diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 5046506e0542..fe2cfd18c89e 100644 --- a/.buildkite/test-amd.yaml +++ b/.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: diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 7e43b123f12e..aab6e52b8eb9 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.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 diff --git a/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py b/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py new file mode 100644 index 000000000000..ae34ffcb8058 --- /dev/null +++ b/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" diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 807a5b9b1a2a..ccf1c60f0fe2 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -399,6 +399,7 @@ def __init__( 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 @@ def _forward_core_rocm( """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 @@ def _forward_core_rocm( 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 @@ def _forward_core_decode_aiter( : attn_metadata.num_actual_tokens ], validate_data=True, + qkvz_layout=self.qkvz_layout, ) )