[DCP] Run fi_a2a on FlashInfer's fused all-to-all + LSE reduce

docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx

@@ -205,7 +205,7 @@ Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request 
 **Blackwell DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles:
 
 - DCP is the only axis that shards the TP-replicated MLA KV; Low-Latency skips it.
-- Leave `--dcp-comm-backend` unset (fabric-resolved: `fi_a2a` on GB200/GB300, `a2a` on B200/B300).
+- Leave `--dcp-comm-backend` unset: it resolves to `fi_a2a` (FlashInfer's fused all-to-all + LSE merge) where the DCP group shares one NVLink domain, GB200/GB300 fabric or one node, and the installed `flashinfer-python` provides `decode_cp_a2a_lse_reduce`; otherwise to `a2a` (the startup log names the missing requirement).
 - No `--enable-symm-mem` under DCP (force-disabled for decode-graph correctness).
 - Explicit `tokenspeed_mla` force-rewrites `--kv-cache-dtype` to fp8; the default `cutedsl_mla` serves either dtype.
 - Calculator ratios run well above 1 here (`r > 1` is legal): `bfloat16` state buys admission, `fp8` KV buys context.

docs/docs/advanced_features/dcp.mdx

@@ -15,7 +15,7 @@ On Kimi K3, DCP applies only to MLA layers. Request-indexed KDA state is unchang
 | Setting | Behavior |
 | --- | --- |
 | `--dcp-size N` | Enables DCP and widens virtual capacity and page size by `N`. Alias: `--decode-context-parallel-size`. |
-| `--dcp-comm-backend` | Selects the MLA partial-output merge: `ag_rs`, `a2a`, or `fi_a2a`. Unset resolves to `fi_a2a` where the DCP group shares an MNNVL domain on Blackwell, otherwise `a2a` on CUDA/ROCm. |
+| `--dcp-comm-backend` | Selects the MLA partial-output merge: `ag_rs`, `a2a`, or `fi_a2a`. Unset resolves to `fi_a2a` where FlashInfer's fused reduce can run (see [Communication backends](#communication-backends)), otherwise `a2a` on CUDA/ROCm. |
 | `--dcp-replicate-q-proj` | Removes the query all-gather for supported MLA weights. Use `--no-dcp-replicate-q-proj` to disable a model-specific default. |
 | `--attn-dp-size N` | Composes only when DCP groups nest inside attention TP. |
 
@@ -44,7 +44,7 @@ sglang serve \
   --port 30000
 ```
 
-Kimi K3 on 32 GPUs with DPA and DCP 8. The model override enables replicated Q by default; the communication backend resolves generically to `fi_a2a` on MNNVL systems, otherwise `a2a`:
+Kimi K3 on 32 GPUs with DPA and DCP 8. The model override enables replicated Q by default; the communication backend resolves generically to `fi_a2a` where the fused reduce can run, otherwise `a2a`:
 
 ```bash Command
 sglang serve \
@@ -132,9 +132,11 @@ Normally each layer all-gathers its absorbed and rotary query components. `--dcp
 | `ag_rs` | Query AG + LSE AG + FP32 output RS | Fallback for platforms without the a2a path; opt in with `--dcp-comm-backend ag_rs` |
 | `a2a`, gathered Q | Query AG + packed NCCL A2A | Two collectives |
 | `a2a`, replicated Q | One packed NCCL A2A | Canonical low-latency path |
-| `fi_a2a`, replicated Q | One FlashInfer MNNVL A2A | Requires Blackwell and a DCP group inside one MNNVL domain (GB200/GB300 fabric, or a single node) |
+| `fi_a2a`, replicated Q | One FlashInfer kernel: all-to-all and LSE merge fused | Requires Blackwell, an MLA model, a DCP group inside one NVLink domain (GB200/GB300 fabric, or one node), and a `flashinfer-python` with the fused `decode_cp_a2a_lse_reduce` |
 
-When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the DCP group shares one MNNVL domain on Blackwell, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`.
+When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the fused reduce can run, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. On CUDA, the startup log names the missing requirement when `fi_a2a` is not selected. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`.
+
+`fi_a2a` keeps two fused-reduce workspaces, one for CUDA graph replays and one for eager decode, each sized for the widest decode CUDA graph batch (`--cuda-graph-max-bs-decode`, times the tokens each request verifies under speculative decoding), also with `--disable-decode-cuda-graph`; eager batches wider than that run as several fused calls. Each workspace takes about `2 × dcp_size × tokens × heads × 1.2 KB` for the 512-wide BF16 latent, where `heads` counts the attention heads per attention-TP rank: 150 MiB at 512 tokens, 16 heads and DCP 8. Both are allocated after KV-cache sizing, out of the memory `--mem-fraction-static` leaves free, and the startup log reports their size. A workspace serves only one CUDA stream, so `--enable-pdmux`, and `SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP` with `--pp-size` above 1, resolve to `a2a`.
 
 DCP adds a context-independent decode collective while cutting context-dependent KV storage and reads by about `c`. Extend is outside that decode cost model: it may gather cached prefix shards, restore token order, and append new tokens, and that work grows with prefix length.
 

docs/docs/advanced_features/server_arguments.mdx

@@ -668,7 +668,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
     </tr>
         <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dcp-comm-backend`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Communication backend for the DCP attention reduction: AllGather + ReduceScatter, fused NCCL All-to-All, or FlashInfer MNNVL All-to-All. Unset resolves to <code>fi_a2a</code> where the DCP group shares an MNNVL domain on Blackwell, else <code>a2a</code> on CUDA/ROCm, else <code>ag_rs</code>.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Communication backend for the DCP attention reduction: AllGather + ReduceScatter, fused NCCL All-to-All, or FlashInfer's fused all-to-all + LSE merge (Blackwell, an MLA model, one NVLink domain, and FlashInfer with <code>decode_cp_a2a_lse_reduce</code>). Unset resolves to <code>fi_a2a</code> where it can run, else <code>a2a</code> on CUDA/ROCm, else <code>ag_rs</code>.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>None</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>ag_rs</code>, <code>a2a</code>, <code>fi_a2a</code></td>
     </tr>

python/sglang/kernels/ops/attention/dcp_kernels.py

@@ -385,7 +385,7 @@ def correct_attn_out(
 
 
 # A2A DCP reduce: LSE-weighted combine of N partial attention outputs
-# (used by the a2a / fi_a2a communication backends, see comm.py).
+# (used by the a2a communication backend, see comm.py).
 
 
 def _lse_pack_dim(output_dtype: torch.dtype) -> int:

python/sglang/srt/arg_groups/fields/parallel.py

@@ -130,9 +130,11 @@ class Parallel(msgspec.Struct):
             help="Communication backend for the decode context-parallel (DCP) "
             "attention reduction: 'ag_rs' (AllGather + ReduceScatter), 'a2a' "
             "(fused NCCL All-to-All exchange of output+LSE + local Triton LSE "
-            "combine), or 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires "
-            "Blackwell and a DCP group within one MNNVL domain). Unset resolves "
-            "to 'fi_a2a' where supported, else 'a2a' on CUDA/ROCm, else 'ag_rs'.",
+            "combine), or 'fi_a2a' (FlashInfer's decode_cp_a2a_lse_reduce, which "
+            "exchanges and merges the partials in one kernel; requires Blackwell, "
+            "an MLA model, the DCP group inside one NVLink domain and a FlashInfer "
+            "build with that op). Unset resolves to 'fi_a2a' where it can run, "
+            "else 'a2a' on CUDA/ROCm, else 'ag_rs'.",
             choices=["ag_rs", "a2a", "fi_a2a"],
             resolvable=True,
         ),

python/sglang/srt/arg_groups/overrides.py

@@ -68,8 +68,8 @@
     num_dp_ranks_of,
 )
 from sglang.srt.utils.common import (
+    fi_a2a_platform_blocker,
     get_quantization_config,
-    is_fi_a2a_supported,
     is_gfx95_supported,
     xpu_has_xmx_support,
 )
@@ -1316,28 +1316,50 @@ def _data_parallelism_defaults(view: Any) -> dict:
     return {}
 
 
+def fi_a2a_blocker(view: Any) -> Optional[str]:
+    """Why --dcp-comm-backend fi_a2a cannot run with this configuration, or None."""
+    reason = fi_a2a_platform_blocker(
+        dcp_size=view.dcp_size,
+        tp_size=view.tp_size,
+        pp_size=view.pp_size,
+        nnodes=view.nnodes,
+    )
+    if reason is not None:
+        return reason
+    if not use_mla_backend(view):
+        return "applies only to MLA models (the MHA DCP reduction ignores it)"
+    if view.enable_pdmux:
+        return (
+            "does not support --enable-pdmux (every captured stream group would "
+            "need its own fused-reduce workspace)"
+        )
+    if envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get() and view.pp_size > 1:
+        return (
+            "cannot be combined with SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP at "
+            "--pp-size > 1 (its decode warmup runs on a third CUDA stream)"
+        )
+    return None
+
+
 @register_post_process
 def _dcp_comm_backend_default(view: Any) -> dict:
     if view.dcp_comm_backend is not None:
         return {}
     if view.dcp_size <= 1:
         return {"dcp_comm_backend": "ag_rs"}
     platform = get_platform()
-    if is_fi_a2a_supported(
-        dcp_size=view.dcp_size,
-        tp_size=view.tp_size,
-        pp_size=view.pp_size,
-        nnodes=view.nnodes,
-    ):
+    blocker = fi_a2a_blocker(view)
+    if blocker is None:
         backend = "fi_a2a"
     elif platform.is_cuda or platform.is_hip:
         backend = "a2a"
     else:
         backend = "ag_rs"
     logger.info(
-        "DCP (dcp_size=%d) selects communication backend %r.",
+        "DCP (dcp_size=%d) selects communication backend %r%s.",
         view.dcp_size,
         backend,
+        f" (fi_a2a {blocker})" if blocker is not None and platform.is_cuda else "",
     )
     return {"dcp_comm_backend": backend}
 

python/sglang/srt/arg_groups/parallel_hook.py

@@ -13,6 +13,7 @@
     _dp_lm_head_validation,
     _tp_lm_head_all_to_all_default,
     declare_resolution,
+    fi_a2a_blocker,
     model_config_of,
     resolved_view,
     resolving_view,
@@ -240,13 +241,13 @@ def handle_decode_context_parallelism(server_args: Any):
             "requires --dcp-size / --decode-context-parallel-size > 1, but "
             f"got dcp_size={cfg.dcp_size}."
         )
-    if cfg.dcp_comm_backend == "fi_a2a" and not get_platform().is_cuda:
-        raise ValueError(
-            "--dcp-comm-backend fi_a2a delegates the exchange to FlashInfer's "
-            "MNNVL All-to-All kernel, which requires Blackwell and a DCP group "
-            "within one MNNVL domain. Use 'a2a' or 'ag_rs' elsewhere, or leave "
-            "the flag unset to resolve it."
-        )
+    if cfg.dcp_comm_backend == "fi_a2a":
+        blocker = fi_a2a_blocker(cfg)
+        if blocker is not None:
+            raise ValueError(
+                f"--dcp-comm-backend fi_a2a {blocker}. Use --dcp-comm-backend a2a, "
+                "or leave the flag unset to resolve a supported backend."
+            )
     if cfg.dcp_replicate_q_proj:
         if cfg.dcp_size <= 1:
             raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.")