[ModelRunnerV2] Support prompt embeds (#42963)

Signed-off-by: gcanlin <[email protected]>
Signed-off-by: Canlin Guo <[email protected]>
Signed-off-by: Nick Hill <[email protected]>
Co-authored-by: Nick Hill <[email protected]>

tests/v1/worker/test_encoder_runner.py

@@ -15,7 +15,13 @@ import numpy as np
 import pytest
 import torch
 
-from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    MultiModalSharedField,
+    PlaceholderRange,
+)
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
 from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -25,6 +31,25 @@ pytestmark = pytest.mark.cpu_test
 HIDDEN = 4
 
 
+def _model_state(cache: EncoderCache) -> MagicMock:
+    """A mock ModelState backed by a real EncoderCache."""
+    state = MagicMock()
+    state.encoder_cache = cache
+    state.device = torch.device("cpu")
+    return state
+
+
+def _embeds_item(embeds: torch.Tensor) -> MultiModalKwargsItem:
+    """A `prompt_embeds` kwargs item, as the HF renderer builds it."""
+    return MultiModalKwargsItem(
+        {
+            "embedding": MultiModalFieldElem(
+                data=embeds, field=MultiModalSharedField(batch_size=1)
+            )
+        }
+    )
+
+
 def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec:
     return MultiModalFeatureSpec(
         data=None,
@@ -197,8 +222,7 @@ def test_execute_mm_encoder_caches_outputs_without_gathering():
     items the connector already holds, and a producer has no load path).
     """
     cache = EncoderCache()
-    state = MagicMock()
-    state.encoder_cache = cache
+    state = _model_state(cache)
     embedding = torch.ones(2, HIDDEN)
     # (mm_hashes, [(modality, kwargs item), ...]), as prepare_mm_inputs returns.
     state.encoder_runner.prepare_mm_inputs.return_value = (
@@ -216,8 +240,7 @@ def test_execute_mm_encoder_caches_outputs_without_gathering():
 def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
     """A step that schedules no encoder input must not touch the encoder."""
     cache = EncoderCache()
-    state = MagicMock()
-    state.encoder_cache = cache
+    state = _model_state(cache)
     state.encoder_runner.prepare_mm_inputs.return_value = ([], [])
 
     ModelState.execute_mm_encoder(state, {})
@@ -226,6 +249,68 @@ def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
     state.encoder_runner.execute_mm_encoder.assert_not_called()
 
 
+def _pe_feature(identifier: str, embeds: torch.Tensor, offset: int = 0):
+    return MultiModalFeatureSpec(
+        data=_embeds_item(embeds),
+        modality="prompt_embeds",
+        identifier=identifier,
+        mm_position=PlaceholderRange(offset=offset, length=embeds.shape[0]),
+    )
+
+
+def test_prepare_mm_inputs_passes_prompt_embeds_through():
+    """`prompt_embeds` is already in embedding space, so no encoder may run.
+
+    The renderer delivers prompt_embeds mixed with real media as an ordinary MM
+    modality. prepare_mm_inputs must cache the tensor directly and keep it out
+    of the encoder batch -- the vision encoder cannot consume it, and a missing
+    cache entry makes the subsequent gather raise "Encoder cache miss".
+    """
+    prompt_embeds = torch.arange(2 * HIDDEN, dtype=torch.float32).view(2, HIDDEN)
+    image_feature = MultiModalFeatureSpec(
+        data=MagicMock(),
+        modality="image",
+        identifier="hash_img",
+        mm_position=PlaceholderRange(offset=2, length=2),
+    )
+    runner = _make_runner(
+        [_pe_feature("hash_pe", prompt_embeds), image_feature], cached=[]
+    )
+
+    mm_hashes, mm_kwargs = runner.prepare_mm_inputs({"req0": [0, 1]})
+
+    # Only the image remains for the encoder; the embeds are already cached.
+    assert mm_hashes == ["hash_img"]
+    assert [modality for modality, _ in mm_kwargs] == ["image"]
+    assert torch.equal(runner.encoder_cache.encoder_outputs["hash_pe"], prompt_embeds)
+
+
+def test_prepare_mm_inputs_skips_cached_prompt_embeds():
+    """A prompt_embeds item already in the cache must not be re-uploaded."""
+    prompt_embeds = torch.ones(3, HIDDEN)
+    feature = _pe_feature("hash_pe", prompt_embeds)
+    runner = _make_runner([feature], cached=[feature])
+    sentinel = runner.encoder_cache.encoder_outputs["hash_pe"]
+
+    mm_hashes, mm_kwargs = runner.prepare_mm_inputs({"req0": [0]})
+
+    assert mm_hashes == [] and mm_kwargs == []
+    assert runner.encoder_cache.encoder_outputs["hash_pe"] is sentinel
+
+
+def test_execute_mm_encoder_skips_encoder_for_prompt_embeds_only():
+    """A batch of nothing but prompt_embeds must not invoke the encoder."""
+    prompt_embeds = torch.ones(3, HIDDEN)
+    runner = _make_runner([_pe_feature("hash_pe", prompt_embeds)], cached=[])
+    state = _model_state(runner.encoder_cache)
+    state.encoder_runner.prepare_mm_inputs.side_effect = runner.prepare_mm_inputs
+
+    ModelState.execute_mm_encoder(state, {"req0": [0]})
+
+    state.encoder_runner.execute_mm_encoder.assert_not_called()
+    assert torch.equal(runner.encoder_cache.encoder_outputs["hash_pe"], prompt_embeds)
+
+
 def test_encoder_timing_stats_registry():
     runner = _make_runner([], [])
     runner.enable_timing = True

tests/v1/worker/test_gpu_model_runner.py

@@ -68,6 +68,16 @@ NUM_BLOCKS = 10
 DEVICE_TYPE = current_platform.device_type
 
 
+@pytest.fixture(autouse=True)
+def _restore_default_dtype():
+    """Several tests here set the process-wide default dtype to float16 and
+    previously leaked it, corrupting later float-sensitive tests in the same
+    pytest process (torch.randn silently produced fp16)."""
+    old = torch.get_default_dtype()
+    yield
+    torch.set_default_dtype(old)
+
+
 def initialize_kv_cache(runner: GPUModelRunner):
     """
     Only perform necessary steps in GPUModelRunner.initialize_kv_cache()

tests/v1/worker/test_prompt_embeds_state.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the Model Runner V2 prompt-embeds overlay (PromptEmbedsState).
+
+The overlay kernel reads each request's GPU-resident prompt embeddings through
+a per-request pointer table and writes the rows scheduled this step into
+`inputs_embeds`, honoring chunked prefill (`num_computed_tokens` offset), the
+prompt/decode boundary (rows past the embeds length untouched), and the
+mixed-mode `prompt_is_token_ids` mask (token-id rows keep the base embedding).
+"""
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+pytest.importorskip("triton")
+if not torch.cuda.is_available():
+    pytest.skip(
+        "CUDA required for prompt-embeds overlay tests", allow_module_level=True
+    )
+
+from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState
+
+HIDDEN = 24
+MAX_NUM_REQS = 8
+DEVICE = torch.device("cuda")
+
+
+@dataclass
+class _NewReqData:
+    req_id: str
+    prompt_embeds: torch.Tensor | None
+    prompt_is_token_ids: list[bool] | None = None
+
+
+@dataclass
+class _Batch:
+    num_reqs: int
+    num_scheduled_tokens: torch.Tensor  # np-like, only .max() is used
+    idx_mapping: torch.Tensor
+    query_start_loc: torch.Tensor
+
+
+def _make_state() -> PromptEmbedsState:
+    return PromptEmbedsState(MAX_NUM_REQS, HIDDEN, torch.float32, DEVICE)
+
+
+def _batch(num_scheduled: list[int], idx_mapping: list[int]) -> _Batch:
+    query_start_loc = [0]
+    for n in num_scheduled:
+        query_start_loc.append(query_start_loc[-1] + n)
+    return _Batch(
+        num_reqs=len(num_scheduled),
+        num_scheduled_tokens=torch.tensor(num_scheduled, dtype=torch.int32),
+        idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64, device=DEVICE),
+        query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE),
+    )
+
+
+def _apply(
+    state: PromptEmbedsState,
+    batch: _Batch,
+    num_computed: list[int],
+    num_tokens: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run the overlay on a fresh base buffer; return (result, base)."""
+    num_computed_tokens = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE)
+    for batch_idx, req_index in enumerate(batch.idx_mapping.tolist()):
+        num_computed_tokens[req_index] = num_computed[batch_idx]
+    base = torch.randn(num_tokens, HIDDEN, dtype=torch.float32, device=DEVICE)
+    inputs_embeds = base.clone()
+    state.apply(batch, num_computed_tokens, inputs_embeds)
+    torch.accelerator.synchronize()
+    return inputs_embeds, base
+
+
+def test_overlay_chunked_prefill_and_decode():
+    """Rows within the embeds range come from prompt_embeds at the
+    num_computed offset; requests without embeds and requests past their
+    embeds length (decode) keep the base embedding."""
+    state = _make_state()
+    embeds_a = torch.randn(6, HIDDEN, dtype=torch.float32)
+    embeds_b = torch.randn(5, HIDDEN, dtype=torch.float32)
+    state.add_request(0, _NewReqData("a", embeds_a))
+    state.add_request(1, _NewReqData("b", embeds_b))
+    state.add_request(2, _NewReqData("c", None))
+    state.apply_staged_writes()
+
+    # a: chunk [2, 6) of its embeds; b: fully decoded; c: no embeds.
+    batch = _batch(num_scheduled=[4, 1, 3], idx_mapping=[0, 1, 2])
+    out, base = _apply(state, batch, num_computed=[2, 7, 1], num_tokens=8)
+
+    torch.testing.assert_close(out[0:4], embeds_a[2:6].to(DEVICE))
+    torch.testing.assert_close(out[4:8], base[4:8])
+
+
+def test_overlay_clamps_to_embeds_length():
+    """A window straddling the end of the prompt embeds writes only the
+    in-range rows (e.g. final prefill chunk + sampled token)."""
+    state = _make_state()
+    embeds = torch.randn(4, HIDDEN, dtype=torch.float32)
+    state.add_request(3, _NewReqData("a", embeds))
+    state.apply_staged_writes()
+
+    batch = _batch(num_scheduled=[3], idx_mapping=[3])
+    out, base = _apply(state, batch, num_computed=[2], num_tokens=3)
+
+    torch.testing.assert_close(out[0:2], embeds[2:4].to(DEVICE))
+    torch.testing.assert_close(out[2:3], base[2:3])
+
+
+def test_overlay_respects_is_token_ids_mask():
+    """Mixed mode: positions marked as real token ids keep the base
+    embedding; only embed positions are overwritten."""
+    state = _make_state()
+    embeds = torch.randn(5, HIDDEN, dtype=torch.float32)
+    is_token_ids = [True, False, False, True, False]
+    state.add_request(0, _NewReqData("a", embeds, is_token_ids))
+    state.apply_staged_writes()
+
+    batch = _batch(num_scheduled=[5], idx_mapping=[0])
+    out, base = _apply(state, batch, num_computed=[0], num_tokens=5)
+
+    embeds_gpu = embeds.to(DEVICE)
+    for pos, is_token in enumerate(is_token_ids):
+        expected = base[pos] if is_token else embeds_gpu[pos]
+        torch.testing.assert_close(out[pos], expected)
+
+
+def test_index_reuse_clears_stale_entry():
+    """A request added at a previously-used index without embeds must not
+    inherit the prior occupant's pointer-table entry."""
+    state = _make_state()
+    state.add_request(0, _NewReqData("a", torch.randn(4, HIDDEN, dtype=torch.float32)))
+    # A second live embeds request so the kernel actually launches (the
+    # overlay is skipped entirely when no request holds embeds).
+    other = torch.randn(2, HIDDEN, dtype=torch.float32)
+    state.add_request(1, _NewReqData("other", other))
+    state.apply_staged_writes()
+    state.remove_request("a")
+    state.add_request(0, _NewReqData("b", None))
+    state.apply_staged_writes()
+
+    batch = _batch(num_scheduled=[2, 2], idx_mapping=[0, 1])
+    out, base = _apply(state, batch, num_computed=[0, 0], num_tokens=4)
+
+    torch.testing.assert_close(out[0:2], base[0:2])
+    torch.testing.assert_close(out[2:4], other.to(DEVICE))

vllm/config/model.py

@@ -926,6 +926,12 @@ class ModelConfig:
                 f"got {type(self.max_model_len).__name__}: {self.max_model_len!r}. "
                 "Example: max_model_len=2048"
             )
+        if self.enable_prompt_embeds and self.is_encoder_decoder:
+            # No encoder-decoder model accepts `inputs_embeds`; their decoders
+            # embed `input_ids` internally.
+            raise ValueError(
+                "--enable-prompt-embeds is not supported with encoder-decoder models."
+            )
         return self
 
     def _resolve_mm_device_do_normalize(

vllm/config/vllm.py

@@ -2449,9 +2449,6 @@ class VllmConfig:
         ):
             unsupported.append("custom logits processors")
 
-        if model_config is not None and model_config.enable_prompt_embeds:
-            unsupported.append("prompt embeds")
-
         if self.cache_config.kv_sharing_fast_prefill:
             # Will be added by https://github.com/vllm-project/vllm/pull/35045
             unsupported.append("KV sharing fast prefill")

vllm/model_executor/models/diffusion_gemma.py

@@ -776,7 +776,7 @@ class DiffusionGemmaModelState(ModelState):
     ) -> None:
         super().__init__(vllm_config, model, encoder_cache, device)
 
-        # Per-step MM data produced by get_mm_embeddings and consumed by
+        # Per-step MM data produced by prepare_inputs_embeds and consumed by
         # prepare_inputs.  Stored as raw (mm_embeds, is_mm_embed) so that
         # prepare_inputs can call embed_input_ids directly into the
         # persistent _inputs_embeds_buf, avoiding the intermediate copy
@@ -874,7 +874,7 @@ class DiffusionGemmaModelState(ModelState):
         if idx is not None:
             self.diffusion_states.remove_request(idx)
 
-    def get_mm_embeddings(
+    def prepare_inputs_embeds(
         self,
         scheduled_encoder_inputs: dict[str, list[int]],
         input_batch: InputBatch,

vllm/model_executor/models/longcat_flash_ngram.py

@@ -273,6 +273,9 @@ class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
 
 
 class LongcatNgramModelState(DefaultModelState):
+    # prepare_inputs builds its own inputs_embeds from n-gram token embeddings.
+    supports_prompt_embeds = False
+
     """Per-request n-gram token history for LongCat-Flash-Lite.
 
     Maintains a small CPU-side per-slot context (last ``n-1`` processed tokens)

vllm/v1/core/sched/output.py

@@ -68,6 +68,14 @@ class NewRequestData:
             prefill_token_ids=prefill_token_ids,
         )
 
+    @property
+    def prompt_len(self) -> int:
+        if self.prompt_token_ids is not None:
+            return len(self.prompt_token_ids)
+        if self.prompt_embeds is not None:
+            return self.prompt_embeds.shape[0]
+        return 0
+
     def __repr__(self) -> str:
         prompt_embeds_shape = (
             self.prompt_embeds.shape if self.prompt_embeds is not None else None

vllm/v1/request.py

@@ -142,11 +142,19 @@ class Request:
             prompt_token_ids, prompt_embeds
         )
         self._output_token_ids: list[int] = []
-        self._all_token_ids: list[int] = (
-            self.prompt_token_ids.copy()
-            if self.prompt_token_ids is not None
-            else [0] * self.num_prompt_tokens
-        )
+        if self.prompt_token_ids is None:
+            self._all_token_ids: list[int] = [0] * self.num_prompt_tokens
+        elif self.prompt_is_token_ids is None:
+            self._all_token_ids = self.prompt_token_ids.copy()
+        else:
+            # Mixed-mode prompt: positions covered by prompt_embeds hold a sentinel
+            # special token id that may lie outside the embedding. Zero them, matching
+            # the no-token-ids case above, so embedding gathers over these placeholder
+            # ids stay in bounds; the actual inputs come from prompt_embeds.
+            self._all_token_ids = [
+                t if is_tok else 0
+                for t, is_tok in zip(self.prompt_token_ids, self.prompt_is_token_ids)
+            ]
 
         # Used in async scheduling.
         self.num_output_placeholders = 0

vllm/v1/worker/gpu/mm/encoder_runner.py

@@ -18,7 +18,7 @@ from vllm.multimodal.utils import (
     group_and_batch_mm_kwargs,
     set_mm_embedding_modality,
 )
-from vllm.utils.torch_utils import PIN_MEMORY
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.utils import (
     EncoderTimingStats,
@@ -91,6 +91,17 @@ class EncoderRunner:
                     continue
                 if mm_feature.identifier in self.encoder_cache.encoder_outputs:
                     continue
+                if mm_feature.modality == "prompt_embeds":
+                    # Passthrough modality: the tensor is already in the
+                    # model's embedding space, so no encoder runs. Cache it
+                    # directly so gather_mm_embeddings splices it via the
+                    # standard is_mm_embed path.
+                    embeds = mm_feature.data["embedding"].data
+                    assert isinstance(embeds, torch.Tensor)
+                    self.encoder_cache.encoder_outputs[mm_feature.identifier] = (
+                        async_tensor_h2d(embeds, device=self.device)
+                    )
+                    continue
                 mm_hashes.append(mm_feature.identifier)
                 mm_kwargs.append((mm_feature.modality, mm_feature.data))
 

vllm/v1/worker/gpu/model_runner.py

@@ -46,6 +46,7 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     initialize_mamba_ssu_backend,
 )
 from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.models.interfaces import requires_raw_input_tokens
 from vllm.model_executor.offloader import (
     create_offloader,
     get_offloader,
@@ -231,6 +232,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
             self.model_config
         )
+        self.uses_inputs_embeds = (
+            self.supports_mm_inputs or self.model_config.enable_prompt_embeds
+        )
         self.encoder_cache = None
         if self.supports_mm_inputs and self.is_first_pp_rank:
             self.encoder_cache = EncoderCache()
@@ -956,7 +960,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
     def add_requests(self, scheduler_output: SchedulerOutput) -> None:
         for new_req_data in scheduler_output.scheduled_new_reqs:
-            assert new_req_data.prompt_token_ids is not None
             assert new_req_data.prefill_token_ids is not None
             req_id = new_req_data.req_id
 
@@ -965,7 +968,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             # with the updated prompt_token_ids and mm_features.
             self._remove_request(req_id)
 
-            prompt_len = len(new_req_data.prompt_token_ids)
+            prompt_len = new_req_data.prompt_len
             sampling_params = new_req_data.sampling_params
             self.req_states.add_request(
                 req_id=req_id,
@@ -980,6 +983,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
             if self.pooling_runner is not None:
                 assert new_req_data.pooling_params is not None
+                assert new_req_data.prompt_token_ids is not None
                 self.pooling_runner.add_request(
                     req_id,
                     req_index,
@@ -1550,17 +1554,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         input_ids = input_batch.input_ids
         inputs_embeds = None
         ec_connector_output = None
-        if self.supports_mm_inputs and self.is_first_pp_rank:
-            # Run MM encoder (if needed) and get multimodal embeddings.
-            # Only first PP rank prepares multimodal embeddings.
+        if self.uses_inputs_embeds and self.is_first_pp_rank:
+            # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds
+            # overlay). Only first PP rank prepares them.
             if dummy_run:
-                # Obtain mm embeddings of correct shape for compiled model.
+                # Obtain embeddings of correct shape for compiled model.
                 inputs_embeds = self.model_state.dummy_inputs_embeds(
                     input_batch.num_tokens_after_padding
                 )
             else:
                 scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
-                if self.lora_config is not None:
+                if self.supports_mm_inputs and self.lora_config is not None:
                     set_active_mm_loras(
                         model=self.model,
                         lora_manager=self.lora_manager,
@@ -1574,16 +1578,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 ) as ec_connector_output:
                     if self.is_encoder_only:
                         # Encode and publish, nothing else: this instance runs no
-                        # language model, so the gather inside get_mm_embeddings
+                        # language model, so the gather inside prepare_inputs_embeds
                         # would build an inputs_embeds nobody reads -- and it
                         # raises "Encoder cache miss" for any scheduled item this
                         # instance did not encode, taking the engine down with it.
                         self.model_state.execute_mm_encoder(scheduled_encoder_inputs)
                     else:
-                        inputs_embeds = self.model_state.get_mm_embeddings(
+                        inputs_embeds = self.model_state.prepare_inputs_embeds(
                             scheduled_encoder_inputs, input_batch, self.req_states
                         )
-            if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
+            if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
                 input_ids = None
 
         if self.is_encoder_only:

vllm/v1/worker/gpu/model_states/__init__.py

@@ -7,6 +7,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.model_executor.layers.attention import Attention, CrossAttention
 from vllm.v1.attention.backend import AttentionType
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
+from vllm.v1.worker.gpu.model_states.interface import ModelState
 
 
 def init_model_state(
@@ -14,11 +15,22 @@ def init_model_state(
     model: nn.Module,
     encoder_cache: EncoderCache | None,
     device: torch.device,
-):
+) -> ModelState:
+    cls = resolve_model_state_cls(vllm_config, model)
+
+    # Reject enable_prompt_embeds for states that would silently ignore it.
+    if vllm_config.model_config.enable_prompt_embeds and not cls.supports_prompt_embeds:
+        raise ValueError(f"--enable-prompt-embeds not supported with {cls.__name__}.")
+
+    return cls(vllm_config, model, encoder_cache, device)
+
+
+def resolve_model_state_cls(
+    vllm_config: VllmConfig, model: nn.Module
+) -> type[ModelState]:
     # Let the model provide its own ModelState if it defines one.
     if hasattr(model, "get_model_state_cls"):
-        cls = model.get_model_state_cls()
-        return cls(vllm_config, model, encoder_cache, device)
+        return model.get_model_state_cls()
 
     # Cross-attention encoder-decoder models (Whisper, CohereASR, NemotronParse, ...)
     if any(isinstance(m, CrossAttention) for m in model.modules()):
@@ -26,7 +38,7 @@ def init_model_state(
             EncoderDecoderModelState,
         )
 
-        return EncoderDecoderModelState(vllm_config, model, encoder_cache, device)
+        return EncoderDecoderModelState
 
     # Encoder-only attention is non-causal and needs no KV cache.
     if any(
@@ -35,13 +47,13 @@ def init_model_state(
     ):
         from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState
 
-        return EncoderOnlyModelState(vllm_config, model, encoder_cache, device)
+        return EncoderOnlyModelState
 
     if vllm_config.model_config.is_hybrid or vllm_config.model_config.is_attention_free:
         from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
 
-        return MambaHybridModelState(vllm_config, model, encoder_cache, device)
+        return MambaHybridModelState
 
     from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 
-    return DefaultModelState(vllm_config, model, encoder_cache, device)
+    return DefaultModelState

vllm/v1/worker/gpu/model_states/default.py

@@ -18,11 +18,14 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.mm.rope import get_rope_state
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner
+from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState
 from vllm.v1.worker.gpu.states import RequestState
 from vllm.v1.worker.utils import AttentionGroup
 
 
 class DefaultModelState(ModelState):
+    supports_prompt_embeds = True
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -32,6 +35,18 @@ class DefaultModelState(ModelState):
     ):
         super().__init__(vllm_config, model, encoder_cache, device)
 
+        self.prompt_embeds_state: PromptEmbedsState | None = None
+        if self.model_config.enable_prompt_embeds:
+            self.prompt_embeds_state = PromptEmbedsState(
+                self.max_num_reqs, self.inputs_embeds_size, self.dtype, self.device
+            )
+            if not self.supports_mm_inputs:
+                # Persistent buffer analogous to encoder_runner.inputs_embeds.
+                embeds_buffer_size = (self.max_num_tokens, self.inputs_embeds_size)
+                self.inputs_embeds = torch.zeros(
+                    embeds_buffer_size, dtype=self.dtype, device=self.device
+                )
+
         self.rope_state = get_rope_state(
             self.model_config,
             model,
@@ -49,42 +64,69 @@ class DefaultModelState(ModelState):
     def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
         if self.rope_state is not None:
             assert new_req_data.prefill_token_ids is not None
+            # `prompt_embeds` is a passthrough modality with no grid info, but
+            # M-RoPE assumes per-feature grids. Filter it out.
+            mm_features = [
+                f for f in new_req_data.mm_features if f.modality != "prompt_embeds"
+            ]
             self.rope_state.init_prefill_positions(
                 req_index,
                 self.model,
                 new_req_data.prefill_token_ids,
-                mm_features=new_req_data.mm_features,
+                mm_features=mm_features,
             )
+        if self.prompt_embeds_state is not None:
+            self.prompt_embeds_state.add_request(req_index, new_req_data)
+
+    def remove_request(self, req_id: str) -> None:
+        if self.prompt_embeds_state is not None:
+            self.prompt_embeds_state.remove_request(req_id)
 
     def apply_staged_writes(self) -> None:
         if self.rope_state is not None:
             self.rope_state.apply_staged_writes()
+        if self.prompt_embeds_state is not None:
+            self.prompt_embeds_state.apply_staged_writes()
 
     def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor:
         """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
-        return self.encoder_runner.inputs_embeds[:num_tokens]
+        if self.supports_mm_inputs:
+            return self.encoder_runner.inputs_embeds[:num_tokens]
+        return self.inputs_embeds[:num_tokens]
 
-    def get_mm_embeddings(
+    def prepare_inputs_embeds(
         self,
         scheduled_encoder_inputs: dict[str, list[int]],
         input_batch: InputBatch,
         req_states: RequestState,
     ) -> torch.Tensor:
-        self.execute_mm_encoder(scheduled_encoder_inputs)
-
-        mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)
-        if self.mm_pruner is not None and mm_embeds:
-            # EVS: recompute mrope positions for pruned media.
-            mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states)
-            # We must flush the staged rope updates for prepare_inputs() to pick up.
-            self.apply_staged_writes()
-
         # Use unpadded input_ids to match is_mm_embed size (num_tokens).
         # input_batch.input_ids may be padded for CUDA graphs.
         input_ids_unpadded = input_batch.input_ids[: input_batch.num_tokens]
-        inputs_embeds = self.encoder_runner.get_inputs_embeds(
-            input_ids_unpadded, mm_embeds, is_mm_embed
-        )
+
+        if self.supports_mm_inputs:
+            self.execute_mm_encoder(scheduled_encoder_inputs)
+
+            mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)
+            if self.mm_pruner is not None and mm_embeds:
+                # EVS: recompute mrope positions for pruned media.
+                mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states)
+                # We must flush the staged rope updates for prepare_inputs() to pick up.
+                self.apply_staged_writes()
+
+            inputs_embeds = self.encoder_runner.get_inputs_embeds(
+                input_ids_unpadded, mm_embeds, is_mm_embed
+            )
+        else:
+            input_embeddings = self.model.embed_input_ids(input_ids_unpadded)
+            self.inputs_embeds[: input_embeddings.shape[0]] = input_embeddings
+            inputs_embeds = self.inputs_embeds
+
+        if self.prompt_embeds_state is not None:
+            self.prompt_embeds_state.apply(
+                input_batch, req_states.num_computed_tokens.gpu, inputs_embeds
+            )
+
         return inputs_embeds[: input_batch.num_tokens_after_padding]
 
     def gather_mm_embeddings(
@@ -115,9 +157,8 @@ class DefaultModelState(ModelState):
 
     def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
         model_inputs = {}
-        if self.supports_mm_inputs:
-            inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens]
-            model_inputs["inputs_embeds"] = inputs_embeds
+        if self.supports_mm_inputs or self.prompt_embeds_state is not None:
+            model_inputs["inputs_embeds"] = self.dummy_inputs_embeds(num_tokens)
         if self.rope_state is not None:
             model_inputs["positions"] = self.rope_state.get_positions(num_tokens)
         return model_inputs

vllm/v1/worker/gpu/model_states/encoder_decoder.py

@@ -54,6 +54,10 @@ class EncoderDecoderModelState(ModelState):
         device: torch.device,
     ) -> None:
         assert encoder_cache is not None
+        if vllm_config.model_config.enable_prompt_embeds:
+            raise ValueError(
+                "--enable-prompt-embeds is not supported with encoder-decoder models."
+            )
         super().__init__(vllm_config, model, encoder_cache, device)
 
         self.max_encoder_len = getattr(
@@ -67,7 +71,7 @@ class EncoderDecoderModelState(ModelState):
 
         self.encoder_outputs: list[torch.Tensor] = []
 
-    def get_mm_embeddings(
+    def prepare_inputs_embeds(
         self,
         scheduled_encoder_inputs: dict[str, list[int]],
         input_batch: InputBatch,

vllm/v1/worker/gpu/model_states/encoder_only.py

@@ -34,6 +34,9 @@ class EncoderOnlyModelState(DefaultModelState):
     the normal KV-backed path untouched.
     """
 
+    # The V2 pooling path is not wired for prompt embeds.
+    supports_prompt_embeds = False
+
     def __init__(
         self,
         vllm_config: VllmConfig,

vllm/v1/worker/gpu/model_states/interface.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from abc import ABC, abstractmethod
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import torch
 import torch.nn as nn
@@ -43,6 +43,9 @@ class ModelSpecificAttnMetadata:
 
 
 class ModelState(ABC):
+    supports_prompt_embeds: ClassVar[bool] = False
+    """Whether this state implements user-provided prompt embeddings."""
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -153,12 +156,13 @@ class ModelState(ABC):
         return None
 
     @abstractmethod
-    def get_mm_embeddings(
+    def prepare_inputs_embeds(
         self,
         scheduled_encoder_inputs: dict[str, list[int]],
         input_batch: InputBatch,
         req_states: RequestState,
     ) -> torch.Tensor | None:
+        """Prepare the ``inputs_embeds`` tensor for the current forward pass."""
         raise NotImplementedError
 
     def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None:

vllm/v1/worker/gpu/model_states/prompt_embeds.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
+from vllm.v1.worker.gpu.input_batch import InputBatch
+
+TOKEN_BLOCK = 16
+
+
+class PromptEmbedsState:
+    """GPU-side state for user-provided prompt embeddings.
+
+    Each request's embeddings are copied to the GPU once at `add_request`,
+    off the per-step hot path. A per-request pointer table (UVA) then lets a
+    single triton kernel overlay all scheduled prompt-embeds rows onto
+    `inputs_embeds` each step, with no python loops or per-request H2D copies.
+    """
+
+    def __init__(
+        self,
+        max_num_reqs: int,
+        hidden_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ):
+        self.hidden_size = hidden_size
+        self.dtype = dtype
+        self.device = device
+
+        # req_id -> (embeds, is_token_ids mask or None). Holds the references
+        # that keep the pointer table below valid.
+        self.gpu_tensors: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {}
+
+        # Indexed by req_state index. Stale entries after removal are
+        # harmless: add_request rewrites all fields for every index it claims.
+        self.embeds_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
+        self.mask_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
+        self.embeds_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        prompt_embeds = new_req_data.prompt_embeds
+        if prompt_embeds is None:
+            self.gpu_tensors.pop(new_req_data.req_id, None)
+            self.embeds_lens.np[req_index] = 0
+            return
+
+        embeds = async_tensor_h2d(prompt_embeds, device=self.device, dtype=self.dtype)
+        embeds = embeds.contiguous()
+        is_token_ids = new_req_data.prompt_is_token_ids
+        mask = None
+        if is_token_ids is not None:
+            mask = async_tensor_h2d(is_token_ids, device=self.device, dtype=torch.uint8)
+        self.gpu_tensors[new_req_data.req_id] = (embeds, mask)
+        self.embeds_ptrs.np[req_index] = embeds.data_ptr()
+        self.mask_ptrs.np[req_index] = 0 if mask is None else mask.data_ptr()
+        self.embeds_lens.np[req_index] = embeds.shape[0]
+
+    def remove_request(self, req_id: str) -> None:
+        self.gpu_tensors.pop(req_id, None)
+
+    def apply_staged_writes(self) -> None:
+        self.embeds_ptrs.copy_to_uva()
+        self.mask_ptrs.copy_to_uva()
+        self.embeds_lens.copy_to_uva()
+
+    def apply(
+        self,
+        input_batch: InputBatch,
+        num_computed_tokens: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+    ) -> None:
+        """Overlay prompt embeddings onto `inputs_embeds` for the batch."""
+        if not self.gpu_tensors:
+            return
+        # The kernel reinterprets raw source pointers as inputs_embeds' dtype.
+        assert inputs_embeds.dtype == self.dtype
+        num_reqs = input_batch.num_reqs
+        max_query_len = int(input_batch.num_scheduled_tokens.max())
+        grid = (num_reqs, triton.cdiv(max_query_len, TOKEN_BLOCK))
+        _apply_prompt_embeds_kernel[grid](
+            inputs_embeds,
+            inputs_embeds.stride(0),
+            self.embeds_ptrs.gpu,
+            self.mask_ptrs.gpu,
+            self.embeds_lens.gpu,
+            input_batch.idx_mapping,
+            input_batch.query_start_loc,
+            num_computed_tokens,
+            self.hidden_size,
+            TOKEN_BLOCK=TOKEN_BLOCK,
+            BLOCK_SIZE=1024,
+        )
+
+
+@triton.jit
+def _apply_prompt_embeds_kernel(
+    inputs_embeds_ptr,
+    inputs_embeds_stride,
+    embeds_ptrs_ptr,  # int64 [max_num_reqs], device pointers (0-len = unused)
+    mask_ptrs_ptr,  # int64 [max_num_reqs], 0 = no is-token-ids mask
+    embeds_lens_ptr,  # int32 [max_num_reqs]
+    idx_mapping_ptr,
+    query_start_loc_ptr,
+    num_computed_tokens_ptr,
+    hidden_size: tl.constexpr,
+    TOKEN_BLOCK: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    batch_idx = tl.program_id(0)
+    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
+    embeds_len = tl.load(embeds_lens_ptr + req_state_idx)
+    num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
+    if num_computed >= embeds_len:
+        # No prompt embeds for this request, or they are fully consumed.
+        return
+
+    query_start = tl.load(query_start_loc_ptr + batch_idx)
+    query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
+    num_rows = tl.minimum(query_end - query_start, embeds_len - num_computed)
+
+    t_start = tl.program_id(1) * TOKEN_BLOCK
+    if t_start >= num_rows:
+        return
+
+    src_ptr = tl.load(embeds_ptrs_ptr + req_state_idx).to(
+        tl.pointer_type(inputs_embeds_ptr.dtype.element_ty)
+    )
+    mask_int = tl.load(mask_ptrs_ptr + req_state_idx)
+    mask_ptr = mask_int.to(tl.pointer_type(tl.int8))
+
+    for t_offset in tl.static_range(TOKEN_BLOCK):
+        t = t_start + t_offset
+        if t < num_rows:
+            src_row = (num_computed + t).to(tl.int64)
+            is_token_id = 0
+            if mask_int != 0:
+                is_token_id = tl.load(mask_ptr + src_row).to(tl.int32)
+            if is_token_id == 0:
+                dst_row = (query_start + t).to(tl.int64)
+                for h in tl.range(0, hidden_size, BLOCK_SIZE):
+                    offs = h + tl.arange(0, BLOCK_SIZE)
+                    h_mask = offs < hidden_size
+                    row = tl.load(src_ptr + src_row * hidden_size + offs, mask=h_mask)
+                    tl.store(
+                        inputs_embeds_ptr + dst_row * inputs_embeds_stride + offs,
+                        row,
+                        mask=h_mask,
+                    )