[None][fix] Support DeepSeekV4 routing in perfect-router planner by qiaoxj07 · Pull Request #14074 · NVIDIA/TensorRT-LLM

qiaoxj07 · 2026-05-13T06:44:04Z

Summary

DeepSeekV3PerfectRouterPlanner._resolve_group_config() only knew how to introspect DeepSeekV3MoeRoutingMethod, which nests n_group / topk_group on routing_impl. DeepSeekV4MoeRoutingMethod stores them directly on the outer object, so passing it in raised at model-init time:

ValueError: DeepSeekV3 perfect-router planning requires routing_method.routing_impl to provide n_group and topk_group

This crashed every executor rank when ENABLE_PERFECT_ROUTER=1 was set on DSv4 (e.g. bm_deepseek-v4-pro-eplb-perfectrouter-sol-1024-1024-ratio08-... benchmarks). The error path went through DeepseekV4MoE.__init__ -> ConfigurableMoE.__init__ -> _init_perfect_router() -> precompute_common_perfect_router_logits() -> planner.get_cache_signature() -> _resolve_group_config(), hitting the same exception on the fallback model_loader.load() retry.

Change

Fall back to the routing method itself when routing_impl is absent. Both V3 (nested on routing_impl) and V4 (flat on the outer object) layouts now resolve to the same (n_group, topk_group) tuple, and the error message is updated to mention both supported shapes.

Add a CPU-only regression test test_perfect_router_resolve_group_config parametrized over both routing classes that drives the planner directly.

Impact

Unblocks ENABLE_PERFECT_ROUTER=1 on DSv4 (perfect-router load-balancer benchmarks).
Backwards-compatible: V3 still resolves through routing_impl first.
No change to runtime kernels or numerics; init-time config check only.

Test Plan

pre-commit run clean on both modified files
pytest tests/unittest/_torch/modules/test_moe_routing.py::test_perfect_router_resolve_group_config (CI)
Rerun the failing DSv4 DEP=16 perfect-router benchmark to confirm init succeeds.

qiaoxj07 · 2026-05-13T06:49:39Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-13T06:56:40Z

PR_Github #48126 [ run ] triggered by Bot. Commit: e54555b Link to invocation

tensorrt-cicd · 2026-05-13T10:30:48Z

PR_Github #48126 [ run ] completed with state SUCCESS. Commit: e54555b
/LLM/main/L0_MergeRequest_PR pipeline #37952 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

qiaoxj07 · 2026-05-13T10:51:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-13T10:58:19Z

PR_Github #48171 [ run ] triggered by Bot. Commit: f44a681 Link to invocation

qiaoxj07 · 2026-05-13T11:30:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-13T11:38:45Z

PR_Github #48173 [ run ] triggered by Bot. Commit: 998d20a Link to invocation

…ew feedback) Per review feedback on PR NVIDIA#14074: instead of having the planner introspect routing_method.routing_impl vs routing_method directly to resolve n_group/topk_group, push the layout abstraction into the routing-method classes themselves. - DeepSeekV3MoeRoutingMethod gains @Property n_group / topk_group that delegate to self.routing_impl. - DeepSeekV4MoeRoutingMethod already stores them as direct attributes, no change needed. - DeepSeekV3PerfectRouterPlanner.get_cache_signature and create_logits now read routing_method.n_group / routing_method.topk_group directly. - Remove the now-redundant _resolve_group_config helper. This way the planner does not need to "recognize whether V3 or V4"; each routing method publishes the same public surface and the storage layout is an implementation detail. The unit regression test is renamed to test_grouped_routing_exposes_n_group_topk_group and asserts the properties directly on both routing-method classes; the end-to-end get_cached_perfect_router_logits test is unchanged and still covers the full planner pipeline for V3, V4 dense, and V4 hashed gates. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

qiaoxj07 · 2026-05-14T00:59:26Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-14T01:09:01Z

PR_Github #48256 [ run ] triggered by Bot. Commit: b80ff77 Link to invocation

DeepSeekV3PerfectRouterPlanner._resolve_group_config() only knew how to introspect DeepSeekV3MoeRoutingMethod, which nests n_group/topk_group on routing_impl. DeepSeekV4MoeRoutingMethod stores them directly on the outer object, so passing it in raised: ValueError: DeepSeekV3 perfect-router planning requires routing_method.routing_impl to provide n_group and topk_group at MoE construction time and crashed every executor rank when ENABLE_PERFECT_ROUTER=1 was set on DeepSeekV4 (e.g. DSv4 DEP=16 benchmarks). Fall back to the routing method itself when routing_impl is absent so both V3 (nested) and V4 (flat) layouts work, and add a CPU-only regression test that exercises both shapes through the planner. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

DeepseekV4MoE.__init__ called precompute_common_perfect_router_logits() again after self.experts = create_moe(...) returned, but the call was missing the required routing_method argument and crashed: TypeError: precompute_common_perfect_router_logits() missing 1 required positional argument: 'routing_method' This was unreachable before the previous fix because the create_moe() call blew up first; now that create_moe() succeeds on DSv4, the dead call became reachable and broke init. Every MoE backend's base class (MoE in interface.py) already invokes _init_perfect_router() during __init__, which precomputes the cache with the correct dtype (_get_perfect_router_dtype() -> float32 for V3/V4 routing) and the correct ep_rank. The block here used the model dtype (e.g. bfloat16) and the default ep_rank=0, producing useless cache entries under wrong keys even when it did run. Drop the dead block and its now-unused import. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

…oE.forward_MoE DeepseekV4MoE.forward_MoE replaced router_logits with cached "ideal" logits via _create_ideal_expert_load_balanced_logits() when ENABLE_PERFECT_ROUTER=1. The helper called get_cached_perfect_router_logits() without the required ep_rank and routing_method arguments: TypeError: get_cached_perfect_router_logits() missing 2 required positional arguments: 'ep_rank' and 'routing_method' This surfaced during warmup once the init-time precompute fix landed, because the runtime swap on this code path was never exercised before. The FusedMoE backend's own forward (interface.py: MoE.forward) already runs router_logits = self._maybe_get_perfect_router_logits(router_logits) as the very first step, using the correct ep_rank, routing_method, and dtype. ConfigurableMoE (returned by create_moe()) only overrides forward_impl/forward_fake, so MoE.forward still runs. The substitution in DeepseekV4MoE.forward_MoE was therefore both duplicated and broken; remove it together with the unused helper and import. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

DeepseekV4Gate with is_hashed=True returns None from its bias callable (fetch_e_score_correction_bias() at modeling_deepseekv4.py:1304-1308), because hashed gates route via tid2eid instead of a learned correction bias. BasePerfectRouterPlanner._force_zero_routing_bias_in_place then crashed on: TypeError: zeros_like(): argument 'input' (position 1) must be Tensor, not NoneType The init-time precompute swallowed this in its per-batch-size try/except (see Warning: Failed to precompute logits for batch size N: zeros_like() ... entries in benchmark logs; the cache ends up at 0 entries). The runtime path through MoE.forward -> _maybe_get_perfect_router_logits has no try/except, so the first forward call after warmup crashes the executor on every rank. Mask-out is a no-op when there is no learned bias to mask. Return early in _force_zero_routing_bias_in_place when bias is None. Replace the narrow _resolve_group_config unit test with an additional end-to-end CPU test that drives get_cached_perfect_router_logits for V3, V4 dense, and V4 hashed routing methods. The hashed case fails without this fix; the V3 / V4-dense cases keep the prior coverage of the planner -> _resolve_group_config -> _force_zero_routing_bias_in_place chain. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

The deployed DSv4-Pro checkpoint config has num_experts_per_tok=6 and n_routed_experts=384 (verified against /lustre/.../models/dsv4-pro/config.json). The previous test used top_k=8 and num_experts=64, which mirror the older DSv3-R1 shape and do not reflect the configuration that ENABLE_PERFECT_ROUTER=1 is actually exercised under in DSv4-Pro DEP=16 disagg-serving benchmarks. The bug fix itself does not depend on top_k, n_group, or topk_group values; the planner constraint experts_per_token <= min(topk_group, n_group) * (num_experts / n_group) is satisfied at top_k=6 with the DSv4-Pro shape (6 <= 4 * 48 = 192). But routing the tests through production-relevant values lets them catch any future regression that is shape-sensitive. Add a "shape" parametrize axis with two cases: - small (num_experts=64, moe_ep_size=1): fast CPU coverage of the three routing layouts. - dsv4_pro_dep16 (num_experts=384, moe_ep_size=16): the exact shape used by the disagg-serving benchmark that motivated this PR. Crossed with the existing v3 / v4_dense / v4_hashed parameter, the test now exercises both routing-method shapes against both expert counts, covering all four bug surfaces fixed in this PR. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

…ew feedback) Per review feedback on PR NVIDIA#14074: instead of having the planner introspect routing_method.routing_impl vs routing_method directly to resolve n_group/topk_group, push the layout abstraction into the routing-method classes themselves. - DeepSeekV3MoeRoutingMethod gains @Property n_group / topk_group that delegate to self.routing_impl. - DeepSeekV4MoeRoutingMethod already stores them as direct attributes, no change needed. - DeepSeekV3PerfectRouterPlanner.get_cache_signature and create_logits now read routing_method.n_group / routing_method.topk_group directly. - Remove the now-redundant _resolve_group_config helper. This way the planner does not need to "recognize whether V3 or V4"; each routing method publishes the same public surface and the storage layout is an implementation detail. The unit regression test is renamed to test_grouped_routing_exposes_n_group_topk_group and asserts the properties directly on both routing-method classes; the end-to-end get_cached_perfect_router_logits test is unchanged and still covers the full planner pipeline for V3, V4 dense, and V4 hashed gates. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

qiaoxj07 · 2026-05-14T09:02:20Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-14T09:09:45Z

PR_Github #48345 [ run ] triggered by Bot. Commit: 459224a Link to invocation

tensorrt-cicd · 2026-05-14T15:00:47Z

PR_Github #48345 [ run ] completed with state SUCCESS. Commit: 459224a
/LLM/main/L0_MergeRequest_PR pipeline #38152 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

qiaoxj07 · 2026-05-14T15:35:16Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-14T15:42:25Z

PR_Github #48395 [ run ] triggered by Bot. Commit: 459224a Link to invocation

tensorrt-cicd · 2026-05-14T16:16:43Z

PR_Github #48395 [ run ] completed with state FAILURE. Commit: 459224a
/LLM/main/L0_MergeRequest_PR pipeline #38198 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

qiaoxj07 · 2026-05-14T18:38:32Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-14T18:46:05Z

PR_Github #48418 [ run ] triggered by Bot. Commit: 459224a Link to invocation

tensorrt-cicd · 2026-05-14T19:23:20Z

PR_Github #48418 [ run ] completed with state SUCCESS. Commit: 459224a
/LLM/main/L0_MergeRequest_PR pipeline #38219 completed with status: 'SUCCESS'

CI Report

Link to invocation

Falls back to per-head RMS reshape when torch.ops.trtllm.deepseek_v4_q_norm is not registered (e.g., Python-only patched image w/o C++ rebuild). Lets us deploy PRs NVIDIA#14074 + NVIDIA#13975 via build_python_changes.sh without a full Docker rebuild. C++ rebuild later will pick up the perf benefit of the fused CUDA kernel automatically.

…DIA#14074) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> (cherry picked from commit d178370) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

…DIA#14074) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

…DIA#14074) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> (cherry picked from commit d178370) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

qiaoxj07 requested a review from a team as a code owner May 13, 2026 06:44

qiaoxj07 requested review from xxi-nv and removed request for a team May 13, 2026 06:44

github-actions Bot assigned qiaoxj07 May 13, 2026

xxi-nv approved these changes May 13, 2026

View reviewed changes

qiaoxj07 requested a review from a team as a code owner May 13, 2026 07:55

qiaoxj07 requested review from syuoni and removed request for a team May 13, 2026 07:55

xxi-nv reviewed May 13, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/fused_moe/routing.py Outdated

lfr-0531 force-pushed the feat/deepseek_v4 branch from 0a93d10 to 118e7a5 Compare May 14, 2026 07:44

lfr-0531 requested a review from a team as a code owner May 14, 2026 07:44

lfr-0531 requested a review from a team May 14, 2026 07:44

lfr-0531 requested a review from a team as a code owner May 14, 2026 07:44

lfr-0531 requested review from yiqingy0 and removed request for a team May 14, 2026 07:44

qiaoxj07 added 5 commits May 14, 2026 16:08

qiaoxj07 force-pushed the fix/dsv4-perfect-router-routing-impl branch from b80ff77 to 459224a Compare May 14, 2026 08:10

lfr-0531 merged commit d178370 into NVIDIA:feat/deepseek_v4 May 15, 2026
6 checks passed

lfr-0531 added the deepseek-v4 label May 19, 2026

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 10, 2026

[None][fix] Support DeepSeekV4 routing in perfect-router planner (NVI…

b213f08

…DIA#14074) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 11, 2026

[None][fix] Support DeepSeekV4 routing in perfect-router planner (NVI…

cb1a20b

…DIA#14074) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 12, 2026

[None][fix] Support DeepSeekV4 routing in perfect-router planner (NVI…

9c0a0b2

…DIA#14074) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

Uh oh!

Conversation

qiaoxj07 commented May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Change

Impact

Test Plan

Uh oh!

qiaoxj07 commented May 13, 2026

Uh oh!

tensorrt-cicd commented May 13, 2026

Uh oh!

tensorrt-cicd commented May 13, 2026

Uh oh!

qiaoxj07 commented May 13, 2026

Uh oh!

tensorrt-cicd commented May 13, 2026

Uh oh!

qiaoxj07 commented May 13, 2026

Uh oh!

tensorrt-cicd commented May 13, 2026

Uh oh!

Uh oh!

qiaoxj07 commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

qiaoxj07 commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

qiaoxj07 commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

qiaoxj07 commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

qiaoxj07 commented May 13, 2026 •

edited

Loading