diff --git a/python/sglang/jit_kernel/tests/test_fused_qk_gemma_rmsnorm_gate.py b/python/sglang/jit_kernel/tests/test_fused_qk_gemma_rmsnorm_gate.py new file mode 100644 index 000000000..94302b148 --- /dev/null +++ b/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"])) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 3d7cf345b..00801ec1b 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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 diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index b7b08b5ce..6c8b1ca8c 100644 --- a/python/sglang/srt/models/utils.py +++ b/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, ( - ("" if len(parts) == 1 else parts[1], weight_data) - for parts, weight_data in group + 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"])