@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 455edc022b

[ModelRunnerV2] Support prompt embeds (#42963)

Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: Canlin Guo <961750412@qq.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Canlin Guo committed Aug 18, 2026 at 05:58 UTC 455edc022b45bbb5a279cff21cd3a34e6371aadf
17 files changed +551 -54
tests/v1/worker/test_encoder_runner.py
+90 -5
@@ -15,7 +15,13 @@ import numpy as np
15 import pytest
16 import torch
17
18 -from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
18 +from vllm.multimodal.inputs import (
19 + MultiModalFeatureSpec,
20 + MultiModalFieldElem,
21 + MultiModalKwargsItem,
22 + MultiModalSharedField,
23 + PlaceholderRange,
24 +)
25 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
26 from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
27 from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -25,6 +31,25 @@ pytestmark = pytest.mark.cpu_test
31 HIDDEN = 4
32
33
34 +def _model_state(cache: EncoderCache) -> MagicMock:
35 + """A mock ModelState backed by a real EncoderCache."""
36 + state = MagicMock()
37 + state.encoder_cache = cache
38 + state.device = torch.device("cpu")
39 + return state
40 +
41 +
42 +def _embeds_item(embeds: torch.Tensor) -> MultiModalKwargsItem:
43 + """A `prompt_embeds` kwargs item, as the HF renderer builds it."""
44 + return MultiModalKwargsItem(
45 + {
46 + "embedding": MultiModalFieldElem(
47 + data=embeds, field=MultiModalSharedField(batch_size=1)
48 + )
49 + }
50 + )
51 +
52 +
53 def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec:
54 return MultiModalFeatureSpec(
55 data=None,
@@ -197,8 +222,7 @@ def test_execute_mm_encoder_caches_outputs_without_gathering():
222 items the connector already holds, and a producer has no load path).
223 """
224 cache = EncoderCache()
200 - state = MagicMock()
201 - state.encoder_cache = cache
225 + state = _model_state(cache)
226 embedding = torch.ones(2, HIDDEN)
227 # (mm_hashes, [(modality, kwargs item), ...]), as prepare_mm_inputs returns.
228 state.encoder_runner.prepare_mm_inputs.return_value = (
@@ -216,8 +240,7 @@ def test_execute_mm_encoder_caches_outputs_without_gathering():
240 def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
241 """A step that schedules no encoder input must not touch the encoder."""
242 cache = EncoderCache()
219 - state = MagicMock()
220 - state.encoder_cache = cache
243 + state = _model_state(cache)
244 state.encoder_runner.prepare_mm_inputs.return_value = ([], [])
245
246 ModelState.execute_mm_encoder(state, {})
@@ -226,6 +249,68 @@ def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
249 state.encoder_runner.execute_mm_encoder.assert_not_called()
250
251
252 +def _pe_feature(identifier: str, embeds: torch.Tensor, offset: int = 0):
253 + return MultiModalFeatureSpec(
254 + data=_embeds_item(embeds),
255 + modality="prompt_embeds",
256 + identifier=identifier,
257 + mm_position=PlaceholderRange(offset=offset, length=embeds.shape[0]),
258 + )
259 +
260 +
261 +def test_prepare_mm_inputs_passes_prompt_embeds_through():
262 + """`prompt_embeds` is already in embedding space, so no encoder may run.
263 +
264 + The renderer delivers prompt_embeds mixed with real media as an ordinary MM
265 + modality. prepare_mm_inputs must cache the tensor directly and keep it out
266 + of the encoder batch -- the vision encoder cannot consume it, and a missing
267 + cache entry makes the subsequent gather raise "Encoder cache miss".
268 + """
269 + prompt_embeds = torch.arange(2 * HIDDEN, dtype=torch.float32).view(2, HIDDEN)
270 + image_feature = MultiModalFeatureSpec(
271 + data=MagicMock(),
272 + modality="image",
273 + identifier="hash_img",
274 + mm_position=PlaceholderRange(offset=2, length=2),
275 + )
276 + runner = _make_runner(
277 + [_pe_feature("hash_pe", prompt_embeds), image_feature], cached=[]
278 + )
279 +
280 + mm_hashes, mm_kwargs = runner.prepare_mm_inputs({"req0": [0, 1]})
281 +
282 + # Only the image remains for the encoder; the embeds are already cached.
283 + assert mm_hashes == ["hash_img"]
284 + assert [modality for modality, _ in mm_kwargs] == ["image"]
285 + assert torch.equal(runner.encoder_cache.encoder_outputs["hash_pe"], prompt_embeds)
286 +
287 +
288 +def test_prepare_mm_inputs_skips_cached_prompt_embeds():
289 + """A prompt_embeds item already in the cache must not be re-uploaded."""
290 + prompt_embeds = torch.ones(3, HIDDEN)
291 + feature = _pe_feature("hash_pe", prompt_embeds)
292 + runner = _make_runner([feature], cached=[feature])
293 + sentinel = runner.encoder_cache.encoder_outputs["hash_pe"]
294 +
295 + mm_hashes, mm_kwargs = runner.prepare_mm_inputs({"req0": [0]})
296 +
297 + assert mm_hashes == [] and mm_kwargs == []
298 + assert runner.encoder_cache.encoder_outputs["hash_pe"] is sentinel
299 +
300 +
301 +def test_execute_mm_encoder_skips_encoder_for_prompt_embeds_only():
302 + """A batch of nothing but prompt_embeds must not invoke the encoder."""
303 + prompt_embeds = torch.ones(3, HIDDEN)
304 + runner = _make_runner([_pe_feature("hash_pe", prompt_embeds)], cached=[])
305 + state = _model_state(runner.encoder_cache)
306 + state.encoder_runner.prepare_mm_inputs.side_effect = runner.prepare_mm_inputs
307 +
308 + ModelState.execute_mm_encoder(state, {"req0": [0]})
309 +
310 + state.encoder_runner.execute_mm_encoder.assert_not_called()
311 + assert torch.equal(runner.encoder_cache.encoder_outputs["hash_pe"], prompt_embeds)
312 +
313 +
314 def test_encoder_timing_stats_registry():
315 runner = _make_runner([], [])
316 runner.enable_timing = True
tests/v1/worker/test_gpu_model_runner.py
+10
@@ -68,6 +68,16 @@ NUM_BLOCKS = 10
68 DEVICE_TYPE = current_platform.device_type
69
70
71 +@pytest.fixture(autouse=True)
72 +def _restore_default_dtype():
73 + """Several tests here set the process-wide default dtype to float16 and
74 + previously leaked it, corrupting later float-sensitive tests in the same
75 + pytest process (torch.randn silently produced fp16)."""
76 + old = torch.get_default_dtype()
77 + yield
78 + torch.set_default_dtype(old)
79 +
80 +
81 def initialize_kv_cache(runner: GPUModelRunner):
82 """
83 Only perform necessary steps in GPUModelRunner.initialize_kv_cache()
tests/v1/worker/test_prompt_embeds_state.py new
+149
@@ -0,0 +1,149 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the Model Runner V2 prompt-embeds overlay (PromptEmbedsState).
4 +
5 +The overlay kernel reads each request's GPU-resident prompt embeddings through
6 +a per-request pointer table and writes the rows scheduled this step into
7 +`inputs_embeds`, honoring chunked prefill (`num_computed_tokens` offset), the
8 +prompt/decode boundary (rows past the embeds length untouched), and the
9 +mixed-mode `prompt_is_token_ids` mask (token-id rows keep the base embedding).
10 +"""
11 +
12 +from dataclasses import dataclass
13 +
14 +import pytest
15 +import torch
16 +
17 +pytest.importorskip("triton")
18 +if not torch.cuda.is_available():
19 + pytest.skip(
20 + "CUDA required for prompt-embeds overlay tests", allow_module_level=True
21 + )
22 +
23 +from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState
24 +
25 +HIDDEN = 24
26 +MAX_NUM_REQS = 8
27 +DEVICE = torch.device("cuda")
28 +
29 +
30 +@dataclass
31 +class _NewReqData:
32 + req_id: str
33 + prompt_embeds: torch.Tensor | None
34 + prompt_is_token_ids: list[bool] | None = None
35 +
36 +
37 +@dataclass
38 +class _Batch:
39 + num_reqs: int
40 + num_scheduled_tokens: torch.Tensor # np-like, only .max() is used
41 + idx_mapping: torch.Tensor
42 + query_start_loc: torch.Tensor
43 +
44 +
45 +def _make_state() -> PromptEmbedsState:
46 + return PromptEmbedsState(MAX_NUM_REQS, HIDDEN, torch.float32, DEVICE)
47 +
48 +
49 +def _batch(num_scheduled: list[int], idx_mapping: list[int]) -> _Batch:
50 + query_start_loc = [0]
51 + for n in num_scheduled:
52 + query_start_loc.append(query_start_loc[-1] + n)
53 + return _Batch(
54 + num_reqs=len(num_scheduled),
55 + num_scheduled_tokens=torch.tensor(num_scheduled, dtype=torch.int32),
56 + idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64, device=DEVICE),
57 + query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE),
58 + )
59 +
60 +
61 +def _apply(
62 + state: PromptEmbedsState,
63 + batch: _Batch,
64 + num_computed: list[int],
65 + num_tokens: int,
66 +) -> tuple[torch.Tensor, torch.Tensor]:
67 + """Run the overlay on a fresh base buffer; return (result, base)."""
68 + num_computed_tokens = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE)
69 + for batch_idx, req_index in enumerate(batch.idx_mapping.tolist()):
70 + num_computed_tokens[req_index] = num_computed[batch_idx]
71 + base = torch.randn(num_tokens, HIDDEN, dtype=torch.float32, device=DEVICE)
72 + inputs_embeds = base.clone()
73 + state.apply(batch, num_computed_tokens, inputs_embeds)
74 + torch.accelerator.synchronize()
75 + return inputs_embeds, base
76 +
77 +
78 +def test_overlay_chunked_prefill_and_decode():
79 + """Rows within the embeds range come from prompt_embeds at the
80 + num_computed offset; requests without embeds and requests past their
81 + embeds length (decode) keep the base embedding."""
82 + state = _make_state()
83 + embeds_a = torch.randn(6, HIDDEN, dtype=torch.float32)
84 + embeds_b = torch.randn(5, HIDDEN, dtype=torch.float32)
85 + state.add_request(0, _NewReqData("a", embeds_a))
86 + state.add_request(1, _NewReqData("b", embeds_b))
87 + state.add_request(2, _NewReqData("c", None))
88 + state.apply_staged_writes()
89 +
90 + # a: chunk [2, 6) of its embeds; b: fully decoded; c: no embeds.
91 + batch = _batch(num_scheduled=[4, 1, 3], idx_mapping=[0, 1, 2])
92 + out, base = _apply(state, batch, num_computed=[2, 7, 1], num_tokens=8)
93 +
94 + torch.testing.assert_close(out[0:4], embeds_a[2:6].to(DEVICE))
95 + torch.testing.assert_close(out[4:8], base[4:8])
96 +
97 +
98 +def test_overlay_clamps_to_embeds_length():
99 + """A window straddling the end of the prompt embeds writes only the
100 + in-range rows (e.g. final prefill chunk + sampled token)."""
101 + state = _make_state()
102 + embeds = torch.randn(4, HIDDEN, dtype=torch.float32)
103 + state.add_request(3, _NewReqData("a", embeds))
104 + state.apply_staged_writes()
105 +
106 + batch = _batch(num_scheduled=[3], idx_mapping=[3])
107 + out, base = _apply(state, batch, num_computed=[2], num_tokens=3)
108 +
109 + torch.testing.assert_close(out[0:2], embeds[2:4].to(DEVICE))
110 + torch.testing.assert_close(out[2:3], base[2:3])
111 +
112 +
113 +def test_overlay_respects_is_token_ids_mask():
114 + """Mixed mode: positions marked as real token ids keep the base
115 + embedding; only embed positions are overwritten."""
116 + state = _make_state()
117 + embeds = torch.randn(5, HIDDEN, dtype=torch.float32)
118 + is_token_ids = [True, False, False, True, False]
119 + state.add_request(0, _NewReqData("a", embeds, is_token_ids))
120 + state.apply_staged_writes()
121 +
122 + batch = _batch(num_scheduled=[5], idx_mapping=[0])
123 + out, base = _apply(state, batch, num_computed=[0], num_tokens=5)
124 +
125 + embeds_gpu = embeds.to(DEVICE)
126 + for pos, is_token in enumerate(is_token_ids):
127 + expected = base[pos] if is_token else embeds_gpu[pos]
128 + torch.testing.assert_close(out[pos], expected)
129 +
130 +
131 +def test_index_reuse_clears_stale_entry():
132 + """A request added at a previously-used index without embeds must not
133 + inherit the prior occupant's pointer-table entry."""
134 + state = _make_state()
135 + state.add_request(0, _NewReqData("a", torch.randn(4, HIDDEN, dtype=torch.float32)))
136 + # A second live embeds request so the kernel actually launches (the
137 + # overlay is skipped entirely when no request holds embeds).
138 + other = torch.randn(2, HIDDEN, dtype=torch.float32)
139 + state.add_request(1, _NewReqData("other", other))
140 + state.apply_staged_writes()
141 + state.remove_request("a")
142 + state.add_request(0, _NewReqData("b", None))
143 + state.apply_staged_writes()
144 +
145 + batch = _batch(num_scheduled=[2, 2], idx_mapping=[0, 1])
146 + out, base = _apply(state, batch, num_computed=[0, 0], num_tokens=4)
147 +
148 + torch.testing.assert_close(out[0:2], base[0:2])
149 + torch.testing.assert_close(out[2:4], other.to(DEVICE))
vllm/config/model.py
+6
@@ -926,6 +926,12 @@ class ModelConfig:
926 f"got {type(self.max_model_len).__name__}: {self.max_model_len!r}. "
927 "Example: max_model_len=2048"
928 )
929 + if self.enable_prompt_embeds and self.is_encoder_decoder:
930 + # No encoder-decoder model accepts `inputs_embeds`; their decoders
931 + # embed `input_ids` internally.
932 + raise ValueError(
933 + "--enable-prompt-embeds is not supported with encoder-decoder models."
934 + )
935 return self
936
937 def _resolve_mm_device_do_normalize(
vllm/config/vllm.py
-3
@@ -2449,9 +2449,6 @@ class VllmConfig:
2449 ):
2450 unsupported.append("custom logits processors")
2451
2452 - if model_config is not None and model_config.enable_prompt_embeds:
2453 - unsupported.append("prompt embeds")
2454 -
2452 if self.cache_config.kv_sharing_fast_prefill:
2453 # Will be added by https://github.com/vllm-project/vllm/pull/35045
2454 unsupported.append("KV sharing fast prefill")
vllm/model_executor/models/diffusion_gemma.py
+2 -2
@@ -776,7 +776,7 @@ class DiffusionGemmaModelState(ModelState):
776 ) -> None:
777 super().__init__(vllm_config, model, encoder_cache, device)
778
779 - # Per-step MM data produced by get_mm_embeddings and consumed by
779 + # Per-step MM data produced by prepare_inputs_embeds and consumed by
780 # prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that
781 # prepare_inputs can call embed_input_ids directly into the
782 # persistent _inputs_embeds_buf, avoiding the intermediate copy
@@ -874,7 +874,7 @@ class DiffusionGemmaModelState(ModelState):
874 if idx is not None:
875 self.diffusion_states.remove_request(idx)
876
877 - def get_mm_embeddings(
877 + def prepare_inputs_embeds(
878 self,
879 scheduled_encoder_inputs: dict[str, list[int]],
880 input_batch: InputBatch,
vllm/model_executor/models/longcat_flash_ngram.py
+3
@@ -273,6 +273,9 @@ class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
273
274
275 class LongcatNgramModelState(DefaultModelState):
276 + # prepare_inputs builds its own inputs_embeds from n-gram token embeddings.
277 + supports_prompt_embeds = False
278 +
279 """Per-request n-gram token history for LongCat-Flash-Lite.
280
281 Maintains a small CPU-side per-slot context (last ``n-1`` processed tokens)
vllm/v1/core/sched/output.py
+8
@@ -68,6 +68,14 @@ class NewRequestData:
68 prefill_token_ids=prefill_token_ids,
69 )
70
71 + @property
72 + def prompt_len(self) -> int:
73 + if self.prompt_token_ids is not None:
74 + return len(self.prompt_token_ids)
75 + if self.prompt_embeds is not None:
76 + return self.prompt_embeds.shape[0]
77 + return 0
78 +
79 def __repr__(self) -> str:
80 prompt_embeds_shape = (
81 self.prompt_embeds.shape if self.prompt_embeds is not None else None
vllm/v1/request.py
+13 -5
@@ -142,11 +142,19 @@ class Request:
142 prompt_token_ids, prompt_embeds
143 )
144 self._output_token_ids: list[int] = []
145 - self._all_token_ids: list[int] = (
146 - self.prompt_token_ids.copy()
147 - if self.prompt_token_ids is not None
148 - else [0] * self.num_prompt_tokens
149 - )
145 + if self.prompt_token_ids is None:
146 + self._all_token_ids: list[int] = [0] * self.num_prompt_tokens
147 + elif self.prompt_is_token_ids is None:
148 + self._all_token_ids = self.prompt_token_ids.copy()
149 + else:
150 + # Mixed-mode prompt: positions covered by prompt_embeds hold a sentinel
151 + # special token id that may lie outside the embedding. Zero them, matching
152 + # the no-token-ids case above, so embedding gathers over these placeholder
153 + # ids stay in bounds; the actual inputs come from prompt_embeds.
154 + self._all_token_ids = [
155 + t if is_tok else 0
156 + for t, is_tok in zip(self.prompt_token_ids, self.prompt_is_token_ids)
157 + ]
158
159 # Used in async scheduling.
160 self.num_output_placeholders = 0
vllm/v1/worker/gpu/mm/encoder_runner.py
+12 -1
@@ -18,7 +18,7 @@ from vllm.multimodal.utils import (
18 group_and_batch_mm_kwargs,
19 set_mm_embedding_modality,
20 )
21 -from vllm.utils.torch_utils import PIN_MEMORY
21 +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
22 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
23 from vllm.v1.worker.utils import (
24 EncoderTimingStats,
@@ -91,6 +91,17 @@ class EncoderRunner:
91 continue
92 if mm_feature.identifier in self.encoder_cache.encoder_outputs:
93 continue
94 + if mm_feature.modality == "prompt_embeds":
95 + # Passthrough modality: the tensor is already in the
96 + # model's embedding space, so no encoder runs. Cache it
97 + # directly so gather_mm_embeddings splices it via the
98 + # standard is_mm_embed path.
99 + embeds = mm_feature.data["embedding"].data
100 + assert isinstance(embeds, torch.Tensor)
101 + self.encoder_cache.encoder_outputs[mm_feature.identifier] = (
102 + async_tensor_h2d(embeds, device=self.device)
103 + )
104 + continue
105 mm_hashes.append(mm_feature.identifier)
106 mm_kwargs.append((mm_feature.modality, mm_feature.data))
107
vllm/v1/worker/gpu/model_runner.py
+14 -10
@@ -46,6 +46,7 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
46 initialize_mamba_ssu_backend,
47 )
48 from vllm.model_executor.model_loader import get_model_loader
49 +from vllm.model_executor.models.interfaces import requires_raw_input_tokens
50 from vllm.model_executor.offloader import (
51 create_offloader,
52 get_offloader,
@@ -231,6 +232,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
232 self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
233 self.model_config
234 )
235 + self.uses_inputs_embeds = (
236 + self.supports_mm_inputs or self.model_config.enable_prompt_embeds
237 + )
238 self.encoder_cache = None
239 if self.supports_mm_inputs and self.is_first_pp_rank:
240 self.encoder_cache = EncoderCache()
@@ -956,7 +960,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
960
961 def add_requests(self, scheduler_output: SchedulerOutput) -> None:
962 for new_req_data in scheduler_output.scheduled_new_reqs:
959 - assert new_req_data.prompt_token_ids is not None
963 assert new_req_data.prefill_token_ids is not None
964 req_id = new_req_data.req_id
965
@@ -965,7 +968,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
968 # with the updated prompt_token_ids and mm_features.
969 self._remove_request(req_id)
970
968 - prompt_len = len(new_req_data.prompt_token_ids)
971 + prompt_len = new_req_data.prompt_len
972 sampling_params = new_req_data.sampling_params
973 self.req_states.add_request(
974 req_id=req_id,
@@ -980,6 +983,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
983
984 if self.pooling_runner is not None:
985 assert new_req_data.pooling_params is not None
986 + assert new_req_data.prompt_token_ids is not None
987 self.pooling_runner.add_request(
988 req_id,
989 req_index,
@@ -1550,17 +1554,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1554 input_ids = input_batch.input_ids
1555 inputs_embeds = None
1556 ec_connector_output = None
1553 - if self.supports_mm_inputs and self.is_first_pp_rank:
1554 - # Run MM encoder (if needed) and get multimodal embeddings.
1555 - # Only first PP rank prepares multimodal embeddings.
1557 + if self.uses_inputs_embeds and self.is_first_pp_rank:
1558 + # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds
1559 + # overlay). Only first PP rank prepares them.
1560 if dummy_run:
1557 - # Obtain mm embeddings of correct shape for compiled model.
1561 + # Obtain embeddings of correct shape for compiled model.
1562 inputs_embeds = self.model_state.dummy_inputs_embeds(
1563 input_batch.num_tokens_after_padding
1564 )
1565 else:
1566 scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
1563 - if self.lora_config is not None:
1567 + if self.supports_mm_inputs and self.lora_config is not None:
1568 set_active_mm_loras(
1569 model=self.model,
1570 lora_manager=self.lora_manager,
@@ -1574,16 +1578,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1578 ) as ec_connector_output:
1579 if self.is_encoder_only:
1580 # Encode and publish, nothing else: this instance runs no
1577 - # language model, so the gather inside get_mm_embeddings
1581 + # language model, so the gather inside prepare_inputs_embeds
1582 # would build an inputs_embeds nobody reads -- and it
1583 # raises "Encoder cache miss" for any scheduled item this
1584 # instance did not encode, taking the engine down with it.
1585 self.model_state.execute_mm_encoder(scheduled_encoder_inputs)
1586 else:
1583 - inputs_embeds = self.model_state.get_mm_embeddings(
1587 + inputs_embeds = self.model_state.prepare_inputs_embeds(
1588 scheduled_encoder_inputs, input_batch, self.req_states
1589 )
1586 - if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
1590 + if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
1591 input_ids = None
1592
1593 if self.is_encoder_only:
vllm/v1/worker/gpu/model_states/__init__.py
+19 -7
@@ -7,6 +7,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config
7 from vllm.model_executor.layers.attention import Attention, CrossAttention
8 from vllm.v1.attention.backend import AttentionType
9 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
10 +from vllm.v1.worker.gpu.model_states.interface import ModelState
11
12
13 def init_model_state(
@@ -14,11 +15,22 @@ def init_model_state(
15 model: nn.Module,
16 encoder_cache: EncoderCache | None,
17 device: torch.device,
17 -):
18 +) -> ModelState:
19 + cls = resolve_model_state_cls(vllm_config, model)
20 +
21 + # Reject enable_prompt_embeds for states that would silently ignore it.
22 + if vllm_config.model_config.enable_prompt_embeds and not cls.supports_prompt_embeds:
23 + raise ValueError(f"--enable-prompt-embeds not supported with {cls.__name__}.")
24 +
25 + return cls(vllm_config, model, encoder_cache, device)
26 +
27 +
28 +def resolve_model_state_cls(
29 + vllm_config: VllmConfig, model: nn.Module
30 +) -> type[ModelState]:
31 # Let the model provide its own ModelState if it defines one.
32 if hasattr(model, "get_model_state_cls"):
20 - cls = model.get_model_state_cls()
21 - return cls(vllm_config, model, encoder_cache, device)
33 + return model.get_model_state_cls()
34
35 # Cross-attention encoder-decoder models (Whisper, CohereASR, NemotronParse, ...)
36 if any(isinstance(m, CrossAttention) for m in model.modules()):
@@ -26,7 +38,7 @@ def init_model_state(
38 EncoderDecoderModelState,
39 )
40
29 - return EncoderDecoderModelState(vllm_config, model, encoder_cache, device)
41 + return EncoderDecoderModelState
42
43 # Encoder-only attention is non-causal and needs no KV cache.
44 if any(
@@ -35,13 +47,13 @@ def init_model_state(
47 ):
48 from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState
49
38 - return EncoderOnlyModelState(vllm_config, model, encoder_cache, device)
50 + return EncoderOnlyModelState
51
52 if vllm_config.model_config.is_hybrid or vllm_config.model_config.is_attention_free:
53 from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
54
43 - return MambaHybridModelState(vllm_config, model, encoder_cache, device)
55 + return MambaHybridModelState
56
57 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
58
47 - return DefaultModelState(vllm_config, model, encoder_cache, device)
59 + return DefaultModelState
vllm/v1/worker/gpu/model_states/default.py
+59 -18
@@ -18,11 +18,14 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
18 from vllm.v1.worker.gpu.mm.rope import get_rope_state
19 from vllm.v1.worker.gpu.model_states.interface import ModelState
20 from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner
21 +from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState
22 from vllm.v1.worker.gpu.states import RequestState
23 from vllm.v1.worker.utils import AttentionGroup
24
25
26 class DefaultModelState(ModelState):
27 + supports_prompt_embeds = True
28 +
29 def __init__(
30 self,
31 vllm_config: VllmConfig,
@@ -32,6 +35,18 @@ class DefaultModelState(ModelState):
35 ):
36 super().__init__(vllm_config, model, encoder_cache, device)
37
38 + self.prompt_embeds_state: PromptEmbedsState | None = None
39 + if self.model_config.enable_prompt_embeds:
40 + self.prompt_embeds_state = PromptEmbedsState(
41 + self.max_num_reqs, self.inputs_embeds_size, self.dtype, self.device
42 + )
43 + if not self.supports_mm_inputs:
44 + # Persistent buffer analogous to encoder_runner.inputs_embeds.
45 + embeds_buffer_size = (self.max_num_tokens, self.inputs_embeds_size)
46 + self.inputs_embeds = torch.zeros(
47 + embeds_buffer_size, dtype=self.dtype, device=self.device
48 + )
49 +
50 self.rope_state = get_rope_state(
51 self.model_config,
52 model,
@@ -49,42 +64,69 @@ class DefaultModelState(ModelState):
64 def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
65 if self.rope_state is not None:
66 assert new_req_data.prefill_token_ids is not None
67 + # `prompt_embeds` is a passthrough modality with no grid info, but
68 + # M-RoPE assumes per-feature grids. Filter it out.
69 + mm_features = [
70 + f for f in new_req_data.mm_features if f.modality != "prompt_embeds"
71 + ]
72 self.rope_state.init_prefill_positions(
73 req_index,
74 self.model,
75 new_req_data.prefill_token_ids,
56 - mm_features=new_req_data.mm_features,
76 + mm_features=mm_features,
77 )
78 + if self.prompt_embeds_state is not None:
79 + self.prompt_embeds_state.add_request(req_index, new_req_data)
80 +
81 + def remove_request(self, req_id: str) -> None:
82 + if self.prompt_embeds_state is not None:
83 + self.prompt_embeds_state.remove_request(req_id)
84
85 def apply_staged_writes(self) -> None:
86 if self.rope_state is not None:
87 self.rope_state.apply_staged_writes()
88 + if self.prompt_embeds_state is not None:
89 + self.prompt_embeds_state.apply_staged_writes()
90
91 def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor:
92 """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
65 - return self.encoder_runner.inputs_embeds[:num_tokens]
93 + if self.supports_mm_inputs:
94 + return self.encoder_runner.inputs_embeds[:num_tokens]
95 + return self.inputs_embeds[:num_tokens]
96
67 - def get_mm_embeddings(
97 + def prepare_inputs_embeds(
98 self,
99 scheduled_encoder_inputs: dict[str, list[int]],
100 input_batch: InputBatch,
101 req_states: RequestState,
102 ) -> torch.Tensor:
73 - self.execute_mm_encoder(scheduled_encoder_inputs)
74 -
75 - mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)
76 - if self.mm_pruner is not None and mm_embeds:
77 - # EVS: recompute mrope positions for pruned media.
78 - mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states)
79 - # We must flush the staged rope updates for prepare_inputs() to pick up.
80 - self.apply_staged_writes()
81 -
103 # Use unpadded input_ids to match is_mm_embed size (num_tokens).
104 # input_batch.input_ids may be padded for CUDA graphs.
105 input_ids_unpadded = input_batch.input_ids[: input_batch.num_tokens]
85 - inputs_embeds = self.encoder_runner.get_inputs_embeds(
86 - input_ids_unpadded, mm_embeds, is_mm_embed
87 - )
106 +
107 + if self.supports_mm_inputs:
108 + self.execute_mm_encoder(scheduled_encoder_inputs)
109 +
110 + mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)
111 + if self.mm_pruner is not None and mm_embeds:
112 + # EVS: recompute mrope positions for pruned media.
113 + mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states)
114 + # We must flush the staged rope updates for prepare_inputs() to pick up.
115 + self.apply_staged_writes()
116 +
117 + inputs_embeds = self.encoder_runner.get_inputs_embeds(
118 + input_ids_unpadded, mm_embeds, is_mm_embed
119 + )
120 + else:
121 + input_embeddings = self.model.embed_input_ids(input_ids_unpadded)
122 + self.inputs_embeds[: input_embeddings.shape[0]] = input_embeddings
123 + inputs_embeds = self.inputs_embeds
124 +
125 + if self.prompt_embeds_state is not None:
126 + self.prompt_embeds_state.apply(
127 + input_batch, req_states.num_computed_tokens.gpu, inputs_embeds
128 + )
129 +
130 return inputs_embeds[: input_batch.num_tokens_after_padding]
131
132 def gather_mm_embeddings(
@@ -115,9 +157,8 @@ class DefaultModelState(ModelState):
157
158 def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
159 model_inputs = {}
118 - if self.supports_mm_inputs:
119 - inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens]
120 - model_inputs["inputs_embeds"] = inputs_embeds
160 + if self.supports_mm_inputs or self.prompt_embeds_state is not None:
161 + model_inputs["inputs_embeds"] = self.dummy_inputs_embeds(num_tokens)
162 if self.rope_state is not None:
163 model_inputs["positions"] = self.rope_state.get_positions(num_tokens)
164 return model_inputs
vllm/v1/worker/gpu/model_states/encoder_decoder.py
+5 -1
@@ -54,6 +54,10 @@ class EncoderDecoderModelState(ModelState):
54 device: torch.device,
55 ) -> None:
56 assert encoder_cache is not None
57 + if vllm_config.model_config.enable_prompt_embeds:
58 + raise ValueError(
59 + "--enable-prompt-embeds is not supported with encoder-decoder models."
60 + )
61 super().__init__(vllm_config, model, encoder_cache, device)
62
63 self.max_encoder_len = getattr(
@@ -67,7 +71,7 @@ class EncoderDecoderModelState(ModelState):
71
72 self.encoder_outputs: list[torch.Tensor] = []
73
70 - def get_mm_embeddings(
74 + def prepare_inputs_embeds(
75 self,
76 scheduled_encoder_inputs: dict[str, list[int]],
77 input_batch: InputBatch,
vllm/v1/worker/gpu/model_states/encoder_only.py
+3
@@ -34,6 +34,9 @@ class EncoderOnlyModelState(DefaultModelState):
34 the normal KV-backed path untouched.
35 """
36
37 + # The V2 pooling path is not wired for prompt embeds.
38 + supports_prompt_embeds = False
39 +
40 def __init__(
41 self,
42 vllm_config: VllmConfig,
vllm/v1/worker/gpu/model_states/interface.py
+6 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 -from typing import Any, cast
4 +from typing import Any, ClassVar, cast
5
6 import torch
7 import torch.nn as nn
@@ -43,6 +43,9 @@ class ModelSpecificAttnMetadata:
43
44
45 class ModelState(ABC):
46 + supports_prompt_embeds: ClassVar[bool] = False
47 + """Whether this state implements user-provided prompt embeddings."""
48 +
49 def __init__(
50 self,
51 vllm_config: VllmConfig,
@@ -153,12 +156,13 @@ class ModelState(ABC):
156 return None
157
158 @abstractmethod
156 - def get_mm_embeddings(
159 + def prepare_inputs_embeds(
160 self,
161 scheduled_encoder_inputs: dict[str, list[int]],
162 input_batch: InputBatch,
163 req_states: RequestState,
164 ) -> torch.Tensor | None:
165 + """Prepare the ``inputs_embeds`` tensor for the current forward pass."""
166 raise NotImplementedError
167
168 def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None:
vllm/v1/worker/gpu/model_states/prompt_embeds.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.triton_utils import tl, triton
6 +from vllm.utils.torch_utils import async_tensor_h2d
7 +from vllm.v1.core.sched.output import NewRequestData
8 +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
9 +from vllm.v1.worker.gpu.input_batch import InputBatch
10 +
11 +TOKEN_BLOCK = 16
12 +
13 +
14 +class PromptEmbedsState:
15 + """GPU-side state for user-provided prompt embeddings.
16 +
17 + Each request's embeddings are copied to the GPU once at `add_request`,
18 + off the per-step hot path. A per-request pointer table (UVA) then lets a
19 + single triton kernel overlay all scheduled prompt-embeds rows onto
20 + `inputs_embeds` each step, with no python loops or per-request H2D copies.
21 + """
22 +
23 + def __init__(
24 + self,
25 + max_num_reqs: int,
26 + hidden_size: int,
27 + dtype: torch.dtype,
28 + device: torch.device,
29 + ):
30 + self.hidden_size = hidden_size
31 + self.dtype = dtype
32 + self.device = device
33 +
34 + # req_id -> (embeds, is_token_ids mask or None). Holds the references
35 + # that keep the pointer table below valid.
36 + self.gpu_tensors: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {}
37 +
38 + # Indexed by req_state index. Stale entries after removal are
39 + # harmless: add_request rewrites all fields for every index it claims.
40 + self.embeds_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
41 + self.mask_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
42 + self.embeds_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
43 +
44 + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
45 + prompt_embeds = new_req_data.prompt_embeds
46 + if prompt_embeds is None:
47 + self.gpu_tensors.pop(new_req_data.req_id, None)
48 + self.embeds_lens.np[req_index] = 0
49 + return
50 +
51 + embeds = async_tensor_h2d(prompt_embeds, device=self.device, dtype=self.dtype)
52 + embeds = embeds.contiguous()
53 + is_token_ids = new_req_data.prompt_is_token_ids
54 + mask = None
55 + if is_token_ids is not None:
56 + mask = async_tensor_h2d(is_token_ids, device=self.device, dtype=torch.uint8)
57 + self.gpu_tensors[new_req_data.req_id] = (embeds, mask)
58 + self.embeds_ptrs.np[req_index] = embeds.data_ptr()
59 + self.mask_ptrs.np[req_index] = 0 if mask is None else mask.data_ptr()
60 + self.embeds_lens.np[req_index] = embeds.shape[0]
61 +
62 + def remove_request(self, req_id: str) -> None:
63 + self.gpu_tensors.pop(req_id, None)
64 +
65 + def apply_staged_writes(self) -> None:
66 + self.embeds_ptrs.copy_to_uva()
67 + self.mask_ptrs.copy_to_uva()
68 + self.embeds_lens.copy_to_uva()
69 +
70 + def apply(
71 + self,
72 + input_batch: InputBatch,
73 + num_computed_tokens: torch.Tensor,
74 + inputs_embeds: torch.Tensor,
75 + ) -> None:
76 + """Overlay prompt embeddings onto `inputs_embeds` for the batch."""
77 + if not self.gpu_tensors:
78 + return
79 + # The kernel reinterprets raw source pointers as inputs_embeds' dtype.
80 + assert inputs_embeds.dtype == self.dtype
81 + num_reqs = input_batch.num_reqs
82 + max_query_len = int(input_batch.num_scheduled_tokens.max())
83 + grid = (num_reqs, triton.cdiv(max_query_len, TOKEN_BLOCK))
84 + _apply_prompt_embeds_kernel[grid](
85 + inputs_embeds,
86 + inputs_embeds.stride(0),
87 + self.embeds_ptrs.gpu,
88 + self.mask_ptrs.gpu,
89 + self.embeds_lens.gpu,
90 + input_batch.idx_mapping,
91 + input_batch.query_start_loc,
92 + num_computed_tokens,
93 + self.hidden_size,
94 + TOKEN_BLOCK=TOKEN_BLOCK,
95 + BLOCK_SIZE=1024,
96 + )
97 +
98 +
99 +@triton.jit
100 +def _apply_prompt_embeds_kernel(
101 + inputs_embeds_ptr,
102 + inputs_embeds_stride,
103 + embeds_ptrs_ptr, # int64 [max_num_reqs], device pointers (0-len = unused)
104 + mask_ptrs_ptr, # int64 [max_num_reqs], 0 = no is-token-ids mask
105 + embeds_lens_ptr, # int32 [max_num_reqs]
106 + idx_mapping_ptr,
107 + query_start_loc_ptr,
108 + num_computed_tokens_ptr,
109 + hidden_size: tl.constexpr,
110 + TOKEN_BLOCK: tl.constexpr,
111 + BLOCK_SIZE: tl.constexpr,
112 +):
113 + batch_idx = tl.program_id(0)
114 + req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
115 + embeds_len = tl.load(embeds_lens_ptr + req_state_idx)
116 + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
117 + if num_computed >= embeds_len:
118 + # No prompt embeds for this request, or they are fully consumed.
119 + return
120 +
121 + query_start = tl.load(query_start_loc_ptr + batch_idx)
122 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
123 + num_rows = tl.minimum(query_end - query_start, embeds_len - num_computed)
124 +
125 + t_start = tl.program_id(1) * TOKEN_BLOCK
126 + if t_start >= num_rows:
127 + return
128 +
129 + src_ptr = tl.load(embeds_ptrs_ptr + req_state_idx).to(
130 + tl.pointer_type(inputs_embeds_ptr.dtype.element_ty)
131 + )
132 + mask_int = tl.load(mask_ptrs_ptr + req_state_idx)
133 + mask_ptr = mask_int.to(tl.pointer_type(tl.int8))
134 +
135 + for t_offset in tl.static_range(TOKEN_BLOCK):
136 + t = t_start + t_offset
137 + if t < num_rows:
138 + src_row = (num_computed + t).to(tl.int64)
139 + is_token_id = 0
140 + if mask_int != 0:
141 + is_token_id = tl.load(mask_ptr + src_row).to(tl.int32)
142 + if is_token_id == 0:
143 + dst_row = (query_start + t).to(tl.int64)
144 + for h in tl.range(0, hidden_size, BLOCK_SIZE):
145 + offs = h + tl.arange(0, BLOCK_SIZE)
146 + h_mask = offs < hidden_size
147 + row = tl.load(src_ptr + src_row * hidden_size + offs, mask=h_mask)
148 + tl.store(
149 + inputs_embeds_ptr + dst_row * inputs_embeds_stride + offs,
150 + row,
151 + mask=h_mask,
152 + )