Support spec v2 for Frozen-KV MTP; remove v1 worker by hnyls2002 · Pull Request #27607 · sgl-project/sglang

hnyls2002 · 2026-06-08T23:38:45Z

Summary

Run Frozen-KV MTP on the spec v2 worker — FrozenKVMTPWorkerV2 (orchestrator, subclasses EAGLEWorkerV2) + FrozenKVMTPDraftWorker (draft layer) — supporting both overlap and non-overlap scheduling
Remove the v1 FrozenKVMTPWorker
Add spec_stage_span nvtx markers on the orchestrator (parity with the v1 worker / EAGLEWorkerV2)

Mechanical verification

The draft layer is a byte-faithful move of the v1 FrozenKVMTPWorker methods into FrozenKVMTPDraftWorker. AST move-equivalence (reproducible script + results): https://gist.github.com/hnyls2002/08fb69f9857f011750ba2455816701ad

19/19 pure-move methods byte-equal (incl. draft_forward and every frozen-KV metadata/view helper)
__init__: only super().__init__ -> TpModelWorker.__init__, two None inits, and assert-message text
draft: only the penalty-cumulate block removed (now handled by the inherited EagleDraftInputV2Mixin.prepare_for_decode, matching EAGLE v2)

Validation

GSM8K first-200, gemma-4-E4B-it + assistant: v1 0.725 vs v2 overlap 0.730 vs v2 sync 0.730; overall accept length ~2.83 across all three paths
26B-A4B, deterministic, few-shot GSM8K (topk=1): v2 recovers to v1's 0.445. This relies on the spec-v2 stop-string finish fix (fix: stop-string check misses early matches during speculative decoding #23802, Fix spec v2 stop output boundary #25980, now on main): under v2 a step commits multiple tokens, so a stop string landing mid-chunk was missed and the model over-generated past it — the few-shot eval then extracted the wrong final number (~0.365). With those merged, v2 matches v1.
Covers topk=1; topk>1 tree path still to validate

Test plan

/rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_26b_a4b_extra.py test_gemma4_mtp_31b_extra.py
test_gemma4_mtp_26b_a4b_extra.py is currently disabled in CI; with the stop-string fix on main it can be re-enabled in a follow-up

CI States

Latest PR Test (Base): ❌ Run #27235719910
Latest PR Test (Extra): ❌ Run #27235719724

Co-Authored-By: Khoa Pham <khoa.pham@radixark.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

gemini-code-assist · 2026-06-08T23:38:48Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

Co-Authored-By: Khoa Pham <khoa.pham@radixark.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

hnyls2002 · 2026-06-08T23:57:19Z

/rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_26b_a4b_extra.py test_gemma4_mtp_31b_extra.py

github-actions · 2026-06-08T23:57:40Z

Results for /rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_26b_a4b_extra.py test_gemma4_mtp_31b_extra.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/test_frozen_kv_mtp.py

🚀 2-gpu-h100 (2 tests): ❌ View workflow run

cd test/ && python3 registered/spec/test_gemma4_mtp_26b_a4b_extra.py
cd test/ && python3 registered/spec/test_gemma4_mtp_31b_extra.py

hnyls2002 · 2026-06-09T00:06:31Z

/rerun-test registered/spec/test_gemma4_mtp_31b_extra.py

github-actions · 2026-06-09T00:06:54Z

Results for /rerun-test registered/spec/test_gemma4_mtp_31b_extra.py:

🚀 2-gpu-h100 (1 test): ❌ View workflow run

cd test/ && python3 registered/spec/test_gemma4_mtp_31b_extra.py

# Conflicts: # python/sglang/srt/speculative/frozen_kv_mtp_worker.py

hnyls2002 · 2026-06-09T07:40:04Z

/rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_31b_extra.py test_gemma4_mtp_26b_a4b_extra.py

github-actions · 2026-06-09T07:40:29Z

Results for /rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_31b_extra.py test_gemma4_mtp_26b_a4b_extra.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/test_frozen_kv_mtp.py

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/test_gemma4_mtp_31b_extra.py
cd test/ && python3 registered/spec/test_gemma4_mtp_26b_a4b_extra.py

github-actions · 2026-06-09T20:20:56Z

Results for /rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_31b_extra.py test_gemma4_mtp_26b_a4b_extra.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/test_frozen_kv_mtp.py

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/test_gemma4_mtp_31b_extra.py
cd test/ && python3 registered/spec/test_gemma4_mtp_26b_a4b_extra.py

hnyls2002 · 2026-06-09T21:05:57Z

/rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_26b_a4b_extra.py test_gemma4_mtp_31b_extra.py test/registered/attention/unittests/dense/test_triton.py

github-actions · 2026-06-09T21:06:22Z

Results for /rerun-test test_frozen_kv_mtp.py test_gemma4_mtp_26b_a4b_extra.py test_gemma4_mtp_31b_extra.py test/registered/attention/unittests/dense/test_triton.py:

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/test_frozen_kv_mtp.py
cd test/ && python3 registered/attention/unittests/dense/test_triton.py

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/test_gemma4_mtp_26b_a4b_extra.py
cd test/ && python3 registered/spec/test_gemma4_mtp_31b_extra.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/dense/test_triton.py

Three conflicts resolved + four-site API-migration port: 1. piecewise_cuda_graph_runner.py (deleted in HEAD, modified in main). Accept the deletion (cg-refactor's unified runner/prefill_cuda_graph_runner.py replaces it). Port main's sgl-project#27468 dflash-aux-hidden-states change to the new runner: capture_hidden_mode = FULL when spec_algorithm.is_dflash(), and use self.capture_hidden_mode (not NULL) in capture_prepare so dflash spec captures with the right variant. 2. frozen_kv_mtp_worker.py (deleted in main, modified in HEAD). Accept main's deletion (V1 -> V2 worker migration, sgl-project#27607). Apply the four cuda_graph_config API migration sites that cg-refactor had in V1 to V2 (frozen_kv_mtp_worker_v2.py): - backup_disable_cuda_graph -> backup_decode_mode (using new cuda_graph_config.decode.backend field) - server_args.disable_cuda_graph = True -> server_args.cuda_graph_config.decode.backend = Backend.DISABLED - matching restoration on draft_model_runner.server_args - init_cuda_graphs early-exit: self.server_args.disable_cuda_graph -> check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED) 3. test/registered/cuda_graph/piecewise/ test_pcg_with_speculative_decoding_dflash.py (file location). Accept at the cg-refactor-renamed location (new test/registered/ cuda_graph/piecewise/ tree, not the legacy test/registered/piecewise_cuda_graph/). Audited the auto-merged code-side files (scheduler.py, schedule_batch.py, memory_pool.py, cuda_graph_buffer_registry.py, arg_groups/speculative_hook.py, spec_info.py, speculative_draft_runner.py) for unused-arg artifacts introduced by the merge. None found — added kwargs (e.g. mamba_indices on DSATokenToKVPool.get_cpu_copy/load_cpu_copy, padding_policy on positions/mrope_positions GraphSlots) are all referenced at the call sites and in the function bodies. Sanity-checked imports on touched files (frozen_kv_mtp_worker_v2, prefill_cuda_graph_runner) and AST-parsed cleanly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Resolve conflicts from spec v2 refactor: - frozen_kv_mtp_worker.py: removed (replaced by frozen_kv_mtp_worker_v2.py in #27607); ported trtllm_mha draft-attn-backend branch into the v2 worker. - frozen_kv_mtp_cuda_graph_runner.py: adopt shared DecodeCudaGraphRunner surface (_replay_graph + _make_graph_key) from #28081 / #28384. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Co-authored-by: Khoa Pham <khoa.pham@radixark.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

frozen mtp spec v2; drop v1 worker

ed77c5b

Co-Authored-By: Khoa Pham <khoa.pham@radixark.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

hnyls2002 requested review from Qiaolin-Yu, Ying1123, kpham-sgl, merrymercy, pyc96 and xiezhq-hermann as code owners June 8, 2026 23:38

github-actions Bot added the speculative-decoding label Jun 8, 2026

hnyls2002 and others added 2 commits June 8, 2026 16:44

rename FrozenDraftWorker -> FrozenKVMTPDraftWorker

a7266a4

Co-Authored-By: Khoa Pham <khoa.pham@radixark.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

trim frozen_kv_mtp_worker_v2 module docstring

f45b55f

Co-Authored-By: Khoa Pham <khoa.pham@radixark.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

kpham-sgl self-assigned this Jun 9, 2026

kpham-sgl mentioned this pull request Jun 9, 2026

[Roadmap] Gemma4 #26596

Closed

8 tasks

hnyls2002 added 2 commits June 9, 2026 00:35

Merge remote-tracking branch 'origin/main' into lsyin/frozen-mtp-spec-v2

ec696e9

# Conflicts: # python/sglang/srt/speculative/frozen_kv_mtp_worker.py

nvtx spans on frozen v2 orchestrator

2e97fa7

fix double bookkeeping in frozen v2 draft

d7ae6c1

hnyls2002 mentioned this pull request Jun 9, 2026

Add UT guarding per-request bookkeeping clock ownership #27710

Merged

drop v1 dead code; un-pin overlap in tests; assert no adaptive

2ac1434

kpham-sgl approved these changes Jun 9, 2026

View reviewed changes

hnyls2002 merged commit decb88e into main Jun 9, 2026
98 of 110 checks passed

hnyls2002 deleted the lsyin/frozen-mtp-spec-v2 branch June 9, 2026 22:30

hnyls2002 mentioned this pull request Jun 10, 2026

[Spec] Remove dead prepare_for_verify / prepare_extend_after_decode + extend-decode kernel #27761

Merged

whybeyoung mentioned this pull request Jun 12, 2026

fix(pd): disable overlap for spec+grammar in disagg decode loop #28039

Merged

amd-bot mentioned this pull request Jun 13, 2026

[CI Monitor] Daily Report - 2026-06-13 bingxche/sglang-ci-bot#103

Open

wirybeaver mentioned this pull request Jul 23, 2026

[Feature][MLX] Gemma 4 MTP speculative decoding on Apple Silicon #32264

Open

6 tasks

jakki-amd pushed a commit to jakki-amd/sglang that referenced this pull request Sep 9, 2026

Support spec v2 for Frozen-KV MTP; remove v1 worker (sgl-project#27607)

a6425f7

Co-authored-by: Khoa Pham <khoa.pham@radixark.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Navigation Menu

Uh oh!

Support spec v2 for Frozen-KV MTP; remove v1 worker - #27607

Support spec v2 for Frozen-KV MTP; remove v1 worker#27607
hnyls2002 merged 7 commits into
mainfrom
lsyin/frozen-mtp-spec-v2

hnyls2002 commented Jun 8, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jun 8, 2026

Uh oh!

hnyls2002 commented Jun 8, 2026

Uh oh!

github-actions Bot commented Jun 8, 2026 •

edited

Loading

Uh oh!

hnyls2002 commented Jun 9, 2026

Uh oh!

github-actions Bot commented Jun 9, 2026 •

edited

Loading

Uh oh!

hnyls2002 commented Jun 9, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 9, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 9, 2026 •

edited

Loading

Uh oh!

hnyls2002 commented Jun 9, 2026

Uh oh!

github-actions Bot commented Jun 9, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

hnyls2002 commented Jun 8, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Mechanical verification

Validation

Test plan

CI States

Uh oh!

gemini-code-assist Bot commented Jun 8, 2026

Uh oh!

hnyls2002 commented Jun 8, 2026

Uh oh!

github-actions Bot commented Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hnyls2002 commented Jun 9, 2026

Uh oh!

github-actions Bot commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hnyls2002 commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hnyls2002 commented Jun 9, 2026

Uh oh!

github-actions Bot commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

hnyls2002 commented Jun 8, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Jun 8, 2026 •

edited

Loading

github-actions Bot commented Jun 9, 2026 •

edited

Loading

hnyls2002 commented Jun 9, 2026 •

edited

Loading

github-actions Bot commented Jun 9, 2026 •

edited

Loading

github-actions Bot commented Jun 9, 2026 •

edited

Loading

github-actions Bot commented Jun 9, 2026 •

edited

Loading