[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>
Mikko Tukiainen committed
Sep 19, 2026 at 16:05 UTC
fbe8a157fbdd4b2f813ae6395397232de97d5c5d
4 files changed
+132
-8
.buildkite/test-amd.yaml
+1
@@ -1834,6 +1834,7 @@ steps:
1834
source_file_dependencies:
1835
- csrc/mamba/
1836
- tests/kernels/mamba
1837
+ - vllm/model_executor/layers/mamba/gdn
1838
- vllm/model_executor/layers/mamba/ops
1839
- vllm/platforms/rocm.py
1840
commands:
.buildkite/test_areas/kernels.yaml
+1
@@ -316,6 +316,7 @@ steps:
316
source_file_dependencies:
317
- csrc/mamba/
318
- tests/kernels/mamba
319
+ - vllm/model_executor/layers/mamba/gdn
320
- vllm/model_executor/layers/mamba/ops
321
commands:
322
- pytest -v -s kernels/mamba
tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py
new
+124
@@ -0,0 +1,124 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Layout dispatch tests for ``QwenGatedDeltaNetAttention._forward_core_rocm``.
4
+
5
+The AITER fused reshape+conv kernel learned Qwen3.5's flat ``[q|k|v|z]``
6
+packing in https://github.com/ROCm/aiter/pull/3251, so flat-layout models take
7
+the decode fast path too instead of falling back to the generic path.
8
+
9
+These tests run host-side on CPU: ``_forward_core_rocm`` is bound to a stub
10
+layer whose ``_forward_core_decode_aiter``/``_forward_core`` record which
11
+branch ran, so no GPU or AITER install is needed.
12
+"""
13
+
14
+from __future__ import annotations
15
+
16
+import types
17
+from unittest.mock import patch
18
+
19
+import pytest
20
+import torch
21
+
22
+from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn
23
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
24
+ QwenGatedDeltaNetAttention,
25
+)
26
+from vllm.platforms import current_platform
27
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
28
+
29
+pytestmark = pytest.mark.skipif(
30
+ not current_platform.is_rocm(), reason="ROCm-specific tests"
31
+)
32
+
33
+PREFIX = "model.layers.0.linear_attn"
34
+H = 2 # num key heads
35
+HV = 4 # num value heads
36
+K = 8 # head_k_dim
37
+V = 8 # head_v_dim
38
+
39
+
40
+def _make_metadata(
41
+ *,
42
+ num_prefills: int = 0,
43
+ num_decodes: int = 2,
44
+ spec_sequence_masks: torch.Tensor | None = None,
45
+):
46
+ """Only the fields the dispatch condition reads carry meaningful values."""
47
+ return GDNAttentionMetadata(
48
+ num_prefills=num_prefills,
49
+ num_prefill_tokens=num_prefills,
50
+ num_decodes=num_decodes,
51
+ num_decode_tokens=num_decodes,
52
+ num_spec_decodes=0,
53
+ num_spec_decode_tokens=0,
54
+ num_actual_tokens=num_prefills + num_decodes,
55
+ spec_sequence_masks=spec_sequence_masks,
56
+ )
57
+
58
+
59
+def _make_layer(gqa_interleaved_layout: bool):
60
+ """Stub layer running the real ``_forward_core_rocm``, recording dispatch."""
61
+ layer = types.SimpleNamespace()
62
+ layer.prefix = PREFIX
63
+ layer.gqa_interleaved_layout = gqa_interleaved_layout
64
+ layer.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat"
65
+
66
+ layer.calls = []
67
+ layer._forward_core_decode_aiter = lambda **kw: layer.calls.append("aiter")
68
+ layer._forward_core = lambda **kw: layer.calls.append("generic")
69
+ layer.prepare_gdn_attention_core_inputs = lambda qkvz, ba, n: (
70
+ torch.zeros(n, H * K * 2 + HV * V),
71
+ torch.zeros(n, HV, V),
72
+ torch.zeros(n, HV),
73
+ torch.zeros(n, HV),
74
+ )
75
+ layer._forward_core_rocm = types.MethodType(
76
+ QwenGatedDeltaNetAttention._forward_core_rocm, layer
77
+ )
78
+ return layer
79
+
80
+
81
+def _run(layer, meta) -> str:
82
+ num_tokens = max(meta.num_actual_tokens, 1)
83
+ ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta})
84
+ with patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=ctx):
85
+ layer._forward_core_rocm(
86
+ qkvz=torch.zeros(num_tokens, 2 * H * K + 2 * HV * V),
87
+ ba=torch.zeros(num_tokens, 2 * HV),
88
+ z_out=torch.zeros(num_tokens, HV, V),
89
+ core_attn_out=torch.zeros(num_tokens, HV, V),
90
+ )
91
+ assert len(layer.calls) == 1
92
+ return layer.calls[0]
93
+
94
+
95
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
96
+def test_pure_decode_takes_the_fast_path(gqa_interleaved_layout: bool) -> None:
97
+ """Both packings reach the fast path; flat used to be guarded out."""
98
+ layer = _make_layer(gqa_interleaved_layout)
99
+ assert _run(layer, _make_metadata()) == "aiter"
100
+
101
+
102
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
103
+def test_layout_string_matches_the_packing(gqa_interleaved_layout: bool) -> None:
104
+ """The value handed to the kernel's ``qkvz_layout`` parameter."""
105
+ layer = _make_layer(gqa_interleaved_layout)
106
+ expected = "interleaved" if gqa_interleaved_layout else "flat"
107
+ assert layer.qkvz_layout == expected
108
+
109
+
110
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
111
+@pytest.mark.parametrize(
112
+ "meta_kwargs",
113
+ [
114
+ {"num_prefills": 1},
115
+ {"num_decodes": 0},
116
+ {"spec_sequence_masks": torch.zeros(1, dtype=torch.bool)},
117
+ ],
118
+ ids=["has_prefill", "no_decode", "spec_decode"],
119
+)
120
+def test_non_pure_decode_batches_use_the_generic_path(
121
+ gqa_interleaved_layout: bool, meta_kwargs: dict
122
+) -> None:
123
+ layer = _make_layer(gqa_interleaved_layout)
124
+ assert _run(layer, _make_metadata(**meta_kwargs)) == "generic"
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+6
-8
@@ -399,6 +399,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
399
self.key_dim = self.head_k_dim * self.num_k_heads
400
self.value_dim = self.head_v_dim * self.num_v_heads
401
self.gqa_interleaved_layout = gqa_interleaved_layout
402
+ self.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat"
403
if current_platform.is_xpu():
404
self._forward_method = self.forward_xpu
405
elif current_platform.is_cpu():
@@ -1206,9 +1207,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1207
"""ROCm AITER fast path: conv1d + recurrent attention from packed
1208
qkvz/ba layout.
1209
1209
- For decode-only (no spec, no prefill) interleaved-GQA layouts,
1210
- dispatches directly to ``_forward_core_decode_aiter``. Otherwise unpacks
1211
- the packed layout and falls through to ``_forward_core``.
1210
+ For decode-only (no spec, no prefill) batches, dispatches directly to
1211
+ ``_forward_core_decode_aiter``. Otherwise unpacks the packed layout and
1212
+ falls through to ``_forward_core``.
1213
1214
Args:
1215
qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim)
@@ -1231,12 +1232,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1232
1233
assert isinstance(attn_metadata, GDNAttentionMetadata)
1234
1234
- # The AITER fused reshape/conv kernel expects Qwen3-Next's interleaved
1235
- # GQA layout. Qwen3.5 uses a non-interleaved q/k/v/z layout and must use
1236
- # the generic path below to split/rearrange inputs correctly.
1235
if (
1238
- self.gqa_interleaved_layout
1239
- and attn_metadata.spec_sequence_masks is None
1236
+ attn_metadata.spec_sequence_masks is None
1237
and attn_metadata.num_prefills == 0
1238
and attn_metadata.num_decodes > 0
1239
):
@@ -1625,6 +1622,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1622
: attn_metadata.num_actual_tokens
1623
],
1624
validate_data=True,
1625
+ qkvz_layout=self.qkvz_layout,
1626
)
1627
)
1628