[None][feat] Refactor DWDP from CUDA IPC to CUDA VMM + MNNVL composite VA - #14453
Merged
Conversation
Commit 1 of the DWDP IPC->VA refactor. Purely additive: no existing
code paths reference the new package yet (wired up in Commit 2).
Package tensorrt_llm/_torch/modules/dwdp/ (migrated from tekit):
* vmm.py CUDA VMM wrappers (RAII handle + VA region, fabric
handle export/import, granularity via lru_cache)
* specs.py WeightSpec / MnnvlHandleSet / PageAlignedLayout
* page_pool.py Double-buffer pool of fabric page handles
* weight_buffer.py Composite VA layout per (layer, weight)
* weight_manager.py Runtime prefetch / wait_and_bind / events
* transport.py MNNVL alloc + MPI-based peer handle exchange
* setup.py Orchestration: Transport -> WeightBuffer ->
WeightManager -> fixup_moe_backends
* __init__.py Package exports
MPI replaces tekit's TCPDWDPStore:
* Transport uses per-pair comm.allgather for handle bytes; the
allgather is itself a sync point (no explicit barrier required
between Phase 1 iterations).
* fixup_moe_backends uses comm.allgather for small scale params
(bias, fc31_alpha, fc2_alpha), replacing tekit's per-sender
serialized put/get/barrier pattern.
* tekit's store.py is intentionally NOT migrated.
collect_moe_params takes layer_indices as an input parameter (SSOT
from DwdpManager._registered_layers; wired up in Commit 2). This
replaces tekit's model-tree walking and makes the MoE layer set a
single source of truth.
Mapping (tensorrt_llm/mapping.py):
* Add dwdp_size / dwdp_rank kwargs (at tail, strictly additive)
* Validate range and store on _dwdp_size / _dwdp_rank
* Override moe_tp_size=1 / moe_ep_size=dwdp_size / cluster=1 when
DWDP enabled
* moe_ep_rank returns dwdp_rank when DWDP enabled
* Expose dwdp_size / dwdp_rank / dwdp_enabled properties
* Include in __eq__ / __hash__ / to_dict
Quality fixes applied during migration:
* vmm.py: logger.debug on cuda-bindings import fallback; thread-safe
granularity cache via functools.lru_cache; logger.info / warning
on tensor_from_ptr fallback paths
* page_pool.py: __del__ uses logger.debug instead of silent pass
* weight_manager.py: drop hardcoded "DeepSeek R1" in docstring
* transport.py: logger.error on create() exception (preserves root
cause before cleanup)
Tests (registered in tests/integration/test_lists/test-db/l0_a10.yml):
* tests/unittest/others/test_mapping_dwdp.py — construction
validation, MoE override, moe_ep_rank branch, to_dict roundtrip,
__eq__/__hash__
* tests/unittest/_torch/modules/test_dwdp_fixup_moe_backends.py —
mock MPI comm; verify allgather semantic equivalence to tekit
TCPDWDPStore version for _allgather_e_score_correction_bias and
_allgather_expert_scales
VA infrastructure originally authored by @dongxuy04.
Co-Authored-By: dongxuy04 <dongxuy@nvidia.com>
Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Commit 2 of the DWDP IPC->VA refactor. This rewires the integration
layer to use the VA infrastructure added in the previous commit. Net
-487 lines across the modified files: 204 insertions, 691 deletions.
pyexecutor/dwdp.py (585 -> 208 lines):
* DwdpManager becomes a pure facade:
- __init__(config, dist, mapping) stores mapping, creates DWDP MPI
sub-communicator (COMM_WORLD.Create_group on rank // dwdp_size).
- __enter__/__exit__ manage the global singleton; duplicate __enter__
now raises instead of silently replacing.
- add_layer(layer_idx) appends to self._registered_layers, which is
the Single Source of Truth for MoE layer indices (setup_dwdp takes
them as input — no more model-tree walking).
- setup(model) delegates to setup_dwdp(model, mapping, device_id,
comm, layer_indices=sorted(_registered_layers)) and caches the
returned DWDPWeightManager.
- prefetch_first_layers / wait_and_bind / record_compute_and_prefetch_next
forward to the DWDPWeightManager.
* Removed: DwdpLayerHandleCollector (~90 lines), DwdpPrefetchBuffer
(~90 lines), all IPC handle plumbing (exchange_all_handles,
initialize_prefetch_buffer, build_weight_view, peer_expert_ranges,
prefetch_layer internals).
pyexecutor/py_executor_creator.py:
* DwdpConfig -> Mapping bridge: when dwdp_config.dwdp_size > 1, rebuild
the Mapping with dwdp_size/dwdp_rank injected (ParallelConfig.to_mapping
doesn't know about DWDP).
* DwdpManager ctor now gets mapping=mapping.
* Replace exchange_all_handles + initialize_prefetch_buffer with
dwdp_manager.setup(model_engine.model) — a single orchestration call.
modules/fused_moe/configurable_moe.py:
* __init__ order fix (F4): DWDP init block now runs BEFORE
_create_comm_strategy_auto so the factory can see self.enable_dwdp.
* _create_comm_strategy_auto returns None when DWDP enabled (fixup
moe backends ep_size=1, no alltoall needed).
* DWDP init simplified: only add_layer; removed dwdp_handle_collector,
dwdp_rank, backend.dwdp_handle_collector.
* wait_and_bind added to forward_impl at Step 3 entry (per-layer,
not per-chunk — correctness fix for multi-chunk forward).
* Removed dwdp_weight_view kwargs injection in _prepare_backend_kwargs.
modules/fused_moe/fused_moe_cute_dsl.py:
* Deleted run_moe_nvfp4_impl_dwdp (~100 lines, multi-B kernel path).
* run_moe_nvfp4 no longer branches on is_dwdp; single-tensor path
handles all cases (VA swaps param.data so kernel sees full weights).
* Removed dwdp_weight_view read in forward hook and
dwdp_handle_collector.register_weights call in load_weights.
modules/fused_moe/interface.py:
* Deleted _init_dwdp_expert_layout and its __init__ call; dropped
get_global_dwdp_manager import. fixup_moe_backends (run from
setup_dwdp) is now the single source of truth for DWDP expert
layout, eliminating the transient EP-sharded state window between
__init__ and setup().
llmapi/llm_args.py:
* DwdpConfig docstring: CUDA IPC -> CUDA VMM + MNNVL fabric handles.
* Fields and status (prototype) unchanged — api_stability passes.
Tests:
* tests/unittest/_torch/executor/test_dwdp_manager.py (new, ~280 lines):
mocks COMM_WORLD / global_mpi_rank / setup_dwdp; verifies
construction validation, global singleton, duplicate __enter__
rejection, add_layer SSOT, setup() forwards sorted layer_indices,
prefetch_first_layers double-depth warmup, record_compute_and_prefetch_next
last-layer no-op, wait_and_bind delegation.
* Registered in tests/integration/test_lists/test-db/l0_a10.yml.
Smoke tested on GB200: 78 pytest passes in 1.02s (Commit 1 + 2 unit
tests + full api_stability suite).
Co-Authored-By: dongxuy04 <dongxuy@nvidia.com>
Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
The VA refactor (a4efa38) inadvertently dropped the IPC version's config-driven expert range computation, replacing it with a hard-coded `num_experts // dwdp_size` in setup.py and transport.py. As a result, DwdpConfig.num_experts_per_worker and num_prefetch_experts became dead fields: read into DwdpManager but never consumed by the runtime. The schema suggested the contract supported by IPC (size = range length, stride = rank-to-rank offset) while the runtime silently assumed uniform integer-division partitioning. Restore the IPC contract by passing both fields through DwdpManager.setup -> setup_dwdp -> DWDPTransport.create. For uniform configs (dwdp=4: 64/64, dwdp=8: 32/32) the formula is mathematically identical to the old integer division; behavior on tested workloads is unchanged. Add `_validate_partition_config` helper enforcing four invariants before the range is computed: 1. positive `num_experts_per_worker` and `num_prefetch_experts` 2. stride <= size (no gaps between consecutive ranks) 3. coverage `(dwdp_size - 1) * stride + size >= num_experts_total` (last rank reaches the final expert) 4. chunk shape parity with `num_experts_per_worker` (DwdpConfig and the fused MoE loader agree on local expert count) These catch the most common schema/runtime mismatches early with specific error messages, and provide the plumbing needed to support non-uniform / redundant partitioning in the future. Redundancy itself is still blocked by the fused MoE backend's `num_experts % ep_size == 0` assertion (interface.py:391); that is a separate, larger change. Verification: - Unit tests: 57 passed + 4 subtests, including 14 new validation cases in tests/unittest/_torch/modules/test_dwdp_setup_validation.py - Accuracy (DSv3-Lite + dwdp=2 + GSM8K): 63.95% (above 61.537 threshold; ~0.7pp above the 63.23% pre-fix data point, within sampling noise) - Perf DEP baseline: 12.88 req/s (vs historical 12.92, -0.3%) - Perf DWDP=4: 14.11 req/s (matches the historical 14.11 to two decimal places, +9.6% over DEP baseline) - Perf DWDP=8: 27.53 req/s on 8 GPUs across 2 trays (1.95x 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>
This is the first of two follow-up commits restoring IPC-era support for non-uniform partitioning (dwdp_size that does not divide num_experts) and redundancy storage on the VA-based DWDP path. Pre-Phase-1 the VA refactor coupled DWDP to fused MoE EP by setting mapping.moe_ep_size = dwdp_size. That triggered the ``num_experts % ep_size == 0`` assert deep in fused_moe/interface.py and locked DWDP to integer-divisible values (2/4/8 only on DSv3's 256 experts). Phase 1 decouples the two: * mapping.py: when dwdp_size > 1, set moe_ep_size = 1 (and moe_ep_rank = 0). The fused MoE backend now sees a single full-table partition; DWDP installs the per-rank layout itself. * pyexecutor/dwdp.py: DwdpManager exposes start_expert_id / end_expert_id derived from the rank-to-rank stride (num_prefetch_experts) and storage size (num_experts_per_worker) in DwdpConfig. These are read by the override hook below. * fused_moe/interface.py: restore the IPC-era _init_dwdp_expert_layout method, called at the end of MoE.__init__ after _init_load_balancer. When DWDP is enabled (get_global_dwdp_manager() is not None) it overrides expert_size_per_partition / slot_start / slot_end / initial_local_expert_ids before create_weights() runs, so the backend allocates num_experts_per_worker slots per rank. For uniform integer-divisible configs (dwdp=4: 64/64, dwdp=8: 32/32) this is mathematically equivalent to the legacy ep_size=dwdp_size layout — verified by the existing 96 unit tests staying green. Phase 2 (next commit) wires the runtime side so non-uniform / redundant configs actually work end-to-end. Tests: * test_mapping_dwdp.py: assert moe_ep_size=1 / moe_ep_rank=0 under DWDP and verify Mapping construction accepts dwdp_size=3/5 unconditionally (Phase 1 contract: rejection moves to DwdpManager / setup_dwdp). * test_dwdp_manager.py: new tests for start_expert_id / end_expert_id under uniform and redundancy configs. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…omposite VA Second of two follow-up commits — Phase 2 + 3 of the restoration plan in the prior commit. Phase 1 made model-load flexible (per-rank storage = num_experts_per_worker regardless of ep_size); this commit makes the runtime side handle non-uniform / overlapping peer ranges end-to-end. Core abstraction (specs.py): * compute_peer_ranges(dwdp_size, size, stride, num_experts) -> PeerRanges: deterministic, no allgather. Returns ``[(start, end_capped)]`` per rank where ``end_capped = min(start + size, num_experts)``. * lookup_owner(expert_id, peer_ranges) -> int: first-match owner lookup, picking the lowest-rank owner under redundancy overlap. Equivalent to ``expert_id // (num_experts // dwdp_size)`` when the partition is uniform. Wire-through (transport.py / weight_manager.py / setup.py): * Transport publishes peer_ranges via ``get_peer_ranges()``; setup_dwdp threads them to fill_edge_bytes (replacing ``prev_expert // experts_per_rank``), DWDPWeightManager.prefetch_layer (replacing ``cursor // _experts_per_rank``), and fixup_moe_backends. The legacy ``num_experts % dwdp_size == 0`` assert in ``DWDPTransport.create()`` is removed. * setup.py:_scatter_shards_to_full is the new common reconstruction helper for fixup_moe_backends's MPI-allgathered scale params and the gate's e_score_correction_bias. Each peer's shard is sized num_experts_per_worker (uniform); ``shard[:end - start]`` is placed at ``full[start:end]``. For redundancy, overlapping writes hit the same expert multiple times but values agree (every rank that owns expert ``e`` loaded ``e`` from the same checkpoint). Validation (Phase 3, setup.py:_validate_partition_config): * New check rejecting tail-padding configurations: ``(dwdp_size-1)*stride + size <= num_experts``. Combined with the existing coverage check (``>=``) this enforces strict equality. Tail padding is rejected because GB200 fabric handles cannot be partially mapped into a num_experts-sized composite VA (cuMemMap requires ``size == handle_phys_size``; partial mapping returns CUDA_ERROR_NOT_SUPPORTED). * Error message lists Mode B recipes for common configs: ``dwdp=3 + 256 -> size=86 stride=85``; ``dwdp=5 + 256 -> 52 / 51``; ``dwdp=7 + 256 -> 40 / 36``. Tests: * tests/unittest/_torch/modules/test_dwdp_peer_ranges.py (new): compute_peer_ranges + lookup_owner under uniform (dwdp=4/8), Mode B overlap (dwdp=3/7), and out-of-range / negative inputs; defensive cap tests for inputs that fail validation upstream. * test_dwdp_setup_validation.py: Mode B overlap pass cases (dwdp=3/5/7) + tail-padding rejection cases (dwdp=3/5). * test_dwdp_fixup_moe_backends.py: ``_scatter_shards_to_full`` direct tests (uniform, Mode B overlap, higher-dim trailing shape) + an end-to-end ``_allgather_expert_scales`` test for Mode B dwdp=3 reconstructing fc31_alpha to a full 256-element vector. * l0_a10.yml: registers the new test file for pre-merge CI. For uniform configs (dwdp=4 / dwdp=8), behaviour is unchanged — ``lookup_owner`` is mathematically equivalent to ``expert_id // (num_experts // dwdp_size)`` and ``compute_peer_ranges`` produces the same start/end tuples the formula would. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
The Phase 2 commit (``63908b7c``) propagated peer_ranges built from ``compute_peer_ranges(num_experts_total=spec.num_experts, ...)``, where ``spec.num_experts`` was derived in ``build_weight_specs`` as ``experts_per_rank * dwdp_size``. Under uniform divisible configurations (dwdp=2 + 36/36 → 72, dwdp=4 + 64/64 → 256) this happened to equal the model's gate-side ``num_experts`` and validation passed. Under Mode B redundancy (e.g. dwdp=3 + size=26/stride=23 on DSv3-Lite's 72 experts), ``experts_per_rank * dwdp_size = 78 ≠ 72``, so the strict-equality validation falsely reported coverage 72 < storage 78 and rejected the configuration. The fix queries the actual gate-side expert count from the first MoE layer's ``experts_module.num_experts`` via the new ``_get_model_num_experts(model, layer_indices)`` helper, and threads it through ``build_weight_specs`` as an explicit ``num_experts_total`` parameter. Downstream, this becomes ``spec.num_experts`` and is used consistently by: * ``_validate_partition_config`` — strict equality is now checked against the real expert count. * ``compute_peer_ranges`` — peer ranges cap at the real expert count. * ``PageAlignedLayout`` / ``WeightBuffer._compute_remote_slices`` — composite VA's full tensor shape is ``(num_experts, ...)`` and remote slices end at ``num_experts``. For uniform divisible cases the new path is mathematically equivalent to the prior code (since storage extent == gate's num_experts when ``size == stride == num_experts // dwdp_size``). Verified by re-running the dwdp=4 accuracy regression (passes with GSM8K = 64.519) and by the Mode B unit tests in ``test_dwdp_peer_ranges.py`` / ``test_dwdp_setup_validation.py``, which already use gate-side ``num_experts_total`` values directly. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Adds ``test_dwdp_accuracy_mode_b_overlap`` to TestDwdpDeepSeekV3Lite, covering the Phase 2 redundancy (overlap) path end-to-end on a real DSv3-Lite + GSM8K workload. Config: ``dwdp_size=2``, ``num_experts_per_worker=40``, ``num_prefetch_experts=32``. The strict-equality check evaluates to ``1*32 + 40 == 72`` (== model's gate-side num_experts), with 8-expert overlap between rank 0's ``[0, 40)`` and rank 1's ``[32, 72)``. This exercises: * ``_init_dwdp_expert_layout`` overriding ``expert_size_per_partition`` to 40 (not the uniform ``72 // 2 == 36``). * ``lookup_owner`` first-match policy in ``fill_edge_bytes`` / ``weight_manager.prefetch_layer`` for the 8 shared experts. * ``_scatter_shards_to_full`` writing each overlap expert from two peers; values must agree because every rank that owns expert ``e`` loaded ``e`` from the same checkpoint. Resource footprint matches the existing ``test_dwdp_accuracy`` (2 CTX TP=1 + 1 GEN TP=2 = 4 GPUs on GB200). ``dwdp_size=2`` is kept intentionally — a ``dwdp_size=3`` variant triggers an unrelated UCX bus error in ``libtensorrt_llm_ucx_wrapper.so``'s ``WorkerProgressThread::progressUntilSync`` when 3 CTX workers simultaneously exchange KV cache. The bug is in the KV transceiver layer (not DWDP); the existing 2-CTX baseline runs cleanly and exercises the overlap path equally well. Verified: PASSED, GSM8K = 63.268 (threshold 61.537, reference 64.740); runtime 234s. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Port the per-rank hostfile + gpu_map launcher mechanism from PR NVIDIA#13888 (Jianbo Hu, "support non-divisible EP in MoE alltoall and slurm benchmark") into the in-tree DWDP examples launcher trio (submit_dwdp.py, disaggr_torch_dwdp.slurm, start_worker_dwdp.sh). Motivation: SLURM block distribution requires (n_ctx + n_gen * gen_tp) to be a multiple of gpus_per_node, otherwise GEN tensor-parallel ranks get split across nodes/trays and incur a large allreduce penalty. Previously worked around by adding empty CTX slots (over-provisioning) or picking dwdp_group multiples that happen to divide cleanly — both fragile and incompatible with Mode B non-uniform expert ranges where num_ctx is e.g. 5 or 6. Dual-path design: * Divisible case (num_ctx_gpus % gpus_per_node == 0): Use the legacy --nodelist + -N + --ntasks-per-node srun command. Block distribution gives the natural rank-to-node mapping; trtllm-serve picks its CUDA device from SLURM_LOCALID. No hostfile / gpu_map files are emitted. Perf-optimal path. * Non-divisible case (Mode B dwdp3 dg=2, dwdp5 dg=1, etc.): Emit hostfile_mpi_worker_base.txt and gpu_map_mpi_worker_base.txt under log_dir, iterating CTX servers then GEN servers so global rank ordering matches split_world_comm's ctx_cfgs + gen_cfgs. disaggr_torch_dwdp.slurm rewrites <nodeN_placeholder> tokens to real hostnames at runtime. srun uses SLURM_HOSTFILE + --distribution=arbitrary to pin every rank's host placement. DWDP-specific deviation from PR NVIDIA#13888: start_worker_dwdp.sh does NOT export CUDA_VISIBLE_DEVICES. DWDP relies on intra-node peer GPU discovery (VA composite cuMemMap of peer MNNVL fabric handles, UCX cuda_ipc / cuda_copy intra-node KV transports, PyTorch peer device enumeration); restricting CUDA visibility to a single GPU breaks these paths and causes a 15% per-CTX-GPU regression (TPOT unchanged, TTFT std blows up 3x). With our allocate_gpus's sequential cursor, gpu_id == SLURM_LOCALID for every rank, so trtllm-serve's internal LOCALID-based device selection already lands each rank on the correct GPU. The gpu_map file is kept for diagnostics and audit logging only. Empirical per-CTX-GPU req/s (DSv3-FP4-V2.1, ISL=8192/OSL=1, conc=256): dwdp4 dg=1 (8 GPU divisible) 3.46 dwdp3 dg=4 (16 GPU divisible) 3.41 dwdp5 dg=4 (24 GPU divisible) ~3.3 (extrapolated from R1b) dwdp3 dg=2 (10 GPU non-divisible) 3.51 dwdp5 dg=1 (9 GPU non-divisible) 3.42 TPOT 14-18 ms across all configs (kernel-neutral). Enables the non-uniform / Mode B perf configurations used in the follow-up dwdp_size=3/5 experiments in this refactor. Co-authored-by: Jianbo Hu <jacohu@nvidia.com> Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…er.wait_and_bind Explain the O(N_moe)→O(1) event simplification that the VA refactor introduces in wait_and_bind. The IPC-era predecessor carried per-layer compute_events to signal 'kernel(L) finished, slot is reusable'; the VA implementation replaces them with per-slot consume_events recorded inside this very method, relying on CUDA stream in-order semantics to provide the same WAR ordering at O(1) bookkeeping cost. Documents the assumption (event recorded on compute_stream AFTER bind and BEFORE next layer's kernel launch) so future maintainers don't silently break it when extending to deeper pipelines (num_buffers > 2). 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>
Re-implements the DWDP IPC-era contention optimization (PR NVIDIA#12974) on top of the VA-based ``DWDPWeightManager.prefetch_layer``. The ``DwdpConfig.contention_opt`` field that PR NVIDIA#12974 added is already carried through the schema; this commit makes it functional again on the VA path. When ``contention_opt=True``, ``prefetch_layer`` issues a single ``cudaMemcpyBatchAsync`` per layer with sub-slices interleaved across peers in 2 MiB chunks (``MEMCPY_BATCH_SLICE_BYTES``), so that any moment in flight touches every peer's NVLink link rather than saturating one. When ``contention_opt=False`` (default) the existing per-slice ``torch.Tensor.copy_`` path is unchanged. Adaptations from the IPC original onto the VA composite-VA structure: * Plan construction walks the same ``get_remote_slices`` / per-peer chunking that ``_prefetch_layer_per_slice`` uses (so Mode B's non-uniform / overlapping ranges resolved by ``lookup_owner`` interleave correctly, not just uniform partitions). * Each peer's per-layer contribution is split into 2 MiB sub-slices, then a round-robin emitter across peers concatenates them into the final ``(dst_ptrs, src_ptrs, sizes)`` triple consumed by ``cudaMemcpyBatchAsync``. * Plans are cached per ``layer_idx`` in ``self._batched_copy_plans`` since peer ``data_ptr()`` and destination slice ``data_ptr()`` are stable after ``setup_dwdp`` returns; this avoids host-side rebuild on every prefetch and matches the cache discipline used by PR NVIDIA#12974. Plumbing: * ``DwdpConfig.contention_opt`` -> ``DwdpManager.setup()`` -> ``setup_dwdp(contention_opt=...)`` -> ``DWDPWeightManager(contention_opt=...)``. * ``DWDPWeightManager.__slots__`` extended with ``_use_batched_prefetch`` and ``_batched_copy_plans``. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…r Mode B Adds ``TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_contention_opt``: a uniform-partition DSv3-Lite + GSM8K end-to-end run with ``contention_opt=True``, mirroring the existing ``test_dwdp_accuracy`` configuration but exercising the batched ``cudaMemcpyBatchAsync`` prefetch path from the previous commit (``Port DWDP contention optimization to VA prefetch path``). The GSM8K threshold is unchanged — the batched path issues the same set of D2D copies as the per-slice default, only with a different launch schedule, so accuracy must match. Also registers ``test_dwdp_accuracy_mode_b_overlap`` (commit ``Add Mode B overlap accuracy integration test``) in ``l0_gb200_multi_gpus.yml`` and ``qa/llm_function_core.txt``; it was added with code coverage but never wired into the CI / QA test lists. After this commit the GB200 multi-GPU CI lane and the QA core function list both cover all three DWDP accuracy scenarios: uniform partition (default contention_opt), uniform partition with contention_opt=True, and Mode B 8-overlap redundancy. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
…regression Adds ``test_dwdp_accuracy_contention_opt`` and ``test_dwdp_accuracy_mode_b_overlap`` to the SKIP list alongside the existing ``test_dwdp_accuracy`` waiver. All three are blocked by the same upstream disagg-serving infrastructure regression tracked in nvbug 6094102 (a 2026-04-28 batch waive that paused nine disagg test configurations across ``test_disaggregated.py``, ``test_disaggregated_serving.py``, and ``test_workers.py`` with matching nvbug numbers 6114139-6114612). Locally reproduced (twice, identical signature on different allocations): pytest hangs at the ``Waiting for servers...`` HTTP polling loop. GEN-side TP=2 worker rank 1 leaves GPU 3 stuck at 2.7 GB while all other GPUs load to ~12 GB; no kernel ever launches and the deterministic memory fingerprint matches between runs. The same pytest harness has been used in the DSv3-Lite + GSM8K accuracy recipe historically (recipe documents ``test_dwdp_accuracy`` GSM8K 64.519, ``test_dwdp_accuracy_mode_b_overlap`` 63.268) — the regression sits in the disagg launch path itself, not in DWDP code. Re-enable when nvbug 6094102 (and its companions) close. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
tianyuz-nv
force-pushed
the
feat/dwdp-va-refactor
branch
from
May 25, 2026 19:18
2d43359 to
cde75fa
Compare
- isort / yapf / ruff-format / autoflake / codespell auto-fixes applied to the DWDP refactor file set (19 .py files + 1 .slurm helper). - Drop unused 'from cuda.bindings import driver as cuda' import in page_pool.py (F401). - Drop unused local 'storage' variable in vmm._tensor_from_ptr_internal (F841). - Fix docstring formatting in test_mapping_dwdp.py::test_hash_distinguishes_dwdp (D205: insert blank line between summary and description). - Suppress ruff-legacy E741 on 13 GEMM-batch dim 'l = ...' assignments in cute_dsl_custom_ops.py (M/N/K/L is the canonical CuTe DSL GEMM dimensions, renaming would be more confusing than the lint warning). - Replace 'composite VAs' with 'composite VA layout' in 4 docstrings to satisfy codespell which mis-flags VAs (Virtual Addresses) as a misspell of 'was'. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
tianyuz-nv
marked this pull request as ready for review
May 26, 2026 08:42
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51044 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51044 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51194 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51194 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51222 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51222 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51241 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51241 [ run ] completed with state |
QiJune
approved these changes
Jun 1, 2026
dc3671
reviewed
Jun 1, 2026
dc3671
approved these changes
Jun 1, 2026
hyukn
approved these changes
Jun 1, 2026
Per @dc3671's review on PR NVIDIA#14453: the happy-path ("must NOT raise") tests for _validate_partition_config are low value -- a valid config passing validation is the trivial case, and the validator already runs in the live DWDP setup path (exercised by the disagg accuracy tests). Remove the 5 happy-path tests (test_uniform_dwdp4/8, test_coverage_with_overlap_passes, test_dwdp{3,7}_mode_b_overlap_passes); keep all 10 failure-mode/reject tests. setup_validation 15 -> 10. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Per @dc3671's review on PR NVIDIA#14453: _validate_partition_config is a defensive validation guard (not a feature) with simple parameter checks, and it already runs in the live DWDP setup path (the non-divisible configs are exercised by the disagg accuracy tests). A dedicated unit-test file adds little coverage, so remove test_dwdp_setup_validation.py entirely and drop its l0_a10 entry. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Per @dc3671's review on PR NVIDIA#14453 ("simplify tests, only test key features"), apply the same rationale as the _validate_partition_config removal to the remaining DWDP unit tests: - drop validation-guard / reject tests (defensive checks on simple params, already exercised by the live setup path); - drop redundant size-variant + skip-condition duplicates (kept one representative each); - remove the now-unused _make_mock_comm_world helper. Kept all key-feature tests (peer-range computation / owner lookup, weight scatter+allgather reconstruction, mapping->MoE wiring, manager prefetch orchestration) and contract plumbing (serialization, eq/hash, lifecycle idempotency). 66 -> 43 tests across 4 files. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Collaborator
Author
|
/bot skip --comment "Test-only change since the last green pipeline (L0_MergeRequest_PR #40664 / PR_Github #51241 on commit 22028b5). These commits only delete and trim DWDP unit tests per @dc3671's review: removed tests/unittest/_torch/modules/dwdp/test_dwdp_setup_validation.py and guard/redundant cases in the other dwdp/ test files. The diff is 5 test files plus 1 test-list line, with no source/functional code touched, and the kept tests are a strict subset of what #40664 already ran, so the existing pipeline result still covers all functional code." |
Collaborator
Author
|
/bot help |
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. |
Collaborator
Author
|
/bot skip --comment "Test-only change: this PR only removes and trims DWDP unit tests per reviewer feedback. No source or functional code changed since the last green pre-merge pipeline 40664, so the existing pipeline result still covers all functional code." |
Collaborator
Author
|
/bot skip --comment "Please see above." |
Collaborator
|
PR_Github #51343 [ skip ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51343 [ skip ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
12 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Migrate the DWDP (Distributed Weight Data Parallelism) implementation from CUDA IPC to a CUDA VMM + MNNVL fabric-handle based design. DWDP accelerates the context (prefill) phase of disaggregated MoE serving by combining data parallelism with NVLink-based expert weight sharing — each worker holds a subset of experts locally and asynchronously prefetches the rest from peers, eliminating both alltoall latency and synchronization costs on DeepSeek-V3-class models.
The original IPC-based DWDP (PR #12136) used
cudaIpcGetMemHandle/cudaIpcOpenMemHandleand a tensor-list interface to the MoE kernel. This refactor moves tocuMemCreate(HANDLE_TYPE_FABRIC)+ a composite virtual address layout that interleaves zero-copy local expert data with double-buffered remote regions. The kernel now sees a single[num_experts, ...]tensor view per layer instead of a per-rank list, removing the need for DWDP-specific multi-buffer kernel paths.Key properties preserved:
DwdpConfigschema,submit_dwdp.py,disaggr_torch_dwdp.slurm, and YAML config layout are all identical to the IPC version. No user-facing breakage.New capabilities unlocked:
dwdp_size=8across 2 trays.dwdp_size+ redundancy storage —dwdp_sizevalues that don't dividenum_experts(e.g. 3, 5, 7 on DSv3's 256 experts) and configurations where adjacent peer ranges intentionally overlap. The IPC implementation supported this via_init_dwdp_expert_layout; the early VA refactor inadvertently dropped it. This family of layouts is referred to as "Mode B" in test names (test_*_mode_b_*) and below — defined formally aspeer[r] = [r × num_prefetch_experts, r × num_prefetch_experts + num_experts_per_worker)clamped tonum_experts. Whennum_experts_per_worker > num_prefetch_expertsadjacent ranges overlap (redundancy); the uniform / non-overlap case (num_experts_per_worker == num_prefetch_experts == num_experts / dwdp_size) is "Mode A".VA infrastructure was originally written by @dongxuy04 on an internal branch; this PR migrates it into TRT-LLM main with the integration glue (DwdpManager facade, MoE backend hookup,
fixup_moe_backendsMPI allgather replacing the internal TCPDWDPStore).Companion housekeeping captured in commits 11–15: with the IPC path gone the DWDP-IPC-era multi-B kernel infrastructure (introduced by PR #12136 alongside the IPC implementation) has no remaining caller and is removed across four phased commits — public custom op, runner / helper plumbing, cute.jit kernels themselves, and the standalone benchmark scripts. Commit 15 then re-applies the IPC-side contention optimization from PR #12974 onto the VA prefetch path so
DwdpConfig.contention_opt=Truebecomes functional again (singlecudaMemcpyBatchAsyncper layer with 2 MiB sub-slices round-robined across peers to limit single-link NVLink contention).Commits
92fd50cb4193612897e96e3fnum_experts_per_workersemantics in DWDP setupf03a1611DwdpConfig.num_groupsagainst the MPI launch7330e2b04d994fe416ed91acnum_expertsfor DWDP setup validation49ffe7ffad745075--distribution=arbitrary85681b49DWDPWeightManager.wait_and_bind67715e02f145063ac349266de3bc13fc8dbe940e66e103d9contention_optpath + register Mode B in CI / QA listscde75fa8e7aecde3cuda.bindings.driverimport and unusedstoragelocal; fix docstring D205; suppress 13 E741 on canonical CuTe GEMM-batch dimlassignments; rewritecomposite VAs→composite VA layoutto placate codespell0a766453@skip_post_blackwell_ultra). Root cause for the earlier failure was build-dir absolute-path symlinks (cpp/build/.../trtllmGenKernels/fmha/trtllm/dev→/lustre/fs1/...) that fail to resolve in a dev SLURM container whose mounts only expose/lustre/fsw. CI wheel install bypasses this entirely (build_wheel.pycallscopy_resolving_symlinkand ships real files undertensorrt_llm/include/...), so the tests can be active in CI.Verification
New unit tests cover DwdpManager lifecycle, Mapping integration with the new DWDP fields, MPI allgather semantics in
fixup_moe_backends, peer-range computation under uniform and Mode B layouts, and partition-config validation. All registered inl0_a10.yml.Three integration tests in
TestDwdpDeepSeekV3Liteexercise the end-to-end DWDP path on DSv3-Lite + GSM8K, all registered inl0_gb200_multi_gpus.ymlandqa/llm_function_core.txt:test_dwdp_accuracy(uniform, defaultcontention_opt=False);test_dwdp_accuracy_contention_opt(uniform, exercises the batchedcudaMemcpyBatchAsyncprefetch path withcontention_opt=True);test_dwdp_accuracy_mode_b_overlap(Mode B 8-overlap, coverscompute_peer_ranges,lookup_owner,_scatter_shards_to_full, thefixup_moe_backendsrewrite, and composite VA stitching).End-to-end performance (DeepSeek-R1-0528-FP4-V2.1, GEN TP=4, conc=256)
5 configurations measured on
scripts/gb200-disagg-scripts/perf_minimal_dwdp{3,4,5,8}.yaml(+ contention_opt variant), ISL=8192 / OSL=1, 4096 prompts via--exec-mode configloadgen. Pass criterion is per-CTX-GPU within ±3% of the IPC-era baseline; TPOT in the 14.7–18.2 ms band serves as a kernel-sanity check.contention_opt=TrueAll 5/5 within ±3% baseline and TPOT band, matching the IPC-era envelope. Coverage axes:
dwdp_size ∈ {3, 4, 5, 8}, both launcher paths (block + arbitrary), divisible + non-divisible layouts, single-tray + cross-tray,contention_optOFF + ON.TP=2 disagg accuracy (3 pytest tests, all PASSED)
The three
TestDwdpDeepSeekV3Litepytest tests run on a 4-GPU GB200 single-tray layout (2 CTX TP=1 + 1 GEN TP=2,dwdp_size=2) and are now active in CI (gated to GB200 via@skip_post_blackwell_ultra). End-to-end GSM8K accuracy on DSv3-Lite NVFP4:test_dwdp_accuracy(uniform, default)test_dwdp_accuracy_contention_opt(uniform,contention_opt=True)test_dwdp_accuracy_mode_b_overlap(Mode B 8-overlap)All three within the test-internal pass threshold (≥ 61.537) and within ±5 pp of the historical baseline.
Cross-node coverage
Cross-node DWDP (where partner-expert prefetch traverses MNNVL fabric instead of intra-node NVLink) is not added as a pytest in this PR for the same reason PR #13888 doesn't — the
disaggregated_mpi_workertest pattern uses nestedmpirunand is incompatible with the multi-node CI infrastructure's outer-srunMPI world. The cross-node code path is logically equivalent to the intra-node path (model-side code doesn't branch on locality; onlycuMemMapof partner expert weights uses a fabric handle instead of an IPC handle), and is transitively covered by the unit tests and end-to-end perf runs above.Compatibility
DwdpConfig(tensorrt_llm/llmapi/llm_args.py): four base fields (dwdp_size,num_groups,num_experts_per_worker,num_prefetch_experts) plus the existingcontention_optfield added by PR [None][feat] Move DWDP contention optimization into DwdpConfig #12974, all preserved. Only the class docstring is updated to mention CUDA VMM + MNNVL.examples/disaggregated/slurm/benchmark/submit_dwdp.py/disaggr_torch_dwdp.slurm/start_worker_dwdp.sh: launcher logic upgraded for non-divisibledwdp_sizeand cross-tray (commit 10); CLI flags unchanged.ctx_config.dwdp_configblock): identical structure.Acknowledgments
VA infrastructure (
_torch/modules/dwdp/) was originally authored by @dongxuy04 on an internal branch. This PR is the migration into TRT-LLM main with the integration glue.The hostfile +
srun --distribution=arbitrarylauncher mechanism in commit 10 is ported from @JacobHu-NV's upstream PR #13888 ("support non-divisible EP in MoE alltoall and slurm benchmark"), adapted for the DWDP examples launcher trio.PR Checklist
DwdpConfigdocstring updated; AGENTS.md not touched.GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.