diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index dec438a7be..fea76ebadf 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -243,6 +243,9 @@ steps: DECODE_BLOCK_SIZE: "256" MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" + # The V4 indexer only supports block-outermost layouts; let it pick. + PREFILLER_KV_LAYOUT: "" + DECODER_KV_LAYOUT: "" commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index 8cd20dced1..01fd558043 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -12,6 +12,7 @@ import logging import statistics import types from contextlib import contextmanager +from math import prod import torch from batch_spec import parse_batch_spec, reorder_for_flashinfer @@ -35,12 +36,20 @@ from vllm.config import ( VllmConfig, set_current_vllm_config, ) +from vllm.platforms import current_platform from vllm.v1.attention.backends.utils import ( CommonAttentionMetadata, - get_kv_cache_layout, - set_kv_cache_layout, + get_supported_kv_cache_layouts, + resolve_kv_cache_layout, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheLayout, + KVCacheTensor, + compute_layer_kv_cache_shape_bytes, + compute_layout_strides, + create_kv_cache_views, ) -from vllm.v1.kv_cache_interface import FullAttentionSpec # ============================================================================ # Backend Configuration @@ -340,49 +349,42 @@ def _create_kv_cache( backend_class, device: torch.device, dtype: torch.dtype, + layout: KVCacheLayout, ) -> list: - """Create KV cache tensors for all layers using the backend's methods. - - Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order() - to create the cache with the correct shape and memory layout. - """ - # Get the logical shape from the backend - cache_shape = backend_class.get_kv_cache_shape( - num_blocks=max_num_blocks, + """Create KV cache tensors for all layers using the standard allocator.""" + if config.kv_cache_dtype.startswith("fp8"): + cache_dtype = current_platform.fp8_dtype() + else: + cache_dtype = dtype + spec = FullAttentionSpec( block_size=config.block_size, num_kv_heads=config.num_kv_heads, head_size=config.head_dim, + dtype=cache_dtype, + ) + # Apply the backend's page customization, as the worker does for the real spec. + spec = backend_class.customize_spec(spec) + total_bytes = ( + prod(compute_layer_kv_cache_shape_bytes(spec, max_num_blocks)) + * config.num_layers + ) + buf = torch.zeros(total_bytes, device=device, dtype=torch.int8) + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, max_num_blocks, config.num_layers, layout + ) + tensor = KVCacheTensor( + size=total_bytes, + layers=[str(i) for i in range(config.num_layers)], + layer_stride=layer_stride, + block_stride=block_stride, + ) + return create_kv_cache_views( + buf, + spec, + max_num_blocks, + layout, + tensor, ) - - # Get the stride order for custom memory layout - try: - stride_order = backend_class.get_kv_cache_stride_order() - assert len(stride_order) == len(cache_shape) - except (AttributeError, NotImplementedError): - stride_order = tuple(range(len(cache_shape))) - - # Permute shape to physical layout order - physical_shape = tuple(cache_shape[i] for i in stride_order) - - # Compute inverse permutation to get back to logical view - inv_order = [stride_order.index(i) for i in range(len(stride_order))] - - # Use fp8 dtype for cache when requested. - cache_dtype = dtype - if config.kv_cache_dtype == "fp8": - from vllm.platforms import current_platform - - cache_dtype = current_platform.fp8_dtype() - - cache_list = [] - for _ in range(config.num_layers): - # Allocate in physical layout order (contiguous in memory) - cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype) - # Permute to logical view - cache = cache.permute(*inv_order) - cache_list.append(cache) - - return cache_list # ============================================================================ @@ -499,13 +501,10 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: backend_class, impl, layer = _create_backend_impl( backend_cfg, config, device, dtype ) - # Set KV cache layout if the backend requires a specific one - # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention) - required_layout = backend_class.get_required_kv_cache_layout() - if required_layout is not None: - set_kv_cache_layout(required_layout) - get_kv_cache_layout.cache_clear() + # (e.g., FlashInfer requires LBHNC on SM100/Blackwell for TRTLLM attention) + supported = get_supported_kv_cache_layouts([backend_class]) + layout = resolve_kv_cache_layout(vllm_config, [[m.name for m in supported]]) common_metadata = _build_common_attn_metadata( q_lens, kv_lens, config.block_size, device @@ -542,7 +541,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: ) cache_list = _create_kv_cache( - config, max_num_blocks, backend_class, device, dtype + config, max_num_blocks, backend_class, device, dtype, layout ) timing_stats, mem_stats = _run_single_benchmark( diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index be2a4cc4f3..d62ab626db 100644 Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md index 800b814250..a35a709793 100644 --- a/docs/features/mooncake_store_connector_usage.md +++ b/docs/features/mooncake_store_connector_usage.md @@ -246,7 +246,6 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`. - `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`. -- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`. - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`. - `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). - `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`. diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index 36b8876c69..031fba5b2f 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -61,7 +61,7 @@ th:not(:first-child) { See [Configuration Notes](#configuration-notes) below for what must match and what may differ. -2 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"}}'`. +2 Cross-layer contiguity is achieved by using a `BLHNC` layout (set via `VLLM_KV_CACHE_LAYOUT=BLHNC`). 3 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. @@ -100,9 +100,9 @@ By default, a **compatibility hash** is checked during handshake. P and D instan ### KV cache layout -- NixlConnector defaults to **`HND`** layout for optimal transfer performance (non-MLA models). -- `NHD` layout is supported but does **not** allow heterogeneous TP head splitting. -- Experimental `HND` ↔ `NHD` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA. +- NixlConnector defaults to **`LBHNC`** (head-major, formerly `HND`) layout for optimal transfer performance (non-MLA models). +- `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting. +- Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA. ### Quantized KV cache diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index a84e3de15b..9799eb9cf1 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -404,21 +404,12 @@ GB-series GPUs support multi-node NVLink. NIXL supports this capability, but KVC ### Heterogeneous KV Layout support -Support use case: Prefill with 'HND' and decode with 'NHD' with experimental configuration +Support use case: Prefill with `LBHNC` and decode with `LBNHC` with experimental configuration ```bash --kv-transfer-config '{..., "enable_permute_local_kv":"True"}' ``` -### Cross layers blocks - -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. -To enable this feature: - -```bash ---kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}' -``` - ## Metrics Reference vLLM periodically logs a `KV Transfer metrics` line summarising NIXL transfer diff --git a/rust/Cargo.lock b/rust/Cargo.lock index ff24484a63..2adb17dc51 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2340,15 +2340,6 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" -[[package]] -name = "lz4_flex" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" -dependencies = [ - "twox-hash", -] - [[package]] name = "logos" version = "0.15.1" @@ -2415,6 +2406,15 @@ dependencies = [ "logos-codegen 0.16.1", ] +[[package]] +name = "lz4_flex" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" +dependencies = [ + "twox-hash", +] + [[package]] name = "macro_rules_attribute" version = "0.2.2" diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index da335b9155..c5f65fe4eb 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -7,7 +7,11 @@ import torch._dynamo from tests.compile.backend import LazyInitPass, TestBackend from tests.utils import TestFP8Layer, flat_product -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + dense_kv_cache_views, +) from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant from vllm.compilation.passes.fusion.attn_quant_fusion import ( ATTN_OP, @@ -39,7 +43,14 @@ from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode +from vllm.v1.attention.backends.utils import ( + get_supported_kv_cache_layouts, + resolve_kv_cache_layout, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + get_kv_quant_mode, +) DEVICE_TYPE = current_platform.device_type FP8_DTYPE = current_platform.fp8_dtype() @@ -108,32 +119,29 @@ class AttentionQuantPatternModel(torch.nn.Module): max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size num_blocks = batch_size * max_blocks - # Fetch the attention backend and kv cache shape and stride order - attn_backend = self.attn.attn_backend - kv_cache_shape = attn_backend.get_kv_cache_shape( - num_blocks, - self.block_size, - self.num_kv_heads, - self.head_size, - cache_dtype_str=self.attn.kv_cache_dtype, - ) - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - inv_order = [ - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) - ] - - # Create dummy KV cache - raw_tensor = torch.zeros( - kv_cache_shape, + spec = AttentionSpec( + block_size=self.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, dtype=self.attn.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype), + ) + supported = get_supported_kv_cache_layouts([self.attn.attn_backend]) + layout = resolve_kv_cache_layout( + self.vllm_config, [[m.name for m in supported]] + ) + raw_tensor = torch.zeros( + num_blocks * spec.page_size_bytes, + dtype=torch.int8, device=self.device, ) - kv_cache = raw_tensor.permute(*inv_order) + kv_cache = dense_kv_cache_views( + raw_tensor, + spec, + num_blocks, + num_layers=1, + layout=layout, + )[0] self.attn.kv_cache = kv_cache diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 106bb8e80a..11c455bfa2 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -150,27 +150,14 @@ class MLAAttentionQuantPatternModel(torch.nn.Module): max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size num_blocks = batch_size * max_blocks - # MLA KV cache is 3D: (num_blocks, block_size, head_size) - attn_backend = self.mla_attn.attn_backend - kv_cache_shape = attn_backend.get_kv_cache_shape( - num_blocks, self.block_size, 1, self.head_size + # MLA KV cache is 4D: (num_blocks, num_heads=1, block_size, head_size) + kv_cache = torch.zeros( + (num_blocks, 1, self.block_size, self.head_size), + dtype=self.kv_cache_dtype, + device=self.device, ) - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - ordered_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - inv_order = [ - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) - ] - - raw_tensor = torch.zeros( - ordered_shape, dtype=self.kv_cache_dtype, device=self.device - ) - kv_cache = raw_tensor.permute(*inv_order) - - self.mla_attn.kv_cache = kv_cache + self.mla_attn.bind_kv_cache(kv_cache) self.attn_metadata = self.builder.build( common_prefix_len=0, common_attn_metadata=common_attn_metadata diff --git a/tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py b/tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py index cc3bfb7693..3c894ee141 100644 --- a/tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py +++ b/tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py @@ -165,29 +165,14 @@ class MLARoPEKVCacheCatTestModel(torch.nn.Module): max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size num_blocks = batch_size * max_blocks - # Fetch the attention backend and kv cache shape and stride order - kv_cache_shape = self.attn_backend.get_kv_cache_shape( - num_blocks, self.block_size, self.num_kv_heads, self.head_size - ) - try: - kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order() - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - inv_order = [ - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) - ] - - raw_tensor = torch.zeros( - num_blocks * self.block_size * self.num_kv_heads * self.head_size, + # MLA KV cache is 4D: (num_blocks, num_heads=1, block_size, head_size) + kv_cache = torch.zeros( + (num_blocks, 1, self.block_size, self.head_size), dtype=self.kv_cache_dtype, device=self.device, ) - raw_tensor = raw_tensor.view(kv_cache_shape) - kv_cache = raw_tensor.permute(*inv_order) - self.mla_attn.kv_cache = kv_cache + self.mla_attn.bind_kv_cache(kv_cache) # Build attn metadata attn_metadata = self.builder.build( diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py index e385999db3..6965346a58 100644 --- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py @@ -10,7 +10,11 @@ from packaging.version import Version import vllm.config from tests.compile.backend import TestBackend -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + dense_kv_cache_views, +) from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import ( @@ -40,7 +44,7 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout pytestmark = pytest.mark.skip_global_cleanup @@ -147,20 +151,25 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module): self.attn._k_scale = self.attn._k_scale.to(device) self.attn._v_scale = self.attn._v_scale.to(device) - self.builder = self.attn.attn_backend.get_builder_cls()( - kv_cache_spec=AttentionSpec( + # Mirror the worker spec-collection loop: the backend publishes the + # packing it needs. + self.kv_cache_spec = self.attn.attn_backend.customize_spec( + AttentionSpec( block_size=self.block_size, num_kv_heads=self.num_kv_heads, head_size=head_size, dtype=self.kv_cache_dtype, - ), + ) + ) + self.builder = self.attn.attn_backend.get_builder_cls()( + kv_cache_spec=self.kv_cache_spec, layer_names=[self.attn.layer_name], vllm_config=vllm_config, device=device, ) def build_attn_metadata( - self, batch_size: int, kv_stride_order: tuple[int, ...] | None = None + self, batch_size: int, layout: KVCacheLayout ) -> CommonAttentionMetadata: batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) common_attn_metadata = create_common_attn_metadata( @@ -170,27 +179,18 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module): max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size num_blocks = batch_size * max_blocks - attn_backend = self.attn.attn_backend - kv_cache_shape = attn_backend.get_kv_cache_shape( - num_blocks, self.block_size, self.num_kv_heads, self.head_size - ) - # Caller can force a physical layout; else use the backend's. - if kv_stride_order is None: - try: - kv_stride_order = attn_backend.get_kv_cache_stride_order() - except (AttributeError, NotImplementedError): - kv_stride_order = tuple(range(len(kv_cache_shape))) - - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_stride_order) - inv_order = [kv_stride_order.index(i) for i in range(len(kv_stride_order))] - raw_tensor = torch.zeros( - 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size, - dtype=self.kv_cache_dtype, + num_blocks * self.kv_cache_spec.page_size_bytes, + dtype=torch.int8, device=self.device, ) - raw_tensor = raw_tensor.view(kv_cache_shape) - kv_cache = raw_tensor.permute(*inv_order) + kv_cache = dense_kv_cache_views( + raw_tensor, + self.kv_cache_spec, + num_blocks, + num_layers=1, + layout=layout, + )[0] # Store as a bare tensor (not wrapped in a list) to match production # `bind_kv_cache` behavior. `get_attention_context` returns this @@ -272,7 +272,7 @@ def _run_qk_norm_rope_kvcache_fusion_test( block_size: int, is_neox: bool, use_shuffle_kv_layout: str, - kv_stride_order: tuple[int, ...], + kv_layout: KVCacheLayout, dtype: torch.dtype, kv_cache_dtype: str, rms_norm_eps: float, @@ -345,7 +345,7 @@ def _run_qk_norm_rope_kvcache_fusion_test( # Run unfused (eager) forward with set_forward_context(None, vllm_config): forward_context = get_forward_context() - attn_metadata = model.build_attn_metadata(num_tokens, kv_stride_order) + attn_metadata = model.build_attn_metadata(num_tokens, kv_layout) forward_context.slot_mapping = { model.layer_name: attn_metadata.slot_mapping } @@ -360,7 +360,7 @@ def _run_qk_norm_rope_kvcache_fusion_test( with set_forward_context(None, vllm_config): model_fused = torch.compile(model, backend=backend) forward_context = get_forward_context() - attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_stride_order) + attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_layout) forward_context.slot_mapping = { model.layer_name: attn_metadata.slot_mapping } @@ -437,10 +437,11 @@ _FUSION_CONFIGS = [ ) @pytest.mark.parametrize("num_tokens", [5, 2048]) @pytest.mark.parametrize( - "kv_stride_order", - # FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size, - # 2*head_size); no separate K/V dim to permute. - [pytest.param((0, 1, 2, 3), id="packed_4d")], + "kv_layout", + [ + pytest.param(KVCacheLayout.LBHNC, id="head_major"), + pytest.param(KVCacheLayout.LBNHC, id="token_major"), + ], ) @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False]) @pytest.mark.parametrize("block_size", [16, 32, 64]) @@ -467,7 +468,7 @@ def test_qk_norm_rope_kvcache_fusion( is_neox: bool, attn_backend: AttentionBackendEnum, enable_aiter_triton_rope: bool, - kv_stride_order: tuple[int, ...], + kv_layout: KVCacheLayout, block_size: int, dtype: torch.dtype, kv_cache_dtype: str, @@ -486,7 +487,7 @@ def test_qk_norm_rope_kvcache_fusion( block_size=block_size, is_neox=is_neox, use_shuffle_kv_layout="0", - kv_stride_order=kv_stride_order, + kv_layout=kv_layout, dtype=dtype, kv_cache_dtype=kv_cache_dtype, rms_norm_eps=rms_norm_eps, diff --git a/tests/compile/passes/test_rope_kvcache_fusion.py b/tests/compile/passes/test_rope_kvcache_fusion.py index 2f68061749..ea3454227b 100644 --- a/tests/compile/passes/test_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rope_kvcache_fusion.py @@ -7,7 +7,11 @@ from torch._higher_order_ops import auto_functionalized import vllm.config from tests.compile.backend import TestBackend -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + dense_kv_cache_views, +) from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.compilation.passes.fusion import rope_kvcache_fusion from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP @@ -38,6 +42,10 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheLayout, +) INDEX_SELECT_OP = torch.ops.aten.index.Tensor VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update @@ -149,9 +157,20 @@ class QKRoPEKVCacheTestModel(torch.nn.Module): FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype ) + # Mirror the worker spec-collection loop: the backend publishes its + # packing (e.g. ROCM_ATTN's separate K/V head groups). + self.kv_cache_spec = self.attn_backend.customize_spec( + FullAttentionSpec( + block_size=self.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + dtype=self.kv_cache_dtype, + ) + ) + # Initialize attn MetadataBuilder self.builder = self.attn_backend.get_builder_cls()( - kv_cache_spec=self.attn.get_kv_cache_spec(vllm_config), + kv_cache_spec=self.kv_cache_spec, layer_names=[self.attn.layer_name], vllm_config=vllm_config, device=device, @@ -168,30 +187,25 @@ class QKRoPEKVCacheTestModel(torch.nn.Module): max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size num_blocks = batch_size * max_blocks - # Fetch the attention backend and kv cache shape and stride order - kv_cache_shape = self.attn_backend.get_kv_cache_shape( - num_blocks, self.block_size, self.num_kv_heads, self.head_size - ) - try: - kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order() - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) + # A backend's preferred supported layout wins over the default, mirroring + # selector-time resolution (e.g. ROCM_ATTN's head groups force LHBNC). + layout = KVCacheLayout.LBNHC + supported = self.attn_backend.supported_kv_cache_layouts() + if supported: + layout = supported[0] - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - inv_order = [ - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) - ] - - # Create dummy KV cache raw_tensor = torch.zeros( - 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size, - dtype=self.kv_cache_dtype, + num_blocks * self.kv_cache_spec.page_size_bytes, + dtype=torch.int8, device=self.device, ) - raw_tensor = raw_tensor.view(kv_cache_shape) - kv_cache = raw_tensor.permute(*inv_order) - - self.attn.kv_cache = kv_cache + self.attn.kv_cache = dense_kv_cache_views( + raw_tensor, + self.kv_cache_spec, + num_blocks, + num_layers=1, + layout=layout, + )[0] # Build attn metadata attn_metadata = self.builder.build( diff --git a/tests/distributed/test_kvlayout.py b/tests/distributed/test_kvlayout.py index c8177f1c7c..122cbf9eeb 100644 --- a/tests/distributed/test_kvlayout.py +++ b/tests/distributed/test_kvlayout.py @@ -21,7 +21,7 @@ def test_get_kv_connector_cache_layout_without_kv_connector(): with set_current_vllm_config(vllm_config): # Test with default settings layout = get_kv_connector_cache_layout() - assert layout == "NHD" + assert layout is None def test_get_kv_connector_cache_layout_with_lmcache_connector(): @@ -35,7 +35,7 @@ def test_get_kv_connector_cache_layout_with_lmcache_connector(): with set_current_vllm_config(vllm_config): # Test with default settings layout = get_kv_connector_cache_layout() - assert layout == "NHD" + assert layout is None def test_get_kv_connector_cache_layout_with_nixl_connector(): @@ -52,7 +52,7 @@ def test_get_kv_connector_cache_layout_with_nixl_connector(): with set_current_vllm_config(vllm_config): # Test with default settings layout = get_kv_connector_cache_layout() - assert layout == "HND" + assert layout == "LBHNC" def test_get_kv_connector_cache_layout_with_multi_connector(): @@ -75,4 +75,4 @@ def test_get_kv_connector_cache_layout_with_multi_connector(): with set_current_vllm_config(vllm_config): # Test with default settings layout = get_kv_connector_cache_layout() - assert layout == "HND" + assert layout == "LBHNC" diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 7b7317ebaa..3e3bfe9dac 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -356,21 +356,21 @@ def test_reshape_and_cache_flash( dequant_nvfp4_kv_cache, ) - def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale): - # data_cache: [N, T, H, data_dim] NHD (contiguous inner dims) - # scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims) - # Permute to HND layout for the dequant utility. + def dequant_nvfp4_cache(data_cache, scale_cache, global_scale): + # data_cache: [B, N, H, data_dim] logical view (layout in strides) + # scale_cache: [B, N, H, scale_dim] logical view (layout in strides) + # Permute to [B, H, N, dim] for the dequant utility. data_hnd = data_cache.permute(0, 2, 1, 3) scale_hnd = scale_cache.permute(0, 2, 1, 3) result_hnd = dequant_nvfp4_kv_cache( data_hnd, scale_hnd, global_scale, head_size, block_size ) - return result_hnd.permute(0, 2, 1, 3) # back to [N, T, H, D] + return result_hnd.permute(0, 2, 1, 3) # back to [B, N, H, dim] - result_key_cache = dequant_nvfp4_cache_nhd( + result_key_cache = dequant_nvfp4_cache( nvfp4_key_data, key_scale_cache, k_scale.item() ) - result_value_cache = dequant_nvfp4_cache_nhd( + result_value_cache = dequant_nvfp4_cache( nvfp4_value_data, value_scale_cache, v_scale.item() ) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 8fe2342808..42883d9899 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -5,10 +5,10 @@ import pytest import torch +from tests.v1.attention.utils import dense_kv_cache_views from vllm import _custom_ops as ops -from vllm.models.minimax_m3.common.indexer import ( - MiniMaxM3IndexerBackend, -) +from vllm.config import CacheConfig +from vllm.config.cache import _layout_from_name from vllm.models.minimax_m3.common.ops.index_topk import ( minimax_m3_index_decode, minimax_m3_index_score, @@ -19,15 +19,15 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import ( minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode, ) -from vllm.models.minimax_m3.common.sparse_attention import ( - MiniMaxM3SparseBackend, - MiniMaxM3SparseTritonImpl, -) +from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl from vllm.platforms import current_platform -from vllm.v1.attention.backends.utils import set_kv_cache_layout -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec -from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache -from vllm.v1.worker.utils import AttentionGroup +from vllm.v1.attention.backends.utils import record_kv_cache_layout +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheLayout, + MLAAttentionSpec, + compute_layer_kv_cache_shape_bytes, +) if not (current_platform.is_cuda() or current_platform.is_rocm()): pytest.skip( @@ -37,35 +37,46 @@ if not (current_platform.is_cuda() or current_platform.is_rocm()): @pytest.fixture -def kv_layout(request): - """Set the global KV cache layout for one test and restore it after.""" - set_kv_cache_layout(request.param) - try: - yield request.param - finally: - set_kv_cache_layout(None) +def kv_layout(request) -> KVCacheLayout: + """Resolve the requested layout name (legacy NHD/HND aliases included).""" + return _layout_from_name(request.param) -def _stride_order_for(backend: type[MiniMaxM3SparseBackend], ndim: int) -> tuple: - """Mirror the allocator's stride-order resolution (identity fallback).""" - try: - stride_order = backend.get_kv_cache_stride_order() - assert len(stride_order) == ndim - except (AttributeError, NotImplementedError): - stride_order = tuple(range(ndim)) +def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]: + """Per-layer physical stride order for the given layout; the 3-dim + indexer side cache (H=1) is contiguous, so identity.""" + if ndim == 3: + return (0, 1, 2) + stride_order = layout.layer_view_order + assert len(stride_order) == ndim return stride_order +def _main_spec() -> FullAttentionSpec: + return FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_DIM, + head_size_v=HEAD_DIM, + dtype=DTYPE, + ) + + +def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]: + """Standardized per-layer logical shape [B, H, N, C] for the main cache, + derived the same way the production allocator does.""" + shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages) + return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize) + + def _allocate_main_kv_via_contract( - num_pages: int, device: torch.device | str = "cuda" + num_pages: int, layout: KVCacheLayout, device: torch.device | str = "cuda" ) -> torch.Tensor: """Build the main KV cache exactly as the production allocator does for the - currently active layout: allocate the physical (permuted) tensor, then - expose the inverse-permuted logical-NHD view the backend sees.""" - logical_shape = MiniMaxM3SparseBackend.get_kv_cache_shape( - num_pages, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM - ) - stride_order = _stride_order_for(MiniMaxM3SparseBackend, len(logical_shape)) + given layout: allocate the physical (permuted) tensor, then expose the + inverse-permuted logical [B, H, N, C] view the backend sees.""" + logical_shape = _main_kv_logical_shape(num_pages) + stride_order = _layer_stride_order(layout, len(logical_shape)) physical_shape = tuple(logical_shape[i] for i in stride_order) inv_order = [stride_order.index(i) for i in range(len(stride_order))] raw = torch.randn(physical_shape, device=device, dtype=DTYPE) @@ -817,7 +828,7 @@ def _reference_sparse_attn( ], ) def test_prefill_sparse_attention_correctness( - kv_layout: str, + kv_layout: KVCacheLayout, q_lens: tuple[int, ...], kv_lens: tuple[int, ...], ): @@ -851,7 +862,7 @@ def test_prefill_sparse_attention_correctness( # Allocate the main KV cache through the backend layout contract so the # physical storage matches the active layout (contiguous NHD or strided # HND), while the kernels and reference see the logical-NHD view. - kv_cache = _allocate_main_kv_via_contract(num_pages) + kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout) # Build sparse block indices with the same contract as the real M3 indexer: # one forced local block, then score-selected older causal blocks. @@ -904,38 +915,31 @@ def test_prefill_sparse_attention_correctness( assert error.max().item() < 1.7e-2 -def test_main_backend_layout_contract(): - """The main sparse backend exposes the logical-NHD shape and the - flash_attn-style stride order for each layout.""" +def test_main_cache_layout_contract(): + """The standardized per-layer logical shape is [B, H, N, C] with packed + K/V content, and the legacy layout aliases resolve to the expected + per-layer stride orders.""" nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM - logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d) + logical = _main_kv_logical_shape(nb) assert logical == (nb, h, bs, 2 * d) # The old separate K/V-axis shape is no longer the logical shape. assert logical != (nb, 2, bs, h, d) - try: - set_kv_cache_layout("HND") - assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3) - set_kv_cache_layout("NHD") - assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3) - finally: - set_kv_cache_layout(None) + assert KVCacheLayout.LBHNC.layer_view_order == (0, 1, 2, 3) + assert KVCacheLayout.LBNHC.layer_view_order == (0, 2, 1, 3) - for layout in ("NHD", "HND"): - try: - set_kv_cache_layout(layout) - order = MiniMaxM3SparseBackend.get_kv_cache_stride_order() - finally: - set_kv_cache_layout(None) + cache_config = CacheConfig() + record_kv_cache_layout(cache_config, "HND") + assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBHNC + cache_config = CacheConfig() + record_kv_cache_layout(cache_config, "NHD") + assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBNHC + + for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC): + order = layout.layer_view_order # Valid permutation: no duplicates, covers every axis. assert set(order) == set(range(len(order))) - # M3 has no cross-layer KV blocks. - with pytest.raises(NotImplementedError): - MiniMaxM3SparseBackend.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - def test_aiter_sparse_pa_layout_contract(monkeypatch): """The shuffle-only AITER path retains separately contiguous K/V storage.""" @@ -948,75 +952,53 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch): lambda: True, ) - nb, bs, h, d = 7, BLOCK_SIZE, 1, HEAD_DIM - logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d) - order = MiniMaxM3SparseBackend.get_kv_cache_stride_order() - assert logical == (nb, 2, bs, h, d) - assert order == (1, 0, 2, 3, 4) + backend = sparse_attn_mod.MiniMaxM3SparseBackend + assert KVCacheLayout.LHBNC in backend.supported_kv_cache_layouts() - physical_shape = tuple(logical[i] for i in order) - inv_order = [order.index(i) for i in range(len(order))] - raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE) - logical_view = raw.permute(*inv_order) - key_cache, value_cache = logical_view.unbind(1) + nb, h = 7, 1 + spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=h, + head_size=HEAD_DIM, + dtype=DTYPE, + num_head_slots=2, + state_content_bytes=h * HEAD_DIM * DTYPE.itemsize, + ) + raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8) + view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LHBNC)[0] + key_cache, value_cache = view.unbind(1) assert key_cache.is_contiguous() assert value_cache.is_contiguous() -def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch): - """Do not pair AITER's separated cache layout with the Triton fallback.""" - import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod - - monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True) - monkeypatch.setattr( - sparse_attn_mod.rocm_aiter_ops, - "is_shuffle_kv_cache_enabled", - lambda: True, +def test_indexer_cache_squeezes_to_contiguous_3d(): + """The indexer side cache is standardized 4D with H=1: under both layouts + the allocator's logical view stays contiguous and squeezes (as + `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim + [num_blocks, block_size, head_dim] cache the kernels consume.""" + nb = 5 + ispec = MLAAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE ) + shape_bytes = compute_layer_kv_cache_shape_bytes(ispec, nb) + assert shape_bytes == (nb, 1, BLOCK_SIZE, HEAD_DIM * DTYPE.itemsize) + assert _layer_stride_order(KVCacheLayout.LBNHC, 3) == (0, 1, 2) - with pytest.raises(ValueError, match="num_kv_heads == 1"): - MiniMaxM3SparseBackend.get_kv_cache_shape(7, BLOCK_SIZE, 2, HEAD_DIM) - - -def test_main_backend_unknown_layout_raises(monkeypatch): - """An unrecognized layout (injected past env-var validation) is rejected.""" - import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod - - monkeypatch.setattr(sparse_attn_mod, "get_kv_cache_layout", lambda: "BOGUS") - with pytest.raises(ValueError, match="Unknown cache layout format"): - MiniMaxM3SparseBackend.get_kv_cache_stride_order() - - -def test_indexer_backend_stride_order_is_identity(): - """The 3-dim indexer cache must not inherit the parent's 4-element stride - order; it overrides to the 3-element identity so the allocator keeps the - contiguous layout.""" - assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2) - - # Cross-layer (per-layer-stacked) KV blocks are not supported. - with pytest.raises(NotImplementedError): - MiniMaxM3IndexerBackend.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - - # The stride order matches the 3-dim indexer shape rank. - indexer_shape = MiniMaxM3IndexerBackend.get_kv_cache_shape( - 5, BLOCK_SIZE, 1, HEAD_DIM - ) - assert len(indexer_shape) == 3 - assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2) + for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC): + iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8) + view = dense_kv_cache_views(iraw, ispec, nb, 1, layout, BLOCK_SIZE)[0] + assert tuple(view.shape) == (nb, 1, BLOCK_SIZE, HEAD_DIM) + indexer_cache = view.squeeze(1) + assert tuple(indexer_cache.shape) == (nb, BLOCK_SIZE, HEAD_DIM) + assert indexer_cache.is_contiguous() def test_hnd_allocation_is_packed_head_major(): """Under HND the backend-visible logical view is the packed head-major physical allocation.""" nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM - logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d) - try: - set_kv_cache_layout("HND") - stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order() - finally: - set_kv_cache_layout(None) + logical = _main_kv_logical_shape(nb) + stride_order = KVCacheLayout.LBHNC.layer_view_order physical_shape = tuple(logical[i] for i in stride_order) assert physical_shape == (nb, h, bs, 2 * d) @@ -1042,27 +1024,13 @@ def test_main_cache_is_block_first_and_unpadded(): """The allocator's contiguous-view branch (not the padded-strided branch) is used for the main GQA cache: its spec is unpadded and the physical layout keeps num_blocks as the first dimension under both layouts.""" - from vllm.v1.kv_cache_interface import FullAttentionSpec - - spec = FullAttentionSpec( - block_size=BLOCK_SIZE, - num_kv_heads=NUM_KV_HEADS, - head_size=HEAD_DIM, - head_size_v=HEAD_DIM, - dtype=DTYPE, - ) + spec = _main_spec() # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided(). assert spec.page_size_padded is None - logical = MiniMaxM3SparseBackend.get_kv_cache_shape( - 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM - ) - for layout in ("NHD", "HND"): - try: - set_kv_cache_layout(layout) - order = MiniMaxM3SparseBackend.get_kv_cache_stride_order() - finally: - set_kv_cache_layout(None) + logical = _main_kv_logical_shape(4) + for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC): + order = layout.layer_view_order inv_order = [order.index(i) for i in range(len(order))] # Physical first dim is num_blocks (block-first); required by the # 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(): @pytest.mark.parametrize("decode_query_len", [1, 4]) @pytest.mark.parametrize("num_padded_reqs", [0, 2]) def test_decode_sparse_attention_correctness( - kv_layout: str, + kv_layout: KVCacheLayout, seq_lens_list: tuple[int, ...], decode_query_len: int, num_padded_reqs: int, @@ -1287,7 +1255,7 @@ def test_decode_sparse_attention_correctness( q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs( seq_lens_list, decode_query_len, num_padded_reqs ) - kv_cache = _allocate_main_kv_via_contract(num_pages) + kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout) actual = torch.empty_like(q) minimax_m3_sparse_attn_decode( @@ -1364,93 +1332,20 @@ def test_decode_wrong_layout_breaks_parity(): assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2 -def _make_attn_group(backend, spec): - return AttentionGroup( - backend=backend, - layer_names=["main"], - kv_cache_spec=spec, - kv_cache_group_id=0, - ) - - def test_main_cache_byte_identical_through_production_allocator(): - """AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main + """AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main `FullAttentionSpec` under HND and assert the backend-visible view has the - same shape, stride, and storage offset as the packed-HND allocation; the - indexer `MLAAttentionSpec` allocates through the same path to its 3-dim - shape.""" + same shape, stride, and storage offset as the packed-HND allocation.""" nb = 4 - spec = FullAttentionSpec( - block_size=BLOCK_SIZE, - num_kv_heads=NUM_KV_HEADS, - head_size=HEAD_DIM, - head_size_v=HEAD_DIM, - dtype=DTYPE, - ) + spec = _main_spec() raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8) - group = _make_attn_group(MiniMaxM3SparseBackend, spec) - try: - set_kv_cache_layout("HND") - kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {}) - finally: - set_kv_cache_layout(None) - view = kv_caches["main"] + view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LBHNC, BLOCK_SIZE)[0] oracle = raw.view(DTYPE).view((nb, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM)) assert tuple(view.shape) == tuple(oracle.shape) assert view.stride() == oracle.stride() assert view.storage_offset() == oracle.storage_offset() - # Indexer cache allocates through the same path under both layouts. - ispec = MLAAttentionSpec( - block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE - ) - for layout in ("NHD", "HND"): - iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8) - igroup = AttentionGroup( - backend=MiniMaxM3IndexerBackend, - layer_names=["idx"], - kv_cache_spec=ispec, - kv_cache_group_id=0, - ) - try: - set_kv_cache_layout(layout) - iout = _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {}) - finally: - set_kv_cache_layout(None) - assert tuple(iout["idx"].shape) == (nb, BLOCK_SIZE, HEAD_DIM) - - -def test_indexer_inherited_stride_order_trips_allocator_assert(): - """AC-4 negative: without the indexer override, the inherited 4-element - stride order trips the allocator's `len(stride_order) == len(shape)` assert - for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the - allocator's `(AttributeError, NotImplementedError)` fallback.""" - - class _BrokenIndexerBackend(MiniMaxM3IndexerBackend): - # Simulate inheriting the parent's 4-element stride order. - get_kv_cache_stride_order = staticmethod( - MiniMaxM3SparseBackend.get_kv_cache_stride_order - ) - - nb = 4 - ispec = MLAAttentionSpec( - block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE - ) - iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8) - igroup = AttentionGroup( - backend=_BrokenIndexerBackend, - layer_names=["idx"], - kv_cache_spec=ispec, - kv_cache_group_id=0, - ) - try: - set_kv_cache_layout("HND") - with pytest.raises(AssertionError): - _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {}) - finally: - set_kv_cache_layout(None) - def test_padded_main_cache_is_flagged(): """AC-2.1 negative: the M3 main cache relies on the allocator's @@ -1465,11 +1360,7 @@ def test_padded_main_cache_is_flagged(): ) assert inv_order[0] == 0, "main GQA cache must remain block-first" - try: - set_kv_cache_layout("HND") - stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order() - finally: - set_kv_cache_layout(None) + stride_order = KVCacheLayout.LBHNC.layer_view_order good = FullAttentionSpec( block_size=BLOCK_SIZE, @@ -1493,13 +1384,13 @@ def test_padded_main_cache_is_flagged(): @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True) -def test_reshape_and_cache_flash_write_persists(kv_layout: str): +def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout): """AC-5 write path: the `reshape_and_cache_flash` write site now consumes packed-content K/V split views. Writing through those views must persist into the bound storage under both layouts.""" torch.manual_seed(0) num_pages = 4 - kv_cache = _allocate_main_kv_via_contract(num_pages) + kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout) with torch.no_grad(): kv_cache.zero_() diff --git a/tests/kernels/attention/test_rocm_aiter_fa.py b/tests/kernels/attention/test_rocm_aiter_fa.py index 566643f43f..423c7c3352 100644 --- a/tests/kernels/attention/test_rocm_aiter_fa.py +++ b/tests/kernels/attention/test_rocm_aiter_fa.py @@ -310,18 +310,23 @@ def test_aiter_mha_backend_contract(): def test_aiter_mha_backend_validates_kv_cache_block_size(): - """The backend should reject KV cache shapes that cannot be gathered + """The backend should reject KV cache specs that cannot be gathered correctly by the ROCm kernel.""" from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend + from vllm.v1.kv_cache_interface import FullAttentionSpec - assert AiterFlashAttentionBackend.get_kv_cache_shape(8, 16, 8, 128) == ( - 8, - 8, - 16, - 256, - ) with pytest.raises(ValueError, match="Block size must be a multiple of 16"): - AiterFlashAttentionBackend.get_kv_cache_shape(8, 15, 8, 128) + AiterFlashAttentionBackend.customize_spec( + FullAttentionSpec( + block_size=15, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ) + ) + # block_size == 1 is the per-token page-size probe and must pass through. + AiterFlashAttentionBackend.customize_spec( + FullAttentionSpec( + block_size=1, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ) + ) def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe(): diff --git a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py index c49ceb03f5..bd267c43be 100644 --- a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py +++ b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py @@ -11,6 +11,7 @@ import pytest import torch from vllm.platforms import current_platform +from vllm.v1.kv_cache_interface import KVCacheLayout if current_platform.is_rocm(): pytest.skip( @@ -34,51 +35,40 @@ def to_float8(x, dtype=None): return x_scl_sat.to(dtype), scale.float().reciprocal() -def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size): - """Create a standard contiguous fp8 KV cache (HND layout).""" - raw = torch.randn( - num_blocks, - 2, - num_kv_heads, - block_size, - head_size, - dtype=torch.bfloat16, - device="cuda", - ) - kv_cache, scale = to_float8(raw) - return kv_cache, scale - - -def make_cross_layer_kv_cache( +def make_random_kv_cache( num_blocks, num_kv_heads, block_size, head_size, - num_layers=4, + layout=KVCacheLayout.LBHNC, + num_layers=None, ): """ - Create a non-contiguous per-layer view mimicking cross-layer allocation. + Create a random fp8 KV cache mimicking the production allocation. - Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size) + Physical layout: [L, B, H, N, 2 * hs] permuted by ``layout.stride_order`` Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size) - with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers). + with non-contiguous strides under cross-layer layouts (they skip over + num_layers), matching the actual forward path. """ - raw = torch.randn( + logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size) + if num_layers is None: + num_layers = 1 if layout.is_layer_compact else 2 + logical_5d = (num_layers, *logical_4d) + physical_5d = tuple(logical_5d[i] for i in layout.stride_order) + inv_order = [layout.stride_order.index(i) for i in range(5)] + + raw_phys = torch.randn(*physical_5d, dtype=torch.bfloat16, device="cuda") + fp8_phys, scale = to_float8(raw_phys) + fp8_4d = fp8_phys.permute(*inv_order)[0] + kv_5d = fp8_4d.view( num_blocks, - 2, num_kv_heads, - num_layers, block_size, + 2, head_size, - dtype=torch.bfloat16, - device="cuda", - ) - fp8_full, scale = to_float8(raw) - layer_view = fp8_full[:, :, :, 0, :, :] - assert not layer_view.is_contiguous(), ( - f"Expected non-contiguous view, got strides {layer_view.stride()}" - ) - return layer_view, scale + ).permute(0, 3, 1, 2, 4) + return kv_5d, scale 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): @pytest.mark.parametrize("block_size", [16, 32]) @pytest.mark.parametrize("batch_size", [1, 4]) @pytest.mark.parametrize("num_pages_per_seq", [3, 8]) -@pytest.mark.parametrize("contiguous", [True, False]) +@pytest.mark.parametrize("layout", list(KVCacheLayout)) @torch.inference_mode() def test_trtllm_kvfp8_dequant( num_kv_heads: int, @@ -122,7 +112,7 @@ def test_trtllm_kvfp8_dequant( block_size: int, batch_size: int, num_pages_per_seq: int, - contiguous: bool, + layout: KVCacheLayout, ): from vllm.v1.attention.backends.flashinfer import ( trtllm_prefill_attn_kvfp8_dequant, @@ -130,20 +120,13 @@ def test_trtllm_kvfp8_dequant( torch.set_default_device("cuda") - if contiguous: - kv_cache, scale = make_contiguous_kv_cache( - NUM_BLOCKS, - num_kv_heads, - block_size, - head_size, - ) - else: - kv_cache, scale = make_cross_layer_kv_cache( - NUM_BLOCKS, - num_kv_heads, - block_size, - head_size, - ) + kv_cache, scale = make_random_kv_cache( + NUM_BLOCKS, + num_kv_heads, + block_size, + head_size, + layout=layout, + ) k_scale = scale.clone() v_scale = scale.clone() @@ -187,7 +170,7 @@ def test_block_tables_with_zero_pages(): torch.set_default_device("cuda") num_kv_heads, block_size, head_size = 8, 16, 64 - kv_cache, scale = make_contiguous_kv_cache( + kv_cache, scale = make_random_kv_cache( NUM_BLOCKS, num_kv_heads, block_size, @@ -234,7 +217,7 @@ def test_all_zero_block_tables(): torch.set_default_device("cuda") num_kv_heads, block_size, head_size = 4, 16, 64 - kv_cache, scale = make_contiguous_kv_cache( + kv_cache, scale = make_random_kv_cache( NUM_BLOCKS, num_kv_heads, block_size, @@ -266,7 +249,7 @@ def test_different_k_v_scales(): torch.set_default_device("cuda") num_kv_heads, block_size, head_size = 8, 16, 64 - kv_cache, _ = make_contiguous_kv_cache( + kv_cache, _ = make_random_kv_cache( NUM_BLOCKS, num_kv_heads, block_size, @@ -299,7 +282,7 @@ def test_single_page_per_seq(): torch.set_default_device("cuda") num_kv_heads, block_size, head_size = 8, 16, 128 - kv_cache, scale = make_contiguous_kv_cache( + kv_cache, scale = make_random_kv_cache( NUM_BLOCKS, num_kv_heads, block_size, @@ -332,7 +315,7 @@ def test_large_page_indices(): num_kv_heads, block_size, head_size = 8, 16, 128 large_num_blocks = 32768 - kv_cache, scale = make_contiguous_kv_cache( + kv_cache, scale = make_random_kv_cache( large_num_blocks, num_kv_heads, block_size, @@ -369,7 +352,7 @@ def test_large_block_size(): torch.set_default_device("cuda") num_kv_heads, block_size, head_size = 4, 64, 128 - kv_cache, scale = make_contiguous_kv_cache( + kv_cache, scale = make_random_kv_cache( NUM_BLOCKS, num_kv_heads, block_size, @@ -401,7 +384,7 @@ def test_large_block_size(): def test_cross_layer_many_layers(): """ Non-contiguous with 36 layers -- matches real gpt-oss-120b. - Strides are far from contiguous (factor of 36 in the gaps). + Cross-layer (BLHNC) strides are far from contiguous (factor of 36 in the gaps). """ from vllm.v1.attention.backends.flashinfer import ( trtllm_prefill_attn_kvfp8_dequant, @@ -409,14 +392,14 @@ def test_cross_layer_many_layers(): torch.set_default_device("cuda") num_kv_heads, block_size, head_size = 8, 16, 64 - num_layers = 36 - kv_cache, scale = make_cross_layer_kv_cache( + kv_cache, scale = make_random_kv_cache( NUM_BLOCKS, num_kv_heads, block_size, head_size, - num_layers=num_layers, + layout=KVCacheLayout.BLHNC, + num_layers=36, ) k_scale = v_scale = scale.clone() diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 86fe732378..0ef3c33aeb 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from fractions import Fraction + import pytest import pytest_asyncio from mistral_common.protocol.transcription.request import ( @@ -75,7 +77,8 @@ def assert_encoder_kv_cache_spec(engine: LLM) -> None: assert audio_config.block_pool_size == 4 assert isinstance(spec, SlidingWindowSpec) assert spec.block_size == 16 - assert spec.num_kv_heads == 128 + assert spec.num_kv_heads == 32 + assert spec.tokens_per_state == Fraction(1, 4) # cdiv(750, 4) == 188 pooled tokens cover the model's window; the extra # +1 is an eviction margin (see whisper_causal.py get_kv_cache_spec). assert spec.sliding_window == cdiv(750, 4) + 1 == 189 diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index f36c4f1fc4..914d8d95ea 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -8,6 +8,7 @@ import pytest from vllm import LLM from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.attention.backends.utils import resolve_kv_cache_layout from vllm.v1.core.kv_cache_utils import ( generate_scheduler_kv_cache_config, get_kv_cache_configs, @@ -82,6 +83,12 @@ def can_initialize( # Avoid calling model.forward() def _initialize_kv_caches_v1(self, vllm_config): kv_cache_specs = self.model_executor.get_kv_cache_specs() + layout = resolve_kv_cache_layout( + vllm_config, + self.model_executor.get_supported_kv_cache_layouts(), + [spec for worker_specs in kv_cache_specs for spec in worker_specs.values()], + ) + self.model_executor.set_kv_cache_layout(layout.name) kv_cache_configs = get_kv_cache_configs( vllm_config, kv_cache_specs, diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index 04a365b48d..9f6beb3eb8 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -294,6 +294,7 @@ class TestTurboQuantKVCacheSpec: head_size_v=128, num_kv_heads=4, sliding_window=None, + get_attn_backend=lambda: TurboQuantAttentionBackend, ) vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32)) diff --git a/tests/test_attention_backend_registry.py b/tests/test_attention_backend_registry.py index 034749874d..6c5e0d7fa5 100644 --- a/tests/test_attention_backend_registry.py +++ b/tests/test_attention_backend_registry.py @@ -38,11 +38,6 @@ class CustomAttentionBackend(AttentionBackend): """Mock builder class.""" return None - @staticmethod - def get_required_kv_cache_layout(): - """Mock KV cache layout.""" - return None - class CustomMambaAttentionImpl(AttentionImpl): """Mock custom mamba attention implementation for testing.""" @@ -71,11 +66,6 @@ class CustomMambaAttentionBackend(AttentionBackend): """Mock builder class.""" return None - @staticmethod - def get_required_kv_cache_layout(): - """Mock KV cache layout.""" - return None - def test_custom_is_not_alias_of_any_backend(): # Get all members of AttentionBackendEnum diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index cb35864b70..cd78911ce8 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -31,10 +31,7 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.attention.backends.utils import ( - set_kv_cache_layout, -) -from vllm.v1.kv_cache_interface import FullAttentionSpec +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout BACKENDS_TO_TEST = [ AttentionBackendEnum.FLASH_ATTN, @@ -44,6 +41,14 @@ BACKENDS_TO_TEST = [ "FLEX_ATTENTION_SLOW", ] + +def _actual_backend(backend) -> AttentionBackendEnum: + """Resolve pseudo-backends (FLEX_ATTENTION_SLOW) to their real enum.""" + if backend == "FLEX_ATTENTION_SLOW": + return AttentionBackendEnum.FLEX_ATTENTION + return backend + + DEVICE_TYPE = current_platform.device_type # Use the platform's preferred FP8 type so the stored cache matches what the @@ -118,6 +123,7 @@ def create_and_prepopulate_kv_cache( device: torch.device, num_blocks: int, common_attn_metadata: CommonAttentionMetadata, + layout: KVCacheLayout, randomize_blocks: bool = True, kv_cache_dtype: str = "auto", ) -> torch.Tensor: @@ -126,19 +132,22 @@ def create_and_prepopulate_kv_cache( Args: k_contexts: List of key context tensors for each sequence v_contexts: List of value context tensors for each sequence - seq_lens: List of sequence lengths block_size: Size of each block num_kv_heads: Number of KV heads head_size: Size of each head dtype: Data type for the cache device: Device to create the cache on num_blocks: Total number of blocks in the cache - block_table: Block table tensor to populate + common_attn_metadata: Provides seq lens, block table and slot mapping + layout: Physical layout to allocate in; the cache is returned as the + logical ``[B, H, N, C]`` view (as ``create_kv_cache_views`` does) randomize_blocks: Whether to randomly permute blocks or use sequential order + kv_cache_dtype: Cache dtype string; fp8 caches use fp8 storage Returns: - Tuple of (kv_cache, updated_block_table) + A 4D tensor in logical ``(num_blocks, num_kv_heads, block_size, + 2 * head_size)`` order with strides determined by ``layout``. """ batch_size = len(k_contexts) seq_lens = common_attn_metadata.seq_lens.cpu() @@ -155,26 +164,33 @@ def create_and_prepopulate_kv_cache( fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype) storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype - kv_cache = torch.zeros( - num_blocks, - block_size, - num_kv_heads, - 2 * head_size, - dtype=storage_dtype, - device=device, - ) - kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size) + # Logical 5D shape is always [L, B, H, N, C]. Cross-layer layouts need + # at least two layers to reproduce the inter-layer gaps in a layer view. + logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size) + num_layers = 1 if layout.is_layer_compact else 2 + logical_5d = (num_layers, *logical_4d) + physical_5d = tuple(logical_5d[i] for i in layout.stride_order) + inv_order = [layout.stride_order.index(i) for i in range(5)] - # Populate the cache with the context tokens + kv_cache_physical = torch.zeros(physical_5d, dtype=storage_dtype, device=device) + # Permute to logical [L, B, H, N, C], then select a layer. This mirrors + # create_kv_cache_views and retains cross-layer strides in the 4D view. + kv_cache = kv_cache_physical.permute(*inv_order)[0] + + # Write context tokens into the cache via the logical view: + # kv_cache[block, :, token_in_block, :] routes correctly regardless + # of physical layout. # Start from block_id=1 since block_id=0 is considered the null block start_block_idx = 1 for i in range(batch_size): k_context, v_context = k_contexts[i], v_contexts[i] - start = start_block_idx * block_size - end = start + k_context.shape[0] - kv_cache_flat[start:end, :, :head_size] = k_context - kv_cache_flat[start:end, :, head_size:] = v_context - + t = torch.arange(k_context.shape[0], device=device) + blk = start_block_idx + t // block_size + off = t % block_size + # Advanced indexing on (blk, off) yields [T, H, hs] destinations. + # index_put is dtype-strict; cast like the scalar path would. + kv_cache[blk, :, off, :head_size] = k_context.to(kv_cache.dtype) + kv_cache[blk, :, off, head_size:] = v_context.to(kv_cache.dtype) # Stay block aligned and allocate enough blocks for the new tokens start_block_idx += cdiv(int(seq_lens[i]), block_size) @@ -214,9 +230,6 @@ def create_and_prepopulate_kv_cache( i, block_indices ] * block_size + token_inter_block_offsets.to(device) - # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs) - kv_cache = kv_cache.transpose(1, 2).contiguous() - if fp8_kv_cache: kv_cache = kv_cache.view(torch.uint8) @@ -236,6 +249,22 @@ class MockAttentionLayer: self._v_scale_float = 1.0 +def _clone_kv_cache_in_layout( + kv_cache: torch.Tensor, layout: KVCacheLayout +) -> torch.Tensor: + """Copy a logical [B, H, N, C] cache into a fresh allocation in `layout`.""" + logical_5d = (1, *kv_cache.shape) + physical = torch.zeros( + tuple(logical_5d[i] for i in layout.stride_order), + dtype=kv_cache.dtype, + device=kv_cache.device, + ) + inv_order = [layout.stride_order.index(i) for i in range(5)] + view = physical.permute(*inv_order)[0] + view.copy_(kv_cache) + return view + + def run_attention_backend( backend: AttentionBackendEnum, kv_cache_spec: FullAttentionSpec, @@ -254,18 +283,17 @@ def run_attention_backend( ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" - # Handle special case for FLEX_ATTENTION_SLOW - actual_backend = backend - - use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0") + use_direct_block_mask = not current_platform.is_rocm() and is_torch_equal_or_newer( + "2.9.0.dev0" + ) if backend == "FLEX_ATTENTION_SLOW": - actual_backend = AttentionBackendEnum.FLEX_ATTENTION use_direct_block_mask = False + backend = _actual_backend(backend) - builder_cls, impl_cls = try_get_attention_backend(actual_backend) + builder_cls, impl_cls = try_get_attention_backend(backend) # Mock flashinfer's get_per_layer_parameters if needed - if actual_backend == AttentionBackendEnum.FLASHINFER: + if backend == AttentionBackendEnum.FLASHINFER: import unittest.mock from vllm.v1.attention.backends.utils import PerLayerParameters @@ -295,7 +323,7 @@ def run_attention_backend( else: # Build metadata builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device) - if actual_backend == AttentionBackendEnum.FLEX_ATTENTION: + if backend == AttentionBackendEnum.FLEX_ATTENTION: builder.direct_build = use_direct_block_mask attn_metadata = builder.build( common_prefix_len=0, @@ -311,17 +339,19 @@ def run_attention_backend( ) head_size = vllm_config.model_config.get_head_size() scale = 1.0 / (head_size**0.5) - impl = impl_cls( - num_heads=num_heads, - head_size=head_size, - scale=scale, - num_kv_heads=num_kv_heads, - alibi_slopes=None, - sliding_window=sliding_window, - attn_type=attn_type, - kv_cache_dtype=kv_cache_dtype, - **({"sinks": sinks} if sinks is not None else {}), - ) + # Impls capture the current vllm config at construction, as in model loading. + with set_current_vllm_config(vllm_config): + impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + scale=scale, + num_kv_heads=num_kv_heads, + alibi_slopes=None, + sliding_window=sliding_window, + attn_type=attn_type, + kv_cache_dtype=kv_cache_dtype, + **({"sinks": sinks} if sinks is not None else {}), + ) # Create mock layer and output buffer mock_layer = MockAttentionLayer(device) @@ -333,7 +363,7 @@ def run_attention_backend( # Run forward pass # NOTE: The query, key, and value are already shaped correctly # in the calling test function. - if not try_backend_includes_kv_cache_update(actual_backend): + if not try_backend_includes_kv_cache_update(backend): impl.do_kv_cache_update( mock_layer, key, value, kv_cache, attn_metadata.slot_mapping ) @@ -347,7 +377,7 @@ def run_attention_backend( def _test_backend_correctness( batch_spec: BatchSpec, model: str, - backend_to_test: list[AttentionBackendEnum | str], + backend_to_test: list[AttentionBackendEnum], mask_mod, *, causal: bool = True, @@ -358,6 +388,7 @@ def _test_backend_correctness( tensor_parallel_size: int = 1, kv_cache_dtype: str = "auto", use_sinks: bool = False, + layout: KVCacheLayout | None = None, ): """ Test that all backends produce similar outputs to a reference implementation @@ -535,6 +566,22 @@ def _test_backend_correctness( common_attn_metadata.causal = causal # 3. Simulate Paged KV Cache and a realistic slot_mapping + # Mirror selector-time resolution locally (caller-requested layout, then + # any backend-required layout, then the default chain). + declared = [ + supported + for backend in backend_to_test + if ( + supported := _actual_backend(backend) + .get_class() + .supported_kv_cache_layouts() + ) + is not None + ] + if layout is None: + # Mirror the resolver: the shared cache uses the declared sets' preferred + # layout; backends that don't support it get a per-backend copy below. + layout = min(declared, key=len)[0] if declared else KVCacheLayout.LBNHC kv_cache = create_and_prepopulate_kv_cache( k_contexts=k_contexts, v_contexts=v_contexts, @@ -545,6 +592,7 @@ def _test_backend_correctness( device=device, num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000, common_attn_metadata=common_attn_metadata, + layout=layout, randomize_blocks=True, kv_cache_dtype=kv_cache_dtype, ) @@ -553,62 +601,44 @@ def _test_backend_correctness( # Note: flex_attention has known Triton kernel compatibility issues # with test infrastructures for backend_name in backend_to_test: - reset_kv_cache_layout = False - - # Resolve backend class for both enum and string names. - actual_backend = backend_name - if backend_name == "FLEX_ATTENTION_SLOW": - actual_backend = AttentionBackendEnum.FLEX_ATTENTION - if hasattr(actual_backend, "get_class"): - backend_cls = actual_backend.get_class() - else: - backend_cls = None + backend_cls = _actual_backend(backend_name).get_class() if is_quantized_kv_cache(kv_cache_dtype) and ( - backend_cls is None - or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype) + not backend_cls.supports_kv_cache_dtype(kv_cache_dtype) ): continue - if backend_name == AttentionBackendEnum.FLASHINFER: - set_kv_cache_layout("HND") - reset_kv_cache_layout = True - kv_cache_for_backend = kv_cache - if backend_cls is not None: - try: - stride_order = backend_cls.get_kv_cache_stride_order() - except (AttributeError, NotImplementedError): - stride_order = tuple(range(kv_cache.ndim)) - if stride_order != tuple(range(kv_cache.ndim)): - # Apply stride order like runtime does in - # _reshape_kv_cache (attn_utils.py:182-210): permute to physical - # layout, make contiguous, then permute to logical layout. - inv_order = [stride_order.index(i) for i in range(len(stride_order))] - kv_cache_for_backend = ( - kv_cache.permute(*stride_order).contiguous().permute(*inv_order) - ) + backend_layout = layout - try: - backend_output = run_attention_backend( - backend_name, - kv_cache_spec, - ["placeholder"], - vllm_config, - device, - common_attn_metadata, - query_vllm, - key_vllm, - value_vllm, - kv_cache_for_backend, - sliding_window=sliding_window, - attn_type=attn_type, - kv_cache_dtype=kv_cache_dtype, - sinks=sinks, - ) - finally: - if reset_kv_cache_layout: - set_kv_cache_layout(None) + backend_supported = backend_cls.supported_kv_cache_layouts() + if backend_supported is not None and backend_layout not in backend_supported: + # Production resolution never pairs this backend with the shared layout + # (e.g. flex is LBNHC-only next to FlashInfer's LBHNC); give it a copy + # of the cache in a layout it supports. + backend_layout = backend_supported[0] + kv_cache_for_backend = _clone_kv_cache_in_layout(kv_cache, backend_layout) + + # FlashInfer reads the layout at plan time; set it to match + # the physical order of the test cache. + vllm_config.cache_config.kv_cache_layout = backend_layout.name + + backend_output = run_attention_backend( + backend_name, + kv_cache_spec, + ["placeholder"], + vllm_config, + device, + common_attn_metadata, + query_vllm, + key_vllm, + value_vllm, + kv_cache_for_backend, + sliding_window=sliding_window, + attn_type=attn_type, + kv_cache_dtype=kv_cache_dtype, + sinks=sinks, + ) # Check shape and dtype consistency assert backend_output.shape == sdpa_output.shape, ( @@ -637,6 +667,38 @@ def _test_backend_correctness( ) +@pytest.mark.parametrize("layout", ["BLHNC", "BHLNC"]) +@pytest.mark.parametrize("batch_spec_name", ["small_decode", "small_prefill"]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_flashinfer_cross_layer_layout( + default_vllm_config, + layout: str, + batch_spec_name: str, + kv_cache_dtype: str, +): + if AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST: + pytest.skip("FlashInfer is not installed") + + def causal_mask_mod( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + kv_idx: torch.Tensor, + *, + context_len: int, + ): + return (q_idx + context_len) >= kv_idx + + _test_backend_correctness( + batch_spec=BATCH_SPECS[batch_spec_name], + model="meta-llama/Meta-Llama-3-8B", + backend_to_test=[AttentionBackendEnum.FLASHINFER], + mask_mod=causal_mask_mod, + kv_cache_dtype=kv_cache_dtype, + layout=KVCacheLayout[layout], + ) + + @pytest.mark.parametrize( "batch_spec_name", [ @@ -984,14 +1046,12 @@ if current_platform.is_rocm(): SLIDING_WINDOW_BACKENDS_TO_TEST = [ AttentionBackendEnum.FLEX_ATTENTION, AttentionBackendEnum.TRITON_ATTN, - "FLEX_ATTENTION_SLOW", ] else: SLIDING_WINDOW_BACKENDS_TO_TEST = [ AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.FLEX_ATTENTION, AttentionBackendEnum.TRITON_ATTN, - "FLEX_ATTENTION_SLOW", ] @@ -1009,7 +1069,10 @@ else: @pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) def test_sliding_window_backend_correctness( - default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int + default_vllm_config, + batch_spec_name: str, + model: str, + tensor_parallel_size: int, ): """Test backend's correctness with sliding window attention.""" @@ -1071,7 +1134,10 @@ def test_sliding_window_backend_correctness( @pytest.mark.parametrize("model", ["google/embeddinggemma-300m"]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2]) def test_sliding_window_encoder_backend_correctness( - default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int + default_vllm_config, + batch_spec_name: str, + model: str, + tensor_parallel_size: int, ): """Test backend's correctness with sliding window attention.""" @@ -1107,7 +1173,6 @@ def test_sliding_window_encoder_backend_correctness( NON_CAUSAL_BACKENDS_TO_TEST = [ AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.FLEX_ATTENTION, - "FLEX_ATTENTION_SLOW", ] if current_platform.is_rocm(): @@ -1128,7 +1193,9 @@ if current_platform.is_rocm(): ) @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"]) def test_non_causal_backend_correctness( - default_vllm_config, batch_spec_name: str, model: str + default_vllm_config, + batch_spec_name: str, + model: str, ): """Test backend's correctness with non-causal (bidirectional) decoder attention, as used by DFlash speculative decoding.""" diff --git a/tests/v1/attention/test_group_head_counts.py b/tests/v1/attention/test_group_head_counts.py index 9691926819..b799bccc29 100644 --- a/tests/v1/attention/test_group_head_counts.py +++ b/tests/v1/attention/test_group_head_counts.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Attention metadata geometry must come from the builder's own group.""" +"""Scheduler metadata sizes a scratchpad from the query head count, so it must +come from the builder's own group: the model-wide ``get_num_attention_heads()`` +is wrong for models that vary it per layer (e.g. Laguna), and too small a +scratchpad is indexed past its end. +""" from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -13,9 +17,8 @@ from vllm.v1.attention.backends.cpu_attn import ( CPUAttentionBackendImpl, CPUAttentionMetadataBuilder, ) -from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder -requires_cpu = pytest.mark.skipif( +pytestmark = pytest.mark.skipif( not current_platform.is_cpu(), reason="CPU attention backend" ) @@ -66,7 +69,6 @@ def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder: ) -@requires_cpu @pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16]) def test_num_heads_comes_from_the_group(group_num_heads): """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): assert builder.num_heads == group_num_heads -@requires_cpu def test_mixed_head_counts_in_one_group_are_rejected(): """Grouping guarantees uniformity; a mixed group means that broke.""" with pytest.raises(AssertionError, match="share num_heads"): _build([MODEL_WIDE_NUM_HEADS, 64]) - - -def test_flash_attention_geometry_comes_from_the_group(): - """FA3 must use group geometry for layers without an ``impl`` wrapper.""" - layers = {f"layer_{i}": SimpleNamespace(num_heads=16) for i in range(2)} - vllm_config = MagicMock() - vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS - vllm_config.model_config.get_num_kv_heads.return_value = NUM_KV_HEADS - vllm_config.model_config.get_head_size.return_value = 128 - vllm_config.model_config.rswa_window = None - vllm_config.model_config.is_mm_prefix_lm = False - vllm_config.parallel_config.cp_kv_cache_interleave_size = 1 - vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs.return_value = ( - False - ) - vllm_config.compilation_config.max_cudagraph_capture_size = None - kv_cache_spec = SimpleNamespace( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.bfloat16, - ) - - with ( - patch( - "vllm.v1.attention.backends.utils.get_layers_from_vllm_config", - return_value=layers, - ), - patch( - "vllm.distributed.parallel_state.get_dcp_group", - side_effect=AssertionError, - ), - patch( - "vllm.v1.attention.backends.flash_attn.get_flash_attn_version", - return_value=3, - ), - ): - builder = FlashAttentionMetadataBuilder( - kv_cache_spec=kv_cache_spec, - layer_names=list(layers), - vllm_config=vllm_config, - device=torch.device("cpu"), - ) - - assert builder.aot_schedule - assert builder.num_heads_q == 16 - assert builder.num_heads_kv == 2 - assert builder.headdim == 64 diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 03cbd0dc36..06625e29e2 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -34,19 +34,19 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_size(): +def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_num_states(): """Regression test: DeepseekV4 compression path must compute slot_mapping from compressed positions, not reuse the uncompressed common metadata mapping. """ device = torch.device("cuda") - # storage_block_size = block_size // compress_ratio = 256 // 4 = 64 + # num_states = block_size // tokens_per_state = 256 // 4 = 64 kv_cache_spec = MLAAttentionSpec( block_size=256, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, - compress_ratio=4, + tokens_per_state=4, ) vllm_config = create_vllm_config(max_model_len=1024) 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_ valid_slots = md.slot_mapping[md.slot_mapping >= 0] assert valid_slots.numel() == 10 # 40 tokens / compress_ratio 4 - storage_bs = kv_cache_spec.storage_block_size # 64 + storage_bs = kv_cache_spec.num_states # 64 # Compressed positions 60..63 land in block 5, positions 64..69 in block 7. expected = torch.tensor( [ diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index e5c32aa490..5c621cf6ad 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -250,6 +250,15 @@ def _prefill_backend_dimension_params(): ) ) ) + elif ( + current_platform.is_rocm() + and prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN + ): + # Pre-existing on main: flash-attn MLA prefill returns wrong + # outputs and faults on ROCm, so it fails on the base commit too. + marks.append( + pytest.mark.skip(reason="FLASH_ATTN MLA prefill broken on ROCm") + ) params.append( pytest.param( prefill_backend, @@ -392,7 +401,7 @@ def create_and_prepopulate_kv_cache( kv_entry_size = head_size kv_cache = torch.zeros( - num_blocks, block_size, kv_entry_size, dtype=torch.uint8, device=device + num_blocks, 1, block_size, kv_entry_size, dtype=torch.uint8, device=device ) scale_tensor = ( scale @@ -401,9 +410,9 @@ def create_and_prepopulate_kv_cache( ) scale_tensor = scale_tensor.to(device=device, dtype=torch.float32) else: - # Create MLA KV cache: (num_blocks, block_size, head_size) + # Create MLA KV cache: (num_blocks, num_heads=1, block_size, head_size) kv_cache = torch.zeros( - num_blocks, block_size, head_size, dtype=dtype, device=device + num_blocks, 1, block_size, head_size, dtype=dtype, device=device ) kv_cache_flat = kv_cache.view(-1, head_size) @@ -424,7 +433,7 @@ def create_and_prepopulate_kv_cache( ops.concat_and_cache_mla( kv_c_context, k_pe_context.squeeze(1), - kv_cache, + kv_cache.squeeze(1), slots, kv_cache_dtype=kv_cache_dtype, scale=scale_tensor, @@ -543,6 +552,10 @@ class MockSparseMLAAttentionLayer: kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto") fp8_attention = kv_cache_dtype.startswith("fp8") + # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache. + if kv_cache.ndim == 4: + kv_cache = kv_cache.squeeze(1) + # Write to KV cache if kv_cache.numel() > 0: ops.concat_and_cache_mla( @@ -678,6 +691,10 @@ class MockMLAAttentionLayer(MLAAttention): output: torch.Tensor, ) -> torch.Tensor: """Replicates MLAAttention.forward_impl logic for testing.""" + # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache. + if kv_cache.ndim == 4: + kv_cache = kv_cache.squeeze(1) + # Write to KV cache kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto") fp8_attention = kv_cache_dtype.startswith("fp8") diff --git a/tests/v1/attention/test_mm_prefix.py b/tests/v1/attention/test_mm_prefix.py index e837121212..6b8717389f 100644 --- a/tests/v1/attention/test_mm_prefix.py +++ b/tests/v1/attention/test_mm_prefix.py @@ -20,7 +20,10 @@ import pytest import torch from vllm.platforms import current_platform -from vllm.v1.attention.backends.utils import fill_mm_prefix_query_ranges +from vllm.v1.attention.backends.utils import ( + fill_mm_prefix_query_ranges, +) +from vllm.v1.kv_cache_interface import KVCacheLayout def _fa4_available() -> bool: @@ -585,6 +588,7 @@ def test_mm_prefix_kv_cache_path(head_size: int): device=DEVICE, num_blocks=2048, common_attn_metadata=common, + layout=KVCacheLayout.LBNHC, randomize_blocks=True, ) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index ac9a87d08c..be58cf8b46 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -813,7 +813,7 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens q = torch.zeros(num_mqa_tokens, 4, 576, dtype=torch.bfloat16, device=device) - kv_cache = torch.zeros(40 * block_size, 576, dtype=torch.bfloat16, device=device) + kv_cache = torch.zeros(40, block_size, 576, dtype=torch.bfloat16, device=device) topk_indices = torch.randint( 0, block_size * 10, @@ -1157,7 +1157,8 @@ def test_sparse_backend_prefill_correctness( q=query_cat, kv_c_normed=kv_c_cat, k_pe=k_pe_cat, - kv_c_and_k_pe_cache=kv_cache, + # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache. + kv_c_and_k_pe_cache=kv_cache.squeeze(1), attn_metadata=metadata, k_scale=torch.tensor(1.0, device=device), output=out_buffer, diff --git a/tests/v1/attention/test_trtllm_attention_integration.py b/tests/v1/attention/test_trtllm_attention_integration.py index c8a080e6d1..d81bc2de4e 100644 --- a/tests/v1/attention/test_trtllm_attention_integration.py +++ b/tests/v1/attention/test_trtllm_attention_integration.py @@ -9,6 +9,7 @@ import pytest import torch from torch.nn.attention.flex_attention import create_block_mask, flex_attention +from tests.v1.attention.test_attention_backends import create_and_prepopulate_kv_cache from tests.v1.attention.utils import ( BatchSpec, create_common_attn_metadata, @@ -16,14 +17,9 @@ from tests.v1.attention.utils import ( ) from vllm.config import set_current_vllm_config from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim, set_random_seed -from vllm.v1.attention.backends.utils import ( - PerLayerParameters, - get_kv_cache_layout, - set_kv_cache_layout, -) -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode +from vllm.v1.attention.backends.utils import PerLayerParameters +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode if not current_platform.is_device_capability_family(100): pytest.skip( @@ -87,96 +83,6 @@ def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): } -def _create_hnd_kv_cache( - k_contexts, - v_contexts, - block_size, - num_kv_heads, - head_size, - dtype, - device, - num_blocks, - common_attn_metadata, - kv_in_head_dim=False, -): - """Create and populate a packed KV cache with HND-compatible strides. - - When kv_in_head_dim=False (default), returns (B, H, N, 2*hs) with K/V - packed in the content dim. When kv_in_head_dim=True, returns - (B, 2*H, N, hs) with K/V as separate head groups. - """ - seq_lens = common_attn_metadata.seq_lens.cpu() - query_lens = ( - common_attn_metadata.query_start_loc_cpu[1:] - - common_attn_metadata.query_start_loc_cpu[:-1] - ) - block_table = common_attn_metadata.block_table_tensor - slot_mapping = common_attn_metadata.slot_mapping - batch_size = len(k_contexts) - - # kv_in_head_dim: (B, N, 2*H, hs) — K/V as separate head groups - # else: (B, N, H, 2*hs) — K/V packed in content dim - n_heads, content = ( - (2 * num_kv_heads, head_size) - if kv_in_head_dim - else (num_kv_heads, 2 * head_size) - ) - kv_cache = torch.zeros( - num_blocks, - block_size, - n_heads, - content, - dtype=dtype, - device=device, - ) - kv_cache_flat = kv_cache.view(-1, n_heads, content) - - start_block_idx = 1 - for i in range(batch_size): - k_ctx, v_ctx = k_contexts[i], v_contexts[i] - start = start_block_idx * block_size - end = start + k_ctx.shape[0] - if kv_in_head_dim: - kv_cache_flat[start:end, :num_kv_heads] = k_ctx - kv_cache_flat[start:end, num_kv_heads:] = v_ctx - else: - kv_cache_flat[start:end, :, :head_size] = k_ctx - kv_cache_flat[start:end, :, head_size:] = v_ctx - start_block_idx += cdiv(int(seq_lens[i]), block_size) - - blocks_end = start_block_idx - - # Randomly permute blocks (starting from block 1; block 0 is null). - perm = torch.randperm(blocks_end - 1) + 1 - inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device) - inv_perm[1:] = torch.argsort(perm) + 1 - kv_cache[1:blocks_end] = kv_cache[perm] - - # Build block table. - start_block_idx = 1 - for i in range(batch_size): - n_blocks = cdiv(int(seq_lens[i]), block_size) - block_table[i, :n_blocks] = inv_perm[ - start_block_idx : start_block_idx + n_blocks - ] - start_block_idx += n_blocks - - # Build slot mapping that is consistent with the block table. - for i in range(batch_size): - ctx_len = int(seq_lens[i]) - int(query_lens[i]) - token_offsets = torch.arange(int(query_lens[i])) + ctx_len - block_indices = token_offsets // block_size - intra_block_offsets = token_offsets % block_size - start = common_attn_metadata.query_start_loc_cpu[i] - end = common_attn_metadata.query_start_loc_cpu[i + 1] - slot_mapping[start:end] = block_table[ - i, block_indices - ] * block_size + intra_block_offsets.to(device) - - # Transpose to canonical: (B, H, N, 2*hs) or (B, 2*H, N, hs) - return kv_cache.transpose(1, 2).contiguous() - - def _create_nvfp4_hnd_kv_cache( k_contexts, v_contexts, @@ -191,7 +97,7 @@ def _create_nvfp4_hnd_kv_cache( ): """Create an nvfp4 KV cache by quantizing bf16 context via reshape_and_cache_flash, using the same block-table layout as - _create_hnd_kv_cache. + create_and_prepopulate_kv_cache. The returned tensor is dtype ``uint8`` with head-group layout ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)`` @@ -217,11 +123,10 @@ def _create_nvfp4_hnd_kv_cache( during quantization. Returns: - ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided). + ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, LBHNC-strided). """ - # First create a bf16 HND cache so block tables are populated. - # Use kv_in_head_dim=True so K/V are separate head groups (B, 2*H, N, hs). - bf16_cache = _create_hnd_kv_cache( + # First create a bf16 cache so block tables are populated. + bf16_cache = create_and_prepopulate_kv_cache( k_contexts, v_contexts, block_size, @@ -231,7 +136,7 @@ def _create_nvfp4_hnd_kv_cache( device, num_blocks, common_attn_metadata, - kv_in_head_dim=True, + layout=KVCacheLayout.LBHNC, ) # (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( k_cache, v_cache = nvfp4_cache.split(num_kv_heads, dim=1) # Flatten bf16 context into tokens and quantize via reshape_and_cache_flash. - # bf16_cache is (B, 2*H, N, hs); split K/V on head dim. + # bf16_cache is (B, H, N, 2*hs); split K/V on the content dim. block_table = common_attn_metadata.block_table_tensor seq_lens = common_attn_metadata.seq_lens.cpu() query_lens = ( @@ -260,10 +165,16 @@ def _create_nvfp4_hnd_kv_cache( # Gather context tokens from the bf16 cache using block table. n_ctx_blocks = (ctx_len + block_size - 1) // block_size blocks = block_table[i, :n_ctx_blocks] - # bf16_cache is (B, 2*H, N, hs); split K and V head groups. - k_bf16, v_bf16 = bf16_cache[blocks].split(num_kv_heads, dim=1) - k_ctx = k_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len] - v_ctx = v_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len] + k_ctx = ( + bf16_cache[blocks, :, :, :head_size] + .transpose(1, 2) + .reshape(-1, num_kv_heads, head_size)[:ctx_len] + ) + v_ctx = ( + bf16_cache[blocks, :, :, head_size:] + .transpose(1, 2) + .reshape(-1, num_kv_heads, head_size)[:ctx_len] + ) # Build slot mapping for these context tokens. token_offsets = torch.arange(ctx_len, device=device) block_indices = token_offsets // block_size @@ -383,7 +294,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL) ) else: kv_scale_val = 1.0 - kv_cache = _create_hnd_kv_cache( + kv_cache = create_and_prepopulate_kv_cache( k_contexts, v_contexts, BLOCK_SIZE, @@ -393,127 +304,117 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL) device, NUM_GPU_BLOCKS, common_attn_metadata, + layout=KVCacheLayout.LBHNC, ) # 3. Run through FlashInfer with TRTLLM enabled - set_kv_cache_layout("HND") - get_kv_cache_layout.cache_clear() + vllm_config.cache_config.kv_cache_layout = "LBHNC" - try: - is_nvfp4 = kv_cache_dtype == "nvfp4" - kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE - spec_dtype = torch.uint8 if is_nvfp4 else dtype - kv_cache_spec = FullAttentionSpec( - block_size=BLOCK_SIZE, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=spec_dtype, - kv_quant_mode=kv_quant_mode, + is_nvfp4 = kv_cache_dtype == "nvfp4" + kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE + spec_dtype = torch.uint8 if is_nvfp4 else dtype + kv_cache_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=spec_dtype, + kv_quant_mode=kv_quant_mode, + ) + layer_names = ["test_layer_0"] + + with ( + set_current_vllm_config(vllm_config), + unittest.mock.patch( + "vllm.utils.flashinfer.supports_trtllm_attention", + return_value=True, + ), + unittest.mock.patch( + "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters", + _mock_get_per_layer_parameters, + ), + ): + builder = FlashInferMetadataBuilder( + kv_cache_spec, layer_names, vllm_config, device + ) + attn_metadata = builder.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, ) - layer_names = ["test_layer_0"] - with ( - set_current_vllm_config(vllm_config), - unittest.mock.patch( - "vllm.utils.flashinfer.supports_trtllm_attention", - return_value=True, - ), - unittest.mock.patch( - "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters", - _mock_get_per_layer_parameters, - ), - ): - builder = FlashInferMetadataBuilder( - kv_cache_spec, layer_names, vllm_config, device + # Verify the correct TRTLLM metadata types were produced. + has_prefills = any(ql > 1 for ql in batch_spec.query_lens) + has_decodes = any(ql == 1 for ql in batch_spec.query_lens) + + if has_prefills: + assert isinstance(attn_metadata.prefill, TRTLLMPrefill), ( + f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}" ) - attn_metadata = builder.build( - common_prefix_len=0, - common_attn_metadata=common_attn_metadata, + if has_decodes: + assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), ( + f"Expected FlashInferTrtllmAPIDecode, got {type(attn_metadata.decode)}" ) + assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN - # Verify the correct TRTLLM metadata types were produced. - has_prefills = any(ql > 1 for ql in batch_spec.query_lens) - has_decodes = any(ql == 1 for ql in batch_spec.query_lens) + impl = FlashInferImpl( + num_heads=num_q_heads, + head_size=head_size, + scale=scale, + num_kv_heads=num_kv_heads, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype=kv_cache_dtype, + ) - if has_prefills: - assert isinstance(attn_metadata.prefill, TRTLLMPrefill), ( - f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}" - ) - if has_decodes: - assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), ( - "Expected FlashInferTrtllmAPIDecode, got " - f"{type(attn_metadata.decode)}" - ) - assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN - - impl = FlashInferImpl( - num_heads=num_q_heads, - head_size=head_size, - scale=scale, - num_kv_heads=num_kv_heads, - alibi_slopes=None, - sliding_window=None, - kv_cache_dtype=kv_cache_dtype, - ) - - mock_layer = MockAttentionLayer(device) - if is_nvfp4: - # For nvfp4, k_scale/v_scale are the global quantization - # scales (amax/448) used by reshape_and_cache_flash. - kv_scale_t = torch.tensor( - kv_scale_val, dtype=torch.float32, device=device - ) - mock_layer._k_scale = kv_scale_t - mock_layer._v_scale = kv_scale_t - mock_layer._k_scale_float = kv_scale_val - mock_layer._v_scale_float = kv_scale_val - output = torch.empty_like(query_vllm) - - impl.do_kv_cache_update( - mock_layer, - key_vllm, - value_vllm, - kv_cache, - attn_metadata.slot_mapping, - ) - - # nvfp4 trtllm kernel requires FP8 queries. In the real - # pipeline the attention layer handles this; here we - # quantize manually. - if is_nvfp4: - finfo = torch.finfo(torch.float8_e4m3fn) - q_amax = query_vllm.abs().amax().clamp(min=1e-12) - q_s = (finfo.max / q_amax * 0.1).item() - query_vllm = ( - (query_vllm * q_s) - .clamp(finfo.min, finfo.max) - .to(torch.float8_e4m3fn) - ) - mock_layer._q_scale = torch.tensor( - 1.0 / q_s, dtype=torch.float32, device=device - ) - mock_layer._q_scale_float = 1.0 / q_s - - output = impl.forward( - mock_layer, - query_vllm, - key_vllm, - value_vllm, - kv_cache, - attn_metadata, - output=output, - ) - - # 4. Compare against SDPA reference + mock_layer = MockAttentionLayer(device) if is_nvfp4: - atol, rtol = 1.0, 1.0 # nvfp4 has higher quantization error - else: - atol, rtol = 1e-2, 1e-2 - torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol) + # For nvfp4, k_scale/v_scale are the global quantization + # scales (amax/448) used by reshape_and_cache_flash. + kv_scale_t = torch.tensor(kv_scale_val, dtype=torch.float32, device=device) + mock_layer._k_scale = kv_scale_t + mock_layer._v_scale = kv_scale_t + mock_layer._k_scale_float = kv_scale_val + mock_layer._v_scale_float = kv_scale_val + output = torch.empty_like(query_vllm) - finally: - set_kv_cache_layout(None) - get_kv_cache_layout.cache_clear() + impl.do_kv_cache_update( + mock_layer, + key_vllm, + value_vllm, + kv_cache, + attn_metadata.slot_mapping, + ) + + # nvfp4 trtllm kernel requires FP8 queries. In the real + # pipeline the attention layer handles this; here we + # quantize manually. + if is_nvfp4: + finfo = torch.finfo(torch.float8_e4m3fn) + q_amax = query_vllm.abs().amax().clamp(min=1e-12) + q_s = (finfo.max / q_amax * 0.1).item() + query_vllm = ( + (query_vllm * q_s).clamp(finfo.min, finfo.max).to(torch.float8_e4m3fn) + ) + mock_layer._q_scale = torch.tensor( + 1.0 / q_s, dtype=torch.float32, device=device + ) + mock_layer._q_scale_float = 1.0 / q_s + + output = impl.forward( + mock_layer, + query_vllm, + key_vllm, + value_vllm, + kv_cache, + attn_metadata, + output=output, + ) + + # 4. Compare against SDPA reference + if is_nvfp4: + atol, rtol = 1.0, 1.0 # nvfp4 has higher quantization error + else: + atol, rtol = 1e-2, 1e-2 + torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol) @pytest.mark.parametrize( diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py index 20562cfe85..6aeeb12b4b 100644 --- a/tests/v1/attention/utils.py +++ b/tests/v1/attention/utils.py @@ -32,7 +32,12 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import ( EncoderOnlyAttentionSpec, FullAttentionSpec, + KVCacheLayout, + KVCacheSpec, + KVCacheTensor, MambaSpec, + compute_layout_strides, + create_kv_cache_views, get_kv_quant_mode, ) @@ -215,6 +220,9 @@ def create_vllm_config( # (these may be set during initialization normally) cache_config.num_gpu_blocks = num_gpu_blocks cache_config.num_cpu_blocks = 0 + # Builders read the resolved layout from the config; tests that exercise a + # different layout overwrite this field. + cache_config.kv_cache_layout = "LBNHC" parallel_config = ParallelConfig( tensor_parallel_size=tensor_parallel_size, @@ -417,3 +425,31 @@ class MockMambaBuilder(BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetad is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool) ) return builder.build(0, common_metadata) + + +def dense_kv_cache_views( + raw: torch.Tensor, + spec: KVCacheSpec, + num_blocks: int, + num_layers: int, + layout: KVCacheLayout, + kernel_block_size: int | None = None, +) -> list[torch.Tensor]: + """``create_kv_cache_views`` for a dense allocation of ``num_layers`` layers.""" + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, num_blocks, num_layers, layout + ) + tensor = KVCacheTensor( + size=raw.numel(), + layers=[str(i) for i in range(num_layers)], + layer_stride=layer_stride, + block_stride=block_stride, + ) + return create_kv_cache_views( + raw, + spec, + num_blocks, + layout, + tensor, + kernel_block_size=kernel_block_size, + ) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 40906f6213..1be6078043 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -1,380 +1,237 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for contiguous KV cache packing.""" +"""Tests for contiguous KV cache packing. + +Every cache group packs its layers densely into one block; groups overlay each other +(a block ID is owned by one group at a time), so the packed block stride is the +largest group's packing. The layout decides whether the layer dim sits outside the +block dim (a contiguous region per layer) or inside it (all layers' pages within each +block); the allocation is the same either way. +""" from unittest.mock import MagicMock import pytest import torch +from vllm.config import CacheConfig from vllm.v1.core.kv_cache_utils import ( - _get_kv_cache_config_packed, - _get_kv_cache_groups_uniform_groups, + _get_kv_cache_bytes_per_block, + _pool_bytes_per_block, get_kv_cache_config_from_groups, ) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, - KVCacheTensor, + KVCacheLayout, MLAAttentionSpec, - SlidingWindowMLASpec, - SlidingWindowSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.utils import allocate_kv_cache + +MEMORY = 8 * 1024 * 1024 -def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec: +def _mla(head_size: int) -> MLAAttentionSpec: return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=512, - dtype=torch.uint8, - page_size_padded=page_size, - cache_dtype_str="fp8_ds_mla", - model_version="deepseek_v4", - alignment=576, + block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8 ) -def _make_full_spec() -> FullAttentionSpec: +def _full() -> FullAttentionSpec: return FullAttentionSpec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float16, + block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float16 ) -def _make_sw_spec() -> SlidingWindowSpec: - return SlidingWindowSpec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float16, - sliding_window=128, +def _uniform_group(specs: dict) -> KVCacheGroupSpec: + return KVCacheGroupSpec( + list(specs), + UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs), ) -def _make_groups(n_c4, n_c128, n_swa): - PS_C4_MLA = 37440 - PS_C4_IDX = 8640 - PS_C128 = 1728 - PS_SWA = 37440 - - mla_specs = {} - for i in range(n_c4): - mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA) - mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX) - for i in range(n_c128): - mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128) - - mla_group = KVCacheGroupSpec( - layer_names=list(mla_specs.keys()), - kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs), - ) - - swa_specs = {} - for i in range(n_swa): - swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA) - - swa_group = KVCacheGroupSpec( - layer_names=list(swa_specs.keys()), - kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs), - ) - - return [mla_group, swa_group] +def _mixed_page_groups(n_mla=3, n_idx=3, n_swa=5): + """A DeepSeek-V4-style hybrid: two groups with different page mixes.""" + g1 = {f"mla.{i}": _mla(512) for i in range(n_mla)} + g1.update({f"idx.{i}": _mla(128) for i in range(n_idx)}) + g2 = {f"swa.{i}": _mla(512) for i in range(n_swa)} + return [_uniform_group(g1), _uniform_group(g2)], g1, g2 -def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None): +def _mock_vllm_config(layout: str | None): config = MagicMock() + config.cache_config = CacheConfig() config.cache_config.num_gpu_blocks_override = None - config.cache_config.prefix_cache_retention_interval = 0 - config.kv_transfer_config = None - if kv_connector_extra_config is not None: - config.kv_transfer_config = MagicMock() - config.kv_transfer_config.kv_connector_extra_config = kv_connector_extra_config + config.cache_config.kv_cache_layout = layout return config -def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024): - groups = _make_groups(n_c4, n_c128, n_swa) - return _get_kv_cache_config_packed(_mock_vllm_config(), groups, mem) - - -def _page_sizes_by_layer( - groups: list[KVCacheGroupSpec], -) -> dict[str, int]: - page_sizes = {} - for group in groups: - specs = group.kv_cache_spec.kv_cache_specs - for layer_name in group.layer_names: - page_sizes[layer_name] = specs[layer_name].page_size_bytes - return page_sizes - - -def _packing_by_layer( - tensors: list[KVCacheTensor], -) -> dict[str, tuple[int, int]]: +def _pages(groups) -> dict[str, int]: return { - layer_name: (tensor.offset, tensor.block_stride) - for tensor in tensors - for layer_name in tensor.shared_by + name: group.kv_cache_spec.kv_cache_specs[name].page_size_bytes + if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) + else group.kv_cache_spec.page_size_bytes + for group in groups + for name in group.layer_names } -def _make_views( - groups: list[KVCacheGroupSpec], - num_blocks: int, - tensors: list[KVCacheTensor], -) -> dict[str, torch.Tensor]: - page_sizes = _page_sizes_by_layer(groups) - packing = _packing_by_layer(tensors) - backing = torch.zeros(tensors[0].size, dtype=torch.uint8) - return { - layer_name: torch.as_strided( - backing, - size=(num_blocks, page_size), - stride=(packing[layer_name][1], 1), - storage_offset=packing[layer_name][0], +def _expected_bytes_per_block(groups) -> int: + pages = _pages(groups) + return max(sum(pages[n] for n in g.layer_names) for g in groups) + + +def _bind(config, layout: str): + return allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout[layout], None) + + +class TestDensePacking: + def test_bytes_per_block_is_largest_group(self): + groups, g1, g2 = _mixed_page_groups() + assert _get_kv_cache_bytes_per_block(groups) == _expected_bytes_per_block( + groups ) - for layer_name, page_size in page_sizes.items() - } - -def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec: - specs = { - f"{prefix}.{i}": MagicMock(page_size_bytes=page_size) - for i, page_size in enumerate(page_sizes) - } - return KVCacheGroupSpec( - layer_names=list(specs), - kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs), - ) - - -class TestInterleavedPacking: - def test_compact_cache_overlays_fp32_state_group(self): - full_specs = {} - state_specs = {} - for i in range(2): - full_specs[f"mla.{i}"] = MLAAttentionSpec( - block_size=256, - num_kv_heads=1, - head_size=512, - dtype=torch.uint8, - page_size_padded=32768, - indexes_kv_by_block_stride=True, - compress_ratio=4, - ) - full_specs[f"indexer.{i}"] = MLAAttentionSpec( - block_size=256, - num_kv_heads=1, - head_size=68, - dtype=torch.uint8, - page_size_padded=4608, - compress_ratio=4, - ) - state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec( - block_size=4, - num_kv_heads=1, - head_size=2048, - dtype=torch.float32, - sliding_window=8, - indexes_kv_by_block_stride=True, - ) - state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec( - block_size=4, - num_kv_heads=1, - head_size=512, - dtype=torch.float32, - sliding_window=8, - indexes_kv_by_block_stride=True, - ) - - grouped_specs = [ - UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs), - UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs), + config = get_kv_cache_config_from_groups( + _mock_vllm_config("BLHNC"), groups, MEMORY + ) + # Groups overlay: both start at offset 0. + assert [tensor.offset for tensor in config.kv_cache_tensors].count(0) == 2 + assert [tensor.layers for tensor in config.kv_cache_tensors] == [ + list(g1)[:3], + list(g1)[3:], + list(g2), ] - groups = _get_kv_cache_groups_uniform_groups(grouped_specs) - assert len(groups) == 2 - assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608} - assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == { - 4352 + def test_layers_within_a_group_are_dense(self): + groups, _, _ = _mixed_page_groups() + pages = _pages(groups) + config = get_kv_cache_config_from_groups( + _mock_vllm_config("BLHNC"), groups, MEMORY + ) + offsets = { + name: tensor.offset + i * tensor.layer_stride + for tensor in config.kv_cache_tensors + for i, name in enumerate(tensor.layers) } - assert { - state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2) - } == {8192} - - full_group_bytes = 2 * (32768 + 4608) - state_group_bytes = 2 * (32768 + 8192) - bytes_per_block = max(full_group_bytes, state_group_bytes) - num_blocks, tensors = _get_kv_cache_config_packed( - _mock_vllm_config(), groups, bytes_per_block * 32 - ) - assert num_blocks == 32 - assert {tensor.block_stride for tensor in tensors} == {bytes_per_block} - - packing = _packing_by_layer(tensors) - assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0 - assert packing["indexer.0"][0] == 32768 - assert packing["indexer_state.0"][0] == 32768 - - def test_deepseek_v4_pro_stride(self): - groups = [ - _make_page_group("full", [32768, 4608] * 30 + [1024] * 31), - _make_page_group("c4_state", [32768, 8192] * 30), - _make_page_group("c128_state", [32768] * 31), - _make_page_group("swa.0", [32768] * 31), - _make_page_group("swa.1", [32768] * 30), - ] - expected_stride = 1_228_800 - - num_blocks, tensors = _get_kv_cache_config_packed( - _mock_vllm_config(), groups, expected_stride * 32 - ) - - assert num_blocks == 32 - assert {tensor.block_stride for tensor in tensors} == {expected_stride} - assert {tensor.size for tensor in tensors} == {expected_stride * 32} - - def test_all_tensors_have_block_stride(self): - _, tensors = _run() - for t in tensors: - assert t.block_stride > 0 - - def test_all_tensors_share_same_size(self): - _, tensors = _run() - sizes = set(t.size for t in tensors) - assert len(sizes) == 1 - assert sizes.pop() > 0 - - def test_offsets_within_one_block(self): - groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) - _, tensors = _get_kv_cache_config_packed( - _mock_vllm_config(), groups, 100 * 1024 * 1024 - ) - page_sizes = _page_sizes_by_layer(groups) - packing = _packing_by_layer(tensors) - for layer_name, page_size in page_sizes.items(): - offset, block_stride = packing[layer_name] - assert offset + page_size <= block_stride - - def test_layouts_are_disjoint_within_each_group(self): - groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) - _, tensors = _get_kv_cache_config_packed( - _mock_vllm_config(), groups, 100 * 1024 * 1024 - ) - page_sizes = _page_sizes_by_layer(groups) - packing = _packing_by_layer(tensors) - for group in groups: - ranges = sorted( - (packing[name][0], packing[name][0] + page_sizes[name]) - for name in group.layer_names + expected = 0 + for name in group.layer_names: + assert offsets[name] == expected + expected += pages[name] + + @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"]) + def test_allocation_is_layout_invariant(self, layout): + specs = {f"l.{i}": _full() for i in range(4)} + groups = [KVCacheGroupSpec(list(specs), _full())] + config = get_kv_cache_config_from_groups( + _mock_vllm_config(layout), groups, MEMORY + ) + (tensor,) = config.kv_cache_tensors + page = _full().page_size_bytes + assert tensor.layers == list(specs) + assert config.num_blocks == MEMORY // (4 * page) + assert tensor.size == 4 * page * config.num_blocks + if layout == "LBNHC": + assert (tensor.layer_stride, tensor.block_stride) == ( + page * config.num_blocks, + page, ) - assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:])) - - def test_all_layers_accounted_for(self): - n_c4, n_c128, n_swa = 5, 4, 7 - _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa) - all_names = set() - for t in tensors: - all_names.update(t.shared_by) - expected = n_c4 * 2 + n_c128 + n_swa - assert len(all_names) == expected - - def test_group_owned_blocks_do_not_alias(self): - groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) - num_blocks, tensors = _get_kv_cache_config_packed( - _mock_vllm_config(), groups, 100 * 1024 * 1024 - ) - views = _make_views(groups, num_blocks, tensors) - - expected = {} - value = 1 - for block_id, group in enumerate(groups): - for layer_name in group.layer_names: - views[layer_name][block_id].fill_(value) - expected[layer_name] = (block_id, value) - value += 1 - - for layer_name, (block_id, value) in expected.items(): - assert (views[layer_name][block_id] == value).all() - - # Once the first group releases its block, another group may reuse it. - for layer_name in groups[1].layer_names: - views[layer_name][0].fill_(255) - for layer_name in groups[1].layer_names: - assert (views[layer_name][0] == 255).all() - - def test_hma_attention_groups_keep_default_backing(self): - full = _make_full_spec() - sw = _make_sw_spec() - page_size = full.page_size_bytes - groups = [ - KVCacheGroupSpec(["full.0", "full.1"], full), - KVCacheGroupSpec(["sw.0", "sw.2"], sw), - KVCacheGroupSpec(["sw.1", "sw.3"], sw), - ] + else: + assert (tensor.layer_stride, tensor.block_stride) == (page, 4 * page) + @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"]) + def test_single_group_mixed_pages_follows_layout(self, layout): + specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)} + groups = [_uniform_group(specs)] config = get_kv_cache_config_from_groups( - _mock_vllm_config(), groups, available_memory=page_size * 2 * 32 + _mock_vllm_config(layout), groups, MEMORY ) + block_stride = _expected_bytes_per_block(groups) + mla_tensor, idx_tensor = config.kv_cache_tensors + assert mla_tensor.layers == ["mla.0", "mla.1"] + assert idx_tensor.layers == ["idx.0"] + assert {t.size for t in config.kv_cache_tensors} == { + block_stride * config.num_blocks + } + if layout == "LBNHC": + assert ( + idx_tensor.offset == 2 * _mla(512).page_size_bytes * config.num_blocks + ) + assert mla_tensor.block_stride == _mla(512).page_size_bytes + else: + assert idx_tensor.offset == 2 * _mla(512).page_size_bytes + assert mla_tensor.block_stride == block_stride - assert config.num_blocks == 32 - assert config.prefix_cache_retention_interval == 0 - assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32 - assert config.kv_cache_tensors == [ - KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]), - KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]), - ] - - def test_hma_attention_groups_use_packed_backing_with_enable_cross_layers(self): - full = _make_full_spec() - sw = _make_sw_spec() - page_size = full.page_size_bytes - groups = [ - KVCacheGroupSpec(["full.0", "full.1"], full), - KVCacheGroupSpec(["sw.0", "sw.2"], sw), - KVCacheGroupSpec(["sw.1", "sw.3"], sw), - ] - + def test_overlaid_groups_alias_and_stay_isolated(self): + groups, g1, g2 = _mixed_page_groups() + # Overlay models resolve to a block-outer layout at backend selection (the + # model's backend declares it); mirror that here. config = get_kv_cache_config_from_groups( - _mock_vllm_config({"enable_cross_layers_blocks": "True"}), - groups, - available_memory=page_size * 2 * 32, + _mock_vllm_config("BLNHC"), groups, MEMORY ) + assert config.num_blocks == MEMORY // _expected_bytes_per_block(groups) + assert _pool_bytes_per_block(groups) == _expected_bytes_per_block(groups) - assert config.num_blocks == 32 - assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32} - assert config.kv_cache_tensors == [ - KVCacheTensor( - size=page_size * 2 * 32, - shared_by=["full.0", "sw.0", "sw.1"], - offset=0, - block_stride=page_size * 2, - ), - KVCacheTensor( - size=page_size * 2 * 32, - shared_by=["full.1", "sw.2", "sw.3"], - offset=page_size, - block_stride=page_size * 2, - ), - ] + views = _bind(config, "BLNHC") + assert set(views) == set(g1) | set(g2) + assert views["swa.0"].data_ptr() == views["mla.0"].data_ptr() - def test_single_group_attention_keeps_unpacked_layout(self): - spec = _make_full_spec() - groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)] + # A block is owned by one group at a time: writes to group-owned + # blocks never disturb the other group's blocks. + for i, name in enumerate(g1): + views[name][0].fill_(i + 1) + views[name][2].fill_(i + 1) + for i, name in enumerate(g2): + views[name][1].fill_(100 + i) + views[name][3].fill_(100 + i) + for i, name in enumerate(g1): + assert (views[name][0].to(torch.int32) == i + 1).all() + assert (views[name][2].to(torch.int32) == i + 1).all() + for i, name in enumerate(g2): + assert (views[name][1].to(torch.int32) == 100 + i).all() + assert (views[name][3].to(torch.int32) == 100 + i).all() + # Layers within a group are disjoint. + views["mla.0"][0].fill_(77) + for i, name in enumerate(list(g1)[1:], start=1): + assert (views[name][0].to(torch.int32) == i + 1).all() + def test_layer_compact_layout_rejected_for_overlaid_groups(self): + # The layout has a single writer (backend-selection resolution); a layer- + # compact layout reaching an overlay model's allocation is an error, not a + # silent flip. + groups, _, _ = _mixed_page_groups() + with pytest.raises( + ValueError, match="cannot express this model's mixed page sizes" + ): + get_kv_cache_config_from_groups(_mock_vllm_config("LBNHC"), groups, MEMORY) + + def test_unresolved_layout_rejected(self): + groups, _, _ = _mixed_page_groups() + with pytest.raises(ValueError, match="has not been resolved"): + get_kv_cache_config_from_groups(_mock_vllm_config(None), groups, MEMORY) + + def test_head_outer_layout_rejected_for_mixed_pages(self): + groups, _, _ = _mixed_page_groups() + with pytest.raises( + ValueError, match="cannot express this model's mixed page sizes" + ): + get_kv_cache_config_from_groups(_mock_vllm_config("LHBNC"), groups, MEMORY) + + @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"]) + def test_bound_views_round_trip(self, layout): + specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)} + groups = [_uniform_group(specs)] config = get_kv_cache_config_from_groups( - _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32 + _mock_vllm_config(layout), groups, MEMORY ) - - assert sum(t.size for t in config.kv_cache_tensors) == ( - spec.page_size_bytes * 2 * 32 - ) - assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0] + views = _bind(config, layout) + for i, name in enumerate(specs): + views[name].fill_(i + 1) + for i, name in enumerate(specs): + assert (views[name].to(torch.int32) == i + 1).all(), name + assert views[name].shape[0] == config.num_blocks if __name__ == "__main__": diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 48943aeb60..870c139ae0 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -124,7 +124,6 @@ def new_kv_cache_spec( page_size_padded=None, sliding_window=None, attention_chunk_size=None, - indexes_kv_by_block_stride=False, kv_quant_mode=KVQuantMode.NONE, ): return FullAttentionSpec( @@ -135,7 +134,6 @@ def new_kv_cache_spec( page_size_padded=page_size_padded, sliding_window=sliding_window, attention_chunk_size=attention_chunk_size, - indexes_kv_by_block_stride=indexes_kv_by_block_stride, kv_quant_mode=kv_quant_mode, ) @@ -147,7 +145,6 @@ def new_sliding_window_spec( dtype=torch.float32, page_size_padded=None, sliding_window=1, - indexes_kv_by_block_stride=False, ): return SlidingWindowSpec( block_size=block_size, @@ -156,7 +153,6 @@ def new_sliding_window_spec( dtype=dtype, page_size_padded=page_size_padded, sliding_window=sliding_window, - indexes_kv_by_block_stride=indexes_kv_by_block_stride, ) @@ -875,6 +871,7 @@ def test_metrics_empty_stats(): def test_get_kv_cache_configs_multiple_workers(): model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() @@ -898,36 +895,21 @@ def test_get_kv_cache_configs_multiple_workers(): ref_kv_cache_spec.page_size_bytes * 2 * 10, ], ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ), - ] + expected = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [expected, expected] # Different available memory. This is the case for TP. # Use the smallest memory available. @@ -939,36 +921,7 @@ def test_get_kv_cache_configs_multiple_workers(): ref_kv_cache_spec.page_size_bytes * 2 * 20, ], ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ), - ] + assert kv_cache_configs == [expected, expected] # Different KV cache specs. This is the case for PP. different_layer_specs = [ @@ -995,7 +948,10 @@ def test_get_kv_cache_configs_multiple_workers(): num_blocks=10, kv_cache_tensors=[ KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ @@ -1006,10 +962,10 @@ def test_get_kv_cache_configs_multiple_workers(): num_blocks=10, kv_cache_tensors=[ KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"] + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer2", "layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ @@ -1046,57 +1002,39 @@ def test_get_kv_cache_configs_multiple_workers(): ref_kv_cache_spec.page_size_bytes * 2 * 10, ], ) + expected_12 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + expected_3 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), + ], + ) assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"] - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), - ], - ), + expected_12, + expected_12, + expected_3, + expected_3, ] # Different workers have different types of layers. This is the case for @@ -1124,10 +1062,10 @@ def test_get_kv_cache_configs_multiple_workers(): num_blocks=10, kv_cache_tensors=[ KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"] + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ @@ -1139,10 +1077,10 @@ def test_get_kv_cache_configs_multiple_workers(): num_blocks=10, kv_cache_tensors=[ KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"] - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer4"] + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer3", "layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ @@ -1180,7 +1118,21 @@ def test_get_kv_cache_configs_multiple_workers(): kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, - shared_by=["layer1", "layer2", "layer3"], + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ @@ -1194,7 +1146,21 @@ def test_get_kv_cache_configs_multiple_workers(): kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * 10, - shared_by=["layer4", "layer5", "layer6"], + layers=["layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer5"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer6"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[ @@ -1233,6 +1199,7 @@ def test_get_kv_cache_configs_multiple_workers(): def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): model_config = ModelConfig(max_model_len=512) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() @@ -1262,7 +1229,10 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - shared_by=["layer1"], + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], @@ -1272,7 +1242,10 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): kv_cache_tensors=[ KVCacheTensor( size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - shared_by=["layer2"], + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, ), ], kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], @@ -1628,78 +1601,18 @@ def test_get_max_concurrency_for_kv_cache_config(): ) -def test_get_max_concurrency_packed_kv_cache_config(): - from vllm.v1.core.kv_cache_utils import ( - _get_kv_cache_config_packed, - _use_packed_kv_cache_config, - ) - - model_config = ModelConfig( - "Qwen/Qwen1.5-7B", - runner="generate", - dtype="float16", - max_model_len=16384, - ) - scheduler_config = SchedulerConfig( - max_num_batched_tokens=1024, - enable_chunked_prefill=True, - max_model_len=model_config.max_model_len, - is_encoder_decoder=model_config.is_encoder_decoder, - async_scheduling=False, - ) - vllm_config = VllmConfig( - model_config=model_config, - scheduler_config=scheduler_config, - ) - - # All-UniformTypeKVCacheSpecs groups select the packed layout. - mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)} - swa_specs = { - f"layer_{i}": SlidingWindowMLASpec( - block_size=16, - num_kv_heads=1, - head_size=576, - dtype=torch.float32, - sliding_window=128, - ) - for i in range(4, 6) - } - kv_cache_groups = [ - KVCacheGroupSpec( - list(mla_specs), - UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs), - ), - KVCacheGroupSpec( - list(swa_specs), - UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs), - ), - ] - assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups) - num_blocks, kv_cache_tensors = _get_kv_cache_config_packed( - vllm_config, kv_cache_groups, 2 * GiB_bytes - ) - assert num_blocks > 0 - kv_cache_config_packed = KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=kv_cache_groups, - ) - # Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages; - # the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The - # previous formula normalized by the first group's page size and gave - # 1061 blocks per request instead of 1097. - assert get_max_concurrency_for_kv_cache_config( - vllm_config, kv_cache_config_packed - ) == num_blocks / (1024 + 73) - - def test_allocate_with_lookahead(): """Verify that lookahead tokens correctly affect block allocation""" block_size = 4 config = KVCacheConfig( num_blocks=10, kv_cache_tensors=[ - KVCacheTensor(size=100, shared_by=["layer1"]), + KVCacheTensor( + size=100, + layers=["layer1"], + layer_stride=100, + block_stride=10, + ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer1"], new_kv_cache_spec(block_size=block_size)), @@ -1763,6 +1676,7 @@ def test_get_kv_cache_config_one_worker(): # pass max_model_len to pass check_enough_kv_cache_memory model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" vllm_config.cache_config.prefix_cache_retention_interval = None mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 @@ -1778,8 +1692,12 @@ def test_get_kv_cache_config_one_worker(): assert kv_cache_config_full == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]), - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), ], kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) @@ -1795,8 +1713,12 @@ def test_get_kv_cache_config_one_worker(): assert kv_cache_config_sliding == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]), - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), ], kv_cache_groups=[ KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) @@ -1815,8 +1737,12 @@ def test_get_kv_cache_config_one_worker(): assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]), - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), ], kv_cache_groups=[ KVCacheGroupSpec( @@ -1838,7 +1764,16 @@ def test_get_kv_cache_config_one_worker(): num_blocks=64, kv_cache_tensors=[ KVCacheTensor( - size=mem_per_block_per_layer * 64, shared_by=["layer_1", "layer_2"] + size=mem_per_block_per_layer * 64, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ @@ -1863,12 +1798,22 @@ def test_get_kv_cache_config_one_worker(): num_blocks=32, kv_cache_tensors=[ KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_1", "layer_3", "layer_4"], + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_2", "layer_5", "layer_6"], + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_3", "layer_5"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_4", "layer_6"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ @@ -1898,15 +1843,28 @@ def test_get_kv_cache_config_one_worker(): num_blocks=32, kv_cache_tensors=[ KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_1", "layer_4", "layer_5", "layer_6"], + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1", "layer_2", "layer_3"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_2", "layer_7", "layer_8", "layer_9"], + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_4", "layer_7", "layer_10"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), KVCacheTensor( - size=mem_per_block_per_layer * 32, shared_by=["layer_3", "layer_10"] + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_5", "layer_8"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_6", "layer_9"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ @@ -1943,28 +1901,23 @@ def test_get_kv_cache_config_one_worker(): num_blocks=32, kv_cache_tensors=[ KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_1", "layer_7"], + size=mem_per_block_per_layer * 32 * 6, + layers=[ + "layer_1", + "layer_2", + "layer_3", + "layer_4", + "layer_5", + "layer_6", + ], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_2", "layer_8"], - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_3", "layer_9"], - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_4", "layer_10"], - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_5", "layer_11"], - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - shared_by=["layer_6"], + size=mem_per_block_per_layer * 32 * 6, + layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ @@ -1987,11 +1940,24 @@ def test_get_kv_cache_config_one_worker(): kv_cache_config_hybrid = get_kv_cache_configs( vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] )[0] + # Layers of different page sizes pack densely into one allocation: the + # 2x-larger layer_1 takes the first region, layer_2 the next. assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=32, kv_cache_tensors=[ - KVCacheTensor(size=mem_per_block_per_layer * 32 * 2, shared_by=["layer_1"]), - KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 2 * 32, + block_stride=mem_per_block_per_layer * 2, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + offset=mem_per_block_per_layer * 2 * 32, + ), ], kv_cache_groups=[ KVCacheGroupSpec( @@ -2015,7 +1981,16 @@ def test_get_kv_cache_config_one_worker(): num_blocks=32, kv_cache_tensors=[ KVCacheTensor( - size=mem_per_block_per_layer * 32, shared_by=["layer_1", "layer_2"] + size=mem_per_block_per_layer * 32, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, ), ], kv_cache_groups=[ @@ -2028,9 +2003,9 @@ def test_get_kv_cache_config_one_worker(): # different hidden size that cannot be aligned by using different block size, # but can be aligned by padding the smaller physical page. - swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True) + swa_spec = new_sliding_window_spec(head_size=96) kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True), + "layer_1": new_kv_cache_spec(head_size=64), "layer_2": swa_spec, } @@ -2041,7 +2016,18 @@ def test_get_kv_cache_config_one_worker(): assert kv_cache_config_hybrid == KVCacheConfig( num_blocks=42, kv_cache_tensors=[ - KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]), + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_1"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_2"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), ], kv_cache_groups=[ KVCacheGroupSpec( @@ -2049,12 +2035,11 @@ def test_get_kv_cache_config_one_worker(): new_kv_cache_spec( head_size=64, page_size_padded=padded_page_size, - indexes_kv_by_block_stride=True, ), ), KVCacheGroupSpec( ["layer_2"], - new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True), + new_sliding_window_spec(head_size=96), ), ], ) @@ -2067,8 +2052,12 @@ def test_get_kv_cache_config_one_worker(): assert kv_cache_config_override_blocks == KVCacheConfig( num_blocks=16, kv_cache_tensors=[ - KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_1"]), - KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_2"]), + KVCacheTensor( + size=mem_per_block_per_layer * 16 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 16, + block_stride=mem_per_block_per_layer, + ), ], kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) @@ -2319,6 +2308,45 @@ def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): ) == (544, 136) +def test_multi_run_layer_compact_strides_place_hoisted_heads(): + """A layer-compact run region is its own dense allocation: under LHBNC the head + groups sit between the layers and the blocks, so a run's block stride is one head + group's slice, not the whole page (regression test for setStorage out-of-bounds + on ROCm hybrid models).""" + full = new_kv_cache_spec() + swa = new_sliding_window_spec(sliding_window=full.block_size * 2) + page = full.page_size_bytes + assert swa.page_size_bytes == page + + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.kv_cache_layout = "LHBNC" + config = kv_cache_utils.get_kv_cache_config_from_groups( + vllm_config, + [ + KVCacheGroupSpec(["full.0", "full.1"], full), + KVCacheGroupSpec(["swa.0", "swa.1"], swa), + ], + available_memory=4 * page * 8, + ) + + num_blocks = config.num_blocks + head_group = full.block_size * full.state_content_size_bytes + for tensor in config.kv_cache_tensors: + # One head group per block, head groups between the layers and blocks. + assert tensor.block_stride == head_group + assert tensor.layer_stride == full.num_heads * num_blocks * head_group + # The last block of the last layer's last head group must stay within the + # allocation (the old page-based stride overran it). + end = ( + tensor.offset + + tensor.layer_stride * (len(tensor.layers) - 1) + + (full.num_heads - 1) * num_blocks * head_group + + tensor.block_stride * (num_blocks - 1) + + head_group + ) + assert end <= tensor.size + + def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): specs = { "target.0.attn": new_mla_spec(), @@ -2640,6 +2668,7 @@ def test_auto_fit_max_model_len(): # Simulate the user passing -1 by setting original_max_model_len model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer kv_cache_specs = { @@ -2658,6 +2687,7 @@ def test_auto_fit_max_model_len(): model_config = ModelConfig(max_model_len=1024) model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" # With limited memory, max_model_len should be reduced # Need memory for at least max_model_len tokens @@ -2678,6 +2708,7 @@ def test_auto_fit_max_model_len_with_hybrid(): # Simulate the user passing -1 by setting original_max_model_len model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer gamma = 2 @@ -2700,6 +2731,7 @@ def test_auto_fit_max_model_len_not_triggered(): model_config = ModelConfig(max_model_len=16) # original_max_model_len should be None by default, not -1 vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 kv_cache_specs = { @@ -2722,6 +2754,7 @@ def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): model_config = ModelConfig(max_model_len=16384) model_config.original_max_model_len = -1 # request auto-fit vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" # Cap the cache to 32 blocks regardless of available memory. vllm_config.cache_config.num_gpu_blocks_override = 32 @@ -2747,6 +2780,7 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): """ model_config = ModelConfig(max_model_len=16384) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" # 32 blocks is far too small for max_model_len=16384 (would need 1024). vllm_config.cache_config.num_gpu_blocks_override = 32 @@ -2769,14 +2803,13 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer block-size multiple, so the smaller page must be padded instead. """ - # Both layers' backends opt into the padded-page strided view (e.g. - # FlashAttention / its DiffKV subclass), so padding is allowed. + # Attention layers read padded pages through a strided view, so padding is + # allowed. target_spec = new_kv_cache_spec( block_size=16, num_kv_heads=1, head_size=192, dtype=torch.bfloat16, - indexes_kv_by_block_stride=True, ) draft_spec = new_sliding_window_spec( block_size=16, @@ -2784,7 +2817,6 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): head_size=128, dtype=torch.bfloat16, sliding_window=1024, - indexes_kv_by_block_stride=True, ) 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(): assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes -def test_unify_kv_cache_page_size_padding_requires_backend_support(): - """Padding is gated on the backend declaring ``indexes_kv_by_block_stride``. - - A backend that does not support the strided padded-page view must raise - rather than silently padding (and misreading KV at runtime). - """ - target_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=1, - head_size=192, - dtype=torch.bfloat16, - indexes_kv_by_block_stride=True, - ) - # The non-divisible draft layer needs padding but its backend does not - # support the strided padded-page view -> must raise, not silently pad. - draft_spec = new_sliding_window_spec( - block_size=16, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - sliding_window=1024, - indexes_kv_by_block_stride=False, - ) - specs = {"target_attn": target_spec, "draft_attn": draft_spec} - - with pytest.raises(NotImplementedError): - kv_cache_utils.unify_kv_cache_spec_page_size(specs) - - def test_unpadded_page_size_includes_per_token_head_scales(): # Per-token-head quant carries inline fp32 scales that are carved from the # 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(): def test_page_size_padded_wins(): # An explicit padded page size takes precedence over the unpadded size. - spec = new_kv_cache_spec( - dtype=torch.uint8, - kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD, - page_size_padded=65536, + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend + + spec = TritonAttentionBackend.customize_spec( + new_kv_cache_spec( + dtype=torch.uint8, + kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD, + page_size_padded=65536, + ) ) assert spec.page_size_bytes == 65536 @@ -2996,14 +3003,18 @@ def test_unify_kv_cache_spec_page_size_mamba(): ) assert unified["mamba_layer"].page_size_bytes == 32768 - # 4. Attention layers with non-divisible page sizes still raise. - with pytest.raises(NotImplementedError): - kv_cache_utils.unify_kv_cache_spec_page_size( - { - "attn_layer": new_kv_cache_spec(block_size=24), # 24576 - "draft_attn_layer": draft_attn_spec, # 32768 - } - ) + # 4. Attention layers with non-divisible page sizes are padded too: every + # backend reads a padded page through the view's block stride, so there is + # no longer a case that must raise. + unified = kv_cache_utils.unify_kv_cache_spec_page_size( + { + "attn_layer": new_kv_cache_spec(block_size=24), # 24576 + "draft_attn_layer": draft_attn_spec, # 32768 + } + ) + assert unified["attn_layer"].page_size_padded == 32768 + assert unified["attn_layer"].page_size_bytes == 32768 + assert unified["attn_layer"].block_size == 24 # 5. Uniform page sizes are returned unchanged. specs = { @@ -3135,6 +3146,7 @@ def test_auto_fit_max_model_len_reserves_null_block(): model_config = ModelConfig(max_model_len=1024) model_config.original_max_model_len = -1 vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" spec = new_kv_cache_spec(block_size=block_size) # Exactly the 1024 / 16 = 64 blocks a full-length request would need. diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 6cd117b1fa..038d45be30 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -154,7 +154,7 @@ def make_kv_cache_config_hybrid_model( elif second_spec_type == "mamba": second_spec = MambaSpec( block_size=block_size, - shapes=(1, 1), + shapes=((1, 1),), dtypes=(torch.float32,), ) @@ -189,7 +189,7 @@ def make_kv_cache_config_three_types( if third_spec_type == "mamba": third_spec = MambaSpec( block_size=block_size, - shapes=(1, 1), + shapes=((1, 1),), dtypes=(torch.float32,), ) elif third_spec_type == "sliding_window": @@ -768,12 +768,12 @@ def _make_hybrid_kv_cache_config( ), "mamba": lambda: MambaSpec( block_size=block_size, - shapes=(1, 1), + shapes=((1, 1),), dtypes=(torch.float32,), ), "mamba_align": lambda: MambaSpec( block_size=block_size, - shapes=(1, 1), + shapes=((1, 1),), dtypes=(torch.float32,), mamba_cache_mode="align", ), @@ -3273,7 +3273,7 @@ def test_hybrid_local_kv_retention_interval_survives_recycling(): num_kv_heads=1, head_size=1, dtype=torch.uint8, - compress_ratio=4, + tokens_per_state=4, ), ), KVCacheGroupSpec( diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 7424b76223..ab28b95669 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -1329,6 +1329,7 @@ def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPat ) mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}] + mock_executor.get_supported_kv_cache_layouts.return_value = [["LBNHC"]] mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024] mock_executor.initialize_from_config.return_value = None mock_executor.supports_draft_weight_updates.return_value = False diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index b8a37e1e96..b9f461db17 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -36,7 +36,9 @@ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS" fi -DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD +# Empty leaves the layout to the backends, for models that constrain it. +PREFILLER_KV_LAYOUT=${PREFILLER_KV_LAYOUT-"HND"} +DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT-"HND"} # Default to HND, optional NHD if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"' else @@ -44,9 +46,9 @@ else fi if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then - KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}' -else - KV_EXTRA_CONFIG='' + # Cross-layer blocks are exercised via the block-outermost layout. + PREFILLER_KV_LAYOUT="BLHNC" + DECODER_KV_LAYOUT="BLHNC" fi # Connector: default pull NixlConnector; NixlPushConnector enables PP prefill. @@ -54,11 +56,11 @@ KV_CONNECTOR=${KV_CONNECTOR:-NixlConnector} # Build the kv-transfer-config for P and D if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then - KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}' - KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}' + KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}'}' + KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}'}' else - KV_CONFIG_P="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}" - KV_CONFIG_D="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}" + KV_CONFIG_P="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}" + KV_CONFIG_D="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}" fi # Models to run @@ -176,7 +178,7 @@ run_tests_for_model() { # Build the command with or without model-specific args BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ - VLLM_KV_CACHE_LAYOUT='HND' \ + ${PREFILLER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$PREFILLER_KV_LAYOUT} \ VLLM_PORT=$INTERNAL_PORT \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ @@ -239,7 +241,7 @@ run_tests_for_model() { # Build the command with or without model-specific args BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \ - VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \ + ${DECODER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT} \ $DECODER_INTERNAL_PORT_ENV \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh index 8ce688a4a2..83fa17d7fb 100755 --- a/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh @@ -9,8 +9,8 @@ # test_accuracy.py. # # By default runs two configurations: -# 1. Normal KV layout (NixlConnector without cross-layer blocks) -# 2. Cross-layer KV layout (NixlConnector with enable_cross_layers_blocks) +# 1. Normal KV layout (LBHNC) +# 2. Cross-layer KV layout (BLHNC) via VLLM_KV_CACHE_LAYOUT # # Usage: # bash tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh @@ -55,8 +55,6 @@ fi # ── KV transfer configs ───────────────────────────────────────────────── -# Normal layout: OffloadingConnector prefers cross-layer but NixlConnector -# does not, so MultiConnector.prefer_cross_layer_blocks = False. KV_CONFIG_NORMAL='{ "kv_connector":"MultiConnector", "kv_role":"kv_both", @@ -72,22 +70,6 @@ KV_CONFIG_NORMAL='{ # Remove whitespace for CLI safety KV_CONFIG_NORMAL=$(echo "$KV_CONFIG_NORMAL" | tr -d '[:space:]') -# Cross-layer layout: both connectors prefer cross-layer blocks. -KV_CONFIG_CROSS_LAYERS='{ - "kv_connector":"MultiConnector", - "kv_role":"kv_both", - "kv_connector_extra_config":{ - "connectors":[ - {"kv_connector":"NixlConnector","kv_role":"kv_both", - "kv_buffer_device":"'$KV_BUFFER_DEVICE'", - "kv_connector_extra_config":{"enable_cross_layers_blocks":"True"}}, - {"kv_connector":"OffloadingConnector","kv_role":"kv_both", - "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}} - ] - } -}' -KV_CONFIG_CROSS_LAYERS=$(echo "$KV_CONFIG_CROSS_LAYERS" | tr -d '[:space:]') - # ── Helpers ────────────────────────────────────────────────────────────── PROCESS_GROUPS=() @@ -218,7 +200,7 @@ run_tests_for_model() { # ── Start prefill instance ── echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT" BASE_CMD="$DEVICE_VISIBILITY_ENV=$PREFILL_GPU \ - VLLM_KV_CACHE_LAYOUT='HND' \ + VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \ vllm serve $model_name \ @@ -245,7 +227,7 @@ run_tests_for_model() { # ── Start decode instance ── echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT" BASE_CMD="$DEVICE_VISIBILITY_ENV=$DECODE_GPU \ - VLLM_KV_CACHE_LAYOUT='HND' \ + VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \ vllm serve $model_name \ @@ -305,7 +287,8 @@ for model in "${MODELS[@]}"; do fi if [[ -z "${SKIP_CROSS_LAYERS:-}" ]]; then - run_tests_for_model "$model" "$KV_CONFIG_CROSS_LAYERS" "MultiConnector cross-layer layout" + VLLM_KV_CACHE_LAYOUT=BLHNC \ + run_tests_for_model "$model" "$KV_CONFIG_NORMAL" "MultiConnector cross-layer layout" fi done diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh index de6c9abcc6..9b9afc4dd9 100755 --- a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh @@ -95,7 +95,7 @@ run_tests_for_model() { # ── Start prefill instance ── echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT" BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \ - VLLM_KV_CACHE_LAYOUT='HND' \ + VLLM_KV_CACHE_LAYOUT='LBHNC' \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \ vllm serve \"$model_name\" \ @@ -121,7 +121,7 @@ run_tests_for_model() { # ── Start decode instance ── echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT" BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \ - VLLM_KV_CACHE_LAYOUT='HND' \ + VLLM_KV_CACHE_LAYOUT='LBHNC' \ UCX_NET_DEVICES=all \ VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \ vllm serve \"$model_name\" \ diff --git a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh index ec98ed3008..28bc22b37d 100755 --- a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh +++ b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh @@ -249,7 +249,7 @@ run_test_for_device() { echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT" env \ ${GPU_DEVICE_VAR}="$GPU_ID" \ - VLLM_KV_CACHE_LAYOUT='HND' \ + VLLM_KV_CACHE_LAYOUT='LBHNC' \ UCX_NET_DEVICES=all \ ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \ VLLM_NIXL_SIDE_CHANNEL_HOST="$NIXL_SIDE_CHANNEL_HOST" \ @@ -288,7 +288,7 @@ run_test_for_device() { echo "Starting decode instance $i on GPU $GPU_ID, port $PORT" env \ ${GPU_DEVICE_VAR}="$GPU_ID" \ - VLLM_KV_CACHE_LAYOUT='HND' \ + VLLM_KV_CACHE_LAYOUT='LBHNC' \ UCX_NET_DEVICES=all \ ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \ VLLM_NIXL_SIDE_CHANNEL_HOST="$NIXL_SIDE_CHANNEL_HOST" \ diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py index aefb536c4c..9c3ba6d555 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py @@ -12,6 +12,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping i _verify_tiling, derive_canonical_mappings, ) +from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -37,12 +38,14 @@ def _ctx(rank, tp=1, dcp=1, pcp=1, interleave=1, total=None, heads=2): def _full_spec(num_kv_heads: int = 2, **kwargs) -> FullAttentionSpec: # block_size=4, head_dim=64, int8 - return FullAttentionSpec( - block_size=4, - num_kv_heads=num_kv_heads, - head_size=64, - dtype=torch.int8, - **kwargs, + return TritonAttentionBackend.customize_spec( + FullAttentionSpec( + block_size=4, + num_kv_heads=num_kv_heads, + head_size=64, + dtype=torch.int8, + **kwargs, + ) ) @@ -381,7 +384,9 @@ def test_fail_closed_cases(): quant_spec = _full_spec(kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD) assert _try_mapping(quant_spec, _split_nhd_cache(quant_spec), _ctx(0, tp=4)) is None # Compressed MLA slots are not 1:1 with tokens - assert _try_mapping(_mla_spec(compress_ratio=2), None, _ctx(0, tp=2, dcp=2)) is None + assert ( + _try_mapping(_mla_spec(tokens_per_state=2), None, _ctx(0, tp=2, dcp=2)) is None + ) # Unrecognized physical layouts swapped = torch.zeros( NUM_BLOCKS, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 091b50fa1a..afb0a9bda5 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -72,43 +72,54 @@ def _make_vllm_config( def _make_kv_cache_config() -> KVCacheConfig: num_blocks = 16 + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) kv_tensor = KVCacheTensor( - size=num_blocks * 8, - shared_by=["layer"], - block_stride=0, + size=spec.page_size_bytes * num_blocks, + layers=["layer"], + layer_stride=spec.page_size_bytes * num_blocks, + block_stride=spec.page_size_bytes, ) return KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=[kv_tensor], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer"], - FullAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ) - ], + kv_cache_groups=[KVCacheGroupSpec(["layer"], spec)], ) def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig: + """One 16 byte-per-block allocation, described two ways. + + Packed: both layers are one dense run. Unpacked: the same bytes as two runs, the + second starting after the first layer's region. Either way the connector accounts + for 16 KV bytes per block. + """ num_blocks = 4 + page = 8 + size = 2 * page * num_blocks if packed: kv_cache_tensors = [ KVCacheTensor( - size=64, - shared_by=[layer_name], - block_stride=16, + size=size, + layers=["layer0", "layer1"], + layer_stride=page * num_blocks, + block_stride=page, ) - for layer_name in ("layer0", "layer1") ] else: kv_cache_tensors = [ - KVCacheTensor(size=40, shared_by=["layer0"]), - KVCacheTensor(size=24, shared_by=["layer1"]), + KVCacheTensor( + size=size, + layers=[layer], + layer_stride=page * num_blocks, + block_stride=page, + offset=i * page * num_blocks, + ) + for i, layer in enumerate(("layer0", "layer1")) ] return KVCacheConfig( @@ -150,6 +161,18 @@ def _mla_spec( ) +_MAMBA_SPEC = MambaSpec( + block_size=16, + shapes=((16, 1),), + dtypes=(torch.float32,), +) +# Page sizes of the specs the replicated-layout cases below are built from. +_MLA_PAGE = _mla_spec().page_size_bytes +_HALF_MLA_PAGE = _mla_spec(head_size=256).page_size_bytes +_FULL_PAGE = _full_attention_spec().page_size_bytes +_MAMBA_PAGE = _MAMBA_SPEC.page_size_bytes + + def _make_mla_kv_cache_config( layer_names: list[str] | None = None, head_size: int = 512, @@ -159,12 +182,14 @@ def _make_mla_kv_cache_config( if layer_names is None: layer_names = ["layer0", "layer1"] spec = _mla_spec(head_size=head_size, dtype=dtype) + layer_stride = spec.page_size_bytes * num_blocks kv_cache_tensors = [ KVCacheTensor( - size=spec.page_size_bytes * num_blocks, - shared_by=[layer_name], + size=layer_stride * len(layer_names), + layers=layer_names, + layer_stride=layer_stride, + block_stride=spec.page_size_bytes, ) - for layer_name in layer_names ] return KVCacheConfig( num_blocks=num_blocks, @@ -174,15 +199,31 @@ def _make_mla_kv_cache_config( def _make_hybrid_kv_cache_config() -> KVCacheConfig: + num_blocks = 4 + full_spec = _full_attention_spec(block_size=12) + mla_spec = _mla_spec() + # Mixed page sizes across overlaying groups: a block is a window of the largest + # group's packing, so the layer dim sits inside the block dim. + window = max(full_spec.page_size_bytes, mla_spec.page_size_bytes) return KVCacheConfig( - num_blocks=4, + num_blocks=num_blocks, kv_cache_tensors=[ - KVCacheTensor(size=40, shared_by=["full_layer"]), - KVCacheTensor(size=24, shared_by=["mla_layer"]), + KVCacheTensor( + size=window * num_blocks, + layers=["full_layer"], + layer_stride=full_spec.page_size_bytes, + block_stride=window, + ), + KVCacheTensor( + size=window * num_blocks, + layers=["mla_layer"], + layer_stride=mla_spec.page_size_bytes, + block_stride=window, + ), ], kv_cache_groups=[ - KVCacheGroupSpec(["full_layer"], _full_attention_spec(block_size=12)), - KVCacheGroupSpec(["mla_layer"], _mla_spec()), + KVCacheGroupSpec(["full_layer"], full_spec), + KVCacheGroupSpec(["mla_layer"], mla_spec), ], ) @@ -263,18 +304,9 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool): assert offloading_config.cache.blocks_per_chunk == 2 -def test_rejects_partially_packed_tensor_layout(): - kv_cache_config = _make_sizing_kv_cache_config(packed=False) - kv_cache_config.kv_cache_tensors[0].block_stride = 16 - - with pytest.raises(AssertionError): - build_offloading_config(_make_vllm_config(), kv_cache_config) - - def test_zero_blocks_skips_tensor_layout_validation(): kv_cache_config = _make_sizing_kv_cache_config(packed=False) kv_cache_config.num_blocks = 0 - kv_cache_config.kv_cache_tensors[0].block_stride = 16 offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config) @@ -374,8 +406,10 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( num_blocks=4, kv_cache_tensors=[ KVCacheTensor( - size=_mla_spec().page_size_bytes * 4, - shared_by=["layer"], + size=_MLA_PAGE * 4, + layers=["layer"], + layer_stride=_MLA_PAGE * 4, + block_stride=_MLA_PAGE, ) ], kv_cache_groups=[ @@ -398,8 +432,10 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( num_blocks=4, kv_cache_tensors=[ KVCacheTensor( - size=_mla_spec().page_size_bytes * 4, - shared_by=["layer"], + size=_MLA_PAGE * 4, + layers=["layer"], + layer_stride=_MLA_PAGE * 4, + block_stride=_MLA_PAGE, ) ], kv_cache_groups=[ @@ -417,16 +453,22 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( "hidden-state", ), ( + # One group, two page sizes: layer1's run starts past layer0's region. KVCacheConfig( num_blocks=4, kv_cache_tensors=[ KVCacheTensor( - size=_mla_spec().page_size_bytes * 4, - shared_by=["layer0"], + size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4, + layers=["layer0"], + layer_stride=_MLA_PAGE * 4, + block_stride=_MLA_PAGE, ), KVCacheTensor( - size=_mla_spec(head_size=256).page_size_bytes * 4, - shared_by=["layer1"], + size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4, + layers=["layer1"], + layer_stride=_HALF_MLA_PAGE * 4, + block_stride=_HALF_MLA_PAGE, + offset=_MLA_PAGE * 4, ), ], kv_cache_groups=[ @@ -445,16 +487,22 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( "uniform-wrapper", ), ( + # Overlaid groups with different page sizes: a block is a window of + # the largest group's packing. KVCacheConfig( num_blocks=4, kv_cache_tensors=[ KVCacheTensor( - size=_mla_spec().page_size_bytes * 4, - shared_by=["mla"], + size=_FULL_PAGE * 4, + layers=["mla"], + layer_stride=_MLA_PAGE, + block_stride=_FULL_PAGE, ), KVCacheTensor( - size=_full_attention_spec().page_size_bytes * 4, - shared_by=["full"], + size=_FULL_PAGE * 4, + layers=["full"], + layer_stride=_FULL_PAGE, + block_stride=_FULL_PAGE, ), ], kv_cache_groups=[ @@ -469,21 +517,21 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( num_blocks=4, kv_cache_tensors=[ KVCacheTensor( - size=_mla_spec().page_size_bytes * 4, - shared_by=["mla"], + size=_MLA_PAGE * 4, + layers=["mla"], + layer_stride=_MLA_PAGE, + block_stride=_MLA_PAGE, + ), + KVCacheTensor( + size=_MLA_PAGE * 4, + layers=["mamba"], + layer_stride=_MAMBA_PAGE, + block_stride=_MLA_PAGE, ), - KVCacheTensor(size=64 * 4, shared_by=["mamba"]), ], kv_cache_groups=[ KVCacheGroupSpec(["mla"], _mla_spec()), - KVCacheGroupSpec( - ["mamba"], - MambaSpec( - block_size=16, - shapes=((16, 1),), - dtypes=(torch.float32,), - ), - ), + KVCacheGroupSpec(["mamba"], _MAMBA_SPEC), ], ), "mla-mamba-hybrid", @@ -493,13 +541,12 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( num_blocks=4, kv_cache_tensors=[ KVCacheTensor( - size=_mla_spec().page_size_bytes * 4, - shared_by=["layer0"], - ), - KVCacheTensor( - size=_mla_spec().page_size_bytes * 4, - shared_by=["layer1"], - ), + size=_MLA_PAGE * 4, + layers=[layer], + layer_stride=_MLA_PAGE * 4, + block_stride=_MLA_PAGE, + ) + for layer in ("layer0", "layer1") ], kv_cache_groups=[ KVCacheGroupSpec(["layer0"], _mla_spec()), @@ -531,16 +578,26 @@ def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting(): indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8) main_layers = [f"main_{i}" for i in range(61)] indexer_layers = [f"indexer_{i}" for i in range(61)] + # A DSA-style group: the main pages and the smaller indexer pages are packed one + # after the other, so a block holds more than 61 MLA pages. + main_bytes = main_spec.page_size_bytes * len(main_layers) + indexer_bytes = indexer_spec.page_size_bytes * len(indexer_layers) + size = (main_bytes + indexer_bytes) * num_blocks kv_cache_config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=[ KVCacheTensor( - size=main_spec.page_size_bytes * len(main_layers) * num_blocks, - shared_by=main_layers, + size=size, + layers=main_layers, + layer_stride=main_spec.page_size_bytes * num_blocks, + block_stride=main_spec.page_size_bytes, ), KVCacheTensor( - size=indexer_spec.page_size_bytes * len(indexer_layers) * num_blocks, - shared_by=indexer_layers, + size=size, + layers=indexer_layers, + layer_stride=indexer_spec.page_size_bytes * num_blocks, + block_stride=indexer_spec.page_size_bytes, + offset=main_bytes * num_blocks, ), ], kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)], @@ -665,34 +722,6 @@ def test_canonical_layout_certifies_v2_model_runner(): ).parallel.is_parallelism_agnostic -def test_prefer_cross_layer_blocks_yields_to_canonical_layout(): - """The connector must not request cross-layer blocks under - canonical_layout: cross-layer slabs have no per-layer refs to certify.""" - from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole - from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import ( - OffloadingConnector, - ) - - kv_cache_config = _make_kv_cache_config() - connector_module = "vllm.distributed.kv_transfer.kv_connector.v1" - - def make_connector(extra_config: dict[str, Any] | None) -> OffloadingConnector: - with ( - patch(f"{connector_module}.offloading_connector.OffloadingSpecFactory"), - patch( - f"{connector_module}.offloading_connector.OffloadingConnectorScheduler" - ), - ): - return OffloadingConnector( - _make_vllm_config(extra_config=extra_config), - KVConnectorRole.SCHEDULER, - kv_cache_config, - ) - - assert make_connector(None).prefer_cross_layer_blocks - assert not make_connector({"canonical_layout": True}).prefer_cross_layer_blocks - - def test_parallelism_agnostic_disabled_on_v2_model_runner(): config = _make_vllm_config() config.use_v2_model_runner = True diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 8a1f3723c4..366624e1d3 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections import defaultdict from unittest.mock import MagicMock import pytest @@ -13,11 +12,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( from vllm.platforms import current_platform from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.attention.backends.utils import set_kv_cache_layout from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheLayout, KVCacheTensor, MambaSpec, MLAAttentionSpec, @@ -65,40 +64,18 @@ elif current_platform.is_xpu(): # --------------------------------------------------------------------------- -def _allocate_and_reshape_kv_caches( +def _allocate_kv_caches( kv_cache_config: KVCacheConfig, attn_groups: list[list], device: torch.device, ): - """ - Use the real GPUModelRunner allocation and reshape methods to produce - kv_caches, just like the model runner does during initialization. - """ - from vllm.v1.worker.gpu_model_runner import GPUModelRunner + """Allocate kv_caches exactly as the model runner does at startup.""" + from vllm.v1.worker.utils import allocate_kv_cache - # Some backends (e.g. FlashAttention) query the KV cache layout during - # reshape, which ultimately calls get_current_vllm_config(). Setting - # the layout override avoids needing a full VllmConfig context. - set_kv_cache_layout("NHD") - try: - runner = object.__new__(GPUModelRunner) - runner.device = device - runner.runner_only_attn_layers = set() - runner.attn_groups = attn_groups - runner.kv_cache_config = kv_cache_config - runner.cache_config = MagicMock(cache_dtype="auto") - runner.shared_kv_cache_layers = {} - runner.model_config = MagicMock() - runner.model_config.hf_config.model_type = "" - runner.compilation_config = MagicMock( - static_forward_context=defaultdict(MagicMock) - ) - runner.kv_caches = [] - - kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups) - return runner.initialize_kv_cache_tensors(kv_cache_config, kernel_block_sizes) - finally: - set_kv_cache_layout(None) + kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups) + return allocate_kv_cache( + kv_cache_config, device, KVCacheLayout.LBNHC, kernel_block_sizes + ) def _single_rank_vllm_config(total_kv_heads: int): @@ -321,10 +298,9 @@ def test_register_kv_caches(backend): Creates one FullAttention group, one MLA group, one Mamba group, and one Mamba-padded group. Each group has GROUP_SIZE layers. - KVCacheTensors are shared across all groups mirroring the real allocation - in kv_cache_utils.py: tensor i is shared by layer i from every group. - The padded-mamba group has a different page size so its layers get their - own dedicated tensors. + KVCacheTensors mirror the real allocation in kv_cache_utils.py: each group is one + run of the shared allocation starting at offset 0, so layer i of every group + aliases the same bytes. Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce the raw per-layer kv_caches registered by the connector. @@ -413,18 +389,15 @@ def test_register_kv_caches(backend): aligned_mamba_layer_names, ] - kv_cache_tensors: list[KVCacheTensor] = [] - for i in range(GROUP_SIZE): - shared_by: list[str] = [] - for group_layer_names in layer_groups: - if len(group_layer_names) > i: - shared_by.append(group_layer_names[i]) - kv_cache_tensors.append( - KVCacheTensor( - size=PAGE_SIZE_BYTES * NUM_BLOCKS, - shared_by=shared_by, - ) + kv_cache_tensors: list[KVCacheTensor] = [ + KVCacheTensor( + size=PAGE_SIZE_BYTES * NUM_BLOCKS * GROUP_SIZE, + layers=group_layer_names, + layer_stride=PAGE_SIZE_BYTES * NUM_BLOCKS, + block_stride=PAGE_SIZE_BYTES, ) + for group_layer_names in layer_groups + ] kv_cache_groups = [ KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec), @@ -472,7 +445,7 @@ def test_register_kv_caches(backend): kv_cache_groups=kv_cache_groups, ) - kv_caches = _allocate_and_reshape_kv_caches( + kv_caches = _allocate_kv_caches( kv_cache_config, attn_groups, device=torch.device(f"{DEVICE_TYPE}:0"), @@ -580,16 +553,24 @@ def test_register_kv_caches_uniform_type(backend): kv_cache_specs={layer_a: spec_a, layer_b: spec_b}, ) + # The group packs its two layers densely, one run each as their pages differ: + # layer_a occupies the first page of every block, layer_b the rest. + window = spec_a.page_size_bytes + spec_b.page_size_bytes kv_cache_config = KVCacheConfig( num_blocks=NUM_BLOCKS, kv_cache_tensors=[ KVCacheTensor( - size=spec_a.page_size_bytes * NUM_BLOCKS, - shared_by=[layer_a], + size=window * NUM_BLOCKS, + layers=[layer_a], + layer_stride=spec_a.page_size_bytes * NUM_BLOCKS, + block_stride=spec_a.page_size_bytes, ), KVCacheTensor( - size=spec_b.page_size_bytes * NUM_BLOCKS, - shared_by=[layer_b], + size=window * NUM_BLOCKS, + layers=[layer_b], + layer_stride=spec_b.page_size_bytes * NUM_BLOCKS, + block_stride=spec_b.page_size_bytes, + offset=spec_a.page_size_bytes * NUM_BLOCKS, ), ], kv_cache_groups=[ @@ -617,7 +598,7 @@ def test_register_kv_caches_uniform_type(backend): ] ] - kv_caches = _allocate_and_reshape_kv_caches( + kv_caches = _allocate_kv_caches( kv_cache_config, attn_groups, device=torch.device(f"{DEVICE_TYPE}:0"), diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 9acf5f5d49..3397a75c3c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -241,13 +241,19 @@ class RequestRunner: ) ] + # Groups overlay one allocation, sized by the largest group. + block_stride = max( + group.kv_cache_spec.page_size_bytes * len(group.layer_names) + for group in kv_cache_groups + ) kv_cache_tensors = [ KVCacheTensor( - size=group.kv_cache_spec.page_size_bytes * num_gpu_blocks, - shared_by=[layer_name], + size=block_stride * num_gpu_blocks, + layers=list(group.layer_names), + layer_stride=group.kv_cache_spec.page_size_bytes * num_gpu_blocks, + block_stride=group.kv_cache_spec.page_size_bytes, ) for group in kv_cache_groups - for layer_name in group.layer_names ] kv_cache_config = KVCacheConfig( @@ -277,7 +283,6 @@ class RequestRunner: ) # register worker kv_caches to enable OffloadingWorker creations - # set_current_vllm_config is needed for get_kv_cache_layout() to work kv_caches: dict[str, torch.Tensor] = {} for group in kv_cache_groups: spec = group.kv_cache_spec diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 13183042ee..0801e3817e 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -98,7 +98,7 @@ def _make_connector_with_fake_worker( ) worker = connector.connector_worker assert isinstance(worker.nixl_wrapper, FakeNixlWrapper) - worker.kv_cache_layout = "HND" + worker.kv_cache_layout = "LBHNC" if do_handshake: remote_agents, _ = worker._nixl_handshake( host="localhost", diff --git a/tests/v1/kv_connector/unit/test_kv_cache_layout.py b/tests/v1/kv_connector/unit/test_kv_cache_layout.py deleted file mode 100644 index 28eb3c61a0..0000000000 --- a/tests/v1/kv_connector/unit/test_kv_cache_layout.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -from vllm.platforms import current_platform - - -def test_mla_common_backend_rejects_cross_layer_kv_cache(): - """MLACommonBackend defaults to the identity permutation (layers dim - first) so MLA backends whose decode kernels are not verified to honor - the cache's block-dim stride stay opted out of cross-layer KV cache.""" - from vllm.model_executor.layers.attention.mla_attention import ( - MLACommonBackend, - ) - - stride_order = MLACommonBackend.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - assert stride_order == (0, 1, 2, 3) - assert stride_order[0] == 0 # layers dim first => no cross-layer - assert MLACommonBackend.get_kv_cache_stride_order( - include_num_layers_dimension=False - ) == (0, 1, 2) - - -@pytest.mark.parametrize( - "backend_path", - # See: https://github.com/vllm-project/vllm/issues/46411 - [ - "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend", - ] - if current_platform.is_rocm() or current_platform.is_xpu() - else [ - "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend", - "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend", - "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend", - "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend", - "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend", - ], -) -def test_verified_mla_backends_support_cross_layer_kv_cache(backend_path): - """Backends whose decode kernels honor the cache's block-dim stride opt - in to the cross-layer layout with a non-identity permutation placing - num_blocks first in physical layout.""" - module_path, name = backend_path.rsplit(".", 1) - backend = getattr( - pytest.importorskip(module_path, reason="backend deps unavailable"), name - ) - - stride_order = backend.get_kv_cache_stride_order(include_num_layers_dimension=True) - assert stride_order == (1, 0, 2, 3) - assert stride_order[0] != 0 # num_blocks first => cross-layer supported - assert backend.get_kv_cache_stride_order(include_num_layers_dimension=False) == ( - 0, - 1, - 2, - ) - - -def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache(): - """DeepseekV32Indexer returns identity permutation (layers dim first) - to signal cross-layer KV cache is unsupported.""" - from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV32IndexerBackend, - ) - - stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - assert stride_order == (0, 1, 2, 3) - assert stride_order[0] == 0 # layers dim first => no cross-layer - assert DeepseekV32IndexerBackend.get_kv_cache_stride_order( - include_num_layers_dimension=False - ) == (0, 1, 2) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index 4847b956b1..74382051da 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -11,6 +11,7 @@ import pytest import torch import zmq.asyncio +from tests.v1.attention.utils import dense_kv_cache_views from vllm import envs from vllm.config import set_current_vllm_config 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 MooncakeBootstrapServer, ) from vllm.utils.network_utils import get_open_port -from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheLayout, ) from vllm.v1.request import RequestStatus @@ -1093,7 +1094,19 @@ async def test_worker_get_finished_timeout(monkeypatch): assert "tx-active" in prefill_worker.reqs_need_send -def test_register_kv_caches(): +@pytest.mark.parametrize( + ("layout", "separate_kv_head_groups"), + [ + (KVCacheLayout.LBHNC, False), + (KVCacheLayout.BLHNC, False), + (KVCacheLayout.BHLNC, False), + # LHBNC gives each head group its own region; the K/V split doubles the + # head count but the registration shape is driven by the layout. + (KVCacheLayout.LHBNC, False), + (KVCacheLayout.LHBNC, True), + ], +) +def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool): """Tests the memory registration logic with the underlying Mooncake engine.""" vllm_config = create_vllm_config( @@ -1118,15 +1131,23 @@ def test_register_kv_caches(): worker = connector.connector_worker mock_thread.return_value.is_alive.return_value = False - kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape( - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64 + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + num_head_slots=2 if separate_kv_head_groups else None, + state_content_bytes=4 * 64 * 2 if separate_kv_head_groups else None, ) - tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16) - tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16) - kv_caches = { - "model.layers.0.self_attn": tensor1, - "model.layers.1.self_attn": tensor2, - } + layer_names = [ + "model.layers.0.self_attn", + "model.layers.1.self_attn", + ] + for layer_name in layer_names: + worker._layer_specs[layer_name] = spec + raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8) + tensor1, tensor2 = dense_kv_cache_views(raw, spec, 2, 2, layout) + kv_caches = dict(zip(layer_names, (tensor1, tensor2))) with patch.object( worker.engine, "batch_register_memory", return_value=0 @@ -1135,16 +1156,35 @@ def test_register_kv_caches(): mock_batch_register.assert_called_once() registered_ptrs, registered_lens = mock_batch_register.call_args[0] - expected_ptrs = {tensor.data_ptr() for tensor in kv_caches.values()} - assert set(registered_ptrs) == expected_ptrs - assert set(registered_lens) == {tensor1.nbytes} + assert registered_ptrs == [raw.data_ptr()] + assert registered_lens == [raw.nbytes] - # Verify block_len_per_layer is set correctly. - assert len(worker.block_len_per_layer) == len(registered_ptrs) - for bl in worker.block_len_per_layer: - assert bl == tensor1.nbytes // tensor1.shape[0] - assert worker.registered_layer_names == list(kv_caches) - assert worker.registered_layer_indices == [0, 1] + if not layout.is_block_compact: + expected_addrs = [ + cache[:, head_idx].data_ptr() + for cache in (tensor1, tensor2) + for head_idx in range(cache.shape[1]) + ] + head_block_bytes = tensor1.stride(0) * tensor1.element_size() + assert worker.kv_caches_base_addr == expected_addrs + assert worker.block_len_per_layer == [head_block_bytes] * len( + expected_addrs + ) + assert worker.kv_block_len_per_layer == [head_block_bytes] * len( + expected_addrs + ) + assert worker.registered_layer_names == [ + layer_name + for layer_name in layer_names + for _ in range(tensor1.shape[1]) + ] + else: + assert len(worker.block_len_per_layer) == len(kv_caches) + for bl in worker.block_len_per_layer: + assert bl == tensor1.stride(0) * tensor1.element_size() + assert worker.kv_block_len_per_layer == [spec.page_size_bytes] * 2 + assert worker.registered_layer_names == list(kv_caches) + assert worker.registered_layer_indices == [0, 1] def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py index 6811d60ee6..fa6482781a 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py @@ -14,6 +14,7 @@ from unittest.mock import patch import pytest import torch +from tests.v1.attention.utils import dense_kv_cache_views from vllm.config import set_current_vllm_config from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import ( KVConnectorRole, @@ -30,6 +31,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheLayout, MambaSpec, ) @@ -138,14 +140,22 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): ) worker = connector.connector_worker - fa_cache = torch.empty((2, 2, 11), dtype=torch.float16) - gdn_conv_state = torch.empty((2, 22), dtype=torch.float16) - gdn_ssm_state = torch.empty((2, 4), dtype=torch.float16) + num_blocks = kv_cache_config.num_blocks + fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec + gdn_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec + fa_raw = torch.empty(num_blocks * fa_spec.page_size_bytes, dtype=torch.int8) + gdn_raw = torch.empty(num_blocks * gdn_spec.page_size_bytes, dtype=torch.int8) + (fa_cache,) = dense_kv_cache_views( + fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC + ) + (gdn_cache,) = dense_kv_cache_views( + gdn_raw, gdn_spec, num_blocks, 1, KVCacheLayout.LBHNC + ) worker.register_kv_caches( { "model.layers.0.self_attn": fa_cache, - "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state), + "model.layers.1.linear_attn": gdn_cache, } ) @@ -154,10 +164,14 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): "model.layers.0.self_attn", "model.layers.1.linear_attn", ] + assert worker.block_len_per_layer == [ + fa_spec.page_size_bytes, + gdn_spec.page_size_bytes, + ] assert worker.registered_group_indices == [0, 1] assert worker.kv_caches_base_addr == [ fa_cache.data_ptr(), - gdn_conv_state.data_ptr(), + gdn_cache.data_ptr(), ] worker.shutdown() @@ -185,8 +199,7 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch): backing = torch.empty((4, 64), dtype=torch.float16) fa_cache = backing[:2, :16] - gdn_conv_state = backing[:3] - gdn_ssm_state = torch.empty((3, 4), dtype=torch.float16) + gdn_cache = backing[:3] with patch.object( worker.engine, "batch_register_memory", return_value=0 @@ -194,13 +207,13 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch): worker.register_kv_caches( { "model.layers.0.self_attn": fa_cache, - "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state), + "model.layers.1.linear_attn": gdn_cache, } ) assert worker.kv_caches_base_addr == [ fa_cache.data_ptr(), - gdn_conv_state.data_ptr(), + gdn_cache.data_ptr(), ] batch_register_memory.assert_called_once() 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(): assert kernel_block_ids == [list(range(34, 51)), [2]] -def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole( +def test_hybrid_gdn_keeps_packed_fa_and_gdn_regions_whole( monkeypatch, ): monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") @@ -359,7 +372,7 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole( ) worker = connector.connector_worker - worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=True) + worker.transfer_topo = SimpleNamespace(is_kv_layout_blocks_first=False) regions = worker._get_transfer_regions( base_addrs=[0x1000, 0x2000], block_lens=[0x100, 0x100], @@ -377,7 +390,6 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole( for region in regions ] == [ (0, 0x1000, 0x40), - (0, 0x1040, 0x40), (1, 0x2000, 0x100), ] diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 4dcddcaa30..729c75970d 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -48,7 +48,14 @@ def _make_kv_cache_config() -> KVCacheConfig: spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) return KVCacheConfig( num_blocks=4, - kv_cache_tensors=[KVCacheTensor(size=8192, shared_by=["layer0"])], + kv_cache_tensors=[ + KVCacheTensor( + size=8192, + layers=["layer0"], + layer_stride=8192, + block_stride=2048, + ) + ], kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)], ) @@ -210,64 +217,6 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events(): assert kv_events.get_all_events() == [event] -def test_prefer_cross_layer_blocks_from_config(): - # Default: disabled - vllm_config = _make_vllm_config() - kv_cache_config = _make_kv_cache_config() - with ( - set_current_vllm_config(vllm_config), - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." - "connector.MooncakeStoreScheduler" - ), - ): - connector = mooncake_store_connector.MooncakeStoreConnector( - vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config - ) - assert connector.prefer_cross_layer_blocks is False - - # Enabled via config - vllm_config_enabled = create_vllm_config( - kv_connector="MooncakeStoreConnector", - kv_role="kv_both", - kv_connector_extra_config={"enable_cross_layers_blocks": "true"}, - ) - with ( - set_current_vllm_config(vllm_config_enabled), - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." - "connector.MooncakeStoreScheduler" - ), - ): - connector_enabled = mooncake_store_connector.MooncakeStoreConnector( - vllm_config_enabled, KVConnectorRole.SCHEDULER, kv_cache_config - ) - assert connector_enabled.prefer_cross_layer_blocks is True - - -def test_register_cross_layers_kv_cache_delegates_to_worker(): - vllm_config = _make_vllm_config() - kv_cache_config = _make_kv_cache_config() - - with ( - set_current_vllm_config(vllm_config), - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." - "connector.MooncakeStoreWorker" - ) as mock_worker_cls, - ): - connector = mooncake_store_connector.MooncakeStoreConnector( - vllm_config, KVConnectorRole.WORKER, kv_cache_config - ) - - fake_tensor = MagicMock() - fake_backend = MagicMock() - connector.register_cross_layers_kv_cache(fake_tensor, fake_backend) - - worker = mock_worker_cls.return_value - worker.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor) - - def test_update_connector_output_and_take_events(): vllm_config = _make_vllm_config() kv_cache_config = _make_kv_cache_config() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index cd23186906..43b123261d 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch import torch +from tests.v1.attention.utils import dense_kv_cache_views from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import ( worker as mooncake_store_worker, ) @@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import ( MambaSpec, SlidingWindowSpec, ) +from vllm.v1.kv_cache_layout import KVCacheLayout class _DictStore: @@ -160,8 +162,18 @@ def test_e2e_swa_plus_full_save_then_lookup_hits(): cfg = KVCacheConfig( num_blocks=4, kv_cache_tensors=[ - KVCacheTensor(size=8192, shared_by=["L0"]), - KVCacheTensor(size=8192, shared_by=["L1"]), + KVCacheTensor( + size=4 * full.page_size_bytes, + layers=["L0"], + layer_stride=4 * full.page_size_bytes, + block_stride=full.page_size_bytes, + ), + KVCacheTensor( + size=4 * swa.page_size_bytes, + layers=["L1"], + layer_stride=4 * swa.page_size_bytes, + block_stride=swa.page_size_bytes, + ), ], kv_cache_groups=[ KVCacheGroupSpec(["L0"], full), @@ -178,9 +190,11 @@ def test_e2e_swa_plus_full_save_then_lookup_hits(): # Register kv_caches using mocked thread classes so register_kv_caches # doesn't try to start real background threads (which set ready_event). + raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8) + raw_swa = torch.zeros(4 * swa.page_size_bytes, dtype=torch.int8) kv_caches = { - "L0": torch.zeros(2, 4, 8, 8, 64), - "L1": torch.zeros(2, 4, 8, 8, 64), + "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0], + "L1": dense_kv_cache_views(raw_swa, swa, 4, 1, KVCacheLayout.LBNHC)[0], } def _fake_thread_init(*args, **kwargs): diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 55081f90bf..7db71c6847 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -16,6 +16,7 @@ from unittest.mock import MagicMock, patch import pytest import torch +from tests.v1.attention.utils import dense_kv_cache_views from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import ( rdma_utils, ) @@ -38,6 +39,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import ) from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + KVCacheLayout, +) class _RecordingBlockHashes: @@ -2961,112 +2967,83 @@ def test_putting_consumer_queues_decode_save(): assert req.current_event is event -def test_register_kv_caches_blocks_first_single_segment(): - """Blocks-first layout (FlashInfer/MLA): one segment per layer.""" +@pytest.mark.parametrize("layout", list(KVCacheLayout)) +def test_register_kv_caches_shared_storage(layout: KVCacheLayout): num_blocks = 10 - page_size_elements = 64 + num_layers = 2 worker = _make_bare_worker(num_gpu_blocks=num_blocks) - - # Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims - tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16) - _register_with_mocked_threads(worker, {"layer0": tensor}) - - db = worker.token_dbs[0] - assert db.kv_caches_base_addr == [tensor.untyped_storage().data_ptr()] - assert db.block_len == [tensor.untyped_storage().nbytes() // num_blocks] - worker.store.register_buffer.assert_called_once_with( - tensor.untyped_storage().data_ptr(), - tensor.untyped_storage().nbytes(), - ) - - -def test_register_kv_caches_kv_first_two_segments(): - """K/V-first layout (FlashAttn): two segments (K, V) per layer.""" - num_blocks = 10 - block_size_tokens = 16 - num_kv_heads = 4 - head_size = 8 - worker = _make_bare_worker(num_gpu_blocks=num_blocks) - - # Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost - tensor = torch.zeros( - 2, - num_blocks, - block_size_tokens, - num_kv_heads, - head_size, + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=2, + head_size=8, dtype=torch.float16, ) - _register_with_mocked_threads(worker, {"layer0": tensor}) - - db = worker.token_dbs[0] - seg_stride = tensor.stride(0) * tensor.element_size() - base = tensor.untyped_storage().data_ptr() - assert db.kv_caches_base_addr == [base, base + seg_stride] - assert db.block_len == [seg_stride // num_blocks] * 2 - - -def test_register_kv_caches_cross_layer_single_segment(): - """Cross-layer tensor: single segment with block_len = page_size * num_layers.""" - num_blocks = 10 - num_layers = 4 - per_layer_page_elements = 64 # elements per layer per block - - worker = _make_bare_worker(num_gpu_blocks=num_blocks) - - # Cross-layer blocks-first tensor: all layers packed into a single - # contiguous block. Shape (num_blocks, num_layers * per_layer_page) - # mimics the physical layout after stride reordering. - total_page_elements = num_layers * per_layer_page_elements - tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16) - - with ( - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." - "worker.KVCacheStoreSendingThread", - side_effect=_auto_set_ready_event, - ), - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." - "worker.KVCacheStoreRecvingThread", - side_effect=_auto_set_ready_event, - ), - ): - # Use the cross-layer wrapper key, same as register_cross_layers_kv_caches - worker.register_kv_caches({"__cross_layer__": tensor}) - - db = worker.token_dbs[0] - assert len(db.kv_caches_base_addr) == 1 - assert db.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr() - - expected_block_len = tensor.untyped_storage().nbytes() // num_blocks - # block_len should be per_layer_page_size * num_layers - assert ( - expected_block_len - == num_layers * per_layer_page_elements * tensor.element_size() + raw = torch.zeros( + num_blocks * num_layers * spec.page_size_bytes, + dtype=torch.int8, ) - assert len(db.block_len) == 1 - assert db.block_len[0] == expected_block_len + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout) - # Also verify via register_cross_layers_kv_caches wrapper - worker2 = _make_bare_worker(num_gpu_blocks=num_blocks) - with ( - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." - "worker.KVCacheStoreSendingThread", - side_effect=_auto_set_ready_event, - ), - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." - "worker.KVCacheStoreRecvingThread", - side_effect=_auto_set_ready_event, - ), - ): - worker2.register_cross_layers_kv_caches(tensor) + _register_with_mocked_threads( + worker, + {"layer0": caches[0], "__cross_layer__": caches[1]}, + ) - db2 = worker2.token_dbs[0] - assert db2.kv_caches_base_addr == db.kv_caches_base_addr - assert db2.block_len == db.block_len + db = worker.token_dbs[0] + if not layout.is_block_compact: + # Each head group is its own region, so a block spans L*H regions. + head_cache = caches[0][:, 0] + assert db.kv_caches_base_addr == [ + cache[:, head_idx].data_ptr() + for cache in caches + for head_idx in range(cache.shape[1]) + ] + assert db.block_len == [head_cache.stride(0) * head_cache.element_size()] * ( + num_layers * caches[0].shape[1] + ) + elif layout.is_layer_compact: + assert db.kv_caches_base_addr == [cache.data_ptr() for cache in caches] + assert db.block_len == [spec.page_size_bytes] * num_layers + else: + assert db.kv_caches_base_addr == [raw.data_ptr()] + assert db.block_len == [num_layers * spec.page_size_bytes] + worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes) + + +def test_register_kv_caches_separate_head_groups(): + # LHBNC gives each head group its own region; the K/V split doubles heads. + layout = KVCacheLayout.LHBNC + num_blocks = 3 + num_layers = 2 + worker = _make_bare_worker(num_gpu_blocks=num_blocks) + spec = FullAttentionSpec( + block_size=4, + num_kv_heads=2, + head_size=8, + dtype=torch.float16, + num_head_slots=2, + state_content_bytes=2 * 8 * 2, + ) + layer_names = ["layer0", "__cross_layer__"] + worker._kv_cache_groups = [KVCacheGroupSpec(layer_names, spec)] + raw = torch.zeros( + num_blocks * num_layers * spec.page_size_bytes, + dtype=torch.int8, + ) + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout) + + _register_with_mocked_threads(worker, dict(zip(layer_names, caches))) + + head_block_bytes = caches[0].stride(0) * caches[0].element_size() + expected_addrs = [ + cache[:, head_idx].data_ptr() + for cache in caches + for head_idx in range(cache.shape[1]) + ] + db = worker.token_dbs[0] + assert db.kv_caches_base_addr == expected_addrs + assert db.block_len == [head_block_bytes] * len(expected_addrs) + worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes) # --------------------------------------------------------------------------- diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index ae62313d57..d169779c5d 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -43,6 +43,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + compute_layer_kv_cache_shape_bytes, ) from .utils import create_request, create_scheduler @@ -56,20 +57,25 @@ def _find_free_port() -> int: def _make_test_kv_cache_config() -> KVCacheConfig: layer_names = ["layer0", "layer1", "layer2"] + num_blocks = 2 + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + ) + page = spec.page_size_bytes return KVCacheConfig( - num_blocks=2, - kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)], - kv_cache_groups=[ - KVCacheGroupSpec( - layer_names=layer_names, - kv_cache_spec=FullAttentionSpec( - block_size=16, - num_kv_heads=4, - head_size=64, - dtype=torch.float16, - ), + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=len(layer_names) * num_blocks * page, + layers=layer_names, + layer_stride=num_blocks * page, + block_stride=page, ) ], + kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)], ) @@ -225,7 +231,7 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker): engine_id, *args, hand_shake_latency: float = 1.8, - kv_cache_layout="HND", + kv_cache_layout="LBHNC", kv_cache_config=None, **kwargs, ): @@ -531,16 +537,15 @@ def test_register_kv_caches(mock_parallel_groups): DEFAULT_PORT = 6301 TP_RANK = 0 DP_RANK = 0 - from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend - - backend_cls = AiterFlashAttentionBackend - - # Create test kv cache tensors using proper backend shape - kv_cache_shape = backend_cls.get_kv_cache_shape( - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64 + # Create test kv cache tensors using KVCacheSpec layout + shape = compute_layer_kv_cache_shape_bytes( + FullAttentionSpec( + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + ), + 2, ) - shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) + shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16) + unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16) kv_caches = { "layer0": shared_tensor, "layer1": unique_tensor, @@ -627,16 +632,15 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups): ROLE = "kv_consumer" vllm_config = create_vllm_config(role=ROLE) - from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend - - backend_cls = AiterFlashAttentionBackend - - # Create test kv cache tensors using proper backend shape - kv_cache_shape = backend_cls.get_kv_cache_shape( - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64 + # Create test kv cache tensors using KVCacheSpec layout + shape = compute_layer_kv_cache_shape_bytes( + FullAttentionSpec( + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + ), + 2, ) - shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16) + shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16) + unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16) kv_caches = { "layer0": shared_tensor, "layer1": unique_tensor, diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 6155932ea7..bceb574b8d 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -12,7 +12,10 @@ import pytest import torch from vllm.platforms import current_platform -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MLAAttentionSpec, +) aiter_available = importlib.util.find_spec("aiter") is not None 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: id="interleaved-kernel-axis-from-spec", ), pytest.param( - (8, 4, 3), + (8, 1, 4, 3), _mla_spec(), 16, { @@ -297,7 +300,7 @@ def test_mixed_layers_compute_distinct_offsets_per_layer(): kv_caches = { "separated": torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16), "interleaved": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16), - "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16), + "indexer": torch.empty((8, 1, 4, 3), dtype=torch.bfloat16), } worker = _worker( kv_caches, @@ -342,7 +345,7 @@ def test_write_transfer_plan_caches_offsets_per_geometry(): kv_caches = { "dense0": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16), "dense1": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16), - "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16), + "indexer": torch.empty((8, 1, 4, 3), dtype=torch.bfloat16), } calls: list[str] = [] @@ -692,7 +695,7 @@ def test_empty_local_block_ids_is_free_only_noop(): def test_registration_regions_do_not_split_interleaved_or_mla_cache(): separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16) interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) - indexer = torch.empty((8, 4, 3), dtype=torch.bfloat16) + indexer = torch.empty((8, 1, 4, 3), dtype=torch.bfloat16) worker = _worker( { "separated": separated, @@ -745,7 +748,9 @@ def test_registration_regions_use_layer_num_blocks(): def test_unsupported_shape_raises_value_error(): - cache = torch.empty((8, 4, 2, 3), dtype=torch.bfloat16) + # A 4-D view whose head/state/content dims all disagree with the spec (a + # standardized [B, H, N, C] view for _full_spec would be (8, 2, 4, 6)). + cache = torch.empty((8, 3, 5, 7), dtype=torch.bfloat16) worker = _worker({"layer": cache}, {"layer": _full_spec()}) with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"): diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 28d2d0dc13..6dce07834e 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -889,16 +889,6 @@ Options: """) -def test_multi_connector_prefer_cross_layer_blocks(mc): - mc._connectors[0].prefer_cross_layer_blocks = False - mc._connectors[1].prefer_cross_layer_blocks = True - assert mc.prefer_cross_layer_blocks is False - - mc._connectors[0].prefer_cross_layer_blocks = True - mc._connectors[1].prefer_cross_layer_blocks = True - assert mc.prefer_cross_layer_blocks is True - - def test_multi_connector_worker_metadata(mc): class MockConnectorWorkerMetadata(KVConnectorWorkerMetadata): def __init__(self, data: set[str]): diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 4b2ae31ba5..d4064d6f6f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -19,6 +19,7 @@ import pytest import ray import torch +from tests.v1.attention.utils import dense_kv_cache_views from vllm import LLM from vllm.config import KVTransferConfig, set_current_vllm_config from vllm.distributed.kv_transfer.kv_connector.utils import ( @@ -54,8 +55,6 @@ from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.platforms.interface import Platform from vllm.sampling_params import SamplingParams -from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend -from vllm.v1.attention.backends.utils import set_kv_cache_layout from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.output_processor import OutputProcessor from vllm.v1.kv_cache_interface import ( @@ -63,12 +62,12 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheTensor, + KVCacheLayout, + MLAAttentionSpec, + compute_layer_kv_cache_shape_bytes, ) from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus -from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin -from vllm.v1.worker.utils import AttentionGroup from .utils import ( create_request, @@ -96,9 +95,6 @@ def clear_kv_transfer(): yield if has_kv_transfer_group(): ensure_kv_transfer_shutdown() - # Reset any KV cache layout override set during tests so it doesn't - # leak into tests in other modules. - set_kv_cache_layout(None) def get_default_xfer_telemetry( @@ -350,12 +346,12 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): ) def test_kv_transfer_handshake(dist_init): """Unit test for basic NixlConnector interface functionality.""" - from vllm.config import set_current_vllm_config # Test setup, we creates a scheduler that contains a NixlConnector # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from # all workers of the instance. vllm_config = create_vllm_config() + vllm_config.cache_config.kv_cache_layout = "BLHNC" # in case the test runs on non-GPU machine vllm_config.kv_transfer_config.kv_buffer_device = "cpu" scheduler = create_scheduler(vllm_config) @@ -386,18 +382,19 @@ def test_kv_transfer_handshake(dist_init): kv_cache_spec = cast( AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec ) - kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape( - num_blocks=kv_cache_config.num_blocks, - block_size=kv_cache_spec.block_size, - num_kv_heads=kv_cache_spec.num_kv_heads, - head_size=kv_cache_spec.head_size, + raw = torch.zeros( + kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 3, + dtype=torch.int8, + ) + caches = dense_kv_cache_views( + raw, + kv_cache_spec, + kv_cache_config.num_blocks, + num_layers=3, + layout=KVCacheLayout.BLHNC, ) - shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) - unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) kv_caches = { - "layer0": shared_tensor, - "layer1": unique_tensor, - "layer2": shared_tensor, + f"layer{layer_idx}": cache for layer_idx, cache in enumerate(caches) } prefill_connector.register_kv_caches(kv_caches) @@ -474,7 +471,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): self, *args, hand_shake_latency: float = 1.8, - kv_cache_layout="HND", + kv_cache_layout="LBHNC", kv_cache_config=None, **kwargs, ): @@ -486,9 +483,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): # Mock register_kv_caches attributes needed for tests that do not call it. self.src_xfer_handles_by_block_size = {self.block_size: 1} self.src_blocks_data = np.empty((0, 3), dtype=np.uint64) - test_shape = self.attn_backends[0].get_kv_cache_shape( - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 - ) + rep_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec + test_shape = compute_layer_kv_cache_shape_bytes(rep_spec, 1) self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, tp_size=self.world_size, @@ -502,7 +498,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): ) self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks + self.vllm_config, self.backend_name ) def _nixl_handshake( @@ -553,9 +549,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): device_id=remote_tp_rank, num_blocks=1, block_lens=remote_block_lens, - # `self.kv_cache_layout` is only forced to HND when vllm engine - # is started. We mock HND here. - kv_cache_layout="HND", + block_strides=remote_block_lens, + kv_cache_layout="LBHNC", block_size=self.block_size, ssm_sizes=(0, 0), attn_backend_name=self.backend_name, @@ -579,6 +574,8 @@ class TestNixlHandshake: self, default_vllm_config, dist_init, pcp_rank ): """Only PCP rank zero publishes and reports sending completion.""" + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + vllm_config = create_vllm_config(kv_role="kv_producer") vllm_config.parallel_config.prefill_context_parallel_size = 2 with ( @@ -656,7 +653,7 @@ class TestNixlHandshake: worker.dst_xfer_side_handles = { FakeNixlConnectorWorker.REMOTE_ENGINE_ID: {0: 1} } - worker.kv_cache_layout = "HND" + worker.kv_cache_layout = "LBHNC" num_xfers = 4 while True: # For the same request_id, initiate multiple xfers across different @@ -1053,7 +1050,9 @@ class TestNixlHandshake: worker.dst_num_blocks[worker.engine_id] = worker.num_blocks # Metadata with different kv_cache_layout than local worker - mismatched_layout = "HND" if worker.kv_cache_layout != "HND" else "NHD" + mismatched_layout = ( + "LBHNC" if worker.kv_cache_layout != "LBHNC" else "LBNHC" + ) meta = NixlAgentMetadata( engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, agent_metadata=FakeNixlWrapper.AGENT_METADATA, @@ -1061,6 +1060,7 @@ class TestNixlHandshake: device_id=0, num_blocks=1, block_lens=worker.block_len_per_layer, + block_strides=worker.block_len_per_layer, kv_cache_layout=mismatched_layout, block_size=worker.block_size, ssm_sizes=(0, 0), @@ -1100,7 +1100,7 @@ class TestNixlHandshake: vllm_config, connector.engine_id, hand_shake_latency=0, - kv_cache_layout="NHD", + kv_cache_layout="LBNHC", ) worker = connector.connector_worker @@ -1111,15 +1111,17 @@ class TestNixlHandshake: worker.dst_num_blocks[worker.engine_id] = worker.num_blocks # Metadata with different kv_cache_layout than local worker + # prefill TP=1, decode TP=2, remote block_lens is double to local + remote_block_lens = [i * 2 for i in worker.block_len_per_layer] meta = NixlAgentMetadata( engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, agent_metadata=FakeNixlWrapper.AGENT_METADATA, kv_caches_base_addr=[0], device_id=0, num_blocks=1, - # prefill TP=1, decode TP=2, remote block_lens is double to local - block_lens=[i * 2 for i in worker.block_len_per_layer], - kv_cache_layout="HND", + block_lens=remote_block_lens, + block_strides=remote_block_lens, + kv_cache_layout="LBHNC", block_size=worker.block_size, ssm_sizes=(0, 0), attn_backend_name=worker.backend_name, @@ -1160,8 +1162,6 @@ class TestNixlHandshake: attn_backends=worker.attn_backends, tensor_shape=None, ) - assert worker.transfer_topo.virtually_split_kv_in_blocks - plan = MagicMock( source_ranks_per_group=((0,), (0,)), rank_offset_factor=1, @@ -1171,6 +1171,7 @@ class TestNixlHandshake: device_id=0, num_blocks=1, block_lens=[remote_block_len], + block_strides=[remote_block_len], ) assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ @@ -1226,6 +1227,7 @@ class TestNixlHandshake: device_id=0, num_blocks=1, block_lens=[fa_len * tp_ratio, idx_len], + block_strides=[fa_len * tp_ratio, idx_len], kv_cache_layout=worker.kv_cache_layout, block_size=worker.block_size, ssm_sizes=(0, 0), @@ -1252,6 +1254,7 @@ class TestNixlHandshake: num_blocks=1, # WRONG: MLA region scaled by tp_ratio (it should be replicated). block_lens=[fa_len * tp_ratio, idx_len * tp_ratio], + block_strides=[fa_len * tp_ratio, idx_len * tp_ratio], kv_cache_layout=worker2.kv_cache_layout, block_size=worker2.block_size, ssm_sizes=(0, 0), @@ -1295,7 +1298,7 @@ class TestNixlHandshake: worker.transfer_topo.total_num_kv_heads = 8 worker.transfer_topo.local_physical_heads = 1 - worker.kv_cache_layout = "HND" + worker.kv_cache_layout = "LBHNC" worker.slot_size_per_layer = [4096] worker.block_len_per_layer = [4096 * worker.block_size] @@ -1311,7 +1314,8 @@ class TestNixlHandshake: device_id=0, num_blocks=1, block_lens=list(worker.block_len_per_layer), - kv_cache_layout="HND", + block_strides=list(worker.block_len_per_layer), + kv_cache_layout="LBHNC", block_size=worker.block_size, ssm_sizes=(0, 0), attn_backend_name=worker.backend_name, @@ -1350,7 +1354,7 @@ class TestNixlHandshake: worker.transfer_topo.total_num_kv_heads = 32 worker.transfer_topo.local_physical_heads = 8 # 32 // 4 - worker.kv_cache_layout = "HND" + worker.kv_cache_layout = "LBHNC" slot_size = 4096 worker.slot_size_per_layer = [slot_size] @@ -1368,7 +1372,8 @@ class TestNixlHandshake: device_id=0, num_blocks=1, block_lens=list(worker.block_len_per_layer), - kv_cache_layout="HND", + block_strides=list(worker.block_len_per_layer), + kv_cache_layout="LBHNC", block_size=worker.block_size, ssm_sizes=(0, 0), attn_backend_name=worker.backend_name, @@ -1839,7 +1844,19 @@ def _run_abort_timeout_test(llm: LLM, timeout: int): llm.llm_engine.engine_core.shutdown() -@pytest.mark.parametrize("enable_cross_layers", ["False", "True"]) +def element_byte_addrs(view: torch.Tensor) -> list[int]: + """Absolute byte addresses of every element of a (possibly strided) view.""" + offsets = torch.zeros(view.shape, dtype=torch.int64) + for dim, (size, stride) in enumerate(zip(view.shape, view.stride())): + shape = [1] * view.ndim + shape[dim] = size + offsets = offsets + torch.arange(size, dtype=torch.int64).view(shape) * stride + esize = view.element_size() + byte_offsets = (offsets.flatten() * esize).tolist() + base = view.data_ptr() + return [base + offset + i for offset in byte_offsets for i in range(esize)] + + @pytest.mark.parametrize( "attn_backend", [ @@ -1853,8 +1870,14 @@ def _run_abort_timeout_test(llm: LLM, timeout: int): "TRITON_ATTN", ], ) +@pytest.mark.parametrize("layout", [layout.name for layout in KVCacheLayout]) +@pytest.mark.parametrize("separate_kv_head_groups", [False, True]) def test_register_kv_caches( - default_vllm_config, dist_init, attn_backend, enable_cross_layers + default_vllm_config, + dist_init, + attn_backend, + layout, + separate_kv_head_groups, ): """ Test that register_kv_caches() properly calls nixl_wrapper methods with @@ -1868,12 +1891,7 @@ def test_register_kv_caches( """ vllm_config = create_vllm_config(attention_backend=attn_backend) - - # Enable cross layers blocks - vllm_config.kv_transfer_config.kv_connector_extra_config[ - "enable_cross_layers_blocks" - ] = enable_cross_layers - set_kv_cache_layout("HND") + vllm_config.cache_config.kv_cache_layout = layout # Import the appropriate backend based on the parameter if attn_backend == "FLASH_ATTN": @@ -1890,59 +1908,37 @@ def test_register_kv_caches( backend_cls = TritonAttentionBackend nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker" - nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector" with ( patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper, patch(f"{nixl_worker}.threading.Event"), patch(f"{nixl_worker}.threading.Thread") as mock_thread, - patch(f"{nixl_connector}.get_current_attn_backend") as mock_get_attn_backend, patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends, ): - # Ensure get_attn_backend returns the correct value due to - # _cached_get_attn_backend returning the backend from previous - # test run if not mocking. - mock_get_attn_backend.return_value = backend_cls mock_get_attn_backends.return_value = [backend_cls] - num_layers = 32 block_size = 16 num_blocks = 8 num_heads = 4 head_size = 16 - # TODO (NickLucche) the fact that connector depends on kv_cache_config for init - # but cross-layer preference cant be inferred prior to creating kv_cache_config - # is a bit awkward. - dummy_connector = NixlConnector( - vllm_config, - KVConnectorRole.WORKER, - make_kv_cache_config(block_size=block_size), - ) kv_cache_spec = FullAttentionSpec( block_size=block_size, num_kv_heads=num_heads, head_size=head_size, dtype=torch.float16, + num_head_slots=2 if separate_kv_head_groups else None, + state_content_bytes=num_heads * head_size * 2 + if separate_kv_head_groups + else None, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0", "layer1", "layer2", "layer3"], kv_cache_spec + ) + ], ) - if dummy_connector.prefer_cross_layer_blocks: - kv_cache_config = KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=kv_cache_spec.page_size_bytes * num_blocks, - shared_by=["all-layers"], - ) - for _ in range(num_layers) - ], - kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)], - ) - else: - kv_cache_config = KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec) - ], - ) # Create connector connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( @@ -1962,93 +1958,36 @@ def test_register_kv_caches( # Reassure the shutdown() check that the thread is terminated mock_thread.return_value.is_alive.return_value = False - expected_tensor_size: int - expected_base_addrs: list[int] - expected_num_entries: int - kv_caches: dict[str, torch.Tensor] - if str(enable_cross_layers).lower() == "true": - assert connector.prefer_cross_layer_blocks == ( - attn_backend in ("FLASH_ATTN", "FLASHINFER", "TRITON_ATTN") - ) - else: - assert not connector.prefer_cross_layer_blocks - - test_shape = backend_cls.get_kv_cache_shape( - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 + raw0 = torch.zeros( + kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 2, + dtype=torch.int8, + device=current_platform.device_type, ) - is_blocks_first = len(test_shape) == 4 and test_shape[0] == 1 - - if connector.prefer_cross_layer_blocks: - with set_current_vllm_config(vllm_config): - _, cross_layers_kv_cache, _ = ( - KVConnectorModelRunnerMixin.allocate_uniform_kv_caches( - kv_cache_config=kv_cache_config, - attn_groups=[ - [ - AttentionGroup( - backend=backend_cls, - layer_names=[], - kv_cache_spec=kv_cache_spec, - kv_cache_group_id=0, - ) - ] - ], - cache_dtype="bfloat16", - device=torch.accelerator.current_device_index(), - kernel_block_sizes=[block_size], - ) - ) - # Store tensor info for validation - expected_tensor_size = ( - cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel() - ) - expected_base_addrs = [ - cross_layers_kv_cache.data_ptr(), - ] - expected_num_entries = 1 - - expected_blocks_count = num_blocks - - kv_caches = {"all-layers": cross_layers_kv_cache} - else: - # Create test kv cache tensors using proper backend shape - kv_cache_shape = backend_cls.get_kv_cache_shape( - num_blocks=kv_cache_config.num_blocks, - block_size=kv_cache_spec.block_size, - num_kv_heads=kv_cache_spec.num_kv_heads, - head_size=kv_cache_spec.head_size, - ) - shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) - unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) - kv_caches = { - "layer0": shared_tensor, - "layer1": unique_tensor, - "layer2": shared_tensor, - } - - # Store tensor info for validation - if is_blocks_first: - expected_tensor_size = ( - shared_tensor.element_size() * shared_tensor.numel() - ) - expected_base_addrs = [ - shared_tensor.data_ptr(), - unique_tensor.data_ptr(), - ] - expected_num_entries = 2 - expected_blocks_count = kv_cache_config.num_blocks * 2 - else: - expected_tensor_size = ( - shared_tensor[0].element_size() * shared_tensor[0].numel() - ) - expected_base_addrs = [ - shared_tensor[0].data_ptr(), - shared_tensor[1].data_ptr(), - unique_tensor[0].data_ptr(), - unique_tensor[1].data_ptr(), - ] - expected_num_entries = 4 - expected_blocks_count = kv_cache_config.num_blocks * 4 + raw1 = torch.zeros( + kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks, + dtype=torch.int8, + device=current_platform.device_type, + ) + tensor0, tensor1 = dense_kv_cache_views( + raw0, + kv_cache_spec, + kv_cache_config.num_blocks, + num_layers=2, + layout=KVCacheLayout[layout], + ) + (tensor2,) = dense_kv_cache_views( + raw1, + kv_cache_spec, + kv_cache_config.num_blocks, + num_layers=1, + layout=KVCacheLayout[layout], + ) + kv_caches = { + "layer0": tensor0, + "layer1": tensor1, + "layer2": tensor2, + "layer3": tensor0, + } # Execute register_kv_caches connector.register_kv_caches(kv_caches) @@ -2056,44 +1995,123 @@ def test_register_kv_caches( # Verify get_reg_descs was called with caches_data assert mock_wrapper_instance.get_reg_descs.called caches_data, _ = mock_wrapper_instance.get_reg_descs.call_args[0] - assert len(caches_data) == expected_num_entries + assert len(caches_data) == 2 - for i, cache_entry in enumerate(caches_data): + for cache_entry, raw in zip(caches_data, (raw0, raw1)): base_addr, size, _tp_rank, _ = cache_entry - assert size == expected_tensor_size, ( - f"Entry {i}: Expected tensor size {expected_tensor_size}, got {size}" - ) - assert base_addr == expected_base_addrs[i], ( - f"Entry {i}: Expected base address {expected_base_addrs[i]}, " - f"got {base_addr}" - ) + assert size == raw.nbytes + assert base_addr == raw.data_ptr() # Verify get_xfer_descs was called with blocks_data assert mock_wrapper_instance.get_xfer_descs.called blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0] - # Validate blocks_data structure and size - assert len(blocks_data) == expected_blocks_count, ( - f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}" - ) + # Layout-blind contract: whatever regions the worker carves out, + # transferring "block b" must move exactly logical block b's bytes for + # every layer. Map each registered byte to the block that owns it, then + # require every descriptor window to hold bytes of a single block and + # the windows of block b to cover exactly block b's bytes. + owner: dict[int, int] = {} + for cache in (tensor0, tensor1, tensor2): + for blk in range(num_blocks): + for addr in element_byte_addrs(cache[blk]): + assert owner.setdefault(addr, blk) == blk + block_bytes = defaultdict(set) + for addr, blk in owner.items(): + block_bytes[blk].add(addr) - if connector.prefer_cross_layer_blocks: - num_blocks = 8 - else: - num_blocks = kv_cache_config.num_blocks + covered: defaultdict[int, set[int]] = defaultdict(set) + for block_start_addr, block_len, _tp_rank in blocks_data: + window = range(block_start_addr, block_start_addr + block_len) + owners = {owner[addr] for addr in window} + assert len(owners) == 1, "descriptor window spans logical blocks" + covered[owners.pop()].update(window) + assert covered == block_bytes - expected_block_len = expected_tensor_size // num_blocks - - for i, block_entry in enumerate(blocks_data): - block_start_addr, block_len, tp_rank = block_entry - assert block_len == expected_block_len, ( - f"Block entry {i}: Expected block len {expected_block_len}, " - f"got {block_len}" - ) + # Region bases are exactly the block-0 window starts, one per region. + base_addrs = connector.connector_worker.kv_caches_base_addr[ + connector.connector_worker.engine_id + ][0] + if layout == "BLHNC": + assert len(base_addrs) == 3 + assert set(base_addrs) == { + start for start, _len, _tp in blocks_data if owner[start] == 0 + } + assert len(blocks_data) == num_blocks * len(base_addrs) assert connector.connector_worker.block_size == 16 +def test_register_packed_dsv4_mla_cache_as_single_region( + default_vllm_config, dist_init +): + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend + + nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker" + with ( + patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper, + patch(f"{nixl_worker}.threading.Event"), + patch(f"{nixl_worker}.threading.Thread") as mock_thread, + patch(f"{nixl_worker}.get_current_attn_backends") as mock_backends, + ): + mock_backends.return_value = [TritonAttentionBackend] + num_blocks = 2 + num_layers = 4 + block_size = 256 + spec = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + cache_dtype_str="fp8_ds_mla", + tokens_per_state=4, + alignment=576, + model_version="deepseek_v4", + state_content_bytes=584, + ) + layer_names = [f"layer{idx}" for idx in range(num_layers)] + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)], + ) + vllm_config = create_vllm_config(attention_backend="TRITON_ATTN") + vllm_config.cache_config.block_size = block_size + vllm_config.cache_config.kv_cache_layout = "BLHNC" + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, + connector.engine_id, + hand_shake_latency=0, + kv_cache_layout="BLHNC", + kv_cache_config=kv_cache_config, + ) + wrapper = mock_nixl_wrapper.return_value + connector.connector_worker.nixl_wrapper = wrapper + wrapper.get_agent_metadata.return_value = b"fake_agent_metadata" + mock_thread.return_value.is_alive.return_value = False + + raw = torch.zeros( + spec.page_size_bytes * num_blocks * num_layers, + dtype=torch.int8, + device=current_platform.device_type, + ) + views = dense_kv_cache_views( + raw, spec, num_blocks, num_layers, KVCacheLayout.BLHNC + ) + connector.register_kv_caches(dict(zip(layer_names, views))) + + blocks_data, _ = wrapper.get_xfer_descs.call_args[0] + packed_block_len = num_layers * spec.page_size_bytes + assert connector.connector_worker.kv_caches_base_addr[ + connector.connector_worker.engine_id + ][0] == [raw.data_ptr()] + assert blocks_data.tolist() == [ + [raw.data_ptr() + block_idx * packed_block_len, packed_block_len, 0] + for block_idx in range(num_blocks) + ] + + class FakePlatform(Platform): device_type: str = "oot" @@ -2318,6 +2336,8 @@ def test_engine_ttl_disabled(default_vllm_config, dist_init): def test_transfer_topology_unregister(): """TransferTopology.unregister_remote_engine removes the engine.""" + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + topo = TransferTopology( tp_rank=0, tp_size=1, @@ -2908,8 +2928,8 @@ def test_speculative_config_compatibility_hash( _set_test_speculative_config(local_config) _set_test_speculative_config(remote_config, **remote_overrides) - local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False) - remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False) + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN") + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN") assert (local_hash == remote_hash) is should_match @@ -2920,10 +2940,8 @@ def test_missing_speculative_config_changes_compatibility_hash(): speculative_config = create_vllm_config() _set_test_speculative_config(speculative_config) - regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN", False) - speculative_hash = compute_nixl_compatibility_hash( - speculative_config, "FLASH_ATTN", False - ) + regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN") + speculative_hash = compute_nixl_compatibility_hash(speculative_config, "FLASH_ATTN") assert regular_hash != speculative_hash @@ -2938,8 +2956,8 @@ def test_speculative_kv_cache_dtype_resolves_to_target(): _set_test_speculative_config(local_config, kv_cache_dtype="fp8") # explicit _set_test_speculative_config(remote_config, kv_cache_dtype=None) # inherits - local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False) - remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False) + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN") + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN") assert local_hash == remote_hash @@ -2955,8 +2973,8 @@ def test_speculative_attention_backend_not_in_compatibility_hash(): _set_test_speculative_config(local_config, attention_backend=None) _set_test_speculative_config(remote_config, attention_backend="FLASHINFER") - local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False) - remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False) + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN") + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN") assert local_hash == remote_hash @@ -2969,17 +2987,17 @@ def test_transfer_mode_changes_compatibility_hash(): config = create_vllm_config() pull_hash = compute_nixl_compatibility_hash( - config, "FLASH_ATTN", False, transfer_mode="pull" + config, "FLASH_ATTN", transfer_mode="pull" ) push_hash = compute_nixl_compatibility_hash( - config, "FLASH_ATTN", False, transfer_mode="push" + config, "FLASH_ATTN", transfer_mode="push" ) assert pull_hash != push_hash assert pull_hash == compute_nixl_compatibility_hash( - config, "FLASH_ATTN", False, transfer_mode="pull" + config, "FLASH_ATTN", transfer_mode="pull" ) - assert compute_nixl_compatibility_hash(config, "FLASH_ATTN", False) == pull_hash + assert compute_nixl_compatibility_hash(config, "FLASH_ATTN") == pull_hash @pytest.mark.skip_global_cleanup @@ -3057,14 +3075,11 @@ def test_compatibility_hash_validation( kv_cache_spec = cast( AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec ) - kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape( - num_blocks=kv_cache_config.num_blocks, - block_size=kv_cache_spec.block_size, - num_kv_heads=kv_cache_spec.num_kv_heads, - head_size=kv_cache_spec.head_size, + shape = compute_layer_kv_cache_shape_bytes( + kv_cache_spec, kv_cache_config.num_blocks ) - shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) - unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) + shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype) + unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype) # Build kv_caches from the actual layer names in kv_cache_config so that # _layer_specs lookups in register_kv_caches always find a matching key. layer_names = [ @@ -3099,18 +3114,19 @@ def test_compatibility_hash_validation( remote_hash = compute_nixl_compatibility_hash( remote_vllm_config, decode_worker.backend_name, - decode_worker.transfer_topo.cross_layers_blocks, ) prefill_block_size = config_overrides.get("block_size", 16) + prefill_block_lens = [4096 * prefill_block_size] prefill_metadata = NixlAgentMetadata( engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, agent_metadata=FakeNixlWrapper.AGENT_METADATA, kv_caches_base_addr=[0], device_id=0, num_blocks=1, - block_lens=[4096 * prefill_block_size], # slot_size * block_size - kv_cache_layout="HND", + block_lens=prefill_block_lens, + block_strides=prefill_block_lens, + kv_cache_layout="LBHNC", block_size=prefill_block_size, ssm_sizes=(0, 0), attn_backend_name=decode_worker.backend_name, @@ -3187,9 +3203,8 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) decode_worker = decode_connector.connector_worker backend = get_current_attn_backend(local_vllm_config) - test_shape = backend.get_kv_cache_shape( - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 - ) + probe_spec = decode_worker.kv_cache_config.kv_cache_groups[0].kv_cache_spec + test_shape = compute_layer_kv_cache_shape_bytes(probe_spec, 1) decode_worker.transfer_topo = TransferTopology( tp_rank=decode_worker.tp_rank, tp_size=decode_worker.world_size, @@ -3205,7 +3220,6 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) decode_worker.compat_hash = compute_nixl_compatibility_hash( decode_worker.vllm_config, decode_worker.backend_name, - decode_worker.transfer_topo.cross_layers_blocks, ) if error_scenario == "handshake_decode_error": diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 11538f0219..66a90d64e5 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1505,14 +1505,19 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): mamba_type=MambaAttentionBackendEnum.GDN_ATTN, ) assert kda_spec.page_size_bytes == unified_page + # The three groups overlay each other, so layer i of every group aliases the same + # region: mla.i, kda_a.i and kda_b.i all live at i * layer_stride. + layer_stride = num_blocks * unified_page return KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=[ KVCacheTensor( - size=num_blocks * unified_page, - shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + size=2 * layer_stride, + layers=[f"{prefix}.0", f"{prefix}.1"], + layer_stride=layer_stride, + block_stride=unified_page, ) - for i in range(2) + for prefix in ("mla", "kda_a", "kda_b") ], kv_cache_groups=[ KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 6ac9ce030d..daf6251dce 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -129,14 +129,19 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo page_size_padded=unified_page, mamba_type=MambaAttentionBackendEnum.GDN_ATTN, ) + # The three groups overlay each other, so layer i of every group aliases the same + # region: mla.i, kda_a.i and kda_b.i all live at i * layer_stride. + layer_stride = num_logical_blocks * unified_page kv_cache_config = KVCacheConfig( num_blocks=num_logical_blocks, kv_cache_tensors=[ KVCacheTensor( - size=num_logical_blocks * unified_page, - shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + size=2 * layer_stride, + layers=[f"{prefix}.0", f"{prefix}.1"], + layer_stride=layer_stride, + block_stride=unified_page, ) - for i in range(2) + for prefix in ("mla", "kda_a", "kda_b") ], kv_cache_groups=[ KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), @@ -230,6 +235,8 @@ def _make_remote_meta( kv_caches_base_addr=[0x10_000_000, 0x20_000_000], num_blocks=remote_num_logical * remote_ppl, block_lens=[kernel_page, kernel_page], + # Non-interleaved remote: consecutive blocks abut, so stride == page length. + block_strides=[kernel_page, kernel_page], kv_cache_layout=worker.kv_cache_layout, block_size=remote_kernel_block_size, ssm_sizes=remote_ssm_sizes, diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index e7326eddcb..d7b6edf07c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -1081,6 +1081,8 @@ class TestPushPipelineParallel: device_id=0, num_blocks=4, block_lens=[block_len] * 4, + # Non-interleaved: consecutive blocks abut, so stride == block length. + block_strides=[block_len] * 4, kv_cache_layout="HND", block_size=16, ssm_sizes=(0, 0), diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py index 78e9e1196f..16d0087673 100644 --- a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py +++ b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py @@ -145,10 +145,15 @@ def _make_kv_cache_config( ), ) ] + # Each group packs its two layers densely; groups overlay one allocation. + layer_stride = _BYTES_PER_BLOCK * num_blocks + size = 2 * layer_stride tensors = [ KVCacheTensor( - size=_BYTES_PER_BLOCK * num_blocks, - shared_by=fa_layers, + size=size, + layers=fa_layers, + layer_stride=layer_stride, + block_stride=_BYTES_PER_BLOCK, ) ] if swa_enabled: @@ -167,8 +172,10 @@ def _make_kv_cache_config( ) tensors.append( KVCacheTensor( - size=_BYTES_PER_BLOCK * num_blocks, - shared_by=sw_layers, + size=size, + layers=sw_layers, + layer_stride=layer_stride, + block_stride=_BYTES_PER_BLOCK, ) ) return KVCacheConfig( diff --git a/tests/v1/kv_connector/unit/test_transfer_topology_sharded.py b/tests/v1/kv_connector/unit/test_transfer_topology_sharded.py index 1fed8ff3b7..c5f1f01088 100644 --- a/tests/v1/kv_connector/unit/test_transfer_topology_sharded.py +++ b/tests/v1/kv_connector/unit/test_transfer_topology_sharded.py @@ -12,14 +12,7 @@ pytestmark = pytest.mark.cpu_test class _FakeAttentionBackend: - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - ) -> tuple[int, int, int, int]: - return (num_blocks, num_kv_heads, block_size, 2 * head_size) + pass def _make_topology( diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index b1fb433533..1bc09995a3 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -130,6 +130,8 @@ def create_vllm_config( cache_dtype=cache_dtype, enable_prefix_caching=True, ) + # Connectors are constructed after layout resolution; mirror that here. + cache_config.kv_cache_layout = "LBNHC" kv_transfer_config = KVTransferConfig( kv_connector=kv_connector, kv_connector_module_path=kv_connector_module_path, diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 9ba1e56487..26da4079b4 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -55,6 +55,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: ), replicated_layout=kwargs.get("replicated_layout", False), canonical_layout=kwargs.get("canonical_layout", False), + kv_cache_layout=kwargs.get("kv_cache_layout", "LBNHC"), ) spec = MagicMock(spec=OffloadingSpec) spec.config = config @@ -214,14 +215,8 @@ def test_parallel_agnostic_separates_persistent_layouts(): def test_canonical_layout_changes_storage_namespace(): # Canonical bytes are not interchangeable with the direct layout, so the # format id must fork the storage namespace. - from vllm.v1.attention.backends.utils import set_kv_cache_layout - - set_kv_cache_layout("NHD") - try: - direct = make_mapper_from_offloading_spec() - canonical = make_mapper_from_offloading_spec(canonical_layout=True) - finally: - set_kv_cache_layout(None) + direct = make_mapper_from_offloading_spec() + canonical = make_mapper_from_offloading_spec(canonical_layout=True) assert "canonical_format" not in direct.fields assert canonical.fields["canonical_format"] == "v1-nhd" assert direct.base_path != canonical.base_path diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 4f7761cf35..714617a4f6 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -91,7 +91,9 @@ def _make_kv_cache_config( tensors.append( KVCacheTensor( size=_BYTES_PER_BLOCK * num_blocks, - shared_by=layer_names, + layers=layer_names, + layer_stride=_BYTES_PER_BLOCK * num_blocks, + block_stride=_BYTES_PER_BLOCK, ) ) return KVCacheConfig( diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 5856120fe7..645cd5daba 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -10,6 +10,7 @@ read partially written / stale blocks and silently corrupt the CPU cache. from __future__ import annotations import time +from unittest.mock import MagicMock import pytest import torch @@ -19,6 +20,8 @@ from vllm.platforms import current_platform if not current_platform.is_cuda_alike(): pytest.skip("Requires CUDA or ROCm", allow_module_level=True) +from tests.v1.attention.utils import dense_kv_cache_views +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend from vllm.v1.simple_kv_offload.cuda_mem_ops import ( CU_MEMCPY_SRC_ACCESS_ORDER_ANY, @@ -188,3 +191,90 @@ def test_build_params_src_access_order(): gpu, cpu, stream, src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM ) assert ordered.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_STREAM + + +@pytest.mark.parametrize("layout", list(KVCacheLayout)) +def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout): + num_blocks = 4 + num_layers = 2 + spec = FullAttentionSpec( + block_size=2, + num_kv_heads=2, + head_size=2, + dtype=torch.float16, + ) + raw = torch.zeros( + num_blocks * num_layers * spec.page_size_bytes, + dtype=torch.int8, + device="cuda", + ) + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout) + cache_config = MagicMock(num_blocks=num_blocks) + worker = SimpleCPUOffloadWorker( + vllm_config=None, + kv_cache_config=cache_config, + cpu_capacity_bytes=raw.nbytes, + ) + worker._backend = MagicMock() + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False) + + worker.register_kv_caches( + {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)} + ) + + assert worker.gpu_kv_caches is not None + if layout.is_layer_compact and not layout.is_block_compact: + expected_regions = num_layers * spec.num_heads + expected_block_bytes = spec.page_size_bytes // spec.num_heads + elif layout.is_layer_compact: + expected_regions = num_layers + expected_block_bytes = spec.page_size_bytes + else: + expected_regions = 1 + expected_block_bytes = spec.page_size_bytes * num_layers + assert len(worker.gpu_kv_caches) == expected_regions + assert {cache.shape for cache in worker.gpu_kv_caches.values()} == { + (num_blocks, expected_block_bytes) + } + + +def test_register_separate_kv_head_groups(monkeypatch): + # LHBNC hoists the K/V head groups outside the block dim, so each layer's + # blocks are registered as one region per group (K, V). + layout = KVCacheLayout.LHBNC + num_blocks = 4 + num_layers = 2 + spec = FullAttentionSpec( + block_size=2, + num_kv_heads=2, + head_size=2, + dtype=torch.float16, + num_head_slots=2, + state_content_bytes=2 * 2 * 2, + ) + raw = torch.zeros( + num_blocks * num_layers * spec.page_size_bytes, + dtype=torch.int8, + device="cuda", + ) + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout) + worker = SimpleCPUOffloadWorker( + vllm_config=None, + kv_cache_config=MagicMock(num_blocks=num_blocks), + cpu_capacity_bytes=raw.nbytes, + ) + worker._backend = MagicMock() + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False) + + worker.register_kv_caches( + {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)} + ) + + assert worker.gpu_kv_caches is not None + assert len(worker.gpu_kv_caches) == num_layers * spec.num_heads + per_group_block_bytes = ( + spec.num_kv_heads * spec.block_size * spec.head_size * spec.dtype.itemsize + ) + assert {cache.shape for cache in worker.gpu_kv_caches.values()} == { + (num_blocks, per_group_block_bytes) + } diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index b446c5a377..01ad759bcd 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -1,48 +1,33 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Padded-page handling in create_kv_cache_views. +Guards that a page_size_padded spec strides the block dimension by the padded page +while keeping per-block content compact, so padding bytes at the end of each page are +never addressed by the logical view. +""" + +import pytest import torch +from tests.v1.attention.utils import dense_kv_cache_views +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, + KVCacheGroupSpec, + KVCacheLayout, KVCacheTensor, - KVQuantMode, - MambaSpec, + MLAAttentionSpec, + compute_layout_strides, +) +from vllm.v1.worker.utils import ( + allocate_kv_cache, + copy_kv_cache_blocks_inplace, ) -from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache -from vllm.v1.worker.utils import AttentionGroup -class FakeFlashAttentionBackend: - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, 2, block_size, num_kv_heads, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - assert not include_num_layers_dimension - return (0, 1, 2, 3, 4) - - -class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend): - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - assert not include_num_layers_dimension - return (0, 1, 3, 2, 4) - - -def test_reshape_padded_flash_attention_kv_cache_strides_by_page(): +def test_reshape_padded_kv_cache_strides_by_padded_page(): num_blocks = 3 spec = FullAttentionSpec( block_size=16, @@ -53,292 +38,192 @@ def test_reshape_padded_flash_attention_kv_cache_strides_by_page(): ) assert spec.real_page_size_bytes == 256 - raw_tensors = { - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) - } - attn_groups = [ - AttentionGroup( - backend=FakeFlashAttentionBackend, - layer_names=["layer"], - kv_cache_spec=spec, - kv_cache_group_id=0, - ) - ] + raw = torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) + (kv_cache,) = dense_kv_cache_views(raw, spec, num_blocks, 1, KVCacheLayout.LBHNC) - kv_cache = _reshape_kv_cache( - attn_groups, - raw_tensors, - "auto", - [spec.block_size], - {}, - )["layer"] - - assert kv_cache.shape == (num_blocks, 2, 16, 1, 2) - assert kv_cache.stride(0) == spec.page_size_bytes // 4 - assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4 - assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4 - assert ( - kv_cache[1, 1].storage_offset() - == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4 - ) + elem_size = 4 # float32 + # Content dim packs K and V: 2 * head_size. + assert kv_cache.shape == (num_blocks, 1, 16, 2 * spec.head_size) + assert kv_cache.dtype == spec.dtype + assert kv_cache.stride(0) == spec.page_size_padded // elem_size + assert kv_cache[1].storage_offset() == spec.page_size_padded // elem_size + # Within one block the (unpadded) content stays compact. + assert kv_cache[0].is_contiguous() -def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page(): - num_blocks = 3 - spec = FullAttentionSpec( - block_size=16, - num_kv_heads=3, - head_size=2, - dtype=torch.float32, - page_size_padded=1024, - ) - assert spec.real_page_size_bytes == 768 - - raw_tensors = { - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) - } - attn_groups = [ - AttentionGroup( - backend=FakeHNDFlashAttentionBackend, - layer_names=["layer"], - kv_cache_spec=spec, - kv_cache_group_id=0, - ) - ] - - kv_cache = _reshape_kv_cache( - attn_groups, - raw_tensors, - "auto", - [spec.block_size], - {}, - )["layer"] - - assert kv_cache.shape == (num_blocks, 2, 16, 3, 2) - assert kv_cache.stride(0) == spec.page_size_bytes // 4 - assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4 - assert kv_cache.stride(2) == 2 - assert kv_cache.stride(3) == spec.block_size * spec.head_size - assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4 - assert ( - kv_cache[1, 1].storage_offset() - == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4 - ) - assert ( - kv_cache[1, 1, 3, 2].storage_offset() - == ( - spec.page_size_bytes - + spec.real_page_size_bytes // 2 - + 3 * spec.head_size * 4 - + 2 * spec.block_size * spec.head_size * 4 - ) - // 4 - ) - - -class FakeDiffKVBackend: - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, block_size, num_kv_heads, head_size * 2) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - assert not include_num_layers_dimension - return (0, 1, 2, 3) - - -def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim(): - num_blocks = 3 - spec = FullAttentionSpec( - block_size=16, +@pytest.mark.parametrize( + ("kernel_block_sizes", "expected_num_blocks", "expected_num_states"), + [ + (None, 4, 64), + ([256], 4, 64), + ([64], 16, 16), + ], +) +def test_allocate_compressed_mla_cache( + kernel_block_sizes: list[int] | None, + expected_num_blocks: int, + expected_num_states: int, +): + spec = MLAAttentionSpec( + block_size=256, num_kv_heads=1, - head_size=2, - dtype=torch.float32, - page_size_padded=384, + head_size=128, + dtype=torch.bfloat16, + tokens_per_state=4, ) - - raw_tensors = { - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) - } - attn_groups = [ - AttentionGroup( - backend=FakeDiffKVBackend, - layer_names=["layer"], - kv_cache_spec=spec, - kv_cache_group_id=0, - ) - ] - - kv_cache = _reshape_kv_cache( - attn_groups, - raw_tensors, - "auto", - [spec.block_size], - {}, - )["layer"] - - assert kv_cache.shape == (num_blocks, 16, 1, 4) - assert kv_cache.stride(0) == spec.page_size_bytes // 4 - assert kv_cache.stride(1) == 4 - - -class FakePerTokenScaleBackend: - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, 2, block_size, num_kv_heads, head_size + 4) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - assert not include_num_layers_dimension - return (0, 1, 2, 3, 4) - - -def test_reshape_padded_quantized_kv_cache_preserves_scale_stride(): - num_blocks = 3 - spec = FullAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=4, - dtype=torch.int8, - kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD, - page_size_padded=384, - ) - assert spec.real_page_size_bytes == 128 - assert spec.page_size_bytes == 384 - - raw_tensors = { - "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) - } - attn_groups = [ - AttentionGroup( - backend=FakePerTokenScaleBackend, - layer_names=["layer"], - kv_cache_spec=spec, - kv_cache_group_id=0, - ) - ] - - kv_cache = _reshape_kv_cache( - attn_groups, - raw_tensors, - "int8_per_token_head", - [spec.block_size], - {}, - )["layer"] - - assert kv_cache.shape == (num_blocks, 2, 16, 1, 8) - assert kv_cache.stride(0) == spec.page_size_bytes - assert kv_cache.stride(1) == 16 * 1 * 8 - assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8 - - -class FakeKVFirstBackend: - """ROCm-style backend that puts K and V ahead of the block dim.""" - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (2, num_blocks, block_size, num_kv_heads, head_size) - - @staticmethod - def get_kv_cache_block_dim( - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> int: - return 1 - - -def _kv_first_setup(shared_by: list[str]): - num_blocks = 3 - attn_spec = FullAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=2, - dtype=torch.float32, - ) - mamba_spec = MambaSpec( - block_size=16, - shapes=((64,),), - dtypes=(torch.float32,), - ) - assert attn_spec.page_size_bytes == mamba_spec.page_size_bytes == 256 - - raw_tensor = torch.zeros(attn_spec.page_size_bytes * num_blocks, dtype=torch.int8) - raw_tensors = {name: raw_tensor for name in shared_by} - attn_groups = [ - AttentionGroup( - backend=FakeKVFirstBackend, - layer_names=["attn"], - kv_cache_spec=attn_spec, - kv_cache_group_id=0, - ) - ] - if "mamba" in shared_by: - attn_groups.append( - AttentionGroup( - backend=FakeKVFirstBackend, - layer_names=["mamba"], - kv_cache_spec=mamba_spec, - kv_cache_group_id=1, - ) - ) - kv_cache_config = KVCacheConfig( - num_blocks=num_blocks, + num_pages = 4 + config = KVCacheConfig( + num_blocks=num_pages, kv_cache_tensors=[ - KVCacheTensor(size=raw_tensor.numel(), shared_by=list(shared_by)) + KVCacheTensor( + size=num_pages * spec.page_size_bytes, + layers=["layer.0"], + layer_stride=num_pages * spec.page_size_bytes, + block_stride=spec.page_size_bytes, + ) ], - kv_cache_groups=[], + kv_cache_groups=[KVCacheGroupSpec(["layer.0"], spec)], ) - kv_caches = _reshape_kv_cache( - attn_groups, - raw_tensors, - "auto", - [attn_spec.block_size] * len(attn_groups), - {}, - kv_cache_config, + caches = allocate_kv_cache( + config, torch.device("cpu"), KVCacheLayout.LBHNC, kernel_block_sizes ) - return num_blocks, kv_caches["attn"] + + assert caches["layer.0"].shape == (expected_num_blocks, 1, expected_num_states, 128) -def test_reshape_kv_first_kv_cache_pages_blocks_when_shared_with_mamba(): - num_blocks, kv_cache = _kv_first_setup(["attn", "mamba"]) +@pytest.mark.parametrize("layout", list(KVCacheLayout)) +def test_copy_kv_cache_blocks_shared_storage(layout: KVCacheLayout): + num_blocks = 4 + num_layers = 2 + spec = FullAttentionSpec( + block_size=2, + num_kv_heads=2, + head_size=2, + dtype=torch.float32, + ) + raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8) + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout) - assert kv_cache.shape == (2, num_blocks, 16, 1, 2) - # Block b has to own page b, so its K and V sit side by side within it. - page = 16 * 1 * 2 * 2 - for block in range(num_blocks): - assert kv_cache[0, block].storage_offset() == block * page - assert kv_cache[1, block].storage_offset() == block * page + page // 2 + for layer_idx, cache in enumerate(caches): + for block_idx in range(num_blocks): + cache[block_idx].fill_(10 * layer_idx + block_idx) + + expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches] + copies = [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)] + + copy_kv_cache_blocks_inplace(caches, num_blocks, copies) + + for layer_idx, cache in enumerate(caches): + torch.testing.assert_close(cache[2], expected[layer_idx][0]) + torch.testing.assert_close(cache[1], expected[layer_idx][1]) -def test_reshape_kv_first_kv_cache_keeps_layout_without_mamba(): - num_blocks, kv_cache = _kv_first_setup(["attn"]) +def test_fixed_block_stride_propagates_outward_in_lhbnc(): + num_blocks = 3 + num_layers = 2 + spec = FullAttentionSpec( + block_size=2, + num_kv_heads=2, + head_size=2, + dtype=torch.float32, + ) + natural = compute_layout_strides(spec, num_blocks, num_layers, KVCacheLayout.LHBNC) + block_stride = natural[1] + 8 - assert kv_cache.shape == (2, num_blocks, 16, 1, 2) - # Nothing else indexes this allocation by page, so K and V stay split into - # one contiguous half each. - assert kv_cache[1, 0].storage_offset() == num_blocks * 16 * 1 * 2 + strides = compute_layout_strides( + spec, + num_blocks, + num_layers, + KVCacheLayout.LHBNC, + fixed_strides=(None, block_stride, None, None, None), + ) + + assert strides[1] == block_stride + assert strides[2] == block_stride * num_blocks + assert strides[0] == strides[2] * spec.num_heads + + +def test_copy_kv_cache_blocks_separate_head_groups(): + # LHBNC stores each head group separately, so a block's bytes are scattered + # across L*H regions. + layout = KVCacheLayout.LHBNC + num_blocks = 4 + num_layers = 2 + spec = FullAttentionSpec( + block_size=2, + num_kv_heads=2, + head_size=2, + dtype=torch.float32, + num_head_slots=2, + state_content_bytes=2 * 2 * 4, + ) + raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8) + caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout) + + for layer_idx, cache in enumerate(caches): + for block_idx in range(num_blocks): + for head_idx in range(cache.shape[1]): + cache[block_idx, head_idx].fill_( + 100 * layer_idx + 10 * head_idx + block_idx + ) + + expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches] + copy_kv_cache_blocks_inplace( + caches, + num_blocks, + [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)], + ) + + for layer_idx, cache in enumerate(caches): + torch.testing.assert_close(cache[2], expected[layer_idx][0]) + torch.testing.assert_close(cache[1], expected[layer_idx][1]) + + +@pytest.mark.parametrize( + "layout,num_layers", + [ + (KVCacheLayout.LBHNC, 2), + # Splitting needs a manager block to be one dense page, which a + # block-outermost layout only gives when the block holds one layer. + (KVCacheLayout.BLHNC, 1), + ], +) +def test_copy_kv_cache_blocks_with_virtual_block_splitting( + layout: KVCacheLayout, num_layers: int +): + num_blocks = 4 + physical_per_logical = 2 + spec = FullAttentionSpec( + block_size=4, + num_kv_heads=1, + head_size=2, + dtype=torch.float32, + ) + raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8) + caches = dense_kv_cache_views( + raw, + spec, + num_blocks, + num_layers, + layout, + kernel_block_size=spec.block_size // physical_per_logical, + ) + + for layer_idx, cache in enumerate(caches): + for block_idx in range(cache.shape[0]): + cache[block_idx].fill_(100 * layer_idx + block_idx) + expected = [[cache[i].clone() for i in range(cache.shape[0])] for cache in caches] + + copy_kv_cache_blocks_inplace( + caches, + num_blocks, + [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)], + ) + + dst_start = 2 * physical_per_logical + for layer_idx, cache in enumerate(caches): + for physical_idx in range(physical_per_logical): + torch.testing.assert_close( + cache[dst_start + physical_idx], expected[layer_idx][physical_idx] + ) diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index d5a698433c..400cfbfc2e 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -1,35 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU regression for the packed-DSV4 KV zeroer geometry. +"""CPU regression for the packed-DSV4 KV zeroer geometry (upstream #50276). -Constructor-crossing regression for the packed-KV zeroing stride fix (upstream -PR #50276): every number below comes from instantiating the REAL production -code (no arithmetic is reimplemented here): +Four DSV4 fp8_ds_mla layers share one cross-layer allocation: each block row packs +the four alignment-padded pages side by side (layout BLHNC). The zeroer must step by +the full packed row per block while zeroing only each layer's meaningful page — never +the alignment padding, never an adjacent layer's page. - - ``vllm.v1.core.kv_cache_utils._get_packed_kv_cache_layout`` (packed - offsets / block stride) - - ``vllm.v1.worker.gpu.attn_utils._reshape_attention_kv_cache`` - (per-layer column views) - - ``vllm.models.deepseek_v4.sparse_mla.DeepseekV4FlashMLABackend`` - (``get_kv_cache_shape`` / ``get_kv_cache_block_dim``) - - ``vllm.v1.worker.utils.AttentionGroup`` / ``KVBlockZeroer`` - (``KVBlockZeroer.__init__`` precomputes the segment tables) - -CPU only: no Triton launch, no CUDA tensor, no CUDA call is made. The test -asserts the packed-DSV4 constructor metadata: - - * block stride = the full packed row (4 aligned pages) - * per-layer zero = the meaningful page (the spec's unpadded - ``real_page_size_bytes``), so adjacent packed layers are - never wiped - * 4 packed layer views; block 99 of the highest-offset layer is the last - written region and stays within the packed backing. - -The pre-fix constructor exposed a five-field ``_meta`` with no block stride, so -its zero width equals the FULL packed row; the extraction below records that -legacy defect and then asserts the fixed semantics, so failure on the old code -is a semantic one (wrong zero width / out-of-bounds block 99), not a -schema-shape accident. +CPU only: the zeroer's ``__init__`` precomputes the segment tables; no kernel is +launched. """ from types import SimpleNamespace @@ -37,160 +16,154 @@ from types import SimpleNamespace import pytest import torch -from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLABackend -from vllm.v1.core.kv_cache_utils import _get_packed_kv_cache_layout -from vllm.v1.kv_cache_interface import MLAAttentionSpec, UniformTypeKVCacheSpecs -from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache -from vllm.v1.worker.utils import AttentionGroup, KVBlockZeroer +from tests.v1.attention.utils import dense_kv_cache_views +from vllm.v1.kv_cache_interface import KVCacheLayout, MLAAttentionSpec +from vllm.v1.worker.utils import ( + AttentionGroup, + KVBlockZeroer, + allocate_kv_cache, +) pytestmark = pytest.mark.cpu_test NUM_BLOCKS = 100 -BLOCK_SIZE_CFG = 256 # vllm_config.cache_config.block_size for DSV4 (group 0) -COMPRESS = 4 # main-MLA compress_ratio -STORAGE_BS = BLOCK_SIZE_CFG // COMPRESS # 64 -KERNEL_BS = 256 # DeepseekV4FlashMLABackend supports [256] -LAST_BLOCK = NUM_BLOCKS - 1 # block 99 +NUM_LAYERS = 4 +ALIGNMENT = 576 -def _extract_segments(meta, base): - """Read (offset_bytes, block_stride_bytes, zero_span_bytes) per segment. - - Accepts both the pre-fix five-field ``_meta`` (no block stride; the zero - width is the FULL packed row -- the defect) and the fixed six-field ``_meta``. - All values come from the real constructor output; nothing is recomputed. Returns - ``(segs, legacy_why)`` where ``legacy_why`` describes the pre-fix defect - when detected. - """ - legacy_why = None - if len(meta) == 6: - seg_addrs, seg_block_strides, seg_page_sizes, _max_chunks, _blk, n_segs = meta - else: - seg_addrs, seg_page_sizes, _max_chunks, _blk, n_segs = meta - seg_block_strides = seg_page_sizes - legacy_why = ( - f"legacy {len(meta)}-field _meta (no block stride): zero width is " - f"the FULL packed row ({int(seg_page_sizes[0]) * 4} B) -- the defect" - ) - segs = [ - (int(a) - base, int(bs) * 4, int(ps) * 4) - for a, bs, ps in zip(seg_addrs, seg_block_strides, seg_page_sizes) - ] - return segs, legacy_why, n_segs - - -def test_dsv4_packed_zeroer_geometry(): - """Real DSV4 constructor metadata separates block stride from zero span.""" - layer_names = [f"model.layers.{i}.self_attn" for i in range(4)] - - specs = [ - MLAAttentionSpec( - block_size=BLOCK_SIZE_CFG, - num_kv_heads=1, - head_size=512, - dtype=torch.uint8, - compress_ratio=COMPRESS, - cache_dtype_str="fp8_ds_mla", - alignment=576, - model_version="deepseek_v4", - state_content_bytes=584, # >576 to allocate room for scales at back of page - ) - for _ in layer_names - ] - uniform = UniformTypeKVCacheSpecs.from_specs( - {ln: sp for ln, sp in zip(layer_names, specs)} +def test_packed_dsv4_zeroer_zeroes_only_each_layers_page(): + spec = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + tokens_per_state=4, + alignment=ALIGNMENT, + # DeepseekV4 fp8_ds_mla: 584B per token, published by the layer. + state_content_bytes=584, ) - assert uniform is not None + # DSV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per stored state. + assert spec.state_content_size_bytes == 584 + assert spec.num_states == 64 + unpadded_page = spec.unpadded_page_size_bytes + padded_page = spec.page_size_bytes + assert unpadded_page == 64 * 584 + assert padded_page > unpadded_page, "alignment must pad the page" - # --- Real packed layout planner ------------------------------------------- - block_stride, layers_by_offset = _get_packed_kv_cache_layout( - [ - AttentionGroup( - backend=DeepseekV4FlashMLABackend, - layer_names=layer_names, - kv_cache_spec=uniform, - kv_cache_group_id=0, - ) - ] - ) - page_bytes = specs[0].page_size_bytes # DSV4 fp8_ds_mla page after 576B alignment - real_page_bytes = specs[0].real_page_size_bytes # unpadded meaningful page - alignment_gap = page_bytes - real_page_bytes # trailing planner padding - total_size = block_stride * NUM_BLOCKS - offsets = sorted(layers_by_offset) - assert block_stride == 4 * page_bytes, (block_stride, 4 * page_bytes) - assert offsets == [i * page_bytes for i in range(4)] + raw = torch.zeros(NUM_BLOCKS * NUM_LAYERS * padded_page, dtype=torch.int8) + views = dense_kv_cache_views(raw, spec, NUM_BLOCKS, NUM_LAYERS, KVCacheLayout.BLHNC) + base = raw.data_ptr() + block_row = NUM_LAYERS * padded_page + for i, view in enumerate(views): + assert view.data_ptr() - base == i * padded_page + assert view.stride(0) * view.element_size() == block_row - # --- Real per-layer views (packed branch) -------------------------------- - backing = torch.zeros(total_size, dtype=torch.uint8) - base = backing.data_ptr() - views = {} - for ln, off in zip(layer_names, offsets): - packing = (off, block_stride) - kv_shape = DeepseekV4FlashMLABackend.get_kv_cache_shape( - NUM_BLOCKS, STORAGE_BS, 1, 512, cache_dtype_str="fp8_ds_mla" - ) - stride_order = tuple(range(len(kv_shape))) - view = _reshape_attention_kv_cache( - backing, specs[0], kv_shape, stride_order, NUM_BLOCKS, packing - ) - views[ln] = view - assert view.data_ptr() - base == off, "data_ptr must be base+offset" - assert view.stride(0) * view.element_size() == block_stride, ( - "row stride == block_stride (logical block stride)" - ) - - # --- Real KVBlockZeroer.__init__ (production arithmetic, CPU) ---------- - sctx = {ln: SimpleNamespace(kv_cache=v) for ln, v in views.items()} zeroer = KVBlockZeroer( torch.device("cpu"), - attn_groups_iter=[ - AttentionGroup( - backend=DeepseekV4FlashMLABackend, - layer_names=layer_names, - kv_cache_spec=specs[0], - kv_cache_group_id=0, - ) - ], - kernel_block_sizes=[KERNEL_BS], - cache_dtype="fp8_ds_mla", - static_forward_context=sctx, + attn_groups_iter=iter( + [ + AttentionGroup( + backend=None, + layer_names=[f"layer.{i}" for i in range(NUM_LAYERS)], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + ), + kernel_block_sizes=[spec.block_size], + static_forward_context={ + f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(NUM_LAYERS) + }, ) - segs, legacy_why, n_segs = _extract_segments(zeroer._meta, base) + seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta - # --- Fixed semantics (proven by metadata values, not by schema shape) ---- - assert n_segs == 4, "all 4 packed layers must register as segments" - assert [off for off, _, _ in segs] == offsets, "one segment per layer offset" - assert all(bs == block_stride for _, bs, _ in segs), ( - "block stride must be the full packed row (4 aligned pages)" + assert n_segs == NUM_LAYERS + # Segments step by the full packed row per block... + assert (seg_block_strides * 4 == block_row).all() + # ...but zero only the layer's meaningful page: no alignment padding, no + # adjacent layer's page. + assert (seg_page_sizes * 4 == unpadded_page).all() + assert sorted(a - base for a in seg_addrs.tolist()) == [ + i * padded_page for i in range(NUM_LAYERS) + ] + # Block 99 of the highest-offset layer stays within the packed backing. + last_end = ( + max(seg_addrs.tolist()) + + (NUM_BLOCKS - 1) * int(seg_block_strides[0]) * 4 + + int(seg_page_sizes[0]) * 4 ) - assert all(zs == real_page_bytes for _, _, zs in segs), ( - f"zero span must be the meaningful page ({real_page_bytes} B), got " - f"{[zs for _, _, zs in segs]} B" + (f"; {legacy_why}" if legacy_why else "") + assert last_end <= base + raw.numel() + + +def test_overlaid_zeroer_dedups_segments_with_max_span(): + """Two groups overlay one allocation; the zeroer must emit one segment per distinct + byte offset, spanning the widest overlaid page, so a newly allocated block is fully + zeroed no matter which group owns it.""" + from unittest.mock import MagicMock + + from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups + from vllm.v1.kv_cache_interface import ( + KVCacheGroupSpec, + KVCacheLayout, + UniformTypeKVCacheSpecs, ) - # --- Block 99 (the real final block) bounds, via constructor metadata --- - # Production write for block b of a segment is - # [off + b*block_stride_bytes, off + b*block_stride_bytes + zero_span_bytes). - worst = 0 - for off, bs, zs in segs: - start = off + LAST_BLOCK * bs - end = start + zs - oob = end - total_size - worst = max(worst, oob) - assert end <= total_size, ( - f"segment@off={off} block {LAST_BLOCK} writes [{start}, {end}) " - f"-> {oob} bytes past backing ({total_size})" + def make_spec(head_size): + return MLAAttentionSpec( + block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8 ) - assert worst == 0, f"highest-offset final block goes {worst} bytes OOB" - # Highest-offset segment's block 99 is the last written region and ends exactly at - # total_size - alignment_gap: the trailing planner padding stays untouched. - off_last, bs_last, zs_last = segs[-1] - assert off_last + LAST_BLOCK * bs_last + zs_last == total_size - alignment_gap, ( - "highest-offset final block must end at total_size - alignment_gap " - f"({alignment_gap}-byte gap intact)" + g1_specs = {"g1.big": make_spec(512), "g1.small": make_spec(128)} + g2_specs = {"g2.huge": make_spec(1024)} + groups = [ + KVCacheGroupSpec( + list(g1_specs), + UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g1_specs), + ), + KVCacheGroupSpec( + list(g2_specs), + UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g2_specs), + ), + ] + from vllm.config import CacheConfig + + vllm_config = MagicMock() + vllm_config.cache_config = CacheConfig() + vllm_config.cache_config.num_gpu_blocks_override = None + vllm_config.cache_config.kv_cache_layout = "BLHNC" + config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024) + views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None) + buf_ptr = views["g1.big"].data_ptr() + + attn_groups = [ + AttentionGroup( + backend=None, + layer_names=list(specs), + kv_cache_spec=next(iter(specs.values())), + kv_cache_group_id=gid, + ) + for gid, specs in enumerate((g1_specs, g2_specs)) + ] + zeroer = KVBlockZeroer( + torch.device("cpu"), + attn_groups_iter=iter(attn_groups), + kernel_block_sizes=[64, 64], + static_forward_context={ + name: SimpleNamespace(kv_cache=views[name]) for name in views + }, ) + seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta - # Zero span < block stride proves adjacent packed layers are never wiped. - assert max(zs for _, _, zs in segs) < block_stride + # g1.big and g2.huge overlay at offset 0 -> one segment with g2's wider + # span; g1.small keeps its own segment. + assert n_segs == 2 + pages = {n: s.page_size_bytes for n, s in (g1_specs | g2_specs).items()} + by_offset = { + a - buf_ptr: p * 4 for a, p in zip(seg_addrs.tolist(), seg_page_sizes.tolist()) + } + assert by_offset[0] == max(pages["g1.big"], pages["g2.huge"]) + assert by_offset[pages["g1.big"]] == pages["g1.small"] + packed_block_stride = max(sum(pages[n] for n in g) for g in (g1_specs, g2_specs)) + assert (seg_block_strides * 4 == packed_block_stride).all() diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 6eaa074cc4..94eb805484 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -92,7 +92,12 @@ def initialize_kv_cache(runner: GPUModelRunner): kv_cache_config = KVCacheConfig( num_blocks=NUM_BLOCKS, kv_cache_tensors=[ - KVCacheTensor(size=tensor_size, shared_by=["layer.0"]), + KVCacheTensor( + size=tensor_size, + layers=["layer.0"], + layer_stride=tensor_size, + block_stride=attn_spec.page_size_bytes, + ), ], kv_cache_groups=[ KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec) @@ -131,6 +136,7 @@ def get_vllm_config(): gpu_memory_utilization=0.9, cache_dtype="auto", ) + cache_config.kv_cache_layout = "LBNHC" parallel_config = ParallelConfig() vllm_config = VllmConfig( model_config=model_config, @@ -839,56 +845,6 @@ def test_update_states_pp_async_multi_request_keeps_rank_state_consistent( ) -def test_kv_cache_stride_order(monkeypatch, model_runner): - # This test checks if GPUModelRunner initializes correctly when an attention - # backend enforces a non-default KV cache stride order. - n_heads = model_runner.model_config.get_num_kv_heads(model_runner.parallel_config) - head_size = model_runner.model_config.get_head_size() - - # Get the expected shape from the backend's get_kv_cache_shape method - # to ensure compatibility with different backends (triton vs flexattention) - attn_backend = None - for attn_group in model_runner._attn_group_iterator(): - attn_backend = attn_group.backend - break - - assert attn_backend is not None, "No attention backend found" - expected_kv_cache_shape = list( - attn_backend.get_kv_cache_shape(NUM_BLOCKS, BLOCK_SIZE, n_heads, head_size) - ) - - # TODO mla test - default_stride = tuple(range(len(expected_kv_cache_shape))) - non_default_stride = (*default_stride[1:], default_stride[0]) - # Permutation that gets you back to expected kv shape - for test_stride in (non_default_stride, default_stride): - - def rnd_stride_order( - include_num_layers_dimension: bool = False, test_stride=test_stride - ): - assert not include_num_layers_dimension - return test_stride - - # Patch the attention backend class and re-trigger the KV cache creation - for attn_group in model_runner._attn_group_iterator(): - attn_backend = attn_group.backend - monkeypatch.setattr( - attn_backend, "get_kv_cache_stride_order", rnd_stride_order - ) - - model_runner.attn_groups = [] - model_runner.kv_caches = [] - model_runner.initialize_kv_cache(model_runner.kv_cache_config) - - # Shape is unchanged, but layout may differ - kv_cache_shape = model_runner.kv_caches[0].shape - assert list(kv_cache_shape) == expected_kv_cache_shape - if default_stride == test_stride: - assert all(kv.is_contiguous() for kv in model_runner.kv_caches) - else: - assert all(not kv.is_contiguous() for kv in model_runner.kv_caches) - - def test_update_config(model_runner): # Simple update model_runner.update_config({"load_config": {"load_format": "dummy"}}) @@ -1124,21 +1080,22 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config): vllm_config, [kv_cache_spec], [available_memory] )[0] assert kv_cache_config.num_blocks == num_expected_blocks - assert len(kv_cache_config.kv_cache_tensors) == 2 - assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2 - assert kv_cache_config.kv_cache_tensors[1].size == available_memory // 2 + assert len(kv_cache_config.kv_cache_tensors) == 1 + assert kv_cache_config.kv_cache_tensors[0].size == available_memory max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes) # max context len with KV sharing should be 2x as large as without assert max_context_len == 1310720 # important: override tensor size to prevent large mem alloc during test - # this will only allocate 2 block worth of memory (2 * 32kb) + # this will only allocate 1 block worth of memory per layer (2 layers * 32kb) kv_cache_config.num_blocks = 1 for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - kv_cache_tensor.size = kv_cache_spec[ - kv_cache_tensor.shared_by[0] - ].page_size_bytes + page_size = kv_cache_spec[kv_cache_tensor.layers[0]].page_size_bytes + kv_cache_tensor.size = page_size * len(kv_cache_tensor.layers) + # One block per layer: a layer's region is exactly one page. + kv_cache_tensor.layer_stride = page_size + kv_cache_tensor.block_stride = page_size runner.initialize_kv_cache(kv_cache_config) @@ -1207,7 +1164,9 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config): # important: override tensor size to prevent large mem alloc during test # this will only allocate 1 block worth of memory (32kb) kv_cache_config.num_blocks = 1 - kv_cache_config.kv_cache_tensors[0].size = kv_cache_spec[layer_0].page_size_bytes + page_size = kv_cache_spec[layer_0].page_size_bytes + kv_cache_config.kv_cache_tensors[0].size = page_size + kv_cache_config.kv_cache_tensors[0].layer_stride = page_size runner.initialize_kv_cache(kv_cache_config) kv_cache_config_after_init = runner.kv_cache_config @@ -1232,7 +1191,7 @@ def test_hybrid_attention_mamba_tensor_shapes(): """ The GPU model runner creates different views into the KVCacheTensors for the attention and mamba layers - (via _reshape_kv_cache_tensors function). This test verifies + (via _allocate_kv_caches). This test verifies that the views are compatible: writing a mamba block will not corrupt an attention block and vice versa """ @@ -1270,6 +1229,7 @@ def test_hybrid_attention_mamba_tensor_shapes(): gpu_memory_utilization=0.9, cache_dtype="auto", ) + cache_config.kv_cache_layout = "LBNHC" parallel_config = ParallelConfig() attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER) vllm_config = VllmConfig( @@ -1561,7 +1521,12 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init): kv_cache_config = KVCacheConfig( num_blocks=NUM_BLOCKS, kv_cache_tensors=[ - KVCacheTensor(size=tensor_size, shared_by=["layer.0"]), + KVCacheTensor( + size=tensor_size, + layers=["layer.0"], + layer_stride=tensor_size, + block_stride=attn_spec.page_size_bytes, + ), ], kv_cache_groups=[ 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(): gpu_memory_utilization=0.9, cache_dtype="auto", ) + cache_config.kv_cache_layout = "LBNHC" parallel_config = ParallelConfig() attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER) vllm_config = VllmConfig( diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 3502898872..945046ad6f 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -7,8 +7,11 @@ from types import SimpleNamespace import pytest import torch +from tests.v1.attention.utils import dense_kv_cache_views from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + FullAttentionSpec, + KVCacheLayout, SlidingWindowSpec, ) from vllm.v1.worker import utils as worker_utils @@ -19,12 +22,6 @@ from vllm.v1.worker.utils import ( ) -class _BlockFirstBackend: - @staticmethod - def get_kv_cache_block_dim(*args, **kwargs): - return 0 - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( "spec", @@ -52,11 +49,8 @@ def test_attention_blocks_are_zeroed(spec): layer_name = "draft.self_attn" zeroer = KVBlockZeroer( device, - attn_groups_iter=[ - AttentionGroup(_BlockFirstBackend, [layer_name], spec, 0) # type: ignore[arg-type] - ], + attn_groups_iter=[AttentionGroup(None, [layer_name], spec, 0)], kernel_block_sizes=[2], - cache_dtype="fp8", static_forward_context={ layer_name: SimpleNamespace(kv_cache=storage), }, @@ -205,16 +199,10 @@ def test_large_dsv4_launch_geometry(monkeypatch): zeroer = KVBlockZeroer( device, attn_groups_iter=[ - AttentionGroup( - _BlockFirstBackend, # type: ignore[arg-type] - [name], - spec, - group_id, - ) + AttentionGroup(None, [name], spec, group_id) for group_id, name in enumerate(layer_names) ], kernel_block_sizes=[1] * n_segs, - cache_dtype="auto", static_forward_context={ name: SimpleNamespace(kv_cache=storage) for name, storage in storages.items() @@ -311,3 +299,46 @@ def test_warmup_respects_available_block_count(): torch.accelerator.synchronize() assert torch.all(storage == 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("layout", list(KVCacheLayout)) +def test_zeroes_exactly_one_block_per_layer(layout: KVCacheLayout): + """The zeroer must zero every byte of the target block in every layer and nothing + outside it — per head-group region under LHBNC, and never past the target block's + tile under block-major layouts (no out-of-bounds writes, no clobbering).""" + device = torch.device("cuda") + num_blocks, num_layers = 4, 2 + spec = FullAttentionSpec( + block_size=4, num_kv_heads=2, head_size=8, dtype=torch.float32 + ) + raw = torch.empty( + num_blocks * num_layers * spec.page_size_bytes, + dtype=torch.int8, + device=device, + ).fill_(1) + views = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout) + groups = [ + AttentionGroup( + backend=None, + layer_names=[f"layer.{i}" for i in range(num_layers)], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + ctx = {f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(num_layers)} + zeroer = KVBlockZeroer( + device, + attn_groups_iter=iter(groups), + kernel_block_sizes=[spec.block_size], + static_forward_context=ctx, + ) + zeroer.zero_block_ids([2]) + torch.accelerator.synchronize() + + for view in views: + assert (view[2] == 0).all(), layout + for b in (0, 1, 3): + assert (view[b].view(torch.int8) == 1).all(), layout + zero_bytes = int((raw == 0).sum().item()) + assert zero_bytes == num_layers * spec.page_size_bytes, layout diff --git a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py index 5a493149a9..5998a03138 100644 --- a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py +++ b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py @@ -45,7 +45,7 @@ def fused_rope_unified_mla_kv_cache_update_impl( cos_sin_cache, is_neox, layer_slot_mapping, - kv_cache, + kv_cache.squeeze(1), kv_cache_dtype, kv_cache_scale, ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 27e46534dc..dc4e6faf3c 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -3,6 +3,7 @@ from collections.abc import Callable from dataclasses import field +from functools import cache from typing import Any, ClassVar, Literal from pydantic import Field, field_validator, model_validator @@ -13,9 +14,28 @@ from vllm.utils.torch_utils import ( is_quantized_kv_cache, kv_cache_uses_per_token_head_scales, ) +from vllm.v1.kv_cache_layout import KVCacheLayout logger = init_logger(__name__) +_LAYOUT_COMPAT_ALIASES = { + "NHD": "LBNHC", + "HND": "LBHNC", +} + + +@cache +def _layout_from_name(layout_name: str) -> KVCacheLayout: + layout_name = _LAYOUT_COMPAT_ALIASES.get(layout_name, layout_name) + try: + return KVCacheLayout[layout_name] + except KeyError: + raise ValueError( + f"Unknown KV cache layout {layout_name!r}. " + f"Valid layouts: {[m.name for m in KVCacheLayout]}" + ) from None + + CacheDType = Literal[ "auto", "float16", @@ -65,6 +85,17 @@ class CacheConfig: """Whether block_size was explicitly provided. Derived automatically.""" user_specified_mamba_block_size: bool = field(default=False, init=False) """Whether mamba_block_size was explicitly provided. Derived automatically.""" + kv_cache_layout: str | None = field(default=None, init=False) + """Resolved physical KV cache layout name (a ``KVCacheLayout`` member). + + ``None`` means the layout has not been resolved yet. The engine core + resolves it once (``resolve_kv_cache_layout``) before memory profiling, and + every worker process adopts the resolved name — via the + ``set_kv_cache_layout`` RPC, or ``KVCacheConfig.kv_cache_layout`` for + workers spawned after resolution — before the KV cache is allocated. Once + set the value is final and read with ``get_resolved_kv_cache_layout``, + which raises on ``None``. Tests and standalone tools may pre-set a value, + which resolution then honors as-is.""" prefix_match_unit: int | None = Field(default=None, gt=0) """The finest token boundary (in tokens) a prefix-cache hit can land on. @@ -309,3 +340,12 @@ class CacheConfig: str(cache_dtype), ) return cache_dtype + + def get_resolved_kv_cache_layout(self) -> KVCacheLayout: + if self.kv_cache_layout is None: + raise ValueError( + "KV cache layout has not been resolved yet; it is resolved once " + "by the engine core (resolve_kv_cache_layout) unless explicitly " + "set by the user." + ) + return _layout_from_name(self.kv_cache_layout) diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 7b068adf8b..2e4899ebdf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -12,7 +12,7 @@ import torch from vllm.config import ( VllmConfig, - get_current_vllm_config, + get_current_vllm_config_or_none, get_layers_from_vllm_config, set_current_vllm_config, ) @@ -35,20 +35,20 @@ EngineId = str BlockIds = tuple[list[int], ...] | list[list[int]] -def get_kv_connector_cache_layout(): - # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is +def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None): + # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout is # used for faster transfer. - vllm_config = get_current_vllm_config() + if vllm_config is None: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return None kv_config = vllm_config.kv_transfer_config if kv_config is not None: connector_cls = KVConnectorFactory.get_connector_class(kv_config) required_kvcache_layout = connector_cls.get_required_kvcache_layout(vllm_config) if required_kvcache_layout is not None: return required_kvcache_layout - logger.info_once( - "Connectors do not specify a kv cache layout, defaulting to NHD." - ) - return "NHD" + return None class KVOutputAggregator: @@ -288,11 +288,11 @@ def kv_postprocess_layout_on_receive(cache, indices): def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio): """ - Transforms the layout of received KV cache to the local block_size and HND. + Transforms the layout of received KV cache to the local block_size and LBHNC. (Only works for local blocksize > remote blocksize) - prefill is HND, smaller block_size - decode(local) is NHD, larger block_size + prefill is LBHNC, smaller block_size + decode(local) is LBNHC, larger block_size """ blocks_to_update = cache.index_select(0, indices) @@ -427,48 +427,6 @@ class TransferTopology: self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {} - # Figure out whether the first dimension of the cache is K/V - # or num_blocks. - attn_backend = self.attn_backends[0] - if not self.is_mamba: - _MOCK_BLOCK_SIZE = 16 - kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape( - num_blocks=1, - block_size=_MOCK_BLOCK_SIZE, - num_kv_heads=1, - head_size=1, - ) - logger.debug("Test kv_cache_shape: %s", kv_cache_shape) - assert kv_cache_shape[0] == 1, ( - "KV cache layout must be blocks-first; expected mocked " - f"num_blocks=1 in leading dim, got shape {kv_cache_shape}." - ) - if not self.is_mla: - assert len(kv_cache_shape) == 4, ( - "Attention KV cache layout must be standardized as " - "[num_blocks, num_kv_heads, block_size, content_size], " - f"got shape {kv_cache_shape}." - ) - - self._cross_layers_blocks = False - if self.tensor_shape is not None: - self._cross_layers_blocks = ( - len(self.tensor_shape) == len(kv_cache_shape) + 1 - ) - - if self._cross_layers_blocks: - logger.debug("Using cross-layer KV cache") - _MOCK_NUM_LAYERS = 80 - kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order( - include_num_layers_dimension=self._cross_layers_blocks - ) - except (AttributeError, NotImplementedError): - assert self.tensor_shape is not None - kv_cache_stride_order = tuple(range(len(self.tensor_shape))) - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - # ============================================================ # Engine registration # ============================================================ @@ -503,24 +461,6 @@ class TransferTopology: for key in [k for k in self._engines if k[0] == remote_engine_id]: del self._engines[key] - # ============================================================ - # Layout properties - # ============================================================ - - @property - def cross_layers_blocks(self) -> bool: - return self._cross_layers_blocks - - @property - def virtually_split_kv_in_blocks(self) -> bool: - # Whether to logically split each block into two separately-indexable - # sub-regions. With K and V packed into the content dim, an attention - # block transfers as a single unit — no K/V sub-split is needed. Only - # Mamba still needs this, to index its two state regions (conv/ssm) - # separately. Not applicable to cross-layer blocks (per-layer - # interleaving means a simple half-split does not separate the parts). - return self.is_mamba and not self._cross_layers_blocks - # ============================================================ # Common methods # ============================================================ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index e49dbebb12..b0d7cf2388 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -48,7 +48,7 @@ from typing import TYPE_CHECKING, Any, Literal import torch from vllm.logger import init_logger -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import KVConnectorOutput @@ -182,14 +182,6 @@ class KVConnectorBase_V1(ABC): """ return False - @property - def prefer_cross_layer_blocks(self) -> bool: - """ - Indicates whether this connector prefers KV blocks that hold KV data for all - layers, which can speed up KV data transfers. Defaults to False. - """ - return False - @property def requires_kv_delivery(self) -> bool: """Whether this connector hands off KV that must be reliably delivered. @@ -279,23 +271,6 @@ class KVConnectorBase_V1(ABC): """ return - def register_cross_layers_kv_cache( - self, kv_cache: torch.Tensor, attn_backend: type["AttentionBackend"] - ): - """ - Initialize with a single KV cache tensor used by all layers. - The first dimension should be num_layers. - This function will only be called for models with uniform layers, - and only if the prefers_cross_layer_blocks is set to True. - Only one of the functions - {register_kv_caches, register_cross_layers_kv_cache} will be called. - - Args: - kv_cache: a cross-layers kv cache tensor - attn_backend: The attention backend that corresponds to all layers - """ - return - def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): """ Set the xPU-specific ops for copying KV between host and device. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py index 782972b03a..3517d55ecd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py @@ -129,20 +129,17 @@ class ExampleConnector(KVConnectorBase_V1): """Inject the KV cache into the layer. Args: - dst_kv_cache_layer (torch.Tensor): the destination KV cache - layer. In shape [num_pages, page_size, xxx] for MLA, - [num_pages, 2, page_size, xxx] otherwise. + dst_kv_cache_layer (torch.Tensor): the destination KV cache layer, + a standardized [B, H, N, C] per-layer view (H == 1 for MLA). src_kv_cache (torch.Tensor): the source KV cache. slot_mapping (torch.Tensor): the slot mapping. In shape [num_tokens]. """ slot_mapping = slot_mapping.to(dst_kv_cache_layer.device, non_blocking=True) if isinstance(attn_metadata, MLACommonMetadata): - dst_kv_cache_layer_shape = dst_kv_cache_layer.shape - num_pages = dst_kv_cache_layer_shape[0] - page_size = dst_kv_cache_layer_shape[1] + # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots. dst_kv_cache_layer = dst_kv_cache_layer.reshape( - num_pages * page_size, -1 + -1, dst_kv_cache_layer.shape[-1] ) dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache else: @@ -225,13 +222,12 @@ class ExampleConnector(KVConnectorBase_V1): ) -> torch.Tensor: """Extract the KV cache from the layer. - Assume the shape of the layer is (num_pages, page_size, xxx) - for MLA, and (num_pages, 2, page_size, xxx) otherwise. + The layer is a standardized [B, H, N, C] per-layer view (H == 1 for MLA). """ slot_mapping = slot_mapping.to(layer.device, non_blocking=True) if isinstance(attn_metadata, MLACommonMetadata): - num_pages, page_size = layer.shape[0], layer.shape[1] - return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...] + # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots. + return layer.reshape(-1, layer.shape[-1])[slot_mapping, ...] block_idxs = slot_mapping // self._block_size offsets = slot_mapping % self._block_size return layer[block_idxs, :, offsets] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index 299ff037ad..c94bd2d195 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -38,8 +38,10 @@ def extract_from_kv_cache( num_tokens: int, ) -> torch.Tensor: """Extract data from KV cache.""" - block_size = kv_cache.shape[1] - return kv_cache[slot_mapping // block_size, slot_mapping % block_size][:num_tokens] + block_size = kv_cache.shape[2] + return kv_cache[slot_mapping // block_size, :, slot_mapping % block_size][ + :num_tokens + ] def load_hidden_states(path: str) -> dict[str, torch.Tensor]: @@ -100,15 +102,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): Must be used in conjunction with the `extract_hidden_states` spec decoding method. """ - @property - def prefer_cross_layer_blocks(self) -> bool: - """ - Indicates whether this connector prefers KV blocks that hold KV data for all - layers, which can speed up KV data transfers. Defaults to False. - """ - # Must be False so that drafter kv cache isn't merged with verifier's - return False - @classmethod def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int: """Index of the KV cache group holding the extracted hidden states. @@ -310,10 +303,12 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): # Block size must match the indexed buffer, else reads hit the wrong # slots. Raise (not assert) so the check survives `python -O`. - if self._block_size != self._kv_cache.shape[1]: + # Views are [num_blocks, num_heads, block_size, head_size], matching what + # extract_from_kv_cache() indexes. + if self._block_size != self._kv_cache.shape[2]: raise ValueError( f"Hidden-states block-size mismatch: derived {self._block_size} " - f"but buffer block size is {self._kv_cache.shape[1]}; read slots " + f"but buffer block size is {self._kv_cache.shape[2]}; read slots " "would be wrong (likely a hybrid block-size resolution bug)." ) @@ -605,9 +600,9 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): "get_required_kvcache_layout should not be called " "on the abstract base class" ) - # NHD means we have (num_tokens, num_heads) - # HND means we have (num_heads, num_tokens) - # For now, we only support NHD layout since this keeps the + # LBNHC means we have (num_tokens, num_heads) + # LBHNC means we have (num_heads, num_tokens) + # For now, we only support LBNHC layout since this keeps the # hidden states for each token together in memory. - # HND is primarily used when sharding heads across devices. - return "NHD" + # LBHNC is primarily used when sharding heads across devices. + return "LBNHC" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py index 526375952f..7294d77cef 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py @@ -541,35 +541,32 @@ class HF3FSKVConnector(KVConnectorBase_V1): self._dtype = first_cache.dtype element_size = first_cache.element_size() - if self._use_mla: - assert len(first_cache.shape) == 3, "MLA format should have 3 dimensions" - # MLA format: [num_blocks, block_size, head_size] - num_blocks, block_size, head_size = first_cache.shape - num_heads = 1 - else: - # MHA format: [2, num_blocks, block_size, num_heads, head_size] - _, num_blocks, block_size, num_heads, head_size = first_cache.shape + # Standardized per-layer [B, H, N, C] view; MLA is just H == 1 with the + # latent vector as the content dim. + assert len(first_cache.shape) == 4, ( + f"expected a [B, H, N, C] KV cache view, got {tuple(first_cache.shape)}" + ) + num_blocks, num_heads, block_size, content_dim = first_cache.shape + cache_strides = first_cache.stride() + assert all( + cache.shape == first_cache.shape and cache.stride() == cache_strides + for cache in self._kv_caches.values() + ), "HF3FS requires uniform KV cache shapes and strides." self._local_total_tokens = num_blocks * block_size self._local_block_size = block_size + self._num_heads = num_heads + self._content_dim = content_dim + self._kv_cache_strides = cache_strides - if self._use_mla: - layer_block_size = block_size * head_size * element_size - self._bytes_per_page = layer_block_size * len(self._kv_caches) - self._shape_per_page = [ - len(self._kv_caches), - block_size, - head_size, - ] - else: - layer_block_size = 2 * block_size * num_heads * head_size * element_size - self._bytes_per_page = layer_block_size * len(self._kv_caches) - self._shape_per_page = [ - len(self._kv_caches), - 2, - block_size, - num_heads * head_size, - ] + layer_block_size = num_heads * block_size * content_dim * element_size + self._bytes_per_page = layer_block_size * len(self._kv_caches) + self._shape_per_page = [ + len(self._kv_caches), + num_heads, + block_size, + content_dim, + ] self._kvcache_ptrs = torch.tensor( [cache.data_ptr() for cache in self._kv_caches.values()], @@ -993,7 +990,10 @@ class HF3FSKVConnector(KVConnectorBase_V1): self._local_total_tokens, buffer_tensor, token_indices, - is_mla=self._use_mla, + self._local_block_size, + self._num_heads, + self._content_dim, + self._kv_cache_strides, ) else: gather_scatter_helper.scatter_kv_caches( @@ -1001,7 +1001,10 @@ class HF3FSKVConnector(KVConnectorBase_V1): self._local_total_tokens, buffer_tensor, token_indices, - is_mla=self._use_mla, + self._local_block_size, + self._num_heads, + self._content_dim, + self._kv_cache_strides, ) def _compute_prefix_hash( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py index 39d852dae6..aebaecd40b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py @@ -12,10 +12,15 @@ def kv_cache_scatter_kernel( source_ptr, token_indices_ptr, num_tokens_in_block, - hidden_size, total_token_in_kvcache, num_layers, - is_mla, + tokens_per_block, + block_stride, + head_stride, + token_stride, + content_stride, + num_heads: tl.constexpr, + content_size: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): layer_idx = tl.program_id(0) @@ -30,39 +35,25 @@ def kv_cache_scatter_kernel( if token_idx >= total_token_in_kvcache: return - if is_mla: - # MLA format: source [num_layers, num_tokens_in_block, hidden_size] - # MLA format: target [total_token_in_kvcache, hidden_size] (per layer) - source_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size - target_offset = token_idx * hidden_size - - for i in range(0, hidden_size, BLOCK_SIZE): - offset = i + tl.arange(0, BLOCK_SIZE) - mask = offset < hidden_size - val = tl.load(source_ptr + source_offset + offset, mask=mask) - tl.store(kv_cache_ptr + target_offset + offset, val, mask=mask) - else: - # MHA format: source [num_layers, 2, num_tokens_in_block, hidden_size] - # MHA format: target [2, total_token_in_kvcache, hidden_size] - source_offset_k = ( - layer_idx * num_tokens_in_block * 2 + token_pos - ) * hidden_size - source_offset_v = ( - layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos - ) * hidden_size - - target_offset_k = token_idx * hidden_size - target_offset_v = (total_token_in_kvcache + token_idx) * hidden_size - - for i in range(0, hidden_size, BLOCK_SIZE): - offset = i + tl.arange(0, BLOCK_SIZE) - mask = offset < hidden_size - - val_k = tl.load(source_ptr + source_offset_k + offset, mask=mask) - val_v = tl.load(source_ptr + source_offset_v + offset, mask=mask) - - tl.store(kv_cache_ptr + target_offset_k + offset, val_k, mask=mask) - tl.store(kv_cache_ptr + target_offset_v + offset, val_v, mask=mask) + block_idx = token_idx // tokens_per_block + state_idx = token_idx % tokens_per_block + hidden_size = num_heads * content_size + for i in range(0, hidden_size, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + head_idx = offset // content_size + content_idx = offset % content_size + source_offset = ( + (layer_idx * num_heads + head_idx) * num_tokens_in_block + token_pos + ) * content_size + content_idx + target_offset = ( + block_idx * block_stride + + head_idx * head_stride + + state_idx * token_stride + + content_idx * content_stride + ) + value = tl.load(source_ptr + source_offset, mask=mask) + tl.store(kv_cache_ptr + target_offset, value, mask=mask) @triton.jit @@ -71,10 +62,15 @@ def kv_cache_gather_kernel( dst_ptr, token_indices_ptr, num_tokens_in_block, - hidden_size, total_token_in_kvcache, num_layers, - is_mla, + tokens_per_block, + block_stride, + head_stride, + token_stride, + content_stride, + num_heads: tl.constexpr, + content_size: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): layer_idx = tl.program_id(0) @@ -89,37 +85,25 @@ def kv_cache_gather_kernel( if token_idx >= total_token_in_kvcache: return - if is_mla: - # MLA format: source [total_token_in_kvcache, hidden_size] (per layer) - # MLA format: dst [num_layers, num_tokens_in_block, hidden_size] - kvcache_offset = token_idx * hidden_size - dst_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size - - for i in range(0, hidden_size, BLOCK_SIZE): - offset = i + tl.arange(0, BLOCK_SIZE) - mask = offset < hidden_size - val = tl.load(kv_cache_ptr + kvcache_offset + offset, mask=mask) - tl.store(dst_ptr + dst_offset + offset, val, mask=mask) - else: - # MHA format: source [2, total_token_in_kvcache, hidden_size] - # MHA format: dst [num_layers, 2, num_tokens_in_block, hidden_size] - dst_offset_k = (layer_idx * num_tokens_in_block * 2 + token_pos) * hidden_size - dst_offset_v = ( - layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos - ) * hidden_size - - kvcache_offset_k = token_idx * hidden_size - kvcache_offset_v = (total_token_in_kvcache + token_idx) * hidden_size - - for i in range(0, hidden_size, BLOCK_SIZE): - offset = i + tl.arange(0, BLOCK_SIZE) - mask = offset < hidden_size - - val_k = tl.load(kv_cache_ptr + kvcache_offset_k + offset, mask=mask) - val_v = tl.load(kv_cache_ptr + kvcache_offset_v + offset, mask=mask) - - tl.store(dst_ptr + dst_offset_k + offset, val_k, mask=mask) - tl.store(dst_ptr + dst_offset_v + offset, val_v, mask=mask) + block_idx = token_idx // tokens_per_block + state_idx = token_idx % tokens_per_block + hidden_size = num_heads * content_size + for i in range(0, hidden_size, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + head_idx = offset // content_size + content_idx = offset % content_size + kvcache_offset = ( + block_idx * block_stride + + head_idx * head_stride + + state_idx * token_stride + + content_idx * content_stride + ) + dst_offset = ( + (layer_idx * num_heads + head_idx) * num_tokens_in_block + token_pos + ) * content_size + content_idx + value = tl.load(kv_cache_ptr + kvcache_offset, mask=mask) + tl.store(dst_ptr + dst_offset, value, mask=mask) def scatter_kv_caches( @@ -127,34 +111,33 @@ def scatter_kv_caches( total_token_in_kvcache: int, src_tensor: torch.Tensor, token_indices: list[int], - is_mla: bool = False, + tokens_per_block: int, + num_heads: int, + content_size: int, + kv_cache_strides: tuple[int, ...], ) -> None: """Scatter KV cache data from source tensor to KV cache storage. Args: kv_caches_ptrs: Tensor of KV cache pointers (one per layer) total_token_in_kvcache: Total number of tokens in KV cache - src_tensor: Source tensor containing data to scatter - - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size] - - MLA format: [num_layers, num_tokens_in_block, hidden_size] + src_tensor: Source ``[L, H, N, C]`` tensor containing data to scatter token_indices: List of token positions to update - is_mla: Whether using MLA model format + tokens_per_block: Number of stored states in each cache block + num_heads: Size of the H axis + content_size: Size of the C axis + kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view """ num_layers = len(kv_caches_ptrs) num_tokens_in_block = len(token_indices) - if is_mla: - # MLA: src_tensor is [num_layers, num_tokens_in_block, hidden_size] - assert len(src_tensor.shape) == 3, ( - f"MLA src_tensor should be 3D, got {src_tensor.shape}" - ) - hidden_size = src_tensor.shape[2] - else: - # MHA: src_tensor is [num_layers, 2, num_tokens_in_block, hidden_size] - assert len(src_tensor.shape) == 4, ( - f"MHA src_tensor should be 4D, got {src_tensor.shape}" - ) - hidden_size = src_tensor.shape[3] + assert src_tensor.shape == ( + num_layers, + num_heads, + num_tokens_in_block, + content_size, + ) + assert len(kv_cache_strides) == 4 device = src_tensor.device token_indices_tensor = torch.tensor( @@ -169,10 +152,12 @@ def scatter_kv_caches( src_tensor, token_indices_tensor, num_tokens_in_block, - hidden_size, total_token_in_kvcache, num_layers, - is_mla, + tokens_per_block, + *kv_cache_strides, + num_heads=num_heads, + content_size=content_size, BLOCK_SIZE=BLOCK_SIZE, ) @@ -182,49 +167,33 @@ def gather_kv_caches( total_token_in_kvcache: int, dst_tensor: torch.Tensor, token_indices: list[int], - is_mla: bool = False, + tokens_per_block: int, + num_heads: int, + content_size: int, + kv_cache_strides: tuple[int, ...], ) -> None: """Gather KV cache data from KV cache storage to destination tensor. Args: kv_caches_ptrs: Tensor of KV cache pointers (one per layer) total_token_in_kvcache: Total number of tokens in KV cache - dst_tensor: Destination tensor to store gathered data - - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size] - - MLA format: [num_layers, num_tokens_in_block, hidden_size] + dst_tensor: Destination ``[L, H, N, C]`` tensor token_indices: List of token positions to gather - is_mla: Whether using MLA model format + tokens_per_block: Number of stored states in each cache block + num_heads: Size of the H axis + content_size: Size of the C axis + kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view """ num_layers = kv_caches_ptrs.shape[0] num_tokens_in_block = len(token_indices) - if is_mla: - # MLA: dst_tensor is [num_layers, num_tokens_in_block, hidden_size] - assert len(dst_tensor.shape) == 3, ( - f"MLA dst_tensor should be 3D, got {dst_tensor.shape}" - ) - assert dst_tensor.shape[0] == num_layers, ( - f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}" - ) - assert dst_tensor.shape[1] == num_tokens_in_block, ( - f"Token count mismatch: {dst_tensor.shape[1]} vs {num_tokens_in_block}" - ) - hidden_size = dst_tensor.shape[2] - else: - # MHA: dst_tensor is [num_layers, 2, num_tokens_in_block, hidden_size] - assert len(dst_tensor.shape) == 4, ( - f"MHA dst_tensor should be 4D, got {dst_tensor.shape}" - ) - assert dst_tensor.shape[0] == num_layers, ( - f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}" - ) - assert dst_tensor.shape[1] == 2, ( - f"MHA should have 2 (K,V) components, got {dst_tensor.shape[1]}" - ) - assert dst_tensor.shape[2] == num_tokens_in_block, ( - f"Token count mismatch: {dst_tensor.shape[2]} vs {num_tokens_in_block}" - ) - hidden_size = dst_tensor.shape[3] + assert dst_tensor.shape == ( + num_layers, + num_heads, + num_tokens_in_block, + content_size, + ) + assert len(kv_cache_strides) == 4 device = dst_tensor.device token_indices_tensor = torch.tensor( @@ -239,10 +208,12 @@ def gather_kv_caches( dst_tensor, token_indices_tensor, num_tokens_in_block, - hidden_size, total_token_in_kvcache, num_layers, - is_mla, + tokens_per_block, + *kv_cache_strides, + num_heads=num_heads, + content_size=content_size, BLOCK_SIZE=BLOCK_SIZE, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index dd94ffb2f9..0d4afbd56d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -49,15 +49,15 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket +from vllm.utils.torch_utils import is_non_overlapping_and_dense from vllm.v1.attention.backend import AttentionMetadata -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, get_kv_cache_layout +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( + AttentionSpec, FullAttentionSpec, KVCacheSpec, MambaSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, SlidingWindowSpec, ) from vllm.v1.request import RequestStatus @@ -122,9 +122,7 @@ def _expand_transfer_regions( kv_block_lens: list[int], layer_names: list[str], layer_indices: list[int], - is_kv_layout_blocks_first: bool, group_indices: list[int] | None = None, - split_kv_regions: list[bool] | None = None, ) -> list[TransferRegion]: """Expand registered KV tensors into the regions transferred by Mooncake.""" assert ( @@ -146,13 +144,6 @@ def _expand_transfer_regions( "Mooncake transfer regions require matching group metadata lengths, " f"got group_indices={len(group_indices)}, layer_names={len(layer_names)}." ) - if split_kv_regions is None: - split_kv_regions = [is_kv_layout_blocks_first] * len(layer_names) - assert len(split_kv_regions) == len(layer_names), ( - "Mooncake transfer regions require matching split metadata, " - f"got split_kv_regions={len(split_kv_regions)}, " - f"layer_names={len(layer_names)}." - ) regions: list[TransferRegion] = [] for ( base_addr, @@ -161,7 +152,6 @@ def _expand_transfer_regions( layer_name, layer_index, group_index, - split_kv_region, ) in zip( base_addrs, block_lens, @@ -169,7 +159,6 @@ def _expand_transfer_regions( layer_names, layer_indices, group_indices, - split_kv_regions, ): regions.append( TransferRegion( @@ -181,17 +170,6 @@ def _expand_transfer_regions( group_index=group_index, ) ) - if split_kv_region: - regions.append( - TransferRegion( - layer_name=layer_name, - layer_index=layer_index, - base_addr=base_addr + kv_block_len, - block_len=block_len, - kv_block_len=kv_block_len, - group_index=group_index, - ) - ) return regions @@ -506,10 +484,10 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA): if vllm_config.model_config.use_mla: return None logger.info_once( - "MooncakeConnector setting KV cache layout to HND for " + "MooncakeConnector setting KV cache layout to LBHNC for " "heterogeneous TP-safe KV transfer." ) - return "HND" + return "LBHNC" ############################################################ # Scheduler Side Methods @@ -1023,12 +1001,10 @@ class MooncakeConnectorWorker: self._sync_block_size_with_kernel() self.attn_backends = get_current_attn_backends(vllm_config) - self.kv_cache_layout = get_kv_cache_layout() logger.debug( "Detected attention backends %s", [backend.get_name() for backend in self.attn_backends], ) - logger.debug("Detected kv cache layout %s", self.kv_cache_layout) self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size} self._layer_specs: dict[str, KVCacheSpec] = {} @@ -1665,7 +1641,7 @@ class MooncakeConnectorWorker: self.registered_layer_indices = [] self.registered_group_indices = [] - for layer_name, cache_or_caches in kv_caches.items(): + for layer_name, cache in kv_caches.items(): layer_index = extract_layer_index(layer_name) layer_spec = self._layer_specs.get(layer_name) if layer_spec is None: @@ -1674,34 +1650,30 @@ class MooncakeConnectorWorker: layer_name, ) continue - if isinstance(layer_spec, MambaSpec): - conv, _ = cache_or_caches - cache_list = [conv] + # One raw page tensor per layer; for Mamba that page holds all the + # recurrent states, unpacked only when binding the cache for execution. + self._log_debug_cache_registration(layer_name, cache) + block_is_contiguous = is_non_overlapping_and_dense(cache[0]) + if not block_is_contiguous: + # Non-block-compact layouts scatter a block across per-head + # regions; each region's blocks are contiguous. + region_caches = [cache[:, head] for head in range(cache.shape[1])] + assert all( + is_non_overlapping_and_dense(region[0]) for region in region_caches + ) else: - # K and V are packed into one blocks-first tensor per layer, - # so each layer registers as a single region. - cache_list = [cache_or_caches] + region_caches = [cache] - logger.debug( - "registering layer %s with %d cache tensor(s)", - layer_name, - len(cache_list), - ) - - for cache in cache_list: - self._log_debug_cache_registration(layer_name, cache) - base_addr = cache.data_ptr() - block_len = cache.stride(0) * cache.element_size() + for region_cache in region_caches: + base_addr = region_cache.data_ptr() + block_len = region_cache.stride(0) * region_cache.element_size() region_base_addresses.append(base_addr) - if isinstance(layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)): - kv_block_len = layer_spec.page_size_bytes - elif self.transfer_topo.virtually_split_kv_in_blocks and not isinstance( - layer_spec, MambaSpec - ): - kv_block_len = block_len // 2 - else: - kv_block_len = block_len + kv_block_len = ( + layer_spec.page_size_bytes + if isinstance(layer_spec, AttentionSpec) and block_is_contiguous + else block_len + ) self.block_len_per_layer.append(block_len) self.kv_block_len_per_layer.append(kv_block_len) self.registered_layer_names.append(layer_name) @@ -1709,12 +1681,12 @@ class MooncakeConnectorWorker: self.registered_group_indices.append( self._layer_group_indices[layer_name] ) - storage = cache.untyped_storage() - storage_addr = storage.data_ptr() - if storage_addr not in seen_storage_ptrs: - seen_storage_ptrs.add(storage_addr) - kv_data_ptrs.append(storage_addr) - kv_data_lens.append(storage.nbytes()) + storage = cache.untyped_storage() + storage_addr = storage.data_ptr() + if storage_addr not in seen_storage_ptrs: + seen_storage_ptrs.add(storage_addr) + kv_data_ptrs.append(storage_addr) + kv_data_lens.append(storage.nbytes()) self.kv_caches_base_addr = region_base_addresses self.seen_base_addresses = kv_data_ptrs @@ -2055,24 +2027,13 @@ class MooncakeConnectorWorker: self._layer_group_indices.get(layer_name, 0) for layer_name in layer_names ] - split_kv_regions = None - if self.transfer_topo.virtually_split_kv_in_blocks: - split_kv_regions = [ - not isinstance( - self._layer_specs[layer_name], - (MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec), - ) - for layer_name in layer_names - ] return _expand_transfer_regions( base_addrs=base_addrs, block_lens=block_lens, kv_block_lens=kv_block_lens, layer_names=layer_names, layer_indices=layer_indices, - is_kv_layout_blocks_first=self.transfer_topo.virtually_split_kv_in_blocks, group_indices=group_indices, - split_kv_regions=split_kv_regions, ) def _get_sender_transfer_plan( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index e35eb6adb3..77adc2d1bc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -89,14 +89,6 @@ class MooncakeStoreKVEvents(KVConnectorKVEvents): class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): """KV connector using MooncakeDistributedStore as shared KV pool.""" - @property - def prefer_cross_layer_blocks(self) -> bool: - extra_config = self._kv_transfer_config.kv_connector_extra_config - return ( - str(extra_config.get("enable_cross_layers_blocks", "False")).lower() - == "true" - ) - @staticmethod def _validate_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig @@ -290,16 +282,6 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): assert self.connector_worker is not None self.connector_worker.register_kv_caches(kv_caches) - def register_cross_layers_kv_cache( - self, kv_cache: torch.Tensor, attn_backend: type - ): - assert self.connector_worker is not None - assert ( - self._kv_cache_config is not None - and len(self._kv_cache_config.kv_cache_groups) == 1 - ), "Cross-layer KV cache does not supported with hybrid models" - self.connector_worker.register_cross_layers_kv_caches(kv_cache) - def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None: # No-op: loads are issued in get_finished() for compute overlap. pass diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index c9fba1c699..4b73614ae5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -58,6 +58,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket +from vllm.utils.torch_utils import is_non_overlapping_and_dense from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import ( BlockHash, @@ -72,6 +73,7 @@ from vllm.v1.kv_cache_interface import ( MLAAttentionSpec, SlidingWindowMLASpec, UniformTypeKVCacheSpecs, + group_kernel_blocks, ) from .metrics import MooncakeStoreConnectorStats @@ -1617,18 +1619,9 @@ class MooncakeStoreWorker: for g_idx, db in enumerate(self.token_dbs) ) - def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: - """Register a cross-layers KV cache tensor. - - Wraps the unified tensor in a single-entry dict so that the - existing stride-based logic in register_kv_caches() produces - the correct single-segment result (block_len = page_size * num_layers). - """ - self.register_kv_caches({"__cross_layer__": kv_cache}) - def register_kv_caches( self, - kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], + kv_caches: dict[str, torch.Tensor], ) -> None: """Register KV cache tensors and start transfer threads.""" if self._capacity_only: @@ -1637,58 +1630,62 @@ class MooncakeStoreWorker: logger.warning("No KV caches to offload.") return - # Resolve each entry to a representative tensor for storage - # deduplication. For attention layers the value is already a tensor; - # for Mamba layers it is a list of tensors that all share the same - # underlying raw storage, so we take the first one. - def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: - assert isinstance(v, torch.Tensor | list) - return v if isinstance(v, torch.Tensor) else v[0] - assert self.cache_config.num_gpu_blocks is not None self.num_blocks = self.cache_config.num_gpu_blocks - seen_ptrs: set[int] = set() + seen_storage_ptrs: set[int] = set() + seen_region_ptrs: set[int] = set() addrs: list[int] = [] block_lens: list[int] = [] - for value in kv_caches.values(): - cache = _repr_tensor(value) + for cache in kv_caches.values(): + cache = group_kernel_blocks(cache, self.num_blocks) cache_storage = cache.untyped_storage() base_addr = cache_storage.data_ptr() - if base_addr in seen_ptrs: - continue - seen_ptrs.add(base_addr) region_len = cache_storage.nbytes() - ret = self.store.register_buffer(base_addr, region_len) - if ret != 0: - logger.error( - "register_buffer failed for addr %#x len %d: %d", - base_addr, - region_len, - ret, - ) + if base_addr not in seen_storage_ptrs: + seen_storage_ptrs.add(base_addr) + ret = self.store.register_buffer(base_addr, region_len) + if ret != 0: + logger.error( + "register_buffer failed for addr %#x len %d: %d", + base_addr, + region_len, + ret, + ) - # Detect layout via stride: a dim whose byte-stride exceeds - # page_size_bytes is an outer segment dim (e.g. the K/V dim of - # FlashAttn's (2, num_blocks, ...)). FlashInfer/MLA's blocks- - # outermost layout has no such dim and yields a single segment. - el = cache.element_size() - page_size_bytes = region_len // self.num_blocks - outer_dims = [ - d for d in range(cache.ndim) if cache.stride(d) * el > page_size_bytes - ] - if not outer_dims: - # Blocks-first layout (FlashInfer / MLA): one segment. + if not is_non_overlapping_and_dense(cache[0]): + # A block is scattered across per-head regions; each region's + # blocks are contiguous. + for head_idx in range(cache.shape[1]): + head_cache = cache[:, head_idx] + assert is_non_overlapping_and_dense(head_cache[0]) + region_addr = head_cache.data_ptr() + if region_addr in seen_region_ptrs: + continue + seen_region_ptrs.add(region_addr) + addrs.append(region_addr) + block_lens.append(head_cache.stride(0) * head_cache.element_size()) + elif cache.stride(0) * cache.element_size() * self.num_blocks == region_len: + # The block stride spans the whole per-block window + # (block-outermost and packed layouts, and single-layer + # tensors), which may hold other layers' pages at higher + # offsets. Register the storage once as one whole-window + # region; per-layer regions would copy overlapping windows and + # run past the storage for offset layers. + if base_addr in seen_region_ptrs: + continue + seen_region_ptrs.add(base_addr) addrs.append(base_addr) - block_lens.append(page_size_bytes) + block_lens.append(region_len // self.num_blocks) else: - # K/V-first layout (FlashAttn / ROCm): split segments. - seg_stride = cache.stride(outer_dims[0]) * el - for idx in range(cache.shape[outer_dims[0]]): - addrs.append(base_addr + idx * seg_stride) - block_lens.append(seg_stride // self.num_blocks) + region_addr = cache.data_ptr() + if region_addr in seen_region_ptrs: + continue + seen_region_ptrs.add(region_addr) + addrs.append(region_addr) + block_lens.append(cache.stride(0) * cache.element_size()) logger.info( "Registered KV caches: num_groups=%d, num_segments=%d, num_blocks=%d", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index 90aa5f823a..0ee2721139 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -57,11 +57,6 @@ def is_mla_cache_layer( return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) -def _content_packed_dim(spec: AttentionSpec) -> int: - head_size_v = getattr(spec, "head_size_v", spec.head_size) - return spec.head_size + head_size_v - - def _spec_dim_matches(value: int, expected: int | None) -> bool: return expected is None or value == expected @@ -110,23 +105,6 @@ def get_layer_transfer_geometry( spec = layer_to_spec[layer_name] is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name) - if is_mla_cache and len(shape) == 3: - num_blocks, block_size, latent_dim = shape - slot_size_bytes = latent_dim * element_size - block_len = block_size * slot_size_bytes - return LayerTransferGeometry( - num_blocks=num_blocks, - block_size=block_size, - block_len=block_len, - slot_size_bytes=slot_size_bytes, - block_stride=stride[0], - local_kv_stride=None, - remote_kv_stride=None, - transfers_per_block=1, - regions_per_block=1, - split_kv_regions=False, - ) - if not is_mla_cache and len(shape) == 5 and shape[0] == 2: _, num_blocks = shape[:2] kernel_blocks_per_block = 1 @@ -230,19 +208,19 @@ def get_layer_transfer_geometry( ) if ( - not is_mla_cache - and isinstance(spec, AttentionSpec) + isinstance(spec, AttentionSpec) and len(shape) == 4 - and shape[1] == spec.num_kv_heads - and shape[2] == spec.block_size - and shape[3] == _content_packed_dim(spec) + and shape[1] == spec.num_heads + and shape[2] == spec.num_states + and shape[3] * element_size == spec.state_content_size_bytes ): - num_blocks, num_kv_heads, block_size, packed_dim = shape - slot_size_bytes = num_kv_heads * packed_dim * element_size - block_len = block_size * slot_size_bytes + # Standardized per-layer [B, H, N, C] view (MLA is just H == 1). + num_blocks, num_heads, num_states, content_dim = shape + slot_size_bytes = num_heads * content_dim * element_size + block_len = num_states * slot_size_bytes return LayerTransferGeometry( num_blocks=num_blocks, - block_size=block_size, + block_size=spec.block_size, block_len=block_len, slot_size_bytes=slot_size_bytes, block_stride=stride[0], @@ -267,6 +245,16 @@ def iter_layer_registration_regions( ) -> list[tuple[torch.Tensor, int]]: geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec) region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len + if geometry.regions_per_block == 1: + # With padded or interleaved pages the block stride exceeds the + # meaningful block_len; register the strided span. The span ends with + # the last block's meaningful bytes, not a whole stride, so it never + # runs past the backing allocation of a strided layer view. + block_stride_bytes = geometry.block_stride * kv_cache.element_size() + region_len = max( + region_len, + (geometry.num_blocks - 1) * block_stride_bytes + geometry.block_len, + ) if geometry.split_kv_regions: return [(cache, region_len) for cache in kv_cache] return [(kv_cache, region_len)] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 743a23e791..2b567da792 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -27,7 +27,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( PromMetricT, ) from vllm.logger import init_logger -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import KVConnectorOutput @@ -209,12 +209,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): c.supports_divergent_local_hybrid_hits for c in self._connectors ) - @property - def prefer_cross_layer_blocks(self) -> bool: - if not self._connectors: - return False - return all(c.prefer_cross_layer_blocks for c in self._connectors) - @property def requires_kv_delivery(self) -> bool: return any(c.requires_kv_delivery for c in self._connectors) @@ -245,13 +239,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): ) return ret - def register_cross_layers_kv_cache( - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] - ): - # Register on all connectors - for c in self._connectors: - c.register_cross_layers_kv_cache(kv_cache, attn_backend) - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): for c in self._connectors: c.register_kv_caches(kv_caches) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 7334a87466..e7267090c8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -71,9 +71,9 @@ from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.network_utils import make_zmq_path from vllm.utils.torch_utils import async_tensor_h2d -from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheLayout, MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec, @@ -89,6 +89,13 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool: + """Return whether all views share storage and a block stride.""" + block_strides = {cache.stride(0) * cache.element_size() for cache in caches} + storage_ptrs = {cache.untyped_storage().data_ptr() for cache in caches} + return len(block_strides) == len(storage_ptrs) == 1 + + class NixlBaseConnectorWorker: """Base implementation of Worker side methods shared by pull and push.""" @@ -302,7 +309,6 @@ class NixlBaseConnectorWorker: for group in kv_cache_config.kv_cache_groups for layer in group.layer_names } - self.hma_group_size = len(kv_cache_config.kv_cache_tensors) # ---- Model state (derived from model config) ---- mamba_ssm_size = (0, 0) @@ -523,7 +529,9 @@ class NixlBaseConnectorWorker: self.attn_backends = get_current_attn_backends(vllm_config) self.backend_name = self.attn_backends[0].get_name() - self.kv_cache_layout = get_kv_cache_layout() + self.kv_cache_layout = ( + vllm_config.cache_config.get_resolved_kv_cache_layout().name + ) self.host_buffer_kv_cache_layout = self.kv_cache_layout logger.info( "Detected attention backend(s) %s", @@ -558,6 +566,11 @@ class NixlBaseConnectorWorker: # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. self.block_len_per_layer = list[int]() + # Per-region block stride in bytes. Taken from the registered tensor's + # stride(0) so it stays correct under layouts that interleave layers + # within a block (BLHNC/BHLNC), where stride > block_len. + self.block_stride_per_layer = list[int]() + # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} @@ -788,16 +801,16 @@ class NixlBaseConnectorWorker: if not self.use_mla: assert kv_cache.ndim == 4 - if self.kv_cache_layout == "NHD": + if self.kv_cache_layout == "LBNHC": if self.kv_transfer_config.enable_permute_local_kv: logger.info_once( "'enable_permute_local_kv' flag is enabled while " - "device KV Layout is NHD. Init host buffer with" - " HND to better support Decode/Prefill TP_ratio > 1." + "device KV Layout is LBNHC. Init host buffer with" + " LBHNC to better support Decode/Prefill TP_ratio > 1." ) - # Since NHD will not support Decode/Prefill TP_ratio > 1, + # Since LBNHC will not support Decode/Prefill TP_ratio > 1, # we can leverage host_buffer for permute. - self.host_buffer_kv_cache_layout = "HND" + self.host_buffer_kv_cache_layout = "LBHNC" else: # Packed KV layout is logical (B, H, N, 2*D). Allocate # (B, N, H, 2*D) and view it as logical (B, H, N, 2*D) @@ -969,120 +982,9 @@ class NixlBaseConnectorWorker: fut.add_done_callback(request_ready) - def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: - """Register a cross-layers KV cache tensor with NIXL. - - `use_uniform_kv_cache()` guarantees a single KV cache group whose - layers all share the same `AttentionSpec`, so any layer name from - `_layer_specs` yields the correct per-layer spec for `page_size_bytes`. - """ - first_layer = next(iter(self._layer_specs)) - # Forwarding a real layer name rather than a synthetic key - self.register_kv_caches({first_layer: kv_cache}) - - def _register_packed_kv_cache( - self, - storage: torch.UntypedStorage, - ) -> None: - """Register a packed KV cache as a single NIXL region. - - The packed allocation interleaves all layers per block, so each - block_stride-byte chunk is one logical block. We register 1 - NIXL region and create 1 descriptor per block. - """ - self.transfer_topo = TransferTopology( - tp_rank=self.tp_rank, - tp_size=self.world_size, - block_size=self.block_size, - engine_id=self.engine_id, - is_mla=self.use_mla, - total_num_kv_heads=self.model_config.get_total_num_kv_heads(), - attn_backends=self.attn_backends, - tensor_shape=None, - is_mamba=self._has_mamba, - ) - self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, - self.backend_name, - self.transfer_topo.cross_layers_blocks, - transfer_mode=self._TRANSFER_MODE, - ) - - total_size = storage.nbytes() - block_stride = total_size // self.num_blocks - base_addr = storage.data_ptr() - device_id = storage.device.index - assert device_id is not None - - logger.info( - "Registering packed KV cache: total_size=%s, block_stride=%s, " - "num_blocks=%s, num_regions=1", - total_size, - block_stride, - self.num_blocks, - ) - - self.device_id = device_id - caches_data = [(base_addr, total_size, self.device_id, "")] - - self.block_len_per_layer = [block_stride] - self.num_regions = 1 - self.num_descs = self.num_blocks - self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr] - - descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type) - self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) - self._registered_descs.append(descs) - - self.dst_num_blocks[self.engine_id] = self.num_blocks - - self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = ( - self.register_local_xfer_handler(self.block_size) - ) - - agent_metadata = NixlAgentMetadata( - engine_id=self.engine_id, - agent_metadata=self.nixl_wrapper.get_agent_metadata(), - device_id=self.device_id, - kv_caches_base_addr=( - self.kv_caches_base_addr[self.engine_id][self.tp_rank] - ), - num_blocks=self.num_blocks, - block_lens=self.block_len_per_layer, - kv_cache_layout=self.kv_cache_layout, - block_size=self.block_size, - ssm_sizes=self._mamba_ssm_size, - attn_backend_name=self.backend_name, - physical_blocks_per_logical_kv_block=( - self._physical_blocks_per_logical_kv_block - ), - dcp_size=self.dcp_size, - pcp_size=self.pcp_size, - ) - assert self.compat_hash is not None - encoder = msgspec.msgpack.Encoder() - self.xfer_handshake_metadata = NixlHandshakePayload( - compatibility_hash=self.compat_hash, - agent_metadata_bytes=encoder.encode(agent_metadata), - ) - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" - # Detect packed allocation: all tensors are strided views into the - # same backing storage (different data_ptr but same storage). - # This happens with DSv4-style contiguous per-block packing. - if len(kv_caches) > 1 and not self._has_mamba: - storage = next(iter(kv_caches.values())).untyped_storage() - storage_ptrs = { - cache.untyped_storage().data_ptr() for cache in kv_caches.values() - } - data_ptrs = {cache.data_ptr() for cache in kv_caches.values()} - if len(storage_ptrs) == 1 and len(data_ptrs) > 1: - self._register_packed_kv_cache(storage) - self.device_kv_caches = kv_caches - return - self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, tp_size=self.world_size, @@ -1100,7 +1002,6 @@ class NixlBaseConnectorWorker: self.compat_hash = compute_nixl_compatibility_hash( self.vllm_config, self.backend_name, - self.transfer_topo.cross_layers_blocks, transfer_mode=self._TRANSFER_MODE, ) @@ -1127,14 +1028,14 @@ class NixlBaseConnectorWorker: ) caches_data = [] - # With hybrid allocator, layers can share a kv cache tensor + seen_storage_addresses: set[int] = set() seen_base_addresses: list[int] = [] + packed_storage = _share_storage_and_block_stride(list(xfer_buffers.values())) + # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead # register separate conv/ssm sub-regions (see `_build_mamba_local`). - tensor_size_bytes = None - for layer_name, cache in xfer_buffers.items(): # NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when # kernel requires a specific block size. This leads to SSM and FA layers @@ -1159,55 +1060,98 @@ class NixlBaseConnectorWorker: else layer_spec.page_size_bytes // self._physical_blocks_per_logical_kv_block ) - if self.transfer_topo._cross_layers_blocks: - # When cross-layers blocks are used, multiply by number of layers - physical_page_size = physical_page_size * len( - self.kv_cache_config.kv_cache_tensors - ) num_blocks = ( self._logical_num_blocks if isinstance(layer_spec, MambaSpec) else self.num_blocks ) - # `page_size` accounts for physical blocks, st KVCache is always - # [`num_blocks` * `page_size`] - curr_tensor_size_bytes = num_blocks * physical_page_size - - base_addr = cache.data_ptr() - is_mla_region = isinstance( - layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) - ) - if base_addr in seen_base_addresses: - # NOTE (NickLucche) HMA employs memory pooling to share tensors - # across groups. This results in skipping all tensors but the ones - # pointed to by group0. Also, generally we will have more blocks - # per tensor but fewer regions. - # A shared tensor may back both SSM and attention layers (e.g. - # KDA+MLA in KimiLinear); the region's FA view is MLA whichever - # layer registered it first. - idx = seen_base_addresses.index(base_addr) - self._region_is_mla[idx] |= is_mla_region - logger.debug("Skipping %s because it's already seen", layer_name) - continue logger.debug( "Registering layer %s with cache shape: %s", layer_name, cache.shape ) - seen_base_addresses.append(base_addr) - # Only record non-Mamba page sizes. - if isinstance(layer_spec, MambaSpec): - self.block_len_per_layer.append( - physical_page_size // self._physical_blocks_per_logical_kv_block - ) - else: - self.block_len_per_layer.append(physical_page_size) - self._region_is_mla.append(is_mla_region) + storage = cache.untyped_storage() + storage_addr = storage.data_ptr() + # Memory registration follows allocations, while transfer regions follow + # logical layers (or contiguous head segments). This keeps strided + # cross-layer views inside their registered allocation. + if storage_addr not in seen_storage_addresses: + seen_storage_addresses.add(storage_addr) + self.device_id = max(cache.get_device(), 0) + caches_data.append((storage_addr, storage.nbytes(), self.device_id, "")) - if not is_mla_region: - if tensor_size_bytes is None: - tensor_size_bytes = curr_tensor_size_bytes - assert tensor_size_bytes == curr_tensor_size_bytes, ( - "All non-MLA kv cache tensors must have the same size" + is_mla_region = isinstance( + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) + ) + + if isinstance(layer_spec, MambaSpec): + region_specs = [ + ( + cache.data_ptr(), + physical_page_size + // self._physical_blocks_per_logical_kv_block, + physical_page_size + // self._physical_blocks_per_logical_kv_block, + ) + ] + else: + if cache.ndim == 1: + # Flat byte view: HMA tensors shared between layer types carry + # no block dimension, so stride(0) is 1 byte. Blocks abut, so + # the stride is the page size. + block_stride = physical_page_size + else: + block_stride = cache.stride(0) * cache.element_size() + storage_is_block_major = num_blocks * block_stride == storage.nbytes() + # A layer whose [H, N, C] interior is dense addresses its own page + # as one chunk, so it can own a region. When it does not, or when + # every layer is packed into this one storage, the block's whole + # row is the only contiguous unit. + hnc_contiguous = ( + cache.ndim == 4 + and cache.stride(2) == cache.shape[3] + and cache.stride(1) == cache.shape[2] * cache.shape[3] ) + if storage_is_block_major and ( + (packed_storage and is_mla_region) or not hnc_contiguous + ): + # TODO(Lucas): handle TP slicing for packed_storage; for now + # restrict to MLA (DSv4) where kv is replicated. + storage_block_len = storage.nbytes() // num_blocks + region_specs = [ + (storage_addr, storage_block_len, storage_block_len) + ] + elif storage_is_block_major: + region_specs = [ + (cache.data_ptr(), physical_page_size, block_stride) + ] + else: + segment_bytes = num_blocks * block_stride + assert cache.nbytes % segment_bytes == 0 + num_segments = cache.nbytes // segment_bytes + region_block_len = ( + block_stride if num_segments > 1 else physical_page_size + ) + region_specs = [ + ( + cache.data_ptr() + segment_idx * segment_bytes, + region_block_len, + block_stride, + ) + for segment_idx in range(num_segments) + ] + + for base_addr, block_len, block_stride in region_specs: + if base_addr in seen_base_addresses: + if is_mla_region: + # Dual-purpose HMA tensor: an MLA layer shares a region + # that a non-MLA layer registered first. MLA is not + # head-sharded, so the region must be flagged MLA. + idx = seen_base_addresses.index(base_addr) + self._region_is_mla[idx] = True + continue + seen_base_addresses.append(base_addr) + self.block_len_per_layer.append(block_len) + self.block_stride_per_layer.append(block_stride) + self._region_is_mla.append(is_mla_region) # When there's a mismatch between kbs<>bs, we rely on HMA to ensure # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs. @@ -1228,11 +1172,6 @@ class NixlBaseConnectorWorker: f"kv_cache_layout={self.kv_cache_layout}" ) - # Need to make sure the device ID is non-negative for NIXL, - # Torch uses -1 to indicate CPU tensors. - self.device_id = max(cache.get_device(), 0) - caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, "")) - logger.debug( "Different block lengths collected: %s", set(self.block_len_per_layer) ) @@ -1240,10 +1179,11 @@ class NixlBaseConnectorWorker: len(self.block_len_per_layer) == len(seen_base_addresses) == len(self._region_is_mla) + == len(self.block_stride_per_layer) ) self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses - self.num_regions = len(caches_data) + self.num_regions = len(seen_base_addresses) if self.pp_size > 1: start_layer, end_layer = self.model_config.get_layers_start_end_indices( @@ -1293,6 +1233,7 @@ class NixlBaseConnectorWorker: kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank], num_blocks=self.num_blocks, block_lens=self.block_len_per_layer, + block_strides=self.block_stride_per_layer, kv_cache_layout=self.kv_cache_layout if not self.use_host_buffer else self.host_buffer_kv_cache_layout, @@ -1431,10 +1372,9 @@ class NixlBaseConnectorWorker: block_arange = np.arange(num_blocks, dtype=np.uint64) parts: list[np.ndarray] = [] for i, base_addr in enumerate(base_addresses): - # K/V are packed into the content dim, so the whole block transfers - # as one unit: desc length equals the block stride. block_len = self.block_len_per_layer[i] // block_size_ratio - addrs = base_addr + block_arange * block_len + block_stride = self.block_stride_per_layer[i] // block_size_ratio + addrs = base_addr + block_arange * block_stride parts.append(self._stack_descs(addrs, block_len, device_id)) return np.concatenate(parts) @@ -1476,7 +1416,7 @@ class NixlBaseConnectorWorker: ) local_block_len = local_block_len // num_reads - page_size = nixl_agent_meta.block_lens[i] + page_size = nixl_agent_meta.block_strides[i] addrs = base_addr + rank_offset + block_arange * page_size parts.append(self._stack_descs(addrs, local_block_len, device_id)) return np.concatenate(parts) @@ -1560,7 +1500,7 @@ class NixlBaseConnectorWorker: tp_ratio = 4 // 2 = 2 Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim] - then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format. + then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "LBHNC" layout format. Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split 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: start:end ] nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end] + nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end] ### Register remote engine in TransferTopology (idempotent). assert self.transfer_topo is not None @@ -1790,10 +1731,10 @@ class NixlBaseConnectorWorker: if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout: if ( self.kv_transfer_config.enable_permute_local_kv - and nixl_agent_meta.kv_cache_layout == "HND" + and nixl_agent_meta.kv_cache_layout == "LBHNC" ): logger.info( - "Remote is HND and local is NHD, enabled additional permute " + "Remote is LBHNC and local is LBNHC, enabled additional permute " "on local device KV." ) assert not self._is_hma_required, ( @@ -1823,18 +1764,18 @@ class NixlBaseConnectorWorker: self.enable_heterogeneous_attn_post_process = True # Heterogeneous TP requires head-splitting, which only works with - # HND layout. MLA and replicated-KV cases don't split on heads. - # Mamba doesn't support heterogeneous TP. + # block-contiguous layouts (e.g. LBHNC). MLA and replicated-KV cases + # don't split on heads. Mamba doesn't support heterogeneous TP. if ( abs(tp_ratio) != 1 and not self.use_mla and not self.transfer_topo.is_kv_replicated(remote_engine_id) - and kv_cache_layout != "HND" + and not KVCacheLayout[kv_cache_layout].is_block_contiguous and not self.enable_permute_local_kv ): raise RuntimeError( "Heterogeneous TP head-dimension splitting requires contiguous heads. " - "Use HND layout on the prefill side." + "Use a block-contiguous layout (e.g. LBHNC) on the prefill side." ) # Per-region block_len validation enforcing the P/D invariant. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 77caec602c..fa5537c325 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -19,7 +19,6 @@ import torch from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.utils import ( EngineId, - get_current_attn_backend, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( CopyBlocksOp, @@ -56,10 +55,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( ) from vllm.forward_context import ForwardContext from vllm.logger import init_logger -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata -from vllm.v1.attention.backends.utils import get_kv_cache_layout +from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -83,36 +80,6 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): def supports_divergent_local_hybrid_hits(self) -> bool: return True - @property - def prefer_cross_layer_blocks(self) -> bool: - if any( - [ - isinstance(group.kv_cache_spec, MambaSpec) - for group in self.kv_cache_config.kv_cache_groups - ] - ): - # Hybrid SSM models do not yet support cross-layer layout - return False - - backend = get_current_attn_backend(self._vllm_config) - if backend.get_name() not in ( - "FLASH_ATTN", - "FLASHINFER", - "TRITON_ATTN", - ): - return False - - # For now there is no benefit to run cross layers when backend - # does not support on HND - if get_kv_cache_layout() != "HND": - return False - - extra_config = self.kv_transfer_config.kv_connector_extra_config - return ( - str(extra_config.get("enable_cross_layers_blocks", "False")).lower() - == "true" - ) - def __init__( self, vllm_config: VllmConfig, @@ -179,9 +146,10 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): # which fallback to the default behavior. return None logger.info_once( - "NixlConnector setting KV cache layout to HND for better xfer performance." + "NixlConnector setting KV cache layout to LBHNC for " + "better xfer performance." ) - return "HND" + return "LBHNC" ############################################################ # Scheduler Side Methods @@ -266,12 +234,6 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): assert self.connector_worker is not None self.connector_worker.register_kv_caches(kv_caches) - def register_cross_layers_kv_cache( - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] - ): - assert self.connector_worker is not None - self.connector_worker.register_cross_layers_kv_caches(kv_cache) - def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): assert self.connector_worker is not None self.connector_worker.set_host_xfer_buffer_ops(copy_operation) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 12cfa51d21..0fca03c93b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -44,8 +44,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:" # 6: Validate EAGLE/MTP speculative configuration compatibility # 7: Include NIXL transfer mode (push vs pull) in the compatibility hash # 8: Add dcp_size and pcp_size to NixlAgentMetadata +# 9: Add block_strides # -NIXL_CONNECTOR_VERSION: int = 8 +NIXL_CONNECTOR_VERSION: int = 9 @dataclass @@ -56,6 +57,7 @@ class NixlAgentMetadata: device_id: int num_blocks: int block_lens: list[int] + block_strides: list[int] kv_cache_layout: str block_size: int ssm_sizes: tuple[int, int] @@ -129,7 +131,6 @@ def _get_speculative_compatibility_factors( def compute_nixl_compatibility_hash( vllm_config: VllmConfig, attn_backend_name: str, - cross_layers_blocks: bool, transfer_mode: str = "pull", ) -> str: """ @@ -180,7 +181,6 @@ def compute_nixl_compatibility_hash( # Attention backend and KV cache dtype affect memory layout "attn_backend_name": attn_backend_name, "cache_dtype": str(cache_config.cache_dtype), - "cross_layers_blocks": cross_layers_blocks, "is_hma_enabled": is_hma_enabled, "speculative_config": _get_speculative_compatibility_factors(vllm_config), # push (WRITE) and pull (READ) connectors are protocol-incompatible diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py index 209616bcc9..27f08c1fd8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py @@ -17,6 +17,7 @@ import torch from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, + KVCacheLayout, KVCacheSpec, MambaSpec, MLAAttentionSpec, @@ -31,13 +32,16 @@ if TYPE_CHECKING: CANONICAL_FORMAT_VERSION = 1 -def canonical_format_id() -> str: +def canonical_format_id(kv_cache_layout: str) -> str: """Identity of the canonical byte format, for namespacing persisted KV. Canonical pages keep the worker's KV layout family, so the id couples the - format version with that family; consumers must match it exactly.""" - from vllm.v1.attention.backends.utils import get_kv_cache_layout - - return f"v{CANONICAL_FORMAT_VERSION}-{get_kv_cache_layout().lower()}" + format version with that family; consumers must match it exactly. The family + keeps its historical NHD/HND spelling so ids stay stable for KV persisted + before the layout enum existed.""" + layout = KVCacheLayout[kv_cache_layout] + legacy = {KVCacheLayout.LBNHC: "nhd", KVCacheLayout.LBHNC: "hnd"} + family = legacy.get(layout, layout.name.lower()) + return f"v{CANONICAL_FORMAT_VERSION}-{family}" @dataclass(frozen=True) @@ -278,7 +282,7 @@ def _layer_mapping( if isinstance(spec, MLAAttentionSpec): # TP-replicated latent; CP shards its tokens across the DCP groups if ( - spec.compress_ratio != 1 + spec.tokens_per_state != 1 or page % bs or ctx.tp_size % ctx.dcp_size or spec.kv_quant_mode.is_per_token_head diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index a38b2ca7a4..70919ce24e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -20,12 +20,7 @@ from vllm.v1.kv_offload.config import ( if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheTensor - - -def is_kv_cache_tensor_packed(kv_cache_tensor: "KVCacheTensor") -> bool: - """Return whether a KV cache tensor uses a packed block stride.""" - return bool(kv_cache_tensor.block_stride) + from vllm.v1.kv_cache_interface import KVCacheConfig def build_offloading_config( @@ -96,18 +91,10 @@ def build_offloading_config( blocks_per_chunk = tokens_per_chunk_int // tokens_per_block worker_kv_bytes_per_block = 0 - if kv_cache_config.num_blocks > 0: - packed_tensors = tuple( - is_kv_cache_tensor_packed(tensor) - for tensor in kv_cache_config.kv_cache_tensors - ) - is_packed = any(packed_tensors) - assert not is_packed or all(packed_tensors) - total_gpu_kv_bytes = ( - kv_cache_config.kv_cache_tensors[0].size - if is_packed - else sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors) - ) + if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors: + # Every KVCacheTensor describes placement within the same backing allocation, + # so its size is the total, not a per-tensor share. + total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks single_group_spec = ( @@ -164,7 +151,7 @@ def build_offloading_config( if isinstance(single_group_spec, FullAttentionSpec): if type(single_group_spec) is MLAAttentionSpec: spec_certifiable = ( - single_group_spec.compress_ratio == 1 + single_group_spec.tokens_per_state == 1 and single_group_spec.real_page_size_bytes % single_group_spec.block_size == 0 @@ -219,4 +206,5 @@ def build_offloading_config( ), replicated_layout=replicated_layout, canonical_layout=canonical_layout, + kv_cache_layout=vllm_config.cache_config.kv_cache_layout, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 211903242e..ecdae9c4f2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -14,16 +14,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( OffloadingWorkerMetadata, ReqId, ) -from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( - is_kv_cache_tensor_packed, -) from vllm.logger import init_logger -from vllm.v1.attention.backend import AttentionBackend from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, + group_kernel_blocks, ) from vllm.v1.kv_offload.base import ( CanonicalKVCacheRef, @@ -73,16 +70,6 @@ class OffloadingConnectorWorker: self.vllm_config, kv_cache_config, kv_caches ) - # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use - # stride(0) as the manager-block stride (equals total_num_bytes_per_block). - # General (non-packed) layouts size the tensor at page_size_bytes per - # manager block, so page_size_bytes is the correct offloading stride. - layer_is_packed: dict[str, bool] = { - ln: is_kv_cache_tensor_packed(kv_tensor) - for kv_tensor in kv_cache_config.kv_cache_tensors - for ln in kv_tensor.shared_by - } - # layer_name -> (num_blocks, page_size_bytes) tensor tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {} # layer_name -> size of (un-padded) page in bytes @@ -100,63 +87,52 @@ class OffloadingConnectorWorker: layer_kv_cache_spec = per_layer_specs.get( layer_name, group_kv_cache_spec ) - if isinstance(layer_kv_cache_spec, AttentionSpec): - layer_kv_cache = kv_caches[layer_name] - assert isinstance(layer_kv_cache, torch.Tensor) + ref = group_kernel_blocks(kv_caches[layer_name], num_blocks) + page = layer_kv_cache_spec.page_size_bytes + elem_size = ref.element_size() + byte_offset = ref.storage_offset() * elem_size + # Packed layouts (e.g. DSv4) interleave layers per block, so the + # attention tensor's stride(0) (the manager-block stride) exceeds + # page_size_bytes; other layouts have stride(0) == page_size_bytes. + block_stride_bytes = ( + ref.stride(0) * elem_size + if isinstance(layer_kv_cache_spec, AttentionSpec) + else page + ) + tensors_per_block[layer_name] = ( + torch.tensor([], dtype=torch.int8, device=ref.device).set_( + ref.untyped_storage(), + byte_offset, + (num_blocks, page), + (block_stride_bytes, 1), + ), + ) + page_size_bytes[layer_name] = page - page = layer_kv_cache_spec.page_size_bytes - elem_size = layer_kv_cache.element_size() - byte_offset = layer_kv_cache.storage_offset() * elem_size - block_stride_bytes = ( - layer_kv_cache.stride(0) * elem_size - if layer_is_packed[layer_name] - else page - ) - raw = torch.empty( - 0, - dtype=torch.int8, - device=layer_kv_cache.device, - ).set_(layer_kv_cache.untyped_storage()) - tensors_per_block[layer_name] = ( - torch.as_strided( - raw, - (num_blocks, page), - (block_stride_bytes, 1), - byte_offset, - ), - ) - page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes + if isinstance(layer_kv_cache_spec, AttentionSpec): unpadded_page_size_bytes[layer_name] = ( layer_kv_cache_spec.unpadded_page_size_bytes ) - elif isinstance(layer_kv_cache_spec, MambaSpec): - layer_kv_cache = kv_caches[layer_name] - assert layer_kv_cache.dtype == torch.int8 - tensors_per_block[layer_name] = ( - layer_kv_cache.view( - num_blocks, layer_kv_cache_spec.page_size_bytes - ), - ) - - page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes unpadded_page_size_bytes[layer_name] = replace( layer_kv_cache_spec, page_size_padded=None ).page_size_bytes - else: raise NotImplementedError - packed_kv_cache_tensor = next( + # Packed layouts (e.g. DSv4) interleave all layers within each manager + # block: a layer view's block stride exceeds its page size. Offload the + # whole packed block as a single transfer region. + packed_layer_name = next( ( - t - for t in kv_cache_config.kv_cache_tensors - if is_kv_cache_tensor_packed(t) and t.shared_by + layer_name + for layer_name, (tensor,) in tensors_per_block.items() + if tensor.stride(0) != tensor.shape[1] ), None, ) - if packed_kv_cache_tensor is not None: - (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]] + if packed_layer_name is not None: + (tensor,) = tensors_per_block[packed_layer_name] block_stride = tensor.stride(0) packed_tensor = tensor.as_strided( (num_blocks, block_stride), @@ -176,29 +152,16 @@ class OffloadingConnectorWorker: block_tensors: list[CanonicalKVCacheTensor] = [] block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list) - for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - # Filter to layers that were actually processed above. - # Packed KV allocation emits KVCacheTensor entries for - # every (tuple_idx, page_size) slot; slots where no group has a - # layer at that index produce an empty shared_by (reserved memory - # with no corresponding model layer). - tensor_layer_names = [ - n for n in kv_cache_tensor.shared_by if n in tensors_per_block - ] - if not tensor_layer_names: - continue + # Layers that alias the same bytes (cache groups overlay each other) share + # one transfer region: exactly the views with equal address and strides. + aliased_layers: dict[tuple[int, tuple[int, ...]], list[str]] = defaultdict(list) + for layer_name, layer_tensors in tensors_per_block.items(): + view = layer_tensors[0] + aliased_layers[(view.data_ptr(), view.stride())].append(layer_name) - # verify all layers in the group reference the exact same tensors + for tensor_layer_names in aliased_layers.values(): assert len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1 - assert ( - len({tensors_per_block[n][0].data_ptr() for n in tensor_layer_names}) - == 1 - ) - assert ( - len({tensors_per_block[n][0].stride() for n in tensor_layer_names}) == 1 - ) - # pick the first layer to represent the group first_layer_name = tensor_layer_names[0] for tensor in tensors_per_block[first_layer_name]: block_tensors.append( @@ -242,53 +205,6 @@ class OffloadingConnectorWorker: self._init_worker(canonical_kv_caches) - def register_cross_layers_kv_cache( - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] - ): - # verify that num_blocks is at physical position 0 in the cross-layers - # tensor layout. - test_shape = attn_backend.get_kv_cache_shape( - num_blocks=1234, block_size=16, num_kv_heads=1, head_size=256 - ) - num_blocks_logical_dim = test_shape.index(1234) + 1 - physical_to_logical = attn_backend.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim) - assert num_blocks_physical_dim == 0 - - kv_cache_groups = self.kv_cache_config.kv_cache_groups - assert len(kv_cache_groups) == 1 - kv_cache_spec = kv_cache_groups[0].kv_cache_spec - num_layers = len(kv_cache_groups[0].layer_names) - page_size_bytes = kv_cache_spec.page_size_bytes * num_layers - - assert kv_cache.storage_offset() == 0 - storage = kv_cache.untyped_storage() - assert len(storage) % page_size_bytes == 0 - num_blocks = len(storage) // page_size_bytes - tensor = ( - torch.tensor( - [], - dtype=torch.int8, - device=kv_cache.device, - ) - .set_(storage) - .view(num_blocks, page_size_bytes) - ) - kv_cache_tensor = CanonicalKVCacheTensor( - tensor=tensor, page_size_bytes=page_size_bytes - ) - # in cross layers layout, there's currently only a single group - kv_cache_data_ref = CanonicalKVCacheRef( - tensor_idx=0, page_size_bytes=page_size_bytes - ) - canonical_kv_caches = CanonicalKVCaches( - tensors=[kv_cache_tensor], group_data_refs=[[kv_cache_data_ref]] - ) - - self._init_worker(canonical_kv_caches) - def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata): assert self.worker is not None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index f003d0d6f8..a64f97805e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -37,7 +37,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( OffloadingConnectorWorker, ) from vllm.forward_context import ForwardContext -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig @@ -47,11 +47,6 @@ from vllm.v1.request import Request class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): - @property - def prefer_cross_layer_blocks(self) -> bool: - # Cross-layer slabs have no per-layer refs to certify canonical mappings - return not self._canonical_layout - @property def requires_kv_delivery(self) -> bool: # 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): assert self.connector_worker is not None self.connector_worker.register_kv_caches(kv_caches) - def register_cross_layers_kv_cache( - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] - ): - assert self.connector_worker is not None - self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend) - def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata): assert self.connector_worker is not None assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata) @@ -196,7 +185,7 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): @classmethod def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None: - return "HND" + return "LBHNC" def reset_cache(self) -> bool | None: assert self.connector_scheduler is not None diff --git a/vllm/envs.py b/vllm/envs.py index 2eba6649c5..7ef9c06592 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -242,7 +242,10 @@ if TYPE_CHECKING: VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5 - VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None + VLLM_KV_CACHE_LAYOUT: ( + Literal["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"] + | None + ) = None VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None VLLM_COMPUTE_NANS_IN_LOGITS: bool = False VLLM_RAISE_ON_LOGIT_NANS: bool = False @@ -1784,18 +1787,22 @@ environment_variables: dict[str, Callable[[], Any]] = { ), # KV Cache layout used throughout vllm. # Some common values are: - # - NHD - # - HND - # Where N=num_blocks, H=num_heads and D=head_size. The default value will + # - LBNHC + # - LBHNC + # Where N=num_states, H=num_heads and C=state_content. The default value will # leave the layout choice to the backend. Mind that backends may only # implement and support a subset of all possible layouts. + # LHBNC hoists the head dim outside the block dim; backends must opt in via + # AttentionBackend.supported_kv_cache_layouts(). "VLLM_KV_CACHE_LAYOUT": env_with_choices( - "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"] + "VLLM_KV_CACHE_LAYOUT", + None, + ["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"], ), # SSM conv state layout used for Mamba models. # - SD: (state_len, dim) — dim contiguous (default) # - DS: (dim, state_len) — TP-sharded dim on dim1, - # consistent with SSM temporal state and HND KV cache layout. + # consistent with SSM temporal state and LBHNC KV cache layout. "VLLM_SSM_CONV_STATE_LAYOUT": env_with_choices( "VLLM_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"] ), diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index fe732e2850..030d9b4d40 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -649,6 +649,10 @@ class MLAAttention(nn.Module, AttentionLayerBase): compile_native=True, ) + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + @property def chunked_prefill_workspace_size(self) -> int: if self._chunked_prefill_workspace_size is None: @@ -1451,27 +1455,6 @@ class MLACommonBackend(AttentionBackend): def get_builder_cls() -> type["MLACommonMetadataBuilder"]: return MLACommonMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, # assumed to be 1 for MLA - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - # Default to identity permutation to signal cross-layer allocation - # is unsupported. Each MLA backend must opt in to support cross-layer - # allocation by overriding this method. - return (0, 1, 2, 3) - return (0, 1, 2) - @classmethod def get_supported_head_sizes(cls) -> list[int]: return [320, 576] diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index 825bf3b1e0..2a129d27bf 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -61,16 +61,16 @@ def hpc_rope_norm_forward( return attn_layer = forward_context.no_compile_layers[layer_name] - # bind_kv_cache stores the per-layer KV cache as a single 5D tensor - # (num_blocks, 2, block_size, num_kv_heads, head_size), so use it directly. + # bind_kv_cache stores the per-layer KV cache as a single tensor + # (num_blocks, num_kv_heads, block_size, 2 * head_size), so use it directly. kv_cache = attn_layer.kv_cache if kv_cache.numel() == 0: output.zero_() return - assert kv_cache.dim() == 5, ( - f"Expected kv_cache to have 5 dims, got {tuple(kv_cache.shape)}" + assert kv_cache.dim() == 4, ( + f"Expected kv_cache to have 4 dims, got {tuple(kv_cache.shape)}" ) rope_norm = _hpc_rope_norm_instances[layer_name] @@ -295,6 +295,9 @@ class HpcRopeNorm(CustomOp, HpcModule): """ import hpc + # (B, H, N, 2*hs) -> two (B, N, H, hs) K/V views, as in HpcAttentionImpl. + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_dim, dim=-1) + num_actual_tokens = attn_metadata.num_actual_tokens num_prefill_reqs = attn_metadata.num_prefills num_decode_reqs = attn_metadata.num_decodes @@ -333,8 +336,8 @@ class HpcRopeNorm(CustomOp, HpcModule): if self.use_fp8: _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8( - key_cache=kv_cache[:, 0], - value_cache=kv_cache[:, 1], + key_cache=key_cache, + value_cache=value_cache, qkv=qkv_prefill, cos_sin=self.cos_sin_cache, num_seqlen_per_req=seq_lens_prefill, @@ -353,8 +356,8 @@ class HpcRopeNorm(CustomOp, HpcModule): attn_metadata.hpc_prefill_q_scale = q_scale else: hpc.rope_norm_store_kv( - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, qkv_prefill, self.cos_sin_cache, seq_lens_prefill, @@ -379,8 +382,8 @@ class HpcRopeNorm(CustomOp, HpcModule): if self.use_fp8: _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8( - key_cache=kv_cache[:, 0], - value_cache=kv_cache[:, 1], + key_cache=key_cache, + value_cache=value_cache, qkv=qkv_decode, cos_sin=self.cos_sin_cache, num_seqlen_per_req=num_seq_kvcache, @@ -401,8 +404,8 @@ class HpcRopeNorm(CustomOp, HpcModule): attn_metadata.hpc_split_k_flag = split_k_flag else: hpc.rope_norm_store_kv( - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, qkv_decode, self.cos_sin_cache, num_seq_kvcache, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 93ce2c7f3d..36f9d0bb7b 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -645,6 +645,11 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C]: the indexer kernels and + # kv_cache_as_quant_view index a 3-D block-major cache. + self.kv_cache = kv_cache.squeeze(1) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: return MLAAttentionSpec( block_size=self.cache_config.block_size, diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index 87dba7b5c9..d94db2c692 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -80,14 +80,14 @@ def dummy_attention(layer_name, _placeholder): def basic_cache( to_cache: torch.Tensor, # shape: [seq_len, num_heads, head_size] - kv_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size] + kv_cache: torch.Tensor, # shape: [num_blocks, num_heads, block_size, head_size] slot_mapping: torch.Tensor, # shape: [seq_len] ): # Padding slots are -1; redirect them to the null block (block 0, never # allocated to a request) so the scatter stays branch-free and sync-free. - block_size = kv_cache.shape[1] + block_size = kv_cache.shape[2] slot_mapping = slot_mapping.clamp_min(0) - kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache + kv_cache[slot_mapping // block_size, :, slot_mapping % block_size] = to_cache ######### CacheOnlyAttentionBackend ######## @@ -123,18 +123,6 @@ class CacheOnlyAttentionBackend(AttentionBackend): def get_impl_cls() -> type["CacheOnlyAttentionImpl"]: return CacheOnlyAttentionImpl - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - # We set `num_kv_heads = num_hidden_layers` and `head_size = hidden_size` - # We also don't use a k/v (2) dim - return (num_blocks, block_size, num_kv_heads, head_size) - @staticmethod def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]: return CacheOnlyAttentionMetadataBuilder diff --git a/vllm/model_executor/models/whisper_causal.py b/vllm/model_executor/models/whisper_causal.py index 91b378e0dc..93f2aa1741 100644 --- a/vllm/model_executor/models/whisper_causal.py +++ b/vllm/model_executor/models/whisper_causal.py @@ -5,6 +5,7 @@ import functools import logging import math from dataclasses import replace +from fractions import Fraction from functools import partial import torch @@ -128,14 +129,10 @@ def create_whisper_attention_backend_with_block_pooling( vllm_config: VllmConfig, device: torch.device, ): - assert kv_cache_spec.num_kv_heads % block_pool_size == 0 - # Scale pooled-unit quantities back to unpooled (encoder-token) - # units for the underlying attention metadata: the KV cache spec - # counts blocks in pooled units, but the kernel runs on the - # `block_pool_size`x-expanded sequence built below. - spec_overrides = dict( + kv_cache_spec = replace( + kv_cache_spec, block_size=kv_cache_spec.block_size * block_pool_size, - num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size, + tokens_per_state=1, ) if isinstance(kv_cache_spec, SlidingWindowSpec): # The manager keeps `sliding_window` in pooled units; the kernel @@ -143,8 +140,7 @@ def create_whisper_attention_backend_with_block_pooling( # unpooled units (`get_kv_cache_spec` sizes the pooled window so # this window always stays resident). assert sliding_window is not None - spec_overrides["sliding_window"] = sliding_window - kv_cache_spec = replace(kv_cache_spec, **spec_overrides) + kv_cache_spec = replace(kv_cache_spec, sliding_window=sliding_window) super().__init__(kv_cache_spec, layer_names, vllm_config, device) # Override model_config-derived values with the actual # encoder values from kv_cache_spec @@ -261,18 +257,6 @@ def create_whisper_attention_backend_with_block_pooling( overrides={ "get_builder_cls": lambda: WhisperCausalAttentionWithBlockPoolingBuilder, "get_impl_cls": lambda: WhisperCausalAttentionWithBlockPoolingImpl, - "get_kv_cache_shape": lambda num_blocks, - block_size, - num_kv_heads, - head_size, - cache_dtype_str: underlying_attn_backend.get_kv_cache_shape( - num_blocks, - # we stretch each block by `block_pool_size` - block_size * block_pool_size, - num_kv_heads // block_pool_size, - head_size, - cache_dtype_str, - ), "forward_includes_kv_cache_update": True, }, ) @@ -341,7 +325,7 @@ class WhisperCausalAttentionWithBlockPooling(Attention): assert isinstance(kv_cache_spec, AttentionSpec) kv_cache_spec = replace( kv_cache_spec, - num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads, + tokens_per_state=Fraction(1, self.block_pool_size), ) if isinstance(kv_cache_spec, SlidingWindowSpec): # The manager counts blocks in pooled units, so express the window in diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index c06185be93..a7e8783873 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -653,7 +653,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the # KV row. block_size = swa_metadata.block_size - swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim) + assert swa_kv_cache.shape[1:] == (block_size, self.head_dim) + swa_kv_cache_3d = swa_kv_cache if cache_dtype == torch.bfloat16: torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( q, @@ -684,6 +685,17 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ) return q_fp8 + def _global_topk_output_buffers( + self, topk_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor] | None: + if self.compress_ratio != 4 or self.eager_scratch_pool is None: + return None + return self.eager_scratch_pool.global_topk_outputs(topk_indices) + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + def get_attn_backend(self) -> type[AttentionBackend]: return self.backend_cls @@ -701,7 +713,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): num_kv_heads=1, head_size=self.head_dim, dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, - compress_ratio=self.compress_ratio, + tokens_per_state=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, alignment=576 if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", @@ -733,16 +745,20 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # head_dim already carries the fp8 scale padding - # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. + # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, - compress_ratio=self.compress_ratio, + tokens_per_state=self.compress_ratio, # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). alignment=576 if uses_fp8_ds_mla_layout else 512, ) diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 590e748145..ccfdbeaa43 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -75,25 +75,6 @@ class CompressorBackend(AttentionBackend): def get_builder_cls() -> type["CompressorMetadataBuilder"]: return CompressorMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - assert num_kv_heads == 1 - return (num_blocks, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (0, 1, 2, 3) - return (0, 1, 2) - @dataclass class CompressorMetadata: @@ -185,6 +166,10 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase): else: raise ValueError(f"Invalid compress ratio: {compress_ratio}") + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # fp8_ds_mla is the UE8M0 paged layout and needs 576B alignment. Plain # full-cache rows share state pages with contiguous KV pages, so padding diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 1250934b8c..bca71abbad 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -22,7 +22,6 @@ from vllm.models.deepseek_v4.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, ) -from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 from vllm.v1.attention.backend import MultipleOf @@ -165,26 +164,6 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend): return None return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x" - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - device_capability = current_platform.get_device_capability() - if device_capability is not None and device_capability.major == 12: - return DeepseekV4SparseMLABackend.get_kv_cache_shape( - num_blocks, - block_size, - num_kv_heads, - head_size, - cache_dtype_str, - ) - assert num_kv_heads == 1 - return (num_blocks, block_size, head_size) - class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4.""" diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 30e386df32..ebff031cea 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -88,21 +88,6 @@ class DeepseekV4SparseMLABackend(AttentionBackend): def supports_compute_capability(cls, capability: DeviceCapability) -> bool: return capability.major in [9, 10] - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if cache_dtype_str == "fp8_ds_mla": - # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale). - # head_size passed in is the semantic head_dim (512). - return (num_blocks, block_size, 584) - else: - return (num_blocks, block_size, head_size) - @dataclass class DeepseekV4FlashMLAMetadata(AttentionMetadata): @@ -151,8 +136,8 @@ class DeepseekV4SparseMLAMetadataBuilder( (max_num_batched_tokens,), dtype=torch.int32, device=device ) - assert hasattr(self.kv_cache_spec, "compress_ratio") - self.compress_ratio = self.kv_cache_spec.compress_ratio + assert isinstance(self.kv_cache_spec.tokens_per_state, int) + self.compress_ratio = self.kv_cache_spec.tokens_per_state # Pre-allocate compressed slot mapping buffer for CUDA graph address # stability when compress_ratio > 1. @@ -206,7 +191,7 @@ class DeepseekV4SparseMLAMetadataBuilder( cm.query_start_loc, cm.seq_lens, cm.block_table_tensor.clamp_(min=0), - int(self.kv_cache_spec.storage_block_size), + int(self.kv_cache_spec.num_states), self.compress_ratio, out=self.compressed_slot_mapping_buffer, ) diff --git a/vllm/models/inkling/amd/sconv_swa_attn.py b/vllm/models/inkling/amd/sconv_swa_attn.py index 7b8958d556..ac8cd073be 100644 --- a/vllm/models/inkling/amd/sconv_swa_attn.py +++ b/vllm/models/inkling/amd/sconv_swa_attn.py @@ -107,32 +107,6 @@ class InklingSconvBackend(AttentionBackend): def get_name() -> str: return "INKLING_SCONV_SWA" - @classmethod - def indexes_kv_by_block_stride(cls) -> bool: - # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the - # padded conv page is read through a strided view. - return True - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - # HND, num-blocks-first, head-major: [num_blocks, H, N, D]. - return (num_blocks, num_kv_heads, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # Identity: physical layout == logical [num_blocks, H, N, D]. - if include_num_layers_dimension: - return (0, 1, 2, 3, 4) - return (0, 1, 2, 3) - @staticmethod def get_impl_cls(): raise NotImplementedError( diff --git a/vllm/models/inkling/nvidia/sconv_swa_attn.py b/vllm/models/inkling/nvidia/sconv_swa_attn.py index 8a2b288d5e..0e3865cc6d 100644 --- a/vllm/models/inkling/nvidia/sconv_swa_attn.py +++ b/vllm/models/inkling/nvidia/sconv_swa_attn.py @@ -107,32 +107,6 @@ class InklingSconvBackend(AttentionBackend): def get_name() -> str: return "INKLING_SCONV_SWA" - @classmethod - def indexes_kv_by_block_stride(cls) -> bool: - # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the - # padded conv page is read through a strided view. - return True - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - # HND, num-blocks-first, head-major: [num_blocks, H, N, D]. - return (num_blocks, num_kv_heads, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # Identity: physical layout == logical [num_blocks, H, N, D]. - if include_num_layers_dimension: - return (0, 1, 2, 3, 4) - return (0, 1, 2, 3) - @staticmethod def get_impl_cls(): raise NotImplementedError( diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 9eb89028bb..21be0c398f 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -384,6 +384,10 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): # ------------------------------------------------------------------ # AttentionLayerBase interface # ------------------------------------------------------------------ + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + def get_attn_backend(self) -> type[AttentionBackend]: return self.attn_backend diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 798f9a475b..9f0e4bad57 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -728,6 +728,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # Main GQA K/V cache. Block size may change after load, refresh it. + # AITER sparse PA wants K and V as two separate head groups + # (H folded into the content dim). + sparse_pa = self.use_aiter_sparse_pa + kv_bytes = ( + self.num_kv_heads * self.head_dim * self.kv_cache_torch_dtype.itemsize + ) return FullAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=self.num_kv_heads, @@ -735,6 +741,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): head_size_v=self.head_dim, dtype=self.kv_cache_torch_dtype, kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + num_head_slots=2 if sparse_pa else None, + state_content_bytes=kv_bytes if sparse_pa else None, ) def _ensure_aiter_sparse_pa_kv_cache(self) -> None: diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 767aa2d19b..49c1fc83b8 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -97,25 +97,6 @@ class MiniMaxM3IndexerBackend(AttentionBackend): def is_sparse(cls) -> bool: return True - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - # M3 does not use cross-layer (per-layer-stacked) KV blocks. - raise NotImplementedError - return (0, 1, 2) - class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase): """Side KV cache for the indexer's per-token index keys (key-only). @@ -156,6 +137,10 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # Key-only: MLAAttentionSpec budgets one vector/token (not 2x for K+V). return MLAAttentionSpec( diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index e8154459dd..82256add46 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -51,10 +51,13 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.utils import ( - get_kv_cache_layout, split_decodes_and_prefills, ) -from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheLayout, + is_quantized_kv_cache, +) logger = init_logger(__name__) @@ -112,47 +115,13 @@ class MiniMaxM3SparseBackend(AttentionBackend): def is_sparse(cls) -> bool: return True - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if minimax_m3_use_aiter_sparse_pa(num_kv_heads): - # AITER's assembly paged-attention kernels require independently - # contiguous K and V storage. Keep that specialized layout behind - # the shuffle flag while every other implementation uses the - # packed-content contract introduced by #44455. - return (num_blocks, 2, block_size, num_kv_heads, head_size) - # K and V are packed into the content dim: logical (B, H, N, 2*hs). - return (num_blocks, num_kv_heads, block_size, 2 * head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets us from - # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory - # layout we want. - if include_num_layers_dimension: - raise NotImplementedError # no cross-layer KV blocks in M3 + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: + # AITER sparse PA unbinds K and V as separate contiguous regions, which + # requires the head dimension to sit outside the block dimension. if _minimax_m3_aiter_sparse_pa_requested(): - # The AITER page-16 sparse PA path reinterprets the K and V slices - # as separate SHUFFLE caches. Keep K/V physically separated so - # kv_cache[:, 0] and kv_cache[:, 1] are contiguous byte ranges. - return (1, 0, 2, 3, 4) - cache_layout = get_kv_cache_layout() - if cache_layout == "NHD": - # (num_blocks, block_size, num_kv_heads, 2*head_size) - stride_order = (0, 2, 1, 3) - elif cache_layout == "HND": - # (num_blocks, num_kv_heads, block_size, 2*head_size) - stride_order = (0, 1, 2, 3) - else: - raise ValueError(f"Unknown cache layout format {cache_layout}.") - return stride_order + return (KVCacheLayout.LHBNC,) + return super().supported_kv_cache_layouts() @dataclass diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 0e133a3773..faee4305c6 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -117,6 +117,21 @@ def is_strictly_contiguous(t: torch.Tensor) -> bool: return True +def is_non_overlapping_and_dense(t: torch.Tensor) -> bool: + """Check if the tensor's elements cover one gapless, non-overlapping byte + range, in any dimension order (i.e. a permuted view of a contiguous + tensor); ``is_contiguous()`` additionally requires row-major order. + """ + expected_stride = 1 + for size, stride in sorted(zip(t.shape, t.stride()), key=lambda p: p[1]): + if size == 1: + continue + if stride != expected_stride: + return False + expected_stride *= size + return True + + def canonicalize_singleton_dim_strides(t: torch.Tensor) -> torch.Tensor: """Fix degenerate strides on size=1 dimensions for CUDA TMA compatibility. diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 04e91abb03..4fdfafaaea 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -23,10 +23,14 @@ if TYPE_CHECKING: from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms.interface import DeviceCapability - from vllm.v1.attention.backends.utils import KVCacheLayoutType - from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, KVQuantMode + from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheLayout, + KVCacheSpec, + KVQuantMode, + ) -from vllm.v1.kv_cache_interface import get_kv_quant_mode +from vllm.v1.kv_cache_interface import KVCacheLayout, get_kv_quant_mode class AttentionType(str, Enum): @@ -84,68 +88,6 @@ class AttentionBackend(ABC): def get_builder_cls(): # -> Type["AttentionMetadataBuilder"]: raise NotImplementedError - @staticmethod - @abstractmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - raise NotImplementedError - - @classmethod - def get_kv_cache_block_dim( - cls, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> int: - """Discover which tensor dim is the block index, since different - backends lay out dims differently.""" - _S = 1234567 - shape = cls.get_kv_cache_shape( - _S, - block_size, - num_kv_heads, - head_size, - cache_dtype_str=cache_dtype_str, - ) - return shape.index(_S) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - """ - Get the physical (memory layout) ordering of the kv cache dimensions. - Standard attention backends pack K and V into the content dim, giving - the logical shape [num_blocks, num_heads, block_size, 2 * head_size]. - e.g. if get_kv_cache_stride_order returns (0, 2, 1, 3) then the physical - ordering of dimensions is - [num_blocks, block_size, num_heads, 2 * head_size]. - - If this function is unimplemented / raises NotImplementedError, - the physical layout of the KV cache will match the logical shape. - - Args: - include_num_layers_dimension: if True, includes an additional - num_layers dimension, which is assumed to be prepended - to the logical KV cache shape. - With the above example, a return value (1, 0, 3, 2, 4) - corresponds to - [num_blocks, num_layers, block_size, num_heads, 2 * head_size]. - - If an additional dimension is NOT included in the returned - tuple, the physical layout will not include a layers dimension. - - Returns: - A tuple of ints which is a permutation of range(len(shape)). - """ - raise NotImplementedError - @classmethod def full_cls_name(cls) -> tuple[str, str]: return (cls.__module__, cls.__qualname__) @@ -215,38 +157,6 @@ class AttentionBackend(ABC): return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes) - @classmethod - def indexes_kv_by_block_stride(cls) -> bool: - """Whether the backend reads KV pages by the runtime block stride. - - True when ``num_blocks`` is the outermost physical dimension of the KV - cache, so the backend tolerates a non-contiguous block dim. This gates - page size padding and cross-layer uniform KV layout. - - Returns: - True if the backend's physical KV layout is num-blocks-first. False - otherwise, including when the backend does not define a layered - stride order. - """ - try: - kv_cache_stride_order = cls.get_kv_cache_stride_order( - include_num_layers_dimension=False - ) - layered_kv_cache_stride_order = cls.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - except (AttributeError, NotImplementedError): - return False - - # Check that attention backend includes a layers dimension. - if len(layered_kv_cache_stride_order) != len(kv_cache_stride_order) + 1: - return False - - # stride_order[0] == 0 means num_layers stays first in physical - # layout (identity permutation), so indexing by block stride is - # not supported. - return layered_kv_cache_stride_order[0] != 0 - @classmethod def is_mla(cls) -> bool: return False @@ -438,7 +348,9 @@ class AttentionBackend(ABC): return invalid_reasons @classmethod - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: + """Layouts this backend's kernels can consume, most preferred first, or + None when the kernels consume any layout and express no preference.""" return None @classmethod diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index a6cd254fea..7f06ec1392 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -30,13 +30,13 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.utils import ( - KVCacheLayoutType, get_num_attention_heads_from_layers, ) from vllm.v1.kv_cache_interface import ( AttentionSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, + KVCacheLayout, ) logger = init_logger(__name__) @@ -65,6 +65,11 @@ class CPUAttentionBackend(AttentionBackend): def get_supported_head_sizes(cls) -> list[int]: return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512] + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # The CPU backend only reads head-major block interiors. + return (KVCacheLayout.LBHNC,) + @staticmethod def get_name() -> str: return "CPU_ATTN" @@ -96,20 +101,6 @@ class CPUAttentionBackend(AttentionBackend): def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]: return CPUAttentionMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return num_blocks, num_kv_heads, block_size, 2 * head_size - - @classmethod - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": - return "HND" - @staticmethod def use_cascade_attention(*args, **kwargs) -> bool: return False diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 76aaed4c5b..141aa45c20 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -33,7 +33,6 @@ from vllm.v1.attention.backends.fa_utils import ( from vllm.v1.attention.backends.utils import ( fill_mm_prefix_query_ranges, get_dcp_local_seq_lens, - get_num_attention_heads_from_layers, ) from vllm.v1.attention.ops.dcp import ( cp_lse_ag_out_rs, @@ -65,7 +64,6 @@ from vllm.v1.attention.backend import ( AttentionMetadataBuilder, CommonAttentionMetadata, ) -from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec from vllm.v1.worker.cp_utils import ( run_split_fa2_dcp_context_attention, @@ -138,43 +136,6 @@ class FlashAttentionBackend(AttentionBackend): def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]: return FlashAttentionMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - # K and V are packed into the content dim: logical (B, H, N, 2*D). - return (num_blocks, num_kv_heads, block_size, 2 * head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets us from - # `get_kv_cache_shape` (logical (B, H, N, 2*D)) to the actual memory - # layout we want. - cache_layout = get_kv_cache_layout() - if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size) - return (1, 0, 3, 2, 4) - elif cache_layout == "NHD": - # (num_blocks, block_size, num_kv_heads, 2*head_size) - stride_order = (0, 2, 1, 3) - elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size) - return (1, 2, 0, 3, 4) - elif cache_layout == "HND": - # (num_blocks, num_kv_heads, block_size, 2*head_size) - stride_order = (0, 1, 2, 3) - else: - raise ValueError(f"Unknown cache layout format {cache_layout}.") - return stride_order - @classmethod def supports_head_size(cls, head_size: int) -> bool: if head_size % 8 != 0: @@ -436,12 +397,12 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad self.compilation_config = vllm_config.compilation_config self.attention_config = vllm_config.attention_config - self.num_heads_q = get_num_attention_heads_from_layers( - vllm_config, layer_names - ) or self.model_config.get_num_attention_heads(self.parallel_config) - self.num_heads_kv = kv_cache_spec.num_kv_heads + self.num_heads_q = self.model_config.get_num_attention_heads( + self.parallel_config + ) + self.num_heads_kv = self.model_config.get_num_kv_heads(self.parallel_config) self.kv_cache_dtype = kv_cache_spec.dtype - self.headdim = kv_cache_spec.head_size + self.headdim = self.model_config.get_head_size() self.block_size = kv_cache_spec.block_size self.max_num_splits = 0 # No upper bound on the number of splits. diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index 02c88b4b36..326b21c2c1 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -21,7 +21,6 @@ from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( if is_flash_attn_varlen_func_available(): from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func -from vllm.v1.attention.backends.utils import get_kv_cache_layout from .flash_attn import ( FlashAttentionBackend, @@ -73,49 +72,6 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend): def get_impl_cls() -> type["FlashAttentionImpl"]: return FlashAttentionDiffKVImpl - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - # Logical (blocks-first, head-major) layout: K and V (with their - # different head sizes) packed in the content dim. - return ( - num_blocks, - num_kv_heads, - block_size, - head_size + FlashAttentionDiffKVBackend.head_size_v, - ) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets us from - # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual - # memory layout we want. - cache_layout = get_kv_cache_layout() - if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v) - return (1, 0, 3, 2, 4) - elif cache_layout == "NHD": - # (num_blocks, block_size, num_kv_heads, C_k+C_v) - stride_order = (0, 2, 1, 3) - elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v) - return (1, 2, 0, 3, 4) - elif cache_layout == "HND": - # (num_blocks, num_kv_heads, block_size, C_k+C_v) - stride_order = (0, 1, 2, 3) - else: - raise ValueError(f"Unknown cache layout format {cache_layout}.") - return stride_order - class FlashAttentionDiffKVImpl(FlashAttentionImpl): vllm_flash_attn_version: int | None diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 5d94ecef2c..792056f61d 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -67,9 +67,8 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.utils import ( - KVCacheLayoutType, get_dcp_local_seq_lens, - get_kv_cache_layout, + get_flashinfer_layout_string, get_num_attention_heads_from_layers, get_per_layer_parameters, infer_global_hyperparameters, @@ -83,6 +82,7 @@ from vllm.v1.attention.ops.dcp import ( from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.kv_cache_interface import ( AttentionSpec, + KVCacheLayout, KVCacheSpec, KVQuantMode, UniformTypeKVCacheSpecs, @@ -282,6 +282,7 @@ def trtllm_prefill_attn_kvfp8_dequant( class BatchDCPPrefillWrapper: def __init__( self, + kv_layout: str, workspace_buffer: torch.Tensor | None = None, dcp_a2a: bool = False, ): @@ -289,9 +290,7 @@ class BatchDCPPrefillWrapper: self._dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False) else: self._dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False) - self._context = BatchPrefillWithPagedKVCacheWrapper( - workspace_buffer, get_kv_cache_layout() - ) + self._context = BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, kv_layout) self._new_tokens = BatchPrefillWithRaggedKVCacheWrapper(workspace_buffer) def plan( @@ -467,44 +466,6 @@ class FlashInferBackend(AttentionBackend): def get_builder_cls() -> type["FlashInferMetadataBuilder"]: return FlashInferMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if cache_dtype_str.startswith("nvfp4"): - full_dim = nvfp4_kv_cache_full_dim(head_size) - return (num_blocks, 2 * num_kv_heads, block_size, full_dim) - # Pack K and V in the content dim (B, H, N, 2*hs). - return (num_blocks, num_kv_heads, block_size, 2 * head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets us from - # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory - # layout we want. - cache_layout = get_kv_cache_layout() - if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size) - return (1, 0, 3, 2, 4) - elif cache_layout == "NHD": - # (num_blocks, block_size, num_kv_heads, 2*head_size) - stride_order = (0, 2, 1, 3) - elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size) - return (1, 2, 0, 3, 4) - elif cache_layout == "HND": - # (num_blocks, num_kv_heads, block_size, 2*head_size) - stride_order = (0, 1, 2, 3) - else: - raise ValueError(f"Unknown cache layout format {cache_layout}.") - return stride_order - @staticmethod def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype: if kv_cache_dtype in ("fp8", "fp8_e4m3"): @@ -563,11 +524,13 @@ class FlashInferBackend(AttentionBackend): ) and supports_trtllm_attention(is_prefill=True) @classmethod - def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None: + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: capability = current_platform.get_device_capability() if capability is not None and capability.major == 10: - return "HND" - return None + # The trtllm-gen kernels consume head-major block interiors; the L/B + # nesting outside the block is immaterial to them. + return (KVCacheLayout.LBHNC, KVCacheLayout.BLHNC) + return super().supported_kv_cache_layouts() forward_includes_kv_cache_update: bool = False @@ -945,6 +908,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): self.paged_kv_indices = self._make_buffer(max_num_pages) self.paged_kv_last_page_len = self._make_buffer(max_num_reqs) + @property + def kv_cache_layout(self) -> KVCacheLayout: + return self.cache_config.get_resolved_kv_cache_layout() + # Keep SM90 prefill/decode Q dtype selection in one place. def get_q_data_type(self, is_prefill: bool) -> torch.dtype: # The user sets --attention-config.disable_flashinfer_q_quantization @@ -1155,7 +1122,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): self._noncausal_prefill_wrapper = ( BatchAttentionWithAttentionSinkWrapper( self._get_workspace_buffer(), - get_kv_cache_layout(), + get_flashinfer_layout_string(self.kv_cache_layout), backend="auto", q_data_type=self.q_data_type_prefill, kv_data_type=self.kv_cache_dtype, @@ -1168,7 +1135,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): self._noncausal_prefill_wrapper = ( BatchPrefillWithPagedKVCacheWrapper( self._get_workspace_buffer(), - get_kv_cache_layout(), + get_flashinfer_layout_string(self.kv_cache_layout), backend="auto", ) ) @@ -1177,6 +1144,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): if self._prefill_wrapper is None: if self.use_dcp: self._prefill_wrapper = BatchDCPPrefillWrapper( + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout), workspace_buffer=self._get_workspace_buffer(), dcp_a2a=self.dcp_a2a, ) @@ -1185,7 +1153,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): assert not self.is_kvcache_nvfp4 self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper( self._get_workspace_buffer(), - get_kv_cache_layout(), + get_flashinfer_layout_string(self.kv_cache_layout), backend="auto", q_data_type=self.q_data_type_prefill, kv_data_type=self.kv_cache_dtype, @@ -1199,7 +1167,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( self._get_workspace_buffer(), - get_kv_cache_layout(), + get_flashinfer_layout_string(self.kv_cache_layout), backend=backend, ) assert self._prefill_wrapper is not None @@ -1225,7 +1193,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( self._get_workspace_buffer(), - get_kv_cache_layout(), + get_flashinfer_layout_string(self.kv_cache_layout), use_cuda_graph=use_cudagraph, paged_kv_indptr_buffer=paged_kv_indptr, paged_kv_indices_buffer=paged_kv_indices, @@ -1248,7 +1216,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): def _get_cascade_wrapper(self): if self._cascade_wrapper is None: self._cascade_wrapper = MultiLevelCascadeAttentionWrapper( - 2, self._get_workspace_buffer(), get_kv_cache_layout() + 2, + self._get_workspace_buffer(), + get_flashinfer_layout_string(self.kv_cache_layout), ) return self._cascade_wrapper @@ -1813,6 +1783,8 @@ class FlashInferImpl(AttentionImpl): num_heads, num_kv_heads, is_prefill=False ) vllm_config = get_current_vllm_config_or_none() + # The layout is resolved after model construction, so read it lazily. + self.cache_config = vllm_config.cache_config if vllm_config else None # Query pre-quantization needs a single dtype for the whole query tensor. # SM90 XQA needs BF16/FP16-Q for decode and FP8 for prefill, # so only enable this for SM100 trtllm-gen where both use FP8-Q. @@ -1848,6 +1820,11 @@ class FlashInferImpl(AttentionImpl): else: self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False) + @property + def kv_cache_layout(self) -> KVCacheLayout: + assert self.cache_config is not None + return self.cache_config.get_resolved_kv_cache_layout() + def fused_output_quant_supported(self, quant_key: QuantKey): # XQA does not support FP8/NVFP4 output, so require trtllm-gen # (SM100+) here. Without that we cannot fuse the output quant. @@ -1919,9 +1896,7 @@ class FlashInferImpl(AttentionImpl): query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] - kv_cache: KV cache tensor with different possible shapes: - - NHD: [num_blocks, 2, block_size, num_kv_heads, head_size] - - HND: [num_blocks, 2, num_kv_heads, block_size, head_size] + kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size] attn_metadata: Metadata for attention. Returns: shape = [num_tokens, num_heads * head_size] @@ -2019,7 +1994,7 @@ class FlashInferImpl(AttentionImpl): if attn_metadata.use_cascade: # Cascade attention (rare case). assert attn_metadata.cascade_wrapper is not None - stride_order = FlashInferBackend.get_kv_cache_stride_order() + stride_order = self.kv_cache_layout.layer_view_order if self.is_kvcache_nvfp4: kv_cache_views = tuple( cache.permute(*stride_order) @@ -2039,7 +2014,7 @@ class FlashInferImpl(AttentionImpl): num_decode_tokens = attn_metadata.num_decode_tokens num_prefill_tokens = attn_metadata.num_prefill_tokens - stride_order = FlashInferBackend.get_kv_cache_stride_order() + stride_order = self.kv_cache_layout.layer_view_order kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1 # with TP=8). PyTorch permits non-canonical strides on size-1 dims; @@ -2195,7 +2170,7 @@ class FlashInferImpl(AttentionImpl): seq_lens_prefill = attn_metadata.prefill.seq_lens # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND - assert get_kv_cache_layout() == "HND" + assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND" assert is_strictly_contiguous(prefill_query) assert is_strictly_contiguous(workspace_buffer) assert is_strictly_contiguous(block_tables_prefill) @@ -2380,7 +2355,7 @@ class FlashInferImpl(AttentionImpl): # trtllm-gen needs HND layout on SM100. XQA is selected # separately on SM90 and does not use this SM100 layout gate. if decode_with_trtllm_gen: - assert get_kv_cache_layout() == "HND" + assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND" else: assert decode_with_xqa assert is_strictly_contiguous(decode_query) @@ -2426,7 +2401,7 @@ class FlashInferImpl(AttentionImpl): window_left=self.window_left, out=output[:num_decode_tokens], sinks=self.sinks, - kv_layout=get_kv_cache_layout(), + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout), q_len_per_req=q_len_per_req, mask=attn_metadata.decode.mask, q_cu_seq_lens=attn_metadata.decode.q_cu_seq_lens, @@ -2499,7 +2474,7 @@ class FlashInferImpl(AttentionImpl): sinks=self.sinks, o_sf_scale=self.o_sf_scale, out=out, - kv_layout=get_kv_cache_layout(), + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout), backend=attn_metadata.decode.kernel.value, q_len_per_req=q_len_per_req, kv_cache_sf=( diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index b7bcac54bb..61c6344373 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -43,7 +43,11 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, MultipleOf, ) -from vllm.v1.kv_cache_interface import AttentionSpec, EncoderOnlyAttentionSpec +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + EncoderOnlyAttentionSpec, + KVCacheLayout, +) logger = init_logger(__name__) @@ -118,6 +122,12 @@ class FlexAttentionBackend(AttentionBackend): def supports_batch_invariance(cls) -> bool: return True + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # Flex flattens the (B, N) block/token axes into a single token dim, which + # only LBNHC's strides allow as a zero-copy view of the layer cache. + return (KVCacheLayout.LBNHC,) + @classmethod def supports_mm_prefix(cls) -> bool: """FlexAttention supports full attention for image tokens.""" @@ -127,25 +137,6 @@ class FlexAttentionBackend(AttentionBackend): def get_impl_cls() -> type["FlexAttentionImpl"]: return FlexAttentionImpl - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - # K and V are packed into the content dim: logical (B, H, N, 2*hs). - return (num_blocks, num_kv_heads, block_size, 2 * head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (1, 0, 3, 2, 4) - return (0, 2, 1, 3) - @staticmethod def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]: return FlexAttentionMetadataBuilder diff --git a/vllm/v1/attention/backends/hpc_attn.py b/vllm/v1/attention/backends/hpc_attn.py index f2ba460e2b..f83d599455 100644 --- a/vllm/v1/attention/backends/hpc_attn.py +++ b/vllm/v1/attention/backends/hpc_attn.py @@ -30,10 +30,9 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.utils import ( - KVCacheLayoutType, split_decodes_and_prefills, ) -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout, KVCacheSpec logger = init_logger(__name__) @@ -265,7 +264,7 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): class HpcAttentionBackend(AttentionBackend): """HPC attention backend (pure attention, no RoPE/Norm). - KV cache layout: NHD (num_blocks, 2, block_size, num_kv_heads, head_size). + KV cache layout: LBNHC (num_blocks, num_kv_heads, block_size, 2 * head_size). """ accept_output_buffer: bool = True @@ -298,24 +297,6 @@ class HpcAttentionBackend(AttentionBackend): def get_builder_cls() -> type["HpcAttnMetadataBuilder"]: return HpcAttnMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, 2, block_size, num_kv_heads, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (1, 0, 2, 3, 4, 5) - return (0, 1, 2, 3, 4) - @classmethod def get_supported_head_sizes(cls) -> list[int]: return [128] @@ -331,8 +312,8 @@ class HpcAttentionBackend(AttentionBackend): return kv_cache_dtype in cls.supported_kv_cache_dtypes @classmethod - def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None: - return "NHD" + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + return (KVCacheLayout.LBNHC,) class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): @@ -452,6 +433,9 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): num_decode_reqs = attn_metadata.num_decodes num_decode_tokens = attn_metadata.num_decode_tokens + # Logical (B, H, N, 2*head_size) -> (B, N, H, head_size) K/V views. + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + # Standard-path KV write: when no upstream fused op wrote KV cache # (i.e. the model didn't wire HpcRopeNorm), fall back to the generic # reshape_and_cache_flash. This is what lets non-HunYuan-V3 models use @@ -466,8 +450,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): torch.ops._C_cache_ops.reshape_and_cache_flash( key, value, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, attn_metadata.slot_mapping, self.kv_cache_dtype, layer._k_scale, @@ -485,7 +469,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): f" (layer={getattr(layer, 'layer_name', '?')})" ) torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype) - kv_cache = kv_cache.view(torch_dtype) + key_cache = key_cache.view(torch_dtype) + value_cache = value_cache.view(torch_dtype) k_scale = layer._k_scale.reshape(1) v_scale = layer._v_scale.reshape(1) @@ -508,8 +493,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): if self.use_fp8: hpc.attention_with_kvcache_prefill_fp8( q_prefill, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, hpc_prefill_q_scale, k_scale, v_scale, @@ -523,8 +508,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): else: hpc.attention_with_kvcache_prefill_bf16( q_prefill, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, cu_seqlens_prefill, block_table_prefill, seq_lens_prefill, @@ -545,8 +530,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): if self.use_fp8: hpc.attention_decode_fp8( q_decode, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, block_table_decode, num_seq_kvcache, hpc_decode_q_scale, @@ -575,8 +560,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): ) hpc.attention_decode_bf16( q_decode, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, block_table_decode, num_seq_kvcache, mtp=mtp, diff --git a/vllm/v1/attention/backends/mla/cutlass_mla.py b/vllm/v1/attention/backends/mla/cutlass_mla.py index 832acbc7ce..8815bd9340 100644 --- a/vllm/v1/attention/backends/mla/cutlass_mla.py +++ b/vllm/v1/attention/backends/mla/cutlass_mla.py @@ -49,14 +49,6 @@ class CutlassMLABackend(MLACommonBackend): def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [128] - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (1, 0, 2, 3) - return (0, 1, 2) - @staticmethod def get_name() -> str: return "CUTLASS_MLA" diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py index 80833dfba6..63daa860fd 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla.py @@ -52,14 +52,6 @@ class FlashAttnMLABackend(MLACommonBackend): def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(16)] - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (1, 0, 2, 3) - return (0, 1, 2) - @staticmethod def get_name() -> str: return "FLASH_ATTN_MLA" diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 7c320b1ea3..4a78a700cf 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -24,6 +24,7 @@ from vllm.v1.attention.backend import ( ) from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, triton_convert_req_index_to_global_index, ) from vllm.v1.kv_cache_interface import AttentionSpec @@ -103,16 +104,6 @@ class FlashAttnMLASparseBackend(AttentionBackend): return "FlashAttention MLA Sparse requires model with index_topk" return None - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, block_size, head_size) - @dataclass class FlashAttnMLASparseMetadata(AttentionMetadata): @@ -223,11 +214,15 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + kv_rows, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) topk_indices, valid_counts = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) @@ -235,13 +230,8 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): cu_seqlens_q = torch.arange( 0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device ) - kv_cache = kv_c_and_k_pe_cache.view( - -1, attn_metadata.block_size, self.head_size - ) - k_cache = kv_cache[:, :, self.kv_lora_rank :].view( - -1, 1, 1, self.qk_rope_head_dim - ) - v_cache = kv_cache[:, :, : self.kv_lora_rank].view(-1, 1, 1, self.kv_lora_rank) + k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1) + v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1) out = flash_attn_varlen_func( q=q_rope, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index 1712f06caf..99dd7f1625 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -28,7 +28,6 @@ from vllm.v1.attention.backend import ( AttentionType, MultipleOf, ) -from vllm.v1.attention.backends.utils import KVCacheLayoutType logger = init_logger(__name__) @@ -133,14 +132,6 @@ class FlashInferMLABackend(MLACommonBackend): def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [32, 64] - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (1, 0, 2, 3) - return (0, 1, 2) - @staticmethod def get_name() -> str: return "FLASHINFER_MLA" @@ -194,10 +185,6 @@ class FlashInferMLABackend(MLACommonBackend): ) return None - @classmethod - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": - return "HND" - class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): can_return_lse_for_decode: bool = True diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index a6b9f4abb0..3cf3d9541d 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -31,7 +31,6 @@ from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, triton_filter_and_convert_dcp_index, ) -from vllm.v1.attention.backends.utils import KVCacheLayoutType from vllm.v1.kv_cache_interface import AttentionSpec if TYPE_CHECKING: @@ -123,20 +122,6 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase): return "FlashInfer MLA Sparse requires model with index_topk config" return None - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, # assumed to be 1 for MLA - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, block_size, head_size) - - @classmethod - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": - return "HND" - class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): """FlashInfer sparse MLA backend for SM120.""" @@ -216,23 +201,6 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): ) return None - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, # assumed to be 1 for MLA - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"): - # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE. - return (num_blocks, block_size, 656) - return (num_blocks, block_size, head_size) - - @classmethod - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": - return None - @dataclass class FlashInferMLASparseMetadata(AttentionMetadata): diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py index a8c7b2555d..f27c5b3dbe 100644 --- a/vllm/v1/attention/backends/mla/flashmla.py +++ b/vllm/v1/attention/backends/mla/flashmla.py @@ -58,14 +58,6 @@ class FlashMLABackend(MLACommonBackend): def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [64] - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (1, 0, 2, 3) - return (0, 1, 2) - @staticmethod def get_name() -> str: return "FLASHMLA" diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 7b0be952ca..2ff9a3013d 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -27,6 +27,7 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, triton_convert_req_index_to_global_index, triton_filter_and_convert_dcp_index, ) @@ -128,20 +129,6 @@ class FlashMLASparseBackend(AttentionBackend): def supports_compute_capability(cls, capability: DeviceCapability) -> bool: return capability.major in [9, 10] - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, # assumed to be 1 for MLA - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if cache_dtype_str == "fp8_ds_mla": - # V3.2 main MLA: 656-byte custom storage format. See module docstring. - return (num_blocks, block_size, 656) - else: - return (num_blocks, block_size, head_size) - @dataclass class FlashMLASparseMetadata(AttentionMetadata): @@ -643,18 +630,22 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): # Convert per-request indices to global slots (decode) or workspace # offsets (prefill). req_id_per_token covers the whole batch; slice it # to the MQA tokens (q may exclude prefill tokens routed to dense MHA). + kv_rows, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) topk_indices, topk_length = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[: topk_indices.shape[0]], attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) return self._bf16_flash_mla_kernel( q, - kv_c_and_k_pe_cache, + kv_rows, topk_indices, topk_length, ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2a21445f5d..4700fbea06 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -36,7 +36,7 @@ from vllm.v1.attention.backends.utils import ( get_dcp_local_seq_lens, split_decodes_and_prefills, ) -from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MLAAttentionSpec logger = init_logger(__name__) @@ -179,34 +179,18 @@ class DeepseekV32IndexerBackend(AttentionBackend): def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]: return DeepseekV32IndexerMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - assert num_kv_heads == 1 - return (num_blocks, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - # DeepseekV32Indexer kernels do not support cross-layer - # KV cache layout. Identity permutation keeps num_layers - # first, signaling incompatibility. - return (0, 1, 2, 3) - return (0, 1, 2) - class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend): @staticmethod def get_name() -> str: return "DEEPSEEK_V4_INDEXER" + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # DeepSeek-V4 packs the indexer pages beside the MLA latent pages inside + # each block, so the layer dim must sit inside the block dim. + return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC) + @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [256] @@ -638,7 +622,10 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): self.compress_ratio = 1 # Get compress_ratio for DeepseekV4 support if isinstance(self.kv_cache_spec, MLAAttentionSpec): - self.compress_ratio = self.kv_cache_spec.compress_ratio + # MLA compression is a whole number of tokens per state (fractions + # are whisper block pooling and never reach MLA). + assert isinstance(self.kv_cache_spec.tokens_per_state, int) + self.compress_ratio = self.kv_cache_spec.tokens_per_state if self.dcp_world_size > 1 and self.compress_ratio > 1: raise NotImplementedError( "DCP is not supported with sparse indexer KV compression " @@ -885,7 +872,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): query_start_loc, seq_lens, block_table, - self.kv_cache_spec.storage_block_size, + self.kv_cache_spec.num_states, self.compress_ratio, out=self.compressed_slot_mapping_buffer, ) @@ -1053,7 +1040,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): if current_platform.is_cuda() and has_deep_gemm(): metadata = get_paged_mqa_logits_metadata( seq_lens, - self.kv_cache_spec.storage_block_size, + self.kv_cache_spec.num_states, self.num_sms, indices=decode_indices, ) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 6ae880251f..3736a38f5f 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -293,16 +293,6 @@ class ROCMAiterMLASparseBackend(AttentionBackend): def get_impl_cls() -> type["ROCMAiterMLASparseImpl"]: return ROCMAiterMLASparseImpl - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, # assumed to be 1 for MLA - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, block_size, head_size) - @classmethod def is_mla(cls) -> bool: return True diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index d858eb677c..83b1afef1f 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -89,6 +89,10 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): # contiguous full-cache layout. assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn) + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + # [B, H=1, N, C] -> [B, N, C] + self.kv_cache = kv_cache.squeeze(1) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous # bf16/fp8 cache uses the natural element-size page. @@ -141,30 +145,6 @@ class DeepseekSparseSWABackend(AttentionBackend): return DeepseekV4ROCMAiterSparseSWAMetadataBuilder return DeepseekSparseSWAMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - assert num_kv_heads == 1 - if cache_dtype_str == "fp8_ds_mla": - # DeepseekV4 SWA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale). - # head_size passed in is the semantic head_dim (512). - return (num_blocks, block_size, 584) - else: - return (num_blocks, block_size, head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (0, 1, 2, 3) - return (0, 1, 2) - @dataclass class DeepseekSparseSWAMetadata: @@ -417,7 +397,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) self.head_size = mla_spec.head_size # Already considered quantization. - self.compress_ratio = mla_spec.compress_ratio + assert isinstance(mla_spec.tokens_per_state, int) + self.compress_ratio = mla_spec.tokens_per_state self.block_size = mla_spec.block_size self.max_model_len = self.vllm_config.model_config.max_model_len self.max_num_batched_tokens = ( diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index fc340a2246..8ffc31ca3b 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -7,6 +7,28 @@ import torch from vllm.triton_utils import tl, triton +def flat_kv_row_view( + kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim] + block_size: int, +) -> tuple[torch.Tensor, int]: + """Flat [row, head_dim] view of a paged cache and its physical rows per block. + + Token offset is ``block_idx * block_stride_rows + offset_in_block``. + When other layers' pages sit between consecutive blocks of this cache, + ``block_stride_rows`` exceeds ``block_size``; those in-between rows are never + indexed (`triton_convert_req_index_to_global_index` ensures this). + """ + num_blocks, _, head_dim = kv_cache.shape + assert kv_cache.stride(0) % head_dim == 0, ( + "block stride is not a whole number of rows; flat row indexing would " + "silently misaddress" + ) + block_stride_rows = kv_cache.stride(0) // head_dim + num_rows = (num_blocks - 1) * block_stride_rows + block_size + rows = kv_cache.as_strided((num_rows, head_dim), (head_dim, 1)) + return rows, block_stride_rows + + # Kernel with prefill workspace support and valid count tracking @triton.jit def _convert_req_index_to_global_index_kernel( @@ -20,6 +42,7 @@ def _convert_req_index_to_global_index_kernel( # shapes (compile-time where possible) max_num_blocks_per_req: tl.constexpr, BLOCK_SIZE: tl.constexpr, + BLOCK_STRIDE_ROWS: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns HAS_PREFILL: tl.constexpr, COUNT_VALID: tl.constexpr, # whether to count valid indices @@ -86,7 +109,7 @@ def _convert_req_index_to_global_index_kernel( bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 is_invalid_tok |= ~valid_block | is_remote base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0) - out_val = base * BLOCK_SIZE + inblock_off + out_val = base * BLOCK_STRIDE_ROWS + inblock_off # Override with prefill output if prefill is enabled if HAS_PREFILL: @@ -156,6 +179,7 @@ def triton_convert_req_index_to_global_index( block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] BLOCK_SIZE: int = 64, + BLOCK_STRIDE_ROWS: int | None = None, NUM_TOPK_TOKENS: int = 2048, BLOCK_N: int = 128, # tile width along columns HAS_PREFILL_WORKSPACE: bool = False, @@ -249,6 +273,7 @@ def triton_convert_req_index_to_global_index( # shapes / constexprs max_num_blocks_per_req, BLOCK_SIZE, + BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE, block_n, HAS_PREFILL_WORKSPACE, return_valid_counts, @@ -365,6 +390,7 @@ def triton_filter_and_convert_dcp_index( None, max_num_blocks_per_req, BLOCK_SIZE, + BLOCK_SIZE, # dense caches on the DCP path block_n, False, # HAS_PREFILL count_valid, diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 57a628b82b..31c16bd2b6 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -23,7 +23,6 @@ from vllm.v1.attention.backend import ( AttentionType, MultipleOf, ) -from vllm.v1.attention.backends.utils import KVCacheLayoutType if TYPE_CHECKING: from vllm.config import VllmConfig @@ -152,10 +151,6 @@ class TokenspeedMLABackend(MLACommonBackend): ) return None - @classmethod - def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": - return "HND" - class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): can_return_lse_for_decode: bool = True diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index e55076700f..6275778902 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -121,14 +121,6 @@ class TritonMLABackend(MLACommonBackend): return True return block_size % 16 == 0 - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - if include_num_layers_dimension: - return (1, 0, 2, 3) - return (0, 1, 2) - @staticmethod def get_name() -> str: return "TRITON_MLA" diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index d24147bfc3..0f493c74da 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -23,7 +23,8 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, MLAAttentionImpl, ) -from vllm.v1.attention.backends.mla.flashmla_sparse import ( +from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, triton_convert_req_index_to_global_index, ) from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface @@ -66,16 +67,6 @@ class XPUMLASparseBackend(AttentionBackend): def is_sparse(cls) -> bool: return True - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, # assumed to be 1 for MLA - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - return (num_blocks, block_size, head_size) - @classmethod def get_supported_head_sizes(cls) -> list[int]: return [576] @@ -249,16 +240,18 @@ class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]): assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + kv_rows, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) topk_indices_global = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token, attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=attn_metadata.topk_tokens, ) - attn_out = self._forward_bf16_kv( - q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata - ) + attn_out = self._forward_bf16_kv(q, kv_rows, topk_indices_global, attn_metadata) return attn_out, None diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 96dae0c266..0aba6ee8f4 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -31,7 +31,7 @@ from vllm.v1.attention.backends.utils import ( split_decodes_prefills_and_extends, ) from vllm.v1.attention.ops.merge_attn_states import merge_attn_states -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout _PARTITION_SIZE_ROCM = 256 _CP_TOKENS_PER_ITER_ROCM = 32 * 1024 @@ -761,18 +761,22 @@ class AiterFlashAttentionBackend(AttentionBackend): def get_builder_cls() -> type["AiterFlashAttentionMetadataBuilder"]: return AiterFlashAttentionMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """Validate the block size the ROCm gather kernels require.""" + # block_size == 1 is the per-token page-size probe (see + # Platform.get_page_size_bytes); real blocks must be gatherable in + # 16-token units by the ROCm kernel. + if spec.block_size != 1 and spec.block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") - # K and V are packed into the content dim: logical (B, H, N, 2*hs). - return (num_blocks, num_kv_heads, block_size, 2 * head_size) + return spec + + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # K and V come out of the content dim as transposed views rather than + # copies, so the head dim may sit on either side of the block dim, but + # the layer must stay outermost. + return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC) @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 2ce0cf4585..9edae3f58f 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -22,6 +22,7 @@ from vllm.v1.attention.backends.rocm_attn import ( RocmAttentionMetadata, RocmAttentionMetadataBuilder, ) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout logger = init_logger(__name__) @@ -77,18 +78,23 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend): def get_impl_cls() -> type["RocmAiterUnifiedAttentionImpl"]: return RocmAiterUnifiedAttentionImpl - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """Keep K and V packed in the content dim, unlike the native HIP + kernels the base class targets.""" + # block_size == 1 is the per-token page-size probe (see + # Platform.get_page_size_bytes); real blocks must be gatherable in + # 16-token units by the ROCm kernel. + if spec.block_size != 1 and spec.block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") - # K and V are packed into the content dim: logical (B, H, N, 2*hs). - return (num_blocks, num_kv_heads, block_size, 2 * head_size) + return spec + + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # K and V come out of the content dim as transposed views rather than + # copies, so the head dim may sit on either side of the block dim, but + # the layer must stay outermost. + return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC) @staticmethod def use_cascade_attention(*args, **kwargs) -> bool: diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index f73e558ec2..f3bb3b3514 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Attention layer with PagedAttention and Triton prefix prefill.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import ClassVar import torch @@ -16,7 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, ) from vllm.platforms import current_platform -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -35,7 +35,7 @@ from vllm.v1.attention.ops.paged_attn import PagedAttention from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( triton_reshape_and_cache_flash, ) -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout logger = init_logger(__name__) @@ -241,17 +241,30 @@ class RocmAttentionBackend(AttentionBackend): AttentionType.ENCODER_ONLY, ) - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - return (2, num_blocks, block_size, num_kv_heads, head_size) + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """K and V as two head groups so the native HIP kernels address each side + as one contiguous region (x-packed interior applied in split_kv_cache).""" + if spec.state_content_bytes is not None: + return spec + assert spec.head_size == spec.head_size_v, ( + "Separate K/V head groups require symmetric K/V head sizes." + ) + return replace( + spec, + num_head_slots=2, + state_content_bytes=spec.num_kv_heads + * spec.head_size + * get_dtype_size(spec.dtype), + ) + + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # The native HIP kernels hardcode group-contiguous block addressing, so + # they need the K and V groups to span all blocks (H outermost within the layer) + # The Triton fallbacks are stride-aware, so block-interior LBHNC is also + # supported, just without the native kernels. + return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC) @staticmethod def use_cascade_attention(*args, **kwargs) -> bool: @@ -380,8 +393,8 @@ class RocmAttentionImpl(AttentionImpl): query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] - kv_cache: shape = - [2, num_blocks, block_size, num_kv_heads, head_size] + kv_cache: logical [num_blocks, 2, block_size, num_kv_heads * + head_size] under LHBNC (physically K/V-group-first) attn_metadata: Metadata for attention. Returns: shape = [num_tokens, num_heads * head_size] @@ -419,8 +432,10 @@ class RocmAttentionImpl(AttentionImpl): layer, ) + # The bound view is logical [B, 2, N, H*hs]; split_kv_cache expects + # the K/V groups first. key_cache, value_cache = PagedAttention.split_kv_cache( - kv_cache, self.num_kv_heads, self.head_size + kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size ) if is_quantized_kv_cache(self.kv_cache_dtype): @@ -477,7 +492,7 @@ class RocmAttentionImpl(AttentionImpl): if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): return key_cache, value_cache = PagedAttention.split_kv_cache( - kv_cache, self.num_kv_heads, self.head_size + kv_cache.transpose(0, 1), self.num_kv_heads, self.head_size ) # Reshape the input keys and values and store them in the cache. @@ -532,7 +547,7 @@ class RocmAttentionImpl(AttentionImpl): if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): return key_cache, value_cache = PagedAttention.split_kv_cache( - kv_cache, + kv_cache.transpose(0, 1), layer.num_kv_heads, # type: ignore[attr-defined] layer.head_size, # type: ignore[attr-defined] ) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 4b1c167e7f..48dd30d666 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -32,7 +32,6 @@ from vllm.v1.attention.backend import ( ) from vllm.v1.attention.backends.utils import ( compute_mm_prefix_range_tensor, - get_kv_cache_layout, get_num_attention_heads_from_layers, ) from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd @@ -45,7 +44,6 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, KVQuantMode, get_kv_quant_mode, - kv_cache_uses_per_token_head_scales, ) logger = init_logger(__name__) @@ -337,62 +335,6 @@ class TritonAttentionBackend(AttentionBackend): def get_impl_cls() -> type["TritonAttentionImpl"]: return TritonAttentionImpl - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - # K and V are packed into the content dim: logical (B, H, N, 2*hs). - if kv_cache_uses_per_token_head_scales(cache_dtype_str): - # Pad the head dim by sizeof(float32)/sizeof(cache_dtype) so the - # per-(token, head) scale fits inline after the quantized data; - # the backend extracts data[:head_size] and scale[head_size:] via - # typed views (see _ensure_scale_caches). INT4 packs two values - # per byte, so the data occupies only head_size // 2 bytes. - from vllm.utils.torch_utils import ( - STR_DTYPE_TO_TORCH_DTYPE, - get_dtype_size, - ) - - cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype_str] - scale_pad = get_dtype_size(torch.float32) // get_dtype_size(cache_dtype) - if get_kv_quant_mode(cache_dtype_str) == KVQuantMode.INT4_PER_TOKEN_HEAD: - data_head_size = head_size // 2 - else: - data_head_size = head_size - padded_hs = data_head_size + scale_pad - return (num_blocks, num_kv_heads, block_size, 2 * padded_hs) - return (num_blocks, num_kv_heads, block_size, 2 * head_size) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets us from - # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory - # layout we want. - cache_layout = get_kv_cache_layout() - if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size) - return (1, 0, 3, 2, 4) - elif cache_layout == "NHD": - # (num_blocks, block_size, num_kv_heads, 2*head_size) - stride_order = (0, 2, 1, 3) - elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size) - return (1, 2, 0, 3, 4) - elif cache_layout == "HND": - # (num_blocks, num_kv_heads, block_size, 2*head_size) - stride_order = (0, 1, 2, 3) - else: - raise ValueError(f"Unknown cache layout: {cache_layout}") - return stride_order - @staticmethod def use_cascade_attention(*args, **kwargs) -> bool: return False diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py index 02ea8401cb..ddbd066206 100644 --- a/vllm/v1/attention/backends/triton_attn_diffkv.py +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -23,7 +23,6 @@ from vllm.v1.attention.backends.triton_attn import ( TritonAttentionMetadata, TritonAttentionMetadataBuilder, ) -from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( triton_reshape_and_cache_flash_diffkv, ) @@ -93,48 +92,6 @@ class TritonAttentionDiffKVBackend(TritonAttentionBackend): def get_builder_cls() -> type["TritonAttentionDiffKVMetadataBuilder"]: return TritonAttentionDiffKVMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - # Logical (blocks-first, head-major) layout: K and V (with their - # different head sizes) packed in the content dim. - return ( - num_blocks, - num_kv_heads, - block_size, - head_size + TritonAttentionDiffKVBackend.head_size_v, - ) - - @staticmethod - def get_kv_cache_stride_order( - include_num_layers_dimension: bool = False, - ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets us from - # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual - # memory layout we want. - cache_layout = get_kv_cache_layout() - if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v) - return (1, 0, 3, 2, 4) - elif cache_layout == "NHD": - # (num_blocks, block_size, num_kv_heads, C_k+C_v) - return (0, 2, 1, 3) - elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v) - return (1, 2, 0, 3, 4) - elif cache_layout == "HND": - # (num_blocks, num_kv_heads, block_size, C_k+C_v) - return (0, 1, 2, 3) - else: - raise ValueError(f"Unknown cache layout format {cache_layout}.") - @classmethod def supports_head_size(cls, head_size: int) -> bool: # DiffKV K head sizes (e.g. 192 for MiMo-V2.5) need to be allowed. diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py index 5d743be886..8542e6164f 100644 --- a/vllm/v1/attention/backends/turboquant_attn.py +++ b/vllm/v1/attention/backends/turboquant_attn.py @@ -64,7 +64,7 @@ from vllm.v1.attention.ops.triton_turboquant_decode import ( triton_turboquant_decode_attention, ) from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout from vllm.v1.worker.workspace import ( current_workspace_manager, is_workspace_manager_initialized, @@ -156,6 +156,10 @@ class TurboQuantAttentionBackend(AttentionBackend): def get_name() -> str: return "TURBOQUANT" + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + return (KVCacheLayout.LBNHC,) + @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [16, 32, 64, 128] @@ -176,39 +180,6 @@ class TurboQuantAttentionBackend(AttentionBackend): def get_builder_cls() -> type["TurboQuantMetadataBuilder"]: return TurboQuantMetadataBuilder - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "turboquant_4bit_nc", - ) -> tuple[int, ...]: - """Combined K+V cache shape — no leading 2 dimension. - - Standard attention backends use (2, num_blocks, block_size, num_kv_heads, - head_dim) with a leading 2 to separate K and V. TurboQuant packs K+V - into a single interleaved slot per head per position. The logical - (blocks-first, head-major) shape is: - - (num_blocks, num_kv_heads, block_size, slot_size_aligned) - - Each slot = [key_packed | value_packed | padding]. - This is safe because TQ has its own get_kv_cache_shape override and - never shares cache tensors with other backends. Layers that fall back - to native dtype via kv_cache_dtype_skip_layers get their own - standard-shaped cache allocation. - - head_size is the model's real head_dim. slot_size_aligned is computed - from the TQ config to ensure correct cache allocation for all head dims. - """ - from vllm.model_executor.layers.quantization.turboquant.config import ( - TurboQuantConfig, - ) - - tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype_str, head_size) - return (num_blocks, num_kv_heads, block_size, tq_config.slot_size_aligned) - @classmethod def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: if kv_cache_dtype is None: diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 9a5224238a..be543bd04b 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -1,26 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import functools import math -from collections.abc import Callable +from collections import defaultdict +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass, field, fields, make_dataclass from typing import ( TYPE_CHECKING, Any, - Literal, Protocol, - cast, - get_args, ) import numpy as np import torch from typing_extensions import runtime_checkable -from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config import CacheConfig, VllmConfig, get_layers_from_vllm_config +from vllm.config.cache import _layout_from_name from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor -from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec +from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MambaSpec if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -41,8 +39,6 @@ from vllm.v1.attention.backend import ( ) logger = init_logger(__name__) -KVCacheLayoutType = Literal["NHD", "HND"] -_KV_CACHE_LAYOUT_OVERRIDE: KVCacheLayoutType | None = None PAD_SLOT_ID = -1 NULL_BLOCK_ID = 0 @@ -155,44 +151,159 @@ def fill_mm_prefix_query_ranges( return num_actual_tokens -def is_valid_kv_cache_layout(value: str) -> bool: - return value in get_args(KVCacheLayoutType) +_FLASHINFER_LAYOUT_NAMES = { + "LBNHC": "NHD", + "LBHNC": "HND", + "BLHNC": "HND", + "BLNHC": "NHD", + "BHLNC": "HND", +} -@functools.lru_cache -def get_kv_cache_layout(): - # Format specified by the code. - global _KV_CACHE_LAYOUT_OVERRIDE +def get_flashinfer_layout_string(layout: KVCacheLayout) -> str: + """Return the layout name in FlashInfer's convention (NHD/HND).""" + assert layout.name in _FLASHINFER_LAYOUT_NAMES, ( + f"KV cache layout {layout.name} has no FlashInfer equivalent; FlashInfer " + "rejects it in supported_kv_cache_layouts" + ) + return _FLASHINFER_LAYOUT_NAMES[layout.name] - cache_layout: Literal["NHD", "HND"] | None = None - if _KV_CACHE_LAYOUT_OVERRIDE is not None: - cache_layout = _KV_CACHE_LAYOUT_OVERRIDE - logger.debug_once( - "`_KV_CACHE_LAYOUT_OVERRIDE` variable detected. " - "Setting KV cache layout to %s.", - cache_layout, + +# Preference order when no backend declares a supported set; LBNHC (NHD) +# first to match main's default. +_DEFAULT_LAYOUT_PREFERENCE = ( + KVCacheLayout.LBNHC, + KVCacheLayout.LBHNC, + KVCacheLayout.BLNHC, + KVCacheLayout.BLHNC, + KVCacheLayout.BHLNC, + KVCacheLayout.LHBNC, +) + + +def _layout_names(layouts: Iterable[KVCacheLayout]) -> list[str]: + return [layout.name for layout in layouts] + + +def get_supported_kv_cache_layouts( + backends: Iterable[type[AttentionBackend]], +) -> list[KVCacheLayout]: + """Layouts every one of the worker's backends supports, most preferred first. + + Every backend declares the layouts its kernels support, most preferred first + (``supported_kv_cache_layouts``), or None when any layout works; workers where + nothing declares follow the default preference. Identical declarations keep + their order; otherwise the layout the most backends put first wins, ties + keeping the enum order. An empty intersection is a hard error. + """ + supported_layouts_lists: list[Sequence[KVCacheLayout]] = [ + layouts + for backend in backends + if (layouts := backend.supported_kv_cache_layouts()) is not None + ] or [_DEFAULT_LAYOUT_PREFERENCE] + + first = supported_layouts_lists[0] + if all(layouts == first for layouts in supported_layouts_lists[1:]): + return list(first) + + priorities: dict[KVCacheLayout, int] = defaultdict(int) + for preferred_layout, *_ in supported_layouts_lists: + priorities[preferred_layout] += 1 + supported_layouts = set.intersection(*map(set, supported_layouts_lists)) + candidates = sorted( + (layout for layout in KVCacheLayout if layout in supported_layouts), + key=lambda layout: priorities[layout], + reverse=True, + ) + if not candidates: + raise ValueError( + "No KV cache layout satisfies every supported set: " + f"{list(map(_layout_names, supported_layouts_lists))}." ) - return cache_layout + return candidates - # Format specified by the user. - cache_layout = envs.VLLM_KV_CACHE_LAYOUT - # When neither the user nor the override specified a layout, get default - if cache_layout is None: - cache_layout = get_kv_connector_cache_layout() + +def record_kv_cache_layout(cache_config: CacheConfig, layout_name: str) -> None: + """Adopt a layout resolved elsewhere (the engine core) in this process.""" + layout = _layout_from_name(layout_name) + existing = cache_config.kv_cache_layout + if existing is not None and existing != layout.name: + raise ValueError( + f"KV cache layout is already resolved to {existing}; " + f"cannot change it to {layout.name}." + ) + cache_config.kv_cache_layout = layout.name + + +def resolve_kv_cache_layout( + vllm_config: VllmConfig, + supported_layouts: list[list[str]], + kv_cache_specs: Iterable[KVCacheSpec] | None = None, +) -> KVCacheLayout: + """Resolve one KV cache layout for the whole model. + + Runs once in the engine core. Every worker reports the layouts its backends + support, most preferred first (``get_supported_kv_cache_layouts``); all + ranks run the same backends, so their lists must agree. Specs mixing HNC + shapes narrow the candidates to block-compact layouts. An explicit + ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails, + with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the + connector's preference is used when compatible and dropped with a warning + otherwise. A layout already present on ``cache_config`` wins outright, and + the result is recorded there (see ``CacheConfig.kv_cache_layout``); it + reaches workers through the ``set_kv_cache_layout`` RPC and + ``KVCacheConfig.kv_cache_layout``. + """ + cache_config = vllm_config.cache_config + if cache_config.kv_cache_layout is not None: + return cache_config.get_resolved_kv_cache_layout() + + assert supported_layouts and all(supported_layouts), ( + "No worker reported supported KV cache layouts." + ) + assert all(names == supported_layouts[0] for names in supported_layouts[1:]), ( + f"Workers disagree on supported KV cache layouts: {supported_layouts}." + ) + candidates = [_layout_from_name(name) for name in supported_layouts[0]] + + # A block-compact layout means the block is densely packed in memory, so any mix of + # specs can re-interpret HNC with different sizes as long as the total number of + # bytes is the same. If not block-compact, each spec must agree on HNC to alias + # the same page (this aliasing is done by the Hybrid Memory Allocator, HMA). + hnc_shapes = { + (spec.num_heads, spec.num_states, spec.page_size_bytes) + for spec in kv_cache_specs or () + } + if len(hnc_shapes) > 1: + candidates = [m for m in candidates if m.is_block_compact] + if not candidates: + raise ValueError( + "Specs with mixed HNC shapes need a block-compact layout, but " + f"none is in every supported set: {supported_layouts}." + ) + + if (requested := envs.VLLM_KV_CACHE_LAYOUT) is not None: + layout = _layout_from_name(requested) + if layout not in candidates: + raise ValueError( + f"VLLM_KV_CACHE_LAYOUT={requested} does not satisfy every " + f"supported set; valid layouts: {_layout_names(candidates)}." + ) + elif (connector := get_kv_connector_cache_layout(vllm_config)) is not None: + layout = _layout_from_name(connector) + if layout not in candidates: + logger.warning_once( + f"KV connector cache layout {connector} does not satisfy every " + f"supported set; valid layouts: {_layout_names(candidates)}. " + f"Using {candidates[0].name} instead." + ) + layout = candidates[0] else: - assert is_valid_kv_cache_layout(cache_layout) - logger.info_once( - "`VLLM_KV_CACHE_LAYOUT` environment variable " - "detected. Setting KV cache layout to %s.", - cache_layout, - ) - return cache_layout + layout = candidates[0] - -def set_kv_cache_layout(cache_layout: KVCacheLayoutType | None): - global _KV_CACHE_LAYOUT_OVERRIDE - _KV_CACHE_LAYOUT_OVERRIDE = cache_layout - get_kv_cache_layout.cache_clear() + logger.info_once("Using %s KV cache layout.", layout.name) + cache_config.kv_cache_layout = layout.name + return layout @dataclass @@ -264,10 +375,7 @@ def get_num_attention_heads_from_layers( ) if not attn_layers: return None - heads = { - cast(Any, getattr(layer, "impl", layer)).num_heads - for layer in attn_layers.values() - } + heads = {layer.impl.num_heads for layer in attn_layers.values()} assert len(heads) == 1, ( f"All layers in one attention group must share num_heads; " f"got {heads} for {layer_names}." diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index b10cb34977..c9631fd457 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -8,7 +8,6 @@ import torch import vllm.envs as envs from vllm.config.cache import CacheDType -from vllm.logger import init_logger from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.attention.backend import AttentionBackend, AttentionType from vllm.v1.attention.backends.registry import ( @@ -18,8 +17,6 @@ from vllm.v1.attention.backends.registry import ( if TYPE_CHECKING: from vllm.v1.kv_cache_interface import KVCacheSpecKind -logger = init_logger(__name__) - class AttentionSelectorConfig(NamedTuple): head_size: int @@ -185,6 +182,8 @@ def get_attn_backend( ) backend = attention_config.backend_per_kind.get(kind.value, backend) + # The KV cache layout is resolved across all of the model's backends at once + # in get_kv_cache_spec(); a single selection cannot see its peers. return _cached_get_attn_backend( backend=backend, attn_selector_config=attn_selector_config, @@ -210,19 +209,6 @@ def _cached_get_attn_backend( f"Invalid attention backend for {current_platform.device_name}" ) backend = resolve_obj_by_qualname(attention_cls) - - # Adjust kv cache layout if the selected backend requires a specific one - required_layout = backend.get_required_kv_cache_layout() - if required_layout is not None: - from vllm.v1.attention.backends.utils import set_kv_cache_layout - - set_kv_cache_layout(required_layout) - logger.info_once( - "Using %s KV cache layout for %s backend.", - required_layout, - backend.get_name(), - ) - return backend diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 46e9b8e96b..337a55aa6c 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -10,7 +10,7 @@ from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, replace from functools import partial -from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload +from typing import Any, NamedTuple, NewType, TypeAlias, overload from vllm import envs from vllm.config import VllmConfig @@ -26,6 +26,7 @@ from vllm.v1.kv_cache_interface import ( HiddenStateCacheSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheLayout, KVCacheSpec, KVCacheTensor, MambaSpec, @@ -33,6 +34,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowMLASpec, SlidingWindowSpec, UniformTypeKVCacheSpecs, + compute_layout_strides, replace_as, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry @@ -897,7 +899,7 @@ def check_enough_kv_cache_memory( # of the specs since grouping may unify them in-place. groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) check_memory = ( - available_memory - _pool_bytes_per_block(vllm_config, groups) + available_memory - _pool_bytes_per_block(groups) if groups else available_memory ) @@ -999,45 +1001,14 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: return num_blocks -def _pool_bytes_per_block( - vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] -) -> int: +def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: """ Bytes consumed by one block in the worker's shared KV cache pool, mirroring the divisor used by `get_kv_cache_config_from_groups` to convert `available_memory` into `num_blocks`. Used to compute the effective KV cache capacity once `num_gpu_blocks_override` is applied. """ - if len(kv_cache_groups) == 1 and isinstance( - kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs - ): - return kv_cache_groups[0].kv_cache_spec.page_size_bytes - if _use_packed_kv_cache_config(vllm_config, kv_cache_groups): - block_stride, _ = _get_packed_kv_cache_layout(kv_cache_groups) - return block_stride - group_size = max(len(g.layer_names) for g in kv_cache_groups) - page_size = get_uniform_page_size([g.kv_cache_spec for g in kv_cache_groups]) - return page_size * group_size - - -def get_num_blocks( - vllm_config: VllmConfig, - num_layers: int, - available_memory: int, - page_size: int, -) -> int: - """ - Get the number of kv cache blocks. - - Args: - vllm_config: The global VllmConfig - num_layers: The number of layers - available_memory: Memory available for KV cache in bytes. - page_size: The page size of the KV cache. - """ - num_blocks = int(available_memory // page_size // num_layers) - num_blocks = max(num_blocks, 0) - return may_override_num_blocks(vllm_config, num_blocks) + return _get_kv_cache_bytes_per_block(kv_cache_groups) def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: @@ -1092,11 +1063,14 @@ def unify_kv_cache_spec_page_size( size by increasing the block size of layers with smaller page size. Two cases cannot be unified by block size alone and pad their physical page to the maximum instead: Mamba layers, whose page size comes from state shapes - and is independent of block size; and attention layers whose page does not - evenly divide the maximum and whose backend opts in via - ``AttentionSpec.indexes_kv_by_block_stride`` (the padded page is read through - a strided view, which not every backend handles). Raise NotImplementedError - if failed to unify the page size. + and is independent of block size; and non-MLA attention layers whose page + does not evenly divide the maximum (the padded page is read through a + strided view). MLA is excluded because sparse MLA indexes the cache in + whole token rows (see ``flat_kv_row_view``), so its block stride can only + be padded by its own row-aligned ``alignment``, not to an arbitrary page + size. Raise NotImplementedError if failed to unify the page size; + ``get_kv_cache_groups`` catches it to try the full-allocation fallback + (e.g. MLA next to an incompatible sliding-window draft). Args: kv_cache_spec: The KVCacheSpec of each attention layer in the model @@ -1130,18 +1104,15 @@ def unify_kv_cache_spec_page_size( ratio = max_page_size // layer_page_size new_block_size = layer_spec.block_size * ratio new_spec = replace(layer_spec, block_size=new_block_size) - elif ( - isinstance(layer_spec, AttentionSpec) - and layer_spec.indexes_kv_by_block_stride + elif isinstance(layer_spec, AttentionSpec) and not isinstance( + layer_spec, MLAAttentionSpec ): new_spec = replace(layer_spec, page_size_padded=max_page_size) else: raise NotImplementedError( f"Layer {layer_name}: page size is not divisible by the " "maximum page size and cannot be padded. Padding is only " - "supported for attention layers whose backend indexes KV " - "pages by the block stride (indexes_kv_by_block_stride is " - "True)." + "supported for non-MLA attention layers." ) assert new_spec.page_size_bytes == max_page_size new_kv_cache_spec[layer_name] = new_spec @@ -1296,83 +1267,64 @@ def _get_kv_cache_groups_uniform_page_size( return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) -def _get_packed_kv_cache_layout( +def _get_per_layer_spec( + group: KVCacheGroupSpec, + layer_name: str, +) -> KVCacheSpec: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + return spec.kv_cache_specs[layer_name] + return spec + + +def _get_kv_cache_bytes_per_block( kv_cache_groups: list[KVCacheGroupSpec], -) -> tuple[int, dict[int, list[str]]]: - """Lay out each cache group densely in one shared block slab. - - A block ID is owned by one cache group at a time, so layouts from different - groups may overlap. Layers within a group remain disjoint. - """ - layers_by_offset: dict[int, list[str]] = defaultdict(list) - block_stride = 0 - for group in kv_cache_groups: - spec = group.kv_cache_spec - byte_offset = 0 - for layer_name in group.layer_names: - if isinstance(spec, UniformTypeKVCacheSpecs): - page_size = spec.kv_cache_specs[layer_name].page_size_bytes - else: - page_size = spec.page_size_bytes - layers_by_offset[byte_offset].append(layer_name) - byte_offset += page_size - block_stride = max(block_stride, byte_offset) - assert block_stride > 0 - return block_stride, layers_by_offset - - -def _use_packed_kv_cache_config( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> bool: - is_dsv4 = all( - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) +) -> int: + """Return the largest cache group's bytes per block.""" + bytes_per_block = max( + sum( + _get_per_layer_spec(group, layer_name).page_size_bytes + for layer_name in group.layer_names + ) for group in kv_cache_groups ) - kv_transfer_config = vllm_config.kv_transfer_config - extra_config = ( - kv_transfer_config.kv_connector_extra_config - if kv_transfer_config is not None - else {} - ) - # NOTE: enable_cross_layers_blocks is an experimental API and subject to change with - # https://github.com/vllm-project/vllm/issues/42082 - enable_cross_layers = ( - str(extra_config.get("enable_cross_layers_blocks", "False")).lower() == "true" - ) - return is_dsv4 or (enable_cross_layers and len(kv_cache_groups) > 1) + assert bytes_per_block > 0 + return bytes_per_block -def _get_kv_cache_config_packed( - vllm_config: VllmConfig, +def validate_kv_cache_layout( + layout: KVCacheLayout, kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> tuple[int, list[KVCacheTensor]]: - """Plan a packed per-block KV cache tensor layout. +) -> None: + """Validate that the resolved layout can express this model's packing. - Cache groups use dense, overlapping layouts within one block slab. Each - emitted tensor aliases the same physical backing allocation. + The layout was chosen once in the engine core from the backends' supported + sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 + indexer) declares block-outermost layouts there, so an inexpressible + layout reaching this point is an error. """ - block_stride, layers_by_offset = _get_packed_kv_cache_layout(kv_cache_groups) + page_sizes = { + _get_per_layer_spec(group, layer_name).page_size_bytes + for group in kv_cache_groups + for layer_name in group.layer_names + } + if len(page_sizes) == 1: + # A rectangular layer dim exists; every layout can express it. + return - num_blocks = available_memory // block_stride - num_blocks = may_override_num_blocks(vllm_config, num_blocks) - - total_size = block_stride * num_blocks - - kv_cache_tensors: list[KVCacheTensor] = [] - for byte_offset in sorted(layers_by_offset): - kv_cache_tensors.append( - KVCacheTensor( - size=total_size, - shared_by=layers_by_offset[byte_offset], - offset=byte_offset, - block_stride=block_stride, - ) + # Mixed page sizes pack pages side by side within a block, which needs each page + # to be one contiguous chunk inside its block (a block-compact layout) and, with + # multiple KV cache groups, the layer dim inside the block dim. + if not layout.is_block_compact or ( + len(kv_cache_groups) > 1 and layout.is_layer_compact + ): + raise ValueError( + f"KV cache layout {layout.name} cannot express this model's " + f"mixed page sizes ({sorted(page_sizes)}); a backend should " + "declare block-outermost supported layouts (e.g. BLHNC), or " + "set VLLM_KV_CACHE_LAYOUT=BLHNC." ) - return num_blocks, kv_cache_tensors - def get_kv_cache_config_from_groups( vllm_config: VllmConfig, @@ -1402,58 +1354,60 @@ def get_kv_cache_config_from_groups( ), ) - # Determine how model runners should initialize the KV cache tensors. - if len(kv_cache_groups) == 1 and isinstance( - kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs - ): - # Special case: all layers have the same type of KV cache but with - # different hidden sizes. Allocate different amount of memory for each - # layer based on its hidden size. - num_blocks = ( - available_memory // kv_cache_groups[0].kv_cache_spec.page_size_bytes - ) - num_blocks = may_override_num_blocks(vllm_config, num_blocks) - per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs - kv_cache_tensors = [ - KVCacheTensor( - size=per_layer_specs[layer_name].page_size_bytes * num_blocks, - shared_by=[layer_name], - ) - for layer_name in kv_cache_groups[0].layer_names - ] - elif _use_packed_kv_cache_config(vllm_config, kv_cache_groups): - # DeepSeek V4 uses the packed layout by default. Other multi-group - # layouts can opt in with --enable-cross-layers. - num_blocks, kv_cache_tensors = _get_kv_cache_config_packed( - vllm_config, kv_cache_groups, available_memory - ) - else: - # General case: - # We will have group_size memory pools, each is shared by one layer from - # each group. As layers of different groups have different block table, - # they will use different parts of the shared Tensor. - # The memory layout for 3 groups (full.0, full.1), (sw.0, sw.2), - # (sw.1, padding) will be: (group_size = 2) - # full.0, sw.0, sw.1: share a Tensor with size=available_memory//2 - # full.1, sw.2: share another Tensor with size=available_memory//2 - group_size = max(len(group.layer_names) for group in kv_cache_groups) + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + validate_kv_cache_layout(layout, kv_cache_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups) + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None - page_size = get_uniform_page_size( - [group.kv_cache_spec for group in kv_cache_groups] - ) - assert group_size > 0, "group_size must be greater than 0" - num_blocks = get_num_blocks( - vllm_config, group_size, available_memory, page_size - ) - kv_cache_tensors = [] - for i in range(group_size): - shared_by = [] - for j in range(len(kv_cache_groups)): - if i < len(kv_cache_groups[j].layer_names): - shared_by.append(kv_cache_groups[j].layer_names[i]) - kv_cache_tensors.append( - KVCacheTensor(size=page_size * num_blocks, shared_by=shared_by) + num_blocks = available_memory // bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + size = bytes_per_block * num_blocks + + # Groups alias from byte 0. Spec regions are laid out differently: + # + # block-outer (the same packing repeats for every block): + # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... + # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... + # |<--- bytes_per_block -->| + # + # layer-outer (only supported for uniform page sizes or single-group models): + # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | + # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | + + kv_cache_tensors = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, interleaved_block_stride, None, None, None), ) + offset = ( + byte_offset + * max(layer_stride, spec.page_size_bytes) + // spec.page_size_bytes + ) + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=offset, + ) + ) + byte_offset += len(layer_names) * spec.page_size_bytes return KVCacheConfig( num_blocks=num_blocks, @@ -1941,57 +1895,26 @@ def _max_memory_usage_bytes_from_groups( This correctly accounts for padding in hybrid models. For example, if a model has 8 full attention layers and 9 sliding window layers, they will be padded to 9 full + 9 sliding window for uniform group sizes. + + Each group independently claims blocks from the shared pool, so a request consumes + the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``. """ if not kv_cache_groups: return 0 - if len(kv_cache_groups) == 1 and isinstance( - kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs - ): - # UniformTypeKVCacheSpecs special case (single group, per-layer specs) - per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs - return sum( - spec.max_memory_usage_bytes(vllm_config) - for spec in per_layer_specs.values() - ) - elif all( - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - for group in kv_cache_groups - ): - # Special case (only DeepseekV4 for now): all groups are - # UniformTypeKVCacheSpecs. - # They must already be page_size aligned and share a common padded - # layer-tuple layout. Even groups with fewer actual tuples still reserve - # the global number of tuple slots in the shared tensor layout. - full_mla_spec = cast(UniformTypeKVCacheSpecs, kv_cache_groups[0].kv_cache_spec) - layer_tuple_bytes = sum(full_mla_spec.get_page_sizes()) - num_layer_tuples = max( - cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).get_num_layer_tuples() - for group in kv_cache_groups - ) - - total_max_mem_usage_bytes = 0 - for group in kv_cache_groups: - group_spec = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec) - g_max_mem_usage_pages = group_spec.max_memory_usage_pages(vllm_config) - g_max_mem_usage_page_bytes = ( - num_layer_tuples * g_max_mem_usage_pages * layer_tuple_bytes + bytes_per_block = _pool_bytes_per_block(kv_cache_groups) + total_blocks = 0 + for group in kv_cache_groups: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + total_blocks += spec.max_memory_usage_pages(vllm_config) + else: + total_blocks += cdiv( + spec.max_memory_usage_bytes(vllm_config), + spec.page_size_bytes, ) - total_max_mem_usage_bytes += g_max_mem_usage_page_bytes - return total_max_mem_usage_bytes - # General case: group_size pools, each shared by one layer per group - # Memory = group_size * page_size * blocks_for_max_len - group_size = max(len(group.layer_names) for group in kv_cache_groups) - page_size = get_uniform_page_size( - [group.kv_cache_spec for group in kv_cache_groups] - ) - blocks_needed = sum( - cdiv(group.kv_cache_spec.max_memory_usage_bytes(vllm_config), page_size) - for group in kv_cache_groups - ) - - return group_size * page_size * blocks_needed + return bytes_per_block * total_blocks def _estimate_max_model_len_from_groups( @@ -2228,7 +2151,7 @@ def get_kv_cache_configs( if not groups: adjusted_memory.append(avail_mem) continue - bytes_per_block = _pool_bytes_per_block(vllm_config, groups) + bytes_per_block = _pool_bytes_per_block(groups) logger.info( "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", avail_mem // bytes_per_block, @@ -2241,7 +2164,7 @@ def get_kv_cache_configs( # the capacity check both plan against usable blocks. Allocation below # still uses the full memory. check_memory = [ - avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem + avail_mem - _pool_bytes_per_block(groups) if groups else avail_mem for groups, avail_mem in zip(projected_groups_per_worker, available_memory) ] @@ -2278,14 +2201,15 @@ def get_kv_cache_configs( min_num_blocks = min( kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs ) - for kv_cache_config in kv_cache_configs: - num_blocks_old = kv_cache_config.num_blocks - kv_cache_config.num_blocks = min_num_blocks - - # Shrink tensor size proportionally - for tensor in kv_cache_config.kv_cache_tensors: - assert tensor.size % num_blocks_old == 0 - tensor.size = tensor.size // num_blocks_old * min_num_blocks + for i, kv_cache_config in enumerate(kv_cache_configs): + if kv_cache_config.num_blocks == min_num_blocks: + continue + # Re-plan with exactly the memory the smallest rank can afford, so + # strides and offsets stay consistent with the shrunken allocation. + groups = kv_cache_config.kv_cache_groups + kv_cache_configs[i] = get_kv_cache_config_from_groups( + vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups) + ) return kv_cache_configs diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 96010465f3..77de51a69a 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -44,6 +44,7 @@ from vllm.utils.gc_utils import ( from vllm.utils.hashing import get_hash_fn_by_name from vllm.utils.network_utils import make_zmq_socket from vllm.utils.system_utils import decorate_logs, set_process_title +from vllm.v1.attention.backends.utils import resolve_kv_cache_layout from vllm.v1.core.kv_cache_utils import ( BlockHash, generate_scheduler_kv_cache_config, @@ -281,6 +282,17 @@ class EngineCore: ) vllm_config.cache_config.enable_prefix_caching = False + # Resolve the KV cache layout before memory profiling: workers that + # capture full cudagraphs initialize a minimal KV cache during it. + # Attention-free models resolve the default so layout reads never precede + # resolution. + layout = resolve_kv_cache_layout( + vllm_config, + self.model_executor.get_supported_kv_cache_layouts(), + [s for specs in kv_cache_specs for s in specs.values()], + ) + self.model_executor.set_kv_cache_layout(layout.name) + has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs) if has_kv_cache: if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: @@ -307,6 +319,8 @@ class EngineCore: kv_cache_configs = get_kv_cache_configs( vllm_config, kv_cache_specs, available_gpu_memory ) + for kv_cache_config in kv_cache_configs: + kv_cache_config.kv_cache_layout = vllm_config.cache_config.kv_cache_layout # If auto-fit reduced max_model_len, sync the new value to workers. # This is needed because workers were spawned before memory profiling diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 6fff5f032d..964f16ed24 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -151,6 +151,14 @@ class Executor(ABC): def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]: return self.collective_rpc("get_kv_cache_spec") + def get_supported_kv_cache_layouts(self) -> list[list[str]]: + """Layouts each worker's backends support, most preferred first.""" + return self.collective_rpc("get_supported_kv_cache_layouts") + + def set_kv_cache_layout(self, layout_name: str) -> None: + """Publish the resolved KV cache layout to the workers.""" + self.collective_rpc("set_kv_cache_layout", args=(layout_name,)) + @overload def collective_rpc( self, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index cbead6d3c8..efc284bcbb 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -8,6 +8,7 @@ from collections import Counter from collections.abc import Collection from dataclasses import dataclass, fields, replace from enum import Enum, IntEnum +from fractions import Fraction from math import prod from typing import TYPE_CHECKING, TypeVar @@ -18,6 +19,7 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_layout import _DIM_B, _DIM_L, KVCacheLayout from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry if TYPE_CHECKING: @@ -147,6 +149,18 @@ class KVCacheSpec: # number of tokens in a block block_size: int + @property + def num_heads(self) -> int: + raise NotImplementedError + + @property + def tokens_per_state(self) -> int | Fraction: + raise NotImplementedError + + @property + def state_content_size_bytes(self) -> int: + raise NotImplementedError + @property def page_size_bytes(self) -> int: """ @@ -158,8 +172,13 @@ class KVCacheSpec: raise NotImplementedError @property - def storage_block_size(self) -> int: - return self.block_size + def num_states(self) -> int: + return self.get_num_kernel_states(self.block_size) + + def get_num_kernel_states(self, kernel_block_size: int) -> int: + if self.tokens_per_state > 0: + return kernel_block_size // self.tokens_per_state + return 1 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: """ @@ -214,6 +233,143 @@ class KVCacheSpec: ) +def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor: + """View a kernel-block-granular layer cache with manager blocks as dim 0. + + Kernel block splitting subdivides each manager block into uniformly strided + kernel blocks, so grouping is a pure view: ``(num_blocks * ratio, ...)`` + """ + if cache.shape[0] == num_blocks: + return cache + assert cache.shape[0] % num_blocks == 0 + return cache.unflatten(0, (num_blocks, -1)) + + +def compute_layer_kv_cache_shape_bytes( + spec: KVCacheSpec, + num_blocks: int, + kernel_block_size: int | None = None, +) -> tuple[int, ...]: + """Return the 4D logical shape ``(B, H, N, C)`` where C is in bytes.""" + bs = kernel_block_size if kernel_block_size is not None else spec.block_size + assert spec.block_size % bs == 0, ( + f"Kernel block size {bs} must divide KV cache block size {spec.block_size}." + ) + blocks_per_page = spec.block_size // bs + return ( + num_blocks * blocks_per_page, + spec.num_heads, + spec.get_num_kernel_states(bs), + spec.state_content_size_bytes, + ) + + +def compute_layout_strides( + spec: KVCacheSpec, + num_blocks: int, + num_layers: int, + layout: KVCacheLayout, + kernel_block_size: int | None = None, + fixed_strides: tuple[int | None, ...] = (None,) * 5, +) -> tuple[int, ...]: + """Byte strides in logical ``[L, B, H, N, C]`` axis order.""" + assert len(fixed_strides) == 5 + assert all(stride is None or stride > 0 for stride in fixed_strides) + shape = ( + num_layers, + *compute_layer_kv_cache_shape_bytes(spec, num_blocks, kernel_block_size), + ) + order = layout.stride_order + padded_page_size = getattr(spec, "page_size_padded", None) + if padded_page_size is not None: + assert kernel_block_size is None or kernel_block_size == spec.block_size, ( + "Padded KV pages do not support kernel block splitting." + ) + page_grid_end = max(order.index(_DIM_L), order.index(_DIM_B)) + 1 + page_grid_shape = tuple(shape[dim] for dim in order[:page_grid_end]) + assert prod(page_grid_shape) == num_layers * num_blocks, ( + "Page padding requires dimensions outside the page tail to be L, B, " + f"or singleton; got {layout.name} with shape {shape}." + ) + + strides = [0] * 5 + current_stride = 1 + for physical_idx, dim in reversed(tuple(enumerate(order))): + if padded_page_size is not None and physical_idx == page_grid_end - 1: + current_stride = max(current_stride, padded_page_size) + assert current_stride % padded_page_size == 0 + strides[dim] = fixed_strides[dim] or current_stride + current_stride = strides[dim] * shape[dim] + return tuple(strides) + + +def create_kv_cache_views( + raw: torch.Tensor, + spec: KVCacheSpec, + num_blocks: int, + layout: KVCacheLayout, + kv_cache_tensor: KVCacheTensor, + kernel_block_size: int | None = None, +) -> list[torch.Tensor]: + """View a flat int8 buffer as one 4D ``[B, H, N, C]`` view per layer. + + Block ``b`` of layer ``l`` starts at the tensor offset plus its layer and + block stride contributions. + """ + num_layers = len(kv_cache_tensor.layers) + layer_stride = kv_cache_tensor.layer_stride + block_stride = kv_cache_tensor.block_stride + shape_bytes = compute_layer_kv_cache_shape_bytes( + spec, num_blocks, kernel_block_size + ) + ratio = shape_bytes[0] // num_blocks + if ratio > 1: + # Kernel blocks subdivide a manager block into `ratio` equal pieces, so + # they sit a constant stride apart only if a block is one dense page: no + # padding at its end, and no other layer's page before the next block. + dense_page_size = prod(compute_layer_kv_cache_shape_bytes(spec, 1)[1:]) + if block_stride != dense_page_size: + raise ValueError( + f"The resolved KV cache layout ({layout.name}) does not store " + "blocks as dense, unpadded pages (block stride " + f"{block_stride} != page {dense_page_size}), so a manager " + f"block cannot be split into {ratio} kernel blocks of " + f"{kernel_block_size} tokens. Reduce --block-size to " + f"{kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a " + "layer-compact layout (e.g. LBNHC)." + ) + assert block_stride % ratio == 0, ( + f"Block stride {block_stride} must divide into {ratio} equal kernel blocks." + ) + block_stride //= ratio + + logical_shape = (num_layers, *shape_bytes) + strides = compute_layout_strides( + spec, + num_blocks, + num_layers, + layout, + kernel_block_size, + fixed_strides=(layer_stride, block_stride, None, None, None), + ) + dtype = getattr(spec, "dtype", None) + + view_5d = torch.as_strided( + raw, + size=logical_shape, + stride=strides, + storage_offset=raw.storage_offset() + kv_cache_tensor.offset, + ) + + views = [] + for layer_idx in range(num_layers): + cache_logical = view_5d[layer_idx] + if dtype is not None: + cache_logical = cache_logical.view(dtype) + views.append(cache_logical) + return views + + @dataclass(frozen=True, kw_only=True) class AttentionSpec(KVCacheSpec): num_kv_heads: int @@ -222,13 +378,16 @@ class AttentionSpec(KVCacheSpec): head_size_v: int = None # type: ignore[assignment] kv_quant_mode: KVQuantMode = KVQuantMode.NONE page_size_padded: int | None = None - indexes_kv_by_block_stride: bool = False num_head_slots: int | None = None """H of the logical ``[B, H, N, C]`` page when packing diverges from one slot per KV head. None means one slot per KV head. Published by the backend. """ state_content_bytes: int | None = None """C in bytes when packed; None means dense K/V content.""" + tokens_per_state: int | Fraction = 1 + """Tokens covered by one stored state. Ints > 1 compress multiple tokens + into one state (DSv4 sparse MLA); fractions < 1 store multiple states per + token (Whisper block pooling: ``Fraction(1, block_pool_size)``).""" def __post_init__(self): if self.head_size_v is None: @@ -249,7 +408,7 @@ class AttentionSpec(KVCacheSpec): @property def unpadded_page_size_bytes(self) -> int: - return self.num_heads * self.storage_block_size * self.state_content_size_bytes + return self.num_heads * self.num_states * self.state_content_size_bytes @property def page_size_bytes(self) -> int: @@ -345,9 +504,9 @@ class FullAttentionSpec(AttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, - indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, + tokens_per_state=specs[0].tokens_per_state, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), # If any layer in the group is non-causal, treat the group as @@ -384,7 +543,6 @@ class MLAAttentionSpec(FullAttentionSpec): cache_dtype_str: str | None = None # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults. alignment: int | None = None # Default to None for no padding. - compress_ratio: int = 1 # Default to 1 for no compression. model_version: str | None = None # Marks draft groups that flatten a non-causal query block into decode rows. non_causal_multi_token_decode: bool = False @@ -395,28 +553,21 @@ class MLAAttentionSpec(FullAttentionSpec): super().__post_init__() _apply_alignment_padding(self) - @property - def storage_block_size(self) -> int: - return self.block_size // self.compress_ratio - @classmethod def merge(cls, specs: list[Self]) -> Self: assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), ( "All attention layers in the same KV cache group must be MLAAttentionSpec." ) cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) - compress_ratio_set = set(spec.compress_ratio for spec in specs) + tokens_per_state_set = set(spec.tokens_per_state for spec in specs) model_version_set = set(spec.model_version for spec in specs) - block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs) assert ( len(cache_dtype_str_set) == 1 - and len(compress_ratio_set) == 1 + and len(tokens_per_state_set) == 1 and len(model_version_set) == 1 - and len(block_stride_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " - "quantization method, compress ratio, model version, and KV block " - "stride indexing." + "quantization method, tokens per state, and model version." ) merged_spec = cls( block_size=specs[0].block_size, @@ -427,9 +578,8 @@ class MLAAttentionSpec(FullAttentionSpec): page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, - indexes_kv_by_block_stride=block_stride_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), - compress_ratio=compress_ratio_set.pop(), + tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs @@ -483,9 +633,9 @@ class RSWASpec(FullAttentionSpec): dtype=base.dtype, kv_quant_mode=base.kv_quant_mode, page_size_padded=base.page_size_padded, - indexes_kv_by_block_stride=base.indexes_kv_by_block_stride, num_head_slots=base.num_head_slots, state_content_bytes=base.state_content_bytes, + tokens_per_state=base.tokens_per_state, sliding_window=base.sliding_window, attention_chunk_size=base.attention_chunk_size, non_causal=base.non_causal, @@ -598,8 +748,8 @@ class SlidingWindowMLASpec(SlidingWindowSpec): cache_dtype_str: str | None = None # DeepseekV4-only: see MLAAttentionSpec.model_version. alignment: int | None = None # Default to None for no padding. - compress_ratio: int = 1 model_version: str | None = None + # MLA stores a single latent vector per state; there is no separate V. head_size_v: int = 0 @@ -610,10 +760,6 @@ class SlidingWindowMLASpec(SlidingWindowSpec): super().__post_init__() _apply_alignment_padding(self) - @property - def storage_block_size(self) -> int: - return self.block_size // self.compress_ratio - @classmethod def merge(cls, specs: list[Self]) -> Self: assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), ( @@ -621,22 +767,20 @@ class SlidingWindowMLASpec(SlidingWindowSpec): "SlidingWindowMLASpec." ) cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) - compress_ratio_set = set(spec.compress_ratio for spec in specs) + tokens_per_state_set = set(spec.tokens_per_state for spec in specs) model_version_set = set(spec.model_version for spec in specs) sliding_window_set = set(spec.sliding_window for spec in specs) - block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs) extra_retained_set = set(spec.extra_retained_tokens for spec in specs) assert ( len(cache_dtype_str_set) == 1 - and len(compress_ratio_set) == 1 + and len(tokens_per_state_set) == 1 and len(model_version_set) == 1 and len(sliding_window_set) == 1 - and len(block_stride_set) == 1 and len(extra_retained_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " - "quantization method, compress ratio, model version, sliding " - "window size, KV block stride indexing, and retained token count." + "quantization method, tokens per state, model version, sliding " + "window size, and retained token count." ) return cls( block_size=specs[0].block_size, @@ -646,11 +790,10 @@ class SlidingWindowMLASpec(SlidingWindowSpec): page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, - indexes_kv_by_block_stride=block_stride_set.pop(), sliding_window=sliding_window_set.pop(), extra_retained_tokens=extra_retained_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), - compress_ratio=compress_ratio_set.pop(), + tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), ) @@ -672,6 +815,15 @@ class MambaSpec(KVCacheSpec): mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2 mamba_cache_mode: str = "none" num_speculative_blocks: int = 0 + num_heads: int = 1 + tokens_per_state: int = -1 + + @property + def state_content_size_bytes(self) -> int: + return sum( + prod(shape) * get_dtype_size(dtype) + for (shape, dtype) in zip(self.shapes, self.dtypes) + ) @property def page_size_bytes(self) -> int: @@ -771,7 +923,6 @@ class SinkFullAttentionSpec(FullAttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, - indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, sliding_window=cls.merge_window_sizes(sliding_window), @@ -857,9 +1008,6 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): return None # NOTE: below util functions are only used by DeepseekV4 for now. - def get_page_sizes(self) -> list[int]: - return list(set(spec.page_size_bytes for spec in self.kv_cache_specs.values())) - def get_num_layer_tuples(self) -> int: return Counter( 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: class KVCacheTensor: """ A class for specifying how the workers should initialize the KV cache. + + Placement of a set of same-shaped layers in the KV cache allocation. + Layer ``layers[l]``'s page for block ``b`` starts at + ``offset + l * layer_stride + b * block_stride`` bytes into the backing + allocation of ``size`` bytes. Layer-outermost layouts give each layer a + contiguous region (``layer_stride = page * num_blocks``, + ``block_stride = page``); block-outermost layouts make each block a + block of all layers' pages (``layer_stride = page``, ``block_stride`` = + the packed block). Tensors whose address ranges overlap + alias the same bytes: cache groups overlay each other, which is sound + because a block ID is owned by one group at a time. """ - size: int # size of the KV cache tensor in bytes - shared_by: list[str] # layer names that share the same KV cache tensor - offset: int = 0 # byte offset of this layer within a contiguous block - block_stride: int = 0 # total bytes per block in a packed layout (0 = not packed) + size: int # total size of the backing allocation in bytes + layers: list[str] # layer names in L order + layer_stride: int + block_stride: int + offset: int = 0 # byte offset of layers[0]'s block 0 @dataclass @@ -972,6 +1132,8 @@ class KVCacheConfig: """ prefix_cache_retention_interval: int | None = None """Resolved retention policy for local prefix-cache checkpoints.""" + kv_cache_layout: str | None = None + """The KV cache layout resolved by the engine core, adopted by all workers.""" @property def has_mamba_layers(self) -> bool: diff --git a/vllm/v1/kv_cache_layout.py b/vllm/v1/kv_cache_layout.py new file mode 100644 index 0000000000..7c1decd393 --- /dev/null +++ b/vllm/v1/kv_cache_layout.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Physical KV cache layout descriptor. + +A leaf module so ``vllm.config`` can import the enum without pulling in the +full KV cache interface. +""" + +from enum import Enum + +# Logical dim indices in the 5D stride permutation [L, B, H, N, C] (see: RFC #42082). +_DIM_L, _DIM_B, _DIM_H, _DIM_N, _DIM_C = 0, 1, 2, 3, 4 + + +class KVCacheLayout(Enum): + """Physical layout descriptor for a KV cache group. + + The logical shape is always [L, B, H, N, ] (RFC #42082). + Each member's value is a stride permutation that maps logical axes + to physical (memory) order. + """ + + LBHNC = (0, 1, 2, 3, 4) # [L, B, H, N, C] (identity) + LBNHC = (0, 1, 3, 2, 4) # [L, B, N, H, C] + LHBNC = (0, 2, 1, 3, 4) # [L, H, B, N, C] + BLHNC = (1, 0, 2, 3, 4) # [B, L, H, N, C] + BLNHC = (1, 0, 3, 2, 4) # [B, L, N, H, C] + BHLNC = (1, 2, 0, 3, 4) # [B, H, L, N, C] + + @property + def stride_order(self) -> tuple[int, ...]: + return self.value + + @property + def layer_view_order(self) -> tuple[int, ...]: + """Physical axis order of a logical 4D per-layer cache view.""" + return tuple(i - 1 for i in self.value if i != _DIM_L) + + @property + def is_layer_compact(self) -> bool: + """True when the layer is compact; i.e. the L dimension is outermost.""" + return self.value[_DIM_L] == 0 + + @property + def is_block_contiguous(self) -> bool: + """True when [H, N, C] is contiguous within a block.""" + return self.value[-3:] == (_DIM_H, _DIM_N, _DIM_C) + + @property + def is_block_compact(self) -> bool: + """True when each page's [H, N, C] bytes form one contiguous run; i.e. + the L and B dimensions are outermost.""" + return set(self.value[:2]) == {_DIM_L, _DIM_B} + + @property + def is_block_outermost(self) -> bool: + """True when B is the outermost physical dimension.""" + return self.value[0] == _DIM_B diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 6c81d15b39..90138fc250 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -445,10 +445,8 @@ class CanonicalKVCacheTensor: """ A canonicalized KV cache tensor whose first dimension is num_blocks. - For attention backends where the raw tensor has num_blocks at a - non-leading physical dimension (e.g. FlashAttention's - (2, num_blocks, ...) layout), the tensor is split so that each - resulting CanonicalKVCacheTensor starts with (num_blocks, ...). + With standardized layouts (RFC #42082) num_blocks is always the leading + logical dimension. """ # The KV cache tensor with shape (num_blocks, ...) diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index 26408d2992..637d3bb664 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -84,3 +84,5 @@ class OffloadingConfig: # True when the canonical per-layer host byte layout was requested via # kv_connector_extra_config; certified per-layer at worker registration. canonical_layout: bool = False + # Resolved KVCacheLayout name of the worker KV cache. + kv_cache_layout: str | None = None diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index 01c31551bc..82ebe31e2d 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -94,7 +94,10 @@ class FileMapper: for group in config.groups ] parallel = config.parallel - canonical_format = canonical_format_id() if config.canonical_layout else None + canonical_format = None + if config.canonical_layout: + assert config.kv_cache_layout is not None + canonical_format = canonical_format_id(config.kv_cache_layout) return cls( root_dir=root_dir, model_name=config.model.name, diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 1e4fcae68c..0a78c5e9c6 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -194,22 +194,19 @@ class SimpleCPUOffloadScheduler: assert len(gpu_config.kv_cache_tensors) > 0 - is_packed = any(t.block_stride for t in gpu_config.kv_cache_tensors) - assert not is_packed or all(t.block_stride for t in gpu_config.kv_cache_tensors) - gpu_total_bytes = ( - gpu_config.kv_cache_tensors[0].size - if is_packed - else sum(t.size for t in gpu_config.kv_cache_tensors) - ) + # Every KVCacheTensor describes placement within the same backing allocation, + # so its size is the total GPU KV cache size. + gpu_total_bytes = gpu_config.kv_cache_tensors[0].size num_gpu_blocks = gpu_config.num_blocks num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes) # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally. cpu_tensors = [ KVCacheTensor( size=t.size // num_gpu_blocks * num_cpu_blocks, - shared_by=list(t.shared_by), - offset=t.offset, + layers=list(t.layers), + layer_stride=t.layer_stride, block_stride=t.block_stride, + offset=t.offset, ) for t in gpu_config.kv_cache_tensors ] diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index e519bca261..f2925915d8 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -95,57 +95,33 @@ class SimpleCPUOffloadWorker: logger.warning("No KV caches to offload.") return - # Resolve each entry to a representative tensor for storage - # deduplication. For attention layers the value is already a tensor; - # for Mamba layers it is a list of tensors that all share the same - # underlying raw storage, so we take the first one. - def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: - assert isinstance(v, torch.Tensor | list) - return v if isinstance(v, torch.Tensor) else v[0] - - any_tensor = _repr_tensor(next(iter(kv_caches.values()))) - self.device = any_tensor.device + self.device = next(iter(kv_caches.values())).device assert self.kv_cache_config is not None num_blocks = self.kv_cache_config.num_blocks - # Deduplicate: multiple layers may share the same backing storage. - seen_ptrs: dict[int, tuple[str, torch.Tensor]] = {} - for name, value in kv_caches.items(): - tensor = _repr_tensor(value) - ptr = tensor.untyped_storage().data_ptr() - if ptr not in seen_ptrs: - seen_ptrs[ptr] = (name, tensor) - - # Build [num_blocks, block_bytes] int8 views from each unique - # storage so that stride(0) gives block_bytes for the copy op. - # - # The physical layout varies across attention backends: - # FlashAttn/ROCm: (2, num_blocks, ...) -> K/V outermost, 2 segments - # FlashInfer/MLA: (num_blocks, ...) -> blocks outermost, 1 segment - # We derive page_size_bytes = storage.nbytes() // num_blocks, then - # classify dims: any dim whose byte-stride exceeds page_size_bytes - # must be an outer segment dim (e.g. the K/V dim of size 2). A less - # hacky way is to update the interface with the layout. + # The DMA backend copies whole blocks as base + block_id * stride(0), + # so view each unique allocation as [num_blocks, block_bytes]. unique_gpu_caches: dict[str, torch.Tensor] = {} - for name, tensor in seen_ptrs.values(): + seen: set[tuple[torch.device, int]] = set() + for name, tensor in kv_caches.items(): storage = tensor.untyped_storage() - raw = torch.empty(0, dtype=torch.int8, device=self.device).set_( - storage, 0, (storage.nbytes(),) + key = (tensor.device, storage.data_ptr()) + if key in seen: + continue + seen.add(key) + + physical_per_block, remainder = divmod(tensor.shape[0], num_blocks) + assert remainder == 0, ( + f"KV cache {name!r} has {tensor.shape[0]} physical blocks, which " + f"is not divisible by {num_blocks} scheduler blocks" ) - el = tensor.element_size() - page_size_bytes = storage.nbytes() // num_blocks - outer_dims = [ - d for d in range(tensor.ndim) if tensor.stride(d) * el > page_size_bytes - ] - if not outer_dims: - unique_gpu_caches[name] = raw.view(num_blocks, -1) - else: - seg_stride = tensor.stride(outer_dims[0]) * el - for idx in range(tensor.shape[outer_dims[0]]): - offset = idx * seg_stride - chunk = raw[offset : offset + seg_stride] - unique_gpu_caches[f"{name}.{idx}"] = chunk.view(num_blocks, -1) + block_bytes = tensor.stride(0) * tensor.element_size() * physical_per_block + raw = torch.empty(0, dtype=torch.int8, device=tensor.device).set_(storage) + regions = raw.view(-1, num_blocks, block_bytes) + for idx, region in enumerate(regions): + key_name = name if len(regions) == 1 else f"{name}.{idx}" + unique_gpu_caches[key_name] = region # Compute per-tensor bytes_per_block. Tensors may have different # page_size_bytes (e.g., UniformTypeKVCacheSpecs with varying head_size). diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 77d56ad882..f603487ace 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,22 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, replace -from math import prod +from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import Any, cast import torch -from vllm.config import ( - VllmConfig, - get_layers_from_vllm_config, - set_current_vllm_config, -) -from vllm.logger import init_logger +from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.multimodal.inputs import MultiModalFeatureSpec -from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backend import ( AttentionCGSupport, CommonAttentionMetadata, @@ -25,20 +18,17 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, KVCacheSpec, - KVQuantMode, - MambaSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.utils import ( AttentionGroup, add_kv_sharing_layers_to_kv_cache_groups, + allocate_kv_cache, bind_kv_cache, prepare_kernel_block_sizes, ) -logger = init_logger(__name__) - @dataclass(frozen=True) class AttentionCGSupportInfo: @@ -69,13 +59,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(vllm_config): if isinstance(spec, AttentionSpec): - backend = attn_module.get_attn_backend() - # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() -> - # get_kv_cache_layout() needs the current vLLM config. - with set_current_vllm_config(vllm_config): - indexes = backend.indexes_kv_by_block_stride() - spec = replace(spec, indexes_kv_by_block_stride=indexes) - spec = backend.customize_spec(spec) + spec = attn_module.get_attn_backend().customize_spec(spec) kv_cache_spec[layer_name] = spec return kv_cache_spec @@ -197,374 +181,22 @@ def get_query_lens_mismatch_unsupported_backend( return None -def _allocate_kv_cache( - kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device -): - kv_cache_raw_tensors: dict[str, torch.Tensor] = {} - packed_backing: torch.Tensor | None = None - for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - if kv_cache_tensor.block_stride > 0: - # Allocate once; all packed tensors alias the same backing. - if packed_backing is None: - packed_backing = torch.zeros( - kv_cache_tensor.size, dtype=torch.int8, device=device - ) - tensor = packed_backing - else: - tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) - for layer_name in kv_cache_tensor.shared_by: - kv_cache_raw_tensors[layer_name] = tensor - - layer_names = set() - for group in kv_cache_config.kv_cache_groups: - for layer_name in group.layer_names: - layer_names.add(layer_name) - assert layer_names == (kv_cache_raw_tensors.keys() | shared_layers.keys()), ( - "Some layers are not correctly initialized" - ) - return kv_cache_raw_tensors - - -def _kv_first_layers_sharing_pool_with_mamba( - attn_groups: Sequence[AttentionGroup], - kernel_block_sizes: list[int], - cache_dtype: str, - kv_cache_config: "KVCacheConfig | None", -) -> set[str]: - """KV-first attention layers that share an allocation with Mamba state. - - Mamba addresses its state by page, so these layers have to lay their blocks - out page by page as well; otherwise block ids from the two groups resolve - to overlapping bytes. - """ - if kv_cache_config is None: - return set() - - kv_first_layers: set[str] = set() - mamba_layers: set[str] = set() - for group in attn_groups: - kv_cache_spec = group.kv_cache_spec - if isinstance(kv_cache_spec, MambaSpec): - mamba_layers.update(group.layer_names) - continue - if not isinstance(kv_cache_spec, AttentionSpec): - continue - if group.kv_cache_group_id >= len(kernel_block_sizes): - continue - block_dim = group.backend.get_kv_cache_block_dim( - kernel_block_sizes[group.kv_cache_group_id], - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=cache_dtype, - ) - if block_dim != 0: - kv_first_layers.update(group.layer_names) - - page_aligned: set[str] = set() - for kv_tensor in kv_cache_config.kv_cache_tensors: - if mamba_layers.isdisjoint(kv_tensor.shared_by): - continue - page_aligned.update(kv_first_layers.intersection(kv_tensor.shared_by)) - return page_aligned - - -def _reshape_attention_kv_cache( - kv_raw_tensor: torch.Tensor, - kv_cache_spec: AttentionSpec, - kv_cache_shape: tuple[int, ...], - kv_cache_stride_order: tuple[int, ...], - num_blocks: int, - packing: tuple[int, int] | None, - page_aligned_blocks: bool = False, -) -> torch.Tensor: - permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - inv_order = [ - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) - ] - dtype = kv_cache_spec.dtype - - if packing is not None: - offset, block_stride = packing - assert inv_order[0] == 0 - page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) - kv_cache = ( - kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes] - .view(dtype) - .view(permuted_kv_cache_shape) - ) - elif kv_cache_spec.page_size_padded is not None: - # Use a strided view to skip the padding between physical pages. - # - # Only num-blocks-first layouts are supported (the block dimension is - # dim 0 of the unpermuted shape). kv-first layouts such as ROCm's - # ``(2, num_blocks, ...)`` are intentionally not supported here. For a - # num-blocks-first layout the only stride that must change is the block - # stride: every other (contiguous) stride already steps within the - # unpadded region of a page, so no further adjustment is needed. - assert kv_cache_shape[0] == num_blocks, ( - "Padded KV pages require a num-blocks-first KV cache layout (got " - f"shape {kv_cache_shape} with num_blocks={num_blocks}); " - "kv-first layouts are not supported." - ) - dtype_size = get_dtype_size(kv_cache_spec.dtype) - page_stride = kv_cache_spec.page_size_bytes // dtype_size - - num_blocks_dim = inv_order[0] - strides = list(torch.empty(permuted_kv_cache_shape, device="meta").stride()) - strides[num_blocks_dim] = page_stride - - kv_cache = torch.as_strided( - kv_raw_tensor.view(dtype), - size=permuted_kv_cache_shape, - stride=tuple(strides), - ) - elif page_aligned_blocks: - # A KV-first layout such as ROCm's ``(2, num_blocks, ...)`` puts block - # ``b``'s K and V in two far-apart halves of the allocation, so block - # ``b`` does not cover page ``b``. Mamba layers sharing the allocation - # do address their state by page, so the two would resolve the same - # bytes. Build the view page-first instead, then swap the dims back. - assert kv_cache_shape[1] == num_blocks and kv_cache_stride_order == tuple( - range(len(kv_cache_shape)) - ), ( - "Page-aligned KV blocks expect a default-strided (kv, num_blocks, " - f"...) layout, got shape {kv_cache_shape} with stride order " - f"{kv_cache_stride_order} and num_blocks={num_blocks}." - ) - kv_cache = ( - kv_raw_tensor.view(dtype) - .view(num_blocks, kv_cache_shape[0], *kv_cache_shape[2:]) - .transpose(0, 1) - ) - else: - # No padding — safe to use a contiguous view. - kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape) - - return kv_cache.permute(*inv_order) - - -def _reshape_kv_cache( - attn_groups: Sequence[AttentionGroup], - kv_cache_raw_tensors: dict[str, torch.Tensor], - cache_dtype: str, - kernel_block_sizes: list[int], - shared_kv_cache_layers: dict[str, str], - kv_cache_config: "KVCacheConfig | None" = None, -) -> dict[str, Any]: - kv_caches: dict[str, Any] = {} - has_attn = False - - layer_packing: dict[str, tuple[int, int]] = {} - if kv_cache_config is not None: - for kv_tensor in kv_cache_config.kv_cache_tensors: - if kv_tensor.block_stride > 0: - for ln in kv_tensor.shared_by: - layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride) - - page_aligned_layers = _kv_first_layers_sharing_pool_with_mamba( - attn_groups, kernel_block_sizes, cache_dtype, kv_cache_config - ) - - for group in attn_groups: - if group.kv_cache_group_id >= len(kernel_block_sizes): - continue - - kv_cache_spec = group.kv_cache_spec - if kv_cache_spec.storage_block_size != kv_cache_spec.block_size: - # use storage_block_size as the kernel block size for groups - # that apply a compression on block size (eg. DeepSeek V4). - kernel_block_size = kv_cache_spec.storage_block_size - else: - kernel_block_size = kernel_block_sizes[group.kv_cache_group_id] - - for layer_name in group.layer_names: - if layer_name in shared_kv_cache_layers: - # Shared layer — tensor will be aliased to its target later. - continue - - kv_raw_tensor = kv_cache_raw_tensors[layer_name] - packing = layer_packing.get(layer_name) - if packing is not None: - _, blk_stride = packing - num_blocks = kv_raw_tensor.numel() // blk_stride - else: - assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 - num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes - - if isinstance(kv_cache_spec, AttentionSpec): - has_attn = True - # Use storage_block_size: it equals block_size for uncompressed - # specs but is smaller for compressed ones (DeepSeek V4), which - # store block_size tokens in block_size // compress_ratio slots. - num_blocks_per_kv_block = ( - kv_cache_spec.storage_block_size // kernel_block_size - ) - kernel_num_blocks = num_blocks * num_blocks_per_kv_block - # Skipped layers (--kv-cache-dtype-skip-layers) keep the - # unquantized shape; only the quantized primary uses the - # quantized cache dtype's (possibly packed) layout. - layer_cache_dtype = ( - "auto" - if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE - else cache_dtype - ) - kv_cache_shape = group.backend.get_kv_cache_shape( - kernel_num_blocks, - kernel_block_size, - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=layer_cache_dtype, - ) - - # FIXME(woosuk): Add kv_cache_stride_order to all attention backends. - try: - kv_cache_stride_order = group.backend.get_kv_cache_stride_order() - assert len(kv_cache_stride_order) == len(kv_cache_shape) - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - - kv_caches[layer_name] = _reshape_attention_kv_cache( - kv_raw_tensor, - kv_cache_spec, - kv_cache_shape, - kv_cache_stride_order, - kernel_num_blocks, - packing, - page_aligned_blocks=layer_name in page_aligned_layers, - ) - - elif isinstance(kv_cache_spec, MambaSpec): - page_size_bytes = kv_cache_spec.page_size_bytes - # Hold a single contiguous [num_blocks, 1, 1, page_size_bytes] - # int8 page view per layer; the layer's bind_kv_cache unpacks - # each block's bytes into its conv/ssm state views. Keeping - # one tensor per layer lets the KV connector register it - # without special-casing Mamba. - kv_caches[layer_name] = kv_raw_tensor[ - : num_blocks * page_size_bytes - ].view(num_blocks, 1, 1, page_size_bytes) - else: - raise NotImplementedError( - f"Unsupported KV cache spec type: {type(kv_cache_spec)}" - ) - - if has_attn and kv_cache_config is not None: - _align_mixed_attention_kv_cache_views( - attn_groups=attn_groups, - kv_caches=kv_caches, - kernel_block_sizes=kernel_block_sizes, - cache_dtype=cache_dtype, - kv_cache_config=kv_cache_config, - page_aligned_layers=page_aligned_layers, - ) - - # Map any sharing layers to their target layer's KV cache. - for layer_name, target_layer_name in shared_kv_cache_layers.items(): - kv_caches[layer_name] = kv_caches[target_layer_name] - - return kv_caches - - -def _align_mixed_attention_kv_cache_views( - attn_groups: Iterable[AttentionGroup], - kv_caches: dict[str, Any], - kernel_block_sizes: list[int], - cache_dtype: str, - kv_cache_config: KVCacheConfig, - page_aligned_layers: set[str], -) -> None: - """Align shared attention KV views when backends disagree on layout. - - Encoder-decoder models can share one raw allocation between decoder - self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention - (blocks-first backends, block dim 0). Keep the physical storage in the - K/V-first layout expected by ROCM_ATTN, and restride the blocks-first - logical views so block IDs address the same bytes. - """ - block_dims_by_layer: dict[str, int] = {} - for group in attn_groups: - kv_cache_spec = group.kv_cache_spec - if not isinstance(kv_cache_spec, AttentionSpec): - continue - if group.kv_cache_group_id >= len(kernel_block_sizes): - continue - block_dim = group.backend.get_kv_cache_block_dim( - kernel_block_sizes[group.kv_cache_group_id], - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=cache_dtype, - ) - for layer_name in group.layer_names: - if layer_name in kv_caches: - block_dims_by_layer[layer_name] = block_dim - - for kv_tensor in kv_cache_config.kv_cache_tensors: - if kv_tensor.block_stride > 0: - continue - if not page_aligned_layers.isdisjoint(kv_tensor.shared_by): - # These KV-first views already address blocks by page, just like - # blocks-first views, so they must not be restrided onto the - # unaligned KV-first storage layout. - continue - shared_block_dims = { - block_dims_by_layer[layer_name] - for layer_name in kv_tensor.shared_by - if layer_name in block_dims_by_layer - } - if 0 not in shared_block_dims or 1 not in shared_block_dims: - continue - - for layer_name in kv_tensor.shared_by: - if block_dims_by_layer.get(layer_name) == 0: - _restride_blocks_first_kv_cache_to_kv_first_storage( - kv_caches[layer_name] - ) - - -def _restride_blocks_first_kv_cache_to_kv_first_storage( - kv_cache: torch.Tensor, -) -> None: - assert kv_cache.ndim >= 3 - assert kv_cache.shape[1] == 2 - page_size = kv_cache.shape[2:].numel() - num_blocks = kv_cache.shape[0] - expected_tail_stride = torch.empty(kv_cache.shape[2:]).stride() - if kv_cache.stride()[2:] != expected_tail_stride: - logger.warning_once( - "Skipping mixed KV-cache layout alignment for a non-NHD " - "blocks-first attention view with stride %s.", - kv_cache.stride(), - ) - return - kv_cache.as_strided_( - size=kv_cache.shape, - stride=(page_size, num_blocks * page_size, *expected_tail_stride), - ) - - def init_kv_cache( runner_kv_caches: list[torch.Tensor | list[torch.Tensor]], forward_context: dict[str, Any], kv_cache_config: KVCacheConfig, - attn_groups: list[list[AttentionGroup]], device: torch.device, - cache_dtype: str, kernel_block_sizes: list[int], vllm_config: VllmConfig, ) -> dict[str, Any]: - shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config) - kv_cache_raw_tensors = _allocate_kv_cache( - kv_cache_config, shared_kv_cache_layers, device - ) - flattened_attn_groups = list(group for groups in attn_groups for group in groups) - kv_caches = _reshape_kv_cache( - attn_groups=flattened_attn_groups, - kv_cache_raw_tensors=kv_cache_raw_tensors, - kernel_block_sizes=kernel_block_sizes, - cache_dtype=cache_dtype, - shared_kv_cache_layers=shared_kv_cache_layers, - kv_cache_config=kv_cache_config, + kv_caches = allocate_kv_cache( + kv_cache_config, + device, + vllm_config.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, ) + for layer_name, target in get_shared_kv_cache_layers(vllm_config).items(): + kv_caches[layer_name] = kv_caches[target] # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per # decoder layer, so a layer name carries two integers (layer + module index). num_attn_module = ( diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index cdacb36e58..317d9fb903 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -51,8 +51,6 @@ class ActiveKVConnector(KVConnector): self.vllm_config = vllm_config self.kv_connector = get_kv_transfer_group() # Register kv caches with KV Connector if applicable. - # TODO: support cross_layers_kv_cache - # (see https://github.com/vllm-project/vllm/pull/27743) self.kv_connector.register_kv_caches(kv_caches_dict) self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 93a8af4f06..6d8a28933f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -618,9 +618,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.kv_caches, self.compilation_config.static_forward_context, self.kv_cache_config, - self.attn_groups, self.device, - self.cache_config.cache_dtype, self.kernel_block_sizes, self.vllm_config, ) @@ -632,7 +630,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.device, attn_groups_iter=(g for groups in self.attn_groups for g in groups), kernel_block_sizes=self.kernel_block_sizes, - cache_dtype=self.cache_config.cache_dtype, static_forward_context=self.compilation_config.static_forward_context, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 684a3d3a44..92837afd68 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import copy from collections.abc import Mapping from typing import Any @@ -104,12 +103,13 @@ class DFlashSpeculator(DraftModelSpeculator): @property def attn_vllm_config(self) -> VllmConfig: # The draft's attention differs from the target's in causality. - config = copy.copy(super().attn_vllm_config) - config.attention_config = replace( - self.vllm_config.attention_config, - use_non_causal=self.requires_non_causal, + return replace( + self.vllm_config, + attention_config=replace( + self.vllm_config.attention_config, + use_non_causal=self.requires_non_causal, + ), ) - return config def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index ae11b2b602..4ef7f5860a 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -42,7 +42,9 @@ from vllm.config.model import PROCESSED_LOGPROBS_MODES from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.eplb.eplb_state import EplbState from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group -from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks +from vllm.distributed.kv_transfer.kv_connector.utils import ( + copy_kv_blocks, +) from vllm.distributed.parallel_state import ( GraphCaptureContext, get_dcp_group, @@ -164,8 +166,6 @@ from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, KVCacheSpec, KVCacheSpecKind, - KVQuantMode, - MambaSpec, SlidingWindowSpec, UniformTypeKVCacheSpecs, get_kv_cache_spec_kind, @@ -221,7 +221,6 @@ from vllm.v1.worker.cp_utils import ( ) from vllm.v1.worker.dp_utils import coordinate_batch_across_dp from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin -from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin @@ -243,6 +242,7 @@ from .utils import ( AttentionGroup, KVBlockZeroer, add_kv_sharing_layers_to_kv_cache_groups, + allocate_kv_cache, bind_kv_cache, copy_kv_cache_blocks_inplace, prepare_kernel_block_sizes, @@ -612,9 +612,6 @@ class GPUModelRunner( # self.model: nn.Module # Set after load_model # Initialize in initialize_kv_cache self.kv_caches: list[torch.Tensor] = [] - # Initialize in initialize_kv_cache_tensors - self.cross_layers_kv_cache: torch.Tensor | None = None - self.cross_layers_attn_backend: type[AttentionBackend] | None = None # indexes: [kv_cache_group_id][attn_group] self.attn_groups: list[list[AttentionGroup]] = [] # self.kv_cache_config: KVCacheConfig @@ -1200,7 +1197,6 @@ class GPUModelRunner( self.device, attn_groups_iter=self._kv_cache_spec_attn_group_iterator(), kernel_block_sizes=self._kernel_block_sizes, - cache_dtype=self.cache_config.cache_dtype, runner_only_attn_layers=self.runner_only_attn_layers, static_forward_context=self.compilation_config.static_forward_context, ) @@ -6707,9 +6703,6 @@ class GPUModelRunner( for i in range(len(self.kv_caches)): self.kv_caches[i] = None # type: ignore self.kv_caches.clear() - if hasattr(self, "cross_layers_kv_cache"): - self.cross_layers_kv_cache = None - self.cross_layers_attn_backend = None if hasattr(self, "attn_groups"): self.attn_groups.clear() if hasattr(self, "kv_cache_config"): @@ -7441,49 +7434,6 @@ class GPUModelRunner( f"!= kv_cache kernel_block_sizes {kernel_block_sizes}" ) - def _allocate_kv_cache_tensors( - self, kv_cache_config: KVCacheConfig - ) -> dict[str, torch.Tensor]: - """ - Initializes the KV cache buffer with the correct size. The buffer needs - to be reshaped to the desired shape before being used by the models. - - Args: - kv_cache_config: The KV cache config - Returns: - dict[str, torch.Tensor]: A map between layer names to their - corresponding memory buffer for KV cache. - """ - kv_cache_raw_tensors: dict[str, torch.Tensor] = {} - packed_backing: torch.Tensor | None = None - for kv_cache_tensor in kv_cache_config.kv_cache_tensors: - if kv_cache_tensor.block_stride > 0: - # Allocate once; all packed tensors alias the same backing. - if packed_backing is None: - packed_backing = torch.zeros( - kv_cache_tensor.size, - dtype=torch.int8, - device=self.device, - ) - tensor = packed_backing - else: - tensor = torch.zeros( - kv_cache_tensor.size, dtype=torch.int8, device=self.device - ) - for layer_name in kv_cache_tensor.shared_by: - kv_cache_raw_tensors[layer_name] = tensor - - layer_names = set() - for group in kv_cache_config.kv_cache_groups: - for layer_name in group.layer_names: - if layer_name in self.runner_only_attn_layers: - continue - layer_names.add(layer_name) - assert layer_names == set(kv_cache_raw_tensors.keys()), ( - "Some layers are not correctly initialized" - ) - return kv_cache_raw_tensors - def _attn_group_iterator(self) -> Iterator[AttentionGroup]: return itertools.chain.from_iterable(self.attn_groups) @@ -7493,183 +7443,6 @@ class GPUModelRunner( for attn_groups in self.attn_groups: yield from attn_groups - def _reshape_kv_cache_tensors( - self, - kv_cache_raw_tensors: dict[str, torch.Tensor], - kernel_block_sizes: list[int], - ) -> dict[str, torch.Tensor]: - """ - Reshape the KV cache tensors to the desired shape and dtype. - - Args: - kv_cache_raw_tensors: The KV cache buffer of each layer, with - correct size but uninitialized shape. - kernel_block_sizes: The kernel block sizes for each KV cache group. - Returns: - Dict[str, torch.Tensor]: A map between layer names to their - corresponding memory buffer for KV cache. - """ - kv_caches: dict[str, torch.Tensor] = {} - has_attn, has_mamba = False, False - - # Map layer names to (offset, block_stride) within the packed - # backing tensor so we can create strided views per layer. - layer_packing: dict[str, tuple[int, int]] = {} - for kv_tensor in self.kv_cache_config.kv_cache_tensors: - if kv_tensor.block_stride > 0: - for ln in kv_tensor.shared_by: - layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride) - for group in self._kv_cache_spec_attn_group_iterator(): - kv_cache_spec = group.kv_cache_spec - attn_backend = group.backend - if group.kv_cache_group_id == len(kernel_block_sizes): - # There may be a last group for layers without kv cache. - continue - kernel_block_size = kernel_block_sizes[group.kv_cache_group_id] - for layer_name in group.layer_names: - if layer_name in self.runner_only_attn_layers: - continue - raw_tensor = kv_cache_raw_tensors[layer_name] - packing = layer_packing.get(layer_name) - if packing is not None: - _, blk_stride = packing - num_blocks = raw_tensor.numel() // blk_stride - else: - assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 - num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes - if isinstance(kv_cache_spec, AttentionSpec): - has_attn = True - num_blocks_per_kv_block = ( - kv_cache_spec.block_size // kernel_block_size - ) - kernel_num_blocks = num_blocks * num_blocks_per_kv_block - - # For MLA with compression, storage_block_size != block_size - if kv_cache_spec.storage_block_size != kv_cache_spec.block_size: - shape_block_size = kv_cache_spec.storage_block_size - else: - shape_block_size = kernel_block_size - - # Skipped layers (--kv-cache-dtype-skip-layers) need - # the unquantized shape. - layer_cache_dtype_str = ( - "auto" - if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE - else getattr( - kv_cache_spec, - "cache_dtype_str", - None, - ) - or self.cache_config.cache_dtype - ) - kv_cache_shape = attn_backend.get_kv_cache_shape( - kernel_num_blocks, - shape_block_size, - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=layer_cache_dtype_str, - ) - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() - assert len(kv_cache_stride_order) == len(kv_cache_shape) - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - raw_tensor = kv_cache_raw_tensors[layer_name] - kv_caches[layer_name] = _reshape_attention_kv_cache( - raw_tensor, - kv_cache_spec, - kv_cache_shape, - kv_cache_stride_order, - kernel_num_blocks, - packing, - ) - - elif isinstance(kv_cache_spec, MambaSpec): - has_mamba = True - raw_tensor = kv_cache_raw_tensors[layer_name] - page_size_bytes = kv_cache_spec.page_size_bytes - # Hold a single contiguous [num_blocks, 1, 1, page_size_bytes] - # int8 page view per layer; the layer's bind_kv_cache unpacks - # each block's bytes into its conv/ssm state views. Keeping - # one tensor per layer lets the KV connector register it - # without special-casing Mamba. - kv_caches[layer_name] = raw_tensor[ - : num_blocks * page_size_bytes - ].view(num_blocks, 1, 1, page_size_bytes) - else: - raise NotImplementedError - - # Reconcile divergent KV layouts to blocks-first. Triggered by hybrid - # attention/mamba models, and by encoder-decoder models whose shared - # decoder/cross-attention allocation mixes K/V-first and blocks-first - # backends (see _has_mixed_attention_kv_layout). - if has_attn and ( - has_mamba or self._has_mixed_attention_kv_layout(kernel_block_sizes) - ): - self._update_hybrid_attention_mamba_layout(kv_caches, kernel_block_sizes) - - return kv_caches - - def _has_mixed_attention_kv_layout(self, kernel_block_sizes: list[int]) -> bool: - """Whether attention groups disagree on the physical KV cache layout. - - Encoder-decoder models (e.g. Whisper) share one raw KV allocation - between a decoder self-attention layer (K/V-first ROCM_ATTN, block dim - 1) and a cross-attention layer (blocks-first, block dim 0). Mixed block - dims mean a block ID maps to different bytes per layer, so the shared - buffer must be normalized to a single (blocks-first) layout. - """ - block_dims: set[int] = set() - for group in self._kv_cache_spec_attn_group_iterator(): - kv_cache_spec = group.kv_cache_spec - if not isinstance(kv_cache_spec, AttentionSpec): - continue - if group.kv_cache_group_id == len(kernel_block_sizes): - continue - block_dims.add( - group.backend.get_kv_cache_block_dim( - kernel_block_sizes[group.kv_cache_group_id], - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=self.cache_config.cache_dtype, - ) - ) - return len(block_dims) > 1 - - def _update_hybrid_attention_mamba_layout( - self, kv_caches: dict[str, torch.Tensor], kernel_block_sizes: list[int] - ) -> None: - """ - Update the layout of attention layers from (2, num_blocks, ...) to - (num_blocks, 2, ...). - - Args: - kv_caches: The KV cache buffer of each layer. - kernel_block_sizes: The kernel block sizes for each KV cache group. - """ - - for group in self._kv_cache_spec_attn_group_iterator(): - kv_cache_spec = group.kv_cache_spec - if not isinstance(kv_cache_spec, AttentionSpec): - continue - block_dim = group.backend.get_kv_cache_block_dim( - kernel_block_sizes[group.kv_cache_group_id], - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=self.cache_config.cache_dtype, - ) - # block_dim: 0 means (num_blocks, 2, ...); 1 means (2, num_blocks, ...). - if block_dim == 0: - continue - assert block_dim == 1 - for layer_name in group.layer_names: - kv_cache = kv_caches[layer_name] - hidden_size = kv_cache.shape[2:].numel() - kv_cache.as_strided_( - size=kv_cache.shape, - stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]), - ) - def initialize_kv_cache_tensors( self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] ) -> dict[str, torch.Tensor]: @@ -7685,29 +7458,12 @@ class GPUModelRunner( corresponding memory buffer for KV cache. """ - # Try creating KV caches optimized for kv-connector transfers - cache_dtype = self.cache_config.cache_dtype - if self.use_uniform_kv_cache(self.attn_groups): - kv_caches, cross_layers_kv_cache, attn_backend = ( - self.allocate_uniform_kv_caches( - kv_cache_config, - self.attn_groups, - cache_dtype, - self.device, - kernel_block_sizes, - ) - ) - self.cross_layers_kv_cache = cross_layers_kv_cache - self.cross_layers_attn_backend = attn_backend - else: - # Fallback to the general case - # Initialize the memory buffer for KV cache - kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config) - - # Change the memory buffer to the desired shape - kv_caches = self._reshape_kv_cache_tensors( - kv_cache_raw_tensors, kernel_block_sizes - ) + kv_caches = allocate_kv_cache( + kv_cache_config, + self.device, + self.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + ) # Set up cross-layer KV cache sharing for layer_name, target_layer_name in self.shared_kv_cache_layers.items(): @@ -7803,13 +7559,7 @@ class GPUModelRunner( if has_kv_transfer_group() and not is_profiling: kv_transfer_group = get_kv_transfer_group() - if self.cross_layers_kv_cache is not None: - assert self.cross_layers_attn_backend is not None - kv_transfer_group.register_cross_layers_kv_cache( - self.cross_layers_kv_cache, self.cross_layers_attn_backend - ) - else: - kv_transfer_group.register_kv_caches(kv_caches) + kv_transfer_group.register_kv_caches(kv_caches) kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks) def get_routed_experts( @@ -7922,13 +7672,7 @@ class GPUModelRunner( # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(self.vllm_config): if isinstance(spec, AttentionSpec): - backend = attn_module.get_attn_backend() - # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() - # -> get_kv_cache_layout() needs the current vLLM config. - with set_current_vllm_config(self.vllm_config): - indexes = backend.indexes_kv_by_block_stride() - spec = replace(spec, indexes_kv_by_block_stride=indexes) - spec = backend.customize_spec(spec) + spec = attn_module.get_attn_backend().customize_spec(spec) kv_cache_spec[layer_name] = spec return kv_cache_spec diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 958972f3c5..c35e22643f 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -68,6 +68,7 @@ from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime +from vllm.v1.attention.backends.utils import record_kv_cache_layout from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.outputs import ( @@ -669,6 +670,11 @@ class Worker(WorkerBase): # so that it's available to the warmup stage. self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks + # Adopt the engine core's layout; workers spawned after resolution + # (e.g. elastic EP scale-up) only see it through the config. + if kv_cache_config.kv_cache_layout is not None: + record_kv_cache_layout(self.cache_config, kv_cache_config.kv_cache_layout) + # Init kv cache connector here, because it requires # `kv_cache_config`. # NOTE(Kuntai): This need to be done before `initialize_kv_cache`, diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 7a640b875f..ceae77e2c0 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -8,27 +8,18 @@ from collections.abc import Generator from contextlib import AbstractContextManager, contextmanager, nullcontext from typing import TYPE_CHECKING -import torch - from vllm.config import VllmConfig -from vllm.config.cache import CacheDType from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase from vllm.forward_context import get_forward_context, set_forward_context -from vllm.logger import init_logger -from vllm.v1.attention.backend import AttentionBackend -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig from vllm.v1.outputs import ( KVConnectorOutput, ModelRunnerOutput, ) -from vllm.v1.worker.utils import AttentionGroup if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput -logger = init_logger(__name__) - # Defined as a kv connector functionality mixin for ModelRunner (GPU, TPU) class KVConnectorModelRunnerMixin: @@ -110,146 +101,3 @@ class KVConnectorModelRunnerMixin: if not defer_finalize: kv_connector.clear_connector_metadata() - - @staticmethod - def use_uniform_kv_cache( - attn_groups: list[list[AttentionGroup]], - ) -> bool: - """ - Determines whether a uniform KV layout should be used. - A uniform layout means all layers KV caches will share the same - underlying tensor, where for a given block number, the respective - KV data for all layers will be contiguous. - This will allow efficient KV transfer of per-block KV data for all - layers at once. - Note this layout will only be applied given 3 conditions: - 1. The KV Cache config contains just a single group where all layers - have the same page size. - 2. A KV connector is configured, and the KV connector instance prefers - to use this layout (prefer_cross_layer_blocks() returns True) - 3. The attention backend indexes KV by the block stride - (kv_cache_spec.indexes_kv_by_block_stride), i.e. num_blocks is the - outermost physical dim so per-block all-layers data is contiguous. - - Note that the actual placement of the num_layers dimensions - in the unified layers tensors will be determined by the attention - backend. - Thus, the layers KV data may still not be contiguous per block - if the attention backend does not support it. - - Args: - attn_groups: The list of attention groups for this model - Returns: - True if we should use a uniform KV cache layout. - """ - - if not has_kv_transfer_group(): - return False - if not get_kv_transfer_group().prefer_cross_layer_blocks: - return False - - if len(attn_groups) != 1 or len(attn_groups[0]) != 1: - return False - - attn_group = attn_groups[0][0] - kv_cache_spec = attn_group.kv_cache_spec - if not isinstance(kv_cache_spec, AttentionSpec): - return False - # Per-token-head quant carves inline-scale views that assume per-layer - # contiguous KV buffers; the cross-layer layout breaks this and corrupts KV. - if kv_cache_spec.kv_quant_mode.is_per_token_head: - return False - return kv_cache_spec.indexes_kv_by_block_stride - - @staticmethod - def allocate_uniform_kv_caches( - kv_cache_config: KVCacheConfig, - attn_groups: list[list[AttentionGroup]], - cache_dtype: CacheDType, - device: torch.device, - kernel_block_sizes: list[int], - ) -> tuple[dict[str, torch.Tensor], torch.Tensor, type[AttentionBackend]]: - """ - Initializes and reshapes KV caches for the simple case where all - layers have the same layout. - - This function assumes use_uniform_kv_cache() returned True. - - Args: - kv_cache_config: The KV cache config - attn_groups: The list of attention groups for this model - cache_dtype: The KV cache dtype - device: The torch device to allocate on. - kernel_block_sizes: The kernel block sizes for each KV cache group. - Returns: - A tuple (kv_caches, cross_layers_kv_cache, attn_backend) where: - kv_caches is a dict mapping between layer names to their - corresponding memory buffer for KV cache. - cross_layers_kv_cache is the cross layers kv cache tensor - attn_backend is the attention backend matching this tensor - """ - attn_group = attn_groups[0][0] - kv_cache_spec = attn_group.kv_cache_spec - assert isinstance(kv_cache_spec, AttentionSpec) - - tensor_sizes = set( - kv_cache_tensor.size for kv_cache_tensor in kv_cache_config.kv_cache_tensors - ) - assert len(tensor_sizes) == 1 - tensor_size = tensor_sizes.pop() - - page_size = kv_cache_spec.page_size_bytes - assert tensor_size % page_size == 0 - num_blocks = tensor_size // page_size - num_layers = len(kv_cache_config.kv_cache_tensors) - total_size = tensor_size * num_layers - - assert len(kernel_block_sizes) == 1 - kernel_block_size = kernel_block_sizes[0] - num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size - kernel_num_blocks = num_blocks * num_blocks_per_kv_block - - attn_backend = attn_group.backend - kv_cache_shape = attn_backend.get_kv_cache_shape( - kernel_num_blocks, - kernel_block_size, - kv_cache_spec.num_kv_heads, - kv_cache_spec.head_size, - cache_dtype_str=cache_dtype, - ) - - # prepend a num_layers dimension into the shape - kv_cache_shape = (num_layers,) + kv_cache_shape - - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order( - include_num_layers_dimension=True - ) - assert len(kv_cache_stride_order) == len(kv_cache_shape) - except (AttributeError, NotImplementedError): - kv_cache_stride_order = tuple(range(len(kv_cache_shape))) - - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - - logger.info("Allocating a cross layer KV cache of shape %s", kv_cache_shape) - - # allocate one contiguous buffer for all layers - cross_layers_kv_cache = ( - torch.zeros(total_size, dtype=torch.int8, device=device) - .view(kv_cache_spec.dtype) - .view(kv_cache_shape) - ) - - # Maintain original KV shape view. - inv_order = [ - kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) - ] - permuted_kv_cache = cross_layers_kv_cache.permute(*inv_order) - - kv_caches = {} - for i, kv_cache_tensor in enumerate(kv_cache_config.kv_cache_tensors): - tensor = permuted_kv_cache[i] - for layer_name in kv_cache_tensor.shared_by: - kv_caches[layer_name] = tensor - - return kv_caches, cross_layers_kv_cache, attn_backend diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index d8f36b1d15..5360adaf74 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -30,9 +30,11 @@ from vllm.v1.kv_cache_interface import ( EncoderOnlyAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheLayout, KVCacheSpec, MambaSpec, UniformTypeKVCacheSpecs, + create_kv_cache_views, ) from vllm.v1.worker.block_table import get_block_table_width @@ -61,10 +63,10 @@ def _zero_kv_blocks_kernel( ): """Zero KV cache blocks across all segments in a single launch. - Each segment is a contiguous region of one block's data. For backends - where blocks are outermost (block_dim=0) there is one segment per - buffer. For backends where K/V is outermost (block_dim=1) there are - two segments per buffer (one for K, one for V). + Each segment is a contiguous region of one block's data. Layer-compact + layouts have one segment per layer buffer; dimensions physically outside + the block dim (separate head groups under LHBNC) and virtual block splits + each get their own segment. Segments may have different block strides and page sizes (e.g. packed KV views or models with multiple KV cache groups like MLA + DSA @@ -110,7 +112,6 @@ class KVBlockZeroer: device: torch.device, attn_groups_iter: Iterable["AttentionGroup"], kernel_block_sizes: list[int], - cache_dtype: str, static_forward_context: dict[str, Any], runner_only_attn_layers: set[str] | None = None, ) -> None: @@ -119,6 +120,12 @@ class KVBlockZeroer: Each entry is the absolute byte address of a segment start on the GPU, so segments in different CUDA allocations work correctly. + Per-layer views are standardized ``[B, H, N, C]`` with blocks at dim 0; dims + physically outside B (separate head groups under LHBNC) each get their own + segment. A segment's page spans everything inside its block, so under BHLNC it + also covers the block's other layers -- safe, since block IDs are global pool + indices and a newly allocated block owns its whole tile. + Block IDs from the scheduler reference logical blocks whose size may differ from the kernel block size (virtual block splitting). Each virtual block is represented as an independent segment so its @@ -133,7 +140,10 @@ class KVBlockZeroer: if runner_only_attn_layers is None: runner_only_attn_layers = set() - seen_ptrs: set[int] = set() + # Overlaid layers (packed layouts) share a base address but may have + # different page sizes; keep the widest span per address so newly + # allocated blocks are fully zeroed for every overlaying group. + seen_ptrs: dict[int, int] = {} seg_addrs: list[int] = [] seg_block_strides: list[int] = [] seg_page_sizes: list[int] = [] @@ -147,12 +157,6 @@ class KVBlockZeroer: kernel_bs = kernel_block_sizes[group.kv_cache_group_id] assert spec.block_size % kernel_bs == 0 ratio = spec.block_size // kernel_bs - block_dim = group.backend.get_kv_cache_block_dim( - kernel_bs, - spec.num_kv_heads, - spec.head_size, - cache_dtype_str=cache_dtype, - ) for layer_name in group.layer_names: if layer_name in runner_only_attn_layers: @@ -161,23 +165,18 @@ class KVBlockZeroer: if not isinstance(kv, torch.Tensor): continue dp = kv.data_ptr() - if dp in seen_ptrs: - continue - seen_ptrs.add(dp) el = kv.element_size() - block_stride_bytes = kv.stride(block_dim) * el + block_stride_bytes = kv.stride(0) * el assert block_stride_bytes % 4 == 0 - assert kv.shape[block_dim] % ratio == 0 + assert kv.shape[0] % ratio == 0 outer_dims = [ d - for d in range(block_dim) + for d in range(1, kv.ndim) if kv.stride(d) * el > block_stride_bytes ] outer_strides = [kv.stride(d) * el for d in outer_dims] - inner_dims = [ - d for d in range(kv.ndim) if d != block_dim and d not in outer_dims - ] + inner_dims = [d for d in range(1, kv.ndim) if d not in outer_dims] kernel_page_bytes = el + sum( (kv.shape[d] - 1) * kv.stride(d) * el for d in inner_dims ) @@ -187,9 +186,18 @@ class KVBlockZeroer: off_bytes = sum(i * s for i, s in zip(outer, outer_strides)) assert (dp + off_bytes) % 4 == 0 for virtual_index in range(ratio): - seg_addrs.append( - dp + off_bytes + virtual_index * block_stride_bytes - ) + addr = dp + off_bytes + virtual_index * block_stride_bytes + if (idx := seen_ptrs.get(addr)) is not None: + assert ( + seg_block_strides[idx] + == logical_block_stride_bytes // 4 + ) + seg_page_sizes[idx] = max( + seg_page_sizes[idx], kernel_page_bytes // 4 + ) + continue + seen_ptrs[addr] = len(seg_addrs) + seg_addrs.append(addr) seg_block_strides.append(logical_block_stride_bytes // 4) seg_page_sizes.append(kernel_page_bytes // 4) @@ -367,6 +375,54 @@ def select_common_block_size( raise ValueError(f"No common block size for {kv_manager_block_size}. ") +def allocate_kv_cache( + kv_cache_config: KVCacheConfig, + device: torch.device, + layout: KVCacheLayout, + kernel_block_sizes: list[int] | None = None, +) -> dict[str, torch.Tensor]: + """Allocate the KV cache and view it as ``[B, H, N, C]`` per layer. + + Every KVCacheTensor places its layers in the same backing allocation: layer ``l`` of + block ``b`` starts at ``offset + l * layer_stride + b * block_stride``. Cache + groups overlay each other, so tensors may address the same bytes. + """ + if not kv_cache_config.kv_cache_tensors: + return {} + + sizes = {tensor.size for tensor in kv_cache_config.kv_cache_tensors} + assert len(sizes) == 1, "KV cache tensors must share one backing allocation." + buf = torch.zeros(sizes.pop(), dtype=torch.int8, device=device) + + kv_caches: dict[str, torch.Tensor] = {} + for tensor in kv_cache_config.kv_cache_tensors: + layer_name = tensor.layers[0] + group_id, group = next( + (group_id, group) + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if layer_name in group.layer_names + ) + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + spec = spec.kv_cache_specs[layer_name] + + num_blocks = kv_cache_config.num_blocks + kernel_block_size = None + if kernel_block_sizes is not None and group_id < len(kernel_block_sizes): + kernel_block_size = kernel_block_sizes[group_id] + + views = create_kv_cache_views( + buf, + spec, + num_blocks, + layout, + tensor, + kernel_block_size=kernel_block_size, + ) + kv_caches.update(zip(tensor.layers, views)) + return kv_caches + + def prepare_kernel_block_sizes( kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]] ) -> list[int]: @@ -551,42 +607,52 @@ def bind_kv_cache( def copy_kv_cache_blocks_inplace( - kv_caches: Iterable[torch.Tensor | list[torch.Tensor]], + kv_caches: Iterable[torch.Tensor], num_blocks: int, kv_cache_block_copies: Sequence[KVCacheBlockCopy], ) -> None: if not kv_cache_block_copies: return - storage_tensors: list[torch.Tensor] = [] - seen_storage: set[int] = set() - for entry in kv_caches: - # Mamba layers hold a list of state tensors; attention layers a single - # tensor. Both alias the shared block-major backing storage. - tensors = entry if isinstance(entry, (list, tuple)) else (entry,) - for tensor in tensors: - ptr = tensor.untyped_storage().data_ptr() - if ptr in seen_storage: - continue - seen_storage.add(ptr) - storage_tensors.append(tensor) - - if not storage_tensors: - return - device = storage_tensors[0].device indices_np = np.array(kv_cache_block_copies, dtype=np.int64) - indices = async_tensor_h2d(indices_np, device=device) - src_indices, dst_indices = indices.unbind(dim=1) + indices: torch.Tensor | None = None + seen: set[tuple[torch.device, int]] = set() + copied_storages: set[tuple[torch.device, int]] = set() + for cache in kv_caches: + # Layers sharing KV (cross-layer sharing) alias the same view; copy it + # once. data_ptr distinguishes per-layer views of a shared allocation. + key = (cache.device, cache.data_ptr()) + if key in seen: + continue + seen.add(key) - for tensor in storage_tensors: - assert tensor.device == device - blocks = torch.empty(0, dtype=torch.uint8, device=device) - blocks.set_(tensor.untyped_storage()) - # Block-major backing storage: block i owns the contiguous byte range - # [i * page_size, (i + 1) * page_size). - assert blocks.numel() % num_blocks == 0 - blocks = blocks.view(num_blocks, -1) - blocks[dst_indices] = blocks[src_indices] + if indices is None: + indices = async_tensor_h2d(indices_np, device=cache.device) + assert cache.device == indices.device + src, dst = indices.unbind(dim=1) + + kernel_blocks_per_block, remainder = divmod(cache.shape[0], num_blocks) + assert remainder == 0, ( + f"{cache.shape[0]} kernel blocks not divisible by " + f"{num_blocks} scheduler blocks" + ) + storage = cache.untyped_storage() + storage_key = (cache.device, storage.data_ptr()) + scheduler_block_stride = ( + cache.stride(0) * cache.element_size() * kernel_blocks_per_block + ) + if storage.nbytes() == num_blocks * scheduler_block_stride: + if storage_key in copied_storages: + continue + copied_storages.add(storage_key) + blocks = torch.empty(0, dtype=torch.uint8, device=cache.device) + blocks.set_(storage) + blocks = blocks.view(num_blocks, -1) + else: + # Fold virtual block splitting into the shape so that dim 0 counts + # scheduler blocks; unflatten of dim 0 is always a view. + blocks = cache.unflatten(0, (num_blocks, kernel_blocks_per_block)) + blocks[dst] = blocks[src] def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) -> bool: diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 998380a877..7a17e10586 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -9,12 +9,17 @@ import torch.nn as nn import vllm.ir from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed.kv_transfer.kv_connector.utils import get_current_attn_backends from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.tracing import instrument from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.system_utils import update_environment_variables +from vllm.v1.attention.backends.utils import ( + get_supported_kv_cache_layouts, + record_kv_cache_layout, +) from vllm.v1.kv_cache_interface import KVCacheSpec if TYPE_CHECKING: @@ -99,6 +104,15 @@ class WorkerBase: """Get specifications for KV cache implementation.""" raise NotImplementedError + def get_supported_kv_cache_layouts(self) -> list[str]: + """Layout names every attention backend supports, most preferred first.""" + backends = get_current_attn_backends(self.vllm_config) + return [layout.name for layout in get_supported_kv_cache_layouts(backends)] + + def set_kv_cache_layout(self, kv_cache_layout: str) -> None: + """Adopt the KV cache layout resolved by the engine core.""" + record_kv_cache_layout(self.vllm_config.cache_config, kv_cache_layout) + def compile_or_warm_up_model(self) -> CompilationTimes: """Prepare model for execution through compilation/warmup.