[ROCm][Perf] Enable the AITER GDN decode fast path for flat qkvz layouts (#53623)

Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

.buildkite/test-amd.yaml

@@ -1834,6 +1834,7 @@ steps:
   source_file_dependencies:
   - csrc/mamba/
   - tests/kernels/mamba
+  - vllm/model_executor/layers/mamba/gdn
   - vllm/model_executor/layers/mamba/ops
   - vllm/platforms/rocm.py
   commands:

.buildkite/test_areas/kernels.yaml

@@ -316,6 +316,7 @@ steps:
   source_file_dependencies:
   - csrc/mamba/
   - tests/kernels/mamba
+  - vllm/model_executor/layers/mamba/gdn
   - vllm/model_executor/layers/mamba/ops
   commands:
     - pytest -v -s kernels/mamba

tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Layout dispatch tests for ``QwenGatedDeltaNetAttention._forward_core_rocm``.
+
+The AITER fused reshape+conv kernel learned Qwen3.5's flat ``[q|k|v|z]``
+packing in https://github.com/ROCm/aiter/pull/3251, so flat-layout models take
+the decode fast path too instead of falling back to the generic path.
+
+These tests run host-side on CPU: ``_forward_core_rocm`` is bound to a stub
+layer whose ``_forward_core_decode_aiter``/``_forward_core`` record which
+branch ran, so no GPU or AITER install is needed.
+"""
+
+from __future__ import annotations
+
+import types
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific tests"
+)
+
+PREFIX = "model.layers.0.linear_attn"
+H = 2  # num key heads
+HV = 4  # num value heads
+K = 8  # head_k_dim
+V = 8  # head_v_dim
+
+
+def _make_metadata(
+    *,
+    num_prefills: int = 0,
+    num_decodes: int = 2,
+    spec_sequence_masks: torch.Tensor | None = None,
+):
+    """Only the fields the dispatch condition reads carry meaningful values."""
+    return GDNAttentionMetadata(
+        num_prefills=num_prefills,
+        num_prefill_tokens=num_prefills,
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decodes,
+        num_spec_decodes=0,
+        num_spec_decode_tokens=0,
+        num_actual_tokens=num_prefills + num_decodes,
+        spec_sequence_masks=spec_sequence_masks,
+    )
+
+
+def _make_layer(gqa_interleaved_layout: bool):
+    """Stub layer running the real ``_forward_core_rocm``, recording dispatch."""
+    layer = types.SimpleNamespace()
+    layer.prefix = PREFIX
+    layer.gqa_interleaved_layout = gqa_interleaved_layout
+    layer.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat"
+
+    layer.calls = []
+    layer._forward_core_decode_aiter = lambda **kw: layer.calls.append("aiter")
+    layer._forward_core = lambda **kw: layer.calls.append("generic")
+    layer.prepare_gdn_attention_core_inputs = lambda qkvz, ba, n: (
+        torch.zeros(n, H * K * 2 + HV * V),
+        torch.zeros(n, HV, V),
+        torch.zeros(n, HV),
+        torch.zeros(n, HV),
+    )
+    layer._forward_core_rocm = types.MethodType(
+        QwenGatedDeltaNetAttention._forward_core_rocm, layer
+    )
+    return layer
+
+
+def _run(layer, meta) -> str:
+    num_tokens = max(meta.num_actual_tokens, 1)
+    ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta})
+    with patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=ctx):
+        layer._forward_core_rocm(
+            qkvz=torch.zeros(num_tokens, 2 * H * K + 2 * HV * V),
+            ba=torch.zeros(num_tokens, 2 * HV),
+            z_out=torch.zeros(num_tokens, HV, V),
+            core_attn_out=torch.zeros(num_tokens, HV, V),
+        )
+    assert len(layer.calls) == 1
+    return layer.calls[0]
+
+
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
+def test_pure_decode_takes_the_fast_path(gqa_interleaved_layout: bool) -> None:
+    """Both packings reach the fast path; flat used to be guarded out."""
+    layer = _make_layer(gqa_interleaved_layout)
+    assert _run(layer, _make_metadata()) == "aiter"
+
+
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
+def test_layout_string_matches_the_packing(gqa_interleaved_layout: bool) -> None:
+    """The value handed to the kernel's ``qkvz_layout`` parameter."""
+    layer = _make_layer(gqa_interleaved_layout)
+    expected = "interleaved" if gqa_interleaved_layout else "flat"
+    assert layer.qkvz_layout == expected
+
+
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
+@pytest.mark.parametrize(
+    "meta_kwargs",
+    [
+        {"num_prefills": 1},
+        {"num_decodes": 0},
+        {"spec_sequence_masks": torch.zeros(1, dtype=torch.bool)},
+    ],
+    ids=["has_prefill", "no_decode", "spec_decode"],
+)
+def test_non_pure_decode_batches_use_the_generic_path(
+    gqa_interleaved_layout: bool, meta_kwargs: dict
+) -> None:
+    layer = _make_layer(gqa_interleaved_layout)
+    assert _run(layer, _make_metadata(**meta_kwargs)) == "generic"

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -399,6 +399,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
         self.key_dim = self.head_k_dim * self.num_k_heads
         self.value_dim = self.head_v_dim * self.num_v_heads
         self.gqa_interleaved_layout = gqa_interleaved_layout
+        self.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat"
         if current_platform.is_xpu():
             self._forward_method = self.forward_xpu
         elif current_platform.is_cpu():
@@ -1206,9 +1207,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
         """ROCm AITER fast path: conv1d + recurrent attention from packed
         qkvz/ba layout.
 
-        For decode-only (no spec, no prefill) interleaved-GQA layouts,
-        dispatches directly to ``_forward_core_decode_aiter``. Otherwise unpacks
-        the packed layout and falls through to ``_forward_core``.
+        For decode-only (no spec, no prefill) batches, dispatches directly to
+        ``_forward_core_decode_aiter``. Otherwise unpacks the packed layout and
+        falls through to ``_forward_core``.
 
         Args:
             qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim)
@@ -1231,12 +1232,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
 
         assert isinstance(attn_metadata, GDNAttentionMetadata)
 
-        # The AITER fused reshape/conv kernel expects Qwen3-Next's interleaved
-        # GQA layout. Qwen3.5 uses a non-interleaved q/k/v/z layout and must use
-        # the generic path below to split/rearrange inputs correctly.
         if (
-            self.gqa_interleaved_layout
-            and attn_metadata.spec_sequence_masks is None
+            attn_metadata.spec_sequence_masks is None
             and attn_metadata.num_prefills == 0
             and attn_metadata.num_decodes > 0
         ):
@@ -1625,6 +1622,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
                     : attn_metadata.num_actual_tokens
                 ],
                 validate_data=True,
+                qkvz_layout=self.qkvz_layout,
             )
         )