[feat][PCP] Enable PCP MoE-Merge+TBO for DeepSeek V4

atom/distributed/pcp_utils.py

@@ -13,7 +13,8 @@
 `layers/utils/cp_utils.py:cp_all_gather_rerange_output`).
 """
 
-from typing import Optional
+import logging
+from typing import NamedTuple, Optional
 
 import torch
 
@@ -24,6 +25,28 @@
 )
 
 
+class PcpBalGroup(NamedTuple):
+    """One request group for PCP+TBO request-boundary split prefill.
+
+    A prefill batch is split into request groups at request boundaries (never
+    inside a sequence); each group is processed as an independent non-TBO PCP
+    mini-batch (padded to a pcp multiple, round-robin striped, reindexed on its
+    own). Consumed by ModelRunner.run_model (per-group stripe / restore) and the
+    attn builder's `_build_ubatch_prefill_metadata_balanced` (slice + reindex).
+    """
+
+    req_start: int  # first request index of this group (inclusive)
+    req_stop: int  # last request index of this group (exclusive)
+    tok_start: int  # global token offset of the group's first token
+    tok_end: int  # global token offset past the group's last REAL token
+    pad_total: (
+        int  # tok count padded to a pcp multiple = pcp_pad_len(tok_end-tok_start, pcp)
+    )
+
+
+logger = logging.getLogger("atom")
+
+
 def get_pcp_world_size() -> int:
     return get_prefill_context_model_parallel_world_size()
 
@@ -39,22 +62,27 @@ def pcp_is_enabled() -> bool:
 def pcp_pad_len(
     total_tokens: int,
     pcp_size: Optional[int] = None,
+    multiple: int = 1,
 ) -> int:
-    """Padded token count so the global sequence is divisible by pcp_size.
+    """Padded token count so the global sequence is divisible by pcp_size * multiple.
 
     Round-robin split requires the global token count to be divisible by pcp_size
-    (see SGLang `can_dsa_cp_split` assert / HIP `apply_cp_reindex`). Returns the
-    padded length (>= total_tokens); callers pad per-token tensors to this
-    length with dummy tokens (KV length 0) before splitting.
+    (see SGLang `can_dsa_cp_split` assert / HIP `apply_cp_reindex`). `multiple` is
+    an extra factor applied on top of pcp_size when the sequence must additionally
+    be evenly divisible by some multiplier. Returns the padded length
+    (>= total_tokens); callers pad per-token tensors to this length with dummy
+    tokens (KV length 0) before splitting.
+
     """
     if pcp_size is None:
         pcp_size = get_pcp_world_size()
-    if pcp_size <= 1:
+    divisor = pcp_size * max(multiple, 1)
+    if divisor <= 1:
         return total_tokens
-    rem = total_tokens % pcp_size
+    rem = total_tokens % divisor
     if rem == 0:
         return total_tokens
-    return total_tokens + (pcp_size - rem)
+    return total_tokens + (divisor - rem)
 
 
 def pcp_round_robin_split(
@@ -111,6 +139,55 @@ def pcp_allgather_rerange(
     return out
 
 
+# ==== MoE-path PCP collectives (rank-major gather + reduce_scatter) ====
+# Rank-major all_gather + reduce_scatter are a mutually-inverse pair:
+#   - gather (1/W -> full): all_gather(dim=0) concats rank-major, so rank r's
+#     1/W stripe lands at rows [r*L:(r+1)*L]. MoE is per-token so the rank-major
+#     (not global) order is fine.
+#   - reduce_scatter (full partial-sum -> 1/W): sums the pcp-half across ranks
+#     AND scatters dim0 back so rank r receives the summed chunk r == its own
+#     original stripe tokens. No rerange/slice needed.
+
+
+def pcp_allgather_rankmajor(
+    input_: torch.Tensor, pcp_size: Optional[int] = None
+) -> torch.Tensor:
+    """Gather this rank's 1/W stripe shard into the full rank-major sequence
+    via a plain all_gather (dim=0). Inverse of pcp_reduce_scatter."""
+    if pcp_size is None:
+        pcp_size = get_pcp_world_size()
+    if pcp_size <= 1:
+        return input_
+    return get_pcp_group().all_gather(input_.contiguous(), dim=0)
+
+
+def pcp_reduce_scatter(
+    input_: torch.Tensor, pcp_size: Optional[int] = None
+) -> torch.Tensor:
+    """Sum the pcp-half across ranks and scatter dim0 back to this rank's 1/W
+    stripe via a plain reduce_scatter (dim=0). Inverse of pcp_allgather_rankmajor."""
+    if pcp_size is None:
+        pcp_size = get_pcp_world_size()
+    if pcp_size <= 1:
+        return input_
+    return get_pcp_group().reduce_scatter(input_.contiguous(), dim=0)
+
+
+def pcp_all_reduce(
+    input_: torch.Tensor, pcp_size: Optional[int] = None
+) -> torch.Tensor:
+    """All-reduce (sum) over the PCP group, no token reshaping. DECODE path:
+    tokens are pcp-redundant (every rank holds the same full batch), so just sum
+    the pcp-half of the intermediate that combine_outputs' tp all_reduce missed.
+    Uses aiter's compile-safe custom-op all_reduce.
+    """
+    if pcp_size is None:
+        pcp_size = get_pcp_world_size()
+    if pcp_size <= 1:
+        return input_
+    return get_pcp_group().all_reduce(input_)
+
+
 def pcp_round_robin_query_indices(
     n_global_q: int, pcp_size: Optional[int] = None, pcp_rank: Optional[int] = None
 ) -> torch.Tensor:

atom/model_engine/llm_engine.py

@@ -70,20 +70,58 @@ def __init__(self, model, tokenizer=None, **kwargs):
                 "them (use -tp N -pcp M without DP-attention, or -dp N "
                 "--enable-dp-attention without -pcp)."
             )
-        # PCP and TBO (two-batch overlap) are not yet compatible: TBO's
-        # UBatchWrapper calls ForCausalLM.forward once per micro-batch with a
-        # sub-slice of tokens, and PCP stripe-splits at that entry — so each
-        # ubatch would be split independently (double-split), giving each rank
-        # ~1/(2*pcp) tokens and a corrupted all-gather restore.
+        # PCP + TBO prefill: supported via coordinated splitting.
+        # PCP + TBO decode: not yet supported (pcp_all_reduce semantics under
+        # per-request ubatch split are unverified).
         if config.prefill_context_parallel_size > 1 and config.enable_tbo:
-            raise ValueError(
-                "prefill_context_parallel_size > 1 (-pcp) combined with "
-                "--enable-tbo is not supported yet (may be supported in a "
-                "future release): TBO calls ForCausalLM.forward per micro-batch "
-                "with a token sub-slice, so PCP would stripe-split each ubatch "
-                "independently (double-split) and corrupt the output. For now, "
-                "disable one of them."
-            )
+            # TBO overlaps compute with the attn<->MoE PCP collectives, which
+            # only exist in MoE merge mode (ATOM_PCP_MOE_MERGE=1). With
+            # ATOM_PCP_MOE_MERGE=0, MoE runs on each rank's 1/W token shard
+            # with NO extra comm between attn and MoE, so TBO has nothing
+            # to overlap and only adds ubatch-splitting overhead. Force it off.
+            if not envs.ATOM_PCP_MOE_MERGE:
+                logger.warning(
+                    "Disabling TBO because ATOM_PCP_MOE_MERGE=0: in this "
+                    "situation it runs MoE on each rank's 1/W token shard with "
+                    "no extra attn<->MoE communication, so TBO has nothing to "
+                    "overlap and only adds overhead."
+                )
+                config.enable_tbo = False
+                config.enable_tbo_decode = False
+            elif config.tensor_parallel_size > 1:
+                # Cross-communicator (PCP x TP) RCCL deadlock:
+                # under TBO the PCP collectives (comm_stream) run concurrently
+                # and UNORDERED with the TP-group all_reduces (compute_stream,
+                # from attention wo_b / MoE RowParallelLinear). On the TPxPCP
+                # rank grid with large collectives this forms a cross-rank
+                # circular wait -> hang (reproduced MI355 TP4PCP2/TP2PCP4 merge
+                # +TBO, 64k/c32). Serialized (non-TBO single stream) is fine, and
+                # TP=1 has no TP communicator so it cannot form the cycle. Only
+                # TP=1 + PCP keeps TBO; TP>1 falls back to non-TBO.
+                logger.warning(
+                    "Disabling TBO: prefill_context_parallel_size > 1 (-pcp) "
+                    "with tensor_parallel_size > 1 (-tp) hangs under TBO due to "
+                    "a PCP<->TP cross-communicator RCCL deadlock (concurrent "
+                    "unordered collectives on comm/compute streams; see "
+                    "PCP_TBO.md 14.4). TBO with PCP is only supported at -tp 1 "
+                    "(e.g. -tp 1 -pcp 8)."
+                )
+                config.enable_tbo = False
+                config.enable_tbo_decode = False
+            elif config.enable_tbo_decode:
+                raise ValueError(
+                    "prefill_context_parallel_size > 1 (-pcp) combined with "
+                    "--enable-tbo all (decode TBO) is not supported yet. "
+                    "Use --enable-tbo (prefill only) with -pcp."
+                )
+            # Under PCP, TBO prefill uses a request-boundary split (the
+            # non-default TBO split mode; never token-midpoint split), so
+            # ATOM_TBO_PREFILL_TOKEN_SPLIT is ignored.
+            if config.enable_tbo and envs.ATOM_TBO_PREFILL_TOKEN_SPLIT:
+                logger.warning(
+                    "ATOM_TBO_PREFILL_TOKEN_SPLIT is ignored under PCP: TBO "
+                    "prefill uses request-boundary balanced grouping."
+                )
         self.rquest_ids = set()
         self.io_processor = InputOutputProcessor(
             config, self.tokenizer, config.kv_cache_block_size