From 8e00cb479621591dc259bd3a697e043dc6e4af8d Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 8 Sep 2026 09:31:08 -0700 Subject: [PATCH] [Core] MRV2 support for fast-prefill Signed-off-by: Nick Hill --- .buildkite/test_areas/engine.yaml | 2 +- .../general/test_kv_sharing_fast_prefill.py | 16 ++- tests/v1/worker/test_attn_utils.py | 64 ++++++++++ tests/v1/worker/test_gpu_model_runner_v2.py | 2 +- vllm/config/cache.py | 1 - vllm/config/vllm.py | 4 - vllm/v1/attention/backends/utils.py | 9 +- vllm/v1/worker/gpu/attn_utils.py | 116 +++++++++++++++++- vllm/v1/worker/gpu/input_batch.py | 5 + vllm/v1/worker/gpu/model_runner.py | 24 +++- vllm/v1/worker/gpu/model_states/default.py | 1 + vllm/v1/worker/gpu/ubatch_utils.py | 1 + 12 files changed, 233 insertions(+), 12 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 4384484c0954..a6b98bf9313b 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.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/ diff --git a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py index 3155da82ea85..d92aa5c77572 100644 --- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py +++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py @@ -52,12 +52,23 @@ def test_prompts(): @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, diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 9751865794b7..906525494f4b 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -7,6 +7,8 @@ never addressed by the logical view. """ +from types import SimpleNamespace + import pytest import torch @@ -22,6 +24,7 @@ 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( diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index eb05b9858af4..1118ce589ca5 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -84,7 +84,7 @@ def narrow(self, *args): 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, diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 4c48fc291ebb..93114ab1d150 100644 --- a/vllm/config/cache.py +++ b/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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 0564cd420ed5..ebb3bdc90d4e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2638,10 +2638,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: ): 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'") diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 716f3689592f..99716ca911e6 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 05949c25ac21..4a859924256e 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/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 @@ 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 @@ 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 @@ def narrow( 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, diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 2a38c3d71061..30b9b8089c04 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -11,6 +11,7 @@ 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, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7dec923c0a8f..fad60d9103fc 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -87,6 +87,7 @@ StepTimingCollector, ) from vllm.v1.worker.gpu.attn_utils import ( + FastPrefillHelper, build_slot_mappings_by_layer, get_kv_cache_spec, init_attn_backend, @@ -319,6 +320,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): 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, @@ -602,10 +604,14 @@ def initialize_kv_cache( 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) @@ -657,7 +663,7 @@ def initialize_kv_cache( 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 @@ -679,6 +685,10 @@ def initialize_kv_cache( 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) @@ -1344,6 +1354,16 @@ def prepare_inputs( 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) @@ -1390,6 +1410,7 @@ def prepare_inputs( 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 @@ -2117,6 +2138,7 @@ def shutdown(self) -> None: 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"): diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 7d42745d1064..6baceac9bc52 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -226,5 +226,6 @@ def prepare_attn( 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 diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 350805707d65..5c9fde3e9356 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -174,6 +174,7 @@ def _slice_input_batch( if input_batch.prompt_lens is None else input_batch.prompt_lens[req_start:req_stop] ), + fast_prefill=None, )