diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index c21fc15a52..f7c47fe6a3 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -4157,6 +4157,22 @@ def test_fp8_dcp_head_envelope_guard(local_heads, dcp_world_size, should_raise): assert local_pad == gathered_pad +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="FlashMLASparseBackend requires CUDA 9.0 or higher", +) +@pytest.mark.parametrize("comm_backend", ["ag_rs", "a2a"]) +def test_fp8_dcp_accepts_comm_backend(comm_backend): + """The mixed-batch path returns (out, lse) for every token in the layout + both DCP combines consume, so neither comm backend is rejected.""" + device = torch.device(DEVICE_TYPE) + vllm_config = _build_sparse_dcp_vllm_config(16, 4, comm_backend) + kv_cache_spec = create_standard_kv_cache_spec(vllm_config) + builder_cls = FlashMLASparseBackend.get_builder_cls() + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + assert builder.fp8_use_mixed_batch + + def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): """A decode row whose top-k shard holds no local candidates (all -1) has undefined kernel out/lse; it must come back as (0, -inf), the identity of diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 3cf0c33c53..80b8da1797 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -446,12 +446,6 @@ class FlashMLASparseMetadataBuilder( ) if parallel_config.decode_context_parallel_size > 1: - if parallel_config.dcp_comm_backend != "ag_rs": - raise NotImplementedError( - "DCP for FlashMLA sparse is only validated with the " - "default 'ag_rs' DCP comm backend; got " - f"'{parallel_config.dcp_comm_backend}'" - ) if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla": raise NotImplementedError( "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla " @@ -1268,7 +1262,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) # The head-padding slice above can leave `out` non-contiguous, and the - # merge feeds it to reduce_scatter. + # DCP combines (reduce_scatter, direct A2A) need packed heads. return out.contiguous(), lse def _fp8_flash_mla_kernel(