Repository navigation
[Gemma4] Optimize Gemm4 with fused Q/K/V RMSNorm + per-expert FP8 ckpt loader - #24696
Conversation
|
/tag-and-rerun-ci |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fused Triton kernel for QKV RMSNorm in Gemma4 models, improving efficiency by replacing three separate normalization passes with a single kernel. It also updates the weight loading logic in the multimodal Gemma4 implementation to support per-expert FP8 checkpoint formats, ensuring correct mapping of expert projections to fused parameters. I have no feedback to provide.
pyc96
left a comment
pyc96
left a comment
There was a problem hiding this comment.
LGTM
Could you also verify quality scores? e.g. MMLU
| # routed-expert weight, leading to all-`<pad>` generation. | ||
| # Warn loudly so a future ckpt-format drift surfaces here | ||
| # instead of silently producing garbage outputs again. | ||
| logger.warning( |
There was a problem hiding this comment.
Do we need the logs here? If the params were not loaded, we should log it in https://github.com/sgl-project/sglang/pull/24696/changes#diff-1131c6a30a9629df571ecba59b1afafdcc4ce164cc45a41f95d78e8837c99c8bL837
There was a problem hiding this comment.
Good point. Removed the inline logger.warning; misses now fall through to the existing centralized unloaded_params check, same as the other loader paths.
@pyc96 Updated mmlu scores. |
* main: (87 commits) [Fix] Disable FlashInfer allreduce fusion under deterministic inference (sgl-project#24629) fix: STANDALONE spec-decode hidden-size mismatch crash (sgl-project#24217) Followup fix for Custom AR V2 in non NVL scenarios (sgl-project#24742) Fix reduce_scatterv producer contract for SUM_LEN (sgl-project#24785) [NPU]Documentation update for communications quantization feature (sgl-project#24668) [Session R3] Add routed_experts_start_len for absolute routing slice control (sgl-project#24851) [Model] Add MiniCPM-V 4.6 support (sgl-project#24855) Support Intern-S2-Preview (sgl-project#24875) [PD] Unify dsv4 dispatch with swa (sgl-project#24888) Optimize MHC pipeline: DeepGemm, fused norm, fused hc_head (sgl-project#24775) Fix PD bootstrap failure handling (sgl-project#24772) [Spec] Cleanup idle stub and shape-check patterns (sgl-project#24881) [Bug] Add dsv4 state_type branch to mooncake disaggregation (sgl-project#24878) [Spec V1] Split draft-extend phase from `EagleDraftInput` into new `EagleDraftExtendInput` (sgl-project#24859) [Gemma4] Optimize Gemm4 with fused Q/K/V RMSNorm + per-expert FP8 ckpt loader (sgl-project#24696) [spec decoding] support kimi-k2.5-eagle3-mla (sgl-project#24826) [SPEC V2] fix: skip stale state updates in spec-v2 overlap (sgl-project#23456) [RL] Call torch.cuda.empty_cache() for `in-place` pause mode to avoid OOM (sgl-project#24854) [diffusion] CI: add cache-dit CI tests (sgl-project#19213) [Utils] Make request dump robust to unpicklable server_args and large meta_info (sgl-project#24767) ... # Conflicts: # python/sglang/srt/utils/common.py
…t loader (sgl-project#24696) Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Motivation
The original motivation is to speedup Gemm4 model. Two unrelated issues on the Gemma4 FP8-MoE path:
Attention RMSNorm is launch-bound. Each Gemma4 attention layer issues three separate sgl-kernel
rmsnormlaunches (q/k/v). At small token counts the per-launch host overhead dominates — three back-to-back launches cost ~60μs regardless of M, even though the actual compute is trivial.compressed-tensors per-expert FP8 checkpoints fail silently. Models like
RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamicstore one weight per expert (experts.<id>.{gate,up,down}_proj.{weight,weight_scale}). The existing loader only recognizes the stackedexperts.gate_up_proj/experts.down_projlayout, so every routed expert is silently skipped → zero-init expert weights → server launches fine but generation collapses to all-<pad>and GSM8K = 0.0.Design
1. Fused Q/K/V RMSNorm (
gemma_qkv_rmsnorm)x · rsqrt(mean(x²) + eps) · wfor all q/k/v heads in a single launch, in-place into the qkv buffer.(M,),BLOCK = next_pow2(head_dim), withtl.static_rangeover heads.qkv.split(...)strided views viastride_*_m— no.contiguous()copy.with_scale=False→ weight is identically 1).Gemma4Attention.forward: only takes the fused path when CUDA +q_norm.scale_shift == 0 == k_norm.scale_shift+not v_norm.with_scale. Non-standard configs and KV-shared layers retain the original three-normpath.
2. Per-expert FP8 loader
Gemma4ForConditionalGeneration.load_weightsfor^.*\.moe\.experts\.<id>\.(gate|up|down)_proj\.(weight|weight_scale)$.param.weight_loader(standard FusedMoE signature).logger.warninginstead of silently skipping — this prevents reproducing the very failure mode the patch fixes.3. Test
test/registered/models/test_gemma4_fp8_per_expert_loading.py(CI:stage-c-test-4-gpu-h100, est 120s) launches the FP8 model with TP=4 and asserts: server health, non-<pad>chat output, and GSM8K-20 ≥ 0.80 (BF16 baseline ~0.95). All three checks fail pre-fix.Accuracy
MMLU 500 examples
BF16 （500 examples, TP=4）：0.878
RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamic
Benchmark
H100 ×4, TP=4,
RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamic, batch=1, in=2048, out=64.Microbench (single H100, bf16, Q=8h, KV=2h, head_dim=128):
Bit-exact vs unfused at M=1/8/128 (max abs diff = 0).
End-to-end (4-GPU, RedHatAI FP8):
Decode benefit is muted because the SGLang decode path runs inside CUDA Graph, which amortizes the launch overhead the fused kernel saves. Prefill (no graph) realizes the full 2% win. Net change is neutral-to-positive across regimes; never negative within the tested M range.
The per-expert FP8 loader is a correctness fix: pre-fix GSM8K-20 = 0.00, post-fix ≈ 0.95 — gated at ≥ 0.80 in the new test.
Test plan
test_gemma4_fp8_per_expert_loading.pyexercises the loader fix end-to-end on 4 GPUs