diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 76b741df3..48b7f7e1d 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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) diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 13a0e70f0..eb4391fab 100644 --- a/python/sglang/srt/models/utils.py +++ b/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"])