@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2cac89f9da

[Spec Decode] Support mixed KV page sizes for DFlash (#45181)

Signed-off-by: Alex Steiner <asteiner@nvidia.com> Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

Alex Steiner committed Jun 21, 2026 at 07:45 UTC 2cac89f9da865dfaceb6d337d97aaff5c9195e48
8 files changed +511 -141
tests/v1/core/test_kv_cache_utils.py
+102 -7
@@ -117,6 +117,7 @@ def new_kv_cache_spec(
117 page_size_padded=None,
118 sliding_window=None,
119 attention_chunk_size=None,
120 + indexes_kv_by_block_stride=False,
121 ):
122 return FullAttentionSpec(
123 block_size=block_size,
@@ -126,6 +127,7 @@ def new_kv_cache_spec(
127 page_size_padded=page_size_padded,
128 sliding_window=sliding_window,
129 attention_chunk_size=attention_chunk_size,
130 + indexes_kv_by_block_stride=indexes_kv_by_block_stride,
131 )
132
133
@@ -136,6 +138,7 @@ def new_sliding_window_spec(
138 dtype=torch.float32,
139 page_size_padded=None,
140 sliding_window=1,
141 + indexes_kv_by_block_stride=False,
142 ):
143 return SlidingWindowSpec(
144 block_size=block_size,
@@ -144,6 +147,7 @@ def new_sliding_window_spec(
147 dtype=dtype,
148 page_size_padded=page_size_padded,
149 sliding_window=sliding_window,
150 + indexes_kv_by_block_stride=indexes_kv_by_block_stride,
151 )
152
153
@@ -1799,16 +1803,38 @@ def test_get_kv_cache_config_one_worker():
1803 ],
1804 )
1805
1802 - # different hidden size that cannot be aligned by using different block size
1806 + # different hidden size that cannot be aligned by using different block size,
1807 + # but can be aligned by padding the smaller physical page.
1808 + swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True)
1809 kv_cache_specs_hybrid = {
1804 - "layer_1": new_kv_cache_spec(head_size=64),
1805 - "layer_2": new_sliding_window_spec(head_size=96),
1810 + "layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True),
1811 + "layer_2": swa_spec,
1812 }
1813
1808 - with pytest.raises(NotImplementedError):
1809 - get_kv_cache_configs(
1810 - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
1811 - )[0]
1814 + kv_cache_config_hybrid = get_kv_cache_configs(
1815 + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
1816 + )[0]
1817 + padded_page_size = swa_spec.page_size_bytes
1818 + assert kv_cache_config_hybrid == KVCacheConfig(
1819 + num_blocks=42,
1820 + kv_cache_tensors=[
1821 + KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]),
1822 + ],
1823 + kv_cache_groups=[
1824 + KVCacheGroupSpec(
1825 + ["layer_1"],
1826 + new_kv_cache_spec(
1827 + head_size=64,
1828 + page_size_padded=padded_page_size,
1829 + indexes_kv_by_block_stride=True,
1830 + ),
1831 + ),
1832 + KVCacheGroupSpec(
1833 + ["layer_2"],
1834 + new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True),
1835 + ),
1836 + ],
1837 + )
1838
1839 # Test num_gpu_blocks_override
1840 vllm_config.cache_config.num_gpu_blocks_override = 16
@@ -2322,6 +2348,75 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
2348 get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
2349
2350
2351 +def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
2352 + """DFlash drafters can have a smaller head size than the target model.
2353 +
2354 + For example, MiMo uses 192-dim target KV heads while its DFlash draft uses
2355 + 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer
2356 + block-size multiple, so the smaller page must be padded instead.
2357 + """
2358 + # Both layers' backends opt into the padded-page strided view (e.g.
2359 + # FlashAttention / its DiffKV subclass), so padding is allowed.
2360 + target_spec = new_kv_cache_spec(
2361 + block_size=16,
2362 + num_kv_heads=1,
2363 + head_size=192,
2364 + dtype=torch.bfloat16,
2365 + indexes_kv_by_block_stride=True,
2366 + )
2367 + draft_spec = new_sliding_window_spec(
2368 + block_size=16,
2369 + num_kv_heads=1,
2370 + head_size=128,
2371 + dtype=torch.bfloat16,
2372 + sliding_window=1024,
2373 + indexes_kv_by_block_stride=True,
2374 + )
2375 +
2376 + unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size(
2377 + {
2378 + "target_attn": target_spec,
2379 + "draft_attn": draft_spec,
2380 + }
2381 + )
2382 +
2383 + assert unified_specs["target_attn"] == target_spec
2384 + unified_draft_spec = unified_specs["draft_attn"]
2385 + assert unified_draft_spec.block_size == draft_spec.block_size
2386 + assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes
2387 + assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes
2388 + assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes
2389 +
2390 +
2391 +def test_unify_kv_cache_page_size_padding_requires_backend_support():
2392 + """Padding is gated on the backend declaring ``indexes_kv_by_block_stride``.
2393 +
2394 + A backend that does not support the strided padded-page view must raise
2395 + rather than silently padding (and misreading KV at runtime).
2396 + """
2397 + target_spec = new_kv_cache_spec(
2398 + block_size=16,
2399 + num_kv_heads=1,
2400 + head_size=192,
2401 + dtype=torch.bfloat16,
2402 + indexes_kv_by_block_stride=True,
2403 + )
2404 + # The non-divisible draft layer needs padding but its backend does not
2405 + # support the strided padded-page view -> must raise, not silently pad.
2406 + draft_spec = new_sliding_window_spec(
2407 + block_size=16,
2408 + num_kv_heads=1,
2409 + head_size=128,
2410 + dtype=torch.bfloat16,
2411 + sliding_window=1024,
2412 + indexes_kv_by_block_stride=False,
2413 + )
2414 + specs = {"target_attn": target_spec, "draft_attn": draft_spec}
2415 +
2416 + with pytest.raises(NotImplementedError):
2417 + kv_cache_utils.unify_kv_cache_spec_page_size(specs)
2418 +
2419 +
2420 def test_unify_hybrid_kv_cache_specs():
2421 # 1. has_full_attention and has_sliding_window
2422 before_spec_1 = new_kv_cache_spec()
tests/v1/worker/test_attn_utils.py new
+242
@@ -0,0 +1,242 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
7 +from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
8 +from vllm.v1.worker.utils import AttentionGroup
9 +
10 +
11 +class FakeFlashAttentionBackend:
12 + @staticmethod
13 + def get_kv_cache_shape(
14 + num_blocks: int,
15 + block_size: int,
16 + num_kv_heads: int,
17 + head_size: int,
18 + cache_dtype_str: str = "auto",
19 + ) -> tuple[int, ...]:
20 + return (num_blocks, 2, block_size, num_kv_heads, head_size)
21 +
22 + @staticmethod
23 + def get_kv_cache_stride_order(
24 + include_num_layers_dimension: bool = False,
25 + ) -> tuple[int, ...]:
26 + assert not include_num_layers_dimension
27 + return (0, 1, 2, 3, 4)
28 +
29 +
30 +class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend):
31 + @staticmethod
32 + def get_kv_cache_stride_order(
33 + include_num_layers_dimension: bool = False,
34 + ) -> tuple[int, ...]:
35 + assert not include_num_layers_dimension
36 + return (0, 1, 3, 2, 4)
37 +
38 +
39 +def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
40 + num_blocks = 3
41 + spec = FullAttentionSpec(
42 + block_size=16,
43 + num_kv_heads=1,
44 + head_size=2,
45 + dtype=torch.float32,
46 + page_size_padded=384,
47 + )
48 + assert spec.real_page_size_bytes == 256
49 +
50 + raw_tensors = {
51 + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
52 + }
53 + attn_groups = [
54 + AttentionGroup(
55 + backend=FakeFlashAttentionBackend,
56 + layer_names=["layer"],
57 + kv_cache_spec=spec,
58 + kv_cache_group_id=0,
59 + )
60 + ]
61 +
62 + kv_cache = _reshape_kv_cache(
63 + attn_groups,
64 + raw_tensors,
65 + "auto",
66 + [spec.block_size],
67 + {},
68 + )["layer"]
69 +
70 + assert kv_cache.shape == (num_blocks, 2, 16, 1, 2)
71 + assert kv_cache.stride(0) == spec.page_size_bytes // 4
72 + assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
73 + assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
74 + assert (
75 + kv_cache[1, 1].storage_offset()
76 + == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
77 + )
78 +
79 +
80 +def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page():
81 + num_blocks = 3
82 + spec = FullAttentionSpec(
83 + block_size=16,
84 + num_kv_heads=3,
85 + head_size=2,
86 + dtype=torch.float32,
87 + page_size_padded=1024,
88 + )
89 + assert spec.real_page_size_bytes == 768
90 +
91 + raw_tensors = {
92 + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
93 + }
94 + attn_groups = [
95 + AttentionGroup(
96 + backend=FakeHNDFlashAttentionBackend,
97 + layer_names=["layer"],
98 + kv_cache_spec=spec,
99 + kv_cache_group_id=0,
100 + )
101 + ]
102 +
103 + kv_cache = _reshape_kv_cache(
104 + attn_groups,
105 + raw_tensors,
106 + "auto",
107 + [spec.block_size],
108 + {},
109 + )["layer"]
110 +
111 + assert kv_cache.shape == (num_blocks, 2, 16, 3, 2)
112 + assert kv_cache.stride(0) == spec.page_size_bytes // 4
113 + assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
114 + assert kv_cache.stride(2) == 2
115 + assert kv_cache.stride(3) == spec.block_size * spec.head_size
116 + assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
117 + assert (
118 + kv_cache[1, 1].storage_offset()
119 + == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
120 + )
121 + assert (
122 + kv_cache[1, 1, 3, 2].storage_offset()
123 + == (
124 + spec.page_size_bytes
125 + + spec.real_page_size_bytes // 2
126 + + 3 * spec.head_size * 4
127 + + 2 * spec.block_size * spec.head_size * 4
128 + )
129 + // 4
130 + )
131 +
132 +
133 +class FakeDiffKVBackend:
134 + @staticmethod
135 + def get_kv_cache_shape(
136 + num_blocks: int,
137 + block_size: int,
138 + num_kv_heads: int,
139 + head_size: int,
140 + cache_dtype_str: str = "auto",
141 + ) -> tuple[int, ...]:
142 + return (num_blocks, block_size, num_kv_heads, head_size * 2)
143 +
144 + @staticmethod
145 + def get_kv_cache_stride_order(
146 + include_num_layers_dimension: bool = False,
147 + ) -> tuple[int, ...]:
148 + assert not include_num_layers_dimension
149 + return (0, 1, 2, 3)
150 +
151 +
152 +def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim():
153 + num_blocks = 3
154 + spec = FullAttentionSpec(
155 + block_size=16,
156 + num_kv_heads=1,
157 + head_size=2,
158 + dtype=torch.float32,
159 + page_size_padded=384,
160 + )
161 +
162 + raw_tensors = {
163 + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
164 + }
165 + attn_groups = [
166 + AttentionGroup(
167 + backend=FakeDiffKVBackend,
168 + layer_names=["layer"],
169 + kv_cache_spec=spec,
170 + kv_cache_group_id=0,
171 + )
172 + ]
173 +
174 + kv_cache = _reshape_kv_cache(
175 + attn_groups,
176 + raw_tensors,
177 + "auto",
178 + [spec.block_size],
179 + {},
180 + )["layer"]
181 +
182 + assert kv_cache.shape == (num_blocks, 16, 1, 4)
183 + assert kv_cache.stride(0) == spec.page_size_bytes // 4
184 + assert kv_cache.stride(1) == 4
185 +
186 +
187 +class FakePerTokenScaleBackend:
188 + @staticmethod
189 + def get_kv_cache_shape(
190 + num_blocks: int,
191 + block_size: int,
192 + num_kv_heads: int,
193 + head_size: int,
194 + cache_dtype_str: str = "auto",
195 + ) -> tuple[int, ...]:
196 + return (num_blocks, 2, block_size, num_kv_heads, head_size + 4)
197 +
198 + @staticmethod
199 + def get_kv_cache_stride_order(
200 + include_num_layers_dimension: bool = False,
201 + ) -> tuple[int, ...]:
202 + assert not include_num_layers_dimension
203 + return (0, 1, 2, 3, 4)
204 +
205 +
206 +def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
207 + num_blocks = 3
208 + spec = FullAttentionSpec(
209 + block_size=16,
210 + num_kv_heads=1,
211 + head_size=4,
212 + dtype=torch.int8,
213 + kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
214 + page_size_padded=384,
215 + )
216 + assert spec.real_page_size_bytes == 128
217 + assert spec.page_size_bytes == 384
218 +
219 + raw_tensors = {
220 + "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
221 + }
222 + attn_groups = [
223 + AttentionGroup(
224 + backend=FakePerTokenScaleBackend,
225 + layer_names=["layer"],
226 + kv_cache_spec=spec,
227 + kv_cache_group_id=0,
228 + )
229 + ]
230 +
231 + kv_cache = _reshape_kv_cache(
232 + attn_groups,
233 + raw_tensors,
234 + "int8_per_token_head",
235 + [spec.block_size],
236 + {},
237 + )["layer"]
238 +
239 + assert kv_cache.shape == (num_blocks, 2, 16, 1, 8)
240 + assert kv_cache.stride(0) == spec.page_size_bytes
241 + assert kv_cache.stride(1) == 16 * 1 * 8
242 + assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8
vllm/v1/attention/backend.py
+32
@@ -201,6 +201,38 @@ class AttentionBackend(ABC):
201
202 return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)
203
204 + @classmethod
205 + def indexes_kv_by_block_stride(cls) -> bool:
206 + """Whether the backend reads KV pages by the runtime block stride.
207 +
208 + True when ``num_blocks`` is the outermost physical dimension of the KV
209 + cache, so the backend tolerates a non-contiguous block dim. This gates
210 + page size padding and cross-layer uniform KV layout.
211 +
212 + Returns:
213 + True if the backend's physical KV layout is num-blocks-first. False
214 + otherwise, including when the backend does not define a layered
215 + stride order.
216 + """
217 + try:
218 + kv_cache_stride_order = cls.get_kv_cache_stride_order(
219 + include_num_layers_dimension=False
220 + )
221 + layered_kv_cache_stride_order = cls.get_kv_cache_stride_order(
222 + include_num_layers_dimension=True
223 + )
224 + except (AttributeError, NotImplementedError):
225 + return False
226 +
227 + # Check that attention backend includes a layers dimension.
228 + if len(layered_kv_cache_stride_order) != len(kv_cache_stride_order) + 1:
229 + return False
230 +
231 + # stride_order[0] == 0 means num_layers stays first in physical
232 + # layout (identity permutation), so indexing by block stride is
233 + # not supported.
234 + return layered_kv_cache_stride_order[0] != 0
235 +
236 @classmethod
237 def is_mla(cls) -> bool:
238 return False
vllm/v1/core/kv_cache_utils.py
+24 -9
@@ -20,6 +20,7 @@ from vllm.utils.math_utils import cdiv, round_up
20 from vllm.utils.mem_utils import format_gib
21 from vllm.utils.torch_utils import get_dtype_size
22 from vllm.v1.kv_cache_interface import (
23 + AttentionSpec,
24 ChunkedLocalAttentionSpec,
25 FullAttentionSpec,
26 HiddenStateCacheSpec,
@@ -1029,9 +1030,14 @@ def unify_kv_cache_spec_page_size(
1030 ) -> dict[str, KVCacheSpec]:
1031 """
1032 Unify the page size of the given KVCacheSpec. If the page size of all layers
1032 - are the same, return the original KVCacheSpec. If not same, unify the page
1033 - size by increasing the block size of layers with smaller page size. Raise
1034 - NotImplementedError if failed to unify the page size.
1033 + are the same, return the original KVCacheSpec. If not same, first try to
1034 + unify page size by increasing the block size of layers with smaller page
1035 + size. If a smaller attention page does not evenly divide the maximum page
1036 + size, keep its logical block size and pad its physical page instead --- but
1037 + only for attention layers whose backend opts in via
1038 + ``AttentionSpec.indexes_kv_by_block_stride`` (the padded page is read through
1039 + a strided view, which not every backend handles). Raise NotImplementedError
1040 + if failed to unify the page size.
1041
1042 Args:
1043 kv_cache_spec: The KVCacheSpec of each attention layer in the model
@@ -1051,14 +1057,23 @@ def unify_kv_cache_spec_page_size(
1057 new_kv_cache_spec[layer_name] = layer_spec
1058 else:
1059 layer_page_size = layer_spec.page_size_bytes
1054 - if max_page_size % layer_page_size != 0:
1060 + if max_page_size % layer_page_size == 0:
1061 + ratio = max_page_size // layer_page_size
1062 + new_block_size = layer_spec.block_size * ratio
1063 + new_spec = replace(layer_spec, block_size=new_block_size)
1064 + elif (
1065 + isinstance(layer_spec, AttentionSpec)
1066 + and layer_spec.indexes_kv_by_block_stride
1067 + ):
1068 + new_spec = replace(layer_spec, page_size_padded=max_page_size)
1069 + else:
1070 raise NotImplementedError(
1056 - "The page size of the layer is not divisible by the "
1057 - "maximum page size. Cannot unify by adjusting block_size."
1071 + f"Layer {layer_name}: page size is not divisible by the "
1072 + "maximum page size and cannot be padded. Padding is only "
1073 + "supported for attention layers whose backend indexes KV "
1074 + "pages by the block stride (indexes_kv_by_block_stride is "
1075 + "True)."
1076 )
1059 - ratio = max_page_size // layer_page_size
1060 - new_block_size = layer_spec.block_size * ratio
1061 - new_spec = replace(layer_spec, block_size=new_block_size)
1077 assert new_spec.page_size_bytes == max_page_size
1078 new_kv_cache_spec[layer_name] = new_spec
1079 return new_kv_cache_spec
vllm/v1/kv_cache_interface.py
+13 -3
@@ -163,6 +163,7 @@ class AttentionSpec(KVCacheSpec):
163 dtype: torch.dtype
164 kv_quant_mode: KVQuantMode = KVQuantMode.NONE
165 page_size_padded: int | None = None
166 + indexes_kv_by_block_stride: bool = False
167
168 @property
169 def page_size_bytes(self) -> int:
@@ -283,6 +284,7 @@ class FullAttentionSpec(AttentionSpec):
284 dtype=specs[0].dtype,
285 kv_quant_mode=specs[0].kv_quant_mode,
286 page_size_padded=specs[0].page_size_padded,
287 + indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
288 sliding_window=cls.merge_window_sizes(sliding_window),
289 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
290 # If any layer in the group is non-causal, treat the group as
@@ -403,13 +405,16 @@ class MLAAttentionSpec(FullAttentionSpec):
405 cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
406 compress_ratio_set = set(spec.compress_ratio for spec in specs)
407 model_version_set = set(spec.model_version for spec in specs)
408 + block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
409 assert (
410 len(cache_dtype_str_set) == 1
411 and len(compress_ratio_set) == 1
412 and len(model_version_set) == 1
413 + and len(block_stride_set) == 1
414 ), (
415 "All attention layers in the same KV cache group must use the same "
412 - "quantization method, compress ratio, and model version."
416 + "quantization method, compress ratio, model version, and KV block "
417 + "stride indexing."
418 )
419 return cls(
420 block_size=specs[0].block_size,
@@ -418,6 +423,7 @@ class MLAAttentionSpec(FullAttentionSpec):
423 dtype=specs[0].dtype,
424 kv_quant_mode=specs[0].kv_quant_mode,
425 page_size_padded=specs[0].page_size_padded,
426 + indexes_kv_by_block_stride=block_stride_set.pop(),
427 cache_dtype_str=cache_dtype_str_set.pop(),
428 compress_ratio=compress_ratio_set.pop(),
429 model_version=model_version_set.pop(),
@@ -584,15 +590,17 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
590 compress_ratio_set = set(spec.compress_ratio for spec in specs)
591 model_version_set = set(spec.model_version for spec in specs)
592 sliding_window_set = set(spec.sliding_window for spec in specs)
593 + block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
594 assert (
595 len(cache_dtype_str_set) == 1
596 and len(compress_ratio_set) == 1
597 and len(model_version_set) == 1
598 and len(sliding_window_set) == 1
599 + and len(block_stride_set) == 1
600 ), (
601 "All attention layers in the same KV cache group must use the same "
594 - "quantization method, compress ratio, model version and sliding "
595 - "window size."
602 + "quantization method, compress ratio, model version, sliding "
603 + "window size, and KV block stride indexing."
604 )
605 return cls(
606 block_size=specs[0].block_size,
@@ -600,6 +608,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
608 head_size=specs[0].head_size,
609 dtype=specs[0].dtype,
610 page_size_padded=specs[0].page_size_padded,
611 + indexes_kv_by_block_stride=block_stride_set.pop(),
612 sliding_window=sliding_window_set.pop(),
613 cache_dtype_str=cache_dtype_str_set.pop(),
614 compress_ratio=compress_ratio_set.pop(),
@@ -711,6 +720,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
720 dtype=specs[0].dtype,
721 kv_quant_mode=specs[0].kv_quant_mode,
722 page_size_padded=specs[0].page_size_padded,
723 + indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
724 sliding_window=cls.merge_window_sizes(sliding_window),
725 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
726 non_causal=any(spec.non_causal for spec in specs),
vllm/v1/worker/gpu/attn_utils.py
+77 -41
@@ -1,13 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Iterable, Sequence
4 -from dataclasses import dataclass
4 +from dataclasses import dataclass, replace
5 from math import prod
6 from typing import Any, cast
7
8 import torch
9
10 -from vllm.config import VllmConfig, get_layers_from_vllm_config
10 +from vllm.config import (
11 + VllmConfig,
12 + get_layers_from_vllm_config,
13 + set_current_vllm_config,
14 +)
15 from vllm.model_executor.layers.attention import Attention
16 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
17 from vllm.utils.torch_utils import get_dtype_size
@@ -47,6 +51,13 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
51 continue
52 # Skip modules that don't need KV cache (eg encoder-only attention)
53 if spec := attn_module.get_kv_cache_spec(vllm_config):
54 + if isinstance(spec, AttentionSpec):
55 + backend = attn_module.get_attn_backend()
56 + # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() ->
57 + # get_kv_cache_layout() needs the current vLLM config.
58 + with set_current_vllm_config(vllm_config):
59 + indexes = backend.indexes_kv_by_block_stride()
60 + spec = replace(spec, indexes_kv_by_block_stride=indexes)
61 kv_cache_spec[layer_name] = spec
62 return kv_cache_spec
63
@@ -180,6 +191,62 @@ def _allocate_kv_cache(
191 return kv_cache_raw_tensors
192
193
194 +def _reshape_attention_kv_cache(
195 + kv_raw_tensor: torch.Tensor,
196 + kv_cache_spec: AttentionSpec,
197 + kv_cache_shape: tuple[int, ...],
198 + kv_cache_stride_order: tuple[int, ...],
199 + num_blocks: int,
200 + packing: tuple[int, int] | None,
201 +) -> torch.Tensor:
202 + permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
203 + inv_order = [
204 + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
205 + ]
206 + dtype = kv_cache_spec.dtype
207 +
208 + if packing is not None:
209 + offset, block_stride = packing
210 + assert inv_order[0] == 0
211 + page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
212 + kv_cache = (
213 + kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes]
214 + .view(dtype)
215 + .view(kv_cache_shape)
216 + )
217 + elif kv_cache_spec.page_size_padded is not None:
218 + # Use a strided view to skip the padding between physical pages.
219 + #
220 + # Only num-blocks-first layouts are supported (the block dimension is
221 + # dim 0 of the unpermuted shape). kv-first layouts such as ROCm's
222 + # ``(2, num_blocks, ...)`` are intentionally not supported here. For a
223 + # num-blocks-first layout the only stride that must change is the block
224 + # stride: every other (contiguous) stride already steps within the
225 + # unpadded region of a page, so no further adjustment is needed.
226 + assert kv_cache_shape[0] == num_blocks, (
227 + "Padded KV pages require a num-blocks-first KV cache layout (got "
228 + f"shape {kv_cache_shape} with num_blocks={num_blocks}); "
229 + "kv-first layouts are not supported."
230 + )
231 + dtype_size = get_dtype_size(kv_cache_spec.dtype)
232 + page_stride = kv_cache_spec.page_size_bytes // dtype_size
233 +
234 + num_blocks_dim = inv_order[0]
235 + strides = list(torch.empty(permuted_kv_cache_shape).stride())
236 + strides[num_blocks_dim] = page_stride
237 +
238 + kv_cache = torch.as_strided(
239 + kv_raw_tensor.view(dtype),
240 + size=permuted_kv_cache_shape,
241 + stride=tuple(strides),
242 + )
243 + else:
244 + # No padding — safe to use a contiguous view.
245 + kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape)
246 +
247 + return kv_cache.permute(*inv_order)
248 +
249 +
250 def _reshape_kv_cache(
251 attn_groups: Sequence[AttentionGroup],
252 kv_cache_raw_tensors: dict[str, torch.Tensor],
@@ -248,45 +315,14 @@ def _reshape_kv_cache(
315 except (AttributeError, NotImplementedError):
316 kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
317
251 - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
252 - inv_order = [
253 - kv_cache_stride_order.index(i)
254 - for i in range(len(kv_cache_stride_order))
255 - ]
256 -
257 - dtype = kv_cache_spec.dtype
258 - if packing is not None:
259 - offset, block_stride = packing
260 - assert inv_order[0] == 0
261 - page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
262 - kv_cache = (
263 - kv_raw_tensor.view(-1, block_stride)[
264 - :, offset : offset + page_bytes
265 - ]
266 - .view(dtype)
267 - .view(kv_cache_shape)
268 - )
269 - elif kv_cache_spec.page_size_padded is not None:
270 - # Use strided view to handle page_size_bytes that
271 - # include padding. This follows the same pattern as
272 - # MambaSpec handling in gpu_model_runner.py.
273 - # NOTE: This assumes kv_cache_shape[0] == num_blocks
274 - # (i.e. the first physical dimension is the block
275 - # index), which holds for all current backends
276 - # (MLA, FlashAttention, TritonAttention, etc.).
277 - dtype_size = get_dtype_size(dtype)
278 - page_stride = kv_cache_spec.page_size_bytes // dtype_size
279 - strides = list(torch.empty(kv_cache_shape).stride())
280 - strides[inv_order[0]] = page_stride
281 - kv_cache = torch.as_strided(
282 - kv_raw_tensor.view(dtype),
283 - size=kv_cache_shape,
284 - stride=tuple(strides),
285 - )
286 - else:
287 - # No padding — safe to use a contiguous view.
288 - kv_cache = kv_raw_tensor.view(dtype).view(kv_cache_shape)
289 - kv_caches[layer_name] = kv_cache.permute(*inv_order)
318 + kv_caches[layer_name] = _reshape_attention_kv_cache(
319 + kv_raw_tensor,
320 + kv_cache_spec,
321 + kv_cache_shape,
322 + kv_cache_stride_order,
323 + kernel_num_blocks,
324 + packing,
325 + )
326
327 elif isinstance(kv_cache_spec, MambaSpec):
328 has_mamba = True
vllm/v1/worker/gpu_model_runner.py
+17 -52
@@ -12,7 +12,6 @@ from contextlib import contextmanager
12 from copy import copy, deepcopy
13 from dataclasses import dataclass, replace
14 from functools import reduce
15 -from math import prod
15 from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast
16
17 import numpy as np
@@ -203,6 +202,7 @@ from vllm.v1.worker.cp_utils import (
202 )
203 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
204 from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
205 +from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache
206 from vllm.v1.worker.gpu.pool.late_interaction_runner import LateInteractionRunner
207 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
208 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
@@ -7125,62 +7125,20 @@ class GPUModelRunner(
7125 kv_cache_spec.head_size,
7126 cache_dtype_str=self.cache_config.cache_dtype,
7127 )
7128 - dtype = kv_cache_spec.dtype
7128 try:
7129 kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
7130 assert len(kv_cache_stride_order) == len(kv_cache_shape)
7131 except (AttributeError, NotImplementedError):
7132 kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
7134 - # The allocation respects the backend-defined stride order
7135 - # to ensure the semantic remains consistent for each
7136 - # backend. We first obtain the generic kv cache shape and
7137 - # then permute it according to the stride order which could
7138 - # result in a non-contiguous tensor.
7139 - kv_cache_shape = tuple(
7140 - kv_cache_shape[i] for i in kv_cache_stride_order
7133 + raw_tensor = kv_cache_raw_tensors[layer_name]
7134 + kv_caches[layer_name] = _reshape_attention_kv_cache(
7135 + raw_tensor,
7136 + kv_cache_spec,
7137 + kv_cache_shape,
7138 + kv_cache_stride_order,
7139 + kernel_num_blocks,
7140 + packing,
7141 )
7142 - # Maintain original KV shape view.
7143 - inv_order = [
7144 - kv_cache_stride_order.index(i)
7145 - for i in range(len(kv_cache_stride_order))
7146 - ]
7147 -
7148 - if packing is not None:
7149 - offset, block_stride = packing
7150 - assert inv_order[0] == 0
7151 - page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
7152 - kv_cache = (
7153 - kv_cache_raw_tensors[layer_name]
7154 - .view(-1, block_stride)[:, offset : offset + page_bytes]
7155 - .view(dtype)
7156 - .view(kv_cache_shape)
7157 - )
7158 - elif kv_cache_spec.page_size_padded is not None:
7159 - # Use strided view to handle page_size_bytes that
7160 - # include padding. This follows
7161 - # the same pattern as MambaSpec handling below.
7162 - # NOTE: This assumes kv_cache_shape[0] == num_blocks
7163 - # (i.e. the first physical dimension is the block
7164 - # index), which holds for MLA backends but NOT for
7165 - # standard attention backends whose shape starts with
7166 - # a K/V dimension of size 2.
7167 - dtype_size = get_dtype_size(dtype)
7168 - page_stride = kv_cache_spec.page_size_bytes // dtype_size
7169 - strides = list(torch.empty(kv_cache_shape).stride())
7170 - strides[inv_order[0]] = page_stride
7171 - kv_cache = torch.as_strided(
7172 - kv_cache_raw_tensors[layer_name].view(dtype),
7173 - size=kv_cache_shape,
7174 - stride=tuple(strides),
7175 - )
7176 - else:
7177 - # No padding — safe to use a contiguous view.
7178 - kv_cache = (
7179 - kv_cache_raw_tensors[layer_name]
7180 - .view(dtype)
7181 - .view(kv_cache_shape)
7182 - )
7183 - kv_caches[layer_name] = kv_cache.permute(*inv_order)
7142
7143 elif isinstance(kv_cache_spec, MambaSpec):
7144 has_mamba = True
@@ -7265,7 +7223,7 @@ class GPUModelRunner(
7223
7224 # Try creating KV caches optimized for kv-connector transfers
7225 cache_dtype = self.cache_config.cache_dtype
7268 - if self.use_uniform_kv_cache(self.attn_groups, cache_dtype):
7226 + if self.use_uniform_kv_cache(self.attn_groups):
7227 kv_caches, cross_layers_kv_cache, attn_backend = (
7228 self.allocate_uniform_kv_caches(
7229 kv_cache_config,
@@ -7515,6 +7473,13 @@ class GPUModelRunner(
7473 continue
7474 # Skip modules that don't need KV cache (eg encoder-only attention)
7475 if spec := attn_module.get_kv_cache_spec(self.vllm_config):
7476 + if isinstance(spec, AttentionSpec):
7477 + backend = attn_module.get_attn_backend()
7478 + # indexes_kv_by_block_stride() -> get_kv_cache_stride_order()
7479 + # -> get_kv_cache_layout() needs the current vLLM config.
7480 + with set_current_vllm_config(self.vllm_config):
7481 + indexes = backend.indexes_kv_by_block_stride()
7482 + spec = replace(spec, indexes_kv_by_block_stride=indexes)
7483 kv_cache_spec[layer_name] = spec
7484
7485 return kv_cache_spec
vllm/v1/worker/kv_connector_model_runner_mixin.py
+4 -29
@@ -114,7 +114,6 @@ class KVConnectorModelRunnerMixin:
114 @staticmethod
115 def use_uniform_kv_cache(
116 attn_groups: list[list[AttentionGroup]],
117 - cache_dtype: CacheDType,
117 ) -> bool:
118 """
119 Determines whether a uniform KV layout should be used.
@@ -128,9 +127,9 @@ class KVConnectorModelRunnerMixin:
127 have the same page size.
128 2. A KV connector is configured, and the KV connector instance prefers
129 to use this layout (prefer_cross_layer_blocks() returns True)
131 - 2. The flash attention backend supports this layout
132 - (get_kv_cache_stride_order(True) includes a placement for a
133 - num_layers dimension)
130 + 3. The attention backend indexes KV by the block stride
131 + (kv_cache_spec.indexes_kv_by_block_stride), i.e. num_blocks is the
132 + outermost physical dim so per-block all-layers data is contiguous.
133
134 Note that the actual placement of the num_layers dimensions
135 in the unified layers tensors will be determined by the attention
@@ -140,7 +139,6 @@ class KVConnectorModelRunnerMixin:
139
140 Args:
141 attn_groups: The list of attention groups for this model
143 - cache_dtype: The KV cache dtype
142 Returns:
143 True if we should use a uniform KV cache layout.
144 """
@@ -157,30 +155,7 @@ class KVConnectorModelRunnerMixin:
155 kv_cache_spec = attn_group.kv_cache_spec
156 if not isinstance(kv_cache_spec, AttentionSpec):
157 return False
160 -
161 - attn_backend = attn_group.backend
162 - kv_cache_shape = attn_backend.get_kv_cache_shape(
163 - 1234,
164 - kv_cache_spec.block_size,
165 - kv_cache_spec.num_kv_heads,
166 - kv_cache_spec.head_size,
167 - cache_dtype_str=cache_dtype,
168 - )
169 -
170 - try:
171 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
172 - include_num_layers_dimension=True
173 - )
174 - except (AttributeError, NotImplementedError):
175 - return False
176 -
177 - # check that attention backend includes a layers dimension
178 - if len(kv_cache_stride_order) != len(kv_cache_shape) + 1:
179 - return False
180 -
181 - # stride_order[0] == 0 means num_layers stays first in physical
182 - # layout (identity permutation), so cross-layer is unsupported.
183 - return kv_cache_stride_order[0] != 0
158 + return kv_cache_spec.indexes_kv_by_block_stride
159
160 @staticmethod
161 def allocate_uniform_kv_caches(