[NVBUG: 6562021] Fix vLLM FlashAttention KV cache layout handling

modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py

@@ -47,6 +47,19 @@
 from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite
 
 
+@functools.cache
+def _flash_attention_kv_cache_layout() -> str:
+    """Return the installed vLLM backend's K/V packing contract."""
+    cache_shape = FlashAttentionBackend.get_kv_cache_shape(3, 16, 1, 16)
+    if cache_shape == (2, 3, 16, 1, 16):
+        return "kv-first"
+    if cache_shape == (3, 2, 16, 1, 16):
+        return "blocks-first"
+    if cache_shape == (3, 1, 16, 32):
+        return "packed"
+    raise RuntimeError(f"Unsupported vLLM FlashAttention KV cache shape {cache_shape}")
+
+
 def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float:
     """Return target sparsity for a phase, defaulting old checkpoint metadata."""
     if isinstance(target_sparse_ratio, float | int):
@@ -514,7 +527,13 @@ def native_forward():
         if resolved is None:
             return native_forward()
 
-        key_cache, value_cache = kv_cache.unbind(0)
+        cache_layout = _flash_attention_kv_cache_layout()
+        if cache_layout == "kv-first":
+            key_cache, value_cache = kv_cache.unbind(0)
+        elif cache_layout == "blocks-first":
+            key_cache, value_cache = kv_cache.unbind(1)
+        else:
+            key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
         is_decode_only = attn_metadata.max_query_len <= 1
         common_kw = {
             "layer": layer,

tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py

@@ -28,7 +28,7 @@
 import vllm
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends import flashinfer as flashinfer_backend
-from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl
 from vllm.v1.attention.backends.flashinfer import (
     FlashInferBackend,
     FlashInferImpl,
@@ -541,6 +541,76 @@ def _make_flash_attention_impl(*, sparse=False, quantized=False):
     return impl
 
 
+def _flash_attention_kv_cache(num_blocks, page_size, num_kv_heads, head_size):
+    layout = vllm_plugin._flash_attention_kv_cache_layout()
+    if layout == "packed":
+        shape = [num_blocks, num_kv_heads, page_size, 2 * head_size]
+    else:
+        shape = [num_blocks, page_size, num_kv_heads, head_size]
+        shape.insert(0 if layout == "kv-first" else 1, 2)
+    return torch.zeros(shape, dtype=torch.float16)
+
+
+@pytest.mark.parametrize(
+    ("layout", "backend_shape"),
+    [
+        ("kv-first", (2, 3, 16, 1, 16)),
+        ("blocks-first", (3, 2, 16, 1, 16)),
+        ("packed", (3, 1, 16, 32)),
+    ],
+)
+def test_flash_attention_forward_follows_backend_kv_cache_layout(
+    monkeypatch, layout, backend_shape
+):
+    impl = _make_flash_attention_impl(sparse=True)
+    if layout == "packed":
+        shape = [3, impl.num_kv_heads, 16, 2 * impl.head_size]
+    else:
+        shape = [3, 16, impl.num_kv_heads, impl.head_size]
+        shape.insert(0 if layout == "kv-first" else 1, 2)
+    monkeypatch.setattr(
+        FlashAttentionBackend, "get_kv_cache_shape", staticmethod(lambda *_args: backend_shape)
+    )
+    vllm_plugin._flash_attention_kv_cache_layout.cache_clear()
+    kv_cache = torch.zeros(shape, dtype=torch.float16)
+    query = torch.zeros(4, impl.num_heads, impl.head_size, dtype=torch.float16)
+    metadata = _flash_attention_metadata(query.shape[0], 16)
+    captured = {}
+
+    def fake_attention(query, **kwargs):
+        captured.update(kwargs)
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention)
+
+    try:
+        impl.forward(
+            layer=None,
+            query=query,
+            key=query,
+            value=query,
+            kv_cache=kv_cache,
+            attn_metadata=metadata,
+            output=torch.empty_like(query),
+        )
+    finally:
+        vllm_plugin._flash_attention_kv_cache_layout.cache_clear()
+
+    if layout == "packed":
+        expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split(
+            impl.head_size, dim=-1
+        )
+    else:
+        expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1)
+    assert captured["k_cache"].shape == expected_key_cache.shape
+    assert captured["v_cache"].shape == expected_value_cache.shape
+    assert captured["k_cache"].stride() == expected_key_cache.stride()
+    assert captured["v_cache"].stride() == expected_value_cache.stride()
+    assert captured["k_cache"].data_ptr() == expected_key_cache.data_ptr()
+    assert captured["v_cache"].data_ptr() == expected_value_cache.data_ptr()
+    assert captured["page_size"] == 16
+
+
 def _flash_attention_mixed_metadata(decode_len=1, prefill_len=17):
     query_lens = (decode_len, prefill_len)
     seq_lens = (16, 34)
@@ -568,7 +638,7 @@ def test_flash_attention_mixed_batch_splits_decode_and_prefill(monkeypatch, quan
     prefill_tokens = 17
     impl = _make_flash_attention_impl(sparse=True, quantized=quantized)
     query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16)
-    kv_cache = torch.zeros(2, 4, 16, 2, 64, dtype=torch.float16)
+    kv_cache = _flash_attention_kv_cache(4, 16, 2, 64)
     metadata = _flash_attention_mixed_metadata(decode_len=1, prefill_len=prefill_tokens)
     layer = SimpleNamespace(
         _query_quant_in_kernel=quantized,
@@ -761,7 +831,7 @@ def test_forward_delegates_cascade_metadata_to_vllm(monkeypatch):
     """Cascade/prefix-cache metadata should use vLLM's native implementation."""
     impl = _clone_sparse_impl(_make_old_impl())
     q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16)
-    kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16)
+    kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size)
     output = torch.empty_like(q)
     attn_metadata = type("AttnMetadata", (), {"use_cascade": True})()
     called = {}
@@ -838,9 +908,7 @@ def test_forward_delegates_launches_without_effective_sparse_work(
     impl = _clone_sparse_impl(_make_old_impl())
     impl.sparse_kw = sparse_kw
     q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16)
-    kv_cache = torch.zeros(
-        2, 1, max_seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16
-    )
+    kv_cache = _flash_attention_kv_cache(1, max_seq_len, impl.num_kv_heads, impl.head_size)
     output = torch.empty_like(q)
     attn_metadata = _flash_attention_metadata(max_query_len, max_seq_len)
     called = {}
@@ -903,7 +971,7 @@ def test_forward_resolves_calibrated_skip_softmax_threshold(monkeypatch):
         "target_sparse_ratio": {"prefill": 0.4, "decode": 0.6},
     }
     q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16)
-    kv_cache = torch.zeros(2, 1, seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16)
+    kv_cache = _flash_attention_kv_cache(1, seq_len, impl.num_kv_heads, impl.head_size)
     attn_metadata = _flash_attention_metadata(max_query_len, seq_len)
     captured = {}
 
@@ -980,7 +1048,7 @@ def quantize_q(query):
     }
     q = torch.full((4, impl.num_heads, impl.head_size), 2.0, dtype=torch.float16)
     q[2:] = 10_000
-    kv_cache = torch.zeros(2, 4, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16)
+    kv_cache = _flash_attention_kv_cache(4, 16, impl.num_kv_heads, impl.head_size)
     metadata = SimpleNamespace(
         num_actual_tokens=q.shape[0],
         max_query_len=1,
@@ -1023,8 +1091,15 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs):
         "v_qdq_scale": 1.0,
     }
     key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"]
-    assert key_cache.data_ptr() == kv_cache[0].data_ptr()
-    assert value_cache.data_ptr() == kv_cache[1].data_ptr()
+    layout = vllm_plugin._flash_attention_kv_cache_layout()
+    if layout == "packed":
+        expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split(
+            impl.head_size, dim=-1
+        )
+    else:
+        expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1)
+    assert key_cache.data_ptr() == expected_key_cache.data_ptr()
+    assert value_cache.data_ptr() == expected_value_cache.data_ptr()
     assert block_table is metadata.block_table
     assert seq_lens is metadata.seq_lens
     assert calls["query"].shape[0] == metadata.seq_lens.shape[0]
@@ -1048,7 +1123,7 @@ def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch):
     }
     impl.sparse_kw = {"skip_softmax_threshold": 0.001}
     q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16)
-    kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16)
+    kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size)
     metadata = _flash_attention_metadata(1, 16)
     captured = {}
 
@@ -1140,7 +1215,7 @@ def test_forward_allows_chunked_prefill_metadata(monkeypatch):
     q_len = 4
     kv_len = 10
     q = torch.zeros(q_len, impl.num_heads, impl.head_size, dtype=torch.float16)
-    kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16)
+    kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size)
     attn_metadata = _flash_attention_metadata(q_len, kv_len)
     captured = {}
 

tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py

@@ -20,7 +20,7 @@
 
 * ``query_start_loc``       -> ``b_start_loc`` / ``b_seq_len``
 * ``seq_lens``              -> ``b_seq_len_k``
-* ``kv_cache.unbind(0)``    -> key_cache / value_cache (axis order)
+* backend-declared K/V axis -> key_cache / value_cache
 * ``k_cache.shape[1]``      -> ``page_size``
 
 Asserted against a contiguous reference call to the underlying Triton kernel.
@@ -33,7 +33,7 @@
 from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl
 
 from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE
-from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ModelOptSparseAttentionImpl
+from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as vllm_plugin
 
 if TRITON_KERNEL_AVAILABLE:
     from modelopt.torch.kernels.common.attention import attention as triton_attention
@@ -46,11 +46,19 @@
 }
 
 
+def _make_backend_paged_cache(k_cache, v_cache):
+    layout = vllm_plugin._flash_attention_kv_cache_layout()
+    if layout == "kv-first":
+        return torch.stack([k_cache, v_cache], dim=0)
+    if layout == "blocks-first":
+        return torch.stack([k_cache, v_cache], dim=1)
+    return torch.cat([k_cache, v_cache], dim=-1).transpose(1, 2)
+
+
 def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size):
-    """Scatter contiguous K/V into a paged KV cache stacked as [2, ...].
+    """Scatter contiguous K/V into the installed vLLM paged-cache layout.
 
-    Returns a single ``kv_cache`` tensor (matching vLLM's layout that
-    ``ModelOptSparseAttentionImpl`` consumes via ``kv_cache.unbind(0)``).
+    Returns a single ``kv_cache`` tensor with the backend-declared K/V axis.
     """
     batch = b_seq_len.shape[0]
     device, dtype = k.device, k.dtype
@@ -76,14 +84,13 @@ def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page
             v_cache[g, :n] = v[start + ts : start + te]
             g += 1
 
-    # Stack on a new leading axis so kv_cache.unbind(0) recovers (k_cache, v_cache).
-    kv_cache = torch.stack([k_cache, v_cache], dim=0)
+    kv_cache = _make_backend_paged_cache(k_cache, v_cache)
     return kv_cache, block_table
 
 
 def _make_impl(num_heads, head_dim, num_kv_heads):
     """Construct ModelOptSparseAttentionImpl with minimal valid kwargs."""
-    return ModelOptSparseAttentionImpl(
+    return vllm_plugin.ModelOptSparseAttentionImpl(
         num_heads=num_heads,
         head_size=head_dim,
         scale=1.0 / (head_dim**0.5),
@@ -132,7 +139,7 @@ def test_prefill_matches_contiguous(self):
             **_ACTIVE_PREFILL_SPARSE_KW,
         )
 
-        # Build paged kv_cache shaped [2, num_blocks, page_size, num_kv_heads, head_dim].
+        # Build the paged cache using the installed backend's K/V axis.
         kv_cache, block_table = _make_paged_cache(
             k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size
         )
@@ -174,7 +181,8 @@ def test_chunked_prefill_is_forwarded_to_kernel(self):
             block_table=torch.zeros(1, 1, device="cuda", dtype=torch.int32),
         )
         q = torch.zeros(4, 2, 64, device="cuda", dtype=torch.float16)
-        kv_cache = torch.zeros(2, 1, 16, 2, 64, device="cuda", dtype=torch.float16)
+        k_cache = torch.zeros(1, 16, 2, 64, device="cuda", dtype=torch.float16)
+        kv_cache = _make_backend_paged_cache(k_cache, torch.zeros_like(k_cache))
         out = impl.forward(
             layer=None,
             query=q,
@@ -346,8 +354,7 @@ def test_page_size_inferred_from_k_cache(self):
         kv_cache, block_table = _make_paged_cache(
             k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size
         )
-        # Sanity: kv_cache axis 1 is page_size.
-        assert kv_cache.shape == (2, seq_len // page_size, page_size, num_kv_heads, head_dim)
+        assert kv_cache.shape[2] == page_size
 
         attn_metadata = SimpleNamespace(
             num_actual_tokens=seq_len,