A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 22:35 UTC
Override attention metadata for fast prefill in some KV sharing setups (#21590)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Jul 30, 2025 at 08:54 UTC
ad510309ee10e5182b99ee94ddc5ace716c65050
6 files changed
+287
-26
tests/v1/e2e/test_kv_sharing_fast_prefill.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import gc
5
+import random
6
+from typing import Optional, Union
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm import LLM, SamplingParams
12
+from vllm.config import CompilationConfig, CompilationLevel
13
+from vllm.forward_context import get_forward_context
14
+from vllm.model_executor.models.gemma3n import Gemma3nForConditionalGeneration
15
+from vllm.model_executor.models.registry import ModelRegistry
16
+from vllm.model_executor.models.utils import extract_layer_index
17
+from vllm.sequence import IntermediateTensors
18
+
19
+from ...utils import fork_new_process_for_each_test
20
+
21
+
22
+class TestGemma3nForConditionalGeneration(Gemma3nForConditionalGeneration):
23
+
24
+ def forward(
25
+ self,
26
+ input_ids: torch.Tensor,
27
+ positions: torch.Tensor,
28
+ intermediate_tensors: Optional[IntermediateTensors] = None,
29
+ inputs_embeds: Optional[torch.Tensor] = None,
30
+ **kwargs,
31
+ ) -> Union[torch.Tensor, IntermediateTensors]:
32
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
33
+ inputs_embeds, **kwargs)
34
+ attn_metadata = get_forward_context().attn_metadata
35
+ # attn_metadata is None during dummy runs
36
+ if (attn_metadata is not None
37
+ and self.cache_config.kv_sharing_fast_prefill):
38
+ assert isinstance(attn_metadata, dict) # true in V1
39
+ # Gemma3n-E2B has 30 layers, with last 20 layers being
40
+ # cross-decoder layers. Check attention metadata is correct
41
+ for layer_name, metadata in attn_metadata.items():
42
+ layer_idx = extract_layer_index(layer_name)
43
+ if layer_idx >= 20:
44
+ assert hasattr(metadata, 'logits_indices_padded')
45
+ assert hasattr(metadata, 'num_logits_indices')
46
+ else:
47
+ assert not hasattr(metadata, 'logits_indices_padded')
48
+ assert not hasattr(metadata, 'num_logits_indices')
49
+
50
+ # Last layer will be a KV sharing layer
51
+ layer_attn_metadata = attn_metadata[
52
+ self.model.language_model.layers[-1].self_attn.attn.layer_name]
53
+ logits_indices_padded = (layer_attn_metadata.logits_indices_padded)
54
+ assert logits_indices_padded is not None
55
+ num_logits_indices = layer_attn_metadata.num_logits_indices
56
+ assert num_logits_indices > 0
57
+ # Reset hidden states to random values and
58
+ # only set logits at logits_indices to valid values
59
+ # Because logits_indices are the only positions that are used
60
+ # for output token sampling, this still produces same outputs
61
+ logits_hs = hidden_states[logits_indices_padded]
62
+ hidden_states = torch.randn_like(hidden_states)
63
+ gen_indices = logits_indices_padded[:num_logits_indices]
64
+ hidden_states[gen_indices] = logits_hs[:num_logits_indices]
65
+
66
+ return hidden_states
67
+
68
+
69
+@pytest.fixture
70
+def test_prompts():
71
+ """
72
+ Adapted from tests/v1/e2e/test_spec_decode.py
73
+ """
74
+ prompt_types = ["repeat", "sentence"]
75
+ # Setting higher num prompts increases the chance of numerics mismatch
76
+ # due to matrix multiplication numerics depending on batch dimension
77
+ num_prompts = 10
78
+ prompts = []
79
+
80
+ random.seed(0)
81
+ random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
82
+
83
+ for kind in random_prompt_type_choices:
84
+ word_choices = ["test", "temp", "hello", "where"]
85
+ word = random.choice(word_choices)
86
+ if kind == "repeat":
87
+ prompt = f"""please repeat the word '{word}' 10 times."""
88
+ elif kind == "sentence":
89
+ prompt = f"""please give a ten-word sentence that
90
+ uses the word {word} at least once."""
91
+ else:
92
+ raise ValueError(f"Unknown prompt type: {kind}")
93
+ prompts.append(prompt)
94
+
95
+ return prompts
96
+
97
+
98
+@fork_new_process_for_each_test
99
+@pytest.mark.parametrize("enforce_eager", [True, False])
100
+def test_kv_sharing_fast_prefill(
101
+ monkeypatch: pytest.MonkeyPatch,
102
+ enforce_eager: bool,
103
+ test_prompts: list[str],
104
+):
105
+ ModelRegistry.register_model("Gemma3nForConditionalGeneration",
106
+ TestGemma3nForConditionalGeneration)
107
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
108
+ compilation_config = CompilationConfig(
109
+ # This allows vLLM compilation backend to handle allocating and
110
+ # managing buffers for cudagraph
111
+ cudagraph_copy_inputs=True,
112
+ level=CompilationLevel.PIECEWISE
113
+ if not enforce_eager else CompilationLevel.NO_COMPILATION)
114
+
115
+ with monkeypatch.context() as m:
116
+ m.setenv("VLLM_USE_V1", "1")
117
+
118
+ llm = LLM(
119
+ model="google/gemma-3n-E2B-it",
120
+ enforce_eager=enforce_eager,
121
+ compilation_config=compilation_config,
122
+ )
123
+ ref_responses = llm.generate(test_prompts, sampling_params)
124
+
125
+ del llm
126
+ gc.collect()
127
+ torch.cuda.empty_cache()
128
+
129
+ llm = LLM(model="google/gemma-3n-E2B-it",
130
+ enforce_eager=enforce_eager,
131
+ compilation_config=compilation_config,
132
+ kv_sharing_fast_prefill=True)
133
+ optimized_responses = llm.generate(test_prompts, sampling_params)
134
+
135
+ misses = 0
136
+
137
+ for ref_response, optimized_response in zip(ref_responses,
138
+ optimized_responses):
139
+ if ref_response.outputs[0].text != optimized_response.outputs[
140
+ 0].text:
141
+ misses += 1
142
+
143
+ assert misses == 0
vllm/config.py
+15
@@ -1795,6 +1795,16 @@ class CacheConfig:
1795
num_cpu_blocks: Optional[int] = field(default=None, init=False)
1796
"""The number of blocks to allocate for CPU memory."""
1797
1798
+ kv_sharing_fast_prefill: bool = False
1799
+ """This feature is work in progress and no prefill optimization takes place
1800
+ with this flag enabled currently.
1801
+
1802
+ In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254),
1803
+ some layers can skip tokens corresponding to prefill. This flag enables
1804
+ attention metadata for eligible layers to be overriden with metadata
1805
+ necessary for implementating this optimization in some models (e.g. Gemma3n)
1806
+ """
1807
+
1808
def compute_hash(self) -> str:
1809
"""
1810
WARNING: Whenever a new field is added to this config,
@@ -1836,6 +1846,11 @@ class CacheConfig:
1846
"GPU memory utilization must be less than 1.0. Got "
1847
f"{self.gpu_memory_utilization}.")
1848
1849
+ if self.kv_sharing_fast_prefill:
1850
+ logger.warning_once(
1851
+ "--kv-sharing-fast-prefill is currently work in progress "
1852
+ "and not functional yet (i.e. no prefill savings)")
1853
+
1854
return self
1855
1856
def _verify_cache_dtype(self) -> None:
vllm/engine/arg_utils.py
+6
@@ -445,6 +445,9 @@ class EngineArgs:
445
# DEPRECATED
446
enable_prompt_adapter: bool = False
447
448
+ kv_sharing_fast_prefill: bool = \
449
+ CacheConfig.kv_sharing_fast_prefill
450
+
451
def __post_init__(self):
452
# support `EngineArgs(compilation_config={...})`
453
# without having to manually construct a
@@ -697,6 +700,8 @@ class EngineArgs:
700
**cache_kwargs["cpu_offload_gb"])
701
cache_group.add_argument("--calculate-kv-scales",
702
**cache_kwargs["calculate_kv_scales"])
703
+ cache_group.add_argument("--kv-sharing-fast-prefill",
704
+ **cache_kwargs["kv_sharing_fast_prefill"])
705
706
# Multimodal related configs
707
multimodal_kwargs = get_kwargs(MultiModalConfig)
@@ -1069,6 +1074,7 @@ class EngineArgs:
1074
prefix_caching_hash_algo=self.prefix_caching_hash_algo,
1075
cpu_offload_gb=self.cpu_offload_gb,
1076
calculate_kv_scales=self.calculate_kv_scales,
1077
+ kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
1078
)
1079
1080
# Get the current placement group if Ray is initialized and
vllm/model_executor/models/gemma3n.py
+1
@@ -793,6 +793,7 @@ class Gemma3nForConditionalGeneration(nn.Module):
793
del lora_config # Unused.
794
super().__init__()
795
self.config = config
796
+ self.cache_config = vllm_config.cache_config
797
self.model = Gemma3nModel(vllm_config=vllm_config,
798
prefix=maybe_prefix(prefix, "model"))
799
self.logits_processor = LogitsProcessor(
vllm/v1/attention/backends/utils.py
+33
-2
@@ -3,8 +3,8 @@
3
import abc
4
import functools
5
from abc import abstractmethod
6
-from dataclasses import dataclass
7
-from typing import TYPE_CHECKING, ClassVar, Generic, Optional, TypeVar
6
+from dataclasses import dataclass, make_dataclass
7
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar
8
9
import numpy as np
10
import torch
@@ -508,3 +508,34 @@ def reorder_batch_to_split_decodes_and_prefills(
508
modified_batch = True
509
510
return modified_batch
511
+
512
+
513
+KV_SHARING_FAST_PREFILL_METADATA_FIELDS = [
514
+ ('logits_indices_padded', Optional[torch.Tensor], None),
515
+ ('num_logits_indices', int, 0),
516
+]
517
+
518
+
519
+def subclass_attention_metadata(
520
+ name_prefix: str,
521
+ metadata_cls: Any,
522
+ fields: list[tuple[str, Any, Any]],
523
+) -> Any:
524
+ """
525
+ Return a new subclass of `metadata_cls` with additional fields
526
+ """
527
+ name: str = name_prefix + metadata_cls.__name__ # type: ignore
528
+ Wrapped = make_dataclass(name, fields, bases=(metadata_cls, ))
529
+ return Wrapped
530
+
531
+
532
+def make_kv_sharing_fast_prefill_attention_metadata(
533
+ metadata_cls: Any, ) -> Any:
534
+ """
535
+ Return a new subclass of `metadata_cls` for fast prefill
536
+ """
537
+ return subclass_attention_metadata(
538
+ name_prefix="KVSharingFastPrefill",
539
+ metadata_cls=metadata_cls,
540
+ fields=KV_SHARING_FAST_PREFILL_METADATA_FIELDS,
541
+ )
vllm/v1/worker/gpu_model_runner.py
+89
-24
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import dataclasses
5
import gc
6
import time
7
from contextlib import contextmanager
@@ -47,6 +48,7 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
48
from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend
49
from vllm.v1.attention.backends.utils import (
50
AttentionMetadataBuilder, CommonAttentionMetadata,
51
+ make_kv_sharing_fast_prefill_attention_metadata,
52
make_local_attention_virtual_batches)
53
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
54
from vllm.v1.kv_cache_interface import (AttentionSpec,
@@ -320,6 +322,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
322
# means this layer will perform attention using the keys and values
323
# from the KV cache of `shared_kv_cache_layers[layer_name]`.
324
self.shared_kv_cache_layers: dict[str, str] = {}
325
+ self.kv_sharing_fast_prefill_eligible_layers: set[str] = set()
326
+
327
+ self.kv_sharing_fast_prefill_logits_indices = None
328
+ if self.cache_config.kv_sharing_fast_prefill:
329
+ self.kv_sharing_fast_prefill_logits_indices = torch.zeros(
330
+ self.max_num_tokens, dtype=torch.int32, device=self.device)
331
332
def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
333
"""
@@ -735,6 +743,55 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
743
744
spec_decode_common_attn_metadata = None
745
746
+ use_spec_decode = len(
747
+ scheduler_output.scheduled_spec_decode_tokens) > 0
748
+ if not use_spec_decode:
749
+ # NOTE(woosuk): Due to chunked prefills, the batch may contain
750
+ # partial requests. While we should not sample any token
751
+ # from these partial requests, we do so for simplicity.
752
+ # We will ignore the sampled tokens from the partial requests.
753
+ # TODO: Support prompt logprobs.
754
+ logits_indices = query_start_loc[1:] - 1
755
+ spec_decode_metadata = None
756
+ else:
757
+ # Get the number of draft tokens for each request.
758
+ # Iterate over the dictionary rather than all requests since not all
759
+ # requests have draft tokens.
760
+ num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
761
+ for req_id, draft_token_ids in (
762
+ scheduler_output.scheduled_spec_decode_tokens.items()):
763
+ req_idx = self.input_batch.req_id_to_index[req_id]
764
+ num_draft_tokens[req_idx] = len(draft_token_ids)
765
+
766
+ spec_decode_metadata = self._calc_spec_decode_metadata(
767
+ num_draft_tokens, cu_num_tokens)
768
+ logits_indices = spec_decode_metadata.logits_indices
769
+
770
+ logits_indices_padded = None
771
+ if self.cache_config.kv_sharing_fast_prefill:
772
+ assert self.kv_sharing_fast_prefill_logits_indices is not None
773
+ num_logits = logits_indices.shape[0]
774
+ assert num_logits > 0
775
+ self.kv_sharing_fast_prefill_logits_indices[:num_logits].copy_(
776
+ logits_indices)
777
+ # There might have leftover indices in logits_indices[num_logits:]
778
+ # from previous iterations, whose values may be greater than the
779
+ # batch size in the current iteration. To ensure indices are always
780
+ # valid, we fill the padded indices with the last index.
781
+ self.kv_sharing_fast_prefill_logits_indices[num_logits:].fill_(
782
+ logits_indices[-1].item())
783
+ if (self.use_cuda_graph
784
+ and num_logits <= self.cudagraph_batch_sizes[-1]):
785
+ # Use piecewise CUDA graphs.
786
+ # Add padding to the batch size.
787
+ num_logits_padded = self.vllm_config.pad_for_cudagraph(
788
+ num_logits)
789
+ else:
790
+ num_logits_padded = num_logits
791
+ logits_indices_padded = (
792
+ self.kv_sharing_fast_prefill_logits_indices[:num_logits_padded]
793
+ )
794
+
795
attn_metadata: dict[str, Any] = {}
796
797
# Prepare encoder attention metadata separately
@@ -806,7 +863,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
863
common_attn_metadata=common_attn_metadata,
864
))
865
866
+ fast_prefill_metadata = attn_metadata_i
867
+ if (self.cache_config.kv_sharing_fast_prefill
868
+ and self.kv_sharing_fast_prefill_eligible_layers):
869
+ # Dynamically create a a dataclass type that inherits
870
+ # from attention metadata type but includes additional
871
+ # fields logits_indices_padded and num_logits_indices
872
+ # which are required for prefill truncation
873
+ fast_prefill_metadata_type = (
874
+ make_kv_sharing_fast_prefill_attention_metadata(
875
+ metadata_cls=type(attn_metadata_i), ))
876
+ fast_prefill_metadata = fast_prefill_metadata_type(
877
+ **dataclasses.asdict(attn_metadata_i),
878
+ logits_indices_padded=logits_indices_padded,
879
+ num_logits_indices=logits_indices.size(0),
880
+ )
881
+
882
for layer_name in kv_cache_group_spec.layer_names:
883
+ if (self.cache_config.kv_sharing_fast_prefill and layer_name
884
+ in self.kv_sharing_fast_prefill_eligible_layers):
885
+ attn_metadata[layer_name] = fast_prefill_metadata
886
+ continue
887
+
888
attn_metadata[layer_name] = attn_metadata_i
889
890
# Hack for now to fix chunked local attention + no hybrid kv cache
@@ -838,30 +916,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
916
b.can_run_in_cudagraph(common_attn_metadata)
917
for b in self.attn_metadata_builders)
918
841
- use_spec_decode = len(
842
- scheduler_output.scheduled_spec_decode_tokens) > 0
843
- if not use_spec_decode:
844
- # NOTE(woosuk): Due to chunked prefills, the batch may contain
845
- # partial requests. While we should not sample any token
846
- # from these partial requests, we do so for simplicity.
847
- # We will ignore the sampled tokens from the partial requests.
848
- # TODO: Support prompt logprobs.
849
- logits_indices = query_start_loc[1:] - 1
850
- spec_decode_metadata = None
851
- else:
852
- # Get the number of draft tokens for each request.
853
- # Iterate over the dictionary rather than all requests since not all
854
- # requests have draft tokens.
855
- num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
856
- for req_id, draft_token_ids in (
857
- scheduler_output.scheduled_spec_decode_tokens.items()):
858
- req_idx = self.input_batch.req_id_to_index[req_id]
859
- num_draft_tokens[req_idx] = len(draft_token_ids)
860
-
861
- spec_decode_metadata = self._calc_spec_decode_metadata(
862
- num_draft_tokens, cu_num_tokens)
863
- logits_indices = spec_decode_metadata.logits_indices
864
-
919
# Hot-Swap lora model
920
if self.lora_config:
921
self.set_active_loras(self.input_batch, num_scheduled_tokens)
@@ -1433,6 +1487,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1487
spec_decode_metadata, num_scheduled_tokens_np,
1488
spec_decode_common_attn_metadata) = (
1489
self._prepare_inputs(scheduler_output))
1490
+
1491
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1492
if (self.use_cuda_graph
1493
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
@@ -2814,6 +2869,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2869
kv_cache_config.kv_cache_groups,
2870
kv_caches,
2871
)
2872
+ attn_layers = get_layers_from_vllm_config(self.vllm_config,
2873
+ Attention)
2874
+ # Iterate in reversed order and add layers that re-use KV cache
2875
+ # e.g. in YOCO-like KV sharing setups (e.g. Gemma3n)
2876
+ for layer_name in reversed(attn_layers):
2877
+ if layer_name in self.shared_kv_cache_layers:
2878
+ self.kv_sharing_fast_prefill_eligible_layers.add(
2879
+ layer_name)
2880
+ else:
2881
+ break
2882
2883
bind_kv_cache(kv_caches,
2884
self.compilation_config.static_forward_context,