[ROCm] Guard AITER GDN decode fast path by layout (#42880)

Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
This commit is contained in:
Tuukka Sarvi
2026-05-18 11:56:22 -07:00
committed by GitHub
parent ce88f01c9a
commit 8fc1c284b9
@@ -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
):