[DCP][Spec] EAGLE/EAGLE3 support for decode context parallelism - #31785
thanhhao98 wants to merge 22 commits into
Conversation
thanhhao98
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Decode Context Parallelism (DCP) optimizations, benchmarks, and tests, including a 2-pass target-verify cascade (non-causal prefix fold and local causal draft fold) and an alternative All-to-All communication backend. The review feedback highlights several critical runtime bugs where torch.dtype objects are incorrectly accessed with .itemsize (which will raise an AttributeError in PyTorch) across activation, kernel, benchmark, and test files, suggesting the use of .element_size() instead. Additionally, the reviewer recommends asserting contiguity for tensors in the new all_to_all_single NCCL wrapper to prevent potential silent data corruption.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
c93cf54 to
264324e
Compare
github-actions
Bot
added
the
documentation
36bc9d2 to
a8f6ef5
Compare
thanhhao98
changed the title
|
Per-GPU aggregate view of the same fix-tip EAGLE3 matrix (complements the per-user table above;
Aggregate ratios sit slightly above the per-user ones (e.g. 0.88 vs 0.84 at 50K/cc4) — run-to-completion aggregates fold in ramp/tail effects; same underlying runs. The regime story is unchanged: below the KV-pool edge DCP pays the verify addon with no capacity benefit; past the edge (overcommit ≥ ~1.0) the off arm loses admission first and DCP wins on both axes. 🤖 Generated with Claude Code |
thanhhao98
changed the title
68998b1 to
f4583ee
Compare
bda5795 to
82e00b0
Compare
thanhhao98
changed the title
…dget Both behaviours are load-bearing and both fail quietly, so they get CPU unit tests. draft_forward_guard is now the ONLY mechanism that unshards a draft forward. Upstream removed ModelRunner.dcp_size/dcp_rank and made the attention-facing accessors derive from dcp_enabled (attn_dcp_size = dcp_size if dcp_enabled else 1), so the guard alone drives attn_dcp_size to 1. If that stops holding, draft forwards silently take DCP branches against an unsharded pool -- the chain-decode corruption this exists to prevent, with no error at the point of failure. The test pins size/rank inside and outside the guard, across nesting, exceptions, and the DCP-off case. The draft-pool budget under-allocates rather than erroring: the widened draft pool eats the transient headroom and the server OOMs at the first large prefill, far from the cause. The test pins the dcp=1 case to the exact upstream ratio, the dcp scaling, monotonicity, that the draft worker budgets nothing, and -- explicitly -- that an unknown draft depth yields 0 rather than a guessed floor.
…alidation Three review findings. 1. SGLANG_DCP_DRAFT_CHAIN_GUARD could produce the corruption it was meant to let us study. The graph CAPTURE in _capture_cuda_graphs was guarded unconditionally while the chain draft (which replays those graphs) was guarded only when the flag was set, so setting it to 0 -- documented as "revert for A/B" -- left capture and replay disagreeing about dcp_enabled. That is exactly the capture/replay divergence the capture-site comment describes as causing an IMA. The flag is removed and all four guard sites are now unconditional. It was introduced in this same PR, so nothing released depended on it. 2. _validate_dcp_spec_topk defaulted its algorithm predicates through getattr. Those come from SpeculativeAlgorithm, an Enum that always defines is_eagle/is_dflash, so the defensiveness bought nothing and would have silently let tree drafts through under DCP if a predicate were ever renamed -- the guard failing open on the exact configuration it exists to reject. Now called directly. The getattr on dcp_size and speculative_eagle_topk is kept and explained: test_spec_registry drives this hook with a deliberately minimal SimpleNamespace. The topk comparison is also explicit about None rather than relying on `or 1`. 3. prepare_decode_context_parallel_metadata took Optional[torch.Tensor] while the five model-side wrappers that forward into it still declared torch.Tensor. Annotations propagated so the contract is consistent along the whole call path. Adds test_dcp_spec_topk_guard.py, which had no coverage at all: rejection for EAGLE and EAGLE3 tree drafts, acceptance for chain drafts and with DCP off, and -- the scoping most likely to regress -- that DSPARK is NOT gated, since is_dflash_family() would pull it in and newly reject the Kimi-Linear + DSPARK + DCP path that already ships. Also cuts the comments back by two thirds (63 -> 21 lines). What is left is only the non-obvious why: the is_dflash_family trap, why the predicates must not be getattr, why banded_depths needs the raw draft depth rather than the DCP-replicated one, and the capture-must-match-replay invariant. Everything that merely restated the code, or recorded history that belongs in a commit message, is gone.
6728a71 to
e23befe
Compare
| def verify_mask(self) -> Optional[VerifyMask]: | ||
| return self._verify_mask | ||
|
|
||
| def update_verify_buffers_to_fill_after_draft( |
There was a problem hiding this comment.
why do we need this no-op function
There was a problem hiding this comment.
without this func, we will get NotImplementedError in
| def draft_forward_guard(is_draft: bool): | ||
| """Run a draft forward with DCP disabled; no-op when is_draft is False. | ||
|
|
||
| The draft KV pool is replicated, not sharded, so every DCP branch must see | ||
| dcp_enabled == False for the whole draft forward (this also drives | ||
| attn_dcp_size to 1 and attn_dcp_rank to 0). | ||
| """ | ||
| if not is_draft: | ||
| return contextlib.nullcontext() | ||
| return get_parallel().override(dcp_enabled=False) |
There was a problem hiding this comment.
I considered this approach many times but decided to go against it because get_parallel() is a runtime static parallel state. .override() is only meant to be used for unit testing.
Let's brainstorm how to best do this offline
There was a problem hiding this comment.
Btw for this part, I considered many previous approaches but ultimately rejected them all
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| def _eagle_draft_layers(kvc: KVCacheConfigurator) -> int: |
There was a problem hiding this comment.
Self note: acceptable currently but might need better approach in the future
| def prepare_context_parallel_metadata_for_dcp( | ||
| self, | ||
| seq_lens: torch.Tensor, | ||
| extend_prefix_lens: torch.Tensor, |
There was a problem hiding this comment.
Curious why make this Optional?
There was a problem hiding this comment.
I removed those Optional.
|
@thanhhao98 could you address the comments and fix the conflicts? thanks! |
…enspeed # Conflicts: # python/sglang/srt/model_executor/pool_configurator.py
…comments Review actions on sgl-project#31785: 1. _validate_dcp_spec_topk -> _validate_dcp_spec, the single entry point for DCP x speculative validation (EAGLE / DFlash / DSpark). 2. Remove or shorten the generated-sounding docstrings and comments flagged in review; the no-op update_verify_buffers_to_fill_after_draft override now matches the bare-pass style of trtllm_mha / triton.
…bing Target-verify is excluded at the single call site (EagerRunner), so the DCP chunked-prefix planner never receives None; the planner None-return and the Optional annotations on the six pass-through signatures were redundant defenses. Reverts planner.py and the four model files to upstream state.
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
…enspeed # Conflicts: # python/sglang/srt/model_executor/pool_configurator.py
STANDALONE inherits the EAGLE V2 tree-draft path, so the chain-only DCP validation must cover it. The draft-extend CUDA graph was captured and replayed outside draft_forward_guard (unlike the draft-decode graph and the eager fallback), so its out-of-graph metadata re-plan saw live DCP state against the replicated draft pool. Also hoist the guard imports to module level (layers.dcp has no reverse dependency; forward() is a hot path).
thanhhao98
changed the title
|
@thanhhao98 could you fix Lint failure? |
nvpohanh
left a comment
nvpohanh
left a comment
There was a problem hiding this comment.
[by Codex] Three inline review findings.
| def _validate_dcp_spec(server_args: ServerArgs) -> None: | ||
| if ( | ||
| server_args.speculative_algorithm is None | ||
| or getattr(server_args, "dcp_size", 1) <= 1 |
There was a problem hiding this comment.
[by Codex] Severity: style | Confidence: High
Please avoid the new dynamic getattr reads here and at line 197. _validate_dcp_spec is typed to receive ServerArgs, where both dcp_size and speculative_eagle_topk are defined, so read those fields directly and make any lightweight test fixture satisfy that contract explicitly. This also prevents a renamed or missing production field from silently disabling the validation.
| model_runner.model, "prepare_context_parallel_metadata_for_dcp" | ||
| if ( | ||
| model_runner.ps.attn_dcp_size > 1 | ||
| and hasattr( |
There was a problem hiding this comment.
[by Codex] Severity: style | Confidence: High
Please remove the retained hasattr check from this modified condition. Model capability should be represented explicitly by the model interface/base implementation or by a typed capability flag selected during initialization, then this path can call prepare_context_parallel_metadata_for_dcp structurally. Runtime method probing makes unsupported models silently skip required DCP metadata preparation.
|
|
||
| algo = SpeculativeAlgorithm.from_string("EAGLE3") | ||
| for name in ("is_eagle", "is_dflash", "is_dspark"): | ||
| self.assertTrue(hasattr(algo, name), f"SpeculativeAlgorithm lost {name}()") |
There was a problem hiding this comment.
[by Codex] Severity: style | Confidence: High
Please avoid hasattr in the new test. Call each required predicate directly (and assert its expected result) so a missing or renamed method raises naturally; that checks the stated contract without dynamic attribute probing.
|
@thanhhao98 could you fix lint issue? Thanks |
isort orders layers.cp.utils before layers.dcp; the branch had them swapped, which failed the pinned pre-commit hook in CI.
TRTLLMMLABackend gained update_verify_buffers_to_fill_after_draft on main in sgl-project#33561, so the copy this branch added is a duplicate definition in the same class. Removing it also retires the review question about the no-op.
attn_dcp_size is now a derived parallel width resolved at publish time (sgl-project#36790, sgl-project#38113), so it no longer follows an overridden dcp_enabled and a draft forward kept the target's DCP width. State the widths the guard means, the way the other production override sites do.
_validate_dcp_spec reads the resolving view instead of defaulting the two ServerArgs fields through getattr; the view also answers with the topk the speculative hook itself resolves, which a raw field read now misses since declarations stopped writing back to the record (sgl-project#36253). The spec-registry fixture states the two fields the validator reads. The eager DCP-metadata branch selects on a runtime_checkable SupportsDecodeContextParallelMetadata protocol rather than probing the model for the method name. The branch stays: a non-MLA DCP target (triton backend) carries no attn_dcp_metadata and must keep skipping it.
register_cpu_ci takes the (stage, runner_config) pair or a suite, never a bare stage: the three files were the only ones in the tree that failed to parse, and collect_tests aborts every CI suite job on the first one. With the registration valid the taxonomy gate applies, so they move under test/registered/unit/dcp/.
The suite fed _validate_dcp_spec an already-final topk, so nothing turned red when the guard read the record instead of the resolving view - the state the speculative hook is actually in once it has declared an auto-chosen topk.
3 participants
Makes EAGLE / EAGLE3 work with decode context parallelism (
--dcp-size > 1). Stacked on #21637.DFlash × DCP already worked (see #33912,
_dflash_draft_cell_size); it gets one shared fix here, not new support.Root cause: the target KV pool is sharded across DCP ranks, but the draft pool is replicated. Draft-side code that reads the DCP topology therefore builds rank-local metadata against a full pool. Every change below follows from that.
Results
8×B200, K2.5-NVFP4 + EAGLE3 (ns2/topk1/ndt3),
--dcp-comm-backend a2a --dcp-replicate-q-proj,jobs 1496709–1496996. Accept length is the correctness gate: the DCP arm must match the off arm,
since a draft chain fed rank-local metadata degenerates after the first token.
GSM8K 0.940–0.960 on every arm. Below the KV-pool edge DCP pays the verify addon for no benefit; past it the off arm loses admission first and DCP wins. Attention FLOPs are conserved under context sharding, so there is no headroom to beat the off arm at shallow overcommit — the ~16% cost at 50K/cc4 is the design point, not a regression.
Tests
test_dcp_draft_guard.py,test_dcp_draft_pool_bounds.py,test_dcp_spec_topk_guard.py— all three cover failures that are otherwise silent: a guard that stops unsharding draft forwards corrupts the chain with no error, an under-budgeted draft pool OOMs at the first large prefill far from its cause, and a mis-scoped topk validator returns wrong tokens rather than raising. The topk tests pin that DSPARK is not gated (is_dflash_family()would pull it in and reject the shipping Kimi-Linear + DSPARK + DCP path).Registered DCP/spec suites pass on 8×B200, including main's
test_kimi_linear_dcp4.py.CI States
Latest PR Test (Base): ❌ Run #34432196662
Latest PR Test (Extra): ❌ Run #34432196449
Latest PR Test (AMD ROCm 10): ❌ Run #34432196657