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