[ROCm] Guard AITER GDN decode fast path by layout (#42880)
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
This commit is contained in:
@@ -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
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user