From e6cb56337b49e606f55fde1870adbbb051e23f9f Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Thu, 10 Sep 2026 10:36:04 -0700 Subject: [PATCH] [Core] MRV2 support for fast-prefill (#56145) 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 4384484c09..a6b98bf931 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 3155da82ea..d92aa5c775 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 @@ 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, diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 9751865794..906525494f 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/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( diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index eb05b9858a..1118ce589c 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 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, diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 4c48fc291e..93114ab1d1 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 294f80b143..f5bee6fa2c 100644 --- a/vllm/config/vllm.py +++ b/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'") diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 716f368959..99716ca911 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 05949c25ac..4a85992425 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 @@ 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, diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 2a38c3d710..30b9b8089c 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/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, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ec5ce63503..533fcbe1e7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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"): diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 7d42745d10..6baceac9bc 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/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 diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 012ce6c471..8007e2dbb4 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/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, )