[Core] MRV2 support for fast-prefill (#56145)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

.buildkite/test_areas/engine.yaml

@@ -77,7 +77,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) E2E Core"
   device: h200_35gb
   key: e2e-core-1-gpu
-  timeout_in_minutes: 40
+  timeout_in_minutes: 50
   source_file_dependencies:
     - vllm/v1/
     - tests/v1/e2e/general/

tests/v1/e2e/general/test_kv_sharing_fast_prefill.py

@@ -52,12 +52,23 @@ use_fork_for_test = (
 
 
 @use_fork_for_test
-@pytest.mark.parametrize("kv_sharing_fast_prefill", [False, True])
+@pytest.mark.parametrize(
+    "kv_sharing_fast_prefill,use_v2_model_runner",
+    [
+        (False, False),
+        (True, False),
+        (True, True),
+        # (False, True) omitted: fast-prefill-off behavior is runner-generic,
+        # and skipping it saves ~9 min of mostly torch.compile time in CI. The
+        # no-silent-fallback assertion below is exercised by the V2 cases.
+    ],
+)
 @pytest.mark.parametrize("enforce_eager", [True, False])
 def test_kv_sharing_fast_prefill(
     monkeypatch: pytest.MonkeyPatch,
     kv_sharing_fast_prefill: bool,
     enforce_eager: bool,
+    use_v2_model_runner: bool,
 ):
     if not enforce_eager and current_platform.is_rocm():
         # Relevant context: https://github.com/vllm-project/vllm/pull/29244
@@ -84,6 +95,7 @@ def test_kv_sharing_fast_prefill(
             m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
         else:
             m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+        m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_model_runner else "0")
 
         prompts, answer, indices = prep_prompts(batch_size)
 
@@ -95,6 +107,8 @@ def test_kv_sharing_fast_prefill(
             kv_sharing_fast_prefill=kv_sharing_fast_prefill,
             attention_backend="TRITON_ATTN",
         )
+        # Guard against a silent fallback to the other model runner.
+        assert llm.llm_engine.vllm_config.use_v2_model_runner == use_v2_model_runner
         responses = llm.generate(prompts, sampling_params)
         check_answers(
             indices,

tests/v1/worker/test_attn_utils.py

@@ -7,6 +7,8 @@ while keeping per-block content compact, so padding bytes at the end of each pag
 never addressed by the logical view.
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -22,6 +24,7 @@ from vllm.v1.kv_cache_interface import (
     MLAAttentionSpec,
     compute_layout_strides,
 )
+from vllm.v1.worker.gpu import attn_utils
 from vllm.v1.worker.gpu.attn_utils import (
     get_attn_cg_support,
     get_query_lens_mismatch_unsupported_backend,
@@ -118,6 +121,67 @@ def test_attention_checks_preserve_global_and_target_scoped_support():
     )
 
 
+def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyPatch):
+    """Fast prefill applies to the contiguous suffix of KV-sharing layers.
+
+    Draft-model layers register after the target model's and may share KV, so
+    they must not extend (or break) the target's eligible suffix.
+    """
+
+    def check(
+        layer_names: list[str],
+        shared: dict[str, str],
+        draft_layer_names: set[str] | None = None,
+    ) -> set[str]:
+        monkeypatch.setattr(
+            attn_utils,
+            "get_layers_from_vllm_config",
+            lambda *a, **k: {name: None for name in layer_names},
+        )
+        monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda *a: shared)
+        vllm_config = SimpleNamespace(
+            cache_config=SimpleNamespace(kv_sharing_fast_prefill=True)
+        )
+        return attn_utils.get_kv_sharing_fast_prefill_eligible_layers(
+            vllm_config, draft_layer_names
+        )
+
+    # No KV sharing: nothing is eligible.
+    assert check(["t0", "t1"], {}) == set()
+
+    # Trailing run of sharing layers (YOCO-style second half).
+    assert check(["t0", "t1", "t2", "t3"], {"t2": "t1", "t3": "t1"}) == {"t2", "t3"}
+
+    # A non-sharing layer after a sharing one breaks the suffix.
+    assert check(["t0", "t1", "t2", "t3"], {"t1": "t0", "t3": "t0"}) == {"t3"}
+
+    # KV-sharing draft layers at the end are collected without an exclusion...
+    assert check(
+        ["t0", "t1", "t2", "t3", "d0", "d1"],
+        {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"},
+    ) == {"t2", "t3", "d0", "d1"}
+
+    # ...so the runner excludes them: skipped, not collected, and they do not
+    # break the target's trailing run.
+    assert check(
+        ["t0", "t1", "t2", "t3", "d0", "d1"],
+        {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"},
+        draft_layer_names={"d0", "d1"},
+    ) == {"t2", "t3"}
+
+    # Feature flag off: nothing is eligible even with sharing layers.
+    monkeypatch.setattr(
+        attn_utils, "get_layers_from_vllm_config", lambda *a, **k: {"t0": None}
+    )
+    monkeypatch.setattr(
+        attn_utils, "get_shared_kv_cache_layers", lambda *a: {"t0": "t0"}
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(kv_sharing_fast_prefill=False)
+    )
+    assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set()
+
+
 def test_reshape_padded_kv_cache_strides_by_padded_page():
     num_blocks = 3
     spec = FullAttentionSpec(

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -84,7 +84,7 @@ def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
     monkeypatch.setattr(
         model_runner_module,
         "init_attn_backend",
-        lambda *args: ([], attn_cg_support, [8, 262144]),
+        lambda *args, **kwargs: ([], attn_cg_support, [8, 262144]),
     )
     monkeypatch.setattr(
         model_runner_module,

vllm/config/cache.py

@@ -234,7 +234,6 @@ class CacheConfig:
     some layers can skip tokens corresponding to prefill. This flag enables
     attention metadata for eligible layers to be overridden with metadata
     necessary for implementing this optimization in some models (e.g. Gemma3n)
-    NOTE: KV cache sharing is not supported for MRv2 (v2 model runner).
     """
 
     kv_cache_memory_bytes: int | None = None

vllm/config/vllm.py

@@ -2674,10 +2674,6 @@ class VllmConfig:
         ):
             unsupported.append("custom logits processors")
 
-        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")
-
         if self.cache_config.mamba_cache_mode == "all":
             unsupported.append("mamba cache mode 'all'")
 

vllm/v1/attention/backends/utils.py

@@ -626,8 +626,13 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
         # Skip computing fast prefill path
         return common_attn_metadata
 
-    assert common_attn_metadata.logits_indices_padded is not None
-    assert common_attn_metadata.num_logits_indices is not None
+    if (
+        common_attn_metadata.logits_indices_padded is None
+        or common_attn_metadata.num_logits_indices is None
+    ):
+        # Fast prefill not armed for this step (e.g. cudagraph capture, or a
+        # pure-decode step): run the KV-sharing layers on the full batch.
+        return common_attn_metadata
 
     logits_indices_padded = common_attn_metadata.logits_indices_padded
     num_logits_indices = common_attn_metadata.num_logits_indices

vllm/v1/worker/gpu/attn_utils.py

@@ -3,11 +3,13 @@
 from collections.abc import Mapping, Sequence
 from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
+import numpy as np
 import torch
 
 from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config.compilation import CUDAGraphMode
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -15,6 +17,7 @@ from vllm.v1.attention.backend import (
     AttentionCGSupport,
     CommonAttentionMetadata,
 )
+from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     KVCacheConfig,
@@ -31,6 +34,12 @@ from vllm.v1.worker.utils import (
     prepare_kernel_block_sizes,
 )
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.cudagraph_utils import (
+        BatchExecutionDescriptor,
+        CudaGraphManager,
+    )
+
 
 @dataclass(frozen=True)
 class AttentionCGSupportInfo:
@@ -50,6 +59,66 @@ class AttentionCGSupportInfo:
         return self
 
 
+@dataclass(frozen=True)
+class FastPrefillBatchMetadata:
+    """Per-step inputs for the KV-sharing fast prefill path."""
+
+    logits_indices_padded: torch.Tensor
+    num_logits_indices: int
+    max_logits_per_req: int
+
+
+class FastPrefillHelper:
+    """Decides per step whether to arm the KV-sharing fast prefill path, and
+    stages the logits indices it runs on.
+    """
+
+    def __init__(self, cudagraph_manager: "CudaGraphManager", max_num_tokens: int):
+        self.max_num_tokens = max_num_tokens
+        self.cudagraph_manager = cudagraph_manager
+        self.logits_indices_buf = torch.zeros(
+            max_num_tokens, dtype=torch.int32, device=cudagraph_manager.device
+        )
+
+    def prepare(
+        self,
+        logits_indices: torch.Tensor,
+        num_reqs: int,
+        cu_num_logits_np: np.ndarray,
+        has_prefill: bool,
+        batch_desc: "BatchExecutionDescriptor",
+    ) -> FastPrefillBatchMetadata | None:
+        if (
+            not has_prefill
+            or batch_desc.cg_mode == CUDAGraphMode.FULL
+            or batch_desc.num_ubatches > 1
+        ):
+            return None
+        buf = self.logits_indices_buf
+        num_logits = logits_indices.shape[0]
+        assert num_logits > 0
+        buf[:num_logits].copy_(logits_indices)
+        # There might be leftover indices in buf[num_logits:] from previous
+        # iterations. Broadcast the scalar GPU-side to keep them valid.
+        buf[num_logits:] = logits_indices[-1]
+        # Pad so the model's KV-sharing layers run their reduced (logits-only)
+        # batch at a captured piecewise cudagraph size. Pre-capture, or with
+        # cudagraphs off, dispatch returns the unpadded count.
+        desc = self.cudagraph_manager.dispatch(
+            num_reqs=num_reqs,
+            num_tokens=num_logits,
+            uniform_token_count=None,
+            num_active_loras=0,
+        )
+        num_logits_padded = min(desc.num_tokens, self.max_num_tokens)
+        return FastPrefillBatchMetadata(
+            logits_indices_padded=buf[:num_logits_padded],
+            num_logits_indices=num_logits,
+            # Largest per-request logits count, known on the host.
+            max_logits_per_req=int(np.diff(cu_num_logits_np).max()),
+        )
+
+
 def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
     kv_cache_spec: dict[str, KVCacheSpec] = {}
     layer_type = cast(type[Any], AttentionLayerBase)
@@ -75,11 +144,42 @@ def get_shared_kv_cache_layers(vllm_config: VllmConfig):
     }
 
 
+def get_kv_sharing_fast_prefill_eligible_layers(
+    vllm_config: VllmConfig, draft_layer_names: set[str] | None = None
+) -> set[str]:
+    """Trailing run of KV-sharing layers, eligible for fast prefill.
+
+    In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other similar
+    KV sharing setups, only the layers that generate KV caches are involved in
+    the prefill phase, enabling prefill to early exit. Layers are registered in
+    execution order, so the eligible layers are the contiguous suffix of
+    KV-sharing layers.
+
+    Speculator draft layers register after the target model's layers (and may
+    themselves share KV), so they are excluded from the walk.
+    """
+    if not vllm_config.cache_config.kv_sharing_fast_prefill:
+        return set()
+    shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
+    if not shared_kv_cache_layers:
+        return set()
+    eligible_layers: set[str] = set()
+    attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
+    for layer_name in reversed(attn_layers):
+        if draft_layer_names is not None and layer_name in draft_layer_names:
+            continue
+        if layer_name not in shared_kv_cache_layers:
+            break
+        eligible_layers.add(layer_name)
+    return eligible_layers
+
+
 def init_attn_backend(
     kv_cache_config: KVCacheConfig,
     vllm_config: VllmConfig,
     device: torch.device,
     active_layer_names: set[str] | None = None,
+    draft_layer_names: set[str] | None = None,
 ) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]:
     # Phase 1: discover attention groups for each kv cache group.
     attn_groups: list[list[AttentionGroup]] = []
@@ -89,6 +189,9 @@ def init_attn_backend(
     add_kv_sharing_layers_to_kv_cache_groups(
         get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups
     )
+    fast_prefill_eligible_layers = get_kv_sharing_fast_prefill_eligible_layers(
+        vllm_config, draft_layer_names
+    )
 
     # Phase 1: discover attention groups for each kv cache group.
     for kv_cache_group_id, kv_cache_group_spec in enumerate(
@@ -106,6 +209,10 @@ def init_attn_backend(
 
         for layer_name in layer_names:
             attn_backend = attn_layers[layer_name].get_attn_backend()
+            if layer_name in fast_prefill_eligible_layers:
+                attn_backend = create_fast_prefill_custom_backend(
+                    "FastPrefill", attn_backend
+                )
 
             layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec
             if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -278,6 +385,7 @@ def build_attn_metadata(
     causal: bool | torch.Tensor | Mapping[int, bool] = True,
     rswa_prefix_lens: torch.Tensor | None = None,
     ubatch_idx: int = 0,
+    fast_prefill: FastPrefillBatchMetadata | None = None,
 ) -> dict[str, Any]:
     seq_lens = seq_lens[:num_reqs]
     if dcp_local_seq_lens is not None:
@@ -305,6 +413,12 @@ def build_attn_metadata(
         group_is_prefilling = common_attn_metadata_extra_kwargs.pop(
             "is_prefilling", is_prefilling
         )
+        if fast_prefill is not None:
+            common_attn_metadata_extra_kwargs.update(
+                logits_indices_padded=fast_prefill.logits_indices_padded,
+                num_logits_indices=fast_prefill.num_logits_indices,
+                max_logits_per_req=fast_prefill.max_logits_per_req,
+            )
         common_attn_metadata = CommonAttentionMetadata(
             query_start_loc=query_start_loc_gpu,
             query_start_loc_cpu=query_start_loc_cpu,

vllm/v1/worker/gpu/input_batch.py

@@ -11,6 +11,7 @@ from vllm.utils import random_uuid
 from vllm.utils.math_utils import cdiv
 
 if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.attn_utils import FastPrefillBatchMetadata
     from vllm.v1.worker.gpu.block_table import BlockTables
 
 
@@ -109,6 +110,10 @@ class InputBatch:
     # stays valid for every replay the graph serves.
     max_query_len: int | None = None
 
+    # Arms the KV-sharing fast prefill path for this step. Absent for dummy
+    # (cudagraph capture) batches, which run the KV-sharing layers in full.
+    fast_prefill: "FastPrefillBatchMetadata | None" = None
+
     @classmethod
     def make_dummy(
         cls,

vllm/v1/worker/gpu/model_runner.py

@@ -89,6 +89,7 @@ from vllm.v1.worker.gpu.async_utils import (
     StepTimingCollector,
 )
 from vllm.v1.worker.gpu.attn_utils import (
+    FastPrefillHelper,
     build_slot_mappings_by_layer,
     get_kv_cache_spec,
     init_attn_backend,
@@ -321,6 +322,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             max_num_tokens=self.max_num_tokens,
             device=self.device,
         )
+        self.fast_prefill: FastPrefillHelper | None = None
         if self.use_pp:
             self.pp_handler = PPHandler(
                 max_num_reqs=self.max_num_reqs,
@@ -610,10 +612,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 for group in self.kv_cache_config.kv_cache_groups
                 for layer_name in group.layer_names
             } - self.speculator.draft_attn_layer_names
+        draft_attn_layer_names = None
+        if isinstance(self.speculator, DraftModelSpeculator):
+            draft_attn_layer_names = self.speculator.draft_attn_layer_names
         self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
             self.kv_cache_config,
             self.vllm_config,
             self.device,
+            draft_layer_names=draft_attn_layer_names,
         )
         additional_attn_cg_support = self.model_state.get_additional_cg_support()
         attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
@@ -665,7 +671,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         initialize_mamba_ssu_backend(
             self.vllm_config.mamba_config,
             self.kv_cache_config,
-            use_replayssm=self.vllm_config.cache_config.use_replayssm,
+            use_replayssm=self.cache_config.use_replayssm,
         )
         if self.adaptive_verification is not None:
             self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
@@ -687,6 +693,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             lora_capture_cases=self.lora_capture_cases,
             varlen_decode=self.adaptive_verification is not None,
         )
+        if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None:
+            self.fast_prefill = FastPrefillHelper(
+                self.cudagraph_manager, self.max_num_tokens
+            )
         check_attention_cp_compatibility(self.vllm_config)
         if isinstance(self.speculator, DraftModelSpeculator):
             # HACK(woosuk)
@@ -1339,6 +1349,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.model_state.num_new_sampled_tokens_per_step,
         )
 
+        fast_prefill = None
+        if self.fast_prefill is not None:
+            fast_prefill = self.fast_prefill.prepare(
+                logits_indices,
+                num_reqs,
+                cu_num_logits_np,
+                batch_req_state.has_prefill,
+                batch_desc,
+            )
+
         # CPU upper bound on seq_lens; padded entries left at zero.
         num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
         seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
@@ -1385,6 +1405,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             cu_num_logits_np=cu_num_logits_np,
             has_structured_output_reqs=scheduler_output.has_structured_output_requests,
             prompt_lens=prompt_lens,
+            fast_prefill=fast_prefill,
             max_query_len=(
                 int(num_scheduled_tokens_upper_bound.max())
                 if adaptive_verification is not None
@@ -2127,6 +2148,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         memory is reclaimable when running in the same process."""
         torch.accelerator.synchronize()
         self.cudagraph_manager = None
+        self.fast_prefill = None
         if hasattr(self, "kv_caches"):
             self.kv_caches.clear()
         if hasattr(self, "attn_groups"):

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

@@ -226,5 +226,6 @@ class DefaultModelState(ModelState):
             for_cudagraph_capture=for_capture,
             rswa_prefix_lens=input_batch.prompt_lens,
             ubatch_idx=ubatch_idx,
+            fast_prefill=input_batch.fast_prefill,
         )
         return attn_metadata

vllm/v1/worker/gpu/ubatch_utils.py

@@ -187,6 +187,7 @@ def _slice_input_batch(
             if input_batch.prompt_lens is None
             else input_batch.prompt_lens[req_start:req_stop]
         ),
+        fast_prefill=None,
     )