From 8fc1c284b94668b60c30737e178cb7e6cd651e89 Mon Sep 17 00:00:00 2001 From: Tuukka Sarvi Date: Mon, 18 May 2026 21:56:22 +0300 Subject: [PATCH] [ROCm] Guard AITER GDN decode fast path by layout (#42880) Signed-off-by: Tuukka Sarvi --- vllm/model_executor/layers/mamba/gdn_linear_attn.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index db55e8ba48..71f3e3d312 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -995,9 +995,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): """ROCm AITER fast path: conv1d + recurrent attention from packed qkvz/ba layout. - For decode-only (no spec, no prefill), dispatches directly to - ``_forward_core_decode_fast``. Otherwise unpacks the packed - layout and falls through to ``_forward_core``. + For decode-only (no spec, no prefill) interleaved-GQA layouts, + dispatches directly to ``_forward_core_decode_fast``. Otherwise unpacks + the packed layout and falls through to ``_forward_core``. Args: qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim) @@ -1018,8 +1018,12 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index] 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 ( - attn_metadata.spec_sequence_masks is None + self.gqa_interleaved_layout + and attn_metadata.spec_sequence_masks is None and attn_metadata.num_prefills == 0 and attn_metadata.num_decodes > 0 ):