A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 09:15 UTC
[ROCm] Guard AITER GDN decode fast path by layout (#42880)
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
Tuukka Sarvi committed
May 18, 2026 at 21:56 UTC
8fc1c284b94668b60c30737e178cb7e6cd651e89
1 file changed
+8
-4
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+8
-4
@@ -995,9 +995,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
995
"""ROCm AITER fast path: conv1d + recurrent attention from packed
996
qkvz/ba layout.
997
998
- For decode-only (no spec, no prefill), dispatches directly to
999
- ``_forward_core_decode_fast``. Otherwise unpacks the packed
1000
- layout and falls through to ``_forward_core``.
998
+ For decode-only (no spec, no prefill) interleaved-GQA layouts,
999
+ dispatches directly to ``_forward_core_decode_fast``. Otherwise unpacks
1000
+ the packed layout and falls through to ``_forward_core``.
1001
1002
Args:
1003
qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim)
@@ -1018,8 +1018,12 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1018
attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]
1019
assert isinstance(attn_metadata, GDNAttentionMetadata)
1020
1021
+ # The AITER fused reshape/conv kernel expects Qwen3-Next's interleaved
1022
+ # GQA layout. Qwen3.5 uses a non-interleaved q/k/v/z layout and must use
1023
+ # the generic path below to split/rearrange inputs correctly.
1024
if (
1022
- attn_metadata.spec_sequence_masks is None
1025
+ self.gqa_interleaved_layout
1026
+ and attn_metadata.spec_sequence_masks is None
1027
and attn_metadata.num_prefills == 0
1028
and attn_metadata.num_decodes > 0
1029
):