[ROCm][Test] Check GDN prefill numerics and output ownership (#58091)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Sep 23, 2026 at 04:19 UTC
a22c637ed97d81aa342789f7c56c2690ee118bfc
3 files changed
+132
-3
.buildkite/test-amd.yaml
+1
@@ -1830,6 +1830,7 @@ steps:
1830
- tests/kernels/mamba
1831
- vllm/model_executor/layers/mamba/gdn
1832
- vllm/model_executor/layers/mamba/ops
1833
+ - vllm/third_party/flash_linear_attention/
1834
- vllm/platforms/rocm.py
1835
commands:
1836
- pytest -v -s kernels/mamba
.buildkite/test_areas/kernels.yaml
+1
@@ -394,6 +394,7 @@ steps:
394
- image-build-amd
395
source_file_dependencies:
396
- vllm/platforms/rocm.py
397
+ - vllm/third_party/flash_linear_attention/
398
399
- label: ":nvidia: (H100) DeepGEMM Kernels"
400
key: kernels-deepgemm-test-h100
tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py
+130
-3
@@ -6,9 +6,9 @@ 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.
9
+The dispatch tests run host-side with recording stubs. The prefill numerical
10
+tests exercise the production ChunkGatedDeltaRule wrapper on ROCm against a
11
+token-wise Torch reference, including ragged state and output-buffer layouts.
12
"""
13
14
from __future__ import annotations
@@ -122,3 +122,130 @@ def test_non_pure_decode_batches_use_the_generic_path(
122
) -> None:
123
layer = _make_layer(gqa_interleaved_layout)
124
assert _run(layer, _make_metadata(**meta_kwargs)) == "generic"
125
+
126
+
127
+def _prefill_reference(q, k, v, g, beta, initial_state, lengths):
128
+ """Independent token-wise recurrence in the layer's [value, key] layout."""
129
+ q, k, v, g, beta, initial_state = (
130
+ x.cpu().float() for x in (q, k, v, g, beta, initial_state)
131
+ )
132
+ q = q.repeat_interleave(v.shape[2] // q.shape[2], dim=2)
133
+ k = k.repeat_interleave(v.shape[2] // k.shape[2], dim=2)
134
+ q = q * q.shape[-1] ** -0.5
135
+ output = torch.empty_like(v)
136
+ final_state = torch.empty_like(initial_state)
137
+ previous_threads = torch.get_num_threads()
138
+ torch.set_num_threads(1)
139
+ try:
140
+ start = 0
141
+ for sequence, length in enumerate(lengths):
142
+ state = initial_state[sequence].clone()
143
+ for token in range(start, start + length):
144
+ state = state * g[0, token].exp()[:, None, None]
145
+ key = k[0, token]
146
+ prediction = (state * key[:, None, :]).sum(dim=-1)
147
+ delta = (v[0, token] - prediction) * beta[0, token, :, None]
148
+ state = state + delta[:, :, None] * key[:, None, :]
149
+ output[0, token] = (state * q[0, token, :, None, :]).sum(dim=-1)
150
+ final_state[sequence] = state
151
+ start += length
152
+ finally:
153
+ torch.set_num_threads(previous_threads)
154
+ return output, final_state
155
+
156
+
157
+@pytest.mark.parametrize(
158
+ "state_dtype,correlated",
159
+ [(torch.bfloat16, False), (torch.float32, False), (torch.float32, True)],
160
+ ids=["ragged-bf16-state", "ragged-fp32-state", "correlated-keys"],
161
+)
162
+@torch.inference_mode()
163
+def test_rocm_prefill_numerics_and_output_layout(state_dtype, correlated):
164
+ from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
165
+ from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
166
+ ChunkGatedDeltaRule,
167
+ )
168
+ from vllm.third_party.flash_linear_attention.ops.index import (
169
+ prepare_chunk_indices,
170
+ prepare_chunk_offsets,
171
+ )
172
+
173
+ if not torch.cuda.is_available():
174
+ pytest.skip("ROCm prefill numerics require a GPU")
175
+ # Only the config is needed for production backend selection; no weights.
176
+ config = VllmConfig(
177
+ model_config=ModelConfig(
178
+ model="Qwen/Qwen3.5-0.8B",
179
+ revision="2fc06364715b967f1860aea9cf38778875588b17",
180
+ max_model_len=1024,
181
+ )
182
+ )
183
+ with set_current_vllm_config(config):
184
+ prefill = ChunkGatedDeltaRule()
185
+ assert prefill.gdn_prefill_backend == "triton"
186
+ assert prefill._forward_method == prefill.forward_native
187
+
188
+ torch.manual_seed(1234)
189
+ lengths = [64] if correlated else [1, 63, 64, 65, 129]
190
+ tokens, key_heads, value_heads, dim = sum(lengths), 4, 8, 128
191
+ q = torch.randn(1, tokens, key_heads, dim, device="cuda")
192
+ q = torch.nn.functional.normalize(q, dim=-1).to(torch.bfloat16)
193
+ k = torch.randn_like(q)
194
+ k = torch.nn.functional.normalize(k.float(), dim=-1).to(q.dtype)
195
+ v = torch.randn(1, tokens, value_heads, dim, device="cuda", dtype=q.dtype)
196
+ g = -torch.rand(1, tokens, value_heads, device="cuda") * 0.1
197
+ beta = torch.rand_like(g)
198
+ state = (torch.randn(len(lengths), value_heads, dim, dim, device="cuda") * 0.05).to(
199
+ state_dtype
200
+ )
201
+ if correlated:
202
+ q.zero_()
203
+ k.zero_()
204
+ q[..., 0] = 1
205
+ k[..., 0] = 1
206
+ g.zero_()
207
+ beta.fill_(0.9)
208
+ state.zero_()
209
+ expected, expected_state = _prefill_reference(q, k, v, g, beta, state, lengths)
210
+ initial_state_copy = state.clone()
211
+ cu_seqlens = torch.tensor(
212
+ [0, *torch.tensor(lengths).cumsum(0).tolist()],
213
+ device="cuda",
214
+ dtype=torch.int32,
215
+ )
216
+ kwargs = dict(
217
+ q=q,
218
+ k=k,
219
+ v=v,
220
+ g=g,
221
+ beta=beta,
222
+ initial_state=state,
223
+ output_final_state=True,
224
+ cu_seqlens=cu_seqlens,
225
+ chunk_indices=prepare_chunk_indices(cu_seqlens, 64),
226
+ chunk_offsets=prepare_chunk_offsets(cu_seqlens, 64),
227
+ use_qk_l2norm_in_kernel=False,
228
+ )
229
+ storage = torch.full(
230
+ (tokens + 3, value_heads, dim), float("nan"), device="cuda", dtype=q.dtype
231
+ )
232
+ output, final_state = prefill(**kwargs, core_attn_out=storage[:tokens])
233
+ assert output.data_ptr() == storage.data_ptr()
234
+ assert output.shape == v.shape
235
+ assert final_state.shape == state.shape
236
+ assert torch.isnan(storage[tokens:]).all()
237
+ torch.testing.assert_close(state, initial_state_copy, atol=0, rtol=0)
238
+
239
+ # Preserve the upstream prefill suite's numerical budgets. The reference
240
+ # here is a token-wise recurrence, independent of the chunk implementation.
241
+ output_error = (output.cpu().float() - expected).abs()
242
+ state_error = (final_state.cpu().float() - expected_state).abs()
243
+ assert torch.isfinite(output).all() and torch.isfinite(final_state).all()
244
+ assert output_error.max() < (1e-1 if correlated else 2e-3)
245
+ assert output_error.mean() < (2e-3 if correlated else 6e-5)
246
+ assert state_error.max() < (2e-1 if correlated else 2e-2)
247
+ assert state_error.mean() < 6e-4
248
+
249
+ no_buffer_output, no_buffer_state = prefill(**kwargs)
250
+ torch.testing.assert_close(no_buffer_output, output, atol=0, rtol=0)
251
+ torch.testing.assert_close(no_buffer_state, final_state, atol=0, rtol=0)