feat(spec-v2): Frozen-KV MTP overlap-scheduling worker (experimental opt-in) + V1 None-safety fixes

[pyc96]

Summary

Stacked on PR #21 (ULT v2). Adds FrozenKVMTPWorkerV2 spec-V2 infrastructure for FROZEN_KV_MTP behind an experimental env-var opt-in (SGLANG_FROZEN_KV_MTP_EXPERIMENTAL_V2=1). Default behavior unchanged: V1 worker + disable_overlap_schedule=True.

State of the V2 implementation

Component	Status
Boots cleanly with overlap enabled	OK
Dispatcher routes to V2 worker	OK
Inherits all V1 correctness-tested methods (draft, verify, _run_assistant_seed_step, etc.)	OK
Implements BaseSpecWorker interface (target_worker, draft_worker, clear_cache_pool)	OK
forward_batch_generation(batch, on_publish=None) with verify-end fence	OK
Ferries next_draft_input with correct shape ([B] bonus_tokens, [B, topk] topk_p/index, [B, H] hidden_states)	OK
Produces accept_lens for batch result processor	OK
Routes prefill -> decode -> seed -> verify -> next iter end-to-end	OK
Request completes and returns correct output to client	OK (single request validated)
on_idle invariant_checker passes (KV pool leak check)	FAIL -- 24-token leak per finished req

Bugs fixed during V2 development

1. V1 draft() and verify() None-safety on batch.sampling_info.penalizer_orchestrator (real fix benefiting V1 too). Required because V2 defers sampling to the schedule stream; the worker side gets penalizer_orchestrator=None.

2. V2 ferry size: was using verify_output.accept_tokens.shape[0] (= flat total across all reqs); now uses batch.req_pool_indices.shape[0] (= per-req batch size).

3. V2 per-req bonus token: was passing the flat accept_tokens; now picks the last accepted token per req via cumsum(num_correct_drafts_per_req_cpu + 1) - 1, mirroring EAGLE V2's fill_bonus_tokens kernel.

4. V2 accept_lens in batch result: _resolve_spec_overlap_tokens at batch_result_processor.py:488 asserts result.accept_lens.is_cpu. Now populated.

Known limitation (NOT fixed)

The on_idle invariant_checker detects a 24-token KV pool leak per finished request:

ValueError: pool memory leak detected! [full] total=106600, available=106570,
evictable=0, protected=6, session_held=0, uncached=0

24 ~= 20 (max_tokens) + 4 (speculative_num_draft_tokens), suggesting the V1 verify path's free of rejected draft slots races with the schedule-stream tear-down of the batch under overlap.

Fixing this requires architectural understanding of SGLang's pool-lifecycle ownership across overlap stream boundaries. Identified next-step direction in the audit (agent-pod/runs/20260525_mtp_v2/analysis/eagle_v2_pattern.md section 14.1 "The target KV pool swap must NOT race with overlap"). The proper fix is likely to either:

• Wrap the verify-side free in a stream-event sync that ensures completion before iter N+1's schedule-side cleanup runs
• Or migrate to the full composition-based V2 pattern (BaseDraftWorker split per the audit's section 15 skeleton, ~600 LOC)

What this PR ships (value)

Item	Value
V1 None-safety in frozen_kv_mtp_worker.py + eagle_info.py	Real bug fixes for V1 paths under future scenarios where penalizer becomes None
FrozenKVMTPWorkerV2 with correct ferry contract	Removes the NotImplementedError stub; next implementer can focus exclusively on the KV-leak fix
Dispatcher + opt-in plumbing	Once the leak is fixed, flip the default by removing the env-var check
2409-line audit at agent-pod/runs/20260525_mtp_v2/analysis/eagle_v2_pattern.md	Comprehensive EAGLE V2 walkthrough + frozen-KV invariants + porting matrix

Reproducer

SGLANG_FROZEN_KV_MTP_EXPERIMENTAL_V2=1 python -m sglang.launch_server \
  --model-path google/gemma-4-31B-it --dtype bfloat16 --trust-remote-code --tp-size 2 \
  --host 0.0.0.0 --port 30000 --attention-backend triton \
  --speculative-algorithm NEXTN \
  --speculative-draft-model-path google/gemma-4-31B-it-assistant \
  --speculative-num-steps 3 --speculative-num-draft-tokens 4 --speculative-eagle-topk 1 \
  --max-running-requests 80 --random-seed 1

# Single request works:
curl -X POST http://127.0.0.1:30000/v1/chat/completions \
  -d '{"model":"google/gemma-4-31B-it","messages":[{"role":"user","content":"What is 2+2?"}],"temperature":0,"max_tokens":15}'

# After the request finishes, the server crashes on on_idle invariant check.
# That is the documented known limitation -- the 24-token KV leak.

Production recommendation

Use SGLang no-MTP (per PR #21). The MTP gap is structural and this PR does not close it yet.

Stack base

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

Files

• python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py (+363 lines vs original stub)
• python/sglang/srt/speculative/frozen_kv_mtp_worker.py (+8/-3 None-safety)
• python/sglang/srt/speculative/eagle_info.py (+11/-4 None-safety)
• python/sglang/srt/speculative/spec_info.py (+13/-3 dispatcher)
• python/sglang/srt/arg_groups/speculative_hook.py (+18/-7 opt-in)
• agent-pod/runs/20260525_mtp_v2/analysis/eagle_v2_pattern.md (2409-line audit)