Motivation

Optimize Gemma4 26B-A4B prefill performance through two complementary approaches:

1. Fused Triton kernels for Gemma4 decoder layers — Reduces kernel launch overhead by fusing multiple operations into single kernels.
2. Enable Piecewise CUDA Graph (PCG) for VLM models — Fixes PCG support for multimodal models that use self.language_model instead of self.model to reference their text backbone.

Main:

PR:

Modifications

Fused Triton Kernels (sglang/srt/layers/gemma4_fused_ops.py)

• gemma_rmsnorm_residual_scalar: Fuses RMSNorm + residual add + scalar multiply into a single Triton kernel (replaces 3 separate ops).
• gemma_dual_rmsnorm_residual_scalar: Fuses the full MoE post-processing pipeline — rmsnorm(rmsnorm(dense_out, w1) + rmsnorm(moe_out, w2), w3) + residual) * scalar — into a single kernel (replaces 5+ separate ops per MoE layer).

Gemma4 Model Integration (sglang/srt/models/gemma4_causal.py)

• Updated Gemma4DecoderLayer.forward() to use the fused kernels in the MoE block, eliminating redundant kernel launches.

PCG VLM Compatibility (sglang/srt/model_executor/model_runner.py)

• resolve_language_model(): Added fallback to model.language_model when model.model is not present, supporting VLM architectures like Gemma4ForConditionalGeneration.
• PCG gate check: Extended the guard condition to also check for language_model attribute, preventing false disabling of PCG for VLMs.
• Layer resolution: Added flexible layer discovery that handles both language_model.model.layers (standard) and language_model.layers (Gemma4 VLM where TextModel is the language_model directly).

PCG Runner Fix (sglang/srt/model_executor/piecewise_cuda_graph_runner.py)

• Updated patch_model target resolution to handle models where the TextModel (with .layers) is the language_model itself, not nested under .model.

Accuracy Tests

MMMU match official 54.9%.

Speed Tests and Profiling

PCG eliminates CPU dispatch overhead by capturing per-layer CUDA graphs for prefill. The benefit is most significant at small-to-medium token counts where kernel launch latency dominates compute time.

Compatibility

• Non-VLM models: No change — existing model.model path is tried first.
• VLM models with self.model (e.g., Qwen2.5-VL): No change — hasattr(model, "model") succeeds as before.
• VLM models with self.language_model (e.g., Gemma4): Now correctly supported.
• Verified correctness on both Gemma4 (text + image) and Qwen2.5-VL (text + image) with PCG enabled.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.