[Gemma4] Add test for MTP models

[kpham-sgl]

Motivation

Add test for PR #24436. Require transformer v5.8.0

Modifications

Accuracy Tests

Speed Tests and Profiling

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.

---

CI States

Latest PR Test (Base): ❌ Run #26135220858
Latest PR Test (Extra): ❌ Run #26135220756

[gemini-code-assist]

Code Review

This pull request implements Frozen-KV MTP (Multi-Token Prediction) speculative decoding for Gemma 4 models, introducing a new assistant model class, worker, and CUDA graph runner. The review highlights critical issues regarding Tensor Parallelism compatibility in embedding and linear layers, suggests avoiding torch.cuda.empty_cache() to prevent performance degradation, recommends copying server_args for isolation, and advises using specific exceptions for better error handling.

I am having trouble creating individual review comments. Click here to see my feedback.

python/sglang/srt/models/gemma4_mtp.py (234)

Using torch.nn.functional.embedding directly on self.target_embed_weight will fail when Tensor Parallelism (TP) is enabled (TP > 1). In SGLang, target_embed_weight is typically a partitioned weight from a VocabParallelEmbedding. A direct lookup will only work for tokens within the local GPU's vocabulary range and will not perform the necessary communication (all-reduce) to produce the full embedding vector on all GPUs. This will lead to incorrect results or IndexError on GPUs where the token IDs are out of the local partition range.

python/sglang/srt/models/gemma4_mtp.py (111)

Using nn.Linear for lm_head bypasses Tensor Parallelism. For large vocabulary sizes (e.g., Gemma's 256k), this replicates a large weight matrix (approx. 1.5GB for bf16) on every GPU, which is memory-inefficient and inconsistent with SGLang's standard use of ColumnParallelLinear for output heads. This also applies to the centroids layer on line 125.

python/sglang/srt/models/gemma4_mtp.py (211)

Calling torch.cuda.empty_cache() inside a model method is generally discouraged as it can cause significant performance overhead due to GPU synchronization and fragmentation. It is better to let the high-level memory manager or the user handle cache clearing if necessary.

python/sglang/srt/speculative/frozen_kv_mtp_worker.py (117)

Modifying server_args.context_length in-place can have unintended side effects if the server_args object is shared with other components (like the target worker). It is safer to create a copy of the arguments for the draft worker to ensure isolation.

        import copy
        server_args = copy.copy(server_args)
        server_args.context_length = target_worker.model_runner.model_config.context_len

python/sglang/srt/server_args.py (3465)

Hardcoding max_running_requests to 48 when using FROZEN_KV_MTP seems arbitrary and might be too restrictive for some hardware configurations. Consider making this a configurable default or providing a more detailed rationale for this specific limit.

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py (146-149)

Raising a generic Exception is discouraged. Using a more specific exception like RuntimeError is preferred for better error handling and clarity.

        except RuntimeError as e:
            raise RuntimeError(
                f"Capture frozen-KV MTP cuda graph failed: {e}\n"
                f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
            ) from e

[kpham-sgl]

/tag-run-ci-label

[kpham-sgl]

/rerun-failed-ci