[Gemma4] Optimize Gemm4 with fused Q/K/V RMSNorm + per-expert FP8 ckpt loader (#24696)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user