[None][feat] Add DWDP (Distributed Weight Data Parallelism) support for MoE inference - #12136
Conversation
svc-trtllm-gh-bot
added
the
Community want to contribute
syuoni
left a comment
syuoni
left a comment
There was a problem hiding this comment.
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!
pengbowang-nv
removed
the
Community want to contribute
svc-trtllm-gh-bot
added
the
Community want to contribute
tianyuz-nv
changed the title
tianyuz-nv
changed the title
1d3e509 to
3532338
Compare
…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>
3532338 to
bbe48fa
Compare
|
/bot run |
|
/bot run --disable-fail-fast |
|
PR_Github #39672 [ run ] triggered by Bot. Commit: |
|
PR_Github #39672 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
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
|
/bot run --disable-fail-fast |
|
PR_Github #39925 [ run ] triggered by Bot. Commit: |
|
PR_Github #39925 [ run ] completed with state
|
|
PR_Github #41281 [ run ] completed with state
|
|
/bot run --stage-list "DGX_B300-4_GPUs-PyTorch-1" --disable-fail-fast |
|
PR_Github #41300 [ run ] triggered by Bot. Commit: |
|
PR_Github #41300 [ run ] completed with state |
|
/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. |
|
/bot skip |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. 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. |
|
/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.' |
|
PR_Github #41320 Bot args parsing error: Failed to parse bot args |
|
/bot skip --comment "L0 MergeRequest PR pipelines 32181 and 32254 (partly tested) already covered CI; all success." |
|
PR_Github #41325 [ skip ] triggered by Bot. Commit: |
|
PR_Github #41325 [ skip ] completed with state |
…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>
… unused by VA path) Commit 3 of the DWDP IPC->VA refactor. Three files (custom op wrapper + two blackwell kernels) are restored verbatim to their pre-PR-NVIDIA#12136 state. These multi-B paths were introduced purely to support DWDP's IPC scheme, which passed N peer expert shards as separate B tensors into each kernel call. The VA pipeline swaps param.data to a single [num_experts, ...] tensor via cuMemMap, so the standard single-B kernel path handles every case — the multi-B parameters, MAX_B_TENSORS branches, and tuple-ified b/sfb/alpha signatures become dead code. Files reverted to the commit before e92ee4f (PR NVIDIA#12136): * _torch/custom_ops/cute_dsl_custom_ops.py - Removes *_multi_b custom op registrations - Restores GatherGroupedGemmInputsHelper to single-tensor layout * _torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py - Removes b_tensor_l_sizes param, MAX_B_TENSORS, num_b_tensors const_expr branches for 1/2/3/4 B tensors, _make_tma_b helper, kernel-side gB_nkl_0..3 / gSFB_nkl_0..3 expansions, tuple'd signatures * _torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py - Same pattern Verified no commits touched these files between PR NVIDIA#12136 and HEAD, so the revert is surgical and does not risk clobbering unrelated work: $ git log --oneline e92ee4f..HEAD -- <file> # empty for all 3 Net change: +371 / -1519 = -1148 lines in the three files. Smoke tested: 78 pytest passes (DWDP units + full api_stability), plus runtime import checks confirm both kernel modules and the custom ops module load cleanly after the revert. Co-Authored-By: dongxuy04 <dongxuy@nvidia.com> Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
`num_groups` has been a dead schema field in DwdpConfig from PR NVIDIA#12136 (the original IPC implementation) through commit 1fbc0d4: read into DwdpManager but never consumed by the runtime. Mis-configured YAMLs where the user's declared topology disagrees with the launch were left to fail mysteriously deeper in MPI sub-communicator creation, or to silently accept a launch that didn't match the schema. Convert the field into runtime validation in DwdpManager.__init__: 1. num_groups must be positive. 2. This rank's computed group_id (`rank // dwdp_size`) must be less than num_groups, so the launch hasn't started more CTX workers than the declared topology can hold. 3. `num_groups * dwdp_size <= MPI world size`, so the world is large enough to fit all declared groups. The three checks together catch over-subscription, world-size under-allocation, and obviously bogus values, while remaining local (no extra inter-rank communication required since DwdpConfig is identical on every rank). Verification: - Unit tests: 61 passed + 4 subtests (4 new num_groups cases in tests/unittest/_torch/executor/test_dwdp_manager.py) - Accuracy (DSv3-Lite + dwdp=2 + GSM8K): PASS, above the 61.5% threshold (3min 52s) - Perf DEP baseline: 12.95 req/s - Perf DWDP=4: 14.26 req/s (+10.1% over DEP) - Perf DWDP=8 cross-tray: 27.21 req/s (1.91x DWDP=4 scaling) Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
`num_groups` has been a dead schema field in DwdpConfig from PR NVIDIA#12136 (the original IPC implementation) through commit 1fbc0d4: read into DwdpManager but never consumed by the runtime. Mis-configured YAMLs where the user's declared topology disagrees with the launch were left to fail mysteriously deeper in MPI sub-communicator creation, or to silently accept a launch that didn't match the schema. Convert the field into runtime validation in DwdpManager.__init__: 1. num_groups must be positive. 2. This rank's computed group_id (`rank // dwdp_size`) must be less than num_groups, so the launch hasn't started more CTX workers than the declared topology can hold. 3. `num_groups * dwdp_size <= MPI world size`, so the world is large enough to fit all declared groups. The three checks together catch over-subscription, world-size under-allocation, and obviously bogus values, while remaining local (no extra inter-rank communication required since DwdpConfig is identical on every rank). Verification: - Unit tests: 61 passed + 4 subtests (4 new num_groups cases in tests/unittest/_torch/executor/test_dwdp_manager.py) - Accuracy (DSv3-Lite + dwdp=2 + GSM8K): PASS, above the 61.5% threshold (3min 52s) - Perf DEP baseline: 12.95 req/s - Perf DWDP=4: 14.26 req/s (+10.1% over DEP) - Perf DWDP=8 cross-tray: 27.21 req/s (1.91x DWDP=4 scaling) Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Phase 1 of removing the DWDP IPC-era multi-B kernel infrastructure that the VA-based DWDP path no longer needs. Three follow-up phases will clean the runner-class plumbing and the kernel-side multi-B branches. Multi-B (multiple B tensors passed as a list to a single fused MoE NVFP4 gather kernel) was introduced in PR NVIDIA#12136 to let the original CUDA-IPC DWDP path GEMM N peer expert shards in one kernel call. In the VA-based pipeline (commit "Switch DWDP integration layer from IPC to VA") the composite VA region exposes peer experts as one contiguous [num_experts, ...] tensor via cuMemMap, so single-B handles every case. After that commit the only callers of the multi-B custom op were: * The single-B custom op itself, which had been rewired to wrap single tensors into single-element lists and forward to multi-B (see PR NVIDIA#12136 commits "make multi_b the inner op, original op as wrapper" and "unify single-B and multi-B kernel wrapper into single entry point"). * test_cute_dsl_moe.py, which exercised both paths. This commit: * Removes the trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion _blackwell_multi_b custom op and its register_fake registration. * Rewires the single-B trtllm::cute_dsl_nvfp4_gather_grouped_gemm _act_fusion_blackwell to call Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner directly, restoring it as the primary entry point. The runner / helper still expect list-form for weight, weight_scale, alpha; a one-line wrap-in-single-element-list keeps Phase 1 minimal. Phase 2 will convert the runner back to bare-tensor inputs. * Updates test_cute_dsl_moe.py to call the single-B op directly and drops the multi-B (num_local_experts > 1) test branch. Verified that nemotron-h, the only model that consumes the new ``activation_type=Relu2`` path added by PR NVIDIA#12884, calls the single-B op (fused_moe_cute_dsl.py:631), not multi-B — so multi-B has no remaining production callers. Both renames + ``activation_type`` parameter introduced by PR NVIDIA#12884 (rename swiglu_fusion → act_fusion, add Relu2 support) are preserved. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…r classes Phase 2 of removing the DWDP IPC-era multi-B kernel infrastructure that the VA-based DWDP path no longer needs. Phase 1 dropped the public multi-B custom op; this phase reverts the wrapper-level plumbing back to bare-tensor inputs (matching the pre-PR-NVIDIA#12136 shape) while preserving every unrelated upstream evolution — notably the swiglu_fusion -> act_fusion rename and ``activation_type`` parameter introduced by PR NVIDIA#12884 (CUTEDSL MoE backend for nemotron-h, the only production consumer of ``activation_type=Relu2``). Changes in tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py: * GatherGroupedGemmInputsHelper.inputs_pre_hook: unpack inputs as bare tensors (``a, b, a_sf, b_sf, alpha, ...``) instead of list form (``a, b_list, a_sf, b_sf_list, alpha_list, ...``); update class docstring accordingly. * Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner: drop the ``b_tensor_l_sizes`` parameter from __init__, ``unique_id`` and the kernel cache key; rewrite ``get_valid_tactics`` / ``forward`` to take bare ``b`` / ``b_sf`` / ``alpha`` tensors and pass single pointers (not tuples) to the kernel. * Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner: same revert; in addition, remove the ``MAX_B_TENSORS = 4`` class constant. ``activation_type`` and the gated/non-gated branches added by PR NVIDIA#12884 are preserved. * cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell and cute_dsl_nvfp4_grouped_gemm_finalize_blackwell custom ops: drop the ``weight: List[torch.Tensor]`` / ``weight_scale: List[...]`` / ``alpha: List[...]`` type annotations in favour of bare ``torch.Tensor``; drop the runtime ``b_tensor_l_sizes`` derivation in finalize_inplace; both fake registrations follow. * cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell: drop the Phase 1 single-element-list wrapping that was carrying the runner through; the call site now passes bare tensors. Changes in tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py: * NvFp4WeightView: six ``List[torch.Tensor]`` fields become ``torch.Tensor`` since the VA composite-VA tensor swap leaves a single full ``[num_experts, ...]`` tensor in every field. ``_build_local_weight_view`` constructs accordingly; six call sites drop their ``[0]`` indexing. * cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell call site drops single-element list wrapping for weight / weight_scale / alpha. Changes in tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py: * Update finalize-blackwell test (around line 563) to call the bare- tensor signature. * Drop the ``num_local_experts > 1`` multi-B branch of the finalize test (the API is gone). Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…nels
Phase 3 of removing the DWDP IPC-era multi-B kernel infrastructure
that the VA-based DWDP path no longer needs. Phase 1 dropped the
public multi-B custom op; Phase 2 reverted the wrapper-level plumbing
to bare-tensor inputs. This phase reverts the cute.jit ``wrapper`` /
``__init__`` of the two affected kernels to a single-B-only shape and
drops the runner-side single-element-tuple bridging that Phase 2 had
left in place.
Affected kernels:
* ``blockscaled_contiguous_gather_grouped_gemm_act_fusion`` (the
nemotron-h activation-fusion gather GEMM)
* ``blockscaled_contiguous_grouped_gemm_finalize_fusion``
Changes per kernel (mirrors the diff in PR NVIDIA#12136-era commit
``b3cfd89065``, adapted onto the post-PR-NVIDIA#12884 ``act_fusion`` rename
and ``activation_type`` parameter):
* ``__init__``: drop the ``b_tensor_l_sizes`` parameter, the
``MAX_B_TENSORS = 4`` class constant, and the if/else block that
populated ``self.b_tensor_l_sizes`` / ``self.b_tensor_l_offsets``.
Set ``self.num_b_tensors = 1`` unconditionally so existing
downstream ``cutlass.const_expr(self.num_b_tensors >= 2/3/4)``
branches in ``__call__`` statically fold away at cute.jit compile
time.
* ``wrapper``: switch from ``b_ptr_tuple`` / ``b_sf_ptr_tuple`` /
``alpha_ptr_tuple`` to bare ``b_ptr`` / ``b_sf_ptr`` / ``alpha_ptr``
pointers and accept ``l: cutlass.Int64`` directly (instead of
reading ``self.b_tensor_l_sizes[0]``). Drop the up-to-four-element
unroll that built ``b_tuple`` / ``b_sf_tuple`` / ``alpha_tuple``;
the wrapper now constructs single ``b`` / ``b_sf`` / ``alpha``
tensors and forwards them to ``__call__``. ``__call__`` retains
its ``isinstance(b, tuple)`` defensive wrap that promotes a bare
tensor to a single-element tuple, so its multi-B-shaped internals
are untouched.
Corresponding changes in
``tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py``:
* Both ``forward`` methods (Sm100BlockScaledContiguousGroupedGemm
FinalizeFusionRunner, Sm100BlockScaledContiguousGatherGroupedGemm
ActFusionRunner) drop the Phase 2 single-element-tuple wrapping for
b_ptr / b_sf_ptr / alpha_ptr.
* Both kernel constructors no longer receive ``b_tensor_l_sizes=(l,)``.
* Both ``compile_args`` / ``exec_args`` lists pass ``l`` to the
wrapper as an explicit argument (since the kernel no longer caches
it inside ``self``).
Multi-B kernel code (Sm100BlockScaledContiguous*GroupedGemm runners
and the cute.jit kernels) was authored entirely by Zongfei Jing
inside PR NVIDIA#12136; with the IPC path gone there is no remaining
caller, so the cleanup is safe. Verified the only consumer of the
``activation_type`` extension introduced by PR NVIDIA#12884 (CUTEDSL MoE for
nemotron-h) takes the single-B code path
(fused_moe_cute_dsl.py:631), not the multi-B one.
Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…rk scripts
Phase 4 of removing the DWDP IPC-era multi-B kernel infrastructure
that the VA-based DWDP path no longer needs. Phase 1-3 removed
multi-B from the public custom op, the runner classes, and the
cute.jit kernels themselves; this phase strips the matching
multi-B mode from the standalone benchmark / regression scripts that
exercised those kernels.
Affected scripts:
* tests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_
gather_grouped_gemm_act_fusion.py
* tests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_
grouped_gemm_finalize_fusion.py
Both files keep their single-B benchmark / reference-check path
(originally introduced by zhichenj from a CUTLASS example). Removed
exclusively the multi-B add-ons that Zongfei Jing layered on top via
PR NVIDIA#12136:
* ``split_groups_to_b_tensors`` helper (gather-side; the finalize
script had its own copy inside main()).
* ``--num_b_tensors`` and ``--b_tensor_l_sizes`` CLI options plus
the argparse validation that derived multi-B sizes from them.
* ``multi_b_mode = b_tensor_l_sizes is not None`` branching inside
``create_tensors`` (multi-B B/sfb/alpha list construction) and
``run`` (multi-B compile / execute / cold-L2 byte accounting /
reference checking).
* ``b_tensor_l_sizes=`` arguments threaded through
``create_tensors``, the kernel constructor, ``cute.compile``, and
``cute.testing.JitArguments``.
The kernels' ``__call__`` methods still carry ``cutlass.const_expr(
self.num_b_tensors >= 2/3/4)`` branches; with the kernel now hard-
coding ``self.num_b_tensors = 1`` those branches statically fold away
at cute.jit compile time, so the standalone scripts driving the
kernel see only the single-B path.
Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
16 participants
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests
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:
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 sharingDwdpHandleCollector: Per-layer collector that gathers CUDA IPC handles for weight/scale/bias tensorsMoE integration (
configurable_moe.py,fused_moe_cute_dsl.py,interface.py):ConfigurableMoEwith automatic detection of compatible backends (CuteDSL + NVFP4)NvFp4WeightViewdataclass for clean separation of DWDP vs non-DWDP weight access patternsCuteDSL kernel extensions (
cute_dsl_custom_ops.py,blockscaled_contiguous_*_fusion.py):Executor integration (
py_executor.py,py_executor_creator.py,_util.py,llm_args.py):DwdpConfigdataclass toLlmArgsfor YAML-based DWDP configurationDwdpManagerduring executor creation when DWDP is enabledDisaggregated serving scripts (
examples/disaggregated/slurm/benchmark/):start_worker_dwdp.shfor launching DWDP workers viampirunsubmit.pywith 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 GEMMtests/scripts/cute_dsl_kernels/run_blockscaled_contiguous_grouped_gemm_finalize_fusion.py: Standalone kernel correctness tests for DWDP finalize GEMMtests/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.