[AMD] fused qk gemma norm kernels to reduce four kernels (#23575)

Co-authored-by: root <root@smci355-ccs-aus-g12-26.cs-aus.dcgpu>

python/sglang/srt/models/qwen3_5.py

@@ -79,6 +79,7 @@ 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.server_args import get_global_server_args
 
 # Utils
@@ -106,7 +107,6 @@ _is_hip = is_hip()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 _is_amx_available = cpu_has_amx_support()
 
-
 cached_get_processor = lru_cache(get_processor)
 
 
@@ -791,6 +791,15 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
                 k_by_head = k.reshape(-1, self.head_dim)
                 k_by_head = self.k_norm(k_by_head)
             current_stream.wait_stream(self.alt_stream)
+        elif _is_hip:
+            q_by_head, k_by_head = fused_qk_gemma_rmsnorm(
+                q,
+                k,
+                self.q_norm.weight.data,
+                self.k_norm.weight.data,
+                self.q_norm.variance_epsilon,
+                self.head_dim,
+            )
         else:
             q_by_head = q.reshape(-1, self.head_dim)
             q_by_head = self.q_norm(q_by_head)

python/sglang/srt/models/utils.py

@@ -21,6 +21,8 @@ from typing import TYPE_CHECKING, Any, Optional, Tuple
 
 import numpy as np
 import torch
+import triton
+import triton.language as tl
 
 from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
 from sglang.srt.environ import envs
@@ -453,5 +455,98 @@ def apply_qk_norm(
     return q, k
 
 
+# ---------------------------------------------------------------------------
+# Fused QK GemmaRMSNorm Triton kernel
+# grid = q_rows (the larger dimension in GQA).  Every block computes Q norm
+# for its row; the first k_rows blocks also compute K norm.  No torch.cat,
+# no tl.where for weight selection, no output slice.
+# ---------------------------------------------------------------------------
+@triton.jit
+def _fused_qk_gemma_rmsnorm_kernel(
+    Q_ptr,
+    K_ptr,
+    Q_out_ptr,
+    K_out_ptr,
+    QW_ptr,
+    KW_ptr,
+    q_stride,
+    k_stride,
+    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
+
+    # Q norm (every block) — use q_stride to handle non-contiguous input
+    q_off = pid * q_stride + cols
+    q = tl.load(Q_ptr + q_off, 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)
+    # output is always contiguous
+    q_out_off = pid * HEAD_DIM + cols
+    tl.store(Q_out_ptr + q_out_off, q_normed, mask=mask)
+
+    # K norm (first k_rows blocks only) — use k_stride for input
+    if pid < k_rows:
+        k_off = pid * k_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(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+    head_dim: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fused QK GemmaRMSNorm — single Triton kernel for both q_norm and k_norm.
+
+    grid = q_rows; every block processes its Q row, and the first k_rows
+    blocks also process K.  No torch.cat, no slice, no tl.where.
+    Passes input strides to the kernel so non-contiguous tensors (e.g. from
+    qkv.split()) are read correctly without an extra .contiguous() copy.
+    """
+    q_flat = q.reshape(-1, head_dim)
+    k_flat = k.reshape(-1, head_dim)
+
+    q_rows = q_flat.shape[0]
+    k_rows = k_flat.shape[0]
+
+    q_out = torch.empty(q_rows, head_dim, dtype=q.dtype, device=q.device)
+    k_out = torch.empty(k_rows, head_dim, dtype=k.dtype, device=k.device)
+
+    BLOCK_HD = triton.next_power_of_2(head_dim)
+
+    _fused_qk_gemma_rmsnorm_kernel[(q_rows,)](
+        q_flat,
+        k_flat,
+        q_out,
+        k_out,
+        q_weight,
+        k_weight,
+        q_flat.stride(0),
+        k_flat.stride(0),
+        k_rows,
+        HEAD_DIM=head_dim,
+        BLOCK_HD=BLOCK_HD,
+        EPS=eps,
+        FP16=(q.dtype == torch.float16),
+    )
+
+    return q_out, k_out
+
+
 # Register the inplace op
 fused_inplace_qknorm = register_custom_op(fused_inplace_qknorm, mutates_args=["q", "k"])