feat(spec-v2): Option-alpha scaffolding -- FROZEN_KV_MTP via EAGLE V2 (experimental, env-gated, NOT yet functional)

[pyc96]

Summary

Scaffolding for Option-α from the vLLM-vs-SGLang MTP comparison
(agent-pod/runs/20260525_mtp_v2/analysis/vllm_vs_sglang_mtp.md):
route FROZEN_KV_MTP to a thin subclass of EAGLEWorkerV2 instead
of the 787-LOC FrozenKVMTPWorker, mirroring vLLM's architectural
pattern (vllm/v1/spec_decode/gemma4.py:31-340 — also 335 LOC).

Behavior is unchanged by default: FROZEN_KV_MTP still uses V1
with disable_overlap_schedule=True. The new EAGLE-V2 path is
gated behind SGLANG_GEMMA4_MTP_VIA_EAGLE=1.

The EAGLE-V2 path is NOT yet functional — it crashes during CUDA
graph capture on a hidden-state shape mismatch. This PR ships the
scaffolding + the detailed analysis of what remains; it is
documentation-quality code for the next implementer, not a working
overlap path.

Stack base

pyc/feat-gemma4-ultimate-v2 (PR #21).

What ships

File	Change
python/sglang/srt/speculative/gemma4_mtp_via_eagle.py (NEW)	Gemma4MTPEagleWorker (subclass of EAGLEWorkerV2): calls bind_frozen_kv_context on the draft after init; monkey-patches draft_forward to snapshot+restore positions around each draft loop (Gemma-4 MTP analog of vLLM's constant_draft_positions=True).
python/sglang/srt/speculative/spec_info.py	Dispatcher routes FROZEN_KV_MTP to Gemma4MTPEagleWorker when SGLANG_GEMMA4_MTP_VIA_EAGLE=1 AND overlap enabled. supports_spec_v2() now returns True for FROZEN_KV_MTP.
python/sglang/srt/arg_groups/speculative_hook.py	_handle_frozen_kv_mtp still forces disable_overlap_schedule=True by default; only skips the override when the env var is set.

The remaining work (what stops this from being ship-ready)

Three sequential blockers:

Blocker 1 — spec_hidden_size mismatch (the immediate crash)

The Gemma-4 MTP assistant model expects forward_batch.spec_info.hidden_states (the recurrent prev_hidden ferried across draft iterations) to be sized [N, backbone_hidden_size] = [N, 5376] for the 31B-it target. EAGLE V2 sizes its hidden_states buffers via EagleDraftInput.hidden_size_for(worker) which returns draft_runner.model_config.spec_hidden_size = the draft model's hidden_size = 1024.

Crash during CUDA graph capture:

File "gemma4_mtp.py", line 253, in forward
    raise ValueError(
ValueError: Frozen-KV MTP forward: token_embed and prev_hidden must
have the same shape (got torch.Size([80, 5376]) vs torch.Size([80, 1024])).

Fix options (any one):

1. Override draft_runner.model_config.spec_hidden_size = target_backbone_hidden_size before EagleDraftWorker.__init__ runs (currently it runs inside EAGLEWorkerV2.__init__ so the subclass can't intercept). Requires touching how the parent constructs the draft worker.
2. Patch Gemma4AssistantForCausalLM to expose a spec_hidden_size class attribute that ModelConfig reads. ~10 LOC in two files.
3. Add an EagleDraftInput.hidden_size_for override hook so the model can declare its required hidden state size.

Blocker 2 — return value contract

Even with the input shape fixed, the assistant's forward() returns LogitsProcessorOutput.hidden_states sized at the internal hidden_size=1024 (the inner Gemma4TextModel output, before post_projection). EAGLE V2 takes this output and stores it as the next iter's spec_info.hidden_states. We need the assistant to instead return post_projection(hidden_states) = [N, 5376] so the recurrent ferry stays at backbone size.

Currently the assistant passes hidden_states_before_norm=projected_states (=post_projection output) to LogitsProcessor, which preserves it as output.hidden_states_before_norm. The fix: swap which one goes into output.hidden_states for the EAGLE-V2 path.

Blocker 3 — constant_draft_positions correctness validation

My current implementation monkey-patches draft_worker.draft_forward to snapshot positions before the EAGLE draft loop and restore after. But INSIDE the loop, EAGLE V2 still does forward_batch.positions.add_(1) at line 493 of eagle_worker_v2.py between draft steps. The attention metadata is built once for draft_index=0 (outside the loop) so each draft step still calls into a kernel that sees incremented positions. Whether this produces correct outputs depends on whether the attention backend reads positions directly or only via the cached metadata. Needs validation once Blocker 1 is past.

The full vLLM approach (SpecDecodeBaseProposer.propose at llm_base_proposer.py:570-587) skips the position-update entirely when constant_draft_positions=True. Equivalent SGLang fix: add a self.constant_draft_positions: bool = False flag to EagleDraftWorker.draft_forward and gate the positions.add_(1) call (~3 LOC change in eagle_worker_v2.py:493). Low-risk if confined to spec-v2 path with explicit opt-in.

Why we shipped scaffolding instead of the full fix

Original time estimate (from the comparison report): 5-10 days for Option-α. Investigation today confirmed this is accurate — the architectural mismatch between EAGLE V2 and Gemma-4 MTP runs three layers deep (init-time buffer sizing, runtime recurrent-state shape, scheduler position-advancement semantics).

Doing this surgery correctly in this session would require:

• Hands-on familiarity with EAGLE V2's CUDA graph capture lifecycle
• Validation matrix: V1 EAGLE (Llama-3/Qwen) + V1 EAGLE3 + V1 FROZEN_KV_MTP + the new V2-EAGLE-MTP path, since changes touch shared code
• 2-3 days minimum of careful iteration

Decision: ship the scaffolding + analysis as infrastructure for the next implementer, preserve V1 default behavior, and document the three remaining blockers concretely.

Verified V1 default path still works

$ curl -sf -X POST http://127.0.0.1:30000/v1/chat/completions \
    -d '{"model":"google/gemma-4-31B-it",...,"What is 2+2?"...}'
Response: 2 + 2 = 4

Server log:

[warn] Overlap scheduler is disabled for Frozen-KV MTP (V1 path;
       set SGLANG_GEMMA4_MTP_VIA_EAGLE=1 to try the experimental EAGLE-V2 path).

Reference

• vLLM's Gemma4Proposer (the architecture this PR mirrors): vllm/v1/spec_decode/gemma4.py:31-335
• vLLM's SpecDecodeBaseProposer.propose with constant_draft_positions: vllm/v1/spec_decode/llm_base_proposer.py:421-600
• Full architectural comparison: agent-pod/runs/20260525_mtp_v2/analysis/vllm_vs_sglang_mtp.md

For production today

Stick with the recommendation from PR #21: use SGLang no-MTP. The MTP gap to vLLM remains until either:

• PR feat(spec-v2): Frozen-KV MTP overlap-scheduling worker (experimental opt-in) + V1 None-safety fixes #23's V2 worker leak is fixed (Option γ, ~1-3 days), or
• The 3 blockers above are resolved on this branch (Option α, ~2-5 days from here)

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.