[Bugfix][ROCm] Give KV-first attention blocks their own page in hybrid models (#51837)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
stefankoncarevic committed
Aug 11, 2026 at 22:52 UTC
3e372c5ff23438eeeafc86c7d8d51026f3dacb6a
2 files changed
+177
-1
tests/v1/worker/test_attn_utils.py
+103
-1
@@ -3,7 +3,13 @@
3
4
import torch
5
6
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
6
+from vllm.v1.kv_cache_interface import (
7
+ FullAttentionSpec,
8
+ KVCacheConfig,
9
+ KVCacheTensor,
10
+ KVQuantMode,
11
+ MambaSpec,
12
+)
13
from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
14
from vllm.v1.worker.utils import AttentionGroup
15
@@ -240,3 +246,99 @@ def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
246
assert kv_cache.stride(0) == spec.page_size_bytes
247
assert kv_cache.stride(1) == 16 * 1 * 8
248
assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8
249
+
250
+
251
+class FakeKVFirstBackend:
252
+ """ROCm-style backend that puts K and V ahead of the block dim."""
253
+
254
+ @staticmethod
255
+ def get_kv_cache_shape(
256
+ num_blocks: int,
257
+ block_size: int,
258
+ num_kv_heads: int,
259
+ head_size: int,
260
+ cache_dtype_str: str = "auto",
261
+ ) -> tuple[int, ...]:
262
+ return (2, num_blocks, block_size, num_kv_heads, head_size)
263
+
264
+ @staticmethod
265
+ def get_kv_cache_block_dim(
266
+ block_size: int,
267
+ num_kv_heads: int,
268
+ head_size: int,
269
+ cache_dtype_str: str = "auto",
270
+ ) -> int:
271
+ return 1
272
+
273
+
274
+def _kv_first_setup(shared_by: list[str]):
275
+ num_blocks = 3
276
+ attn_spec = FullAttentionSpec(
277
+ block_size=16,
278
+ num_kv_heads=1,
279
+ head_size=2,
280
+ dtype=torch.float32,
281
+ )
282
+ mamba_spec = MambaSpec(
283
+ block_size=16,
284
+ shapes=((64,),),
285
+ dtypes=(torch.float32,),
286
+ )
287
+ assert attn_spec.page_size_bytes == mamba_spec.page_size_bytes == 256
288
+
289
+ raw_tensor = torch.zeros(attn_spec.page_size_bytes * num_blocks, dtype=torch.int8)
290
+ raw_tensors = {name: raw_tensor for name in shared_by}
291
+ attn_groups = [
292
+ AttentionGroup(
293
+ backend=FakeKVFirstBackend,
294
+ layer_names=["attn"],
295
+ kv_cache_spec=attn_spec,
296
+ kv_cache_group_id=0,
297
+ )
298
+ ]
299
+ if "mamba" in shared_by:
300
+ attn_groups.append(
301
+ AttentionGroup(
302
+ backend=FakeKVFirstBackend,
303
+ layer_names=["mamba"],
304
+ kv_cache_spec=mamba_spec,
305
+ kv_cache_group_id=1,
306
+ )
307
+ )
308
+ kv_cache_config = KVCacheConfig(
309
+ num_blocks=num_blocks,
310
+ kv_cache_tensors=[
311
+ KVCacheTensor(size=raw_tensor.numel(), shared_by=list(shared_by))
312
+ ],
313
+ kv_cache_groups=[],
314
+ )
315
+
316
+ kv_caches = _reshape_kv_cache(
317
+ attn_groups,
318
+ raw_tensors,
319
+ "auto",
320
+ [attn_spec.block_size] * len(attn_groups),
321
+ {},
322
+ kv_cache_config,
323
+ )
324
+ return num_blocks, kv_caches["attn"]
325
+
326
+
327
+def test_reshape_kv_first_kv_cache_pages_blocks_when_shared_with_mamba():
328
+ num_blocks, kv_cache = _kv_first_setup(["attn", "mamba"])
329
+
330
+ assert kv_cache.shape == (2, num_blocks, 16, 1, 2)
331
+ # Block b has to own page b, so its K and V sit side by side within it.
332
+ page = 16 * 1 * 2 * 2
333
+ for block in range(num_blocks):
334
+ assert kv_cache[0, block].storage_offset() == block * page
335
+ assert kv_cache[1, block].storage_offset() == block * page + page // 2
336
+
337
+
338
+def test_reshape_kv_first_kv_cache_keeps_layout_without_mamba():
339
+ num_blocks, kv_cache = _kv_first_setup(["attn"])
340
+
341
+ assert kv_cache.shape == (2, num_blocks, 16, 1, 2)
342
+ # Nothing else indexes this allocation by page, so K and V stay split into
343
+ # one contiguous half each.
344
+ assert kv_cache[1, 0].storage_offset() == num_blocks * 16 * 1 * 2
vllm/v1/worker/gpu/attn_utils.py
+74
@@ -208,6 +208,49 @@ def _allocate_kv_cache(
208
return kv_cache_raw_tensors
209
210
211
+def _kv_first_layers_sharing_pool_with_mamba(
212
+ attn_groups: Sequence[AttentionGroup],
213
+ kernel_block_sizes: list[int],
214
+ cache_dtype: str,
215
+ kv_cache_config: "KVCacheConfig | None",
216
+) -> set[str]:
217
+ """KV-first attention layers that share an allocation with Mamba state.
218
+
219
+ Mamba addresses its state by page, so these layers have to lay their blocks
220
+ out page by page as well; otherwise block ids from the two groups resolve
221
+ to overlapping bytes.
222
+ """
223
+ if kv_cache_config is None:
224
+ return set()
225
+
226
+ kv_first_layers: set[str] = set()
227
+ mamba_layers: set[str] = set()
228
+ for group in attn_groups:
229
+ kv_cache_spec = group.kv_cache_spec
230
+ if isinstance(kv_cache_spec, MambaSpec):
231
+ mamba_layers.update(group.layer_names)
232
+ continue
233
+ if not isinstance(kv_cache_spec, AttentionSpec):
234
+ continue
235
+ if group.kv_cache_group_id >= len(kernel_block_sizes):
236
+ continue
237
+ block_dim = group.backend.get_kv_cache_block_dim(
238
+ kernel_block_sizes[group.kv_cache_group_id],
239
+ kv_cache_spec.num_kv_heads,
240
+ kv_cache_spec.head_size,
241
+ cache_dtype_str=cache_dtype,
242
+ )
243
+ if block_dim != 0:
244
+ kv_first_layers.update(group.layer_names)
245
+
246
+ page_aligned: set[str] = set()
247
+ for kv_tensor in kv_cache_config.kv_cache_tensors:
248
+ if mamba_layers.isdisjoint(kv_tensor.shared_by):
249
+ continue
250
+ page_aligned.update(kv_first_layers.intersection(kv_tensor.shared_by))
251
+ return page_aligned
252
+
253
+
254
def _reshape_attention_kv_cache(
255
kv_raw_tensor: torch.Tensor,
256
kv_cache_spec: AttentionSpec,
@@ -215,6 +258,7 @@ def _reshape_attention_kv_cache(
258
kv_cache_stride_order: tuple[int, ...],
259
num_blocks: int,
260
packing: tuple[int, int] | None,
261
+ page_aligned_blocks: bool = False,
262
) -> torch.Tensor:
263
permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
264
inv_order = [
@@ -257,6 +301,24 @@ def _reshape_attention_kv_cache(
301
size=permuted_kv_cache_shape,
302
stride=tuple(strides),
303
)
304
+ elif page_aligned_blocks:
305
+ # A KV-first layout such as ROCm's ``(2, num_blocks, ...)`` puts block
306
+ # ``b``'s K and V in two far-apart halves of the allocation, so block
307
+ # ``b`` does not cover page ``b``. Mamba layers sharing the allocation
308
+ # do address their state by page, so the two would resolve the same
309
+ # bytes. Build the view page-first instead, then swap the dims back.
310
+ assert kv_cache_shape[1] == num_blocks and kv_cache_stride_order == tuple(
311
+ range(len(kv_cache_shape))
312
+ ), (
313
+ "Page-aligned KV blocks expect a default-strided (kv, num_blocks, "
314
+ f"...) layout, got shape {kv_cache_shape} with stride order "
315
+ f"{kv_cache_stride_order} and num_blocks={num_blocks}."
316
+ )
317
+ kv_cache = (
318
+ kv_raw_tensor.view(dtype)
319
+ .view(num_blocks, kv_cache_shape[0], *kv_cache_shape[2:])
320
+ .transpose(0, 1)
321
+ )
322
else:
323
# No padding — safe to use a contiguous view.
324
kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape)
@@ -282,6 +344,10 @@ def _reshape_kv_cache(
344
for ln in kv_tensor.shared_by:
345
layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride)
346
347
+ page_aligned_layers = _kv_first_layers_sharing_pool_with_mamba(
348
+ attn_groups, kernel_block_sizes, cache_dtype, kv_cache_config
349
+ )
350
+
351
for group in attn_groups:
352
if group.kv_cache_group_id >= len(kernel_block_sizes):
353
continue
@@ -347,6 +413,7 @@ def _reshape_kv_cache(
413
kv_cache_stride_order,
414
kernel_num_blocks,
415
packing,
416
+ page_aligned_blocks=layer_name in page_aligned_layers,
417
)
418
419
elif isinstance(kv_cache_spec, MambaSpec):
@@ -371,6 +438,7 @@ def _reshape_kv_cache(
438
kernel_block_sizes=kernel_block_sizes,
439
cache_dtype=cache_dtype,
440
kv_cache_config=kv_cache_config,
441
+ page_aligned_layers=page_aligned_layers,
442
)
443
444
# Map any sharing layers to their target layer's KV cache.
@@ -386,6 +454,7 @@ def _align_mixed_attention_kv_cache_views(
454
kernel_block_sizes: list[int],
455
cache_dtype: str,
456
kv_cache_config: KVCacheConfig,
457
+ page_aligned_layers: set[str],
458
) -> None:
459
"""Align shared attention KV views when backends disagree on layout.
460
@@ -415,6 +484,11 @@ def _align_mixed_attention_kv_cache_views(
484
for kv_tensor in kv_cache_config.kv_cache_tensors:
485
if kv_tensor.block_stride > 0:
486
continue
487
+ if not page_aligned_layers.isdisjoint(kv_tensor.shared_by):
488
+ # These KV-first views already address blocks by page, just like
489
+ # blocks-first views, so they must not be restrided onto the
490
+ # unaligned KV-first storage layout.
491
+ continue
492
shared_block_dims = {
493
block_dims_by_layer[layer_name]
494
for layer_name in kv_tensor.shared_by