[None][feat] Add DWDP (Distributed Weight Data Parallelism) support for MoE inference by tianyuz-nv · Pull Request #12136 · NVIDIA/TensorRT-LLM

tianyuz-nv · 2026-03-12T03:36:37Z

Summary by CodeRabbit

Release Notes

New Features
- Added DWDP (Double-Weighted Distributed Prefetching) support for disaggregated MoE inference, enabling efficient cross-rank weight prefetching with ping-pong buffering for improved serving performance.
- Extended MoE kernels to support multiple weight tensor configurations, improving flexibility for advanced quantization and distributed weight management scenarios.
Bug Fixes
- Fixed rank synchronization in disaggregated MoE initialization to use session-scoped barriers for improved compatibility.
Tests
- Added DWDP accuracy test for DeepSeek-V3-Lite model with GSM8K benchmark.

Description

In the context phase of LLM inference, workload imbalances and communication bottlenecks often lead to excessive synchronization overhead, limiting GPU utilization. To resolve this, we propose Distributed Weight Data Parallelism (DWDP), a strategy that leverages Data Parallelism combined with NVLink-based weight offloading to enable fully asynchronous execution across ranks.

Key properties of DWDP:

Context-phase acceleration: DWDP is designed specifically for the context (prefill) phase of disaggregated serving.
Hardware-aware: DWDP performance gains rely on the GB200 NVL72 architecture for high-bandwidth NVLink weight transfer.
Fully asynchronous execution: DWDP eliminates synchronization barriers across ranks, providing significant performance advantages under imbalanced workloads.
Fine-grained resource management: DWDP enables flexible expert-to-worker assignment, allowing more efficient GPU memory utilization.

A detailed technical report on DWDP internals and optimizations will be published separately. We welcome discussions and feedback.

Changes

Core DWDP runtime (tensorrt_llm/_torch/pyexecutor/dwdp.py — new file):

DwdpManager: Orchestrates IPC handle exchange across MPI ranks for zero-copy expert weight sharing
DwdpHandleCollector: Per-layer collector that gathers CUDA IPC handles for weight/scale/bias tensors
Expert weight prefetching with double-buffering via CUDA streams

MoE integration (configurable_moe.py, fused_moe_cute_dsl.py, interface.py):

Integrate DWDP into ConfigurableMoE with automatic detection of compatible backends (CuteDSL + NVFP4)
Add NvFp4WeightView dataclass for clean separation of DWDP vs non-DWDP weight access patterns
Add contiguous gather/scatter grouped GEMM kernels for DWDP expert subset computation

CuteDSL kernel extensions (cute_dsl_custom_ops.py, blockscaled_contiguous_*_fusion.py):

Extend blockscaled contiguous grouped GEMM kernels to support DWDP's gather-based expert selection
Unify single-B and multi-B kernel wrappers into a single entry point

Executor integration (py_executor.py, py_executor_creator.py, _util.py, llm_args.py):

Add DwdpConfig dataclass to LlmArgs for YAML-based DWDP configuration
Initialize DwdpManager during executor creation when DWDP is enabled
Wire DWDP expert prefetching into the per-step execution loop

Disaggregated serving scripts (examples/disaggregated/slurm/benchmark/):

Add start_worker_dwdp.sh for launching DWDP workers via mpirun
Extend submit.py with DWDP-specific configuration (dwdp_size, num_group, experts_per_worker, etc.)

Test Coverage

tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py: Unit tests for DWDP-specific CuteDSL MoE kernels (contiguous gather grouped GEMM with SwiGLU fusion, finalize fusion)
tests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py: Standalone kernel correctness tests for DWDP gather GEMM
tests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_grouped_gemm_finalize_fusion.py: Standalone kernel correctness tests for DWDP finalize GEMM
tests/integration/defs/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_dwdp_accuracy: End-to-end DWDP disaggregated serving accuracy test with DeepSeek-V3-Lite (NVFP4, 4 GPUs, GSM8K)

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)
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.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

syuoni

This is a nice feature, great to see it's going to production.

Regarding the CuTeDSL MoE interface, overall I would suggest:

Use two different ops for single-b and multi-b cases.
Above the op level, we use two different code paths; this avoids confusion to users -- most users would use the single-b op only.
Below the op level, we unify them by multi-b implementation; this simplifies implementation and avoids code duplication.

Thanks!

…or MoE inference Core DWDP runtime (dwdp.py): - DwdpManager: IPC handle exchange across MPI ranks - DwdpHandleCollector: per-layer weight/scale/bias handle collection - Expert weight prefetching with double-buffering MoE integration (configurable_moe.py, fused_moe_cute_dsl.py, interface.py): - DWDP support in ConfigurableMoE with CuteDSL+NVFP4 backend - NvFp4WeightView for DWDP weight access patterns - Contiguous gather/scatter grouped GEMM kernels CuteDSL kernel extensions: - Blockscaled contiguous gather grouped GEMM with SwiGLU fusion - Blockscaled contiguous grouped GEMM finalize fusion Executor integration (py_executor.py, py_executor_creator.py, llm_args.py): - DwdpConfig dataclass for YAML-based configuration - DwdpManager initialization and per-step prefetching Disaggregated serving scripts: - start_worker_dwdp.sh for MPI-based worker launch - submit.py DWDP configuration support CI test: - DWDP accuracy test with DeepSeek-V3-Lite (NVFP4, 4 GPUs, GSM8K) Co-authored-by: wanqian-nv <221923321+wanqian-nv@users.noreply.github.com> Co-authored-by: zongfeijing <20381269+zongfeijing@users.noreply.github.com> Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>

tianyuz-nv · 2026-03-19T08:17:29Z

/bot run

JintaoPengCS · 2026-03-20T02:25:02Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-20T02:31:39Z

PR_Github #39672 [ run ] triggered by Bot. Commit: bbe48fa Link to invocation

tensorrt-cicd · 2026-03-20T08:46:40Z

PR_Github #39672 [ run ] completed with state SUCCESS. Commit: bbe48fa
/LLM/main/L0_MergeRequest_PR pipeline #30874 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

JintaoPengCS · 2026-03-23T00:23:04Z

/bot run --disable-fail-fast

Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>

…matting - Remove DwdpConfig.from_dict() to comply with Pydantic best practices test - Initialize GroupedGemmInputsHelper in test_nvfp4_gather_grouped_gemm_swiglu_blackwell - Apply pre-commit formatting (isort, yapf, ruff, autoflake, trailing whitespace) Signed-off-by: Tianyu Zhang <tianyuz@nvidia.com> Signed-off-by: tianyuz-nv <tianyuz@nvidia.com> Made-with: Cursor

Kefeng-Duan · 2026-03-23T11:25:11Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-23T11:30:53Z

PR_Github #39925 [ run ] triggered by Bot. Commit: f153099 Link to invocation

tensorrt-cicd · 2026-03-23T19:55:37Z

PR_Github #39925 [ run ] completed with state SUCCESS. Commit: f153099
/LLM/main/L0_MergeRequest_PR pipeline #31092 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

tianyuz-nv · 2026-03-31T17:42:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-31T17:51:04Z

PR_Github #40992 [ run ] triggered by Bot. Commit: afcd735 Link to invocation

tensorrt-cicd · 2026-04-01T02:18:53Z

PR_Github #40992 [ run ] completed with state FAILURE. Commit: afcd735
/LLM/main/L0_MergeRequest_PR pipeline #31975 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

tianyuz-nv · 2026-04-01T15:28:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-01T15:34:52Z

PR_Github #41218 [ run ] triggered by Bot. Commit: be12482 Link to invocation

tensorrt-cicd · 2026-04-01T23:32:16Z

PR_Github #41218 [ run ] completed with state SUCCESS. Commit: be12482
/LLM/main/L0_MergeRequest_PR pipeline #32181 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

Kefeng-Duan · 2026-04-01T23:46:55Z

/bot run --stage-list "DGX_B300-4_GPUs-PyTorch-1"

tensorrt-cicd · 2026-04-01T23:53:28Z

PR_Github #41281 [ run ] triggered by Bot. Commit: be12482 Link to invocation

tensorrt-cicd · 2026-04-02T01:39:22Z

PR_Github #41281 [ run ] completed with state FAILURE. Commit: be12482
/LLM/main/L0_MergeRequest_PR pipeline #32239 (Partly Tested) 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

Kefeng-Duan · 2026-04-02T01:54:11Z

/bot run --stage-list "DGX_B300-4_GPUs-PyTorch-1" --disable-fail-fast

tensorrt-cicd · 2026-04-02T02:00:45Z

PR_Github #41300 [ run ] triggered by Bot. Commit: be12482 Link to invocation

tensorrt-cicd · 2026-04-02T03:01:40Z

PR_Github #41300 [ run ] completed with state SUCCESS. Commit: be12482
/LLM/main/L0_MergeRequest_PR pipeline #32254 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

Kefeng-Duan · 2026-04-02T03:23:39Z

/LLM/main/L0_MergeRequest_PR pipeline #32181 and /LLM/main/L0_MergeRequest_PR pipeline #32254 (Partly Tested) have covered all CI tests, and all success.

Kefeng-Duan · 2026-04-02T03:23:49Z

/bot skip

github-actions · 2026-04-02T03:23:57Z

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental) --high-priority]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

Kefeng-Duan · 2026-04-02T03:25:45Z

/bot skip --comment '/LLM/main/L0_MergeRequest_PR pipeline #32181 and /LLM/main/L0_MergeRequest_PR pipeline #32254 (Partly Tested) have covered all CI tests, and all success.'

tensorrt-cicd · 2026-04-02T03:31:24Z

PR_Github #41320 Bot args parsing error: Failed to parse bot args

Link to invocation

tianyuz-nv · 2026-04-02T03:35:28Z

/bot skip --comment "L0 MergeRequest PR pipelines 32181 and 32254 (partly tested) already covered CI; all success."

tensorrt-cicd · 2026-04-02T03:40:56Z

PR_Github #41325 [ skip ] triggered by Bot. Commit: be12482 Link to invocation

tensorrt-cicd · 2026-04-02T03:51:56Z

PR_Github #41325 [ skip ] completed with state SUCCESS. Commit: be12482
Skipping testing for commit be12482

Link to invocation

…nfigurableMoE load_weights PR NVIDIA#12136 (DWDP) added a load_weights override in CuteDslFusedMoE that dropped the allow_partial_loading parameter from the base class signature. ConfigurableMoE.load_weights also lacked this parameter. This causes TypeError when qwen2_moe_weight_mapper calls module.load_weights(weights=..., allow_partial_loading=...) on models using the CuteDSL or ConfigurableMoE backend (e.g., Qwen3 MoE). Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

…or MoE inference (NVIDIA#12136) Signed-off-by: tianyuz-nv <tianyuz@nvidia.com> Signed-off-by: Tianyu Zhang <tianyuz@nvidia.com> Signed-off-by: Kefeng-Duan <176893526+Kefeng-Duan@users.noreply.github.com> Co-authored-by: wanqian-nv <221923321+wanqian-nv@users.noreply.github.com> Co-authored-by: zongfeijing <20381269+zongfeijing@users.noreply.github.com> Co-authored-by: Kefeng-Duan <176893526+Kefeng-Duan@users.noreply.github.com>

svc-trtllm-gh-bot added the Community want to contribute PRs initiated from Community label Mar 12, 2026

xxi-nv reviewed Mar 12, 2026

View reviewed changes

syuoni reviewed Mar 12, 2026

View reviewed changes

pengbowang-nv removed the Community want to contribute PRs initiated from Community label Mar 12, 2026

svc-trtllm-gh-bot added the Community want to contribute PRs initiated from Community label Mar 12, 2026

wanqian-nv reviewed Mar 12, 2026

View reviewed changes

tianyuz-nv changed the title ~~Dwdp productization~~ [None][feat] Add DWDP (Disaggregated Weight Data Parallelism) support for MoE inference Mar 17, 2026

tianyuz-nv changed the title ~~[None][feat] Add DWDP (Disaggregated Weight Data Parallelism) support for MoE inference~~ [None][feat] Add DWDP (Distributed Weight Data Parallelism) support for MoE inference Mar 17, 2026

tianyuz-nv force-pushed the dwdp_productization branch 2 times, most recently from 1d3e509 to 3532338 Compare March 19, 2026 05:55

tianyuz-nv force-pushed the dwdp_productization branch from 3532338 to bbe48fa Compare March 19, 2026 08:09

tianyuz-nv added 3 commits March 22, 2026 22:47

Merge upstream/main into dwdp_productization

7e4a855

Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>

Register DWDP accuracy test in CI test lists

f76d026

Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>

tianyuz-nv marked this pull request as ready for review March 24, 2026 02:39

tianyuz-nv requested review from a team as code owners March 24, 2026 02:39

Merge remote-tracking branch 'upstream/main' into dwdp_productization

be12482

Kefeng-Duan enabled auto-merge (squash) April 2, 2026 03:29

Kefeng-Duan merged commit e92ee4f into NVIDIA:main Apr 2, 2026
5 checks passed

tianyuz-nv mentioned this pull request Apr 2, 2026

[None][fix] Fix CuteDslFusedMoE.load_weights signature to accept allow_partial_loading #12690

Closed

hlu1 mentioned this pull request Apr 4, 2026

[Feature] Distributed Weight Data Parallelism (DWDP) for Sparse MoE Models sgl-project/sglang#22084

Open

qiaoxj07 mentioned this pull request Apr 4, 2026

[None][fix] Add missing allow_partial_loading param to CuteDSL and ConfigurableMoE load_weights #12761

Merged

2 tasks

hlu1 mentioned this pull request Apr 11, 2026

[CuteDSL MoE] Update the CuteDSL MoE kernels support DWDP flashinfer-ai/flashinfer#3036

Open

yhyang201 mentioned this pull request Apr 13, 2026

[WIP] feat: support multi-B weight tensors (DWDP) in CuTe DSL NVFP4 MoE flashinfer-ai/flashinfer#3041

Open

5 tasks

Conversation

tianyuz-nv commented Mar 12, 2026 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Release Notes

Description

Changes

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

syuoni left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tianyuz-nv commented Mar 19, 2026

Uh oh!

JintaoPengCS commented Mar 20, 2026

Uh oh!

tensorrt-cicd commented Mar 20, 2026

Uh oh!

tensorrt-cicd commented Mar 20, 2026

Uh oh!

JintaoPengCS commented Mar 23, 2026

Uh oh!

Kefeng-Duan commented Mar 23, 2026

Uh oh!

tensorrt-cicd commented Mar 23, 2026

Uh oh!

tensorrt-cicd commented Mar 23, 2026

Uh oh!

tianyuz-nv commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

tianyuz-nv commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

Kefeng-Duan commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

Kefeng-Duan commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

Kefeng-Duan commented Apr 2, 2026

Uh oh!

tianyuz-nv commented Mar 12, 2026 •

edited by coderabbitai bot

Loading