[DCP] Resolve --dcp-comm-backend to fi_a2a/a2a by default for every model (#39165)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

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

@@ -121,17 +121,18 @@ class Parallel(msgspec.Struct):
         ),
     ] = 1
     dcp_comm_backend: A[
-        str,
+        Optional[str],
         Arg(
             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 "
-            "SM90+ and MNNVL fabric memory, e.g. GB200 NVL72).",
+            "Blackwell and a DCP group within one MNNVL domain). Unset resolves "
+            "to 'fi_a2a' where supported, else 'a2a' on CUDA/ROCm, else 'ag_rs'.",
             choices=["ag_rs", "a2a", "fi_a2a"],
             resolvable=True,
         ),
-    ] = "ag_rs"
+    ] = None
     dcp_replicate_q_proj: A[
         Optional[bool],
         Arg(

python/sglang/srt/arg_groups/model_overrides/kimi_k3.py

@@ -16,7 +16,6 @@ from sglang.srt.arg_groups.model_override_base import (
     resolving_view,
 )
 from sglang.srt.runtime_context import get_platform
-from sglang.srt.utils.common import get_device_name
 
 logger = logging.getLogger(__name__)
 
@@ -130,29 +129,12 @@ def _kimi_k3_overrides(server_args: Any, hf_config: Any) -> dict:
                 f"Decode attention backend for Kimi-K3 DCP must be 'cutedsl_mla', 'tokenspeed_mla' or 'aiter', got {decode_backend!r}."
             )
 
-        if cfg.dcp_replicate_q_proj is None:
+        if cfg.dcp_replicate_q_proj is None and cfg.dcp_comm_backend in (
+            "a2a",
+            "fi_a2a",
+        ):
             logger.info("Kimi-K3 DCP enables replicated Q projection by default.")
             overrides["dcp_replicate_q_proj"] = True
-
-        from sglang.srt.layers.dcp.comm import is_fi_a2a_supported
-
-        device_name = get_device_name()
-        dcp_comm_backend = (
-            "fi_a2a"
-            if is_fi_a2a_supported(
-                dcp_size=cfg.dcp_size,
-                tp_size=cfg.tp_size,
-                pp_size=cfg.pp_size,
-                nnodes=cfg.nnodes,
-            )
-            else "a2a"
-        )
-        logger.info(
-            "Kimi-K3 DCP selects communication backend on "
-            f"{device_name!r}: {cfg.dcp_comm_backend!r} -> "
-            f"{dcp_comm_backend!r}."
-        )
-        overrides["dcp_comm_backend"] = dcp_comm_backend
         return overrides
 
     if not (get_platform().is_sm100 and get_platform().device_sm in (100, 103)):

python/sglang/srt/arg_groups/overrides.py

@@ -78,6 +78,7 @@ from sglang.srt.runtime_context import (
 )
 from sglang.srt.utils.common import (
     get_quantization_config,
+    is_fi_a2a_supported,
     is_gfx95_supported,
     xpu_has_xmx_support,
 )
@@ -1424,6 +1425,32 @@ def _data_parallelism_defaults(view: Any) -> dict:
     return {}
 
 
+@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,
+    ):
+        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.",
+        view.dcp_size,
+        backend,
+    )
+    return {"dcp_comm_backend": backend}
+
+
 @register_post_process
 def _tp_lm_head_all_to_all_default(view: Any) -> dict:
     """Enable the TP LM-head all-to-all path only for pure-DP decode nodes.

python/sglang/srt/arg_groups/parallel_hook.py

@@ -9,6 +9,7 @@ from typing import Any
 
 from sglang.srt.arg_groups.overrides import (
     _data_parallelism_defaults,
+    _dcp_comm_backend_default,
     _dp_lm_head_validation,
     _tp_lm_head_all_to_all_default,
     declare_resolution,
@@ -114,7 +115,8 @@ def handle_context_parallelism(server_args: Any):
     )
 
 
-def handle_dcp_validation(server_args: Any):
+def handle_decode_context_parallelism(server_args: Any):
+    run_post_process_pass(server_args, _dcp_comm_backend_default)
     cfg = resolving_view(server_args)
     if cfg.dcp_size < 1:
         raise ValueError(
@@ -132,10 +134,9 @@ def handle_dcp_validation(server_args: Any):
     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 an NVIDIA CUDA platform "
-            "with SM90+ and MNNVL fabric memory (e.g. GB200 NVL72). The "
-            "authoritative fabric probe runs at model-runner init; use 'a2a' "
-            "or 'ag_rs' on clusters without MNNVL."
+            "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_replicate_q_proj:
         if cfg.dcp_size <= 1:

python/sglang/srt/arg_groups/pipeline.py

@@ -163,7 +163,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
     from sglang.srt.arg_groups.parallel_hook import (
         handle_context_parallelism,
         handle_data_parallelism,
-        handle_dcp_validation,
+        handle_decode_context_parallelism,
         handle_dwdp,
         handle_elastic_ep,
         handle_eplb_and_dispatch,
@@ -171,7 +171,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
     )
 
     validate_prefill_only_disable_kv_cache_args(server_args)
-    handle_dcp_validation(server_args)
+    handle_decode_context_parallelism(server_args)
 
     # Model-arch prefill CUDA-graph default must land before cuda-graph
     # resolution (the declarative registry materializes too late to affect

python/sglang/srt/layers/dcp/comm.py

@@ -36,9 +36,9 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
 )
 from sglang.srt.distributed.parallel_state import GroupCoordinator
-from sglang.srt.runtime_context import get_parallel, get_platform
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import is_hip
-from sglang.srt.utils.common import is_mnnvl_fabric_device
+from sglang.srt.utils.common import is_fi_a2a_supported
 
 _is_hip = is_hip()
 
@@ -390,17 +390,6 @@ def all_gather_kv_cache_for_dcp(
 _FI_A2A_STATE: Optional[dict] = None
 
 
-def is_fi_a2a_supported(
-    *, dcp_size: int, tp_size: int, pp_size: int, nnodes: int
-) -> bool:
-    if not get_platform().is_sm100:
-        return False
-    if is_mnnvl_fabric_device():
-        return True
-    tp_size_per_node = tp_size // max(nnodes // pp_size, 1)
-    return tp_size_per_node % dcp_size == 0
-
-
 def init_fi_a2a_workspace(cp_group: "GroupCoordinator") -> None:
     # Call once per process BEFORE CUDA-graph capture: the FlashInfer init syncs
     # the stream and barriers cross-rank, neither of which is capturable.

python/sglang/srt/utils/common.py

@@ -106,6 +106,7 @@ from sglang.srt.runtime_context import (
     get_flags,
     get_model,
     get_parallel,
+    get_platform,
     get_spec,
 )
 from sglang.srt.utils.video_decoder import _BACKEND, VideoDecoderWrapper
@@ -875,6 +876,17 @@ def is_mnnvl_fabric_device() -> bool:
     return any(tag in name for tag in ("GB200", "GB300"))
 
 
+def is_fi_a2a_supported(
+    *, dcp_size: int, tp_size: int, pp_size: int, nnodes: int
+) -> bool:
+    if not get_platform().is_sm100:
+        return False
+    if is_mnnvl_fabric_device():
+        return True
+    tp_size_per_node = tp_size // max(nnodes // pp_size, 1)
+    return tp_size_per_node % dcp_size == 0
+
+
 @lru_cache(maxsize=1)
 def is_habana_available() -> bool:
     return find_spec("habana_frameworks") is not None