From 325ab245a182d07c62c51cc655e7b2dc2f0a31d6 Mon Sep 17 00:00:00 2001
From: Thanhhao <31717833+thanhhao98@users.noreply.github.com>
Date: Tue, 8 Sep 2026 21:55:36 +0700
Subject: [PATCH] [DCP] Allow fi_a2a on single-node systems Blackwell without
 MNNVL fabric ( ex B200 B300) (#37767)

Co-authored-by: Hao Phan <htphan@nvidia.com>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
---
 .../srt/arg_groups/model_overrides/kimi_k3.py | 15 ++++++-
 python/sglang/srt/layers/dcp/comm.py          | 40 ++++++++++++++-----
 2 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py b/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py
index d1579a31a..e8e3b58ad 100644
--- a/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py
+++ b/python/sglang/srt/arg_groups/model_overrides/kimi_k3.py
@@ -16,7 +16,7 @@ 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, is_mnnvl_fabric_device
+from sglang.srt.utils.common import get_device_name
 
 logger = logging.getLogger(__name__)
 
@@ -110,8 +110,19 @@ def _kimi_k3_overrides(server_args: Any, hf_config: Any) -> dict:
             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_mnnvl_fabric_device() else "a2a"
+        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} -> "
diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py
index 915a29b7a..8f2e1ef20 100644
--- a/python/sglang/srt/layers/dcp/comm.py
+++ b/python/sglang/srt/layers/dcp/comm.py
@@ -36,7 +36,8 @@ 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
+from sglang.srt.runtime_context import get_parallel, get_platform
+from sglang.srt.utils.common import is_mnnvl_fabric_device
 
 
 def _warn_deprecated_dcp_accessor(name: str, replacement: str) -> None:
@@ -384,6 +385,17 @@ 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.
@@ -401,7 +413,7 @@ def init_fi_a2a_workspace(cp_group: "GroupCoordinator") -> None:
             decode_cp_a2a_init_workspace,
         )
         from flashinfer.comm.mapping import Mapping
-        from flashinfer.comm.mnnvl import MnnvlConfig, is_mnnvl_fabric_supported
+        from flashinfer.comm.mnnvl import MnnvlConfig
     except ImportError as e:
         raise ImportError(
             "--dcp-comm-backend fi_a2a requires FlashInfer with the DCP "
@@ -416,15 +428,25 @@ def init_fi_a2a_workspace(cp_group: "GroupCoordinator") -> None:
         TorchDistributedCommBackend,
     )
 
-    if not is_mnnvl_fabric_supported(torch.cuda.current_device()):
-        raise RuntimeError(
-            "--dcp-comm-backend fi_a2a requires MNNVL fabric memory (e.g. "
-            "GB200 NVL72); is_mnnvl_fabric_supported() returned False. Use "
-            "--dcp-comm-backend a2a or ag_rs on clusters without MNNVL."
-        )
-
     cp_size = cp_group.world_size
     cp_rank = cp_group.rank_in_group
+    parallel = get_parallel()
+
+    if not is_fi_a2a_supported(
+        dcp_size=cp_size,
+        tp_size=parallel.tp_size,
+        pp_size=parallel.pp_size,
+        nnodes=parallel.nnodes,
+    ):
+        raise RuntimeError(
+            "--dcp-comm-backend fi_a2a needs a Blackwell system whose DCP group "
+            "shares one MNNVL domain: either MNNVL fabric memory (GB200/GB300) "
+            f"or a DCP group inside one node (got dcp_size={cp_size}, "
+            f"tp_size={parallel.tp_size}, pp_size={parallel.pp_size}, "
+            f"nnodes={parallel.nnodes}). Use --dcp-comm-backend a2a or ag_rs "
+            "otherwise."
+        )
+
     mapping = Mapping(
         world_size=cp_size,
         rank=cp_rank,