From 59944f4477d0c1e90132734109db6779b523d106 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 6 Jul 2026 23:57:44 +0000 Subject: [PATCH] dcp: don't fail boot when B12X DCP collectives don't support the world size The B12X PCIe DCP channel exists only for world sizes 2/4/8, and the runtime dispatchers already fall back to NCCL collectives per call. The dedicated DCP warmup added in 'Optimize B12X DCP collectives and warmup' turned that graceful fallback into a boot failure: with TP6 + DCP3/DCP6 (GLM head66 configs that worked on the v13 stack) warmup_b12x_dcp_a2a raised 'B12X PCIe DCP query all-gather is unavailable for the configured attention geometry' and EngineCore died. - warmup_b12x_dcp_a2a: log once and skip when world size is not 2/4/8 (still raises for genuinely broken geometries at supported sizes). - MLAAttention: don't set dcp_b12x for unsupported DCP sizes, so neither warmup nor the per-step dispatch attempts the PCIe channel at all. Verified on GLM-5.2 NVFP4 TP6/DCP6/MTP3 (v13 wiki shape: GMU 0.957, max_model_len 128000, max_num_batched_tokens 2048): boots clean, KV cache 989k tokens, test.py c0/c3000 coherent (CJK 0), 73-75 tok/s. Co-Authored-By: Claude Fable 5 --- tests/distributed/test_dcp_a2a.py | 26 +++++++++++++++++++ .../layers/attention/mla_attention.py | 5 ++++ vllm/v1/attention/ops/dcp_alltoall.py | 10 +++++++ 3 files changed, 41 insertions(+) diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index 93ae30112db8..e8d375799861 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/tests/distributed/test_dcp_a2a.py @@ -438,6 +438,32 @@ def test_b12x_query_gather_requires_env(monkeypatch: pytest.MonkeyPatch): assert actual is expected +def test_warmup_skips_unsupported_world_size(monkeypatch: pytest.MonkeyPatch): + from vllm.v1.attention.ops import dcp_alltoall + + monkeypatch.setenv("VLLM_USE_B12X_DCP_A2A", "1") + monkeypatch.setattr( + dcp_alltoall, + "_try_b12x_dcp_all_gather_heads", + lambda *args, **kwargs: pytest.fail( + "warmup must not touch the B12X channel for world size 6" + ), + ) + group = _FakeCPGroup(6, None) # type: ignore[arg-type] + + # Must log-and-return instead of raising: the runtime dispatchers fall + # back to NCCL for DCP world sizes without a B12X channel (e.g. TP6). + dcp_alltoall.warmup_b12x_dcp_a2a( + group, # type: ignore[arg-type] + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_batch_size=8192, + total_heads=66, + head_dim=512, + query_head_dim=576, + ) + + class TestPackedA2AKernels: @pytest.mark.skipif( torch.accelerator.device_count() < 1, reason="CUDA is required." diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 38c911bd5971..a2b759a2bca5 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -556,6 +556,11 @@ def __init__( self.dcp_a2a and envs.VLLM_USE_B12X_DCP_A2A and self.attn_backend.get_name() == "B12X_MLA_SPARSE" + # The B12X PCIe DCP channel only exists for world sizes 2/4/8; + # other DCP sizes (e.g. TP6 with DCP3/DCP6) use NCCL collectives. + and _vllm_config is not None + and _vllm_config.parallel_config.decode_context_parallel_size + in (2, 4, 8) ) self.dcp_max_batch_size = ( int(_vllm_config.scheduler_config.max_num_batched_tokens) diff --git a/vllm/v1/attention/ops/dcp_alltoall.py b/vllm/v1/attention/ops/dcp_alltoall.py index a0e85bc90758..065e21b73050 100644 --- a/vllm/v1/attention/ops/dcp_alltoall.py +++ b/vllm/v1/attention/ops/dcp_alltoall.py @@ -280,6 +280,16 @@ def warmup_b12x_dcp_a2a( """Create and exercise the B12X DCP channel before CUDA graph capture.""" if not envs.VLLM_USE_B12X_DCP_A2A: return + if cp_group.world_size not in (2, 4, 8): + # The PCIe channel only exists for these world sizes. The runtime + # dispatchers already fall back to NCCL collectives per call, so an + # unsupported DCP size (e.g. TP6 with DCP3/DCP6) must not fail boot. + logger.warning_once( + "B12X PCIe DCP collectives support world sizes 2/4/8; " + "DCP world size %d uses NCCL collectives instead.", + cp_group.world_size, + ) + return if query_head_dim is None: query_head_dim = head_dim local_query = torch.empty(