[Gemma4] Optimize Gemm4 with fused Q/K/V RMSNorm + per-expert FP8 ckpt loader (#24696)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

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,

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,20 +341,62 @@ 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)
 
+        # 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:
             q = q.unflatten(-1, (self.num_heads, self.head_dim))
             q = self.q_norm(q)
             q = q.flatten(-2, -1)
-
+            if is_kv_shared:
-        # 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
             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)
 

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.<id>.gate_proj.{weight,weight_scale}
+            #   experts.<id>.up_proj.{weight,weight_scale}
+            #   experts.<id>.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

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.<id>.{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 `<pad>` 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-`<pad>` 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 `<pad>` 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(
+            "<pad>", 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()