[AMD][Perf] Fuse GatedDeltaNet QKVZBA split/reshape/cat into a single Triton kernel for Qwen3.5-architecture MoE on HIP (#34421)

python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py

@@ -4,6 +4,11 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.utils import get_bool_env_var, is_hip
+
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+
 # =============================================================================
 # Fused kernel — reads INTERLEAVED input format
 # Used by Qwen3-Next whose checkpoint stores fused in_proj_qkvz weights
@@ -293,6 +298,16 @@ def fused_qkvzba_split_reshape_cat_contiguous(
     )
     a = torch.empty_like(b)
     grid = (batch * seq_len, num_heads_qk)
+    # Each program moves `v_per_group * head_v` elements for both v and z. For
+    # the small head-group ratios (<= 512 elements) a single warp is the best
+    # fit; wider ratios (e.g. 8 v-heads per k-head) need more lanes so the
+    # per-program vector load/store does not serialize. The threshold was tuned
+    # on MI355X, so it is confined to the HIP/aiter path; every other backend
+    # keeps the original `num_warps=1`.
+    num_warps = 1
+    if _use_aiter:
+        v_elems_per_program = (num_heads_v // num_heads_qk) * head_v
+        num_warps = 1 if v_elems_per_program <= 512 else 4
     fused_qkvzba_split_reshape_cat_contiguous_kernel[grid](
         mixed_qkv,
         z,
@@ -304,7 +319,7 @@ def fused_qkvzba_split_reshape_cat_contiguous(
         num_heads_v,
         head_qk,
         head_v,
-        num_warps=1,
+        num_warps=num_warps,
         num_stages=3,
     )
     return mixed_qkv, z, b, a

python/sglang/srt/models/qwen3_5.py

@@ -136,6 +136,13 @@ _qknorm_use_alt_stream = _is_cuda or (
 )
 _is_amx_available = cpu_has_amx_support()
 
+# Head-group ratios (num_v_heads // num_k_heads) served by the fused
+# split/reshape/cat Triton kernel. On AMD/aiter the ratio-8 layout is also
+# covered by the fused kernel, which removes the two `.contiguous()` copies
+# plus the `torch.cat` of the unfused fallback. Other backends keep the
+# original tuple so their control flow is unchanged.
+_GDN_FUSED_QKVZBA_RATIOS = (1, 2, 4, 8) if _use_aiter else (1, 2, 4)
+
 cached_get_processor = lru_cache(get_processor)
 
 
@@ -636,7 +643,10 @@ class Qwen3_5GatedDeltaNet(nn.Module):
             hidden_states
         )
 
-        if self.num_v_heads // self.num_k_heads in [1, 2, 4] and not _is_npu:
+        if (
+            self.num_v_heads // self.num_k_heads in _GDN_FUSED_QKVZBA_RATIOS
+            and not _is_npu
+        ):
             if _is_cpu:
                 num_k_heads_tp = self.num_k_heads // self.attn_tp_size
                 num_v_heads_tp = self.num_v_heads // self.attn_tp_size