[None][fix] enable static EPLB for the one-model DSpark drafter - #16938
Merged
Conversation
The one-model external-drafter branch builds the draft ModelConfig from the
speculative checkpoint but never propagated moe_load_balancer. With EPLB
enabled, the engine-wide MoeLoadBalancer is live during DSpark construction
while the draft config's moe_load_balancer is None, so every DSpark stage MoE
trips `assert moe_load_balancer_config is not None` in
MoE._init_load_balancer and the gen worker dies at model init.
Propagate moe_load_balancer for DSpark only. DSpark's draft stages are full
DeepSeek-V4 blocks sharing the target's expert topology and layer-index
namespace (layer_idx = num_hidden_layers + stage_id), so they register into
the target's balancer as additional EPLB layers. PARD, DFlash and
draft-target are independent checkpoints whose topology and layer numbering
need not match the target's; their kwargs are unchanged and a test pins that.
Add two config-time validations, both gated on an actually-active balancer so
the non-EPLB path is untouched:
- initial_global_assignments must cover the DSpark stage indices. A config
generated before DSpark (or with 1-layer MTP) lacks them, which otherwise
surfaces as a bare KeyError from deep inside MoE init; now it reports every
missing index at once and points at the regeneration docs.
- reject layer_updates_per_iter > 0. DSpark supports static EPLB only:
online EPLB requires every registered MoE layer to run exactly once per
iteration, but the draft MoE is skipped on iterations with no generation
requests (context-only batches, warmup), and the balancer's CPU worker
then spins forever in its untimed waitCpuStage(). Failing at startup beats
a silent deadlock.
Also require the draft config's num_hidden_layers to match the target's when
EPLB is active, since the stage layer indices are derived from it and must
line up with the namespace initial_global_assignments is keyed by.
Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
longlee0622
force-pushed
the
dev/dspark-static-eplb-pr
branch
from
July 28, 2026 08:11
0b5730b to
f1a8aa8
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #62157 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62157 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62170 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62170 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
WalkthroughChangesDSpark EPLB integration
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant SpecDecOneEngineForCausalLM
participant external_drafter_config_kwargs
participant ModelConfig
SpecDecOneEngineForCausalLM->>external_drafter_config_kwargs: Build external drafter kwargs
external_drafter_config_kwargs->>ModelConfig: Pass copied settings with spec_config=None
external_drafter_config_kwargs-->>ModelConfig: Include moe_load_balancer for DSpark
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_dspark.py`:
- Around line 82-120: Annotate every affected function without using Any: in
tensorrt_llm/_torch/models/modeling_dspark.py lines 82-120, add concrete
parameter and return types to _active_moe_load_balancer,
validate_dspark_eplb_layer_base, and validate_dspark_eplb_stage_layers, plus
Google-style argument documentation for the public validators; in
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py lines
46-76, type the helpers and fixture generator; in the same file lines 87-210,
annotate test parameters and None returns; and in
tensorrt_llm/_torch/models/modeling_speculative.py lines 1842-1869, replace the
bare dict return annotation with precise input and mapping types.
- Around line 156-159: Update the assignment guard in the DSpark
placement-validation logic to bypass validation only when
initial_global_assignments is None. Preserve validation for an explicitly
supplied empty mapping so missing layer assignments are rejected before EPLB
initialization.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 55c3d84a-deb4-454e-8b54-e4cce547704d
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_dspark.pytensorrt_llm/_torch/models/modeling_speculative.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py
Collaborator
|
PR_Github #62308 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62308 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62385 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62385 [ run ] completed with state |
longlee0622
enabled auto-merge (squash)
July 29, 2026 04:16
1 task
JunyiXu-nv
approved these changes
Jul 29, 2026
WeiHaocheng
approved these changes
Jul 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Enable static MoE EPLB for the one-model DSpark drafter. With EPLB configured, the
DSpark gen worker currently dies at model init:
The target model is built inside
maybe_create_moe_load_balancer(...), so theengine-wide
MoeLoadBalanceris live while the DSpark stages are constructed. Butthe generic external-drafter branch builds
draft_configfrom the speculativecheckpoint without propagating
moe_load_balancer, so every DSpark stage MoE seesDSparkDraftModel._derive_draft_model_config()is not the loss point — its shallowcopy preserves whatever it receives; the config is already
Noneby then.Changes
1. Propagate
moe_load_balancer— DSpark only.external_drafter_config_kwargs()factors out the kwargs the one-model externaldrafter inherits, and adds
moe_load_balanceronly whenspec_dec_mode.is_dspark(). DSpark's stages are full DeepSeek-V4 blocks sharing thetarget's expert topology and layer-index namespace (
layer_idx = num_hidden_layers + stage_id), so they can register into the target's balancer. PARD / DFlash /draft-target are independent checkpoints whose topology and layer numbering need not
match the target's, and whose EPLB configs would be keyed against a different
namespace — their kwargs are byte-for-byte unchanged, and a test pins that so a
future refactor cannot silently generalize this.
2. Validate the EPLB config covers the DSpark stages.
A config generated before DSpark (or in the 1-layer-MTP era) lacks the stage
indices, which today surfaces as a bare
KeyError: layer_idx 62 not foundfrom deepinside MoE init. The new check runs before any DSpark MoE is built and reports every
missing index at once, with the layer-index convention and a pointer to the
regeneration docs.
3. Reject online EPLB (
layer_updates_per_iter > 0) at config time.DSpark supports static EPLB only. Online EPLB requires every registered MoE layer to
run exactly once per iteration, but the DSpark draft MoE is skipped on iterations
with no generation requests anywhere (context-only batches, warmup). The balancer's
CPU worker then spins forever in its untimed
waitCpuStage(), waiting for a GPUsignal only an MoE forward emits — a silent deadlock. Failing at startup beats
hanging in production.
Also: when EPLB is active, require the draft config's
num_hidden_layersto matchthe target's, since the stage indices derive from it and must line up with the
namespace
initial_global_assignmentsis keyed by.All three validations are gated on
get_moe_load_balancer()being non-None — theexact condition under which
MoE._init_load_balancerconsumes the config — so thenon-EPLB path is untouched.
Test Coverage
Unit —
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py,14 CPU-only cases: propagation, PARD/DFlash/draft-target isolation,
spec_config=Nonepreservation, stage-layer coverage, auto-placement, EPLB-inactive no-op, online-EPLB
rejection, layer-base match/mismatch.
Multi-GPU smoke (no CI stage covers DSpark + EPLB + real weights, so this was run
by hand):
ValueError: initial_global_assignments is missing DSpark layer(s) [62, 63]. ... indices [61, 64).on all 8 ranks — the oldAssertionErroris gonenum_hidden_layers(43) + num_stages(3), i.e. the DSpark stages land on 43/44/45 as designed; weight load,register_weight_slots_after_to_cuda()andfinalize_model()all passThe full-coverage run stops in
_run_attention_warmupon a CuTe-DSL codegen error(
nvgpu.cvt_fptruncoperand type), so it does not reach token generation. Thatfailure is unrelated to this change: an otherwise identical run with EPLB fully
disabled fails identically.
Notes for reviewers
per-drafter EPLB config domain and layer identity design first.
does not change.
MoeLoadBalancer(...)exists leavesMpiPoolSession.shutdownblocked joining itsworkers forever, while they still hold GPU contexts. Confirmed with a control that
fails from pre-existing code (
interface.pyassignment-length check) on targetlayer 0, with these validators inactive — it hangs identically. Worth a separate
issue: it turns a fast, clear config error into a wedged node.
PR Checklist
[JIRA/NVBUG/None][type] descriptionDev Engineer Review
QA Engineer Review
Added 14 CPU-only unit tests covering propagation, drafter isolation, layer coverage, inactive/online EPLB behavior, and layer-count validation.
No corresponding entries were added to
tests/integration/test_lists/test-db/orqa/. Existing DSpark performance entries inqa/llm_perf_core.ymldo not cover this unit-test file.Verdict: needs follow-up