Files
sglang/python
leavelet 2c7353d061 Fix spec_v2 CP draft extend hidden-state shape mismatch
EAGLEWorkerV2's target prefill always captured the global FULL hidden,
so the bs>1 NSA CP shared-KV draft extend received a global hidden whose
row count (real total tokens) matched neither the per-rank CP-local count
nor the MLP-sync CP-aligned padded count, tripping
[CP_SHARED_KV_FAIL_FAST][draft_batch_gt1_spec_hidden_shape_mismatch]
under SGLANG_ENABLE_SPEC_V2=1.

Mirror the legacy EAGLEWorker contract: add _can_use_cp_draft_shared_kv
and, when CP draft shared-KV applies, capture the CP-local hidden side
channel (CaptureHiddenMode.NULL + capture_draft_hidden_states) instead of
FULL. The v2 consumer already prefers draft_hidden_states when present
(commit 5e22279670 added the consumer but not this producer side).

Fixes the shape at the source rather than loosening the NextN fail-fast.
CP-off keeps FULL; non-CP / bs=1 unaffected; legacy v1 path untouched.
Validated on g0033 PD (prefill + 2-node decode): bs>1 up to 17 incl. the
former-crashing bs=6, 0 non-400 errors, GSM8K 0.965 / 0 invalid.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-13 09:06:00 +00:00
..
2026-06-10 05:54:43 +08:00