[None][feat] Expert DWDP: a balanced expert partition when dwdp_size does not divide the experts

Expert DWDP gives every rank one contiguous range of routed experts, set by two DwdpConfig scalars,
num_experts_per_worker (size) and num_prefetch_experts (stride). Setup required
(dwdp_size - 1) * stride + size == num_experts with stride <= size, so when dwdp_size does not
divide the expert count only overlapping configs passed. Kimi-K3 DWDP12 runs 82/74 and stores 88
duplicate experts per layer. On a K3 DWDP12 context server (attached dense DWDP, max_num_tokens
32768, KV cache fraction 0.70), the balanced partition below lowers each rank's profiling peak by
10.64-12.42 GiB and raises its KV cache by 7.45-8.70 GiB against 82/74. The VMM does not need equal
ranges: a handle maps only whole, whatever its size. The software assumed equal ranges in these
places:
- the range formula, computed four times;
- each MoE backend, which allocated num_experts_per_worker rows;
- every peer's handle size and view rows, taken from this rank's;
- the fixup, which allgathered a small parameter only when its first dimension was
  num_experts_per_worker, so with unequal counts some ranks would skip a collective.

ExpertPartition (dwdp/specs.py) is one table of every rank's [start, end). Every rank computes it
from DwdpConfig, dwdp_size and the model's expert count N. It is the only source of any rank's
range for DwdpManager, MoE._init_dwdp_expert_layout, setup_dwdp and the transport.
- Both fields 0: balanced. With divmod(N, dwdp_size) == (base, rem), the ranks below rem store
  base + 1 experts and the others base, as _compute_ep_partition splits EP ranks (K3 DWDP12:
  8 x 75 + 4 x 74). Needs N >= dwdp_size.
- Both positive: strided, where rank r stores [r * stride, min(r * stride + size, N)). Accepted:
  - uniform: stride == size and dwdp_size * size == N;
  - redundant: stride < size and (dwdp_size - 1) * stride + size == N; the lowest rank storing an
    expert serves it;
  - capped: stride == size and (dwdp_size - 1) * size < N < dwdp_size * size; the last rank
    allocates only the experts left, with no padding.
- Anything else raises ValueError: one field set or a negative value (DwdpConfig also rejects
  these when the config is parsed), a stride larger than the size, ranges that end before N,
  overlap past N, or an empty last rank.
No config field is added. Both fields become NonNegativeInt, which the golden manifest records as
int, so the manifest and the API references are unchanged. The two field descriptions change, and
the serve JSON schemas (tensorrt_llm/schemas/trtllm-serve-config.schema.json and
trtllm-serve-disagg-config.schema.json, regenerated with scripts/generate_trtllm_serve_schemas.py)
gain minimum 0 for both.

What reads the table:
- DwdpManager.expert_partition(N), cached per N with no collective, and local_expert_range(N).
  MoE._init_dwdp_expert_layout sizes each backend with exactly its rank's rows, so an invalid
  config raises at model construction, before create_weights. initial_global_assignments is
  range(N) for every partition: only a load balancer reads it, and DWDP excludes one.
- DWDPTransport.create takes the ExpertPartition. The transport allocates
  handle_geometry(expert_bytes, start, end, granularity) for the rank's range. It maps every peer's
  handle whole, with that peer's size and offset, and gives each peer view that peer's rows.
- From the handshake to the vote, setup_dwdp collects first, then decides the same way on every
  rank. A failure on one rank in that span thus raises on every rank after the same collective,
  instead of leaving its peers blocked in the next one.
  - _agree_on_partition makes one allgather: N, every rank's range, the kind, the granularity, and
    for each (layer, weight) expert chunk its rows, per-expert shape and dtype. Each rank is
    compared with the most common report, so the error names the ranks that differ, rank 0
    included (it lists at most 8 problems); a rank that sends no handshake is named too.
  - Each Phase-1 handle allgather carries (status, payload), and every rank checks the statuses
    before any import. The handles need no geometry check: every rank computes each handle's size
    and data offset with handle_geometry from what the handshake agreed on. Phase 3's Barrier
    becomes allgather(status). On POSIX_FD, Phase 2 opens each peer's pidfd inside the import, so
    that failure is reported in Phase 3 as well; a failed map or access grant releases its VA
    before the error is reported.
  - Steps 5-7 (the weight buffer and ring, the edge bytes, the weight manager) only record a local
    error. The fixup's one collective, _exchange_expert_scalars, is the group's vote. If any rank
    failed, that rank raises its own error and the others raise a RuntimeError naming it.
    Otherwise it assembles the per-expert alphas (and a per-expert gate bias) along the partition.
    It also requires overlapping copies, and per-layer parameters such as fc31_input_scale, to be
    bitwise equal on every rank. fixup_moe_backends makes no collective, so K3's 184 fixup
    allgathers become one.
  - If a verdict collective itself fails (a peer died, the store timed out), a rank that had
    failed raises its own error, with the collective's error as its context.
  The backend fixup after the vote is local: it re-materializes the activation constants
  (install_activation_params) for the full expert count, installs the exchanged tensors, and a
  failure there raises on that rank only. A rank that dies, or fails before the handshake, leaves
  its peers waiting until the rendezvous times out.
- Removed: _validate_partition_config, _allgather_e_score_correction_bias,
  _allgather_expert_scales, _get_first_spec, WeightBuffer.compute_peer_data_offset,
  compute_peer_ranges, DWDPTransport.get_peer_ranges, and DwdpManager's start_expert_id and
  end_expert_id. Stale comments are fixed: ring pages are local NONE handles, the whole-handle rule
  applies to every VMM handle, the mapping limit, the composite-VA padding, and the edge bytes.

Compatibility: explicit size/stride configs keep the same ranges, rows, handle sizes and offsets,
peer views, ring pages, copy plans and fixup results, byte for byte. ExpertPartition.strided
accepts exactly the old validator's configs plus capped tables, so the uniform and redundant
configs of test_dwdp_aggregated.py, test_dwdp_disaggregated_serving.py and the Kimi-K3
perf-sanity configs (56/56) stay valid. For these configs only the following change: the
collectives (19 -> 15 per rank on a 3-layer model), the log wording, the point where a config
error is raised, and initial_global_assignments before the fixup (range(N), where a redundant
config had a round-robin list); ranks that disagree now raise together. 0/0, which was rejected
before (Kimi-K3 while the model is built, with an IndexError; DeepSeek-V3 in setup, after loading),
now means balanced.

Tests:
- test_dwdp_peer_ranges.py: balanced equals _compute_ep_partition. On a 60,836-config grid,
  strided accepts exactly the configs whose ranges end at N (the old validator's set) plus capped
  tables, with the strided ranges. Also covers from_config accept and reject cases, invalid
  tables, and handle_geometry equal to PageAlignedLayout.
- test_dwdp_manager.py: DwdpConfig validation; 0/0 resolved and cached per expert count; backend
  rows from _init_dwdp_expert_layout (74 on a balanced rank, 71 on the capped last rank, 82 on a
  redundant one); a config that does not fit raises at layout time. The existing range checks
  read local_expert_range and expert_partition instead of the removed start_expert_id and
  end_expert_id.
- test_prefetch_copy_plan.py: with peers of 75/74, 45/44 and 38/37 experts, every descriptor reads
  its owner's rows, the remote experts are covered once, the batched plan moves the same bytes, and
  the gather plan has one row per peer group with all of the group's bytes.
- tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py: ported to the exchange. Every rank
  makes one allgather, and every inconsistency raises on every rank. A ConfigurableMoE-shaped
  wrapper around a backend set up by the real NVFP4CuteDslFusedMoEMethod.setup_quant_scales, on a
  4 + 3 split: after the fixup both see every expert, quant_scales holds the full alphas, and a
  full-sized gate bias stays as it is.
- test_dwdp_setup_lifecycle.py: patches the handshake, the granularity query and the vote.
- test_trtllm_serve_config_schemas.py: the serve schema accepts 0/0 and rejects a negative
  num_experts_per_worker or num_prefetch_experts; the disaggregated schema rejects a negative one
  under context_servers.
- New test_dwdp_setup_protocol.py (cpu_only): one thread per rank runs the real setup_dwdp over
  StoreDwdpGroup on a HashStore, with a fake VMM that maps handles only whole. The weight buffer,
  edge-byte fill and weight manager are stand-ins; those of the buffer and the manager record
  their arguments.
  - Every peer is mapped with its own size and rows (balanced, capped, uniform, redundant),
    steps 5-7 get the partition's local range and peer ranges, and the exchange gives every rank
    the full fc31_alpha (from 75/74 shards on the balanced DWDP12).
  - A failure on one rank raises on every rank after the same collective, and nothing leaks.
    Covered failures: the handshake (rows, partition, dtype), an export, the access grant of an
    imported handle, and the ring.
  - When rank 0 alone computed another partition, the handshake error names rank 0.

Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com>

Author: zongfeijing

tensorrt_llm/_torch/modules/dwdp/__init__.py

@@ -45,13 +45,14 @@
 from .setup import setup_dwdp, teardown_dwdp
 from .specs import (
     EdgeInfo,
+    ExpertPartition,
     LayerWeightSpecs,
     MnnvlHandleSet,
     PageAlignedLayout,
     PeerRange,
     PeerRanges,
     WeightSpec,
-    compute_peer_ranges,
+    handle_geometry,
     lookup_owner,
 )
 from .transport import DWDPTransport
@@ -87,7 +88,8 @@
     "PageAlignedLayout",
     "PeerRange",
     "PeerRanges",
-    "compute_peer_ranges",
+    "ExpertPartition",
+    "handle_geometry",
     "lookup_owner",
     # VMM utilities
     "align_up",

tensorrt_llm/_torch/modules/dwdp/page_pool.py

@@ -15,22 +15,22 @@
 
 """Page pool for DWDP remote double buffer.
 
-The page pool manages per-page fabric handles for the remote regions of the
+The page pool manages per-page handles for the remote regions of the
 composite buffer. Two pools (slot 0, slot 1) support double buffering.
 
-Each pool is a list of page-sized fabric handles that can be mapped at different
-VA positions for different layers. Pages are reused across layers in the same slot:
+Each pool is a list of page-sized local (``CU_MEM_HANDLE_TYPE_NONE``) handles
+that can be mapped at different VA positions for different layers. Pages are
+reused across layers in the same slot:
     - Layer 3 (buf_idx=0): pool[0] pages 0..K for pre/post regions
     - Layer 5 (buf_idx=0): SAME pool[0] pages (double buffer reuse)
     - Layer 4 (buf_idx=1): pool[1] pages 0..M
 
 This naturally handles heterogeneous layers: each layer maps exactly the pages
 it needs, even if remote sizes differ across layers.
 
-Verified on GB200:
-    - Per-page fabric handles composited into one VA: PASS
-    - Same page handles reused across layers: PASS
-    - Non-fabric handles cannot be mapped: FAIL (not supported on GB200)
+The pages are never exported: they are mapped next to the fabric handle of a
+weight's local experts in one VA reservation, and one set_access covers the
+whole reservation.
 """
 
 from __future__ import annotations
@@ -50,11 +50,11 @@
 
 
 class PagePool:
-    """Pool of per-page fabric handles for remote double buffer.
+    """Pool of per-page local handles for remote double buffer.
 
     Two pools (slot 0, slot 1) for double buffering. Each pool is a list
-    of page-sized fabric handles. Different layers map these pages at
-    different VA positions within their composite VA.
+    of page-sized ``CU_MEM_HANDLE_TYPE_NONE`` handles. Different layers map
+    these pages at different VA positions within their composite VA.
 
     Pages are reused across layers in the same slot:
         - Layer 3 (buf_idx=0): pool[0] pages 0..K for pre/post regions
@@ -64,9 +64,8 @@ class PagePool:
     Pool page size vs VMM granularity:
         Each pool handle is ``page_size`` bytes (default 8 * granularity =
         16MB on GB200).  Larger pool pages reduce the number of cuMemMap
-        calls by 8x, avoiding the CUDA driver internal tracking limit
-        (~130K mappings per process) that causes OOM when 4 processes share
-        a GB200 node with 2MB pages.
+        calls and live mappings by 8x.  With 2MB pages, 4 processes on a
+        GB200 node run out of memory at ~130K live mappings per process.
 
     Attributes:
         device_id: CUDA device ordinal.