diff --git a/python/sglang/srt/layers/gemma4_fused_ops.py b/python/sglang/srt/layers/gemma4_fused_ops.py index eeb9ebcea..ad6f01d98 100644 --- a/python/sglang/srt/layers/gemma4_fused_ops.py +++ b/python/sglang/srt/layers/gemma4_fused_ops.py @@ -4,6 +4,8 @@ Fuses standard RMSNorm + residual-add (+ optional scalar multiply) into a single kernel pass to reduce kernel launch overhead. """ +from typing import Optional + import torch import triton import triton.language as tl @@ -130,6 +132,119 @@ def _gemma_dual_rmsnorm_residual_kernel( tl.store(Out_ptr + row * stride_o + cols, out.to(x1.dtype), mask=mask) +@triton.jit +def _gemma_qkv_rmsnorm_kernel( + Q_ptr, + K_ptr, + V_ptr, + Q_w_ptr, + K_w_ptr, + stride_q_m, + stride_k_m, + stride_v_m, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + eps, + HAS_KV: tl.constexpr, + BLOCK: tl.constexpr, +): + """Per-token fused RMSNorm of Q (with q_w), K (with k_w), V (no scale). + + Layout assumption: each tensor's last dim packs (num_heads, head_dim) contiguously + so per-head offset is `h * HEAD_DIM`. The token (M) stride is taken from + stride_*_m so the kernel works on strided views (e.g. slices of a larger + qkv buffer produced by `qkv.split`) without requiring `.contiguous()` copies. + V uses `weight=ones` semantics so the multiply-by-weight is omitted. + """ + m = tl.program_id(0) + cols = tl.arange(0, BLOCK) + mask = cols < HEAD_DIM + + qw = tl.load(Q_w_ptr + cols, mask=mask, other=0.0).to(tl.float32) + + # Q heads + for h in tl.static_range(NUM_Q_HEADS): + off = m * stride_q_m + h * HEAD_DIM + cols + x = tl.load(Q_ptr + off, mask=mask, other=0.0).to(tl.float32) + rrms = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps) + out = x * rrms * qw + tl.store(Q_ptr + off, out.to(Q_ptr.dtype.element_ty), mask=mask) + + if HAS_KV: + kw = tl.load(K_w_ptr + cols, mask=mask, other=0.0).to(tl.float32) + + # K heads + for h in tl.static_range(NUM_KV_HEADS): + off = m * stride_k_m + h * HEAD_DIM + cols + x = tl.load(K_ptr + off, mask=mask, other=0.0).to(tl.float32) + rrms = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps) + out = x * rrms * kw + tl.store(K_ptr + off, out.to(K_ptr.dtype.element_ty), mask=mask) + + # V heads (no scaling: V-norm uses weight=ones) + for h in tl.static_range(NUM_KV_HEADS): + off = m * stride_v_m + h * HEAD_DIM + cols + x = tl.load(V_ptr + off, mask=mask, other=0.0).to(tl.float32) + rrms = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps) + out = x * rrms + tl.store(V_ptr + off, out.to(V_ptr.dtype.element_ty), mask=mask) + + +def gemma_qkv_rmsnorm( + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + q_weight: torch.Tensor, + k_weight: Optional[torch.Tensor], + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + eps: float = 1e-6, +) -> None: + """In-place fused RMSNorm on Q, K, V for Gemma4 attention. + + All three norms compute `x * rsqrt(mean(x^2) + eps)` independently per head. + Q is scaled by `q_weight`, K by `k_weight`, V by 1 (Gemma4's V-norm has + `with_scale=False`). + + Inputs may be 2D `(M, num_heads * head_dim)` or strided views of a larger + buffer (such as q/k/v slices from `qkv.split`). The kernel uses the actual + `stride(0)` so no `.contiguous()` copy is required. Within a token, the + last dim must be contiguous so heads pack as `h * head_dim` offsets. + + If k and v are both None (KV-shared layer), only Q is normalized. + """ + assert q.is_cuda + assert q.stride(-1) == 1, "Q's last dim must be contiguous" + assert q_weight.shape[-1] == head_dim + M = q.shape[0] if q.dim() >= 2 else 1 + BLOCK = triton.next_power_of_2(head_dim) + + has_kv = k is not None and v is not None + if has_kv: + assert k.is_cuda and v.is_cuda + assert k.stride(-1) == 1 and v.stride(-1) == 1 + assert k_weight is not None and k_weight.shape[-1] == head_dim + + _gemma_qkv_rmsnorm_kernel[(M,)]( + q, + k if has_kv else q, + v if has_kv else q, + q_weight, + k_weight if has_kv else q_weight, + q.stride(0), + k.stride(0) if has_kv else 0, + v.stride(0) if has_kv else 0, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads if has_kv else 0, + HEAD_DIM=head_dim, + eps=eps, + HAS_KV=has_kv, + BLOCK=BLOCK, + ) + + def gemma_dual_rmsnorm_residual_scalar( x1: torch.Tensor, weight1: torch.Tensor, diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index 901d1a7de..9c04746d2 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -30,6 +30,7 @@ from sglang.srt.distributed import ( ) from sglang.srt.layers.gemma4_fused_ops import ( gemma_dual_rmsnorm_residual_scalar, + gemma_qkv_rmsnorm, gemma_rmsnorm_residual_scalar, ) from sglang.srt.layers.layernorm import Gemma4RMSNorm, RMSNorm @@ -340,22 +341,64 @@ class Gemma4Attention(nn.Module): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = q.unflatten(-1, (self.num_heads, self.head_dim)) - q = self.q_norm(q) - q = q.flatten(-2, -1) - - # Check if we should use shared KV cache - if self.is_kv_shared_layer and self.kv_shared_layer_index is not None: - # For KV shared layers, we skip K/V computation and normalization - # The RadixAttention will handle retrieving shared KV from cache - k = None - v = None + # Fused Q/K/V RMSNorm: replaces three separate norm kernels with one. + # Preconditions for the fused path: tensors on CUDA, q_norm/k_norm use + # the standard norm*weight (scale_shift==0) and v_norm has weight=ones + # (with_scale=False) — the canonical Gemma4 attention configuration. + is_kv_shared = ( + self.is_kv_shared_layer and self.kv_shared_layer_index is not None + ) + can_fuse_qkv_norm = ( + q.is_cuda + and self.q_norm.scale_shift == 0.0 + and self.k_norm.scale_shift == 0.0 + and not self.v_norm.with_scale + ) + if can_fuse_qkv_norm: + if is_kv_shared: + gemma_qkv_rmsnorm( + q, + None, + None, + self.q_norm.weight.data, + None, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + eps=self.q_norm.eps, + ) + k = None + v = None + else: + gemma_qkv_rmsnorm( + q, + k, + v, + self.q_norm.weight.data, + self.k_norm.weight.data, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + eps=self.q_norm.eps, + ) + # Match the original norm path's output shapes: q stays 2D, + # k/v become 3D so the subsequent `.flatten(-2, -1)` works. + # Use reshape (not view) since k/v are strided slice views of + # the qkv buffer and may not satisfy view's contiguity rules. + k = k.reshape(-1, self.num_kv_heads, self.head_dim) + v = v.reshape(-1, self.num_kv_heads, self.head_dim) else: - k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) - k = self.k_norm(k) - - v = v.unflatten(-1, (self.num_kv_heads, self.head_dim)) - v = self.v_norm(v) + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + q = self.q_norm(q) + q = q.flatten(-2, -1) + if is_kv_shared: + k = None + v = None + else: + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) + k = self.k_norm(k) + v = v.unflatten(-1, (self.num_kv_heads, self.head_dim)) + v = self.v_norm(v) # Apply rotary embedding if k is not None: diff --git a/python/sglang/srt/models/gemma4_mm.py b/python/sglang/srt/models/gemma4_mm.py index 479043b73..4ce8a5909 100644 --- a/python/sglang/srt/models/gemma4_mm.py +++ b/python/sglang/srt/models/gemma4_mm.py @@ -802,6 +802,41 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): and int(m.group(1)) in k_eq_v_layers ) + # Per-expert checkpoint format used by compressed-tensors / FP8 + # (e.g. RedHatAI/*-FP8-Dynamic). Each expert is stored as a + # separate key with shape (out, in): + # experts..gate_proj.{weight,weight_scale} + # experts..up_proj.{weight,weight_scale} + # experts..down_proj.{weight,weight_scale} + # These need to be folded into sglang's fused FusedMoE params: + # experts.w13_weight[_scale] (gate->shard "w1", up->shard "w3") + # experts.w2_weight[_scale] (down->shard "w2") + per_expert_match = re.match( + r"^(.*?\.moe\.experts\.)(\d+)\.(gate_proj|up_proj|down_proj)" + r"\.(weight|weight_scale)$", + name, + ) + if per_expert_match: + prefix = per_expert_match.group(1) + expert_id = int(per_expert_match.group(2)) + proj = per_expert_match.group(3) + suffix = per_expert_match.group(4) + if proj == "gate_proj": + base, sid = "w13_weight", "w1" + elif proj == "up_proj": + base, sid = "w13_weight", "w3" + else: # down_proj + base, sid = "w2_weight", "w2" + if suffix == "weight_scale": + base += "_scale" + fused_name = prefix + base + if fused_name in params_dict: + param = params_dict[fused_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, fused_name, sid, expert_id) + loaded_params.add(fused_name) + continue + # MoE expert weights checked first (gate_up_proj contains "up_proj" # which would false-match the stacked dense MLP mapping). orig_name = name diff --git a/test/registered/models/test_gemma4_fp8_per_expert_loading.py b/test/registered/models/test_gemma4_fp8_per_expert_loading.py new file mode 100644 index 000000000..3565263fb --- /dev/null +++ b/test/registered/models/test_gemma4_fp8_per_expert_loading.py @@ -0,0 +1,109 @@ +"""End-to-end test for compressed-tensors per-expert FP8 MoE checkpoint +loading on Gemma4 (e.g. RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamic). + +Regression coverage for the load_weights path that recognises +`experts..{gate,up,down}_proj.{weight,weight_scale}` keys and folds +them into SGLang's fused FusedMoE parameters. Without that path, all +routed-expert weights are silently skipped at load time and the model +emits only `` tokens at inference (GSM8K collapses to 0.0). +""" + +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.utils import get_device_sm, kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +# Compressed-tensors per-expert FP8 MoE checkpoint that exercises the +# loader path (gated repo + ~27 GB download + 4 GPUs at TP=4). +register_cuda_ci(est_time=120, suite="stage-c-test-4-gpu-h100") + + +@unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher") +class TestGemma4FP8PerExpertLoading(CustomTestCase): + """Three-stage check that catches the silent-skip failure mode: + 1. server health + 2. completion is not the all-`` garbage state + 3. GSM8K accuracy matches the BF16 baseline + """ + + model = "RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamic" + base_url = DEFAULT_URL_FOR_TEST + + @classmethod + def setUpClass(cls): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp", + "4", + "--trust-remote-code", + "--random-seed", + "42", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_health(self): + r = requests.get(self.base_url + "/health") + self.assertEqual(r.status_code, 200) + + def test_basic_generation_not_garbage(self): + """Pre-fix the server starts but every routed expert is zero-init, + which leads chat completions to deterministic `` spam.""" + r = requests.post( + self.base_url + "/v1/chat/completions", + json={ + "model": self.model, + "messages": [{"role": "user", "content": "What is 7 + 5?"}], + "temperature": 0, + "max_tokens": 32, + }, + ) + self.assertEqual(r.status_code, 200) + text = r.json()["choices"][0]["message"]["content"] + self.assertNotIn( + "", text, f"Output looks like the pre-fix garbage state: {text!r}" + ) + self.assertGreater(len(text.strip()), 0, "Empty completion") + self.assertIn("12", text, f"Expected the answer to mention '12': {text!r}") + + def test_gsm8k_accuracy(self): + """Pre-fix this scores exactly 0.00 (zero routed-expert weights); + post-fix it matches the BF16 baseline (~0.95 on 20 samples).""" + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + num_examples=20, + num_threads=16, + ) + metrics = run_eval(args) + score = float(metrics["score"]) + print(f"Gemma4 FP8 per-expert GSM8K-20 score: {score:.3f}") + # Threshold rules out the failure mode (0.00) while leaving ample + # margin under the BF16 baseline (~0.95). + self.assertGreaterEqual( + score, + 0.80, + f"Per-expert FP8 ckpt accuracy collapsed: {score} " + "(pre-fix value is 0.00; BF16 baseline is ~0.95).", + ) + + +if __name__ == "__main__": + unittest.main()