diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 7e65d650f7..b446c5a377 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -3,7 +3,13 @@ import torch -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheTensor, + KVQuantMode, + MambaSpec, +) from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache from vllm.v1.worker.utils import AttentionGroup @@ -240,3 +246,99 @@ def test_reshape_padded_quantized_kv_cache_preserves_scale_stride(): 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, + kv_cache_tensors=[ + KVCacheTensor(size=raw_tensor.numel(), shared_by=list(shared_by)) + ], + kv_cache_groups=[], + ) + + kv_caches = _reshape_kv_cache( + attn_groups, + raw_tensors, + "auto", + [attn_spec.block_size] * len(attn_groups), + {}, + kv_cache_config, + ) + return num_blocks, kv_caches["attn"] + + +def test_reshape_kv_first_kv_cache_pages_blocks_when_shared_with_mamba(): + num_blocks, kv_cache = _kv_first_setup(["attn", "mamba"]) + + 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 + + +def test_reshape_kv_first_kv_cache_keeps_layout_without_mamba(): + num_blocks, kv_cache = _kv_first_setup(["attn"]) + + 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 diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 598bddf47a..4c137a8fe8 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -208,6 +208,49 @@ def _allocate_kv_cache( 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, @@ -215,6 +258,7 @@ def _reshape_attention_kv_cache( 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 = [ @@ -257,6 +301,24 @@ def _reshape_attention_kv_cache( 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) @@ -282,6 +344,10 @@ def _reshape_kv_cache( 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 @@ -347,6 +413,7 @@ def _reshape_kv_cache( kv_cache_stride_order, kernel_num_blocks, packing, + page_aligned_blocks=layer_name in page_aligned_layers, ) elif isinstance(kv_cache_spec, MambaSpec): @@ -371,6 +438,7 @@ def _reshape_kv_cache( 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. @@ -386,6 +454,7 @@ def _align_mixed_attention_kv_cache_views( 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. @@ -415,6 +484,11 @@ def _align_mixed_attention_kv_cache_views( 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