Skip to content

[PD] Bootstrap EAGLE draft state on decode after target-only prefill - #40955

Draft
roborluo wants to merge 7 commits into
sgl-project:mainfrom
roborluo:pd-decode-draft-bootstrap
Draft

roborluo wants to merge 7 commits into
sgl-project:mainfrom
roborluo:pd-decode-draft-bootstrap

Conversation

@roborluo

@roborluo roborluo commented Sep 23, 2026 •

Copy link
Copy Markdown

Motivation

Target-only prefill does not initialize the draft model's prefix KV or proposal state. Decode-side EAGLE therefore needs real initialization rather than accepting absent metadata or synthesizing probabilities.

Dependency

Depends on #40953.

This stacked PR targets main, so its full diff includes that prerequisite. Review the decode-bootstrap increment separately; it will need rebasing after the prerequisite merges. It does not depend on the workload-routing PR.

Modifications

  • Add default-off --disaggregation-decode-draft-bootstrap and an explicit prefix-token budget.
  • Advertise draft-state capability in Python and Rust bootstrap registries. Validate the actual handoff capabilities, not merely whether a rejection-sampling flag is set.
  • Select a target-only receive layout before metadata registration. Do not mix seeded and target-only peers on one listener.
  • Replay the committed prefix on decode through target and draft models while preserving the already sampled handoff token. Reuse allocated KV slots and drain overlap before replay.
  • Publish draft state only after successful initialization; reject unsupported/oversized requests without silently truncating them.
  • Expose additional target replay work as pd_draft_bootstrap_tokens, without altering prefill cache-hit accounting.

This is an experimental correctness fallback, not a throughput optimization: it repeats full-prefix target and draft compute. Supported scope is single-layer EAGLE, TP-only decode, supported DeepSeek/GLM MoE text architectures, Mooncake, and Python egress. Decode prefix caching, HiCache, offload, HiSparse, constrained decoding, embeddings and multimodal requests are excluded. Rust egress is rejected for this mode because it does not expose replay accounting.

Accuracy Tests

  • Passed: 13 source-isolated CPU bootstrap tests with real CPU tensors and mocked runners: capability matrix, prefix boundaries, unchanged handoff, failure-before-publish and incomplete-KV rejection.
  • Passed: 7 prerequisite CPU proposal-metadata tests (plus 7 subtests).
  • Passed: changed Python compilation, focused Ruff checks and git diff --check.
  • Full SGLang unit collection is blocked locally by missing Triton. Rust tests were not run because the Rust toolchain is unavailable.
  • GPU transport, real runner correctness, overlap, stochastic sampling accuracy and performance are unvalidated on this upstream port. This PR remains Draft.
python test/manual/disaggregation/pd_draft_bootstrap/run_cpu_tests.py
python -m pytest test/registered/unit/disaggregation/test_pd_draft_metadata_cpu.py

The manual README documents the full-environment test commands and outstanding GPU validation matrix.

Speed Tests and Profiling

No speedup or benchmark result is claimed. Account for the replay counter separately in computed-token metrics. No internal model fixtures, cluster configs, logs or results are included.

Checklist

  • Add configuration, capability negotiation, replay accounting and regression tests.
  • Apply the repository's Python import/lint/format tools to changed Python files.
  • Merge/rebase the prerequisite.
  • Run full SGLang and Rust tests.
  • Validate GPU execution and sampling quality before merge.

CI States

Latest PR Test (Base): ❌ Run #35906910682
Latest PR Test (Extra): ❌ Run #35906910092
Latest PR Test (AMD ROCm 10): ❌ Run #35906910780

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant