[https://nvbugs/5451740][fix] Add DP padding back on SM120

[peaceh-nv]

Description

1. Add DP padding back on the devices that have limited communication bw, current just enabled for SM120. Without DP padding, NCCL reduce_scatter + all_gather would fallback to reduce + broadcast, which can cause 30% context phase comm perf regression.
2. Another non arch-specific issue caused by dp_padding missing is that when num_chunks>1 for chunked MoE in IFB scenario (e.g. input_seqlen=4k, moe_max_num_tokens=3k, num_chunks=2), input tokens for 8 ranks are [4096, 1, 1, 1, 1, 1, 1, 1], all ranks needs to be split into 2 chunks, which makes rank1-7 have input tensor with zero dim, which leads to NCCL kernel failure.

Summary by CodeRabbit

• New Features

  • Added automatic DP padding for attention on SM120 in specific MoE + tensor-parallel setups to ensure consistent routing and outputs.
  • Exposed additional MoE backends for use in configurations, enabling broader deployment options.

• Bug Fixes

  • Resolved potential shape/routing inconsistencies in MoE pipelines when using tensor parallelism on SM120 by applying conditional padding.

[peaceh-nv]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds SM120-specific DP padding logic to compute_routed_output paths, conditionally padding hidden_states and propagating a use_dp_padding flag through routing/experts. Updates imports to expose CutlassFusedMoE and WideEPMoE and drop MoEWeightLoadingMode.

Changes

Cohort / File(s)	Summary
MoE DP padding and imports tensorrt_llm/_torch/models/modeling_deepseekv3.py	- Update imports: add CutlassFusedMoE, WideEPMoE; remove MoEWeightLoadingMode - Add SM120 DP-padding decision for attention when TP>1 and specific MoE/FP8/QDQ conditions - Pad hidden_states by token dimension to max rank tokens when enabled - Pass use_dp_padding into routing/experts - Mirror the same logic in Deepseekv3MoE.compute_routed_output

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor Caller
  participant Model as compute_routed_output
  participant Router as Routing Method
  participant Experts as MoE Experts

  Caller->>Model: forward(hidden_states, ...)
  Note over Model: Determine use_dp_padding if SM120 && TP>1 &&\n(MoE type/FP8/QDQ conditions)
  alt use_dp_padding = true
    Model->>Model: Pad hidden_states (token dim) to max tokens across ranks
  end
  Model->>Router: route(hidden_states, use_dp_padding)
  Router-->>Model: routing logits / assignments
  Model->>Experts: dispatch(hidden_states, assignments, use_dp_padding)
  Experts-->>Model: expert outputs (padded if enabled)
  alt use_dp_padding = true
    Model->>Model: Depad outputs to original token lengths
  end
  Model-->>Caller: routed output
  Note right of Model: Same flow mirrored in Deepseekv3MoE.compute_routed_output

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description Check	⚠️ Warning	The PR description provides a clear explanation of the issue and solution but omits the required “## Test Coverage” section detailing relevant tests and the “## PR Checklist” section from the repository’s template.	Please add a “## Test Coverage” section outlining the tests that verify the new DP padding logic and include the “## PR Checklist” section with the appropriate checklist items as specified in the repository template.
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The title correctly uses a valid NVBugs identifier and a clear “[fix]” tag and succinctly summarizes the main change of reintroducing DP padding on SM120 devices without unnecessary detail or noise.

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between b1dc84b and 28053d7.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/models/modeling_deepseekv3.py (2 hunks)

🧰 Additional context used

📓 Path-based instructions (3)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use only spaces, no tabs; indent with 4 spaces.

Files:

• tensorrt_llm/_torch/models/modeling_deepseekv3.py

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.

Files:

• tensorrt_llm/_torch/models/modeling_deepseekv3.py

**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).

Files:

• tensorrt_llm/_torch/models/modeling_deepseekv3.py

🧠 Learnings (1)

📚 Learning: 2025-09-19T21:28:13.751Z

Learnt from: jhaotingc
PR: NVIDIA/TensorRT-LLM#7856
File: cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp:159-166
Timestamp: 2025-09-19T21:28:13.751Z
Learning: In TensorRT-LLM blockScaleMoe routing (cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu), the DeepSeek routing method performs reinterpret_cast<float*>(routingLogits) at line 89, which could cause issues if routing_logits are BF16. However, Qwen3-FP8 models use RenormalizeNaive routing method and are not affected by this dtype casting issue.

Applied to files:

• tensorrt_llm/_torch/models/modeling_deepseekv3.py

🧬 Code graph analysis (1)

tensorrt_llm/_torch/models/modeling_deepseekv3.py (4)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)

• CutlassFusedMoE (25-600)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (1)

• WideEPMoE (41-1019)

tensorrt_llm/_utils.py (1)

• get_sm_version (689-691)

tensorrt_llm/_torch/modules/fused_moe/interface.py (2)

• has_fp8_qdq (284-287)
• has_nvfp4 (296-299)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Pre-commit Check

🔇 Additional comments (3)

tensorrt_llm/_torch/models/modeling_deepseekv3.py (3)

60-61: Import changes look good and align with PR objectives.

The updated imports correctly add CutlassFusedMoE and WideEPMoE while removing MoEWeightLoadingMode, which matches the AI summary and supports the DP padding functionality being added.

---

856-866: LGTM! SM120-specific DP padding logic correctly implemented.

The DP padding logic is well-implemented with appropriate conditions:

• Targets SM120 specifically (get_sm_version() == 120)
• Applies only when DP is enabled and TP size > 1
• Correctly excludes CutlassFusedMoE and WideEPMoE from padding or requires specific quantization configurations
• Properly pads hidden_states to match maximum rank token count across all ranks

The padding prevents the communication performance regression and chunked MoE failures described in the PR objectives.

---

876-876: DP padding flag correctly propagated to experts.

The use_dp_padding flag is appropriately passed to the experts forward call, ensuring the padding behavior is propagated through the MoE computation pipeline.

Tip

👮 Agentic pre-merge checks are now available in preview!

Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.

• Built-in checks – Quickly apply ready-made checks to enforce title conventions, require pull request descriptions that follow templates, validate linked issues for compliance, and more.
• Custom agentic checks – Define your own rules using CodeRabbit’s advanced agentic capabilities to enforce organization-specific policies and workflows. For example, you can instruct CodeRabbit’s agent to verify that API documentation is updated whenever API schema files are modified in a PR. Note: Upto 5 custom checks are currently allowed during the preview period. Pricing for this feature will be announced in a few weeks.

Please see the documentation for more information.

Example:

reviews:
  pre_merge_checks:
    custom_checks:
      - name: "Undocumented Breaking Changes"
        mode: "warning"
        instructions: |
          Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).

Please share your feedback with us on this Discord post.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[tensorrt-cicd]

PR_Github #19824 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #19824 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #14916 completed with status: 'FAILURE'

[peaceh-nv]

/bot run

[peaceh-nv]

/bot run

[tensorrt-cicd]

PR_Github #19864 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #19865 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #19864 [ run ] completed with state ABORTED
LLM/main/L0_MergeRequest_PR #14947 (Blue Ocean) completed with status: ABORTED

[peaceh-nv]

/bot run

[tensorrt-cicd]

PR_Github #19872 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #19865 [ run ] completed with state ABORTED
LLM/main/L0_MergeRequest_PR #14949 (Blue Ocean) completed with status: ABORTED

[peaceh-nv]

/bot run

[tensorrt-cicd]

PR_Github #19875 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #19872 [ run ] completed with state ABORTED
LLM/main/L0_MergeRequest_PR #14955 (Blue Ocean) completed with status: ABORTED

[peaceh-nv]

/bot run

[tensorrt-cicd]

PR_Github #19890 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #19875 [ run ] completed with state ABORTED
LLM/main/L0_MergeRequest_PR #14957 (Blue Ocean) completed with status: ABORTED

[peaceh-nv]

/bot run

[tensorrt-cicd]

PR_Github #19893 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #19890 [ run ] completed with state ABORTED
LLM/main/L0_MergeRequest_PR #14969 (Blue Ocean) completed with status: ABORTED

[tensorrt-cicd]

PR_Github #19893 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #14972 completed with status: 'FAILURE'

[peaceh-nv]

/bot run

[tensorrt-cicd]

PR_Github #20008 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #20008 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #15069 completed with status: 'SUCCESS'