[Core] MRV2 support for fast-prefill (#56145)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Sep 10, 2026 at 10:36 UTC
e6cb56337b49e606f55fde1870adbbb051e23f9f
12 files changed
+233
-12
.buildkite/test_areas/engine.yaml
+1
-1
@@ -77,7 +77,7 @@ steps:
77
- label: ":nvidia: (H200 MIG 35GB) E2E Core"
78
device: h200_35gb
79
key: e2e-core-1-gpu
80
- timeout_in_minutes: 40
80
+ timeout_in_minutes: 50
81
source_file_dependencies:
82
- vllm/v1/
83
- tests/v1/e2e/general/
tests/v1/e2e/general/test_kv_sharing_fast_prefill.py
+15
-1
@@ -52,12 +52,23 @@ use_fork_for_test = (
52
53
54
@use_fork_for_test
55
-@pytest.mark.parametrize("kv_sharing_fast_prefill", [False, True])
55
+@pytest.mark.parametrize(
56
+ "kv_sharing_fast_prefill,use_v2_model_runner",
57
+ [
58
+ (False, False),
59
+ (True, False),
60
+ (True, True),
61
+ # (False, True) omitted: fast-prefill-off behavior is runner-generic,
62
+ # and skipping it saves ~9 min of mostly torch.compile time in CI. The
63
+ # no-silent-fallback assertion below is exercised by the V2 cases.
64
+ ],
65
+)
66
@pytest.mark.parametrize("enforce_eager", [True, False])
67
def test_kv_sharing_fast_prefill(
68
monkeypatch: pytest.MonkeyPatch,
69
kv_sharing_fast_prefill: bool,
70
enforce_eager: bool,
71
+ use_v2_model_runner: bool,
72
):
73
if not enforce_eager and current_platform.is_rocm():
74
# Relevant context: https://github.com/vllm-project/vllm/pull/29244
@@ -84,6 +95,7 @@ def test_kv_sharing_fast_prefill(
95
m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
96
else:
97
m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
98
+ m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_model_runner else "0")
99
100
prompts, answer, indices = prep_prompts(batch_size)
101
@@ -95,6 +107,8 @@ def test_kv_sharing_fast_prefill(
107
kv_sharing_fast_prefill=kv_sharing_fast_prefill,
108
attention_backend="TRITON_ATTN",
109
)
110
+ # Guard against a silent fallback to the other model runner.
111
+ assert llm.llm_engine.vllm_config.use_v2_model_runner == use_v2_model_runner
112
responses = llm.generate(prompts, sampling_params)
113
check_answers(
114
indices,
tests/v1/worker/test_attn_utils.py
+64
@@ -7,6 +7,8 @@ while keeping per-block content compact, so padding bytes at the end of each pag
7
never addressed by the logical view.
8
"""
9
10
+from types import SimpleNamespace
11
+
12
import pytest
13
import torch
14
@@ -22,6 +24,7 @@ from vllm.v1.kv_cache_interface import (
24
MLAAttentionSpec,
25
compute_layout_strides,
26
)
27
+from vllm.v1.worker.gpu import attn_utils
28
from vllm.v1.worker.gpu.attn_utils import (
29
get_attn_cg_support,
30
get_query_lens_mismatch_unsupported_backend,
@@ -118,6 +121,67 @@ def test_attention_checks_preserve_global_and_target_scoped_support():
121
)
122
123
124
+def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyPatch):
125
+ """Fast prefill applies to the contiguous suffix of KV-sharing layers.
126
+
127
+ Draft-model layers register after the target model's and may share KV, so
128
+ they must not extend (or break) the target's eligible suffix.
129
+ """
130
+
131
+ def check(
132
+ layer_names: list[str],
133
+ shared: dict[str, str],
134
+ draft_layer_names: set[str] | None = None,
135
+ ) -> set[str]:
136
+ monkeypatch.setattr(
137
+ attn_utils,
138
+ "get_layers_from_vllm_config",
139
+ lambda *a, **k: {name: None for name in layer_names},
140
+ )
141
+ monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda *a: shared)
142
+ vllm_config = SimpleNamespace(
143
+ cache_config=SimpleNamespace(kv_sharing_fast_prefill=True)
144
+ )
145
+ return attn_utils.get_kv_sharing_fast_prefill_eligible_layers(
146
+ vllm_config, draft_layer_names
147
+ )
148
+
149
+ # No KV sharing: nothing is eligible.
150
+ assert check(["t0", "t1"], {}) == set()
151
+
152
+ # Trailing run of sharing layers (YOCO-style second half).
153
+ assert check(["t0", "t1", "t2", "t3"], {"t2": "t1", "t3": "t1"}) == {"t2", "t3"}
154
+
155
+ # A non-sharing layer after a sharing one breaks the suffix.
156
+ assert check(["t0", "t1", "t2", "t3"], {"t1": "t0", "t3": "t0"}) == {"t3"}
157
+
158
+ # KV-sharing draft layers at the end are collected without an exclusion...
159
+ assert check(
160
+ ["t0", "t1", "t2", "t3", "d0", "d1"],
161
+ {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"},
162
+ ) == {"t2", "t3", "d0", "d1"}
163
+
164
+ # ...so the runner excludes them: skipped, not collected, and they do not
165
+ # break the target's trailing run.
166
+ assert check(
167
+ ["t0", "t1", "t2", "t3", "d0", "d1"],
168
+ {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"},
169
+ draft_layer_names={"d0", "d1"},
170
+ ) == {"t2", "t3"}
171
+
172
+ # Feature flag off: nothing is eligible even with sharing layers.
173
+ monkeypatch.setattr(
174
+ attn_utils, "get_layers_from_vllm_config", lambda *a, **k: {"t0": None}
175
+ )
176
+ monkeypatch.setattr(
177
+ attn_utils, "get_shared_kv_cache_layers", lambda *a: {"t0": "t0"}
178
+ )
179
+ vllm_config = SimpleNamespace(
180
+ cache_config=SimpleNamespace(kv_sharing_fast_prefill=False)
181
+ )
182
+ assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set()
183
+
184
+
185
def test_reshape_padded_kv_cache_strides_by_padded_page():
186
num_blocks = 3
187
spec = FullAttentionSpec(
tests/v1/worker/test_gpu_model_runner_v2.py
+1
-1
@@ -84,7 +84,7 @@ def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
84
monkeypatch.setattr(
85
model_runner_module,
86
"init_attn_backend",
87
- lambda *args: ([], attn_cg_support, [8, 262144]),
87
+ lambda *args, **kwargs: ([], attn_cg_support, [8, 262144]),
88
)
89
monkeypatch.setattr(
90
model_runner_module,
vllm/config/cache.py
-1
@@ -234,7 +234,6 @@ class CacheConfig:
234
some layers can skip tokens corresponding to prefill. This flag enables
235
attention metadata for eligible layers to be overridden with metadata
236
necessary for implementing this optimization in some models (e.g. Gemma3n)
237
- NOTE: KV cache sharing is not supported for MRv2 (v2 model runner).
237
"""
238
239
kv_cache_memory_bytes: int | None = None
vllm/config/vllm.py
-4
@@ -2674,10 +2674,6 @@ class VllmConfig:
2674
):
2675
unsupported.append("custom logits processors")
2676
2677
- if self.cache_config.kv_sharing_fast_prefill:
2678
- # Will be added by https://github.com/vllm-project/vllm/pull/35045
2679
- unsupported.append("KV sharing fast prefill")
2680
-
2677
if self.cache_config.mamba_cache_mode == "all":
2678
unsupported.append("mamba cache mode 'all'")
2679
vllm/v1/attention/backends/utils.py
+7
-2
@@ -626,8 +626,13 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
626
# Skip computing fast prefill path
627
return common_attn_metadata
628
629
- assert common_attn_metadata.logits_indices_padded is not None
630
- assert common_attn_metadata.num_logits_indices is not None
629
+ if (
630
+ common_attn_metadata.logits_indices_padded is None
631
+ or common_attn_metadata.num_logits_indices is None
632
+ ):
633
+ # Fast prefill not armed for this step (e.g. cudagraph capture, or a
634
+ # pure-decode step): run the KV-sharing layers on the full batch.
635
+ return common_attn_metadata
636
637
logits_indices_padded = common_attn_metadata.logits_indices_padded
638
num_logits_indices = common_attn_metadata.num_logits_indices
vllm/v1/worker/gpu/attn_utils.py
+115
-1
@@ -3,11 +3,13 @@
3
from collections.abc import Mapping, Sequence
4
from contextlib import AbstractContextManager, nullcontext
5
from dataclasses import dataclass
6
-from typing import Any, cast
6
+from typing import TYPE_CHECKING, Any, cast
7
8
+import numpy as np
9
import torch
10
11
from vllm.config import VllmConfig, get_layers_from_vllm_config
12
+from vllm.config.compilation import CUDAGraphMode
13
from vllm.model_executor.layers.attention import Attention
14
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
15
from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -15,6 +17,7 @@ from vllm.v1.attention.backend import (
17
AttentionCGSupport,
18
CommonAttentionMetadata,
19
)
20
+from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend
21
from vllm.v1.kv_cache_interface import (
22
AttentionSpec,
23
KVCacheConfig,
@@ -31,6 +34,12 @@ from vllm.v1.worker.utils import (
34
prepare_kernel_block_sizes,
35
)
36
37
+if TYPE_CHECKING:
38
+ from vllm.v1.worker.gpu.cudagraph_utils import (
39
+ BatchExecutionDescriptor,
40
+ CudaGraphManager,
41
+ )
42
+
43
44
@dataclass(frozen=True)
45
class AttentionCGSupportInfo:
@@ -50,6 +59,66 @@ class AttentionCGSupportInfo:
59
return self
60
61
62
+@dataclass(frozen=True)
63
+class FastPrefillBatchMetadata:
64
+ """Per-step inputs for the KV-sharing fast prefill path."""
65
+
66
+ logits_indices_padded: torch.Tensor
67
+ num_logits_indices: int
68
+ max_logits_per_req: int
69
+
70
+
71
+class FastPrefillHelper:
72
+ """Decides per step whether to arm the KV-sharing fast prefill path, and
73
+ stages the logits indices it runs on.
74
+ """
75
+
76
+ def __init__(self, cudagraph_manager: "CudaGraphManager", max_num_tokens: int):
77
+ self.max_num_tokens = max_num_tokens
78
+ self.cudagraph_manager = cudagraph_manager
79
+ self.logits_indices_buf = torch.zeros(
80
+ max_num_tokens, dtype=torch.int32, device=cudagraph_manager.device
81
+ )
82
+
83
+ def prepare(
84
+ self,
85
+ logits_indices: torch.Tensor,
86
+ num_reqs: int,
87
+ cu_num_logits_np: np.ndarray,
88
+ has_prefill: bool,
89
+ batch_desc: "BatchExecutionDescriptor",
90
+ ) -> FastPrefillBatchMetadata | None:
91
+ if (
92
+ not has_prefill
93
+ or batch_desc.cg_mode == CUDAGraphMode.FULL
94
+ or batch_desc.num_ubatches > 1
95
+ ):
96
+ return None
97
+ buf = self.logits_indices_buf
98
+ num_logits = logits_indices.shape[0]
99
+ assert num_logits > 0
100
+ buf[:num_logits].copy_(logits_indices)
101
+ # There might be leftover indices in buf[num_logits:] from previous
102
+ # iterations. Broadcast the scalar GPU-side to keep them valid.
103
+ buf[num_logits:] = logits_indices[-1]
104
+ # Pad so the model's KV-sharing layers run their reduced (logits-only)
105
+ # batch at a captured piecewise cudagraph size. Pre-capture, or with
106
+ # cudagraphs off, dispatch returns the unpadded count.
107
+ desc = self.cudagraph_manager.dispatch(
108
+ num_reqs=num_reqs,
109
+ num_tokens=num_logits,
110
+ uniform_token_count=None,
111
+ num_active_loras=0,
112
+ )
113
+ num_logits_padded = min(desc.num_tokens, self.max_num_tokens)
114
+ return FastPrefillBatchMetadata(
115
+ logits_indices_padded=buf[:num_logits_padded],
116
+ num_logits_indices=num_logits,
117
+ # Largest per-request logits count, known on the host.
118
+ max_logits_per_req=int(np.diff(cu_num_logits_np).max()),
119
+ )
120
+
121
+
122
def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
123
kv_cache_spec: dict[str, KVCacheSpec] = {}
124
layer_type = cast(type[Any], AttentionLayerBase)
@@ -75,11 +144,42 @@ def get_shared_kv_cache_layers(vllm_config: VllmConfig):
144
}
145
146
147
+def get_kv_sharing_fast_prefill_eligible_layers(
148
+ vllm_config: VllmConfig, draft_layer_names: set[str] | None = None
149
+) -> set[str]:
150
+ """Trailing run of KV-sharing layers, eligible for fast prefill.
151
+
152
+ In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other similar
153
+ KV sharing setups, only the layers that generate KV caches are involved in
154
+ the prefill phase, enabling prefill to early exit. Layers are registered in
155
+ execution order, so the eligible layers are the contiguous suffix of
156
+ KV-sharing layers.
157
+
158
+ Speculator draft layers register after the target model's layers (and may
159
+ themselves share KV), so they are excluded from the walk.
160
+ """
161
+ if not vllm_config.cache_config.kv_sharing_fast_prefill:
162
+ return set()
163
+ shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
164
+ if not shared_kv_cache_layers:
165
+ return set()
166
+ eligible_layers: set[str] = set()
167
+ attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
168
+ for layer_name in reversed(attn_layers):
169
+ if draft_layer_names is not None and layer_name in draft_layer_names:
170
+ continue
171
+ if layer_name not in shared_kv_cache_layers:
172
+ break
173
+ eligible_layers.add(layer_name)
174
+ return eligible_layers
175
+
176
+
177
def init_attn_backend(
178
kv_cache_config: KVCacheConfig,
179
vllm_config: VllmConfig,
180
device: torch.device,
181
active_layer_names: set[str] | None = None,
182
+ draft_layer_names: set[str] | None = None,
183
) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]:
184
# Phase 1: discover attention groups for each kv cache group.
185
attn_groups: list[list[AttentionGroup]] = []
@@ -89,6 +189,9 @@ def init_attn_backend(
189
add_kv_sharing_layers_to_kv_cache_groups(
190
get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups
191
)
192
+ fast_prefill_eligible_layers = get_kv_sharing_fast_prefill_eligible_layers(
193
+ vllm_config, draft_layer_names
194
+ )
195
196
# Phase 1: discover attention groups for each kv cache group.
197
for kv_cache_group_id, kv_cache_group_spec in enumerate(
@@ -106,6 +209,10 @@ def init_attn_backend(
209
210
for layer_name in layer_names:
211
attn_backend = attn_layers[layer_name].get_attn_backend()
212
+ if layer_name in fast_prefill_eligible_layers:
213
+ attn_backend = create_fast_prefill_custom_backend(
214
+ "FastPrefill", attn_backend
215
+ )
216
217
layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec
218
if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -278,6 +385,7 @@ def build_attn_metadata(
385
causal: bool | torch.Tensor | Mapping[int, bool] = True,
386
rswa_prefix_lens: torch.Tensor | None = None,
387
ubatch_idx: int = 0,
388
+ fast_prefill: FastPrefillBatchMetadata | None = None,
389
) -> dict[str, Any]:
390
seq_lens = seq_lens[:num_reqs]
391
if dcp_local_seq_lens is not None:
@@ -305,6 +413,12 @@ def build_attn_metadata(
413
group_is_prefilling = common_attn_metadata_extra_kwargs.pop(
414
"is_prefilling", is_prefilling
415
)
416
+ if fast_prefill is not None:
417
+ common_attn_metadata_extra_kwargs.update(
418
+ logits_indices_padded=fast_prefill.logits_indices_padded,
419
+ num_logits_indices=fast_prefill.num_logits_indices,
420
+ max_logits_per_req=fast_prefill.max_logits_per_req,
421
+ )
422
common_attn_metadata = CommonAttentionMetadata(
423
query_start_loc=query_start_loc_gpu,
424
query_start_loc_cpu=query_start_loc_cpu,
vllm/v1/worker/gpu/input_batch.py
+5
@@ -11,6 +11,7 @@ from vllm.utils import random_uuid
11
from vllm.utils.math_utils import cdiv
12
13
if TYPE_CHECKING:
14
+ from vllm.v1.worker.gpu.attn_utils import FastPrefillBatchMetadata
15
from vllm.v1.worker.gpu.block_table import BlockTables
16
17
@@ -109,6 +110,10 @@ class InputBatch:
110
# stays valid for every replay the graph serves.
111
max_query_len: int | None = None
112
113
+ # Arms the KV-sharing fast prefill path for this step. Absent for dummy
114
+ # (cudagraph capture) batches, which run the KV-sharing layers in full.
115
+ fast_prefill: "FastPrefillBatchMetadata | None" = None
116
+
117
@classmethod
118
def make_dummy(
119
cls,
vllm/v1/worker/gpu/model_runner.py
+23
-1
@@ -89,6 +89,7 @@ from vllm.v1.worker.gpu.async_utils import (
89
StepTimingCollector,
90
)
91
from vllm.v1.worker.gpu.attn_utils import (
92
+ FastPrefillHelper,
93
build_slot_mappings_by_layer,
94
get_kv_cache_spec,
95
init_attn_backend,
@@ -321,6 +322,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
322
max_num_tokens=self.max_num_tokens,
323
device=self.device,
324
)
325
+ self.fast_prefill: FastPrefillHelper | None = None
326
if self.use_pp:
327
self.pp_handler = PPHandler(
328
max_num_reqs=self.max_num_reqs,
@@ -610,10 +612,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
612
for group in self.kv_cache_config.kv_cache_groups
613
for layer_name in group.layer_names
614
} - self.speculator.draft_attn_layer_names
615
+ draft_attn_layer_names = None
616
+ if isinstance(self.speculator, DraftModelSpeculator):
617
+ draft_attn_layer_names = self.speculator.draft_attn_layer_names
618
self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
619
self.kv_cache_config,
620
self.vllm_config,
621
self.device,
622
+ draft_layer_names=draft_attn_layer_names,
623
)
624
additional_attn_cg_support = self.model_state.get_additional_cg_support()
625
attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
@@ -665,7 +671,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
671
initialize_mamba_ssu_backend(
672
self.vllm_config.mamba_config,
673
self.kv_cache_config,
668
- use_replayssm=self.vllm_config.cache_config.use_replayssm,
674
+ use_replayssm=self.cache_config.use_replayssm,
675
)
676
if self.adaptive_verification is not None:
677
self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
@@ -687,6 +693,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
693
lora_capture_cases=self.lora_capture_cases,
694
varlen_decode=self.adaptive_verification is not None,
695
)
696
+ if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None:
697
+ self.fast_prefill = FastPrefillHelper(
698
+ self.cudagraph_manager, self.max_num_tokens
699
+ )
700
check_attention_cp_compatibility(self.vllm_config)
701
if isinstance(self.speculator, DraftModelSpeculator):
702
# HACK(woosuk)
@@ -1339,6 +1349,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1349
self.model_state.num_new_sampled_tokens_per_step,
1350
)
1351
1352
+ fast_prefill = None
1353
+ if self.fast_prefill is not None:
1354
+ fast_prefill = self.fast_prefill.prepare(
1355
+ logits_indices,
1356
+ num_reqs,
1357
+ cu_num_logits_np,
1358
+ batch_req_state.has_prefill,
1359
+ batch_desc,
1360
+ )
1361
+
1362
# CPU upper bound on seq_lens; padded entries left at zero.
1363
num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
1364
seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
@@ -1385,6 +1405,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1405
cu_num_logits_np=cu_num_logits_np,
1406
has_structured_output_reqs=scheduler_output.has_structured_output_requests,
1407
prompt_lens=prompt_lens,
1408
+ fast_prefill=fast_prefill,
1409
max_query_len=(
1410
int(num_scheduled_tokens_upper_bound.max())
1411
if adaptive_verification is not None
@@ -2127,6 +2148,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2148
memory is reclaimable when running in the same process."""
2149
torch.accelerator.synchronize()
2150
self.cudagraph_manager = None
2151
+ self.fast_prefill = None
2152
if hasattr(self, "kv_caches"):
2153
self.kv_caches.clear()
2154
if hasattr(self, "attn_groups"):
vllm/v1/worker/gpu/model_states/default.py
+1
@@ -226,5 +226,6 @@ class DefaultModelState(ModelState):
226
for_cudagraph_capture=for_capture,
227
rswa_prefix_lens=input_batch.prompt_lens,
228
ubatch_idx=ubatch_idx,
229
+ fast_prefill=input_batch.fast_prefill,
230
)
231
return attn_metadata
vllm/v1/worker/gpu/ubatch_utils.py
+1
@@ -187,6 +187,7 @@ def _slice_input_batch(
187
if input_batch.prompt_lens is None
188
else input_batch.prompt_lens[req_start:req_stop]
189
),
190
+ fast_prefill=None,
191
)
192
193