diff --git a/README.md b/README.md index 3db2ef832b2c..7a3775d9a318 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,9 @@ state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs.< ## Tech Blogs +* [04/03] DWDP: Distributed Weight Data Parallelism for High-Performance LLM Inference on NVL72 +✨ [➡️ link](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog19_DWDP_Distributed_Weight_Data_Parallelism_for_High_Performance_LLM_Inference_on_NVL72.md) + * [03/16] Optimizing MoE Communication with One-Sided AlltoAll Over NVLink ✨ [➡️ link](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog18_Optimizing_MoE_Communication_with_One_Sided_AlltoAll_Over_NVLink.md) diff --git a/docs/source/blogs/media/tech_blog19_async_comm_contention.png b/docs/source/blogs/media/tech_blog19_async_comm_contention.png new file mode 100644 index 000000000000..71d15a541c72 Binary files /dev/null and b/docs/source/blogs/media/tech_blog19_async_comm_contention.png differ diff --git a/docs/source/blogs/media/tech_blog19_dwdp_overview.png b/docs/source/blogs/media/tech_blog19_dwdp_overview.png new file mode 100644 index 000000000000..592c6797f9ef Binary files /dev/null and b/docs/source/blogs/media/tech_blog19_dwdp_overview.png differ diff --git a/docs/source/blogs/media/tech_blog19_dwdp_runtime_flow.png b/docs/source/blogs/media/tech_blog19_dwdp_runtime_flow.png new file mode 100644 index 000000000000..e156c575e428 Binary files /dev/null and b/docs/source/blogs/media/tech_blog19_dwdp_runtime_flow.png differ diff --git a/docs/source/blogs/media/tech_blog19_e2e_pareto_frontier.png b/docs/source/blogs/media/tech_blog19_e2e_pareto_frontier.png new file mode 100644 index 000000000000..332392791089 Binary files /dev/null and b/docs/source/blogs/media/tech_blog19_e2e_pareto_frontier.png differ diff --git a/docs/source/blogs/media/tech_blog19_sync_overhead_in_dep.png b/docs/source/blogs/media/tech_blog19_sync_overhead_in_dep.png new file mode 100644 index 000000000000..35372b2eaa37 Binary files /dev/null and b/docs/source/blogs/media/tech_blog19_sync_overhead_in_dep.png differ diff --git a/docs/source/blogs/tech_blog/blog19_DWDP_Distributed_Weight_Data_Parallelism_for_High_Performance_LLM_Inference_on_NVL72.md b/docs/source/blogs/tech_blog/blog19_DWDP_Distributed_Weight_Data_Parallelism_for_High_Performance_LLM_Inference_on_NVL72.md new file mode 100644 index 000000000000..0e412474827d --- /dev/null +++ b/docs/source/blogs/tech_blog/blog19_DWDP_Distributed_Weight_Data_Parallelism_for_High_Performance_LLM_Inference_on_NVL72.md @@ -0,0 +1,357 @@ +# DWDP: Distributed Weight Data Parallelism for High-Performance LLM Inference on NVL72 + +By NVIDIA TensorRT LLM Team + +In LLM inference, workload imbalances and communication bottlenecks often lead to excessive synchronization overhead, limiting GPU utilization. We present DWDP (Distributed Weight Data Parallelism), an inference parallelization strategy that preserves data-parallel execution while offloading MoE weights across peer GPUs. By removing collective inter-rank synchronization, DWDP allows each GPU to progress independently. Implemented in TensorRT LLM and evaluated with DeepSeek-R1 on GB200 NVL72, DWDP improves end-to-end output TPS/GPU by 8.8% at comparable TPS/user in the 20-100 TPS/user serving range under 8K input sequence length and 1K output sequence length. The DWDP implementation has been merged into TensorRT LLM ([PR #12136](https://github.com/NVIDIA/TensorRT-LLM/pull/12136)). A more detailed technical introduction is also available on arXiv ([arXiv:2604.01621](https://arxiv.org/abs/2604.01621)). + +## Table of Contents + +- [Motivation](#motivation) +- [DWDP Overview](#dwdp-overview) + - [High-Level Design](#high-level-design) + - [Roofline Analysis](#roofline-analysis) +- [DWDP Implementation](#dwdp-implementation) + - [Key Components](#key-components) + - [Runtime Flow](#runtime-flow) + - [Current Code-Level Constraints](#current-code-level-constraints) +- [Key Optimizations](#key-optimizations) + - [Eliminating Split-Weight Merge Overhead](#eliminating-split-weight-merge-overhead) + - [Mitigating Asynchronous Communication Contention](#mitigating-asynchronous-communication-contention) +- [Evaluation](#evaluation) + - [Experimental Setup](#experimental-setup) + - [Context-Only Evaluation](#context-only-evaluation) + - [End-to-End Evaluation](#end-to-end-evaluation) + - [Reproducing Steps](#reproducing-steps) +- [Summary](#summary) +- [Future Work](#future-work) +- [Acknowledgment](#acknowledgment) + +## Motivation + +Most existing inference parallelism strategies introduce layer-wise inter-rank synchronization. That synchronization becomes increasingly problematic in real-world LLM serving, where per-rank workloads are rarely balanced. At the request level, different ranks often see different sequence lengths and KV-cache hit rates. At the weight level, activated computation can also vary across ranks, especially for MoE models. Together, these effects create substantial per-rank latency variation during inference. Once the execution model synchronizes at layer boundaries, end-to-end throughput becomes bounded by the slowest rank. + +This effect can be quantified using a DEP configuration for DeepSeek-R1 on GB200 with `ISL/OSL = 8K/1` and input ratio `0.8`. In that setup, synchronization overhead reaches approximately `10%` when the coefficient of variation of per-rank sequence lengths is `20%`, which is well within the range observed in production workloads. In other words, synchronization overhead is not a corner case. Under realistic imbalance, it can materially reduce end-to-end inference throughput. + +This leads to the key design question behind DWDP: can we remove collective synchronization and let each rank progress independently? + +
+
+ Synchronization overhead caused by workload imbalance in DEP +
+
+

Figure 1. Synchronization overhead caused by workload imbalance in DEP for DeepSeek-R1 on GB200 with ISL/OSL = 8K/1 and input ratio 0.8.

+ +## DWDP Overview + +### High-Level Design + +Figure 2 shows the core idea of DWDP on an MoE model such as DeepSeek-R1. DWDP preserves data-parallel execution across ranks while offloading MoE weights across peer GPUs. This design specifically targets MoE weights because they dominate the model memory footprint, whereas attention weights account for a much smaller share. + +Within a DWDP group, attention weights are fully replicated on each rank, while the experts in every MoE layer are partitioned across ranks. As a result, each rank permanently stores only its local experts, and the remaining experts reside on peer GPUs. Before executing an MoE layer, the rank fetches the missing remote experts it needs for that layer. + +At runtime, DWDP overlaps the asynchronous prefetch of remote experts for layer `l+1` with the MoE block of layer `l` and the attention block of layer `l+1`. Together, these two blocks create the compute window that hides remote weight prefetch. Before the MoE block of layer `l+1` begins, the rank waits only for its own prefetched experts to arrive. After the layer finishes, those prefetched remote experts are released. To sustain this pipeline across layers, DWDP uses double buffering with prefetching. + +To eliminate collective inter-rank synchronization during inference, DWDP avoids NCCL-based collective remote-weight gathering such as all-gather. Instead, each rank pulls remote experts from peer GPUs through copy-engine-based `cudaMemcpyAsync`, which does not consume SM resources. These transfers are issued as serial peer-to-peer pulls, so they do not introduce synchronization across the group. Once a rank has the experts it needs for the next MoE block, it can continue independently. + +DWDP also provides greater flexibility in expert placement. Because each rank only needs to fetch the weights for one layer before executing its MoE block, DWDP does not require the number of experts to be exactly divisible by the DWDP group size, and it does not require a perfectly disjoint expert partition across ranks. Instead, ranks can be configured with the same number of local experts while allowing redundant expert placement when necessary, for example to support group sizes that do not evenly divide the number of experts. This weaker placement constraint enables resource provisioning at single-rank granularity. When memory permits, the same redundancy can also reduce remote prefetch overhead by increasing the number of local experts on each rank. + +
+
+ Overview of DWDP with DWDP group size 4 +
+
+

Figure 2. Overview of DWDP with DWDP group size 4.

+ +### Roofline Analysis + +We use a simple layer-wise roofline-style model to identify when DWDP can outperform DEP and what fundamentally limits its gain. This analysis focuses on the context phase of DeepSeek-R1 on GB200 and compares DWDP4 against DEP4, where both methods use a four-rank execution group. + +We focus on two derived metrics in Table 1: `T_compute / T_prefetch`, which indicates whether DWDP can hide remote weight prefetch, and `T_DEP / T_DWDP`, which captures DWDP's expected advantage over DEP. + +Table 1 shows that DWDP begins to outperform DEP at around 16K input tokens at batch size 1. As input sequence length increases, `T_compute / T_prefetch` grows from below `1` to above `1`, indicating that longer contexts provide a sufficiently large compute window to amortize and eventually hide remote prefetch overhead. This reveals that DWDP needs enough computation per layer to cover remote weight prefetch. The 16K crossover is specific to the batch-size-1 setting. Increasing the batch size enlarges the compute window and can make DWDP beneficial even for shorter contexts. + +| Input sequence length | `T_compute / T_prefetch` | `T_DEP / T_DWDP` | +| --- | ---: | ---: | +| 1024 | 0.19 | 0.10 | +| 8192 | 0.62 | 0.73 | +| 16384 | 1.52 | 1.27 | +| 32768 | 4.77 | 1.17 | + +*Table 1. Roofline-style analysis data for DeepSeek-R1 context on GB200. The crossover around 16K tokens is where DWDP begins to outperform DEP at batch size 1.* + +DWDP's advantage over DEP comes from eliminating synchronized all-to-all communication from the critical path. This advantage, however, is not monotonic in input sequence length. Once the sequence becomes very long, computation dominates both methods, so synchronized all-to-all overhead accounts for a smaller fraction of DEP's latency. +Accordingly, the marginal speedup of DWDP decreases as ISL grows further. + +Importantly, this is a conservative analysis: it assumes perfectly balanced workloads and therefore does not capture the additional benefit DWDP can deliver under real-world imbalance, where avoiding synchronization overhead matters even more. + + +## DWDP Implementation + +In this section, we focus on the main DWDP runtime components and the runtime flow during inference. + + +### Key Components + +#### `DwdpConfig` + +The configuration surface lives in [`tensorrt_llm/llmapi/llm_args.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/be12482/tensorrt_llm/llmapi/llm_args.py). DWDP is off by default. In the current productized flow, this config is used on the context server of disaggregated serving. + +The four fields are: + +- `dwdp_size`: the number of GPUs in each DWDP group +- `num_groups`: the number of DWDP groups; total context workers = `num_groups * dwdp_size` +- `num_experts_per_worker`: the number of experts each worker keeps locally +- `num_prefetch_experts`: the number of experts each worker fetches from each peer rank + +Together, these fields define the DWDP group structure and how experts are split between local residency and remote prefetch before inference starts. + + +#### `DwdpLayerHandleCollector` + +Each DWDP-enabled MoE layer registers a `DwdpLayerHandleCollector`. During model initialization, it serves as the per-layer metadata carrier that later enables runtime prefetch. + +- record the CUDA IPC handles for that layer's local MoE weights and related tensors +- record tensor shapes, dtypes, and allocation offsets +- hold peer pointers to that layer's remote MoE weights and related tensors on peer GPUs + + +#### `DwdpPrefetchBuffer` + +`DwdpPrefetchBuffer` is the runtime buffer that stores prefetched remote experts. Its role is to keep the next layer's remote experts ready without overwriting the data still needed by the current layer. + +- two prefetch buffers in ping-pong form +- a dedicated prefetch stream +- prefetch-completion events that tell compute when prefetched data is ready +- compute-completion events that tell the next prefetch when a buffer can be safely reused + +#### `DwdpManager` + +`DwdpManager`, implemented in [`tensorrt_llm/_torch/pyexecutor/dwdp.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/be12482/tensorrt_llm/_torch/pyexecutor/dwdp.py), is the control center of the DWDP runtime. It owns the DWDP lifecycle and orchestrates when prefetch happens, while `DwdpPrefetchBuffer` provides the storage, stream, and events used by that pipeline. `DwdpManager` is responsible for: + +- forming the DWDP group from the global MPI world +- creating and tracking one `DwdpLayerHandleCollector` for each DWDP-enabled MoE layer +- all-gathering metadata across the DWDP group, where the metadata records the local MoE parameter information on each GPU +- allocating and initializing the prefetch buffer +- triggering layer-by-layer prefetch at the right time + + +### Runtime Flow + +
+
+ DWDP runtime flow +
+
+ +In the current DWDP code path, the runtime flow can be summarized by the diagram above. + +1. Configuration and group formation + + If `dwdp_config` is present, `py_executor_creator.py` creates a `DwdpManager`. At that moment, DWDP forms the worker's DWDP group, determines its local DWDP rank within that group, and determines its local expert range. + +2. Per-layer metadata registration during model initialization + + As each DWDP-enabled MoE layer is initialized, it calls `DwdpManager.add_layer(...)` and gets back a `DwdpLayerHandleCollector`. At this stage, DWDP is not moving any weights yet. It is only creating the per-layer objects that will later record the local MoE parameter metadata needed for remote prefetch. + +3. Local metadata registration and handle exchange + + After `load_weights()` completes for a DWDP-enabled MoE layer, `DwdpLayerHandleCollector.register_weights(...)` records the local metadata for that layer. This metadata mainly includes the CUDA IPC handles for the local MoE weights, along with the tensor information needed for peer access. After all relevant layers have registered their local metadata, `DwdpManager.exchange_all_handles()` all-gathers that metadata within the DWDP group so each rank knows which peer-GPU tensors it can pull during prefetch. + +4. Prefetch buffer initialization + + After handle exchange, `DwdpManager.initialize_prefetch_buffer()` allocates the `DwdpPrefetchBuffer` and initializes the events used to coordinate prefetch and compute. At this point, the runtime has everything it needs to start asynchronous layer-by-layer prefetch. + +5. Warmup at the start of each forward step + + At the start of each forward step, `PyExecutor` calls `prefetch_first_layers()`. This primes the first DWDP prefetches so that the first DWDP-enabled MoE layers in that step do not enter with an empty pipeline. In architectures such as DeepSeek-R1, the dense and attention work between consecutive MoE blocks is the compute window that DWDP tries to use to hide remote prefetch. + +6. Layer-by-layer prefetch during inference + + During inference, when a DWDP-enabled MoE layer is about to run, DWDP first waits for that layer's prefetched remote experts to be ready. After the layer finishes, `DwdpManager.record_compute_and_prefetch_next(...)` records compute completion for that layer and immediately triggers prefetch for the next layer that will reuse the same ping-pong slot. This is the steady-state loop that keeps prefetch and compute overlapped. + +### Current Code-Level Constraints + +The current implementation has the following constraints: + +- DWDP only supports the `CuteDSL` MoE backend with `NVFP4`. +- DWDP only supports `TP = 1` inside each DWDP group. +- DWDP only supports the MPI worker launch flow used by `trtllm-serve disaggregated_mpi_worker`. +- DWDP does not support overlap scheduler. +- DWDP does not support EPLB on the same MoE path. +- DWDP requires fused-finalize-enabled FC2. + +## Key Optimizations + +### Eliminating Split-Weight Merge Overhead + +DWDP naturally produces split weights for each MoE layer: local experts stay in the model weights, while remote experts arrive in prefetch buffers. Existing groupedGEMM kernels usually assume that all required weights already live in one contiguous buffer. A straightforward implementation would therefore merge local and remote experts through a device-to-device (D2D) copy before every MoE call. + +That extra merge is expensive because it inserts another bandwidth-heavy step directly on the critical path. In a baseline context-only profiling case with DeepSeek-R1 on GB200x4 under `ISL = 8K`, `ratio = 0.8`, and `max_num_tokens = 32768`, the baseline DWDP pays an additional `34 us` of D2D copy for this pre-launch merge, which accounts for about `3%` of iteration latency. + +To remove that overhead, we extend the cuteDSL groupedGEMM kernels to support TensorList-based inputs so the groupedGEMM kernel can consume multiple weight buffers directly. Instead of first materializing a merged expert-weight buffer, the kernel performs the required indexing and address calculation internally while remaining compatible with the existing layout and sharding scheme. Although this design introduces a small amount of additional instruction overhead, including extra address computations and descriptor loads, profiling and end-to-end evaluation show no meaningful performance regression. In practice, the dominant bottlenecks remain the main compute workload and memory traffic, indicating that the proposed approach effectively removes pre-merge D2D overhead without negatively affecting overall performance. + +The relevant code changes in [PR #12136](https://github.com/NVIDIA/TensorRT-LLM/pull/12136) are in: + +- [`tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/be12482/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py) +- [`tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/be12482/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py) +- [`tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/be12482/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py) +- [`tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/be12482/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py) + +### Mitigating Asynchronous Communication Contention + +Another practical challenge in DWDP is that asynchronous remote-weight pulls can create many-to-one contention at the source-side copy engine. As Figure 3 shows, in each MoE layer multiple ranks may simultaneously pull missing remote experts from the same peer rank. When the layer-wise compute window is only comparable to the remote-weight prefetch time, this source-side serialization stretches the communication window and exposes visible compute bubbles before the next compute region can begin. + +One mitigation is to split each remote-weight transfer into fixed-size slices and schedule those slices in a round-robin order across active destination ranks. This reduces random communication delay by time-multiplexing the source-side copy engine more evenly across competing pulls. + +
+
+ Nsight Systems trace showing many-to-one source-side communication contention in DWDP +
+
+

Figure 3. Nsight Systems trace showing many-to-one source-side communication contention in DWDP under a short compute-window setting.

+ +Our experiments show that the additional gain is most visible when the compute window is short. For example, under the `ISL = 8K` context-only workload with `ISL ratio = 0.5` and `max_num_tokens (MNT) = 16384`, contention mitigation delivers an `8%` TPS/GPU gain over the DWDP version without this optimization. + +**Important:** This optimization is part of the broader DWDP design exploration, but it is not yet included in the current productized DWDP code path. + + + +## Evaluation + +### Experimental Setup + +The experiments in this section use the following setup. +Unless otherwise stated, the results in this section do not include the additional performance gain from the contention-mitigation optimization described above. + +- Hardware: GB200 NVL72 +- Commit: the DWDP implementation evaluated in this section was developed based on TensorRT LLM commit `3a89495` +- Model: [DeepSeek-R1-0528-NVFP4-v2](https://huggingface.co/nvidia/DeepSeek-R1-0528-NVFP4-v2) +- Serving mode: disaggregated serving, with DWDP applied on the context server + +We split the discussion into context-only and end-to-end results. + +### Context-Only Evaluation + +The context-only study isolates the context phase, uses the Artificial Analysis dataset, and compares DWDP against a DEP baseline. + +We first examine a context-only iteration-latency breakdown of DEP4 and DWDP4 for DeepSeek-R1 under `ISL = 8K`, `ratio = 0.8`, and `max_num_tokens = 32768` on GB200x4. The last column reports per-category deltas normalized to the DEP4 iteration latency. + + +| Category | DEP4 (`us`) | DWDP4 (`us`) | `Delta / T_DEP4` | +| --- | ---: | ---: | ---: | +| Attention | 269.67 | 320.56 | -3.86% | +| GroupedGEMM | 342.40 | 337.42 | 0.38% | +| DenseGEMM | 177.50 | 189.28 | -0.89% | +| Others | 241.69 | 284.32 | -3.23% | +| Communication | 126.74 | 0.00 | 9.60% | +| P2P Copy | 0.00 | 429.00 | -- | +| Synchronization Cost | 161.85 | 0.00 | 12.26% | +| **Iteration Latency** | **1319.85** | **1131.58** | **14.26%** | + +*Table 3. Context-only iteration-latency breakdown of DEP4 and DWDP4 for DeepSeek-R1 under `ISL = 8K`, `ratio = 0.8`, and `max_num_tokens = 32768` on GB200x4.* + +The breakdown highlights both the promise and the remaining inefficiencies of DWDP. Relative to DEP, DWDP removes synchronization cost entirely and takes communication off the critical path. Together, these two effects correspond to a `21.86%` gross reduction in iteration latency. + +At the same time, compute categories such as Attention and Others become slower. This slowdown reduces the realized gain to a net `14.26%` improvement. Our follow-up analysis shows that it comes from communication-computation interference, and the dominant cause is power-induced frequency throttling. + +### End-to-End Evaluation + +The end-to-end study uses the SemiAnalysis dataset with ISL=`8K`, OSL=`1K`, and input ratio `0.8`. The generation server configuration is kept fixed, and DWDP is applied only to the context server. The comparison is made against Pareto points from the DEP baseline. + +
+
+ End-to-end Pareto frontier comparison between baseline and DWDP +
+
+

Figure 4. End-to-end Pareto frontier comparison between baseline and DWDP.

+ +Figure 4 shows that DWDP pushes the end-to-end Pareto points toward better serving efficiency: at similar TPS/user, it achieves higher output TPS/GPU than the baseline across most of the target range. + +Table 4 summarizes the average speedup in each TPS/user range. The gain is most pronounced at lower TPS/user. +Comparing Pareto points with similar TPS/user, we find that DWDP typically uses fewer context GPUs than the baseline. This suggests that the gain primarily comes from reduced context GPU demand. + +The serving-efficiency benefit becomes smaller at high TPS/user. In this region, the system is more heavily generation-bottlenecked, and the context stage cannot accumulate enough tokens to amortize DWDP's prefetch overhead. + +| TPS/user range | Avg. DWDP TPS/user speedup | Avg. DWDP TPS/GPU speedup | +| --- | ---: | ---: | +| 20-30 | 1.15 | 1.10 | +| 40-50 | 1.16 | 1.08 | +| 60-70 | 1.00 | 1.10 | +| 80-90 | 1.00 | 1.06 | +| 170-180 | 1.00 | 0.97 | + +*Table 4. End-to-end performance summary of DWDP across target TPS/user ranges.* + +We also evaluate median TTFT, including queueing time. The results are summarized in Table 5. +Compared with the baseline, DWDP increases TTFT across the evaluated TPS/user ranges. At low TPS/user, TTFT can increase substantially for pairs with more aggressive reductions in context GPU count. These regressions come from lowering the aggregate service rate of the context stage and worsening rate matching between the context and generation stages. We expect this issue to be mitigated by better request matching in future work, especially because DWDP enables finer-grained context configurations. + +| TPS/user range | TPS/GPU speedup | Baseline TTFT (ms) | DWDP TTFT (ms) | +| --- | ---: | ---: | ---: | +| 20-30 | 1.10 | 2538 | 8314 | +| 40-50 | 1.08 | 1919 | 7012 | +| 60-70 | 1.12 | 965 | 1640 | +| 80-90 | 1.06 | 1669 | 2280 | +| 170-180 | 0.97 | 494 | 660 | + +*Table 5. Median TTFT comparison across target TPS/user ranges.* + + +### Reproducing Steps + +The reproduction files are provided in `examples/dwdp/`: + +- `env.yaml`: cluster/container/model/dataset inputs +- `dwdp_reproduce.yaml`: DWDP experiment matrix +- `reproduce.py`: config generator and launcher for `submit_dwdp.py` + +The above experiments use: + +- GB200 NVL72 +- DeepSeek-R1 NVFP4 checkpoint +- SemiAnalysis dataset with maximum input length `8K`, output length `1K`, and input ratio `0.8` + +The end-to-end results reported in this blog are Pareto points selected from these reproduced experiments. + +Before running, edit `examples/dwdp/env.yaml` for your environment (for example `partition`, `account`, `container_image`, `model_path`, and dataset paths), then launch: + +```bash +python3 -m pip install pyyaml +python3 examples/dwdp/reproduce.py \ + --env-config examples/dwdp/env.yaml \ + --reproduce-config examples/dwdp/dwdp_reproduce.yaml \ + --output-dir examples/dwdp/generated +``` + +## Summary + +- DWDP's first advantage is that it removes the synchronization overhead caused by imbalanced workloads, which makes it a better fit for real LLM serving. +- DWDP's second advantage is flexibility: it gives the system finer-grained freedom when provisioning context GPUs in disaggregated serving. +- DWDP needs a sufficiently large compute window to hide remote expert prefetch, which is why it is best matched to the context side. +- DWDP depends on strong hardware support. High-bandwidth peer GPU connectivity such as GB200 NVL72 is what makes DWDP practical. +- DWDP introduces new engineering challenges, especially around split-weight handling and asynchronous remote-weight prefetch. +- Today, DWDP supports only a subset of code paths and deployment assumptions. Expanding that support remains future work. + + +## Future Work + + +### Integrate Contention Mitigation into the Productized Path + +The contention-mitigation optimization discussed earlier is not yet included in the current productized DWDP path. Integrating it into the production runtime is a natural next step, especially for workloads with short compute windows where many-to-one copy-engine contention is more likely to surface. + +### Decouple Launch-Time Coordination from MPI + +DWDP relies on the `trtllm-serve disaggregated_mpi_worker` launch path and separate launch scripts because handle exchange and group formation currently depend on MPI communication across context workers. We will replace this MPI-based launch-time coordination with a TCP-based method. + +### Move Beyond CUDA IPC for Broader Topologies + +Remote expert access is built on CUDA IPC handles, which are not suitable for cross-node deployment. We will replace them with a fabric-capable remote-memory mechanism so that DWDP can support broader topologies. + +### Reduce Reliance on Kernel-Specialized Split-Weight Handling + +Longer term, it may be worth exploring memory-management approaches that present a more unified weight view to the kernel, for example through virtual-memory-based assembly. This could reduce the reliance on kernel-specialized handling of split weights. + + +## Acknowledgment + +We would like to thank everyone who contributed to this work. diff --git a/examples/dwdp/README.md b/examples/dwdp/README.md new file mode 100644 index 000000000000..28b9038eed11 --- /dev/null +++ b/examples/dwdp/README.md @@ -0,0 +1,101 @@ +# DWDP Reproduction + +This directory provides a thin reproduction layer on top of +`examples/disaggregated/slurm/benchmark/submit_dwdp.py`. +It does not modify that launcher. Instead, it combines: + +- `env.yaml`: cluster, container, model, and dataset inputs provided by the user +- `dwdp_reproduce.yaml`: the DWDP reproduction matrix +- `reproduce.py`: the script that merges both files, generates full benchmark + configs, and forwards them to `submit_dwdp.py` + +## Files + +- `env.yaml` + Holds environment-specific inputs such as Slurm settings, container image, + mount list, model path, and dataset mapping. +- `dwdp_reproduce.yaml` + Holds only experiment parameters such as `isl`, `osl`, `ctx_tp`, `gen_tp`, + `batch`, `prefetch`, and DWDP settings. +- `generated/` + Output directory for the generated full configs that are passed to + `submit_dwdp.py`. + +## How It Works + +`reproduce.py` reads `env.yaml` and `dwdp_reproduce.yaml`, generates one full +benchmark config per experiment, writes the config by default to `generated/`, then +invokes: + +```bash +python examples/disaggregated/slurm/benchmark/submit_dwdp.py -c +``` + +## Configure `env.yaml` + +Update these sections before running: + +- `slurm` + Set `partition`, `account`, `time`, and any cluster-specific `extra_args`. +- `hardware` + Set `gpus_per_node` for your cluster. +- `environment` + Set `container_image`, `container_mount`, `model_path`, and usually + `trtllm_repo`. + Leave `log_dir` unset unless you intentionally want a fixed log location. + When `log_dir` is omitted, `submit_dwdp.py` creates a unique per-run log + directory automatically. +- `datasets` + Map short dataset keys to concrete dataset files. + +`environment.work_dir` is optional. If omitted, `reproduce.py` automatically +points it to `examples/disaggregated/slurm/benchmark`, which is what +`submit_dwdp.py` expects for locating the benchmark shell scripts. + +## Configure `dwdp_reproduce.yaml` + +This file can define both context-only and end-to-end reproduction experiments. + +The reproduction matrix is split into: + +- `experiment_defaults` + Common fields shared across many experiments. +- `experiments` + One entry per benchmark case. + +Each experiment may reference datasets in two ways: + +- `dataset_key`: resolves through `env.yaml -> datasets` +- `dataset_file`: directly provides the full dataset path + +`dataset_key` is the preferred path when several experiments share the same +dataset file. + +## Usage + +Install required Python dependency first: + +```bash +python3 -m pip install pyyaml +``` + + +```bash +python3 examples/dwdp/reproduce.py \ + --env-config /path/to/env.yaml \ + --reproduce-config /path/to/dwdp_reproduce.yaml \ + --output-dir /path/to/generated +``` + +Before running, update `dwdp_reproduce.yaml` as needed so it includes the +reproduction experiments you want to launch. + +## Generated Configs + +Generated configs are written by default to `examples/dwdp/generated/`. +The filenames include both the experiment name and the generated benchmark +identifier so they can be inspected or reused directly with +`submit_dwdp.py`. + +> **IMPORTANT:** Leave `environment.log_dir` unset by default. Logs are written +> under `examples/disaggregated/slurm/benchmark/logs/`. diff --git a/examples/dwdp/dwdp_reproduce.yaml b/examples/dwdp/dwdp_reproduce.yaml new file mode 100644 index 000000000000..b2fe21afd893 --- /dev/null +++ b/examples/dwdp/dwdp_reproduce.yaml @@ -0,0 +1,396 @@ +# DWDP reproduction matrix. +# +# This file defines only experiment-level benchmark settings. Cluster-, model-, +# and dataset-specific paths belong in ``env.yaml``. + +experiment_defaults: + ratio: 0.8 + isl_std: 0 + enable_dp: true + +experiments: + # Naming convention: + # - `*_minus_1` means fewer context servers than the main case. + # Context-only reproduction. + - name: "context_only_dep" + dataset_key: "context_only_ratio_08" + isl: 8192 + osl: 1 + num_ctx_servers: 1 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 4 + batch: 64 + gen_max_tokens: 256 + ctx_max_bs: 64 + ctx_max_num_tokens: 33000 + mtp: 0 + eplb: 0 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "context_only_dwdp" + dataset_key: "context_only_ratio_08" + isl: 8192 + osl: 1 + num_ctx_servers: 4 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 4 + batch: 64 + gen_max_tokens: 256 + ctx_max_bs: 64 + ctx_max_num_tokens: 33000 + mtp: 0 + eplb: 0 + prefetch: 64 + dwdp: true + dwdp_group: 1 + + # End-to-end reproduction: 20-40 TPS / user. + - name: "e2e_20_40_dep" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 5 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 8 + batch: 256 + gen_max_tokens: 512 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_20_40_dep_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 4 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 8 + batch: 256 + gen_max_tokens: 512 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_20_40_dwdp" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 20 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 8 + batch: 256 + gen_max_tokens: 512 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 64 + dwdp: true + dwdp_group: 5 + + - name: "e2e_20_40_dwdp_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 16 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 8 + batch: 256 + gen_max_tokens: 512 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 64 + dwdp: true + dwdp_group: 4 + + # End-to-end reproduction: 40-60 TPS / user. + - name: "e2e_40_60_dep" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 4 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 8 + batch: 128 + gen_max_tokens: 256 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 0 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_40_60_dep_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 3 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 8 + batch: 128 + gen_max_tokens: 256 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 0 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_40_60_dwdp" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 16 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 8 + batch: 128 + gen_max_tokens: 256 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 0 + prefetch: 64 + dwdp: true + dwdp_group: 4 + + - name: "e2e_40_60_dwdp_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 12 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 8 + batch: 128 + gen_max_tokens: 256 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 0 + prefetch: 64 + dwdp: true + dwdp_group: 3 + + # End-to-end reproduction: 60-80 TPS / user. + - name: "e2e_60_80_dep" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 7 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 16 + batch: 64 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_60_80_dep_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 6 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 16 + batch: 64 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_60_80_dwdp" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 28 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 16 + batch: 64 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 64 + dwdp: true + dwdp_group: 7 + + - name: "e2e_60_80_dwdp_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 24 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 16 + batch: 64 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 64 + dwdp: true + dwdp_group: 6 + + - name: "e2e_60_80_dwdp_minus_2" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 20 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 16 + batch: 64 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 1 + eplb: 256 + prefetch: 64 + dwdp: true + dwdp_group: 5 + + # End-to-end reproduction: 80-100 TPS / user. + - name: "e2e_80_100_dep" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 4 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 16 + batch: 32 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 3 + eplb: 256 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_80_100_dwdp" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 16 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 16 + batch: 32 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 3 + eplb: 256 + prefetch: 64 + dwdp: true + dwdp_group: 4 + + - name: "e2e_80_100_dwdp_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 12 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 16 + batch: 32 + gen_max_tokens: 128 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 3 + eplb: 256 + prefetch: 64 + dwdp: true + dwdp_group: 3 + + # End-to-end reproduction: 160-180 TPS / user. + - name: "e2e_160_180_dep" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 2 + ctx_tp: 4 + num_gen_servers: 1 + gen_tp: 32 + batch: 2 + gen_max_tokens: 8 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 3 + eplb: 0 + prefetch: 0 + dwdp: false + dwdp_group: 1 + + - name: "e2e_160_180_dwdp" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 8 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 32 + batch: 2 + gen_max_tokens: 8 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 3 + eplb: 0 + prefetch: 64 + dwdp: true + dwdp_group: 2 + + - name: "e2e_160_180_dwdp_minus_1" + dataset_key: "e2e_ratio_08" + isl: 8192 + osl: 1024 + num_ctx_servers: 4 + ctx_tp: 1 + num_gen_servers: 1 + gen_tp: 32 + batch: 2 + gen_max_tokens: 8 + ctx_max_bs: 4 + ctx_max_num_tokens: 33000 + mtp: 3 + eplb: 0 + prefetch: 64 + dwdp: true + dwdp_group: 1 diff --git a/examples/dwdp/env.yaml b/examples/dwdp/env.yaml new file mode 100644 index 000000000000..53ae724c96e2 --- /dev/null +++ b/examples/dwdp/env.yaml @@ -0,0 +1,60 @@ +# Cluster and environment inputs for DWDP reproduction. +# +# ``reproduce.py`` combines this file with ``dwdp_reproduce.yaml`` to generate +# full benchmark configs for ``examples/disaggregated/slurm/benchmark/submit_dwdp.py``. +# +# Use this file with ``examples/dwdp/reproduce.py`` as the value of +# ``--env-config``. If ``--env-config`` is omitted, the bundled ``env.yaml`` +# next to ``reproduce.py`` is used. + +slurm: + partition: "" + account: "" + time: "04:00:00" + job_name_prefix: "dwdp-repro" + extra_args: "--gres=gpu:4" + set_segment: true + numa_bind: true + +hardware: + gpus_per_node: 4 + +environment: + container_image: "" + container_mount: "" + model_path: "" + trtllm_repo: "" + build_wheel: false + cuda_architectures: "" + trtllm_wheel_path: "" + # Optional override. Leave this unset to let submit_dwdp.py create + # a unique per-run log directory automatically. + # log_dir: "/path/to/custom/logs" + # If omitted, reproduce.py points work_dir to + # examples/disaggregated/slurm/benchmark automatically. + # work_dir: "/path/to/examples/disaggregated/slurm/benchmark" + worker_env_var: >- + TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 + TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 + ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0 + server_env_var: "TRTLLM_SERVER_DISABLE_GC=1" + ctx_worker_env_var: "" + gen_worker_env_var: "" + +benchmark_defaults: + mode: "e2e" + use_nv_sa_benchmark: false + streaming: true + +profiling: + nsys_on: false + ctx_profile_range: "10-15" + gen_profile_range: "200-250" + +accuracy: + enable_accuracy_test: false + tasks: {} + +datasets: + context_only_ratio_08: "" + e2e_ratio_08: "" diff --git a/examples/dwdp/reproduce.py b/examples/dwdp/reproduce.py new file mode 100644 index 000000000000..7d652d7a65c3 --- /dev/null +++ b/examples/dwdp/reproduce.py @@ -0,0 +1,526 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate and submit DWDP reproduction configs. + +This script combines a user-provided environment YAML and a reproduction +matrix YAML, writes full benchmark configs, and forwards them to +``examples/disaggregated/slurm/benchmark/submit_dwdp.py``. +""" + +import argparse +import math +import re +import subprocess +import sys +from copy import deepcopy +from pathlib import Path +from typing import Any, Dict, List + +try: + import yaml +except ModuleNotFoundError as exc: + yaml = None + YAML_IMPORT_ERROR = exc +else: + YAML_IMPORT_ERROR = None + + +SCRIPT_DIR = Path(__file__).resolve().parent +BENCHMARK_DIR = SCRIPT_DIR.parent / "disaggregated" / "slurm" / "benchmark" +SUBMIT_DWDP_SCRIPT = BENCHMARK_DIR / "submit_dwdp.py" +DEFAULT_ENV_CONFIG = SCRIPT_DIR / "env.yaml" +DEFAULT_REPRODUCE_CONFIG = SCRIPT_DIR / "dwdp_reproduce.yaml" +DEFAULT_OUTPUT_DIR = SCRIPT_DIR / "generated" +TOTAL_EXPERTS = 256 + +DEFAULT_WORKER_ENV_VAR = ( + "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 " + "TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 " + "ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0" +) +DEFAULT_SERVER_ENV_VAR = "TRTLLM_SERVER_DISABLE_GC=1" + +REQUIRED_EXPERIMENT_FIELDS = { + "isl", + "osl", + "num_ctx_servers", + "ctx_tp", + "num_gen_servers", + "gen_tp", + "batch", + "gen_max_tokens", + "ctx_max_bs", + "ctx_max_num_tokens", + "mtp", + "eplb", + "prefetch", + "dwdp", + "dwdp_group", + "ratio", +} + + +def parse_args() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Generate DWDP reproduction configs and submit them" + ) + parser.add_argument( + "--env-config", + default=str(DEFAULT_ENV_CONFIG), + help="Path to the environment YAML", + ) + parser.add_argument( + "--reproduce-config", + default=str(DEFAULT_REPRODUCE_CONFIG), + help="Path to the DWDP reproduction YAML", + ) + parser.add_argument( + "--output-dir", + default=str(DEFAULT_OUTPUT_DIR), + help="Directory for generated benchmark configs", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Generate configs and invoke submit_dwdp.py with --dry-run", + ) + return parser.parse_args() + + +def load_yaml_file(path: Path) -> Dict[str, Any]: + """Load a YAML file into a dictionary.""" + with open(path, "r", encoding="utf-8") as file_obj: + data = yaml.safe_load(file_obj) + + if data is None: + return {} + if not isinstance(data, dict): + raise ValueError(f"Expected mapping at top level in {path}") + return data + + +def merge_nested_dicts(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]: + """Recursively merge two dictionaries.""" + merged = deepcopy(base) + for key, value in override.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = merge_nested_dicts(merged[key], value) + else: + merged[key] = deepcopy(value) + return merged + + +def _as_bool(value: Any) -> bool: + """Convert booleans and boolean-like strings to bool.""" + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.lower() == "true" + return bool(value) + + +def _slugify(value: str) -> str: + """Convert arbitrary text into a file-safe slug.""" + slug = re.sub(r"[^a-zA-Z0-9]+", "_", value).strip("_").lower() + return slug or "experiment" + + +def _cuda_graph_batch_sizes(gen_max_num_tokens: int) -> List[int]: + """Build CUDA graph batch sizes following the internal benchmark defaults.""" + sizes = [1, 2, 4] + stages = [(8, 1024, 8), (1040, 2048, 16)] + for start, limit, step in stages: + if gen_max_num_tokens >= start: + end = min(gen_max_num_tokens, limit) + sizes.extend(range(start, end + 1, step)) + if gen_max_num_tokens > 2048: + sizes.append(gen_max_num_tokens) + return sizes + + +def calc_seq_lens(isl: int, osl: int, ratio: float = 0) -> List[int]: + """Return ``[ctx_max_seq_len, gen_max_seq_len]``.""" + if ratio != 0: + if isl == 1024: + ctx_max_seq_len = 1044 + gen_max_seq_len = 2068 + else: + ctx_max_seq_len = isl + 1024 + gen_max_seq_len = ctx_max_seq_len + osl + else: + ctx_max_seq_len = isl + 1024 + gen_max_seq_len = ctx_max_seq_len + osl + return [ctx_max_seq_len, gen_max_seq_len] + + +def build_worker_config(experiment: Dict[str, Any]) -> Dict[str, Any]: + """Build the ``worker_config`` section for the benchmark launcher.""" + isl = int(experiment["isl"]) + osl = int(experiment["osl"]) + ratio = float(experiment.get("ratio", 0)) + ctx_tp = int(experiment["ctx_tp"]) + gen_tp = int(experiment["gen_tp"]) + batch = int(experiment["batch"]) + gen_max_tokens = int(experiment["gen_max_tokens"]) + ctx_max_bs = int(experiment.get("ctx_max_bs", 4)) + ctx_max_num_tokens = int(experiment.get("ctx_max_num_tokens", isl + 1024)) + mtp = int(experiment.get("mtp", 0)) + eplb = int(experiment.get("eplb", 0)) + enable_dp = _as_bool(experiment.get("enable_dp", True)) + prefetch = int(experiment.get("prefetch", 0)) + dwdp_enabled = _as_bool(experiment.get("dwdp", False)) + dwdp_group = int(experiment.get("dwdp_group", 1)) + num_ctx_servers = int(experiment["num_ctx_servers"]) + + if dwdp_group <= 0: + raise ValueError("dwdp_group must be greater than zero") + if prefetch < 0: + raise ValueError("prefetch must be greater than or equal to zero") + if dwdp_enabled and num_ctx_servers % dwdp_group != 0: + raise ValueError("num_ctx_servers must be divisible by dwdp_group when DWDP is enabled") + + ctx_max_seq_len, gen_max_seq_len = calc_seq_lens(isl, osl, ratio) + max_tokens_in_buffer = math.ceil(ctx_max_seq_len / 64) * 64 + dwdp_size = num_ctx_servers // dwdp_group if dwdp_enabled else 1 + experts_per_worker = int(TOTAL_EXPERTS - (dwdp_size - 1) * prefetch) + if dwdp_enabled and experts_per_worker <= 0: + raise ValueError( + "Invalid DWDP configuration: " + f"TOTAL_EXPERTS={TOTAL_EXPERTS}, dwdp_size={dwdp_size}, and prefetch={prefetch} " + "produce a non-positive num_experts_per_worker. " + "Reduce prefetch or use a smaller DWDP size." + ) + + gen_cfg: Dict[str, Any] = { + "tensor_parallel_size": gen_tp, + "moe_expert_parallel_size": gen_tp, + "enable_attention_dp": True, + "enable_lm_head_tp_in_adp": True, + "pipeline_parallel_size": 1, + "context_parallel_size": 1, + "max_batch_size": batch, + "max_num_tokens": gen_max_tokens, + "max_seq_len": gen_max_seq_len, + "cuda_graph_config": { + "enable_padding": True, + "batch_sizes": _cuda_graph_batch_sizes(gen_max_tokens), + }, + "print_iter_log": True, + "trust_remote_code": True, + "kv_cache_config": { + "enable_block_reuse": False, + "dtype": "fp8", + "free_gpu_memory_fraction": 0.75, + }, + "stream_interval": 100, + "moe_config": { + "backend": "CUTEDSL", + }, + "cache_transceiver_config": { + "backend": "UCX", + "max_tokens_in_buffer": max_tokens_in_buffer, + }, + "num_postprocess_workers": 4, + } + if mtp > 0: + gen_cfg["speculative_config"] = { + "decoding_type": "MTP", + "num_nextn_predict_layers": mtp, + } + if eplb > 0: + gen_cfg["moe_config"]["load_balancer"] = { + "num_slots": eplb, + "layer_updates_per_iter": 1, + } + + ctx_cfg: Dict[str, Any] = { + "max_batch_size": ctx_max_bs, + "max_num_tokens": ctx_max_num_tokens, + "max_seq_len": ctx_max_seq_len, + "tensor_parallel_size": ctx_tp, + "context_parallel_size": 1, + "moe_expert_parallel_size": ctx_tp, + "enable_attention_dp": enable_dp, + "pipeline_parallel_size": 1, + "print_iter_log": True, + "trust_remote_code": True, + "cuda_graph_config": None, + "disable_overlap_scheduler": True, + "kv_cache_config": { + "enable_block_reuse": False, + "dtype": "fp8", + "free_gpu_memory_fraction": 0.3, + }, + "cache_transceiver_config": { + "backend": "UCX", + "max_tokens_in_buffer": max_tokens_in_buffer, + }, + "moe_config": { + "backend": "CUTEDSL", + }, + } + if dwdp_enabled: + ctx_cfg["dwdp_config"] = { + "dwdp_size": dwdp_size, + "num_groups": dwdp_group, + "num_experts_per_worker": experts_per_worker, + "num_prefetch_experts": prefetch, + } + if mtp > 0: + ctx_cfg["speculative_config"] = { + "decoding_type": "MTP", + "num_nextn_predict_layers": mtp, + } + + return {"gen": gen_cfg, "ctx": ctx_cfg} + + +def _build_sub_file(experiment: Dict[str, Any]) -> str: + """Build the benchmark config identifier.""" + batch = int(experiment["batch"]) + gen_tp = int(experiment["gen_tp"]) + isl = int(experiment["isl"]) + osl = int(experiment["osl"]) + mtp = int(experiment.get("mtp", 0)) + eplb = int(experiment.get("eplb", 0)) + ratio = experiment.get("ratio", 0) + num_ctx_servers = int(experiment["num_ctx_servers"]) + ctx_tp = int(experiment["ctx_tp"]) + num_gen_servers = int(experiment["num_gen_servers"]) + dwdp = str(experiment.get("dwdp", "false")).lower() + return ( + f"isl{isl}_osl{osl}_sa{ratio}_lbz{batch}_" + f"{num_ctx_servers}ctx{ctx_tp}_{num_gen_servers}gen{gen_tp}_" + f"dwdp{dwdp}_mtp{mtp}_eplb{eplb}" + ) + + +def resolve_dataset_file(env_config: Dict[str, Any], experiment: Dict[str, Any]) -> str: + """Resolve the dataset path for an experiment.""" + dataset_file = experiment.get("dataset_file") + if dataset_file: + return str(dataset_file) + + dataset_key = experiment.get("dataset_key") + datasets = env_config.get("datasets", {}) + if dataset_key: + if dataset_key not in datasets: + raise ValueError(f"dataset_key '{dataset_key}' is not defined in env.yaml datasets") + return str(datasets[dataset_key]) + + raise ValueError("Each experiment must define either dataset_file or dataset_key") + + +def build_job_name(slurm_config: Dict[str, Any], experiment: Dict[str, Any], sub_file: str) -> str: + """Build a short, descriptive Slurm job name.""" + prefix = str(slurm_config.get("job_name_prefix", "dwdp-reproduce")) + experiment_name = experiment.get("name") + if experiment_name: + return f"{prefix}_{_slugify(str(experiment_name))}" + return f"{prefix}_{sub_file}" + + +def build_accuracy_config(env_config: Dict[str, Any], experiment: Dict[str, Any]) -> Dict[str, Any]: + """Build the accuracy section.""" + base_accuracy = env_config.get("accuracy", {}) + experiment_accuracy = experiment.get("accuracy", {}) + accuracy = merge_nested_dicts(base_accuracy, experiment_accuracy) + # Use env.yaml as the default; experiment-level override takes precedence. + if "enable_accuracy_test" in experiment_accuracy: + accuracy["enable_accuracy_test"] = _as_bool(experiment_accuracy["enable_accuracy_test"]) + else: + accuracy["enable_accuracy_test"] = _as_bool( + base_accuracy.get("enable_accuracy_test", False) + ) + accuracy.setdefault("tasks", {}) + return accuracy + + +def build_full_config(env_config: Dict[str, Any], experiment: Dict[str, Any]) -> Dict[str, Any]: + """Build a full config for submit_dwdp.py.""" + slurm_config = env_config.get("slurm", {}) + hardware_config = env_config.get("hardware", {}) + environment_config = env_config.get("environment", {}) + benchmark_defaults = env_config.get("benchmark_defaults", {}) + profiling_defaults = env_config.get("profiling", {}) + + missing_slurm = [key for key in ("partition", "account", "time") if key not in slurm_config] + if missing_slurm: + raise ValueError( + f"env.yaml missing required slurm keys: {', '.join(sorted(missing_slurm))}" + ) + + if "gpus_per_node" not in hardware_config: + raise ValueError("env.yaml missing required hardware.gpus_per_node") + + missing_environment = [ + key + for key in ("container_image", "container_mount", "model_path") + if key not in environment_config + ] + if missing_environment: + raise ValueError( + "env.yaml missing required environment keys: " + ", ".join(sorted(missing_environment)) + ) + + missing_fields = sorted(REQUIRED_EXPERIMENT_FIELDS - set(experiment)) + if missing_fields: + raise ValueError("Experiment is missing required fields: " + ", ".join(missing_fields)) + + isl = int(experiment["isl"]) + osl = int(experiment["osl"]) + batch = int(experiment["batch"]) + gen_tp = int(experiment["gen_tp"]) + if batch <= 0 or gen_tp <= 0: + raise ValueError("batch and gen_tp must be greater than zero") + num_prompts = int(experiment.get("num_prompts", 4096 if osl == 1 else 20000)) + concurrency_list = str(experiment.get("concurrency_list", batch * gen_tp)) + multi_round = int(experiment.get("multi_round", max(1, num_prompts // (batch * gen_tp)))) + sub_file = _build_sub_file(experiment) + dataset_file = resolve_dataset_file(env_config, experiment) + launcher_work_dir = str(environment_config.get("work_dir", BENCHMARK_DIR)) + + full_environment = { + "container_mount": environment_config["container_mount"], + "container_image": environment_config["container_image"], + "model_path": environment_config["model_path"], + "trtllm_repo": environment_config.get("trtllm_repo", ""), + "build_wheel": _as_bool(environment_config.get("build_wheel", False)), + "cuda_architectures": environment_config.get("cuda_architectures", ""), + "trtllm_wheel_path": environment_config.get("trtllm_wheel_path", ""), + "work_dir": launcher_work_dir, + "worker_env_var": environment_config.get("worker_env_var", DEFAULT_WORKER_ENV_VAR), + "server_env_var": environment_config.get("server_env_var", DEFAULT_SERVER_ENV_VAR), + } + for key in ("ctx_worker_env_var", "gen_worker_env_var"): + if key in environment_config: + full_environment[key] = environment_config[key] + full_environment = merge_nested_dicts(full_environment, experiment.get("environment", {})) + if not full_environment.get("log_dir"): + full_environment.pop("log_dir", None) + benchmark_config = { + "mode": benchmark_defaults.get("mode", "e2e"), + "use_nv_sa_benchmark": _as_bool(benchmark_defaults.get("use_nv_sa_benchmark", False)), + "multi_round": multi_round, + "benchmark_ratio": float(experiment.get("ratio", 0)), + "streaming": _as_bool(benchmark_defaults.get("streaming", True)), + "concurrency_list": concurrency_list, + "input_length": isl, + "output_length": osl, + "dataset_file": dataset_file, + } + benchmark_config = merge_nested_dicts(benchmark_config, experiment.get("benchmark", {})) + + profiling_config = { + "nsys_on": _as_bool(profiling_defaults.get("nsys_on", False)), + "ctx_profile_range": profiling_defaults.get("ctx_profile_range", "10-15"), + "gen_profile_range": profiling_defaults.get("gen_profile_range", "200-250"), + } + profiling_config = merge_nested_dicts(profiling_config, experiment.get("profiling", {})) + + config = { + "slurm": { + "script_file": "disaggr_torch_dwdp.slurm", + "partition": slurm_config["partition"], + "account": slurm_config["account"], + "job_time": slurm_config["time"], + "job_name": build_job_name(slurm_config, experiment, sub_file), + "extra_args": slurm_config.get("extra_args", ""), + "set_segment": _as_bool(slurm_config.get("set_segment", True)), + "numa_bind": _as_bool(slurm_config.get("numa_bind", True)), + }, + "benchmark": benchmark_config, + "hardware": { + "gpus_per_node": int(hardware_config["gpus_per_node"]), + "num_ctx_servers": int(experiment["num_ctx_servers"]), + "num_gen_servers": int(experiment["num_gen_servers"]), + }, + "environment": full_environment, + "profiling": profiling_config, + "accuracy": build_accuracy_config(env_config, experiment), + "worker_config": build_worker_config(experiment), + } + config["slurm"] = merge_nested_dicts(config["slurm"], experiment.get("slurm", {})) + return config + + +def write_config_file(config: Dict[str, Any], output_path: Path) -> None: + """Write the generated config to disk.""" + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", encoding="utf-8") as file_obj: + yaml.safe_dump(config, file_obj, default_flow_style=False, sort_keys=False) + + +def build_output_path(output_dir: Path, experiment: Dict[str, Any]) -> Path: + """Build the config output path for one experiment.""" + sub_file = _build_sub_file(experiment) + name = experiment.get("name") + if name: + filename = f"{_slugify(str(name))}_{sub_file}_config.yaml" + else: + filename = f"{sub_file}_config.yaml" + return output_dir / filename + + +def load_experiments(reproduce_config: Dict[str, Any]) -> List[Dict[str, Any]]: + """Load experiments after applying top-level defaults.""" + experiment_defaults = reproduce_config.get("experiment_defaults", {}) + if not isinstance(experiment_defaults, dict): + raise ValueError("experiment_defaults must be a mapping when provided") + experiments = reproduce_config.get("experiments", []) + if not isinstance(experiments, list) or not experiments: + raise ValueError("dwdp_reproduce.yaml must contain a non-empty experiments list") + + merged_experiments = [] + for experiment in experiments: + if not isinstance(experiment, dict): + raise ValueError("Each experiment entry must be a mapping") + merged_experiment = merge_nested_dicts(experiment_defaults, experiment) + merged_experiment.setdefault("isl_std", 0) + merged_experiments.append(merged_experiment) + return merged_experiments + + +def submit_config(submit_dwdp_script: Path, config_path: Path, dry_run: bool) -> None: + """Forward a generated config to submit_dwdp.py.""" + command = [sys.executable, str(submit_dwdp_script), "-c", str(config_path)] + if dry_run: + command.append("--dry-run") + + print(" " + " ".join(command)) + subprocess.run(command, check=True) + + +def main() -> None: + """Entry point.""" + args = parse_args() + + if YAML_IMPORT_ERROR is not None: + raise SystemExit( + "Missing Python dependency: 'pyyaml'. Install it with: python3 -m pip install pyyaml" + ) from YAML_IMPORT_ERROR + + if not SUBMIT_DWDP_SCRIPT.is_file(): + raise FileNotFoundError(f"submit_dwdp.py not found at expected path: {SUBMIT_DWDP_SCRIPT}") + + env_config = load_yaml_file(Path(args.env_config)) + reproduce_config = load_yaml_file(Path(args.reproduce_config)) + experiments = load_experiments(reproduce_config) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + for index, experiment in enumerate(experiments, start=1): + name = experiment.get("name", _build_sub_file(experiment)) + print(f"[{index}/{len(experiments)}] {name}") + config = build_full_config(env_config, experiment) + output_path = build_output_path(output_dir, experiment) + write_config_file(config, output_path) + print(f" Generated config: {output_path}") + submit_config(SUBMIT_DWDP_SCRIPT, output_path, dry_run=args.dry_run) + + +if __name__ == "__main__": + main()