| 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 | 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 |
| 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" |
| 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) |