[None][test] Replace disaggregated DWDP accuracy tests with aggregated coverage

tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py

@@ -487,6 +487,11 @@ def calculate_num_chunks(self, all_rank_num_tokens: List[int]) -> int:
         """
         if self.use_dp and self.comm is not None:
             num_rows = self._dp_padded_num_rows(all_rank_num_tokens)
+        elif self.enable_dwdp:
+            # DWDP prefetches expert weights instead of dispatching tokens, so a
+            # rank only ever processes its own tokens, never more. Keyed off
+            # ``enable_dwdp`` so no non-DWDP path changes the branch it takes.
+            num_rows = max(all_rank_num_tokens)
         else:
             # non-DP: no cross-rank dispatch. The scheduler fills all_rank_num_tokens
             # from [x.shape[0]] before calling here, so it must be a single-element list.

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -536,10 +536,34 @@ def create_py_executor(
     )
     logger.info("ATTENTION RUNTIME FEATURES: ", attn_runtime_features)
 
-    # Initialize DWDP Manager (only for context workers in disaggregated serving)
+    # Initialize DWDP Manager.
+    #
+    # DWDP needs every global MPI rank to be a complete, unsharded model replica
+    # owning one expert slice -- `dwdp_rank = global_mpi_rank() % dwdp_size` is
+    # only meaningful under that bijection. Disaggregated context workers satisfy
+    # this with TP=1, aggregated serving with full attention DP; everything that
+    # shards a replica across ranks is rejected. Pipeline and context parallelism
+    # are rejected even under attention DP, since pairing ranks that hold
+    # different layers as DWDP peers gives wrong expert weights rather than an
+    # error. The TP check tests `dp_size == tp_size` rather than
+    # `enable_attention_dp` so a future partial attention DP cannot slip through.
     dwdp_manager: Optional[DwdpManager] = None
     if llm_args.dwdp_config is not None:
-        assert mapping.tp_size == 1 and llm_args.dwdp_config.dwdp_size > 1, "DWDP requires TP=1 and dwdp_size > 1"
+        if llm_args.dwdp_config.dwdp_size <= 1:
+            raise ValueError(
+                f"DWDP requires dwdp_size > 1, got {llm_args.dwdp_config.dwdp_size}."
+            )
+        if mapping.pp_size > 1 or mapping.cp_size > 1:
+            raise ValueError(
+                "DWDP requires each rank to be a complete model replica, so "
+                "pipeline and context parallelism are not supported, but got "
+                f"pp_size={mapping.pp_size}, cp_size={mapping.cp_size}.")
+        if mapping.tp_size > 1 and mapping.dp_size != mapping.tp_size:
+            raise ValueError(
+                "DWDP requires each rank to be a complete model replica: use "
+                "tp_size=1 (disaggregated context worker) or "
+                "enable_attention_dp=True (aggregated serving), but got "
+                f"tp_size={mapping.tp_size}, dp_size={mapping.dp_size}.")
         dwdp_manager = DwdpManager(config=llm_args.dwdp_config,
                                    dist=dist,
                                    mapping=mapping)

tests/integration/defs/accuracy/test_dwdp_aggregated.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""DWDP accuracy tests in aggregated serving.
+
+DWDP was originally built for the context phase of disaggregated serving, where
+each context worker is its own TP=1 instance and the workers are joined into one
+MPI world by ``trtllm-serve disaggregated_mpi_worker``. The same invariant --
+every rank is a complete model replica owning one expert slice -- also holds for
+a single aggregated instance running attention DP, because ``Mapping.dp_size ==
+tp_size`` there and attention is replicated rather than tensor-sharded.
+
+Running DWDP aggregated keeps the expert-sharing paths under test while removing
+the disaggregated KV cache transceiver from the picture, so these tests do not
+depend on the UCX transport configuration of the cluster they run on.
+"""
+
+import pytest
+
+from tensorrt_llm import LLM
+from tensorrt_llm.llmapi import KvCacheConfig, MoeConfig
+from tensorrt_llm.llmapi.llm_args import DwdpConfig
+
+from ..conftest import llm_models_root, skip_post_blackwell_ultra, skip_pre_blackwell
+from .accuracy_core import GSM8K, LlmapiAccuracyTestHarness
+
+# DeepSeek-V3-Lite has 72 routed experts, partitioned across ``DWDP_SIZE``
+# workers. Mode A is the uniform partition (``size == stride == 72 //
+# DWDP_SIZE``); Mode B uses ``size > stride`` so adjacent peer ranges overlap,
+# and ``(DWDP_SIZE - 1) * stride + size == 72`` must hold exactly.
+#
+# DWDP_SIZE is 4 rather than the 2 the disaggregated tests used: aggregated
+# serving has no generation server, so the whole allocation goes to DWDP peers.
+# Three remote peers per rank instead of one is also what makes
+# ``contention_opt`` meaningful, since it interleaves prefetch slices across
+# peers to spread them over several NVLink links.
+#
+# The MPI world must be exactly ``num_groups * dwdp_size`` ranks -- a rank
+# computes ``group_id = rank // dwdp_size`` and DwdpManager rejects
+# ``group_id >= num_groups``. ``tensor_parallel_size`` below is that world size,
+# so it has to track DWDP_SIZE and the ``num_groups=1`` passed to DwdpConfig.
+DWDP_SIZE = 4
+
+
+class TestDwdpAggDeepSeekV3Lite(LlmapiAccuracyTestHarness):
+    MODEL_NAME = "deepseek-ai/DeepSeek-V3-Lite"
+
+    @pytest.mark.skip_less_device(DWDP_SIZE)
+    @skip_pre_blackwell
+    @skip_post_blackwell_ultra
+    @pytest.mark.parametrize(
+        "num_experts_per_worker,num_prefetch_experts,contention_opt",
+        [
+            (18, 18, False),
+            (24, 16, False),
+            (18, 18, True),
+        ],
+        ids=["mode_a_uniform", "mode_b_overlap", "mode_a_uniform_contention_opt"],
+    )
+    def test_dwdp_agg_accuracy(
+        self,
+        num_experts_per_worker: int,
+        num_prefetch_experts: int,
+        contention_opt: bool,
+    ) -> None:
+        dwdp_config = DwdpConfig(
+            dwdp_size=DWDP_SIZE,
+            num_groups=1,
+            num_experts_per_worker=num_experts_per_worker,
+            num_prefetch_experts=num_prefetch_experts,
+            contention_opt=contention_opt,
+        )
+
+        # Attention DP makes each of the DWDP_SIZE ranks an independent replica
+        # serving its own requests; DWDP supplies the expert weights a rank does
+        # not hold locally. Overlap scheduling is not supported with DWDP.
+        with LLM(
+            f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp",
+            tensor_parallel_size=DWDP_SIZE,
+            enable_attention_dp=True,
+            dwdp_config=dwdp_config,
+            moe_config=MoeConfig(backend="CUTEDSL"),
+            disable_overlap_scheduler=True,
+            enable_autotuner=False,
+            enable_chunked_prefill=False,
+            cuda_graph_config=None,
+            max_batch_size=16,
+            max_num_tokens=8192,
+            kv_cache_config=KvCacheConfig(
+                free_gpu_memory_fraction=0.4,
+                enable_block_reuse=False,
+                enable_partial_reuse=False,
+                tokens_per_block=32,
+            ),
+        ) as llm:
+            # Guard against DWDP silently not running at all. If the config were
+            # dropped before create_py_executor, MoE would fall back to the normal
+            # parallel path, which is also correct and scores the same, so the
+            # accuracy check below could not tell the two apart. create_py_executor
+            # either honours dwdp_config or raises -- it has no branch that ignores
+            # it -- so an LLM that constructed successfully while still carrying the
+            # config had a DwdpManager built for it.
+            assert llm.args.dwdp_config == dwdp_config
+
+            task = GSM8K(self.MODEL_NAME)
+            task.evaluate(llm)

tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py

@@ -2,6 +2,23 @@
 
 Separated from test_disaggregated_serving.py to isolate MPI-dependent test
 infrastructure for easier maintenance.
+
+NOTE: these tests no longer gate pre-merge CI. DWDP accuracy is gated by
+test_dwdp_aggregated.py, which exercises the same expert-sharing paths without
+the disaggregated KV cache transceiver and is therefore not exposed to
+per-cluster transport configuration. The tests here stay registered in the QA
+list, waived under nvbugs/6276923, so that the file keeps being collected --
+it imports helpers from test_disaggregated_serving.py and would otherwise rot
+unnoticed -- and so that dropping the waive is the natural signal once the
+disaggregated path is healthy again. They also remain useful as a manual
+reproduction of the disaggregated DWDP path.
+
+When running it manually, check the launcher's UCX settings first: SLURM
+enroot/pyxis injects ``UCX_TLS=tcp`` from the host MPI stack on some clusters,
+which pins the KV cache transceiver to a transport that can fail there and hang
+the run in ``check_gen_transfer_status``. Clear or pin ``UCX_TLS`` for the
+cluster before running -- see jenkins/scripts/slurm_env_setup.sh and
+examples/disaggregated/slurm/benchmark/start_worker_dwdp.sh.
 """
 
 import contextlib

tests/integration/test_lists/qa/llm_function_core.txt

@@ -67,6 +67,9 @@ accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_chunked_prefill
 accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first
 accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first_kv_cache_v1
 accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_nixl_backend
+accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform]
+accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_b_overlap]
+accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform_contention_opt]
 accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy
 accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_contention_opt
 accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_mode_b_overlap

tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

@@ -49,9 +49,9 @@ l0_gb200_multi_gpus:
   - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-overlap_scheduler]
   - accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] TIMEOUT (90)
   - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60)
-  - accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy
-  - accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_contention_opt
-  - accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_mode_b_overlap
+  - accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform]
+  - accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_b_overlap]
+  - accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform_contention_opt]
   - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm
   - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing
   - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant