@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8bdc70ec7b

[6/N][KV-Cache Layout Refactor] Standardize KV cache layout (#51718)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Aug 21, 2026 at 20:47 UTC 8bdc70ec7b379279ec0152343239c2d50aced687
150 files changed +3993 -5802
.buildkite/test_areas/disaggregated.yaml
+3
@@ -243,6 +243,9 @@ steps:
243 DECODE_BLOCK_SIZE: "256"
244 MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
245 VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
246 + # The V4 indexer only supports block-outermost layouts; let it pick.
247 + PREFILLER_KV_LAYOUT: ""
248 + DECODER_KV_LAYOUT: ""
249 commands:
250 - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
251 - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
benchmarks/attention_benchmarks/runner.py
+47 -48
@@ -12,6 +12,7 @@ import logging
12 import statistics
13 import types
14 from contextlib import contextmanager
15 +from math import prod
16
17 import torch
18 from batch_spec import parse_batch_spec, reorder_for_flashinfer
@@ -35,12 +36,20 @@ from vllm.config import (
36 VllmConfig,
37 set_current_vllm_config,
38 )
39 +from vllm.platforms import current_platform
40 from vllm.v1.attention.backends.utils import (
41 CommonAttentionMetadata,
40 - get_kv_cache_layout,
41 - set_kv_cache_layout,
42 + get_supported_kv_cache_layouts,
43 + resolve_kv_cache_layout,
44 +)
45 +from vllm.v1.kv_cache_interface import (
46 + FullAttentionSpec,
47 + KVCacheLayout,
48 + KVCacheTensor,
49 + compute_layer_kv_cache_shape_bytes,
50 + compute_layout_strides,
51 + create_kv_cache_views,
52 )
43 -from vllm.v1.kv_cache_interface import FullAttentionSpec
53
54 # ============================================================================
55 # Backend Configuration
@@ -340,49 +349,42 @@ def _create_kv_cache(
349 backend_class,
350 device: torch.device,
351 dtype: torch.dtype,
352 + layout: KVCacheLayout,
353 ) -> list:
344 - """Create KV cache tensors for all layers using the backend's methods.
345 -
346 - Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
347 - to create the cache with the correct shape and memory layout.
348 - """
349 - # Get the logical shape from the backend
350 - cache_shape = backend_class.get_kv_cache_shape(
351 - num_blocks=max_num_blocks,
354 + """Create KV cache tensors for all layers using the standard allocator."""
355 + if config.kv_cache_dtype.startswith("fp8"):
356 + cache_dtype = current_platform.fp8_dtype()
357 + else:
358 + cache_dtype = dtype
359 + spec = FullAttentionSpec(
360 block_size=config.block_size,
361 num_kv_heads=config.num_kv_heads,
362 head_size=config.head_dim,
363 + dtype=cache_dtype,
364 + )
365 + # Apply the backend's page customization, as the worker does for the real spec.
366 + spec = backend_class.customize_spec(spec)
367 + total_bytes = (
368 + prod(compute_layer_kv_cache_shape_bytes(spec, max_num_blocks))
369 + * config.num_layers
370 + )
371 + buf = torch.zeros(total_bytes, device=device, dtype=torch.int8)
372 + layer_stride, block_stride, _, _, _ = compute_layout_strides(
373 + spec, max_num_blocks, config.num_layers, layout
374 + )
375 + tensor = KVCacheTensor(
376 + size=total_bytes,
377 + layers=[str(i) for i in range(config.num_layers)],
378 + layer_stride=layer_stride,
379 + block_stride=block_stride,
380 + )
381 + return create_kv_cache_views(
382 + buf,
383 + spec,
384 + max_num_blocks,
385 + layout,
386 + tensor,
387 )
356 -
357 - # Get the stride order for custom memory layout
358 - try:
359 - stride_order = backend_class.get_kv_cache_stride_order()
360 - assert len(stride_order) == len(cache_shape)
361 - except (AttributeError, NotImplementedError):
362 - stride_order = tuple(range(len(cache_shape)))
363 -
364 - # Permute shape to physical layout order
365 - physical_shape = tuple(cache_shape[i] for i in stride_order)
366 -
367 - # Compute inverse permutation to get back to logical view
368 - inv_order = [stride_order.index(i) for i in range(len(stride_order))]
369 -
370 - # Use fp8 dtype for cache when requested.
371 - cache_dtype = dtype
372 - if config.kv_cache_dtype == "fp8":
373 - from vllm.platforms import current_platform
374 -
375 - cache_dtype = current_platform.fp8_dtype()
376 -
377 - cache_list = []
378 - for _ in range(config.num_layers):
379 - # Allocate in physical layout order (contiguous in memory)
380 - cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
381 - # Permute to logical view
382 - cache = cache.permute(*inv_order)
383 - cache_list.append(cache)
384 -
385 - return cache_list
388
389
390 # ============================================================================
@@ -499,13 +501,10 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
501 backend_class, impl, layer = _create_backend_impl(
502 backend_cfg, config, device, dtype
503 )
502 -
504 # Set KV cache layout if the backend requires a specific one
504 - # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
505 - required_layout = backend_class.get_required_kv_cache_layout()
506 - if required_layout is not None:
507 - set_kv_cache_layout(required_layout)
508 - get_kv_cache_layout.cache_clear()
505 + # (e.g., FlashInfer requires LBHNC on SM100/Blackwell for TRTLLM attention)
506 + supported = get_supported_kv_cache_layouts([backend_class])
507 + layout = resolve_kv_cache_layout(vllm_config, [[m.name for m in supported]])
508
509 common_metadata = _build_common_attn_metadata(
510 q_lens, kv_lens, config.block_size, device
@@ -542,7 +541,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
541 )
542
543 cache_list = _create_kv_cache(
545 - config, max_num_blocks, backend_class, device, dtype
544 + config, max_num_blocks, backend_class, device, dtype, layout
545 )
546
547 timing_stats, mem_stats = _run_single_benchmark(
docs/assets/contributing/dockerfile-stages-dependency.png
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ
docs/features/mooncake_store_connector_usage.md
-1
@@ -246,7 +246,6 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants
246
247 - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
248 - `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.
249 -- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
249 - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
250 - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
251 - `save_decode_cache` (bool): Enable offloading decode tokens' KV cache. A `kv_consumer` does not save during prefill; when decode starts, it fills any missing block-aligned prompt prefix before appending completed decode blocks. Default: `false`.
docs/features/nixl_connector_compatibility.md
+4 -4
@@ -61,7 +61,7 @@ th:not(:first-child) {
61 See [Configuration Notes](#configuration-notes) below for what must match and
62 what may differ.
63
64 -<sup>2</sup> Requires `FLASH_ATTN` or `FLASHINFER` backend **and** `HND` KV cache layout. Enable via `--kv-transfer-config '{"kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'`.
64 +<sup>2</sup> Cross-layer contiguity is achieved by using a `BLHNC` layout (set via `VLLM_KV_CACHE_LAYOUT=BLHNC`).
65
66 <sup>3</sup> Supported only when HMA is **not** required (i.e., non-hybrid models). Block IDs are remapped automatically. Only P block size < D block size is supported.
67
@@ -100,9 +100,9 @@ By default, a **compatibility hash** is checked during handshake. P and D instan
100
101 ### KV cache layout
102
103 -- NixlConnector defaults to **`HND`** layout for optimal transfer performance (non-MLA models).
104 -- `NHD` layout is supported but does **not** allow heterogeneous TP head splitting.
105 -- Experimental `HND` ↔ `NHD` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA.
103 +- NixlConnector defaults to **`LBHNC`** (head-major, formerly `HND`) layout for optimal transfer performance (non-MLA models).
104 +- `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting.
105 +- Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA.
106
107 ### Quantized KV cache
108
docs/features/nixl_connector_usage.md
+1 -10
@@ -404,21 +404,12 @@ GB-series GPUs support multi-node NVLink. NIXL supports this capability, but KVC
404
405 ### Heterogeneous KV Layout support
406
407 -Support use case: Prefill with 'HND' and decode with 'NHD' with experimental configuration
407 +Support use case: Prefill with `LBHNC` and decode with `LBNHC` with experimental configuration
408
409 ```bash
410 --kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
411 ```
412
413 -### Cross layers blocks
414 -
415 -By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
416 -To enable this feature:
417 -
418 -```bash
419 ---kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
420 -```
421 -
413 ## Metrics Reference
414
415 vLLM periodically logs a `KV Transfer metrics` line summarising NIXL transfer
rust/Cargo.lock
+9 -9
@@ -2340,15 +2340,6 @@ version = "0.4.29"
2340 source = "registry+https://github.com/rust-lang/crates.io-index"
2341 checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
2342
2343 -[[package]]
2344 -name = "lz4_flex"
2345 -version = "0.14.0"
2346 -source = "registry+https://github.com/rust-lang/crates.io-index"
2347 -checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226"
2348 -dependencies = [
2349 - "twox-hash",
2350 -]
2351 -
2343 [[package]]
2344 name = "logos"
2345 version = "0.15.1"
@@ -2415,6 +2406,15 @@ dependencies = [
2406 "logos-codegen 0.16.1",
2407 ]
2408
2409 +[[package]]
2410 +name = "lz4_flex"
2411 +version = "0.14.0"
2412 +source = "registry+https://github.com/rust-lang/crates.io-index"
2413 +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226"
2414 +dependencies = [
2415 + "twox-hash",
2416 +]
2417 +
2418 [[package]]
2419 name = "macro_rules_attribute"
2420 version = "0.2.2"
tests/compile/passes/test_fusion_attn.py
+32 -24
@@ -7,7 +7,11 @@ import torch._dynamo
7
8 from tests.compile.backend import LazyInitPass, TestBackend
9 from tests.utils import TestFP8Layer, flat_product
10 -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
10 +from tests.v1.attention.utils import (
11 + BatchSpec,
12 + create_common_attn_metadata,
13 + dense_kv_cache_views,
14 +)
15 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
16 from vllm.compilation.passes.fusion.attn_quant_fusion import (
17 ATTN_OP,
@@ -39,7 +43,14 @@ from vllm.platforms import current_platform
43 from vllm.utils.flashinfer import has_flashinfer
44 from vllm.v1.attention.backend import AttentionMetadata
45 from vllm.v1.attention.backends.registry import AttentionBackendEnum
42 -from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
46 +from vllm.v1.attention.backends.utils import (
47 + get_supported_kv_cache_layouts,
48 + resolve_kv_cache_layout,
49 +)
50 +from vllm.v1.kv_cache_interface import (
51 + AttentionSpec,
52 + get_kv_quant_mode,
53 +)
54
55 DEVICE_TYPE = current_platform.device_type
56 FP8_DTYPE = current_platform.fp8_dtype()
@@ -108,32 +119,29 @@ class AttentionQuantPatternModel(torch.nn.Module):
119 max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
120 num_blocks = batch_size * max_blocks
121
111 - # Fetch the attention backend and kv cache shape and stride order
112 - attn_backend = self.attn.attn_backend
113 - kv_cache_shape = attn_backend.get_kv_cache_shape(
114 - num_blocks,
115 - self.block_size,
116 - self.num_kv_heads,
117 - self.head_size,
118 - cache_dtype_str=self.attn.kv_cache_dtype,
122 + spec = AttentionSpec(
123 + block_size=self.block_size,
124 + num_kv_heads=self.num_kv_heads,
125 + head_size=self.head_size,
126 + dtype=self.attn.kv_cache_torch_dtype,
127 + kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
128 + )
129 + supported = get_supported_kv_cache_layouts([self.attn.attn_backend])
130 + layout = resolve_kv_cache_layout(
131 + self.vllm_config, [[m.name for m in supported]]
132 )
120 - try:
121 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
122 - except (AttributeError, NotImplementedError):
123 - kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
124 -
125 - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
126 - inv_order = [
127 - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
128 - ]
129 -
130 - # Create dummy KV cache
133 raw_tensor = torch.zeros(
132 - kv_cache_shape,
133 - dtype=self.attn.kv_cache_torch_dtype,
134 + num_blocks * spec.page_size_bytes,
135 + dtype=torch.int8,
136 device=self.device,
137 )
136 - kv_cache = raw_tensor.permute(*inv_order)
138 + kv_cache = dense_kv_cache_views(
139 + raw_tensor,
140 + spec,
141 + num_blocks,
142 + num_layers=1,
143 + layout=layout,
144 + )[0]
145
146 self.attn.kv_cache = kv_cache
147
tests/compile/passes/test_mla_attn_quant_fusion.py
+6 -19
@@ -150,27 +150,14 @@ class MLAAttentionQuantPatternModel(torch.nn.Module):
150 max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
151 num_blocks = batch_size * max_blocks
152
153 - # MLA KV cache is 3D: (num_blocks, block_size, head_size)
154 - attn_backend = self.mla_attn.attn_backend
155 - kv_cache_shape = attn_backend.get_kv_cache_shape(
156 - num_blocks, self.block_size, 1, self.head_size
157 - )
158 - try:
159 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
160 - except (AttributeError, NotImplementedError):
161 - kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
162 -
163 - ordered_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
164 - inv_order = [
165 - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
166 - ]
167 -
168 - raw_tensor = torch.zeros(
169 - ordered_shape, dtype=self.kv_cache_dtype, device=self.device
153 + # MLA KV cache is 4D: (num_blocks, num_heads=1, block_size, head_size)
154 + kv_cache = torch.zeros(
155 + (num_blocks, 1, self.block_size, self.head_size),
156 + dtype=self.kv_cache_dtype,
157 + device=self.device,
158 )
171 - kv_cache = raw_tensor.permute(*inv_order)
159
173 - self.mla_attn.kv_cache = kv_cache
160 + self.mla_attn.bind_kv_cache(kv_cache)
161
162 self.attn_metadata = self.builder.build(
163 common_prefix_len=0, common_attn_metadata=common_attn_metadata
tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py
+4 -19
@@ -165,29 +165,14 @@ class MLARoPEKVCacheCatTestModel(torch.nn.Module):
165 max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
166 num_blocks = batch_size * max_blocks
167
168 - # Fetch the attention backend and kv cache shape and stride order
169 - kv_cache_shape = self.attn_backend.get_kv_cache_shape(
170 - num_blocks, self.block_size, self.num_kv_heads, self.head_size
171 - )
172 - try:
173 - kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
174 - except (AttributeError, NotImplementedError):
175 - kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
176 -
177 - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
178 - inv_order = [
179 - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
180 - ]
181 -
182 - raw_tensor = torch.zeros(
183 - num_blocks * self.block_size * self.num_kv_heads * self.head_size,
168 + # MLA KV cache is 4D: (num_blocks, num_heads=1, block_size, head_size)
169 + kv_cache = torch.zeros(
170 + (num_blocks, 1, self.block_size, self.head_size),
171 dtype=self.kv_cache_dtype,
172 device=self.device,
173 )
187 - raw_tensor = raw_tensor.view(kv_cache_shape)
188 - kv_cache = raw_tensor.permute(*inv_order)
174
190 - self.mla_attn.kv_cache = kv_cache
175 + self.mla_attn.bind_kv_cache(kv_cache)
176
177 # Build attn metadata
178 attn_metadata = self.builder.build(
tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py
+34 -33
@@ -10,7 +10,11 @@ from packaging.version import Version
10
11 import vllm.config
12 from tests.compile.backend import TestBackend
13 -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
13 +from tests.v1.attention.utils import (
14 + BatchSpec,
15 + create_common_attn_metadata,
16 + dense_kv_cache_views,
17 +)
18 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
19 from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
20 from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import (
@@ -40,7 +44,7 @@ from vllm.v1.attention.backend import (
44 CommonAttentionMetadata,
45 )
46 from vllm.v1.attention.backends.registry import AttentionBackendEnum
43 -from vllm.v1.kv_cache_interface import AttentionSpec
47 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
48
49 pytestmark = pytest.mark.skip_global_cleanup
50
@@ -147,20 +151,25 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
151 self.attn._k_scale = self.attn._k_scale.to(device)
152 self.attn._v_scale = self.attn._v_scale.to(device)
153
150 - self.builder = self.attn.attn_backend.get_builder_cls()(
151 - kv_cache_spec=AttentionSpec(
154 + # Mirror the worker spec-collection loop: the backend publishes the
155 + # packing it needs.
156 + self.kv_cache_spec = self.attn.attn_backend.customize_spec(
157 + AttentionSpec(
158 block_size=self.block_size,
159 num_kv_heads=self.num_kv_heads,
160 head_size=head_size,
161 dtype=self.kv_cache_dtype,
156 - ),
162 + )
163 + )
164 + self.builder = self.attn.attn_backend.get_builder_cls()(
165 + kv_cache_spec=self.kv_cache_spec,
166 layer_names=[self.attn.layer_name],
167 vllm_config=vllm_config,
168 device=device,
169 )
170
171 def build_attn_metadata(
163 - self, batch_size: int, kv_stride_order: tuple[int, ...] | None = None
172 + self, batch_size: int, layout: KVCacheLayout
173 ) -> CommonAttentionMetadata:
174 batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
175 common_attn_metadata = create_common_attn_metadata(
@@ -170,27 +179,18 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
179 max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
180 num_blocks = batch_size * max_blocks
181
173 - attn_backend = self.attn.attn_backend
174 - kv_cache_shape = attn_backend.get_kv_cache_shape(
175 - num_blocks, self.block_size, self.num_kv_heads, self.head_size
176 - )
177 - # Caller can force a physical layout; else use the backend's.
178 - if kv_stride_order is None:
179 - try:
180 - kv_stride_order = attn_backend.get_kv_cache_stride_order()
181 - except (AttributeError, NotImplementedError):
182 - kv_stride_order = tuple(range(len(kv_cache_shape)))
183 -
184 - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_stride_order)
185 - inv_order = [kv_stride_order.index(i) for i in range(len(kv_stride_order))]
186 -
182 raw_tensor = torch.zeros(
188 - 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
189 - dtype=self.kv_cache_dtype,
183 + num_blocks * self.kv_cache_spec.page_size_bytes,
184 + dtype=torch.int8,
185 device=self.device,
186 )
192 - raw_tensor = raw_tensor.view(kv_cache_shape)
193 - kv_cache = raw_tensor.permute(*inv_order)
187 + kv_cache = dense_kv_cache_views(
188 + raw_tensor,
189 + self.kv_cache_spec,
190 + num_blocks,
191 + num_layers=1,
192 + layout=layout,
193 + )[0]
194
195 # Store as a bare tensor (not wrapped in a list) to match production
196 # `bind_kv_cache` behavior. `get_attention_context` returns this
@@ -272,7 +272,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
272 block_size: int,
273 is_neox: bool,
274 use_shuffle_kv_layout: str,
275 - kv_stride_order: tuple[int, ...],
275 + kv_layout: KVCacheLayout,
276 dtype: torch.dtype,
277 kv_cache_dtype: str,
278 rms_norm_eps: float,
@@ -345,7 +345,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
345 # Run unfused (eager) forward
346 with set_forward_context(None, vllm_config):
347 forward_context = get_forward_context()
348 - attn_metadata = model.build_attn_metadata(num_tokens, kv_stride_order)
348 + attn_metadata = model.build_attn_metadata(num_tokens, kv_layout)
349 forward_context.slot_mapping = {
350 model.layer_name: attn_metadata.slot_mapping
351 }
@@ -360,7 +360,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
360 with set_forward_context(None, vllm_config):
361 model_fused = torch.compile(model, backend=backend)
362 forward_context = get_forward_context()
363 - attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_stride_order)
363 + attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_layout)
364 forward_context.slot_mapping = {
365 model.layer_name: attn_metadata.slot_mapping
366 }
@@ -437,10 +437,11 @@ _FUSION_CONFIGS = [
437 )
438 @pytest.mark.parametrize("num_tokens", [5, 2048])
439 @pytest.mark.parametrize(
440 - "kv_stride_order",
441 - # FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size,
442 - # 2*head_size); no separate K/V dim to permute.
443 - [pytest.param((0, 1, 2, 3), id="packed_4d")],
440 + "kv_layout",
441 + [
442 + pytest.param(KVCacheLayout.LBHNC, id="head_major"),
443 + pytest.param(KVCacheLayout.LBNHC, id="token_major"),
444 + ],
445 )
446 @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
447 @pytest.mark.parametrize("block_size", [16, 32, 64])
@@ -467,7 +468,7 @@ def test_qk_norm_rope_kvcache_fusion(
468 is_neox: bool,
469 attn_backend: AttentionBackendEnum,
470 enable_aiter_triton_rope: bool,
470 - kv_stride_order: tuple[int, ...],
471 + kv_layout: KVCacheLayout,
472 block_size: int,
473 dtype: torch.dtype,
474 kv_cache_dtype: str,
@@ -486,7 +487,7 @@ def test_qk_norm_rope_kvcache_fusion(
487 block_size=block_size,
488 is_neox=is_neox,
489 use_shuffle_kv_layout="0",
489 - kv_stride_order=kv_stride_order,
490 + kv_layout=kv_layout,
491 dtype=dtype,
492 kv_cache_dtype=kv_cache_dtype,
493 rms_norm_eps=rms_norm_eps,
tests/compile/passes/test_rope_kvcache_fusion.py
+36 -22
@@ -7,7 +7,11 @@ from torch._higher_order_ops import auto_functionalized
7
8 import vllm.config
9 from tests.compile.backend import TestBackend
10 -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
10 +from tests.v1.attention.utils import (
11 + BatchSpec,
12 + create_common_attn_metadata,
13 + dense_kv_cache_views,
14 +)
15 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
16 from vllm.compilation.passes.fusion import rope_kvcache_fusion
17 from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
@@ -38,6 +42,10 @@ from vllm.v1.attention.backend import (
42 CommonAttentionMetadata,
43 )
44 from vllm.v1.attention.backends.registry import AttentionBackendEnum
45 +from vllm.v1.kv_cache_interface import (
46 + FullAttentionSpec,
47 + KVCacheLayout,
48 +)
49
50 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
51 VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
@@ -149,9 +157,20 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
157 FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
158 )
159
160 + # Mirror the worker spec-collection loop: the backend publishes its
161 + # packing (e.g. ROCM_ATTN's separate K/V head groups).
162 + self.kv_cache_spec = self.attn_backend.customize_spec(
163 + FullAttentionSpec(
164 + block_size=self.block_size,
165 + num_kv_heads=self.num_kv_heads,
166 + head_size=self.head_size,
167 + dtype=self.kv_cache_dtype,
168 + )
169 + )
170 +
171 # Initialize attn MetadataBuilder
172 self.builder = self.attn_backend.get_builder_cls()(
154 - kv_cache_spec=self.attn.get_kv_cache_spec(vllm_config),
173 + kv_cache_spec=self.kv_cache_spec,
174 layer_names=[self.attn.layer_name],
175 vllm_config=vllm_config,
176 device=device,
@@ -168,30 +187,25 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
187 max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
188 num_blocks = batch_size * max_blocks
189
171 - # Fetch the attention backend and kv cache shape and stride order
172 - kv_cache_shape = self.attn_backend.get_kv_cache_shape(
173 - num_blocks, self.block_size, self.num_kv_heads, self.head_size
174 - )
175 - try:
176 - kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
177 - except (AttributeError, NotImplementedError):
178 - kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
179 -
180 - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
181 - inv_order = [
182 - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
183 - ]
190 + # A backend's preferred supported layout wins over the default, mirroring
191 + # selector-time resolution (e.g. ROCM_ATTN's head groups force LHBNC).
192 + layout = KVCacheLayout.LBNHC
193 + supported = self.attn_backend.supported_kv_cache_layouts()
194 + if supported:
195 + layout = supported[0]
196
185 - # Create dummy KV cache
197 raw_tensor = torch.zeros(
187 - 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
188 - dtype=self.kv_cache_dtype,
198 + num_blocks * self.kv_cache_spec.page_size_bytes,
199 + dtype=torch.int8,
200 device=self.device,
201 )
191 - raw_tensor = raw_tensor.view(kv_cache_shape)
192 - kv_cache = raw_tensor.permute(*inv_order)
193 -
194 - self.attn.kv_cache = kv_cache
202 + self.attn.kv_cache = dense_kv_cache_views(
203 + raw_tensor,
204 + self.kv_cache_spec,
205 + num_blocks,
206 + num_layers=1,
207 + layout=layout,
208 + )[0]
209
210 # Build attn metadata
211 attn_metadata = self.builder.build(
tests/distributed/test_kvlayout.py
+4 -4
@@ -21,7 +21,7 @@ def test_get_kv_connector_cache_layout_without_kv_connector():
21 with set_current_vllm_config(vllm_config):
22 # Test with default settings
23 layout = get_kv_connector_cache_layout()
24 - assert layout == "NHD"
24 + assert layout is None
25
26
27 def test_get_kv_connector_cache_layout_with_lmcache_connector():
@@ -35,7 +35,7 @@ def test_get_kv_connector_cache_layout_with_lmcache_connector():
35 with set_current_vllm_config(vllm_config):
36 # Test with default settings
37 layout = get_kv_connector_cache_layout()
38 - assert layout == "NHD"
38 + assert layout is None
39
40
41 def test_get_kv_connector_cache_layout_with_nixl_connector():
@@ -52,7 +52,7 @@ def test_get_kv_connector_cache_layout_with_nixl_connector():
52 with set_current_vllm_config(vllm_config):
53 # Test with default settings
54 layout = get_kv_connector_cache_layout()
55 - assert layout == "HND"
55 + assert layout == "LBHNC"
56
57
58 def test_get_kv_connector_cache_layout_with_multi_connector():
@@ -75,4 +75,4 @@ def test_get_kv_connector_cache_layout_with_multi_connector():
75 with set_current_vllm_config(vllm_config):
76 # Test with default settings
77 layout = get_kv_connector_cache_layout()
78 - assert layout == "HND"
78 + assert layout == "LBHNC"
tests/kernels/attention/test_cache.py
+7 -7
@@ -356,21 +356,21 @@ def test_reshape_and_cache_flash(
356 dequant_nvfp4_kv_cache,
357 )
358
359 - def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
360 - # data_cache: [N, T, H, data_dim] NHD (contiguous inner dims)
361 - # scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
362 - # Permute to HND layout for the dequant utility.
359 + def dequant_nvfp4_cache(data_cache, scale_cache, global_scale):
360 + # data_cache: [B, N, H, data_dim] logical view (layout in strides)
361 + # scale_cache: [B, N, H, scale_dim] logical view (layout in strides)
362 + # Permute to [B, H, N, dim] for the dequant utility.
363 data_hnd = data_cache.permute(0, 2, 1, 3)
364 scale_hnd = scale_cache.permute(0, 2, 1, 3)
365 result_hnd = dequant_nvfp4_kv_cache(
366 data_hnd, scale_hnd, global_scale, head_size, block_size
367 )
368 - return result_hnd.permute(0, 2, 1, 3) # back to [N, T, H, D]
368 + return result_hnd.permute(0, 2, 1, 3) # back to [B, N, H, dim]
369
370 - result_key_cache = dequant_nvfp4_cache_nhd(
370 + result_key_cache = dequant_nvfp4_cache(
371 nvfp4_key_data, key_scale_cache, k_scale.item()
372 )
373 - result_value_cache = dequant_nvfp4_cache_nhd(
373 + result_value_cache = dequant_nvfp4_cache(
374 nvfp4_value_data, value_scale_cache, v_scale.item()
375 )
376
tests/kernels/attention/test_minimax_m3.py
+111 -220
@@ -5,10 +5,10 @@
5 import pytest
6 import torch
7
8 +from tests.v1.attention.utils import dense_kv_cache_views
9 from vllm import _custom_ops as ops
9 -from vllm.models.minimax_m3.common.indexer import (
10 - MiniMaxM3IndexerBackend,
11 -)
10 +from vllm.config import CacheConfig
11 +from vllm.config.cache import _layout_from_name
12 from vllm.models.minimax_m3.common.ops.index_topk import (
13 minimax_m3_index_decode,
14 minimax_m3_index_score,
@@ -19,15 +19,15 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import (
19 minimax_m3_sparse_attn,
20 minimax_m3_sparse_attn_decode,
21 )
22 -from vllm.models.minimax_m3.common.sparse_attention import (
23 - MiniMaxM3SparseBackend,
24 - MiniMaxM3SparseTritonImpl,
25 -)
22 +from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl
23 from vllm.platforms import current_platform
27 -from vllm.v1.attention.backends.utils import set_kv_cache_layout
28 -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
29 -from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
30 -from vllm.v1.worker.utils import AttentionGroup
24 +from vllm.v1.attention.backends.utils import record_kv_cache_layout
25 +from vllm.v1.kv_cache_interface import (
26 + FullAttentionSpec,
27 + KVCacheLayout,
28 + MLAAttentionSpec,
29 + compute_layer_kv_cache_shape_bytes,
30 +)
31
32 if not (current_platform.is_cuda() or current_platform.is_rocm()):
33 pytest.skip(
@@ -37,35 +37,46 @@ if not (current_platform.is_cuda() or current_platform.is_rocm()):
37
38
39 @pytest.fixture
40 -def kv_layout(request):
41 - """Set the global KV cache layout for one test and restore it after."""
42 - set_kv_cache_layout(request.param)
43 - try:
44 - yield request.param
45 - finally:
46 - set_kv_cache_layout(None)
47 -
48 -
49 -def _stride_order_for(backend: type[MiniMaxM3SparseBackend], ndim: int) -> tuple:
50 - """Mirror the allocator's stride-order resolution (identity fallback)."""
51 - try:
52 - stride_order = backend.get_kv_cache_stride_order()
53 - assert len(stride_order) == ndim
54 - except (AttributeError, NotImplementedError):
55 - stride_order = tuple(range(ndim))
40 +def kv_layout(request) -> KVCacheLayout:
41 + """Resolve the requested layout name (legacy NHD/HND aliases included)."""
42 + return _layout_from_name(request.param)
43 +
44 +
45 +def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]:
46 + """Per-layer physical stride order for the given layout; the 3-dim
47 + indexer side cache (H=1) is contiguous, so identity."""
48 + if ndim == 3:
49 + return (0, 1, 2)
50 + stride_order = layout.layer_view_order
51 + assert len(stride_order) == ndim
52 return stride_order
53
54
55 +def _main_spec() -> FullAttentionSpec:
56 + return FullAttentionSpec(
57 + block_size=BLOCK_SIZE,
58 + num_kv_heads=NUM_KV_HEADS,
59 + head_size=HEAD_DIM,
60 + head_size_v=HEAD_DIM,
61 + dtype=DTYPE,
62 + )
63 +
64 +
65 +def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]:
66 + """Standardized per-layer logical shape [B, H, N, C] for the main cache,
67 + derived the same way the production allocator does."""
68 + shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages)
69 + return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize)
70 +
71 +
72 def _allocate_main_kv_via_contract(
60 - num_pages: int, device: torch.device | str = "cuda"
73 + num_pages: int, layout: KVCacheLayout, device: torch.device | str = "cuda"
74 ) -> torch.Tensor:
75 """Build the main KV cache exactly as the production allocator does for the
63 - currently active layout: allocate the physical (permuted) tensor, then
64 - expose the inverse-permuted logical-NHD view the backend sees."""
65 - logical_shape = MiniMaxM3SparseBackend.get_kv_cache_shape(
66 - num_pages, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
67 - )
68 - stride_order = _stride_order_for(MiniMaxM3SparseBackend, len(logical_shape))
76 + given layout: allocate the physical (permuted) tensor, then expose the
77 + inverse-permuted logical [B, H, N, C] view the backend sees."""
78 + logical_shape = _main_kv_logical_shape(num_pages)
79 + stride_order = _layer_stride_order(layout, len(logical_shape))
80 physical_shape = tuple(logical_shape[i] for i in stride_order)
81 inv_order = [stride_order.index(i) for i in range(len(stride_order))]
82 raw = torch.randn(physical_shape, device=device, dtype=DTYPE)
@@ -817,7 +828,7 @@ def _reference_sparse_attn(
828 ],
829 )
830 def test_prefill_sparse_attention_correctness(
820 - kv_layout: str,
831 + kv_layout: KVCacheLayout,
832 q_lens: tuple[int, ...],
833 kv_lens: tuple[int, ...],
834 ):
@@ -851,7 +862,7 @@ def test_prefill_sparse_attention_correctness(
862 # Allocate the main KV cache through the backend layout contract so the
863 # physical storage matches the active layout (contiguous NHD or strided
864 # HND), while the kernels and reference see the logical-NHD view.
854 - kv_cache = _allocate_main_kv_via_contract(num_pages)
865 + kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
866
867 # Build sparse block indices with the same contract as the real M3 indexer:
868 # one forced local block, then score-selected older causal blocks.
@@ -904,38 +915,31 @@ def test_prefill_sparse_attention_correctness(
915 assert error.max().item() < 1.7e-2
916
917
907 -def test_main_backend_layout_contract():
908 - """The main sparse backend exposes the logical-NHD shape and the
909 - flash_attn-style stride order for each layout."""
918 +def test_main_cache_layout_contract():
919 + """The standardized per-layer logical shape is [B, H, N, C] with packed
920 + K/V content, and the legacy layout aliases resolve to the expected
921 + per-layer stride orders."""
922 nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
911 - logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
923 + logical = _main_kv_logical_shape(nb)
924 assert logical == (nb, h, bs, 2 * d)
925 # The old separate K/V-axis shape is no longer the logical shape.
926 assert logical != (nb, 2, bs, h, d)
927
916 - try:
917 - set_kv_cache_layout("HND")
918 - assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3)
919 - set_kv_cache_layout("NHD")
920 - assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3)
921 - finally:
922 - set_kv_cache_layout(None)
923 -
924 - for layout in ("NHD", "HND"):
925 - try:
926 - set_kv_cache_layout(layout)
927 - order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
928 - finally:
929 - set_kv_cache_layout(None)
928 + assert KVCacheLayout.LBHNC.layer_view_order == (0, 1, 2, 3)
929 + assert KVCacheLayout.LBNHC.layer_view_order == (0, 2, 1, 3)
930 +
931 + cache_config = CacheConfig()
932 + record_kv_cache_layout(cache_config, "HND")
933 + assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBHNC
934 + cache_config = CacheConfig()
935 + record_kv_cache_layout(cache_config, "NHD")
936 + assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBNHC
937 +
938 + for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
939 + order = layout.layer_view_order
940 # Valid permutation: no duplicates, covers every axis.
941 assert set(order) == set(range(len(order)))
942
933 - # M3 has no cross-layer KV blocks.
934 - with pytest.raises(NotImplementedError):
935 - MiniMaxM3SparseBackend.get_kv_cache_stride_order(
936 - include_num_layers_dimension=True
937 - )
938 -
943
944 def test_aiter_sparse_pa_layout_contract(monkeypatch):
945 """The shuffle-only AITER path retains separately contiguous K/V storage."""
@@ -948,75 +952,53 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch):
952 lambda: True,
953 )
954
951 - nb, bs, h, d = 7, BLOCK_SIZE, 1, HEAD_DIM
952 - logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
953 - order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
954 - assert logical == (nb, 2, bs, h, d)
955 - assert order == (1, 0, 2, 3, 4)
955 + backend = sparse_attn_mod.MiniMaxM3SparseBackend
956 + assert KVCacheLayout.LHBNC in backend.supported_kv_cache_layouts()
957
957 - physical_shape = tuple(logical[i] for i in order)
958 - inv_order = [order.index(i) for i in range(len(order))]
959 - raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
960 - logical_view = raw.permute(*inv_order)
961 - key_cache, value_cache = logical_view.unbind(1)
958 + nb, h = 7, 1
959 + spec = FullAttentionSpec(
960 + block_size=BLOCK_SIZE,
961 + num_kv_heads=h,
962 + head_size=HEAD_DIM,
963 + dtype=DTYPE,
964 + num_head_slots=2,
965 + state_content_bytes=h * HEAD_DIM * DTYPE.itemsize,
966 + )
967 + raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
968 + view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LHBNC)[0]
969 + key_cache, value_cache = view.unbind(1)
970 assert key_cache.is_contiguous()
971 assert value_cache.is_contiguous()
972
973
966 -def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
967 - """Do not pair AITER's separated cache layout with the Triton fallback."""
968 - import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
969 -
970 - monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
971 - monkeypatch.setattr(
972 - sparse_attn_mod.rocm_aiter_ops,
973 - "is_shuffle_kv_cache_enabled",
974 - lambda: True,
974 +def test_indexer_cache_squeezes_to_contiguous_3d():
975 + """The indexer side cache is standardized 4D with H=1: under both layouts
976 + the allocator's logical view stays contiguous and squeezes (as
977 + `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim
978 + [num_blocks, block_size, head_dim] cache the kernels consume."""
979 + nb = 5
980 + ispec = MLAAttentionSpec(
981 + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
982 )
983 + shape_bytes = compute_layer_kv_cache_shape_bytes(ispec, nb)
984 + assert shape_bytes == (nb, 1, BLOCK_SIZE, HEAD_DIM * DTYPE.itemsize)
985 + assert _layer_stride_order(KVCacheLayout.LBNHC, 3) == (0, 1, 2)
986
977 - with pytest.raises(ValueError, match="num_kv_heads == 1"):
978 - MiniMaxM3SparseBackend.get_kv_cache_shape(7, BLOCK_SIZE, 2, HEAD_DIM)
979 -
980 -
981 -def test_main_backend_unknown_layout_raises(monkeypatch):
982 - """An unrecognized layout (injected past env-var validation) is rejected."""
983 - import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
984 -
985 - monkeypatch.setattr(sparse_attn_mod, "get_kv_cache_layout", lambda: "BOGUS")
986 - with pytest.raises(ValueError, match="Unknown cache layout format"):
987 - MiniMaxM3SparseBackend.get_kv_cache_stride_order()
988 -
989 -
990 -def test_indexer_backend_stride_order_is_identity():
991 - """The 3-dim indexer cache must not inherit the parent's 4-element stride
992 - order; it overrides to the 3-element identity so the allocator keeps the
993 - contiguous layout."""
994 - assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2)
995 -
996 - # Cross-layer (per-layer-stacked) KV blocks are not supported.
997 - with pytest.raises(NotImplementedError):
998 - MiniMaxM3IndexerBackend.get_kv_cache_stride_order(
999 - include_num_layers_dimension=True
1000 - )
1001 -
1002 - # The stride order matches the 3-dim indexer shape rank.
1003 - indexer_shape = MiniMaxM3IndexerBackend.get_kv_cache_shape(
1004 - 5, BLOCK_SIZE, 1, HEAD_DIM
1005 - )
1006 - assert len(indexer_shape) == 3
1007 - assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
987 + for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
988 + iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
989 + view = dense_kv_cache_views(iraw, ispec, nb, 1, layout, BLOCK_SIZE)[0]
990 + assert tuple(view.shape) == (nb, 1, BLOCK_SIZE, HEAD_DIM)
991 + indexer_cache = view.squeeze(1)
992 + assert tuple(indexer_cache.shape) == (nb, BLOCK_SIZE, HEAD_DIM)
993 + assert indexer_cache.is_contiguous()
994
995
996 def test_hnd_allocation_is_packed_head_major():
997 """Under HND the backend-visible logical view is the packed head-major
998 physical allocation."""
999 nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
1014 - logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
1015 - try:
1016 - set_kv_cache_layout("HND")
1017 - stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
1018 - finally:
1019 - set_kv_cache_layout(None)
1000 + logical = _main_kv_logical_shape(nb)
1001 + stride_order = KVCacheLayout.LBHNC.layer_view_order
1002
1003 physical_shape = tuple(logical[i] for i in stride_order)
1004 assert physical_shape == (nb, h, bs, 2 * d)
@@ -1042,27 +1024,13 @@ def test_main_cache_is_block_first_and_unpadded():
1024 """The allocator's contiguous-view branch (not the padded-strided branch)
1025 is used for the main GQA cache: its spec is unpadded and the physical
1026 layout keeps num_blocks as the first dimension under both layouts."""
1045 - from vllm.v1.kv_cache_interface import FullAttentionSpec
1046 -
1047 - spec = FullAttentionSpec(
1048 - block_size=BLOCK_SIZE,
1049 - num_kv_heads=NUM_KV_HEADS,
1050 - head_size=HEAD_DIM,
1051 - head_size_v=HEAD_DIM,
1052 - dtype=DTYPE,
1053 - )
1027 + spec = _main_spec()
1028 # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
1029 assert spec.page_size_padded is None
1030
1057 - logical = MiniMaxM3SparseBackend.get_kv_cache_shape(
1058 - 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
1059 - )
1060 - for layout in ("NHD", "HND"):
1061 - try:
1062 - set_kv_cache_layout(layout)
1063 - order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
1064 - finally:
1065 - set_kv_cache_layout(None)
1031 + logical = _main_kv_logical_shape(4)
1032 + for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
1033 + order = layout.layer_view_order
1034 inv_order = [order.index(i) for i in range(len(order))]
1035 # Physical first dim is num_blocks (block-first); required by the
1036 # padded-strided branch's block-first assumption if it were ever taken.
@@ -1275,7 +1243,7 @@ def test_aiter_decode_sparse_block_table_supports_spec_decode():
1243 @pytest.mark.parametrize("decode_query_len", [1, 4])
1244 @pytest.mark.parametrize("num_padded_reqs", [0, 2])
1245 def test_decode_sparse_attention_correctness(
1278 - kv_layout: str,
1246 + kv_layout: KVCacheLayout,
1247 seq_lens_list: tuple[int, ...],
1248 decode_query_len: int,
1249 num_padded_reqs: int,
@@ -1287,7 +1255,7 @@ def test_decode_sparse_attention_correctness(
1255 q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(
1256 seq_lens_list, decode_query_len, num_padded_reqs
1257 )
1290 - kv_cache = _allocate_main_kv_via_contract(num_pages)
1258 + kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
1259
1260 actual = torch.empty_like(q)
1261 minimax_m3_sparse_attn_decode(
@@ -1364,93 +1332,20 @@ def test_decode_wrong_layout_breaks_parity():
1332 assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2
1333
1334
1367 -def _make_attn_group(backend, spec):
1368 - return AttentionGroup(
1369 - backend=backend,
1370 - layer_names=["main"],
1371 - kv_cache_spec=spec,
1372 - kv_cache_group_id=0,
1373 - )
1374 -
1375 -
1335 def test_main_cache_byte_identical_through_production_allocator():
1377 - """AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main
1336 + """AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main
1337 `FullAttentionSpec` under HND and assert the backend-visible view has the
1379 - same shape, stride, and storage offset as the packed-HND allocation; the
1380 - indexer `MLAAttentionSpec` allocates through the same path to its 3-dim
1381 - shape."""
1338 + same shape, stride, and storage offset as the packed-HND allocation."""
1339 nb = 4
1383 - spec = FullAttentionSpec(
1384 - block_size=BLOCK_SIZE,
1385 - num_kv_heads=NUM_KV_HEADS,
1386 - head_size=HEAD_DIM,
1387 - head_size_v=HEAD_DIM,
1388 - dtype=DTYPE,
1389 - )
1340 + spec = _main_spec()
1341 raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
1391 - group = _make_attn_group(MiniMaxM3SparseBackend, spec)
1392 - try:
1393 - set_kv_cache_layout("HND")
1394 - kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {})
1395 - finally:
1396 - set_kv_cache_layout(None)
1397 - view = kv_caches["main"]
1342 + view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LBHNC, BLOCK_SIZE)[0]
1343
1344 oracle = raw.view(DTYPE).view((nb, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM))
1345 assert tuple(view.shape) == tuple(oracle.shape)
1346 assert view.stride() == oracle.stride()
1347 assert view.storage_offset() == oracle.storage_offset()
1348
1404 - # Indexer cache allocates through the same path under both layouts.
1405 - ispec = MLAAttentionSpec(
1406 - block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
1407 - )
1408 - for layout in ("NHD", "HND"):
1409 - iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
1410 - igroup = AttentionGroup(
1411 - backend=MiniMaxM3IndexerBackend,
1412 - layer_names=["idx"],
1413 - kv_cache_spec=ispec,
1414 - kv_cache_group_id=0,
1415 - )
1416 - try:
1417 - set_kv_cache_layout(layout)
1418 - iout = _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
1419 - finally:
1420 - set_kv_cache_layout(None)
1421 - assert tuple(iout["idx"].shape) == (nb, BLOCK_SIZE, HEAD_DIM)
1422 -
1423 -
1424 -def test_indexer_inherited_stride_order_trips_allocator_assert():
1425 - """AC-4 negative: without the indexer override, the inherited 4-element
1426 - stride order trips the allocator's `len(stride_order) == len(shape)` assert
1427 - for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the
1428 - allocator's `(AttributeError, NotImplementedError)` fallback."""
1429 -
1430 - class _BrokenIndexerBackend(MiniMaxM3IndexerBackend):
1431 - # Simulate inheriting the parent's 4-element stride order.
1432 - get_kv_cache_stride_order = staticmethod(
1433 - MiniMaxM3SparseBackend.get_kv_cache_stride_order
1434 - )
1435 -
1436 - nb = 4
1437 - ispec = MLAAttentionSpec(
1438 - block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
1439 - )
1440 - iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
1441 - igroup = AttentionGroup(
1442 - backend=_BrokenIndexerBackend,
1443 - layer_names=["idx"],
1444 - kv_cache_spec=ispec,
1445 - kv_cache_group_id=0,
1446 - )
1447 - try:
1448 - set_kv_cache_layout("HND")
1449 - with pytest.raises(AssertionError):
1450 - _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
1451 - finally:
1452 - set_kv_cache_layout(None)
1453 -
1349
1350 def test_padded_main_cache_is_flagged():
1351 """AC-2.1 negative: the M3 main cache relies on the allocator's
@@ -1465,11 +1360,7 @@ def test_padded_main_cache_is_flagged():
1360 )
1361 assert inv_order[0] == 0, "main GQA cache must remain block-first"
1362
1468 - try:
1469 - set_kv_cache_layout("HND")
1470 - stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
1471 - finally:
1472 - set_kv_cache_layout(None)
1363 + stride_order = KVCacheLayout.LBHNC.layer_view_order
1364
1365 good = FullAttentionSpec(
1366 block_size=BLOCK_SIZE,
@@ -1493,13 +1384,13 @@ def test_padded_main_cache_is_flagged():
1384
1385
1386 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
1496 -def test_reshape_and_cache_flash_write_persists(kv_layout: str):
1387 +def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout):
1388 """AC-5 write path: the `reshape_and_cache_flash` write site now consumes
1389 packed-content K/V split views. Writing through those views must persist
1390 into the bound storage under both layouts."""
1391 torch.manual_seed(0)
1392 num_pages = 4
1502 - kv_cache = _allocate_main_kv_via_contract(num_pages)
1393 + kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
1394 with torch.no_grad():
1395 kv_cache.zero_()
1396
tests/kernels/attention/test_rocm_aiter_fa.py
+13 -8
@@ -310,18 +310,23 @@ def test_aiter_mha_backend_contract():
310
311
312 def test_aiter_mha_backend_validates_kv_cache_block_size():
313 - """The backend should reject KV cache shapes that cannot be gathered
313 + """The backend should reject KV cache specs that cannot be gathered
314 correctly by the ROCm kernel."""
315 from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
316 + from vllm.v1.kv_cache_interface import FullAttentionSpec
317
317 - assert AiterFlashAttentionBackend.get_kv_cache_shape(8, 16, 8, 128) == (
318 - 8,
319 - 8,
320 - 16,
321 - 256,
322 - )
318 with pytest.raises(ValueError, match="Block size must be a multiple of 16"):
324 - AiterFlashAttentionBackend.get_kv_cache_shape(8, 15, 8, 128)
319 + AiterFlashAttentionBackend.customize_spec(
320 + FullAttentionSpec(
321 + block_size=15, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
322 + )
323 + )
324 + # block_size == 1 is the per-token page-size probe and must pass through.
325 + AiterFlashAttentionBackend.customize_spec(
326 + FullAttentionSpec(
327 + block_size=1, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
328 + )
329 + )
330
331
332 def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe():
tests/kernels/attention/test_trtllm_kvfp8_dequant.py
+41 -58
@@ -11,6 +11,7 @@ import pytest
11 import torch
12
13 from vllm.platforms import current_platform
14 +from vllm.v1.kv_cache_interface import KVCacheLayout
15
16 if current_platform.is_rocm():
17 pytest.skip(
@@ -34,51 +35,40 @@ def to_float8(x, dtype=None):
35 return x_scl_sat.to(dtype), scale.float().reciprocal()
36
37
37 -def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
38 - """Create a standard contiguous fp8 KV cache (HND layout)."""
39 - raw = torch.randn(
40 - num_blocks,
41 - 2,
42 - num_kv_heads,
43 - block_size,
44 - head_size,
45 - dtype=torch.bfloat16,
46 - device="cuda",
47 - )
48 - kv_cache, scale = to_float8(raw)
49 - return kv_cache, scale
50 -
51 -
52 -def make_cross_layer_kv_cache(
38 +def make_random_kv_cache(
39 num_blocks,
40 num_kv_heads,
41 block_size,
42 head_size,
57 - num_layers=4,
43 + layout=KVCacheLayout.LBHNC,
44 + num_layers=None,
45 ):
46 """
60 - Create a non-contiguous per-layer view mimicking cross-layer allocation.
47 + Create a random fp8 KV cache mimicking the production allocation.
48
62 - Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
49 + Physical layout: [L, B, H, N, 2 * hs] permuted by ``layout.stride_order``
50 Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size)
64 - with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
51 + with non-contiguous strides under cross-layer layouts (they skip over
52 + num_layers), matching the actual forward path.
53 """
66 - raw = torch.randn(
54 + logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size)
55 + if num_layers is None:
56 + num_layers = 1 if layout.is_layer_compact else 2
57 + logical_5d = (num_layers, *logical_4d)
58 + physical_5d = tuple(logical_5d[i] for i in layout.stride_order)
59 + inv_order = [layout.stride_order.index(i) for i in range(5)]
60 +
61 + raw_phys = torch.randn(*physical_5d, dtype=torch.bfloat16, device="cuda")
62 + fp8_phys, scale = to_float8(raw_phys)
63 + fp8_4d = fp8_phys.permute(*inv_order)[0]
64 + kv_5d = fp8_4d.view(
65 num_blocks,
68 - 2,
66 num_kv_heads,
70 - num_layers,
67 block_size,
68 + 2,
69 head_size,
73 - dtype=torch.bfloat16,
74 - device="cuda",
75 - )
76 - fp8_full, scale = to_float8(raw)
77 - layer_view = fp8_full[:, :, :, 0, :, :]
78 - assert not layer_view.is_contiguous(), (
79 - f"Expected non-contiguous view, got strides {layer_view.stride()}"
80 - )
81 - return layer_view, scale
70 + ).permute(0, 3, 1, 2, 4)
71 + return kv_5d, scale
72
73
74 def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
@@ -114,7 +104,7 @@ def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
104 @pytest.mark.parametrize("block_size", [16, 32])
105 @pytest.mark.parametrize("batch_size", [1, 4])
106 @pytest.mark.parametrize("num_pages_per_seq", [3, 8])
117 -@pytest.mark.parametrize("contiguous", [True, False])
107 +@pytest.mark.parametrize("layout", list(KVCacheLayout))
108 @torch.inference_mode()
109 def test_trtllm_kvfp8_dequant(
110 num_kv_heads: int,
@@ -122,7 +112,7 @@ def test_trtllm_kvfp8_dequant(
112 block_size: int,
113 batch_size: int,
114 num_pages_per_seq: int,
125 - contiguous: bool,
115 + layout: KVCacheLayout,
116 ):
117 from vllm.v1.attention.backends.flashinfer import (
118 trtllm_prefill_attn_kvfp8_dequant,
@@ -130,20 +120,13 @@ def test_trtllm_kvfp8_dequant(
120
121 torch.set_default_device("cuda")
122
133 - if contiguous:
134 - kv_cache, scale = make_contiguous_kv_cache(
135 - NUM_BLOCKS,
136 - num_kv_heads,
137 - block_size,
138 - head_size,
139 - )
140 - else:
141 - kv_cache, scale = make_cross_layer_kv_cache(
142 - NUM_BLOCKS,
143 - num_kv_heads,
144 - block_size,
145 - head_size,
146 - )
123 + kv_cache, scale = make_random_kv_cache(
124 + NUM_BLOCKS,
125 + num_kv_heads,
126 + block_size,
127 + head_size,
128 + layout=layout,
129 + )
130
131 k_scale = scale.clone()
132 v_scale = scale.clone()
@@ -187,7 +170,7 @@ def test_block_tables_with_zero_pages():
170 torch.set_default_device("cuda")
171 num_kv_heads, block_size, head_size = 8, 16, 64
172
190 - kv_cache, scale = make_contiguous_kv_cache(
173 + kv_cache, scale = make_random_kv_cache(
174 NUM_BLOCKS,
175 num_kv_heads,
176 block_size,
@@ -234,7 +217,7 @@ def test_all_zero_block_tables():
217 torch.set_default_device("cuda")
218 num_kv_heads, block_size, head_size = 4, 16, 64
219
237 - kv_cache, scale = make_contiguous_kv_cache(
220 + kv_cache, scale = make_random_kv_cache(
221 NUM_BLOCKS,
222 num_kv_heads,
223 block_size,
@@ -266,7 +249,7 @@ def test_different_k_v_scales():
249 torch.set_default_device("cuda")
250 num_kv_heads, block_size, head_size = 8, 16, 64
251
269 - kv_cache, _ = make_contiguous_kv_cache(
252 + kv_cache, _ = make_random_kv_cache(
253 NUM_BLOCKS,
254 num_kv_heads,
255 block_size,
@@ -299,7 +282,7 @@ def test_single_page_per_seq():
282 torch.set_default_device("cuda")
283 num_kv_heads, block_size, head_size = 8, 16, 128
284
302 - kv_cache, scale = make_contiguous_kv_cache(
285 + kv_cache, scale = make_random_kv_cache(
286 NUM_BLOCKS,
287 num_kv_heads,
288 block_size,
@@ -332,7 +315,7 @@ def test_large_page_indices():
315 num_kv_heads, block_size, head_size = 8, 16, 128
316 large_num_blocks = 32768
317
335 - kv_cache, scale = make_contiguous_kv_cache(
318 + kv_cache, scale = make_random_kv_cache(
319 large_num_blocks,
320 num_kv_heads,
321 block_size,
@@ -369,7 +352,7 @@ def test_large_block_size():
352 torch.set_default_device("cuda")
353 num_kv_heads, block_size, head_size = 4, 64, 128
354
372 - kv_cache, scale = make_contiguous_kv_cache(
355 + kv_cache, scale = make_random_kv_cache(
356 NUM_BLOCKS,
357 num_kv_heads,
358 block_size,
@@ -401,7 +384,7 @@ def test_large_block_size():
384 def test_cross_layer_many_layers():
385 """
386 Non-contiguous with 36 layers -- matches real gpt-oss-120b.
404 - Strides are far from contiguous (factor of 36 in the gaps).
387 + Cross-layer (BLHNC) strides are far from contiguous (factor of 36 in the gaps).
388 """
389 from vllm.v1.attention.backends.flashinfer import (
390 trtllm_prefill_attn_kvfp8_dequant,
@@ -409,14 +392,14 @@ def test_cross_layer_many_layers():
392
393 torch.set_default_device("cuda")
394 num_kv_heads, block_size, head_size = 8, 16, 64
412 - num_layers = 36
395
414 - kv_cache, scale = make_cross_layer_kv_cache(
396 + kv_cache, scale = make_random_kv_cache(
397 NUM_BLOCKS,
398 num_kv_heads,
399 block_size,
400 head_size,
419 - num_layers=num_layers,
401 + layout=KVCacheLayout.BLHNC,
402 + num_layers=36,
403 )
404 k_scale = v_scale = scale.clone()
405
tests/models/multimodal/generation/test_voxtral_realtime.py
+4 -1
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from fractions import Fraction
4 +
5 import pytest
6 import pytest_asyncio
7 from mistral_common.protocol.transcription.request import (
@@ -75,7 +77,8 @@ def assert_encoder_kv_cache_spec(engine: LLM) -> None:
77 assert audio_config.block_pool_size == 4
78 assert isinstance(spec, SlidingWindowSpec)
79 assert spec.block_size == 16
78 - assert spec.num_kv_heads == 128
80 + assert spec.num_kv_heads == 32
81 + assert spec.tokens_per_state == Fraction(1, 4)
82 # cdiv(750, 4) == 188 pooled tokens cover the model's window; the extra
83 # +1 is an eviction margin (see whisper_causal.py get_kv_cache_spec).
84 assert spec.sliding_window == cdiv(750, 4) + 1 == 189
tests/models/test_initialization.py
+7
@@ -8,6 +8,7 @@ import pytest
8
9 from vllm import LLM
10 from vllm.utils.mem_constants import GiB_bytes
11 +from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
12 from vllm.v1.core.kv_cache_utils import (
13 generate_scheduler_kv_cache_config,
14 get_kv_cache_configs,
@@ -82,6 +83,12 @@ def can_initialize(
83 # Avoid calling model.forward()
84 def _initialize_kv_caches_v1(self, vllm_config):
85 kv_cache_specs = self.model_executor.get_kv_cache_specs()
86 + layout = resolve_kv_cache_layout(
87 + vllm_config,
88 + self.model_executor.get_supported_kv_cache_layouts(),
89 + [spec for worker_specs in kv_cache_specs for spec in worker_specs.values()],
90 + )
91 + self.model_executor.set_kv_cache_layout(layout.name)
92 kv_cache_configs = get_kv_cache_configs(
93 vllm_config,
94 kv_cache_specs,
tests/quantization/test_turboquant.py
+1
@@ -294,6 +294,7 @@ class TestTurboQuantKVCacheSpec:
294 head_size_v=128,
295 num_kv_heads=4,
296 sliding_window=None,
297 + get_attn_backend=lambda: TurboQuantAttentionBackend,
298 )
299 vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32))
300
tests/test_attention_backend_registry.py
-10
@@ -38,11 +38,6 @@ class CustomAttentionBackend(AttentionBackend):
38 """Mock builder class."""
39 return None
40
41 - @staticmethod
42 - def get_required_kv_cache_layout():
43 - """Mock KV cache layout."""
44 - return None
45 -
41
42 class CustomMambaAttentionImpl(AttentionImpl):
43 """Mock custom mamba attention implementation for testing."""
@@ -71,11 +66,6 @@ class CustomMambaAttentionBackend(AttentionBackend):
66 """Mock builder class."""
67 return None
68
74 - @staticmethod
75 - def get_required_kv_cache_layout():
76 - """Mock KV cache layout."""
77 - return None
78 -
69
70 def test_custom_is_not_alias_of_any_backend():
71 # Get all members of AttentionBackendEnum
tests/v1/attention/test_attention_backends.py
+170 -103
@@ -31,10 +31,7 @@ from vllm.v1.attention.backend import (
31 CommonAttentionMetadata,
32 )
33 from vllm.v1.attention.backends.registry import AttentionBackendEnum
34 -from vllm.v1.attention.backends.utils import (
35 - set_kv_cache_layout,
36 -)
37 -from vllm.v1.kv_cache_interface import FullAttentionSpec
34 +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
35
36 BACKENDS_TO_TEST = [
37 AttentionBackendEnum.FLASH_ATTN,
@@ -44,6 +41,14 @@ BACKENDS_TO_TEST = [
41 "FLEX_ATTENTION_SLOW",
42 ]
43
44 +
45 +def _actual_backend(backend) -> AttentionBackendEnum:
46 + """Resolve pseudo-backends (FLEX_ATTENTION_SLOW) to their real enum."""
47 + if backend == "FLEX_ATTENTION_SLOW":
48 + return AttentionBackendEnum.FLEX_ATTENTION
49 + return backend
50 +
51 +
52 DEVICE_TYPE = current_platform.device_type
53
54 # Use the platform's preferred FP8 type so the stored cache matches what the
@@ -118,6 +123,7 @@ def create_and_prepopulate_kv_cache(
123 device: torch.device,
124 num_blocks: int,
125 common_attn_metadata: CommonAttentionMetadata,
126 + layout: KVCacheLayout,
127 randomize_blocks: bool = True,
128 kv_cache_dtype: str = "auto",
129 ) -> torch.Tensor:
@@ -126,19 +132,22 @@ def create_and_prepopulate_kv_cache(
132 Args:
133 k_contexts: List of key context tensors for each sequence
134 v_contexts: List of value context tensors for each sequence
129 - seq_lens: List of sequence lengths
135 block_size: Size of each block
136 num_kv_heads: Number of KV heads
137 head_size: Size of each head
138 dtype: Data type for the cache
139 device: Device to create the cache on
140 num_blocks: Total number of blocks in the cache
136 - block_table: Block table tensor to populate
141 + common_attn_metadata: Provides seq lens, block table and slot mapping
142 + layout: Physical layout to allocate in; the cache is returned as the
143 + logical ``[B, H, N, C]`` view (as ``create_kv_cache_views`` does)
144 randomize_blocks: Whether to randomly permute blocks
145 or use sequential order
146 + kv_cache_dtype: Cache dtype string; fp8 caches use fp8 storage
147
148 Returns:
141 - Tuple of (kv_cache, updated_block_table)
149 + A 4D tensor in logical ``(num_blocks, num_kv_heads, block_size,
150 + 2 * head_size)`` order with strides determined by ``layout``.
151 """
152 batch_size = len(k_contexts)
153 seq_lens = common_attn_metadata.seq_lens.cpu()
@@ -155,26 +164,33 @@ def create_and_prepopulate_kv_cache(
164 fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
165 storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
166
158 - kv_cache = torch.zeros(
159 - num_blocks,
160 - block_size,
161 - num_kv_heads,
162 - 2 * head_size,
163 - dtype=storage_dtype,
164 - device=device,
165 - )
166 - kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
167 -
168 - # Populate the cache with the context tokens
167 + # Logical 5D shape is always [L, B, H, N, C]. Cross-layer layouts need
168 + # at least two layers to reproduce the inter-layer gaps in a layer view.
169 + logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size)
170 + num_layers = 1 if layout.is_layer_compact else 2
171 + logical_5d = (num_layers, *logical_4d)
172 + physical_5d = tuple(logical_5d[i] for i in layout.stride_order)
173 + inv_order = [layout.stride_order.index(i) for i in range(5)]
174 +
175 + kv_cache_physical = torch.zeros(physical_5d, dtype=storage_dtype, device=device)
176 + # Permute to logical [L, B, H, N, C], then select a layer. This mirrors
177 + # create_kv_cache_views and retains cross-layer strides in the 4D view.
178 + kv_cache = kv_cache_physical.permute(*inv_order)[0]
179 +
180 + # Write context tokens into the cache via the logical view:
181 + # kv_cache[block, :, token_in_block, :] routes correctly regardless
182 + # of physical layout.
183 # Start from block_id=1 since block_id=0 is considered the null block
184 start_block_idx = 1
185 for i in range(batch_size):
186 k_context, v_context = k_contexts[i], v_contexts[i]
173 - start = start_block_idx * block_size
174 - end = start + k_context.shape[0]
175 - kv_cache_flat[start:end, :, :head_size] = k_context
176 - kv_cache_flat[start:end, :, head_size:] = v_context
177 -
187 + t = torch.arange(k_context.shape[0], device=device)
188 + blk = start_block_idx + t // block_size
189 + off = t % block_size
190 + # Advanced indexing on (blk, off) yields [T, H, hs] destinations.
191 + # index_put is dtype-strict; cast like the scalar path would.
192 + kv_cache[blk, :, off, :head_size] = k_context.to(kv_cache.dtype)
193 + kv_cache[blk, :, off, head_size:] = v_context.to(kv_cache.dtype)
194 # Stay block aligned and allocate enough blocks for the new tokens
195 start_block_idx += cdiv(int(seq_lens[i]), block_size)
196
@@ -214,9 +230,6 @@ def create_and_prepopulate_kv_cache(
230 i, block_indices
231 ] * block_size + token_inter_block_offsets.to(device)
232
217 - # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
218 - kv_cache = kv_cache.transpose(1, 2).contiguous()
219 -
233 if fp8_kv_cache:
234 kv_cache = kv_cache.view(torch.uint8)
235
@@ -236,6 +249,22 @@ class MockAttentionLayer:
249 self._v_scale_float = 1.0
250
251
252 +def _clone_kv_cache_in_layout(
253 + kv_cache: torch.Tensor, layout: KVCacheLayout
254 +) -> torch.Tensor:
255 + """Copy a logical [B, H, N, C] cache into a fresh allocation in `layout`."""
256 + logical_5d = (1, *kv_cache.shape)
257 + physical = torch.zeros(
258 + tuple(logical_5d[i] for i in layout.stride_order),
259 + dtype=kv_cache.dtype,
260 + device=kv_cache.device,
261 + )
262 + inv_order = [layout.stride_order.index(i) for i in range(5)]
263 + view = physical.permute(*inv_order)[0]
264 + view.copy_(kv_cache)
265 + return view
266 +
267 +
268 def run_attention_backend(
269 backend: AttentionBackendEnum,
270 kv_cache_spec: FullAttentionSpec,
@@ -254,18 +283,17 @@ def run_attention_backend(
283 ) -> torch.Tensor:
284 """Run attention computation using the specified backend's AttentionImpl."""
285
257 - # Handle special case for FLEX_ATTENTION_SLOW
258 - actual_backend = backend
259 -
260 - use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
286 + use_direct_block_mask = not current_platform.is_rocm() and is_torch_equal_or_newer(
287 + "2.9.0.dev0"
288 + )
289 if backend == "FLEX_ATTENTION_SLOW":
262 - actual_backend = AttentionBackendEnum.FLEX_ATTENTION
290 use_direct_block_mask = False
291 + backend = _actual_backend(backend)
292
265 - builder_cls, impl_cls = try_get_attention_backend(actual_backend)
293 + builder_cls, impl_cls = try_get_attention_backend(backend)
294
295 # Mock flashinfer's get_per_layer_parameters if needed
268 - if actual_backend == AttentionBackendEnum.FLASHINFER:
296 + if backend == AttentionBackendEnum.FLASHINFER:
297 import unittest.mock
298
299 from vllm.v1.attention.backends.utils import PerLayerParameters
@@ -295,7 +323,7 @@ def run_attention_backend(
323 else:
324 # Build metadata
325 builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
298 - if actual_backend == AttentionBackendEnum.FLEX_ATTENTION:
326 + if backend == AttentionBackendEnum.FLEX_ATTENTION:
327 builder.direct_build = use_direct_block_mask
328 attn_metadata = builder.build(
329 common_prefix_len=0,
@@ -311,17 +339,19 @@ def run_attention_backend(
339 )
340 head_size = vllm_config.model_config.get_head_size()
341 scale = 1.0 / (head_size**0.5)
314 - impl = impl_cls(
315 - num_heads=num_heads,
316 - head_size=head_size,
317 - scale=scale,
318 - num_kv_heads=num_kv_heads,
319 - alibi_slopes=None,
320 - sliding_window=sliding_window,
321 - attn_type=attn_type,
322 - kv_cache_dtype=kv_cache_dtype,
323 - **({"sinks": sinks} if sinks is not None else {}),
324 - )
342 + # Impls capture the current vllm config at construction, as in model loading.
343 + with set_current_vllm_config(vllm_config):
344 + impl = impl_cls(
345 + num_heads=num_heads,
346 + head_size=head_size,
347 + scale=scale,
348 + num_kv_heads=num_kv_heads,
349 + alibi_slopes=None,
350 + sliding_window=sliding_window,
351 + attn_type=attn_type,
352 + kv_cache_dtype=kv_cache_dtype,
353 + **({"sinks": sinks} if sinks is not None else {}),
354 + )
355
356 # Create mock layer and output buffer
357 mock_layer = MockAttentionLayer(device)
@@ -333,7 +363,7 @@ def run_attention_backend(
363 # Run forward pass
364 # NOTE: The query, key, and value are already shaped correctly
365 # in the calling test function.
336 - if not try_backend_includes_kv_cache_update(actual_backend):
366 + if not try_backend_includes_kv_cache_update(backend):
367 impl.do_kv_cache_update(
368 mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
369 )
@@ -347,7 +377,7 @@ def run_attention_backend(
377 def _test_backend_correctness(
378 batch_spec: BatchSpec,
379 model: str,
350 - backend_to_test: list[AttentionBackendEnum | str],
380 + backend_to_test: list[AttentionBackendEnum],
381 mask_mod,
382 *,
383 causal: bool = True,
@@ -358,6 +388,7 @@ def _test_backend_correctness(
388 tensor_parallel_size: int = 1,
389 kv_cache_dtype: str = "auto",
390 use_sinks: bool = False,
391 + layout: KVCacheLayout | None = None,
392 ):
393 """
394 Test that all backends produce similar outputs to a reference implementation
@@ -535,6 +566,22 @@ def _test_backend_correctness(
566 common_attn_metadata.causal = causal
567
568 # 3. Simulate Paged KV Cache and a realistic slot_mapping
569 + # Mirror selector-time resolution locally (caller-requested layout, then
570 + # any backend-required layout, then the default chain).
571 + declared = [
572 + supported
573 + for backend in backend_to_test
574 + if (
575 + supported := _actual_backend(backend)
576 + .get_class()
577 + .supported_kv_cache_layouts()
578 + )
579 + is not None
580 + ]
581 + if layout is None:
582 + # Mirror the resolver: the shared cache uses the declared sets' preferred
583 + # layout; backends that don't support it get a per-backend copy below.
584 + layout = min(declared, key=len)[0] if declared else KVCacheLayout.LBNHC
585 kv_cache = create_and_prepopulate_kv_cache(
586 k_contexts=k_contexts,
587 v_contexts=v_contexts,
@@ -545,6 +592,7 @@ def _test_backend_correctness(
592 device=device,
593 num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
594 common_attn_metadata=common_attn_metadata,
595 + layout=layout,
596 randomize_blocks=True,
597 kv_cache_dtype=kv_cache_dtype,
598 )
@@ -553,62 +601,44 @@ def _test_backend_correctness(
601 # Note: flex_attention has known Triton kernel compatibility issues
602 # with test infrastructures
603 for backend_name in backend_to_test:
556 - reset_kv_cache_layout = False
557 -
558 - # Resolve backend class for both enum and string names.
559 - actual_backend = backend_name
560 - if backend_name == "FLEX_ATTENTION_SLOW":
561 - actual_backend = AttentionBackendEnum.FLEX_ATTENTION
562 - if hasattr(actual_backend, "get_class"):
563 - backend_cls = actual_backend.get_class()
564 - else:
565 - backend_cls = None
604 + backend_cls = _actual_backend(backend_name).get_class()
605
606 if is_quantized_kv_cache(kv_cache_dtype) and (
568 - backend_cls is None
569 - or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
607 + not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
608 ):
609 continue
610
573 - if backend_name == AttentionBackendEnum.FLASHINFER:
574 - set_kv_cache_layout("HND")
575 - reset_kv_cache_layout = True
576 -
611 kv_cache_for_backend = kv_cache
578 - if backend_cls is not None:
579 - try:
580 - stride_order = backend_cls.get_kv_cache_stride_order()
581 - except (AttributeError, NotImplementedError):
582 - stride_order = tuple(range(kv_cache.ndim))
583 - if stride_order != tuple(range(kv_cache.ndim)):
584 - # Apply stride order like runtime does in
585 - # _reshape_kv_cache (attn_utils.py:182-210): permute to physical
586 - # layout, make contiguous, then permute to logical layout.
587 - inv_order = [stride_order.index(i) for i in range(len(stride_order))]
588 - kv_cache_for_backend = (
589 - kv_cache.permute(*stride_order).contiguous().permute(*inv_order)
590 - )
591 -
592 - try:
593 - backend_output = run_attention_backend(
594 - backend_name,
595 - kv_cache_spec,
596 - ["placeholder"],
597 - vllm_config,
598 - device,
599 - common_attn_metadata,
600 - query_vllm,
601 - key_vllm,
602 - value_vllm,
603 - kv_cache_for_backend,
604 - sliding_window=sliding_window,
605 - attn_type=attn_type,
606 - kv_cache_dtype=kv_cache_dtype,
607 - sinks=sinks,
608 - )
609 - finally:
610 - if reset_kv_cache_layout:
611 - set_kv_cache_layout(None)
612 + backend_layout = layout
613 +
614 + backend_supported = backend_cls.supported_kv_cache_layouts()
615 + if backend_supported is not None and backend_layout not in backend_supported:
616 + # Production resolution never pairs this backend with the shared layout
617 + # (e.g. flex is LBNHC-only next to FlashInfer's LBHNC); give it a copy
618 + # of the cache in a layout it supports.
619 + backend_layout = backend_supported[0]
620 + kv_cache_for_backend = _clone_kv_cache_in_layout(kv_cache, backend_layout)
621 +
622 + # FlashInfer reads the layout at plan time; set it to match
623 + # the physical order of the test cache.
624 + vllm_config.cache_config.kv_cache_layout = backend_layout.name
625 +
626 + backend_output = run_attention_backend(
627 + backend_name,
628 + kv_cache_spec,
629 + ["placeholder"],
630 + vllm_config,
631 + device,
632 + common_attn_metadata,
633 + query_vllm,
634 + key_vllm,
635 + value_vllm,
636 + kv_cache_for_backend,
637 + sliding_window=sliding_window,
638 + attn_type=attn_type,
639 + kv_cache_dtype=kv_cache_dtype,
640 + sinks=sinks,
641 + )
642
643 # Check shape and dtype consistency
644 assert backend_output.shape == sdpa_output.shape, (
@@ -637,6 +667,38 @@ def _test_backend_correctness(
667 )
668
669
670 +@pytest.mark.parametrize("layout", ["BLHNC", "BHLNC"])
671 +@pytest.mark.parametrize("batch_spec_name", ["small_decode", "small_prefill"])
672 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
673 +def test_flashinfer_cross_layer_layout(
674 + default_vllm_config,
675 + layout: str,
676 + batch_spec_name: str,
677 + kv_cache_dtype: str,
678 +):
679 + if AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST:
680 + pytest.skip("FlashInfer is not installed")
681 +
682 + def causal_mask_mod(
683 + b: torch.Tensor,
684 + h: torch.Tensor,
685 + q_idx: torch.Tensor,
686 + kv_idx: torch.Tensor,
687 + *,
688 + context_len: int,
689 + ):
690 + return (q_idx + context_len) >= kv_idx
691 +
692 + _test_backend_correctness(
693 + batch_spec=BATCH_SPECS[batch_spec_name],
694 + model="meta-llama/Meta-Llama-3-8B",
695 + backend_to_test=[AttentionBackendEnum.FLASHINFER],
696 + mask_mod=causal_mask_mod,
697 + kv_cache_dtype=kv_cache_dtype,
698 + layout=KVCacheLayout[layout],
699 + )
700 +
701 +
702 @pytest.mark.parametrize(
703 "batch_spec_name",
704 [
@@ -984,14 +1046,12 @@ if current_platform.is_rocm():
1046 SLIDING_WINDOW_BACKENDS_TO_TEST = [
1047 AttentionBackendEnum.FLEX_ATTENTION,
1048 AttentionBackendEnum.TRITON_ATTN,
987 - "FLEX_ATTENTION_SLOW",
1049 ]
1050 else:
1051 SLIDING_WINDOW_BACKENDS_TO_TEST = [
1052 AttentionBackendEnum.FLASH_ATTN,
1053 AttentionBackendEnum.FLEX_ATTENTION,
1054 AttentionBackendEnum.TRITON_ATTN,
994 - "FLEX_ATTENTION_SLOW",
1055 ]
1056
1057
@@ -1009,7 +1069,10 @@ else:
1069 @pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])
1070 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
1071 def test_sliding_window_backend_correctness(
1012 - default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
1072 + default_vllm_config,
1073 + batch_spec_name: str,
1074 + model: str,
1075 + tensor_parallel_size: int,
1076 ):
1077 """Test backend's correctness with sliding window attention."""
1078
@@ -1071,7 +1134,10 @@ def test_sliding_window_backend_correctness(
1134 @pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
1135 @pytest.mark.parametrize("tensor_parallel_size", [1, 2])
1136 def test_sliding_window_encoder_backend_correctness(
1074 - default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
1137 + default_vllm_config,
1138 + batch_spec_name: str,
1139 + model: str,
1140 + tensor_parallel_size: int,
1141 ):
1142 """Test backend's correctness with sliding window attention."""
1143
@@ -1107,7 +1173,6 @@ def test_sliding_window_encoder_backend_correctness(
1173 NON_CAUSAL_BACKENDS_TO_TEST = [
1174 AttentionBackendEnum.FLASH_ATTN,
1175 AttentionBackendEnum.FLEX_ATTENTION,
1110 - "FLEX_ATTENTION_SLOW",
1176 ]
1177
1178 if current_platform.is_rocm():
@@ -1128,7 +1193,9 @@ if current_platform.is_rocm():
1193 )
1194 @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
1195 def test_non_causal_backend_correctness(
1131 - default_vllm_config, batch_spec_name: str, model: str
1196 + default_vllm_config,
1197 + batch_spec_name: str,
1198 + model: str,
1199 ):
1200 """Test backend's correctness with non-causal (bidirectional) decoder
1201 attention, as used by DFlash speculative decoding."""
tests/v1/attention/test_group_head_counts.py
+6 -53
@@ -1,6 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Attention metadata geometry must come from the builder's own group."""
3 +"""Scheduler metadata sizes a scratchpad from the query head count, so it must
4 +come from the builder's own group: the model-wide ``get_num_attention_heads()``
5 +is wrong for models that vary it per layer (e.g. Laguna), and too small a
6 +scratchpad is indexed past its end.
7 +"""
8
9 from types import SimpleNamespace
10 from unittest.mock import MagicMock, patch
@@ -13,9 +17,8 @@ from vllm.v1.attention.backends.cpu_attn import (
17 CPUAttentionBackendImpl,
18 CPUAttentionMetadataBuilder,
19 )
16 -from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder
20
18 -requires_cpu = pytest.mark.skipif(
21 +pytestmark = pytest.mark.skipif(
22 not current_platform.is_cpu(), reason="CPU attention backend"
23 )
24
@@ -66,7 +69,6 @@ def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:
69 )
70
71
69 -@requires_cpu
72 @pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16])
73 def test_num_heads_comes_from_the_group(group_num_heads):
74 """The group's own count wins, even when it is not the model-wide one."""
@@ -74,56 +76,7 @@ def test_num_heads_comes_from_the_group(group_num_heads):
76 assert builder.num_heads == group_num_heads
77
78
77 -@requires_cpu
79 def test_mixed_head_counts_in_one_group_are_rejected():
80 """Grouping guarantees uniformity; a mixed group means that broke."""
81 with pytest.raises(AssertionError, match="share num_heads"):
82 _build([MODEL_WIDE_NUM_HEADS, 64])
82 -
83 -
84 -def test_flash_attention_geometry_comes_from_the_group():
85 - """FA3 must use group geometry for layers without an ``impl`` wrapper."""
86 - layers = {f"layer_{i}": SimpleNamespace(num_heads=16) for i in range(2)}
87 - vllm_config = MagicMock()
88 - vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS
89 - vllm_config.model_config.get_num_kv_heads.return_value = NUM_KV_HEADS
90 - vllm_config.model_config.get_head_size.return_value = 128
91 - vllm_config.model_config.rswa_window = None
92 - vllm_config.model_config.is_mm_prefix_lm = False
93 - vllm_config.parallel_config.cp_kv_cache_interleave_size = 1
94 - vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs.return_value = (
95 - False
96 - )
97 - vllm_config.compilation_config.max_cudagraph_capture_size = None
98 - kv_cache_spec = SimpleNamespace(
99 - block_size=16,
100 - num_kv_heads=2,
101 - head_size=64,
102 - dtype=torch.bfloat16,
103 - )
104 -
105 - with (
106 - patch(
107 - "vllm.v1.attention.backends.utils.get_layers_from_vllm_config",
108 - return_value=layers,
109 - ),
110 - patch(
111 - "vllm.distributed.parallel_state.get_dcp_group",
112 - side_effect=AssertionError,
113 - ),
114 - patch(
115 - "vllm.v1.attention.backends.flash_attn.get_flash_attn_version",
116 - return_value=3,
117 - ),
118 - ):
119 - builder = FlashAttentionMetadataBuilder(
120 - kv_cache_spec=kv_cache_spec,
121 - layer_names=list(layers),
122 - vllm_config=vllm_config,
123 - device=torch.device("cpu"),
124 - )
125 -
126 - assert builder.aot_schedule
127 - assert builder.num_heads_q == 16
128 - assert builder.num_heads_kv == 2
129 - assert builder.headdim == 64
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+4 -4
@@ -34,19 +34,19 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
34
35
36 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
37 -def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_size():
37 +def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_num_states():
38 """Regression test: DeepseekV4 compression path must compute slot_mapping from
39 compressed positions, not reuse the uncompressed common metadata mapping.
40 """
41 device = torch.device("cuda")
42
43 - # storage_block_size = block_size // compress_ratio = 256 // 4 = 64
43 + # num_states = block_size // tokens_per_state = 256 // 4 = 64
44 kv_cache_spec = MLAAttentionSpec(
45 block_size=256,
46 num_kv_heads=1,
47 head_size=128,
48 dtype=torch.bfloat16,
49 - compress_ratio=4,
49 + tokens_per_state=4,
50 )
51 vllm_config = create_vllm_config(max_model_len=1024)
52 max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024)
@@ -95,7 +95,7 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_
95 valid_slots = md.slot_mapping[md.slot_mapping >= 0]
96 assert valid_slots.numel() == 10 # 40 tokens / compress_ratio 4
97
98 - storage_bs = kv_cache_spec.storage_block_size # 64
98 + storage_bs = kv_cache_spec.num_states # 64
99 # Compressed positions 60..63 land in block 5, positions 64..69 in block 7.
100 expected = torch.tensor(
101 [
tests/v1/attention/test_mla_backends.py
+21 -4
@@ -250,6 +250,15 @@ def _prefill_backend_dimension_params():
250 )
251 )
252 )
253 + elif (
254 + current_platform.is_rocm()
255 + and prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN
256 + ):
257 + # Pre-existing on main: flash-attn MLA prefill returns wrong
258 + # outputs and faults on ROCm, so it fails on the base commit too.
259 + marks.append(
260 + pytest.mark.skip(reason="FLASH_ATTN MLA prefill broken on ROCm")
261 + )
262 params.append(
263 pytest.param(
264 prefill_backend,
@@ -392,7 +401,7 @@ def create_and_prepopulate_kv_cache(
401 kv_entry_size = head_size
402
403 kv_cache = torch.zeros(
395 - num_blocks, block_size, kv_entry_size, dtype=torch.uint8, device=device
404 + num_blocks, 1, block_size, kv_entry_size, dtype=torch.uint8, device=device
405 )
406 scale_tensor = (
407 scale
@@ -401,9 +410,9 @@ def create_and_prepopulate_kv_cache(
410 )
411 scale_tensor = scale_tensor.to(device=device, dtype=torch.float32)
412 else:
404 - # Create MLA KV cache: (num_blocks, block_size, head_size)
413 + # Create MLA KV cache: (num_blocks, num_heads=1, block_size, head_size)
414 kv_cache = torch.zeros(
406 - num_blocks, block_size, head_size, dtype=dtype, device=device
415 + num_blocks, 1, block_size, head_size, dtype=dtype, device=device
416 )
417 kv_cache_flat = kv_cache.view(-1, head_size)
418
@@ -424,7 +433,7 @@ def create_and_prepopulate_kv_cache(
433 ops.concat_and_cache_mla(
434 kv_c_context,
435 k_pe_context.squeeze(1),
427 - kv_cache,
436 + kv_cache.squeeze(1),
437 slots,
438 kv_cache_dtype=kv_cache_dtype,
439 scale=scale_tensor,
@@ -543,6 +552,10 @@ class MockSparseMLAAttentionLayer:
552 kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
553 fp8_attention = kv_cache_dtype.startswith("fp8")
554
555 + # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache.
556 + if kv_cache.ndim == 4:
557 + kv_cache = kv_cache.squeeze(1)
558 +
559 # Write to KV cache
560 if kv_cache.numel() > 0:
561 ops.concat_and_cache_mla(
@@ -678,6 +691,10 @@ class MockMLAAttentionLayer(MLAAttention):
691 output: torch.Tensor,
692 ) -> torch.Tensor:
693 """Replicates MLAAttention.forward_impl logic for testing."""
694 + # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache.
695 + if kv_cache.ndim == 4:
696 + kv_cache = kv_cache.squeeze(1)
697 +
698 # Write to KV cache
699 kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
700 fp8_attention = kv_cache_dtype.startswith("fp8")
tests/v1/attention/test_mm_prefix.py
+5 -1
@@ -20,7 +20,10 @@ import pytest
20 import torch
21
22 from vllm.platforms import current_platform
23 -from vllm.v1.attention.backends.utils import fill_mm_prefix_query_ranges
23 +from vllm.v1.attention.backends.utils import (
24 + fill_mm_prefix_query_ranges,
25 +)
26 +from vllm.v1.kv_cache_interface import KVCacheLayout
27
28
29 def _fa4_available() -> bool:
@@ -585,6 +588,7 @@ def test_mm_prefix_kv_cache_path(head_size: int):
588 device=DEVICE,
589 num_blocks=2048,
590 common_attn_metadata=common,
591 + layout=KVCacheLayout.LBNHC,
592 randomize_blocks=True,
593 )
594
tests/v1/attention/test_sparse_mla_backends.py
+3 -2
@@ -813,7 +813,7 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
813 assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens
814
815 q = torch.zeros(num_mqa_tokens, 4, 576, dtype=torch.bfloat16, device=device)
816 - kv_cache = torch.zeros(40 * block_size, 576, dtype=torch.bfloat16, device=device)
816 + kv_cache = torch.zeros(40, block_size, 576, dtype=torch.bfloat16, device=device)
817 topk_indices = torch.randint(
818 0,
819 block_size * 10,
@@ -1157,7 +1157,8 @@ def test_sparse_backend_prefill_correctness(
1157 q=query_cat,
1158 kv_c_normed=kv_c_cat,
1159 k_pe=k_pe_cat,
1160 - kv_c_and_k_pe_cache=kv_cache,
1160 + # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache.
1161 + kv_c_and_k_pe_cache=kv_cache.squeeze(1),
1162 attn_metadata=metadata,
1163 k_scale=torch.tensor(1.0, device=device),
1164 output=out_buffer,
tests/v1/attention/test_trtllm_attention_integration.py
+120 -219
@@ -9,6 +9,7 @@ import pytest
9 import torch
10 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
11
12 +from tests.v1.attention.test_attention_backends import create_and_prepopulate_kv_cache
13 from tests.v1.attention.utils import (
14 BatchSpec,
15 create_common_attn_metadata,
@@ -16,14 +17,9 @@ from tests.v1.attention.utils import (
17 )
18 from vllm.config import set_current_vllm_config
19 from vllm.platforms import current_platform
19 -from vllm.utils.math_utils import cdiv
20 from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim, set_random_seed
21 -from vllm.v1.attention.backends.utils import (
22 - PerLayerParameters,
23 - get_kv_cache_layout,
24 - set_kv_cache_layout,
25 -)
26 -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
21 +from vllm.v1.attention.backends.utils import PerLayerParameters
22 +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode
23
24 if not current_platform.is_device_capability_family(100):
25 pytest.skip(
@@ -87,96 +83,6 @@ def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
83 }
84
85
90 -def _create_hnd_kv_cache(
91 - k_contexts,
92 - v_contexts,
93 - block_size,
94 - num_kv_heads,
95 - head_size,
96 - dtype,
97 - device,
98 - num_blocks,
99 - common_attn_metadata,
100 - kv_in_head_dim=False,
101 -):
102 - """Create and populate a packed KV cache with HND-compatible strides.
103 -
104 - When kv_in_head_dim=False (default), returns (B, H, N, 2*hs) with K/V
105 - packed in the content dim. When kv_in_head_dim=True, returns
106 - (B, 2*H, N, hs) with K/V as separate head groups.
107 - """
108 - seq_lens = common_attn_metadata.seq_lens.cpu()
109 - query_lens = (
110 - common_attn_metadata.query_start_loc_cpu[1:]
111 - - common_attn_metadata.query_start_loc_cpu[:-1]
112 - )
113 - block_table = common_attn_metadata.block_table_tensor
114 - slot_mapping = common_attn_metadata.slot_mapping
115 - batch_size = len(k_contexts)
116 -
117 - # kv_in_head_dim: (B, N, 2*H, hs) — K/V as separate head groups
118 - # else: (B, N, H, 2*hs) — K/V packed in content dim
119 - n_heads, content = (
120 - (2 * num_kv_heads, head_size)
121 - if kv_in_head_dim
122 - else (num_kv_heads, 2 * head_size)
123 - )
124 - kv_cache = torch.zeros(
125 - num_blocks,
126 - block_size,
127 - n_heads,
128 - content,
129 - dtype=dtype,
130 - device=device,
131 - )
132 - kv_cache_flat = kv_cache.view(-1, n_heads, content)
133 -
134 - start_block_idx = 1
135 - for i in range(batch_size):
136 - k_ctx, v_ctx = k_contexts[i], v_contexts[i]
137 - start = start_block_idx * block_size
138 - end = start + k_ctx.shape[0]
139 - if kv_in_head_dim:
140 - kv_cache_flat[start:end, :num_kv_heads] = k_ctx
141 - kv_cache_flat[start:end, num_kv_heads:] = v_ctx
142 - else:
143 - kv_cache_flat[start:end, :, :head_size] = k_ctx
144 - kv_cache_flat[start:end, :, head_size:] = v_ctx
145 - start_block_idx += cdiv(int(seq_lens[i]), block_size)
146 -
147 - blocks_end = start_block_idx
148 -
149 - # Randomly permute blocks (starting from block 1; block 0 is null).
150 - perm = torch.randperm(blocks_end - 1) + 1
151 - inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
152 - inv_perm[1:] = torch.argsort(perm) + 1
153 - kv_cache[1:blocks_end] = kv_cache[perm]
154 -
155 - # Build block table.
156 - start_block_idx = 1
157 - for i in range(batch_size):
158 - n_blocks = cdiv(int(seq_lens[i]), block_size)
159 - block_table[i, :n_blocks] = inv_perm[
160 - start_block_idx : start_block_idx + n_blocks
161 - ]
162 - start_block_idx += n_blocks
163 -
164 - # Build slot mapping that is consistent with the block table.
165 - for i in range(batch_size):
166 - ctx_len = int(seq_lens[i]) - int(query_lens[i])
167 - token_offsets = torch.arange(int(query_lens[i])) + ctx_len
168 - block_indices = token_offsets // block_size
169 - intra_block_offsets = token_offsets % block_size
170 - start = common_attn_metadata.query_start_loc_cpu[i]
171 - end = common_attn_metadata.query_start_loc_cpu[i + 1]
172 - slot_mapping[start:end] = block_table[
173 - i, block_indices
174 - ] * block_size + intra_block_offsets.to(device)
175 -
176 - # Transpose to canonical: (B, H, N, 2*hs) or (B, 2*H, N, hs)
177 - return kv_cache.transpose(1, 2).contiguous()
178 -
179 -
86 def _create_nvfp4_hnd_kv_cache(
87 k_contexts,
88 v_contexts,
@@ -191,7 +97,7 @@ def _create_nvfp4_hnd_kv_cache(
97 ):
98 """Create an nvfp4 KV cache by quantizing bf16 context via
99 reshape_and_cache_flash, using the same block-table layout as
194 - _create_hnd_kv_cache.
100 + create_and_prepopulate_kv_cache.
101
102 The returned tensor is dtype ``uint8`` with head-group layout
103 ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)``
@@ -217,11 +123,10 @@ def _create_nvfp4_hnd_kv_cache(
123 during quantization.
124
125 Returns:
220 - ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided).
126 + ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, LBHNC-strided).
127 """
222 - # First create a bf16 HND cache so block tables are populated.
223 - # Use kv_in_head_dim=True so K/V are separate head groups (B, 2*H, N, hs).
224 - bf16_cache = _create_hnd_kv_cache(
128 + # First create a bf16 cache so block tables are populated.
129 + bf16_cache = create_and_prepopulate_kv_cache(
130 k_contexts,
131 v_contexts,
132 block_size,
@@ -231,7 +136,7 @@ def _create_nvfp4_hnd_kv_cache(
136 device,
137 num_blocks,
138 common_attn_metadata,
234 - kv_in_head_dim=True,
139 + layout=KVCacheLayout.LBHNC,
140 )
141
142 # (num_blocks, 2 * num_kv_heads, block_size, full_dim) — K heads first, then V heads
@@ -244,7 +149,7 @@ def _create_nvfp4_hnd_kv_cache(
149 k_cache, v_cache = nvfp4_cache.split(num_kv_heads, dim=1)
150
151 # Flatten bf16 context into tokens and quantize via reshape_and_cache_flash.
247 - # bf16_cache is (B, 2*H, N, hs); split K/V on head dim.
152 + # bf16_cache is (B, H, N, 2*hs); split K/V on the content dim.
153 block_table = common_attn_metadata.block_table_tensor
154 seq_lens = common_attn_metadata.seq_lens.cpu()
155 query_lens = (
@@ -260,10 +165,16 @@ def _create_nvfp4_hnd_kv_cache(
165 # Gather context tokens from the bf16 cache using block table.
166 n_ctx_blocks = (ctx_len + block_size - 1) // block_size
167 blocks = block_table[i, :n_ctx_blocks]
263 - # bf16_cache is (B, 2*H, N, hs); split K and V head groups.
264 - k_bf16, v_bf16 = bf16_cache[blocks].split(num_kv_heads, dim=1)
265 - k_ctx = k_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
266 - v_ctx = v_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
168 + k_ctx = (
169 + bf16_cache[blocks, :, :, :head_size]
170 + .transpose(1, 2)
171 + .reshape(-1, num_kv_heads, head_size)[:ctx_len]
172 + )
173 + v_ctx = (
174 + bf16_cache[blocks, :, :, head_size:]
175 + .transpose(1, 2)
176 + .reshape(-1, num_kv_heads, head_size)[:ctx_len]
177 + )
178 # Build slot mapping for these context tokens.
179 token_offsets = torch.arange(ctx_len, device=device)
180 block_indices = token_offsets // block_size
@@ -383,7 +294,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
294 )
295 else:
296 kv_scale_val = 1.0
386 - kv_cache = _create_hnd_kv_cache(
297 + kv_cache = create_and_prepopulate_kv_cache(
298 k_contexts,
299 v_contexts,
300 BLOCK_SIZE,
@@ -393,127 +304,117 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
304 device,
305 NUM_GPU_BLOCKS,
306 common_attn_metadata,
307 + layout=KVCacheLayout.LBHNC,
308 )
309
310 # 3. Run through FlashInfer with TRTLLM enabled
399 - set_kv_cache_layout("HND")
400 - get_kv_cache_layout.cache_clear()
401 -
402 - try:
403 - is_nvfp4 = kv_cache_dtype == "nvfp4"
404 - kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE
405 - spec_dtype = torch.uint8 if is_nvfp4 else dtype
406 - kv_cache_spec = FullAttentionSpec(
407 - block_size=BLOCK_SIZE,
408 - num_kv_heads=num_kv_heads,
409 - head_size=head_size,
410 - dtype=spec_dtype,
411 - kv_quant_mode=kv_quant_mode,
311 + vllm_config.cache_config.kv_cache_layout = "LBHNC"
312 +
313 + is_nvfp4 = kv_cache_dtype == "nvfp4"
314 + kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE
315 + spec_dtype = torch.uint8 if is_nvfp4 else dtype
316 + kv_cache_spec = FullAttentionSpec(
317 + block_size=BLOCK_SIZE,
318 + num_kv_heads=num_kv_heads,
319 + head_size=head_size,
320 + dtype=spec_dtype,
321 + kv_quant_mode=kv_quant_mode,
322 + )
323 + layer_names = ["test_layer_0"]
324 +
325 + with (
326 + set_current_vllm_config(vllm_config),
327 + unittest.mock.patch(
328 + "vllm.utils.flashinfer.supports_trtllm_attention",
329 + return_value=True,
330 + ),
331 + unittest.mock.patch(
332 + "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
333 + _mock_get_per_layer_parameters,
334 + ),
335 + ):
336 + builder = FlashInferMetadataBuilder(
337 + kv_cache_spec, layer_names, vllm_config, device
338 + )
339 + attn_metadata = builder.build(
340 + common_prefix_len=0,
341 + common_attn_metadata=common_attn_metadata,
342 )
413 - layer_names = ["test_layer_0"]
414 -
415 - with (
416 - set_current_vllm_config(vllm_config),
417 - unittest.mock.patch(
418 - "vllm.utils.flashinfer.supports_trtllm_attention",
419 - return_value=True,
420 - ),
421 - unittest.mock.patch(
422 - "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
423 - _mock_get_per_layer_parameters,
424 - ),
425 - ):
426 - builder = FlashInferMetadataBuilder(
427 - kv_cache_spec, layer_names, vllm_config, device
343 +
344 + # Verify the correct TRTLLM metadata types were produced.
345 + has_prefills = any(ql > 1 for ql in batch_spec.query_lens)
346 + has_decodes = any(ql == 1 for ql in batch_spec.query_lens)
347 +
348 + if has_prefills:
349 + assert isinstance(attn_metadata.prefill, TRTLLMPrefill), (
350 + f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
351 )
429 - attn_metadata = builder.build(
430 - common_prefix_len=0,
431 - common_attn_metadata=common_attn_metadata,
352 + if has_decodes:
353 + assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), (
354 + f"Expected FlashInferTrtllmAPIDecode, got {type(attn_metadata.decode)}"
355 )
356 + assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN
357
434 - # Verify the correct TRTLLM metadata types were produced.
435 - has_prefills = any(ql > 1 for ql in batch_spec.query_lens)
436 - has_decodes = any(ql == 1 for ql in batch_spec.query_lens)
437 -
438 - if has_prefills:
439 - assert isinstance(attn_metadata.prefill, TRTLLMPrefill), (
440 - f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
441 - )
442 - if has_decodes:
443 - assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), (
444 - "Expected FlashInferTrtllmAPIDecode, got "
445 - f"{type(attn_metadata.decode)}"
446 - )
447 - assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN
448 -
449 - impl = FlashInferImpl(
450 - num_heads=num_q_heads,
451 - head_size=head_size,
452 - scale=scale,
453 - num_kv_heads=num_kv_heads,
454 - alibi_slopes=None,
455 - sliding_window=None,
456 - kv_cache_dtype=kv_cache_dtype,
457 - )
358 + impl = FlashInferImpl(
359 + num_heads=num_q_heads,
360 + head_size=head_size,
361 + scale=scale,
362 + num_kv_heads=num_kv_heads,
363 + alibi_slopes=None,
364 + sliding_window=None,
365 + kv_cache_dtype=kv_cache_dtype,
366 + )
367
459 - mock_layer = MockAttentionLayer(device)
460 - if is_nvfp4:
461 - # For nvfp4, k_scale/v_scale are the global quantization
462 - # scales (amax/448) used by reshape_and_cache_flash.
463 - kv_scale_t = torch.tensor(
464 - kv_scale_val, dtype=torch.float32, device=device
465 - )
466 - mock_layer._k_scale = kv_scale_t
467 - mock_layer._v_scale = kv_scale_t
468 - mock_layer._k_scale_float = kv_scale_val
469 - mock_layer._v_scale_float = kv_scale_val
470 - output = torch.empty_like(query_vllm)
471 -
472 - impl.do_kv_cache_update(
473 - mock_layer,
474 - key_vllm,
475 - value_vllm,
476 - kv_cache,
477 - attn_metadata.slot_mapping,
478 - )
368 + mock_layer = MockAttentionLayer(device)
369 + if is_nvfp4:
370 + # For nvfp4, k_scale/v_scale are the global quantization
371 + # scales (amax/448) used by reshape_and_cache_flash.
372 + kv_scale_t = torch.tensor(kv_scale_val, dtype=torch.float32, device=device)
373 + mock_layer._k_scale = kv_scale_t
374 + mock_layer._v_scale = kv_scale_t
375 + mock_layer._k_scale_float = kv_scale_val
376 + mock_layer._v_scale_float = kv_scale_val
377 + output = torch.empty_like(query_vllm)
378 +
379 + impl.do_kv_cache_update(
380 + mock_layer,
381 + key_vllm,
382 + value_vllm,
383 + kv_cache,
384 + attn_metadata.slot_mapping,
385 + )
386
480 - # nvfp4 trtllm kernel requires FP8 queries. In the real
481 - # pipeline the attention layer handles this; here we
482 - # quantize manually.
483 - if is_nvfp4:
484 - finfo = torch.finfo(torch.float8_e4m3fn)
485 - q_amax = query_vllm.abs().amax().clamp(min=1e-12)
486 - q_s = (finfo.max / q_amax * 0.1).item()
487 - query_vllm = (
488 - (query_vllm * q_s)
489 - .clamp(finfo.min, finfo.max)
490 - .to(torch.float8_e4m3fn)
491 - )
492 - mock_layer._q_scale = torch.tensor(
493 - 1.0 / q_s, dtype=torch.float32, device=device
494 - )
495 - mock_layer._q_scale_float = 1.0 / q_s
496 -
497 - output = impl.forward(
498 - mock_layer,
499 - query_vllm,
500 - key_vllm,
501 - value_vllm,
502 - kv_cache,
503 - attn_metadata,
504 - output=output,
387 + # nvfp4 trtllm kernel requires FP8 queries. In the real
388 + # pipeline the attention layer handles this; here we
389 + # quantize manually.
390 + if is_nvfp4:
391 + finfo = torch.finfo(torch.float8_e4m3fn)
392 + q_amax = query_vllm.abs().amax().clamp(min=1e-12)
393 + q_s = (finfo.max / q_amax * 0.1).item()
394 + query_vllm = (
395 + (query_vllm * q_s).clamp(finfo.min, finfo.max).to(torch.float8_e4m3fn)
396 + )
397 + mock_layer._q_scale = torch.tensor(
398 + 1.0 / q_s, dtype=torch.float32, device=device
399 )
400 + mock_layer._q_scale_float = 1.0 / q_s
401 +
402 + output = impl.forward(
403 + mock_layer,
404 + query_vllm,
405 + key_vllm,
406 + value_vllm,
407 + kv_cache,
408 + attn_metadata,
409 + output=output,
410 + )
411
507 - # 4. Compare against SDPA reference
508 - if is_nvfp4:
509 - atol, rtol = 1.0, 1.0 # nvfp4 has higher quantization error
510 - else:
511 - atol, rtol = 1e-2, 1e-2
512 - torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol)
513 -
514 - finally:
515 - set_kv_cache_layout(None)
516 - get_kv_cache_layout.cache_clear()
412 + # 4. Compare against SDPA reference
413 + if is_nvfp4:
414 + atol, rtol = 1.0, 1.0 # nvfp4 has higher quantization error
415 + else:
416 + atol, rtol = 1e-2, 1e-2
417 + torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol)
418
419
420 @pytest.mark.parametrize(
tests/v1/attention/utils.py
+36
@@ -32,7 +32,12 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
32 from vllm.v1.kv_cache_interface import (
33 EncoderOnlyAttentionSpec,
34 FullAttentionSpec,
35 + KVCacheLayout,
36 + KVCacheSpec,
37 + KVCacheTensor,
38 MambaSpec,
39 + compute_layout_strides,
40 + create_kv_cache_views,
41 get_kv_quant_mode,
42 )
43
@@ -215,6 +220,9 @@ def create_vllm_config(
220 # (these may be set during initialization normally)
221 cache_config.num_gpu_blocks = num_gpu_blocks
222 cache_config.num_cpu_blocks = 0
223 + # Builders read the resolved layout from the config; tests that exercise a
224 + # different layout overwrite this field.
225 + cache_config.kv_cache_layout = "LBNHC"
226
227 parallel_config = ParallelConfig(
228 tensor_parallel_size=tensor_parallel_size,
@@ -417,3 +425,31 @@ class MockMambaBuilder(BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetad
425 is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool)
426 )
427 return builder.build(0, common_metadata)
428 +
429 +
430 +def dense_kv_cache_views(
431 + raw: torch.Tensor,
432 + spec: KVCacheSpec,
433 + num_blocks: int,
434 + num_layers: int,
435 + layout: KVCacheLayout,
436 + kernel_block_size: int | None = None,
437 +) -> list[torch.Tensor]:
438 + """``create_kv_cache_views`` for a dense allocation of ``num_layers`` layers."""
439 + layer_stride, block_stride, _, _, _ = compute_layout_strides(
440 + spec, num_blocks, num_layers, layout
441 + )
442 + tensor = KVCacheTensor(
443 + size=raw.numel(),
444 + layers=[str(i) for i in range(num_layers)],
445 + layer_stride=layer_stride,
446 + block_stride=block_stride,
447 + )
448 + return create_kv_cache_views(
449 + raw,
450 + spec,
451 + num_blocks,
452 + layout,
453 + tensor,
454 + kernel_block_size=kernel_block_size,
455 + )
tests/v1/core/test_contiguous_kv_packing.py
+175 -318
@@ -1,380 +1,237 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for contiguous KV cache packing."""
3 +"""Tests for contiguous KV cache packing.
4 +
5 +Every cache group packs its layers densely into one block; groups overlay each other
6 +(a block ID is owned by one group at a time), so the packed block stride is the
7 +largest group's packing. The layout decides whether the layer dim sits outside the
8 +block dim (a contiguous region per layer) or inside it (all layers' pages within each
9 +block); the allocation is the same either way.
10 +"""
11
12 from unittest.mock import MagicMock
13
14 import pytest
15 import torch
16
17 +from vllm.config import CacheConfig
18 from vllm.v1.core.kv_cache_utils import (
11 - _get_kv_cache_config_packed,
12 - _get_kv_cache_groups_uniform_groups,
19 + _get_kv_cache_bytes_per_block,
20 + _pool_bytes_per_block,
21 get_kv_cache_config_from_groups,
22 )
23 from vllm.v1.kv_cache_interface import (
24 FullAttentionSpec,
25 KVCacheGroupSpec,
18 - KVCacheTensor,
26 + KVCacheLayout,
27 MLAAttentionSpec,
20 - SlidingWindowMLASpec,
21 - SlidingWindowSpec,
28 UniformTypeKVCacheSpecs,
29 )
30 +from vllm.v1.worker.utils import allocate_kv_cache
31 +
32 +MEMORY = 8 * 1024 * 1024
33
34
26 -def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
35 +def _mla(head_size: int) -> MLAAttentionSpec:
36 return MLAAttentionSpec(
28 - block_size=block_size,
29 - num_kv_heads=1,
30 - head_size=512,
31 - dtype=torch.uint8,
32 - page_size_padded=page_size,
33 - cache_dtype_str="fp8_ds_mla",
34 - model_version="deepseek_v4",
35 - alignment=576,
37 + block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8
38 )
39
40
39 -def _make_full_spec() -> FullAttentionSpec:
41 +def _full() -> FullAttentionSpec:
42 return FullAttentionSpec(
41 - block_size=16,
42 - num_kv_heads=2,
43 - head_size=64,
44 - dtype=torch.float16,
43 + block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float16
44 )
45
46
48 -def _make_sw_spec() -> SlidingWindowSpec:
49 - return SlidingWindowSpec(
50 - block_size=16,
51 - num_kv_heads=2,
52 - head_size=64,
53 - dtype=torch.float16,
54 - sliding_window=128,
55 - )
56 -
57 -
58 -def _make_groups(n_c4, n_c128, n_swa):
59 - PS_C4_MLA = 37440
60 - PS_C4_IDX = 8640
61 - PS_C128 = 1728
62 - PS_SWA = 37440
63 -
64 - mla_specs = {}
65 - for i in range(n_c4):
66 - mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
67 - mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
68 - for i in range(n_c128):
69 - mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
70 -
71 - mla_group = KVCacheGroupSpec(
72 - layer_names=list(mla_specs.keys()),
73 - kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
47 +def _uniform_group(specs: dict) -> KVCacheGroupSpec:
48 + return KVCacheGroupSpec(
49 + list(specs),
50 + UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs),
51 )
52
76 - swa_specs = {}
77 - for i in range(n_swa):
78 - swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
79 -
80 - swa_group = KVCacheGroupSpec(
81 - layer_names=list(swa_specs.keys()),
82 - kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
83 - )
53
85 - return [mla_group, swa_group]
54 +def _mixed_page_groups(n_mla=3, n_idx=3, n_swa=5):
55 + """A DeepSeek-V4-style hybrid: two groups with different page mixes."""
56 + g1 = {f"mla.{i}": _mla(512) for i in range(n_mla)}
57 + g1.update({f"idx.{i}": _mla(128) for i in range(n_idx)})
58 + g2 = {f"swa.{i}": _mla(512) for i in range(n_swa)}
59 + return [_uniform_group(g1), _uniform_group(g2)], g1, g2
60
61
88 -def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None):
62 +def _mock_vllm_config(layout: str | None):
63 config = MagicMock()
64 + config.cache_config = CacheConfig()
65 config.cache_config.num_gpu_blocks_override = None
91 - config.cache_config.prefix_cache_retention_interval = 0
92 - config.kv_transfer_config = None
93 - if kv_connector_extra_config is not None:
94 - config.kv_transfer_config = MagicMock()
95 - config.kv_transfer_config.kv_connector_extra_config = kv_connector_extra_config
66 + config.cache_config.kv_cache_layout = layout
67 return config
68
69
99 -def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
100 - groups = _make_groups(n_c4, n_c128, n_swa)
101 - return _get_kv_cache_config_packed(_mock_vllm_config(), groups, mem)
102 -
103 -
104 -def _page_sizes_by_layer(
105 - groups: list[KVCacheGroupSpec],
106 -) -> dict[str, int]:
107 - page_sizes = {}
108 - for group in groups:
109 - specs = group.kv_cache_spec.kv_cache_specs
110 - for layer_name in group.layer_names:
111 - page_sizes[layer_name] = specs[layer_name].page_size_bytes
112 - return page_sizes
113 -
114 -
115 -def _packing_by_layer(
116 - tensors: list[KVCacheTensor],
117 -) -> dict[str, tuple[int, int]]:
70 +def _pages(groups) -> dict[str, int]:
71 return {
119 - layer_name: (tensor.offset, tensor.block_stride)
120 - for tensor in tensors
121 - for layer_name in tensor.shared_by
72 + name: group.kv_cache_spec.kv_cache_specs[name].page_size_bytes
73 + if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
74 + else group.kv_cache_spec.page_size_bytes
75 + for group in groups
76 + for name in group.layer_names
77 }
78
79
125 -def _make_views(
126 - groups: list[KVCacheGroupSpec],
127 - num_blocks: int,
128 - tensors: list[KVCacheTensor],
129 -) -> dict[str, torch.Tensor]:
130 - page_sizes = _page_sizes_by_layer(groups)
131 - packing = _packing_by_layer(tensors)
132 - backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
133 - return {
134 - layer_name: torch.as_strided(
135 - backing,
136 - size=(num_blocks, page_size),
137 - stride=(packing[layer_name][1], 1),
138 - storage_offset=packing[layer_name][0],
139 - )
140 - for layer_name, page_size in page_sizes.items()
141 - }
80 +def _expected_bytes_per_block(groups) -> int:
81 + pages = _pages(groups)
82 + return max(sum(pages[n] for n in g.layer_names) for g in groups)
83
84
144 -def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec:
145 - specs = {
146 - f"{prefix}.{i}": MagicMock(page_size_bytes=page_size)
147 - for i, page_size in enumerate(page_sizes)
148 - }
149 - return KVCacheGroupSpec(
150 - layer_names=list(specs),
151 - kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs),
152 - )
153 -
154 -
155 -class TestInterleavedPacking:
156 - def test_compact_cache_overlays_fp32_state_group(self):
157 - full_specs = {}
158 - state_specs = {}
159 - for i in range(2):
160 - full_specs[f"mla.{i}"] = MLAAttentionSpec(
161 - block_size=256,
162 - num_kv_heads=1,
163 - head_size=512,
164 - dtype=torch.uint8,
165 - page_size_padded=32768,
166 - indexes_kv_by_block_stride=True,
167 - compress_ratio=4,
168 - )
169 - full_specs[f"indexer.{i}"] = MLAAttentionSpec(
170 - block_size=256,
171 - num_kv_heads=1,
172 - head_size=68,
173 - dtype=torch.uint8,
174 - page_size_padded=4608,
175 - compress_ratio=4,
176 - )
177 - state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec(
178 - block_size=4,
179 - num_kv_heads=1,
180 - head_size=2048,
181 - dtype=torch.float32,
182 - sliding_window=8,
183 - indexes_kv_by_block_stride=True,
184 - )
185 - state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec(
186 - block_size=4,
187 - num_kv_heads=1,
188 - head_size=512,
189 - dtype=torch.float32,
190 - sliding_window=8,
191 - indexes_kv_by_block_stride=True,
192 - )
85 +def _bind(config, layout: str):
86 + return allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout[layout], None)
87
194 - grouped_specs = [
195 - UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs),
196 - UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs),
197 - ]
198 - groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
88
200 - assert len(groups) == 2
201 - assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608}
202 - assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == {
203 - 4352
204 - }
205 - assert {
206 - state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2)
207 - } == {8192}
208 -
209 - full_group_bytes = 2 * (32768 + 4608)
210 - state_group_bytes = 2 * (32768 + 8192)
211 - bytes_per_block = max(full_group_bytes, state_group_bytes)
212 - num_blocks, tensors = _get_kv_cache_config_packed(
213 - _mock_vllm_config(), groups, bytes_per_block * 32
89 +class TestDensePacking:
90 + def test_bytes_per_block_is_largest_group(self):
91 + groups, g1, g2 = _mixed_page_groups()
92 + assert _get_kv_cache_bytes_per_block(groups) == _expected_bytes_per_block(
93 + groups
94 )
215 - assert num_blocks == 32
216 - assert {tensor.block_stride for tensor in tensors} == {bytes_per_block}
217 -
218 - packing = _packing_by_layer(tensors)
219 - assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0
220 - assert packing["indexer.0"][0] == 32768
221 - assert packing["indexer_state.0"][0] == 32768
222 -
223 - def test_deepseek_v4_pro_stride(self):
224 - groups = [
225 - _make_page_group("full", [32768, 4608] * 30 + [1024] * 31),
226 - _make_page_group("c4_state", [32768, 8192] * 30),
227 - _make_page_group("c128_state", [32768] * 31),
228 - _make_page_group("swa.0", [32768] * 31),
229 - _make_page_group("swa.1", [32768] * 30),
230 - ]
231 - expected_stride = 1_228_800
95
233 - num_blocks, tensors = _get_kv_cache_config_packed(
234 - _mock_vllm_config(), groups, expected_stride * 32
96 + config = get_kv_cache_config_from_groups(
97 + _mock_vllm_config("BLHNC"), groups, MEMORY
98 )
99 + # Groups overlay: both start at offset 0.
100 + assert [tensor.offset for tensor in config.kv_cache_tensors].count(0) == 2
101 + assert [tensor.layers for tensor in config.kv_cache_tensors] == [
102 + list(g1)[:3],
103 + list(g1)[3:],
104 + list(g2),
105 + ]
106
237 - assert num_blocks == 32
238 - assert {tensor.block_stride for tensor in tensors} == {expected_stride}
239 - assert {tensor.size for tensor in tensors} == {expected_stride * 32}
240 -
241 - def test_all_tensors_have_block_stride(self):
242 - _, tensors = _run()
243 - for t in tensors:
244 - assert t.block_stride > 0
245 -
246 - def test_all_tensors_share_same_size(self):
247 - _, tensors = _run()
248 - sizes = set(t.size for t in tensors)
249 - assert len(sizes) == 1
250 - assert sizes.pop() > 0
251 -
252 - def test_offsets_within_one_block(self):
253 - groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
254 - _, tensors = _get_kv_cache_config_packed(
255 - _mock_vllm_config(), groups, 100 * 1024 * 1024
256 - )
257 - page_sizes = _page_sizes_by_layer(groups)
258 - packing = _packing_by_layer(tensors)
259 - for layer_name, page_size in page_sizes.items():
260 - offset, block_stride = packing[layer_name]
261 - assert offset + page_size <= block_stride
262 -
263 - def test_layouts_are_disjoint_within_each_group(self):
264 - groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
265 - _, tensors = _get_kv_cache_config_packed(
266 - _mock_vllm_config(), groups, 100 * 1024 * 1024
107 + def test_layers_within_a_group_are_dense(self):
108 + groups, _, _ = _mixed_page_groups()
109 + pages = _pages(groups)
110 + config = get_kv_cache_config_from_groups(
111 + _mock_vllm_config("BLHNC"), groups, MEMORY
112 )
268 - page_sizes = _page_sizes_by_layer(groups)
269 - packing = _packing_by_layer(tensors)
270 -
113 + offsets = {
114 + name: tensor.offset + i * tensor.layer_stride
115 + for tensor in config.kv_cache_tensors
116 + for i, name in enumerate(tensor.layers)
117 + }
118 for group in groups:
272 - ranges = sorted(
273 - (packing[name][0], packing[name][0] + page_sizes[name])
274 - for name in group.layer_names
275 - )
276 - assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
277 -
278 - def test_all_layers_accounted_for(self):
279 - n_c4, n_c128, n_swa = 5, 4, 7
280 - _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
281 - all_names = set()
282 - for t in tensors:
283 - all_names.update(t.shared_by)
284 - expected = n_c4 * 2 + n_c128 + n_swa
285 - assert len(all_names) == expected
286 -
287 - def test_group_owned_blocks_do_not_alias(self):
288 - groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
289 - num_blocks, tensors = _get_kv_cache_config_packed(
290 - _mock_vllm_config(), groups, 100 * 1024 * 1024
291 - )
292 - views = _make_views(groups, num_blocks, tensors)
293 -
294 - expected = {}
295 - value = 1
296 - for block_id, group in enumerate(groups):
297 - for layer_name in group.layer_names:
298 - views[layer_name][block_id].fill_(value)
299 - expected[layer_name] = (block_id, value)
300 - value += 1
301 -
302 - for layer_name, (block_id, value) in expected.items():
303 - assert (views[layer_name][block_id] == value).all()
304 -
305 - # Once the first group releases its block, another group may reuse it.
306 - for layer_name in groups[1].layer_names:
307 - views[layer_name][0].fill_(255)
308 - for layer_name in groups[1].layer_names:
309 - assert (views[layer_name][0] == 255).all()
310 -
311 - def test_hma_attention_groups_keep_default_backing(self):
312 - full = _make_full_spec()
313 - sw = _make_sw_spec()
314 - page_size = full.page_size_bytes
315 - groups = [
316 - KVCacheGroupSpec(["full.0", "full.1"], full),
317 - KVCacheGroupSpec(["sw.0", "sw.2"], sw),
318 - KVCacheGroupSpec(["sw.1", "sw.3"], sw),
319 - ]
320 -
119 + expected = 0
120 + for name in group.layer_names:
121 + assert offsets[name] == expected
122 + expected += pages[name]
123 +
124 + @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"])
125 + def test_allocation_is_layout_invariant(self, layout):
126 + specs = {f"l.{i}": _full() for i in range(4)}
127 + groups = [KVCacheGroupSpec(list(specs), _full())]
128 config = get_kv_cache_config_from_groups(
322 - _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
129 + _mock_vllm_config(layout), groups, MEMORY
130 )
131 + (tensor,) = config.kv_cache_tensors
132 + page = _full().page_size_bytes
133 + assert tensor.layers == list(specs)
134 + assert config.num_blocks == MEMORY // (4 * page)
135 + assert tensor.size == 4 * page * config.num_blocks
136 + if layout == "LBNHC":
137 + assert (tensor.layer_stride, tensor.block_stride) == (
138 + page * config.num_blocks,
139 + page,
140 + )
141 + else:
142 + assert (tensor.layer_stride, tensor.block_stride) == (page, 4 * page)
143
325 - assert config.num_blocks == 32
326 - assert config.prefix_cache_retention_interval == 0
327 - assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
328 - assert config.kv_cache_tensors == [
329 - KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
330 - KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
331 - ]
332 -
333 - def test_hma_attention_groups_use_packed_backing_with_enable_cross_layers(self):
334 - full = _make_full_spec()
335 - sw = _make_sw_spec()
336 - page_size = full.page_size_bytes
337 - groups = [
338 - KVCacheGroupSpec(["full.0", "full.1"], full),
339 - KVCacheGroupSpec(["sw.0", "sw.2"], sw),
340 - KVCacheGroupSpec(["sw.1", "sw.3"], sw),
341 - ]
342 -
144 + @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"])
145 + def test_single_group_mixed_pages_follows_layout(self, layout):
146 + specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)}
147 + groups = [_uniform_group(specs)]
148 config = get_kv_cache_config_from_groups(
344 - _mock_vllm_config({"enable_cross_layers_blocks": "True"}),
345 - groups,
346 - available_memory=page_size * 2 * 32,
149 + _mock_vllm_config(layout), groups, MEMORY
150 )
348 -
349 - assert config.num_blocks == 32
350 - assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
351 - assert config.kv_cache_tensors == [
352 - KVCacheTensor(
353 - size=page_size * 2 * 32,
354 - shared_by=["full.0", "sw.0", "sw.1"],
355 - offset=0,
356 - block_stride=page_size * 2,
357 - ),
358 - KVCacheTensor(
359 - size=page_size * 2 * 32,
360 - shared_by=["full.1", "sw.2", "sw.3"],
361 - offset=page_size,
362 - block_stride=page_size * 2,
363 - ),
364 - ]
365 -
366 - def test_single_group_attention_keeps_unpacked_layout(self):
367 - spec = _make_full_spec()
368 - groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
369 -
151 + block_stride = _expected_bytes_per_block(groups)
152 + mla_tensor, idx_tensor = config.kv_cache_tensors
153 + assert mla_tensor.layers == ["mla.0", "mla.1"]
154 + assert idx_tensor.layers == ["idx.0"]
155 + assert {t.size for t in config.kv_cache_tensors} == {
156 + block_stride * config.num_blocks
157 + }
158 + if layout == "LBNHC":
159 + assert (
160 + idx_tensor.offset == 2 * _mla(512).page_size_bytes * config.num_blocks
161 + )
162 + assert mla_tensor.block_stride == _mla(512).page_size_bytes
163 + else:
164 + assert idx_tensor.offset == 2 * _mla(512).page_size_bytes
165 + assert mla_tensor.block_stride == block_stride
166 +
167 + def test_overlaid_groups_alias_and_stay_isolated(self):
168 + groups, g1, g2 = _mixed_page_groups()
169 + # Overlay models resolve to a block-outer layout at backend selection (the
170 + # model's backend declares it); mirror that here.
171 config = get_kv_cache_config_from_groups(
371 - _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
172 + _mock_vllm_config("BLNHC"), groups, MEMORY
173 )
373 -
374 - assert sum(t.size for t in config.kv_cache_tensors) == (
375 - spec.page_size_bytes * 2 * 32
174 + assert config.num_blocks == MEMORY // _expected_bytes_per_block(groups)
175 + assert _pool_bytes_per_block(groups) == _expected_bytes_per_block(groups)
176 +
177 + views = _bind(config, "BLNHC")
178 + assert set(views) == set(g1) | set(g2)
179 + assert views["swa.0"].data_ptr() == views["mla.0"].data_ptr()
180 +
181 + # A block is owned by one group at a time: writes to group-owned
182 + # blocks never disturb the other group's blocks.
183 + for i, name in enumerate(g1):
184 + views[name][0].fill_(i + 1)
185 + views[name][2].fill_(i + 1)
186 + for i, name in enumerate(g2):
187 + views[name][1].fill_(100 + i)
188 + views[name][3].fill_(100 + i)
189 + for i, name in enumerate(g1):
190 + assert (views[name][0].to(torch.int32) == i + 1).all()
191 + assert (views[name][2].to(torch.int32) == i + 1).all()
192 + for i, name in enumerate(g2):
193 + assert (views[name][1].to(torch.int32) == 100 + i).all()
194 + assert (views[name][3].to(torch.int32) == 100 + i).all()
195 + # Layers within a group are disjoint.
196 + views["mla.0"][0].fill_(77)
197 + for i, name in enumerate(list(g1)[1:], start=1):
198 + assert (views[name][0].to(torch.int32) == i + 1).all()
199 +
200 + def test_layer_compact_layout_rejected_for_overlaid_groups(self):
201 + # The layout has a single writer (backend-selection resolution); a layer-
202 + # compact layout reaching an overlay model's allocation is an error, not a
203 + # silent flip.
204 + groups, _, _ = _mixed_page_groups()
205 + with pytest.raises(
206 + ValueError, match="cannot express this model's mixed page sizes"
207 + ):
208 + get_kv_cache_config_from_groups(_mock_vllm_config("LBNHC"), groups, MEMORY)
209 +
210 + def test_unresolved_layout_rejected(self):
211 + groups, _, _ = _mixed_page_groups()
212 + with pytest.raises(ValueError, match="has not been resolved"):
213 + get_kv_cache_config_from_groups(_mock_vllm_config(None), groups, MEMORY)
214 +
215 + def test_head_outer_layout_rejected_for_mixed_pages(self):
216 + groups, _, _ = _mixed_page_groups()
217 + with pytest.raises(
218 + ValueError, match="cannot express this model's mixed page sizes"
219 + ):
220 + get_kv_cache_config_from_groups(_mock_vllm_config("LHBNC"), groups, MEMORY)
221 +
222 + @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"])
223 + def test_bound_views_round_trip(self, layout):
224 + specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)}
225 + groups = [_uniform_group(specs)]
226 + config = get_kv_cache_config_from_groups(
227 + _mock_vllm_config(layout), groups, MEMORY
228 )
377 - assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
229 + views = _bind(config, layout)
230 + for i, name in enumerate(specs):
231 + views[name].fill_(i + 1)
232 + for i, name in enumerate(specs):
233 + assert (views[name].to(torch.int32) == i + 1).all(), name
234 + assert views[name].shape[0] == config.num_blocks
235
236
237 if __name__ == "__main__":
tests/v1/core/test_kv_cache_utils.py
+300 -288
@@ -124,7 +124,6 @@ def new_kv_cache_spec(
124 page_size_padded=None,
125 sliding_window=None,
126 attention_chunk_size=None,
127 - indexes_kv_by_block_stride=False,
127 kv_quant_mode=KVQuantMode.NONE,
128 ):
129 return FullAttentionSpec(
@@ -135,7 +134,6 @@ def new_kv_cache_spec(
134 page_size_padded=page_size_padded,
135 sliding_window=sliding_window,
136 attention_chunk_size=attention_chunk_size,
138 - indexes_kv_by_block_stride=indexes_kv_by_block_stride,
137 kv_quant_mode=kv_quant_mode,
138 )
139
@@ -147,7 +145,6 @@ def new_sliding_window_spec(
145 dtype=torch.float32,
146 page_size_padded=None,
147 sliding_window=1,
150 - indexes_kv_by_block_stride=False,
148 ):
149 return SlidingWindowSpec(
150 block_size=block_size,
@@ -156,7 +153,6 @@ def new_sliding_window_spec(
153 dtype=dtype,
154 page_size_padded=page_size_padded,
155 sliding_window=sliding_window,
159 - indexes_kv_by_block_stride=indexes_kv_by_block_stride,
156 )
157
158
@@ -875,6 +871,7 @@ def test_metrics_empty_stats():
871 def test_get_kv_cache_configs_multiple_workers():
872 model_config = ModelConfig(max_model_len=16)
873 vllm_config = VllmConfig(model_config=model_config)
874 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
875 vllm_config.cache_config.prefix_cache_retention_interval = None
876
877 ref_kv_cache_spec = new_kv_cache_spec()
@@ -898,36 +895,21 @@ def test_get_kv_cache_configs_multiple_workers():
895 ref_kv_cache_spec.page_size_bytes * 2 * 10,
896 ],
897 )
901 - assert kv_cache_configs == [
902 - KVCacheConfig(
903 - num_blocks=10,
904 - kv_cache_tensors=[
905 - KVCacheTensor(
906 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
907 - ),
908 - KVCacheTensor(
909 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
910 - ),
911 - ],
912 - kv_cache_groups=[
913 - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
914 - ],
915 - ),
916 - KVCacheConfig(
917 - num_blocks=10,
918 - kv_cache_tensors=[
919 - KVCacheTensor(
920 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
921 - ),
922 - KVCacheTensor(
923 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
924 - ),
925 - ],
926 - kv_cache_groups=[
927 - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
928 - ],
929 - ),
930 - ]
898 + expected = KVCacheConfig(
899 + num_blocks=10,
900 + kv_cache_tensors=[
901 + KVCacheTensor(
902 + size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
903 + layers=["layer1", "layer2"],
904 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
905 + block_stride=ref_kv_cache_spec.page_size_bytes,
906 + ),
907 + ],
908 + kv_cache_groups=[
909 + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
910 + ],
911 + )
912 + assert kv_cache_configs == [expected, expected]
913
914 # Different available memory. This is the case for TP.
915 # Use the smallest memory available.
@@ -939,36 +921,7 @@ def test_get_kv_cache_configs_multiple_workers():
921 ref_kv_cache_spec.page_size_bytes * 2 * 20,
922 ],
923 )
942 - assert kv_cache_configs == [
943 - KVCacheConfig(
944 - num_blocks=10,
945 - kv_cache_tensors=[
946 - KVCacheTensor(
947 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
948 - ),
949 - KVCacheTensor(
950 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
951 - ),
952 - ],
953 - kv_cache_groups=[
954 - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
955 - ],
956 - ),
957 - KVCacheConfig(
958 - num_blocks=10,
959 - kv_cache_tensors=[
960 - KVCacheTensor(
961 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
962 - ),
963 - KVCacheTensor(
964 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
965 - ),
966 - ],
967 - kv_cache_groups=[
968 - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
969 - ],
970 - ),
971 - ]
924 + assert kv_cache_configs == [expected, expected]
925
926 # Different KV cache specs. This is the case for PP.
927 different_layer_specs = [
@@ -995,7 +948,10 @@ def test_get_kv_cache_configs_multiple_workers():
948 num_blocks=10,
949 kv_cache_tensors=[
950 KVCacheTensor(
998 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
951 + size=ref_kv_cache_spec.page_size_bytes * 10,
952 + layers=["layer1"],
953 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
954 + block_stride=ref_kv_cache_spec.page_size_bytes,
955 ),
956 ],
957 kv_cache_groups=[
@@ -1006,10 +962,10 @@ def test_get_kv_cache_configs_multiple_workers():
962 num_blocks=10,
963 kv_cache_tensors=[
964 KVCacheTensor(
1009 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
1010 - ),
1011 - KVCacheTensor(
1012 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
965 + size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
966 + layers=["layer2", "layer3"],
967 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
968 + block_stride=ref_kv_cache_spec.page_size_bytes,
969 ),
970 ],
971 kv_cache_groups=[
@@ -1046,57 +1002,39 @@ def test_get_kv_cache_configs_multiple_workers():
1002 ref_kv_cache_spec.page_size_bytes * 2 * 10,
1003 ],
1004 )
1005 + expected_12 = KVCacheConfig(
1006 + num_blocks=10,
1007 + kv_cache_tensors=[
1008 + KVCacheTensor(
1009 + size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1010 + layers=["layer1", "layer2"],
1011 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1012 + block_stride=ref_kv_cache_spec.page_size_bytes,
1013 + ),
1014 + ],
1015 + kv_cache_groups=[
1016 + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
1017 + ],
1018 + )
1019 + expected_3 = KVCacheConfig(
1020 + num_blocks=10,
1021 + kv_cache_tensors=[
1022 + KVCacheTensor(
1023 + size=ref_kv_cache_spec.page_size_bytes * 10,
1024 + layers=["layer3"],
1025 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1026 + block_stride=ref_kv_cache_spec.page_size_bytes,
1027 + ),
1028 + ],
1029 + kv_cache_groups=[
1030 + KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
1031 + ],
1032 + )
1033 assert kv_cache_configs == [
1050 - KVCacheConfig(
1051 - num_blocks=10,
1052 - kv_cache_tensors=[
1053 - KVCacheTensor(
1054 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
1055 - ),
1056 - KVCacheTensor(
1057 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
1058 - ),
1059 - ],
1060 - kv_cache_groups=[
1061 - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
1062 - ],
1063 - ),
1064 - KVCacheConfig(
1065 - num_blocks=10,
1066 - kv_cache_tensors=[
1067 - KVCacheTensor(
1068 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
1069 - ),
1070 - KVCacheTensor(
1071 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
1072 - ),
1073 - ],
1074 - kv_cache_groups=[
1075 - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
1076 - ],
1077 - ),
1078 - KVCacheConfig(
1079 - num_blocks=10,
1080 - kv_cache_tensors=[
1081 - KVCacheTensor(
1082 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
1083 - ),
1084 - ],
1085 - kv_cache_groups=[
1086 - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
1087 - ],
1088 - ),
1089 - KVCacheConfig(
1090 - num_blocks=10,
1091 - kv_cache_tensors=[
1092 - KVCacheTensor(
1093 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
1094 - ),
1095 - ],
1096 - kv_cache_groups=[
1097 - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
1098 - ],
1099 - ),
1034 + expected_12,
1035 + expected_12,
1036 + expected_3,
1037 + expected_3,
1038 ]
1039
1040 # Different workers have different types of layers. This is the case for
@@ -1124,10 +1062,10 @@ def test_get_kv_cache_configs_multiple_workers():
1062 num_blocks=10,
1063 kv_cache_tensors=[
1064 KVCacheTensor(
1127 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
1128 - ),
1129 - KVCacheTensor(
1130 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
1065 + size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1066 + layers=["layer1", "layer2"],
1067 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1068 + block_stride=ref_kv_cache_spec.page_size_bytes,
1069 ),
1070 ],
1071 kv_cache_groups=[
@@ -1139,10 +1077,10 @@ def test_get_kv_cache_configs_multiple_workers():
1077 num_blocks=10,
1078 kv_cache_tensors=[
1079 KVCacheTensor(
1142 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
1143 - ),
1144 - KVCacheTensor(
1145 - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer4"]
1080 + size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
1081 + layers=["layer3", "layer4"],
1082 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1083 + block_stride=ref_kv_cache_spec.page_size_bytes,
1084 ),
1085 ],
1086 kv_cache_groups=[
@@ -1180,7 +1118,21 @@ def test_get_kv_cache_configs_multiple_workers():
1118 kv_cache_tensors=[
1119 KVCacheTensor(
1120 size=ref_kv_cache_spec.page_size_bytes * 10,
1183 - shared_by=["layer1", "layer2", "layer3"],
1121 + layers=["layer1"],
1122 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1123 + block_stride=ref_kv_cache_spec.page_size_bytes,
1124 + ),
1125 + KVCacheTensor(
1126 + size=ref_kv_cache_spec.page_size_bytes * 10,
1127 + layers=["layer2"],
1128 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1129 + block_stride=ref_kv_cache_spec.page_size_bytes,
1130 + ),
1131 + KVCacheTensor(
1132 + size=ref_kv_cache_spec.page_size_bytes * 10,
1133 + layers=["layer3"],
1134 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1135 + block_stride=ref_kv_cache_spec.page_size_bytes,
1136 ),
1137 ],
1138 kv_cache_groups=[
@@ -1194,7 +1146,21 @@ def test_get_kv_cache_configs_multiple_workers():
1146 kv_cache_tensors=[
1147 KVCacheTensor(
1148 size=ref_kv_cache_spec.page_size_bytes * 10,
1197 - shared_by=["layer4", "layer5", "layer6"],
1149 + layers=["layer4"],
1150 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1151 + block_stride=ref_kv_cache_spec.page_size_bytes,
1152 + ),
1153 + KVCacheTensor(
1154 + size=ref_kv_cache_spec.page_size_bytes * 10,
1155 + layers=["layer5"],
1156 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1157 + block_stride=ref_kv_cache_spec.page_size_bytes,
1158 + ),
1159 + KVCacheTensor(
1160 + size=ref_kv_cache_spec.page_size_bytes * 10,
1161 + layers=["layer6"],
1162 + layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
1163 + block_stride=ref_kv_cache_spec.page_size_bytes,
1164 ),
1165 ],
1166 kv_cache_groups=[
@@ -1233,6 +1199,7 @@ def test_get_kv_cache_configs_multiple_workers():
1199 def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
1200 model_config = ModelConfig(max_model_len=512)
1201 vllm_config = VllmConfig(model_config=model_config)
1202 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
1203 vllm_config.cache_config.prefix_cache_retention_interval = None
1204
1205 ref_kv_cache_spec = new_kv_cache_spec()
@@ -1262,7 +1229,10 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
1229 kv_cache_tensors=[
1230 KVCacheTensor(
1231 size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
1265 - shared_by=["layer1"],
1232 + layers=["layer1"],
1233 + layer_stride=ref_kv_cache_spec.page_size_bytes
1234 + * expected_num_blocks,
1235 + block_stride=ref_kv_cache_spec.page_size_bytes,
1236 ),
1237 ],
1238 kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)],
@@ -1272,7 +1242,10 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
1242 kv_cache_tensors=[
1243 KVCacheTensor(
1244 size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
1275 - shared_by=["layer2"],
1245 + layers=["layer2"],
1246 + layer_stride=ref_kv_cache_spec.page_size_bytes
1247 + * expected_num_blocks,
1248 + block_stride=ref_kv_cache_spec.page_size_bytes,
1249 ),
1250 ],
1251 kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)],
@@ -1628,78 +1601,18 @@ def test_get_max_concurrency_for_kv_cache_config():
1601 )
1602
1603
1631 -def test_get_max_concurrency_packed_kv_cache_config():
1632 - from vllm.v1.core.kv_cache_utils import (
1633 - _get_kv_cache_config_packed,
1634 - _use_packed_kv_cache_config,
1635 - )
1636 -
1637 - model_config = ModelConfig(
1638 - "Qwen/Qwen1.5-7B",
1639 - runner="generate",
1640 - dtype="float16",
1641 - max_model_len=16384,
1642 - )
1643 - scheduler_config = SchedulerConfig(
1644 - max_num_batched_tokens=1024,
1645 - enable_chunked_prefill=True,
1646 - max_model_len=model_config.max_model_len,
1647 - is_encoder_decoder=model_config.is_encoder_decoder,
1648 - async_scheduling=False,
1649 - )
1650 - vllm_config = VllmConfig(
1651 - model_config=model_config,
1652 - scheduler_config=scheduler_config,
1653 - )
1654 -
1655 - # All-UniformTypeKVCacheSpecs groups select the packed layout.
1656 - mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
1657 - swa_specs = {
1658 - f"layer_{i}": SlidingWindowMLASpec(
1659 - block_size=16,
1660 - num_kv_heads=1,
1661 - head_size=576,
1662 - dtype=torch.float32,
1663 - sliding_window=128,
1664 - )
1665 - for i in range(4, 6)
1666 - }
1667 - kv_cache_groups = [
1668 - KVCacheGroupSpec(
1669 - list(mla_specs),
1670 - UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
1671 - ),
1672 - KVCacheGroupSpec(
1673 - list(swa_specs),
1674 - UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
1675 - ),
1676 - ]
1677 - assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
1678 - num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
1679 - vllm_config, kv_cache_groups, 2 * GiB_bytes
1680 - )
1681 - assert num_blocks > 0
1682 - kv_cache_config_packed = KVCacheConfig(
1683 - num_blocks=num_blocks,
1684 - kv_cache_tensors=kv_cache_tensors,
1685 - kv_cache_groups=kv_cache_groups,
1686 - )
1687 - # Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
1688 - # the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
1689 - # previous formula normalized by the first group's page size and gave
1690 - # 1061 blocks per request instead of 1097.
1691 - assert get_max_concurrency_for_kv_cache_config(
1692 - vllm_config, kv_cache_config_packed
1693 - ) == num_blocks / (1024 + 73)
1694 -
1695 -
1604 def test_allocate_with_lookahead():
1605 """Verify that lookahead tokens correctly affect block allocation"""
1606 block_size = 4
1607 config = KVCacheConfig(
1608 num_blocks=10,
1609 kv_cache_tensors=[
1702 - KVCacheTensor(size=100, shared_by=["layer1"]),
1610 + KVCacheTensor(
1611 + size=100,
1612 + layers=["layer1"],
1613 + layer_stride=100,
1614 + block_stride=10,
1615 + ),
1616 ],
1617 kv_cache_groups=[
1618 KVCacheGroupSpec(["layer1"], new_kv_cache_spec(block_size=block_size)),
@@ -1763,6 +1676,7 @@ def test_get_kv_cache_config_one_worker():
1676 # pass max_model_len to pass check_enough_kv_cache_memory
1677 model_config = ModelConfig(max_model_len=16)
1678 vllm_config = VllmConfig(model_config=model_config)
1679 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
1680 vllm_config.cache_config.prefix_cache_retention_interval = None
1681
1682 mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
@@ -1778,8 +1692,12 @@ def test_get_kv_cache_config_one_worker():
1692 assert kv_cache_config_full == KVCacheConfig(
1693 num_blocks=32,
1694 kv_cache_tensors=[
1781 - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
1782 - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
1695 + KVCacheTensor(
1696 + size=mem_per_block_per_layer * 32 * 2,
1697 + layers=["layer_1", "layer_2"],
1698 + layer_stride=mem_per_block_per_layer * 32,
1699 + block_stride=mem_per_block_per_layer,
1700 + ),
1701 ],
1702 kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
1703 )
@@ -1795,8 +1713,12 @@ def test_get_kv_cache_config_one_worker():
1713 assert kv_cache_config_sliding == KVCacheConfig(
1714 num_blocks=32,
1715 kv_cache_tensors=[
1798 - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
1799 - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
1716 + KVCacheTensor(
1717 + size=mem_per_block_per_layer * 32 * 2,
1718 + layers=["layer_1", "layer_2"],
1719 + layer_stride=mem_per_block_per_layer * 32,
1720 + block_stride=mem_per_block_per_layer,
1721 + ),
1722 ],
1723 kv_cache_groups=[
1724 KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec())
@@ -1815,8 +1737,12 @@ def test_get_kv_cache_config_one_worker():
1737 assert kv_cache_config_hybrid == KVCacheConfig(
1738 num_blocks=32,
1739 kv_cache_tensors=[
1818 - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
1819 - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
1740 + KVCacheTensor(
1741 + size=mem_per_block_per_layer * 32 * 2,
1742 + layers=["layer_1", "layer_2"],
1743 + layer_stride=mem_per_block_per_layer * 32,
1744 + block_stride=mem_per_block_per_layer,
1745 + ),
1746 ],
1747 kv_cache_groups=[
1748 KVCacheGroupSpec(
@@ -1838,7 +1764,16 @@ def test_get_kv_cache_config_one_worker():
1764 num_blocks=64,
1765 kv_cache_tensors=[
1766 KVCacheTensor(
1841 - size=mem_per_block_per_layer * 64, shared_by=["layer_1", "layer_2"]
1767 + size=mem_per_block_per_layer * 64,
1768 + layers=["layer_1"],
1769 + layer_stride=mem_per_block_per_layer * 64,
1770 + block_stride=mem_per_block_per_layer,
1771 + ),
1772 + KVCacheTensor(
1773 + size=mem_per_block_per_layer * 64,
1774 + layers=["layer_2"],
1775 + layer_stride=mem_per_block_per_layer * 64,
1776 + block_stride=mem_per_block_per_layer,
1777 ),
1778 ],
1779 kv_cache_groups=[
@@ -1863,12 +1798,22 @@ def test_get_kv_cache_config_one_worker():
1798 num_blocks=32,
1799 kv_cache_tensors=[
1800 KVCacheTensor(
1866 - size=mem_per_block_per_layer * 32,
1867 - shared_by=["layer_1", "layer_3", "layer_4"],
1801 + size=mem_per_block_per_layer * 32 * 2,
1802 + layers=["layer_1", "layer_2"],
1803 + layer_stride=mem_per_block_per_layer * 32,
1804 + block_stride=mem_per_block_per_layer,
1805 ),
1806 KVCacheTensor(
1870 - size=mem_per_block_per_layer * 32,
1871 - shared_by=["layer_2", "layer_5", "layer_6"],
1807 + size=mem_per_block_per_layer * 32 * 2,
1808 + layers=["layer_3", "layer_5"],
1809 + layer_stride=mem_per_block_per_layer * 32,
1810 + block_stride=mem_per_block_per_layer,
1811 + ),
1812 + KVCacheTensor(
1813 + size=mem_per_block_per_layer * 32 * 2,
1814 + layers=["layer_4", "layer_6"],
1815 + layer_stride=mem_per_block_per_layer * 32,
1816 + block_stride=mem_per_block_per_layer,
1817 ),
1818 ],
1819 kv_cache_groups=[
@@ -1898,15 +1843,28 @@ def test_get_kv_cache_config_one_worker():
1843 num_blocks=32,
1844 kv_cache_tensors=[
1845 KVCacheTensor(
1901 - size=mem_per_block_per_layer * 32,
1902 - shared_by=["layer_1", "layer_4", "layer_5", "layer_6"],
1846 + size=mem_per_block_per_layer * 32 * 3,
1847 + layers=["layer_1", "layer_2", "layer_3"],
1848 + layer_stride=mem_per_block_per_layer * 32,
1849 + block_stride=mem_per_block_per_layer,
1850 ),
1851 KVCacheTensor(
1905 - size=mem_per_block_per_layer * 32,
1906 - shared_by=["layer_2", "layer_7", "layer_8", "layer_9"],
1852 + size=mem_per_block_per_layer * 32 * 3,
1853 + layers=["layer_4", "layer_7", "layer_10"],
1854 + layer_stride=mem_per_block_per_layer * 32,
1855 + block_stride=mem_per_block_per_layer,
1856 + ),
1857 + KVCacheTensor(
1858 + size=mem_per_block_per_layer * 32 * 3,
1859 + layers=["layer_5", "layer_8"],
1860 + layer_stride=mem_per_block_per_layer * 32,
1861 + block_stride=mem_per_block_per_layer,
1862 ),
1863 KVCacheTensor(
1909 - size=mem_per_block_per_layer * 32, shared_by=["layer_3", "layer_10"]
1864 + size=mem_per_block_per_layer * 32 * 3,
1865 + layers=["layer_6", "layer_9"],
1866 + layer_stride=mem_per_block_per_layer * 32,
1867 + block_stride=mem_per_block_per_layer,
1868 ),
1869 ],
1870 kv_cache_groups=[
@@ -1943,28 +1901,23 @@ def test_get_kv_cache_config_one_worker():
1901 num_blocks=32,
1902 kv_cache_tensors=[
1903 KVCacheTensor(
1946 - size=mem_per_block_per_layer * 32,
1947 - shared_by=["layer_1", "layer_7"],
1948 - ),
1949 - KVCacheTensor(
1950 - size=mem_per_block_per_layer * 32,
1951 - shared_by=["layer_2", "layer_8"],
1952 - ),
1953 - KVCacheTensor(
1954 - size=mem_per_block_per_layer * 32,
1955 - shared_by=["layer_3", "layer_9"],
1956 - ),
1957 - KVCacheTensor(
1958 - size=mem_per_block_per_layer * 32,
1959 - shared_by=["layer_4", "layer_10"],
1960 - ),
1961 - KVCacheTensor(
1962 - size=mem_per_block_per_layer * 32,
1963 - shared_by=["layer_5", "layer_11"],
1904 + size=mem_per_block_per_layer * 32 * 6,
1905 + layers=[
1906 + "layer_1",
1907 + "layer_2",
1908 + "layer_3",
1909 + "layer_4",
1910 + "layer_5",
1911 + "layer_6",
1912 + ],
1913 + layer_stride=mem_per_block_per_layer * 32,
1914 + block_stride=mem_per_block_per_layer,
1915 ),
1916 KVCacheTensor(
1966 - size=mem_per_block_per_layer * 32,
1967 - shared_by=["layer_6"],
1917 + size=mem_per_block_per_layer * 32 * 6,
1918 + layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"],
1919 + layer_stride=mem_per_block_per_layer * 32,
1920 + block_stride=mem_per_block_per_layer,
1921 ),
1922 ],
1923 kv_cache_groups=[
@@ -1987,11 +1940,24 @@ def test_get_kv_cache_config_one_worker():
1940 kv_cache_config_hybrid = get_kv_cache_configs(
1941 vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32]
1942 )[0]
1943 + # Layers of different page sizes pack densely into one allocation: the
1944 + # 2x-larger layer_1 takes the first region, layer_2 the next.
1945 assert kv_cache_config_hybrid == KVCacheConfig(
1946 num_blocks=32,
1947 kv_cache_tensors=[
1993 - KVCacheTensor(size=mem_per_block_per_layer * 32 * 2, shared_by=["layer_1"]),
1994 - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
1948 + KVCacheTensor(
1949 + size=mem_per_block_per_layer * 32 * 3,
1950 + layers=["layer_1"],
1951 + layer_stride=mem_per_block_per_layer * 2 * 32,
1952 + block_stride=mem_per_block_per_layer * 2,
1953 + ),
1954 + KVCacheTensor(
1955 + size=mem_per_block_per_layer * 32 * 3,
1956 + layers=["layer_2"],
1957 + layer_stride=mem_per_block_per_layer * 32,
1958 + block_stride=mem_per_block_per_layer,
1959 + offset=mem_per_block_per_layer * 2 * 32,
1960 + ),
1961 ],
1962 kv_cache_groups=[
1963 KVCacheGroupSpec(
@@ -2015,7 +1981,16 @@ def test_get_kv_cache_config_one_worker():
1981 num_blocks=32,
1982 kv_cache_tensors=[
1983 KVCacheTensor(
2018 - size=mem_per_block_per_layer * 32, shared_by=["layer_1", "layer_2"]
1984 + size=mem_per_block_per_layer * 32,
1985 + layers=["layer_1"],
1986 + layer_stride=mem_per_block_per_layer * 32,
1987 + block_stride=mem_per_block_per_layer,
1988 + ),
1989 + KVCacheTensor(
1990 + size=mem_per_block_per_layer * 32,
1991 + layers=["layer_2"],
1992 + layer_stride=mem_per_block_per_layer * 32,
1993 + block_stride=mem_per_block_per_layer,
1994 ),
1995 ],
1996 kv_cache_groups=[
@@ -2028,9 +2003,9 @@ def test_get_kv_cache_config_one_worker():
2003
2004 # different hidden size that cannot be aligned by using different block size,
2005 # but can be aligned by padding the smaller physical page.
2031 - swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True)
2006 + swa_spec = new_sliding_window_spec(head_size=96)
2007 kv_cache_specs_hybrid = {
2033 - "layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True),
2008 + "layer_1": new_kv_cache_spec(head_size=64),
2009 "layer_2": swa_spec,
2010 }
2011
@@ -2041,7 +2016,18 @@ def test_get_kv_cache_config_one_worker():
2016 assert kv_cache_config_hybrid == KVCacheConfig(
2017 num_blocks=42,
2018 kv_cache_tensors=[
2044 - KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]),
2019 + KVCacheTensor(
2020 + size=padded_page_size * 42,
2021 + layers=["layer_1"],
2022 + layer_stride=padded_page_size * 42,
2023 + block_stride=padded_page_size,
2024 + ),
2025 + KVCacheTensor(
2026 + size=padded_page_size * 42,
2027 + layers=["layer_2"],
2028 + layer_stride=padded_page_size * 42,
2029 + block_stride=padded_page_size,
2030 + ),
2031 ],
2032 kv_cache_groups=[
2033 KVCacheGroupSpec(
@@ -2049,12 +2035,11 @@ def test_get_kv_cache_config_one_worker():
2035 new_kv_cache_spec(
2036 head_size=64,
2037 page_size_padded=padded_page_size,
2052 - indexes_kv_by_block_stride=True,
2038 ),
2039 ),
2040 KVCacheGroupSpec(
2041 ["layer_2"],
2057 - new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True),
2042 + new_sliding_window_spec(head_size=96),
2043 ),
2044 ],
2045 )
@@ -2067,8 +2052,12 @@ def test_get_kv_cache_config_one_worker():
2052 assert kv_cache_config_override_blocks == KVCacheConfig(
2053 num_blocks=16,
2054 kv_cache_tensors=[
2070 - KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_1"]),
2071 - KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_2"]),
2055 + KVCacheTensor(
2056 + size=mem_per_block_per_layer * 16 * 2,
2057 + layers=["layer_1", "layer_2"],
2058 + layer_stride=mem_per_block_per_layer * 16,
2059 + block_stride=mem_per_block_per_layer,
2060 + ),
2061 ],
2062 kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
2063 )
@@ -2319,6 +2308,45 @@ def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity():
2308 ) == (544, 136)
2309
2310
2311 +def test_multi_run_layer_compact_strides_place_hoisted_heads():
2312 + """A layer-compact run region is its own dense allocation: under LHBNC the head
2313 + groups sit between the layers and the blocks, so a run's block stride is one head
2314 + group's slice, not the whole page (regression test for setStorage out-of-bounds
2315 + on ROCm hybrid models)."""
2316 + full = new_kv_cache_spec()
2317 + swa = new_sliding_window_spec(sliding_window=full.block_size * 2)
2318 + page = full.page_size_bytes
2319 + assert swa.page_size_bytes == page
2320 +
2321 + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
2322 + vllm_config.cache_config.kv_cache_layout = "LHBNC"
2323 + config = kv_cache_utils.get_kv_cache_config_from_groups(
2324 + vllm_config,
2325 + [
2326 + KVCacheGroupSpec(["full.0", "full.1"], full),
2327 + KVCacheGroupSpec(["swa.0", "swa.1"], swa),
2328 + ],
2329 + available_memory=4 * page * 8,
2330 + )
2331 +
2332 + num_blocks = config.num_blocks
2333 + head_group = full.block_size * full.state_content_size_bytes
2334 + for tensor in config.kv_cache_tensors:
2335 + # One head group per block, head groups between the layers and blocks.
2336 + assert tensor.block_stride == head_group
2337 + assert tensor.layer_stride == full.num_heads * num_blocks * head_group
2338 + # The last block of the last layer's last head group must stay within the
2339 + # allocation (the old page-based stride overran it).
2340 + end = (
2341 + tensor.offset
2342 + + tensor.layer_stride * (len(tensor.layers) - 1)
2343 + + (full.num_heads - 1) * num_blocks * head_group
2344 + + tensor.block_stride * (num_blocks - 1)
2345 + + head_group
2346 + )
2347 + assert end <= tensor.size
2348 +
2349 +
2350 def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified():
2351 specs = {
2352 "target.0.attn": new_mla_spec(),
@@ -2640,6 +2668,7 @@ def test_auto_fit_max_model_len():
2668 # Simulate the user passing -1 by setting original_max_model_len
2669 model_config.original_max_model_len = -1
2670 vllm_config = VllmConfig(model_config=model_config)
2671 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
2672
2673 mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer
2674 kv_cache_specs = {
@@ -2658,6 +2687,7 @@ def test_auto_fit_max_model_len():
2687 model_config = ModelConfig(max_model_len=1024)
2688 model_config.original_max_model_len = -1
2689 vllm_config = VllmConfig(model_config=model_config)
2690 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
2691
2692 # With limited memory, max_model_len should be reduced
2693 # Need memory for at least max_model_len tokens
@@ -2678,6 +2708,7 @@ def test_auto_fit_max_model_len_with_hybrid():
2708 # Simulate the user passing -1 by setting original_max_model_len
2709 model_config.original_max_model_len = -1
2710 vllm_config = VllmConfig(model_config=model_config)
2711 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
2712
2713 mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer
2714 gamma = 2
@@ -2700,6 +2731,7 @@ def test_auto_fit_max_model_len_not_triggered():
2731 model_config = ModelConfig(max_model_len=16)
2732 # original_max_model_len should be None by default, not -1
2733 vllm_config = VllmConfig(model_config=model_config)
2734 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
2735
2736 mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
2737 kv_cache_specs = {
@@ -2722,6 +2754,7 @@ def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
2754 model_config = ModelConfig(max_model_len=16384)
2755 model_config.original_max_model_len = -1 # request auto-fit
2756 vllm_config = VllmConfig(model_config=model_config)
2757 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
2758 # Cap the cache to 32 blocks regardless of available memory.
2759 vllm_config.cache_config.num_gpu_blocks_override = 32
2760
@@ -2747,6 +2780,7 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
2780 """
2781 model_config = ModelConfig(max_model_len=16384)
2782 vllm_config = VllmConfig(model_config=model_config)
2783 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
2784 # 32 blocks is far too small for max_model_len=16384 (would need 1024).
2785 vllm_config.cache_config.num_gpu_blocks_override = 32
2786
@@ -2769,14 +2803,13 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
2803 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer
2804 block-size multiple, so the smaller page must be padded instead.
2805 """
2772 - # Both layers' backends opt into the padded-page strided view (e.g.
2773 - # FlashAttention / its DiffKV subclass), so padding is allowed.
2806 + # Attention layers read padded pages through a strided view, so padding is
2807 + # allowed.
2808 target_spec = new_kv_cache_spec(
2809 block_size=16,
2810 num_kv_heads=1,
2811 head_size=192,
2812 dtype=torch.bfloat16,
2779 - indexes_kv_by_block_stride=True,
2813 )
2814 draft_spec = new_sliding_window_spec(
2815 block_size=16,
@@ -2784,7 +2817,6 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
2817 head_size=128,
2818 dtype=torch.bfloat16,
2819 sliding_window=1024,
2787 - indexes_kv_by_block_stride=True,
2820 )
2821
2822 unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size(
@@ -2802,35 +2834,6 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
2834 assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes
2835
2836
2805 -def test_unify_kv_cache_page_size_padding_requires_backend_support():
2806 - """Padding is gated on the backend declaring ``indexes_kv_by_block_stride``.
2807 -
2808 - A backend that does not support the strided padded-page view must raise
2809 - rather than silently padding (and misreading KV at runtime).
2810 - """
2811 - target_spec = new_kv_cache_spec(
2812 - block_size=16,
2813 - num_kv_heads=1,
2814 - head_size=192,
2815 - dtype=torch.bfloat16,
2816 - indexes_kv_by_block_stride=True,
2817 - )
2818 - # The non-divisible draft layer needs padding but its backend does not
2819 - # support the strided padded-page view -> must raise, not silently pad.
2820 - draft_spec = new_sliding_window_spec(
2821 - block_size=16,
2822 - num_kv_heads=1,
2823 - head_size=128,
2824 - dtype=torch.bfloat16,
2825 - sliding_window=1024,
2826 - indexes_kv_by_block_stride=False,
2827 - )
2828 - specs = {"target_attn": target_spec, "draft_attn": draft_spec}
2829 -
2830 - with pytest.raises(NotImplementedError):
2831 - kv_cache_utils.unify_kv_cache_spec_page_size(specs)
2832 -
2833 -
2837 def test_unpadded_page_size_includes_per_token_head_scales():
2838 # Per-token-head quant carries inline fp32 scales that are carved from the
2839 # raw KV allocation, so they must be budgeted into the offload width. The
@@ -2850,10 +2853,14 @@ def test_unpadded_page_size_includes_per_token_head_scales():
2853
2854 def test_page_size_padded_wins():
2855 # An explicit padded page size takes precedence over the unpadded size.
2853 - spec = new_kv_cache_spec(
2854 - dtype=torch.uint8,
2855 - kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
2856 - page_size_padded=65536,
2856 + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
2857 +
2858 + spec = TritonAttentionBackend.customize_spec(
2859 + new_kv_cache_spec(
2860 + dtype=torch.uint8,
2861 + kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
2862 + page_size_padded=65536,
2863 + )
2864 )
2865 assert spec.page_size_bytes == 65536
2866
@@ -2996,14 +3003,18 @@ def test_unify_kv_cache_spec_page_size_mamba():
3003 )
3004 assert unified["mamba_layer"].page_size_bytes == 32768
3005
2999 - # 4. Attention layers with non-divisible page sizes still raise.
3000 - with pytest.raises(NotImplementedError):
3001 - kv_cache_utils.unify_kv_cache_spec_page_size(
3002 - {
3003 - "attn_layer": new_kv_cache_spec(block_size=24), # 24576
3004 - "draft_attn_layer": draft_attn_spec, # 32768
3005 - }
3006 - )
3006 + # 4. Attention layers with non-divisible page sizes are padded too: every
3007 + # backend reads a padded page through the view's block stride, so there is
3008 + # no longer a case that must raise.
3009 + unified = kv_cache_utils.unify_kv_cache_spec_page_size(
3010 + {
3011 + "attn_layer": new_kv_cache_spec(block_size=24), # 24576
3012 + "draft_attn_layer": draft_attn_spec, # 32768
3013 + }
3014 + )
3015 + assert unified["attn_layer"].page_size_padded == 32768
3016 + assert unified["attn_layer"].page_size_bytes == 32768
3017 + assert unified["attn_layer"].block_size == 24
3018
3019 # 5. Uniform page sizes are returned unchanged.
3020 specs = {
@@ -3135,6 +3146,7 @@ def test_auto_fit_max_model_len_reserves_null_block():
3146 model_config = ModelConfig(max_model_len=1024)
3147 model_config.original_max_model_len = -1
3148 vllm_config = VllmConfig(model_config=model_config)
3149 + vllm_config.cache_config.kv_cache_layout = "LBNHC"
3150 spec = new_kv_cache_spec(block_size=block_size)
3151
3152 # Exactly the 1024 / 16 = 64 blocks a full-length request would need.
tests/v1/core/test_prefix_caching.py
+5 -5
@@ -154,7 +154,7 @@ def make_kv_cache_config_hybrid_model(
154 elif second_spec_type == "mamba":
155 second_spec = MambaSpec(
156 block_size=block_size,
157 - shapes=(1, 1),
157 + shapes=((1, 1),),
158 dtypes=(torch.float32,),
159 )
160
@@ -189,7 +189,7 @@ def make_kv_cache_config_three_types(
189 if third_spec_type == "mamba":
190 third_spec = MambaSpec(
191 block_size=block_size,
192 - shapes=(1, 1),
192 + shapes=((1, 1),),
193 dtypes=(torch.float32,),
194 )
195 elif third_spec_type == "sliding_window":
@@ -768,12 +768,12 @@ def _make_hybrid_kv_cache_config(
768 ),
769 "mamba": lambda: MambaSpec(
770 block_size=block_size,
771 - shapes=(1, 1),
771 + shapes=((1, 1),),
772 dtypes=(torch.float32,),
773 ),
774 "mamba_align": lambda: MambaSpec(
775 block_size=block_size,
776 - shapes=(1, 1),
776 + shapes=((1, 1),),
777 dtypes=(torch.float32,),
778 mamba_cache_mode="align",
779 ),
@@ -3273,7 +3273,7 @@ def test_hybrid_local_kv_retention_interval_survives_recycling():
3273 num_kv_heads=1,
3274 head_size=1,
3275 dtype=torch.uint8,
3276 - compress_ratio=4,
3276 + tokens_per_state=4,
3277 ),
3278 ),
3279 KVCacheGroupSpec(
tests/v1/engine/test_engine_core_client.py
+1
@@ -1329,6 +1329,7 @@ def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPat
1329 )
1330
1331 mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}]
1332 + mock_executor.get_supported_kv_cache_layouts.return_value = [["LBNHC"]]
1333 mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024]
1334 mock_executor.initialize_from_config.return_value = None
1335 mock_executor.supports_draft_weight_updates.return_value = False
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+12 -10
@@ -36,7 +36,9 @@ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
36 echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS"
37 fi
38
39 -DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD
39 +# Empty leaves the layout to the backends, for models that constrain it.
40 +PREFILLER_KV_LAYOUT=${PREFILLER_KV_LAYOUT-"HND"}
41 +DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT-"HND"} # Default to HND, optional NHD
42 if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then
43 KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"'
44 else
@@ -44,9 +46,9 @@ else
46 fi
47
48 if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
47 - KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}'
48 -else
49 - KV_EXTRA_CONFIG=''
49 + # Cross-layer blocks are exercised via the block-outermost layout.
50 + PREFILLER_KV_LAYOUT="BLHNC"
51 + DECODER_KV_LAYOUT="BLHNC"
52 fi
53
54 # Connector: default pull NixlConnector; NixlPushConnector enables PP prefill.
@@ -54,11 +56,11 @@ KV_CONNECTOR=${KV_CONNECTOR:-NixlConnector}
56
57 # Build the kv-transfer-config for P and D
58 if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
57 - KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
58 - KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
59 + KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}'}'
60 + KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}'}'
61 else
60 - KV_CONFIG_P="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
61 - KV_CONFIG_D="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
62 + KV_CONFIG_P="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}"
63 + KV_CONFIG_D="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}"
64 fi
65
66 # Models to run
@@ -176,7 +178,7 @@ run_tests_for_model() {
178
179 # Build the command with or without model-specific args
180 BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
179 - VLLM_KV_CACHE_LAYOUT='HND' \
181 + ${PREFILLER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$PREFILLER_KV_LAYOUT} \
182 VLLM_PORT=$INTERNAL_PORT \
183 UCX_NET_DEVICES=all \
184 VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
@@ -239,7 +241,7 @@ run_tests_for_model() {
241
242 # Build the command with or without model-specific args
243 BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
242 - VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
244 + ${DECODER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT} \
245 $DECODER_INTERNAL_PORT_ENV \
246 UCX_NET_DEVICES=all \
247 VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
+6 -23
@@ -9,8 +9,8 @@
9 # test_accuracy.py.
10 #
11 # By default runs two configurations:
12 -# 1. Normal KV layout (NixlConnector without cross-layer blocks)
13 -# 2. Cross-layer KV layout (NixlConnector with enable_cross_layers_blocks)
12 +# 1. Normal KV layout (LBHNC)
13 +# 2. Cross-layer KV layout (BLHNC) via VLLM_KV_CACHE_LAYOUT
14 #
15 # Usage:
16 # bash tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
@@ -55,8 +55,6 @@ fi
55
56 # ── KV transfer configs ─────────────────────────────────────────────────
57
58 -# Normal layout: OffloadingConnector prefers cross-layer but NixlConnector
59 -# does not, so MultiConnector.prefer_cross_layer_blocks = False.
58 KV_CONFIG_NORMAL='{
59 "kv_connector":"MultiConnector",
60 "kv_role":"kv_both",
@@ -72,22 +70,6 @@ KV_CONFIG_NORMAL='{
70 # Remove whitespace for CLI safety
71 KV_CONFIG_NORMAL=$(echo "$KV_CONFIG_NORMAL" | tr -d '[:space:]')
72
75 -# Cross-layer layout: both connectors prefer cross-layer blocks.
76 -KV_CONFIG_CROSS_LAYERS='{
77 - "kv_connector":"MultiConnector",
78 - "kv_role":"kv_both",
79 - "kv_connector_extra_config":{
80 - "connectors":[
81 - {"kv_connector":"NixlConnector","kv_role":"kv_both",
82 - "kv_buffer_device":"'$KV_BUFFER_DEVICE'",
83 - "kv_connector_extra_config":{"enable_cross_layers_blocks":"True"}},
84 - {"kv_connector":"OffloadingConnector","kv_role":"kv_both",
85 - "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}}
86 - ]
87 - }
88 -}'
89 -KV_CONFIG_CROSS_LAYERS=$(echo "$KV_CONFIG_CROSS_LAYERS" | tr -d '[:space:]')
90 -
73 # ── Helpers ──────────────────────────────────────────────────────────────
74
75 PROCESS_GROUPS=()
@@ -218,7 +200,7 @@ run_tests_for_model() {
200 # ── Start prefill instance ──
201 echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
202 BASE_CMD="$DEVICE_VISIBILITY_ENV=$PREFILL_GPU \
221 - VLLM_KV_CACHE_LAYOUT='HND' \
203 + VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
204 UCX_NET_DEVICES=all \
205 VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
206 vllm serve $model_name \
@@ -245,7 +227,7 @@ run_tests_for_model() {
227 # ── Start decode instance ──
228 echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
229 BASE_CMD="$DEVICE_VISIBILITY_ENV=$DECODE_GPU \
248 - VLLM_KV_CACHE_LAYOUT='HND' \
230 + VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
231 UCX_NET_DEVICES=all \
232 VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
233 vllm serve $model_name \
@@ -305,7 +287,8 @@ for model in "${MODELS[@]}"; do
287 fi
288
289 if [[ -z "${SKIP_CROSS_LAYERS:-}" ]]; then
308 - run_tests_for_model "$model" "$KV_CONFIG_CROSS_LAYERS" "MultiConnector cross-layer layout"
290 + VLLM_KV_CACHE_LAYOUT=BLHNC \
291 + run_tests_for_model "$model" "$KV_CONFIG_NORMAL" "MultiConnector cross-layer layout"
292 fi
293 done
294
tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+2 -2
@@ -95,7 +95,7 @@ run_tests_for_model() {
95 # ── Start prefill instance ──
96 echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
97 BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \
98 - VLLM_KV_CACHE_LAYOUT='HND' \
98 + VLLM_KV_CACHE_LAYOUT='LBHNC' \
99 UCX_NET_DEVICES=all \
100 VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
101 vllm serve \"$model_name\" \
@@ -121,7 +121,7 @@ run_tests_for_model() {
121 # ── Start decode instance ──
122 echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
123 BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \
124 - VLLM_KV_CACHE_LAYOUT='HND' \
124 + VLLM_KV_CACHE_LAYOUT='LBHNC' \
125 UCX_NET_DEVICES=all \
126 VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
127 vllm serve \"$model_name\" \
tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+2 -2
@@ -249,7 +249,7 @@ run_test_for_device() {
249 echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
250 env \
251 ${GPU_DEVICE_VAR}="$GPU_ID" \
252 - VLLM_KV_CACHE_LAYOUT='HND' \
252 + VLLM_KV_CACHE_LAYOUT='LBHNC' \
253 UCX_NET_DEVICES=all \
254 ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
255 VLLM_NIXL_SIDE_CHANNEL_HOST="$NIXL_SIDE_CHANNEL_HOST" \
@@ -288,7 +288,7 @@ run_test_for_device() {
288 echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
289 env \
290 ${GPU_DEVICE_VAR}="$GPU_ID" \
291 - VLLM_KV_CACHE_LAYOUT='HND' \
291 + VLLM_KV_CACHE_LAYOUT='LBHNC' \
292 UCX_NET_DEVICES=all \
293 ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
294 VLLM_NIXL_SIDE_CHANNEL_HOST="$NIXL_SIDE_CHANNEL_HOST" \
tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py
+12 -7
@@ -12,6 +12,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping i
12 _verify_tiling,
13 derive_canonical_mappings,
14 )
15 +from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
16 from vllm.v1.kv_cache_interface import (
17 FullAttentionSpec,
18 KVCacheGroupSpec,
@@ -37,12 +38,14 @@ def _ctx(rank, tp=1, dcp=1, pcp=1, interleave=1, total=None, heads=2):
38
39 def _full_spec(num_kv_heads: int = 2, **kwargs) -> FullAttentionSpec:
40 # block_size=4, head_dim=64, int8
40 - return FullAttentionSpec(
41 - block_size=4,
42 - num_kv_heads=num_kv_heads,
43 - head_size=64,
44 - dtype=torch.int8,
45 - **kwargs,
41 + return TritonAttentionBackend.customize_spec(
42 + FullAttentionSpec(
43 + block_size=4,
44 + num_kv_heads=num_kv_heads,
45 + head_size=64,
46 + dtype=torch.int8,
47 + **kwargs,
48 + )
49 )
50
51
@@ -381,7 +384,9 @@ def test_fail_closed_cases():
384 quant_spec = _full_spec(kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD)
385 assert _try_mapping(quant_spec, _split_nhd_cache(quant_spec), _ctx(0, tp=4)) is None
386 # Compressed MLA slots are not 1:1 with tokens
384 - assert _try_mapping(_mla_spec(compress_ratio=2), None, _ctx(0, tp=2, dcp=2)) is None
387 + assert (
388 + _try_mapping(_mla_spec(tokens_per_state=2), None, _ctx(0, tp=2, dcp=2)) is None
389 + )
390 # Unrecognized physical layouts
391 swapped = torch.zeros(
392 NUM_BLOCKS,
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+128 -99
@@ -72,43 +72,54 @@ def _make_vllm_config(
72
73 def _make_kv_cache_config() -> KVCacheConfig:
74 num_blocks = 16
75 + spec = FullAttentionSpec(
76 + block_size=16,
77 + num_kv_heads=1,
78 + head_size=1,
79 + dtype=torch.float32,
80 + )
81 kv_tensor = KVCacheTensor(
76 - size=num_blocks * 8,
77 - shared_by=["layer"],
78 - block_stride=0,
82 + size=spec.page_size_bytes * num_blocks,
83 + layers=["layer"],
84 + layer_stride=spec.page_size_bytes * num_blocks,
85 + block_stride=spec.page_size_bytes,
86 )
87 return KVCacheConfig(
88 num_blocks=num_blocks,
89 kv_cache_tensors=[kv_tensor],
83 - kv_cache_groups=[
84 - KVCacheGroupSpec(
85 - ["layer"],
86 - FullAttentionSpec(
87 - block_size=16,
88 - num_kv_heads=1,
89 - head_size=1,
90 - dtype=torch.float32,
91 - ),
92 - )
93 - ],
90 + kv_cache_groups=[KVCacheGroupSpec(["layer"], spec)],
91 )
92
93
94 def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
95 + """One 16 byte-per-block allocation, described two ways.
96 +
97 + Packed: both layers are one dense run. Unpacked: the same bytes as two runs, the
98 + second starting after the first layer's region. Either way the connector accounts
99 + for 16 KV bytes per block.
100 + """
101 num_blocks = 4
102 + page = 8
103 + size = 2 * page * num_blocks
104 if packed:
105 kv_cache_tensors = [
106 KVCacheTensor(
102 - size=64,
103 - shared_by=[layer_name],
104 - block_stride=16,
107 + size=size,
108 + layers=["layer0", "layer1"],
109 + layer_stride=page * num_blocks,
110 + block_stride=page,
111 )
106 - for layer_name in ("layer0", "layer1")
112 ]
113 else:
114 kv_cache_tensors = [
110 - KVCacheTensor(size=40, shared_by=["layer0"]),
111 - KVCacheTensor(size=24, shared_by=["layer1"]),
115 + KVCacheTensor(
116 + size=size,
117 + layers=[layer],
118 + layer_stride=page * num_blocks,
119 + block_stride=page,
120 + offset=i * page * num_blocks,
121 + )
122 + for i, layer in enumerate(("layer0", "layer1"))
123 ]
124
125 return KVCacheConfig(
@@ -150,6 +161,18 @@ def _mla_spec(
161 )
162
163
164 +_MAMBA_SPEC = MambaSpec(
165 + block_size=16,
166 + shapes=((16, 1),),
167 + dtypes=(torch.float32,),
168 +)
169 +# Page sizes of the specs the replicated-layout cases below are built from.
170 +_MLA_PAGE = _mla_spec().page_size_bytes
171 +_HALF_MLA_PAGE = _mla_spec(head_size=256).page_size_bytes
172 +_FULL_PAGE = _full_attention_spec().page_size_bytes
173 +_MAMBA_PAGE = _MAMBA_SPEC.page_size_bytes
174 +
175 +
176 def _make_mla_kv_cache_config(
177 layer_names: list[str] | None = None,
178 head_size: int = 512,
@@ -159,12 +182,14 @@ def _make_mla_kv_cache_config(
182 if layer_names is None:
183 layer_names = ["layer0", "layer1"]
184 spec = _mla_spec(head_size=head_size, dtype=dtype)
185 + layer_stride = spec.page_size_bytes * num_blocks
186 kv_cache_tensors = [
187 KVCacheTensor(
164 - size=spec.page_size_bytes * num_blocks,
165 - shared_by=[layer_name],
188 + size=layer_stride * len(layer_names),
189 + layers=layer_names,
190 + layer_stride=layer_stride,
191 + block_stride=spec.page_size_bytes,
192 )
167 - for layer_name in layer_names
193 ]
194 return KVCacheConfig(
195 num_blocks=num_blocks,
@@ -174,15 +199,31 @@ def _make_mla_kv_cache_config(
199
200
201 def _make_hybrid_kv_cache_config() -> KVCacheConfig:
202 + num_blocks = 4
203 + full_spec = _full_attention_spec(block_size=12)
204 + mla_spec = _mla_spec()
205 + # Mixed page sizes across overlaying groups: a block is a window of the largest
206 + # group's packing, so the layer dim sits inside the block dim.
207 + window = max(full_spec.page_size_bytes, mla_spec.page_size_bytes)
208 return KVCacheConfig(
178 - num_blocks=4,
209 + num_blocks=num_blocks,
210 kv_cache_tensors=[
180 - KVCacheTensor(size=40, shared_by=["full_layer"]),
181 - KVCacheTensor(size=24, shared_by=["mla_layer"]),
211 + KVCacheTensor(
212 + size=window * num_blocks,
213 + layers=["full_layer"],
214 + layer_stride=full_spec.page_size_bytes,
215 + block_stride=window,
216 + ),
217 + KVCacheTensor(
218 + size=window * num_blocks,
219 + layers=["mla_layer"],
220 + layer_stride=mla_spec.page_size_bytes,
221 + block_stride=window,
222 + ),
223 ],
224 kv_cache_groups=[
184 - KVCacheGroupSpec(["full_layer"], _full_attention_spec(block_size=12)),
185 - KVCacheGroupSpec(["mla_layer"], _mla_spec()),
225 + KVCacheGroupSpec(["full_layer"], full_spec),
226 + KVCacheGroupSpec(["mla_layer"], mla_spec),
227 ],
228 )
229
@@ -263,18 +304,9 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
304 assert offloading_config.cache.blocks_per_chunk == 2
305
306
266 -def test_rejects_partially_packed_tensor_layout():
267 - kv_cache_config = _make_sizing_kv_cache_config(packed=False)
268 - kv_cache_config.kv_cache_tensors[0].block_stride = 16
269 -
270 - with pytest.raises(AssertionError):
271 - build_offloading_config(_make_vllm_config(), kv_cache_config)
272 -
273 -
307 def test_zero_blocks_skips_tensor_layout_validation():
308 kv_cache_config = _make_sizing_kv_cache_config(packed=False)
309 kv_cache_config.num_blocks = 0
277 - kv_cache_config.kv_cache_tensors[0].block_stride = 16
310
311 offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config)
312
@@ -374,8 +406,10 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
406 num_blocks=4,
407 kv_cache_tensors=[
408 KVCacheTensor(
377 - size=_mla_spec().page_size_bytes * 4,
378 - shared_by=["layer"],
409 + size=_MLA_PAGE * 4,
410 + layers=["layer"],
411 + layer_stride=_MLA_PAGE * 4,
412 + block_stride=_MLA_PAGE,
413 )
414 ],
415 kv_cache_groups=[
@@ -398,8 +432,10 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
432 num_blocks=4,
433 kv_cache_tensors=[
434 KVCacheTensor(
401 - size=_mla_spec().page_size_bytes * 4,
402 - shared_by=["layer"],
435 + size=_MLA_PAGE * 4,
436 + layers=["layer"],
437 + layer_stride=_MLA_PAGE * 4,
438 + block_stride=_MLA_PAGE,
439 )
440 ],
441 kv_cache_groups=[
@@ -417,16 +453,22 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
453 "hidden-state",
454 ),
455 (
456 + # One group, two page sizes: layer1's run starts past layer0's region.
457 KVCacheConfig(
458 num_blocks=4,
459 kv_cache_tensors=[
460 KVCacheTensor(
424 - size=_mla_spec().page_size_bytes * 4,
425 - shared_by=["layer0"],
461 + size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
462 + layers=["layer0"],
463 + layer_stride=_MLA_PAGE * 4,
464 + block_stride=_MLA_PAGE,
465 ),
466 KVCacheTensor(
428 - size=_mla_spec(head_size=256).page_size_bytes * 4,
429 - shared_by=["layer1"],
467 + size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
468 + layers=["layer1"],
469 + layer_stride=_HALF_MLA_PAGE * 4,
470 + block_stride=_HALF_MLA_PAGE,
471 + offset=_MLA_PAGE * 4,
472 ),
473 ],
474 kv_cache_groups=[
@@ -445,16 +487,22 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
487 "uniform-wrapper",
488 ),
489 (
490 + # Overlaid groups with different page sizes: a block is a window of
491 + # the largest group's packing.
492 KVCacheConfig(
493 num_blocks=4,
494 kv_cache_tensors=[
495 KVCacheTensor(
452 - size=_mla_spec().page_size_bytes * 4,
453 - shared_by=["mla"],
496 + size=_FULL_PAGE * 4,
497 + layers=["mla"],
498 + layer_stride=_MLA_PAGE,
499 + block_stride=_FULL_PAGE,
500 ),
501 KVCacheTensor(
456 - size=_full_attention_spec().page_size_bytes * 4,
457 - shared_by=["full"],
502 + size=_FULL_PAGE * 4,
503 + layers=["full"],
504 + layer_stride=_FULL_PAGE,
505 + block_stride=_FULL_PAGE,
506 ),
507 ],
508 kv_cache_groups=[
@@ -469,21 +517,21 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
517 num_blocks=4,
518 kv_cache_tensors=[
519 KVCacheTensor(
472 - size=_mla_spec().page_size_bytes * 4,
473 - shared_by=["mla"],
520 + size=_MLA_PAGE * 4,
521 + layers=["mla"],
522 + layer_stride=_MLA_PAGE,
523 + block_stride=_MLA_PAGE,
524 + ),
525 + KVCacheTensor(
526 + size=_MLA_PAGE * 4,
527 + layers=["mamba"],
528 + layer_stride=_MAMBA_PAGE,
529 + block_stride=_MLA_PAGE,
530 ),
475 - KVCacheTensor(size=64 * 4, shared_by=["mamba"]),
531 ],
532 kv_cache_groups=[
533 KVCacheGroupSpec(["mla"], _mla_spec()),
479 - KVCacheGroupSpec(
480 - ["mamba"],
481 - MambaSpec(
482 - block_size=16,
483 - shapes=((16, 1),),
484 - dtypes=(torch.float32,),
485 - ),
486 - ),
534 + KVCacheGroupSpec(["mamba"], _MAMBA_SPEC),
535 ],
536 ),
537 "mla-mamba-hybrid",
@@ -493,13 +541,12 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
541 num_blocks=4,
542 kv_cache_tensors=[
543 KVCacheTensor(
496 - size=_mla_spec().page_size_bytes * 4,
497 - shared_by=["layer0"],
498 - ),
499 - KVCacheTensor(
500 - size=_mla_spec().page_size_bytes * 4,
501 - shared_by=["layer1"],
502 - ),
544 + size=_MLA_PAGE * 4,
545 + layers=[layer],
546 + layer_stride=_MLA_PAGE * 4,
547 + block_stride=_MLA_PAGE,
548 + )
549 + for layer in ("layer0", "layer1")
550 ],
551 kv_cache_groups=[
552 KVCacheGroupSpec(["layer0"], _mla_spec()),
@@ -531,16 +578,26 @@ def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
578 indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8)
579 main_layers = [f"main_{i}" for i in range(61)]
580 indexer_layers = [f"indexer_{i}" for i in range(61)]
581 + # A DSA-style group: the main pages and the smaller indexer pages are packed one
582 + # after the other, so a block holds more than 61 MLA pages.
583 + main_bytes = main_spec.page_size_bytes * len(main_layers)
584 + indexer_bytes = indexer_spec.page_size_bytes * len(indexer_layers)
585 + size = (main_bytes + indexer_bytes) * num_blocks
586 kv_cache_config = KVCacheConfig(
587 num_blocks=num_blocks,
588 kv_cache_tensors=[
589 KVCacheTensor(
538 - size=main_spec.page_size_bytes * len(main_layers) * num_blocks,
539 - shared_by=main_layers,
590 + size=size,
591 + layers=main_layers,
592 + layer_stride=main_spec.page_size_bytes * num_blocks,
593 + block_stride=main_spec.page_size_bytes,
594 ),
595 KVCacheTensor(
542 - size=indexer_spec.page_size_bytes * len(indexer_layers) * num_blocks,
543 - shared_by=indexer_layers,
596 + size=size,
597 + layers=indexer_layers,
598 + layer_stride=indexer_spec.page_size_bytes * num_blocks,
599 + block_stride=indexer_spec.page_size_bytes,
600 + offset=main_bytes * num_blocks,
601 ),
602 ],
603 kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)],
@@ -665,34 +722,6 @@ def test_canonical_layout_certifies_v2_model_runner():
722 ).parallel.is_parallelism_agnostic
723
724
668 -def test_prefer_cross_layer_blocks_yields_to_canonical_layout():
669 - """The connector must not request cross-layer blocks under
670 - canonical_layout: cross-layer slabs have no per-layer refs to certify."""
671 - from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
672 - from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
673 - OffloadingConnector,
674 - )
675 -
676 - kv_cache_config = _make_kv_cache_config()
677 - connector_module = "vllm.distributed.kv_transfer.kv_connector.v1"
678 -
679 - def make_connector(extra_config: dict[str, Any] | None) -> OffloadingConnector:
680 - with (
681 - patch(f"{connector_module}.offloading_connector.OffloadingSpecFactory"),
682 - patch(
683 - f"{connector_module}.offloading_connector.OffloadingConnectorScheduler"
684 - ),
685 - ):
686 - return OffloadingConnector(
687 - _make_vllm_config(extra_config=extra_config),
688 - KVConnectorRole.SCHEDULER,
689 - kv_cache_config,
690 - )
691 -
692 - assert make_connector(None).prefer_cross_layer_blocks
693 - assert not make_connector({"canonical_layout": True}).prefer_cross_layer_blocks
694 -
695 -
725 def test_parallelism_agnostic_disabled_on_v2_model_runner():
726 config = _make_vllm_config()
727 config.use_v2_model_runner = True
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+33 -52
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections import defaultdict
3 from unittest.mock import MagicMock
4
5 import pytest
@@ -13,11 +12,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
12 from vllm.platforms import current_platform
13 from vllm.utils.torch_utils import get_dtype_size
14 from vllm.v1.attention.backends.registry import AttentionBackendEnum
16 -from vllm.v1.attention.backends.utils import set_kv_cache_layout
15 from vllm.v1.kv_cache_interface import (
16 FullAttentionSpec,
17 KVCacheConfig,
18 KVCacheGroupSpec,
19 + KVCacheLayout,
20 KVCacheTensor,
21 MambaSpec,
22 MLAAttentionSpec,
@@ -65,40 +64,18 @@ elif current_platform.is_xpu():
64 # ---------------------------------------------------------------------------
65
66
68 -def _allocate_and_reshape_kv_caches(
67 +def _allocate_kv_caches(
68 kv_cache_config: KVCacheConfig,
69 attn_groups: list[list],
70 device: torch.device,
71 ):
73 - """
74 - Use the real GPUModelRunner allocation and reshape methods to produce
75 - kv_caches, just like the model runner does during initialization.
76 - """
77 - from vllm.v1.worker.gpu_model_runner import GPUModelRunner
78 -
79 - # Some backends (e.g. FlashAttention) query the KV cache layout during
80 - # reshape, which ultimately calls get_current_vllm_config(). Setting
81 - # the layout override avoids needing a full VllmConfig context.
82 - set_kv_cache_layout("NHD")
83 - try:
84 - runner = object.__new__(GPUModelRunner)
85 - runner.device = device
86 - runner.runner_only_attn_layers = set()
87 - runner.attn_groups = attn_groups
88 - runner.kv_cache_config = kv_cache_config
89 - runner.cache_config = MagicMock(cache_dtype="auto")
90 - runner.shared_kv_cache_layers = {}
91 - runner.model_config = MagicMock()
92 - runner.model_config.hf_config.model_type = ""
93 - runner.compilation_config = MagicMock(
94 - static_forward_context=defaultdict(MagicMock)
95 - )
96 - runner.kv_caches = []
72 + """Allocate kv_caches exactly as the model runner does at startup."""
73 + from vllm.v1.worker.utils import allocate_kv_cache
74
98 - kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
99 - return runner.initialize_kv_cache_tensors(kv_cache_config, kernel_block_sizes)
100 - finally:
101 - set_kv_cache_layout(None)
75 + kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
76 + return allocate_kv_cache(
77 + kv_cache_config, device, KVCacheLayout.LBNHC, kernel_block_sizes
78 + )
79
80
81 def _single_rank_vllm_config(total_kv_heads: int):
@@ -321,10 +298,9 @@ def test_register_kv_caches(backend):
298 Creates one FullAttention group, one MLA group, one Mamba group, and
299 one Mamba-padded group. Each group has GROUP_SIZE layers.
300
324 - KVCacheTensors are shared across all groups mirroring the real allocation
325 - in kv_cache_utils.py: tensor i is shared by layer i from every group.
326 - The padded-mamba group has a different page size so its layers get their
327 - own dedicated tensors.
301 + KVCacheTensors mirror the real allocation in kv_cache_utils.py: each group is one
302 + run of the shared allocation starting at offset 0, so layer i of every group
303 + aliases the same bytes.
304
305 Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
306 the raw per-layer kv_caches registered by the connector.
@@ -413,18 +389,15 @@ def test_register_kv_caches(backend):
389 aligned_mamba_layer_names,
390 ]
391
416 - kv_cache_tensors: list[KVCacheTensor] = []
417 - for i in range(GROUP_SIZE):
418 - shared_by: list[str] = []
419 - for group_layer_names in layer_groups:
420 - if len(group_layer_names) > i:
421 - shared_by.append(group_layer_names[i])
422 - kv_cache_tensors.append(
423 - KVCacheTensor(
424 - size=PAGE_SIZE_BYTES * NUM_BLOCKS,
425 - shared_by=shared_by,
426 - )
392 + kv_cache_tensors: list[KVCacheTensor] = [
393 + KVCacheTensor(
394 + size=PAGE_SIZE_BYTES * NUM_BLOCKS * GROUP_SIZE,
395 + layers=group_layer_names,
396 + layer_stride=PAGE_SIZE_BYTES * NUM_BLOCKS,
397 + block_stride=PAGE_SIZE_BYTES,
398 )
399 + for group_layer_names in layer_groups
400 + ]
401
402 kv_cache_groups = [
403 KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec),
@@ -472,7 +445,7 @@ def test_register_kv_caches(backend):
445 kv_cache_groups=kv_cache_groups,
446 )
447
475 - kv_caches = _allocate_and_reshape_kv_caches(
448 + kv_caches = _allocate_kv_caches(
449 kv_cache_config,
450 attn_groups,
451 device=torch.device(f"{DEVICE_TYPE}:0"),
@@ -580,16 +553,24 @@ def test_register_kv_caches_uniform_type(backend):
553 kv_cache_specs={layer_a: spec_a, layer_b: spec_b},
554 )
555
556 + # The group packs its two layers densely, one run each as their pages differ:
557 + # layer_a occupies the first page of every block, layer_b the rest.
558 + window = spec_a.page_size_bytes + spec_b.page_size_bytes
559 kv_cache_config = KVCacheConfig(
560 num_blocks=NUM_BLOCKS,
561 kv_cache_tensors=[
562 KVCacheTensor(
587 - size=spec_a.page_size_bytes * NUM_BLOCKS,
588 - shared_by=[layer_a],
563 + size=window * NUM_BLOCKS,
564 + layers=[layer_a],
565 + layer_stride=spec_a.page_size_bytes * NUM_BLOCKS,
566 + block_stride=spec_a.page_size_bytes,
567 ),
568 KVCacheTensor(
591 - size=spec_b.page_size_bytes * NUM_BLOCKS,
592 - shared_by=[layer_b],
569 + size=window * NUM_BLOCKS,
570 + layers=[layer_b],
571 + layer_stride=spec_b.page_size_bytes * NUM_BLOCKS,
572 + block_stride=spec_b.page_size_bytes,
573 + offset=spec_a.page_size_bytes * NUM_BLOCKS,
574 ),
575 ],
576 kv_cache_groups=[
@@ -617,7 +598,7 @@ def test_register_kv_caches_uniform_type(backend):
598 ]
599 ]
600
620 - kv_caches = _allocate_and_reshape_kv_caches(
601 + kv_caches = _allocate_kv_caches(
602 kv_cache_config,
603 attn_groups,
604 device=torch.device(f"{DEVICE_TYPE}:0"),
tests/v1/kv_connector/unit/offloading_connector/utils.py
+9 -4
@@ -241,13 +241,19 @@ class RequestRunner:
241 )
242 ]
243
244 + # Groups overlay one allocation, sized by the largest group.
245 + block_stride = max(
246 + group.kv_cache_spec.page_size_bytes * len(group.layer_names)
247 + for group in kv_cache_groups
248 + )
249 kv_cache_tensors = [
250 KVCacheTensor(
246 - size=group.kv_cache_spec.page_size_bytes * num_gpu_blocks,
247 - shared_by=[layer_name],
251 + size=block_stride * num_gpu_blocks,
252 + layers=list(group.layer_names),
253 + layer_stride=group.kv_cache_spec.page_size_bytes * num_gpu_blocks,
254 + block_stride=group.kv_cache_spec.page_size_bytes,
255 )
256 for group in kv_cache_groups
250 - for layer_name in group.layer_names
257 ]
258
259 kv_cache_config = KVCacheConfig(
@@ -277,7 +283,6 @@ class RequestRunner:
283 )
284
285 # register worker kv_caches to enable OffloadingWorker creations
280 - # set_current_vllm_config is needed for get_kv_cache_layout() to work
286 kv_caches: dict[str, torch.Tensor] = {}
287 for group in kv_cache_groups:
288 spec = group.kv_cache_spec
tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py
+1 -1
@@ -98,7 +98,7 @@ def _make_connector_with_fake_worker(
98 )
99 worker = connector.connector_worker
100 assert isinstance(worker.nixl_wrapper, FakeNixlWrapper)
101 - worker.kv_cache_layout = "HND"
101 + worker.kv_cache_layout = "LBHNC"
102 if do_handshake:
103 remote_agents, _ = worker._nixl_handshake(
104 host="localhost",
tests/v1/kv_connector/unit/test_kv_cache_layout.py deleted
-75
@@ -1,75 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -from vllm.platforms import current_platform
7 -
8 -
9 -def test_mla_common_backend_rejects_cross_layer_kv_cache():
10 - """MLACommonBackend defaults to the identity permutation (layers dim
11 - first) so MLA backends whose decode kernels are not verified to honor
12 - the cache's block-dim stride stay opted out of cross-layer KV cache."""
13 - from vllm.model_executor.layers.attention.mla_attention import (
14 - MLACommonBackend,
15 - )
16 -
17 - stride_order = MLACommonBackend.get_kv_cache_stride_order(
18 - include_num_layers_dimension=True
19 - )
20 - assert stride_order == (0, 1, 2, 3)
21 - assert stride_order[0] == 0 # layers dim first => no cross-layer
22 - assert MLACommonBackend.get_kv_cache_stride_order(
23 - include_num_layers_dimension=False
24 - ) == (0, 1, 2)
25 -
26 -
27 -@pytest.mark.parametrize(
28 - "backend_path",
29 - # See: https://github.com/vllm-project/vllm/issues/46411
30 - [
31 - "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
32 - ]
33 - if current_platform.is_rocm() or current_platform.is_xpu()
34 - else [
35 - "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
36 - "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
37 - "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
38 - "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
39 - "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend",
40 - ],
41 -)
42 -def test_verified_mla_backends_support_cross_layer_kv_cache(backend_path):
43 - """Backends whose decode kernels honor the cache's block-dim stride opt
44 - in to the cross-layer layout with a non-identity permutation placing
45 - num_blocks first in physical layout."""
46 - module_path, name = backend_path.rsplit(".", 1)
47 - backend = getattr(
48 - pytest.importorskip(module_path, reason="backend deps unavailable"), name
49 - )
50 -
51 - stride_order = backend.get_kv_cache_stride_order(include_num_layers_dimension=True)
52 - assert stride_order == (1, 0, 2, 3)
53 - assert stride_order[0] != 0 # num_blocks first => cross-layer supported
54 - assert backend.get_kv_cache_stride_order(include_num_layers_dimension=False) == (
55 - 0,
56 - 1,
57 - 2,
58 - )
59 -
60 -
61 -def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
62 - """DeepseekV32Indexer returns identity permutation (layers dim first)
63 - to signal cross-layer KV cache is unsupported."""
64 - from vllm.v1.attention.backends.mla.indexer import (
65 - DeepseekV32IndexerBackend,
66 - )
67 -
68 - stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order(
69 - include_num_layers_dimension=True
70 - )
71 - assert stride_order == (0, 1, 2, 3)
72 - assert stride_order[0] == 0 # layers dim first => no cross-layer
73 - assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
74 - include_num_layers_dimension=False
75 - ) == (0, 1, 2)
tests/v1/kv_connector/unit/test_mooncake_connector.py
+60 -20
@@ -11,6 +11,7 @@ import pytest
11 import torch
12 import zmq.asyncio
13
14 +from tests.v1.attention.utils import dense_kv_cache_views
15 from vllm import envs
16 from vllm.config import set_current_vllm_config
17 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
@@ -32,11 +33,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import
33 MooncakeBootstrapServer,
34 )
35 from vllm.utils.network_utils import get_open_port
35 -from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
36 from vllm.v1.kv_cache_interface import (
37 FullAttentionSpec,
38 KVCacheConfig,
39 KVCacheGroupSpec,
40 + KVCacheLayout,
41 )
42 from vllm.v1.request import RequestStatus
43
@@ -1093,7 +1094,19 @@ async def test_worker_get_finished_timeout(monkeypatch):
1094 assert "tx-active" in prefill_worker.reqs_need_send
1095
1096
1096 -def test_register_kv_caches():
1097 +@pytest.mark.parametrize(
1098 + ("layout", "separate_kv_head_groups"),
1099 + [
1100 + (KVCacheLayout.LBHNC, False),
1101 + (KVCacheLayout.BLHNC, False),
1102 + (KVCacheLayout.BHLNC, False),
1103 + # LHBNC gives each head group its own region; the K/V split doubles the
1104 + # head count but the registration shape is driven by the layout.
1105 + (KVCacheLayout.LHBNC, False),
1106 + (KVCacheLayout.LHBNC, True),
1107 + ],
1108 +)
1109 +def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool):
1110 """Tests the memory registration logic with the underlying Mooncake engine."""
1111
1112 vllm_config = create_vllm_config(
@@ -1118,15 +1131,23 @@ def test_register_kv_caches():
1131 worker = connector.connector_worker
1132 mock_thread.return_value.is_alive.return_value = False
1133
1121 - kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
1122 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
1134 + spec = FullAttentionSpec(
1135 + block_size=16,
1136 + num_kv_heads=4,
1137 + head_size=64,
1138 + dtype=torch.float16,
1139 + num_head_slots=2 if separate_kv_head_groups else None,
1140 + state_content_bytes=4 * 64 * 2 if separate_kv_head_groups else None,
1141 )
1124 - tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1125 - tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1126 - kv_caches = {
1127 - "model.layers.0.self_attn": tensor1,
1128 - "model.layers.1.self_attn": tensor2,
1129 - }
1142 + layer_names = [
1143 + "model.layers.0.self_attn",
1144 + "model.layers.1.self_attn",
1145 + ]
1146 + for layer_name in layer_names:
1147 + worker._layer_specs[layer_name] = spec
1148 + raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
1149 + tensor1, tensor2 = dense_kv_cache_views(raw, spec, 2, 2, layout)
1150 + kv_caches = dict(zip(layer_names, (tensor1, tensor2)))
1151
1152 with patch.object(
1153 worker.engine, "batch_register_memory", return_value=0
@@ -1135,16 +1156,35 @@ def test_register_kv_caches():
1156
1157 mock_batch_register.assert_called_once()
1158 registered_ptrs, registered_lens = mock_batch_register.call_args[0]
1138 - expected_ptrs = {tensor.data_ptr() for tensor in kv_caches.values()}
1139 - assert set(registered_ptrs) == expected_ptrs
1140 - assert set(registered_lens) == {tensor1.nbytes}
1141 -
1142 - # Verify block_len_per_layer is set correctly.
1143 - assert len(worker.block_len_per_layer) == len(registered_ptrs)
1144 - for bl in worker.block_len_per_layer:
1145 - assert bl == tensor1.nbytes // tensor1.shape[0]
1146 - assert worker.registered_layer_names == list(kv_caches)
1147 - assert worker.registered_layer_indices == [0, 1]
1159 + assert registered_ptrs == [raw.data_ptr()]
1160 + assert registered_lens == [raw.nbytes]
1161 +
1162 + if not layout.is_block_compact:
1163 + expected_addrs = [
1164 + cache[:, head_idx].data_ptr()
1165 + for cache in (tensor1, tensor2)
1166 + for head_idx in range(cache.shape[1])
1167 + ]
1168 + head_block_bytes = tensor1.stride(0) * tensor1.element_size()
1169 + assert worker.kv_caches_base_addr == expected_addrs
1170 + assert worker.block_len_per_layer == [head_block_bytes] * len(
1171 + expected_addrs
1172 + )
1173 + assert worker.kv_block_len_per_layer == [head_block_bytes] * len(
1174 + expected_addrs
1175 + )
1176 + assert worker.registered_layer_names == [
1177 + layer_name
1178 + for layer_name in layer_names
1179 + for _ in range(tensor1.shape[1])
1180 + ]
1181 + else:
1182 + assert len(worker.block_len_per_layer) == len(kv_caches)
1183 + for bl in worker.block_len_per_layer:
1184 + assert bl == tensor1.stride(0) * tensor1.element_size()
1185 + assert worker.kv_block_len_per_layer == [spec.page_size_bytes] * 2
1186 + assert worker.registered_layer_names == list(kv_caches)
1187 + assert worker.registered_layer_indices == [0, 1]
1188
1189
1190 def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
+24 -12
@@ -14,6 +14,7 @@ from unittest.mock import patch
14 import pytest
15 import torch
16
17 +from tests.v1.attention.utils import dense_kv_cache_views
18 from vllm.config import set_current_vllm_config
19 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
20 KVConnectorRole,
@@ -30,6 +31,7 @@ from vllm.v1.kv_cache_interface import (
31 FullAttentionSpec,
32 KVCacheConfig,
33 KVCacheGroupSpec,
34 + KVCacheLayout,
35 MambaSpec,
36 )
37
@@ -138,14 +140,22 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
140 )
141 worker = connector.connector_worker
142
141 - fa_cache = torch.empty((2, 2, 11), dtype=torch.float16)
142 - gdn_conv_state = torch.empty((2, 22), dtype=torch.float16)
143 - gdn_ssm_state = torch.empty((2, 4), dtype=torch.float16)
143 + num_blocks = kv_cache_config.num_blocks
144 + fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
145 + gdn_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec
146 + fa_raw = torch.empty(num_blocks * fa_spec.page_size_bytes, dtype=torch.int8)
147 + gdn_raw = torch.empty(num_blocks * gdn_spec.page_size_bytes, dtype=torch.int8)
148 + (fa_cache,) = dense_kv_cache_views(
149 + fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC
150 + )
151 + (gdn_cache,) = dense_kv_cache_views(
152 + gdn_raw, gdn_spec, num_blocks, 1, KVCacheLayout.LBHNC
153 + )
154
155 worker.register_kv_caches(
156 {
157 "model.layers.0.self_attn": fa_cache,
148 - "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
158 + "model.layers.1.linear_attn": gdn_cache,
159 }
160 )
161
@@ -154,10 +164,14 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
164 "model.layers.0.self_attn",
165 "model.layers.1.linear_attn",
166 ]
167 + assert worker.block_len_per_layer == [
168 + fa_spec.page_size_bytes,
169 + gdn_spec.page_size_bytes,
170 + ]
171 assert worker.registered_group_indices == [0, 1]
172 assert worker.kv_caches_base_addr == [
173 fa_cache.data_ptr(),
160 - gdn_conv_state.data_ptr(),
174 + gdn_cache.data_ptr(),
175 ]
176
177 worker.shutdown()
@@ -185,8 +199,7 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
199
200 backing = torch.empty((4, 64), dtype=torch.float16)
201 fa_cache = backing[:2, :16]
188 - gdn_conv_state = backing[:3]
189 - gdn_ssm_state = torch.empty((3, 4), dtype=torch.float16)
202 + gdn_cache = backing[:3]
203
204 with patch.object(
205 worker.engine, "batch_register_memory", return_value=0
@@ -194,13 +207,13 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
207 worker.register_kv_caches(
208 {
209 "model.layers.0.self_attn": fa_cache,
197 - "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
210 + "model.layers.1.linear_attn": gdn_cache,
211 }
212 )
213
214 assert worker.kv_caches_base_addr == [
215 fa_cache.data_ptr(),
203 - gdn_conv_state.data_ptr(),
216 + gdn_cache.data_ptr(),
217 ]
218 batch_register_memory.assert_called_once()
219 registered_ptrs, registered_lens = batch_register_memory.call_args[0]
@@ -339,7 +352,7 @@ def test_logical_to_kernel_block_ids_expands_fa_not_gdn():
352 assert kernel_block_ids == [list(range(34, 51)), [2]]
353
354
342 -def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
355 +def test_hybrid_gdn_keeps_packed_fa_and_gdn_regions_whole(
356 monkeypatch,
357 ):
358 monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
@@ -359,7 +372,7 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
372 )
373 worker = connector.connector_worker
374
362 - worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=True)
375 + worker.transfer_topo = SimpleNamespace(is_kv_layout_blocks_first=False)
376 regions = worker._get_transfer_regions(
377 base_addrs=[0x1000, 0x2000],
378 block_lens=[0x100, 0x100],
@@ -377,7 +390,6 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
390 for region in regions
391 ] == [
392 (0, 0x1000, 0x40),
380 - (0, 0x1040, 0x40),
393 (1, 0x2000, 0x100),
394 ]
395
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+8 -59
@@ -48,7 +48,14 @@ def _make_kv_cache_config() -> KVCacheConfig:
48 spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
49 return KVCacheConfig(
50 num_blocks=4,
51 - kv_cache_tensors=[KVCacheTensor(size=8192, shared_by=["layer0"])],
51 + kv_cache_tensors=[
52 + KVCacheTensor(
53 + size=8192,
54 + layers=["layer0"],
55 + layer_stride=8192,
56 + block_stride=2048,
57 + )
58 + ],
59 kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)],
60 )
61
@@ -210,64 +217,6 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events():
217 assert kv_events.get_all_events() == [event]
218
219
213 -def test_prefer_cross_layer_blocks_from_config():
214 - # Default: disabled
215 - vllm_config = _make_vllm_config()
216 - kv_cache_config = _make_kv_cache_config()
217 - with (
218 - set_current_vllm_config(vllm_config),
219 - patch(
220 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
221 - "connector.MooncakeStoreScheduler"
222 - ),
223 - ):
224 - connector = mooncake_store_connector.MooncakeStoreConnector(
225 - vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
226 - )
227 - assert connector.prefer_cross_layer_blocks is False
228 -
229 - # Enabled via config
230 - vllm_config_enabled = create_vllm_config(
231 - kv_connector="MooncakeStoreConnector",
232 - kv_role="kv_both",
233 - kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
234 - )
235 - with (
236 - set_current_vllm_config(vllm_config_enabled),
237 - patch(
238 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
239 - "connector.MooncakeStoreScheduler"
240 - ),
241 - ):
242 - connector_enabled = mooncake_store_connector.MooncakeStoreConnector(
243 - vllm_config_enabled, KVConnectorRole.SCHEDULER, kv_cache_config
244 - )
245 - assert connector_enabled.prefer_cross_layer_blocks is True
246 -
247 -
248 -def test_register_cross_layers_kv_cache_delegates_to_worker():
249 - vllm_config = _make_vllm_config()
250 - kv_cache_config = _make_kv_cache_config()
251 -
252 - with (
253 - set_current_vllm_config(vllm_config),
254 - patch(
255 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
256 - "connector.MooncakeStoreWorker"
257 - ) as mock_worker_cls,
258 - ):
259 - connector = mooncake_store_connector.MooncakeStoreConnector(
260 - vllm_config, KVConnectorRole.WORKER, kv_cache_config
261 - )
262 -
263 - fake_tensor = MagicMock()
264 - fake_backend = MagicMock()
265 - connector.register_cross_layers_kv_cache(fake_tensor, fake_backend)
266 -
267 - worker = mock_worker_cls.return_value
268 - worker.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
269 -
270 -
220 def test_update_connector_output_and_take_events():
221 vllm_config = _make_vllm_config()
222 kv_cache_config = _make_kv_cache_config()
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+18 -4
@@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch
10
11 import torch
12
13 +from tests.v1.attention.utils import dense_kv_cache_views
14 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
15 worker as mooncake_store_worker,
16 )
@@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
36 MambaSpec,
37 SlidingWindowSpec,
38 )
39 +from vllm.v1.kv_cache_layout import KVCacheLayout
40
41
42 class _DictStore:
@@ -160,8 +162,18 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():
162 cfg = KVCacheConfig(
163 num_blocks=4,
164 kv_cache_tensors=[
163 - KVCacheTensor(size=8192, shared_by=["L0"]),
164 - KVCacheTensor(size=8192, shared_by=["L1"]),
165 + KVCacheTensor(
166 + size=4 * full.page_size_bytes,
167 + layers=["L0"],
168 + layer_stride=4 * full.page_size_bytes,
169 + block_stride=full.page_size_bytes,
170 + ),
171 + KVCacheTensor(
172 + size=4 * swa.page_size_bytes,
173 + layers=["L1"],
174 + layer_stride=4 * swa.page_size_bytes,
175 + block_stride=swa.page_size_bytes,
176 + ),
177 ],
178 kv_cache_groups=[
179 KVCacheGroupSpec(["L0"], full),
@@ -178,9 +190,11 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():
190
191 # Register kv_caches using mocked thread classes so register_kv_caches
192 # doesn't try to start real background threads (which set ready_event).
193 + raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
194 + raw_swa = torch.zeros(4 * swa.page_size_bytes, dtype=torch.int8)
195 kv_caches = {
182 - "L0": torch.zeros(2, 4, 8, 8, 64),
183 - "L1": torch.zeros(2, 4, 8, 8, 64),
196 + "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0],
197 + "L1": dense_kv_cache_views(raw_swa, swa, 4, 1, KVCacheLayout.LBNHC)[0],
198 }
199
200 def _fake_thread_init(*args, **kwargs):
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+72 -95
@@ -16,6 +16,7 @@ from unittest.mock import MagicMock, patch
16 import pytest
17 import torch
18
19 +from tests.v1.attention.utils import dense_kv_cache_views
20 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
21 rdma_utils,
22 )
@@ -38,6 +39,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import
39 )
40 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
41 from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
42 +from vllm.v1.kv_cache_interface import (
43 + FullAttentionSpec,
44 + KVCacheGroupSpec,
45 + KVCacheLayout,
46 +)
47
48
49 class _RecordingBlockHashes:
@@ -2961,112 +2967,83 @@ def test_putting_consumer_queues_decode_save():
2967 assert req.current_event is event
2968
2969
2964 -def test_register_kv_caches_blocks_first_single_segment():
2965 - """Blocks-first layout (FlashInfer/MLA): one segment per layer."""
2970 +@pytest.mark.parametrize("layout", list(KVCacheLayout))
2971 +def test_register_kv_caches_shared_storage(layout: KVCacheLayout):
2972 num_blocks = 10
2967 - page_size_elements = 64
2973 + num_layers = 2
2974 worker = _make_bare_worker(num_gpu_blocks=num_blocks)
2969 -
2970 - # Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
2971 - tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
2972 - _register_with_mocked_threads(worker, {"layer0": tensor})
2973 -
2974 - db = worker.token_dbs[0]
2975 - assert db.kv_caches_base_addr == [tensor.untyped_storage().data_ptr()]
2976 - assert db.block_len == [tensor.untyped_storage().nbytes() // num_blocks]
2977 - worker.store.register_buffer.assert_called_once_with(
2978 - tensor.untyped_storage().data_ptr(),
2979 - tensor.untyped_storage().nbytes(),
2975 + spec = FullAttentionSpec(
2976 + block_size=16,
2977 + num_kv_heads=2,
2978 + head_size=8,
2979 + dtype=torch.float16,
2980 )
2981 + raw = torch.zeros(
2982 + num_blocks * num_layers * spec.page_size_bytes,
2983 + dtype=torch.int8,
2984 + )
2985 + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
2986
2982 -
2983 -def test_register_kv_caches_kv_first_two_segments():
2984 - """K/V-first layout (FlashAttn): two segments (K, V) per layer."""
2985 - num_blocks = 10
2986 - block_size_tokens = 16
2987 - num_kv_heads = 4
2988 - head_size = 8
2989 - worker = _make_bare_worker(num_gpu_blocks=num_blocks)
2990 -
2991 - # Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
2992 - tensor = torch.zeros(
2993 - 2,
2994 - num_blocks,
2995 - block_size_tokens,
2996 - num_kv_heads,
2997 - head_size,
2998 - dtype=torch.float16,
2987 + _register_with_mocked_threads(
2988 + worker,
2989 + {"layer0": caches[0], "__cross_layer__": caches[1]},
2990 )
3000 - _register_with_mocked_threads(worker, {"layer0": tensor})
2991
2992 db = worker.token_dbs[0]
3003 - seg_stride = tensor.stride(0) * tensor.element_size()
3004 - base = tensor.untyped_storage().data_ptr()
3005 - assert db.kv_caches_base_addr == [base, base + seg_stride]
3006 - assert db.block_len == [seg_stride // num_blocks] * 2
3007 -
3008 -
3009 -def test_register_kv_caches_cross_layer_single_segment():
3010 - """Cross-layer tensor: single segment with block_len = page_size * num_layers."""
3011 - num_blocks = 10
3012 - num_layers = 4
3013 - per_layer_page_elements = 64 # elements per layer per block
3014 -
2993 + if not layout.is_block_compact:
2994 + # Each head group is its own region, so a block spans L*H regions.
2995 + head_cache = caches[0][:, 0]
2996 + assert db.kv_caches_base_addr == [
2997 + cache[:, head_idx].data_ptr()
2998 + for cache in caches
2999 + for head_idx in range(cache.shape[1])
3000 + ]
3001 + assert db.block_len == [head_cache.stride(0) * head_cache.element_size()] * (
3002 + num_layers * caches[0].shape[1]
3003 + )
3004 + elif layout.is_layer_compact:
3005 + assert db.kv_caches_base_addr == [cache.data_ptr() for cache in caches]
3006 + assert db.block_len == [spec.page_size_bytes] * num_layers
3007 + else:
3008 + assert db.kv_caches_base_addr == [raw.data_ptr()]
3009 + assert db.block_len == [num_layers * spec.page_size_bytes]
3010 + worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
3011 +
3012 +
3013 +def test_register_kv_caches_separate_head_groups():
3014 + # LHBNC gives each head group its own region; the K/V split doubles heads.
3015 + layout = KVCacheLayout.LHBNC
3016 + num_blocks = 3
3017 + num_layers = 2
3018 worker = _make_bare_worker(num_gpu_blocks=num_blocks)
3016 -
3017 - # Cross-layer blocks-first tensor: all layers packed into a single
3018 - # contiguous block. Shape (num_blocks, num_layers * per_layer_page)
3019 - # mimics the physical layout after stride reordering.
3020 - total_page_elements = num_layers * per_layer_page_elements
3021 - tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16)
3022 -
3023 - with (
3024 - patch(
3025 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
3026 - "worker.KVCacheStoreSendingThread",
3027 - side_effect=_auto_set_ready_event,
3028 - ),
3029 - patch(
3030 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
3031 - "worker.KVCacheStoreRecvingThread",
3032 - side_effect=_auto_set_ready_event,
3033 - ),
3034 - ):
3035 - # Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
3036 - worker.register_kv_caches({"__cross_layer__": tensor})
3037 -
3038 - db = worker.token_dbs[0]
3039 - assert len(db.kv_caches_base_addr) == 1
3040 - assert db.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
3041 -
3042 - expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
3043 - # block_len should be per_layer_page_size * num_layers
3044 - assert (
3045 - expected_block_len
3046 - == num_layers * per_layer_page_elements * tensor.element_size()
3019 + spec = FullAttentionSpec(
3020 + block_size=4,
3021 + num_kv_heads=2,
3022 + head_size=8,
3023 + dtype=torch.float16,
3024 + num_head_slots=2,
3025 + state_content_bytes=2 * 8 * 2,
3026 )
3048 - assert len(db.block_len) == 1
3049 - assert db.block_len[0] == expected_block_len
3027 + layer_names = ["layer0", "__cross_layer__"]
3028 + worker._kv_cache_groups = [KVCacheGroupSpec(layer_names, spec)]
3029 + raw = torch.zeros(
3030 + num_blocks * num_layers * spec.page_size_bytes,
3031 + dtype=torch.int8,
3032 + )
3033 + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
3034
3051 - # Also verify via register_cross_layers_kv_caches wrapper
3052 - worker2 = _make_bare_worker(num_gpu_blocks=num_blocks)
3053 - with (
3054 - patch(
3055 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
3056 - "worker.KVCacheStoreSendingThread",
3057 - side_effect=_auto_set_ready_event,
3058 - ),
3059 - patch(
3060 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
3061 - "worker.KVCacheStoreRecvingThread",
3062 - side_effect=_auto_set_ready_event,
3063 - ),
3064 - ):
3065 - worker2.register_cross_layers_kv_caches(tensor)
3035 + _register_with_mocked_threads(worker, dict(zip(layer_names, caches)))
3036
3067 - db2 = worker2.token_dbs[0]
3068 - assert db2.kv_caches_base_addr == db.kv_caches_base_addr
3069 - assert db2.block_len == db.block_len
3037 + head_block_bytes = caches[0].stride(0) * caches[0].element_size()
3038 + expected_addrs = [
3039 + cache[:, head_idx].data_ptr()
3040 + for cache in caches
3041 + for head_idx in range(cache.shape[1])
3042 + ]
3043 + db = worker.token_dbs[0]
3044 + assert db.kv_caches_base_addr == expected_addrs
3045 + assert db.block_len == [head_block_bytes] * len(expected_addrs)
3046 + worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
3047
3048
3049 # ---------------------------------------------------------------------------
tests/v1/kv_connector/unit/test_moriio_connector.py
+34 -30
@@ -43,6 +43,7 @@ from vllm.v1.kv_cache_interface import (
43 KVCacheConfig,
44 KVCacheGroupSpec,
45 KVCacheTensor,
46 + compute_layer_kv_cache_shape_bytes,
47 )
48
49 from .utils import create_request, create_scheduler
@@ -56,20 +57,25 @@ def _find_free_port() -> int:
57
58 def _make_test_kv_cache_config() -> KVCacheConfig:
59 layer_names = ["layer0", "layer1", "layer2"]
60 + num_blocks = 2
61 + spec = FullAttentionSpec(
62 + block_size=16,
63 + num_kv_heads=4,
64 + head_size=64,
65 + dtype=torch.float16,
66 + )
67 + page = spec.page_size_bytes
68 return KVCacheConfig(
60 - num_blocks=2,
61 - kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)],
62 - kv_cache_groups=[
63 - KVCacheGroupSpec(
64 - layer_names=layer_names,
65 - kv_cache_spec=FullAttentionSpec(
66 - block_size=16,
67 - num_kv_heads=4,
68 - head_size=64,
69 - dtype=torch.float16,
70 - ),
69 + num_blocks=num_blocks,
70 + kv_cache_tensors=[
71 + KVCacheTensor(
72 + size=len(layer_names) * num_blocks * page,
73 + layers=layer_names,
74 + layer_stride=num_blocks * page,
75 + block_stride=page,
76 )
77 ],
78 + kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)],
79 )
80
81
@@ -225,7 +231,7 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker):
231 engine_id,
232 *args,
233 hand_shake_latency: float = 1.8,
228 - kv_cache_layout="HND",
234 + kv_cache_layout="LBHNC",
235 kv_cache_config=None,
236 **kwargs,
237 ):
@@ -531,16 +537,15 @@ def test_register_kv_caches(mock_parallel_groups):
537 DEFAULT_PORT = 6301
538 TP_RANK = 0
539 DP_RANK = 0
534 - from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
535 -
536 - backend_cls = AiterFlashAttentionBackend
537 -
538 - # Create test kv cache tensors using proper backend shape
539 - kv_cache_shape = backend_cls.get_kv_cache_shape(
540 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
540 + # Create test kv cache tensors using KVCacheSpec layout
541 + shape = compute_layer_kv_cache_shape_bytes(
542 + FullAttentionSpec(
543 + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
544 + ),
545 + 2,
546 )
542 - shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
543 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
547 + shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
548 + unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
549 kv_caches = {
550 "layer0": shared_tensor,
551 "layer1": unique_tensor,
@@ -627,16 +632,15 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups):
632
633 ROLE = "kv_consumer"
634 vllm_config = create_vllm_config(role=ROLE)
630 - from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
631 -
632 - backend_cls = AiterFlashAttentionBackend
633 -
634 - # Create test kv cache tensors using proper backend shape
635 - kv_cache_shape = backend_cls.get_kv_cache_shape(
636 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
635 + # Create test kv cache tensors using KVCacheSpec layout
636 + shape = compute_layer_kv_cache_shape_bytes(
637 + FullAttentionSpec(
638 + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
639 + ),
640 + 2,
641 )
638 - shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
639 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
642 + shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
643 + unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
644 kv_caches = {
645 "layer0": shared_tensor,
646 "layer1": unique_tensor,
tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+11 -6
@@ -12,7 +12,10 @@ import pytest
12 import torch
13
14 from vllm.platforms import current_platform
15 -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
15 +from vllm.v1.kv_cache_interface import (
16 + FullAttentionSpec,
17 + MLAAttentionSpec,
18 +)
19
20 aiter_available = importlib.util.find_spec("aiter") is not None
21 mori_available = importlib.util.find_spec("mori") is not None
@@ -241,7 +244,7 @@ def _write_task(layer_name: str, transfer_id: str = "xfer") -> Any:
244 id="interleaved-kernel-axis-from-spec",
245 ),
246 pytest.param(
244 - (8, 4, 3),
247 + (8, 1, 4, 3),
248 _mla_spec(),
249 16,
250 {
@@ -297,7 +300,7 @@ def test_mixed_layers_compute_distinct_offsets_per_layer():
300 kv_caches = {
301 "separated": torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16),
302 "interleaved": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
300 - "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
303 + "indexer": torch.empty((8, 1, 4, 3), dtype=torch.bfloat16),
304 }
305 worker = _worker(
306 kv_caches,
@@ -342,7 +345,7 @@ def test_write_transfer_plan_caches_offsets_per_geometry():
345 kv_caches = {
346 "dense0": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
347 "dense1": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
345 - "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
348 + "indexer": torch.empty((8, 1, 4, 3), dtype=torch.bfloat16),
349 }
350 calls: list[str] = []
351
@@ -692,7 +695,7 @@ def test_empty_local_block_ids_is_free_only_noop():
695 def test_registration_regions_do_not_split_interleaved_or_mla_cache():
696 separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
697 interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
695 - indexer = torch.empty((8, 4, 3), dtype=torch.bfloat16)
698 + indexer = torch.empty((8, 1, 4, 3), dtype=torch.bfloat16)
699 worker = _worker(
700 {
701 "separated": separated,
@@ -745,7 +748,9 @@ def test_registration_regions_use_layer_num_blocks():
748
749
750 def test_unsupported_shape_raises_value_error():
748 - cache = torch.empty((8, 4, 2, 3), dtype=torch.bfloat16)
751 + # A 4-D view whose head/state/content dims all disagree with the spec (a
752 + # standardized [B, H, N, C] view for _full_spec would be (8, 2, 4, 6)).
753 + cache = torch.empty((8, 3, 5, 7), dtype=torch.bfloat16)
754 worker = _worker({"layer": cache}, {"layer": _full_spec()})
755
756 with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"):
tests/v1/kv_connector/unit/test_multi_connector.py
-10
@@ -889,16 +889,6 @@ Options:
889 """)
890
891
892 -def test_multi_connector_prefer_cross_layer_blocks(mc):
893 - mc._connectors[0].prefer_cross_layer_blocks = False
894 - mc._connectors[1].prefer_cross_layer_blocks = True
895 - assert mc.prefer_cross_layer_blocks is False
896 -
897 - mc._connectors[0].prefer_cross_layer_blocks = True
898 - mc._connectors[1].prefer_cross_layer_blocks = True
899 - assert mc.prefer_cross_layer_blocks is True
900 -
901 -
892 def test_multi_connector_worker_metadata(mc):
893 class MockConnectorWorkerMetadata(KVConnectorWorkerMetadata):
894 def __init__(self, data: set[str]):
tests/v1/kv_connector/unit/test_nixl_connector.py
+235 -221
@@ -19,6 +19,7 @@ import pytest
19 import ray
20 import torch
21
22 +from tests.v1.attention.utils import dense_kv_cache_views
23 from vllm import LLM
24 from vllm.config import KVTransferConfig, set_current_vllm_config
25 from vllm.distributed.kv_transfer.kv_connector.utils import (
@@ -54,8 +55,6 @@ from vllm.outputs import RequestOutput
55 from vllm.platforms import current_platform
56 from vllm.platforms.interface import Platform
57 from vllm.sampling_params import SamplingParams
57 -from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
58 -from vllm.v1.attention.backends.utils import set_kv_cache_layout
58 from vllm.v1.engine import EngineCoreRequest
59 from vllm.v1.engine.output_processor import OutputProcessor
60 from vllm.v1.kv_cache_interface import (
@@ -63,12 +62,12 @@ from vllm.v1.kv_cache_interface import (
62 FullAttentionSpec,
63 KVCacheConfig,
64 KVCacheGroupSpec,
66 - KVCacheTensor,
65 + KVCacheLayout,
66 + MLAAttentionSpec,
67 + compute_layer_kv_cache_shape_bytes,
68 )
69 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
70 from vllm.v1.request import RequestStatus
70 -from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
71 -from vllm.v1.worker.utils import AttentionGroup
71
72 from .utils import (
73 create_request,
@@ -96,9 +95,6 @@ def clear_kv_transfer():
95 yield
96 if has_kv_transfer_group():
97 ensure_kv_transfer_shutdown()
99 - # Reset any KV cache layout override set during tests so it doesn't
100 - # leak into tests in other modules.
101 - set_kv_cache_layout(None)
98
99
100 def get_default_xfer_telemetry(
@@ -350,12 +346,12 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():
346 )
347 def test_kv_transfer_handshake(dist_init):
348 """Unit test for basic NixlConnector interface functionality."""
353 - from vllm.config import set_current_vllm_config
349
350 # Test setup, we creates a scheduler that contains a NixlConnector
351 # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from
352 # all workers of the instance.
353 vllm_config = create_vllm_config()
354 + vllm_config.cache_config.kv_cache_layout = "BLHNC"
355 # in case the test runs on non-GPU machine
356 vllm_config.kv_transfer_config.kv_buffer_device = "cpu"
357 scheduler = create_scheduler(vllm_config)
@@ -386,18 +382,19 @@ def test_kv_transfer_handshake(dist_init):
382 kv_cache_spec = cast(
383 AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
384 )
389 - kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
390 - num_blocks=kv_cache_config.num_blocks,
391 - block_size=kv_cache_spec.block_size,
392 - num_kv_heads=kv_cache_spec.num_kv_heads,
393 - head_size=kv_cache_spec.head_size,
385 + raw = torch.zeros(
386 + kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 3,
387 + dtype=torch.int8,
388 + )
389 + caches = dense_kv_cache_views(
390 + raw,
391 + kv_cache_spec,
392 + kv_cache_config.num_blocks,
393 + num_layers=3,
394 + layout=KVCacheLayout.BLHNC,
395 )
395 - shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
396 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
396 kv_caches = {
398 - "layer0": shared_tensor,
399 - "layer1": unique_tensor,
400 - "layer2": shared_tensor,
397 + f"layer{layer_idx}": cache for layer_idx, cache in enumerate(caches)
398 }
399 prefill_connector.register_kv_caches(kv_caches)
400
@@ -474,7 +471,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
471 self,
472 *args,
473 hand_shake_latency: float = 1.8,
477 - kv_cache_layout="HND",
474 + kv_cache_layout="LBHNC",
475 kv_cache_config=None,
476 **kwargs,
477 ):
@@ -486,9 +483,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
483 # Mock register_kv_caches attributes needed for tests that do not call it.
484 self.src_xfer_handles_by_block_size = {self.block_size: 1}
485 self.src_blocks_data = np.empty((0, 3), dtype=np.uint64)
489 - test_shape = self.attn_backends[0].get_kv_cache_shape(
490 - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
491 - )
486 + rep_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
487 + test_shape = compute_layer_kv_cache_shape_bytes(rep_spec, 1)
488 self.transfer_topo = TransferTopology(
489 tp_rank=self.tp_rank,
490 tp_size=self.world_size,
@@ -502,7 +498,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
498 )
499
500 self.compat_hash = compute_nixl_compatibility_hash(
505 - self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
501 + self.vllm_config, self.backend_name
502 )
503
504 def _nixl_handshake(
@@ -553,9 +549,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
549 device_id=remote_tp_rank,
550 num_blocks=1,
551 block_lens=remote_block_lens,
556 - # `self.kv_cache_layout` is only forced to HND when vllm engine
557 - # is started. We mock HND here.
558 - kv_cache_layout="HND",
552 + block_strides=remote_block_lens,
553 + kv_cache_layout="LBHNC",
554 block_size=self.block_size,
555 ssm_sizes=(0, 0),
556 attn_backend_name=self.backend_name,
@@ -579,6 +574,8 @@ class TestNixlHandshake:
574 self, default_vllm_config, dist_init, pcp_rank
575 ):
576 """Only PCP rank zero publishes and reports sending completion."""
577 + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
578 +
579 vllm_config = create_vllm_config(kv_role="kv_producer")
580 vllm_config.parallel_config.prefill_context_parallel_size = 2
581 with (
@@ -656,7 +653,7 @@ class TestNixlHandshake:
653 worker.dst_xfer_side_handles = {
654 FakeNixlConnectorWorker.REMOTE_ENGINE_ID: {0: 1}
655 }
659 - worker.kv_cache_layout = "HND"
656 + worker.kv_cache_layout = "LBHNC"
657 num_xfers = 4
658 while True:
659 # For the same request_id, initiate multiple xfers across different
@@ -1053,7 +1050,9 @@ class TestNixlHandshake:
1050 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1051
1052 # Metadata with different kv_cache_layout than local worker
1056 - mismatched_layout = "HND" if worker.kv_cache_layout != "HND" else "NHD"
1053 + mismatched_layout = (
1054 + "LBHNC" if worker.kv_cache_layout != "LBHNC" else "LBNHC"
1055 + )
1056 meta = NixlAgentMetadata(
1057 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1058 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
@@ -1061,6 +1060,7 @@ class TestNixlHandshake:
1060 device_id=0,
1061 num_blocks=1,
1062 block_lens=worker.block_len_per_layer,
1063 + block_strides=worker.block_len_per_layer,
1064 kv_cache_layout=mismatched_layout,
1065 block_size=worker.block_size,
1066 ssm_sizes=(0, 0),
@@ -1100,7 +1100,7 @@ class TestNixlHandshake:
1100 vllm_config,
1101 connector.engine_id,
1102 hand_shake_latency=0,
1103 - kv_cache_layout="NHD",
1103 + kv_cache_layout="LBNHC",
1104 )
1105 worker = connector.connector_worker
1106
@@ -1111,15 +1111,17 @@ class TestNixlHandshake:
1111 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1112
1113 # Metadata with different kv_cache_layout than local worker
1114 + # prefill TP=1, decode TP=2, remote block_lens is double to local
1115 + remote_block_lens = [i * 2 for i in worker.block_len_per_layer]
1116 meta = NixlAgentMetadata(
1117 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1118 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1119 kv_caches_base_addr=[0],
1120 device_id=0,
1121 num_blocks=1,
1120 - # prefill TP=1, decode TP=2, remote block_lens is double to local
1121 - block_lens=[i * 2 for i in worker.block_len_per_layer],
1122 - kv_cache_layout="HND",
1122 + block_lens=remote_block_lens,
1123 + block_strides=remote_block_lens,
1124 + kv_cache_layout="LBHNC",
1125 block_size=worker.block_size,
1126 ssm_sizes=(0, 0),
1127 attn_backend_name=worker.backend_name,
@@ -1160,8 +1162,6 @@ class TestNixlHandshake:
1162 attn_backends=worker.attn_backends,
1163 tensor_shape=None,
1164 )
1163 - assert worker.transfer_topo.virtually_split_kv_in_blocks
1164 -
1165 plan = MagicMock(
1166 source_ranks_per_group=((0,), (0,)),
1167 rank_offset_factor=1,
@@ -1171,6 +1171,7 @@ class TestNixlHandshake:
1171 device_id=0,
1172 num_blocks=1,
1173 block_lens=[remote_block_len],
1174 + block_strides=[remote_block_len],
1175 )
1176
1177 assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
@@ -1226,6 +1227,7 @@ class TestNixlHandshake:
1227 device_id=0,
1228 num_blocks=1,
1229 block_lens=[fa_len * tp_ratio, idx_len],
1230 + block_strides=[fa_len * tp_ratio, idx_len],
1231 kv_cache_layout=worker.kv_cache_layout,
1232 block_size=worker.block_size,
1233 ssm_sizes=(0, 0),
@@ -1252,6 +1254,7 @@ class TestNixlHandshake:
1254 num_blocks=1,
1255 # WRONG: MLA region scaled by tp_ratio (it should be replicated).
1256 block_lens=[fa_len * tp_ratio, idx_len * tp_ratio],
1257 + block_strides=[fa_len * tp_ratio, idx_len * tp_ratio],
1258 kv_cache_layout=worker2.kv_cache_layout,
1259 block_size=worker2.block_size,
1260 ssm_sizes=(0, 0),
@@ -1295,7 +1298,7 @@ class TestNixlHandshake:
1298
1299 worker.transfer_topo.total_num_kv_heads = 8
1300 worker.transfer_topo.local_physical_heads = 1
1298 - worker.kv_cache_layout = "HND"
1301 + worker.kv_cache_layout = "LBHNC"
1302
1303 worker.slot_size_per_layer = [4096]
1304 worker.block_len_per_layer = [4096 * worker.block_size]
@@ -1311,7 +1314,8 @@ class TestNixlHandshake:
1314 device_id=0,
1315 num_blocks=1,
1316 block_lens=list(worker.block_len_per_layer),
1314 - kv_cache_layout="HND",
1317 + block_strides=list(worker.block_len_per_layer),
1318 + kv_cache_layout="LBHNC",
1319 block_size=worker.block_size,
1320 ssm_sizes=(0, 0),
1321 attn_backend_name=worker.backend_name,
@@ -1350,7 +1354,7 @@ class TestNixlHandshake:
1354
1355 worker.transfer_topo.total_num_kv_heads = 32
1356 worker.transfer_topo.local_physical_heads = 8 # 32 // 4
1353 - worker.kv_cache_layout = "HND"
1357 + worker.kv_cache_layout = "LBHNC"
1358
1359 slot_size = 4096
1360 worker.slot_size_per_layer = [slot_size]
@@ -1368,7 +1372,8 @@ class TestNixlHandshake:
1372 device_id=0,
1373 num_blocks=1,
1374 block_lens=list(worker.block_len_per_layer),
1371 - kv_cache_layout="HND",
1375 + block_strides=list(worker.block_len_per_layer),
1376 + kv_cache_layout="LBHNC",
1377 block_size=worker.block_size,
1378 ssm_sizes=(0, 0),
1379 attn_backend_name=worker.backend_name,
@@ -1839,7 +1844,19 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1844 llm.llm_engine.engine_core.shutdown()
1845
1846
1842 -@pytest.mark.parametrize("enable_cross_layers", ["False", "True"])
1847 +def element_byte_addrs(view: torch.Tensor) -> list[int]:
1848 + """Absolute byte addresses of every element of a (possibly strided) view."""
1849 + offsets = torch.zeros(view.shape, dtype=torch.int64)
1850 + for dim, (size, stride) in enumerate(zip(view.shape, view.stride())):
1851 + shape = [1] * view.ndim
1852 + shape[dim] = size
1853 + offsets = offsets + torch.arange(size, dtype=torch.int64).view(shape) * stride
1854 + esize = view.element_size()
1855 + byte_offsets = (offsets.flatten() * esize).tolist()
1856 + base = view.data_ptr()
1857 + return [base + offset + i for offset in byte_offsets for i in range(esize)]
1858 +
1859 +
1860 @pytest.mark.parametrize(
1861 "attn_backend",
1862 [
@@ -1853,8 +1870,14 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1870 "TRITON_ATTN",
1871 ],
1872 )
1873 +@pytest.mark.parametrize("layout", [layout.name for layout in KVCacheLayout])
1874 +@pytest.mark.parametrize("separate_kv_head_groups", [False, True])
1875 def test_register_kv_caches(
1857 - default_vllm_config, dist_init, attn_backend, enable_cross_layers
1876 + default_vllm_config,
1877 + dist_init,
1878 + attn_backend,
1879 + layout,
1880 + separate_kv_head_groups,
1881 ):
1882 """
1883 Test that register_kv_caches() properly calls nixl_wrapper methods with
@@ -1868,12 +1891,7 @@ def test_register_kv_caches(
1891 """
1892
1893 vllm_config = create_vllm_config(attention_backend=attn_backend)
1871 -
1872 - # Enable cross layers blocks
1873 - vllm_config.kv_transfer_config.kv_connector_extra_config[
1874 - "enable_cross_layers_blocks"
1875 - ] = enable_cross_layers
1876 - set_kv_cache_layout("HND")
1894 + vllm_config.cache_config.kv_cache_layout = layout
1895
1896 # Import the appropriate backend based on the parameter
1897 if attn_backend == "FLASH_ATTN":
@@ -1890,59 +1908,37 @@ def test_register_kv_caches(
1908 backend_cls = TritonAttentionBackend
1909
1910 nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
1893 - nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector"
1911 with (
1912 patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
1913 patch(f"{nixl_worker}.threading.Event"),
1914 patch(f"{nixl_worker}.threading.Thread") as mock_thread,
1898 - patch(f"{nixl_connector}.get_current_attn_backend") as mock_get_attn_backend,
1915 patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends,
1916 ):
1901 - # Ensure get_attn_backend returns the correct value due to
1902 - # _cached_get_attn_backend returning the backend from previous
1903 - # test run if not mocking.
1904 - mock_get_attn_backend.return_value = backend_cls
1917 mock_get_attn_backends.return_value = [backend_cls]
1906 - num_layers = 32
1918 block_size = 16
1919 num_blocks = 8
1920 num_heads = 4
1921 head_size = 16
1922
1912 - # TODO (NickLucche) the fact that connector depends on kv_cache_config for init
1913 - # but cross-layer preference cant be inferred prior to creating kv_cache_config
1914 - # is a bit awkward.
1915 - dummy_connector = NixlConnector(
1916 - vllm_config,
1917 - KVConnectorRole.WORKER,
1918 - make_kv_cache_config(block_size=block_size),
1919 - )
1923 kv_cache_spec = FullAttentionSpec(
1924 block_size=block_size,
1925 num_kv_heads=num_heads,
1926 head_size=head_size,
1927 dtype=torch.float16,
1928 + num_head_slots=2 if separate_kv_head_groups else None,
1929 + state_content_bytes=num_heads * head_size * 2
1930 + if separate_kv_head_groups
1931 + else None,
1932 + )
1933 + kv_cache_config = KVCacheConfig(
1934 + num_blocks=num_blocks,
1935 + kv_cache_tensors=[],
1936 + kv_cache_groups=[
1937 + KVCacheGroupSpec(
1938 + ["layer0", "layer1", "layer2", "layer3"], kv_cache_spec
1939 + )
1940 + ],
1941 )
1926 - if dummy_connector.prefer_cross_layer_blocks:
1927 - kv_cache_config = KVCacheConfig(
1928 - num_blocks=num_blocks,
1929 - kv_cache_tensors=[
1930 - KVCacheTensor(
1931 - size=kv_cache_spec.page_size_bytes * num_blocks,
1932 - shared_by=["all-layers"],
1933 - )
1934 - for _ in range(num_layers)
1935 - ],
1936 - kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)],
1937 - )
1938 - else:
1939 - kv_cache_config = KVCacheConfig(
1940 - num_blocks=num_blocks,
1941 - kv_cache_tensors=[],
1942 - kv_cache_groups=[
1943 - KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec)
1944 - ],
1945 - )
1942 # Create connector
1943 connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
1944 connector.connector_worker = FakeNixlConnectorWorker(
@@ -1962,93 +1958,36 @@ def test_register_kv_caches(
1958 # Reassure the shutdown() check that the thread is terminated
1959 mock_thread.return_value.is_alive.return_value = False
1960
1965 - expected_tensor_size: int
1966 - expected_base_addrs: list[int]
1967 - expected_num_entries: int
1968 - kv_caches: dict[str, torch.Tensor]
1969 - if str(enable_cross_layers).lower() == "true":
1970 - assert connector.prefer_cross_layer_blocks == (
1971 - attn_backend in ("FLASH_ATTN", "FLASHINFER", "TRITON_ATTN")
1972 - )
1973 - else:
1974 - assert not connector.prefer_cross_layer_blocks
1975 -
1976 - test_shape = backend_cls.get_kv_cache_shape(
1977 - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1961 + raw0 = torch.zeros(
1962 + kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 2,
1963 + dtype=torch.int8,
1964 + device=current_platform.device_type,
1965 )
1979 - is_blocks_first = len(test_shape) == 4 and test_shape[0] == 1
1980 -
1981 - if connector.prefer_cross_layer_blocks:
1982 - with set_current_vllm_config(vllm_config):
1983 - _, cross_layers_kv_cache, _ = (
1984 - KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
1985 - kv_cache_config=kv_cache_config,
1986 - attn_groups=[
1987 - [
1988 - AttentionGroup(
1989 - backend=backend_cls,
1990 - layer_names=[],
1991 - kv_cache_spec=kv_cache_spec,
1992 - kv_cache_group_id=0,
1993 - )
1994 - ]
1995 - ],
1996 - cache_dtype="bfloat16",
1997 - device=torch.accelerator.current_device_index(),
1998 - kernel_block_sizes=[block_size],
1999 - )
2000 - )
2001 - # Store tensor info for validation
2002 - expected_tensor_size = (
2003 - cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
2004 - )
2005 - expected_base_addrs = [
2006 - cross_layers_kv_cache.data_ptr(),
2007 - ]
2008 - expected_num_entries = 1
2009 -
2010 - expected_blocks_count = num_blocks
2011 -
2012 - kv_caches = {"all-layers": cross_layers_kv_cache}
2013 - else:
2014 - # Create test kv cache tensors using proper backend shape
2015 - kv_cache_shape = backend_cls.get_kv_cache_shape(
2016 - num_blocks=kv_cache_config.num_blocks,
2017 - block_size=kv_cache_spec.block_size,
2018 - num_kv_heads=kv_cache_spec.num_kv_heads,
2019 - head_size=kv_cache_spec.head_size,
2020 - )
2021 - shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
2022 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
2023 - kv_caches = {
2024 - "layer0": shared_tensor,
2025 - "layer1": unique_tensor,
2026 - "layer2": shared_tensor,
2027 - }
2028 -
2029 - # Store tensor info for validation
2030 - if is_blocks_first:
2031 - expected_tensor_size = (
2032 - shared_tensor.element_size() * shared_tensor.numel()
2033 - )
2034 - expected_base_addrs = [
2035 - shared_tensor.data_ptr(),
2036 - unique_tensor.data_ptr(),
2037 - ]
2038 - expected_num_entries = 2
2039 - expected_blocks_count = kv_cache_config.num_blocks * 2
2040 - else:
2041 - expected_tensor_size = (
2042 - shared_tensor[0].element_size() * shared_tensor[0].numel()
2043 - )
2044 - expected_base_addrs = [
2045 - shared_tensor[0].data_ptr(),
2046 - shared_tensor[1].data_ptr(),
2047 - unique_tensor[0].data_ptr(),
2048 - unique_tensor[1].data_ptr(),
2049 - ]
2050 - expected_num_entries = 4
2051 - expected_blocks_count = kv_cache_config.num_blocks * 4
1966 + raw1 = torch.zeros(
1967 + kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks,
1968 + dtype=torch.int8,
1969 + device=current_platform.device_type,
1970 + )
1971 + tensor0, tensor1 = dense_kv_cache_views(
1972 + raw0,
1973 + kv_cache_spec,
1974 + kv_cache_config.num_blocks,
1975 + num_layers=2,
1976 + layout=KVCacheLayout[layout],
1977 + )
1978 + (tensor2,) = dense_kv_cache_views(
1979 + raw1,
1980 + kv_cache_spec,
1981 + kv_cache_config.num_blocks,
1982 + num_layers=1,
1983 + layout=KVCacheLayout[layout],
1984 + )
1985 + kv_caches = {
1986 + "layer0": tensor0,
1987 + "layer1": tensor1,
1988 + "layer2": tensor2,
1989 + "layer3": tensor0,
1990 + }
1991
1992 # Execute register_kv_caches
1993 connector.register_kv_caches(kv_caches)
@@ -2056,42 +1995,121 @@ def test_register_kv_caches(
1995 # Verify get_reg_descs was called with caches_data
1996 assert mock_wrapper_instance.get_reg_descs.called
1997 caches_data, _ = mock_wrapper_instance.get_reg_descs.call_args[0]
2059 - assert len(caches_data) == expected_num_entries
1998 + assert len(caches_data) == 2
1999
2061 - for i, cache_entry in enumerate(caches_data):
2000 + for cache_entry, raw in zip(caches_data, (raw0, raw1)):
2001 base_addr, size, _tp_rank, _ = cache_entry
2063 - assert size == expected_tensor_size, (
2064 - f"Entry {i}: Expected tensor size {expected_tensor_size}, got {size}"
2065 - )
2066 - assert base_addr == expected_base_addrs[i], (
2067 - f"Entry {i}: Expected base address {expected_base_addrs[i]}, "
2068 - f"got {base_addr}"
2069 - )
2002 + assert size == raw.nbytes
2003 + assert base_addr == raw.data_ptr()
2004
2005 # Verify get_xfer_descs was called with blocks_data
2006 assert mock_wrapper_instance.get_xfer_descs.called
2007 blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
2008
2075 - # Validate blocks_data structure and size
2076 - assert len(blocks_data) == expected_blocks_count, (
2077 - f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
2078 - )
2009 + # Layout-blind contract: whatever regions the worker carves out,
2010 + # transferring "block b" must move exactly logical block b's bytes for
2011 + # every layer. Map each registered byte to the block that owns it, then
2012 + # require every descriptor window to hold bytes of a single block and
2013 + # the windows of block b to cover exactly block b's bytes.
2014 + owner: dict[int, int] = {}
2015 + for cache in (tensor0, tensor1, tensor2):
2016 + for blk in range(num_blocks):
2017 + for addr in element_byte_addrs(cache[blk]):
2018 + assert owner.setdefault(addr, blk) == blk
2019 + block_bytes = defaultdict(set)
2020 + for addr, blk in owner.items():
2021 + block_bytes[blk].add(addr)
2022 +
2023 + covered: defaultdict[int, set[int]] = defaultdict(set)
2024 + for block_start_addr, block_len, _tp_rank in blocks_data:
2025 + window = range(block_start_addr, block_start_addr + block_len)
2026 + owners = {owner[addr] for addr in window}
2027 + assert len(owners) == 1, "descriptor window spans logical blocks"
2028 + covered[owners.pop()].update(window)
2029 + assert covered == block_bytes
2030 +
2031 + # Region bases are exactly the block-0 window starts, one per region.
2032 + base_addrs = connector.connector_worker.kv_caches_base_addr[
2033 + connector.connector_worker.engine_id
2034 + ][0]
2035 + if layout == "BLHNC":
2036 + assert len(base_addrs) == 3
2037 + assert set(base_addrs) == {
2038 + start for start, _len, _tp in blocks_data if owner[start] == 0
2039 + }
2040 + assert len(blocks_data) == num_blocks * len(base_addrs)
2041
2080 - if connector.prefer_cross_layer_blocks:
2081 - num_blocks = 8
2082 - else:
2083 - num_blocks = kv_cache_config.num_blocks
2042 + assert connector.connector_worker.block_size == 16
2043
2085 - expected_block_len = expected_tensor_size // num_blocks
2044
2087 - for i, block_entry in enumerate(blocks_data):
2088 - block_start_addr, block_len, tp_rank = block_entry
2089 - assert block_len == expected_block_len, (
2090 - f"Block entry {i}: Expected block len {expected_block_len}, "
2091 - f"got {block_len}"
2092 - )
2045 +def test_register_packed_dsv4_mla_cache_as_single_region(
2046 + default_vllm_config, dist_init
2047 +):
2048 + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
2049
2094 - assert connector.connector_worker.block_size == 16
2050 + nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
2051 + with (
2052 + patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
2053 + patch(f"{nixl_worker}.threading.Event"),
2054 + patch(f"{nixl_worker}.threading.Thread") as mock_thread,
2055 + patch(f"{nixl_worker}.get_current_attn_backends") as mock_backends,
2056 + ):
2057 + mock_backends.return_value = [TritonAttentionBackend]
2058 + num_blocks = 2
2059 + num_layers = 4
2060 + block_size = 256
2061 + spec = MLAAttentionSpec(
2062 + block_size=block_size,
2063 + num_kv_heads=1,
2064 + head_size=512,
2065 + dtype=torch.uint8,
2066 + cache_dtype_str="fp8_ds_mla",
2067 + tokens_per_state=4,
2068 + alignment=576,
2069 + model_version="deepseek_v4",
2070 + state_content_bytes=584,
2071 + )
2072 + layer_names = [f"layer{idx}" for idx in range(num_layers)]
2073 + kv_cache_config = KVCacheConfig(
2074 + num_blocks=num_blocks,
2075 + kv_cache_tensors=[],
2076 + kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)],
2077 + )
2078 + vllm_config = create_vllm_config(attention_backend="TRITON_ATTN")
2079 + vllm_config.cache_config.block_size = block_size
2080 + vllm_config.cache_config.kv_cache_layout = "BLHNC"
2081 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
2082 + connector.connector_worker = FakeNixlConnectorWorker(
2083 + vllm_config,
2084 + connector.engine_id,
2085 + hand_shake_latency=0,
2086 + kv_cache_layout="BLHNC",
2087 + kv_cache_config=kv_cache_config,
2088 + )
2089 + wrapper = mock_nixl_wrapper.return_value
2090 + connector.connector_worker.nixl_wrapper = wrapper
2091 + wrapper.get_agent_metadata.return_value = b"fake_agent_metadata"
2092 + mock_thread.return_value.is_alive.return_value = False
2093 +
2094 + raw = torch.zeros(
2095 + spec.page_size_bytes * num_blocks * num_layers,
2096 + dtype=torch.int8,
2097 + device=current_platform.device_type,
2098 + )
2099 + views = dense_kv_cache_views(
2100 + raw, spec, num_blocks, num_layers, KVCacheLayout.BLHNC
2101 + )
2102 + connector.register_kv_caches(dict(zip(layer_names, views)))
2103 +
2104 + blocks_data, _ = wrapper.get_xfer_descs.call_args[0]
2105 + packed_block_len = num_layers * spec.page_size_bytes
2106 + assert connector.connector_worker.kv_caches_base_addr[
2107 + connector.connector_worker.engine_id
2108 + ][0] == [raw.data_ptr()]
2109 + assert blocks_data.tolist() == [
2110 + [raw.data_ptr() + block_idx * packed_block_len, packed_block_len, 0]
2111 + for block_idx in range(num_blocks)
2112 + ]
2113
2114
2115 class FakePlatform(Platform):
@@ -2318,6 +2336,8 @@ def test_engine_ttl_disabled(default_vllm_config, dist_init):
2336
2337 def test_transfer_topology_unregister():
2338 """TransferTopology.unregister_remote_engine removes the engine."""
2339 + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
2340 +
2341 topo = TransferTopology(
2342 tp_rank=0,
2343 tp_size=1,
@@ -2908,8 +2928,8 @@ def test_speculative_config_compatibility_hash(
2928 _set_test_speculative_config(local_config)
2929 _set_test_speculative_config(remote_config, **remote_overrides)
2930
2911 - local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
2912 - remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
2931 + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
2932 + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
2933
2934 assert (local_hash == remote_hash) is should_match
2935
@@ -2920,10 +2940,8 @@ def test_missing_speculative_config_changes_compatibility_hash():
2940 speculative_config = create_vllm_config()
2941 _set_test_speculative_config(speculative_config)
2942
2923 - regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN", False)
2924 - speculative_hash = compute_nixl_compatibility_hash(
2925 - speculative_config, "FLASH_ATTN", False
2926 - )
2943 + regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN")
2944 + speculative_hash = compute_nixl_compatibility_hash(speculative_config, "FLASH_ATTN")
2945
2946 assert regular_hash != speculative_hash
2947
@@ -2938,8 +2956,8 @@ def test_speculative_kv_cache_dtype_resolves_to_target():
2956 _set_test_speculative_config(local_config, kv_cache_dtype="fp8") # explicit
2957 _set_test_speculative_config(remote_config, kv_cache_dtype=None) # inherits
2958
2941 - local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
2942 - remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
2959 + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
2960 + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
2961
2962 assert local_hash == remote_hash
2963
@@ -2955,8 +2973,8 @@ def test_speculative_attention_backend_not_in_compatibility_hash():
2973 _set_test_speculative_config(local_config, attention_backend=None)
2974 _set_test_speculative_config(remote_config, attention_backend="FLASHINFER")
2975
2958 - local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
2959 - remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
2976 + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
2977 + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
2978
2979 assert local_hash == remote_hash
2980
@@ -2969,17 +2987,17 @@ def test_transfer_mode_changes_compatibility_hash():
2987 config = create_vllm_config()
2988
2989 pull_hash = compute_nixl_compatibility_hash(
2972 - config, "FLASH_ATTN", False, transfer_mode="pull"
2990 + config, "FLASH_ATTN", transfer_mode="pull"
2991 )
2992 push_hash = compute_nixl_compatibility_hash(
2975 - config, "FLASH_ATTN", False, transfer_mode="push"
2993 + config, "FLASH_ATTN", transfer_mode="push"
2994 )
2995
2996 assert pull_hash != push_hash
2997 assert pull_hash == compute_nixl_compatibility_hash(
2980 - config, "FLASH_ATTN", False, transfer_mode="pull"
2998 + config, "FLASH_ATTN", transfer_mode="pull"
2999 )
2982 - assert compute_nixl_compatibility_hash(config, "FLASH_ATTN", False) == pull_hash
3000 + assert compute_nixl_compatibility_hash(config, "FLASH_ATTN") == pull_hash
3001
3002
3003 @pytest.mark.skip_global_cleanup
@@ -3057,14 +3075,11 @@ def test_compatibility_hash_validation(
3075 kv_cache_spec = cast(
3076 AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
3077 )
3060 - kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape(
3061 - num_blocks=kv_cache_config.num_blocks,
3062 - block_size=kv_cache_spec.block_size,
3063 - num_kv_heads=kv_cache_spec.num_kv_heads,
3064 - head_size=kv_cache_spec.head_size,
3078 + shape = compute_layer_kv_cache_shape_bytes(
3079 + kv_cache_spec, kv_cache_config.num_blocks
3080 )
3066 - shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
3067 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
3081 + shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
3082 + unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
3083 # Build kv_caches from the actual layer names in kv_cache_config so that
3084 # _layer_specs lookups in register_kv_caches always find a matching key.
3085 layer_names = [
@@ -3099,18 +3114,19 @@ def test_compatibility_hash_validation(
3114 remote_hash = compute_nixl_compatibility_hash(
3115 remote_vllm_config,
3116 decode_worker.backend_name,
3102 - decode_worker.transfer_topo.cross_layers_blocks,
3117 )
3118
3119 prefill_block_size = config_overrides.get("block_size", 16)
3120 + prefill_block_lens = [4096 * prefill_block_size]
3121 prefill_metadata = NixlAgentMetadata(
3122 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3123 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
3124 kv_caches_base_addr=[0],
3125 device_id=0,
3126 num_blocks=1,
3112 - block_lens=[4096 * prefill_block_size], # slot_size * block_size
3113 - kv_cache_layout="HND",
3127 + block_lens=prefill_block_lens,
3128 + block_strides=prefill_block_lens,
3129 + kv_cache_layout="LBHNC",
3130 block_size=prefill_block_size,
3131 ssm_sizes=(0, 0),
3132 attn_backend_name=decode_worker.backend_name,
@@ -3187,9 +3203,8 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
3203 decode_worker = decode_connector.connector_worker
3204
3205 backend = get_current_attn_backend(local_vllm_config)
3190 - test_shape = backend.get_kv_cache_shape(
3191 - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
3192 - )
3206 + probe_spec = decode_worker.kv_cache_config.kv_cache_groups[0].kv_cache_spec
3207 + test_shape = compute_layer_kv_cache_shape_bytes(probe_spec, 1)
3208 decode_worker.transfer_topo = TransferTopology(
3209 tp_rank=decode_worker.tp_rank,
3210 tp_size=decode_worker.world_size,
@@ -3205,7 +3220,6 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
3220 decode_worker.compat_hash = compute_nixl_compatibility_hash(
3221 decode_worker.vllm_config,
3222 decode_worker.backend_name,
3208 - decode_worker.transfer_topo.cross_layers_blocks,
3223 )
3224
3225 if error_scenario == "handshake_decode_error":
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+8 -3
@@ -1505,14 +1505,19 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
1505 mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
1506 )
1507 assert kda_spec.page_size_bytes == unified_page
1508 + # The three groups overlay each other, so layer i of every group aliases the same
1509 + # region: mla.i, kda_a.i and kda_b.i all live at i * layer_stride.
1510 + layer_stride = num_blocks * unified_page
1511 return KVCacheConfig(
1512 num_blocks=num_blocks,
1513 kv_cache_tensors=[
1514 KVCacheTensor(
1512 - size=num_blocks * unified_page,
1513 - shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
1515 + size=2 * layer_stride,
1516 + layers=[f"{prefix}.0", f"{prefix}.1"],
1517 + layer_stride=layer_stride,
1518 + block_stride=unified_page,
1519 )
1515 - for i in range(2)
1520 + for prefix in ("mla", "kda_a", "kda_b")
1521 ],
1522 kv_cache_groups=[
1523 KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+10 -3
@@ -129,14 +129,19 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo
129 page_size_padded=unified_page,
130 mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
131 )
132 + # The three groups overlay each other, so layer i of every group aliases the same
133 + # region: mla.i, kda_a.i and kda_b.i all live at i * layer_stride.
134 + layer_stride = num_logical_blocks * unified_page
135 kv_cache_config = KVCacheConfig(
136 num_blocks=num_logical_blocks,
137 kv_cache_tensors=[
138 KVCacheTensor(
136 - size=num_logical_blocks * unified_page,
137 - shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
139 + size=2 * layer_stride,
140 + layers=[f"{prefix}.0", f"{prefix}.1"],
141 + layer_stride=layer_stride,
142 + block_stride=unified_page,
143 )
139 - for i in range(2)
144 + for prefix in ("mla", "kda_a", "kda_b")
145 ],
146 kv_cache_groups=[
147 KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
@@ -230,6 +235,8 @@ def _make_remote_meta(
235 kv_caches_base_addr=[0x10_000_000, 0x20_000_000],
236 num_blocks=remote_num_logical * remote_ppl,
237 block_lens=[kernel_page, kernel_page],
238 + # Non-interleaved remote: consecutive blocks abut, so stride == page length.
239 + block_strides=[kernel_page, kernel_page],
240 kv_cache_layout=worker.kv_cache_layout,
241 block_size=remote_kernel_block_size,
242 ssm_sizes=remote_ssm_sizes,
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+2
@@ -1081,6 +1081,8 @@ class TestPushPipelineParallel:
1081 device_id=0,
1082 num_blocks=4,
1083 block_lens=[block_len] * 4,
1084 + # Non-interleaved: consecutive blocks abut, so stride == block length.
1085 + block_strides=[block_len] * 4,
1086 kv_cache_layout="HND",
1087 block_size=16,
1088 ssm_sizes=(0, 0),
tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py
+11 -4
@@ -145,10 +145,15 @@ def _make_kv_cache_config(
145 ),
146 )
147 ]
148 + # Each group packs its two layers densely; groups overlay one allocation.
149 + layer_stride = _BYTES_PER_BLOCK * num_blocks
150 + size = 2 * layer_stride
151 tensors = [
152 KVCacheTensor(
150 - size=_BYTES_PER_BLOCK * num_blocks,
151 - shared_by=fa_layers,
153 + size=size,
154 + layers=fa_layers,
155 + layer_stride=layer_stride,
156 + block_stride=_BYTES_PER_BLOCK,
157 )
158 ]
159 if swa_enabled:
@@ -167,8 +172,10 @@ def _make_kv_cache_config(
172 )
173 tensors.append(
174 KVCacheTensor(
170 - size=_BYTES_PER_BLOCK * num_blocks,
171 - shared_by=sw_layers,
175 + size=size,
176 + layers=sw_layers,
177 + layer_stride=layer_stride,
178 + block_stride=_BYTES_PER_BLOCK,
179 )
180 )
181 return KVCacheConfig(
tests/v1/kv_connector/unit/test_transfer_topology_sharded.py
+1 -8
@@ -12,14 +12,7 @@ pytestmark = pytest.mark.cpu_test
12
13
14 class _FakeAttentionBackend:
15 - @staticmethod
16 - def get_kv_cache_shape(
17 - num_blocks: int,
18 - block_size: int,
19 - num_kv_heads: int,
20 - head_size: int,
21 - ) -> tuple[int, int, int, int]:
22 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
15 + pass
16
17
18 def _make_topology(
tests/v1/kv_connector/unit/utils.py
+2
@@ -130,6 +130,8 @@ def create_vllm_config(
130 cache_dtype=cache_dtype,
131 enable_prefix_caching=True,
132 )
133 + # Connectors are constructed after layout resolution; mirror that here.
134 + cache_config.kv_cache_layout = "LBNHC"
135 kv_transfer_config = KVTransferConfig(
136 kv_connector=kv_connector,
137 kv_connector_module_path=kv_connector_module_path,
tests/v1/kv_offload/test_file_mapper.py
+3 -8
@@ -55,6 +55,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
55 ),
56 replicated_layout=kwargs.get("replicated_layout", False),
57 canonical_layout=kwargs.get("canonical_layout", False),
58 + kv_cache_layout=kwargs.get("kv_cache_layout", "LBNHC"),
59 )
60 spec = MagicMock(spec=OffloadingSpec)
61 spec.config = config
@@ -214,14 +215,8 @@ def test_parallel_agnostic_separates_persistent_layouts():
215 def test_canonical_layout_changes_storage_namespace():
216 # Canonical bytes are not interchangeable with the direct layout, so the
217 # format id must fork the storage namespace.
217 - from vllm.v1.attention.backends.utils import set_kv_cache_layout
218 -
219 - set_kv_cache_layout("NHD")
220 - try:
221 - direct = make_mapper_from_offloading_spec()
222 - canonical = make_mapper_from_offloading_spec(canonical_layout=True)
223 - finally:
224 - set_kv_cache_layout(None)
218 + direct = make_mapper_from_offloading_spec()
219 + canonical = make_mapper_from_offloading_spec(canonical_layout=True)
220 assert "canonical_format" not in direct.fields
221 assert canonical.fields["canonical_format"] == "v1-nhd"
222 assert direct.base_path != canonical.base_path
tests/v1/simple_kv_offload/test_scheduler.py
+3 -1
@@ -91,7 +91,9 @@ def _make_kv_cache_config(
91 tensors.append(
92 KVCacheTensor(
93 size=_BYTES_PER_BLOCK * num_blocks,
94 - shared_by=layer_names,
94 + layers=layer_names,
95 + layer_stride=_BYTES_PER_BLOCK * num_blocks,
96 + block_stride=_BYTES_PER_BLOCK,
97 )
98 )
99 return KVCacheConfig(
tests/v1/simple_kv_offload/test_worker.py
+90
@@ -10,6 +10,7 @@ read partially written / stale blocks and silently corrupt the CPU cache.
10 from __future__ import annotations
11
12 import time
13 +from unittest.mock import MagicMock
14
15 import pytest
16 import torch
@@ -19,6 +20,8 @@ from vllm.platforms import current_platform
20 if not current_platform.is_cuda_alike():
21 pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
22
23 +from tests.v1.attention.utils import dense_kv_cache_views
24 +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
25 from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
26 from vllm.v1.simple_kv_offload.cuda_mem_ops import (
27 CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
@@ -188,3 +191,90 @@ def test_build_params_src_access_order():
191 gpu, cpu, stream, src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
192 )
193 assert ordered.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
194 +
195 +
196 +@pytest.mark.parametrize("layout", list(KVCacheLayout))
197 +def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout):
198 + num_blocks = 4
199 + num_layers = 2
200 + spec = FullAttentionSpec(
201 + block_size=2,
202 + num_kv_heads=2,
203 + head_size=2,
204 + dtype=torch.float16,
205 + )
206 + raw = torch.zeros(
207 + num_blocks * num_layers * spec.page_size_bytes,
208 + dtype=torch.int8,
209 + device="cuda",
210 + )
211 + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
212 + cache_config = MagicMock(num_blocks=num_blocks)
213 + worker = SimpleCPUOffloadWorker(
214 + vllm_config=None,
215 + kv_cache_config=cache_config,
216 + cpu_capacity_bytes=raw.nbytes,
217 + )
218 + worker._backend = MagicMock()
219 + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
220 +
221 + worker.register_kv_caches(
222 + {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
223 + )
224 +
225 + assert worker.gpu_kv_caches is not None
226 + if layout.is_layer_compact and not layout.is_block_compact:
227 + expected_regions = num_layers * spec.num_heads
228 + expected_block_bytes = spec.page_size_bytes // spec.num_heads
229 + elif layout.is_layer_compact:
230 + expected_regions = num_layers
231 + expected_block_bytes = spec.page_size_bytes
232 + else:
233 + expected_regions = 1
234 + expected_block_bytes = spec.page_size_bytes * num_layers
235 + assert len(worker.gpu_kv_caches) == expected_regions
236 + assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
237 + (num_blocks, expected_block_bytes)
238 + }
239 +
240 +
241 +def test_register_separate_kv_head_groups(monkeypatch):
242 + # LHBNC hoists the K/V head groups outside the block dim, so each layer's
243 + # blocks are registered as one region per group (K, V).
244 + layout = KVCacheLayout.LHBNC
245 + num_blocks = 4
246 + num_layers = 2
247 + spec = FullAttentionSpec(
248 + block_size=2,
249 + num_kv_heads=2,
250 + head_size=2,
251 + dtype=torch.float16,
252 + num_head_slots=2,
253 + state_content_bytes=2 * 2 * 2,
254 + )
255 + raw = torch.zeros(
256 + num_blocks * num_layers * spec.page_size_bytes,
257 + dtype=torch.int8,
258 + device="cuda",
259 + )
260 + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
261 + worker = SimpleCPUOffloadWorker(
262 + vllm_config=None,
263 + kv_cache_config=MagicMock(num_blocks=num_blocks),
264 + cpu_capacity_bytes=raw.nbytes,
265 + )
266 + worker._backend = MagicMock()
267 + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
268 +
269 + worker.register_kv_caches(
270 + {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
271 + )
272 +
273 + assert worker.gpu_kv_caches is not None
274 + assert len(worker.gpu_kv_caches) == num_layers * spec.num_heads
275 + per_group_block_bytes = (
276 + spec.num_kv_heads * spec.block_size * spec.head_size * spec.dtype.itemsize
277 + )
278 + assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
279 + (num_blocks, per_group_block_bytes)
280 + }
tests/v1/worker/test_attn_utils.py
+171 -286
@@ -1,48 +1,33 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Padded-page handling in create_kv_cache_views.
4
5 +Guards that a page_size_padded spec strides the block dimension by the padded page
6 +while keeping per-block content compact, so padding bytes at the end of each page are
7 +never addressed by the logical view.
8 +"""
9 +
10 +import pytest
11 import torch
12
13 +from tests.v1.attention.utils import dense_kv_cache_views
14 +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
15 from vllm.v1.kv_cache_interface import (
16 FullAttentionSpec,
17 KVCacheConfig,
18 + KVCacheGroupSpec,
19 + KVCacheLayout,
20 KVCacheTensor,
10 - KVQuantMode,
11 - MambaSpec,
21 + MLAAttentionSpec,
22 + compute_layout_strides,
23 +)
24 +from vllm.v1.worker.utils import (
25 + allocate_kv_cache,
26 + copy_kv_cache_blocks_inplace,
27 )
13 -from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
14 -from vllm.v1.worker.utils import AttentionGroup
15 -
16 -
17 -class FakeFlashAttentionBackend:
18 - @staticmethod
19 - def get_kv_cache_shape(
20 - num_blocks: int,
21 - block_size: int,
22 - num_kv_heads: int,
23 - head_size: int,
24 - cache_dtype_str: str = "auto",
25 - ) -> tuple[int, ...]:
26 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
27 -
28 - @staticmethod
29 - def get_kv_cache_stride_order(
30 - include_num_layers_dimension: bool = False,
31 - ) -> tuple[int, ...]:
32 - assert not include_num_layers_dimension
33 - return (0, 1, 2, 3, 4)
34 -
35 -
36 -class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend):
37 - @staticmethod
38 - def get_kv_cache_stride_order(
39 - include_num_layers_dimension: bool = False,
40 - ) -> tuple[int, ...]:
41 - assert not include_num_layers_dimension
42 - return (0, 1, 3, 2, 4)
28
29
45 -def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
30 +def test_reshape_padded_kv_cache_strides_by_padded_page():
31 num_blocks = 3
32 spec = FullAttentionSpec(
33 block_size=16,
@@ -53,292 +38,192 @@ def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
38 )
39 assert spec.real_page_size_bytes == 256
40
56 - raw_tensors = {
57 - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
58 - }
59 - attn_groups = [
60 - AttentionGroup(
61 - backend=FakeFlashAttentionBackend,
62 - layer_names=["layer"],
63 - kv_cache_spec=spec,
64 - kv_cache_group_id=0,
65 - )
66 - ]
67 -
68 - kv_cache = _reshape_kv_cache(
69 - attn_groups,
70 - raw_tensors,
71 - "auto",
72 - [spec.block_size],
73 - {},
74 - )["layer"]
41 + raw = torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
42 + (kv_cache,) = dense_kv_cache_views(raw, spec, num_blocks, 1, KVCacheLayout.LBHNC)
43 +
44 + elem_size = 4 # float32
45 + # Content dim packs K and V: 2 * head_size.
46 + assert kv_cache.shape == (num_blocks, 1, 16, 2 * spec.head_size)
47 + assert kv_cache.dtype == spec.dtype
48 + assert kv_cache.stride(0) == spec.page_size_padded // elem_size
49 + assert kv_cache[1].storage_offset() == spec.page_size_padded // elem_size
50 + # Within one block the (unpadded) content stays compact.
51 + assert kv_cache[0].is_contiguous()
52 +
53 +
54 +@pytest.mark.parametrize(
55 + ("kernel_block_sizes", "expected_num_blocks", "expected_num_states"),
56 + [
57 + (None, 4, 64),
58 + ([256], 4, 64),
59 + ([64], 16, 16),
60 + ],
61 +)
62 +def test_allocate_compressed_mla_cache(
63 + kernel_block_sizes: list[int] | None,
64 + expected_num_blocks: int,
65 + expected_num_states: int,
66 +):
67 + spec = MLAAttentionSpec(
68 + block_size=256,
69 + num_kv_heads=1,
70 + head_size=128,
71 + dtype=torch.bfloat16,
72 + tokens_per_state=4,
73 + )
74 + num_pages = 4
75 + config = KVCacheConfig(
76 + num_blocks=num_pages,
77 + kv_cache_tensors=[
78 + KVCacheTensor(
79 + size=num_pages * spec.page_size_bytes,
80 + layers=["layer.0"],
81 + layer_stride=num_pages * spec.page_size_bytes,
82 + block_stride=spec.page_size_bytes,
83 + )
84 + ],
85 + kv_cache_groups=[KVCacheGroupSpec(["layer.0"], spec)],
86 + )
87
76 - assert kv_cache.shape == (num_blocks, 2, 16, 1, 2)
77 - assert kv_cache.stride(0) == spec.page_size_bytes // 4
78 - assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
79 - assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
80 - assert (
81 - kv_cache[1, 1].storage_offset()
82 - == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
88 + caches = allocate_kv_cache(
89 + config, torch.device("cpu"), KVCacheLayout.LBHNC, kernel_block_sizes
90 )
91
92 + assert caches["layer.0"].shape == (expected_num_blocks, 1, expected_num_states, 128)
93
86 -def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page():
87 - num_blocks = 3
94 +
95 +@pytest.mark.parametrize("layout", list(KVCacheLayout))
96 +def test_copy_kv_cache_blocks_shared_storage(layout: KVCacheLayout):
97 + num_blocks = 4
98 + num_layers = 2
99 spec = FullAttentionSpec(
89 - block_size=16,
90 - num_kv_heads=3,
100 + block_size=2,
101 + num_kv_heads=2,
102 head_size=2,
103 dtype=torch.float32,
93 - page_size_padded=1024,
104 )
95 - assert spec.real_page_size_bytes == 768
96 -
97 - raw_tensors = {
98 - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
99 - }
100 - attn_groups = [
101 - AttentionGroup(
102 - backend=FakeHNDFlashAttentionBackend,
103 - layer_names=["layer"],
104 - kv_cache_spec=spec,
105 - kv_cache_group_id=0,
106 - )
107 - ]
105 + raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8)
106 + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
107
109 - kv_cache = _reshape_kv_cache(
110 - attn_groups,
111 - raw_tensors,
112 - "auto",
113 - [spec.block_size],
114 - {},
115 - )["layer"]
108 + for layer_idx, cache in enumerate(caches):
109 + for block_idx in range(num_blocks):
110 + cache[block_idx].fill_(10 * layer_idx + block_idx)
111
117 - assert kv_cache.shape == (num_blocks, 2, 16, 3, 2)
118 - assert kv_cache.stride(0) == spec.page_size_bytes // 4
119 - assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
120 - assert kv_cache.stride(2) == 2
121 - assert kv_cache.stride(3) == spec.block_size * spec.head_size
122 - assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
123 - assert (
124 - kv_cache[1, 1].storage_offset()
125 - == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
126 - )
127 - assert (
128 - kv_cache[1, 1, 3, 2].storage_offset()
129 - == (
130 - spec.page_size_bytes
131 - + spec.real_page_size_bytes // 2
132 - + 3 * spec.head_size * 4
133 - + 2 * spec.block_size * spec.head_size * 4
134 - )
135 - // 4
136 - )
112 + expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches]
113 + copies = [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)]
114
115 + copy_kv_cache_blocks_inplace(caches, num_blocks, copies)
116
139 -class FakeDiffKVBackend:
140 - @staticmethod
141 - def get_kv_cache_shape(
142 - num_blocks: int,
143 - block_size: int,
144 - num_kv_heads: int,
145 - head_size: int,
146 - cache_dtype_str: str = "auto",
147 - ) -> tuple[int, ...]:
148 - return (num_blocks, block_size, num_kv_heads, head_size * 2)
117 + for layer_idx, cache in enumerate(caches):
118 + torch.testing.assert_close(cache[2], expected[layer_idx][0])
119 + torch.testing.assert_close(cache[1], expected[layer_idx][1])
120
150 - @staticmethod
151 - def get_kv_cache_stride_order(
152 - include_num_layers_dimension: bool = False,
153 - ) -> tuple[int, ...]:
154 - assert not include_num_layers_dimension
155 - return (0, 1, 2, 3)
121
157 -
158 -def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim():
122 +def test_fixed_block_stride_propagates_outward_in_lhbnc():
123 num_blocks = 3
124 + num_layers = 2
125 spec = FullAttentionSpec(
161 - block_size=16,
162 - num_kv_heads=1,
126 + block_size=2,
127 + num_kv_heads=2,
128 head_size=2,
129 dtype=torch.float32,
165 - page_size_padded=384,
130 + )
131 + natural = compute_layout_strides(spec, num_blocks, num_layers, KVCacheLayout.LHBNC)
132 + block_stride = natural[1] + 8
133 +
134 + strides = compute_layout_strides(
135 + spec,
136 + num_blocks,
137 + num_layers,
138 + KVCacheLayout.LHBNC,
139 + fixed_strides=(None, block_stride, None, None, None),
140 )
141
168 - raw_tensors = {
169 - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
170 - }
171 - attn_groups = [
172 - AttentionGroup(
173 - backend=FakeDiffKVBackend,
174 - layer_names=["layer"],
175 - kv_cache_spec=spec,
176 - kv_cache_group_id=0,
177 - )
178 - ]
179 -
180 - kv_cache = _reshape_kv_cache(
181 - attn_groups,
182 - raw_tensors,
183 - "auto",
184 - [spec.block_size],
185 - {},
186 - )["layer"]
187 -
188 - assert kv_cache.shape == (num_blocks, 16, 1, 4)
189 - assert kv_cache.stride(0) == spec.page_size_bytes // 4
190 - assert kv_cache.stride(1) == 4
191 -
192 -
193 -class FakePerTokenScaleBackend:
194 - @staticmethod
195 - def get_kv_cache_shape(
196 - num_blocks: int,
197 - block_size: int,
198 - num_kv_heads: int,
199 - head_size: int,
200 - cache_dtype_str: str = "auto",
201 - ) -> tuple[int, ...]:
202 - return (num_blocks, 2, block_size, num_kv_heads, head_size + 4)
203 -
204 - @staticmethod
205 - def get_kv_cache_stride_order(
206 - include_num_layers_dimension: bool = False,
207 - ) -> tuple[int, ...]:
208 - assert not include_num_layers_dimension
209 - return (0, 1, 2, 3, 4)
142 + assert strides[1] == block_stride
143 + assert strides[2] == block_stride * num_blocks
144 + assert strides[0] == strides[2] * spec.num_heads
145
146
212 -def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
213 - num_blocks = 3
147 +def test_copy_kv_cache_blocks_separate_head_groups():
148 + # LHBNC stores each head group separately, so a block's bytes are scattered
149 + # across L*H regions.
150 + layout = KVCacheLayout.LHBNC
151 + num_blocks = 4
152 + num_layers = 2
153 spec = FullAttentionSpec(
215 - block_size=16,
216 - num_kv_heads=1,
217 - head_size=4,
218 - dtype=torch.int8,
219 - kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
220 - page_size_padded=384,
154 + block_size=2,
155 + num_kv_heads=2,
156 + head_size=2,
157 + dtype=torch.float32,
158 + num_head_slots=2,
159 + state_content_bytes=2 * 2 * 4,
160 + )
161 + raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8)
162 + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
163 +
164 + for layer_idx, cache in enumerate(caches):
165 + for block_idx in range(num_blocks):
166 + for head_idx in range(cache.shape[1]):
167 + cache[block_idx, head_idx].fill_(
168 + 100 * layer_idx + 10 * head_idx + block_idx
169 + )
170 +
171 + expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches]
172 + copy_kv_cache_blocks_inplace(
173 + caches,
174 + num_blocks,
175 + [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
176 )
222 - assert spec.real_page_size_bytes == 128
223 - assert spec.page_size_bytes == 384
224 -
225 - raw_tensors = {
226 - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
227 - }
228 - attn_groups = [
229 - AttentionGroup(
230 - backend=FakePerTokenScaleBackend,
231 - layer_names=["layer"],
232 - kv_cache_spec=spec,
233 - kv_cache_group_id=0,
234 - )
235 - ]
236 -
237 - kv_cache = _reshape_kv_cache(
238 - attn_groups,
239 - raw_tensors,
240 - "int8_per_token_head",
241 - [spec.block_size],
242 - {},
243 - )["layer"]
244 -
245 - assert kv_cache.shape == (num_blocks, 2, 16, 1, 8)
246 - assert kv_cache.stride(0) == spec.page_size_bytes
247 - assert kv_cache.stride(1) == 16 * 1 * 8
248 - assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8
249 -
250 -
251 -class FakeKVFirstBackend:
252 - """ROCm-style backend that puts K and V ahead of the block dim."""
253 -
254 - @staticmethod
255 - def get_kv_cache_shape(
256 - num_blocks: int,
257 - block_size: int,
258 - num_kv_heads: int,
259 - head_size: int,
260 - cache_dtype_str: str = "auto",
261 - ) -> tuple[int, ...]:
262 - return (2, num_blocks, block_size, num_kv_heads, head_size)
177
264 - @staticmethod
265 - def get_kv_cache_block_dim(
266 - block_size: int,
267 - num_kv_heads: int,
268 - head_size: int,
269 - cache_dtype_str: str = "auto",
270 - ) -> int:
271 - return 1
178 + for layer_idx, cache in enumerate(caches):
179 + torch.testing.assert_close(cache[2], expected[layer_idx][0])
180 + torch.testing.assert_close(cache[1], expected[layer_idx][1])
181
182
274 -def _kv_first_setup(shared_by: list[str]):
275 - num_blocks = 3
276 - attn_spec = FullAttentionSpec(
277 - block_size=16,
183 +@pytest.mark.parametrize(
184 + "layout,num_layers",
185 + [
186 + (KVCacheLayout.LBHNC, 2),
187 + # Splitting needs a manager block to be one dense page, which a
188 + # block-outermost layout only gives when the block holds one layer.
189 + (KVCacheLayout.BLHNC, 1),
190 + ],
191 +)
192 +def test_copy_kv_cache_blocks_with_virtual_block_splitting(
193 + layout: KVCacheLayout, num_layers: int
194 +):
195 + num_blocks = 4
196 + physical_per_logical = 2
197 + spec = FullAttentionSpec(
198 + block_size=4,
199 num_kv_heads=1,
200 head_size=2,
201 dtype=torch.float32,
202 )
282 - mamba_spec = MambaSpec(
283 - block_size=16,
284 - shapes=((64,),),
285 - dtypes=(torch.float32,),
203 + raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8)
204 + caches = dense_kv_cache_views(
205 + raw,
206 + spec,
207 + num_blocks,
208 + num_layers,
209 + layout,
210 + kernel_block_size=spec.block_size // physical_per_logical,
211 )
287 - assert attn_spec.page_size_bytes == mamba_spec.page_size_bytes == 256
212
289 - raw_tensor = torch.zeros(attn_spec.page_size_bytes * num_blocks, dtype=torch.int8)
290 - raw_tensors = {name: raw_tensor for name in shared_by}
291 - attn_groups = [
292 - AttentionGroup(
293 - backend=FakeKVFirstBackend,
294 - layer_names=["attn"],
295 - kv_cache_spec=attn_spec,
296 - kv_cache_group_id=0,
297 - )
298 - ]
299 - if "mamba" in shared_by:
300 - attn_groups.append(
301 - AttentionGroup(
302 - backend=FakeKVFirstBackend,
303 - layer_names=["mamba"],
304 - kv_cache_spec=mamba_spec,
305 - kv_cache_group_id=1,
306 - )
307 - )
308 - kv_cache_config = KVCacheConfig(
309 - num_blocks=num_blocks,
310 - kv_cache_tensors=[
311 - KVCacheTensor(size=raw_tensor.numel(), shared_by=list(shared_by))
312 - ],
313 - kv_cache_groups=[],
314 - )
213 + for layer_idx, cache in enumerate(caches):
214 + for block_idx in range(cache.shape[0]):
215 + cache[block_idx].fill_(100 * layer_idx + block_idx)
216 + expected = [[cache[i].clone() for i in range(cache.shape[0])] for cache in caches]
217
316 - kv_caches = _reshape_kv_cache(
317 - attn_groups,
318 - raw_tensors,
319 - "auto",
320 - [attn_spec.block_size] * len(attn_groups),
321 - {},
322 - kv_cache_config,
218 + copy_kv_cache_blocks_inplace(
219 + caches,
220 + num_blocks,
221 + [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
222 )
324 - return num_blocks, kv_caches["attn"]
325 -
326 -
327 -def test_reshape_kv_first_kv_cache_pages_blocks_when_shared_with_mamba():
328 - num_blocks, kv_cache = _kv_first_setup(["attn", "mamba"])
329 -
330 - assert kv_cache.shape == (2, num_blocks, 16, 1, 2)
331 - # Block b has to own page b, so its K and V sit side by side within it.
332 - page = 16 * 1 * 2 * 2
333 - for block in range(num_blocks):
334 - assert kv_cache[0, block].storage_offset() == block * page
335 - assert kv_cache[1, block].storage_offset() == block * page + page // 2
223
337 -
338 -def test_reshape_kv_first_kv_cache_keeps_layout_without_mamba():
339 - num_blocks, kv_cache = _kv_first_setup(["attn"])
340 -
341 - assert kv_cache.shape == (2, num_blocks, 16, 1, 2)
342 - # Nothing else indexes this allocation by page, so K and V stay split into
343 - # one contiguous half each.
344 - assert kv_cache[1, 0].storage_offset() == num_blocks * 16 * 1 * 2
224 + dst_start = 2 * physical_per_logical
225 + for layer_idx, cache in enumerate(caches):
226 + for physical_idx in range(physical_per_logical):
227 + torch.testing.assert_close(
228 + cache[dst_start + physical_idx], expected[layer_idx][physical_idx]
229 + )
tests/v1/worker/test_dsv4_packed_zeroer_geometry.py
+146 -173
@@ -1,35 +1,14 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""CPU regression for the packed-DSV4 KV zeroer geometry.
4 -
5 -Constructor-crossing regression for the packed-KV zeroing stride fix (upstream
6 -PR #50276): every number below comes from instantiating the REAL production
7 -code (no arithmetic is reimplemented here):
8 -
9 - - ``vllm.v1.core.kv_cache_utils._get_packed_kv_cache_layout`` (packed
10 - offsets / block stride)
11 - - ``vllm.v1.worker.gpu.attn_utils._reshape_attention_kv_cache``
12 - (per-layer column views)
13 - - ``vllm.models.deepseek_v4.sparse_mla.DeepseekV4FlashMLABackend``
14 - (``get_kv_cache_shape`` / ``get_kv_cache_block_dim``)
15 - - ``vllm.v1.worker.utils.AttentionGroup`` / ``KVBlockZeroer``
16 - (``KVBlockZeroer.__init__`` precomputes the segment tables)
17 -
18 -CPU only: no Triton launch, no CUDA tensor, no CUDA call is made. The test
19 -asserts the packed-DSV4 constructor metadata:
20 -
21 - * block stride = the full packed row (4 aligned pages)
22 - * per-layer zero = the meaningful page (the spec's unpadded
23 - ``real_page_size_bytes``), so adjacent packed layers are
24 - never wiped
25 - * 4 packed layer views; block 99 of the highest-offset layer is the last
26 - written region and stays within the packed backing.
27 -
28 -The pre-fix constructor exposed a five-field ``_meta`` with no block stride, so
29 -its zero width equals the FULL packed row; the extraction below records that
30 -legacy defect and then asserts the fixed semantics, so failure on the old code
31 -is a semantic one (wrong zero width / out-of-bounds block 99), not a
32 -schema-shape accident.
3 +"""CPU regression for the packed-DSV4 KV zeroer geometry (upstream #50276).
4 +
5 +Four DSV4 fp8_ds_mla layers share one cross-layer allocation: each block row packs
6 +the four alignment-padded pages side by side (layout BLHNC). The zeroer must step by
7 +the full packed row per block while zeroing only each layer's meaningful page — never
8 +the alignment padding, never an adjacent layer's page.
9 +
10 +CPU only: the zeroer's ``__init__`` precomputes the segment tables; no kernel is
11 +launched.
12 """
13
14 from types import SimpleNamespace
@@ -37,160 +16,154 @@ from types import SimpleNamespace
16 import pytest
17 import torch
18
40 -from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLABackend
41 -from vllm.v1.core.kv_cache_utils import _get_packed_kv_cache_layout
42 -from vllm.v1.kv_cache_interface import MLAAttentionSpec, UniformTypeKVCacheSpecs
43 -from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache
44 -from vllm.v1.worker.utils import AttentionGroup, KVBlockZeroer
19 +from tests.v1.attention.utils import dense_kv_cache_views
20 +from vllm.v1.kv_cache_interface import KVCacheLayout, MLAAttentionSpec
21 +from vllm.v1.worker.utils import (
22 + AttentionGroup,
23 + KVBlockZeroer,
24 + allocate_kv_cache,
25 +)
26
27 pytestmark = pytest.mark.cpu_test
28
29 NUM_BLOCKS = 100
49 -BLOCK_SIZE_CFG = 256 # vllm_config.cache_config.block_size for DSV4 (group 0)
50 -COMPRESS = 4 # main-MLA compress_ratio
51 -STORAGE_BS = BLOCK_SIZE_CFG // COMPRESS # 64
52 -KERNEL_BS = 256 # DeepseekV4FlashMLABackend supports [256]
53 -LAST_BLOCK = NUM_BLOCKS - 1 # block 99
54 -
55 -
56 -def _extract_segments(meta, base):
57 - """Read (offset_bytes, block_stride_bytes, zero_span_bytes) per segment.
58 -
59 - Accepts both the pre-fix five-field ``_meta`` (no block stride; the zero
60 - width is the FULL packed row -- the defect) and the fixed six-field ``_meta``.
61 - All values come from the real constructor output; nothing is recomputed. Returns
62 - ``(segs, legacy_why)`` where ``legacy_why`` describes the pre-fix defect
63 - when detected.
64 - """
65 - legacy_why = None
66 - if len(meta) == 6:
67 - seg_addrs, seg_block_strides, seg_page_sizes, _max_chunks, _blk, n_segs = meta
68 - else:
69 - seg_addrs, seg_page_sizes, _max_chunks, _blk, n_segs = meta
70 - seg_block_strides = seg_page_sizes
71 - legacy_why = (
72 - f"legacy {len(meta)}-field _meta (no block stride): zero width is "
73 - f"the FULL packed row ({int(seg_page_sizes[0]) * 4} B) -- the defect"
74 - )
75 - segs = [
76 - (int(a) - base, int(bs) * 4, int(ps) * 4)
77 - for a, bs, ps in zip(seg_addrs, seg_block_strides, seg_page_sizes)
78 - ]
79 - return segs, legacy_why, n_segs
80 -
81 -
82 -def test_dsv4_packed_zeroer_geometry():
83 - """Real DSV4 constructor metadata separates block stride from zero span."""
84 - layer_names = [f"model.layers.{i}.self_attn" for i in range(4)]
85 -
86 - specs = [
87 - MLAAttentionSpec(
88 - block_size=BLOCK_SIZE_CFG,
89 - num_kv_heads=1,
90 - head_size=512,
91 - dtype=torch.uint8,
92 - compress_ratio=COMPRESS,
93 - cache_dtype_str="fp8_ds_mla",
94 - alignment=576,
95 - model_version="deepseek_v4",
96 - state_content_bytes=584, # >576 to allocate room for scales at back of page
97 - )
98 - for _ in layer_names
99 - ]
100 - uniform = UniformTypeKVCacheSpecs.from_specs(
101 - {ln: sp for ln, sp in zip(layer_names, specs)}
30 +NUM_LAYERS = 4
31 +ALIGNMENT = 576
32 +
33 +
34 +def test_packed_dsv4_zeroer_zeroes_only_each_layers_page():
35 + spec = MLAAttentionSpec(
36 + block_size=256,
37 + num_kv_heads=1,
38 + head_size=512,
39 + dtype=torch.uint8,
40 + cache_dtype_str="fp8_ds_mla",
41 + model_version="deepseek_v4",
42 + tokens_per_state=4,
43 + alignment=ALIGNMENT,
44 + # DeepseekV4 fp8_ds_mla: 584B per token, published by the layer.
45 + state_content_bytes=584,
46 )
103 - assert uniform is not None
104 -
105 - # --- Real packed layout planner -------------------------------------------
106 - block_stride, layers_by_offset = _get_packed_kv_cache_layout(
107 - [
108 - AttentionGroup(
109 - backend=DeepseekV4FlashMLABackend,
110 - layer_names=layer_names,
111 - kv_cache_spec=uniform,
112 - kv_cache_group_id=0,
113 - )
114 - ]
115 - )
116 - page_bytes = specs[0].page_size_bytes # DSV4 fp8_ds_mla page after 576B alignment
117 - real_page_bytes = specs[0].real_page_size_bytes # unpadded meaningful page
118 - alignment_gap = page_bytes - real_page_bytes # trailing planner padding
119 - total_size = block_stride * NUM_BLOCKS
120 - offsets = sorted(layers_by_offset)
121 - assert block_stride == 4 * page_bytes, (block_stride, 4 * page_bytes)
122 - assert offsets == [i * page_bytes for i in range(4)]
123 -
124 - # --- Real per-layer views (packed branch) --------------------------------
125 - backing = torch.zeros(total_size, dtype=torch.uint8)
126 - base = backing.data_ptr()
127 - views = {}
128 - for ln, off in zip(layer_names, offsets):
129 - packing = (off, block_stride)
130 - kv_shape = DeepseekV4FlashMLABackend.get_kv_cache_shape(
131 - NUM_BLOCKS, STORAGE_BS, 1, 512, cache_dtype_str="fp8_ds_mla"
132 - )
133 - stride_order = tuple(range(len(kv_shape)))
134 - view = _reshape_attention_kv_cache(
135 - backing, specs[0], kv_shape, stride_order, NUM_BLOCKS, packing
136 - )
137 - views[ln] = view
138 - assert view.data_ptr() - base == off, "data_ptr must be base+offset"
139 - assert view.stride(0) * view.element_size() == block_stride, (
140 - "row stride == block_stride (logical block stride)"
141 - )
47 + # DSV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per stored state.
48 + assert spec.state_content_size_bytes == 584
49 + assert spec.num_states == 64
50 + unpadded_page = spec.unpadded_page_size_bytes
51 + padded_page = spec.page_size_bytes
52 + assert unpadded_page == 64 * 584
53 + assert padded_page > unpadded_page, "alignment must pad the page"
54 +
55 + raw = torch.zeros(NUM_BLOCKS * NUM_LAYERS * padded_page, dtype=torch.int8)
56 + views = dense_kv_cache_views(raw, spec, NUM_BLOCKS, NUM_LAYERS, KVCacheLayout.BLHNC)
57 + base = raw.data_ptr()
58 + block_row = NUM_LAYERS * padded_page
59 + for i, view in enumerate(views):
60 + assert view.data_ptr() - base == i * padded_page
61 + assert view.stride(0) * view.element_size() == block_row
62
143 - # --- Real KVBlockZeroer.__init__ (production arithmetic, CPU) ----------
144 - sctx = {ln: SimpleNamespace(kv_cache=v) for ln, v in views.items()}
63 zeroer = KVBlockZeroer(
64 torch.device("cpu"),
147 - attn_groups_iter=[
148 - AttentionGroup(
149 - backend=DeepseekV4FlashMLABackend,
150 - layer_names=layer_names,
151 - kv_cache_spec=specs[0],
152 - kv_cache_group_id=0,
153 - )
154 - ],
155 - kernel_block_sizes=[KERNEL_BS],
156 - cache_dtype="fp8_ds_mla",
157 - static_forward_context=sctx,
65 + attn_groups_iter=iter(
66 + [
67 + AttentionGroup(
68 + backend=None,
69 + layer_names=[f"layer.{i}" for i in range(NUM_LAYERS)],
70 + kv_cache_spec=spec,
71 + kv_cache_group_id=0,
72 + )
73 + ]
74 + ),
75 + kernel_block_sizes=[spec.block_size],
76 + static_forward_context={
77 + f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(NUM_LAYERS)
78 + },
79 )
159 - segs, legacy_why, n_segs = _extract_segments(zeroer._meta, base)
160 -
161 - # --- Fixed semantics (proven by metadata values, not by schema shape) ----
162 - assert n_segs == 4, "all 4 packed layers must register as segments"
163 - assert [off for off, _, _ in segs] == offsets, "one segment per layer offset"
164 - assert all(bs == block_stride for _, bs, _ in segs), (
165 - "block stride must be the full packed row (4 aligned pages)"
80 + seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
81 +
82 + assert n_segs == NUM_LAYERS
83 + # Segments step by the full packed row per block...
84 + assert (seg_block_strides * 4 == block_row).all()
85 + # ...but zero only the layer's meaningful page: no alignment padding, no
86 + # adjacent layer's page.
87 + assert (seg_page_sizes * 4 == unpadded_page).all()
88 + assert sorted(a - base for a in seg_addrs.tolist()) == [
89 + i * padded_page for i in range(NUM_LAYERS)
90 + ]
91 + # Block 99 of the highest-offset layer stays within the packed backing.
92 + last_end = (
93 + max(seg_addrs.tolist())
94 + + (NUM_BLOCKS - 1) * int(seg_block_strides[0]) * 4
95 + + int(seg_page_sizes[0]) * 4
96 )
167 - assert all(zs == real_page_bytes for _, _, zs in segs), (
168 - f"zero span must be the meaningful page ({real_page_bytes} B), got "
169 - f"{[zs for _, _, zs in segs]} B" + (f"; {legacy_why}" if legacy_why else "")
97 + assert last_end <= base + raw.numel()
98 +
99 +
100 +def test_overlaid_zeroer_dedups_segments_with_max_span():
101 + """Two groups overlay one allocation; the zeroer must emit one segment per distinct
102 + byte offset, spanning the widest overlaid page, so a newly allocated block is fully
103 + zeroed no matter which group owns it."""
104 + from unittest.mock import MagicMock
105 +
106 + from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups
107 + from vllm.v1.kv_cache_interface import (
108 + KVCacheGroupSpec,
109 + KVCacheLayout,
110 + UniformTypeKVCacheSpecs,
111 )
112
172 - # --- Block 99 (the real final block) bounds, via constructor metadata ---
173 - # Production write for block b of a segment is
174 - # [off + b*block_stride_bytes, off + b*block_stride_bytes + zero_span_bytes).
175 - worst = 0
176 - for off, bs, zs in segs:
177 - start = off + LAST_BLOCK * bs
178 - end = start + zs
179 - oob = end - total_size
180 - worst = max(worst, oob)
181 - assert end <= total_size, (
182 - f"segment@off={off} block {LAST_BLOCK} writes [{start}, {end}) "
183 - f"-> {oob} bytes past backing ({total_size})"
113 + def make_spec(head_size):
114 + return MLAAttentionSpec(
115 + block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8
116 )
185 - assert worst == 0, f"highest-offset final block goes {worst} bytes OOB"
186 -
187 - # Highest-offset segment's block 99 is the last written region and ends exactly at
188 - # total_size - alignment_gap: the trailing planner padding stays untouched.
189 - off_last, bs_last, zs_last = segs[-1]
190 - assert off_last + LAST_BLOCK * bs_last + zs_last == total_size - alignment_gap, (
191 - "highest-offset final block must end at total_size - alignment_gap "
192 - f"({alignment_gap}-byte gap intact)"
193 - )
117
195 - # Zero span < block stride proves adjacent packed layers are never wiped.
196 - assert max(zs for _, _, zs in segs) < block_stride
118 + g1_specs = {"g1.big": make_spec(512), "g1.small": make_spec(128)}
119 + g2_specs = {"g2.huge": make_spec(1024)}
120 + groups = [
121 + KVCacheGroupSpec(
122 + list(g1_specs),
123 + UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g1_specs),
124 + ),
125 + KVCacheGroupSpec(
126 + list(g2_specs),
127 + UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g2_specs),
128 + ),
129 + ]
130 + from vllm.config import CacheConfig
131 +
132 + vllm_config = MagicMock()
133 + vllm_config.cache_config = CacheConfig()
134 + vllm_config.cache_config.num_gpu_blocks_override = None
135 + vllm_config.cache_config.kv_cache_layout = "BLHNC"
136 + config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024)
137 + views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None)
138 + buf_ptr = views["g1.big"].data_ptr()
139 +
140 + attn_groups = [
141 + AttentionGroup(
142 + backend=None,
143 + layer_names=list(specs),
144 + kv_cache_spec=next(iter(specs.values())),
145 + kv_cache_group_id=gid,
146 + )
147 + for gid, specs in enumerate((g1_specs, g2_specs))
148 + ]
149 + zeroer = KVBlockZeroer(
150 + torch.device("cpu"),
151 + attn_groups_iter=iter(attn_groups),
152 + kernel_block_sizes=[64, 64],
153 + static_forward_context={
154 + name: SimpleNamespace(kv_cache=views[name]) for name in views
155 + },
156 + )
157 + seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
158 +
159 + # g1.big and g2.huge overlay at offset 0 -> one segment with g2's wider
160 + # span; g1.small keeps its own segment.
161 + assert n_segs == 2
162 + pages = {n: s.page_size_bytes for n, s in (g1_specs | g2_specs).items()}
163 + by_offset = {
164 + a - buf_ptr: p * 4 for a, p in zip(seg_addrs.tolist(), seg_page_sizes.tolist())
165 + }
166 + assert by_offset[0] == max(pages["g1.big"], pages["g2.huge"])
167 + assert by_offset[pages["g1.big"]] == pages["g1.small"]
168 + packed_block_stride = max(sum(pages[n] for n in g) for g in (g1_specs, g2_specs))
169 + assert (seg_block_strides * 4 == packed_block_stride).all()
tests/v1/worker/test_gpu_model_runner.py
+27 -61
@@ -92,7 +92,12 @@ def initialize_kv_cache(runner: GPUModelRunner):
92 kv_cache_config = KVCacheConfig(
93 num_blocks=NUM_BLOCKS,
94 kv_cache_tensors=[
95 - KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
95 + KVCacheTensor(
96 + size=tensor_size,
97 + layers=["layer.0"],
98 + layer_stride=tensor_size,
99 + block_stride=attn_spec.page_size_bytes,
100 + ),
101 ],
102 kv_cache_groups=[
103 KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
@@ -131,6 +136,7 @@ def get_vllm_config():
136 gpu_memory_utilization=0.9,
137 cache_dtype="auto",
138 )
139 + cache_config.kv_cache_layout = "LBNHC"
140 parallel_config = ParallelConfig()
141 vllm_config = VllmConfig(
142 model_config=model_config,
@@ -839,56 +845,6 @@ def test_update_states_pp_async_multi_request_keeps_rank_state_consistent(
845 )
846
847
842 -def test_kv_cache_stride_order(monkeypatch, model_runner):
843 - # This test checks if GPUModelRunner initializes correctly when an attention
844 - # backend enforces a non-default KV cache stride order.
845 - n_heads = model_runner.model_config.get_num_kv_heads(model_runner.parallel_config)
846 - head_size = model_runner.model_config.get_head_size()
847 -
848 - # Get the expected shape from the backend's get_kv_cache_shape method
849 - # to ensure compatibility with different backends (triton vs flexattention)
850 - attn_backend = None
851 - for attn_group in model_runner._attn_group_iterator():
852 - attn_backend = attn_group.backend
853 - break
854 -
855 - assert attn_backend is not None, "No attention backend found"
856 - expected_kv_cache_shape = list(
857 - attn_backend.get_kv_cache_shape(NUM_BLOCKS, BLOCK_SIZE, n_heads, head_size)
858 - )
859 -
860 - # TODO mla test
861 - default_stride = tuple(range(len(expected_kv_cache_shape)))
862 - non_default_stride = (*default_stride[1:], default_stride[0])
863 - # Permutation that gets you back to expected kv shape
864 - for test_stride in (non_default_stride, default_stride):
865 -
866 - def rnd_stride_order(
867 - include_num_layers_dimension: bool = False, test_stride=test_stride
868 - ):
869 - assert not include_num_layers_dimension
870 - return test_stride
871 -
872 - # Patch the attention backend class and re-trigger the KV cache creation
873 - for attn_group in model_runner._attn_group_iterator():
874 - attn_backend = attn_group.backend
875 - monkeypatch.setattr(
876 - attn_backend, "get_kv_cache_stride_order", rnd_stride_order
877 - )
878 -
879 - model_runner.attn_groups = []
880 - model_runner.kv_caches = []
881 - model_runner.initialize_kv_cache(model_runner.kv_cache_config)
882 -
883 - # Shape is unchanged, but layout may differ
884 - kv_cache_shape = model_runner.kv_caches[0].shape
885 - assert list(kv_cache_shape) == expected_kv_cache_shape
886 - if default_stride == test_stride:
887 - assert all(kv.is_contiguous() for kv in model_runner.kv_caches)
888 - else:
889 - assert all(not kv.is_contiguous() for kv in model_runner.kv_caches)
890 -
891 -
848 def test_update_config(model_runner):
849 # Simple update
850 model_runner.update_config({"load_config": {"load_format": "dummy"}})
@@ -1124,21 +1080,22 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):
1080 vllm_config, [kv_cache_spec], [available_memory]
1081 )[0]
1082 assert kv_cache_config.num_blocks == num_expected_blocks
1127 - assert len(kv_cache_config.kv_cache_tensors) == 2
1128 - assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2
1129 - assert kv_cache_config.kv_cache_tensors[1].size == available_memory // 2
1083 + assert len(kv_cache_config.kv_cache_tensors) == 1
1084 + assert kv_cache_config.kv_cache_tensors[0].size == available_memory
1085
1086 max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
1087 # max context len with KV sharing should be 2x as large as without
1088 assert max_context_len == 1310720
1089
1090 # important: override tensor size to prevent large mem alloc during test
1136 - # this will only allocate 2 block worth of memory (2 * 32kb)
1091 + # this will only allocate 1 block worth of memory per layer (2 layers * 32kb)
1092 kv_cache_config.num_blocks = 1
1093 for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
1139 - kv_cache_tensor.size = kv_cache_spec[
1140 - kv_cache_tensor.shared_by[0]
1141 - ].page_size_bytes
1094 + page_size = kv_cache_spec[kv_cache_tensor.layers[0]].page_size_bytes
1095 + kv_cache_tensor.size = page_size * len(kv_cache_tensor.layers)
1096 + # One block per layer: a layer's region is exactly one page.
1097 + kv_cache_tensor.layer_stride = page_size
1098 + kv_cache_tensor.block_stride = page_size
1099
1100 runner.initialize_kv_cache(kv_cache_config)
1101
@@ -1207,7 +1164,9 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
1164 # important: override tensor size to prevent large mem alloc during test
1165 # this will only allocate 1 block worth of memory (32kb)
1166 kv_cache_config.num_blocks = 1
1210 - kv_cache_config.kv_cache_tensors[0].size = kv_cache_spec[layer_0].page_size_bytes
1167 + page_size = kv_cache_spec[layer_0].page_size_bytes
1168 + kv_cache_config.kv_cache_tensors[0].size = page_size
1169 + kv_cache_config.kv_cache_tensors[0].layer_stride = page_size
1170
1171 runner.initialize_kv_cache(kv_cache_config)
1172 kv_cache_config_after_init = runner.kv_cache_config
@@ -1232,7 +1191,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
1191 """
1192 The GPU model runner creates different views into the
1193 KVCacheTensors for the attention and mamba layers
1235 - (via _reshape_kv_cache_tensors function). This test verifies
1194 + (via _allocate_kv_caches). This test verifies
1195 that the views are compatible: writing a mamba block
1196 will not corrupt an attention block and vice versa
1197 """
@@ -1270,6 +1229,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
1229 gpu_memory_utilization=0.9,
1230 cache_dtype="auto",
1231 )
1232 + cache_config.kv_cache_layout = "LBNHC"
1233 parallel_config = ParallelConfig()
1234 attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)
1235 vllm_config = VllmConfig(
@@ -1561,7 +1521,12 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
1521 kv_cache_config = KVCacheConfig(
1522 num_blocks=NUM_BLOCKS,
1523 kv_cache_tensors=[
1564 - KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
1524 + KVCacheTensor(
1525 + size=tensor_size,
1526 + layers=["layer.0"],
1527 + layer_stride=tensor_size,
1528 + block_stride=attn_spec.page_size_bytes,
1529 + ),
1530 ],
1531 kv_cache_groups=[
1532 KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
@@ -1726,6 +1691,7 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
1691 gpu_memory_utilization=0.9,
1692 cache_dtype="auto",
1693 )
1694 + cache_config.kv_cache_layout = "LBNHC"
1695 parallel_config = ParallelConfig()
1696 attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)
1697 vllm_config = VllmConfig(
tests/v1/worker/test_kv_block_zeroer.py
+48 -17
@@ -7,8 +7,11 @@ from types import SimpleNamespace
7 import pytest
8 import torch
9
10 +from tests.v1.attention.utils import dense_kv_cache_views
11 from vllm.v1.kv_cache_interface import (
12 ChunkedLocalAttentionSpec,
13 + FullAttentionSpec,
14 + KVCacheLayout,
15 SlidingWindowSpec,
16 )
17 from vllm.v1.worker import utils as worker_utils
@@ -19,12 +22,6 @@ from vllm.v1.worker.utils import (
22 )
23
24
22 -class _BlockFirstBackend:
23 - @staticmethod
24 - def get_kv_cache_block_dim(*args, **kwargs):
25 - return 0
26 -
27 -
25 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
26 @pytest.mark.parametrize(
27 "spec",
@@ -52,11 +49,8 @@ def test_attention_blocks_are_zeroed(spec):
49 layer_name = "draft.self_attn"
50 zeroer = KVBlockZeroer(
51 device,
55 - attn_groups_iter=[
56 - AttentionGroup(_BlockFirstBackend, [layer_name], spec, 0) # type: ignore[arg-type]
57 - ],
52 + attn_groups_iter=[AttentionGroup(None, [layer_name], spec, 0)],
53 kernel_block_sizes=[2],
59 - cache_dtype="fp8",
54 static_forward_context={
55 layer_name: SimpleNamespace(kv_cache=storage),
56 },
@@ -205,16 +199,10 @@ def test_large_dsv4_launch_geometry(monkeypatch):
199 zeroer = KVBlockZeroer(
200 device,
201 attn_groups_iter=[
208 - AttentionGroup(
209 - _BlockFirstBackend, # type: ignore[arg-type]
210 - [name],
211 - spec,
212 - group_id,
213 - )
202 + AttentionGroup(None, [name], spec, group_id)
203 for group_id, name in enumerate(layer_names)
204 ],
205 kernel_block_sizes=[1] * n_segs,
217 - cache_dtype="auto",
206 static_forward_context={
207 name: SimpleNamespace(kv_cache=storage)
208 for name, storage in storages.items()
@@ -311,3 +299,46 @@ def test_warmup_respects_available_block_count():
299 torch.accelerator.synchronize()
300
301 assert torch.all(storage == 1)
302 +
303 +
304 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
305 +@pytest.mark.parametrize("layout", list(KVCacheLayout))
306 +def test_zeroes_exactly_one_block_per_layer(layout: KVCacheLayout):
307 + """The zeroer must zero every byte of the target block in every layer and nothing
308 + outside it — per head-group region under LHBNC, and never past the target block's
309 + tile under block-major layouts (no out-of-bounds writes, no clobbering)."""
310 + device = torch.device("cuda")
311 + num_blocks, num_layers = 4, 2
312 + spec = FullAttentionSpec(
313 + block_size=4, num_kv_heads=2, head_size=8, dtype=torch.float32
314 + )
315 + raw = torch.empty(
316 + num_blocks * num_layers * spec.page_size_bytes,
317 + dtype=torch.int8,
318 + device=device,
319 + ).fill_(1)
320 + views = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
321 + groups = [
322 + AttentionGroup(
323 + backend=None,
324 + layer_names=[f"layer.{i}" for i in range(num_layers)],
325 + kv_cache_spec=spec,
326 + kv_cache_group_id=0,
327 + )
328 + ]
329 + ctx = {f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(num_layers)}
330 + zeroer = KVBlockZeroer(
331 + device,
332 + attn_groups_iter=iter(groups),
333 + kernel_block_sizes=[spec.block_size],
334 + static_forward_context=ctx,
335 + )
336 + zeroer.zero_block_ids([2])
337 + torch.accelerator.synchronize()
338 +
339 + for view in views:
340 + assert (view[2] == 0).all(), layout
341 + for b in (0, 1, 3):
342 + assert (view[b].view(torch.int8) == 1).all(), layout
343 + zero_bytes = int((raw == 0).sum().item())
344 + assert zero_bytes == num_layers * spec.page_size_bytes, layout
vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py
+1 -1
@@ -45,7 +45,7 @@ def fused_rope_unified_mla_kv_cache_update_impl(
45 cos_sin_cache,
46 is_neox,
47 layer_slot_mapping,
48 - kv_cache,
48 + kv_cache.squeeze(1),
49 kv_cache_dtype,
50 kv_cache_scale,
51 )
vllm/config/cache.py
+40
@@ -3,6 +3,7 @@
3
4 from collections.abc import Callable
5 from dataclasses import field
6 +from functools import cache
7 from typing import Any, ClassVar, Literal
8
9 from pydantic import Field, field_validator, model_validator
@@ -13,9 +14,28 @@ from vllm.utils.torch_utils import (
14 is_quantized_kv_cache,
15 kv_cache_uses_per_token_head_scales,
16 )
17 +from vllm.v1.kv_cache_layout import KVCacheLayout
18
19 logger = init_logger(__name__)
20
21 +_LAYOUT_COMPAT_ALIASES = {
22 + "NHD": "LBNHC",
23 + "HND": "LBHNC",
24 +}
25 +
26 +
27 +@cache
28 +def _layout_from_name(layout_name: str) -> KVCacheLayout:
29 + layout_name = _LAYOUT_COMPAT_ALIASES.get(layout_name, layout_name)
30 + try:
31 + return KVCacheLayout[layout_name]
32 + except KeyError:
33 + raise ValueError(
34 + f"Unknown KV cache layout {layout_name!r}. "
35 + f"Valid layouts: {[m.name for m in KVCacheLayout]}"
36 + ) from None
37 +
38 +
39 CacheDType = Literal[
40 "auto",
41 "float16",
@@ -65,6 +85,17 @@ class CacheConfig:
85 """Whether block_size was explicitly provided. Derived automatically."""
86 user_specified_mamba_block_size: bool = field(default=False, init=False)
87 """Whether mamba_block_size was explicitly provided. Derived automatically."""
88 + kv_cache_layout: str | None = field(default=None, init=False)
89 + """Resolved physical KV cache layout name (a ``KVCacheLayout`` member).
90 +
91 + ``None`` means the layout has not been resolved yet. The engine core
92 + resolves it once (``resolve_kv_cache_layout``) before memory profiling, and
93 + every worker process adopts the resolved name — via the
94 + ``set_kv_cache_layout`` RPC, or ``KVCacheConfig.kv_cache_layout`` for
95 + workers spawned after resolution — before the KV cache is allocated. Once
96 + set the value is final and read with ``get_resolved_kv_cache_layout``,
97 + which raises on ``None``. Tests and standalone tools may pre-set a value,
98 + which resolution then honors as-is."""
99 prefix_match_unit: int | None = Field(default=None, gt=0)
100 """The finest token boundary (in tokens) a prefix-cache hit can land on.
101
@@ -309,3 +340,12 @@ class CacheConfig:
340 str(cache_dtype),
341 )
342 return cache_dtype
343 +
344 + def get_resolved_kv_cache_layout(self) -> KVCacheLayout:
345 + if self.kv_cache_layout is None:
346 + raise ValueError(
347 + "KV cache layout has not been resolved yet; it is resolved once "
348 + "by the engine core (resolve_kv_cache_layout) unless explicitly "
349 + "set by the user."
350 + )
351 + return _layout_from_name(self.kv_cache_layout)
vllm/distributed/kv_transfer/kv_connector/utils.py
+11 -71
@@ -12,7 +12,7 @@ import torch
12
13 from vllm.config import (
14 VllmConfig,
15 - get_current_vllm_config,
15 + get_current_vllm_config_or_none,
16 get_layers_from_vllm_config,
17 set_current_vllm_config,
18 )
@@ -35,20 +35,20 @@ EngineId = str
35 BlockIds = tuple[list[int], ...] | list[list[int]]
36
37
38 -def get_kv_connector_cache_layout():
39 - # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is
38 +def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None):
39 + # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout is
40 # used for faster transfer.
41 - vllm_config = get_current_vllm_config()
41 + if vllm_config is None:
42 + vllm_config = get_current_vllm_config_or_none()
43 + if vllm_config is None:
44 + return None
45 kv_config = vllm_config.kv_transfer_config
46 if kv_config is not None:
47 connector_cls = KVConnectorFactory.get_connector_class(kv_config)
48 required_kvcache_layout = connector_cls.get_required_kvcache_layout(vllm_config)
49 if required_kvcache_layout is not None:
50 return required_kvcache_layout
48 - logger.info_once(
49 - "Connectors do not specify a kv cache layout, defaulting to NHD."
50 - )
51 - return "NHD"
51 + return None
52
53
54 class KVOutputAggregator:
@@ -288,11 +288,11 @@ def kv_postprocess_layout_on_receive(cache, indices):
288
289 def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio):
290 """
291 - Transforms the layout of received KV cache to the local block_size and HND.
291 + Transforms the layout of received KV cache to the local block_size and LBHNC.
292 (Only works for local blocksize > remote blocksize)
293
294 - prefill is HND, smaller block_size
295 - decode(local) is NHD, larger block_size
294 + prefill is LBHNC, smaller block_size
295 + decode(local) is LBNHC, larger block_size
296 """
297 blocks_to_update = cache.index_select(0, indices)
298
@@ -427,48 +427,6 @@ class TransferTopology:
427
428 self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {}
429
430 - # Figure out whether the first dimension of the cache is K/V
431 - # or num_blocks.
432 - attn_backend = self.attn_backends[0]
433 - if not self.is_mamba:
434 - _MOCK_BLOCK_SIZE = 16
435 - kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(
436 - num_blocks=1,
437 - block_size=_MOCK_BLOCK_SIZE,
438 - num_kv_heads=1,
439 - head_size=1,
440 - )
441 - logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
442 - assert kv_cache_shape[0] == 1, (
443 - "KV cache layout must be blocks-first; expected mocked "
444 - f"num_blocks=1 in leading dim, got shape {kv_cache_shape}."
445 - )
446 - if not self.is_mla:
447 - assert len(kv_cache_shape) == 4, (
448 - "Attention KV cache layout must be standardized as "
449 - "[num_blocks, num_kv_heads, block_size, content_size], "
450 - f"got shape {kv_cache_shape}."
451 - )
452 -
453 - self._cross_layers_blocks = False
454 - if self.tensor_shape is not None:
455 - self._cross_layers_blocks = (
456 - len(self.tensor_shape) == len(kv_cache_shape) + 1
457 - )
458 -
459 - if self._cross_layers_blocks:
460 - logger.debug("Using cross-layer KV cache")
461 - _MOCK_NUM_LAYERS = 80
462 - kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
463 - try:
464 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
465 - include_num_layers_dimension=self._cross_layers_blocks
466 - )
467 - except (AttributeError, NotImplementedError):
468 - assert self.tensor_shape is not None
469 - kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
470 - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
471 -
430 # ============================================================
431 # Engine registration
432 # ============================================================
@@ -503,24 +461,6 @@ class TransferTopology:
461 for key in [k for k in self._engines if k[0] == remote_engine_id]:
462 del self._engines[key]
463
506 - # ============================================================
507 - # Layout properties
508 - # ============================================================
509 -
510 - @property
511 - def cross_layers_blocks(self) -> bool:
512 - return self._cross_layers_blocks
513 -
514 - @property
515 - def virtually_split_kv_in_blocks(self) -> bool:
516 - # Whether to logically split each block into two separately-indexable
517 - # sub-regions. With K and V packed into the content dim, an attention
518 - # block transfers as a single unit — no K/V sub-split is needed. Only
519 - # Mamba still needs this, to index its two state regions (conv/ssm)
520 - # separately. Not applicable to cross-layer blocks (per-layer
521 - # interleaving means a simple half-split does not separate the parts).
522 - return self.is_mamba and not self._cross_layers_blocks
523 -
464 # ============================================================
465 # Common methods
466 # ============================================================
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+1 -26
@@ -48,7 +48,7 @@ from typing import TYPE_CHECKING, Any, Literal
48 import torch
49
50 from vllm.logger import init_logger
51 -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
51 +from vllm.v1.attention.backend import AttentionMetadata
52 from vllm.v1.core.sched.output import SchedulerOutput
53 from vllm.v1.outputs import KVConnectorOutput
54
@@ -182,14 +182,6 @@ class KVConnectorBase_V1(ABC):
182 """
183 return False
184
185 - @property
186 - def prefer_cross_layer_blocks(self) -> bool:
187 - """
188 - Indicates whether this connector prefers KV blocks that hold KV data for all
189 - layers, which can speed up KV data transfers. Defaults to False.
190 - """
191 - return False
192 -
185 @property
186 def requires_kv_delivery(self) -> bool:
187 """Whether this connector hands off KV that must be reliably delivered.
@@ -279,23 +271,6 @@ class KVConnectorBase_V1(ABC):
271 """
272 return
273
282 - def register_cross_layers_kv_cache(
283 - self, kv_cache: torch.Tensor, attn_backend: type["AttentionBackend"]
284 - ):
285 - """
286 - Initialize with a single KV cache tensor used by all layers.
287 - The first dimension should be num_layers.
288 - This function will only be called for models with uniform layers,
289 - and only if the prefers_cross_layer_blocks is set to True.
290 - Only one of the functions
291 - {register_kv_caches, register_cross_layers_kv_cache} will be called.
292 -
293 - Args:
294 - kv_cache: a cross-layers kv cache tensor
295 - attn_backend: The attention backend that corresponds to all layers
296 - """
297 - return
298 -
274 def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
275 """
276 Set the xPU-specific ops for copying KV between host and device.
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+7 -11
@@ -129,20 +129,17 @@ class ExampleConnector(KVConnectorBase_V1):
129 """Inject the KV cache into the layer.
130
131 Args:
132 - dst_kv_cache_layer (torch.Tensor): the destination KV cache
133 - layer. In shape [num_pages, page_size, xxx] for MLA,
134 - [num_pages, 2, page_size, xxx] otherwise.
132 + dst_kv_cache_layer (torch.Tensor): the destination KV cache layer,
133 + a standardized [B, H, N, C] per-layer view (H == 1 for MLA).
134 src_kv_cache (torch.Tensor): the source KV cache.
135 slot_mapping (torch.Tensor): the slot mapping. In shape
136 [num_tokens].
137 """
138 slot_mapping = slot_mapping.to(dst_kv_cache_layer.device, non_blocking=True)
139 if isinstance(attn_metadata, MLACommonMetadata):
141 - dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
142 - num_pages = dst_kv_cache_layer_shape[0]
143 - page_size = dst_kv_cache_layer_shape[1]
140 + # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots.
141 dst_kv_cache_layer = dst_kv_cache_layer.reshape(
145 - num_pages * page_size, -1
142 + -1, dst_kv_cache_layer.shape[-1]
143 )
144 dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
145 else:
@@ -225,13 +222,12 @@ class ExampleConnector(KVConnectorBase_V1):
222 ) -> torch.Tensor:
223 """Extract the KV cache from the layer.
224
228 - Assume the shape of the layer is (num_pages, page_size, xxx)
229 - for MLA, and (num_pages, 2, page_size, xxx) otherwise.
225 + The layer is a standardized [B, H, N, C] per-layer view (H == 1 for MLA).
226 """
227 slot_mapping = slot_mapping.to(layer.device, non_blocking=True)
228 if isinstance(attn_metadata, MLACommonMetadata):
233 - num_pages, page_size = layer.shape[0], layer.shape[1]
234 - return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]
229 + # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots.
230 + return layer.reshape(-1, layer.shape[-1])[slot_mapping, ...]
231 block_idxs = slot_mapping // self._block_size
232 offsets = slot_mapping % self._block_size
233 return layer[block_idxs, :, offsets]
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+13 -18
@@ -38,8 +38,10 @@ def extract_from_kv_cache(
38 num_tokens: int,
39 ) -> torch.Tensor:
40 """Extract data from KV cache."""
41 - block_size = kv_cache.shape[1]
42 - return kv_cache[slot_mapping // block_size, slot_mapping % block_size][:num_tokens]
41 + block_size = kv_cache.shape[2]
42 + return kv_cache[slot_mapping // block_size, :, slot_mapping % block_size][
43 + :num_tokens
44 + ]
45
46
47 def load_hidden_states(path: str) -> dict[str, torch.Tensor]:
@@ -100,15 +102,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
102 Must be used in conjunction with the `extract_hidden_states` spec decoding method.
103 """
104
103 - @property
104 - def prefer_cross_layer_blocks(self) -> bool:
105 - """
106 - Indicates whether this connector prefers KV blocks that hold KV data for all
107 - layers, which can speed up KV data transfers. Defaults to False.
108 - """
109 - # Must be False so that drafter kv cache isn't merged with verifier's
110 - return False
111 -
105 @classmethod
106 def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
107 """Index of the KV cache group holding the extracted hidden states.
@@ -310,10 +303,12 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
303
304 # Block size must match the indexed buffer, else reads hit the wrong
305 # slots. Raise (not assert) so the check survives `python -O`.
313 - if self._block_size != self._kv_cache.shape[1]:
306 + # Views are [num_blocks, num_heads, block_size, head_size], matching what
307 + # extract_from_kv_cache() indexes.
308 + if self._block_size != self._kv_cache.shape[2]:
309 raise ValueError(
310 f"Hidden-states block-size mismatch: derived {self._block_size} "
316 - f"but buffer block size is {self._kv_cache.shape[1]}; read slots "
311 + f"but buffer block size is {self._kv_cache.shape[2]}; read slots "
312 "would be wrong (likely a hybrid block-size resolution bug)."
313 )
314
@@ -605,9 +600,9 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
600 "get_required_kvcache_layout should not be called "
601 "on the abstract base class"
602 )
608 - # NHD means we have (num_tokens, num_heads)
609 - # HND means we have (num_heads, num_tokens)
610 - # For now, we only support NHD layout since this keeps the
603 + # LBNHC means we have (num_tokens, num_heads)
604 + # LBHNC means we have (num_heads, num_tokens)
605 + # For now, we only support LBNHC layout since this keeps the
606 # hidden states for each token together in memory.
612 - # HND is primarily used when sharding heads across devices.
613 - return "NHD"
607 + # LBHNC is primarily used when sharding heads across devices.
608 + return "LBNHC"
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py
+31 -28
@@ -541,35 +541,32 @@ class HF3FSKVConnector(KVConnectorBase_V1):
541 self._dtype = first_cache.dtype
542 element_size = first_cache.element_size()
543
544 - if self._use_mla:
545 - assert len(first_cache.shape) == 3, "MLA format should have 3 dimensions"
546 - # MLA format: [num_blocks, block_size, head_size]
547 - num_blocks, block_size, head_size = first_cache.shape
548 - num_heads = 1
549 - else:
550 - # MHA format: [2, num_blocks, block_size, num_heads, head_size]
551 - _, num_blocks, block_size, num_heads, head_size = first_cache.shape
544 + # Standardized per-layer [B, H, N, C] view; MLA is just H == 1 with the
545 + # latent vector as the content dim.
546 + assert len(first_cache.shape) == 4, (
547 + f"expected a [B, H, N, C] KV cache view, got {tuple(first_cache.shape)}"
548 + )
549 + num_blocks, num_heads, block_size, content_dim = first_cache.shape
550 + cache_strides = first_cache.stride()
551 + assert all(
552 + cache.shape == first_cache.shape and cache.stride() == cache_strides
553 + for cache in self._kv_caches.values()
554 + ), "HF3FS requires uniform KV cache shapes and strides."
555
556 self._local_total_tokens = num_blocks * block_size
557 self._local_block_size = block_size
555 -
556 - if self._use_mla:
557 - layer_block_size = block_size * head_size * element_size
558 - self._bytes_per_page = layer_block_size * len(self._kv_caches)
559 - self._shape_per_page = [
560 - len(self._kv_caches),
561 - block_size,
562 - head_size,
563 - ]
564 - else:
565 - layer_block_size = 2 * block_size * num_heads * head_size * element_size
566 - self._bytes_per_page = layer_block_size * len(self._kv_caches)
567 - self._shape_per_page = [
568 - len(self._kv_caches),
569 - 2,
570 - block_size,
571 - num_heads * head_size,
572 - ]
558 + self._num_heads = num_heads
559 + self._content_dim = content_dim
560 + self._kv_cache_strides = cache_strides
561 +
562 + layer_block_size = num_heads * block_size * content_dim * element_size
563 + self._bytes_per_page = layer_block_size * len(self._kv_caches)
564 + self._shape_per_page = [
565 + len(self._kv_caches),
566 + num_heads,
567 + block_size,
568 + content_dim,
569 + ]
570
571 self._kvcache_ptrs = torch.tensor(
572 [cache.data_ptr() for cache in self._kv_caches.values()],
@@ -993,7 +990,10 @@ class HF3FSKVConnector(KVConnectorBase_V1):
990 self._local_total_tokens,
991 buffer_tensor,
992 token_indices,
996 - is_mla=self._use_mla,
993 + self._local_block_size,
994 + self._num_heads,
995 + self._content_dim,
996 + self._kv_cache_strides,
997 )
998 else:
999 gather_scatter_helper.scatter_kv_caches(
@@ -1001,7 +1001,10 @@ class HF3FSKVConnector(KVConnectorBase_V1):
1001 self._local_total_tokens,
1002 buffer_tensor,
1003 token_indices,
1004 - is_mla=self._use_mla,
1004 + self._local_block_size,
1005 + self._num_heads,
1006 + self._content_dim,
1007 + self._kv_cache_strides,
1008 )
1009
1010 def _compute_prefix_hash(
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py
+92 -121
@@ -12,10 +12,15 @@ def kv_cache_scatter_kernel(
12 source_ptr,
13 token_indices_ptr,
14 num_tokens_in_block,
15 - hidden_size,
15 total_token_in_kvcache,
16 num_layers,
18 - is_mla,
17 + tokens_per_block,
18 + block_stride,
19 + head_stride,
20 + token_stride,
21 + content_stride,
22 + num_heads: tl.constexpr,
23 + content_size: tl.constexpr,
24 BLOCK_SIZE: tl.constexpr,
25 ):
26 layer_idx = tl.program_id(0)
@@ -30,39 +35,25 @@ def kv_cache_scatter_kernel(
35 if token_idx >= total_token_in_kvcache:
36 return
37
33 - if is_mla:
34 - # MLA format: source [num_layers, num_tokens_in_block, hidden_size]
35 - # MLA format: target [total_token_in_kvcache, hidden_size] (per layer)
36 - source_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
37 - target_offset = token_idx * hidden_size
38 -
39 - for i in range(0, hidden_size, BLOCK_SIZE):
40 - offset = i + tl.arange(0, BLOCK_SIZE)
41 - mask = offset < hidden_size
42 - val = tl.load(source_ptr + source_offset + offset, mask=mask)
43 - tl.store(kv_cache_ptr + target_offset + offset, val, mask=mask)
44 - else:
45 - # MHA format: source [num_layers, 2, num_tokens_in_block, hidden_size]
46 - # MHA format: target [2, total_token_in_kvcache, hidden_size]
47 - source_offset_k = (
48 - layer_idx * num_tokens_in_block * 2 + token_pos
49 - ) * hidden_size
50 - source_offset_v = (
51 - layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
52 - ) * hidden_size
53 -
54 - target_offset_k = token_idx * hidden_size
55 - target_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
56 -
57 - for i in range(0, hidden_size, BLOCK_SIZE):
58 - offset = i + tl.arange(0, BLOCK_SIZE)
59 - mask = offset < hidden_size
60 -
61 - val_k = tl.load(source_ptr + source_offset_k + offset, mask=mask)
62 - val_v = tl.load(source_ptr + source_offset_v + offset, mask=mask)
63 -
64 - tl.store(kv_cache_ptr + target_offset_k + offset, val_k, mask=mask)
65 - tl.store(kv_cache_ptr + target_offset_v + offset, val_v, mask=mask)
38 + block_idx = token_idx // tokens_per_block
39 + state_idx = token_idx % tokens_per_block
40 + hidden_size = num_heads * content_size
41 + for i in range(0, hidden_size, BLOCK_SIZE):
42 + offset = i + tl.arange(0, BLOCK_SIZE)
43 + mask = offset < hidden_size
44 + head_idx = offset // content_size
45 + content_idx = offset % content_size
46 + source_offset = (
47 + (layer_idx * num_heads + head_idx) * num_tokens_in_block + token_pos
48 + ) * content_size + content_idx
49 + target_offset = (
50 + block_idx * block_stride
51 + + head_idx * head_stride
52 + + state_idx * token_stride
53 + + content_idx * content_stride
54 + )
55 + value = tl.load(source_ptr + source_offset, mask=mask)
56 + tl.store(kv_cache_ptr + target_offset, value, mask=mask)
57
58
59 @triton.jit
@@ -71,10 +62,15 @@ def kv_cache_gather_kernel(
62 dst_ptr,
63 token_indices_ptr,
64 num_tokens_in_block,
74 - hidden_size,
65 total_token_in_kvcache,
66 num_layers,
77 - is_mla,
67 + tokens_per_block,
68 + block_stride,
69 + head_stride,
70 + token_stride,
71 + content_stride,
72 + num_heads: tl.constexpr,
73 + content_size: tl.constexpr,
74 BLOCK_SIZE: tl.constexpr,
75 ):
76 layer_idx = tl.program_id(0)
@@ -89,37 +85,25 @@ def kv_cache_gather_kernel(
85 if token_idx >= total_token_in_kvcache:
86 return
87
92 - if is_mla:
93 - # MLA format: source [total_token_in_kvcache, hidden_size] (per layer)
94 - # MLA format: dst [num_layers, num_tokens_in_block, hidden_size]
95 - kvcache_offset = token_idx * hidden_size
96 - dst_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
97 -
98 - for i in range(0, hidden_size, BLOCK_SIZE):
99 - offset = i + tl.arange(0, BLOCK_SIZE)
100 - mask = offset < hidden_size
101 - val = tl.load(kv_cache_ptr + kvcache_offset + offset, mask=mask)
102 - tl.store(dst_ptr + dst_offset + offset, val, mask=mask)
103 - else:
104 - # MHA format: source [2, total_token_in_kvcache, hidden_size]
105 - # MHA format: dst [num_layers, 2, num_tokens_in_block, hidden_size]
106 - dst_offset_k = (layer_idx * num_tokens_in_block * 2 + token_pos) * hidden_size
107 - dst_offset_v = (
108 - layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
109 - ) * hidden_size
110 -
111 - kvcache_offset_k = token_idx * hidden_size
112 - kvcache_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
113 -
114 - for i in range(0, hidden_size, BLOCK_SIZE):
115 - offset = i + tl.arange(0, BLOCK_SIZE)
116 - mask = offset < hidden_size
117 -
118 - val_k = tl.load(kv_cache_ptr + kvcache_offset_k + offset, mask=mask)
119 - val_v = tl.load(kv_cache_ptr + kvcache_offset_v + offset, mask=mask)
120 -
121 - tl.store(dst_ptr + dst_offset_k + offset, val_k, mask=mask)
122 - tl.store(dst_ptr + dst_offset_v + offset, val_v, mask=mask)
88 + block_idx = token_idx // tokens_per_block
89 + state_idx = token_idx % tokens_per_block
90 + hidden_size = num_heads * content_size
91 + for i in range(0, hidden_size, BLOCK_SIZE):
92 + offset = i + tl.arange(0, BLOCK_SIZE)
93 + mask = offset < hidden_size
94 + head_idx = offset // content_size
95 + content_idx = offset % content_size
96 + kvcache_offset = (
97 + block_idx * block_stride
98 + + head_idx * head_stride
99 + + state_idx * token_stride
100 + + content_idx * content_stride
101 + )
102 + dst_offset = (
103 + (layer_idx * num_heads + head_idx) * num_tokens_in_block + token_pos
104 + ) * content_size + content_idx
105 + value = tl.load(kv_cache_ptr + kvcache_offset, mask=mask)
106 + tl.store(dst_ptr + dst_offset, value, mask=mask)
107
108
109 def scatter_kv_caches(
@@ -127,34 +111,33 @@ def scatter_kv_caches(
111 total_token_in_kvcache: int,
112 src_tensor: torch.Tensor,
113 token_indices: list[int],
130 - is_mla: bool = False,
114 + tokens_per_block: int,
115 + num_heads: int,
116 + content_size: int,
117 + kv_cache_strides: tuple[int, ...],
118 ) -> None:
119 """Scatter KV cache data from source tensor to KV cache storage.
120
121 Args:
122 kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
123 total_token_in_kvcache: Total number of tokens in KV cache
137 - src_tensor: Source tensor containing data to scatter
138 - - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
139 - - MLA format: [num_layers, num_tokens_in_block, hidden_size]
124 + src_tensor: Source ``[L, H, N, C]`` tensor containing data to scatter
125 token_indices: List of token positions to update
141 - is_mla: Whether using MLA model format
126 + tokens_per_block: Number of stored states in each cache block
127 + num_heads: Size of the H axis
128 + content_size: Size of the C axis
129 + kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view
130 """
131 num_layers = len(kv_caches_ptrs)
132 num_tokens_in_block = len(token_indices)
133
146 - if is_mla:
147 - # MLA: src_tensor is [num_layers, num_tokens_in_block, hidden_size]
148 - assert len(src_tensor.shape) == 3, (
149 - f"MLA src_tensor should be 3D, got {src_tensor.shape}"
150 - )
151 - hidden_size = src_tensor.shape[2]
152 - else:
153 - # MHA: src_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
154 - assert len(src_tensor.shape) == 4, (
155 - f"MHA src_tensor should be 4D, got {src_tensor.shape}"
156 - )
157 - hidden_size = src_tensor.shape[3]
134 + assert src_tensor.shape == (
135 + num_layers,
136 + num_heads,
137 + num_tokens_in_block,
138 + content_size,
139 + )
140 + assert len(kv_cache_strides) == 4
141
142 device = src_tensor.device
143 token_indices_tensor = torch.tensor(
@@ -169,10 +152,12 @@ def scatter_kv_caches(
152 src_tensor,
153 token_indices_tensor,
154 num_tokens_in_block,
172 - hidden_size,
155 total_token_in_kvcache,
156 num_layers,
175 - is_mla,
157 + tokens_per_block,
158 + *kv_cache_strides,
159 + num_heads=num_heads,
160 + content_size=content_size,
161 BLOCK_SIZE=BLOCK_SIZE,
162 )
163
@@ -182,49 +167,33 @@ def gather_kv_caches(
167 total_token_in_kvcache: int,
168 dst_tensor: torch.Tensor,
169 token_indices: list[int],
185 - is_mla: bool = False,
170 + tokens_per_block: int,
171 + num_heads: int,
172 + content_size: int,
173 + kv_cache_strides: tuple[int, ...],
174 ) -> None:
175 """Gather KV cache data from KV cache storage to destination tensor.
176
177 Args:
178 kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
179 total_token_in_kvcache: Total number of tokens in KV cache
192 - dst_tensor: Destination tensor to store gathered data
193 - - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
194 - - MLA format: [num_layers, num_tokens_in_block, hidden_size]
180 + dst_tensor: Destination ``[L, H, N, C]`` tensor
181 token_indices: List of token positions to gather
196 - is_mla: Whether using MLA model format
182 + tokens_per_block: Number of stored states in each cache block
183 + num_heads: Size of the H axis
184 + content_size: Size of the C axis
185 + kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view
186 """
187 num_layers = kv_caches_ptrs.shape[0]
188 num_tokens_in_block = len(token_indices)
189
201 - if is_mla:
202 - # MLA: dst_tensor is [num_layers, num_tokens_in_block, hidden_size]
203 - assert len(dst_tensor.shape) == 3, (
204 - f"MLA dst_tensor should be 3D, got {dst_tensor.shape}"
205 - )
206 - assert dst_tensor.shape[0] == num_layers, (
207 - f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
208 - )
209 - assert dst_tensor.shape[1] == num_tokens_in_block, (
210 - f"Token count mismatch: {dst_tensor.shape[1]} vs {num_tokens_in_block}"
211 - )
212 - hidden_size = dst_tensor.shape[2]
213 - else:
214 - # MHA: dst_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
215 - assert len(dst_tensor.shape) == 4, (
216 - f"MHA dst_tensor should be 4D, got {dst_tensor.shape}"
217 - )
218 - assert dst_tensor.shape[0] == num_layers, (
219 - f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
220 - )
221 - assert dst_tensor.shape[1] == 2, (
222 - f"MHA should have 2 (K,V) components, got {dst_tensor.shape[1]}"
223 - )
224 - assert dst_tensor.shape[2] == num_tokens_in_block, (
225 - f"Token count mismatch: {dst_tensor.shape[2]} vs {num_tokens_in_block}"
226 - )
227 - hidden_size = dst_tensor.shape[3]
190 + assert dst_tensor.shape == (
191 + num_layers,
192 + num_heads,
193 + num_tokens_in_block,
194 + content_size,
195 + )
196 + assert len(kv_cache_strides) == 4
197
198 device = dst_tensor.device
199 token_indices_tensor = torch.tensor(
@@ -239,10 +208,12 @@ def gather_kv_caches(
208 dst_tensor,
209 token_indices_tensor,
210 num_tokens_in_block,
242 - hidden_size,
211 total_token_in_kvcache,
212 num_layers,
245 - is_mla,
213 + tokens_per_block,
214 + *kv_cache_strides,
215 + num_heads=num_heads,
216 + content_size=content_size,
217 BLOCK_SIZE=BLOCK_SIZE,
218 )
219
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+32 -71
@@ -49,15 +49,15 @@ from vllm.model_executor.models.utils import extract_layer_index
49 from vllm.platforms import current_platform
50 from vllm.utils.math_utils import cdiv
51 from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
52 +from vllm.utils.torch_utils import is_non_overlapping_and_dense
53 from vllm.v1.attention.backend import AttentionMetadata
53 -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, get_kv_cache_layout
54 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
55 from vllm.v1.core.sched.output import SchedulerOutput
56 from vllm.v1.kv_cache_interface import (
57 + AttentionSpec,
58 FullAttentionSpec,
59 KVCacheSpec,
60 MambaSpec,
59 - MLAAttentionSpec,
60 - SlidingWindowMLASpec,
61 SlidingWindowSpec,
62 )
63 from vllm.v1.request import RequestStatus
@@ -122,9 +122,7 @@ def _expand_transfer_regions(
122 kv_block_lens: list[int],
123 layer_names: list[str],
124 layer_indices: list[int],
125 - is_kv_layout_blocks_first: bool,
125 group_indices: list[int] | None = None,
127 - split_kv_regions: list[bool] | None = None,
126 ) -> list[TransferRegion]:
127 """Expand registered KV tensors into the regions transferred by Mooncake."""
128 assert (
@@ -146,13 +144,6 @@ def _expand_transfer_regions(
144 "Mooncake transfer regions require matching group metadata lengths, "
145 f"got group_indices={len(group_indices)}, layer_names={len(layer_names)}."
146 )
149 - if split_kv_regions is None:
150 - split_kv_regions = [is_kv_layout_blocks_first] * len(layer_names)
151 - assert len(split_kv_regions) == len(layer_names), (
152 - "Mooncake transfer regions require matching split metadata, "
153 - f"got split_kv_regions={len(split_kv_regions)}, "
154 - f"layer_names={len(layer_names)}."
155 - )
147 regions: list[TransferRegion] = []
148 for (
149 base_addr,
@@ -161,7 +152,6 @@ def _expand_transfer_regions(
152 layer_name,
153 layer_index,
154 group_index,
164 - split_kv_region,
155 ) in zip(
156 base_addrs,
157 block_lens,
@@ -169,7 +159,6 @@ def _expand_transfer_regions(
159 layer_names,
160 layer_indices,
161 group_indices,
172 - split_kv_regions,
162 ):
163 regions.append(
164 TransferRegion(
@@ -181,17 +170,6 @@ def _expand_transfer_regions(
170 group_index=group_index,
171 )
172 )
184 - if split_kv_region:
185 - regions.append(
186 - TransferRegion(
187 - layer_name=layer_name,
188 - layer_index=layer_index,
189 - base_addr=base_addr + kv_block_len,
190 - block_len=block_len,
191 - kv_block_len=kv_block_len,
192 - group_index=group_index,
193 - )
194 - )
173 return regions
174
175
@@ -506,10 +484,10 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
484 if vllm_config.model_config.use_mla:
485 return None
486 logger.info_once(
509 - "MooncakeConnector setting KV cache layout to HND for "
487 + "MooncakeConnector setting KV cache layout to LBHNC for "
488 "heterogeneous TP-safe KV transfer."
489 )
512 - return "HND"
490 + return "LBHNC"
491
492 ############################################################
493 # Scheduler Side Methods
@@ -1023,12 +1001,10 @@ class MooncakeConnectorWorker:
1001 self._sync_block_size_with_kernel()
1002
1003 self.attn_backends = get_current_attn_backends(vllm_config)
1026 - self.kv_cache_layout = get_kv_cache_layout()
1004 logger.debug(
1005 "Detected attention backends %s",
1006 [backend.get_name() for backend in self.attn_backends],
1007 )
1031 - logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
1008
1009 self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
1010 self._layer_specs: dict[str, KVCacheSpec] = {}
@@ -1665,7 +1641,7 @@ class MooncakeConnectorWorker:
1641 self.registered_layer_indices = []
1642 self.registered_group_indices = []
1643
1668 - for layer_name, cache_or_caches in kv_caches.items():
1644 + for layer_name, cache in kv_caches.items():
1645 layer_index = extract_layer_index(layer_name)
1646 layer_spec = self._layer_specs.get(layer_name)
1647 if layer_spec is None:
@@ -1674,34 +1650,30 @@ class MooncakeConnectorWorker:
1650 layer_name,
1651 )
1652 continue
1677 - if isinstance(layer_spec, MambaSpec):
1678 - conv, _ = cache_or_caches
1679 - cache_list = [conv]
1653 + # One raw page tensor per layer; for Mamba that page holds all the
1654 + # recurrent states, unpacked only when binding the cache for execution.
1655 + self._log_debug_cache_registration(layer_name, cache)
1656 + block_is_contiguous = is_non_overlapping_and_dense(cache[0])
1657 + if not block_is_contiguous:
1658 + # Non-block-compact layouts scatter a block across per-head
1659 + # regions; each region's blocks are contiguous.
1660 + region_caches = [cache[:, head] for head in range(cache.shape[1])]
1661 + assert all(
1662 + is_non_overlapping_and_dense(region[0]) for region in region_caches
1663 + )
1664 else:
1681 - # K and V are packed into one blocks-first tensor per layer,
1682 - # so each layer registers as a single region.
1683 - cache_list = [cache_or_caches]
1684 -
1685 - logger.debug(
1686 - "registering layer %s with %d cache tensor(s)",
1687 - layer_name,
1688 - len(cache_list),
1689 - )
1665 + region_caches = [cache]
1666
1691 - for cache in cache_list:
1692 - self._log_debug_cache_registration(layer_name, cache)
1693 - base_addr = cache.data_ptr()
1694 - block_len = cache.stride(0) * cache.element_size()
1667 + for region_cache in region_caches:
1668 + base_addr = region_cache.data_ptr()
1669 + block_len = region_cache.stride(0) * region_cache.element_size()
1670 region_base_addresses.append(base_addr)
1671
1697 - if isinstance(layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)):
1698 - kv_block_len = layer_spec.page_size_bytes
1699 - elif self.transfer_topo.virtually_split_kv_in_blocks and not isinstance(
1700 - layer_spec, MambaSpec
1701 - ):
1702 - kv_block_len = block_len // 2
1703 - else:
1704 - kv_block_len = block_len
1672 + kv_block_len = (
1673 + layer_spec.page_size_bytes
1674 + if isinstance(layer_spec, AttentionSpec) and block_is_contiguous
1675 + else block_len
1676 + )
1677 self.block_len_per_layer.append(block_len)
1678 self.kv_block_len_per_layer.append(kv_block_len)
1679 self.registered_layer_names.append(layer_name)
@@ -1709,12 +1681,12 @@ class MooncakeConnectorWorker:
1681 self.registered_group_indices.append(
1682 self._layer_group_indices[layer_name]
1683 )
1712 - storage = cache.untyped_storage()
1713 - storage_addr = storage.data_ptr()
1714 - if storage_addr not in seen_storage_ptrs:
1715 - seen_storage_ptrs.add(storage_addr)
1716 - kv_data_ptrs.append(storage_addr)
1717 - kv_data_lens.append(storage.nbytes())
1684 + storage = cache.untyped_storage()
1685 + storage_addr = storage.data_ptr()
1686 + if storage_addr not in seen_storage_ptrs:
1687 + seen_storage_ptrs.add(storage_addr)
1688 + kv_data_ptrs.append(storage_addr)
1689 + kv_data_lens.append(storage.nbytes())
1690
1691 self.kv_caches_base_addr = region_base_addresses
1692 self.seen_base_addresses = kv_data_ptrs
@@ -2055,24 +2027,13 @@ class MooncakeConnectorWorker:
2027 self._layer_group_indices.get(layer_name, 0)
2028 for layer_name in layer_names
2029 ]
2058 - split_kv_regions = None
2059 - if self.transfer_topo.virtually_split_kv_in_blocks:
2060 - split_kv_regions = [
2061 - not isinstance(
2062 - self._layer_specs[layer_name],
2063 - (MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec),
2064 - )
2065 - for layer_name in layer_names
2066 - ]
2030 return _expand_transfer_regions(
2031 base_addrs=base_addrs,
2032 block_lens=block_lens,
2033 kv_block_lens=kv_block_lens,
2034 layer_names=layer_names,
2035 layer_indices=layer_indices,
2073 - is_kv_layout_blocks_first=self.transfer_topo.virtually_split_kv_in_blocks,
2036 group_indices=group_indices,
2075 - split_kv_regions=split_kv_regions,
2037 )
2038
2039 def _get_sender_transfer_plan(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
-18
@@ -89,14 +89,6 @@ class MooncakeStoreKVEvents(KVConnectorKVEvents):
89 class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
90 """KV connector using MooncakeDistributedStore as shared KV pool."""
91
92 - @property
93 - def prefer_cross_layer_blocks(self) -> bool:
94 - extra_config = self._kv_transfer_config.kv_connector_extra_config
95 - return (
96 - str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
97 - == "true"
98 - )
99 -
92 @staticmethod
93 def _validate_kv_cache_config(
94 vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
@@ -290,16 +282,6 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
282 assert self.connector_worker is not None
283 self.connector_worker.register_kv_caches(kv_caches)
284
293 - def register_cross_layers_kv_cache(
294 - self, kv_cache: torch.Tensor, attn_backend: type
295 - ):
296 - assert self.connector_worker is not None
297 - assert (
298 - self._kv_cache_config is not None
299 - and len(self._kv_cache_config.kv_cache_groups) == 1
300 - ), "Cross-layer KV cache does not supported with hybrid models"
301 - self.connector_worker.register_cross_layers_kv_caches(kv_cache)
302 -
285 def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
286 # No-op: loads are issued in get_finished() for compute overlap.
287 pass
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+46 -49
@@ -58,6 +58,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import
58 from vllm.logger import init_logger
59 from vllm.utils.math_utils import cdiv
60 from vllm.utils.network_utils import get_ip, make_zmq_socket
61 +from vllm.utils.torch_utils import is_non_overlapping_and_dense
62 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
63 from vllm.v1.core.kv_cache_utils import (
64 BlockHash,
@@ -72,6 +73,7 @@ from vllm.v1.kv_cache_interface import (
73 MLAAttentionSpec,
74 SlidingWindowMLASpec,
75 UniformTypeKVCacheSpecs,
76 + group_kernel_blocks,
77 )
78
79 from .metrics import MooncakeStoreConnectorStats
@@ -1617,18 +1619,9 @@ class MooncakeStoreWorker:
1619 for g_idx, db in enumerate(self.token_dbs)
1620 )
1621
1620 - def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
1621 - """Register a cross-layers KV cache tensor.
1622 -
1623 - Wraps the unified tensor in a single-entry dict so that the
1624 - existing stride-based logic in register_kv_caches() produces
1625 - the correct single-segment result (block_len = page_size * num_layers).
1626 - """
1627 - self.register_kv_caches({"__cross_layer__": kv_cache})
1628 -
1622 def register_kv_caches(
1623 self,
1631 - kv_caches: dict[str, torch.Tensor | list[torch.Tensor]],
1624 + kv_caches: dict[str, torch.Tensor],
1625 ) -> None:
1626 """Register KV cache tensors and start transfer threads."""
1627 if self._capacity_only:
@@ -1637,58 +1630,62 @@ class MooncakeStoreWorker:
1630 logger.warning("No KV caches to offload.")
1631 return
1632
1640 - # Resolve each entry to a representative tensor for storage
1641 - # deduplication. For attention layers the value is already a tensor;
1642 - # for Mamba layers it is a list of tensors that all share the same
1643 - # underlying raw storage, so we take the first one.
1644 - def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
1645 - assert isinstance(v, torch.Tensor | list)
1646 - return v if isinstance(v, torch.Tensor) else v[0]
1647 -
1633 assert self.cache_config.num_gpu_blocks is not None
1634 self.num_blocks = self.cache_config.num_gpu_blocks
1635
1651 - seen_ptrs: set[int] = set()
1636 + seen_storage_ptrs: set[int] = set()
1637 + seen_region_ptrs: set[int] = set()
1638 addrs: list[int] = []
1639 block_lens: list[int] = []
1640
1655 - for value in kv_caches.values():
1656 - cache = _repr_tensor(value)
1641 + for cache in kv_caches.values():
1642 + cache = group_kernel_blocks(cache, self.num_blocks)
1643 cache_storage = cache.untyped_storage()
1644 base_addr = cache_storage.data_ptr()
1659 - if base_addr in seen_ptrs:
1660 - continue
1661 - seen_ptrs.add(base_addr)
1645 region_len = cache_storage.nbytes()
1646
1664 - ret = self.store.register_buffer(base_addr, region_len)
1665 - if ret != 0:
1666 - logger.error(
1667 - "register_buffer failed for addr %#x len %d: %d",
1668 - base_addr,
1669 - region_len,
1670 - ret,
1671 - )
1647 + if base_addr not in seen_storage_ptrs:
1648 + seen_storage_ptrs.add(base_addr)
1649 + ret = self.store.register_buffer(base_addr, region_len)
1650 + if ret != 0:
1651 + logger.error(
1652 + "register_buffer failed for addr %#x len %d: %d",
1653 + base_addr,
1654 + region_len,
1655 + ret,
1656 + )
1657
1673 - # Detect layout via stride: a dim whose byte-stride exceeds
1674 - # page_size_bytes is an outer segment dim (e.g. the K/V dim of
1675 - # FlashAttn's (2, num_blocks, ...)). FlashInfer/MLA's blocks-
1676 - # outermost layout has no such dim and yields a single segment.
1677 - el = cache.element_size()
1678 - page_size_bytes = region_len // self.num_blocks
1679 - outer_dims = [
1680 - d for d in range(cache.ndim) if cache.stride(d) * el > page_size_bytes
1681 - ]
1682 - if not outer_dims:
1683 - # Blocks-first layout (FlashInfer / MLA): one segment.
1658 + if not is_non_overlapping_and_dense(cache[0]):
1659 + # A block is scattered across per-head regions; each region's
1660 + # blocks are contiguous.
1661 + for head_idx in range(cache.shape[1]):
1662 + head_cache = cache[:, head_idx]
1663 + assert is_non_overlapping_and_dense(head_cache[0])
1664 + region_addr = head_cache.data_ptr()
1665 + if region_addr in seen_region_ptrs:
1666 + continue
1667 + seen_region_ptrs.add(region_addr)
1668 + addrs.append(region_addr)
1669 + block_lens.append(head_cache.stride(0) * head_cache.element_size())
1670 + elif cache.stride(0) * cache.element_size() * self.num_blocks == region_len:
1671 + # The block stride spans the whole per-block window
1672 + # (block-outermost and packed layouts, and single-layer
1673 + # tensors), which may hold other layers' pages at higher
1674 + # offsets. Register the storage once as one whole-window
1675 + # region; per-layer regions would copy overlapping windows and
1676 + # run past the storage for offset layers.
1677 + if base_addr in seen_region_ptrs:
1678 + continue
1679 + seen_region_ptrs.add(base_addr)
1680 addrs.append(base_addr)
1685 - block_lens.append(page_size_bytes)
1681 + block_lens.append(region_len // self.num_blocks)
1682 else:
1687 - # K/V-first layout (FlashAttn / ROCm): split segments.
1688 - seg_stride = cache.stride(outer_dims[0]) * el
1689 - for idx in range(cache.shape[outer_dims[0]]):
1690 - addrs.append(base_addr + idx * seg_stride)
1691 - block_lens.append(seg_stride // self.num_blocks)
1683 + region_addr = cache.data_ptr()
1684 + if region_addr in seen_region_ptrs:
1685 + continue
1686 + seen_region_ptrs.add(region_addr)
1687 + addrs.append(region_addr)
1688 + block_lens.append(cache.stride(0) * cache.element_size())
1689
1690 logger.info(
1691 "Registered KV caches: num_groups=%d, num_segments=%d, num_blocks=%d",
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
+19 -31
@@ -57,11 +57,6 @@ def is_mla_cache_layer(
57 return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
58
59
60 -def _content_packed_dim(spec: AttentionSpec) -> int:
61 - head_size_v = getattr(spec, "head_size_v", spec.head_size)
62 - return spec.head_size + head_size_v
63 -
64 -
60 def _spec_dim_matches(value: int, expected: int | None) -> bool:
61 return expected is None or value == expected
62
@@ -110,23 +105,6 @@ def get_layer_transfer_geometry(
105 spec = layer_to_spec[layer_name]
106 is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
107
113 - if is_mla_cache and len(shape) == 3:
114 - num_blocks, block_size, latent_dim = shape
115 - slot_size_bytes = latent_dim * element_size
116 - block_len = block_size * slot_size_bytes
117 - return LayerTransferGeometry(
118 - num_blocks=num_blocks,
119 - block_size=block_size,
120 - block_len=block_len,
121 - slot_size_bytes=slot_size_bytes,
122 - block_stride=stride[0],
123 - local_kv_stride=None,
124 - remote_kv_stride=None,
125 - transfers_per_block=1,
126 - regions_per_block=1,
127 - split_kv_regions=False,
128 - )
129 -
108 if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
109 _, num_blocks = shape[:2]
110 kernel_blocks_per_block = 1
@@ -230,19 +208,19 @@ def get_layer_transfer_geometry(
208 )
209
210 if (
233 - not is_mla_cache
234 - and isinstance(spec, AttentionSpec)
211 + isinstance(spec, AttentionSpec)
212 and len(shape) == 4
236 - and shape[1] == spec.num_kv_heads
237 - and shape[2] == spec.block_size
238 - and shape[3] == _content_packed_dim(spec)
213 + and shape[1] == spec.num_heads
214 + and shape[2] == spec.num_states
215 + and shape[3] * element_size == spec.state_content_size_bytes
216 ):
240 - num_blocks, num_kv_heads, block_size, packed_dim = shape
241 - slot_size_bytes = num_kv_heads * packed_dim * element_size
242 - block_len = block_size * slot_size_bytes
217 + # Standardized per-layer [B, H, N, C] view (MLA is just H == 1).
218 + num_blocks, num_heads, num_states, content_dim = shape
219 + slot_size_bytes = num_heads * content_dim * element_size
220 + block_len = num_states * slot_size_bytes
221 return LayerTransferGeometry(
222 num_blocks=num_blocks,
245 - block_size=block_size,
223 + block_size=spec.block_size,
224 block_len=block_len,
225 slot_size_bytes=slot_size_bytes,
226 block_stride=stride[0],
@@ -267,6 +245,16 @@ def iter_layer_registration_regions(
245 ) -> list[tuple[torch.Tensor, int]]:
246 geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
247 region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
248 + if geometry.regions_per_block == 1:
249 + # With padded or interleaved pages the block stride exceeds the
250 + # meaningful block_len; register the strided span. The span ends with
251 + # the last block's meaningful bytes, not a whole stride, so it never
252 + # runs past the backing allocation of a strided layer view.
253 + block_stride_bytes = geometry.block_stride * kv_cache.element_size()
254 + region_len = max(
255 + region_len,
256 + (geometry.num_blocks - 1) * block_stride_bytes + geometry.block_len,
257 + )
258 if geometry.split_kv_regions:
259 return [(cache, region_len) for cache in kv_cache]
260 return [(kv_cache, region_len)]
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+1 -14
@@ -27,7 +27,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
27 PromMetricT,
28 )
29 from vllm.logger import init_logger
30 -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
30 +from vllm.v1.attention.backend import AttentionMetadata
31 from vllm.v1.core.sched.output import SchedulerOutput
32 from vllm.v1.outputs import KVConnectorOutput
33
@@ -209,12 +209,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
209 c.supports_divergent_local_hybrid_hits for c in self._connectors
210 )
211
212 - @property
213 - def prefer_cross_layer_blocks(self) -> bool:
214 - if not self._connectors:
215 - return False
216 - return all(c.prefer_cross_layer_blocks for c in self._connectors)
217 -
212 @property
213 def requires_kv_delivery(self) -> bool:
214 return any(c.requires_kv_delivery for c in self._connectors)
@@ -245,13 +239,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
239 )
240 return ret
241
248 - def register_cross_layers_kv_cache(
249 - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
250 - ):
251 - # Register on all connectors
252 - for c in self._connectors:
253 - c.register_cross_layers_kv_cache(kv_cache, attn_backend)
254 -
242 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
243 for c in self._connectors:
244 c.register_kv_caches(kv_caches)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+118 -177
@@ -71,9 +71,9 @@ from vllm.platforms import current_platform
71 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
72 from vllm.utils.network_utils import make_zmq_path
73 from vllm.utils.torch_utils import async_tensor_h2d
74 -from vllm.v1.attention.backends.utils import get_kv_cache_layout
74 from vllm.v1.kv_cache_interface import (
75 FullAttentionSpec,
76 + KVCacheLayout,
77 MambaSpec,
78 MLAAttentionSpec,
79 SlidingWindowMLASpec,
@@ -89,6 +89,13 @@ if TYPE_CHECKING:
89 logger = init_logger(__name__)
90
91
92 +def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
93 + """Return whether all views share storage and a block stride."""
94 + block_strides = {cache.stride(0) * cache.element_size() for cache in caches}
95 + storage_ptrs = {cache.untyped_storage().data_ptr() for cache in caches}
96 + return len(block_strides) == len(storage_ptrs) == 1
97 +
98 +
99 class NixlBaseConnectorWorker:
100 """Base implementation of Worker side methods shared by pull and push."""
101
@@ -302,7 +309,6 @@ class NixlBaseConnectorWorker:
309 for group in kv_cache_config.kv_cache_groups
310 for layer in group.layer_names
311 }
305 - self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
312
313 # ---- Model state (derived from model config) ----
314 mamba_ssm_size = (0, 0)
@@ -523,7 +529,9 @@ class NixlBaseConnectorWorker:
529 self.attn_backends = get_current_attn_backends(vllm_config)
530 self.backend_name = self.attn_backends[0].get_name()
531
526 - self.kv_cache_layout = get_kv_cache_layout()
532 + self.kv_cache_layout = (
533 + vllm_config.cache_config.get_resolved_kv_cache_layout().name
534 + )
535 self.host_buffer_kv_cache_layout = self.kv_cache_layout
536 logger.info(
537 "Detected attention backend(s) %s",
@@ -558,6 +566,11 @@ class NixlBaseConnectorWorker:
566 # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
567 self.block_len_per_layer = list[int]()
568
569 + # Per-region block stride in bytes. Taken from the registered tensor's
570 + # stride(0) so it stays correct under layouts that interleave layers
571 + # within a block (BLHNC/BHLNC), where stride > block_len.
572 + self.block_stride_per_layer = list[int]()
573 +
574 # Per-engine TP mappings. Generated during handshake.
575 self.tp_mappings: dict[EngineId, TPMapping] = {}
576
@@ -788,16 +801,16 @@ class NixlBaseConnectorWorker:
801 if not self.use_mla:
802 assert kv_cache.ndim == 4
803
791 - if self.kv_cache_layout == "NHD":
804 + if self.kv_cache_layout == "LBNHC":
805 if self.kv_transfer_config.enable_permute_local_kv:
806 logger.info_once(
807 "'enable_permute_local_kv' flag is enabled while "
795 - "device KV Layout is NHD. Init host buffer with"
796 - " HND to better support Decode/Prefill TP_ratio > 1."
808 + "device KV Layout is LBNHC. Init host buffer with"
809 + " LBHNC to better support Decode/Prefill TP_ratio > 1."
810 )
798 - # Since NHD will not support Decode/Prefill TP_ratio > 1,
811 + # Since LBNHC will not support Decode/Prefill TP_ratio > 1,
812 # we can leverage host_buffer for permute.
800 - self.host_buffer_kv_cache_layout = "HND"
813 + self.host_buffer_kv_cache_layout = "LBHNC"
814 else:
815 # Packed KV layout is logical (B, H, N, 2*D). Allocate
816 # (B, N, H, 2*D) and view it as logical (B, H, N, 2*D)
@@ -969,120 +982,9 @@ class NixlBaseConnectorWorker:
982
983 fut.add_done_callback(request_ready)
984
972 - def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
973 - """Register a cross-layers KV cache tensor with NIXL.
974 -
975 - `use_uniform_kv_cache()` guarantees a single KV cache group whose
976 - layers all share the same `AttentionSpec`, so any layer name from
977 - `_layer_specs` yields the correct per-layer spec for `page_size_bytes`.
978 - """
979 - first_layer = next(iter(self._layer_specs))
980 - # Forwarding a real layer name rather than a synthetic key
981 - self.register_kv_caches({first_layer: kv_cache})
982 -
983 - def _register_packed_kv_cache(
984 - self,
985 - storage: torch.UntypedStorage,
986 - ) -> None:
987 - """Register a packed KV cache as a single NIXL region.
988 -
989 - The packed allocation interleaves all layers per block, so each
990 - block_stride-byte chunk is one logical block. We register 1
991 - NIXL region and create 1 descriptor per block.
992 - """
993 - self.transfer_topo = TransferTopology(
994 - tp_rank=self.tp_rank,
995 - tp_size=self.world_size,
996 - block_size=self.block_size,
997 - engine_id=self.engine_id,
998 - is_mla=self.use_mla,
999 - total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
1000 - attn_backends=self.attn_backends,
1001 - tensor_shape=None,
1002 - is_mamba=self._has_mamba,
1003 - )
1004 - self.compat_hash = compute_nixl_compatibility_hash(
1005 - self.vllm_config,
1006 - self.backend_name,
1007 - self.transfer_topo.cross_layers_blocks,
1008 - transfer_mode=self._TRANSFER_MODE,
1009 - )
1010 -
1011 - total_size = storage.nbytes()
1012 - block_stride = total_size // self.num_blocks
1013 - base_addr = storage.data_ptr()
1014 - device_id = storage.device.index
1015 - assert device_id is not None
1016 -
1017 - logger.info(
1018 - "Registering packed KV cache: total_size=%s, block_stride=%s, "
1019 - "num_blocks=%s, num_regions=1",
1020 - total_size,
1021 - block_stride,
1022 - self.num_blocks,
1023 - )
1024 -
1025 - self.device_id = device_id
1026 - caches_data = [(base_addr, total_size, self.device_id, "")]
1027 -
1028 - self.block_len_per_layer = [block_stride]
1029 - self.num_regions = 1
1030 - self.num_descs = self.num_blocks
1031 - self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr]
1032 -
1033 - descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
1034 - self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
1035 - self._registered_descs.append(descs)
1036 -
1037 - self.dst_num_blocks[self.engine_id] = self.num_blocks
1038 -
1039 - self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = (
1040 - self.register_local_xfer_handler(self.block_size)
1041 - )
1042 -
1043 - agent_metadata = NixlAgentMetadata(
1044 - engine_id=self.engine_id,
1045 - agent_metadata=self.nixl_wrapper.get_agent_metadata(),
1046 - device_id=self.device_id,
1047 - kv_caches_base_addr=(
1048 - self.kv_caches_base_addr[self.engine_id][self.tp_rank]
1049 - ),
1050 - num_blocks=self.num_blocks,
1051 - block_lens=self.block_len_per_layer,
1052 - kv_cache_layout=self.kv_cache_layout,
1053 - block_size=self.block_size,
1054 - ssm_sizes=self._mamba_ssm_size,
1055 - attn_backend_name=self.backend_name,
1056 - physical_blocks_per_logical_kv_block=(
1057 - self._physical_blocks_per_logical_kv_block
1058 - ),
1059 - dcp_size=self.dcp_size,
1060 - pcp_size=self.pcp_size,
1061 - )
1062 - assert self.compat_hash is not None
1063 - encoder = msgspec.msgpack.Encoder()
1064 - self.xfer_handshake_metadata = NixlHandshakePayload(
1065 - compatibility_hash=self.compat_hash,
1066 - agent_metadata_bytes=encoder.encode(agent_metadata),
1067 - )
1068 -
985 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
986 """Register the KV Cache data in nixl."""
987
1072 - # Detect packed allocation: all tensors are strided views into the
1073 - # same backing storage (different data_ptr but same storage).
1074 - # This happens with DSv4-style contiguous per-block packing.
1075 - if len(kv_caches) > 1 and not self._has_mamba:
1076 - storage = next(iter(kv_caches.values())).untyped_storage()
1077 - storage_ptrs = {
1078 - cache.untyped_storage().data_ptr() for cache in kv_caches.values()
1079 - }
1080 - data_ptrs = {cache.data_ptr() for cache in kv_caches.values()}
1081 - if len(storage_ptrs) == 1 and len(data_ptrs) > 1:
1082 - self._register_packed_kv_cache(storage)
1083 - self.device_kv_caches = kv_caches
1084 - return
1085 -
988 self.transfer_topo = TransferTopology(
989 tp_rank=self.tp_rank,
990 tp_size=self.world_size,
@@ -1100,7 +1002,6 @@ class NixlBaseConnectorWorker:
1002 self.compat_hash = compute_nixl_compatibility_hash(
1003 self.vllm_config,
1004 self.backend_name,
1103 - self.transfer_topo.cross_layers_blocks,
1005 transfer_mode=self._TRANSFER_MODE,
1006 )
1007
@@ -1127,14 +1028,14 @@ class NixlBaseConnectorWorker:
1028 )
1029
1030 caches_data = []
1130 - # With hybrid allocator, layers can share a kv cache tensor
1031 + seen_storage_addresses: set[int] = set()
1032 seen_base_addresses: list[int] = []
1033
1034 + packed_storage = _share_storage_and_block_stride(list(xfer_buffers.values()))
1035 +
1036 # K and V are packed into the content dim, so each attention layer is a
1037 # single NIXL region whose block transfers as one unit. Mamba layers instead
1038 # register separate conv/ssm sub-regions (see `_build_mamba_local`).
1136 - tensor_size_bytes = None
1137 -
1039 for layer_name, cache in xfer_buffers.items():
1040 # NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when
1041 # kernel requires a specific block size. This leads to SSM and FA layers
@@ -1159,55 +1060,98 @@ class NixlBaseConnectorWorker:
1060 else layer_spec.page_size_bytes
1061 // self._physical_blocks_per_logical_kv_block
1062 )
1162 - if self.transfer_topo._cross_layers_blocks:
1163 - # When cross-layers blocks are used, multiply by number of layers
1164 - physical_page_size = physical_page_size * len(
1165 - self.kv_cache_config.kv_cache_tensors
1166 - )
1063 num_blocks = (
1064 self._logical_num_blocks
1065 if isinstance(layer_spec, MambaSpec)
1066 else self.num_blocks
1067 )
1172 - # `page_size` accounts for physical blocks, st KVCache is always
1173 - # [`num_blocks` * `page_size`]
1174 - curr_tensor_size_bytes = num_blocks * physical_page_size
1068 + logger.debug(
1069 + "Registering layer %s with cache shape: %s", layer_name, cache.shape
1070 + )
1071 + storage = cache.untyped_storage()
1072 + storage_addr = storage.data_ptr()
1073 + # Memory registration follows allocations, while transfer regions follow
1074 + # logical layers (or contiguous head segments). This keeps strided
1075 + # cross-layer views inside their registered allocation.
1076 + if storage_addr not in seen_storage_addresses:
1077 + seen_storage_addresses.add(storage_addr)
1078 + self.device_id = max(cache.get_device(), 0)
1079 + caches_data.append((storage_addr, storage.nbytes(), self.device_id, ""))
1080
1176 - base_addr = cache.data_ptr()
1081 is_mla_region = isinstance(
1082 layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
1083 )
1180 - if base_addr in seen_base_addresses:
1181 - # NOTE (NickLucche) HMA employs memory pooling to share tensors
1182 - # across groups. This results in skipping all tensors but the ones
1183 - # pointed to by group0. Also, generally we will have more blocks
1184 - # per tensor but fewer regions.
1185 - # A shared tensor may back both SSM and attention layers (e.g.
1186 - # KDA+MLA in KimiLinear); the region's FA view is MLA whichever
1187 - # layer registered it first.
1188 - idx = seen_base_addresses.index(base_addr)
1189 - self._region_is_mla[idx] |= is_mla_region
1190 - logger.debug("Skipping %s because it's already seen", layer_name)
1191 - continue
1192 - logger.debug(
1193 - "Registering layer %s with cache shape: %s", layer_name, cache.shape
1194 - )
1195 - seen_base_addresses.append(base_addr)
1196 - # Only record non-Mamba page sizes.
1084 +
1085 if isinstance(layer_spec, MambaSpec):
1198 - self.block_len_per_layer.append(
1199 - physical_page_size // self._physical_blocks_per_logical_kv_block
1200 - )
1086 + region_specs = [
1087 + (
1088 + cache.data_ptr(),
1089 + physical_page_size
1090 + // self._physical_blocks_per_logical_kv_block,
1091 + physical_page_size
1092 + // self._physical_blocks_per_logical_kv_block,
1093 + )
1094 + ]
1095 else:
1202 - self.block_len_per_layer.append(physical_page_size)
1203 - self._region_is_mla.append(is_mla_region)
1204 -
1205 - if not is_mla_region:
1206 - if tensor_size_bytes is None:
1207 - tensor_size_bytes = curr_tensor_size_bytes
1208 - assert tensor_size_bytes == curr_tensor_size_bytes, (
1209 - "All non-MLA kv cache tensors must have the same size"
1096 + if cache.ndim == 1:
1097 + # Flat byte view: HMA tensors shared between layer types carry
1098 + # no block dimension, so stride(0) is 1 byte. Blocks abut, so
1099 + # the stride is the page size.
1100 + block_stride = physical_page_size
1101 + else:
1102 + block_stride = cache.stride(0) * cache.element_size()
1103 + storage_is_block_major = num_blocks * block_stride == storage.nbytes()
1104 + # A layer whose [H, N, C] interior is dense addresses its own page
1105 + # as one chunk, so it can own a region. When it does not, or when
1106 + # every layer is packed into this one storage, the block's whole
1107 + # row is the only contiguous unit.
1108 + hnc_contiguous = (
1109 + cache.ndim == 4
1110 + and cache.stride(2) == cache.shape[3]
1111 + and cache.stride(1) == cache.shape[2] * cache.shape[3]
1112 )
1113 + if storage_is_block_major and (
1114 + (packed_storage and is_mla_region) or not hnc_contiguous
1115 + ):
1116 + # TODO(Lucas): handle TP slicing for packed_storage; for now
1117 + # restrict to MLA (DSv4) where kv is replicated.
1118 + storage_block_len = storage.nbytes() // num_blocks
1119 + region_specs = [
1120 + (storage_addr, storage_block_len, storage_block_len)
1121 + ]
1122 + elif storage_is_block_major:
1123 + region_specs = [
1124 + (cache.data_ptr(), physical_page_size, block_stride)
1125 + ]
1126 + else:
1127 + segment_bytes = num_blocks * block_stride
1128 + assert cache.nbytes % segment_bytes == 0
1129 + num_segments = cache.nbytes // segment_bytes
1130 + region_block_len = (
1131 + block_stride if num_segments > 1 else physical_page_size
1132 + )
1133 + region_specs = [
1134 + (
1135 + cache.data_ptr() + segment_idx * segment_bytes,
1136 + region_block_len,
1137 + block_stride,
1138 + )
1139 + for segment_idx in range(num_segments)
1140 + ]
1141 +
1142 + for base_addr, block_len, block_stride in region_specs:
1143 + if base_addr in seen_base_addresses:
1144 + if is_mla_region:
1145 + # Dual-purpose HMA tensor: an MLA layer shares a region
1146 + # that a non-MLA layer registered first. MLA is not
1147 + # head-sharded, so the region must be flagged MLA.
1148 + idx = seen_base_addresses.index(base_addr)
1149 + self._region_is_mla[idx] = True
1150 + continue
1151 + seen_base_addresses.append(base_addr)
1152 + self.block_len_per_layer.append(block_len)
1153 + self.block_stride_per_layer.append(block_stride)
1154 + self._region_is_mla.append(is_mla_region)
1155
1156 # When there's a mismatch between kbs<>bs, we rely on HMA to ensure
1157 # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
@@ -1228,11 +1172,6 @@ class NixlBaseConnectorWorker:
1172 f"kv_cache_layout={self.kv_cache_layout}"
1173 )
1174
1231 - # Need to make sure the device ID is non-negative for NIXL,
1232 - # Torch uses -1 to indicate CPU tensors.
1233 - self.device_id = max(cache.get_device(), 0)
1234 - caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, ""))
1235 -
1175 logger.debug(
1176 "Different block lengths collected: %s", set(self.block_len_per_layer)
1177 )
@@ -1240,10 +1179,11 @@ class NixlBaseConnectorWorker:
1179 len(self.block_len_per_layer)
1180 == len(seen_base_addresses)
1181 == len(self._region_is_mla)
1182 + == len(self.block_stride_per_layer)
1183 )
1184
1185 self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
1246 - self.num_regions = len(caches_data)
1186 + self.num_regions = len(seen_base_addresses)
1187
1188 if self.pp_size > 1:
1189 start_layer, end_layer = self.model_config.get_layers_start_end_indices(
@@ -1293,6 +1233,7 @@ class NixlBaseConnectorWorker:
1233 kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank],
1234 num_blocks=self.num_blocks,
1235 block_lens=self.block_len_per_layer,
1236 + block_strides=self.block_stride_per_layer,
1237 kv_cache_layout=self.kv_cache_layout
1238 if not self.use_host_buffer
1239 else self.host_buffer_kv_cache_layout,
@@ -1431,10 +1372,9 @@ class NixlBaseConnectorWorker:
1372 block_arange = np.arange(num_blocks, dtype=np.uint64)
1373 parts: list[np.ndarray] = []
1374 for i, base_addr in enumerate(base_addresses):
1434 - # K/V are packed into the content dim, so the whole block transfers
1435 - # as one unit: desc length equals the block stride.
1375 block_len = self.block_len_per_layer[i] // block_size_ratio
1437 - addrs = base_addr + block_arange * block_len
1376 + block_stride = self.block_stride_per_layer[i] // block_size_ratio
1377 + addrs = base_addr + block_arange * block_stride
1378 parts.append(self._stack_descs(addrs, block_len, device_id))
1379 return np.concatenate(parts)
1380
@@ -1476,7 +1416,7 @@ class NixlBaseConnectorWorker:
1416 )
1417 local_block_len = local_block_len // num_reads
1418
1479 - page_size = nixl_agent_meta.block_lens[i]
1419 + page_size = nixl_agent_meta.block_strides[i]
1420 addrs = base_addr + rank_offset + block_arange * page_size
1421 parts.append(self._stack_descs(addrs, local_block_len, device_id))
1422 return np.concatenate(parts)
@@ -1560,7 +1500,7 @@ class NixlBaseConnectorWorker:
1500 tp_ratio = 4 // 2 = 2
1501
1502 Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim]
1563 - then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format.
1503 + then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "LBHNC" layout format.
1504 Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio
1505 first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split
1506 along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0.
@@ -1600,6 +1540,7 @@ class NixlBaseConnectorWorker:
1540 start:end
1541 ]
1542 nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
1543 + nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
1544
1545 ### Register remote engine in TransferTopology (idempotent).
1546 assert self.transfer_topo is not None
@@ -1790,10 +1731,10 @@ class NixlBaseConnectorWorker:
1731 if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout:
1732 if (
1733 self.kv_transfer_config.enable_permute_local_kv
1793 - and nixl_agent_meta.kv_cache_layout == "HND"
1734 + and nixl_agent_meta.kv_cache_layout == "LBHNC"
1735 ):
1736 logger.info(
1796 - "Remote is HND and local is NHD, enabled additional permute "
1737 + "Remote is LBHNC and local is LBNHC, enabled additional permute "
1738 "on local device KV."
1739 )
1740 assert not self._is_hma_required, (
@@ -1823,18 +1764,18 @@ class NixlBaseConnectorWorker:
1764 self.enable_heterogeneous_attn_post_process = True
1765
1766 # Heterogeneous TP requires head-splitting, which only works with
1826 - # HND layout. MLA and replicated-KV cases don't split on heads.
1827 - # Mamba doesn't support heterogeneous TP.
1767 + # block-contiguous layouts (e.g. LBHNC). MLA and replicated-KV cases
1768 + # don't split on heads. Mamba doesn't support heterogeneous TP.
1769 if (
1770 abs(tp_ratio) != 1
1771 and not self.use_mla
1772 and not self.transfer_topo.is_kv_replicated(remote_engine_id)
1832 - and kv_cache_layout != "HND"
1773 + and not KVCacheLayout[kv_cache_layout].is_block_contiguous
1774 and not self.enable_permute_local_kv
1775 ):
1776 raise RuntimeError(
1777 "Heterogeneous TP head-dimension splitting requires contiguous heads. "
1837 - "Use HND layout on the prefill side."
1778 + "Use a block-contiguous layout (e.g. LBHNC) on the prefill side."
1779 )
1780
1781 # Per-region block_len validation enforcing the P/D invariant.
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+4 -42
@@ -19,7 +19,6 @@ import torch
19 from vllm.config import VllmConfig
20 from vllm.distributed.kv_transfer.kv_connector.utils import (
21 EngineId,
22 - get_current_attn_backend,
22 )
23 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
24 CopyBlocksOp,
@@ -56,10 +55,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
55 )
56 from vllm.forward_context import ForwardContext
57 from vllm.logger import init_logger
59 -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
60 -from vllm.v1.attention.backends.utils import get_kv_cache_layout
58 +from vllm.v1.attention.backend import AttentionMetadata
59 from vllm.v1.core.sched.output import SchedulerOutput
62 -from vllm.v1.kv_cache_interface import MambaSpec
60 from vllm.v1.outputs import KVConnectorOutput
61
62 if TYPE_CHECKING:
@@ -83,36 +80,6 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
80 def supports_divergent_local_hybrid_hits(self) -> bool:
81 return True
82
86 - @property
87 - def prefer_cross_layer_blocks(self) -> bool:
88 - if any(
89 - [
90 - isinstance(group.kv_cache_spec, MambaSpec)
91 - for group in self.kv_cache_config.kv_cache_groups
92 - ]
93 - ):
94 - # Hybrid SSM models do not yet support cross-layer layout
95 - return False
96 -
97 - backend = get_current_attn_backend(self._vllm_config)
98 - if backend.get_name() not in (
99 - "FLASH_ATTN",
100 - "FLASHINFER",
101 - "TRITON_ATTN",
102 - ):
103 - return False
104 -
105 - # For now there is no benefit to run cross layers when backend
106 - # does not support on HND
107 - if get_kv_cache_layout() != "HND":
108 - return False
109 -
110 - extra_config = self.kv_transfer_config.kv_connector_extra_config
111 - return (
112 - str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
113 - == "true"
114 - )
115 -
83 def __init__(
84 self,
85 vllm_config: VllmConfig,
@@ -179,9 +146,10 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
146 # which fallback to the default behavior.
147 return None
148 logger.info_once(
182 - "NixlConnector setting KV cache layout to HND for better xfer performance."
149 + "NixlConnector setting KV cache layout to LBHNC for "
150 + "better xfer performance."
151 )
184 - return "HND"
152 + return "LBHNC"
153
154 ############################################################
155 # Scheduler Side Methods
@@ -266,12 +234,6 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
234 assert self.connector_worker is not None
235 self.connector_worker.register_kv_caches(kv_caches)
236
269 - def register_cross_layers_kv_cache(
270 - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
271 - ):
272 - assert self.connector_worker is not None
273 - self.connector_worker.register_cross_layers_kv_caches(kv_cache)
274 -
237 def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
238 assert self.connector_worker is not None
239 self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+3 -3
@@ -44,8 +44,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
44 # 6: Validate EAGLE/MTP speculative configuration compatibility
45 # 7: Include NIXL transfer mode (push vs pull) in the compatibility hash
46 # 8: Add dcp_size and pcp_size to NixlAgentMetadata
47 +# 9: Add block_strides
48 #
48 -NIXL_CONNECTOR_VERSION: int = 8
49 +NIXL_CONNECTOR_VERSION: int = 9
50
51
52 @dataclass
@@ -56,6 +57,7 @@ class NixlAgentMetadata:
57 device_id: int
58 num_blocks: int
59 block_lens: list[int]
60 + block_strides: list[int]
61 kv_cache_layout: str
62 block_size: int
63 ssm_sizes: tuple[int, int]
@@ -129,7 +131,6 @@ def _get_speculative_compatibility_factors(
131 def compute_nixl_compatibility_hash(
132 vllm_config: VllmConfig,
133 attn_backend_name: str,
132 - cross_layers_blocks: bool,
134 transfer_mode: str = "pull",
135 ) -> str:
136 """
@@ -180,7 +181,6 @@ def compute_nixl_compatibility_hash(
181 # Attention backend and KV cache dtype affect memory layout
182 "attn_backend_name": attn_backend_name,
183 "cache_dtype": str(cache_config.cache_dtype),
183 - "cross_layers_blocks": cross_layers_blocks,
184 "is_hma_enabled": is_hma_enabled,
185 "speculative_config": _get_speculative_compatibility_factors(vllm_config),
186 # push (WRITE) and pull (READ) connectors are protocol-incompatible
vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py
+10 -6
@@ -17,6 +17,7 @@ import torch
17 from vllm.v1.kv_cache_interface import (
18 AttentionSpec,
19 KVCacheConfig,
20 + KVCacheLayout,
21 KVCacheSpec,
22 MambaSpec,
23 MLAAttentionSpec,
@@ -31,13 +32,16 @@ if TYPE_CHECKING:
32 CANONICAL_FORMAT_VERSION = 1
33
34
34 -def canonical_format_id() -> str:
35 +def canonical_format_id(kv_cache_layout: str) -> str:
36 """Identity of the canonical byte format, for namespacing persisted KV.
37 Canonical pages keep the worker's KV layout family, so the id couples the
37 - format version with that family; consumers must match it exactly."""
38 - from vllm.v1.attention.backends.utils import get_kv_cache_layout
39 -
40 - return f"v{CANONICAL_FORMAT_VERSION}-{get_kv_cache_layout().lower()}"
38 + format version with that family; consumers must match it exactly. The family
39 + keeps its historical NHD/HND spelling so ids stay stable for KV persisted
40 + before the layout enum existed."""
41 + layout = KVCacheLayout[kv_cache_layout]
42 + legacy = {KVCacheLayout.LBNHC: "nhd", KVCacheLayout.LBHNC: "hnd"}
43 + family = legacy.get(layout, layout.name.lower())
44 + return f"v{CANONICAL_FORMAT_VERSION}-{family}"
45
46
47 @dataclass(frozen=True)
@@ -278,7 +282,7 @@ def _layer_mapping(
282 if isinstance(spec, MLAAttentionSpec):
283 # TP-replicated latent; CP shards its tokens across the DCP groups
284 if (
281 - spec.compress_ratio != 1
285 + spec.tokens_per_state != 1
286 or page % bs
287 or ctx.tp_size % ctx.dcp_size
288 or spec.kv_quant_mode.is_per_token_head
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+7 -19
@@ -20,12 +20,7 @@ from vllm.v1.kv_offload.config import (
20
21 if TYPE_CHECKING:
22 from vllm.config import VllmConfig
23 - from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheTensor
24 -
25 -
26 -def is_kv_cache_tensor_packed(kv_cache_tensor: "KVCacheTensor") -> bool:
27 - """Return whether a KV cache tensor uses a packed block stride."""
28 - return bool(kv_cache_tensor.block_stride)
23 + from vllm.v1.kv_cache_interface import KVCacheConfig
24
25
26 def build_offloading_config(
@@ -96,18 +91,10 @@ def build_offloading_config(
91 blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
92
93 worker_kv_bytes_per_block = 0
99 - if kv_cache_config.num_blocks > 0:
100 - packed_tensors = tuple(
101 - is_kv_cache_tensor_packed(tensor)
102 - for tensor in kv_cache_config.kv_cache_tensors
103 - )
104 - is_packed = any(packed_tensors)
105 - assert not is_packed or all(packed_tensors)
106 - total_gpu_kv_bytes = (
107 - kv_cache_config.kv_cache_tensors[0].size
108 - if is_packed
109 - else sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors)
110 - )
94 + if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors:
95 + # Every KVCacheTensor describes placement within the same backing allocation,
96 + # so its size is the total, not a per-tensor share.
97 + total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size
98 worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
99
100 single_group_spec = (
@@ -164,7 +151,7 @@ def build_offloading_config(
151 if isinstance(single_group_spec, FullAttentionSpec):
152 if type(single_group_spec) is MLAAttentionSpec:
153 spec_certifiable = (
167 - single_group_spec.compress_ratio == 1
154 + single_group_spec.tokens_per_state == 1
155 and single_group_spec.real_page_size_bytes
156 % single_group_spec.block_size
157 == 0
@@ -219,4 +206,5 @@ def build_offloading_config(
206 ),
207 replicated_layout=replicated_layout,
208 canonical_layout=canonical_layout,
209 + kv_cache_layout=vllm_config.cache_config.kv_cache_layout,
210 )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+40 -124
@@ -14,16 +14,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
14 OffloadingWorkerMetadata,
15 ReqId,
16 )
17 -from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
18 - is_kv_cache_tensor_packed,
19 -)
17 from vllm.logger import init_logger
21 -from vllm.v1.attention.backend import AttentionBackend
18 from vllm.v1.kv_cache_interface import (
19 AttentionSpec,
20 KVCacheConfig,
21 MambaSpec,
22 UniformTypeKVCacheSpecs,
23 + group_kernel_blocks,
24 )
25 from vllm.v1.kv_offload.base import (
26 CanonicalKVCacheRef,
@@ -73,16 +70,6 @@ class OffloadingConnectorWorker:
70 self.vllm_config, kv_cache_config, kv_caches
71 )
72
76 - # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
77 - # stride(0) as the manager-block stride (equals total_num_bytes_per_block).
78 - # General (non-packed) layouts size the tensor at page_size_bytes per
79 - # manager block, so page_size_bytes is the correct offloading stride.
80 - layer_is_packed: dict[str, bool] = {
81 - ln: is_kv_cache_tensor_packed(kv_tensor)
82 - for kv_tensor in kv_cache_config.kv_cache_tensors
83 - for ln in kv_tensor.shared_by
84 - }
85 -
73 # layer_name -> (num_blocks, page_size_bytes) tensor
74 tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
75 # layer_name -> size of (un-padded) page in bytes
@@ -100,63 +87,52 @@ class OffloadingConnectorWorker:
87 layer_kv_cache_spec = per_layer_specs.get(
88 layer_name, group_kv_cache_spec
89 )
103 - if isinstance(layer_kv_cache_spec, AttentionSpec):
104 - layer_kv_cache = kv_caches[layer_name]
105 - assert isinstance(layer_kv_cache, torch.Tensor)
90 + ref = group_kernel_blocks(kv_caches[layer_name], num_blocks)
91 + page = layer_kv_cache_spec.page_size_bytes
92 + elem_size = ref.element_size()
93 + byte_offset = ref.storage_offset() * elem_size
94 + # Packed layouts (e.g. DSv4) interleave layers per block, so the
95 + # attention tensor's stride(0) (the manager-block stride) exceeds
96 + # page_size_bytes; other layouts have stride(0) == page_size_bytes.
97 + block_stride_bytes = (
98 + ref.stride(0) * elem_size
99 + if isinstance(layer_kv_cache_spec, AttentionSpec)
100 + else page
101 + )
102 + tensors_per_block[layer_name] = (
103 + torch.tensor([], dtype=torch.int8, device=ref.device).set_(
104 + ref.untyped_storage(),
105 + byte_offset,
106 + (num_blocks, page),
107 + (block_stride_bytes, 1),
108 + ),
109 + )
110 + page_size_bytes[layer_name] = page
111
107 - page = layer_kv_cache_spec.page_size_bytes
108 - elem_size = layer_kv_cache.element_size()
109 - byte_offset = layer_kv_cache.storage_offset() * elem_size
110 - block_stride_bytes = (
111 - layer_kv_cache.stride(0) * elem_size
112 - if layer_is_packed[layer_name]
113 - else page
114 - )
115 - raw = torch.empty(
116 - 0,
117 - dtype=torch.int8,
118 - device=layer_kv_cache.device,
119 - ).set_(layer_kv_cache.untyped_storage())
120 - tensors_per_block[layer_name] = (
121 - torch.as_strided(
122 - raw,
123 - (num_blocks, page),
124 - (block_stride_bytes, 1),
125 - byte_offset,
126 - ),
127 - )
128 - page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
112 + if isinstance(layer_kv_cache_spec, AttentionSpec):
113 unpadded_page_size_bytes[layer_name] = (
114 layer_kv_cache_spec.unpadded_page_size_bytes
115 )
132 -
116 elif isinstance(layer_kv_cache_spec, MambaSpec):
134 - layer_kv_cache = kv_caches[layer_name]
135 - assert layer_kv_cache.dtype == torch.int8
136 - tensors_per_block[layer_name] = (
137 - layer_kv_cache.view(
138 - num_blocks, layer_kv_cache_spec.page_size_bytes
139 - ),
140 - )
141 -
142 - page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
117 unpadded_page_size_bytes[layer_name] = replace(
118 layer_kv_cache_spec, page_size_padded=None
119 ).page_size_bytes
146 -
120 else:
121 raise NotImplementedError
122
150 - packed_kv_cache_tensor = next(
123 + # Packed layouts (e.g. DSv4) interleave all layers within each manager
124 + # block: a layer view's block stride exceeds its page size. Offload the
125 + # whole packed block as a single transfer region.
126 + packed_layer_name = next(
127 (
152 - t
153 - for t in kv_cache_config.kv_cache_tensors
154 - if is_kv_cache_tensor_packed(t) and t.shared_by
128 + layer_name
129 + for layer_name, (tensor,) in tensors_per_block.items()
130 + if tensor.stride(0) != tensor.shape[1]
131 ),
132 None,
133 )
158 - if packed_kv_cache_tensor is not None:
159 - (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]
134 + if packed_layer_name is not None:
135 + (tensor,) = tensors_per_block[packed_layer_name]
136 block_stride = tensor.stride(0)
137 packed_tensor = tensor.as_strided(
138 (num_blocks, block_stride),
@@ -176,29 +152,16 @@ class OffloadingConnectorWorker:
152
153 block_tensors: list[CanonicalKVCacheTensor] = []
154 block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
179 - for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
180 - # Filter to layers that were actually processed above.
181 - # Packed KV allocation emits KVCacheTensor entries for
182 - # every (tuple_idx, page_size) slot; slots where no group has a
183 - # layer at that index produce an empty shared_by (reserved memory
184 - # with no corresponding model layer).
185 - tensor_layer_names = [
186 - n for n in kv_cache_tensor.shared_by if n in tensors_per_block
187 - ]
188 - if not tensor_layer_names:
189 - continue
190 -
191 - # verify all layers in the group reference the exact same tensors
155 + # Layers that alias the same bytes (cache groups overlay each other) share
156 + # one transfer region: exactly the views with equal address and strides.
157 + aliased_layers: dict[tuple[int, tuple[int, ...]], list[str]] = defaultdict(list)
158 + for layer_name, layer_tensors in tensors_per_block.items():
159 + view = layer_tensors[0]
160 + aliased_layers[(view.data_ptr(), view.stride())].append(layer_name)
161 +
162 + for tensor_layer_names in aliased_layers.values():
163 assert len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1
193 - assert (
194 - len({tensors_per_block[n][0].data_ptr() for n in tensor_layer_names})
195 - == 1
196 - )
197 - assert (
198 - len({tensors_per_block[n][0].stride() for n in tensor_layer_names}) == 1
199 - )
164
201 - # pick the first layer to represent the group
165 first_layer_name = tensor_layer_names[0]
166 for tensor in tensors_per_block[first_layer_name]:
167 block_tensors.append(
@@ -242,53 +205,6 @@ class OffloadingConnectorWorker:
205
206 self._init_worker(canonical_kv_caches)
207
245 - def register_cross_layers_kv_cache(
246 - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
247 - ):
248 - # verify that num_blocks is at physical position 0 in the cross-layers
249 - # tensor layout.
250 - test_shape = attn_backend.get_kv_cache_shape(
251 - num_blocks=1234, block_size=16, num_kv_heads=1, head_size=256
252 - )
253 - num_blocks_logical_dim = test_shape.index(1234) + 1
254 - physical_to_logical = attn_backend.get_kv_cache_stride_order(
255 - include_num_layers_dimension=True
256 - )
257 - num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim)
258 - assert num_blocks_physical_dim == 0
259 -
260 - kv_cache_groups = self.kv_cache_config.kv_cache_groups
261 - assert len(kv_cache_groups) == 1
262 - kv_cache_spec = kv_cache_groups[0].kv_cache_spec
263 - num_layers = len(kv_cache_groups[0].layer_names)
264 - page_size_bytes = kv_cache_spec.page_size_bytes * num_layers
265 -
266 - assert kv_cache.storage_offset() == 0
267 - storage = kv_cache.untyped_storage()
268 - assert len(storage) % page_size_bytes == 0
269 - num_blocks = len(storage) // page_size_bytes
270 - tensor = (
271 - torch.tensor(
272 - [],
273 - dtype=torch.int8,
274 - device=kv_cache.device,
275 - )
276 - .set_(storage)
277 - .view(num_blocks, page_size_bytes)
278 - )
279 - kv_cache_tensor = CanonicalKVCacheTensor(
280 - tensor=tensor, page_size_bytes=page_size_bytes
281 - )
282 - # in cross layers layout, there's currently only a single group
283 - kv_cache_data_ref = CanonicalKVCacheRef(
284 - tensor_idx=0, page_size_bytes=page_size_bytes
285 - )
286 - canonical_kv_caches = CanonicalKVCaches(
287 - tensors=[kv_cache_tensor], group_data_refs=[[kv_cache_data_ref]]
288 - )
289 -
290 - self._init_worker(canonical_kv_caches)
291 -
208 def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
209 assert self.worker is not None
210
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+2 -13
@@ -37,7 +37,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
37 OffloadingConnectorWorker,
38 )
39 from vllm.forward_context import ForwardContext
40 -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
40 +from vllm.v1.attention.backend import AttentionMetadata
41 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
42 from vllm.v1.core.sched.output import SchedulerOutput
43 from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -47,11 +47,6 @@ from vllm.v1.request import Request
47
48
49 class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
50 - @property
51 - def prefer_cross_layer_blocks(self) -> bool:
52 - # Cross-layer slabs have no per-layer refs to certify canonical mappings
53 - return not self._canonical_layout
54 -
50 @property
51 def requires_kv_delivery(self) -> bool:
52 # Runs as kv_both, but is a best-effort cache: a dropped save is just a
@@ -91,12 +86,6 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
86 assert self.connector_worker is not None
87 self.connector_worker.register_kv_caches(kv_caches)
88
94 - def register_cross_layers_kv_cache(
95 - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
96 - ):
97 - assert self.connector_worker is not None
98 - self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend)
99 -
89 def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
90 assert self.connector_worker is not None
91 assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)
@@ -196,7 +185,7 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
185
186 @classmethod
187 def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None:
199 - return "HND"
188 + return "LBHNC"
189
190 def reset_cache(self) -> bool | None:
191 assert self.connector_scheduler is not None
vllm/envs.py
+13 -6
@@ -242,7 +242,10 @@ if TYPE_CHECKING:
242 VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
243 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
244 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5
245 - VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
245 + VLLM_KV_CACHE_LAYOUT: (
246 + Literal["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"]
247 + | None
248 + ) = None
249 VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
250 VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
251 VLLM_RAISE_ON_LOGIT_NANS: bool = False
@@ -1784,18 +1787,22 @@ environment_variables: dict[str, Callable[[], Any]] = {
1787 ),
1788 # KV Cache layout used throughout vllm.
1789 # Some common values are:
1787 - # - NHD
1788 - # - HND
1789 - # Where N=num_blocks, H=num_heads and D=head_size. The default value will
1790 + # - LBNHC
1791 + # - LBHNC
1792 + # Where N=num_states, H=num_heads and C=state_content. The default value will
1793 # leave the layout choice to the backend. Mind that backends may only
1794 # implement and support a subset of all possible layouts.
1795 + # LHBNC hoists the head dim outside the block dim; backends must opt in via
1796 + # AttentionBackend.supported_kv_cache_layouts().
1797 "VLLM_KV_CACHE_LAYOUT": env_with_choices(
1793 - "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
1798 + "VLLM_KV_CACHE_LAYOUT",
1799 + None,
1800 + ["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"],
1801 ),
1802 # SSM conv state layout used for Mamba models.
1803 # - SD: (state_len, dim) — dim contiguous (default)
1804 # - DS: (dim, state_len) — TP-sharded dim on dim1,
1798 - # consistent with SSM temporal state and HND KV cache layout.
1805 + # consistent with SSM temporal state and LBHNC KV cache layout.
1806 "VLLM_SSM_CONV_STATE_LAYOUT": env_with_choices(
1807 "VLLM_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"]
1808 ),
vllm/model_executor/layers/attention/mla_attention.py
+4 -21
@@ -649,6 +649,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
649 compile_native=True,
650 )
651
652 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
653 + # [B, H=1, N, C] -> [B, N, C]
654 + self.kv_cache = kv_cache.squeeze(1)
655 +
656 @property
657 def chunked_prefill_workspace_size(self) -> int:
658 if self._chunked_prefill_workspace_size is None:
@@ -1451,27 +1455,6 @@ class MLACommonBackend(AttentionBackend):
1455 def get_builder_cls() -> type["MLACommonMetadataBuilder"]:
1456 return MLACommonMetadataBuilder
1457
1454 - @staticmethod
1455 - def get_kv_cache_shape(
1456 - num_blocks: int,
1457 - block_size: int,
1458 - num_kv_heads: int, # assumed to be 1 for MLA
1459 - head_size: int,
1460 - cache_dtype_str: str = "auto",
1461 - ) -> tuple[int, ...]:
1462 - return (num_blocks, block_size, head_size)
1463 -
1464 - @staticmethod
1465 - def get_kv_cache_stride_order(
1466 - include_num_layers_dimension: bool = False,
1467 - ) -> tuple[int, ...]:
1468 - if include_num_layers_dimension:
1469 - # Default to identity permutation to signal cross-layer allocation
1470 - # is unsupported. Each MLA backend must opt in to support cross-layer
1471 - # allocation by overriding this method.
1472 - return (0, 1, 2, 3)
1473 - return (0, 1, 2)
1474 -
1458 @classmethod
1459 def get_supported_head_sizes(cls) -> list[int]:
1460 return [320, 576]
vllm/model_executor/layers/hpc/rope_norm.py
+15 -12
@@ -61,16 +61,16 @@ def hpc_rope_norm_forward(
61 return
62
63 attn_layer = forward_context.no_compile_layers[layer_name]
64 - # bind_kv_cache stores the per-layer KV cache as a single 5D tensor
65 - # (num_blocks, 2, block_size, num_kv_heads, head_size), so use it directly.
64 + # bind_kv_cache stores the per-layer KV cache as a single tensor
65 + # (num_blocks, num_kv_heads, block_size, 2 * head_size), so use it directly.
66 kv_cache = attn_layer.kv_cache
67
68 if kv_cache.numel() == 0:
69 output.zero_()
70 return
71
72 - assert kv_cache.dim() == 5, (
73 - f"Expected kv_cache to have 5 dims, got {tuple(kv_cache.shape)}"
72 + assert kv_cache.dim() == 4, (
73 + f"Expected kv_cache to have 4 dims, got {tuple(kv_cache.shape)}"
74 )
75
76 rope_norm = _hpc_rope_norm_instances[layer_name]
@@ -295,6 +295,9 @@ class HpcRopeNorm(CustomOp, HpcModule):
295 """
296 import hpc
297
298 + # (B, H, N, 2*hs) -> two (B, N, H, hs) K/V views, as in HpcAttentionImpl.
299 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_dim, dim=-1)
300 +
301 num_actual_tokens = attn_metadata.num_actual_tokens
302 num_prefill_reqs = attn_metadata.num_prefills
303 num_decode_reqs = attn_metadata.num_decodes
@@ -333,8 +336,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
336
337 if self.use_fp8:
338 _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
336 - key_cache=kv_cache[:, 0],
337 - value_cache=kv_cache[:, 1],
339 + key_cache=key_cache,
340 + value_cache=value_cache,
341 qkv=qkv_prefill,
342 cos_sin=self.cos_sin_cache,
343 num_seqlen_per_req=seq_lens_prefill,
@@ -353,8 +356,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
356 attn_metadata.hpc_prefill_q_scale = q_scale
357 else:
358 hpc.rope_norm_store_kv(
356 - kv_cache[:, 0],
357 - kv_cache[:, 1],
359 + key_cache,
360 + value_cache,
361 qkv_prefill,
362 self.cos_sin_cache,
363 seq_lens_prefill,
@@ -379,8 +382,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
382
383 if self.use_fp8:
384 _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
382 - key_cache=kv_cache[:, 0],
383 - value_cache=kv_cache[:, 1],
385 + key_cache=key_cache,
386 + value_cache=value_cache,
387 qkv=qkv_decode,
388 cos_sin=self.cos_sin_cache,
389 num_seqlen_per_req=num_seq_kvcache,
@@ -401,8 +404,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
404 attn_metadata.hpc_split_k_flag = split_k_flag
405 else:
406 hpc.rope_norm_store_kv(
404 - kv_cache[:, 0],
405 - kv_cache[:, 1],
407 + key_cache,
408 + value_cache,
409 qkv_decode,
410 self.cos_sin_cache,
411 num_seq_kvcache,
vllm/model_executor/models/deepseek_v2.py
+5
@@ -645,6 +645,11 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
645 raise ValueError(f"Duplicate layer name: {prefix}")
646 compilation_config.static_forward_context[prefix] = self
647
648 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
649 + # [B, H=1, N, C] -> [B, N, C]: the indexer kernels and
650 + # kv_cache_as_quant_view index a 3-D block-major cache.
651 + self.kv_cache = kv_cache.squeeze(1)
652 +
653 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
654 return MLAAttentionSpec(
655 block_size=self.cache_config.block_size,
vllm/model_executor/models/extract_hidden_states.py
+3 -15
@@ -80,14 +80,14 @@ def dummy_attention(layer_name, _placeholder):
80
81 def basic_cache(
82 to_cache: torch.Tensor, # shape: [seq_len, num_heads, head_size]
83 - kv_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size]
83 + kv_cache: torch.Tensor, # shape: [num_blocks, num_heads, block_size, head_size]
84 slot_mapping: torch.Tensor, # shape: [seq_len]
85 ):
86 # Padding slots are -1; redirect them to the null block (block 0, never
87 # allocated to a request) so the scatter stays branch-free and sync-free.
88 - block_size = kv_cache.shape[1]
88 + block_size = kv_cache.shape[2]
89 slot_mapping = slot_mapping.clamp_min(0)
90 - kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache
90 + kv_cache[slot_mapping // block_size, :, slot_mapping % block_size] = to_cache
91
92
93 ######### CacheOnlyAttentionBackend ########
@@ -123,18 +123,6 @@ class CacheOnlyAttentionBackend(AttentionBackend):
123 def get_impl_cls() -> type["CacheOnlyAttentionImpl"]:
124 return CacheOnlyAttentionImpl
125
126 - @staticmethod
127 - def get_kv_cache_shape(
128 - num_blocks: int,
129 - block_size: int,
130 - num_kv_heads: int,
131 - head_size: int,
132 - cache_dtype_str: str = "auto",
133 - ) -> tuple[int, ...]:
134 - # We set `num_kv_heads = num_hidden_layers` and `head_size = hidden_size`
135 - # We also don't use a k/v (2) dim
136 - return (num_blocks, block_size, num_kv_heads, head_size)
137 -
126 @staticmethod
127 def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]:
128 return CacheOnlyAttentionMetadataBuilder
vllm/model_executor/models/whisper_causal.py
+6 -22
@@ -5,6 +5,7 @@ import functools
5 import logging
6 import math
7 from dataclasses import replace
8 +from fractions import Fraction
9 from functools import partial
10
11 import torch
@@ -128,14 +129,10 @@ def create_whisper_attention_backend_with_block_pooling(
129 vllm_config: VllmConfig,
130 device: torch.device,
131 ):
131 - assert kv_cache_spec.num_kv_heads % block_pool_size == 0
132 - # Scale pooled-unit quantities back to unpooled (encoder-token)
133 - # units for the underlying attention metadata: the KV cache spec
134 - # counts blocks in pooled units, but the kernel runs on the
135 - # `block_pool_size`x-expanded sequence built below.
136 - spec_overrides = dict(
132 + kv_cache_spec = replace(
133 + kv_cache_spec,
134 block_size=kv_cache_spec.block_size * block_pool_size,
138 - num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
135 + tokens_per_state=1,
136 )
137 if isinstance(kv_cache_spec, SlidingWindowSpec):
138 # The manager keeps `sliding_window` in pooled units; the kernel
@@ -143,8 +140,7 @@ def create_whisper_attention_backend_with_block_pooling(
140 # unpooled units (`get_kv_cache_spec` sizes the pooled window so
141 # this window always stays resident).
142 assert sliding_window is not None
146 - spec_overrides["sliding_window"] = sliding_window
147 - kv_cache_spec = replace(kv_cache_spec, **spec_overrides)
143 + kv_cache_spec = replace(kv_cache_spec, sliding_window=sliding_window)
144 super().__init__(kv_cache_spec, layer_names, vllm_config, device)
145 # Override model_config-derived values with the actual
146 # encoder values from kv_cache_spec
@@ -261,18 +257,6 @@ def create_whisper_attention_backend_with_block_pooling(
257 overrides={
258 "get_builder_cls": lambda: WhisperCausalAttentionWithBlockPoolingBuilder,
259 "get_impl_cls": lambda: WhisperCausalAttentionWithBlockPoolingImpl,
264 - "get_kv_cache_shape": lambda num_blocks,
265 - block_size,
266 - num_kv_heads,
267 - head_size,
268 - cache_dtype_str: underlying_attn_backend.get_kv_cache_shape(
269 - num_blocks,
270 - # we stretch each block by `block_pool_size`
271 - block_size * block_pool_size,
272 - num_kv_heads // block_pool_size,
273 - head_size,
274 - cache_dtype_str,
275 - ),
260 "forward_includes_kv_cache_update": True,
261 },
262 )
@@ -341,7 +325,7 @@ class WhisperCausalAttentionWithBlockPooling(Attention):
325 assert isinstance(kv_cache_spec, AttentionSpec)
326 kv_cache_spec = replace(
327 kv_cache_spec,
344 - num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
328 + tokens_per_state=Fraction(1, self.block_pool_size),
329 )
330 if isinstance(kv_cache_spec, SlidingWindowSpec):
331 # The manager counts blocks in pooled units, so express the window in
vllm/models/deepseek_v4/attention.py
+20 -4
@@ -653,7 +653,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
653 # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
654 # KV row.
655 block_size = swa_metadata.block_size
656 - swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim)
656 + assert swa_kv_cache.shape[1:] == (block_size, self.head_dim)
657 + swa_kv_cache_3d = swa_kv_cache
658 if cache_dtype == torch.bfloat16:
659 torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
660 q,
@@ -684,6 +685,17 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
685 )
686 return q_fp8
687
688 + def _global_topk_output_buffers(
689 + self, topk_indices: torch.Tensor
690 + ) -> tuple[torch.Tensor, torch.Tensor] | None:
691 + if self.compress_ratio != 4 or self.eager_scratch_pool is None:
692 + return None
693 + return self.eager_scratch_pool.global_topk_outputs(topk_indices)
694 +
695 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
696 + # [B, H=1, N, C] -> [B, N, C]
697 + self.kv_cache = kv_cache.squeeze(1)
698 +
699 def get_attn_backend(self) -> type[AttentionBackend]:
700 return self.backend_cls
701
@@ -701,7 +713,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
713 num_kv_heads=1,
714 head_size=self.head_dim,
715 dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
704 - compress_ratio=self.compress_ratio,
716 + tokens_per_state=self.compress_ratio,
717 cache_dtype_str=self.kv_cache_dtype,
718 alignment=576 if uses_fp8_ds_mla_layout else 512,
719 model_version="deepseek_v4",
@@ -733,16 +745,20 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
745 raise ValueError(f"Duplicate layer name: {prefix}")
746 compilation_config.static_forward_context[prefix] = self
747
748 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
749 + # [B, H=1, N, C] -> [B, N, C]
750 + self.kv_cache = kv_cache.squeeze(1)
751 +
752 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
753 # head_dim already carries the fp8 scale padding
738 - # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout.
754 + # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout.
755 uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
756 return MLAAttentionSpec(
757 block_size=self.cache_config.block_size,
758 num_kv_heads=1,
759 head_size=self.head_dim,
760 dtype=self.dtype,
745 - compress_ratio=self.compress_ratio,
761 + tokens_per_state=self.compress_ratio,
762 # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
763 alignment=576 if uses_fp8_ds_mla_layout else 512,
764 )
vllm/models/deepseek_v4/compressor.py
+4 -19
@@ -75,25 +75,6 @@ class CompressorBackend(AttentionBackend):
75 def get_builder_cls() -> type["CompressorMetadataBuilder"]:
76 return CompressorMetadataBuilder
77
78 - @staticmethod
79 - def get_kv_cache_shape(
80 - num_blocks: int,
81 - block_size: int,
82 - num_kv_heads: int,
83 - head_size: int,
84 - cache_dtype_str: str = "auto",
85 - ) -> tuple[int, ...]:
86 - assert num_kv_heads == 1
87 - return (num_blocks, block_size, head_size)
88 -
89 - @staticmethod
90 - def get_kv_cache_stride_order(
91 - include_num_layers_dimension: bool = False,
92 - ) -> tuple[int, ...]:
93 - if include_num_layers_dimension:
94 - return (0, 1, 2, 3)
95 - return (0, 1, 2)
96 -
78
79 @dataclass
80 class CompressorMetadata:
@@ -185,6 +166,10 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
166 else:
167 raise ValueError(f"Invalid compress ratio: {compress_ratio}")
168
169 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
170 + # [B, H=1, N, C] -> [B, N, C]
171 + self.kv_cache = kv_cache.squeeze(1)
172 +
173 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
174 # fp8_ds_mla is the UE8M0 paged layout and needs 576B alignment. Plain
175 # full-cache rows share state pages with contiguous KV pages, so padding
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
-21
@@ -22,7 +22,6 @@ from vllm.models.deepseek_v4.sparse_mla import (
22 DeepseekV4FlashMLAMetadata,
23 DeepseekV4SparseMLABackend,
24 )
25 -from vllm.platforms import current_platform
25 from vllm.platforms.interface import DeviceCapability
26 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
27 from vllm.v1.attention.backend import MultipleOf
@@ -165,26 +164,6 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
164 return None
165 return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"
166
168 - @staticmethod
169 - def get_kv_cache_shape(
170 - num_blocks: int,
171 - block_size: int,
172 - num_kv_heads: int,
173 - head_size: int,
174 - cache_dtype_str: str = "auto",
175 - ) -> tuple[int, ...]:
176 - device_capability = current_platform.get_device_capability()
177 - if device_capability is not None and device_capability.major == 12:
178 - return DeepseekV4SparseMLABackend.get_kv_cache_shape(
179 - num_blocks,
180 - block_size,
181 - num_kv_heads,
182 - head_size,
183 - cache_dtype_str,
184 - )
185 - assert num_kv_heads == 1
186 - return (num_blocks, block_size, head_size)
187 -
167
168 class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
169 """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
vllm/models/deepseek_v4/sparse_mla.py
+3 -18
@@ -88,21 +88,6 @@ class DeepseekV4SparseMLABackend(AttentionBackend):
88 def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
89 return capability.major in [9, 10]
90
91 - @staticmethod
92 - def get_kv_cache_shape(
93 - num_blocks: int,
94 - block_size: int,
95 - num_kv_heads: int,
96 - head_size: int,
97 - cache_dtype_str: str = "auto",
98 - ) -> tuple[int, ...]:
99 - if cache_dtype_str == "fp8_ds_mla":
100 - # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
101 - # head_size passed in is the semantic head_dim (512).
102 - return (num_blocks, block_size, 584)
103 - else:
104 - return (num_blocks, block_size, head_size)
105 -
91
92 @dataclass
93 class DeepseekV4FlashMLAMetadata(AttentionMetadata):
@@ -151,8 +136,8 @@ class DeepseekV4SparseMLAMetadataBuilder(
136 (max_num_batched_tokens,), dtype=torch.int32, device=device
137 )
138
154 - assert hasattr(self.kv_cache_spec, "compress_ratio")
155 - self.compress_ratio = self.kv_cache_spec.compress_ratio
139 + assert isinstance(self.kv_cache_spec.tokens_per_state, int)
140 + self.compress_ratio = self.kv_cache_spec.tokens_per_state
141
142 # Pre-allocate compressed slot mapping buffer for CUDA graph address
143 # stability when compress_ratio > 1.
@@ -206,7 +191,7 @@ class DeepseekV4SparseMLAMetadataBuilder(
191 cm.query_start_loc,
192 cm.seq_lens,
193 cm.block_table_tensor.clamp_(min=0),
209 - int(self.kv_cache_spec.storage_block_size),
194 + int(self.kv_cache_spec.num_states),
195 self.compress_ratio,
196 out=self.compressed_slot_mapping_buffer,
197 )
vllm/models/inkling/amd/sconv_swa_attn.py
-26
@@ -107,32 +107,6 @@ class InklingSconvBackend(AttentionBackend):
107 def get_name() -> str:
108 return "INKLING_SCONV_SWA"
109
110 - @classmethod
111 - def indexes_kv_by_block_stride(cls) -> bool:
112 - # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the
113 - # padded conv page is read through a strided view.
114 - return True
115 -
116 - @staticmethod
117 - def get_kv_cache_shape(
118 - num_blocks: int,
119 - block_size: int,
120 - num_kv_heads: int,
121 - head_size: int,
122 - cache_dtype_str: str = "auto",
123 - ) -> tuple[int, ...]:
124 - # HND, num-blocks-first, head-major: [num_blocks, H, N, D].
125 - return (num_blocks, num_kv_heads, block_size, head_size)
126 -
127 - @staticmethod
128 - def get_kv_cache_stride_order(
129 - include_num_layers_dimension: bool = False,
130 - ) -> tuple[int, ...]:
131 - # Identity: physical layout == logical [num_blocks, H, N, D].
132 - if include_num_layers_dimension:
133 - return (0, 1, 2, 3, 4)
134 - return (0, 1, 2, 3)
135 -
110 @staticmethod
111 def get_impl_cls():
112 raise NotImplementedError(
vllm/models/inkling/nvidia/sconv_swa_attn.py
-26
@@ -107,32 +107,6 @@ class InklingSconvBackend(AttentionBackend):
107 def get_name() -> str:
108 return "INKLING_SCONV_SWA"
109
110 - @classmethod
111 - def indexes_kv_by_block_stride(cls) -> bool:
112 - # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the
113 - # padded conv page is read through a strided view.
114 - return True
115 -
116 - @staticmethod
117 - def get_kv_cache_shape(
118 - num_blocks: int,
119 - block_size: int,
120 - num_kv_heads: int,
121 - head_size: int,
122 - cache_dtype_str: str = "auto",
123 - ) -> tuple[int, ...]:
124 - # HND, num-blocks-first, head-major: [num_blocks, H, N, D].
125 - return (num_blocks, num_kv_heads, block_size, head_size)
126 -
127 - @staticmethod
128 - def get_kv_cache_stride_order(
129 - include_num_layers_dimension: bool = False,
130 - ) -> tuple[int, ...]:
131 - # Identity: physical layout == logical [num_blocks, H, N, D].
132 - if include_num_layers_dimension:
133 - return (0, 1, 2, 3, 4)
134 - return (0, 1, 2, 3)
135 -
110 @staticmethod
111 def get_impl_cls():
112 raise NotImplementedError(
vllm/models/kimi_k3/nvidia/mla.py
+4
@@ -384,6 +384,10 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
384 # ------------------------------------------------------------------
385 # AttentionLayerBase interface
386 # ------------------------------------------------------------------
387 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
388 + # [B, H=1, N, C] -> [B, N, C]
389 + self.kv_cache = kv_cache.squeeze(1)
390 +
391 def get_attn_backend(self) -> type[AttentionBackend]:
392 return self.attn_backend
393
vllm/models/minimax_m3/amd/model.py
+8
@@ -728,6 +728,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
728
729 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
730 # Main GQA K/V cache. Block size may change after load, refresh it.
731 + # AITER sparse PA wants K and V as two separate head groups
732 + # (H folded into the content dim).
733 + sparse_pa = self.use_aiter_sparse_pa
734 + kv_bytes = (
735 + self.num_kv_heads * self.head_dim * self.kv_cache_torch_dtype.itemsize
736 + )
737 return FullAttentionSpec(
738 block_size=vllm_config.cache_config.block_size,
739 num_kv_heads=self.num_kv_heads,
@@ -735,6 +741,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
741 head_size_v=self.head_dim,
742 dtype=self.kv_cache_torch_dtype,
743 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
744 + num_head_slots=2 if sparse_pa else None,
745 + state_content_bytes=kv_bytes if sparse_pa else None,
746 )
747
748 def _ensure_aiter_sparse_pa_kv_cache(self) -> None:
vllm/models/minimax_m3/common/indexer.py
+4 -19
@@ -97,25 +97,6 @@ class MiniMaxM3IndexerBackend(AttentionBackend):
97 def is_sparse(cls) -> bool:
98 return True
99
100 - @staticmethod
101 - def get_kv_cache_shape(
102 - num_blocks: int,
103 - block_size: int,
104 - num_kv_heads: int,
105 - head_size: int,
106 - cache_dtype_str: str = "auto",
107 - ) -> tuple[int, ...]:
108 - return (num_blocks, block_size, head_size)
109 -
110 - @staticmethod
111 - def get_kv_cache_stride_order(
112 - include_num_layers_dimension: bool = False,
113 - ) -> tuple[int, ...]:
114 - if include_num_layers_dimension:
115 - # M3 does not use cross-layer (per-layer-stacked) KV blocks.
116 - raise NotImplementedError
117 - return (0, 1, 2)
118 -
100
101 class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
102 """Side KV cache for the indexer's per-token index keys (key-only).
@@ -156,6 +137,10 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
137 raise ValueError(f"Duplicate layer name: {prefix}")
138 compilation_config.static_forward_context[prefix] = self
139
140 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
141 + # [B, H=1, N, C] -> [B, N, C]
142 + self.kv_cache = kv_cache.squeeze(1)
143 +
144 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
145 # Key-only: MLAAttentionSpec budgets one vector/token (not 2x for K+V).
146 return MLAAttentionSpec(
vllm/models/minimax_m3/common/sparse_attention.py
+11 -42
@@ -51,10 +51,13 @@ from vllm.v1.attention.backend import (
51 MultipleOf,
52 )
53 from vllm.v1.attention.backends.utils import (
54 - get_kv_cache_layout,
54 split_decodes_and_prefills,
55 )
57 -from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
56 +from vllm.v1.kv_cache_interface import (
57 + AttentionSpec,
58 + KVCacheLayout,
59 + is_quantized_kv_cache,
60 +)
61
62 logger = init_logger(__name__)
63
@@ -112,47 +115,13 @@ class MiniMaxM3SparseBackend(AttentionBackend):
115 def is_sparse(cls) -> bool:
116 return True
117
115 - @staticmethod
116 - def get_kv_cache_shape(
117 - num_blocks: int,
118 - block_size: int,
119 - num_kv_heads: int,
120 - head_size: int,
121 - cache_dtype_str: str = "auto",
122 - ) -> tuple[int, ...]:
123 - if minimax_m3_use_aiter_sparse_pa(num_kv_heads):
124 - # AITER's assembly paged-attention kernels require independently
125 - # contiguous K and V storage. Keep that specialized layout behind
126 - # the shuffle flag while every other implementation uses the
127 - # packed-content contract introduced by #44455.
128 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
129 - # K and V are packed into the content dim: logical (B, H, N, 2*hs).
130 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
131 -
132 - @staticmethod
133 - def get_kv_cache_stride_order(
134 - include_num_layers_dimension: bool = False,
135 - ) -> tuple[int, ...]:
136 - # `stride_order` indicates the permutation that gets us from
137 - # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
138 - # layout we want.
139 - if include_num_layers_dimension:
140 - raise NotImplementedError # no cross-layer KV blocks in M3
118 + @classmethod
119 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
120 + # AITER sparse PA unbinds K and V as separate contiguous regions, which
121 + # requires the head dimension to sit outside the block dimension.
122 if _minimax_m3_aiter_sparse_pa_requested():
142 - # The AITER page-16 sparse PA path reinterprets the K and V slices
143 - # as separate SHUFFLE caches. Keep K/V physically separated so
144 - # kv_cache[:, 0] and kv_cache[:, 1] are contiguous byte ranges.
145 - return (1, 0, 2, 3, 4)
146 - cache_layout = get_kv_cache_layout()
147 - if cache_layout == "NHD":
148 - # (num_blocks, block_size, num_kv_heads, 2*head_size)
149 - stride_order = (0, 2, 1, 3)
150 - elif cache_layout == "HND":
151 - # (num_blocks, num_kv_heads, block_size, 2*head_size)
152 - stride_order = (0, 1, 2, 3)
153 - else:
154 - raise ValueError(f"Unknown cache layout format {cache_layout}.")
155 - return stride_order
123 + return (KVCacheLayout.LHBNC,)
124 + return super().supported_kv_cache_layouts()
125
126
127 @dataclass
vllm/utils/torch_utils.py
+15
@@ -117,6 +117,21 @@ def is_strictly_contiguous(t: torch.Tensor) -> bool:
117 return True
118
119
120 +def is_non_overlapping_and_dense(t: torch.Tensor) -> bool:
121 + """Check if the tensor's elements cover one gapless, non-overlapping byte
122 + range, in any dimension order (i.e. a permuted view of a contiguous
123 + tensor); ``is_contiguous()`` additionally requires row-major order.
124 + """
125 + expected_stride = 1
126 + for size, stride in sorted(zip(t.shape, t.stride()), key=lambda p: p[1]):
127 + if size == 1:
128 + continue
129 + if stride != expected_stride:
130 + return False
131 + expected_stride *= size
132 + return True
133 +
134 +
135 def canonicalize_singleton_dim_strides(t: torch.Tensor) -> torch.Tensor:
136 """Fix degenerate strides on size=1 dimensions for CUDA TMA compatibility.
137
vllm/v1/attention/backend.py
+10 -98
@@ -23,10 +23,14 @@ if TYPE_CHECKING:
23 from vllm.model_executor.layers.linear import ColumnParallelLinear
24 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
25 from vllm.platforms.interface import DeviceCapability
26 - from vllm.v1.attention.backends.utils import KVCacheLayoutType
27 - from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, KVQuantMode
26 + from vllm.v1.kv_cache_interface import (
27 + AttentionSpec,
28 + KVCacheLayout,
29 + KVCacheSpec,
30 + KVQuantMode,
31 + )
32
29 -from vllm.v1.kv_cache_interface import get_kv_quant_mode
33 +from vllm.v1.kv_cache_interface import KVCacheLayout, get_kv_quant_mode
34
35
36 class AttentionType(str, Enum):
@@ -84,68 +88,6 @@ class AttentionBackend(ABC):
88 def get_builder_cls(): # -> Type["AttentionMetadataBuilder"]:
89 raise NotImplementedError
90
87 - @staticmethod
88 - @abstractmethod
89 - def get_kv_cache_shape(
90 - num_blocks: int,
91 - block_size: int,
92 - num_kv_heads: int,
93 - head_size: int,
94 - cache_dtype_str: str = "auto",
95 - ) -> tuple[int, ...]:
96 - raise NotImplementedError
97 -
98 - @classmethod
99 - def get_kv_cache_block_dim(
100 - cls,
101 - block_size: int,
102 - num_kv_heads: int,
103 - head_size: int,
104 - cache_dtype_str: str = "auto",
105 - ) -> int:
106 - """Discover which tensor dim is the block index, since different
107 - backends lay out dims differently."""
108 - _S = 1234567
109 - shape = cls.get_kv_cache_shape(
110 - _S,
111 - block_size,
112 - num_kv_heads,
113 - head_size,
114 - cache_dtype_str=cache_dtype_str,
115 - )
116 - return shape.index(_S)
117 -
118 - @staticmethod
119 - def get_kv_cache_stride_order(
120 - include_num_layers_dimension: bool = False,
121 - ) -> tuple[int, ...]:
122 - """
123 - Get the physical (memory layout) ordering of the kv cache dimensions.
124 - Standard attention backends pack K and V into the content dim, giving
125 - the logical shape [num_blocks, num_heads, block_size, 2 * head_size].
126 - e.g. if get_kv_cache_stride_order returns (0, 2, 1, 3) then the physical
127 - ordering of dimensions is
128 - [num_blocks, block_size, num_heads, 2 * head_size].
129 -
130 - If this function is unimplemented / raises NotImplementedError,
131 - the physical layout of the KV cache will match the logical shape.
132 -
133 - Args:
134 - include_num_layers_dimension: if True, includes an additional
135 - num_layers dimension, which is assumed to be prepended
136 - to the logical KV cache shape.
137 - With the above example, a return value (1, 0, 3, 2, 4)
138 - corresponds to
139 - [num_blocks, num_layers, block_size, num_heads, 2 * head_size].
140 -
141 - If an additional dimension is NOT included in the returned
142 - tuple, the physical layout will not include a layers dimension.
143 -
144 - Returns:
145 - A tuple of ints which is a permutation of range(len(shape)).
146 - """
147 - raise NotImplementedError
148 -
91 @classmethod
92 def full_cls_name(cls) -> tuple[str, str]:
93 return (cls.__module__, cls.__qualname__)
@@ -215,38 +157,6 @@ class AttentionBackend(ABC):
157
158 return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)
159
218 - @classmethod
219 - def indexes_kv_by_block_stride(cls) -> bool:
220 - """Whether the backend reads KV pages by the runtime block stride.
221 -
222 - True when ``num_blocks`` is the outermost physical dimension of the KV
223 - cache, so the backend tolerates a non-contiguous block dim. This gates
224 - page size padding and cross-layer uniform KV layout.
225 -
226 - Returns:
227 - True if the backend's physical KV layout is num-blocks-first. False
228 - otherwise, including when the backend does not define a layered
229 - stride order.
230 - """
231 - try:
232 - kv_cache_stride_order = cls.get_kv_cache_stride_order(
233 - include_num_layers_dimension=False
234 - )
235 - layered_kv_cache_stride_order = cls.get_kv_cache_stride_order(
236 - include_num_layers_dimension=True
237 - )
238 - except (AttributeError, NotImplementedError):
239 - return False
240 -
241 - # Check that attention backend includes a layers dimension.
242 - if len(layered_kv_cache_stride_order) != len(kv_cache_stride_order) + 1:
243 - return False
244 -
245 - # stride_order[0] == 0 means num_layers stays first in physical
246 - # layout (identity permutation), so indexing by block stride is
247 - # not supported.
248 - return layered_kv_cache_stride_order[0] != 0
249 -
160 @classmethod
161 def is_mla(cls) -> bool:
162 return False
@@ -438,7 +348,9 @@ class AttentionBackend(ABC):
348 return invalid_reasons
349
350 @classmethod
441 - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
351 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
352 + """Layouts this backend's kernels can consume, most preferred first, or
353 + None when the kernels consume any layout and express no preference."""
354 return None
355
356 @classmethod
vllm/v1/attention/backends/cpu_attn.py
+6 -15
@@ -30,13 +30,13 @@ from vllm.v1.attention.backend import (
30 MultipleOf,
31 )
32 from vllm.v1.attention.backends.utils import (
33 - KVCacheLayoutType,
33 get_num_attention_heads_from_layers,
34 )
35 from vllm.v1.kv_cache_interface import (
36 AttentionSpec,
37 CrossAttentionSpec,
38 EncoderOnlyAttentionSpec,
39 + KVCacheLayout,
40 )
41
42 logger = init_logger(__name__)
@@ -65,6 +65,11 @@ class CPUAttentionBackend(AttentionBackend):
65 def get_supported_head_sizes(cls) -> list[int]:
66 return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512]
67
68 + @classmethod
69 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
70 + # The CPU backend only reads head-major block interiors.
71 + return (KVCacheLayout.LBHNC,)
72 +
73 @staticmethod
74 def get_name() -> str:
75 return "CPU_ATTN"
@@ -96,20 +101,6 @@ class CPUAttentionBackend(AttentionBackend):
101 def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]:
102 return CPUAttentionMetadataBuilder
103
99 - @staticmethod
100 - def get_kv_cache_shape(
101 - num_blocks: int,
102 - block_size: int,
103 - num_kv_heads: int,
104 - head_size: int,
105 - cache_dtype_str: str = "auto",
106 - ) -> tuple[int, ...]:
107 - return num_blocks, num_kv_heads, block_size, 2 * head_size
108 -
109 - @classmethod
110 - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
111 - return "HND"
112 -
104 @staticmethod
105 def use_cascade_attention(*args, **kwargs) -> bool:
106 return False
vllm/v1/attention/backends/flash_attn.py
+5 -44
@@ -33,7 +33,6 @@ from vllm.v1.attention.backends.fa_utils import (
33 from vllm.v1.attention.backends.utils import (
34 fill_mm_prefix_query_ranges,
35 get_dcp_local_seq_lens,
36 - get_num_attention_heads_from_layers,
36 )
37 from vllm.v1.attention.ops.dcp import (
38 cp_lse_ag_out_rs,
@@ -65,7 +64,6 @@ from vllm.v1.attention.backend import (
64 AttentionMetadataBuilder,
65 CommonAttentionMetadata,
66 )
68 -from vllm.v1.attention.backends.utils import get_kv_cache_layout
67 from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
68 from vllm.v1.worker.cp_utils import (
69 run_split_fa2_dcp_context_attention,
@@ -138,43 +136,6 @@ class FlashAttentionBackend(AttentionBackend):
136 def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]:
137 return FlashAttentionMetadataBuilder
138
141 - @staticmethod
142 - def get_kv_cache_shape(
143 - num_blocks: int,
144 - block_size: int,
145 - num_kv_heads: int,
146 - head_size: int,
147 - cache_dtype_str: str = "auto",
148 - ) -> tuple[int, ...]:
149 - if block_size % 16 != 0:
150 - raise ValueError("Block size must be a multiple of 16.")
151 - # K and V are packed into the content dim: logical (B, H, N, 2*D).
152 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
153 -
154 - @staticmethod
155 - def get_kv_cache_stride_order(
156 - include_num_layers_dimension: bool = False,
157 - ) -> tuple[int, ...]:
158 - # `stride_order` indicates the permutation that gets us from
159 - # `get_kv_cache_shape` (logical (B, H, N, 2*D)) to the actual memory
160 - # layout we want.
161 - cache_layout = get_kv_cache_layout()
162 - if cache_layout == "NHD" and include_num_layers_dimension:
163 - # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
164 - return (1, 0, 3, 2, 4)
165 - elif cache_layout == "NHD":
166 - # (num_blocks, block_size, num_kv_heads, 2*head_size)
167 - stride_order = (0, 2, 1, 3)
168 - elif cache_layout == "HND" and include_num_layers_dimension:
169 - # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
170 - return (1, 2, 0, 3, 4)
171 - elif cache_layout == "HND":
172 - # (num_blocks, num_kv_heads, block_size, 2*head_size)
173 - stride_order = (0, 1, 2, 3)
174 - else:
175 - raise ValueError(f"Unknown cache layout format {cache_layout}.")
176 - return stride_order
177 -
139 @classmethod
140 def supports_head_size(cls, head_size: int) -> bool:
141 if head_size % 8 != 0:
@@ -436,12 +397,12 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
397 self.compilation_config = vllm_config.compilation_config
398 self.attention_config = vllm_config.attention_config
399
439 - self.num_heads_q = get_num_attention_heads_from_layers(
440 - vllm_config, layer_names
441 - ) or self.model_config.get_num_attention_heads(self.parallel_config)
442 - self.num_heads_kv = kv_cache_spec.num_kv_heads
400 + self.num_heads_q = self.model_config.get_num_attention_heads(
401 + self.parallel_config
402 + )
403 + self.num_heads_kv = self.model_config.get_num_kv_heads(self.parallel_config)
404 self.kv_cache_dtype = kv_cache_spec.dtype
444 - self.headdim = kv_cache_spec.head_size
405 + self.headdim = self.model_config.get_head_size()
406 self.block_size = kv_cache_spec.block_size
407
408 self.max_num_splits = 0 # No upper bound on the number of splits.
vllm/v1/attention/backends/flash_attn_diffkv.py
-44
@@ -21,7 +21,6 @@ from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
21
22 if is_flash_attn_varlen_func_available():
23 from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
24 -from vllm.v1.attention.backends.utils import get_kv_cache_layout
24
25 from .flash_attn import (
26 FlashAttentionBackend,
@@ -73,49 +72,6 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
72 def get_impl_cls() -> type["FlashAttentionImpl"]:
73 return FlashAttentionDiffKVImpl
74
76 - @staticmethod
77 - def get_kv_cache_shape(
78 - num_blocks: int,
79 - block_size: int,
80 - num_kv_heads: int,
81 - head_size: int,
82 - cache_dtype_str: str = "auto",
83 - ) -> tuple[int, ...]:
84 - if block_size % 16 != 0:
85 - raise ValueError("Block size must be a multiple of 16.")
86 - # Logical (blocks-first, head-major) layout: K and V (with their
87 - # different head sizes) packed in the content dim.
88 - return (
89 - num_blocks,
90 - num_kv_heads,
91 - block_size,
92 - head_size + FlashAttentionDiffKVBackend.head_size_v,
93 - )
94 -
95 - @staticmethod
96 - def get_kv_cache_stride_order(
97 - include_num_layers_dimension: bool = False,
98 - ) -> tuple[int, ...]:
99 - # `stride_order` indicates the permutation that gets us from
100 - # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual
101 - # memory layout we want.
102 - cache_layout = get_kv_cache_layout()
103 - if cache_layout == "NHD" and include_num_layers_dimension:
104 - # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v)
105 - return (1, 0, 3, 2, 4)
106 - elif cache_layout == "NHD":
107 - # (num_blocks, block_size, num_kv_heads, C_k+C_v)
108 - stride_order = (0, 2, 1, 3)
109 - elif cache_layout == "HND" and include_num_layers_dimension:
110 - # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v)
111 - return (1, 2, 0, 3, 4)
112 - elif cache_layout == "HND":
113 - # (num_blocks, num_kv_heads, block_size, C_k+C_v)
114 - stride_order = (0, 1, 2, 3)
115 - else:
116 - raise ValueError(f"Unknown cache layout format {cache_layout}.")
117 - return stride_order
118 -
75
76 class FlashAttentionDiffKVImpl(FlashAttentionImpl):
77 vllm_flash_attn_version: int | None
vllm/v1/attention/backends/flashinfer.py
+36 -61
@@ -67,9 +67,8 @@ from vllm.v1.attention.backend import (
67 MultipleOf,
68 )
69 from vllm.v1.attention.backends.utils import (
70 - KVCacheLayoutType,
70 get_dcp_local_seq_lens,
72 - get_kv_cache_layout,
71 + get_flashinfer_layout_string,
72 get_num_attention_heads_from_layers,
73 get_per_layer_parameters,
74 infer_global_hyperparameters,
@@ -83,6 +82,7 @@ from vllm.v1.attention.ops.dcp import (
82 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
83 from vllm.v1.kv_cache_interface import (
84 AttentionSpec,
85 + KVCacheLayout,
86 KVCacheSpec,
87 KVQuantMode,
88 UniformTypeKVCacheSpecs,
@@ -282,6 +282,7 @@ def trtllm_prefill_attn_kvfp8_dequant(
282 class BatchDCPPrefillWrapper:
283 def __init__(
284 self,
285 + kv_layout: str,
286 workspace_buffer: torch.Tensor | None = None,
287 dcp_a2a: bool = False,
288 ):
@@ -289,9 +290,7 @@ class BatchDCPPrefillWrapper:
290 self._dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
291 else:
292 self._dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
292 - self._context = BatchPrefillWithPagedKVCacheWrapper(
293 - workspace_buffer, get_kv_cache_layout()
294 - )
293 + self._context = BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, kv_layout)
294 self._new_tokens = BatchPrefillWithRaggedKVCacheWrapper(workspace_buffer)
295
296 def plan(
@@ -467,44 +466,6 @@ class FlashInferBackend(AttentionBackend):
466 def get_builder_cls() -> type["FlashInferMetadataBuilder"]:
467 return FlashInferMetadataBuilder
468
470 - @staticmethod
471 - def get_kv_cache_shape(
472 - num_blocks: int,
473 - block_size: int,
474 - num_kv_heads: int,
475 - head_size: int,
476 - cache_dtype_str: str = "auto",
477 - ) -> tuple[int, ...]:
478 - if cache_dtype_str.startswith("nvfp4"):
479 - full_dim = nvfp4_kv_cache_full_dim(head_size)
480 - return (num_blocks, 2 * num_kv_heads, block_size, full_dim)
481 - # Pack K and V in the content dim (B, H, N, 2*hs).
482 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
483 -
484 - @staticmethod
485 - def get_kv_cache_stride_order(
486 - include_num_layers_dimension: bool = False,
487 - ) -> tuple[int, ...]:
488 - # `stride_order` indicates the permutation that gets us from
489 - # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
490 - # layout we want.
491 - cache_layout = get_kv_cache_layout()
492 - if cache_layout == "NHD" and include_num_layers_dimension:
493 - # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
494 - return (1, 0, 3, 2, 4)
495 - elif cache_layout == "NHD":
496 - # (num_blocks, block_size, num_kv_heads, 2*head_size)
497 - stride_order = (0, 2, 1, 3)
498 - elif cache_layout == "HND" and include_num_layers_dimension:
499 - # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
500 - return (1, 2, 0, 3, 4)
501 - elif cache_layout == "HND":
502 - # (num_blocks, num_kv_heads, block_size, 2*head_size)
503 - stride_order = (0, 1, 2, 3)
504 - else:
505 - raise ValueError(f"Unknown cache layout format {cache_layout}.")
506 - return stride_order
507 -
469 @staticmethod
470 def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:
471 if kv_cache_dtype in ("fp8", "fp8_e4m3"):
@@ -563,11 +524,13 @@ class FlashInferBackend(AttentionBackend):
524 ) and supports_trtllm_attention(is_prefill=True)
525
526 @classmethod
566 - def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
527 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
528 capability = current_platform.get_device_capability()
529 if capability is not None and capability.major == 10:
569 - return "HND"
570 - return None
530 + # The trtllm-gen kernels consume head-major block interiors; the L/B
531 + # nesting outside the block is immaterial to them.
532 + return (KVCacheLayout.LBHNC, KVCacheLayout.BLHNC)
533 + return super().supported_kv_cache_layouts()
534
535 forward_includes_kv_cache_update: bool = False
536
@@ -945,6 +908,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
908 self.paged_kv_indices = self._make_buffer(max_num_pages)
909 self.paged_kv_last_page_len = self._make_buffer(max_num_reqs)
910
911 + @property
912 + def kv_cache_layout(self) -> KVCacheLayout:
913 + return self.cache_config.get_resolved_kv_cache_layout()
914 +
915 # Keep SM90 prefill/decode Q dtype selection in one place.
916 def get_q_data_type(self, is_prefill: bool) -> torch.dtype:
917 # The user sets --attention-config.disable_flashinfer_q_quantization
@@ -1155,7 +1122,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1122 self._noncausal_prefill_wrapper = (
1123 BatchAttentionWithAttentionSinkWrapper(
1124 self._get_workspace_buffer(),
1158 - get_kv_cache_layout(),
1125 + get_flashinfer_layout_string(self.kv_cache_layout),
1126 backend="auto",
1127 q_data_type=self.q_data_type_prefill,
1128 kv_data_type=self.kv_cache_dtype,
@@ -1168,7 +1135,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1135 self._noncausal_prefill_wrapper = (
1136 BatchPrefillWithPagedKVCacheWrapper(
1137 self._get_workspace_buffer(),
1171 - get_kv_cache_layout(),
1138 + get_flashinfer_layout_string(self.kv_cache_layout),
1139 backend="auto",
1140 )
1141 )
@@ -1177,6 +1144,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1144 if self._prefill_wrapper is None:
1145 if self.use_dcp:
1146 self._prefill_wrapper = BatchDCPPrefillWrapper(
1147 + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
1148 workspace_buffer=self._get_workspace_buffer(),
1149 dcp_a2a=self.dcp_a2a,
1150 )
@@ -1185,7 +1153,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1153 assert not self.is_kvcache_nvfp4
1154 self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper(
1155 self._get_workspace_buffer(),
1188 - get_kv_cache_layout(),
1156 + get_flashinfer_layout_string(self.kv_cache_layout),
1157 backend="auto",
1158 q_data_type=self.q_data_type_prefill,
1159 kv_data_type=self.kv_cache_dtype,
@@ -1199,7 +1167,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1167 backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
1168 self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
1169 self._get_workspace_buffer(),
1202 - get_kv_cache_layout(),
1170 + get_flashinfer_layout_string(self.kv_cache_layout),
1171 backend=backend,
1172 )
1173 assert self._prefill_wrapper is not None
@@ -1225,7 +1193,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1193 backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
1194 decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
1195 self._get_workspace_buffer(),
1228 - get_kv_cache_layout(),
1196 + get_flashinfer_layout_string(self.kv_cache_layout),
1197 use_cuda_graph=use_cudagraph,
1198 paged_kv_indptr_buffer=paged_kv_indptr,
1199 paged_kv_indices_buffer=paged_kv_indices,
@@ -1248,7 +1216,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1216 def _get_cascade_wrapper(self):
1217 if self._cascade_wrapper is None:
1218 self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
1251 - 2, self._get_workspace_buffer(), get_kv_cache_layout()
1219 + 2,
1220 + self._get_workspace_buffer(),
1221 + get_flashinfer_layout_string(self.kv_cache_layout),
1222 )
1223 return self._cascade_wrapper
1224
@@ -1813,6 +1783,8 @@ class FlashInferImpl(AttentionImpl):
1783 num_heads, num_kv_heads, is_prefill=False
1784 )
1785 vllm_config = get_current_vllm_config_or_none()
1786 + # The layout is resolved after model construction, so read it lazily.
1787 + self.cache_config = vllm_config.cache_config if vllm_config else None
1788 # Query pre-quantization needs a single dtype for the whole query tensor.
1789 # SM90 XQA needs BF16/FP16-Q for decode and FP8 for prefill,
1790 # so only enable this for SM100 trtllm-gen where both use FP8-Q.
@@ -1848,6 +1820,11 @@ class FlashInferImpl(AttentionImpl):
1820 else:
1821 self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
1822
1823 + @property
1824 + def kv_cache_layout(self) -> KVCacheLayout:
1825 + assert self.cache_config is not None
1826 + return self.cache_config.get_resolved_kv_cache_layout()
1827 +
1828 def fused_output_quant_supported(self, quant_key: QuantKey):
1829 # XQA does not support FP8/NVFP4 output, so require trtllm-gen
1830 # (SM100+) here. Without that we cannot fuse the output quant.
@@ -1919,9 +1896,7 @@ class FlashInferImpl(AttentionImpl):
1896 query: shape = [num_tokens, num_heads, head_size]
1897 key: shape = [num_tokens, num_kv_heads, head_size]
1898 value: shape = [num_tokens, num_kv_heads, head_size]
1922 - kv_cache: KV cache tensor with different possible shapes:
1923 - - NHD: [num_blocks, 2, block_size, num_kv_heads, head_size]
1924 - - HND: [num_blocks, 2, num_kv_heads, block_size, head_size]
1899 + kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size]
1900 attn_metadata: Metadata for attention.
1901 Returns:
1902 shape = [num_tokens, num_heads * head_size]
@@ -2019,7 +1994,7 @@ class FlashInferImpl(AttentionImpl):
1994 if attn_metadata.use_cascade:
1995 # Cascade attention (rare case).
1996 assert attn_metadata.cascade_wrapper is not None
2022 - stride_order = FlashInferBackend.get_kv_cache_stride_order()
1997 + stride_order = self.kv_cache_layout.layer_view_order
1998 if self.is_kvcache_nvfp4:
1999 kv_cache_views = tuple(
2000 cache.permute(*stride_order)
@@ -2039,7 +2014,7 @@ class FlashInferImpl(AttentionImpl):
2014 num_decode_tokens = attn_metadata.num_decode_tokens
2015 num_prefill_tokens = attn_metadata.num_prefill_tokens
2016
2042 - stride_order = FlashInferBackend.get_kv_cache_stride_order()
2017 + stride_order = self.kv_cache_layout.layer_view_order
2018 kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous
2019 # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1
2020 # with TP=8). PyTorch permits non-canonical strides on size-1 dims;
@@ -2195,7 +2170,7 @@ class FlashInferImpl(AttentionImpl):
2170 seq_lens_prefill = attn_metadata.prefill.seq_lens
2171
2172 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
2198 - assert get_kv_cache_layout() == "HND"
2173 + assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND"
2174 assert is_strictly_contiguous(prefill_query)
2175 assert is_strictly_contiguous(workspace_buffer)
2176 assert is_strictly_contiguous(block_tables_prefill)
@@ -2380,7 +2355,7 @@ class FlashInferImpl(AttentionImpl):
2355 # trtllm-gen needs HND layout on SM100. XQA is selected
2356 # separately on SM90 and does not use this SM100 layout gate.
2357 if decode_with_trtllm_gen:
2383 - assert get_kv_cache_layout() == "HND"
2358 + assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND"
2359 else:
2360 assert decode_with_xqa
2361 assert is_strictly_contiguous(decode_query)
@@ -2426,7 +2401,7 @@ class FlashInferImpl(AttentionImpl):
2401 window_left=self.window_left,
2402 out=output[:num_decode_tokens],
2403 sinks=self.sinks,
2429 - kv_layout=get_kv_cache_layout(),
2404 + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
2405 q_len_per_req=q_len_per_req,
2406 mask=attn_metadata.decode.mask,
2407 q_cu_seq_lens=attn_metadata.decode.q_cu_seq_lens,
@@ -2499,7 +2474,7 @@ class FlashInferImpl(AttentionImpl):
2474 sinks=self.sinks,
2475 o_sf_scale=self.o_sf_scale,
2476 out=out,
2502 - kv_layout=get_kv_cache_layout(),
2477 + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
2478 backend=attn_metadata.decode.kernel.value,
2479 q_len_per_req=q_len_per_req,
2480 kv_cache_sf=(
vllm/v1/attention/backends/flex_attention.py
+11 -20
@@ -43,7 +43,11 @@ from vllm.v1.attention.backend import (
43 CommonAttentionMetadata,
44 MultipleOf,
45 )
46 -from vllm.v1.kv_cache_interface import AttentionSpec, EncoderOnlyAttentionSpec
46 +from vllm.v1.kv_cache_interface import (
47 + AttentionSpec,
48 + EncoderOnlyAttentionSpec,
49 + KVCacheLayout,
50 +)
51
52 logger = init_logger(__name__)
53
@@ -118,6 +122,12 @@ class FlexAttentionBackend(AttentionBackend):
122 def supports_batch_invariance(cls) -> bool:
123 return True
124
125 + @classmethod
126 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
127 + # Flex flattens the (B, N) block/token axes into a single token dim, which
128 + # only LBNHC's strides allow as a zero-copy view of the layer cache.
129 + return (KVCacheLayout.LBNHC,)
130 +
131 @classmethod
132 def supports_mm_prefix(cls) -> bool:
133 """FlexAttention supports full attention for image tokens."""
@@ -127,25 +137,6 @@ class FlexAttentionBackend(AttentionBackend):
137 def get_impl_cls() -> type["FlexAttentionImpl"]:
138 return FlexAttentionImpl
139
130 - @staticmethod
131 - def get_kv_cache_shape(
132 - num_blocks: int,
133 - block_size: int,
134 - num_kv_heads: int,
135 - head_size: int,
136 - cache_dtype_str: str = "auto",
137 - ) -> tuple[int, ...]:
138 - # K and V are packed into the content dim: logical (B, H, N, 2*hs).
139 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
140 -
141 - @staticmethod
142 - def get_kv_cache_stride_order(
143 - include_num_layers_dimension: bool = False,
144 - ) -> tuple[int, ...]:
145 - if include_num_layers_dimension:
146 - return (1, 0, 3, 2, 4)
147 - return (0, 2, 1, 3)
148 -
140 @staticmethod
141 def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]:
142 return FlexAttentionMetadataBuilder
vllm/v1/attention/backends/hpc_attn.py
+19 -34
@@ -30,10 +30,9 @@ from vllm.v1.attention.backend import (
30 MultipleOf,
31 )
32 from vllm.v1.attention.backends.utils import (
33 - KVCacheLayoutType,
33 split_decodes_and_prefills,
34 )
36 -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
35 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout, KVCacheSpec
36
37 logger = init_logger(__name__)
38
@@ -265,7 +264,7 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
264 class HpcAttentionBackend(AttentionBackend):
265 """HPC attention backend (pure attention, no RoPE/Norm).
266
268 - KV cache layout: NHD (num_blocks, 2, block_size, num_kv_heads, head_size).
267 + KV cache layout: LBNHC (num_blocks, num_kv_heads, block_size, 2 * head_size).
268 """
269
270 accept_output_buffer: bool = True
@@ -298,24 +297,6 @@ class HpcAttentionBackend(AttentionBackend):
297 def get_builder_cls() -> type["HpcAttnMetadataBuilder"]:
298 return HpcAttnMetadataBuilder
299
301 - @staticmethod
302 - def get_kv_cache_shape(
303 - num_blocks: int,
304 - block_size: int,
305 - num_kv_heads: int,
306 - head_size: int,
307 - cache_dtype_str: str = "auto",
308 - ) -> tuple[int, ...]:
309 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
310 -
311 - @staticmethod
312 - def get_kv_cache_stride_order(
313 - include_num_layers_dimension: bool = False,
314 - ) -> tuple[int, ...]:
315 - if include_num_layers_dimension:
316 - return (1, 0, 2, 3, 4, 5)
317 - return (0, 1, 2, 3, 4)
318 -
300 @classmethod
301 def get_supported_head_sizes(cls) -> list[int]:
302 return [128]
@@ -331,8 +312,8 @@ class HpcAttentionBackend(AttentionBackend):
312 return kv_cache_dtype in cls.supported_kv_cache_dtypes
313
314 @classmethod
334 - def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
335 - return "NHD"
315 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
316 + return (KVCacheLayout.LBNHC,)
317
318
319 class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
@@ -452,6 +433,9 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
433 num_decode_reqs = attn_metadata.num_decodes
434 num_decode_tokens = attn_metadata.num_decode_tokens
435
436 + # Logical (B, H, N, 2*head_size) -> (B, N, H, head_size) K/V views.
437 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
438 +
439 # Standard-path KV write: when no upstream fused op wrote KV cache
440 # (i.e. the model didn't wire HpcRopeNorm), fall back to the generic
441 # reshape_and_cache_flash. This is what lets non-HunYuan-V3 models use
@@ -466,8 +450,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
450 torch.ops._C_cache_ops.reshape_and_cache_flash(
451 key,
452 value,
469 - kv_cache[:, 0],
470 - kv_cache[:, 1],
453 + key_cache,
454 + value_cache,
455 attn_metadata.slot_mapping,
456 self.kv_cache_dtype,
457 layer._k_scale,
@@ -485,7 +469,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
469 f" (layer={getattr(layer, 'layer_name', '?')})"
470 )
471 torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype)
488 - kv_cache = kv_cache.view(torch_dtype)
472 + key_cache = key_cache.view(torch_dtype)
473 + value_cache = value_cache.view(torch_dtype)
474 k_scale = layer._k_scale.reshape(1)
475 v_scale = layer._v_scale.reshape(1)
476
@@ -508,8 +493,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
493 if self.use_fp8:
494 hpc.attention_with_kvcache_prefill_fp8(
495 q_prefill,
511 - kv_cache[:, 0],
512 - kv_cache[:, 1],
496 + key_cache,
497 + value_cache,
498 hpc_prefill_q_scale,
499 k_scale,
500 v_scale,
@@ -523,8 +508,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
508 else:
509 hpc.attention_with_kvcache_prefill_bf16(
510 q_prefill,
526 - kv_cache[:, 0],
527 - kv_cache[:, 1],
511 + key_cache,
512 + value_cache,
513 cu_seqlens_prefill,
514 block_table_prefill,
515 seq_lens_prefill,
@@ -545,8 +530,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
530 if self.use_fp8:
531 hpc.attention_decode_fp8(
532 q_decode,
548 - kv_cache[:, 0],
549 - kv_cache[:, 1],
533 + key_cache,
534 + value_cache,
535 block_table_decode,
536 num_seq_kvcache,
537 hpc_decode_q_scale,
@@ -575,8 +560,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
560 )
561 hpc.attention_decode_bf16(
562 q_decode,
578 - kv_cache[:, 0],
579 - kv_cache[:, 1],
563 + key_cache,
564 + value_cache,
565 block_table_decode,
566 num_seq_kvcache,
567 mtp=mtp,
vllm/v1/attention/backends/mla/cutlass_mla.py
-8
@@ -49,14 +49,6 @@ class CutlassMLABackend(MLACommonBackend):
49 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
50 return [128]
51
52 - @staticmethod
53 - def get_kv_cache_stride_order(
54 - include_num_layers_dimension: bool = False,
55 - ) -> tuple[int, ...]:
56 - if include_num_layers_dimension:
57 - return (1, 0, 2, 3)
58 - return (0, 1, 2)
59 -
52 @staticmethod
53 def get_name() -> str:
54 return "CUTLASS_MLA"
vllm/v1/attention/backends/mla/flashattn_mla.py
-8
@@ -52,14 +52,6 @@ class FlashAttnMLABackend(MLACommonBackend):
52 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
53 return [MultipleOf(16)]
54
55 - @staticmethod
56 - def get_kv_cache_stride_order(
57 - include_num_layers_dimension: bool = False,
58 - ) -> tuple[int, ...]:
59 - if include_num_layers_dimension:
60 - return (1, 0, 2, 3)
61 - return (0, 1, 2)
62 -
55 @staticmethod
56 def get_name() -> str:
57 return "FLASH_ATTN_MLA"
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+7 -17
@@ -24,6 +24,7 @@ from vllm.v1.attention.backend import (
24 )
25 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
26 from vllm.v1.attention.backends.mla.sparse_utils import (
27 + flat_kv_row_view,
28 triton_convert_req_index_to_global_index,
29 )
30 from vllm.v1.kv_cache_interface import AttentionSpec
@@ -103,16 +104,6 @@ class FlashAttnMLASparseBackend(AttentionBackend):
104 return "FlashAttention MLA Sparse requires model with index_topk"
105 return None
106
106 - @staticmethod
107 - def get_kv_cache_shape(
108 - num_blocks: int,
109 - block_size: int,
110 - num_kv_heads: int,
111 - head_size: int,
112 - cache_dtype_str: str = "auto",
113 - ) -> tuple[int, ...]:
114 - return (num_blocks, block_size, head_size)
115 -
107
108 @dataclass
109 class FlashAttnMLASparseMetadata(AttentionMetadata):
@@ -223,11 +214,15 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
214
215 assert self.topk_indices_buffer is not None
216 topk_indices = self.topk_indices_buffer[:num_actual_toks]
217 + kv_rows, block_stride_rows = flat_kv_row_view(
218 + kv_c_and_k_pe_cache, attn_metadata.block_size
219 + )
220 topk_indices, valid_counts = triton_convert_req_index_to_global_index(
221 attn_metadata.req_id_per_token[:num_actual_toks],
222 attn_metadata.block_table,
223 topk_indices,
224 BLOCK_SIZE=attn_metadata.block_size,
225 + BLOCK_STRIDE_ROWS=block_stride_rows,
226 NUM_TOPK_TOKENS=topk_indices.shape[1],
227 return_valid_counts=True,
228 )
@@ -235,13 +230,8 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
230 cu_seqlens_q = torch.arange(
231 0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device
232 )
238 - kv_cache = kv_c_and_k_pe_cache.view(
239 - -1, attn_metadata.block_size, self.head_size
240 - )
241 - k_cache = kv_cache[:, :, self.kv_lora_rank :].view(
242 - -1, 1, 1, self.qk_rope_head_dim
243 - )
244 - v_cache = kv_cache[:, :, : self.kv_lora_rank].view(-1, 1, 1, self.kv_lora_rank)
233 + k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
234 + v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
235
236 out = flash_attn_varlen_func(
237 q=q_rope,
vllm/v1/attention/backends/mla/flashinfer_mla.py
-13
@@ -28,7 +28,6 @@ from vllm.v1.attention.backend import (
28 AttentionType,
29 MultipleOf,
30 )
31 -from vllm.v1.attention.backends.utils import KVCacheLayoutType
31
32 logger = init_logger(__name__)
33
@@ -133,14 +132,6 @@ class FlashInferMLABackend(MLACommonBackend):
132 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
133 return [32, 64]
134
136 - @staticmethod
137 - def get_kv_cache_stride_order(
138 - include_num_layers_dimension: bool = False,
139 - ) -> tuple[int, ...]:
140 - if include_num_layers_dimension:
141 - return (1, 0, 2, 3)
142 - return (0, 1, 2)
143 -
135 @staticmethod
136 def get_name() -> str:
137 return "FLASHINFER_MLA"
@@ -194,10 +185,6 @@ class FlashInferMLABackend(MLACommonBackend):
185 )
186 return None
187
197 - @classmethod
198 - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
199 - return "HND"
200 -
188
189 class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
190 can_return_lse_for_decode: bool = True
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
-32
@@ -31,7 +31,6 @@ from vllm.v1.attention.backends.mla.sparse_utils import (
31 triton_convert_req_index_to_global_index,
32 triton_filter_and_convert_dcp_index,
33 )
34 -from vllm.v1.attention.backends.utils import KVCacheLayoutType
34 from vllm.v1.kv_cache_interface import AttentionSpec
35
36 if TYPE_CHECKING:
@@ -123,20 +122,6 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
122 return "FlashInfer MLA Sparse requires model with index_topk config"
123 return None
124
126 - @staticmethod
127 - def get_kv_cache_shape(
128 - num_blocks: int,
129 - block_size: int,
130 - num_kv_heads: int, # assumed to be 1 for MLA
131 - head_size: int,
132 - cache_dtype_str: str = "auto",
133 - ) -> tuple[int, ...]:
134 - return (num_blocks, block_size, head_size)
135 -
136 - @classmethod
137 - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
138 - return "HND"
139 -
125
126 class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
127 """FlashInfer sparse MLA backend for SM120."""
@@ -216,23 +201,6 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
201 )
202 return None
203
219 - @staticmethod
220 - def get_kv_cache_shape(
221 - num_blocks: int,
222 - block_size: int,
223 - num_kv_heads: int, # assumed to be 1 for MLA
224 - head_size: int,
225 - cache_dtype_str: str = "auto",
226 - ) -> tuple[int, ...]:
227 - if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"):
228 - # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE.
229 - return (num_blocks, block_size, 656)
230 - return (num_blocks, block_size, head_size)
231 -
232 - @classmethod
233 - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
234 - return None
235 -
204
205 @dataclass
206 class FlashInferMLASparseMetadata(AttentionMetadata):
vllm/v1/attention/backends/mla/flashmla.py
-8
@@ -58,14 +58,6 @@ class FlashMLABackend(MLACommonBackend):
58 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
59 return [64]
60
61 - @staticmethod
62 - def get_kv_cache_stride_order(
63 - include_num_layers_dimension: bool = False,
64 - ) -> tuple[int, ...]:
65 - if include_num_layers_dimension:
66 - return (1, 0, 2, 3)
67 - return (0, 1, 2)
68 -
61 @staticmethod
62 def get_name() -> str:
63 return "FLASHMLA"
vllm/v1/attention/backends/mla/flashmla_sparse.py
+6 -15
@@ -27,6 +27,7 @@ from vllm.v1.attention.backend import (
27 MultipleOf,
28 )
29 from vllm.v1.attention.backends.mla.sparse_utils import (
30 + flat_kv_row_view,
31 triton_convert_req_index_to_global_index,
32 triton_filter_and_convert_dcp_index,
33 )
@@ -128,20 +129,6 @@ class FlashMLASparseBackend(AttentionBackend):
129 def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
130 return capability.major in [9, 10]
131
131 - @staticmethod
132 - def get_kv_cache_shape(
133 - num_blocks: int,
134 - block_size: int,
135 - num_kv_heads: int, # assumed to be 1 for MLA
136 - head_size: int,
137 - cache_dtype_str: str = "auto",
138 - ) -> tuple[int, ...]:
139 - if cache_dtype_str == "fp8_ds_mla":
140 - # V3.2 main MLA: 656-byte custom storage format. See module docstring.
141 - return (num_blocks, block_size, 656)
142 - else:
143 - return (num_blocks, block_size, head_size)
144 -
132
133 @dataclass
134 class FlashMLASparseMetadata(AttentionMetadata):
@@ -643,18 +630,22 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
630 # Convert per-request indices to global slots (decode) or workspace
631 # offsets (prefill). req_id_per_token covers the whole batch; slice it
632 # to the MQA tokens (q may exclude prefill tokens routed to dense MHA).
633 + kv_rows, block_stride_rows = flat_kv_row_view(
634 + kv_c_and_k_pe_cache, attn_metadata.block_size
635 + )
636 topk_indices, topk_length = triton_convert_req_index_to_global_index(
637 attn_metadata.req_id_per_token[: topk_indices.shape[0]],
638 attn_metadata.block_table,
639 topk_indices,
640 BLOCK_SIZE=attn_metadata.block_size,
641 + BLOCK_STRIDE_ROWS=block_stride_rows,
642 NUM_TOPK_TOKENS=topk_indices.shape[1],
643 return_valid_counts=True,
644 )
645
646 return self._bf16_flash_mla_kernel(
647 q,
657 - kv_c_and_k_pe_cache,
648 + kv_rows,
649 topk_indices,
650 topk_length,
651 )
vllm/v1/attention/backends/mla/indexer.py
+13 -26
@@ -36,7 +36,7 @@ from vllm.v1.attention.backends.utils import (
36 get_dcp_local_seq_lens,
37 split_decodes_and_prefills,
38 )
39 -from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
39 +from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MLAAttentionSpec
40
41 logger = init_logger(__name__)
42
@@ -179,34 +179,18 @@ class DeepseekV32IndexerBackend(AttentionBackend):
179 def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]:
180 return DeepseekV32IndexerMetadataBuilder
181
182 - @staticmethod
183 - def get_kv_cache_shape(
184 - num_blocks: int,
185 - block_size: int,
186 - num_kv_heads: int,
187 - head_size: int,
188 - cache_dtype_str: str = "auto",
189 - ) -> tuple[int, ...]:
190 - assert num_kv_heads == 1
191 - return (num_blocks, block_size, head_size)
192 -
193 - @staticmethod
194 - def get_kv_cache_stride_order(
195 - include_num_layers_dimension: bool = False,
196 - ) -> tuple[int, ...]:
197 - if include_num_layers_dimension:
198 - # DeepseekV32Indexer kernels do not support cross-layer
199 - # KV cache layout. Identity permutation keeps num_layers
200 - # first, signaling incompatibility.
201 - return (0, 1, 2, 3)
202 - return (0, 1, 2)
203 -
182
183 class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
184 @staticmethod
185 def get_name() -> str:
186 return "DEEPSEEK_V4_INDEXER"
187
188 + @classmethod
189 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
190 + # DeepSeek-V4 packs the indexer pages beside the MLA latent pages inside
191 + # each block, so the layer dim must sit inside the block dim.
192 + return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC)
193 +
194 @staticmethod
195 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
196 return [256]
@@ -638,7 +622,10 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
622 self.compress_ratio = 1
623 # Get compress_ratio for DeepseekV4 support
624 if isinstance(self.kv_cache_spec, MLAAttentionSpec):
641 - self.compress_ratio = self.kv_cache_spec.compress_ratio
625 + # MLA compression is a whole number of tokens per state (fractions
626 + # are whisper block pooling and never reach MLA).
627 + assert isinstance(self.kv_cache_spec.tokens_per_state, int)
628 + self.compress_ratio = self.kv_cache_spec.tokens_per_state
629 if self.dcp_world_size > 1 and self.compress_ratio > 1:
630 raise NotImplementedError(
631 "DCP is not supported with sparse indexer KV compression "
@@ -885,7 +872,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
872 query_start_loc,
873 seq_lens,
874 block_table,
888 - self.kv_cache_spec.storage_block_size,
875 + self.kv_cache_spec.num_states,
876 self.compress_ratio,
877 out=self.compressed_slot_mapping_buffer,
878 )
@@ -1053,7 +1040,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1040 if current_platform.is_cuda() and has_deep_gemm():
1041 metadata = get_paged_mqa_logits_metadata(
1042 seq_lens,
1056 - self.kv_cache_spec.storage_block_size,
1043 + self.kv_cache_spec.num_states,
1044 self.num_sms,
1045 indices=decode_indices,
1046 )
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
-10
@@ -293,16 +293,6 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
293 def get_impl_cls() -> type["ROCMAiterMLASparseImpl"]:
294 return ROCMAiterMLASparseImpl
295
296 - @staticmethod
297 - def get_kv_cache_shape(
298 - num_blocks: int,
299 - block_size: int,
300 - num_kv_heads: int, # assumed to be 1 for MLA
301 - head_size: int,
302 - cache_dtype_str: str = "auto",
303 - ) -> tuple[int, ...]:
304 - return (num_blocks, block_size, head_size)
305 -
296 @classmethod
297 def is_mla(cls) -> bool:
298 return True
vllm/v1/attention/backends/mla/sparse_swa.py
+6 -25
@@ -89,6 +89,10 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
89 # contiguous full-cache layout.
90 assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)
91
92 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
93 + # [B, H=1, N, C] -> [B, N, C]
94 + self.kv_cache = kv_cache.squeeze(1)
95 +
96 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
97 # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
98 # bf16/fp8 cache uses the natural element-size page.
@@ -141,30 +145,6 @@ class DeepseekSparseSWABackend(AttentionBackend):
145 return DeepseekV4ROCMAiterSparseSWAMetadataBuilder
146 return DeepseekSparseSWAMetadataBuilder
147
144 - @staticmethod
145 - def get_kv_cache_shape(
146 - num_blocks: int,
147 - block_size: int,
148 - num_kv_heads: int,
149 - head_size: int,
150 - cache_dtype_str: str = "auto",
151 - ) -> tuple[int, ...]:
152 - assert num_kv_heads == 1
153 - if cache_dtype_str == "fp8_ds_mla":
154 - # DeepseekV4 SWA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
155 - # head_size passed in is the semantic head_dim (512).
156 - return (num_blocks, block_size, 584)
157 - else:
158 - return (num_blocks, block_size, head_size)
159 -
160 - @staticmethod
161 - def get_kv_cache_stride_order(
162 - include_num_layers_dimension: bool = False,
163 - ) -> tuple[int, ...]:
164 - if include_num_layers_dimension:
165 - return (0, 1, 2, 3)
166 - return (0, 1, 2)
167 -
148
149 @dataclass
150 class DeepseekSparseSWAMetadata:
@@ -417,7 +397,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
397 assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec)
398 mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec)
399 self.head_size = mla_spec.head_size # Already considered quantization.
420 - self.compress_ratio = mla_spec.compress_ratio
400 + assert isinstance(mla_spec.tokens_per_state, int)
401 + self.compress_ratio = mla_spec.tokens_per_state
402 self.block_size = mla_spec.block_size
403 self.max_model_len = self.vllm_config.model_config.max_model_len
404 self.max_num_batched_tokens = (
vllm/v1/attention/backends/mla/sparse_utils.py
+27 -1
@@ -7,6 +7,28 @@ import torch
7 from vllm.triton_utils import tl, triton
8
9
10 +def flat_kv_row_view(
11 + kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim]
12 + block_size: int,
13 +) -> tuple[torch.Tensor, int]:
14 + """Flat [row, head_dim] view of a paged cache and its physical rows per block.
15 +
16 + Token offset is ``block_idx * block_stride_rows + offset_in_block``.
17 + When other layers' pages sit between consecutive blocks of this cache,
18 + ``block_stride_rows`` exceeds ``block_size``; those in-between rows are never
19 + indexed (`triton_convert_req_index_to_global_index` ensures this).
20 + """
21 + num_blocks, _, head_dim = kv_cache.shape
22 + assert kv_cache.stride(0) % head_dim == 0, (
23 + "block stride is not a whole number of rows; flat row indexing would "
24 + "silently misaddress"
25 + )
26 + block_stride_rows = kv_cache.stride(0) // head_dim
27 + num_rows = (num_blocks - 1) * block_stride_rows + block_size
28 + rows = kv_cache.as_strided((num_rows, head_dim), (head_dim, 1))
29 + return rows, block_stride_rows
30 +
31 +
32 # Kernel with prefill workspace support and valid count tracking
33 @triton.jit
34 def _convert_req_index_to_global_index_kernel(
@@ -20,6 +42,7 @@ def _convert_req_index_to_global_index_kernel(
42 # shapes (compile-time where possible)
43 max_num_blocks_per_req: tl.constexpr,
44 BLOCK_SIZE: tl.constexpr,
45 + BLOCK_STRIDE_ROWS: tl.constexpr,
46 BLOCK_N: tl.constexpr, # tile width along columns
47 HAS_PREFILL: tl.constexpr,
48 COUNT_VALID: tl.constexpr, # whether to count valid indices
@@ -86,7 +109,7 @@ def _convert_req_index_to_global_index_kernel(
109 bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
110 is_invalid_tok |= ~valid_block | is_remote
111 base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
89 - out_val = base * BLOCK_SIZE + inblock_off
112 + out_val = base * BLOCK_STRIDE_ROWS + inblock_off
113
114 # Override with prefill output if prefill is enabled
115 if HAS_PREFILL:
@@ -156,6 +179,7 @@ def triton_convert_req_index_to_global_index(
179 block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req]
180 token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS]
181 BLOCK_SIZE: int = 64,
182 + BLOCK_STRIDE_ROWS: int | None = None,
183 NUM_TOPK_TOKENS: int = 2048,
184 BLOCK_N: int = 128, # tile width along columns
185 HAS_PREFILL_WORKSPACE: bool = False,
@@ -249,6 +273,7 @@ def triton_convert_req_index_to_global_index(
273 # shapes / constexprs
274 max_num_blocks_per_req,
275 BLOCK_SIZE,
276 + BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,
277 block_n,
278 HAS_PREFILL_WORKSPACE,
279 return_valid_counts,
@@ -365,6 +390,7 @@ def triton_filter_and_convert_dcp_index(
390 None,
391 max_num_blocks_per_req,
392 BLOCK_SIZE,
393 + BLOCK_SIZE, # dense caches on the DCP path
394 block_n,
395 False, # HAS_PREFILL
396 count_valid,
vllm/v1/attention/backends/mla/tokenspeed_mla.py
-5
@@ -23,7 +23,6 @@ from vllm.v1.attention.backend import (
23 AttentionType,
24 MultipleOf,
25 )
26 -from vllm.v1.attention.backends.utils import KVCacheLayoutType
26
27 if TYPE_CHECKING:
28 from vllm.config import VllmConfig
@@ -152,10 +151,6 @@ class TokenspeedMLABackend(MLACommonBackend):
151 )
152 return None
153
155 - @classmethod
156 - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
157 - return "HND"
158 -
154
155 class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
156 can_return_lse_for_decode: bool = True
vllm/v1/attention/backends/mla/triton_mla.py
-8
@@ -121,14 +121,6 @@ class TritonMLABackend(MLACommonBackend):
121 return True
122 return block_size % 16 == 0
123
124 - @staticmethod
125 - def get_kv_cache_stride_order(
126 - include_num_layers_dimension: bool = False,
127 - ) -> tuple[int, ...]:
128 - if include_num_layers_dimension:
129 - return (1, 0, 2, 3)
130 - return (0, 1, 2)
131 -
124 @staticmethod
125 def get_name() -> str:
126 return "TRITON_MLA"
vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+7 -14
@@ -23,7 +23,8 @@ from vllm.v1.attention.backend import (
23 CommonAttentionMetadata,
24 MLAAttentionImpl,
25 )
26 -from vllm.v1.attention.backends.mla.flashmla_sparse import (
26 +from vllm.v1.attention.backends.mla.sparse_utils import (
27 + flat_kv_row_view,
28 triton_convert_req_index_to_global_index,
29 )
30 from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
@@ -66,16 +67,6 @@ class XPUMLASparseBackend(AttentionBackend):
67 def is_sparse(cls) -> bool:
68 return True
69
69 - @staticmethod
70 - def get_kv_cache_shape(
71 - num_blocks: int,
72 - block_size: int,
73 - num_kv_heads: int, # assumed to be 1 for MLA
74 - head_size: int,
75 - cache_dtype_str: str = "auto",
76 - ) -> tuple[int, ...]:
77 - return (num_blocks, block_size, head_size)
78 -
70 @classmethod
71 def get_supported_head_sizes(cls) -> list[int]:
72 return [576]
@@ -249,16 +240,18 @@ class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]):
240 assert self.topk_indices_buffer is not None
241 topk_indices = self.topk_indices_buffer[:num_actual_toks]
242
243 + kv_rows, block_stride_rows = flat_kv_row_view(
244 + kv_c_and_k_pe_cache, attn_metadata.block_size
245 + )
246 topk_indices_global = triton_convert_req_index_to_global_index(
247 attn_metadata.req_id_per_token,
248 attn_metadata.block_table,
249 topk_indices,
250 BLOCK_SIZE=attn_metadata.block_size,
251 + BLOCK_STRIDE_ROWS=block_stride_rows,
252 NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
253 )
254
260 - attn_out = self._forward_bf16_kv(
261 - q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata
262 - )
255 + attn_out = self._forward_bf16_kv(q, kv_rows, topk_indices_global, attn_metadata)
256
257 return attn_out, None
vllm/v1/attention/backends/rocm_aiter_fa.py
+16 -12
@@ -31,7 +31,7 @@ from vllm.v1.attention.backends.utils import (
31 split_decodes_prefills_and_extends,
32 )
33 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
34 -from vllm.v1.kv_cache_interface import AttentionSpec
34 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
35
36 _PARTITION_SIZE_ROCM = 256
37 _CP_TOKENS_PER_ITER_ROCM = 32 * 1024
@@ -761,18 +761,22 @@ class AiterFlashAttentionBackend(AttentionBackend):
761 def get_builder_cls() -> type["AiterFlashAttentionMetadataBuilder"]:
762 return AiterFlashAttentionMetadataBuilder
763
764 - @staticmethod
765 - def get_kv_cache_shape(
766 - num_blocks: int,
767 - block_size: int,
768 - num_kv_heads: int,
769 - head_size: int,
770 - cache_dtype_str: str = "auto",
771 - ) -> tuple[int, ...]:
772 - if block_size % 16 != 0:
764 + @classmethod
765 + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
766 + """Validate the block size the ROCm gather kernels require."""
767 + # block_size == 1 is the per-token page-size probe (see
768 + # Platform.get_page_size_bytes); real blocks must be gatherable in
769 + # 16-token units by the ROCm kernel.
770 + if spec.block_size != 1 and spec.block_size % 16 != 0:
771 raise ValueError("Block size must be a multiple of 16.")
774 - # K and V are packed into the content dim: logical (B, H, N, 2*hs).
775 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
772 + return spec
773 +
774 + @classmethod
775 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
776 + # K and V come out of the content dim as transposed views rather than
777 + # copies, so the head dim may sit on either side of the block dim, but
778 + # the layer must stay outermost.
779 + return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
780
781 @classmethod
782 def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+17 -11
@@ -22,6 +22,7 @@ from vllm.v1.attention.backends.rocm_attn import (
22 RocmAttentionMetadata,
23 RocmAttentionMetadataBuilder,
24 )
25 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
26
27 logger = init_logger(__name__)
28
@@ -77,18 +78,23 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
78 def get_impl_cls() -> type["RocmAiterUnifiedAttentionImpl"]:
79 return RocmAiterUnifiedAttentionImpl
80
80 - @staticmethod
81 - def get_kv_cache_shape(
82 - num_blocks: int,
83 - block_size: int,
84 - num_kv_heads: int,
85 - head_size: int,
86 - cache_dtype_str: str = "auto",
87 - ) -> tuple[int, ...]:
88 - if block_size % 16 != 0:
81 + @classmethod
82 + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
83 + """Keep K and V packed in the content dim, unlike the native HIP
84 + kernels the base class targets."""
85 + # block_size == 1 is the per-token page-size probe (see
86 + # Platform.get_page_size_bytes); real blocks must be gatherable in
87 + # 16-token units by the ROCm kernel.
88 + if spec.block_size != 1 and spec.block_size % 16 != 0:
89 raise ValueError("Block size must be a multiple of 16.")
90 - # K and V are packed into the content dim: logical (B, H, N, 2*hs).
91 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
90 + return spec
91 +
92 + @classmethod
93 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
94 + # K and V come out of the content dim as transposed views rather than
95 + # copies, so the head dim may sit on either side of the block dim, but
96 + # the layer must stay outermost.
97 + return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
98
99 @staticmethod
100 def use_cascade_attention(*args, **kwargs) -> bool:
vllm/v1/attention/backends/rocm_attn.py
+34 -19
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Attention layer with PagedAttention and Triton prefix prefill."""
4
5 -from dataclasses import dataclass
5 +from dataclasses import dataclass, replace
6 from typing import ClassVar
7
8 import torch
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 kFp8StaticTensorSym,
17 )
18 from vllm.platforms import current_platform
19 -from vllm.utils.torch_utils import is_quantized_kv_cache
19 +from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache
20 from vllm.v1.attention.backend import (
21 AttentionBackend,
22 AttentionCGSupport,
@@ -35,7 +35,7 @@ from vllm.v1.attention.ops.paged_attn import PagedAttention
35 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
36 triton_reshape_and_cache_flash,
37 )
38 -from vllm.v1.kv_cache_interface import AttentionSpec
38 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
39
40 logger = init_logger(__name__)
41
@@ -241,17 +241,30 @@ class RocmAttentionBackend(AttentionBackend):
241 AttentionType.ENCODER_ONLY,
242 )
243
244 - @staticmethod
245 - def get_kv_cache_shape(
246 - num_blocks: int,
247 - block_size: int,
248 - num_kv_heads: int,
249 - head_size: int,
250 - cache_dtype_str: str = "auto",
251 - ) -> tuple[int, ...]:
252 - if block_size % 16 != 0:
253 - raise ValueError("Block size must be a multiple of 16.")
254 - return (2, num_blocks, block_size, num_kv_heads, head_size)
244 + @classmethod
245 + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
246 + """K and V as two head groups so the native HIP kernels address each side
247 + as one contiguous region (x-packed interior applied in split_kv_cache)."""
248 + if spec.state_content_bytes is not None:
249 + return spec
250 + assert spec.head_size == spec.head_size_v, (
251 + "Separate K/V head groups require symmetric K/V head sizes."
252 + )
253 + return replace(
254 + spec,
255 + num_head_slots=2,
256 + state_content_bytes=spec.num_kv_heads
257 + * spec.head_size
258 + * get_dtype_size(spec.dtype),
259 + )
260 +
261 + @classmethod
262 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
263 + # The native HIP kernels hardcode group-contiguous block addressing, so
264 + # they need the K and V groups to span all blocks (H outermost within the layer)
265 + # The Triton fallbacks are stride-aware, so block-interior LBHNC is also
266 + # supported, just without the native kernels.
267 + return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC)
268
269 @staticmethod
270 def use_cascade_attention(*args, **kwargs) -> bool:
@@ -380,8 +393,8 @@ class RocmAttentionImpl(AttentionImpl):
393 query: shape = [num_tokens, num_heads, head_size]
394 key: shape = [num_tokens, num_kv_heads, head_size]
395 value: shape = [num_tokens, num_kv_heads, head_size]
383 - kv_cache: shape =
384 - [2, num_blocks, block_size, num_kv_heads, head_size]
396 + kv_cache: logical [num_blocks, 2, block_size, num_kv_heads *
397 + head_size] under LHBNC (physically K/V-group-first)
398 attn_metadata: Metadata for attention.
399 Returns:
400 shape = [num_tokens, num_heads * head_size]
@@ -419,8 +432,10 @@ class RocmAttentionImpl(AttentionImpl):
432 layer,
433 )
434
435 + # The bound view is logical [B, 2, N, H*hs]; split_kv_cache expects
436 + # the K/V groups first.
437 key_cache, value_cache = PagedAttention.split_kv_cache(
423 - kv_cache, self.num_kv_heads, self.head_size
438 + kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size
439 )
440
441 if is_quantized_kv_cache(self.kv_cache_dtype):
@@ -477,7 +492,7 @@ class RocmAttentionImpl(AttentionImpl):
492 if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
493 return
494 key_cache, value_cache = PagedAttention.split_kv_cache(
480 - kv_cache, self.num_kv_heads, self.head_size
495 + kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size
496 )
497
498 # Reshape the input keys and values and store them in the cache.
@@ -532,7 +547,7 @@ class RocmAttentionImpl(AttentionImpl):
547 if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
548 return
549 key_cache, value_cache = PagedAttention.split_kv_cache(
535 - kv_cache,
550 + kv_cache.transpose(0, 1),
551 layer.num_kv_heads, # type: ignore[attr-defined]
552 layer.head_size, # type: ignore[attr-defined]
553 )
vllm/v1/attention/backends/triton_attn.py
-58
@@ -32,7 +32,6 @@ from vllm.v1.attention.backend import (
32 )
33 from vllm.v1.attention.backends.utils import (
34 compute_mm_prefix_range_tensor,
35 - get_kv_cache_layout,
35 get_num_attention_heads_from_layers,
36 )
37 from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
@@ -45,7 +44,6 @@ from vllm.v1.kv_cache_interface import (
44 AttentionSpec,
45 KVQuantMode,
46 get_kv_quant_mode,
48 - kv_cache_uses_per_token_head_scales,
47 )
48
49 logger = init_logger(__name__)
@@ -337,62 +335,6 @@ class TritonAttentionBackend(AttentionBackend):
335 def get_impl_cls() -> type["TritonAttentionImpl"]:
336 return TritonAttentionImpl
337
340 - @staticmethod
341 - def get_kv_cache_shape(
342 - num_blocks: int,
343 - block_size: int,
344 - num_kv_heads: int,
345 - head_size: int,
346 - cache_dtype_str: str = "auto",
347 - ) -> tuple[int, ...]:
348 - if block_size % 16 != 0:
349 - raise ValueError("Block size must be a multiple of 16.")
350 - # K and V are packed into the content dim: logical (B, H, N, 2*hs).
351 - if kv_cache_uses_per_token_head_scales(cache_dtype_str):
352 - # Pad the head dim by sizeof(float32)/sizeof(cache_dtype) so the
353 - # per-(token, head) scale fits inline after the quantized data;
354 - # the backend extracts data[:head_size] and scale[head_size:] via
355 - # typed views (see _ensure_scale_caches). INT4 packs two values
356 - # per byte, so the data occupies only head_size // 2 bytes.
357 - from vllm.utils.torch_utils import (
358 - STR_DTYPE_TO_TORCH_DTYPE,
359 - get_dtype_size,
360 - )
361 -
362 - cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype_str]
363 - scale_pad = get_dtype_size(torch.float32) // get_dtype_size(cache_dtype)
364 - if get_kv_quant_mode(cache_dtype_str) == KVQuantMode.INT4_PER_TOKEN_HEAD:
365 - data_head_size = head_size // 2
366 - else:
367 - data_head_size = head_size
368 - padded_hs = data_head_size + scale_pad
369 - return (num_blocks, num_kv_heads, block_size, 2 * padded_hs)
370 - return (num_blocks, num_kv_heads, block_size, 2 * head_size)
371 -
372 - @staticmethod
373 - def get_kv_cache_stride_order(
374 - include_num_layers_dimension: bool = False,
375 - ) -> tuple[int, ...]:
376 - # `stride_order` indicates the permutation that gets us from
377 - # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
378 - # layout we want.
379 - cache_layout = get_kv_cache_layout()
380 - if cache_layout == "NHD" and include_num_layers_dimension:
381 - # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
382 - return (1, 0, 3, 2, 4)
383 - elif cache_layout == "NHD":
384 - # (num_blocks, block_size, num_kv_heads, 2*head_size)
385 - stride_order = (0, 2, 1, 3)
386 - elif cache_layout == "HND" and include_num_layers_dimension:
387 - # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
388 - return (1, 2, 0, 3, 4)
389 - elif cache_layout == "HND":
390 - # (num_blocks, num_kv_heads, block_size, 2*head_size)
391 - stride_order = (0, 1, 2, 3)
392 - else:
393 - raise ValueError(f"Unknown cache layout: {cache_layout}")
394 - return stride_order
395 -
338 @staticmethod
339 def use_cascade_attention(*args, **kwargs) -> bool:
340 return False
vllm/v1/attention/backends/triton_attn_diffkv.py
-43
@@ -23,7 +23,6 @@ from vllm.v1.attention.backends.triton_attn import (
23 TritonAttentionMetadata,
24 TritonAttentionMetadataBuilder,
25 )
26 -from vllm.v1.attention.backends.utils import get_kv_cache_layout
26 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
27 triton_reshape_and_cache_flash_diffkv,
28 )
@@ -93,48 +92,6 @@ class TritonAttentionDiffKVBackend(TritonAttentionBackend):
92 def get_builder_cls() -> type["TritonAttentionDiffKVMetadataBuilder"]:
93 return TritonAttentionDiffKVMetadataBuilder
94
96 - @staticmethod
97 - def get_kv_cache_shape(
98 - num_blocks: int,
99 - block_size: int,
100 - num_kv_heads: int,
101 - head_size: int,
102 - cache_dtype_str: str = "auto",
103 - ) -> tuple[int, ...]:
104 - if block_size % 16 != 0:
105 - raise ValueError("Block size must be a multiple of 16.")
106 - # Logical (blocks-first, head-major) layout: K and V (with their
107 - # different head sizes) packed in the content dim.
108 - return (
109 - num_blocks,
110 - num_kv_heads,
111 - block_size,
112 - head_size + TritonAttentionDiffKVBackend.head_size_v,
113 - )
114 -
115 - @staticmethod
116 - def get_kv_cache_stride_order(
117 - include_num_layers_dimension: bool = False,
118 - ) -> tuple[int, ...]:
119 - # `stride_order` indicates the permutation that gets us from
120 - # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual
121 - # memory layout we want.
122 - cache_layout = get_kv_cache_layout()
123 - if cache_layout == "NHD" and include_num_layers_dimension:
124 - # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v)
125 - return (1, 0, 3, 2, 4)
126 - elif cache_layout == "NHD":
127 - # (num_blocks, block_size, num_kv_heads, C_k+C_v)
128 - return (0, 2, 1, 3)
129 - elif cache_layout == "HND" and include_num_layers_dimension:
130 - # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v)
131 - return (1, 2, 0, 3, 4)
132 - elif cache_layout == "HND":
133 - # (num_blocks, num_kv_heads, block_size, C_k+C_v)
134 - return (0, 1, 2, 3)
135 - else:
136 - raise ValueError(f"Unknown cache layout format {cache_layout}.")
137 -
95 @classmethod
96 def supports_head_size(cls, head_size: int) -> bool:
97 # DiffKV K head sizes (e.g. 192 for MiMo-V2.5) need to be allowed.
vllm/v1/attention/backends/turboquant_attn.py
+5 -34
@@ -64,7 +64,7 @@ from vllm.v1.attention.ops.triton_turboquant_decode import (
64 triton_turboquant_decode_attention,
65 )
66 from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store
67 -from vllm.v1.kv_cache_interface import AttentionSpec
67 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
68 from vllm.v1.worker.workspace import (
69 current_workspace_manager,
70 is_workspace_manager_initialized,
@@ -156,6 +156,10 @@ class TurboQuantAttentionBackend(AttentionBackend):
156 def get_name() -> str:
157 return "TURBOQUANT"
158
159 + @classmethod
160 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
161 + return (KVCacheLayout.LBNHC,)
162 +
163 @staticmethod
164 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
165 return [16, 32, 64, 128]
@@ -176,39 +180,6 @@ class TurboQuantAttentionBackend(AttentionBackend):
180 def get_builder_cls() -> type["TurboQuantMetadataBuilder"]:
181 return TurboQuantMetadataBuilder
182
179 - @staticmethod
180 - def get_kv_cache_shape(
181 - num_blocks: int,
182 - block_size: int,
183 - num_kv_heads: int,
184 - head_size: int,
185 - cache_dtype_str: str = "turboquant_4bit_nc",
186 - ) -> tuple[int, ...]:
187 - """Combined K+V cache shape — no leading 2 dimension.
188 -
189 - Standard attention backends use (2, num_blocks, block_size, num_kv_heads,
190 - head_dim) with a leading 2 to separate K and V. TurboQuant packs K+V
191 - into a single interleaved slot per head per position. The logical
192 - (blocks-first, head-major) shape is:
193 -
194 - (num_blocks, num_kv_heads, block_size, slot_size_aligned)
195 -
196 - Each slot = [key_packed | value_packed | padding].
197 - This is safe because TQ has its own get_kv_cache_shape override and
198 - never shares cache tensors with other backends. Layers that fall back
199 - to native dtype via kv_cache_dtype_skip_layers get their own
200 - standard-shaped cache allocation.
201 -
202 - head_size is the model's real head_dim. slot_size_aligned is computed
203 - from the TQ config to ensure correct cache allocation for all head dims.
204 - """
205 - from vllm.model_executor.layers.quantization.turboquant.config import (
206 - TurboQuantConfig,
207 - )
208 -
209 - tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype_str, head_size)
210 - return (num_blocks, num_kv_heads, block_size, tq_config.slot_size_aligned)
211 -
183 @classmethod
184 def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
185 if kv_cache_dtype is None:
vllm/v1/attention/backends/utils.py
+151 -43
@@ -1,26 +1,24 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import functools
3 import math
5 -from collections.abc import Callable
4 +from collections import defaultdict
5 +from collections.abc import Callable, Iterable, Sequence
6 from dataclasses import dataclass, field, fields, make_dataclass
7 from typing import (
8 TYPE_CHECKING,
9 Any,
10 - Literal,
10 Protocol,
12 - cast,
13 - get_args,
11 )
12
13 import numpy as np
14 import torch
15 from typing_extensions import runtime_checkable
16
20 -from vllm.config import VllmConfig, get_layers_from_vllm_config
17 +from vllm.config import CacheConfig, VllmConfig, get_layers_from_vllm_config
18 +from vllm.config.cache import _layout_from_name
19 from vllm.utils.math_utils import cdiv
20 from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor
23 -from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
21 +from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MambaSpec
22
23 if TYPE_CHECKING:
24 from vllm.v1.core.sched.output import SchedulerOutput
@@ -41,8 +39,6 @@ from vllm.v1.attention.backend import (
39 )
40
41 logger = init_logger(__name__)
44 -KVCacheLayoutType = Literal["NHD", "HND"]
45 -_KV_CACHE_LAYOUT_OVERRIDE: KVCacheLayoutType | None = None
42
43 PAD_SLOT_ID = -1
44 NULL_BLOCK_ID = 0
@@ -155,44 +151,159 @@ def fill_mm_prefix_query_ranges(
151 return num_actual_tokens
152
153
158 -def is_valid_kv_cache_layout(value: str) -> bool:
159 - return value in get_args(KVCacheLayoutType)
154 +_FLASHINFER_LAYOUT_NAMES = {
155 + "LBNHC": "NHD",
156 + "LBHNC": "HND",
157 + "BLHNC": "HND",
158 + "BLNHC": "NHD",
159 + "BHLNC": "HND",
160 +}
161
162
162 -@functools.lru_cache
163 -def get_kv_cache_layout():
164 - # Format specified by the code.
165 - global _KV_CACHE_LAYOUT_OVERRIDE
163 +def get_flashinfer_layout_string(layout: KVCacheLayout) -> str:
164 + """Return the layout name in FlashInfer's convention (NHD/HND)."""
165 + assert layout.name in _FLASHINFER_LAYOUT_NAMES, (
166 + f"KV cache layout {layout.name} has no FlashInfer equivalent; FlashInfer "
167 + "rejects it in supported_kv_cache_layouts"
168 + )
169 + return _FLASHINFER_LAYOUT_NAMES[layout.name]
170 +
171 +
172 +# Preference order when no backend declares a supported set; LBNHC (NHD)
173 +# first to match main's default.
174 +_DEFAULT_LAYOUT_PREFERENCE = (
175 + KVCacheLayout.LBNHC,
176 + KVCacheLayout.LBHNC,
177 + KVCacheLayout.BLNHC,
178 + KVCacheLayout.BLHNC,
179 + KVCacheLayout.BHLNC,
180 + KVCacheLayout.LHBNC,
181 +)
182 +
183 +
184 +def _layout_names(layouts: Iterable[KVCacheLayout]) -> list[str]:
185 + return [layout.name for layout in layouts]
186 +
187
167 - cache_layout: Literal["NHD", "HND"] | None = None
168 - if _KV_CACHE_LAYOUT_OVERRIDE is not None:
169 - cache_layout = _KV_CACHE_LAYOUT_OVERRIDE
170 - logger.debug_once(
171 - "`_KV_CACHE_LAYOUT_OVERRIDE` variable detected. "
172 - "Setting KV cache layout to %s.",
173 - cache_layout,
188 +def get_supported_kv_cache_layouts(
189 + backends: Iterable[type[AttentionBackend]],
190 +) -> list[KVCacheLayout]:
191 + """Layouts every one of the worker's backends supports, most preferred first.
192 +
193 + Every backend declares the layouts its kernels support, most preferred first
194 + (``supported_kv_cache_layouts``), or None when any layout works; workers where
195 + nothing declares follow the default preference. Identical declarations keep
196 + their order; otherwise the layout the most backends put first wins, ties
197 + keeping the enum order. An empty intersection is a hard error.
198 + """
199 + supported_layouts_lists: list[Sequence[KVCacheLayout]] = [
200 + layouts
201 + for backend in backends
202 + if (layouts := backend.supported_kv_cache_layouts()) is not None
203 + ] or [_DEFAULT_LAYOUT_PREFERENCE]
204 +
205 + first = supported_layouts_lists[0]
206 + if all(layouts == first for layouts in supported_layouts_lists[1:]):
207 + return list(first)
208 +
209 + priorities: dict[KVCacheLayout, int] = defaultdict(int)
210 + for preferred_layout, *_ in supported_layouts_lists:
211 + priorities[preferred_layout] += 1
212 + supported_layouts = set.intersection(*map(set, supported_layouts_lists))
213 + candidates = sorted(
214 + (layout for layout in KVCacheLayout if layout in supported_layouts),
215 + key=lambda layout: priorities[layout],
216 + reverse=True,
217 + )
218 + if not candidates:
219 + raise ValueError(
220 + "No KV cache layout satisfies every supported set: "
221 + f"{list(map(_layout_names, supported_layouts_lists))}."
222 )
175 - return cache_layout
223 + return candidates
224
177 - # Format specified by the user.
178 - cache_layout = envs.VLLM_KV_CACHE_LAYOUT
179 - # When neither the user nor the override specified a layout, get default
180 - if cache_layout is None:
181 - cache_layout = get_kv_connector_cache_layout()
182 - else:
183 - assert is_valid_kv_cache_layout(cache_layout)
184 - logger.info_once(
185 - "`VLLM_KV_CACHE_LAYOUT` environment variable "
186 - "detected. Setting KV cache layout to %s.",
187 - cache_layout,
225 +
226 +def record_kv_cache_layout(cache_config: CacheConfig, layout_name: str) -> None:
227 + """Adopt a layout resolved elsewhere (the engine core) in this process."""
228 + layout = _layout_from_name(layout_name)
229 + existing = cache_config.kv_cache_layout
230 + if existing is not None and existing != layout.name:
231 + raise ValueError(
232 + f"KV cache layout is already resolved to {existing}; "
233 + f"cannot change it to {layout.name}."
234 )
189 - return cache_layout
235 + cache_config.kv_cache_layout = layout.name
236 +
237 +
238 +def resolve_kv_cache_layout(
239 + vllm_config: VllmConfig,
240 + supported_layouts: list[list[str]],
241 + kv_cache_specs: Iterable[KVCacheSpec] | None = None,
242 +) -> KVCacheLayout:
243 + """Resolve one KV cache layout for the whole model.
244 +
245 + Runs once in the engine core. Every worker reports the layouts its backends
246 + support, most preferred first (``get_supported_kv_cache_layouts``); all
247 + ranks run the same backends, so their lists must agree. Specs mixing HNC
248 + shapes narrow the candidates to block-compact layouts. An explicit
249 + ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails,
250 + with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the
251 + connector's preference is used when compatible and dropped with a warning
252 + otherwise. A layout already present on ``cache_config`` wins outright, and
253 + the result is recorded there (see ``CacheConfig.kv_cache_layout``); it
254 + reaches workers through the ``set_kv_cache_layout`` RPC and
255 + ``KVCacheConfig.kv_cache_layout``.
256 + """
257 + cache_config = vllm_config.cache_config
258 + if cache_config.kv_cache_layout is not None:
259 + return cache_config.get_resolved_kv_cache_layout()
260 +
261 + assert supported_layouts and all(supported_layouts), (
262 + "No worker reported supported KV cache layouts."
263 + )
264 + assert all(names == supported_layouts[0] for names in supported_layouts[1:]), (
265 + f"Workers disagree on supported KV cache layouts: {supported_layouts}."
266 + )
267 + candidates = [_layout_from_name(name) for name in supported_layouts[0]]
268 +
269 + # A block-compact layout means the block is densely packed in memory, so any mix of
270 + # specs can re-interpret HNC with different sizes as long as the total number of
271 + # bytes is the same. If not block-compact, each spec must agree on HNC to alias
272 + # the same page (this aliasing is done by the Hybrid Memory Allocator, HMA).
273 + hnc_shapes = {
274 + (spec.num_heads, spec.num_states, spec.page_size_bytes)
275 + for spec in kv_cache_specs or ()
276 + }
277 + if len(hnc_shapes) > 1:
278 + candidates = [m for m in candidates if m.is_block_compact]
279 + if not candidates:
280 + raise ValueError(
281 + "Specs with mixed HNC shapes need a block-compact layout, but "
282 + f"none is in every supported set: {supported_layouts}."
283 + )
284
285 + if (requested := envs.VLLM_KV_CACHE_LAYOUT) is not None:
286 + layout = _layout_from_name(requested)
287 + if layout not in candidates:
288 + raise ValueError(
289 + f"VLLM_KV_CACHE_LAYOUT={requested} does not satisfy every "
290 + f"supported set; valid layouts: {_layout_names(candidates)}."
291 + )
292 + elif (connector := get_kv_connector_cache_layout(vllm_config)) is not None:
293 + layout = _layout_from_name(connector)
294 + if layout not in candidates:
295 + logger.warning_once(
296 + f"KV connector cache layout {connector} does not satisfy every "
297 + f"supported set; valid layouts: {_layout_names(candidates)}. "
298 + f"Using {candidates[0].name} instead."
299 + )
300 + layout = candidates[0]
301 + else:
302 + layout = candidates[0]
303
192 -def set_kv_cache_layout(cache_layout: KVCacheLayoutType | None):
193 - global _KV_CACHE_LAYOUT_OVERRIDE
194 - _KV_CACHE_LAYOUT_OVERRIDE = cache_layout
195 - get_kv_cache_layout.cache_clear()
304 + logger.info_once("Using %s KV cache layout.", layout.name)
305 + cache_config.kv_cache_layout = layout.name
306 + return layout
307
308
309 @dataclass
@@ -264,10 +375,7 @@ def get_num_attention_heads_from_layers(
375 )
376 if not attn_layers:
377 return None
267 - heads = {
268 - cast(Any, getattr(layer, "impl", layer)).num_heads
269 - for layer in attn_layers.values()
270 - }
378 + heads = {layer.impl.num_heads for layer in attn_layers.values()}
379 assert len(heads) == 1, (
380 f"All layers in one attention group must share num_heads; "
381 f"got {heads} for {layer_names}."
vllm/v1/attention/selector.py
+2 -16
@@ -8,7 +8,6 @@ import torch
8
9 import vllm.envs as envs
10 from vllm.config.cache import CacheDType
11 -from vllm.logger import init_logger
11 from vllm.utils.import_utils import resolve_obj_by_qualname
12 from vllm.v1.attention.backend import AttentionBackend, AttentionType
13 from vllm.v1.attention.backends.registry import (
@@ -18,8 +17,6 @@ from vllm.v1.attention.backends.registry import (
17 if TYPE_CHECKING:
18 from vllm.v1.kv_cache_interface import KVCacheSpecKind
19
21 -logger = init_logger(__name__)
22 -
20
21 class AttentionSelectorConfig(NamedTuple):
22 head_size: int
@@ -185,6 +182,8 @@ def get_attn_backend(
182 )
183 backend = attention_config.backend_per_kind.get(kind.value, backend)
184
185 + # The KV cache layout is resolved across all of the model's backends at once
186 + # in get_kv_cache_spec(); a single selection cannot see its peers.
187 return _cached_get_attn_backend(
188 backend=backend,
189 attn_selector_config=attn_selector_config,
@@ -210,19 +209,6 @@ def _cached_get_attn_backend(
209 f"Invalid attention backend for {current_platform.device_name}"
210 )
211 backend = resolve_obj_by_qualname(attention_cls)
213 -
214 - # Adjust kv cache layout if the selected backend requires a specific one
215 - required_layout = backend.get_required_kv_cache_layout()
216 - if required_layout is not None:
217 - from vllm.v1.attention.backends.utils import set_kv_cache_layout
218 -
219 - set_kv_cache_layout(required_layout)
220 - logger.info_once(
221 - "Using %s KV cache layout for %s backend.",
222 - required_layout,
223 - backend.get_name(),
224 - )
225 -
212 return backend
213
214
vllm/v1/core/kv_cache_utils.py
+138 -214
@@ -10,7 +10,7 @@ from collections import defaultdict
10 from collections.abc import Callable, Iterable, Iterator, Sequence
11 from dataclasses import dataclass, replace
12 from functools import partial
13 -from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload
13 +from typing import Any, NamedTuple, NewType, TypeAlias, overload
14
15 from vllm import envs
16 from vllm.config import VllmConfig
@@ -26,6 +26,7 @@ from vllm.v1.kv_cache_interface import (
26 HiddenStateCacheSpec,
27 KVCacheConfig,
28 KVCacheGroupSpec,
29 + KVCacheLayout,
30 KVCacheSpec,
31 KVCacheTensor,
32 MambaSpec,
@@ -33,6 +34,7 @@ from vllm.v1.kv_cache_interface import (
34 SlidingWindowMLASpec,
35 SlidingWindowSpec,
36 UniformTypeKVCacheSpecs,
37 + compute_layout_strides,
38 replace_as,
39 )
40 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
@@ -897,7 +899,7 @@ def check_enough_kv_cache_memory(
899 # of the specs since grouping may unify them in-place.
900 groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec))
901 check_memory = (
900 - available_memory - _pool_bytes_per_block(vllm_config, groups)
902 + available_memory - _pool_bytes_per_block(groups)
903 if groups
904 else available_memory
905 )
@@ -999,45 +1001,14 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
1001 return num_blocks
1002
1003
1002 -def _pool_bytes_per_block(
1003 - vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec]
1004 -) -> int:
1004 +def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
1005 """
1006 Bytes consumed by one block in the worker's shared KV cache pool, mirroring
1007 the divisor used by `get_kv_cache_config_from_groups` to convert
1008 `available_memory` into `num_blocks`. Used to compute the effective KV cache
1009 capacity once `num_gpu_blocks_override` is applied.
1010 """
1011 - if len(kv_cache_groups) == 1 and isinstance(
1012 - kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
1013 - ):
1014 - return kv_cache_groups[0].kv_cache_spec.page_size_bytes
1015 - if _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
1016 - block_stride, _ = _get_packed_kv_cache_layout(kv_cache_groups)
1017 - return block_stride
1018 - group_size = max(len(g.layer_names) for g in kv_cache_groups)
1019 - page_size = get_uniform_page_size([g.kv_cache_spec for g in kv_cache_groups])
1020 - return page_size * group_size
1021 -
1022 -
1023 -def get_num_blocks(
1024 - vllm_config: VllmConfig,
1025 - num_layers: int,
1026 - available_memory: int,
1027 - page_size: int,
1028 -) -> int:
1029 - """
1030 - Get the number of kv cache blocks.
1031 -
1032 - Args:
1033 - vllm_config: The global VllmConfig
1034 - num_layers: The number of layers
1035 - available_memory: Memory available for KV cache in bytes.
1036 - page_size: The page size of the KV cache.
1037 - """
1038 - num_blocks = int(available_memory // page_size // num_layers)
1039 - num_blocks = max(num_blocks, 0)
1040 - return may_override_num_blocks(vllm_config, num_blocks)
1011 + return _get_kv_cache_bytes_per_block(kv_cache_groups)
1012
1013
1014 def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
@@ -1092,11 +1063,14 @@ def unify_kv_cache_spec_page_size(
1063 size by increasing the block size of layers with smaller page size. Two
1064 cases cannot be unified by block size alone and pad their physical page to
1065 the maximum instead: Mamba layers, whose page size comes from state shapes
1095 - and is independent of block size; and attention layers whose page does not
1096 - evenly divide the maximum and whose backend opts in via
1097 - ``AttentionSpec.indexes_kv_by_block_stride`` (the padded page is read through
1098 - a strided view, which not every backend handles). Raise NotImplementedError
1099 - if failed to unify the page size.
1066 + and is independent of block size; and non-MLA attention layers whose page
1067 + does not evenly divide the maximum (the padded page is read through a
1068 + strided view). MLA is excluded because sparse MLA indexes the cache in
1069 + whole token rows (see ``flat_kv_row_view``), so its block stride can only
1070 + be padded by its own row-aligned ``alignment``, not to an arbitrary page
1071 + size. Raise NotImplementedError if failed to unify the page size;
1072 + ``get_kv_cache_groups`` catches it to try the full-allocation fallback
1073 + (e.g. MLA next to an incompatible sliding-window draft).
1074
1075 Args:
1076 kv_cache_spec: The KVCacheSpec of each attention layer in the model
@@ -1130,18 +1104,15 @@ def unify_kv_cache_spec_page_size(
1104 ratio = max_page_size // layer_page_size
1105 new_block_size = layer_spec.block_size * ratio
1106 new_spec = replace(layer_spec, block_size=new_block_size)
1133 - elif (
1134 - isinstance(layer_spec, AttentionSpec)
1135 - and layer_spec.indexes_kv_by_block_stride
1107 + elif isinstance(layer_spec, AttentionSpec) and not isinstance(
1108 + layer_spec, MLAAttentionSpec
1109 ):
1110 new_spec = replace(layer_spec, page_size_padded=max_page_size)
1111 else:
1112 raise NotImplementedError(
1113 f"Layer {layer_name}: page size is not divisible by the "
1114 "maximum page size and cannot be padded. Padding is only "
1142 - "supported for attention layers whose backend indexes KV "
1143 - "pages by the block stride (indexes_kv_by_block_stride is "
1144 - "True)."
1115 + "supported for non-MLA attention layers."
1116 )
1117 assert new_spec.page_size_bytes == max_page_size
1118 new_kv_cache_spec[layer_name] = new_spec
@@ -1296,83 +1267,64 @@ def _get_kv_cache_groups_uniform_page_size(
1267 return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
1268
1269
1299 -def _get_packed_kv_cache_layout(
1300 - kv_cache_groups: list[KVCacheGroupSpec],
1301 -) -> tuple[int, dict[int, list[str]]]:
1302 - """Lay out each cache group densely in one shared block slab.
1303 -
1304 - A block ID is owned by one cache group at a time, so layouts from different
1305 - groups may overlap. Layers within a group remain disjoint.
1306 - """
1307 - layers_by_offset: dict[int, list[str]] = defaultdict(list)
1308 - block_stride = 0
1309 - for group in kv_cache_groups:
1310 - spec = group.kv_cache_spec
1311 - byte_offset = 0
1312 - for layer_name in group.layer_names:
1313 - if isinstance(spec, UniformTypeKVCacheSpecs):
1314 - page_size = spec.kv_cache_specs[layer_name].page_size_bytes
1315 - else:
1316 - page_size = spec.page_size_bytes
1317 - layers_by_offset[byte_offset].append(layer_name)
1318 - byte_offset += page_size
1319 - block_stride = max(block_stride, byte_offset)
1320 - assert block_stride > 0
1321 - return block_stride, layers_by_offset
1270 +def _get_per_layer_spec(
1271 + group: KVCacheGroupSpec,
1272 + layer_name: str,
1273 +) -> KVCacheSpec:
1274 + spec = group.kv_cache_spec
1275 + if isinstance(spec, UniformTypeKVCacheSpecs):
1276 + return spec.kv_cache_specs[layer_name]
1277 + return spec
1278
1279
1324 -def _use_packed_kv_cache_config(
1325 - vllm_config: VllmConfig,
1280 +def _get_kv_cache_bytes_per_block(
1281 kv_cache_groups: list[KVCacheGroupSpec],
1327 -) -> bool:
1328 - is_dsv4 = all(
1329 - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1282 +) -> int:
1283 + """Return the largest cache group's bytes per block."""
1284 + bytes_per_block = max(
1285 + sum(
1286 + _get_per_layer_spec(group, layer_name).page_size_bytes
1287 + for layer_name in group.layer_names
1288 + )
1289 for group in kv_cache_groups
1290 )
1332 - kv_transfer_config = vllm_config.kv_transfer_config
1333 - extra_config = (
1334 - kv_transfer_config.kv_connector_extra_config
1335 - if kv_transfer_config is not None
1336 - else {}
1337 - )
1338 - # NOTE: enable_cross_layers_blocks is an experimental API and subject to change with
1339 - # https://github.com/vllm-project/vllm/issues/42082
1340 - enable_cross_layers = (
1341 - str(extra_config.get("enable_cross_layers_blocks", "False")).lower() == "true"
1342 - )
1343 - return is_dsv4 or (enable_cross_layers and len(kv_cache_groups) > 1)
1291 + assert bytes_per_block > 0
1292 + return bytes_per_block
1293
1294
1346 -def _get_kv_cache_config_packed(
1347 - vllm_config: VllmConfig,
1295 +def validate_kv_cache_layout(
1296 + layout: KVCacheLayout,
1297 kv_cache_groups: list[KVCacheGroupSpec],
1349 - available_memory: int,
1350 -) -> tuple[int, list[KVCacheTensor]]:
1351 - """Plan a packed per-block KV cache tensor layout.
1298 +) -> None:
1299 + """Validate that the resolved layout can express this model's packing.
1300
1353 - Cache groups use dense, overlapping layouts within one block slab. Each
1354 - emitted tensor aliases the same physical backing allocation.
1301 + The layout was chosen once in the engine core from the backends' supported
1302 + sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4
1303 + indexer) declares block-outermost layouts there, so an inexpressible
1304 + layout reaching this point is an error.
1305 """
1356 - block_stride, layers_by_offset = _get_packed_kv_cache_layout(kv_cache_groups)
1357 -
1358 - num_blocks = available_memory // block_stride
1359 - num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1360 -
1361 - total_size = block_stride * num_blocks
1306 + page_sizes = {
1307 + _get_per_layer_spec(group, layer_name).page_size_bytes
1308 + for group in kv_cache_groups
1309 + for layer_name in group.layer_names
1310 + }
1311 + if len(page_sizes) == 1:
1312 + # A rectangular layer dim exists; every layout can express it.
1313 + return
1314
1363 - kv_cache_tensors: list[KVCacheTensor] = []
1364 - for byte_offset in sorted(layers_by_offset):
1365 - kv_cache_tensors.append(
1366 - KVCacheTensor(
1367 - size=total_size,
1368 - shared_by=layers_by_offset[byte_offset],
1369 - offset=byte_offset,
1370 - block_stride=block_stride,
1371 - )
1315 + # Mixed page sizes pack pages side by side within a block, which needs each page
1316 + # to be one contiguous chunk inside its block (a block-compact layout) and, with
1317 + # multiple KV cache groups, the layer dim inside the block dim.
1318 + if not layout.is_block_compact or (
1319 + len(kv_cache_groups) > 1 and layout.is_layer_compact
1320 + ):
1321 + raise ValueError(
1322 + f"KV cache layout {layout.name} cannot express this model's "
1323 + f"mixed page sizes ({sorted(page_sizes)}); a backend should "
1324 + "declare block-outermost supported layouts (e.g. BLHNC), or "
1325 + "set VLLM_KV_CACHE_LAYOUT=BLHNC."
1326 )
1327
1374 - return num_blocks, kv_cache_tensors
1375 -
1328
1329 def get_kv_cache_config_from_groups(
1330 vllm_config: VllmConfig,
@@ -1402,58 +1354,60 @@ def get_kv_cache_config_from_groups(
1354 ),
1355 )
1356
1405 - # Determine how model runners should initialize the KV cache tensors.
1406 - if len(kv_cache_groups) == 1 and isinstance(
1407 - kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
1408 - ):
1409 - # Special case: all layers have the same type of KV cache but with
1410 - # different hidden sizes. Allocate different amount of memory for each
1411 - # layer based on its hidden size.
1412 - num_blocks = (
1413 - available_memory // kv_cache_groups[0].kv_cache_spec.page_size_bytes
1414 - )
1415 - num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1416 - per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs
1417 - kv_cache_tensors = [
1418 - KVCacheTensor(
1419 - size=per_layer_specs[layer_name].page_size_bytes * num_blocks,
1420 - shared_by=[layer_name],
1357 + layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
1358 + validate_kv_cache_layout(layout, kv_cache_groups)
1359 + bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups)
1360 + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
1361 +
1362 + num_blocks = available_memory // bytes_per_block
1363 + num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1364 + size = bytes_per_block * num_blocks
1365 +
1366 + # Groups alias from byte 0. Spec regions are laid out differently:
1367 + #
1368 + # block-outer (the same packing repeats for every block):
1369 + # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ...
1370 + # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ...
1371 + # |<--- bytes_per_block -->|
1372 + #
1373 + # layer-outer (only supported for uniform page sizes or single-group models):
1374 + # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] |
1375 + # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] |
1376 +
1377 + kv_cache_tensors = []
1378 + for group in kv_cache_groups:
1379 + group_spec = group.kv_cache_spec
1380 + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list)
1381 + if isinstance(group_spec, UniformTypeKVCacheSpecs):
1382 + for layer_name, spec in group_spec.kv_cache_specs.items():
1383 + layers_by_spec[spec].append(layer_name)
1384 + elif group.layer_names:
1385 + layers_by_spec[group_spec].extend(group.layer_names)
1386 +
1387 + byte_offset = 0
1388 + for spec, layer_names in layers_by_spec.items():
1389 + layer_stride, block_stride, _, _, _ = compute_layout_strides(
1390 + spec,
1391 + num_blocks,
1392 + len(layer_names),
1393 + layout,
1394 + fixed_strides=(None, interleaved_block_stride, None, None, None),
1395 + )
1396 + offset = (
1397 + byte_offset
1398 + * max(layer_stride, spec.page_size_bytes)
1399 + // spec.page_size_bytes
1400 )
1422 - for layer_name in kv_cache_groups[0].layer_names
1423 - ]
1424 - elif _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
1425 - # DeepSeek V4 uses the packed layout by default. Other multi-group
1426 - # layouts can opt in with --enable-cross-layers.
1427 - num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
1428 - vllm_config, kv_cache_groups, available_memory
1429 - )
1430 - else:
1431 - # General case:
1432 - # We will have group_size memory pools, each is shared by one layer from
1433 - # each group. As layers of different groups have different block table,
1434 - # they will use different parts of the shared Tensor.
1435 - # The memory layout for 3 groups (full.0, full.1), (sw.0, sw.2),
1436 - # (sw.1, padding) will be: (group_size = 2)
1437 - # full.0, sw.0, sw.1: share a Tensor with size=available_memory//2
1438 - # full.1, sw.2: share another Tensor with size=available_memory//2
1439 - group_size = max(len(group.layer_names) for group in kv_cache_groups)
1440 -
1441 - page_size = get_uniform_page_size(
1442 - [group.kv_cache_spec for group in kv_cache_groups]
1443 - )
1444 - assert group_size > 0, "group_size must be greater than 0"
1445 - num_blocks = get_num_blocks(
1446 - vllm_config, group_size, available_memory, page_size
1447 - )
1448 - kv_cache_tensors = []
1449 - for i in range(group_size):
1450 - shared_by = []
1451 - for j in range(len(kv_cache_groups)):
1452 - if i < len(kv_cache_groups[j].layer_names):
1453 - shared_by.append(kv_cache_groups[j].layer_names[i])
1401 kv_cache_tensors.append(
1455 - KVCacheTensor(size=page_size * num_blocks, shared_by=shared_by)
1402 + KVCacheTensor(
1403 + size=size,
1404 + layers=layer_names,
1405 + layer_stride=layer_stride,
1406 + block_stride=block_stride,
1407 + offset=offset,
1408 + )
1409 )
1410 + byte_offset += len(layer_names) * spec.page_size_bytes
1411
1412 return KVCacheConfig(
1413 num_blocks=num_blocks,
@@ -1941,57 +1895,26 @@ def _max_memory_usage_bytes_from_groups(
1895 This correctly accounts for padding in hybrid models. For example, if a
1896 model has 8 full attention layers and 9 sliding window layers, they will
1897 be padded to 9 full + 9 sliding window for uniform group sizes.
1898 +
1899 + Each group independently claims blocks from the shared pool, so a request consumes
1900 + the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``.
1901 """
1902 if not kv_cache_groups:
1903 return 0
1904
1948 - if len(kv_cache_groups) == 1 and isinstance(
1949 - kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
1950 - ):
1951 - # UniformTypeKVCacheSpecs special case (single group, per-layer specs)
1952 - per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs
1953 - return sum(
1954 - spec.max_memory_usage_bytes(vllm_config)
1955 - for spec in per_layer_specs.values()
1956 - )
1957 - elif all(
1958 - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1959 - for group in kv_cache_groups
1960 - ):
1961 - # Special case (only DeepseekV4 for now): all groups are
1962 - # UniformTypeKVCacheSpecs.
1963 - # They must already be page_size aligned and share a common padded
1964 - # layer-tuple layout. Even groups with fewer actual tuples still reserve
1965 - # the global number of tuple slots in the shared tensor layout.
1966 - full_mla_spec = cast(UniformTypeKVCacheSpecs, kv_cache_groups[0].kv_cache_spec)
1967 - layer_tuple_bytes = sum(full_mla_spec.get_page_sizes())
1968 - num_layer_tuples = max(
1969 - cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).get_num_layer_tuples()
1970 - for group in kv_cache_groups
1971 - )
1972 -
1973 - total_max_mem_usage_bytes = 0
1974 - for group in kv_cache_groups:
1975 - group_spec = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec)
1976 - g_max_mem_usage_pages = group_spec.max_memory_usage_pages(vllm_config)
1977 - g_max_mem_usage_page_bytes = (
1978 - num_layer_tuples * g_max_mem_usage_pages * layer_tuple_bytes
1905 + bytes_per_block = _pool_bytes_per_block(kv_cache_groups)
1906 + total_blocks = 0
1907 + for group in kv_cache_groups:
1908 + spec = group.kv_cache_spec
1909 + if isinstance(spec, UniformTypeKVCacheSpecs):
1910 + total_blocks += spec.max_memory_usage_pages(vllm_config)
1911 + else:
1912 + total_blocks += cdiv(
1913 + spec.max_memory_usage_bytes(vllm_config),
1914 + spec.page_size_bytes,
1915 )
1980 - total_max_mem_usage_bytes += g_max_mem_usage_page_bytes
1981 - return total_max_mem_usage_bytes
1982 -
1983 - # General case: group_size pools, each shared by one layer per group
1984 - # Memory = group_size * page_size * blocks_for_max_len
1985 - group_size = max(len(group.layer_names) for group in kv_cache_groups)
1986 - page_size = get_uniform_page_size(
1987 - [group.kv_cache_spec for group in kv_cache_groups]
1988 - )
1989 - blocks_needed = sum(
1990 - cdiv(group.kv_cache_spec.max_memory_usage_bytes(vllm_config), page_size)
1991 - for group in kv_cache_groups
1992 - )
1916
1994 - return group_size * page_size * blocks_needed
1917 + return bytes_per_block * total_blocks
1918
1919
1920 def _estimate_max_model_len_from_groups(
@@ -2228,7 +2151,7 @@ def get_kv_cache_configs(
2151 if not groups:
2152 adjusted_memory.append(avail_mem)
2153 continue
2231 - bytes_per_block = _pool_bytes_per_block(vllm_config, groups)
2154 + bytes_per_block = _pool_bytes_per_block(groups)
2155 logger.info(
2156 "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
2157 avail_mem // bytes_per_block,
@@ -2241,7 +2164,7 @@ def get_kv_cache_configs(
2164 # the capacity check both plan against usable blocks. Allocation below
2165 # still uses the full memory.
2166 check_memory = [
2244 - avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem
2167 + avail_mem - _pool_bytes_per_block(groups) if groups else avail_mem
2168 for groups, avail_mem in zip(projected_groups_per_worker, available_memory)
2169 ]
2170
@@ -2278,14 +2201,15 @@ def get_kv_cache_configs(
2201 min_num_blocks = min(
2202 kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs
2203 )
2281 - for kv_cache_config in kv_cache_configs:
2282 - num_blocks_old = kv_cache_config.num_blocks
2283 - kv_cache_config.num_blocks = min_num_blocks
2284 -
2285 - # Shrink tensor size proportionally
2286 - for tensor in kv_cache_config.kv_cache_tensors:
2287 - assert tensor.size % num_blocks_old == 0
2288 - tensor.size = tensor.size // num_blocks_old * min_num_blocks
2204 + for i, kv_cache_config in enumerate(kv_cache_configs):
2205 + if kv_cache_config.num_blocks == min_num_blocks:
2206 + continue
2207 + # Re-plan with exactly the memory the smallest rank can afford, so
2208 + # strides and offsets stay consistent with the shrunken allocation.
2209 + groups = kv_cache_config.kv_cache_groups
2210 + kv_cache_configs[i] = get_kv_cache_config_from_groups(
2211 + vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups)
2212 + )
2213
2214 return kv_cache_configs
2215
vllm/v1/engine/core.py
+14
@@ -44,6 +44,7 @@ from vllm.utils.gc_utils import (
44 from vllm.utils.hashing import get_hash_fn_by_name
45 from vllm.utils.network_utils import make_zmq_socket
46 from vllm.utils.system_utils import decorate_logs, set_process_title
47 +from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
48 from vllm.v1.core.kv_cache_utils import (
49 BlockHash,
50 generate_scheduler_kv_cache_config,
@@ -281,6 +282,17 @@ class EngineCore:
282 )
283 vllm_config.cache_config.enable_prefix_caching = False
284
285 + # Resolve the KV cache layout before memory profiling: workers that
286 + # capture full cudagraphs initialize a minimal KV cache during it.
287 + # Attention-free models resolve the default so layout reads never precede
288 + # resolution.
289 + layout = resolve_kv_cache_layout(
290 + vllm_config,
291 + self.model_executor.get_supported_kv_cache_layouts(),
292 + [s for specs in kv_cache_specs for s in specs.values()],
293 + )
294 + self.model_executor.set_kv_cache_layout(layout.name)
295 +
296 has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)
297 if has_kv_cache:
298 if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
@@ -307,6 +319,8 @@ class EngineCore:
319 kv_cache_configs = get_kv_cache_configs(
320 vllm_config, kv_cache_specs, available_gpu_memory
321 )
322 + for kv_cache_config in kv_cache_configs:
323 + kv_cache_config.kv_cache_layout = vllm_config.cache_config.kv_cache_layout
324
325 # If auto-fit reduced max_model_len, sync the new value to workers.
326 # This is needed because workers were spawned before memory profiling
vllm/v1/executor/abstract.py
+8
@@ -151,6 +151,14 @@ class Executor(ABC):
151 def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
152 return self.collective_rpc("get_kv_cache_spec")
153
154 + def get_supported_kv_cache_layouts(self) -> list[list[str]]:
155 + """Layouts each worker's backends support, most preferred first."""
156 + return self.collective_rpc("get_supported_kv_cache_layouts")
157 +
158 + def set_kv_cache_layout(self, layout_name: str) -> None:
159 + """Publish the resolved KV cache layout to the workers."""
160 + self.collective_rpc("set_kv_cache_layout", args=(layout_name,))
161 +
162 @overload
163 def collective_rpc(
164 self,
vllm/v1/kv_cache_interface.py
+202 -40
@@ -8,6 +8,7 @@ from collections import Counter
8 from collections.abc import Collection
9 from dataclasses import dataclass, fields, replace
10 from enum import Enum, IntEnum
11 +from fractions import Fraction
12 from math import prod
13 from typing import TYPE_CHECKING, TypeVar
14
@@ -18,6 +19,7 @@ from vllm.logger import init_logger
19 from vllm.utils.math_utils import cdiv, round_up
20 from vllm.utils.torch_utils import get_dtype_size
21 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
22 +from vllm.v1.kv_cache_layout import _DIM_B, _DIM_L, KVCacheLayout
23 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
24
25 if TYPE_CHECKING:
@@ -147,6 +149,18 @@ class KVCacheSpec:
149 # number of tokens in a block
150 block_size: int
151
152 + @property
153 + def num_heads(self) -> int:
154 + raise NotImplementedError
155 +
156 + @property
157 + def tokens_per_state(self) -> int | Fraction:
158 + raise NotImplementedError
159 +
160 + @property
161 + def state_content_size_bytes(self) -> int:
162 + raise NotImplementedError
163 +
164 @property
165 def page_size_bytes(self) -> int:
166 """
@@ -158,8 +172,13 @@ class KVCacheSpec:
172 raise NotImplementedError
173
174 @property
161 - def storage_block_size(self) -> int:
162 - return self.block_size
175 + def num_states(self) -> int:
176 + return self.get_num_kernel_states(self.block_size)
177 +
178 + def get_num_kernel_states(self, kernel_block_size: int) -> int:
179 + if self.tokens_per_state > 0:
180 + return kernel_block_size // self.tokens_per_state
181 + return 1
182
183 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
184 """
@@ -214,6 +233,143 @@ class KVCacheSpec:
233 )
234
235
236 +def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor:
237 + """View a kernel-block-granular layer cache with manager blocks as dim 0.
238 +
239 + Kernel block splitting subdivides each manager block into uniformly strided
240 + kernel blocks, so grouping is a pure view: ``(num_blocks * ratio, ...)``
241 + """
242 + if cache.shape[0] == num_blocks:
243 + return cache
244 + assert cache.shape[0] % num_blocks == 0
245 + return cache.unflatten(0, (num_blocks, -1))
246 +
247 +
248 +def compute_layer_kv_cache_shape_bytes(
249 + spec: KVCacheSpec,
250 + num_blocks: int,
251 + kernel_block_size: int | None = None,
252 +) -> tuple[int, ...]:
253 + """Return the 4D logical shape ``(B, H, N, C)`` where C is in bytes."""
254 + bs = kernel_block_size if kernel_block_size is not None else spec.block_size
255 + assert spec.block_size % bs == 0, (
256 + f"Kernel block size {bs} must divide KV cache block size {spec.block_size}."
257 + )
258 + blocks_per_page = spec.block_size // bs
259 + return (
260 + num_blocks * blocks_per_page,
261 + spec.num_heads,
262 + spec.get_num_kernel_states(bs),
263 + spec.state_content_size_bytes,
264 + )
265 +
266 +
267 +def compute_layout_strides(
268 + spec: KVCacheSpec,
269 + num_blocks: int,
270 + num_layers: int,
271 + layout: KVCacheLayout,
272 + kernel_block_size: int | None = None,
273 + fixed_strides: tuple[int | None, ...] = (None,) * 5,
274 +) -> tuple[int, ...]:
275 + """Byte strides in logical ``[L, B, H, N, C]`` axis order."""
276 + assert len(fixed_strides) == 5
277 + assert all(stride is None or stride > 0 for stride in fixed_strides)
278 + shape = (
279 + num_layers,
280 + *compute_layer_kv_cache_shape_bytes(spec, num_blocks, kernel_block_size),
281 + )
282 + order = layout.stride_order
283 + padded_page_size = getattr(spec, "page_size_padded", None)
284 + if padded_page_size is not None:
285 + assert kernel_block_size is None or kernel_block_size == spec.block_size, (
286 + "Padded KV pages do not support kernel block splitting."
287 + )
288 + page_grid_end = max(order.index(_DIM_L), order.index(_DIM_B)) + 1
289 + page_grid_shape = tuple(shape[dim] for dim in order[:page_grid_end])
290 + assert prod(page_grid_shape) == num_layers * num_blocks, (
291 + "Page padding requires dimensions outside the page tail to be L, B, "
292 + f"or singleton; got {layout.name} with shape {shape}."
293 + )
294 +
295 + strides = [0] * 5
296 + current_stride = 1
297 + for physical_idx, dim in reversed(tuple(enumerate(order))):
298 + if padded_page_size is not None and physical_idx == page_grid_end - 1:
299 + current_stride = max(current_stride, padded_page_size)
300 + assert current_stride % padded_page_size == 0
301 + strides[dim] = fixed_strides[dim] or current_stride
302 + current_stride = strides[dim] * shape[dim]
303 + return tuple(strides)
304 +
305 +
306 +def create_kv_cache_views(
307 + raw: torch.Tensor,
308 + spec: KVCacheSpec,
309 + num_blocks: int,
310 + layout: KVCacheLayout,
311 + kv_cache_tensor: KVCacheTensor,
312 + kernel_block_size: int | None = None,
313 +) -> list[torch.Tensor]:
314 + """View a flat int8 buffer as one 4D ``[B, H, N, C]`` view per layer.
315 +
316 + Block ``b`` of layer ``l`` starts at the tensor offset plus its layer and
317 + block stride contributions.
318 + """
319 + num_layers = len(kv_cache_tensor.layers)
320 + layer_stride = kv_cache_tensor.layer_stride
321 + block_stride = kv_cache_tensor.block_stride
322 + shape_bytes = compute_layer_kv_cache_shape_bytes(
323 + spec, num_blocks, kernel_block_size
324 + )
325 + ratio = shape_bytes[0] // num_blocks
326 + if ratio > 1:
327 + # Kernel blocks subdivide a manager block into `ratio` equal pieces, so
328 + # they sit a constant stride apart only if a block is one dense page: no
329 + # padding at its end, and no other layer's page before the next block.
330 + dense_page_size = prod(compute_layer_kv_cache_shape_bytes(spec, 1)[1:])
331 + if block_stride != dense_page_size:
332 + raise ValueError(
333 + f"The resolved KV cache layout ({layout.name}) does not store "
334 + "blocks as dense, unpadded pages (block stride "
335 + f"{block_stride} != page {dense_page_size}), so a manager "
336 + f"block cannot be split into {ratio} kernel blocks of "
337 + f"{kernel_block_size} tokens. Reduce --block-size to "
338 + f"{kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a "
339 + "layer-compact layout (e.g. LBNHC)."
340 + )
341 + assert block_stride % ratio == 0, (
342 + f"Block stride {block_stride} must divide into {ratio} equal kernel blocks."
343 + )
344 + block_stride //= ratio
345 +
346 + logical_shape = (num_layers, *shape_bytes)
347 + strides = compute_layout_strides(
348 + spec,
349 + num_blocks,
350 + num_layers,
351 + layout,
352 + kernel_block_size,
353 + fixed_strides=(layer_stride, block_stride, None, None, None),
354 + )
355 + dtype = getattr(spec, "dtype", None)
356 +
357 + view_5d = torch.as_strided(
358 + raw,
359 + size=logical_shape,
360 + stride=strides,
361 + storage_offset=raw.storage_offset() + kv_cache_tensor.offset,
362 + )
363 +
364 + views = []
365 + for layer_idx in range(num_layers):
366 + cache_logical = view_5d[layer_idx]
367 + if dtype is not None:
368 + cache_logical = cache_logical.view(dtype)
369 + views.append(cache_logical)
370 + return views
371 +
372 +
373 @dataclass(frozen=True, kw_only=True)
374 class AttentionSpec(KVCacheSpec):
375 num_kv_heads: int
@@ -222,13 +378,16 @@ class AttentionSpec(KVCacheSpec):
378 head_size_v: int = None # type: ignore[assignment]
379 kv_quant_mode: KVQuantMode = KVQuantMode.NONE
380 page_size_padded: int | None = None
225 - indexes_kv_by_block_stride: bool = False
381 num_head_slots: int | None = None
382 """H of the logical ``[B, H, N, C]`` page when packing diverges from one
383 slot per KV head. None means one slot per KV head. Published by the backend.
384 """
385 state_content_bytes: int | None = None
386 """C in bytes when packed; None means dense K/V content."""
387 + tokens_per_state: int | Fraction = 1
388 + """Tokens covered by one stored state. Ints > 1 compress multiple tokens
389 + into one state (DSv4 sparse MLA); fractions < 1 store multiple states per
390 + token (Whisper block pooling: ``Fraction(1, block_pool_size)``)."""
391
392 def __post_init__(self):
393 if self.head_size_v is None:
@@ -249,7 +408,7 @@ class AttentionSpec(KVCacheSpec):
408
409 @property
410 def unpadded_page_size_bytes(self) -> int:
252 - return self.num_heads * self.storage_block_size * self.state_content_size_bytes
411 + return self.num_heads * self.num_states * self.state_content_size_bytes
412
413 @property
414 def page_size_bytes(self) -> int:
@@ -345,9 +504,9 @@ class FullAttentionSpec(AttentionSpec):
504 dtype=specs[0].dtype,
505 kv_quant_mode=specs[0].kv_quant_mode,
506 page_size_padded=specs[0].page_size_padded,
348 - indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
507 num_head_slots=specs[0].num_head_slots,
508 state_content_bytes=specs[0].state_content_bytes,
509 + tokens_per_state=specs[0].tokens_per_state,
510 sliding_window=cls.merge_window_sizes(sliding_window),
511 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
512 # If any layer in the group is non-causal, treat the group as
@@ -384,7 +543,6 @@ class MLAAttentionSpec(FullAttentionSpec):
543 cache_dtype_str: str | None = None
544 # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults.
545 alignment: int | None = None # Default to None for no padding.
387 - compress_ratio: int = 1 # Default to 1 for no compression.
546 model_version: str | None = None
547 # Marks draft groups that flatten a non-causal query block into decode rows.
548 non_causal_multi_token_decode: bool = False
@@ -395,28 +553,21 @@ class MLAAttentionSpec(FullAttentionSpec):
553 super().__post_init__()
554 _apply_alignment_padding(self)
555
398 - @property
399 - def storage_block_size(self) -> int:
400 - return self.block_size // self.compress_ratio
401 -
556 @classmethod
557 def merge(cls, specs: list[Self]) -> Self:
558 assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), (
559 "All attention layers in the same KV cache group must be MLAAttentionSpec."
560 )
561 cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
408 - compress_ratio_set = set(spec.compress_ratio for spec in specs)
562 + tokens_per_state_set = set(spec.tokens_per_state for spec in specs)
563 model_version_set = set(spec.model_version for spec in specs)
410 - block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
564 assert (
565 len(cache_dtype_str_set) == 1
413 - and len(compress_ratio_set) == 1
566 + and len(tokens_per_state_set) == 1
567 and len(model_version_set) == 1
415 - and len(block_stride_set) == 1
568 ), (
569 "All attention layers in the same KV cache group must use the same "
418 - "quantization method, compress ratio, model version, and KV block "
419 - "stride indexing."
570 + "quantization method, tokens per state, and model version."
571 )
572 merged_spec = cls(
573 block_size=specs[0].block_size,
@@ -427,9 +578,8 @@ class MLAAttentionSpec(FullAttentionSpec):
578 page_size_padded=specs[0].page_size_padded,
579 num_head_slots=specs[0].num_head_slots,
580 state_content_bytes=specs[0].state_content_bytes,
430 - indexes_kv_by_block_stride=block_stride_set.pop(),
581 cache_dtype_str=cache_dtype_str_set.pop(),
432 - compress_ratio=compress_ratio_set.pop(),
582 + tokens_per_state=tokens_per_state_set.pop(),
583 model_version=model_version_set.pop(),
584 non_causal_multi_token_decode=any(
585 spec.non_causal_multi_token_decode for spec in specs
@@ -483,9 +633,9 @@ class RSWASpec(FullAttentionSpec):
633 dtype=base.dtype,
634 kv_quant_mode=base.kv_quant_mode,
635 page_size_padded=base.page_size_padded,
486 - indexes_kv_by_block_stride=base.indexes_kv_by_block_stride,
636 num_head_slots=base.num_head_slots,
637 state_content_bytes=base.state_content_bytes,
638 + tokens_per_state=base.tokens_per_state,
639 sliding_window=base.sliding_window,
640 attention_chunk_size=base.attention_chunk_size,
641 non_causal=base.non_causal,
@@ -598,8 +748,8 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
748 cache_dtype_str: str | None = None
749 # DeepseekV4-only: see MLAAttentionSpec.model_version.
750 alignment: int | None = None # Default to None for no padding.
601 - compress_ratio: int = 1
751 model_version: str | None = None
752 +
753 # MLA stores a single latent vector per state; there is no separate V.
754 head_size_v: int = 0
755
@@ -610,10 +760,6 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
760 super().__post_init__()
761 _apply_alignment_padding(self)
762
613 - @property
614 - def storage_block_size(self) -> int:
615 - return self.block_size // self.compress_ratio
616 -
763 @classmethod
764 def merge(cls, specs: list[Self]) -> Self:
765 assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), (
@@ -621,22 +767,20 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
767 "SlidingWindowMLASpec."
768 )
769 cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
624 - compress_ratio_set = set(spec.compress_ratio for spec in specs)
770 + tokens_per_state_set = set(spec.tokens_per_state for spec in specs)
771 model_version_set = set(spec.model_version for spec in specs)
772 sliding_window_set = set(spec.sliding_window for spec in specs)
627 - block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
773 extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
774 assert (
775 len(cache_dtype_str_set) == 1
631 - and len(compress_ratio_set) == 1
776 + and len(tokens_per_state_set) == 1
777 and len(model_version_set) == 1
778 and len(sliding_window_set) == 1
634 - and len(block_stride_set) == 1
779 and len(extra_retained_set) == 1
780 ), (
781 "All attention layers in the same KV cache group must use the same "
638 - "quantization method, compress ratio, model version, sliding "
639 - "window size, KV block stride indexing, and retained token count."
782 + "quantization method, tokens per state, model version, sliding "
783 + "window size, and retained token count."
784 )
785 return cls(
786 block_size=specs[0].block_size,
@@ -646,11 +790,10 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
790 page_size_padded=specs[0].page_size_padded,
791 num_head_slots=specs[0].num_head_slots,
792 state_content_bytes=specs[0].state_content_bytes,
649 - indexes_kv_by_block_stride=block_stride_set.pop(),
793 sliding_window=sliding_window_set.pop(),
794 extra_retained_tokens=extra_retained_set.pop(),
795 cache_dtype_str=cache_dtype_str_set.pop(),
653 - compress_ratio=compress_ratio_set.pop(),
796 + tokens_per_state=tokens_per_state_set.pop(),
797 model_version=model_version_set.pop(),
798 )
799
@@ -672,6 +815,15 @@ class MambaSpec(KVCacheSpec):
815 mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2
816 mamba_cache_mode: str = "none"
817 num_speculative_blocks: int = 0
818 + num_heads: int = 1
819 + tokens_per_state: int = -1
820 +
821 + @property
822 + def state_content_size_bytes(self) -> int:
823 + return sum(
824 + prod(shape) * get_dtype_size(dtype)
825 + for (shape, dtype) in zip(self.shapes, self.dtypes)
826 + )
827
828 @property
829 def page_size_bytes(self) -> int:
@@ -771,7 +923,6 @@ class SinkFullAttentionSpec(FullAttentionSpec):
923 dtype=specs[0].dtype,
924 kv_quant_mode=specs[0].kv_quant_mode,
925 page_size_padded=specs[0].page_size_padded,
774 - indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride,
926 num_head_slots=specs[0].num_head_slots,
927 state_content_bytes=specs[0].state_content_bytes,
928 sliding_window=cls.merge_window_sizes(sliding_window),
@@ -857,9 +1008,6 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1008 return None
1009
1010 # NOTE: below util functions are only used by DeepseekV4 for now.
860 - def get_page_sizes(self) -> list[int]:
861 - return list(set(spec.page_size_bytes for spec in self.kv_cache_specs.values()))
862 -
1011 def get_num_layer_tuples(self) -> int:
1012 return Counter(
1013 spec.page_size_bytes for spec in self.kv_cache_specs.values()
@@ -929,12 +1077,24 @@ def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None:
1077 class KVCacheTensor:
1078 """
1079 A class for specifying how the workers should initialize the KV cache.
1080 +
1081 + Placement of a set of same-shaped layers in the KV cache allocation.
1082 + Layer ``layers[l]``'s page for block ``b`` starts at
1083 + ``offset + l * layer_stride + b * block_stride`` bytes into the backing
1084 + allocation of ``size`` bytes. Layer-outermost layouts give each layer a
1085 + contiguous region (``layer_stride = page * num_blocks``,
1086 + ``block_stride = page``); block-outermost layouts make each block a
1087 + block of all layers' pages (``layer_stride = page``, ``block_stride`` =
1088 + the packed block). Tensors whose address ranges overlap
1089 + alias the same bytes: cache groups overlay each other, which is sound
1090 + because a block ID is owned by one group at a time.
1091 """
1092
934 - size: int # size of the KV cache tensor in bytes
935 - shared_by: list[str] # layer names that share the same KV cache tensor
936 - offset: int = 0 # byte offset of this layer within a contiguous block
937 - block_stride: int = 0 # total bytes per block in a packed layout (0 = not packed)
1093 + size: int # total size of the backing allocation in bytes
1094 + layers: list[str] # layer names in L order
1095 + layer_stride: int
1096 + block_stride: int
1097 + offset: int = 0 # byte offset of layers[0]'s block 0
1098
1099
1100 @dataclass
@@ -972,6 +1132,8 @@ class KVCacheConfig:
1132 """
1133 prefix_cache_retention_interval: int | None = None
1134 """Resolved retention policy for local prefix-cache checkpoints."""
1135 + kv_cache_layout: str | None = None
1136 + """The KV cache layout resolved by the engine core, adopted by all workers."""
1137
1138 @property
1139 def has_mamba_layers(self) -> bool:
vllm/v1/kv_cache_layout.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Physical KV cache layout descriptor.
4 +
5 +A leaf module so ``vllm.config`` can import the enum without pulling in the
6 +full KV cache interface.
7 +"""
8 +
9 +from enum import Enum
10 +
11 +# Logical dim indices in the 5D stride permutation [L, B, H, N, C] (see: RFC #42082).
12 +_DIM_L, _DIM_B, _DIM_H, _DIM_N, _DIM_C = 0, 1, 2, 3, 4
13 +
14 +
15 +class KVCacheLayout(Enum):
16 + """Physical layout descriptor for a KV cache group.
17 +
18 + The logical shape is always [L, B, H, N, <content>] (RFC #42082).
19 + Each member's value is a stride permutation that maps logical axes
20 + to physical (memory) order.
21 + """
22 +
23 + LBHNC = (0, 1, 2, 3, 4) # [L, B, H, N, C] (identity)
24 + LBNHC = (0, 1, 3, 2, 4) # [L, B, N, H, C]
25 + LHBNC = (0, 2, 1, 3, 4) # [L, H, B, N, C]
26 + BLHNC = (1, 0, 2, 3, 4) # [B, L, H, N, C]
27 + BLNHC = (1, 0, 3, 2, 4) # [B, L, N, H, C]
28 + BHLNC = (1, 2, 0, 3, 4) # [B, H, L, N, C]
29 +
30 + @property
31 + def stride_order(self) -> tuple[int, ...]:
32 + return self.value
33 +
34 + @property
35 + def layer_view_order(self) -> tuple[int, ...]:
36 + """Physical axis order of a logical 4D per-layer cache view."""
37 + return tuple(i - 1 for i in self.value if i != _DIM_L)
38 +
39 + @property
40 + def is_layer_compact(self) -> bool:
41 + """True when the layer is compact; i.e. the L dimension is outermost."""
42 + return self.value[_DIM_L] == 0
43 +
44 + @property
45 + def is_block_contiguous(self) -> bool:
46 + """True when [H, N, C] is contiguous within a block."""
47 + return self.value[-3:] == (_DIM_H, _DIM_N, _DIM_C)
48 +
49 + @property
50 + def is_block_compact(self) -> bool:
51 + """True when each page's [H, N, C] bytes form one contiguous run; i.e.
52 + the L and B dimensions are outermost."""
53 + return set(self.value[:2]) == {_DIM_L, _DIM_B}
54 +
55 + @property
56 + def is_block_outermost(self) -> bool:
57 + """True when B is the outermost physical dimension."""
58 + return self.value[0] == _DIM_B
vllm/v1/kv_offload/base.py
+2 -4
@@ -445,10 +445,8 @@ class CanonicalKVCacheTensor:
445 """
446 A canonicalized KV cache tensor whose first dimension is num_blocks.
447
448 - For attention backends where the raw tensor has num_blocks at a
449 - non-leading physical dimension (e.g. FlashAttention's
450 - (2, num_blocks, ...) layout), the tensor is split so that each
451 - resulting CanonicalKVCacheTensor starts with (num_blocks, ...).
448 + With standardized layouts (RFC #42082) num_blocks is always the leading
449 + logical dimension.
450 """
451
452 # The KV cache tensor with shape (num_blocks, ...)
vllm/v1/kv_offload/config.py
+2
@@ -84,3 +84,5 @@ class OffloadingConfig:
84 # True when the canonical per-layer host byte layout was requested via
85 # kv_connector_extra_config; certified per-layer at worker registration.
86 canonical_layout: bool = False
87 + # Resolved KVCacheLayout name of the worker KV cache.
88 + kv_cache_layout: str | None = None
vllm/v1/kv_offload/file_mapper.py
+4 -1
@@ -94,7 +94,10 @@ class FileMapper:
94 for group in config.groups
95 ]
96 parallel = config.parallel
97 - canonical_format = canonical_format_id() if config.canonical_layout else None
97 + canonical_format = None
98 + if config.canonical_layout:
99 + assert config.kv_cache_layout is not None
100 + canonical_format = canonical_format_id(config.kv_cache_layout)
101 return cls(
102 root_dir=root_dir,
103 model_name=config.model.name,
vllm/v1/simple_kv_offload/manager.py
+6 -9
@@ -194,22 +194,19 @@ class SimpleCPUOffloadScheduler:
194
195 assert len(gpu_config.kv_cache_tensors) > 0
196
197 - is_packed = any(t.block_stride for t in gpu_config.kv_cache_tensors)
198 - assert not is_packed or all(t.block_stride for t in gpu_config.kv_cache_tensors)
199 - gpu_total_bytes = (
200 - gpu_config.kv_cache_tensors[0].size
201 - if is_packed
202 - else sum(t.size for t in gpu_config.kv_cache_tensors)
203 - )
197 + # Every KVCacheTensor describes placement within the same backing allocation,
198 + # so its size is the total GPU KV cache size.
199 + gpu_total_bytes = gpu_config.kv_cache_tensors[0].size
200 num_gpu_blocks = gpu_config.num_blocks
201 num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes)
202 # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally.
203 cpu_tensors = [
204 KVCacheTensor(
205 size=t.size // num_gpu_blocks * num_cpu_blocks,
210 - shared_by=list(t.shared_by),
211 - offset=t.offset,
206 + layers=list(t.layers),
207 + layer_stride=t.layer_stride,
208 block_stride=t.block_stride,
209 + offset=t.offset,
210 )
211 for t in gpu_config.kv_cache_tensors
212 ]
vllm/v1/simple_kv_offload/worker.py
+20 -44
@@ -95,57 +95,33 @@ class SimpleCPUOffloadWorker:
95 logger.warning("No KV caches to offload.")
96 return
97
98 - # Resolve each entry to a representative tensor for storage
99 - # deduplication. For attention layers the value is already a tensor;
100 - # for Mamba layers it is a list of tensors that all share the same
101 - # underlying raw storage, so we take the first one.
102 - def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
103 - assert isinstance(v, torch.Tensor | list)
104 - return v if isinstance(v, torch.Tensor) else v[0]
105 -
106 - any_tensor = _repr_tensor(next(iter(kv_caches.values())))
107 - self.device = any_tensor.device
98 + self.device = next(iter(kv_caches.values())).device
99
100 assert self.kv_cache_config is not None
101 num_blocks = self.kv_cache_config.num_blocks
102
112 - # Deduplicate: multiple layers may share the same backing storage.
113 - seen_ptrs: dict[int, tuple[str, torch.Tensor]] = {}
114 - for name, value in kv_caches.items():
115 - tensor = _repr_tensor(value)
116 - ptr = tensor.untyped_storage().data_ptr()
117 - if ptr not in seen_ptrs:
118 - seen_ptrs[ptr] = (name, tensor)
119 -
120 - # Build [num_blocks, block_bytes] int8 views from each unique
121 - # storage so that stride(0) gives block_bytes for the copy op.
122 - #
123 - # The physical layout varies across attention backends:
124 - # FlashAttn/ROCm: (2, num_blocks, ...) -> K/V outermost, 2 segments
125 - # FlashInfer/MLA: (num_blocks, ...) -> blocks outermost, 1 segment
126 - # We derive page_size_bytes = storage.nbytes() // num_blocks, then
127 - # classify dims: any dim whose byte-stride exceeds page_size_bytes
128 - # must be an outer segment dim (e.g. the K/V dim of size 2). A less
129 - # hacky way is to update the interface with the layout.
103 + # The DMA backend copies whole blocks as base + block_id * stride(0),
104 + # so view each unique allocation as [num_blocks, block_bytes].
105 unique_gpu_caches: dict[str, torch.Tensor] = {}
131 - for name, tensor in seen_ptrs.values():
106 + seen: set[tuple[torch.device, int]] = set()
107 + for name, tensor in kv_caches.items():
108 storage = tensor.untyped_storage()
133 - raw = torch.empty(0, dtype=torch.int8, device=self.device).set_(
134 - storage, 0, (storage.nbytes(),)
109 + key = (tensor.device, storage.data_ptr())
110 + if key in seen:
111 + continue
112 + seen.add(key)
113 +
114 + physical_per_block, remainder = divmod(tensor.shape[0], num_blocks)
115 + assert remainder == 0, (
116 + f"KV cache {name!r} has {tensor.shape[0]} physical blocks, which "
117 + f"is not divisible by {num_blocks} scheduler blocks"
118 )
136 - el = tensor.element_size()
137 - page_size_bytes = storage.nbytes() // num_blocks
138 - outer_dims = [
139 - d for d in range(tensor.ndim) if tensor.stride(d) * el > page_size_bytes
140 - ]
141 - if not outer_dims:
142 - unique_gpu_caches[name] = raw.view(num_blocks, -1)
143 - else:
144 - seg_stride = tensor.stride(outer_dims[0]) * el
145 - for idx in range(tensor.shape[outer_dims[0]]):
146 - offset = idx * seg_stride
147 - chunk = raw[offset : offset + seg_stride]
148 - unique_gpu_caches[f"{name}.{idx}"] = chunk.view(num_blocks, -1)
119 + block_bytes = tensor.stride(0) * tensor.element_size() * physical_per_block
120 + raw = torch.empty(0, dtype=torch.int8, device=tensor.device).set_(storage)
121 + regions = raw.view(-1, num_blocks, block_bytes)
122 + for idx, region in enumerate(regions):
123 + key_name = name if len(regions) == 1 else f"{name}.{idx}"
124 + unique_gpu_caches[key_name] = region
125
126 # Compute per-tensor bytes_per_block. Tensors may have different
127 # page_size_bytes (e.g., UniformTypeKVCacheSpecs with varying head_size).
vllm/v1/worker/gpu/attn_utils.py
+12 -380
@@ -1,22 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Iterable, Mapping, Sequence
4 -from dataclasses import dataclass, replace
5 -from math import prod
3 +from collections.abc import Mapping, Sequence
4 +from dataclasses import dataclass
5 from typing import Any, cast
6
7 import torch
8
10 -from vllm.config import (
11 - VllmConfig,
12 - get_layers_from_vllm_config,
13 - set_current_vllm_config,
14 -)
15 -from vllm.logger import init_logger
9 +from vllm.config import VllmConfig, get_layers_from_vllm_config
10 from vllm.model_executor.layers.attention import Attention
11 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
12 from vllm.multimodal.inputs import MultiModalFeatureSpec
19 -from vllm.utils.torch_utils import get_dtype_size
13 from vllm.v1.attention.backend import (
14 AttentionCGSupport,
15 CommonAttentionMetadata,
@@ -25,20 +18,17 @@ from vllm.v1.kv_cache_interface import (
18 AttentionSpec,
19 KVCacheConfig,
20 KVCacheSpec,
28 - KVQuantMode,
29 - MambaSpec,
21 UniformTypeKVCacheSpecs,
22 )
23 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
24 from vllm.v1.worker.utils import (
25 AttentionGroup,
26 add_kv_sharing_layers_to_kv_cache_groups,
27 + allocate_kv_cache,
28 bind_kv_cache,
29 prepare_kernel_block_sizes,
30 )
31
40 -logger = init_logger(__name__)
41 -
32
33 @dataclass(frozen=True)
34 class AttentionCGSupportInfo:
@@ -69,13 +59,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
59 # Skip modules that don't need KV cache (eg encoder-only attention)
60 if spec := attn_module.get_kv_cache_spec(vllm_config):
61 if isinstance(spec, AttentionSpec):
72 - backend = attn_module.get_attn_backend()
73 - # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() ->
74 - # get_kv_cache_layout() needs the current vLLM config.
75 - with set_current_vllm_config(vllm_config):
76 - indexes = backend.indexes_kv_by_block_stride()
77 - spec = replace(spec, indexes_kv_by_block_stride=indexes)
78 - spec = backend.customize_spec(spec)
62 + spec = attn_module.get_attn_backend().customize_spec(spec)
63 kv_cache_spec[layer_name] = spec
64 return kv_cache_spec
65
@@ -197,374 +181,22 @@ def get_query_lens_mismatch_unsupported_backend(
181 return None
182
183
200 -def _allocate_kv_cache(
201 - kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device
202 -):
203 - kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
204 - packed_backing: torch.Tensor | None = None
205 - for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
206 - if kv_cache_tensor.block_stride > 0:
207 - # Allocate once; all packed tensors alias the same backing.
208 - if packed_backing is None:
209 - packed_backing = torch.zeros(
210 - kv_cache_tensor.size, dtype=torch.int8, device=device
211 - )
212 - tensor = packed_backing
213 - else:
214 - tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)
215 - for layer_name in kv_cache_tensor.shared_by:
216 - kv_cache_raw_tensors[layer_name] = tensor
217 -
218 - layer_names = set()
219 - for group in kv_cache_config.kv_cache_groups:
220 - for layer_name in group.layer_names:
221 - layer_names.add(layer_name)
222 - assert layer_names == (kv_cache_raw_tensors.keys() | shared_layers.keys()), (
223 - "Some layers are not correctly initialized"
224 - )
225 - return kv_cache_raw_tensors
226 -
227 -
228 -def _kv_first_layers_sharing_pool_with_mamba(
229 - attn_groups: Sequence[AttentionGroup],
230 - kernel_block_sizes: list[int],
231 - cache_dtype: str,
232 - kv_cache_config: "KVCacheConfig | None",
233 -) -> set[str]:
234 - """KV-first attention layers that share an allocation with Mamba state.
235 -
236 - Mamba addresses its state by page, so these layers have to lay their blocks
237 - out page by page as well; otherwise block ids from the two groups resolve
238 - to overlapping bytes.
239 - """
240 - if kv_cache_config is None:
241 - return set()
242 -
243 - kv_first_layers: set[str] = set()
244 - mamba_layers: set[str] = set()
245 - for group in attn_groups:
246 - kv_cache_spec = group.kv_cache_spec
247 - if isinstance(kv_cache_spec, MambaSpec):
248 - mamba_layers.update(group.layer_names)
249 - continue
250 - if not isinstance(kv_cache_spec, AttentionSpec):
251 - continue
252 - if group.kv_cache_group_id >= len(kernel_block_sizes):
253 - continue
254 - block_dim = group.backend.get_kv_cache_block_dim(
255 - kernel_block_sizes[group.kv_cache_group_id],
256 - kv_cache_spec.num_kv_heads,
257 - kv_cache_spec.head_size,
258 - cache_dtype_str=cache_dtype,
259 - )
260 - if block_dim != 0:
261 - kv_first_layers.update(group.layer_names)
262 -
263 - page_aligned: set[str] = set()
264 - for kv_tensor in kv_cache_config.kv_cache_tensors:
265 - if mamba_layers.isdisjoint(kv_tensor.shared_by):
266 - continue
267 - page_aligned.update(kv_first_layers.intersection(kv_tensor.shared_by))
268 - return page_aligned
269 -
270 -
271 -def _reshape_attention_kv_cache(
272 - kv_raw_tensor: torch.Tensor,
273 - kv_cache_spec: AttentionSpec,
274 - kv_cache_shape: tuple[int, ...],
275 - kv_cache_stride_order: tuple[int, ...],
276 - num_blocks: int,
277 - packing: tuple[int, int] | None,
278 - page_aligned_blocks: bool = False,
279 -) -> torch.Tensor:
280 - permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
281 - inv_order = [
282 - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
283 - ]
284 - dtype = kv_cache_spec.dtype
285 -
286 - if packing is not None:
287 - offset, block_stride = packing
288 - assert inv_order[0] == 0
289 - page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
290 - kv_cache = (
291 - kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes]
292 - .view(dtype)
293 - .view(permuted_kv_cache_shape)
294 - )
295 - elif kv_cache_spec.page_size_padded is not None:
296 - # Use a strided view to skip the padding between physical pages.
297 - #
298 - # Only num-blocks-first layouts are supported (the block dimension is
299 - # dim 0 of the unpermuted shape). kv-first layouts such as ROCm's
300 - # ``(2, num_blocks, ...)`` are intentionally not supported here. For a
301 - # num-blocks-first layout the only stride that must change is the block
302 - # stride: every other (contiguous) stride already steps within the
303 - # unpadded region of a page, so no further adjustment is needed.
304 - assert kv_cache_shape[0] == num_blocks, (
305 - "Padded KV pages require a num-blocks-first KV cache layout (got "
306 - f"shape {kv_cache_shape} with num_blocks={num_blocks}); "
307 - "kv-first layouts are not supported."
308 - )
309 - dtype_size = get_dtype_size(kv_cache_spec.dtype)
310 - page_stride = kv_cache_spec.page_size_bytes // dtype_size
311 -
312 - num_blocks_dim = inv_order[0]
313 - strides = list(torch.empty(permuted_kv_cache_shape, device="meta").stride())
314 - strides[num_blocks_dim] = page_stride
315 -
316 - kv_cache = torch.as_strided(
317 - kv_raw_tensor.view(dtype),
318 - size=permuted_kv_cache_shape,
319 - stride=tuple(strides),
320 - )
321 - elif page_aligned_blocks:
322 - # A KV-first layout such as ROCm's ``(2, num_blocks, ...)`` puts block
323 - # ``b``'s K and V in two far-apart halves of the allocation, so block
324 - # ``b`` does not cover page ``b``. Mamba layers sharing the allocation
325 - # do address their state by page, so the two would resolve the same
326 - # bytes. Build the view page-first instead, then swap the dims back.
327 - assert kv_cache_shape[1] == num_blocks and kv_cache_stride_order == tuple(
328 - range(len(kv_cache_shape))
329 - ), (
330 - "Page-aligned KV blocks expect a default-strided (kv, num_blocks, "
331 - f"...) layout, got shape {kv_cache_shape} with stride order "
332 - f"{kv_cache_stride_order} and num_blocks={num_blocks}."
333 - )
334 - kv_cache = (
335 - kv_raw_tensor.view(dtype)
336 - .view(num_blocks, kv_cache_shape[0], *kv_cache_shape[2:])
337 - .transpose(0, 1)
338 - )
339 - else:
340 - # No padding — safe to use a contiguous view.
341 - kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape)
342 -
343 - return kv_cache.permute(*inv_order)
344 -
345 -
346 -def _reshape_kv_cache(
347 - attn_groups: Sequence[AttentionGroup],
348 - kv_cache_raw_tensors: dict[str, torch.Tensor],
349 - cache_dtype: str,
350 - kernel_block_sizes: list[int],
351 - shared_kv_cache_layers: dict[str, str],
352 - kv_cache_config: "KVCacheConfig | None" = None,
353 -) -> dict[str, Any]:
354 - kv_caches: dict[str, Any] = {}
355 - has_attn = False
356 -
357 - layer_packing: dict[str, tuple[int, int]] = {}
358 - if kv_cache_config is not None:
359 - for kv_tensor in kv_cache_config.kv_cache_tensors:
360 - if kv_tensor.block_stride > 0:
361 - for ln in kv_tensor.shared_by:
362 - layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride)
363 -
364 - page_aligned_layers = _kv_first_layers_sharing_pool_with_mamba(
365 - attn_groups, kernel_block_sizes, cache_dtype, kv_cache_config
366 - )
367 -
368 - for group in attn_groups:
369 - if group.kv_cache_group_id >= len(kernel_block_sizes):
370 - continue
371 -
372 - kv_cache_spec = group.kv_cache_spec
373 - if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
374 - # use storage_block_size as the kernel block size for groups
375 - # that apply a compression on block size (eg. DeepSeek V4).
376 - kernel_block_size = kv_cache_spec.storage_block_size
377 - else:
378 - kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
379 -
380 - for layer_name in group.layer_names:
381 - if layer_name in shared_kv_cache_layers:
382 - # Shared layer — tensor will be aliased to its target later.
383 - continue
384 -
385 - kv_raw_tensor = kv_cache_raw_tensors[layer_name]
386 - packing = layer_packing.get(layer_name)
387 - if packing is not None:
388 - _, blk_stride = packing
389 - num_blocks = kv_raw_tensor.numel() // blk_stride
390 - else:
391 - assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
392 - num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes
393 -
394 - if isinstance(kv_cache_spec, AttentionSpec):
395 - has_attn = True
396 - # Use storage_block_size: it equals block_size for uncompressed
397 - # specs but is smaller for compressed ones (DeepSeek V4), which
398 - # store block_size tokens in block_size // compress_ratio slots.
399 - num_blocks_per_kv_block = (
400 - kv_cache_spec.storage_block_size // kernel_block_size
401 - )
402 - kernel_num_blocks = num_blocks * num_blocks_per_kv_block
403 - # Skipped layers (--kv-cache-dtype-skip-layers) keep the
404 - # unquantized shape; only the quantized primary uses the
405 - # quantized cache dtype's (possibly packed) layout.
406 - layer_cache_dtype = (
407 - "auto"
408 - if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
409 - else cache_dtype
410 - )
411 - kv_cache_shape = group.backend.get_kv_cache_shape(
412 - kernel_num_blocks,
413 - kernel_block_size,
414 - kv_cache_spec.num_kv_heads,
415 - kv_cache_spec.head_size,
416 - cache_dtype_str=layer_cache_dtype,
417 - )
418 -
419 - # FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
420 - try:
421 - kv_cache_stride_order = group.backend.get_kv_cache_stride_order()
422 - assert len(kv_cache_stride_order) == len(kv_cache_shape)
423 - except (AttributeError, NotImplementedError):
424 - kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
425 -
426 - kv_caches[layer_name] = _reshape_attention_kv_cache(
427 - kv_raw_tensor,
428 - kv_cache_spec,
429 - kv_cache_shape,
430 - kv_cache_stride_order,
431 - kernel_num_blocks,
432 - packing,
433 - page_aligned_blocks=layer_name in page_aligned_layers,
434 - )
435 -
436 - elif isinstance(kv_cache_spec, MambaSpec):
437 - page_size_bytes = kv_cache_spec.page_size_bytes
438 - # Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]
439 - # int8 page view per layer; the layer's bind_kv_cache unpacks
440 - # each block's bytes into its conv/ssm state views. Keeping
441 - # one tensor per layer lets the KV connector register it
442 - # without special-casing Mamba.
443 - kv_caches[layer_name] = kv_raw_tensor[
444 - : num_blocks * page_size_bytes
445 - ].view(num_blocks, 1, 1, page_size_bytes)
446 - else:
447 - raise NotImplementedError(
448 - f"Unsupported KV cache spec type: {type(kv_cache_spec)}"
449 - )
450 -
451 - if has_attn and kv_cache_config is not None:
452 - _align_mixed_attention_kv_cache_views(
453 - attn_groups=attn_groups,
454 - kv_caches=kv_caches,
455 - kernel_block_sizes=kernel_block_sizes,
456 - cache_dtype=cache_dtype,
457 - kv_cache_config=kv_cache_config,
458 - page_aligned_layers=page_aligned_layers,
459 - )
460 -
461 - # Map any sharing layers to their target layer's KV cache.
462 - for layer_name, target_layer_name in shared_kv_cache_layers.items():
463 - kv_caches[layer_name] = kv_caches[target_layer_name]
464 -
465 - return kv_caches
466 -
467 -
468 -def _align_mixed_attention_kv_cache_views(
469 - attn_groups: Iterable[AttentionGroup],
470 - kv_caches: dict[str, Any],
471 - kernel_block_sizes: list[int],
472 - cache_dtype: str,
473 - kv_cache_config: KVCacheConfig,
474 - page_aligned_layers: set[str],
475 -) -> None:
476 - """Align shared attention KV views when backends disagree on layout.
477 -
478 - Encoder-decoder models can share one raw allocation between decoder
479 - self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention
480 - (blocks-first backends, block dim 0). Keep the physical storage in the
481 - K/V-first layout expected by ROCM_ATTN, and restride the blocks-first
482 - logical views so block IDs address the same bytes.
483 - """
484 - block_dims_by_layer: dict[str, int] = {}
485 - for group in attn_groups:
486 - kv_cache_spec = group.kv_cache_spec
487 - if not isinstance(kv_cache_spec, AttentionSpec):
488 - continue
489 - if group.kv_cache_group_id >= len(kernel_block_sizes):
490 - continue
491 - block_dim = group.backend.get_kv_cache_block_dim(
492 - kernel_block_sizes[group.kv_cache_group_id],
493 - kv_cache_spec.num_kv_heads,
494 - kv_cache_spec.head_size,
495 - cache_dtype_str=cache_dtype,
496 - )
497 - for layer_name in group.layer_names:
498 - if layer_name in kv_caches:
499 - block_dims_by_layer[layer_name] = block_dim
500 -
501 - for kv_tensor in kv_cache_config.kv_cache_tensors:
502 - if kv_tensor.block_stride > 0:
503 - continue
504 - if not page_aligned_layers.isdisjoint(kv_tensor.shared_by):
505 - # These KV-first views already address blocks by page, just like
506 - # blocks-first views, so they must not be restrided onto the
507 - # unaligned KV-first storage layout.
508 - continue
509 - shared_block_dims = {
510 - block_dims_by_layer[layer_name]
511 - for layer_name in kv_tensor.shared_by
512 - if layer_name in block_dims_by_layer
513 - }
514 - if 0 not in shared_block_dims or 1 not in shared_block_dims:
515 - continue
516 -
517 - for layer_name in kv_tensor.shared_by:
518 - if block_dims_by_layer.get(layer_name) == 0:
519 - _restride_blocks_first_kv_cache_to_kv_first_storage(
520 - kv_caches[layer_name]
521 - )
522 -
523 -
524 -def _restride_blocks_first_kv_cache_to_kv_first_storage(
525 - kv_cache: torch.Tensor,
526 -) -> None:
527 - assert kv_cache.ndim >= 3
528 - assert kv_cache.shape[1] == 2
529 - page_size = kv_cache.shape[2:].numel()
530 - num_blocks = kv_cache.shape[0]
531 - expected_tail_stride = torch.empty(kv_cache.shape[2:]).stride()
532 - if kv_cache.stride()[2:] != expected_tail_stride:
533 - logger.warning_once(
534 - "Skipping mixed KV-cache layout alignment for a non-NHD "
535 - "blocks-first attention view with stride %s.",
536 - kv_cache.stride(),
537 - )
538 - return
539 - kv_cache.as_strided_(
540 - size=kv_cache.shape,
541 - stride=(page_size, num_blocks * page_size, *expected_tail_stride),
542 - )
543 -
544 -
184 def init_kv_cache(
185 runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
186 forward_context: dict[str, Any],
187 kv_cache_config: KVCacheConfig,
549 - attn_groups: list[list[AttentionGroup]],
188 device: torch.device,
551 - cache_dtype: str,
189 kernel_block_sizes: list[int],
190 vllm_config: VllmConfig,
191 ) -> dict[str, Any]:
555 - shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
556 - kv_cache_raw_tensors = _allocate_kv_cache(
557 - kv_cache_config, shared_kv_cache_layers, device
558 - )
559 - flattened_attn_groups = list(group for groups in attn_groups for group in groups)
560 - kv_caches = _reshape_kv_cache(
561 - attn_groups=flattened_attn_groups,
562 - kv_cache_raw_tensors=kv_cache_raw_tensors,
563 - kernel_block_sizes=kernel_block_sizes,
564 - cache_dtype=cache_dtype,
565 - shared_kv_cache_layers=shared_kv_cache_layers,
566 - kv_cache_config=kv_cache_config,
192 + kv_caches = allocate_kv_cache(
193 + kv_cache_config,
194 + device,
195 + vllm_config.cache_config.get_resolved_kv_cache_layout(),
196 + kernel_block_sizes,
197 )
198 + for layer_name, target in get_shared_kv_cache_layers(vllm_config).items():
199 + kv_caches[layer_name] = kv_caches[target]
200 # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per
201 # decoder layer, so a layer name carries two integers (layer + module index).
202 num_attn_module = (
vllm/v1/worker/gpu/kv_connector.py
-2
@@ -51,8 +51,6 @@ class ActiveKVConnector(KVConnector):
51 self.vllm_config = vllm_config
52 self.kv_connector = get_kv_transfer_group()
53 # Register kv caches with KV Connector if applicable.
54 - # TODO: support cross_layers_kv_cache
55 - # (see https://github.com/vllm-project/vllm/pull/27743)
54 self.kv_connector.register_kv_caches(kv_caches_dict)
55 self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks)
56
vllm/v1/worker/gpu/model_runner.py
-3
@@ -618,9 +618,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
618 self.kv_caches,
619 self.compilation_config.static_forward_context,
620 self.kv_cache_config,
621 - self.attn_groups,
621 self.device,
623 - self.cache_config.cache_dtype,
622 self.kernel_block_sizes,
623 self.vllm_config,
624 )
@@ -632,7 +630,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
630 self.device,
631 attn_groups_iter=(g for groups in self.attn_groups for g in groups),
632 kernel_block_sizes=self.kernel_block_sizes,
635 - cache_dtype=self.cache_config.cache_dtype,
633 static_forward_context=self.compilation_config.static_forward_context,
634 )
635
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+6 -6
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import copy
3 from collections.abc import Mapping
4 from typing import Any
5
@@ -104,12 +103,13 @@ class DFlashSpeculator(DraftModelSpeculator):
103 @property
104 def attn_vllm_config(self) -> VllmConfig:
105 # The draft's attention differs from the target's in causality.
107 - config = copy.copy(super().attn_vllm_config)
108 - config.attention_config = replace(
109 - self.vllm_config.attention_config,
110 - use_non_causal=self.requires_non_causal,
106 + return replace(
107 + self.vllm_config,
108 + attention_config=replace(
109 + self.vllm_config.attention_config,
110 + use_non_causal=self.requires_non_causal,
111 + ),
112 )
112 - return config
113
114 def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
115 wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
vllm/v1/worker/gpu_model_runner.py
+12 -268
@@ -42,7 +42,9 @@ from vllm.config.model import PROCESSED_LOGPROBS_MODES
42 from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
43 from vllm.distributed.eplb.eplb_state import EplbState
44 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
45 -from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
45 +from vllm.distributed.kv_transfer.kv_connector.utils import (
46 + copy_kv_blocks,
47 +)
48 from vllm.distributed.parallel_state import (
49 GraphCaptureContext,
50 get_dcp_group,
@@ -164,8 +166,6 @@ from vllm.v1.kv_cache_interface import (
166 KVCacheGroupSpec,
167 KVCacheSpec,
168 KVCacheSpecKind,
167 - KVQuantMode,
168 - MambaSpec,
169 SlidingWindowSpec,
170 UniformTypeKVCacheSpecs,
171 get_kv_cache_spec_kind,
@@ -221,7 +221,6 @@ from vllm.v1.worker.cp_utils import (
221 )
222 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
223 from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
224 -from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache
224 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
225 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
226 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
@@ -243,6 +242,7 @@ from .utils import (
242 AttentionGroup,
243 KVBlockZeroer,
244 add_kv_sharing_layers_to_kv_cache_groups,
245 + allocate_kv_cache,
246 bind_kv_cache,
247 copy_kv_cache_blocks_inplace,
248 prepare_kernel_block_sizes,
@@ -612,9 +612,6 @@ class GPUModelRunner(
612 # self.model: nn.Module # Set after load_model
613 # Initialize in initialize_kv_cache
614 self.kv_caches: list[torch.Tensor] = []
615 - # Initialize in initialize_kv_cache_tensors
616 - self.cross_layers_kv_cache: torch.Tensor | None = None
617 - self.cross_layers_attn_backend: type[AttentionBackend] | None = None
615 # indexes: [kv_cache_group_id][attn_group]
616 self.attn_groups: list[list[AttentionGroup]] = []
617 # self.kv_cache_config: KVCacheConfig
@@ -1200,7 +1197,6 @@ class GPUModelRunner(
1197 self.device,
1198 attn_groups_iter=self._kv_cache_spec_attn_group_iterator(),
1199 kernel_block_sizes=self._kernel_block_sizes,
1203 - cache_dtype=self.cache_config.cache_dtype,
1200 runner_only_attn_layers=self.runner_only_attn_layers,
1201 static_forward_context=self.compilation_config.static_forward_context,
1202 )
@@ -6707,9 +6703,6 @@ class GPUModelRunner(
6703 for i in range(len(self.kv_caches)):
6704 self.kv_caches[i] = None # type: ignore
6705 self.kv_caches.clear()
6710 - if hasattr(self, "cross_layers_kv_cache"):
6711 - self.cross_layers_kv_cache = None
6712 - self.cross_layers_attn_backend = None
6706 if hasattr(self, "attn_groups"):
6707 self.attn_groups.clear()
6708 if hasattr(self, "kv_cache_config"):
@@ -7441,49 +7434,6 @@ class GPUModelRunner(
7434 f"!= kv_cache kernel_block_sizes {kernel_block_sizes}"
7435 )
7436
7444 - def _allocate_kv_cache_tensors(
7445 - self, kv_cache_config: KVCacheConfig
7446 - ) -> dict[str, torch.Tensor]:
7447 - """
7448 - Initializes the KV cache buffer with the correct size. The buffer needs
7449 - to be reshaped to the desired shape before being used by the models.
7450 -
7451 - Args:
7452 - kv_cache_config: The KV cache config
7453 - Returns:
7454 - dict[str, torch.Tensor]: A map between layer names to their
7455 - corresponding memory buffer for KV cache.
7456 - """
7457 - kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
7458 - packed_backing: torch.Tensor | None = None
7459 - for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
7460 - if kv_cache_tensor.block_stride > 0:
7461 - # Allocate once; all packed tensors alias the same backing.
7462 - if packed_backing is None:
7463 - packed_backing = torch.zeros(
7464 - kv_cache_tensor.size,
7465 - dtype=torch.int8,
7466 - device=self.device,
7467 - )
7468 - tensor = packed_backing
7469 - else:
7470 - tensor = torch.zeros(
7471 - kv_cache_tensor.size, dtype=torch.int8, device=self.device
7472 - )
7473 - for layer_name in kv_cache_tensor.shared_by:
7474 - kv_cache_raw_tensors[layer_name] = tensor
7475 -
7476 - layer_names = set()
7477 - for group in kv_cache_config.kv_cache_groups:
7478 - for layer_name in group.layer_names:
7479 - if layer_name in self.runner_only_attn_layers:
7480 - continue
7481 - layer_names.add(layer_name)
7482 - assert layer_names == set(kv_cache_raw_tensors.keys()), (
7483 - "Some layers are not correctly initialized"
7484 - )
7485 - return kv_cache_raw_tensors
7486 -
7437 def _attn_group_iterator(self) -> Iterator[AttentionGroup]:
7438 return itertools.chain.from_iterable(self.attn_groups)
7439
@@ -7493,183 +7443,6 @@ class GPUModelRunner(
7443 for attn_groups in self.attn_groups:
7444 yield from attn_groups
7445
7496 - def _reshape_kv_cache_tensors(
7497 - self,
7498 - kv_cache_raw_tensors: dict[str, torch.Tensor],
7499 - kernel_block_sizes: list[int],
7500 - ) -> dict[str, torch.Tensor]:
7501 - """
7502 - Reshape the KV cache tensors to the desired shape and dtype.
7503 -
7504 - Args:
7505 - kv_cache_raw_tensors: The KV cache buffer of each layer, with
7506 - correct size but uninitialized shape.
7507 - kernel_block_sizes: The kernel block sizes for each KV cache group.
7508 - Returns:
7509 - Dict[str, torch.Tensor]: A map between layer names to their
7510 - corresponding memory buffer for KV cache.
7511 - """
7512 - kv_caches: dict[str, torch.Tensor] = {}
7513 - has_attn, has_mamba = False, False
7514 -
7515 - # Map layer names to (offset, block_stride) within the packed
7516 - # backing tensor so we can create strided views per layer.
7517 - layer_packing: dict[str, tuple[int, int]] = {}
7518 - for kv_tensor in self.kv_cache_config.kv_cache_tensors:
7519 - if kv_tensor.block_stride > 0:
7520 - for ln in kv_tensor.shared_by:
7521 - layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride)
7522 - for group in self._kv_cache_spec_attn_group_iterator():
7523 - kv_cache_spec = group.kv_cache_spec
7524 - attn_backend = group.backend
7525 - if group.kv_cache_group_id == len(kernel_block_sizes):
7526 - # There may be a last group for layers without kv cache.
7527 - continue
7528 - kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
7529 - for layer_name in group.layer_names:
7530 - if layer_name in self.runner_only_attn_layers:
7531 - continue
7532 - raw_tensor = kv_cache_raw_tensors[layer_name]
7533 - packing = layer_packing.get(layer_name)
7534 - if packing is not None:
7535 - _, blk_stride = packing
7536 - num_blocks = raw_tensor.numel() // blk_stride
7537 - else:
7538 - assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
7539 - num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
7540 - if isinstance(kv_cache_spec, AttentionSpec):
7541 - has_attn = True
7542 - num_blocks_per_kv_block = (
7543 - kv_cache_spec.block_size // kernel_block_size
7544 - )
7545 - kernel_num_blocks = num_blocks * num_blocks_per_kv_block
7546 -
7547 - # For MLA with compression, storage_block_size != block_size
7548 - if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
7549 - shape_block_size = kv_cache_spec.storage_block_size
7550 - else:
7551 - shape_block_size = kernel_block_size
7552 -
7553 - # Skipped layers (--kv-cache-dtype-skip-layers) need
7554 - # the unquantized shape.
7555 - layer_cache_dtype_str = (
7556 - "auto"
7557 - if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
7558 - else getattr(
7559 - kv_cache_spec,
7560 - "cache_dtype_str",
7561 - None,
7562 - )
7563 - or self.cache_config.cache_dtype
7564 - )
7565 - kv_cache_shape = attn_backend.get_kv_cache_shape(
7566 - kernel_num_blocks,
7567 - shape_block_size,
7568 - kv_cache_spec.num_kv_heads,
7569 - kv_cache_spec.head_size,
7570 - cache_dtype_str=layer_cache_dtype_str,
7571 - )
7572 - try:
7573 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
7574 - assert len(kv_cache_stride_order) == len(kv_cache_shape)
7575 - except (AttributeError, NotImplementedError):
7576 - kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
7577 - raw_tensor = kv_cache_raw_tensors[layer_name]
7578 - kv_caches[layer_name] = _reshape_attention_kv_cache(
7579 - raw_tensor,
7580 - kv_cache_spec,
7581 - kv_cache_shape,
7582 - kv_cache_stride_order,
7583 - kernel_num_blocks,
7584 - packing,
7585 - )
7586 -
7587 - elif isinstance(kv_cache_spec, MambaSpec):
7588 - has_mamba = True
7589 - raw_tensor = kv_cache_raw_tensors[layer_name]
7590 - page_size_bytes = kv_cache_spec.page_size_bytes
7591 - # Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]
7592 - # int8 page view per layer; the layer's bind_kv_cache unpacks
7593 - # each block's bytes into its conv/ssm state views. Keeping
7594 - # one tensor per layer lets the KV connector register it
7595 - # without special-casing Mamba.
7596 - kv_caches[layer_name] = raw_tensor[
7597 - : num_blocks * page_size_bytes
7598 - ].view(num_blocks, 1, 1, page_size_bytes)
7599 - else:
7600 - raise NotImplementedError
7601 -
7602 - # Reconcile divergent KV layouts to blocks-first. Triggered by hybrid
7603 - # attention/mamba models, and by encoder-decoder models whose shared
7604 - # decoder/cross-attention allocation mixes K/V-first and blocks-first
7605 - # backends (see _has_mixed_attention_kv_layout).
7606 - if has_attn and (
7607 - has_mamba or self._has_mixed_attention_kv_layout(kernel_block_sizes)
7608 - ):
7609 - self._update_hybrid_attention_mamba_layout(kv_caches, kernel_block_sizes)
7610 -
7611 - return kv_caches
7612 -
7613 - def _has_mixed_attention_kv_layout(self, kernel_block_sizes: list[int]) -> bool:
7614 - """Whether attention groups disagree on the physical KV cache layout.
7615 -
7616 - Encoder-decoder models (e.g. Whisper) share one raw KV allocation
7617 - between a decoder self-attention layer (K/V-first ROCM_ATTN, block dim
7618 - 1) and a cross-attention layer (blocks-first, block dim 0). Mixed block
7619 - dims mean a block ID maps to different bytes per layer, so the shared
7620 - buffer must be normalized to a single (blocks-first) layout.
7621 - """
7622 - block_dims: set[int] = set()
7623 - for group in self._kv_cache_spec_attn_group_iterator():
7624 - kv_cache_spec = group.kv_cache_spec
7625 - if not isinstance(kv_cache_spec, AttentionSpec):
7626 - continue
7627 - if group.kv_cache_group_id == len(kernel_block_sizes):
7628 - continue
7629 - block_dims.add(
7630 - group.backend.get_kv_cache_block_dim(
7631 - kernel_block_sizes[group.kv_cache_group_id],
7632 - kv_cache_spec.num_kv_heads,
7633 - kv_cache_spec.head_size,
7634 - cache_dtype_str=self.cache_config.cache_dtype,
7635 - )
7636 - )
7637 - return len(block_dims) > 1
7638 -
7639 - def _update_hybrid_attention_mamba_layout(
7640 - self, kv_caches: dict[str, torch.Tensor], kernel_block_sizes: list[int]
7641 - ) -> None:
7642 - """
7643 - Update the layout of attention layers from (2, num_blocks, ...) to
7644 - (num_blocks, 2, ...).
7645 -
7646 - Args:
7647 - kv_caches: The KV cache buffer of each layer.
7648 - kernel_block_sizes: The kernel block sizes for each KV cache group.
7649 - """
7650 -
7651 - for group in self._kv_cache_spec_attn_group_iterator():
7652 - kv_cache_spec = group.kv_cache_spec
7653 - if not isinstance(kv_cache_spec, AttentionSpec):
7654 - continue
7655 - block_dim = group.backend.get_kv_cache_block_dim(
7656 - kernel_block_sizes[group.kv_cache_group_id],
7657 - kv_cache_spec.num_kv_heads,
7658 - kv_cache_spec.head_size,
7659 - cache_dtype_str=self.cache_config.cache_dtype,
7660 - )
7661 - # block_dim: 0 means (num_blocks, 2, ...); 1 means (2, num_blocks, ...).
7662 - if block_dim == 0:
7663 - continue
7664 - assert block_dim == 1
7665 - for layer_name in group.layer_names:
7666 - kv_cache = kv_caches[layer_name]
7667 - hidden_size = kv_cache.shape[2:].numel()
7668 - kv_cache.as_strided_(
7669 - size=kv_cache.shape,
7670 - stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]),
7671 - )
7672 -
7446 def initialize_kv_cache_tensors(
7447 self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
7448 ) -> dict[str, torch.Tensor]:
@@ -7685,29 +7458,12 @@ class GPUModelRunner(
7458 corresponding memory buffer for KV cache.
7459 """
7460
7688 - # Try creating KV caches optimized for kv-connector transfers
7689 - cache_dtype = self.cache_config.cache_dtype
7690 - if self.use_uniform_kv_cache(self.attn_groups):
7691 - kv_caches, cross_layers_kv_cache, attn_backend = (
7692 - self.allocate_uniform_kv_caches(
7693 - kv_cache_config,
7694 - self.attn_groups,
7695 - cache_dtype,
7696 - self.device,
7697 - kernel_block_sizes,
7698 - )
7699 - )
7700 - self.cross_layers_kv_cache = cross_layers_kv_cache
7701 - self.cross_layers_attn_backend = attn_backend
7702 - else:
7703 - # Fallback to the general case
7704 - # Initialize the memory buffer for KV cache
7705 - kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
7706 -
7707 - # Change the memory buffer to the desired shape
7708 - kv_caches = self._reshape_kv_cache_tensors(
7709 - kv_cache_raw_tensors, kernel_block_sizes
7710 - )
7461 + kv_caches = allocate_kv_cache(
7462 + kv_cache_config,
7463 + self.device,
7464 + self.cache_config.get_resolved_kv_cache_layout(),
7465 + kernel_block_sizes,
7466 + )
7467
7468 # Set up cross-layer KV cache sharing
7469 for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
@@ -7803,13 +7559,7 @@ class GPUModelRunner(
7559
7560 if has_kv_transfer_group() and not is_profiling:
7561 kv_transfer_group = get_kv_transfer_group()
7806 - if self.cross_layers_kv_cache is not None:
7807 - assert self.cross_layers_attn_backend is not None
7808 - kv_transfer_group.register_cross_layers_kv_cache(
7809 - self.cross_layers_kv_cache, self.cross_layers_attn_backend
7810 - )
7811 - else:
7812 - kv_transfer_group.register_kv_caches(kv_caches)
7562 + kv_transfer_group.register_kv_caches(kv_caches)
7563 kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
7564
7565 def get_routed_experts(
@@ -7922,13 +7672,7 @@ class GPUModelRunner(
7672 # Skip modules that don't need KV cache (eg encoder-only attention)
7673 if spec := attn_module.get_kv_cache_spec(self.vllm_config):
7674 if isinstance(spec, AttentionSpec):
7925 - backend = attn_module.get_attn_backend()
7926 - # indexes_kv_by_block_stride() -> get_kv_cache_stride_order()
7927 - # -> get_kv_cache_layout() needs the current vLLM config.
7928 - with set_current_vllm_config(self.vllm_config):
7929 - indexes = backend.indexes_kv_by_block_stride()
7930 - spec = replace(spec, indexes_kv_by_block_stride=indexes)
7931 - spec = backend.customize_spec(spec)
7675 + spec = attn_module.get_attn_backend().customize_spec(spec)
7676 kv_cache_spec[layer_name] = spec
7677
7678 return kv_cache_spec
vllm/v1/worker/gpu_worker.py
+6
@@ -68,6 +68,7 @@ from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
68 from vllm.utils.mem_constants import GiB_bytes
69 from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
70 from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime
71 +from vllm.v1.attention.backends.utils import record_kv_cache_layout
72 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
73 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
74 from vllm.v1.outputs import (
@@ -669,6 +670,11 @@ class Worker(WorkerBase):
670 # so that it's available to the warmup stage.
671 self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
672
673 + # Adopt the engine core's layout; workers spawned after resolution
674 + # (e.g. elastic EP scale-up) only see it through the config.
675 + if kv_cache_config.kv_cache_layout is not None:
676 + record_kv_cache_layout(self.cache_config, kv_cache_config.kv_cache_layout)
677 +
678 # Init kv cache connector here, because it requires
679 # `kv_cache_config`.
680 # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
vllm/v1/worker/kv_connector_model_runner_mixin.py
-152
@@ -8,27 +8,18 @@ from collections.abc import Generator
8 from contextlib import AbstractContextManager, contextmanager, nullcontext
9 from typing import TYPE_CHECKING
10
11 -import torch
12 -
11 from vllm.config import VllmConfig
14 -from vllm.config.cache import CacheDType
12 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
13 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
14 from vllm.forward_context import get_forward_context, set_forward_context
18 -from vllm.logger import init_logger
19 -from vllm.v1.attention.backend import AttentionBackend
20 -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
15 from vllm.v1.outputs import (
16 KVConnectorOutput,
17 ModelRunnerOutput,
18 )
25 -from vllm.v1.worker.utils import AttentionGroup
19
20 if TYPE_CHECKING:
21 from vllm.v1.core.sched.output import SchedulerOutput
22
30 -logger = init_logger(__name__)
31 -
23
24 # Defined as a kv connector functionality mixin for ModelRunner (GPU, TPU)
25 class KVConnectorModelRunnerMixin:
@@ -110,146 +101,3 @@ class KVConnectorModelRunnerMixin:
101
102 if not defer_finalize:
103 kv_connector.clear_connector_metadata()
113 -
114 - @staticmethod
115 - def use_uniform_kv_cache(
116 - attn_groups: list[list[AttentionGroup]],
117 - ) -> bool:
118 - """
119 - Determines whether a uniform KV layout should be used.
120 - A uniform layout means all layers KV caches will share the same
121 - underlying tensor, where for a given block number, the respective
122 - KV data for all layers will be contiguous.
123 - This will allow efficient KV transfer of per-block KV data for all
124 - layers at once.
125 - Note this layout will only be applied given 3 conditions:
126 - 1. The KV Cache config contains just a single group where all layers
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)
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
136 - backend.
137 - Thus, the layers KV data may still not be contiguous per block
138 - if the attention backend does not support it.
139 -
140 - Args:
141 - attn_groups: The list of attention groups for this model
142 - Returns:
143 - True if we should use a uniform KV cache layout.
144 - """
145 -
146 - if not has_kv_transfer_group():
147 - return False
148 - if not get_kv_transfer_group().prefer_cross_layer_blocks:
149 - return False
150 -
151 - if len(attn_groups) != 1 or len(attn_groups[0]) != 1:
152 - return False
153 -
154 - attn_group = attn_groups[0][0]
155 - kv_cache_spec = attn_group.kv_cache_spec
156 - if not isinstance(kv_cache_spec, AttentionSpec):
157 - return False
158 - # Per-token-head quant carves inline-scale views that assume per-layer
159 - # contiguous KV buffers; the cross-layer layout breaks this and corrupts KV.
160 - if kv_cache_spec.kv_quant_mode.is_per_token_head:
161 - return False
162 - return kv_cache_spec.indexes_kv_by_block_stride
163 -
164 - @staticmethod
165 - def allocate_uniform_kv_caches(
166 - kv_cache_config: KVCacheConfig,
167 - attn_groups: list[list[AttentionGroup]],
168 - cache_dtype: CacheDType,
169 - device: torch.device,
170 - kernel_block_sizes: list[int],
171 - ) -> tuple[dict[str, torch.Tensor], torch.Tensor, type[AttentionBackend]]:
172 - """
173 - Initializes and reshapes KV caches for the simple case where all
174 - layers have the same layout.
175 -
176 - This function assumes use_uniform_kv_cache() returned True.
177 -
178 - Args:
179 - kv_cache_config: The KV cache config
180 - attn_groups: The list of attention groups for this model
181 - cache_dtype: The KV cache dtype
182 - device: The torch device to allocate on.
183 - kernel_block_sizes: The kernel block sizes for each KV cache group.
184 - Returns:
185 - A tuple (kv_caches, cross_layers_kv_cache, attn_backend) where:
186 - kv_caches is a dict mapping between layer names to their
187 - corresponding memory buffer for KV cache.
188 - cross_layers_kv_cache is the cross layers kv cache tensor
189 - attn_backend is the attention backend matching this tensor
190 - """
191 - attn_group = attn_groups[0][0]
192 - kv_cache_spec = attn_group.kv_cache_spec
193 - assert isinstance(kv_cache_spec, AttentionSpec)
194 -
195 - tensor_sizes = set(
196 - kv_cache_tensor.size for kv_cache_tensor in kv_cache_config.kv_cache_tensors
197 - )
198 - assert len(tensor_sizes) == 1
199 - tensor_size = tensor_sizes.pop()
200 -
201 - page_size = kv_cache_spec.page_size_bytes
202 - assert tensor_size % page_size == 0
203 - num_blocks = tensor_size // page_size
204 - num_layers = len(kv_cache_config.kv_cache_tensors)
205 - total_size = tensor_size * num_layers
206 -
207 - assert len(kernel_block_sizes) == 1
208 - kernel_block_size = kernel_block_sizes[0]
209 - num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size
210 - kernel_num_blocks = num_blocks * num_blocks_per_kv_block
211 -
212 - attn_backend = attn_group.backend
213 - kv_cache_shape = attn_backend.get_kv_cache_shape(
214 - kernel_num_blocks,
215 - kernel_block_size,
216 - kv_cache_spec.num_kv_heads,
217 - kv_cache_spec.head_size,
218 - cache_dtype_str=cache_dtype,
219 - )
220 -
221 - # prepend a num_layers dimension into the shape
222 - kv_cache_shape = (num_layers,) + kv_cache_shape
223 -
224 - try:
225 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
226 - include_num_layers_dimension=True
227 - )
228 - assert len(kv_cache_stride_order) == len(kv_cache_shape)
229 - except (AttributeError, NotImplementedError):
230 - kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
231 -
232 - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
233 -
234 - logger.info("Allocating a cross layer KV cache of shape %s", kv_cache_shape)
235 -
236 - # allocate one contiguous buffer for all layers
237 - cross_layers_kv_cache = (
238 - torch.zeros(total_size, dtype=torch.int8, device=device)
239 - .view(kv_cache_spec.dtype)
240 - .view(kv_cache_shape)
241 - )
242 -
243 - # Maintain original KV shape view.
244 - inv_order = [
245 - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
246 - ]
247 - permuted_kv_cache = cross_layers_kv_cache.permute(*inv_order)
248 -
249 - kv_caches = {}
250 - for i, kv_cache_tensor in enumerate(kv_cache_config.kv_cache_tensors):
251 - tensor = permuted_kv_cache[i]
252 - for layer_name in kv_cache_tensor.shared_by:
253 - kv_caches[layer_name] = tensor
254 -
255 - return kv_caches, cross_layers_kv_cache, attn_backend
vllm/v1/worker/utils.py
+119 -53
@@ -30,9 +30,11 @@ from vllm.v1.kv_cache_interface import (
30 EncoderOnlyAttentionSpec,
31 KVCacheConfig,
32 KVCacheGroupSpec,
33 + KVCacheLayout,
34 KVCacheSpec,
35 MambaSpec,
36 UniformTypeKVCacheSpecs,
37 + create_kv_cache_views,
38 )
39 from vllm.v1.worker.block_table import get_block_table_width
40
@@ -61,10 +63,10 @@ def _zero_kv_blocks_kernel(
63 ):
64 """Zero KV cache blocks across all segments in a single launch.
65
64 - Each segment is a contiguous region of one block's data. For backends
65 - where blocks are outermost (block_dim=0) there is one segment per
66 - buffer. For backends where K/V is outermost (block_dim=1) there are
67 - two segments per buffer (one for K, one for V).
66 + Each segment is a contiguous region of one block's data. Layer-compact
67 + layouts have one segment per layer buffer; dimensions physically outside
68 + the block dim (separate head groups under LHBNC) and virtual block splits
69 + each get their own segment.
70
71 Segments may have different block strides and page sizes (e.g. packed
72 KV views or models with multiple KV cache groups like MLA + DSA
@@ -110,7 +112,6 @@ class KVBlockZeroer:
112 device: torch.device,
113 attn_groups_iter: Iterable["AttentionGroup"],
114 kernel_block_sizes: list[int],
113 - cache_dtype: str,
115 static_forward_context: dict[str, Any],
116 runner_only_attn_layers: set[str] | None = None,
117 ) -> None:
@@ -119,6 +120,12 @@ class KVBlockZeroer:
120 Each entry is the absolute byte address of a segment start on the
121 GPU, so segments in different CUDA allocations work correctly.
122
123 + Per-layer views are standardized ``[B, H, N, C]`` with blocks at dim 0; dims
124 + physically outside B (separate head groups under LHBNC) each get their own
125 + segment. A segment's page spans everything inside its block, so under BHLNC it
126 + also covers the block's other layers -- safe, since block IDs are global pool
127 + indices and a newly allocated block owns its whole tile.
128 +
129 Block IDs from the scheduler reference logical blocks whose size
130 may differ from the kernel block size (virtual block splitting).
131 Each virtual block is represented as an independent segment so its
@@ -133,7 +140,10 @@ class KVBlockZeroer:
140
141 if runner_only_attn_layers is None:
142 runner_only_attn_layers = set()
136 - seen_ptrs: set[int] = set()
143 + # Overlaid layers (packed layouts) share a base address but may have
144 + # different page sizes; keep the widest span per address so newly
145 + # allocated blocks are fully zeroed for every overlaying group.
146 + seen_ptrs: dict[int, int] = {}
147 seg_addrs: list[int] = []
148 seg_block_strides: list[int] = []
149 seg_page_sizes: list[int] = []
@@ -147,12 +157,6 @@ class KVBlockZeroer:
157 kernel_bs = kernel_block_sizes[group.kv_cache_group_id]
158 assert spec.block_size % kernel_bs == 0
159 ratio = spec.block_size // kernel_bs
150 - block_dim = group.backend.get_kv_cache_block_dim(
151 - kernel_bs,
152 - spec.num_kv_heads,
153 - spec.head_size,
154 - cache_dtype_str=cache_dtype,
155 - )
160
161 for layer_name in group.layer_names:
162 if layer_name in runner_only_attn_layers:
@@ -161,23 +165,18 @@ class KVBlockZeroer:
165 if not isinstance(kv, torch.Tensor):
166 continue
167 dp = kv.data_ptr()
164 - if dp in seen_ptrs:
165 - continue
166 - seen_ptrs.add(dp)
168
169 el = kv.element_size()
169 - block_stride_bytes = kv.stride(block_dim) * el
170 + block_stride_bytes = kv.stride(0) * el
171 assert block_stride_bytes % 4 == 0
171 - assert kv.shape[block_dim] % ratio == 0
172 + assert kv.shape[0] % ratio == 0
173 outer_dims = [
174 d
174 - for d in range(block_dim)
175 + for d in range(1, kv.ndim)
176 if kv.stride(d) * el > block_stride_bytes
177 ]
178 outer_strides = [kv.stride(d) * el for d in outer_dims]
178 - inner_dims = [
179 - d for d in range(kv.ndim) if d != block_dim and d not in outer_dims
180 - ]
179 + inner_dims = [d for d in range(1, kv.ndim) if d not in outer_dims]
180 kernel_page_bytes = el + sum(
181 (kv.shape[d] - 1) * kv.stride(d) * el for d in inner_dims
182 )
@@ -187,9 +186,18 @@ class KVBlockZeroer:
186 off_bytes = sum(i * s for i, s in zip(outer, outer_strides))
187 assert (dp + off_bytes) % 4 == 0
188 for virtual_index in range(ratio):
190 - seg_addrs.append(
191 - dp + off_bytes + virtual_index * block_stride_bytes
192 - )
189 + addr = dp + off_bytes + virtual_index * block_stride_bytes
190 + if (idx := seen_ptrs.get(addr)) is not None:
191 + assert (
192 + seg_block_strides[idx]
193 + == logical_block_stride_bytes // 4
194 + )
195 + seg_page_sizes[idx] = max(
196 + seg_page_sizes[idx], kernel_page_bytes // 4
197 + )
198 + continue
199 + seen_ptrs[addr] = len(seg_addrs)
200 + seg_addrs.append(addr)
201 seg_block_strides.append(logical_block_stride_bytes // 4)
202 seg_page_sizes.append(kernel_page_bytes // 4)
203
@@ -367,6 +375,54 @@ def select_common_block_size(
375 raise ValueError(f"No common block size for {kv_manager_block_size}. ")
376
377
378 +def allocate_kv_cache(
379 + kv_cache_config: KVCacheConfig,
380 + device: torch.device,
381 + layout: KVCacheLayout,
382 + kernel_block_sizes: list[int] | None = None,
383 +) -> dict[str, torch.Tensor]:
384 + """Allocate the KV cache and view it as ``[B, H, N, C]`` per layer.
385 +
386 + Every KVCacheTensor places its layers in the same backing allocation: layer ``l`` of
387 + block ``b`` starts at ``offset + l * layer_stride + b * block_stride``. Cache
388 + groups overlay each other, so tensors may address the same bytes.
389 + """
390 + if not kv_cache_config.kv_cache_tensors:
391 + return {}
392 +
393 + sizes = {tensor.size for tensor in kv_cache_config.kv_cache_tensors}
394 + assert len(sizes) == 1, "KV cache tensors must share one backing allocation."
395 + buf = torch.zeros(sizes.pop(), dtype=torch.int8, device=device)
396 +
397 + kv_caches: dict[str, torch.Tensor] = {}
398 + for tensor in kv_cache_config.kv_cache_tensors:
399 + layer_name = tensor.layers[0]
400 + group_id, group = next(
401 + (group_id, group)
402 + for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
403 + if layer_name in group.layer_names
404 + )
405 + spec = group.kv_cache_spec
406 + if isinstance(spec, UniformTypeKVCacheSpecs):
407 + spec = spec.kv_cache_specs[layer_name]
408 +
409 + num_blocks = kv_cache_config.num_blocks
410 + kernel_block_size = None
411 + if kernel_block_sizes is not None and group_id < len(kernel_block_sizes):
412 + kernel_block_size = kernel_block_sizes[group_id]
413 +
414 + views = create_kv_cache_views(
415 + buf,
416 + spec,
417 + num_blocks,
418 + layout,
419 + tensor,
420 + kernel_block_size=kernel_block_size,
421 + )
422 + kv_caches.update(zip(tensor.layers, views))
423 + return kv_caches
424 +
425 +
426 def prepare_kernel_block_sizes(
427 kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]]
428 ) -> list[int]:
@@ -551,42 +607,52 @@ def bind_kv_cache(
607
608
609 def copy_kv_cache_blocks_inplace(
554 - kv_caches: Iterable[torch.Tensor | list[torch.Tensor]],
610 + kv_caches: Iterable[torch.Tensor],
611 num_blocks: int,
612 kv_cache_block_copies: Sequence[KVCacheBlockCopy],
613 ) -> None:
614 if not kv_cache_block_copies:
615 return
616
561 - storage_tensors: list[torch.Tensor] = []
562 - seen_storage: set[int] = set()
563 - for entry in kv_caches:
564 - # Mamba layers hold a list of state tensors; attention layers a single
565 - # tensor. Both alias the shared block-major backing storage.
566 - tensors = entry if isinstance(entry, (list, tuple)) else (entry,)
567 - for tensor in tensors:
568 - ptr = tensor.untyped_storage().data_ptr()
569 - if ptr in seen_storage:
570 - continue
571 - seen_storage.add(ptr)
572 - storage_tensors.append(tensor)
573 -
574 - if not storage_tensors:
575 - return
576 - device = storage_tensors[0].device
617 indices_np = np.array(kv_cache_block_copies, dtype=np.int64)
578 - indices = async_tensor_h2d(indices_np, device=device)
579 - src_indices, dst_indices = indices.unbind(dim=1)
580 -
581 - for tensor in storage_tensors:
582 - assert tensor.device == device
583 - blocks = torch.empty(0, dtype=torch.uint8, device=device)
584 - blocks.set_(tensor.untyped_storage())
585 - # Block-major backing storage: block i owns the contiguous byte range
586 - # [i * page_size, (i + 1) * page_size).
587 - assert blocks.numel() % num_blocks == 0
588 - blocks = blocks.view(num_blocks, -1)
589 - blocks[dst_indices] = blocks[src_indices]
618 + indices: torch.Tensor | None = None
619 + seen: set[tuple[torch.device, int]] = set()
620 + copied_storages: set[tuple[torch.device, int]] = set()
621 + for cache in kv_caches:
622 + # Layers sharing KV (cross-layer sharing) alias the same view; copy it
623 + # once. data_ptr distinguishes per-layer views of a shared allocation.
624 + key = (cache.device, cache.data_ptr())
625 + if key in seen:
626 + continue
627 + seen.add(key)
628 +
629 + if indices is None:
630 + indices = async_tensor_h2d(indices_np, device=cache.device)
631 + assert cache.device == indices.device
632 + src, dst = indices.unbind(dim=1)
633 +
634 + kernel_blocks_per_block, remainder = divmod(cache.shape[0], num_blocks)
635 + assert remainder == 0, (
636 + f"{cache.shape[0]} kernel blocks not divisible by "
637 + f"{num_blocks} scheduler blocks"
638 + )
639 + storage = cache.untyped_storage()
640 + storage_key = (cache.device, storage.data_ptr())
641 + scheduler_block_stride = (
642 + cache.stride(0) * cache.element_size() * kernel_blocks_per_block
643 + )
644 + if storage.nbytes() == num_blocks * scheduler_block_stride:
645 + if storage_key in copied_storages:
646 + continue
647 + copied_storages.add(storage_key)
648 + blocks = torch.empty(0, dtype=torch.uint8, device=cache.device)
649 + blocks.set_(storage)
650 + blocks = blocks.view(num_blocks, -1)
651 + else:
652 + # Fold virtual block splitting into the shape so that dim 0 counts
653 + # scheduler blocks; unflatten of dim 0 is always a view.
654 + blocks = cache.unflatten(0, (num_blocks, kernel_blocks_per_block))
655 + blocks[dst] = blocks[src]
656
657
658 def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) -> bool:
vllm/v1/worker/worker_base.py
+14
@@ -9,12 +9,17 @@ import torch.nn as nn
9
10 import vllm.ir
11 from vllm.config import VllmConfig, set_current_vllm_config
12 +from vllm.distributed.kv_transfer.kv_connector.utils import get_current_attn_backends
13 from vllm.logger import init_logger
14 from vllm.lora.request import LoRARequest
15 from vllm.multimodal import MULTIMODAL_REGISTRY
16 from vllm.tracing import instrument
17 from vllm.utils.import_utils import resolve_obj_by_qualname
18 from vllm.utils.system_utils import update_environment_variables
19 +from vllm.v1.attention.backends.utils import (
20 + get_supported_kv_cache_layouts,
21 + record_kv_cache_layout,
22 +)
23 from vllm.v1.kv_cache_interface import KVCacheSpec
24
25 if TYPE_CHECKING:
@@ -99,6 +104,15 @@ class WorkerBase:
104 """Get specifications for KV cache implementation."""
105 raise NotImplementedError
106
107 + def get_supported_kv_cache_layouts(self) -> list[str]:
108 + """Layout names every attention backend supports, most preferred first."""
109 + backends = get_current_attn_backends(self.vllm_config)
110 + return [layout.name for layout in get_supported_kv_cache_layouts(backends)]
111 +
112 + def set_kv_cache_layout(self, kv_cache_layout: str) -> None:
113 + """Adopt the KV cache layout resolved by the engine core."""
114 + record_kv_cache_layout(self.vllm_config.cache_config, kv_cache_layout)
115 +
116 def compile_or_warm_up_model(self) -> CompilationTimes:
117 """Prepare model for execution through compilation/warmup.
118