minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-05-10 00:24:12 -07:00
committed by GitHub
co-authored by luoyuan.luo
parent a87fb399de
commit d3fd91ed97
4 changed files with 317 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/gemma4_fused_ops.py
+115
View File
@@ -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
+58 -15
View File
@@ -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:
python/sglang/srt/models/gemma4_mm.py
+35
View File
@@ -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
+109
View File
@@ -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()