[#14146][bugfix] Fix CuteDSL MoE ghost-token and global-index bugs

[waynehacking8]

Summary

Fixes #14146 — two bugs in generate_token_selected_experts (cute_dsl_custom_ops.py), the helper that synthesizes token_selected_experts for autotuner profiling (inputs_pre_hook). They affect how representative the profiling distribution is — not inference routing or accuracy — but the generated profile should still match the per-expert allocation:

• Ghost-token assignment: when num_tokens_per_expert[j] == 0, the inner loop still assigns 1 token to that expert before breaking. With DeepSeek-V3 EP=16, top_k=8, 1 token: 7 ghost experts get assigned instead of 0. Fix: if num_tokens_j <= 0: continue.
• Prioritization uses local index: the priority threshold top_k - (num_experts - j) uses the loop variable j (local) instead of j + self.local_expert_offset (global). Whenever num_local_experts < num_experts (EP > 1) the threshold is always negative, so the at-risk-token prioritization is dead code on every rank; the global index restores it near the tail of the global expert range. Fix: compute global_j and use it consistently.

Reproduction

# DeepSeek-V3, ep_size=16, top_k=8, num_tokens=1:
#   num_tokens_per_expert = [1, 0, 0, ..., 0]  (1 nonzero, 15 zeros)
#   Buggy: assigns experts [0,1,2,3,4,5,6,7] — 7 ghost experts
#   Fixed: assigns experts [0] — correct

Test plan

• test_ghost_token_zero_allocation — experts with 0 allocation receive no tokens (reproduces the [1,1,0,…] → assigned [0,1] case above)
• test_prioritization_rescues_at_risk_tokens — deterministic last-rank config (num_experts=4, top_k=2, 2 local experts, offset=2, 4 tokens): global-index prioritization fills 3 tokens to top_k vs 2 under the local-index bug, so reverting global_j → j fails the test (the previous assertion only checked the expert-id range, which is offset-invariant and missed bug 2)
• Existing test_grouped_gemm_inputs_helper and test_moe_sort pass

Closes #14146

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

Summary by CodeRabbit

• Bug Fixes

  • Fixed expert index mapping in mixture-of-experts token selection to use global indices instead of local indices, ensuring correct expert assignment.
  • Added conditional handling for model weight loading to properly support checkpoints with optional bias parameters.

• Tests

  • Added test coverage for zero-token expert handling and global index prioritization validation.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR fixes two distinct issues: CuteDSL MoE token-to-expert selection now correctly skips zero-allocation experts and uses global expert indices in prioritization calculations, and DeepSeek V3 model loading now conditionally imports the e_score_correction_bias parameter when present in checkpoints. Tests validate the MoE fixes.

Changes

CuteDSL MoE Expert Selection Fixes

Layer / File(s)	Summary
Expert index offset and prioritization logic tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py	Loop skips iterations when num_tokens_j <= 0, computes global_j = j + local_expert_offset, updates the prioritized condition to use global_j instead of local j, and stores the global expert index in token_selected_experts instead of recomputing it.
Expert selection test validation tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py	test_ghost_token_zero_allocation verifies experts with zero token allocation receive no assignments; test_prioritization_uses_global_index verifies selected experts fall within the expected local range when local_expert_offset is non-zero.

DeepSeek V3 Weight Loading Conditional

Layer / File(s)	Summary
Conditional e_score_correction_bias loading tensorrt_llm/_torch/models/modeling_deepseekv3.py	DeepseekV3Gate.load_weights now checks for e_score_correction_bias key presence in the checkpoint before copying it into self.e_score_correction_bias, replacing the prior unconditional copy.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and accurately summarizes the main changes: fixing two bugs in CuteDSL MoE ghost-token assignment and global-index prioritization, with a specific issue reference.
Linked Issues check	✅ Passed	The PR implementation directly addresses all coding objectives from issue #14146: prevents ghost-token assignment by skipping zero-allocation experts, corrects prioritization to use global indices, adds unit tests for both fixes, and maintains intended behavior across multi-rank configurations.
Out of Scope Changes check	✅ Passed	All code changes are directly scoped to the stated objectives: fixing ghost-token and global-index bugs in generate_token_selected_experts, adding corresponding tests, and making a conditional fix for optional e_score_correction_bias loading in DeepseekV3Gate.
Description check	✅ Passed	PR description comprehensively explains both bugs, provides reproduction steps, test plan, and clearly documents the fixes with specific examples.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

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_deepseekv3.py`:
- Around line 899-902: The current conditional only copies
weights[0]["e_score_correction_bias"] when present, leaving
self.e_score_correction_bias uninitialized if the key is absent; add an else
branch in the same block that deterministically initializes
self.e_score_correction_bias (e.g., zero-fill or a chosen default) so
checkpoints that omit "e_score_correction_bias" produce deterministic behavior.
Ensure the fallback matches self.e_score_correction_bias's dtype and device (use
methods like zero_() or torch.zeros_like semantics) and reference the same
attribute name self.e_score_correction_bias and the source weights container
weights[0] when implementing the change.

In `@tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py`:
- Around line 75-93: The test test_prioritization_uses_global_index currently
only checks ID bounds; change it to assert actual prioritization output matches
the expected global expert set: use
GroupedGemmInputsHelper.generate_num_tokens_per_expert to get
num_tokens_per_expert, compute the expected set of global expert IDs as indices
where num_tokens_per_expert > 0 shifted by local_expert_offset (or otherwise
derived per the helper's mapping), then call generate_token_selected_experts and
assert that the non-negative entries in its result correspond exactly to that
expected global set (or contain the expected top-k per-expert selection), so
that failures show a prioritization/regression rather than just an out-of-range
ID. Ensure you reference test_prioritization_uses_global_index,
GroupedGemmInputsHelper, generate_num_tokens_per_expert and
generate_token_selected_experts when updating the assertions.

🪄 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: afcacc7d-6420-4184-9fb0-03eb17e585d1

📥 Commits

Reviewing files that changed from the base of the PR and between b852703 and 230c841.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
• tensorrt_llm/_torch/models/modeling_deepseekv3.py
• tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py

[waynehacking8]

Status check: this PR has been mergeable and review-ready since 06-09 (coderabbit pass addressed at open time). Could a maintainer trigger CI (/bot run) when convenient? Happy to rebase or adjust if anything in the recent refactors moved — the fix still applies cleanly to current main.