diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 79f8937c6379..f4b7ee520ad9 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -1,16 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4.""" +"""Tests for contiguous KV cache packing.""" from unittest.mock import MagicMock import pytest import torch -from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4 +from vllm import envs +from vllm.v1.core.kv_cache_utils import ( + _get_kv_cache_config_deepseek_v4, + get_kv_cache_config_from_groups, +) from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, KVCacheGroupSpec, + KVCacheTensor, MLAAttentionSpec, + SlidingWindowSpec, UniformTypeKVCacheSpecs, ) @@ -28,6 +35,25 @@ def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec: ) +def _make_full_spec() -> FullAttentionSpec: + return FullAttentionSpec( + 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 _make_groups(n_c4, n_c128, n_swa): PS_C4_MLA = 37440 PS_C4_IDX = 8640 @@ -130,6 +156,73 @@ def test_strided_views_are_independent(self): for i, v in enumerate(views): assert (v == i + 1).all(), f"View {i} was corrupted" + def test_hma_attention_groups_keep_default_backing(self, monkeypatch): + monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False) + 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), + ] + + config = get_kv_cache_config_from_groups( + _mock_vllm_config(), groups, available_memory=page_size * 2 * 32 + ) + + assert config.num_blocks == 32 + 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_flag(self, monkeypatch): + monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False) + 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), + ] + + config = get_kv_cache_config_from_groups( + _mock_vllm_config(), groups, available_memory=page_size * 2 * 32 + ) + + 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, + ), + ] + + def test_single_group_attention_keeps_unpacked_layout(self): + spec = _make_full_spec() + groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)] + + config = get_kv_cache_config_from_groups( + _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32 + ) + + 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] + if __name__ == "__main__": pytest.main([__file__, "-v"]) 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 8583bb4b1e09..f22d6738b4fd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -50,7 +50,8 @@ def _register_handlers(self, kv_caches: CanonicalKVCaches): def register_kv_caches( self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]] ): - num_blocks = self.spec.kv_cache_config.num_blocks + kv_cache_config = self.spec.kv_cache_config + num_blocks = kv_cache_config.num_blocks # layer_name -> (num_blocks, page_size_bytes) tensor tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {} @@ -58,7 +59,7 @@ def register_kv_caches( unpadded_page_size_bytes: dict[str, int] = {} # layer_name -> size of page in bytes page_size_bytes: dict[str, int] = {} - for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups: + for kv_cache_group in kv_cache_config.kv_cache_groups: group_layer_names = kv_cache_group.layer_names group_kv_cache_spec = kv_cache_group.kv_cache_spec if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs): @@ -122,9 +123,35 @@ def register_kv_caches( else: raise NotImplementedError + packed_kv_cache_tensor = next( + (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None + ) + is_dsv4 = all( + isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) + for group in kv_cache_config.kv_cache_groups + ) + if packed_kv_cache_tensor is not None and not is_dsv4: + (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]] + block_stride = tensor.stride(0) + packed_tensor = tensor.as_strided( + (num_blocks, block_stride), + (block_stride, 1), + storage_offset=0, + ) + self._register_handlers( + CanonicalKVCaches( + [CanonicalKVCacheTensor(packed_tensor, block_stride)], + [ + [CanonicalKVCacheRef(0, block_stride)] + for _ in kv_cache_config.kv_cache_groups + ], + ) + ) + return + block_tensors: list[CanonicalKVCacheTensor] = [] block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list) - for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors: + for kv_cache_tensor in kv_cache_config.kv_cache_tensors: # Filter to layers that were actually processed above. # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for # every (tuple_idx, page_size) slot; slots where no group has a @@ -166,7 +193,7 @@ def register_kv_caches( ) group_data_refs: list[list[CanonicalKVCacheRef]] = [] - for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups: + for kv_cache_group in kv_cache_config.kv_cache_groups: group_refs: list[CanonicalKVCacheRef] = [] for layer_name in kv_cache_group.layer_names: group_refs += block_data_refs[layer_name] diff --git a/vllm/envs.py b/vllm/envs.py index a94e084ab628..d9b10afba20d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -209,6 +209,7 @@ VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5 VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None + VLLM_USE_PACKED_HMA_KV_CACHE: bool = False VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None VLLM_COMPUTE_NANS_IN_LOGITS: bool = False VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[ @@ -1608,6 +1609,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_KV_CACHE_LAYOUT": env_with_choices( "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"] ), + # Opt into packed per-block KV cache allocation for multi-group + # attention-only HMA models (e.g. gpt-oss, Gemma 3/4). + "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool( + int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0")) + ), # SSM conv state layout used for Mamba models. # - SD: (state_len, dim) — dim contiguous (default) # - DS: (dim, state_len) — TP-sharded dim on dim1, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index a1ebe08c0789..4e1d28d7d5da 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -938,9 +938,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ): return kv_cache_groups[0].kv_cache_spec.page_size_bytes - if all( - isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups - ): + if _use_packed_kv_cache_groups(kv_cache_groups): # buckets = {page_size: [[layer_names], [layer_names], ...]} buckets = _bucket_layers_by_page_size(kv_cache_groups) return sum(ps * len(slots) for ps, slots in buckets.items()) @@ -1218,16 +1216,29 @@ def _bucket_layers_by_page_size( return buckets -def _get_kv_cache_config_deepseek_v4( +def _use_packed_kv_cache_groups( + kv_cache_groups: list[KVCacheGroupSpec], +) -> bool: + is_dsv4 = all( + isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) + for group in kv_cache_groups + ) + return is_dsv4 or ( + bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1 + ) + + +def _get_kv_cache_config_packed( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], available_memory: int, ) -> tuple[int, list[KVCacheTensor]]: - """DeepseekV4 KV cache tensor layout planning. + """Plan a packed per-block KV cache tensor layout. Emit one KVCacheTensor per (slot_idx, page_size). Layers from different groups at the same slot share a tensor (they have independent block - tables so block-id namespaces never collide). + tables so block-id namespaces never collide). Each emitted tensor aliases + one physical backing allocation, with per-block data laid out contiguously. """ # buckets = {page_size: [[layer_names], [layer_names], ...]} buckets = _bucket_layers_by_page_size(kv_cache_groups) @@ -1255,6 +1266,9 @@ def _get_kv_cache_config_deepseek_v4( return num_blocks, kv_cache_tensors +_get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_packed + + def get_kv_cache_config_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], @@ -1299,13 +1313,11 @@ def get_kv_cache_config_from_groups( ) for layer_name in kv_cache_groups[0].layer_names ] - elif all( - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - for group in kv_cache_groups - ): - # DeepseekV4: UniformTypeKVCacheSpecs but multiple groups. - # Delegate to the DeepseekV4-specific allocator. - num_blocks, kv_cache_tensors = _get_kv_cache_config_deepseek_v4( + elif _use_packed_kv_cache_groups(kv_cache_groups): + # DeepSeek V4 keeps the existing packed layout. Other multi-group + # attention-only HMA layouts can opt in with + # VLLM_USE_PACKED_HMA_KV_CACHE=1. + num_blocks, kv_cache_tensors = _get_kv_cache_config_packed( vllm_config, kv_cache_groups, available_memory ) else: diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index d65ba9439e19..b8fb893f14dd 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -58,7 +58,15 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): self.cpu_page_size_per_worker = 0 assert kv_cache_config is not None if kv_cache_config.num_blocks > 0 and world_size > 0: - total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors) + is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors) + assert not is_packed or all( + t.block_stride for t in kv_cache_config.kv_cache_tensors + ) + total_gpu_kv_bytes = ( + kv_cache_config.kv_cache_tensors[0].size + if is_packed + else sum(t.size for t in kv_cache_config.kv_cache_tensors) + ) kv_bytes_per_block = ( total_gpu_kv_bytes // kv_cache_config.num_blocks ) * world_size diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index fe984be96a20..5e431e62388e 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -187,7 +187,13 @@ def _derive_cpu_config( assert len(gpu_config.kv_cache_tensors) > 0 - gpu_total_bytes = sum(t.size for t in gpu_config.kv_cache_tensors) + 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) + ) 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. @@ -195,6 +201,8 @@ def _derive_cpu_config( KVCacheTensor( size=t.size // num_gpu_blocks * num_cpu_blocks, shared_by=list(t.shared_by), + offset=t.offset, + block_stride=t.block_stride, ) for t in gpu_config.kv_cache_tensors ]