[AMD][Perf] Fuse QK RMSNorm + gate extraction Triton kernel for Qwen3.5 on HIP (#27656)

python/sglang/jit_kernel/tests/test_fused_qk_gemma_rmsnorm_gate.py

@@ -0,0 +1,151 @@
+import itertools
+import sys
+
+import pytest
+import torch
+
+from sglang.srt.models.utils import fused_qk_gemma_rmsnorm_with_gate
+from sglang.test.ci.ci_register import register_amd_ci
+
+register_amd_ci(est_time=20, suite="jit-kernel-unit-test-amd")
+
+
+def reference_qk_gemma_rmsnorm_with_gate(
+    q_gate: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    num_heads: int,
+):
+    """Pure-PyTorch reference: deinterleave q/gate, GemmaRMSNorm q and k."""
+    seq_len = q_gate.shape[0]
+
+    # Deinterleave q and gate from [q_h0, gate_h0, q_h1, gate_h1, ...]
+    qg_3d = q_gate.view(seq_len, num_heads, 2 * head_dim)
+    q = qg_3d[:, :, :head_dim].contiguous().view(-1, head_dim)
+    gate = qg_3d[:, :, head_dim:].contiguous().view(-1, head_dim)
+
+    k_flat = k.reshape(-1, head_dim)
+
+    # GemmaRMSNorm: x * rsqrt(mean(x^2) + eps) * (weight + 1)
+    def gemma_rmsnorm(x, w):
+        x_fp32 = x.float()
+        var = x_fp32.pow(2).mean(dim=-1, keepdim=True)
+        normed = x_fp32 * (var + eps).rsqrt() * (w.float() + 1.0)
+        return normed.to(x.dtype)
+
+    q_out = gemma_rmsnorm(q, q_weight)
+    k_out = gemma_rmsnorm(k_flat, k_weight)
+
+    return q_out, k_out, gate
+
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+SEQ_LENS = [1, 2, 4, 7, 16, 128]
+NUM_HEADS_LIST = [8, 16, 32]
+NUM_KV_HEADS_LIST = [2, 4, 8]
+HEAD_DIM_LIST = [64, 128]
+
+
+@pytest.mark.parametrize(
+    "seq_len,num_heads,num_kv_heads,head_dim",
+    list(itertools.product(SEQ_LENS, NUM_HEADS_LIST, NUM_KV_HEADS_LIST, HEAD_DIM_LIST)),
+)
+def test_fused_qk_gemma_rmsnorm_with_gate(
+    seq_len: int, num_heads: int, num_kv_heads: int, head_dim: int
+):
+    if num_kv_heads > num_heads:
+        pytest.skip("num_kv_heads > num_heads is not a valid config")
+
+    eps = 1e-6
+    q_size = num_heads * head_dim
+    kv_size = num_kv_heads * head_dim
+
+    # Build a full qkv buffer and split — this gives non-contiguous k,
+    # which is the real usage pattern
+    qkv = torch.randn(
+        seq_len, q_size * 2 + kv_size + kv_size, device=DEVICE, dtype=DTYPE
+    )
+    q_gate, k, v = qkv.split([q_size * 2, kv_size, kv_size], dim=-1)
+
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+
+    # Reference
+    q_ref, k_ref, gate_ref = reference_qk_gemma_rmsnorm_with_gate(
+        q_gate, k, q_weight, k_weight, eps, head_dim, num_heads
+    )
+
+    # Fused kernel
+    q_out, k_out, gate_out = fused_qk_gemma_rmsnorm_with_gate(
+        q_gate, k, q_weight, k_weight, eps, head_dim, num_heads
+    )
+
+    torch.testing.assert_close(q_out, q_ref, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(k_out, k_ref, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(gate_out, gate_ref, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("head_dim", [64, 128, 256])
+def test_gate_is_exact_copy(head_dim: int):
+    """Gate output must be a bitwise-exact copy of the input gate data."""
+    seq_len = 4
+    num_heads = 16
+    num_kv_heads = 4
+    eps = 1e-6
+    q_size = num_heads * head_dim
+    kv_size = num_kv_heads * head_dim
+
+    qkv = torch.randn(
+        seq_len, q_size * 2 + kv_size + kv_size, device=DEVICE, dtype=DTYPE
+    )
+    q_gate, k, v = qkv.split([q_size * 2, kv_size, kv_size], dim=-1)
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+
+    _, _, gate_out = fused_qk_gemma_rmsnorm_with_gate(
+        q_gate, k, q_weight, k_weight, eps, head_dim, num_heads
+    )
+
+    # Extract gate from interleaved buffer manually
+    qg_3d = q_gate.view(seq_len, num_heads, 2 * head_dim)
+    gate_expected = qg_3d[:, :, head_dim:].contiguous().view(-1, head_dim)
+
+    assert torch.equal(gate_out, gate_expected), "Gate must be bitwise exact"
+
+
+@pytest.mark.parametrize("seq_len", [1, 8])
+def test_contiguous_k_also_works(seq_len: int):
+    """Kernel should work even when k is already contiguous."""
+    num_heads = 16
+    num_kv_heads = 4
+    head_dim = 128
+    eps = 1e-6
+    q_size = num_heads * head_dim
+    kv_size = num_kv_heads * head_dim
+
+    q_gate = torch.randn(seq_len, q_size * 2, device=DEVICE, dtype=DTYPE)
+    k = torch.randn(seq_len, kv_size, device=DEVICE, dtype=DTYPE)
+    assert k.is_contiguous()
+
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+
+    q_ref, k_ref, gate_ref = reference_qk_gemma_rmsnorm_with_gate(
+        q_gate, k, q_weight, k_weight, eps, head_dim, num_heads
+    )
+    q_out, k_out, gate_out = fused_qk_gemma_rmsnorm_with_gate(
+        q_gate, k, q_weight, k_weight, eps, head_dim, num_heads
+    )
+
+    torch.testing.assert_close(q_out, q_ref, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(k_out, k_ref, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(gate_out, gate_ref, atol=0, rtol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

python/sglang/srt/models/qwen3_5.py

@@ -84,7 +84,10 @@ from sglang.srt.models.qwen2_moe import Qwen2MoeMLP, Qwen2MoeSparseMoeBlock
 
 # Models
 from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
-from sglang.srt.models.utils import fused_qk_gemma_rmsnorm
+from sglang.srt.models.utils import (
+    fused_qk_gemma_rmsnorm,
+    fused_qk_gemma_rmsnorm_with_gate,
+)
 from sglang.srt.server_args import get_global_server_args
 
 # Utils
@@ -889,6 +892,33 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
         q, k = self.rotary_emb(positions, q, k)
         return q, k, v, gate
 
+    def forward_prepare_hip(self, positions, hidden_states):
+        qkv, _ = self.qkv_proj(hidden_states)
+        if self.attn_output_gate:
+            q_gate, k, v = qkv.split(
+                [self.q_size * 2, self.kv_size, self.kv_size], dim=-1
+            )
+            seq_len = q_gate.shape[0]
+            q_flat, k_flat, gate_flat = fused_qk_gemma_rmsnorm_with_gate(
+                q_gate,
+                k,
+                self.q_norm.weight.data,
+                self.k_norm.weight.data,
+                self.q_norm.variance_epsilon,
+                self.head_dim,
+                self.num_heads,
+            )
+            q = q_flat.view(seq_len, -1)
+            k = k_flat.view(seq_len, -1)
+            gate = gate_flat.view(seq_len, -1)
+        else:
+            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+            gate = None
+            q, k = self._apply_qk_norm(q, k)
+
+        q, k = self.rotary_emb(positions, q, k)
+        return q, k, v, gate
+
     def forward_prepare_npu(self, positions, hidden_states, forward_batch):
         qkv, _ = self.qkv_proj(hidden_states)
         # Calculate first full attention layer ID based on config
@@ -916,7 +946,12 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
         """Full attention forward pass."""
-        if (
+        if _is_hip and self.attn_output_gate:
+            q, k, v, gate = self.forward_prepare_hip(
+                positions=positions,
+                hidden_states=hidden_states,
+            )
+        elif (
             not _is_npu
             or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
             or not self.attn_output_gate
@@ -1512,7 +1547,6 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
 
 
 class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration):
-
     packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping
     hf_to_sglang_mapper = None
 

python/sglang/srt/models/utils.py

@@ -154,9 +154,12 @@ class AutoWeightsLoader:
             for weight_name, weight_data in weights
         )
         for prefix, group in itertools.groupby(weights_by_parts, key=lambda x: x[0][0]):
-            yield prefix, (
+            yield (
+                prefix,
+                (
                     ("" if len(parts) == 1 else parts[1], weight_data)
                     for parts, weight_data in group
+                ),
             )
 
     @staticmethod
@@ -375,7 +378,6 @@ def compute_cu_seqlens_from_grid_numpy(grid_thw: torch.Tensor) -> torch.Tensor:
 
 
 class RotaryPosMixin:
-
     @staticmethod
     @lru_cache(maxsize=1024)
     def rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor:
@@ -594,5 +596,126 @@ def fused_qk_gemma_rmsnorm(
     return q_out, k_out
 
 
+# ---------------------------------------------------------------------------
+# Fused QK GemmaRMSNorm + gate extraction kernel
+# For models with attn_output_gate (e.g. Qwen3.5) where q and gate are
+# interleaved per head: [q_h0, gate_h0, q_h1, gate_h1, ...].
+# Reads q from the interleaved buffer, normalizes it, and copies gate to a
+# contiguous output — all in a single kernel launch.  Eliminates two
+# elementwise copy kernels that would otherwise be needed to deinterleave.
+# ---------------------------------------------------------------------------
+@triton.jit
+def _fused_qk_gemma_rmsnorm_gate_kernel(
+    QG_ptr,
+    K_ptr,
+    Q_out_ptr,
+    K_out_ptr,
+    Gate_out_ptr,
+    QW_ptr,
+    KW_ptr,
+    qg_token_stride,
+    qg_head_stride,
+    k_token_stride,
+    k_head_stride,
+    num_heads,
+    num_kv_heads,
+    k_rows,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_HD: tl.constexpr,
+    EPS: tl.constexpr,
+    FP16: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_HD)
+    mask = cols < HEAD_DIM
+    out_dtype = tl.float16 if FP16 else tl.bfloat16
+
+    token_idx = pid // num_heads
+    head_idx = pid % num_heads
+
+    base = token_idx * qg_token_stride + head_idx * qg_head_stride
+
+    # Q norm
+    q = tl.load(QG_ptr + base + cols, mask=mask, other=0.0).to(tl.float32)
+    w_q = tl.load(QW_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    q_var = tl.sum(q * q, axis=0) / HEAD_DIM
+    q_normed = (q * tl.rsqrt(q_var + EPS) * (w_q + 1.0)).to(out_dtype)
+    out_off = pid * HEAD_DIM + cols
+    tl.store(Q_out_ptr + out_off, q_normed, mask=mask)
+
+    # Gate copy
+    gate = tl.load(QG_ptr + base + HEAD_DIM + cols, mask=mask, other=0.0)
+    tl.store(Gate_out_ptr + out_off, gate, mask=mask)
+
+    # K norm (first k_rows blocks only)
+    if pid < k_rows:
+        token_idx_k = pid // num_kv_heads
+        head_idx_k = pid % num_kv_heads
+        k_off = token_idx_k * k_token_stride + head_idx_k * k_head_stride + cols
+        k = tl.load(K_ptr + k_off, mask=mask, other=0.0).to(tl.float32)
+        w_k = tl.load(KW_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+        k_var = tl.sum(k * k, axis=0) / HEAD_DIM
+        k_normed = (k * tl.rsqrt(k_var + EPS) * (w_k + 1.0)).to(out_dtype)
+        k_out_off = pid * HEAD_DIM + cols
+        tl.store(K_out_ptr + k_out_off, k_normed, mask=mask)
+
+
+def fused_qk_gemma_rmsnorm_with_gate(
+    q_gate: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    num_heads: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fused QK GemmaRMSNorm + gate extraction from interleaved q_gate buffer.
+
+    q_gate: (seq, q_size*2) where q and gate are interleaved per head,
+            i.e. [q_h0, gate_h0, q_h1, gate_h1, ...] with q_size = num_heads * head_dim.
+            Can be a non-contiguous view from qkv.split().
+    k: (seq, kv_size) — same as fused_qk_gemma_rmsnorm.
+
+    Returns (q_out, k_out, gate_out) all contiguous with shape
+    (seq*num_heads, head_dim), (seq*num_kv_heads, head_dim), (seq*num_heads, head_dim).
+    """
+    seq_len = q_gate.shape[0]
+    qg_3d = q_gate.view(seq_len, num_heads, 2 * head_dim)
+    num_kv_heads = k.shape[-1] // head_dim
+    k_3d = k.view(seq_len, num_kv_heads, head_dim)
+
+    q_rows = seq_len * num_heads
+    k_rows = seq_len * num_kv_heads
+
+    q_out = torch.empty(q_rows, head_dim, dtype=q_gate.dtype, device=q_gate.device)
+    k_out = torch.empty(k_rows, head_dim, dtype=k.dtype, device=k.device)
+    gate_out = torch.empty(q_rows, head_dim, dtype=q_gate.dtype, device=q_gate.device)
+
+    BLOCK_HD = triton.next_power_of_2(head_dim)
+
+    _fused_qk_gemma_rmsnorm_gate_kernel[(q_rows,)](
+        qg_3d,
+        k_3d,
+        q_out,
+        k_out,
+        gate_out,
+        q_weight,
+        k_weight,
+        qg_3d.stride(0),
+        qg_3d.stride(1),
+        k_3d.stride(0),
+        k_3d.stride(1),
+        num_heads,
+        num_kv_heads,
+        k_rows,
+        HEAD_DIM=head_dim,
+        BLOCK_HD=BLOCK_HD,
+        EPS=eps,
+        FP16=(q_gate.dtype == torch.float16),
+    )
+
+    return q_out, k_out, gate_out
+
+
 # Register the inplace op
 fused_inplace_qknorm = register_custom_op(fused_inplace_qknorm, mutates_args=["q", "k"])