Fix stale KV-cache API references missed by the standardization

- benchmarks/attention_benchmarks/runner.py: import get_kv_cache_layout
  (resolve_kv_cache_layout was renamed) and pass reshape_kv_cache's
  offset/layer_stride/block_stride. Fixes the non-optional MI355
  "Attention Benchmarks Smoke Test" job. Reported by @okorzh-amd.
- test_rocm_aiter_fa: replace the deleted get_kv_cache_shape assertions
  with the customize_spec packing contract.
- test_sparse_mla_backends: squeeze the head axis before calling
  forward_mha directly, mirroring bind_kv_cache. Without it the
  chunked-context prefill read the cache at the wrong stride.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

benchmarks/attention_benchmarks/runner.py

@@ -39,12 +39,13 @@
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import (
     CommonAttentionMetadata,
+    get_kv_cache_layout,
     initialize_kv_cache_layout,
-    resolve_kv_cache_layout,
 )
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     compute_layer_kv_cache_shape_bytes,
+    layer_kv_cache_strides,
     reshape_kv_cache,
 )
 
@@ -357,13 +358,25 @@ def _create_kv_cache(
         head_size=config.head_dim,
         dtype=cache_dtype,
     )
-    layout = resolve_kv_cache_layout()
+    layout = get_kv_cache_layout()
     total_bytes = (
         prod(compute_layer_kv_cache_shape_bytes(spec, max_num_blocks))
         * config.num_layers
     )
     buf = torch.zeros(total_bytes, device=device, dtype=torch.int8)
-    return reshape_kv_cache(buf, spec, max_num_blocks, config.num_layers, layout)
+    layer_stride, block_stride = layer_kv_cache_strides(
+        spec, max_num_blocks, config.num_layers, layout
+    )
+    return reshape_kv_cache(
+        buf,
+        spec,
+        max_num_blocks,
+        config.num_layers,
+        layout,
+        offset=0,
+        layer_stride=layer_stride,
+        block_stride=block_stride,
+    )
 
 
 # ============================================================================

tests/kernels/attention/test_rocm_aiter_fa.py

@@ -309,19 +309,26 @@ def test_aiter_mha_backend_contract():
     )
 
 
-def test_aiter_mha_backend_validates_kv_cache_block_size():
-    """The backend should reject KV cache shapes that cannot be gathered
-    correctly by the ROCm kernel."""
+def test_aiter_mha_backend_publishes_separate_kv_head_groups(monkeypatch):
+    """K and V must be published as two token-major head groups, which is how
+    the fused QK-norm+RoPE+cache kernel addresses the cache."""
+    from vllm.v1.attention.backends import rocm_aiter_fa
     from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
+    from vllm.v1.kv_cache_interface import FullAttentionSpec
 
-    assert AiterFlashAttentionBackend.get_kv_cache_shape(8, 16, 8, 128) == (
-        8,
-        8,
-        16,
-        256,
+    spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
     )
-    with pytest.raises(ValueError, match="Block size must be a multiple of 16"):
-        AiterFlashAttentionBackend.get_kv_cache_shape(8, 15, 8, 128)
+
+    monkeypatch.setattr(rocm_aiter_fa, "_use_separate_kv_head_groups", lambda: True)
+    published = AiterFlashAttentionBackend.customize_spec(spec)
+    assert (published.num_head_slots, published.state_content_bytes) == (2, 8 * 128 * 2)
+    assert published.page_size_bytes == spec.page_size_bytes
+    assert AiterFlashAttentionBackend.get_required_kv_cache_layout() == "LBHNC"
+
+    monkeypatch.setattr(rocm_aiter_fa, "_use_separate_kv_head_groups", lambda: False)
+    assert AiterFlashAttentionBackend.customize_spec(spec) == spec
+    assert AiterFlashAttentionBackend.get_required_kv_cache_layout() is None
 
 
 def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe():

tests/v1/attention/test_sparse_mla_backends.py

@@ -1157,7 +1157,9 @@ def test_sparse_backend_prefill_correctness(
             q=query_cat,
             kv_c_normed=kv_c_cat,
             k_pe=k_pe_cat,
-            kv_c_and_k_pe_cache=kv_cache,
+            # Impls receive the bind-time-squeezed [B, N, C] cache; mirror
+            # mla_attention.py's bind_kv_cache squeeze here.
+            kv_c_and_k_pe_cache=kv_cache.squeeze(1),
             attn_metadata=metadata,
             k_scale=torch.tensor(1.0, device=device),
             output=out_buffer,