[ROCm] Align mixed encoder-decoder KV cache views in V2 runner (#47685)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 6, 2026 at 23:09 UTC
2f71b2bd9f693f06973a8abf3823f095bd46ef69
1 file changed
+82
vllm/v1/worker/gpu/attn_utils.py
+82
@@ -12,6 +12,7 @@ from vllm.config import (
12
get_layers_from_vllm_config,
13
set_current_vllm_config,
14
)
15
+from vllm.logger import init_logger
16
from vllm.model_executor.layers.attention import Attention
17
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
18
from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -37,6 +38,8 @@ from vllm.v1.worker.utils import (
38
prepare_kernel_block_sizes,
39
)
40
41
+logger = init_logger(__name__)
42
+
43
44
@dataclass(frozen=True)
45
class AttentionCGSupportInfo:
@@ -368,6 +371,14 @@ def _reshape_kv_cache(
371
kernel_block_sizes=kernel_block_sizes,
372
cache_dtype=cache_dtype,
373
)
374
+ elif has_attn and kv_cache_config is not None:
375
+ _align_mixed_attention_kv_cache_views(
376
+ attn_groups=attn_groups,
377
+ kv_caches=kv_caches,
378
+ kernel_block_sizes=kernel_block_sizes,
379
+ cache_dtype=cache_dtype,
380
+ kv_cache_config=kv_cache_config,
381
+ )
382
383
# Map any sharing layers to their target layer's KV cache.
384
for layer_name, target_layer_name in shared_kv_cache_layers.items():
@@ -376,6 +387,77 @@ def _reshape_kv_cache(
387
return kv_caches
388
389
390
+def _align_mixed_attention_kv_cache_views(
391
+ attn_groups: Iterable[AttentionGroup],
392
+ kv_caches: dict[str, Any],
393
+ kernel_block_sizes: list[int],
394
+ cache_dtype: str,
395
+ kv_cache_config: KVCacheConfig,
396
+) -> None:
397
+ """Align shared attention KV views when backends disagree on layout.
398
+
399
+ Encoder-decoder models can share one raw allocation between decoder
400
+ self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention
401
+ (blocks-first backends, block dim 0). Keep the physical storage in the
402
+ K/V-first layout expected by ROCM_ATTN, and restride the blocks-first
403
+ logical views so block IDs address the same bytes.
404
+ """
405
+ block_dims_by_layer: dict[str, int] = {}
406
+ for group in attn_groups:
407
+ kv_cache_spec = group.kv_cache_spec
408
+ if not isinstance(kv_cache_spec, AttentionSpec):
409
+ continue
410
+ if group.kv_cache_group_id >= len(kernel_block_sizes):
411
+ continue
412
+ block_dim = group.backend.get_kv_cache_block_dim(
413
+ kernel_block_sizes[group.kv_cache_group_id],
414
+ kv_cache_spec.num_kv_heads,
415
+ kv_cache_spec.head_size,
416
+ cache_dtype_str=cache_dtype,
417
+ )
418
+ for layer_name in group.layer_names:
419
+ if layer_name in kv_caches:
420
+ block_dims_by_layer[layer_name] = block_dim
421
+
422
+ for kv_tensor in kv_cache_config.kv_cache_tensors:
423
+ if kv_tensor.block_stride > 0:
424
+ continue
425
+ shared_block_dims = {
426
+ block_dims_by_layer[layer_name]
427
+ for layer_name in kv_tensor.shared_by
428
+ if layer_name in block_dims_by_layer
429
+ }
430
+ if 0 not in shared_block_dims or 1 not in shared_block_dims:
431
+ continue
432
+
433
+ for layer_name in kv_tensor.shared_by:
434
+ if block_dims_by_layer.get(layer_name) == 0:
435
+ _restride_blocks_first_kv_cache_to_kv_first_storage(
436
+ kv_caches[layer_name]
437
+ )
438
+
439
+
440
+def _restride_blocks_first_kv_cache_to_kv_first_storage(
441
+ kv_cache: torch.Tensor,
442
+) -> None:
443
+ assert kv_cache.ndim >= 3
444
+ assert kv_cache.shape[1] == 2
445
+ page_size = kv_cache.shape[2:].numel()
446
+ num_blocks = kv_cache.shape[0]
447
+ expected_tail_stride = torch.empty(kv_cache.shape[2:]).stride()
448
+ if kv_cache.stride()[2:] != expected_tail_stride:
449
+ logger.warning_once(
450
+ "Skipping mixed KV-cache layout alignment for a non-NHD "
451
+ "blocks-first attention view with stride %s.",
452
+ kv_cache.stride(),
453
+ )
454
+ return
455
+ kv_cache.as_strided_(
456
+ size=kv_cache.shape,
457
+ stride=(page_size, num_blocks * page_size, *expected_tail_stride),
458
+ )
459
+
460
+
461
def _update_hybrid_attention_layout(
462
attn_groups: Iterable[AttentionGroup],
463
kv_caches: dict[str, Any],