diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 72e504b48f45..31a5cb88477b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/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. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 0b89c079b3bc..f3f257f8465e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/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) diff --git a/tests/integration/defs/accuracy/test_dwdp_aggregated.py b/tests/integration/defs/accuracy/test_dwdp_aggregated.py new file mode 100644 index 000000000000..0c5623096111 --- /dev/null +++ b/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) diff --git a/tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py b/tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py index cb98648cfab8..5b9ed7e2b215 100644 --- a/tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py +++ b/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 diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 5c48a75d648d..758a360d0ef4 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/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 diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 044926cd3f11..dceecb3761be 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/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