@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9a4600e4dc

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[CORE] Prompt Embeddings Support for v1 Engine (#24278)

Signed-off-by: Andrew Sansom <andrew@protopia.ai> Signed-off-by: Andrew Sansom <qthequartermasterman@gmail.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

Andrew Sansom committed Sep 18, 2025 at 19:03 UTC 9a4600e4dcbbd13988c31d5198d3ab8b4172ecca

20 files changed +304 -75

tests/basic_correctness/test_basic_correctness.py

-10

@@ -76,11 +76,6 @@ def test_models(

76 model_executor: str,

77 enable_prompt_embeds: bool,

78 ) -> None:

79 -

80 - if enable_prompt_embeds and envs.is_set(

81 - "VLLM_USE_V1") and envs.VLLM_USE_V1:

82 - pytest.skip("enable_prompt_embeds is not supported in v1.")

83 -

79 if not envs.VLLM_USE_V1:

80 if async_scheduling:

81 pytest.skip("async_scheduling only supported in v1.")

@@ -164,11 +159,6 @@ def test_models_distributed(

159 extra_env: dict[str, str],

160 enable_prompt_embeds: bool,

161 ) -> None:

167 -

168 - if enable_prompt_embeds and envs.is_set(

169 - "VLLM_USE_V1") and envs.VLLM_USE_V1:

170 - pytest.skip("enable_prompt_embeds is not supported in v1.")

171 -

162 if test_suite != TARGET_TEST_SUITE:

163 pytest.skip(f"Skip test for {test_suite}")

164

tests/entrypoints/openai/test_completion_with_prompt_embeds.py

-1

@@ -36,7 +36,6 @@ def default_server_args() -> list[str]:

36 "--enforce-eager",

37 # Prompt Embeds server args

38 "--enable-prompt-embeds",

39 - "--no-enable-chunked-prefill",

39 ]

tests/models/language/generation/test_common.py

-6

@@ -125,12 +125,6 @@ def test_models(hf_runner, vllm_runner, example_prompts, model: str,

125 # in parts of the operators

126 pytest.skip(f"Skipping '{model}' model test with AITER kernel.")

127

128 - # Note: can be removed when

129 - # https://github.com/vllm-project/vllm/pull/24278 finished

130 - if current_platform.is_cpu() and use_prompt_embeds:

131 - pytest.skip("Skipping use_prompt_embeds=True with "

132 - "V1-only CPU backend.")

133 -

128 with hf_runner(model) as hf_model:

129 hf_outputs = hf_model.generate_greedy_logprobs_limit(

130 example_prompts, max_tokens, num_logprobs)

vllm/engine/arg_utils.py

+18 -6

@@ -1513,12 +1513,6 @@ class EngineArgs:

1513 recommend_to_remove=False)

1514 return False

1515

1516 - # No text embedding inputs so far.

1517 - if self.enable_prompt_embeds:

1518 - _raise_or_fallback(feature_name="--enable-prompt-embeds",

1519 - recommend_to_remove=False)

1520 - return False

1521 -

1516 # No Mamba or Encoder-Decoder so far.

1517 if not model_config.is_v1_compatible:

1518 _raise_or_fallback(feature_name=model_config.architectures,

@@ -1651,6 +1645,13 @@ class EngineArgs:

1645 "models in V0 and has been disabled.")

1646 self.enable_prefix_caching = False

1647

1648 + if self.enable_prompt_embeds:

1649 + logger.warning(

1650 + "--enable-prompt-embeds and --enable-prefix-caching "

1651 + "are not supported together in V0. Prefix caching has "

1652 + "been disabled.")

1653 + self.enable_prefix_caching = False

1654 +

1655 # Set max_num_seqs to 256 for VLLM_V0.

1656 if self.max_num_seqs is None:

1657 self.max_num_seqs = 256

@@ -1664,6 +1665,17 @@ class EngineArgs:

1665 # For pooling tasks the default is False

1666 if model_config.runner_type != "pooling":

1667 self.enable_chunked_prefill = True

1668 +

1669 + # TODO: When prefix caching supports prompt embeds inputs, this

1670 + # check can be removed.

1671 + if (self.enable_prompt_embeds

1672 + and self.enable_prefix_caching is not False):

1673 + logger.warning(

1674 + "--enable-prompt-embeds and --enable-prefix-caching "

1675 + "are not supported together in V1. Prefix caching has "

1676 + "been disabled.")

1677 + self.enable_prefix_caching = False

1678 +

1679 if self.enable_prefix_caching is None:

1680 self.enable_prefix_caching = True

1681 else:

vllm/entrypoints/openai/protocol.py

+1 -1

@@ -973,7 +973,6 @@ class CompletionRequest(OpenAIBaseModel):

973 # https://platform.openai.com/docs/api-reference/completions/create

974 model: Optional[str] = None

975 prompt: Optional[Union[list[int], list[list[int]], str, list[str]]] = None

976 - prompt_embeds: Optional[Union[bytes, list[bytes]]] = None

976 best_of: Optional[int] = None

977 echo: Optional[bool] = False

978 frequency_penalty: Optional[float] = 0.0

@@ -1009,6 +1008,7 @@ class CompletionRequest(OpenAIBaseModel):

1008 # --8<-- [end:completion-sampling-params]

1009

1010 # --8<-- [start:completion-extra-params]

1011 + prompt_embeds: Optional[Union[bytes, list[bytes]]] = None

1012 add_special_tokens: bool = Field(

1013 default=True,

1014 description=(

vllm/utils/__init__.py

+27

@@ -3443,3 +3443,30 @@ def decorate_logs(process_name: Optional[str] = None) -> None:

3443 pid = os.getpid()

3444 _add_prefix(sys.stdout, process_name, pid)

3445 _add_prefix(sys.stderr, process_name, pid)

3446 +

3447 +

3448 +def length_from_prompt_token_ids_or_embeds(

3449 + prompt_token_ids: Optional[list[int]],

3450 + prompt_embeds: Optional[torch.Tensor],

3451 +) -> int:

3452 + """Calculate the request length (in number of tokens) give either

3453 + prompt_token_ids or prompt_embeds.

3454 + """

3455 + prompt_token_len = None if prompt_token_ids is None else len(

3456 + prompt_token_ids)

3457 + prompt_embeds_len = \

3458 + None if prompt_embeds is None else len(prompt_embeds)

3459 +

3460 + if prompt_token_len is None:

3461 + if prompt_embeds_len is None:

3462 + raise ValueError(

3463 + "Neither prompt_token_ids nor prompt_embeds were defined.")

3464 + return prompt_embeds_len

3465 + else:

3466 + if (prompt_embeds_len is not None

3467 + and prompt_embeds_len != prompt_token_len):

3468 + raise ValueError(

3469 + "Prompt token ids and prompt embeds had different lengths"

3470 + f" prompt_token_ids={prompt_token_len}"

3471 + f" prompt_embeds={prompt_embeds_len}")

3472 + return prompt_token_len

vllm/v1/core/sched/output.py

+18 -6

@@ -11,6 +11,7 @@ from vllm._bc_linter import bc_linter_include

11 if TYPE_CHECKING:

12 import numpy as np

13 import numpy.typing as npt

14 + import torch

16 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

17 KVConnectorMetadata)

@@ -26,13 +27,14 @@ if TYPE_CHECKING:

27 class NewRequestData:

29 req_id: str

29 - prompt_token_ids: list[int]

30 + prompt_token_ids: Optional[list[int]]

31 mm_features: list[MultiModalFeatureSpec]

32 sampling_params: Optional[SamplingParams]

33 pooling_params: Optional[PoolingParams]

34 block_ids: tuple[list[int], ...]

35 num_computed_tokens: int

36 lora_request: Optional[LoRARequest]

37 + prompt_embeds: Optional[torch.Tensor] = None

39 @classmethod

40 def from_request(

@@ -49,9 +51,12 @@ class NewRequestData:

51 block_ids=block_ids,

52 num_computed_tokens=request.num_computed_tokens,

53 lora_request=request.lora_request,

54 + prompt_embeds=request.prompt_embeds,

55 )

54 - def __repr__(self):

57 + def __repr__(self) -> str:

58 + prompt_embeds_shape = (self.prompt_embeds.shape

59 + if self.prompt_embeds else None)

60 return (f"NewRequestData("

61 f"req_id={self.req_id},"

62 f"prompt_token_ids={self.prompt_token_ids},"

@@ -59,19 +64,26 @@ class NewRequestData:

64 f"sampling_params={self.sampling_params},"

65 f"block_ids={self.block_ids},"

66 f"num_computed_tokens={self.num_computed_tokens},"

62 - f"lora_request={self.lora_request}"

67 + f"lora_request={self.lora_request},"

68 + f"prompt_embeds_shape={prompt_embeds_shape}"

69 ")")

71 # Version of __repr__ with the prompt data obfuscated

66 - def anon_repr(self):

72 + def anon_repr(self) -> str:

73 + prompt_token_ids_len = len(

74 + self.prompt_token_ids

75 + ) if self.prompt_token_ids is not None else None

76 + prompt_embeds_shape = (self.prompt_embeds.shape

77 + if self.prompt_embeds else None)

78 return (f"NewRequestData("

79 f"req_id={self.req_id},"

69 - f"prompt_token_ids_len={len(self.prompt_token_ids)},"

80 + f"prompt_token_ids_len={prompt_token_ids_len},"

81 f"mm_features={self.mm_features},"

82 f"sampling_params={self.sampling_params},"

83 f"block_ids={self.block_ids},"

84 f"num_computed_tokens={self.num_computed_tokens},"

74 - f"lora_request={self.lora_request}"

85 + f"lora_request={self.lora_request},"

86 + f"prompt_embeds_shape={prompt_embeds_shape}"

87 ")")

vllm/v1/engine/__init__.py

+2 -1

@@ -47,7 +47,7 @@ class EngineCoreRequest(

47 gc=False): # type: ignore[call-arg]

49 request_id: str

50 - prompt_token_ids: list[int]

50 + prompt_token_ids: Optional[list[int]]

51 mm_features: Optional[list[MultiModalFeatureSpec]]

52 sampling_params: Optional[SamplingParams]

53 pooling_params: Optional[PoolingParams]

@@ -56,6 +56,7 @@ class EngineCoreRequest(

56 lora_request: Optional[LoRARequest]

57 cache_salt: Optional[str]

58 data_parallel_rank: Optional[int]

59 + prompt_embeds: Optional[torch.Tensor] = None

61 # Index of the client, used to ensure outputs are sent back to the same

62 # client for this request when scaling out the front-end.

vllm/v1/engine/detokenizer.py

+21 -10

@@ -13,6 +13,7 @@ from vllm.engine.output_processor.stop_checker import StopChecker

13 from vllm.logger import init_logger

14 from vllm.transformers_utils.detokenizer_utils import (

15 AnyTokenizer, convert_prompt_ids_to_tokens, detokenize_incrementally)

16 +from vllm.utils import length_from_prompt_token_ids_or_embeds

17 from vllm.v1.engine import EngineCoreRequest

19 logger = init_logger(__name__)

@@ -179,11 +180,12 @@ class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):

180 self.tokenizer: Tokenizer = tokenizer._tokenizer

181

182 # Find a safe place to start.

182 - prompt_suffix = request.prompt_token_ids

183 + prompt_token_ids = request.prompt_token_ids or []

184 + prompt_suffix = prompt_token_ids

185 prompt_len = len(prompt_suffix)

186 if prompt_len > 4:

187 for i in range(4, min(prompt_len + 1, 24)):

186 - suffix = request.prompt_token_ids[-i:]

188 + suffix = prompt_token_ids[-i:]

189 if '�' not in self.tokenizer.decode(suffix):

190 prompt_suffix = suffix

191 break

@@ -260,16 +262,25 @@ class SlowIncrementalDetokenizer(BaseIncrementalDetokenizer):

262 params = request.sampling_params

263 assert params is not None

264

265 + self.prompt_len = length_from_prompt_token_ids_or_embeds(

266 + request.prompt_token_ids, request.prompt_embeds)

267 +

268 # Metadata for incremental detokenization.

264 - self.tokens, self.prefix_offset, self.read_offset = (

265 - convert_prompt_ids_to_tokens(

266 - tokenizer=tokenizer,

267 - prompt_ids=request.prompt_token_ids,

268 - skip_special_tokens=params.skip_special_tokens,

269 - ))

269 + if request.prompt_token_ids is not None:

270 + self.tokens, self.prefix_offset, self.read_offset = (

271 + convert_prompt_ids_to_tokens(

272 + tokenizer=tokenizer,

273 + prompt_ids=request.prompt_token_ids,

274 + skip_special_tokens=params.skip_special_tokens,

275 + ))

276 + else:

277 + # Prompt embedding requests cannot be detokenized, in general.

278 + self.tokens = [""] * self.prompt_len

279 + self.prefix_offset = 0

280 + self.read_offest = 0

281

271 - self.token_ids.extend(request.prompt_token_ids)

272 - self.prompt_len = len(request.prompt_token_ids)

282 + self.token_ids.extend(request.prompt_token_ids

283 + or [0] * self.prompt_len)

284

285 self.skip_special_tokens = params.skip_special_tokens

286 self.spaces_between_special_tokens = (

vllm/v1/engine/output_processor.py

+20 -5

@@ -14,6 +14,7 @@ from vllm.sampling_params import RequestOutputKind

14 from vllm.tracing import (SpanAttributes, SpanKind, Tracer,

15 extract_trace_context)

16 from vllm.transformers_utils.tokenizer import AnyTokenizer

17 +from vllm.utils import length_from_prompt_token_ids_or_embeds

18 from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason

19 from vllm.v1.engine.detokenizer import IncrementalDetokenizer

20 from vllm.v1.engine.logprobs import LogprobsProcessor

@@ -86,7 +87,8 @@ class RequestState:

87 lora_name: Optional[str],

88 output_kind: RequestOutputKind,

89 prompt: Optional[str],

89 - prompt_token_ids: list[int],

90 + prompt_token_ids: Optional[list[int]],

91 + prompt_embeds: Optional[torch.Tensor],

92 logprobs_processor: Optional[LogprobsProcessor],

93 detokenizer: Optional[IncrementalDetokenizer],

94 max_tokens_param: Optional[int],

@@ -104,7 +106,9 @@ class RequestState:

106 self.output_kind = output_kind

107 self.prompt = prompt

108 self.prompt_token_ids = prompt_token_ids

107 - self.prompt_len = len(prompt_token_ids)

109 + self.prompt_embeds = prompt_embeds

110 + self.prompt_len = length_from_prompt_token_ids_or_embeds(

111 + self.prompt_token_ids, self.prompt_embeds)

112 self.logprobs_processor = logprobs_processor

113 self.detokenizer = detokenizer

114 self.max_tokens_param = max_tokens_param

@@ -165,6 +169,7 @@ class RequestState:

169 output_kind=output_kind,

170 prompt=prompt,

171 prompt_token_ids=request.prompt_token_ids,

172 + prompt_embeds=request.prompt_embeds,

173 logprobs_processor=logprobs_processor,

174 detokenizer=detokenizer,

175 max_tokens_param=max_tokens_param,

@@ -223,6 +228,8 @@ class RequestState:

228 first_output = outputs[0]

229 if isinstance(first_output, PoolingOutput):

230 assert len(outputs) == 1

231 + # Prompt embeddings are currently not supported by pooling requests.

232 + assert self.prompt_token_ids is not None

233 return PoolingRequestOutput(

234 request_id=request_id,

235 outputs=first_output,

@@ -236,10 +243,15 @@ class RequestState:

243 else:

244 prompt_logprobs = self.logprobs_processor.prompt_logprobs

245

246 + # If prompt embeds were used, put placeholder prompt token ids

247 + prompt_token_ids = self.prompt_token_ids

248 + if prompt_token_ids is None and self.prompt_embeds is not None:

249 + prompt_token_ids = [0] * len(self.prompt_embeds)

250 +

251 return RequestOutput(

252 request_id=request_id,

253 prompt=self.prompt,

242 - prompt_token_ids=self.prompt_token_ids,

254 + prompt_token_ids=prompt_token_ids,

255 prompt_logprobs=prompt_logprobs,

256 outputs=cast(list[CompletionOutput], outputs),

257 finished=finished,

@@ -469,6 +481,8 @@ class OutputProcessor:

481

482 arrival_time_nano_seconds = int(req_state.stats.arrival_time * 1e9)

483 trace_context = extract_trace_context(engine_core_output.trace_headers)

484 + prompt_length = length_from_prompt_token_ids_or_embeds(

485 + req_state.prompt_token_ids, req_state.prompt_embeds)

486 with (self.tracer.start_as_current_span(

487 "llm_request",

488 kind=SpanKind.SERVER,

@@ -488,7 +502,7 @@ class OutputProcessor:

502 span.set_attribute(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE,

503 queued_time)

504 span.set_attribute(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS,

491 - len(req_state.prompt_token_ids))

505 + prompt_length)

506 span.set_attribute(SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS,

507 metrics.num_generation_tokens)

508 span.set_attribute(

@@ -544,7 +558,8 @@ class OutputProcessor:

558 assert req_state.stats is not None

559 iteration_stats.update_from_finished_request(

560 finish_reason=finish_reason,

547 - num_prompt_tokens=len(req_state.prompt_token_ids),

561 + num_prompt_tokens=length_from_prompt_token_ids_or_embeds(

562 + req_state.prompt_token_ids, req_state.prompt_embeds),

563 max_tokens_param=req_state.max_tokens_param,

564 req_stats=req_state.stats)

565 self.lora_states.finish_request(req_state)

vllm/v1/engine/processor.py

+29 -9

@@ -19,6 +19,7 @@ from vllm.multimodal.utils import argsort_mm_positions

19 from vllm.pooling_params import PoolingParams

20 from vllm.sampling_params import SamplingParams

21 from vllm.transformers_utils.tokenizer import AnyTokenizer

22 +from vllm.utils import length_from_prompt_token_ids_or_embeds

23 from vllm.v1.engine import EngineCoreRequest

24 from vllm.v1.structured_output.backend_guidance import (

25 validate_guidance_grammar)

@@ -390,6 +391,16 @@ class Processor:

391 self._validate_model_inputs(processed_inputs)

392

393 encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)

394 + # Mypy does not always properly infer the types of some elements of

395 + # discriminated unions of TypedDicts, because of how it handles

396 + # inheritance of TypedDict. If we explicitly extract the items we want

397 + # we can avoid type errors from using `dict.get` later in the method.

398 + prompt_str: Optional[str] = None if decoder_inputs[

399 + "type"] == "embeds" else decoder_inputs.get("prompt")

400 + prompt_token_ids = decoder_inputs[

401 + "prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None

402 + prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs[

403 + "type"] == "embeds" else None

404

405 sampling_params = None

406 pooling_params = None

@@ -398,9 +409,10 @@ class Processor:

409 sampling_params = params.clone()

410 # If unset max tokens, then generate up to the max_model_len.

411 if sampling_params.max_tokens is None:

401 - sampling_params.max_tokens = (

402 - self.model_config.max_model_len -

403 - len(decoder_inputs["prompt_token_ids"]))

412 + seq_len = length_from_prompt_token_ids_or_embeds(

413 + prompt_token_ids, prompt_embeds)

414 + sampling_params.max_tokens = \

415 + self.model_config.max_model_len - seq_len

416 sampling_params.update_from_generation_config(

417 self.generation_config_fields, eos_token_id)

418 if self.tokenizer is not None:

@@ -430,9 +442,10 @@ class Processor:

442 identifier=decoder_mm_hashes[modality][idx],

443 mm_position=decoder_mm_positions[modality][idx]))

444

433 - return decoder_inputs.get("prompt"), EngineCoreRequest(

445 + return prompt_str, EngineCoreRequest(

446 request_id=request_id,

435 - prompt_token_ids=decoder_inputs["prompt_token_ids"],

447 + prompt_token_ids=prompt_token_ids,

448 + prompt_embeds=prompt_embeds,

449 mm_features=mm_features,

450 sampling_params=sampling_params,

451 pooling_params=pooling_params,

@@ -461,10 +474,17 @@ class Processor:

474 ):

475 model_config = self.model_config

476

464 - prompt_ids = prompt_inputs["prompt_token_ids"]

477 + prompt_ids = None if prompt_inputs[

478 + "type"] == "embeds" else prompt_inputs["prompt_token_ids"]

479 + prompt_embeds = prompt_inputs["prompt_embeds"] if prompt_inputs[

480 + "type"] == "embeds" else None

481 + prompt_len = length_from_prompt_token_ids_or_embeds(

482 + prompt_ids, prompt_embeds)

483 if not prompt_ids:

484 if prompt_type == "encoder" and model_config.is_multimodal_model:

485 pass # Mllama may have empty encoder inputs for text-only data

486 + elif prompt_inputs["type"] == "embeds":

487 + pass # Prompt embeds should not have prompt_ids.

488 else:

489 raise ValueError(f"The {prompt_type} prompt cannot be empty")

490

@@ -472,7 +492,7 @@ class Processor:

492 tokenizer = None

493 else:

494 tokenizer = self.tokenizer

475 - max_input_id = max(prompt_ids, default=0)

495 + max_input_id = max(prompt_ids or [], default=0)

496

497 # NOTE: tokenizer.max_token_id is the tokenizer’s vocab size while

498 # self.model_config.get_vocab_size() is the model’s vocab size.

@@ -490,7 +510,7 @@ class Processor:

510 f"Token id {max_input_id} is out of vocabulary")

511

512 max_prompt_len = self.model_config.max_model_len

493 - if len(prompt_ids) > max_prompt_len:

513 + if prompt_len > max_prompt_len:

514 if prompt_type == "encoder" and model_config.is_multimodal_model:

515 mm_registry = self.input_preprocessor.mm_registry

516 mm_processor = mm_registry.create_processor(

@@ -514,7 +534,7 @@ class Processor:

534 "number of text tokens.")

535

536 raise ValueError(

517 - f"The {prompt_type} prompt (length {len(prompt_ids)}) is "

537 + f"The {prompt_type} prompt (length {prompt_len}) is "

538 f"longer than the maximum model length of {max_prompt_len}. "

539 f"{suggestion}")

540

vllm/v1/request.py

+12 -3

@@ -7,9 +7,12 @@ from collections.abc import Mapping

7 from functools import partial

8 from typing import TYPE_CHECKING, Any, Callable, Optional, Union

10 +import torch

11 +

12 from vllm.multimodal.inputs import MultiModalFeatureSpec

13 from vllm.pooling_params import PoolingParams

14 from vllm.sampling_params import SamplingParams

15 +from vllm.utils import length_from_prompt_token_ids_or_embeds

16 from vllm.v1.engine import (EngineCoreEvent, EngineCoreEventType,

17 EngineCoreRequest, FinishReason)

18 from vllm.v1.structured_output.request import StructuredOutputRequest

@@ -25,12 +28,13 @@ class Request:

28 def __init__(

29 self,

30 request_id: str,

28 - prompt_token_ids: list[int],

31 + prompt_token_ids: Optional[list[int]],

32 sampling_params: Optional[SamplingParams],

33 pooling_params: Optional[PoolingParams],

34 eos_token_id: Optional[int],

35 client_index: int = 0,

36 arrival_time: Optional[float] = None,

37 + prompt_embeds: Optional[torch.Tensor] = None,

38 mm_features: Optional[list[MultiModalFeatureSpec]] = None,

39 lora_request: Optional["LoRARequest"] = None,

40 structured_output_request: Optional["StructuredOutputRequest"] = None,

@@ -79,9 +83,13 @@ class Request:

83 "sampling_params and pooling_params can't both be unset")

85 self.prompt_token_ids = prompt_token_ids

82 - self.num_prompt_tokens = len(self.prompt_token_ids)

86 + self.prompt_embeds = prompt_embeds

87 + self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

88 + prompt_token_ids, prompt_embeds)

89 self._output_token_ids: list[int] = []

84 - self._all_token_ids: list[int] = self.prompt_token_ids.copy()

90 + self._all_token_ids: list[int] = self.prompt_token_ids.copy(

91 + ) if self.prompt_token_ids is not None else [0

92 + ] * self.num_prompt_tokens

93 self.num_output_placeholders = 0 # Used in async scheduling.

94 self.spec_token_ids: list[int] = []

95 self.num_computed_tokens = 0

@@ -123,6 +131,7 @@ class Request:

131 request_id=request.request_id,

132 client_index=request.client_index,

133 prompt_token_ids=request.prompt_token_ids,

134 + prompt_embeds=request.prompt_embeds,

135 mm_features=request.mm_features,

136 sampling_params=request.sampling_params,

137 pooling_params=request.pooling_params,

vllm/v1/sample/logits_processor/__init__.py

+1 -1

@@ -243,7 +243,7 @@ class AdapterLogitsProcessor(LogitsProcessor):

243 def _new_state(

244 self,

245 params: SamplingParams,

246 - prompt_ids: list[int],

246 + prompt_ids: Optional[list[int]],

247 output_ids: list[int],

248 ) -> Optional[partial[torch.Tensor]]:

249 """Return state representation for new request

vllm/v1/sample/logits_processor/builtin.py

+4 -2

@@ -187,7 +187,8 @@ class MinTokensLogitsProcessor(LogitsProcessor):

187

188 @staticmethod

189 def add_request(

190 - params: SamplingParams, _: list[int], output_tok_ids: list[int]

190 + params: SamplingParams, _: Optional[list[int]],

191 + output_tok_ids: list[int]

192 ) -> Optional[tuple[int, Sequence[int], set[int]]]:

193 min_tokens = params.min_tokens

194 if not min_tokens or len(output_tok_ids) >= min_tokens:

@@ -234,7 +235,8 @@ class MinTokensLogitsProcessor(LogitsProcessor):

235

236 def process_dict_updates(

237 req_entries: dict[int, T], batch_update: Optional[BatchUpdate],

237 - new_state: Callable[[SamplingParams, list[int], list[int]], Optional[T]]

238 + new_state: Callable[[SamplingParams, Optional[list[int]], list[int]],

239 + Optional[T]]

240 ) -> bool:

241 """Utility function to update dict state for sparse LogitsProcessors."""

242

vllm/v1/sample/logits_processor/interface.py

+1 -1

@@ -26,7 +26,7 @@ RemovedRequest = int

27 # (index, params, prompt_tok_ids, output_tok_ids) tuples for new

28 # requests added to the batch.

29 -AddedRequest = tuple[int, SamplingParams, list[int], list[int]]

29 +AddedRequest = tuple[int, SamplingParams, Optional[list[int]], list[int]]

31 # (index 1, index 2, directionality) tuples representing

32 # one-way moves or two-way swaps of requests in batch

vllm/v1/serial_utils.py

+1 -1

@@ -174,7 +174,7 @@ class MsgpackEncoder:

174 ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:

175 assert self.aux_buffers is not None

176 # view the tensor as a contiguous 1D array of bytes

177 - arr = obj.flatten().contiguous().view(torch.uint8).numpy()

177 + arr = obj.flatten().contiguous().cpu().view(torch.uint8).numpy()

178 if obj.nbytes < self.size_threshold:

179 # Smaller tensors are encoded inline, just like ndarrays.

180 data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, arr.data)

vllm/v1/worker/gpu_input_batch.py

+48 -7

@@ -13,7 +13,7 @@ from vllm.lora.request import LoRARequest

13 from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems

14 from vllm.pooling_params import PoolingParams

15 from vllm.sampling_params import SamplingParams, SamplingType

16 -from vllm.utils import swap_dict_values

16 +from vllm.utils import length_from_prompt_token_ids_or_embeds, swap_dict_values

17 from vllm.v1.outputs import LogprobsTensors

18 from vllm.v1.pool.metadata import PoolingMetadata

19 from vllm.v1.sample.logits_processor import (BatchUpdateBuilder,

@@ -29,7 +29,7 @@ from vllm.v1.worker.block_table import MultiGroupBlockTable

29 class CachedRequestState:

31 req_id: str

32 - prompt_token_ids: list[int]

32 + prompt_token_ids: Optional[list[int]]

33 mm_features: list[MultiModalFeatureSpec]

34 sampling_params: Optional[SamplingParams]

35 pooling_params: Optional[PoolingParams]

@@ -43,9 +43,11 @@ class CachedRequestState:

43 mrope_position_delta: Optional[int] = None

45 lora_request: Optional[LoRARequest] = None

46 + prompt_embeds: Optional[torch.Tensor] = None

48 def __post_init__(self):

48 - self.num_prompt_tokens = len(self.prompt_token_ids)

49 + self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

50 + self.prompt_token_ids, self.prompt_embeds)

52 @property

53 def num_tokens(self) -> int:

@@ -63,6 +65,10 @@ class CachedRequestState:

66 def get_token_id(self, idx: int) -> int:

67 if idx < self.num_prompt_tokens:

68 + if self.prompt_token_ids is None:

69 + raise ValueError(

70 + f"Tried to access token index {idx}, but that token was "

71 + "provided via prompt_embeds, and its ID is unknown.")

72 return self.prompt_token_ids[idx]

73 elif idx - self.num_prompt_tokens < len(self.output_token_ids):

74 return self.output_token_ids[idx - self.num_prompt_tokens]

@@ -109,6 +115,14 @@ class InputBatch:

115 pin_memory=False,

116 )

117 self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()

118 + self.is_token_ids = torch.zeros((max_num_reqs, max_model_len),

119 + device="cpu",

120 + dtype=bool,

121 + pin_memory=False)

122 + # Store prompt embeddings per request to avoid OOM from large upfront

123 + # allocation if max_model_len is big.

124 + # Maps req_index -> tensor of shape (num_prompt_tokens, hidden_size)

125 + self.req_prompt_embeds: dict[int, torch.Tensor] = {}

126 self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)

127 self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)

128 self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)

@@ -310,15 +324,23 @@ class InputBatch:

324 self.req_id_to_index[req_id] = req_index

325

326 # Copy the prompt token ids and output token ids.

313 - num_prompt_tokens = len(request.prompt_token_ids)

327 + num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

328 + request.prompt_token_ids, request.prompt_embeds)

329 self.num_prompt_tokens[req_index] = num_prompt_tokens

315 - self.token_ids_cpu[

316 - req_index, :num_prompt_tokens] = request.prompt_token_ids

330 start_idx = num_prompt_tokens

331 end_idx = start_idx + len(request.output_token_ids)

332 + if request.prompt_token_ids is not None:

333 + self.token_ids_cpu[

334 + req_index, :num_prompt_tokens] = request.prompt_token_ids

335 + self.is_token_ids[req_index, :num_prompt_tokens] = True

336 + else:

337 + self.is_token_ids[req_index, :num_prompt_tokens] = False

338 + if request.prompt_embeds is not None:

339 + self.req_prompt_embeds[req_index] = request.prompt_embeds

340 self.token_ids_cpu[req_index,

341 start_idx:end_idx] = request.output_token_ids

321 - # Number of token ids in token_ids_cpu.

342 + self.is_token_ids[req_index, start_idx:end_idx] = True

343 + # Number of token ids in prompt (token_ids_cpu or prompt_embeds).

344 # NOTE(woosuk): This may include spec decode tokens.

345 self.num_tokens[req_index] = request.num_tokens

346 # Number of tokens without spec decode tokens.

@@ -503,6 +525,20 @@ class InputBatch:

525 self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]

526 self.token_ids_cpu[i2, ...] = tmp

527

528 + self.is_token_ids[[i1, i2], ...] = self.is_token_ids[[i2, i1], ...]

529 +

530 + # Swap prompt embeddings if they exist

531 + embeds_i1 = self.req_prompt_embeds.get(i1)

532 + embeds_i2 = self.req_prompt_embeds.get(i2)

533 + if embeds_i1 is not None:

534 + self.req_prompt_embeds[i2] = embeds_i1

535 + else:

536 + self.req_prompt_embeds.pop(i2, None)

537 + if embeds_i2 is not None:

538 + self.req_prompt_embeds[i1] = embeds_i2

539 + else:

540 + self.req_prompt_embeds.pop(i1, None)

541 +

542 self.block_table.swap_row(i1, i2)

543

544 self.request_lora_mapping[i1], self.request_lora_mapping[i2] = \

@@ -592,6 +628,11 @@ class InputBatch:

628 num_tokens = self.num_tokens[last_req_index]

629 self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[

630 last_req_index, :num_tokens]

631 + self.is_token_ids[empty_index, :num_tokens] = self.is_token_ids[

632 + last_req_index, :num_tokens]

633 + if last_req_index in self.req_prompt_embeds:

634 + self.req_prompt_embeds[

635 + empty_index] = self.req_prompt_embeds.pop(last_req_index)

636 self.num_tokens[empty_index] = num_tokens

637 self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[

638 last_req_index]

vllm/v1/worker/gpu_model_runner.py

+96 -3

@@ -56,7 +56,9 @@ from vllm.sequence import IntermediateTensors, PoolerOutput

56 from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

57 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,

58 GiB_bytes, check_use_alibi, get_dtype_size,

59 - is_pin_memory_available, round_up, supports_dynamo)

59 + is_pin_memory_available,

60 + length_from_prompt_token_ids_or_embeds, round_up,

61 + supports_dynamo)

62 from vllm.v1.attention.backends.flash_attn import AttentionMetadata

63 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder

64 from vllm.v1.attention.backends.utils import (

@@ -197,6 +199,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

199 cache_config.cache_dtype]

200

201 self.is_pooling_model = (model_config.runner_type == 'pooling')

202 + self.enable_prompt_embeds = model_config.enable_prompt_embeds

203 self.is_multimodal_raw_input_only_model = (

204 model_config.is_multimodal_raw_input_only_model)

205

@@ -342,6 +345,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

345 self.hidden_size,

346 dtype=self.dtype,

347 numpy=False)

348 + self.is_token_ids = self._make_buffer(self.max_num_tokens,

349 + dtype=torch.bool)

350 self.discard_request_indices = self._make_buffer(self.max_num_reqs,

351 dtype=torch.int64)

352 self.num_discarded_requests = 0

@@ -574,6 +579,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

579 req_state = CachedRequestState(

580 req_id=req_id,

581 prompt_token_ids=new_req_data.prompt_token_ids,

582 + prompt_embeds=new_req_data.prompt_embeds,

583 mm_features=new_req_data.mm_features,

584 sampling_params=sampling_params,

585 pooling_params=pooling_params,

@@ -819,6 +825,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

825 if self.input_batch.prev_sampled_token_ids is None:

826 # Normal scheduling case

827 self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

828 + self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)

829 + self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)

830 return

831

832 # Async scheduling case, where some decode requests from the previous

@@ -844,6 +852,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

852 # If not all requests are decodes from the last iteration,

853 # We need to copy the input_ids_cpu to the GPU first.

854 self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

855 + self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)

856 + self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)

857 if num_commmon_tokens == 0:

858 # No requests in common with the previous iteration

859 # So input_ids_cpu will have all the input ids.

@@ -857,6 +867,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

867 self.input_batch.prev_sampled_token_ids[:num_commmon_tokens,

868 0],

869 non_blocking=True)

870 + self.is_token_ids.gpu[:num_commmon_tokens] = True

871 return

872 # Upload the index tensors asynchronously

873 # so the scatter can be non-blocking.

@@ -947,14 +958,60 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

958 # where M is the max_model_len.

959 token_indices = (positions_np +

960 req_indices * self.input_batch.token_ids_cpu.shape[1])

961 + token_indices_tensor = torch.from_numpy(token_indices)

962

963 # NOTE(woosuk): We use torch.index_select instead of np.take here

964 # because torch.index_select is much faster than np.take for large

965 # tensors.

966 torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(),

967 0,

956 - torch.from_numpy(token_indices),

968 + token_indices_tensor,

969 out=self.input_ids.cpu[:total_num_scheduled_tokens])

970 + is_token_ids = self.input_batch.is_token_ids.flatten()

971 + torch.index_select(

972 + is_token_ids,

973 + 0,

974 + token_indices_tensor,

975 + out=self.is_token_ids.cpu[:total_num_scheduled_tokens])

976 +

977 + # Because we did not pre-allocate a massive prompt_embeds CPU tensor on

978 + # the InputBatch, we need to fill in the prompt embeds into the expected

979 + # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor.

980 + if self.input_batch.req_prompt_embeds:

981 + output_idx = 0

982 + for req_idx in range(num_reqs):

983 + num_sched = num_scheduled_tokens[req_idx]

984 +

985 + # Skip if this request doesn't have embeddings

986 + if req_idx not in self.input_batch.req_prompt_embeds:

987 + output_idx += num_sched

988 + continue

989 +

990 + # Skip if no tokens scheduled

991 + if num_sched <= 0:

992 + output_idx += num_sched

993 + continue

994 +

995 + req_embeds = self.input_batch.req_prompt_embeds[req_idx]

996 + start_pos = self.input_batch.num_computed_tokens_cpu[req_idx]

997 +

998 + # Skip if trying to read beyond available embeddings

999 + if start_pos >= req_embeds.shape[0]:

1000 + output_idx += num_sched

1001 + continue

1002 +

1003 + # Copy available embeddings

1004 + end_pos = start_pos + num_sched

1005 + actual_end = min(end_pos, req_embeds.shape[0])

1006 + actual_num_sched = actual_end - start_pos

1007 +

1008 + if actual_num_sched > 0:

1009 + self.inputs_embeds.cpu[output_idx:output_idx +

1010 + actual_num_sched].copy_(

1011 + req_embeds[start_pos:actual_end]

1012 + )

1013 +

1014 + output_idx += num_sched

1015

1016 self.input_batch.block_table.compute_slot_mapping(

1017 req_indices, positions_np)

@@ -1279,7 +1336,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1336 self.input_batch.num_computed_tokens_cpu[index]

1337 num_scheduled_tokens = \

1338 scheduler_output.num_scheduled_tokens[req_id]

1282 - num_prompt_tokens = len(req.prompt_token_ids)

1339 + num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

1340 + req.prompt_token_ids, req.prompt_embeds)

1341

1342 if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens:

1343 prompt_part_len = max(0,

@@ -1845,6 +1903,32 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1903 **self._init_model_kwargs(num_scheduled_tokens),

1904 **self._extract_mm_kwargs(scheduler_output),

1905 }

1906 + elif (self.enable_prompt_embeds and get_pp_group().is_first_rank):

1907 + # Get the input embeddings for the tokens that are not input embeds,

1908 + # then put them into the appropriate positions.

1909 + # TODO(qthequartermasterman): Since even when prompt embeds are

1910 + # enabled, (a) not all requests will use prompt embeds, and (b)

1911 + # after the initial prompt is processed, the rest of the generated

1912 + # tokens will be token ids, it is not desirable to have the

1913 + # embedding layer outside of the CUDA graph all the time. The v0

1914 + # engine avoids this by "double compiling" the CUDA graph, once

1915 + # with input_ids and again with inputs_embeds, for all num_tokens.

1916 + # If a batch only has token ids, then including the embedding layer

1917 + # in the CUDA graph will be more performant (like in the else case

1918 + # below).

1919 + token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens] \

1920 + .nonzero(as_tuple=False) \

1921 + .squeeze(1)

1922 + # Some tokens ids may need to become embeds

1923 + if token_ids_idx.numel() > 0:

1924 + token_ids = self.input_ids.gpu[token_ids_idx]

1925 + tokens_to_embeds = self.model.get_input_embeddings(

1926 + input_ids=token_ids)

1927 + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds

1928 +

1929 + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]

1930 + model_kwargs = self._init_model_kwargs(num_input_tokens)

1931 + input_ids = None

1932 else:

1933 # For text-only models, we use token ids as input.

1934 # While it is possible to use embeddings as input just like the

@@ -2023,6 +2107,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2107

2108 self.input_batch.token_ids_cpu[req_idx,

2109 start_idx:end_idx] = sampled_ids

2110 + self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True

2111 self.input_batch.num_tokens_no_spec[req_idx] = end_idx

2112 self.input_batch.num_tokens[req_idx] = end_idx

2113

@@ -2570,6 +2655,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2655

2656 # Get metadata for this request.

2657 request = self.requests[req_id]

2658 + if request.prompt_token_ids is None:

2659 + # Prompt logprobs is incompatible with prompt embeddings

2660 + continue

2661 +

2662 num_prompt_tokens = len(request.prompt_token_ids)

2663 prompt_token_ids = torch.tensor(request.prompt_token_ids).to(

2664 self.device, non_blocking=True)

@@ -2922,6 +3011,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3011 **model_kwargs,

3012 **self._dummy_mm_kwargs(num_reqs),

3013 }

3014 + elif self.enable_prompt_embeds:

3015 + input_ids = None

3016 + inputs_embeds = self.inputs_embeds.gpu[:num_tokens]

3017 + model_kwargs = self._init_model_kwargs(num_tokens)

3018 else:

3019 input_ids = self.input_ids.gpu[:num_tokens]

3020 inputs_embeds = None

vllm/v1/worker/tpu_input_batch.py

+4 -2

@@ -9,7 +9,7 @@ import torch

10 from vllm.lora.request import LoRARequest

11 from vllm.sampling_params import SamplingType

12 -from vllm.utils import swap_dict_values

12 +from vllm.utils import length_from_prompt_token_ids_or_embeds, swap_dict_values

13 from vllm.v1.outputs import LogprobsTensors

14 from vllm.v1.worker.block_table import MultiGroupBlockTable

15 from vllm.v1.worker.gpu_input_batch import CachedRequestState

@@ -213,7 +213,9 @@ class InputBatch:

213 self.req_id_to_index[req_id] = req_index

214

215 # Copy the prompt token ids and output token ids.

216 - num_prompt_tokens = len(request.prompt_token_ids)

216 + num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

217 + request.prompt_token_ids, request.prompt_embeds)

218 + # TODO: copy prompt_embeds

219 self.num_prompt_tokens[req_index] = num_prompt_tokens

220 self.token_ids_cpu[

221 req_index, :num_prompt_tokens] = request.prompt_token_ids

vllm/v1/worker/tpu_model_runner.py

@@ -387,6 +387,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

387 self.requests[req_id] = CachedRequestState(

388 req_id=req_id,

389 prompt_token_ids=new_req_data.prompt_token_ids,

390 + prompt_embeds=new_req_data.prompt_embeds,

391 mm_features=new_req_data.mm_features,

392 sampling_params=sampling_params,

393 pooling_params=None,