[Attention][DCP] Support a2a DCP comm backend for FlashMLA sparse (#61285)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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

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(