From e6adeb2408b4add15f7844b46aaaefaabf17ee6c Mon Sep 17 00:00:00 2001 From: Dakai An Date: Tue, 6 Oct 2026 20:05:13 +0800 Subject: [PATCH] [Perf][DSV4.1] Enable the fused all-reduce + mHC boundary at TP2 The CuTe DSL Lamport all-reduce fused with mHC post, collapse and RMSNorm (and the MoE finalize publish) is parametrized by the TP world size, but supports_mhc_all_reduce only admitted TP4. The Blackwell single-node TP recipes now default to TP2, where every sublayer boundary at <=16 tokens fell back to a separate all-reduce plus the unfused mHC post and pre. Admit TP2 and extend the kernel-vs-unfused test to TP2. Signed-off-by: Dakai An --- tests/distributed/test_custom_all_reduce.py | 9 +++++---- tests/kernels/mhc/test_mhc_kernels.py | 4 +++- vllm/models/deepseek_v41/nvidia/ops/mhc.py | 2 +- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index 84ff52449995..4978cdc2fd5e 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -176,10 +176,11 @@ def check(outputs): @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="Requires SM100" ) -def test_all_reduce_mhc_matches_unfused_path(monkeypatch): - if torch.accelerator.device_count() < 4: - pytest.skip("Requires four GPUs with NVLink multicast") - multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc) +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_all_reduce_mhc_matches_unfused_path(monkeypatch, tp_size): + if torch.accelerator.device_count() < tp_size: + pytest.skip(f"Requires {tp_size} GPUs") + multi_process_parallel(monkeypatch, tp_size, 1, _all_reduce_mhc) @pytest.mark.parametrize( diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 00e9834d6215..6219b8486796 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -55,7 +55,9 @@ "tp,ep,hidden,hc,multicast,expected", [ (4, False, 5120, 4, 1, True), - (2, False, 5120, 4, 1, False), + (2, False, 5120, 4, 1, True), + (1, False, 5120, 4, 1, False), + (8, False, 5120, 4, 1, False), (4, True, 5120, 4, 1, False), (4, False, 4096, 4, 1, False), (4, False, 5120, 2, 1, False), diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py index edfdbd1ee40d..b15ab0389ea9 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -60,7 +60,7 @@ def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool: parallel = vllm_config.parallel_config config = vllm_config.model_config.hf_config if ( - parallel.tensor_parallel_size != 4 + parallel.tensor_parallel_size not in (2, 4) or parallel.enable_expert_parallel or config.hidden_size != 5120 or config.hc_mult != 4