From b7f87a2513c2761952ca018c7b98b144cecd3b09 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 13 Aug 2026 15:18:35 +0800 Subject: [PATCH] [AMD][Perf] Fuse GatedDeltaNet QKVZBA split/reshape/cat into a single Triton kernel for Qwen3.5-architecture MoE on HIP (#34421) --- .../ops/attention/triton_gdn_fused_proj.py | 17 ++++++++++++++++- python/sglang/srt/models/qwen3_5.py | 12 +++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py index 076a1cc3f..8e39197ba 100644 --- a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py +++ b/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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 40d1810f9..a2fb07c8a 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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