iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][ROCm] Give KV-first attention blocks their own page in hybrid models (#51837)

Browse Source 
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
This commit is contained in:
stefankoncarevic
2026-08-11 16:52:30 -04:00
committed by GitHub
parent 0f0cb918b7
commit 3e372c5ff2
2 changed files with 177 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_attn_utils.py
+103 -1
View File
@@ -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
vllm/v1/worker/gpu/attn_utils.py
+74
View File
@@ -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