From 69737e283139b3e98548e8449e16468861a1b019 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 12 Jun 2026 11:36:44 +0800 Subject: [PATCH 01/17] [Fea] add moe_pcp_merge --- atom/config.py | 10 ++++++++ atom/model_engine/arg_utils.py | 11 +++++++++ atom/model_engine/llm_engine.py | 9 +++++++ atom/models/deepseek_v4.py | 42 +++++++++++++++++++++++++++++++++ 4 files changed, 72 insertions(+) diff --git a/atom/config.py b/atom/config.py index a651dd9299..789c32481e 100644 --- a/atom/config.py +++ b/atom/config.py @@ -749,6 +749,16 @@ class ParallelConfig: data_parallel_size: int = 1 """Number of data parallel groups. MoE layers will be sharded according to the product of the tensor parallel size and data parallel size.""" + moe_pcp_merge: bool = False + """PCP (Prefill Context Parallel) MoE comm mode. Only meaningful when + prefill_context_parallel_size > 1. + - False (mode A, default): MoE runs on each rank's 1/W token shard with NO + extra comm (per-token independent). Equivalent to SGLang moe_dp_size == + pcp_size. + - True (mode B): all-gather hidden 1/W -> full before MoE and slice full -> + 1/W after, so MoE sees the complete token set (MoE itself is untouched / + PCP-agnostic). Equivalent to SGLang moe_dp_size == 1. Costs one extra + hidden all-gather per layer; see plan §P3.""" data_parallel_size_local: int = 1 """Number of local data parallel groups.""" data_parallel_rank: int = 0 diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index a8533600c7..89fe3270bb 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -31,6 +31,7 @@ class EngineArgs: trust_remote_code: bool = False tensor_parallel_size: int = 1 prefill_context_parallel_size: int = 1 + moe_pcp_merge: bool = False data_parallel_size: int = 1 enforce_eager: bool = False enable_prefix_caching: bool = True @@ -88,6 +89,16 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="Prefill context parallel size. Independent dimension " "(world = tp x pcp); splits the sequence during prefill.", ) + parser.add_argument( + "--moe-pcp-merge", + action=argparse.BooleanOptionalAction, + default=False, + help="PCP MoE comm mode (only when -pcp > 1). Default (False, mode " + "A): MoE runs on each rank's 1/W token shard, no extra comm. " + "True (mode B): all-gather hidden to full tokens before MoE and " + "slice back after, so MoE sees the complete sequence (MoE itself " + "untouched). Adds one hidden all-gather per layer.", + ) parser.add_argument( "--data-parallel-size", "-dp", diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index a20b019dc2..7d41223f83 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -39,6 +39,7 @@ def __init__(self, model, tokenizer=None, **kwargs): config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} data_parallel_size = kwargs.get("data_parallel_size", 1) data_parallel_master_port = kwargs.get("data_parallel_master_port", None) + moe_pcp_merge = kwargs.get("moe_pcp_merge", False) config = Config(model, **config_kwargs) self.config = config self.tokenizer = tokenizer or _load_tokenizer( @@ -55,6 +56,14 @@ def __init__(self, model, tokenizer=None, **kwargs): if data_parallel_master_port is not None: config.parallel_config.data_parallel_master_port = data_parallel_master_port self.data_parallel_size = data_parallel_size + # PCP MoE comm mode (mode B). Only meaningful when PCP is enabled. + if moe_pcp_merge and config.prefill_context_parallel_size <= 1: + raise ValueError( + "moe_pcp_merge=True (PCP MoE mode B) requires " + "prefill_context_parallel_size > 1 (-pcp). With PCP disabled " + "there is no token sharding to all-gather for MoE." + ) + config.parallel_config.moe_pcp_merge = moe_pcp_merge # PCP and DP-attention are not yet compatible: PCP stripe-splits # input_ids to 1/pcp_size in ForCausalLM.forward, but DP-attention's # `_gather_ids_for_dp` all-gathers using dp_metadata sizes computed on diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index cc1f15216a..254b376181 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -2669,9 +2669,24 @@ def forward( self.norm_eps, ) x = hc_state.x_prev + # PCP MoE mode B: all-gather this rank's 1/W token shard back to the + # full (padded) sequence so MoE sees every token (MoE itself is + # untouched / PCP-agnostic), then slice the result back to 1/W to match + # the residual stream (which stays sharded throughout). Mode A leaves x + # at 1/W and MoE runs per-shard with no extra comm. Only `x` is + # gathered/sliced — fuse_hc already produced the 1/W ffn input, and the + # residual in hc_state remains 1/W. The gather length is the fixed + # padded total (not data-dependent), so it is torch.compile-safe like + # the K/V all-gather inside attention.forward_impl. + moe_merge = _moe_pcp_merge_active() + if moe_merge: + pcp_ws = get_pcp_world_size() + x = pcp_allgather_rerange(x, pcp_ws) # [1/W, dim] -> [full, dim] x = self.ffn( x ) # [num_tokens, dim] (input_ids read from forward_context for hash MoE) + if moe_merge: + x = pcp_split_stripe(x, pcp_ws) # [full, dim] -> [1/W, dim] hc_state.x_prev = x return hc_state @@ -2755,6 +2770,19 @@ def _pcp_active() -> bool: return fc.context.is_prefill and not fc.context.is_dummy_run +def _moe_pcp_merge_active() -> bool: + """Whether PCP MoE mode B (all-gather hidden -> full before MoE, slice back + after) applies in this forward. + + True only when this is a real PCP prefill (`_pcp_active`) AND the + `moe_pcp_merge` flag is set. Mode A (default) returns False: MoE runs on the + 1/W shard with no extra comm. + """ + if not _pcp_active(): + return False + return bool(get_current_atom_config().parallel_config.moe_pcp_merge) + + @support_torch_compile class DeepseekV4Model(nn.Module): """Full model: embed -> expand to hc_mult copies -> N blocks -> hc_head -> logits. @@ -2976,6 +3004,15 @@ def forward( # runs entirely on 1/W; the final hidden is all-gathered + un-padded # after self.model(...) returns. use_pcp = _pcp_active() + moe_merge = _moe_pcp_merge_active() + # Mode B is incompatible with DP-attention id-gathering for now: both + # rewrite ctx.context.input_ids with different (full-PCP vs DP-gathered) + # token sets, and stacking them is unverified. Disallow until needed. + assert not (moe_merge and self._need_ids_gather), ( + "moe_pcp_merge (PCP MoE mode B) is not supported together with " + "DP-attention input-id gathering yet." + ) + full_padded_ids = None if use_pcp: pcp_size = get_pcp_world_size() n_global = input_ids.shape[0] @@ -2983,6 +3020,11 @@ def forward( if pad > 0: input_ids = torch.cat([input_ids, input_ids.new_zeros(pad)], dim=0) positions = torch.cat([positions, positions.new_zeros(pad)], dim=0) + # Mode B: stash the full (padded, global-order) ids for hash MoE, + # which sees full tokens after the per-layer all-gather. The token + # order matches pcp_allgather_rerange's reconstruction in Block. + if moe_merge: + full_padded_ids = input_ids input_ids = pcp_round_robin_split(input_ids, pcp_size) positions = pcp_round_robin_split(positions, pcp_size) From 97c87adf6fb8aef9d29b8eed4bc536b9ee0c9f8f Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Wed, 24 Jun 2026 13:54:22 +0800 Subject: [PATCH 02/17] fix moe_pcp_merge --- atom/distributed/pcp_utils.py | 55 ++++++++++ atom/model_ops/moe.py | 24 +++++ atom/models/deepseek_v4.py | 187 +++++++++++++++++++++++++++++----- 3 files changed, 239 insertions(+), 27 deletions(-) diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 79ccf37585..6db969b29c 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -111,6 +111,61 @@ def pcp_allgather_rerange( return out +# ==== MoE-path PCP collectives (scheme B, rank-major gather + reduce_scatter) ==== +# Used by the MoE merge path (moe_pcp_merge_forward custom op). Because that op +# is opaque to Dynamo (the collectives run eagerly inside it, NOT in the compiled +# graph), we use the NATURAL collectives — no all-reduce emulation needed. The +# earlier all-reduce-only scheme was a workaround for "collectives miscompiled +# inside the @support_torch_compile graph"; with the opaque-op fix (方向C) the +# gather/scatter no longer live in the traced graph, so plain all_gather / +# reduce_scatter are correct and cheaper (all_gather moves W× fewer bytes). +# +# 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: diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 37b4bc8b9e..1380ed6861 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -146,6 +146,30 @@ def flatten_tp_across_dp(dp_rank: int): tp_size = tp_size_ tp_rank = 0 if tp_size_ == 1 else get_tp_group().rank_in_group + # PCP moe_pcp_merge: fold the prefill-context-parallel dim into the MoE + # tensor/expert sharding so the W=pcp_size redundant copies become real + # shards (intermediate//W*tp or expert//W*tp). The flattened rank MUST be + # `pcp_rank * tp_size + tp_rank`: this makes each TP group [0..tp-1] own a + # contiguous half and its PCP partner own the other half, so the existing + # tp all_reduce + the new pcp reduce_scatter (two orthogonal groups) sum + # all W*tp shards. The reversed mapping would make PCP partners overlap + # and double-count. ep_size/ep_rank below inherit tp_size/tp_rank, so EP + # is covered by the same flatten. + from aiter.dist.parallel_state import ( + get_prefill_context_model_parallel_rank, + get_prefill_context_model_parallel_world_size, + ) + + pcp_merge = ( + parallel_config.parallel_config.moe_pcp_merge + and get_prefill_context_model_parallel_world_size() > 1 + ) + if pcp_merge: + pcp_size = get_prefill_context_model_parallel_world_size() + pcp_rank = get_prefill_context_model_parallel_rank() + tp_rank = pcp_rank * tp_size + tp_rank + tp_size = pcp_size * tp_size + atom_config = get_current_atom_config() if not use_ep: diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 254b376181..fb273f7b5b 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -43,10 +43,14 @@ ) from atom.distributed.pcp_utils import ( get_pcp_world_size, + pcp_all_reduce, + pcp_allgather_rankmajor, pcp_allgather_rerange, pcp_pad_len, + pcp_reduce_scatter, pcp_round_robin_split, ) +from atom.utils.custom_register import direct_register_custom_op from aiter.ops.topk import top_k_per_row_decode, top_k_per_row_prefill from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits from aiter.ops.triton.fusions.fused_clamp_act_mul import ( @@ -2267,9 +2271,15 @@ def __init__( and self.alt_stream is not None and envs.ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD > 0 ) - if self._use_dual_stream: - # Register self in static_forward_context so the custom op - # dispatcher can look us up by `layer_name` (= self.prefix). + # Register self in static_forward_context so the custom op dispatcher + # can look us up by `layer_name` (= self.prefix). Needed by + # maybe_dual_stream_forward (dual-stream) AND moe_pcp_merge_forward + # (PCP mode B, 阶段19) — the latter requires registration regardless of + # dual-stream, so register whenever either consumer is active. + _merge_on = get_pcp_world_size() > 1 and bool( + get_current_atom_config().parallel_config.moe_pcp_merge + ) + if self._use_dual_stream or _merge_on: get_current_atom_config().compilation_config.static_forward_context[ prefix ] = self @@ -2384,6 +2394,18 @@ def combine_outputs( all-reduce across TP ranks. """ if shared is not None: + # Scheme B (moe_pcp_merge, non-fused shared only): the shared expert + # is NOT pcp-sharded (its MergedColumn/RowParallelLinear bind to the + # 4-card tp group, so every pcp rank holds the same shared weights + # and computes the same full shared output — pcp-redundant). After + # this combine the result rides through Block.forward's pcp + # reduce_scatter, which SUMS the pcp partners. Without correction the + # shared part would be summed pcp_size times (doubled for pcp=2). So + # pre-scale shared by 1/pcp_size: the reduce_scatter then RESTORES it + # to 1x instead of multiplying. routed is genuinely pcp-sharded + # (partial sum) so it must NOT be scaled — only shared. + if _moe_pcp_merge_active() or _moe_pcp_merge_decode_active(): + shared = shared * (1.0 / get_pcp_world_size()) routed = routed + shared if self.tp_size > 1: routed = tensor_model_parallel_all_reduce(routed) @@ -2474,6 +2496,13 @@ def __init__( indexer_stream=indexer_stream, ) self.ffn = MoE(layer_id, args, prefix=f"{prefix}.ffn", alt_stream=alt_stream) + # 阶段19: static (config-only, is_dummy-independent) gate for routing MoE + # through the opaque moe_pcp_merge_forward custom op. Constant across + # warmup/real so Dynamo specializes it identically -> the op call is + # baked into code0 and the gate inside the op stays dynamic. + self._moe_merge_enabled = get_pcp_world_size() > 1 and bool( + get_current_atom_config().parallel_config.moe_pcp_merge + ) self.attn_norm = RMSNorm(args.dim, self.norm_eps) self.ffn_norm = RMSNorm(args.dim, self.norm_eps) self.hc_mult = hc_mult = args.hc_mult @@ -2669,24 +2698,24 @@ def forward( self.norm_eps, ) x = hc_state.x_prev - # PCP MoE mode B: all-gather this rank's 1/W token shard back to the - # full (padded) sequence so MoE sees every token (MoE itself is - # untouched / PCP-agnostic), then slice the result back to 1/W to match - # the residual stream (which stays sharded throughout). Mode A leaves x - # at 1/W and MoE runs per-shard with no extra comm. Only `x` is - # gathered/sliced — fuse_hc already produced the 1/W ffn input, and the - # residual in hc_state remains 1/W. The gather length is the fixed - # padded total (not data-dependent), so it is torch.compile-safe like - # the K/V all-gather inside attention.forward_impl. - moe_merge = _moe_pcp_merge_active() - if moe_merge: - pcp_ws = get_pcp_world_size() - x = pcp_allgather_rerange(x, pcp_ws) # [1/W, dim] -> [full, dim] - x = self.ffn( - x - ) # [num_tokens, dim] (input_ids read from forward_context for hash MoE) - if moe_merge: - x = pcp_split_stripe(x, pcp_ws) # [full, dim] -> [1/W, dim] + # PCP moe_pcp_merge (scheme B): MoE weights are sharded W*tp ways (pcp + # folded into tp/ep at build time), so each rank holds 1/(W*tp) of the + # experts and must see the full token set. 阶段19 (方向 C): the gate + + # merge collectives + MoE all run inside the OPAQUE moe_pcp_merge_forward + # custom op. Being a Dynamo barrier, the gate is read eagerly each call + # and never specialized/baked (the failure mode that killed the in-graph + # gate, 阶段16-18). The op is shape-preserving ([1/W,dim] in/out: + # prefill gathers->full, runs MoE, reduce_scatters->1/W; decode runs MoE + # then pcp all_reduce). `self._moe_merge_enabled` is a static config bool + # (is_dummy-independent) so Dynamo specializes it identically at warmup + # and real -> the op call is baked into code0; the runtime gate lives + # inside the op. Mode A / non-merge keeps the plain self.ffn path. + if self._moe_merge_enabled: + x = torch.ops.aiter.moe_pcp_merge_forward(x, self.ffn.prefix) + else: + x = self.ffn( + x + ) # [num_tokens, dim] (input_ids from forward_context for hash MoE) hc_state.x_prev = x return hc_state @@ -2757,6 +2786,80 @@ def forward( return self.get_logits(norm(x)) # [bs, vocab] +# ===== 阶段19 (方向 C): PCP MoE merge as an OPAQUE custom op ===== +# The whole "gate + merge collectives + MoE" runs inside this custom op. Being +# a Dynamo barrier, the gate (read from forward_context) is evaluated EAGERLY on +# every invocation and is NEVER specialized/baked — which is what defeated the +# in-graph gate (阶段16-18: torch.compile traced the gate to its warmup value, +# dispatch_to_code(0) then froze it). The op is shape-preserving +# ([n_local, dim] in/out: prefill gathers 1/W->full, runs MoE, reduce_scatters +# full->1/W; decode runs MoE then pcp all_reduce), so the fake impl is +# empty_like and Dynamo needs no dynamic-shape reasoning. prefill/decode split +# on the REAL is_prefill read inside the op: decode's cudagraph capture (阶段18: +# is_prefill=False, is_dummy=False) bakes the all_reduce kernel into the decode +# graph; prefill (no cudagraph) runs the gather/scatter eagerly each call. The +# gate KEEPS is_dummy_run (via _moe_pcp_merge_active/_decode_active) so it stays +# consistent with the out-of-graph split gate _pcp_active: at warmup +# (is_dummy=True) it neither splits nor merges. Opacity — NOT removing is_dummy +# — is what avoids the specialization that killed the in-graph gate (阶段16-18). + + +def _run_moe(moe: "MoE", x: torch.Tensor) -> torch.Tensor: + """Replicate MoE.forward's dual/single dispatch (we are already inside an + opaque op, so call the underlying methods directly rather than re-entering + the maybe_dual_stream_forward custom op).""" + threshold = envs.ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD + num_tokens = x.shape[0] + if moe._use_dual_stream and 0 < num_tokens <= threshold: + return moe.dual_stream_moe_forward(x) + return moe.single_stream_moe_forward(x) + + +def moe_pcp_merge_forward( + hidden_states: torch.Tensor, # [n_local, dim] + layer_name: str, +) -> torch.Tensor: # [n_local, dim] (shape-preserving) + moe = get_current_atom_config().compilation_config.static_forward_context[ + layer_name + ] + # Gate is read EAGERLY here (op is opaque to Dynamo), so it is NEVER baked — + # keeping is_dummy_run in the gate is correct and NECESSARY: it keeps this op + # consistent with the out-of-graph split gate `_pcp_active()` (also has + # is_dummy). At warmup (is_dummy=True) neither splits nor merges, so x stays + # full and the hash-MoE input_ids (also un-split at warmup) match. Removing + # is_dummy here would merge a never-split warmup batch -> input_ids/hidden + # length mismatch -> crash. (The in-graph gate needed is_dummy *removed* to + # bake True into code0; the opaque op does NOT — opacity is the fix.) + do_prefill = _moe_pcp_merge_active() + do_decode = _moe_pcp_merge_decode_active() + ws = get_pcp_world_size() + x = hidden_states + if do_prefill: + x = pcp_allgather_rankmajor(x, ws) # [1/W,dim] -> [full,dim] + x = _run_moe(moe, x) + if do_prefill: + x = pcp_reduce_scatter(x, ws) # [full,dim] -> [1/W,dim] + elif do_decode: + x = pcp_all_reduce(x) # sum pcp half (decode tokens are pcp-redundant) + return x + + +def _moe_pcp_merge_forward_fake( + hidden_states: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +direct_register_custom_op( + op_name="moe_pcp_merge_forward", + op_func=moe_pcp_merge_forward, + mutates_args=(), + fake_impl=_moe_pcp_merge_forward_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + def _pcp_active() -> bool: """Whether to apply PCP round-robin-split in this forward. @@ -2783,6 +2886,24 @@ def _moe_pcp_merge_active() -> bool: return bool(get_current_atom_config().parallel_config.moe_pcp_merge) +def _moe_pcp_merge_decode_active() -> bool: + """moe_pcp_merge DECODE path: MoE weights are sharded W*tp ways (pcp folded + into tp at build time), but decode does NOT stripe-split tokens — every pcp + rank holds the same full batch. So decode skips gather/reduce_scatter and + instead does one pcp all_reduce after MoE to sum the pcp-half of the + intermediate that combine_outputs' tp all_reduce misses. + + True only when pcp>1, moe_pcp_merge set, and this is a real DECODE forward + (not prefill — that's `_moe_pcp_merge_active` — and not dummy/warmup). + """ + if get_pcp_world_size() <= 1: + return False + if not bool(get_current_atom_config().parallel_config.moe_pcp_merge): + return False + fc = get_forward_context() + return (not fc.context.is_prefill) and (not fc.context.is_dummy_run) + + @support_torch_compile class DeepseekV4Model(nn.Module): """Full model: embed -> expand to hc_mult copies -> N blocks -> hc_head -> logits. @@ -2871,7 +2992,7 @@ def forward( hc_state = HCState(residual=h, post_mix=None, comb_mix=None, x_prev=None) for layer in self.layers: - hc_state = layer(hc_state, positions) # [num_tokens, hc, dim] + hc_state = layer(hc_state, positions) h = self.layers[-1].hc_post( hc_state.x_prev, hc_state.residual, hc_state.post_mix, hc_state.comb_mix ) @@ -3004,6 +3125,10 @@ def forward( # runs entirely on 1/W; the final hidden is all-gathered + un-padded # after self.model(...) returns. use_pcp = _pcp_active() + # NOTE: moe_merge here is the OUT-OF-GRAPH gate, used only for the + # input_ids gather below (round-robin split + hash-MoE id alignment). + # The MoE merge collectives gate themselves separately inside the opaque + # moe_pcp_merge_forward custom op (called from Block.forward, 阶段19). moe_merge = _moe_pcp_merge_active() # Mode B is incompatible with DP-attention id-gathering for now: both # rewrite ctx.context.input_ids with different (full-PCP vs DP-gathered) @@ -3020,13 +3145,16 @@ def forward( if pad > 0: input_ids = torch.cat([input_ids, input_ids.new_zeros(pad)], dim=0) positions = torch.cat([positions, positions.new_zeros(pad)], dim=0) - # Mode B: stash the full (padded, global-order) ids for hash MoE, - # which sees full tokens after the per-layer all-gather. The token - # order matches pcp_allgather_rerange's reconstruction in Block. - if moe_merge: - full_padded_ids = input_ids input_ids = pcp_round_robin_split(input_ids, pcp_size) positions = pcp_round_robin_split(positions, pcp_size) + # Scheme B: each Block rank-major all-gathers hidden before MoE + # (pcp_allgather_rankmajor = plain all_gather(dim=0), RANK-MAJOR + # order: [rank0's stripe | rank1's stripe | ...]). The hash MoE + # indexes input_ids per hidden row, so the ids must be gathered in + # the SAME rank-major order. pcp_allgather_rankmajor is int-safe + # (plain all_gather), so reuse it for the ids too. + if moe_merge: + full_padded_ids = pcp_allgather_rankmajor(input_ids, pcp_size) if self._need_ids_gather: # DP-attention (no EP) hash routing: input_ids is local but the MoE @@ -3041,6 +3169,11 @@ def forward( # (measured). The ids tensor is [N,1] int (tiny vs hidden [N,7168]), # so inline costs ~nothing in overlap. ctx.context.input_ids = MoE._gather_ids_for_dp(input_ids.flatten(), ctx) + elif moe_merge: + # Scheme B: hash MoE runs on the rank-major full token set (each + # Block rank-major all-gathers hidden before MoE), so the ids it + # indexes must be the matching rank-major full ids. + ctx.context.input_ids = full_padded_ids else: ctx.context.input_ids = input_ids h = self.model(input_ids, positions) From feac9bce5800b6e6efe01248cbeb52b8c4c21a6e Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Tue, 30 Jun 2026 19:40:49 +0800 Subject: [PATCH 03/17] feat(pcp+tbo): implement PCP+TBO Option A (prefill, divisible lengths) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0: Relax PCP+TBO raise in llm_engine.py — only block TBO-decode+PCP; prefill-only TBO now allowed. P1 (coordinated split, b1 approach): - pcp_utils.py: add num_ubatches param to pcp_pad_len (N·pcp alignment) - model_runner.py: _preprocess overrides TBO eligibility with 1/pcp local token count; Option A gate (n_prefill % 2*pcp == 0) avoids dummy tokens; prepare_inputs builds ubatch_slices in 1/pcp space with correct request boundaries; run_model does PCP split before UBatchWrapper, updates context.positions and forward_vars["positions"/"cu_seqlens_q"]; out-of- graph pcp_allgather_rerange after UBatchWrapper. - ubatch_wrapper.py: _make_ubatch_context carries input_ids per-ubatch slice for hash-MoE (moe_pcp_merge mode B). - deepseek_v4.py: ForCausalLM.forward skips PCP split and final pcp_allgather_rerange when tbo_active(); per-ubatch input_ids allgather for hash-MoE. - deepseek_v4_attn.py: build_ubatch_prefill_metadata uses real_num_tokens for _attach_v4_per_fwd_meta; Option B (arbitrary length) planned as per-ubatch reindex — see PCP_TBO.md §11. P2: moe_pcp_merge_forward inserts TBO yield/switch around pcp_allgather_rankmajor and pcp_reduce_scatter. Attention-internal PCP collectives (compressor kv_score, kv, positions allgather) wrapped with tbo_yield_and_switch / tbo_switch_to_compute_sync to serialize RCCL submissions across TBO ubatch threads. Add ATOM_TBO_DEBUG env var and per-collective debug logging in pcp_utils.py for locating RCCL hang root cause (see PCP_TBO.md §9.5). Verified: warmup ✓, gsm8k=0.86 (Option A mixed path), server stable. Option B (arbitrary length) deferred — requires per-ubatch reindex. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- atom/distributed/pcp_utils.py | 58 ++++++- atom/model_engine/llm_engine.py | 22 +-- atom/model_engine/model_runner.py | 162 +++++++++++++++++- atom/model_ops/attentions/deepseek_v4_attn.py | 17 +- atom/models/deepseek_v4.py | 113 +++++++++--- atom/utils/envs.py | 5 + atom/utils/tbo/ubatch_wrapper.py | 9 + 7 files changed, 334 insertions(+), 52 deletions(-) diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 6db969b29c..32140180b3 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -13,6 +13,8 @@ `layers/utils/cp_utils.py:cp_all_gather_rerange_output`). """ +import itertools +import logging from typing import Optional import torch @@ -23,6 +25,43 @@ get_prefill_context_model_parallel_world_size, ) +logger = logging.getLogger("atom") + +# Monotonic counter so debug lines can be matched against RCCL WorkNCCL SeqNum +# when locating a PCP+TBO collective hang. GIL-protected; shared across the 2 +# TBO ubatch threads (the issue order across threads is exactly what we trace). +_pcp_comm_seq = itertools.count() + + +def _tbo_debug_collective(fn: str, input_: torch.Tensor) -> None: + """Emit a WARNING-level trace for one PCP collective when ATOM_TBO_DEBUG=1. + + Logs (in issue order) the call site, input shape/numel, the current TBO + ubatch id, and the PCP rank — so the last line before an RCCL allgather + timeout pinpoints which collective (attention KV/compressor/indexer vs MoE + hidden) deadlocked, and on which ubatch. + """ + from atom.utils import envs + + if not envs.ATOM_TBO_DEBUG: + return + try: + from atom.utils.tbo.ubatching import tbo_active, tbo_current_ubatch_id + + ub = tbo_current_ubatch_id() if tbo_active() else -1 + except Exception: + ub = -1 + seq = next(_pcp_comm_seq) + logger.warning( + "[TBO_DEBUG] seq=%d %s in_shape=%s numel=%d ubatch=%d pcp_rank=%d", + seq, + fn, + tuple(input_.shape), + input_.numel(), + ub, + get_prefill_context_model_parallel_rank(), + ) + def get_pcp_world_size() -> int: return get_prefill_context_model_parallel_world_size() @@ -39,22 +78,29 @@ def pcp_is_enabled() -> bool: def pcp_pad_len( total_tokens: int, pcp_size: Optional[int] = None, + num_ubatches: 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 * num_ubatches. 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. + + When TBO is active with N ubatches, pass num_ubatches=N so each per-rank + shard (total // pcp_size) is also divisible by N — guaranteeing every + ubatch is a pcp_size multiple and the token-midpoint split lands on an + exact pcp_size boundary. """ if pcp_size is None: pcp_size = get_pcp_world_size() - if pcp_size <= 1: + divisor = pcp_size * max(num_ubatches, 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( @@ -96,6 +142,7 @@ def pcp_allgather_rerange( pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ + _tbo_debug_collective("pcp_allgather_rerange", input_) group = get_pcp_group() # aiter all_gather(dim=0) returns rank-major concat: [pcp*L, *rest]. gathered = group.all_gather(input_.contiguous(), dim=0) @@ -138,6 +185,7 @@ def pcp_allgather_rankmajor( pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ + _tbo_debug_collective("pcp_allgather_rankmajor", input_) return get_pcp_group().all_gather(input_.contiguous(), dim=0) @@ -150,6 +198,7 @@ def pcp_reduce_scatter( pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ + _tbo_debug_collective("pcp_reduce_scatter", input_) return get_pcp_group().reduce_scatter(input_.contiguous(), dim=0) @@ -163,6 +212,7 @@ def pcp_all_reduce(input_: torch.Tensor, pcp_size: Optional[int] = None) -> torc pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ + _tbo_debug_collective("pcp_all_reduce", input_) return get_pcp_group().all_reduce(input_) diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index 7d41223f83..bb04a9cacb 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -79,20 +79,16 @@ 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 (P1). + # 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." - ) + if 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." + ) self.rquest_ids = set() self.io_processor = InputOutputProcessor( config, self.tokenizer, config.kv_cache_block_size diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 582926f306..b2c6b8ca05 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -38,11 +38,18 @@ from atom.kv_transfer.disaggregation import KVConnectorOutput from atom.utils.forward_context import get_kvconnector from atom.utils.tbo import ( + UBatchSlice, UBatchWrapper, local_tbo_precompute, maybe_create_ubatch_slices, sync_dp_for_tbo, ) +from atom.distributed.pcp_utils import ( + pcp_allgather_rankmajor, + pcp_allgather_rerange, + pcp_pad_len, + pcp_round_robin_split, +) from atom.utils.forward_context import ( Context, DPMetadata, @@ -1733,6 +1740,27 @@ def _preprocess( self.config, batch, is_prefill, num_scheduled_tokens ) + # PCP+TBO prefill (障碍 A, b1): the PCP split happens in run_model, + # before UBatchWrapper. Each ubatch sees 1/(2*pcp) tokens. Override the + # local TBO eligibility/sizes with the local (1/pcp) token count so + # ub_slices are created in 1/pcp space rather than the global space. + # + # 选项 A（务实先验证）：仅当 total_prefill_tokens 被 2*pcp 整除时启用 TBO。 + # 此时各 PCP rank 真实 token 数相等、每 ubatch = local//2 跨 rank 一致、 + # 无 dummy padding token → 回避跨 rank allgather size mismatch（Bug5）。 + # 不整除时 local_eligible=False → tbo_collective_active=False → 回退非 TBO。 + # 选项 B（任意长度）需 per-ubatch reindex，见 PCP_TBO.md §11。 + pcp_size = self.config.prefill_context_parallel_size + if tbo_on and is_prefill and pcp_size > 1 and not batch.is_dummy_run: + n_prefill = batch.total_tokens_num_prefill + if n_prefill % (2 * pcp_size) == 0: + local_tokens = n_prefill // pcp_size + local_eligible = local_tokens >= 2 + local_ub0 = local_tokens // 2 + local_ub1 = local_tokens - local_ub0 + else: + local_eligible = False + if dp_size <= 1: # Single-rank: TBO decision is purely local; no collective needed. # dp_uniform_decode=True mirrors the DP-disabled case in the @@ -1853,14 +1881,87 @@ def prepare_inputs(self, batch: ScheduledBatch, input_ids: torch.Tensor = None): input_ids, ) - ubatch_slices = self._maybe_create_tbo_slices( - batch, - is_prefill, - scheduled_bs if not is_prefill else 0, - actual_num_tokens, - num_scheduled_tokens, - tbo_collective_active, - ) + pcp_size = self.config.prefill_context_parallel_size + _pcp_tbo_prefill = ( + is_prefill and pcp_size > 1 and tbo_collective_active + and not batch.is_dummy_run + ) + if _pcp_tbo_prefill: + # Compute padded local token count and local cu_seqlens_q first; + # both are needed to determine correct per-ubatch request boundaries. + padded_total = pcp_pad_len(batch.total_tokens_num_prefill, pcp_size, num_ubatches=2) + local_tokens = padded_total // pcp_size + half = local_tokens // 2 + num_prefill_reqs = batch.total_seqs_num_prefill + + # Local cu_seqlens_q: derived from attn_metadata.batch_id_per_token + # which _apply_pcp_reindex already reduced to 1/pcp space. + # build_ubatch_prefill_metadata reads forward_vars["cu_seqlens_q"] (set + # to global values in prepare_prefill); overwrite with local values now. + bid = attn_metadata.batch_id_per_token + valid_bid = bid[bid >= 0] + if valid_bid.numel() > 0: + local_counts_np = ( + torch.bincount(valid_bid.to(torch.int64), minlength=num_prefill_reqs) + .cpu() + .numpy() + ) + else: + local_counts_np = np.zeros(num_prefill_reqs, dtype=np.int64) + local_cu = np.zeros(num_prefill_reqs + 1, dtype=np.int32) + local_cu[1:] = local_counts_np.cumsum() + self.forward_vars["cu_seqlens_q"].np[: num_prefill_reqs + 1] = local_cu + + # token_slice boundary MUST use local_tokens//2 (padded, cross-rank + # consistent via N·pcp pad), NOT real_local//2 (per-rank real count). + # real_local differs across ranks when total is not divisible by pcp + # (e.g. 8193 → rank0=4097, rank1=4096), so using it for token_slice + # makes ubatch1 sizes diverge across ranks → RCCL allgather hang (Bug5). + # Dummy tokens in [real_local:local_tokens] are handled by real_num_tokens + # inside build_ubatch_prefill_metadata (_attach_v4_per_fwd_meta receives + # real_num_tokens = sum(extend_lens) instead of ub_num_tokens). + + # Determine per-ubatch request boundaries from local_cu, using padded + # half as the split point so both ranks see the same boundary. + req_stop_ub0 = int(np.searchsorted(local_cu[:num_prefill_reqs], half, side='left')) + req_stop_ub0 = max(1, min(req_stop_ub0, num_prefill_reqs)) + req_start_ub1 = int(np.searchsorted(local_cu[1:num_prefill_reqs + 1], half + 1, side='left')) + req_start_ub1 = max(0, min(req_start_ub1, num_prefill_reqs - 1)) + + ubatch_slices = [ + UBatchSlice( + request_slice=slice(0, req_stop_ub0), + token_slice=slice(0, half), # padded half, cross-rank consistent + ), + UBatchSlice( + request_slice=slice(req_start_ub1, num_prefill_reqs), + token_slice=slice(half, local_tokens), # padded boundary + ), + ] + + # Local positions: round-robin split of the global positions tensor. + # build_ubatch_prefill_metadata reads forward_vars["positions"].np/gpu; + # overwrite with local (1/pcp) values so per-ubatch indexing is correct. + n_global_pos = positions.shape[0] + pos_padded = positions + if padded_total > n_global_pos: + pos_padded = torch.cat( + [positions, positions.new_zeros(padded_total - n_global_pos)] + ) + local_positions = pcp_round_robin_split(pos_padded, pcp_size) + self.forward_vars["positions"].np[:local_tokens] = ( + local_positions.cpu().numpy() + ) + self.forward_vars["positions"].copy_to_gpu(local_tokens) + else: + ubatch_slices = self._maybe_create_tbo_slices( + batch, + is_prefill, + scheduled_bs if not is_prefill else 0, + actual_num_tokens, + num_scheduled_tokens, + tbo_collective_active, + ) set_forward_context( attn_metadata=attn_metadata, @@ -1965,6 +2066,38 @@ def run_model( # internally short-circuits for prefill / cudagraph. forward_mode.assert_shape_contract(input_ids, forward_context.attn_metadata) + # PCP+TBO prefill (障碍 A, b1): apply PCP split here, before UBatchWrapper, + # so each ubatch receives 1/(2*pcp) tokens and ub_slices (in 1/pcp space) + # align with the metadata already reindexed by _apply_pcp_reindex. + _pcp_size = self.config.prefill_context_parallel_size + _pcp_tbo_prefill = ( + _pcp_size > 1 + and isinstance(self.model, UBatchWrapper) + and forward_context.ubatch_slices is not None + and is_prefill + and not forward_context.context.is_dummy_run + ) + if _pcp_tbo_prefill: + n_global = input_ids.shape[0] + _padded = pcp_pad_len(n_global, _pcp_size, num_ubatches=2) + if _padded > n_global: + _pad = _padded - n_global + input_ids = torch.cat([input_ids, input_ids.new_zeros(_pad)]) + positions = torch.cat([positions, positions.new_zeros(_pad)]) + input_ids = pcp_round_robin_split(input_ids, _pcp_size) + positions = pcp_round_robin_split(positions, _pcp_size) + # Update context.positions to the local (1/pcp) version so that + # _make_ubatch_context slices it correctly in 1/pcp space. + forward_context.context.positions = positions + # Store local (1/pcp) input_ids for hash MoE (mode B, moe_pcp_merge). + # _make_ubatch_context will slice [token_slice] per ubatch; each + # ForCausalLM.forward then allgathers its per-ubatch slice across PCP + # ranks, giving padded_total//2 ids — matching what moe_pcp_merge_forward + # allgathers for hidden states. Pre-allgathering here would give + # padded_total ids, which mismatches MoE's padded_total//2 num_tokens. + if self.config.parallel_config.moe_pcp_merge: + forward_context.context.input_ids = input_ids # local, 1/pcp tokens + if not forward_mode.use_cudagraph: # prefill, or decode forced eager (enforce_eager / DP peer # prefill / bs above the largest captured graph). @@ -2014,6 +2147,19 @@ def run_model( model_output = self.model( input_ids, positions, inputs_embeds=inputs_embeds ) + # PCP+TBO prefill: UBatchWrapper has concatenated the two ubatch + # shards (each 1/(2*pcp) tokens); do the single pcp_allgather_rerange + # here and crop to n_global. ForCausalLM.forward skipped this step + # to avoid 2 concurrent RCCL calls from the two ubatch threads. + if _pcp_tbo_prefill: + if self.use_aux_hidden_state_outputs: + _h, _aux = model_output + _h = pcp_allgather_rerange(_h, _pcp_size)[:n_global] + model_output = (_h, _aux) + else: + model_output = pcp_allgather_rerange( + model_output, _pcp_size + )[:n_global] if self.use_aux_hidden_state_outputs: hidden_states, self._aux_hidden_states = model_output else: diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index d5d90639a5..5a02bf64cb 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1616,6 +1616,12 @@ def build_ubatch_prefill_metadata( ts = ub_slice.token_slice ub_num_reqs = rs.stop - rs.start ub_num_tokens = ts.stop - ts.start + # ub_num_tokens above counts the full token_slice width (= padded size + # when PCP+TBO inserts dummy tokens to keep cross-rank alignment). + # real_num_tokens is computed below from sum(extend_lens_np) after + # clamping, which excludes dummy-padding rows (extend_lens==0 for them). + # Passed as `total_tokens` to _attach_v4_per_fwd_meta so that + # batch_id_unpadded_np[:total_tokens] = ... doesn't broadcast-error. if src.state_slot_mapping is not None: ub_attn.state_slot_mapping = src.state_slot_mapping[rs] @@ -1632,6 +1638,10 @@ def build_ubatch_prefill_metadata( extend_lens_np = (clamped_ends - clamped_starts).astype(np.int32) ub_cu = np.zeros(ub_num_reqs + 1, dtype=np.int32) np.cumsum(extend_lens_np, dtype=np.int32, out=ub_cu[1:]) + # Real (non-dummy) token count = sum of clamped extend lengths. + # Equals ub_num_tokens in standard TBO; may be smaller when PCP+TBO + # pads token_slice to a pcp-aligned boundary (Option B). + real_num_tokens = int(ub_cu[ub_num_reqs]) ub_start_pos_for_ctx = positions_np[ub_cu[:ub_num_reqs]].astype(np.int32) context_lens_np = (ub_start_pos_for_ctx + extend_lens_np).astype(np.int32) from atom.model_ops.v4_kernels import make_compress_plans @@ -1669,9 +1679,14 @@ def build_ubatch_prefill_metadata( extend_lens_np, # ubatch's per-seq token counts ub_attn.state_slot_mapping_cpu, ub_num_reqs, - ub_num_tokens, + real_num_tokens, # actual tokens only (excludes PCP dummy padding) ) + # Option A: real_num_tokens == ub_num_tokens (no dummy, since + # _preprocess gates on n_prefill % (2*pcp)==0). All downstream + # functions receive ub_num_tokens consistently. + # Option B (arbitrary length, with dummy tokens) requires per-ubatch + # reindex — see PCP_TBO.md §11 for the full analysis and remaining work. positions_gpu = var["positions"].gpu[ts.start : ts.stop] self._attach_v4_indexer_meta( ub_attn, diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index fb273f7b5b..491e067b09 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1006,7 +1006,17 @@ def forward( # state_slot_mapping passed to fused_compress_attn are full-sequence # (never split in the builder), so they match the gathered `combined`. if _pcp_active(): + from atom.utils.tbo.ubatching import ( + tbo_active as _tbo_active, + tbo_yield_and_switch_from_compute_to_comm, + tbo_switch_to_compute_sync, + ) + _tbo = _tbo_active() + if _tbo: + tbo_yield_and_switch_from_compute_to_comm() combined = pcp_allgather_rerange(combined, get_pcp_world_size()) + if _tbo: + tbo_switch_to_compute_sync() # TBO decode: copy `combined` into a fixed-address buffer so CUDAGraph # capture/replay see a stable pointer (allocator may re-place it). from atom.utils.tbo.ubatching import tbo_active, tbo_current_ubatch_id @@ -1927,6 +1937,14 @@ def forward_impl( pcp_on = _pcp_active() if pcp_on: pcp_ws = get_pcp_world_size() + from atom.utils.tbo.ubatching import ( + tbo_active as _tbo_active_attn, + tbo_yield_and_switch_from_compute_to_comm, + tbo_switch_to_compute_sync, + ) + _tbo_attn = _tbo_active_attn() + if _tbo_attn: + tbo_yield_and_switch_from_compute_to_comm() kv_full = pcp_allgather_rerange(kv, pcp_ws) # positions must match kv_full's full-sequence coords for the # swa_write ring addressing (`positions[src] % cache_size`). @@ -1935,6 +1953,8 @@ def forward_impl( # the same rerange used for kv (NOT fc.context.positions, which # the builder reindexed to 1/W). positions_full = pcp_allgather_rerange(positions, pcp_ws) + if _tbo_attn: + tbo_switch_to_compute_sync() else: kv_full = kv positions_full = positions @@ -2159,10 +2179,12 @@ class MoE(nn.Module): `FusedMoE.select_experts(scoring_func="sqrtsoftplus", e_score_correction_bias=...)`, which we extended in atom/model_ops/moe.py to add the V4 path. - Hash routing for `layer_id < n_hash_layers` (first 3 V4 layers) is NOT yet - wired through FusedMoE — the `tid2eid` buffer is declared so weight loading - completes, but inference uses the standard sqrtsoftplus path. Hash layers - will produce incorrect routing; correct hash routing lands in PR3+. + Hash routing for `layer_id < n_hash_layers` (first 3 V4 layers) is wired + through FusedMoE via the `custom_routing_function` hook: hash layers load a + `tid2eid` table (token-id -> expert-id) instead of `gate.bias`, and + `select_experts` gives `custom_routing_function` precedence over the + standard sqrtsoftplus path. Expert *selection* comes from `tid2eid[input_ids]` + while expert *weights* still use sqrtsoftplus(gate_logits). Accuracy verified. """ def __init__( @@ -2835,10 +2857,27 @@ def moe_pcp_merge_forward( ws = get_pcp_world_size() x = hidden_states if do_prefill: + from atom.utils.tbo.ubatching import ( + tbo_active as _tbo_active, + tbo_yield_and_switch_from_compute_to_comm, + tbo_switch_to_compute_sync, + ) + _tbo = _tbo_active() + # allgather: when TBO is active, yield to the partner ubatch so its + # compute overlaps this collective on the comm stream (P2 overlap). + if _tbo: + tbo_yield_and_switch_from_compute_to_comm() x = pcp_allgather_rankmajor(x, ws) # [1/W,dim] -> [full,dim] + if _tbo: + tbo_switch_to_compute_sync() x = _run_moe(moe, x) if do_prefill: + # reduce_scatter: same yield pattern. + if _tbo: + tbo_yield_and_switch_from_compute_to_comm() x = pcp_reduce_scatter(x, ws) # [full,dim] -> [1/W,dim] + if _tbo: + tbo_switch_to_compute_sync() elif do_decode: x = pcp_all_reduce(x) # sum pcp half (decode tokens are pcp-redundant) return x @@ -3139,22 +3178,30 @@ def forward( ) full_padded_ids = None if use_pcp: + from atom.utils.tbo.ubatching import tbo_active as _tbo_active pcp_size = get_pcp_world_size() - n_global = input_ids.shape[0] - pad = pcp_pad_len(n_global, pcp_size) - n_global - if pad > 0: - input_ids = torch.cat([input_ids, input_ids.new_zeros(pad)], dim=0) - positions = torch.cat([positions, positions.new_zeros(pad)], dim=0) - input_ids = pcp_round_robin_split(input_ids, pcp_size) - positions = pcp_round_robin_split(positions, pcp_size) - # Scheme B: each Block rank-major all-gathers hidden before MoE - # (pcp_allgather_rankmajor = plain all_gather(dim=0), RANK-MAJOR - # order: [rank0's stripe | rank1's stripe | ...]). The hash MoE - # indexes input_ids per hidden row, so the ids must be gathered in - # the SAME rank-major order. pcp_allgather_rankmajor is int-safe - # (plain all_gather), so reuse it for the ids too. - if moe_merge: - full_padded_ids = pcp_allgather_rankmajor(input_ids, pcp_size) + if _tbo_active(): + # PCP+TBO prefill (障碍 A, b1): the split was done upstream in + # run_model; input_ids is already 1/(2*pcp) tokens. full_padded_ids + # was precomputed there and stored in ctx.context.input_ids (carried + # into this ubatch context by _make_ubatch_context). Nothing to do. + pass + else: + n_global = input_ids.shape[0] + pad = pcp_pad_len(n_global, pcp_size) - n_global + if pad > 0: + input_ids = torch.cat([input_ids, input_ids.new_zeros(pad)], dim=0) + positions = torch.cat([positions, positions.new_zeros(pad)], dim=0) + input_ids = pcp_round_robin_split(input_ids, pcp_size) + positions = pcp_round_robin_split(positions, pcp_size) + # Scheme B: each Block rank-major all-gathers hidden before MoE + # (pcp_allgather_rankmajor = plain all_gather(dim=0), RANK-MAJOR + # order: [rank0's stripe | rank1's stripe | ...]). The hash MoE + # indexes input_ids per hidden row, so the ids must be gathered in + # the SAME rank-major order. pcp_allgather_rankmajor is int-safe + # (plain all_gather), so reuse it for the ids too. + if moe_merge: + full_padded_ids = pcp_allgather_rankmajor(input_ids, pcp_size) if self._need_ids_gather: # DP-attention (no EP) hash routing: input_ids is local but the MoE @@ -3170,19 +3217,33 @@ def forward( # so inline costs ~nothing in overlap. ctx.context.input_ids = MoE._gather_ids_for_dp(input_ids.flatten(), ctx) elif moe_merge: - # Scheme B: hash MoE runs on the rank-major full token set (each - # Block rank-major all-gathers hidden before MoE), so the ids it - # indexes must be the matching rank-major full ids. - ctx.context.input_ids = full_padded_ids + if full_padded_ids is not None: + # Normal PCP path: set ids from the just-computed all-gather. + ctx.context.input_ids = full_padded_ids + else: + # PCP+TBO path: input_ids is the per-ubatch slice of local ids + # (set by _make_ubatch_context from run_model's local ids). + # Allgather across PCP ranks to get padded_total//2 ids, + # matching what moe_pcp_merge_forward allgathers for hidden states. + # Both PCP ranks call this at the same TBO ubatch phase → synchronized. + ctx.context.input_ids = pcp_allgather_rankmajor( + input_ids.flatten(), pcp_size + ) else: ctx.context.input_ids = input_ids h = self.model(input_ids, positions) # ----- PCP: all-gather shards, restore original order, drop pad ----- if use_pcp: - h = pcp_allgather_rerange(h, pcp_size) - if pad > 0: - h = h[:n_global] + if _tbo_active(): + # PCP+TBO: skip the all-gather here. Each ubatch returns its + # local (1/(2*pcp)) shard; run_model does a single + # pcp_allgather_rerange after UBatchWrapper cats both shards. + pass + else: + h = pcp_allgather_rerange(h, pcp_size) + if pad > 0: + h = h[:n_global] return h def compute_logits( diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 3bd2a53abb..c5fa7aee75 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -237,6 +237,11 @@ "ATOM_TBO_PREFILL_TOKEN_SPLIT": lambda: ( os.getenv("ATOM_TBO_PREFILL_TOKEN_SPLIT", "1") == "1" ), + # Debug logging for TBO (+ PCP) collectives and ubatch splitting. Default off. + # When "1", PCP collective sites and TBO ubatch decisions emit WARNING logs + # (function name, input shape, ubatch id) to trace cross-rank collective order + # — used to locate RCCL allgather hangs under PCP+TBO. + "ATOM_TBO_DEBUG": lambda: os.getenv("ATOM_TBO_DEBUG", "0") == "1", # --- NUMA binding --- # Master switch: pin each GPU worker to its GPU-local NUMA node's CPU cores # and preferred memory. Default off so baseline/pinned A/B stays clean. diff --git a/atom/utils/tbo/ubatch_wrapper.py b/atom/utils/tbo/ubatch_wrapper.py index ebc799865f..917789d183 100644 --- a/atom/utils/tbo/ubatch_wrapper.py +++ b/atom/utils/tbo/ubatch_wrapper.py @@ -506,6 +506,15 @@ def _make_ubatch_context( batch_size=ub_num_reqs, graph_bs=graph_bs, is_draft=ctx.context.is_draft, + # Carry over per-ubatch slice of input_ids for hash MoE (PCP+TBO mode). + # run_model stores local (1/pcp) ids; each ubatch takes its token_slice. + # ForCausalLM.forward then allgathers the slice to get ids matching the + # MoE's per-ubatch allgathered hidden states (padded_total//2 tokens). + input_ids=( + ctx.context.input_ids[ub_slice.token_slice] + if ctx.context.input_ids is not None + else None + ), ) return ForwardContext( From b85174468e31f2655f40eaa02b0c6060b278d1cd Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Wed, 1 Jul 2026 17:22:49 +0800 Subject: [PATCH 04/17] Enable PCP+TBO --- atom/distributed/pcp_utils.py | 19 +- atom/model_engine/model_runner.py | 236 ++++++++++-------- atom/model_ops/attentions/deepseek_v4_attn.py | 229 ++++++++++++++++- 3 files changed, 362 insertions(+), 122 deletions(-) diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 32140180b3..906a8ea2d3 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -15,10 +15,27 @@ import itertools import logging -from typing import Optional +from typing import NamedTuple, Optional import torch + +class PcpBalGroup(NamedTuple): + """One request group for PCP+TBO balanced prefill (PCP_TBO.md §12). + + 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) + from aiter.dist.parallel_state import ( get_pcp_group, get_prefill_context_model_parallel_rank, diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index b2c6b8ca05..c998d27c01 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -45,6 +45,7 @@ sync_dp_for_tbo, ) from atom.distributed.pcp_utils import ( + PcpBalGroup, pcp_allgather_rankmajor, pcp_allgather_rerange, pcp_pad_len, @@ -1740,26 +1741,32 @@ def _preprocess( self.config, batch, is_prefill, num_scheduled_tokens ) - # PCP+TBO prefill (障碍 A, b1): the PCP split happens in run_model, - # before UBatchWrapper. Each ubatch sees 1/(2*pcp) tokens. Override the - # local TBO eligibility/sizes with the local (1/pcp) token count so - # ub_slices are created in 1/pcp space rather than the global space. - # - # 选项 A（务实先验证）：仅当 total_prefill_tokens 被 2*pcp 整除时启用 TBO。 - # 此时各 PCP rank 真实 token 数相等、每 ubatch = local//2 跨 rank 一致、 - # 无 dummy padding token → 回避跨 rank allgather size mismatch（Bug5）。 - # 不整除时 local_eligible=False → tbo_collective_active=False → 回退非 TBO。 - # 选项 B（任意长度）需 per-ubatch reindex，见 PCP_TBO.md §11。 + # PCP+TBO prefill — b2 + balanced (PCP_TBO.md §12): split requests into + # two GROUPS at a request boundary (never split a sequence's tokens), so + # each ubatch = "non-TBO PCP on a request subset". This root-fixes the + # token-split R1/R2 structural conflicts (compressor cross-token + # compression / SWA ring), because every sequence stays whole within one + # group. Requires num_reqs >= 2 (balanced needs two non-empty groups); + # bs=1 falls back to non-TBO. pcp_size = self.config.prefill_context_parallel_size + # True for eligible PCP+TBO balanced prefill; read by build_ubatch / + # run_model / prepare_prefill to route the per-group path. + self._pcp_tbo_balanced_active = False + # Per-group descriptors; reset each step, set in prepare_inputs balanced + # branch. Guards run_model/build_ubatch against stale values. + self._pcp_bal_groups = None if tbo_on and is_prefill and pcp_size > 1 and not batch.is_dummy_run: + num_prefill_reqs = batch.total_seqs_num_prefill n_prefill = batch.total_tokens_num_prefill - if n_prefill % (2 * pcp_size) == 0: - local_tokens = n_prefill // pcp_size - local_eligible = local_tokens >= 2 - local_ub0 = local_tokens // 2 - local_ub1 = local_tokens - local_ub0 - else: - local_eligible = False + # Rough local sizing for TBO eligibility. PCP is always dp=1, so the + # dp_size<=1 fast path below returns local_eligible verbatim as + # tbo_collective_active; local_ub0/ub1 are only used by the dp>1 + # sync path (never hit under PCP). + local_tokens = n_prefill // pcp_size + local_eligible = num_prefill_reqs >= 2 and local_tokens >= 2 + local_ub0 = local_tokens // 2 + local_ub1 = local_tokens - local_ub0 + self._pcp_tbo_balanced_active = local_eligible if dp_size <= 1: # Single-rank: TBO decision is purely local; no collective needed. @@ -1882,77 +1889,52 @@ def prepare_inputs(self, batch: ScheduledBatch, input_ids: torch.Tensor = None): ) pcp_size = self.config.prefill_context_parallel_size - _pcp_tbo_prefill = ( + _pcp_tbo_balanced = ( is_prefill and pcp_size > 1 and tbo_collective_active and not batch.is_dummy_run - ) - if _pcp_tbo_prefill: - # Compute padded local token count and local cu_seqlens_q first; - # both are needed to determine correct per-ubatch request boundaries. - padded_total = pcp_pad_len(batch.total_tokens_num_prefill, pcp_size, num_ubatches=2) - local_tokens = padded_total // pcp_size - half = local_tokens // 2 + and getattr(self, "_pcp_tbo_balanced_active", False) + ) + if _pcp_tbo_balanced: + # b2 + balanced (PCP_TBO.md §12): split REQUESTS into two groups at a + # request boundary near the token midpoint. Each group is an + # independent "non-TBO PCP mini-batch": padded to a pcp multiple and + # round-robin striped as a whole, so every sequence stays intact in + # one group (root-fixes token-split R1/R2). forward_vars stay GLOBAL + # here; build_ubatch_prefill_metadata slices the FULL (un-reindexed) + # metadata per group and calls _apply_pcp_reindex on it. num_prefill_reqs = batch.total_seqs_num_prefill - - # Local cu_seqlens_q: derived from attn_metadata.batch_id_per_token - # which _apply_pcp_reindex already reduced to 1/pcp space. - # build_ubatch_prefill_metadata reads forward_vars["cu_seqlens_q"] (set - # to global values in prepare_prefill); overwrite with local values now. - bid = attn_metadata.batch_id_per_token - valid_bid = bid[bid >= 0] - if valid_bid.numel() > 0: - local_counts_np = ( - torch.bincount(valid_bid.to(torch.int64), minlength=num_prefill_reqs) - .cpu() - .numpy() - ) - else: - local_counts_np = np.zeros(num_prefill_reqs, dtype=np.int64) - local_cu = np.zeros(num_prefill_reqs + 1, dtype=np.int32) - local_cu[1:] = local_counts_np.cumsum() - self.forward_vars["cu_seqlens_q"].np[: num_prefill_reqs + 1] = local_cu - - # token_slice boundary MUST use local_tokens//2 (padded, cross-rank - # consistent via N·pcp pad), NOT real_local//2 (per-rank real count). - # real_local differs across ranks when total is not divisible by pcp - # (e.g. 8193 → rank0=4097, rank1=4096), so using it for token_slice - # makes ubatch1 sizes diverge across ranks → RCCL allgather hang (Bug5). - # Dummy tokens in [real_local:local_tokens] are handled by real_num_tokens - # inside build_ubatch_prefill_metadata (_attach_v4_per_fwd_meta receives - # real_num_tokens = sum(extend_lens) instead of ub_num_tokens). - - # Determine per-ubatch request boundaries from local_cu, using padded - # half as the split point so both ranks see the same boundary. - req_stop_ub0 = int(np.searchsorted(local_cu[:num_prefill_reqs], half, side='left')) - req_stop_ub0 = max(1, min(req_stop_ub0, num_prefill_reqs)) - req_start_ub1 = int(np.searchsorted(local_cu[1:num_prefill_reqs + 1], half + 1, side='left')) - req_start_ub1 = max(0, min(req_start_ub1, num_prefill_reqs - 1)) - + per_req = np.asarray( + num_scheduled_tokens[:num_prefill_reqs], dtype=np.int64 + ) + total_tok = int(per_req.sum()) + cum = np.cumsum(per_req) # cum[j] = sum of reqs [0..j] + target = total_tok // 2 + # request boundary whose cumulative token count is closest to target + split_idx = int(np.searchsorted(cum, target, side="left")) + 1 + split_idx = max(1, min(split_idx, num_prefill_reqs - 1)) + B = int(cum[split_idx - 1]) # global token count of group0 (reqs [0:split_idx]) + H0 = pcp_pad_len(B, pcp_size) + H1 = pcp_pad_len(total_tok - B, pcp_size) + l0 = H0 // pcp_size + l1 = H1 // pcp_size + # token_slice is in the LOCAL concat space [g0_local | g1_local] that + # run_model produces (see per-group stripe there). ubatch_slices = [ UBatchSlice( - request_slice=slice(0, req_stop_ub0), - token_slice=slice(0, half), # padded half, cross-rank consistent + request_slice=slice(0, split_idx), + token_slice=slice(0, l0), ), UBatchSlice( - request_slice=slice(req_start_ub1, num_prefill_reqs), - token_slice=slice(half, local_tokens), # padded boundary + request_slice=slice(split_idx, num_prefill_reqs), + token_slice=slice(l0, l0 + l1), ), ] - - # Local positions: round-robin split of the global positions tensor. - # build_ubatch_prefill_metadata reads forward_vars["positions"].np/gpu; - # overwrite with local (1/pcp) values so per-ubatch indexing is correct. - n_global_pos = positions.shape[0] - pos_padded = positions - if padded_total > n_global_pos: - pos_padded = torch.cat( - [positions, positions.new_zeros(padded_total - n_global_pos)] - ) - local_positions = pcp_round_robin_split(pos_padded, pcp_size) - self.forward_vars["positions"].np[:local_tokens] = ( - local_positions.cpu().numpy() - ) - self.forward_vars["positions"].copy_to_gpu(local_tokens) + # Per-group descriptors consumed by run_model (per-group stripe) and + # build_ubatch_prefill_metadata (slice+reindex). See PcpBalGroup. + self._pcp_bal_groups = [ + PcpBalGroup(0, split_idx, 0, B, H0), + PcpBalGroup(split_idx, num_prefill_reqs, B, total_tok, H1), + ] else: ubatch_slices = self._maybe_create_tbo_slices( batch, @@ -2041,6 +2023,28 @@ def prepare_model(self, batch: ScheduledBatch): needs_independent_noise, ) + def _restore_pcp_balanced_output( + self, + mo: torch.Tensor, + groups: "list[PcpBalGroup]", + pcp_size: int, + ) -> torch.Tensor: + """Restore PCP+TBO balanced output (PCP_TBO.md §12). + + UBatchWrapper concatenated the two groups' 1/pcp output shards + [g0_local | g1_local]. Each group was striped independently, so restore + per group: pcp_allgather_rerange its shard back to the group's global + order, crop the per-group pad, then concat to the full global sequence. + """ + outs = [] + off = 0 + for grp in groups: + local_len = grp.pad_total // pcp_size # group's 1/pcp token count + seg = pcp_allgather_rerange(mo[off : off + local_len], pcp_size) + outs.append(seg[: grp.tok_end - grp.tok_start]) # crop per-group pad + off += local_len + return torch.cat(outs) + def run_model( self, input_ids: torch.Tensor, @@ -2066,37 +2070,43 @@ def run_model( # internally short-circuits for prefill / cudagraph. forward_mode.assert_shape_contract(input_ids, forward_context.attn_metadata) - # PCP+TBO prefill (障碍 A, b1): apply PCP split here, before UBatchWrapper, - # so each ubatch receives 1/(2*pcp) tokens and ub_slices (in 1/pcp space) - # align with the metadata already reindexed by _apply_pcp_reindex. + # PCP+TBO prefill — b2 + balanced (PCP_TBO.md §12): PER-GROUP stripe here, + # before UBatchWrapper. Each request group is padded to a pcp multiple and + # round-robin striped as a WHOLE (so sequences stay intact per group), then + # the two groups' 1/pcp shards are concatenated. token_slice (built in + # prepare_inputs) indexes into this [g0_local | g1_local] concat. _pcp_size = self.config.prefill_context_parallel_size - _pcp_tbo_prefill = ( + _pcp_bal_groups = getattr(self, "_pcp_bal_groups", None) + _pcp_tbo_balanced = ( _pcp_size > 1 and isinstance(self.model, UBatchWrapper) and forward_context.ubatch_slices is not None and is_prefill and not forward_context.context.is_dummy_run + and _pcp_bal_groups is not None ) - if _pcp_tbo_prefill: + if _pcp_tbo_balanced: n_global = input_ids.shape[0] - _padded = pcp_pad_len(n_global, _pcp_size, num_ubatches=2) - if _padded > n_global: - _pad = _padded - n_global - input_ids = torch.cat([input_ids, input_ids.new_zeros(_pad)]) - positions = torch.cat([positions, positions.new_zeros(_pad)]) - input_ids = pcp_round_robin_split(input_ids, _pcp_size) - positions = pcp_round_robin_split(positions, _pcp_size) - # Update context.positions to the local (1/pcp) version so that - # _make_ubatch_context slices it correctly in 1/pcp space. + g_ids, g_pos = [], [] + for grp in _pcp_bal_groups: + seg_ids = input_ids[grp.tok_start : grp.tok_end] + seg_pos = positions[grp.tok_start : grp.tok_end] + pad = grp.pad_total - (grp.tok_end - grp.tok_start) + if pad > 0: + seg_ids = torch.cat([seg_ids, seg_ids.new_zeros(pad)]) + seg_pos = torch.cat([seg_pos, seg_pos.new_zeros(pad)]) + g_ids.append(pcp_round_robin_split(seg_ids, _pcp_size)) + g_pos.append(pcp_round_robin_split(seg_pos, _pcp_size)) + input_ids = torch.cat(g_ids) + positions = torch.cat(g_pos) + # context.positions = local per-group concat so _make_ubatch_context + # slices each ubatch's forward positions correctly. forward_context.context.positions = positions - # Store local (1/pcp) input_ids for hash MoE (mode B, moe_pcp_merge). - # _make_ubatch_context will slice [token_slice] per ubatch; each - # ForCausalLM.forward then allgathers its per-ubatch slice across PCP - # ranks, giving padded_total//2 ids — matching what moe_pcp_merge_forward - # allgathers for hidden states. Pre-allgathering here would give - # padded_total ids, which mismatches MoE's padded_total//2 num_tokens. + # Hash MoE (mode B): local per-group-concat ids. Each ForCausalLM.forward + # allgathers its ubatch's slice (g_i local, H_i/pcp) across pcp ranks → + # H_i ids, matching moe_pcp_merge_forward's per-ubatch hidden allgather. if self.config.parallel_config.moe_pcp_merge: - forward_context.context.input_ids = input_ids # local, 1/pcp tokens + forward_context.context.input_ids = input_ids if not forward_mode.use_cudagraph: # prefill, or decode forced eager (enforce_eager / DP peer @@ -2147,19 +2157,25 @@ def run_model( model_output = self.model( input_ids, positions, inputs_embeds=inputs_embeds ) - # PCP+TBO prefill: UBatchWrapper has concatenated the two ubatch - # shards (each 1/(2*pcp) tokens); do the single pcp_allgather_rerange - # here and crop to n_global. ForCausalLM.forward skipped this step - # to avoid 2 concurrent RCCL calls from the two ubatch threads. - if _pcp_tbo_prefill: + # PCP+TBO prefill (balanced): UBatchWrapper concatenated the two + # groups' 1/pcp output shards [g0_local | g1_local]. Restore each + # group independently: pcp_allgather_rerange its shard back to the + # group's global order, crop off the per-group pad, then concat to + # the full global sequence. Per-group (not single global) because + # each group was striped independently. + if _pcp_tbo_balanced: if self.use_aux_hidden_state_outputs: _h, _aux = model_output - _h = pcp_allgather_rerange(_h, _pcp_size)[:n_global] - model_output = (_h, _aux) + model_output = ( + self._restore_pcp_balanced_output( + _h, _pcp_bal_groups, _pcp_size + ), + _aux, + ) else: - model_output = pcp_allgather_rerange( - model_output, _pcp_size - )[:n_global] + model_output = self._restore_pcp_balanced_output( + model_output, _pcp_bal_groups, _pcp_size + ) if self.use_aux_hidden_state_outputs: hidden_states, self._aux_hidden_states = model_output else: diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 5a02bf64cb..0f899c931c 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1469,7 +1469,12 @@ def prepare_prefill(self, batch: ScheduledBatch): # model.forward entry round-robin-splits hidden/positions to 1/W, so the # per-query (per-token) metadata must be reduced to the SAME owned-query # set. Per-seq / KV-write fields stay full (every rank keeps full KV). - if pcp_is_enabled() and not batch.is_dummy_run: + # PCP+TBO balanced (§12): DEFER reindex to per-group in + # build_ubatch_prefill_metadata (each request group reindexed + # independently on its own pcp pad). Keep the FULL un-reindexed metadata + # here so build_ubatch can slice it per group. + _bal = getattr(self.model_runner, "_pcp_tbo_balanced_active", False) + if pcp_is_enabled() and not batch.is_dummy_run and not _bal: # Gate on `not is_dummy_run`: ForCausalLM.forward's round-robin-split is # skipped on dummy/warmup runs (_pcp_active() returns False there), # so reindexing metadata to 1/W here would pair full-size @@ -1515,6 +1520,9 @@ def _apply_pcp_reindex( pcp_size = get_pcp_world_size() device = attn_metadata.batch_id_per_token.device # Pad to a multiple of pcp_size; dummy (pad) queries get zero-length KV. + # This runs on the non-TBO PCP path (full-batch reindex) and, under + # PCP+TBO balanced (§12), per request GROUP (each group reindexed + # independently on its own pcp pad). Either way the divisor is pcp_size. padded_total = pcp_pad_len(total_tokens, pcp_size) n_pad = padded_total - total_tokens owned_q = pcp_round_robin_query_indices(padded_total, pcp_size).to(device) @@ -1605,9 +1613,29 @@ def build_ubatch_prefill_metadata( padded_bs: int, ubatch_idx: int = 0, ) -> AttentionMetaData_DSV4: - """Split prefill AttentionMetaData for V4 TBO micro-batches.""" + """Split prefill AttentionMetaData for V4 TBO micro-batches. + + Two paths: + - PCP+TBO balanced (§12, current PCP+TBO scheme): dispatches to + `_build_ubatch_prefill_metadata_balanced(attn_metadata, ubatch_idx)`, + which derives the group from `model_runner._pcp_bal_groups[ubatch_idx]` + and **ignores `ub_slice` / `padded_bs`** (the group's request/token + ranges come from the PcpBalGroup, not the ub_slice). + - Non-balanced TBO (token-split, §11): uses `ub_slice` / `padded_bs`. + """ from atom.utils.tbo.ubatch_splitting import split_attn_metadata + # PCP+TBO balanced (§12): each ubatch = one request group processed as an + # independent non-TBO PCP mini-batch. Slice the FULL (un-reindexed) + # metadata to the group + call _apply_pcp_reindex on it (reuse the proven + # reindex). Bypasses the token-split rebuild path entirely. + if getattr(self.model_runner, "_pcp_tbo_balanced_active", False) and getattr( + self.model_runner, "_pcp_bal_groups", None + ) is not None: + return self._build_ubatch_prefill_metadata_balanced( + attn_metadata, ubatch_idx + ) + ub_attn = split_attn_metadata(attn_metadata, ub_slice, padded_bs) ub_attn.__class__ = AttentionMetaData_DSV4 @@ -1682,11 +1710,23 @@ def build_ubatch_prefill_metadata( real_num_tokens, # actual tokens only (excludes PCP dummy padding) ) - # Option A: real_num_tokens == ub_num_tokens (no dummy, since - # _preprocess gates on n_prefill % (2*pcp)==0). All downstream - # functions receive ub_num_tokens consistently. - # Option B (arbitrary length, with dummy tokens) requires per-ubatch - # reindex — see PCP_TBO.md §11 for the full analysis and remaining work. + # Option B (PCP+TBO arbitrary length) — Path 3 (four steps). + # See PCP_TBO.md §11.7 for full analysis / path comparison. + # + # Step 1: _attach_v4_per_fwd_meta already received real_num_tokens. + # Extend batch_id_per_token from real_num_tokens to ub_num_tokens + # with -1 sentinel so downstream kernels that use batch_id_per_token + # see the correct padded size matching ub_num_tokens. + if real_num_tokens < ub_num_tokens and ub_attn.batch_id_per_token is not None: + dummy_count = ub_num_tokens - real_num_tokens + ub_attn.batch_id_per_token = torch.cat([ + ub_attn.batch_id_per_token, + ub_attn.batch_id_per_token.new_full((dummy_count,), -1), + ]) + + # Step 2: indexer gets ub_num_tokens so indexer_meta matches Q shape. + # Model forward's Q tensor has ub_num_tokens rows; _score_topk_prefill + # requires seq_base.shape == topk_global.shape == ub_num_tokens. positions_gpu = var["positions"].gpu[ts.start : ts.stop] self._attach_v4_indexer_meta( ub_attn, @@ -1695,26 +1735,55 @@ def build_ubatch_prefill_metadata( positions_gpu=positions_gpu, ) - # start_pos = position of first token of each seq in this ubatch. + # Step 3: paged prefill meta uses real_num_tokens + clipped positions_np + # so that positions_arr and chunk_start_pt have matching sizes inside fn. ub_start_pos_per_seq_np = positions_np[ub_cu[:ub_num_reqs]] - ub_positions_gpu = var["positions"].gpu[ts.start : ts.stop] + ub_positions_gpu = var["positions"].gpu[ts.start : ts.start + real_num_tokens] ub_block_tables_gpu = var["block_tables"].gpu[rs.start : rs.stop] ub_cu_q_per_seq_gpu = torch.from_numpy( np.ascontiguousarray(ub_cu[:ub_num_reqs], dtype=np.int32) ).to(self.device, non_blocking=True) self._build_paged_prefill_meta( ub_attn, - positions_np, + positions_np[:real_num_tokens], # clip to real tokens only ub_cu, extend_lens_np, ub_start_pos_per_seq_np, ub_attn.state_slot_mapping_cpu, ub_num_reqs, - ub_num_tokens, + real_num_tokens, positions_gpu=ub_positions_gpu, cu_q_per_seq_gpu=ub_cu_q_per_seq_gpu, block_tables_gpu=ub_block_tables_gpu, ) + + # Step 4 (Option B): pad paged-prefill per-token ragged/dense + # buffers from real_num_tokens to ub_num_tokens so ALL per-token + # metadata matches the model-forward token count (x.size(0) = + # ub_num_tokens; kernels read T from tensor shapes, see §11.4/§11.7). + # _build_paged_prefill_meta above built them at real_num_tokens + # (it can't take ub_num_tokens directly — its internal + # np.repeat(batch_id) is real_num_tokens while positions_arr[:T] + # would be ub_num_tokens → shape mismatch). Append n_pad dummy + # queries: pcp_pad_indptr repeats the last prefix-sum (zero-length + # KV), pcp_pad_dense appends zeros (skip=0) — identical to + # _apply_pcp_reindex's dummy handling, so downstream kernels skip + # them via batch_id=-1. No-op when real==ub (ubatch0 / divisible). + n_pad_ub = ub_num_tokens - real_num_tokens + if n_pad_ub > 0: + for ind_attr in ( + "kv_indptr_prefix_swa", + "kv_indptr_prefix_csa", + "kv_indptr_prefix_hca", + "kv_indptr_extend", + ): + indptr = getattr(ub_attn, ind_attr, None) + if indptr is not None: + setattr(ub_attn, ind_attr, pcp_pad_indptr(indptr, n_pad_ub)) + if ub_attn.skip_prefix_len_csa is not None: + ub_attn.skip_prefix_len_csa = pcp_pad_dense( + ub_attn.skip_prefix_len_csa, n_pad_ub + ) finally: bt[:ub_num_reqs] = saved_bt var["context_lens"].np[:ub_num_reqs] = saved_ctx @@ -1738,6 +1807,11 @@ def build_ubatch_prefill_metadata( if ub_attn.cu_seqlens_k is not None: ub_attn.cu_seqlens_k = ub_cu_gpu + # NOTE: this token-split (§11) build path is retained for the non-balanced + # TBO+PCP case, but the current PCP+TBO scheme is b2+balanced (§12), which + # dispatches to `_build_ubatch_prefill_metadata_balanced` at the top of + # `build_ubatch_prefill_metadata` and never reaches here. + # Clone all GPU tensors that are views into shared CpuGpuBuffers. # Without this, building the next ubatch overwrites this ubatch's # data via the same underlying buffer. @@ -1756,6 +1830,139 @@ def build_ubatch_prefill_metadata( return ub_attn + def _build_ubatch_prefill_metadata_balanced( + self, + attn_metadata: AttentionMetaData, + ubatch_idx: int, + ) -> AttentionMetaData_DSV4: + """PCP+TBO balanced (§12): build one request group's metadata as an + independent non-TBO PCP mini-batch. + + `attn_metadata` is the FULL, UN-reindexed metadata (global). We slice it + to this group's requests + global token range, then run the proven + `_apply_pcp_reindex` on the group (pads the group to a pcp multiple and + round-robin strides to 1/pcp — matching run_model's per-group stripe). + Per-seq / KV-write fields (cu_seqlens_q, compress_plans, state_slot) stay + GLOBAL for the group (the compressor/swa_write see the group's full + all-gathered tokens), exactly as non-TBO PCP does for the whole batch. + """ + from atom.utils.tbo.ubatch_splitting import UBatchSlice, split_attn_metadata + from atom.model_ops.v4_kernels import make_compress_plans + + mr = self.model_runner + grp = mr._pcp_bal_groups[ubatch_idx] # PcpBalGroup + rs0, rs1 = grp.req_start, grp.req_stop + gts, gte = grp.tok_start, grp.tok_end + H = grp.pad_total + group_bs = rs1 - rs0 + group_total = gte - gts # group's global token count (real, pre-pad) + device = self.device + var = mr.forward_vars + src = cast(AttentionMetaData_DSV4, attn_metadata) + + # ---- base fields via split on the GROUP's GLOBAL token range ---- + # full metadata is global, so a global token_slice slices cu_seqlens_q / + # slot_mapping / context_lens correctly (per-request, rebased). + g_slice = UBatchSlice( + request_slice=slice(rs0, rs1), + token_slice=slice(gts, gte), + ) + ub = split_attn_metadata(attn_metadata, g_slice, group_bs) + ub.__class__ = AttentionMetaData_DSV4 + # split_attn_metadata doesn't carry these: state drives prefill/decode + # dispatch; indexer_meta must be non-None so _apply_pcp_reindex rebuilds + # it for the group (it rebuilds from batch_id+positions, ignoring content). + ub.state = src.state + ub.indexer_meta = src.indexer_meta + + # ---- per-seq DSV4 fields sliced by request ---- + if src.state_slot_mapping is not None: + ub.state_slot_mapping = src.state_slot_mapping[rs0:rs1].contiguous() + if src.state_slot_mapping_cpu is not None: + ub.state_slot_mapping_cpu = src.state_slot_mapping_cpu[rs0:rs1] + if src.n_committed_csa_per_seq is not None: + ub.n_committed_csa_per_seq = src.n_committed_csa_per_seq[rs0:rs1].contiguous() + if src.n_committed_csa_per_seq_cpu is not None: + ub.n_committed_csa_per_seq_cpu = src.n_committed_csa_per_seq_cpu[rs0:rs1] + if src.n_committed_hca_per_seq_cpu is not None: + ub.n_committed_hca_per_seq_cpu = src.n_committed_hca_per_seq_cpu[rs0:rs1] + + # ---- per-token DSV4 fields sliced by the GLOBAL token range [gts,gte) ---- + owned = torch.arange(gts, gte, device=device) + for ind_attr, idx_attr in ( + ("kv_indptr_prefix_swa", "kv_indices_prefix_swa"), + ("kv_indptr_prefix_csa", "kv_indices_prefix_csa"), + ("kv_indptr_prefix_hca", "kv_indices_prefix_hca"), + ("kv_indptr_extend", "kv_indices_extend"), + ): + indptr = getattr(src, ind_attr, None) + indices = getattr(src, idx_attr, None) + if indptr is None or indices is None: + continue + ni, nx = pcp_reindex_ragged(indptr, indices, owned) + # kv_indices_extend are ROW offsets into the per-fwd kv_full tensor. + # In the full metadata they index the WHOLE sequence's kv_full [0,T); + # for this group kv_full only holds the group's tokens (global order + # [gts,gte) → rows [0, gte-gts)), so rebase by gts. (prefix indices + # point into unified_kv by absolute cache slot — no rebase.) Balanced + # splits on request boundaries so each query's SWA window stays within + # its sequence (within the group) → row >= gts, rebased value >= 0. + if idx_attr == "kv_indices_extend" and nx.numel() > 0: + nx = nx - gts + setattr(ub, ind_attr, ni) + setattr(ub, idx_attr, nx) + # batch_id_per_token: slice + rebase global req id → group-local (keep -1). + if src.batch_id_per_token is not None: + bid = src.batch_id_per_token[gts:gte].clone() + ub.batch_id_per_token = torch.where(bid >= 0, bid - rs0, bid) + if src.skip_prefix_len_csa is not None: + ub.skip_prefix_len_csa = src.skip_prefix_len_csa[gts:gte].contiguous() + ub.swa_pages = src.swa_pages + + # ---- compress_plans: group's GLOBAL per-request (compressor all-gathers + # the group to full order). Built from global cu / context_lens slices. ---- + if self._unique_compress_ratios_overlap: + gcu = var["cu_seqlens_q"].np # GLOBAL (not overwritten for balanced) + ext = (gcu[rs0 + 1 : rs1 + 1] - gcu[rs0:rs1]).astype(np.int32) + ctx = np.asarray(var["context_lens"].np[rs0:rs1], dtype=np.int32) + plan_bufs = self._get_ubatch_compress_plan_buffers(ubatch_idx) + ub.compress_plans = make_compress_plans( + np.ascontiguousarray(ext, dtype=np.int32), + np.ascontiguousarray(ctx, dtype=np.int32), + self._unique_compress_ratios_overlap, + plan_buffers=plan_bufs, + decode_capacity_per_ratio=None, + ) + else: + ub.compress_plans = {} + + # ---- reindex the group to 1/pcp (proven path) ---- + # positions: group's GLOBAL positions (forward_vars stay global for + # balanced). _apply_pcp_reindex pads group_total to pcp + strides — + # matching run_model's per-group pcp_round_robin_split. + group_positions = var["positions"].gpu[gts:gte] + self._apply_pcp_reindex(ub, group_positions, group_bs, group_total) + + # max_seqlen_q from the group's per-request extend lengths. + if ub.cu_seqlens_q is not None and group_bs > 0: + per_req_q = ub.cu_seqlens_q[1 : group_bs + 1] - ub.cu_seqlens_q[:group_bs] + if per_req_q.numel() > 0: + ub.max_seqlen_q = int(per_req_q.max().item()) + + # Clone GPU tensors that are slices/views into shared CpuGpuBuffers, so a + # later ubatch (or fwd) reusing the same buffer can't overwrite this + # ubatch's data (mirrors the non-balanced path's clones). + # n_committed_csa_per_seq is a view of src's shared buffer (the [rs0:rs1] + # .contiguous() slice above stays a view when already contiguous). + if ub.n_committed_csa_per_seq is not None: + ub.n_committed_csa_per_seq = ub.n_committed_csa_per_seq.clone() + if ub.indexer_meta is not None: + im = ub.indexer_meta + for k in ("cu_committed_gpu", "batch_id_per_token_gpu", "n_committed_per_seq_gpu"): + if im.get(k) is not None: + im[k] = im[k].clone() + return ub + def _attach_v4_per_fwd_meta( self, attn_metadata: AttentionMetaData_DSV4, From 74227ef04b09c38ebc1f3314fc34e0826e6e3a7b Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 3 Jul 2026 09:57:36 +0000 Subject: [PATCH 05/17] Using ATOM_PCP_MOE_MERGE --- atom/config.py | 10 ---------- atom/model_engine/arg_utils.py | 11 ----------- atom/model_engine/llm_engine.py | 9 --------- atom/model_engine/model_runner.py | 2 +- atom/model_ops/moe.py | 2 +- atom/models/deepseek_v4.py | 14 ++++++-------- atom/utils/envs.py | 8 ++++++++ 7 files changed, 16 insertions(+), 40 deletions(-) diff --git a/atom/config.py b/atom/config.py index 789c32481e..a651dd9299 100644 --- a/atom/config.py +++ b/atom/config.py @@ -749,16 +749,6 @@ class ParallelConfig: data_parallel_size: int = 1 """Number of data parallel groups. MoE layers will be sharded according to the product of the tensor parallel size and data parallel size.""" - moe_pcp_merge: bool = False - """PCP (Prefill Context Parallel) MoE comm mode. Only meaningful when - prefill_context_parallel_size > 1. - - False (mode A, default): MoE runs on each rank's 1/W token shard with NO - extra comm (per-token independent). Equivalent to SGLang moe_dp_size == - pcp_size. - - True (mode B): all-gather hidden 1/W -> full before MoE and slice full -> - 1/W after, so MoE sees the complete token set (MoE itself is untouched / - PCP-agnostic). Equivalent to SGLang moe_dp_size == 1. Costs one extra - hidden all-gather per layer; see plan §P3.""" data_parallel_size_local: int = 1 """Number of local data parallel groups.""" data_parallel_rank: int = 0 diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 89fe3270bb..a8533600c7 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -31,7 +31,6 @@ class EngineArgs: trust_remote_code: bool = False tensor_parallel_size: int = 1 prefill_context_parallel_size: int = 1 - moe_pcp_merge: bool = False data_parallel_size: int = 1 enforce_eager: bool = False enable_prefix_caching: bool = True @@ -89,16 +88,6 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="Prefill context parallel size. Independent dimension " "(world = tp x pcp); splits the sequence during prefill.", ) - parser.add_argument( - "--moe-pcp-merge", - action=argparse.BooleanOptionalAction, - default=False, - help="PCP MoE comm mode (only when -pcp > 1). Default (False, mode " - "A): MoE runs on each rank's 1/W token shard, no extra comm. " - "True (mode B): all-gather hidden to full tokens before MoE and " - "slice back after, so MoE sees the complete sequence (MoE itself " - "untouched). Adds one hidden all-gather per layer.", - ) parser.add_argument( "--data-parallel-size", "-dp", diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index bb04a9cacb..d6df31596b 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -39,7 +39,6 @@ def __init__(self, model, tokenizer=None, **kwargs): config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} data_parallel_size = kwargs.get("data_parallel_size", 1) data_parallel_master_port = kwargs.get("data_parallel_master_port", None) - moe_pcp_merge = kwargs.get("moe_pcp_merge", False) config = Config(model, **config_kwargs) self.config = config self.tokenizer = tokenizer or _load_tokenizer( @@ -56,14 +55,6 @@ def __init__(self, model, tokenizer=None, **kwargs): if data_parallel_master_port is not None: config.parallel_config.data_parallel_master_port = data_parallel_master_port self.data_parallel_size = data_parallel_size - # PCP MoE comm mode (mode B). Only meaningful when PCP is enabled. - if moe_pcp_merge and config.prefill_context_parallel_size <= 1: - raise ValueError( - "moe_pcp_merge=True (PCP MoE mode B) requires " - "prefill_context_parallel_size > 1 (-pcp). With PCP disabled " - "there is no token sharding to all-gather for MoE." - ) - config.parallel_config.moe_pcp_merge = moe_pcp_merge # PCP and DP-attention are not yet compatible: PCP stripe-splits # input_ids to 1/pcp_size in ForCausalLM.forward, but DP-attention's # `_gather_ids_for_dp` all-gathers using dp_metadata sizes computed on diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index c998d27c01..18fcaf01ad 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -2105,7 +2105,7 @@ def run_model( # Hash MoE (mode B): local per-group-concat ids. Each ForCausalLM.forward # allgathers its ubatch's slice (g_i local, H_i/pcp) across pcp ranks → # H_i ids, matching moe_pcp_merge_forward's per-ubatch hidden allgather. - if self.config.parallel_config.moe_pcp_merge: + if envs.ATOM_PCP_MOE_MERGE: forward_context.context.input_ids = input_ids if not forward_mode.use_cudagraph: diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 1380ed6861..ff86faa8a0 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -161,7 +161,7 @@ def flatten_tp_across_dp(dp_rank: int): ) pcp_merge = ( - parallel_config.parallel_config.moe_pcp_merge + envs.ATOM_PCP_MOE_MERGE and get_prefill_context_model_parallel_world_size() > 1 ) if pcp_merge: diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 1ddfbdf6c8..f7283fcdb1 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -2309,9 +2309,7 @@ def __init__( # maybe_dual_stream_forward (dual-stream) AND moe_pcp_merge_forward # (PCP mode B, 阶段19) — the latter requires registration regardless of # dual-stream, so register whenever either consumer is active. - _merge_on = get_pcp_world_size() > 1 and bool( - get_current_atom_config().parallel_config.moe_pcp_merge - ) + _merge_on = get_pcp_world_size() > 1 and bool(envs.ATOM_PCP_MOE_MERGE) if self._use_dual_stream or _merge_on: get_current_atom_config().compilation_config.static_forward_context[ prefix @@ -2534,7 +2532,7 @@ def __init__( # warmup/real so Dynamo specializes it identically -> the op call is # baked into code0 and the gate inside the op stays dynamic. self._moe_merge_enabled = get_pcp_world_size() > 1 and bool( - get_current_atom_config().parallel_config.moe_pcp_merge + envs.ATOM_PCP_MOE_MERGE ) self.attn_norm = RMSNorm(args.dim, self.norm_eps) self.ffn_norm = RMSNorm(args.dim, self.norm_eps) @@ -2928,12 +2926,12 @@ def _moe_pcp_merge_active() -> bool: after) applies in this forward. True only when this is a real PCP prefill (`_pcp_active`) AND the - `moe_pcp_merge` flag is set. Mode A (default) returns False: MoE runs on the + `ATOM_PCP_MOE_MERGE` env is set. Mode A returns False: MoE runs on the 1/W shard with no extra comm. """ if not _pcp_active(): return False - return bool(get_current_atom_config().parallel_config.moe_pcp_merge) + return bool(envs.ATOM_PCP_MOE_MERGE) def _moe_pcp_merge_decode_active() -> bool: @@ -2943,12 +2941,12 @@ def _moe_pcp_merge_decode_active() -> bool: instead does one pcp all_reduce after MoE to sum the pcp-half of the intermediate that combine_outputs' tp all_reduce misses. - True only when pcp>1, moe_pcp_merge set, and this is a real DECODE forward + True only when pcp>1, ATOM_PCP_MOE_MERGE set, and this is a real DECODE forward (not prefill — that's `_moe_pcp_merge_active` — and not dummy/warmup). """ if get_pcp_world_size() <= 1: return False - if not bool(get_current_atom_config().parallel_config.moe_pcp_merge): + if not bool(envs.ATOM_PCP_MOE_MERGE): return False fc = get_forward_context() return (not fc.context.is_prefill) and (not fc.context.is_dummy_run) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index c5fa7aee75..afc33d82ef 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -237,6 +237,14 @@ "ATOM_TBO_PREFILL_TOKEN_SPLIT": lambda: ( os.getenv("ATOM_TBO_PREFILL_TOKEN_SPLIT", "1") == "1" ), + # --- PCP MoE comm mode --- + # Fold the PCP (prefill-context-parallel) dim into the MoE tp/ep sharding. + # Only meaningful when prefill_context_parallel_size > 1; + # Default "1": all-gather hidden 1/W -> full before MoE and slice + # full -> 1/W after, so MoE sees the complete token set (MoE itself is + # untouched / PCP-agnostic). Costs one extra hidden all-gather per layer. + # "0": MoE runs on each rank's 1/W token shard with no extra comm. + "ATOM_PCP_MOE_MERGE": lambda: os.getenv("ATOM_PCP_MOE_MERGE", "1") == "1", # Debug logging for TBO (+ PCP) collectives and ubatch splitting. Default off. # When "1", PCP collective sites and TBO ubatch decisions emit WARNING logs # (function name, input shape, ubatch id) to trace cross-rank collective order From eb01b8b39db96e262742f88658ea4a9ac2017f0c Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Mon, 6 Jul 2026 01:50:33 +0000 Subject: [PATCH 06/17] Remove debug log --- atom/distributed/pcp_utils.py | 56 ++----------------------------- atom/model_engine/llm_engine.py | 2 +- atom/model_engine/model_runner.py | 2 +- atom/utils/envs.py | 5 --- 4 files changed, 4 insertions(+), 61 deletions(-) diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 906a8ea2d3..c7388f3aa7 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -13,7 +13,6 @@ `layers/utils/cp_utils.py:cp_all_gather_rerange_output`). """ -import itertools import logging from typing import NamedTuple, Optional @@ -21,7 +20,7 @@ class PcpBalGroup(NamedTuple): - """One request group for PCP+TBO balanced prefill (PCP_TBO.md §12). + """One request group for PCP+TBO balanced 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 @@ -44,41 +43,6 @@ class PcpBalGroup(NamedTuple): logger = logging.getLogger("atom") -# Monotonic counter so debug lines can be matched against RCCL WorkNCCL SeqNum -# when locating a PCP+TBO collective hang. GIL-protected; shared across the 2 -# TBO ubatch threads (the issue order across threads is exactly what we trace). -_pcp_comm_seq = itertools.count() - - -def _tbo_debug_collective(fn: str, input_: torch.Tensor) -> None: - """Emit a WARNING-level trace for one PCP collective when ATOM_TBO_DEBUG=1. - - Logs (in issue order) the call site, input shape/numel, the current TBO - ubatch id, and the PCP rank — so the last line before an RCCL allgather - timeout pinpoints which collective (attention KV/compressor/indexer vs MoE - hidden) deadlocked, and on which ubatch. - """ - from atom.utils import envs - - if not envs.ATOM_TBO_DEBUG: - return - try: - from atom.utils.tbo.ubatching import tbo_active, tbo_current_ubatch_id - - ub = tbo_current_ubatch_id() if tbo_active() else -1 - except Exception: - ub = -1 - seq = next(_pcp_comm_seq) - logger.warning( - "[TBO_DEBUG] seq=%d %s in_shape=%s numel=%d ubatch=%d pcp_rank=%d", - seq, - fn, - tuple(input_.shape), - input_.numel(), - ub, - get_prefill_context_model_parallel_rank(), - ) - def get_pcp_world_size() -> int: return get_prefill_context_model_parallel_world_size() @@ -104,10 +68,6 @@ def pcp_pad_len( padded length (>= total_tokens); callers pad per-token tensors to this length with dummy tokens (KV length 0) before splitting. - When TBO is active with N ubatches, pass num_ubatches=N so each per-rank - shard (total // pcp_size) is also divisible by N — guaranteeing every - ubatch is a pcp_size multiple and the token-midpoint split lands on an - exact pcp_size boundary. """ if pcp_size is None: pcp_size = get_pcp_world_size() @@ -159,7 +119,6 @@ def pcp_allgather_rerange( pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ - _tbo_debug_collective("pcp_allgather_rerange", input_) group = get_pcp_group() # aiter all_gather(dim=0) returns rank-major concat: [pcp*L, *rest]. gathered = group.all_gather(input_.contiguous(), dim=0) @@ -175,15 +134,7 @@ def pcp_allgather_rerange( return out -# ==== MoE-path PCP collectives (scheme B, rank-major gather + reduce_scatter) ==== -# Used by the MoE merge path (moe_pcp_merge_forward custom op). Because that op -# is opaque to Dynamo (the collectives run eagerly inside it, NOT in the compiled -# graph), we use the NATURAL collectives — no all-reduce emulation needed. The -# earlier all-reduce-only scheme was a workaround for "collectives miscompiled -# inside the @support_torch_compile graph"; with the opaque-op fix (方向C) the -# gather/scatter no longer live in the traced graph, so plain all_gather / -# reduce_scatter are correct and cheaper (all_gather moves W× fewer bytes). -# +# ==== 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 @@ -202,7 +153,6 @@ def pcp_allgather_rankmajor( pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ - _tbo_debug_collective("pcp_allgather_rankmajor", input_) return get_pcp_group().all_gather(input_.contiguous(), dim=0) @@ -215,7 +165,6 @@ def pcp_reduce_scatter( pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ - _tbo_debug_collective("pcp_reduce_scatter", input_) return get_pcp_group().reduce_scatter(input_.contiguous(), dim=0) @@ -229,7 +178,6 @@ def pcp_all_reduce(input_: torch.Tensor, pcp_size: Optional[int] = None) -> torc pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ - _tbo_debug_collective("pcp_all_reduce", input_) return get_pcp_group().all_reduce(input_) diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index d6df31596b..c545f7c5d5 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -70,7 +70,7 @@ 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 + TBO prefill: supported via coordinated splitting (P1). + # 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: diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 18fcaf01ad..b8b79ab525 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1741,7 +1741,7 @@ def _preprocess( self.config, batch, is_prefill, num_scheduled_tokens ) - # PCP+TBO prefill — b2 + balanced (PCP_TBO.md §12): split requests into + # PCP+TBO prefill — b2 + balanced: split requests into # two GROUPS at a request boundary (never split a sequence's tokens), so # each ubatch = "non-TBO PCP on a request subset". This root-fixes the # token-split R1/R2 structural conflicts (compressor cross-token diff --git a/atom/utils/envs.py b/atom/utils/envs.py index afc33d82ef..4b496b1f72 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -245,11 +245,6 @@ # untouched / PCP-agnostic). Costs one extra hidden all-gather per layer. # "0": MoE runs on each rank's 1/W token shard with no extra comm. "ATOM_PCP_MOE_MERGE": lambda: os.getenv("ATOM_PCP_MOE_MERGE", "1") == "1", - # Debug logging for TBO (+ PCP) collectives and ubatch splitting. Default off. - # When "1", PCP collective sites and TBO ubatch decisions emit WARNING logs - # (function name, input shape, ubatch id) to trace cross-rank collective order - # — used to locate RCCL allgather hangs under PCP+TBO. - "ATOM_TBO_DEBUG": lambda: os.getenv("ATOM_TBO_DEBUG", "0") == "1", # --- NUMA binding --- # Master switch: pin each GPU worker to its GPU-local NUMA node's CPU cores # and preferred memory. Default off so baseline/pinned A/B stays clean. From ff2da6f463f4e18dc9c44518f64c8525ec89551e Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Mon, 6 Jul 2026 03:01:21 +0000 Subject: [PATCH 07/17] Refactor comments --- atom/distributed/pcp_utils.py | 14 ++-- atom/model_engine/llm_engine.py | 24 +++++- atom/model_engine/model_runner.py | 35 ++++---- atom/model_ops/attentions/deepseek_v4_attn.py | 81 +++---------------- atom/models/deepseek_v4.py | 54 +++---------- 5 files changed, 67 insertions(+), 141 deletions(-) diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index c7388f3aa7..5221d39cfb 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -59,19 +59,21 @@ def pcp_is_enabled() -> bool: def pcp_pad_len( total_tokens: int, pcp_size: Optional[int] = None, - num_ubatches: int = 1, + multiple: int = 1, ) -> int: - """Padded token count so the global sequence is divisible by pcp_size * num_ubatches. + """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() - divisor = pcp_size * max(num_ubatches, 1) + divisor = pcp_size * max(multiple, 1) if divisor <= 1: return total_tokens rem = total_tokens % divisor diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index c545f7c5d5..b0fd3b009b 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -74,12 +74,34 @@ def __init__(self, model, tokenizer=None, **kwargs): # 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: - if config.enable_tbo_decode: + # 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.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 request-boundary balanced grouping + # (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 diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index b8b79ab525..f11773dacc 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1741,13 +1741,10 @@ def _preprocess( self.config, batch, is_prefill, num_scheduled_tokens ) - # PCP+TBO prefill — b2 + balanced: split requests into - # two GROUPS at a request boundary (never split a sequence's tokens), so - # each ubatch = "non-TBO PCP on a request subset". This root-fixes the - # token-split R1/R2 structural conflicts (compressor cross-token - # compression / SWA ring), because every sequence stays whole within one - # group. Requires num_reqs >= 2 (balanced needs two non-empty groups); - # bs=1 falls back to non-TBO. + # PCP+TBO prefill: split requests into two GROUPS at a request boundary + # (never split a sequence's tokens), so each ubatch = "non-TBO PCP on a + # request subset". Requires num_reqs >= 2 (balanced needs two non-empty + # groups); bs=1 falls back to non-TBO. pcp_size = self.config.prefill_context_parallel_size # True for eligible PCP+TBO balanced prefill; read by build_ubatch / # run_model / prepare_prefill to route the per-group path. @@ -1895,12 +1892,12 @@ def prepare_inputs(self, batch: ScheduledBatch, input_ids: torch.Tensor = None): and getattr(self, "_pcp_tbo_balanced_active", False) ) if _pcp_tbo_balanced: - # b2 + balanced (PCP_TBO.md §12): split REQUESTS into two groups at a - # request boundary near the token midpoint. Each group is an - # independent "non-TBO PCP mini-batch": padded to a pcp multiple and - # round-robin striped as a whole, so every sequence stays intact in - # one group (root-fixes token-split R1/R2). forward_vars stay GLOBAL - # here; build_ubatch_prefill_metadata slices the FULL (un-reindexed) + # split REQUESTS into two groups at a request boundary near the + # token midpoint. Each group is an independent "non-TBO PCP + # mini-batch": padded to a pcp multiple and round-robin striped + # as a whole, so every sequence stays intact in one group + # (root-fixes token-split R1/R2). forward_vars stay GLOBAL here; + # build_ubatch_prefill_metadata slices the FULL (un-reindexed) # metadata per group and calls _apply_pcp_reindex on it. num_prefill_reqs = batch.total_seqs_num_prefill per_req = np.asarray( @@ -2070,11 +2067,11 @@ def run_model( # internally short-circuits for prefill / cudagraph. forward_mode.assert_shape_contract(input_ids, forward_context.attn_metadata) - # PCP+TBO prefill — b2 + balanced (PCP_TBO.md §12): PER-GROUP stripe here, - # before UBatchWrapper. Each request group is padded to a pcp multiple and - # round-robin striped as a WHOLE (so sequences stay intact per group), then - # the two groups' 1/pcp shards are concatenated. token_slice (built in - # prepare_inputs) indexes into this [g0_local | g1_local] concat. + # PCP+TBO prefill: PER-GROUP stripe here, before UBatchWrapper. Each + # request group is padded to a pcp multiple and round-robin striped as a + # WHOLE (so sequences stay intact per group), then the two groups' 1/pcp + # shards are concatenated. token_slice (built in prepare_inputs) indexes + # into this [g0_local | g1_local] concat. _pcp_size = self.config.prefill_context_parallel_size _pcp_bal_groups = getattr(self, "_pcp_bal_groups", None) _pcp_tbo_balanced = ( @@ -2102,7 +2099,7 @@ def run_model( # context.positions = local per-group concat so _make_ubatch_context # slices each ubatch's forward positions correctly. forward_context.context.positions = positions - # Hash MoE (mode B): local per-group-concat ids. Each ForCausalLM.forward + # Hash MoE: local per-group-concat ids. Each ForCausalLM.forward # allgathers its ubatch's slice (g_i local, H_i/pcp) across pcp ranks → # H_i ids, matching moe_pcp_merge_forward's per-ubatch hidden allgather. if envs.ATOM_PCP_MOE_MERGE: diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 0f899c931c..09ace3be74 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1469,7 +1469,7 @@ def prepare_prefill(self, batch: ScheduledBatch): # model.forward entry round-robin-splits hidden/positions to 1/W, so the # per-query (per-token) metadata must be reduced to the SAME owned-query # set. Per-seq / KV-write fields stay full (every rank keeps full KV). - # PCP+TBO balanced (§12): DEFER reindex to per-group in + # PCP+TBO balanced: DEFER reindex to per-group in # build_ubatch_prefill_metadata (each request group reindexed # independently on its own pcp pad). Keep the FULL un-reindexed metadata # here so build_ubatch can slice it per group. @@ -1521,8 +1521,8 @@ def _apply_pcp_reindex( device = attn_metadata.batch_id_per_token.device # Pad to a multiple of pcp_size; dummy (pad) queries get zero-length KV. # This runs on the non-TBO PCP path (full-batch reindex) and, under - # PCP+TBO balanced (§12), per request GROUP (each group reindexed - # independently on its own pcp pad). Either way the divisor is pcp_size. + # PCP+TBO balanced, per request GROUP (each group reindexed independently + # on its own pcp pad). Either way the divisor is pcp_size. padded_total = pcp_pad_len(total_tokens, pcp_size) n_pad = padded_total - total_tokens owned_q = pcp_round_robin_query_indices(padded_total, pcp_size).to(device) @@ -1616,7 +1616,7 @@ def build_ubatch_prefill_metadata( """Split prefill AttentionMetaData for V4 TBO micro-batches. Two paths: - - PCP+TBO balanced (§12, current PCP+TBO scheme): dispatches to + - PCP+TBO balanced: dispatches to `_build_ubatch_prefill_metadata_balanced(attn_metadata, ubatch_idx)`, which derives the group from `model_runner._pcp_bal_groups[ubatch_idx]` and **ignores `ub_slice` / `padded_bs`** (the group's request/token @@ -1625,7 +1625,7 @@ def build_ubatch_prefill_metadata( """ from atom.utils.tbo.ubatch_splitting import split_attn_metadata - # PCP+TBO balanced (§12): each ubatch = one request group processed as an + # PCP+TBO balanced: each ubatch = one request group processed as an # independent non-TBO PCP mini-batch. Slice the FULL (un-reindexed) # metadata to the group + call _apply_pcp_reindex on it (reuse the proven # reindex). Bypasses the token-split rebuild path entirely. @@ -1644,12 +1644,6 @@ def build_ubatch_prefill_metadata( ts = ub_slice.token_slice ub_num_reqs = rs.stop - rs.start ub_num_tokens = ts.stop - ts.start - # ub_num_tokens above counts the full token_slice width (= padded size - # when PCP+TBO inserts dummy tokens to keep cross-rank alignment). - # real_num_tokens is computed below from sum(extend_lens_np) after - # clamping, which excludes dummy-padding rows (extend_lens==0 for them). - # Passed as `total_tokens` to _attach_v4_per_fwd_meta so that - # batch_id_unpadded_np[:total_tokens] = ... doesn't broadcast-error. if src.state_slot_mapping is not None: ub_attn.state_slot_mapping = src.state_slot_mapping[rs] @@ -1666,10 +1660,6 @@ def build_ubatch_prefill_metadata( extend_lens_np = (clamped_ends - clamped_starts).astype(np.int32) ub_cu = np.zeros(ub_num_reqs + 1, dtype=np.int32) np.cumsum(extend_lens_np, dtype=np.int32, out=ub_cu[1:]) - # Real (non-dummy) token count = sum of clamped extend lengths. - # Equals ub_num_tokens in standard TBO; may be smaller when PCP+TBO - # pads token_slice to a pcp-aligned boundary (Option B). - real_num_tokens = int(ub_cu[ub_num_reqs]) ub_start_pos_for_ctx = positions_np[ub_cu[:ub_num_reqs]].astype(np.int32) context_lens_np = (ub_start_pos_for_ctx + extend_lens_np).astype(np.int32) from atom.model_ops.v4_kernels import make_compress_plans @@ -1707,26 +1697,9 @@ def build_ubatch_prefill_metadata( extend_lens_np, # ubatch's per-seq token counts ub_attn.state_slot_mapping_cpu, ub_num_reqs, - real_num_tokens, # actual tokens only (excludes PCP dummy padding) + ub_num_tokens, ) - # Option B (PCP+TBO arbitrary length) — Path 3 (four steps). - # See PCP_TBO.md §11.7 for full analysis / path comparison. - # - # Step 1: _attach_v4_per_fwd_meta already received real_num_tokens. - # Extend batch_id_per_token from real_num_tokens to ub_num_tokens - # with -1 sentinel so downstream kernels that use batch_id_per_token - # see the correct padded size matching ub_num_tokens. - if real_num_tokens < ub_num_tokens and ub_attn.batch_id_per_token is not None: - dummy_count = ub_num_tokens - real_num_tokens - ub_attn.batch_id_per_token = torch.cat([ - ub_attn.batch_id_per_token, - ub_attn.batch_id_per_token.new_full((dummy_count,), -1), - ]) - - # Step 2: indexer gets ub_num_tokens so indexer_meta matches Q shape. - # Model forward's Q tensor has ub_num_tokens rows; _score_topk_prefill - # requires seq_base.shape == topk_global.shape == ub_num_tokens. positions_gpu = var["positions"].gpu[ts.start : ts.stop] self._attach_v4_indexer_meta( ub_attn, @@ -1735,55 +1708,26 @@ def build_ubatch_prefill_metadata( positions_gpu=positions_gpu, ) - # Step 3: paged prefill meta uses real_num_tokens + clipped positions_np - # so that positions_arr and chunk_start_pt have matching sizes inside fn. + # start_pos = position of first token of each seq in this ubatch. ub_start_pos_per_seq_np = positions_np[ub_cu[:ub_num_reqs]] - ub_positions_gpu = var["positions"].gpu[ts.start : ts.start + real_num_tokens] + ub_positions_gpu = var["positions"].gpu[ts.start : ts.stop] ub_block_tables_gpu = var["block_tables"].gpu[rs.start : rs.stop] ub_cu_q_per_seq_gpu = torch.from_numpy( np.ascontiguousarray(ub_cu[:ub_num_reqs], dtype=np.int32) ).to(self.device, non_blocking=True) self._build_paged_prefill_meta( ub_attn, - positions_np[:real_num_tokens], # clip to real tokens only + positions_np, ub_cu, extend_lens_np, ub_start_pos_per_seq_np, ub_attn.state_slot_mapping_cpu, ub_num_reqs, - real_num_tokens, + ub_num_tokens, positions_gpu=ub_positions_gpu, cu_q_per_seq_gpu=ub_cu_q_per_seq_gpu, block_tables_gpu=ub_block_tables_gpu, ) - - # Step 4 (Option B): pad paged-prefill per-token ragged/dense - # buffers from real_num_tokens to ub_num_tokens so ALL per-token - # metadata matches the model-forward token count (x.size(0) = - # ub_num_tokens; kernels read T from tensor shapes, see §11.4/§11.7). - # _build_paged_prefill_meta above built them at real_num_tokens - # (it can't take ub_num_tokens directly — its internal - # np.repeat(batch_id) is real_num_tokens while positions_arr[:T] - # would be ub_num_tokens → shape mismatch). Append n_pad dummy - # queries: pcp_pad_indptr repeats the last prefix-sum (zero-length - # KV), pcp_pad_dense appends zeros (skip=0) — identical to - # _apply_pcp_reindex's dummy handling, so downstream kernels skip - # them via batch_id=-1. No-op when real==ub (ubatch0 / divisible). - n_pad_ub = ub_num_tokens - real_num_tokens - if n_pad_ub > 0: - for ind_attr in ( - "kv_indptr_prefix_swa", - "kv_indptr_prefix_csa", - "kv_indptr_prefix_hca", - "kv_indptr_extend", - ): - indptr = getattr(ub_attn, ind_attr, None) - if indptr is not None: - setattr(ub_attn, ind_attr, pcp_pad_indptr(indptr, n_pad_ub)) - if ub_attn.skip_prefix_len_csa is not None: - ub_attn.skip_prefix_len_csa = pcp_pad_dense( - ub_attn.skip_prefix_len_csa, n_pad_ub - ) finally: bt[:ub_num_reqs] = saved_bt var["context_lens"].np[:ub_num_reqs] = saved_ctx @@ -1807,11 +1751,6 @@ def build_ubatch_prefill_metadata( if ub_attn.cu_seqlens_k is not None: ub_attn.cu_seqlens_k = ub_cu_gpu - # NOTE: this token-split (§11) build path is retained for the non-balanced - # TBO+PCP case, but the current PCP+TBO scheme is b2+balanced (§12), which - # dispatches to `_build_ubatch_prefill_metadata_balanced` at the top of - # `build_ubatch_prefill_metadata` and never reaches here. - # Clone all GPU tensors that are views into shared CpuGpuBuffers. # Without this, building the next ubatch overwrites this ubatch's # data via the same underlying buffer. diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index f7283fcdb1..ab15ff0ca8 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -2307,7 +2307,7 @@ def __init__( # Register self in static_forward_context so the custom op dispatcher # can look us up by `layer_name` (= self.prefix). Needed by # maybe_dual_stream_forward (dual-stream) AND moe_pcp_merge_forward - # (PCP mode B, 阶段19) — the latter requires registration regardless of + # — the latter requires registration regardless of # dual-stream, so register whenever either consumer is active. _merge_on = get_pcp_world_size() > 1 and bool(envs.ATOM_PCP_MOE_MERGE) if self._use_dual_stream or _merge_on: @@ -2425,7 +2425,7 @@ def combine_outputs( all-reduce across TP ranks. """ if shared is not None: - # Scheme B (moe_pcp_merge, non-fused shared only): the shared expert + # PCP with ATOM_PCP_MOE_MERGE=1 (non-fused shared only): the shared expert # is NOT pcp-sharded (its MergedColumn/RowParallelLinear bind to the # 4-card tp group, so every pcp rank holds the same shared weights # and computes the same full shared output — pcp-redundant). After @@ -2527,10 +2527,6 @@ def __init__( indexer_stream=indexer_stream, ) self.ffn = MoE(layer_id, args, prefix=f"{prefix}.ffn", alt_stream=alt_stream) - # 阶段19: static (config-only, is_dummy-independent) gate for routing MoE - # through the opaque moe_pcp_merge_forward custom op. Constant across - # warmup/real so Dynamo specializes it identically -> the op call is - # baked into code0 and the gate inside the op stays dynamic. self._moe_merge_enabled = get_pcp_world_size() > 1 and bool( envs.ATOM_PCP_MOE_MERGE ) @@ -2729,18 +2725,6 @@ def forward( self.norm_eps, ) x = hc_state.x_prev - # PCP moe_pcp_merge (scheme B): MoE weights are sharded W*tp ways (pcp - # folded into tp/ep at build time), so each rank holds 1/(W*tp) of the - # experts and must see the full token set. 阶段19 (方向 C): the gate + - # merge collectives + MoE all run inside the OPAQUE moe_pcp_merge_forward - # custom op. Being a Dynamo barrier, the gate is read eagerly each call - # and never specialized/baked (the failure mode that killed the in-graph - # gate, 阶段16-18). The op is shape-preserving ([1/W,dim] in/out: - # prefill gathers->full, runs MoE, reduce_scatters->1/W; decode runs MoE - # then pcp all_reduce). `self._moe_merge_enabled` is a static config bool - # (is_dummy-independent) so Dynamo specializes it identically at warmup - # and real -> the op call is baked into code0; the runtime gate lives - # inside the op. Mode A / non-merge keeps the plain self.ffn path. if self._moe_merge_enabled: x = torch.ops.aiter.moe_pcp_merge_forward(x, self.ffn.prefix) else: @@ -2817,24 +2801,6 @@ def forward( return self.get_logits(norm(x)) # [bs, vocab] -# ===== 阶段19 (方向 C): PCP MoE merge as an OPAQUE custom op ===== -# The whole "gate + merge collectives + MoE" runs inside this custom op. Being -# a Dynamo barrier, the gate (read from forward_context) is evaluated EAGERLY on -# every invocation and is NEVER specialized/baked — which is what defeated the -# in-graph gate (阶段16-18: torch.compile traced the gate to its warmup value, -# dispatch_to_code(0) then froze it). The op is shape-preserving -# ([n_local, dim] in/out: prefill gathers 1/W->full, runs MoE, reduce_scatters -# full->1/W; decode runs MoE then pcp all_reduce), so the fake impl is -# empty_like and Dynamo needs no dynamic-shape reasoning. prefill/decode split -# on the REAL is_prefill read inside the op: decode's cudagraph capture (阶段18: -# is_prefill=False, is_dummy=False) bakes the all_reduce kernel into the decode -# graph; prefill (no cudagraph) runs the gather/scatter eagerly each call. The -# gate KEEPS is_dummy_run (via _moe_pcp_merge_active/_decode_active) so it stays -# consistent with the out-of-graph split gate _pcp_active: at warmup -# (is_dummy=True) it neither splits nor merges. Opacity — NOT removing is_dummy -# — is what avoids the specialization that killed the in-graph gate (阶段16-18). - - def _run_moe(moe: "MoE", x: torch.Tensor) -> torch.Tensor: """Replicate MoE.forward's dual/single dispatch (we are already inside an opaque op, so call the underlying methods directly rather than re-entering @@ -3176,23 +3142,23 @@ def forward( # NOTE: moe_merge here is the OUT-OF-GRAPH gate, used only for the # input_ids gather below (round-robin split + hash-MoE id alignment). # The MoE merge collectives gate themselves separately inside the opaque - # moe_pcp_merge_forward custom op (called from Block.forward, 阶段19). + # moe_pcp_merge_forward custom op (called from Block.forward). moe_merge = _moe_pcp_merge_active() - # Mode B is incompatible with DP-attention id-gathering for now: both + # PCP with ATOM_PCP_MOE_MERGE=1 (default) is incompatible with DP-attention id-gathering for now: both # rewrite ctx.context.input_ids with different (full-PCP vs DP-gathered) # token sets, and stacking them is unverified. Disallow until needed. assert not (moe_merge and self._need_ids_gather), ( - "moe_pcp_merge (PCP MoE mode B) is not supported together with " - "DP-attention input-id gathering yet." + "PCP with ATOM_PCP_MOE_MERGE=1 (default) is not supported " + "together with DP-attention input-id gathering yet." ) full_padded_ids = None if use_pcp: from atom.utils.tbo.ubatching import tbo_active as _tbo_active pcp_size = get_pcp_world_size() if _tbo_active(): - # PCP+TBO prefill (障碍 A, b1): the split was done upstream in - # run_model; input_ids is already 1/(2*pcp) tokens. full_padded_ids - # was precomputed there and stored in ctx.context.input_ids (carried + # PCP+TBO prefill: the split was done upstream in run_model; + # input_ids is already 1/(2*pcp) tokens. full_padded_ids was + # precomputed there and stored in ctx.context.input_ids (carried # into this ubatch context by _make_ubatch_context). Nothing to do. pass else: @@ -3203,7 +3169,7 @@ def forward( positions = torch.cat([positions, positions.new_zeros(pad)], dim=0) input_ids = pcp_round_robin_split(input_ids, pcp_size) positions = pcp_round_robin_split(positions, pcp_size) - # Scheme B: each Block rank-major all-gathers hidden before MoE + # each Block rank-major all-gathers hidden before MoE # (pcp_allgather_rankmajor = plain all_gather(dim=0), RANK-MAJOR # order: [rank0's stripe | rank1's stripe | ...]). The hash MoE # indexes input_ids per hidden row, so the ids must be gathered in From b5b629fbea4d3ce14ec82b73e119bb85dcee51b3 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Mon, 6 Jul 2026 03:21:47 +0000 Subject: [PATCH 08/17] Refactor comments 2 --- atom/model_ops/attentions/deepseek_v4_attn.py | 2 +- atom/models/deepseek_v4.py | 15 ++++++++------- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 8cd1446815..fbaea9523d 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1780,7 +1780,7 @@ def _build_ubatch_prefill_metadata_balanced( attn_metadata: AttentionMetaData, ubatch_idx: int, ) -> AttentionMetaData_DSV4: - """PCP+TBO balanced (§12): build one request group's metadata as an + """PCP+TBO balanced: build one request group's metadata as an independent non-TBO PCP mini-batch. `attn_metadata` is the FULL, UN-reindexed metadata (global). We slice it diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index ab15ff0ca8..6cc64cae5d 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -2309,8 +2309,8 @@ def __init__( # maybe_dual_stream_forward (dual-stream) AND moe_pcp_merge_forward # — the latter requires registration regardless of # dual-stream, so register whenever either consumer is active. - _merge_on = get_pcp_world_size() > 1 and bool(envs.ATOM_PCP_MOE_MERGE) - if self._use_dual_stream or _merge_on: + _pcp_merge_on = get_pcp_world_size() > 1 and bool(envs.ATOM_PCP_MOE_MERGE) + if self._use_dual_stream or _pcp_merge_on: get_current_atom_config().compilation_config.static_forward_context[ prefix ] = self @@ -2888,8 +2888,8 @@ def _pcp_active() -> bool: def _moe_pcp_merge_active() -> bool: - """Whether PCP MoE mode B (all-gather hidden -> full before MoE, slice back - after) applies in this forward. + """Whether PCP applies `attn with PCP -> all-gather hidden -> full before MoE, + slice back` after in this forward. True only when this is a real PCP prefill (`_pcp_active`) AND the `ATOM_PCP_MOE_MERGE` env is set. Mode A returns False: MoE runs on the @@ -3144,9 +3144,10 @@ def forward( # The MoE merge collectives gate themselves separately inside the opaque # moe_pcp_merge_forward custom op (called from Block.forward). moe_merge = _moe_pcp_merge_active() - # PCP with ATOM_PCP_MOE_MERGE=1 (default) is incompatible with DP-attention id-gathering for now: both - # rewrite ctx.context.input_ids with different (full-PCP vs DP-gathered) - # token sets, and stacking them is unverified. Disallow until needed. + # PCP with ATOM_PCP_MOE_MERGE=1 (default) is incompatible with DP-attention + # id-gathering for now: both rewrite ctx.context.input_ids with different + # (full-PCP vs DP-gathered) token sets, and stacking them is unverified. + # Disallow until needed. assert not (moe_merge and self._need_ids_gather), ( "PCP with ATOM_PCP_MOE_MERGE=1 (default) is not supported " "together with DP-attention input-id gathering yet." From d7fec57048e6ecf953a8cd8cff35c376a5a00b19 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Tue, 7 Jul 2026 04:46:33 +0000 Subject: [PATCH 09/17] Add PCP doc --- atom/distributed/pcp_utils.py | 2 +- atom/model_engine/llm_engine.py | 6 +- atom/model_engine/model_runner.py | 17 +- atom/model_ops/attentions/deepseek_v4_attn.py | 20 +-- docs/context_parallel_guide.md | 161 ++++++++++++++++++ docs/index.rst | 1 + 6 files changed, 185 insertions(+), 22 deletions(-) create mode 100644 docs/context_parallel_guide.md diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 5221d39cfb..7ad365ceac 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -20,7 +20,7 @@ class PcpBalGroup(NamedTuple): - """One request group for PCP+TBO balanced prefill. + """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 diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index b0fd3b009b..c86c7b0269 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -94,9 +94,9 @@ def __init__(self, model, tokenizer=None, **kwargs): "--enable-tbo all (decode TBO) is not supported yet. " "Use --enable-tbo (prefill only) with -pcp." ) - # Under PCP, TBO prefill uses request-boundary balanced grouping - # (never token-midpoint split), so ATOM_TBO_PREFILL_TOKEN_SPLIT is - # ignored. + # 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 " diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index f11773dacc..ff1216c50d 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1743,14 +1743,15 @@ def _preprocess( # PCP+TBO prefill: split requests into two GROUPS at a request boundary # (never split a sequence's tokens), so each ubatch = "non-TBO PCP on a - # request subset". Requires num_reqs >= 2 (balanced needs two non-empty - # groups); bs=1 falls back to non-TBO. + # request subset". Requires num_reqs >= 2 (request-boundary split needs + # two non-empty groups); bs=1 falls back to non-TBO. pcp_size = self.config.prefill_context_parallel_size - # True for eligible PCP+TBO balanced prefill; read by build_ubatch / - # run_model / prepare_prefill to route the per-group path. + # True for eligible PCP+TBO request-boundary split prefill; read by + # build_ubatch / run_model / prepare_prefill to route the per-group path. self._pcp_tbo_balanced_active = False - # Per-group descriptors; reset each step, set in prepare_inputs balanced - # branch. Guards run_model/build_ubatch against stale values. + # Per-group descriptors; reset each step, set in prepare_inputs + # request-boundary-split branch. Guards run_model/build_ubatch against + # stale values. self._pcp_bal_groups = None if tbo_on and is_prefill and pcp_size > 1 and not batch.is_dummy_run: num_prefill_reqs = batch.total_seqs_num_prefill @@ -2026,7 +2027,7 @@ def _restore_pcp_balanced_output( groups: "list[PcpBalGroup]", pcp_size: int, ) -> torch.Tensor: - """Restore PCP+TBO balanced output (PCP_TBO.md §12). + """Restore PCP+TBO request-boundary-split output. UBatchWrapper concatenated the two groups' 1/pcp output shards [g0_local | g1_local]. Each group was striped independently, so restore @@ -2154,7 +2155,7 @@ def run_model( model_output = self.model( input_ids, positions, inputs_embeds=inputs_embeds ) - # PCP+TBO prefill (balanced): UBatchWrapper concatenated the two + # PCP+TBO prefill (request-boundary split): UBatchWrapper concatenated the two # groups' 1/pcp output shards [g0_local | g1_local]. Restore each # group independently: pcp_allgather_rerange its shard back to the # group's global order, crop off the per-group pad, then concat to diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index fbaea9523d..7287f42b93 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1475,7 +1475,7 @@ def prepare_prefill(self, batch: ScheduledBatch): # model.forward entry round-robin-splits hidden/positions to 1/W, so the # per-query (per-token) metadata must be reduced to the SAME owned-query # set. Per-seq / KV-write fields stay full (every rank keeps full KV). - # PCP+TBO balanced: DEFER reindex to per-group in + # PCP+TBO request-boundary split: DEFER reindex to per-group in # build_ubatch_prefill_metadata (each request group reindexed # independently on its own pcp pad). Keep the FULL un-reindexed metadata # here so build_ubatch can slice it per group. @@ -1527,7 +1527,7 @@ def _apply_pcp_reindex( device = attn_metadata.batch_id_per_token.device # Pad to a multiple of pcp_size; dummy (pad) queries get zero-length KV. # This runs on the non-TBO PCP path (full-batch reindex) and, under - # PCP+TBO balanced, per request GROUP (each group reindexed independently + # PCP+TBO request-boundary split, per request GROUP (each group reindexed independently # on its own pcp pad). Either way the divisor is pcp_size. padded_total = pcp_pad_len(total_tokens, pcp_size) n_pad = padded_total - total_tokens @@ -1622,16 +1622,16 @@ def build_ubatch_prefill_metadata( """Split prefill AttentionMetaData for V4 TBO micro-batches. Two paths: - - PCP+TBO balanced: dispatches to + - PCP+TBO request-boundary split: dispatches to `_build_ubatch_prefill_metadata_balanced(attn_metadata, ubatch_idx)`, which derives the group from `model_runner._pcp_bal_groups[ubatch_idx]` and **ignores `ub_slice` / `padded_bs`** (the group's request/token ranges come from the PcpBalGroup, not the ub_slice). - - Non-balanced TBO (token-split, §11): uses `ub_slice` / `padded_bs`. + - Token-split TBO (default, §11): uses `ub_slice` / `padded_bs`. """ from atom.utils.tbo.ubatch_splitting import split_attn_metadata - # PCP+TBO balanced: each ubatch = one request group processed as an + # PCP+TBO request-boundary split: each ubatch = one request group processed as an # independent non-TBO PCP mini-batch. Slice the FULL (un-reindexed) # metadata to the group + call _apply_pcp_reindex on it (reuse the proven # reindex). Bypasses the token-split rebuild path entirely. @@ -1780,7 +1780,7 @@ def _build_ubatch_prefill_metadata_balanced( attn_metadata: AttentionMetaData, ubatch_idx: int, ) -> AttentionMetaData_DSV4: - """PCP+TBO balanced: build one request group's metadata as an + """PCP+TBO request-boundary split: build one request group's metadata as an independent non-TBO PCP mini-batch. `attn_metadata` is the FULL, UN-reindexed metadata (global). We slice it @@ -1867,7 +1867,7 @@ def _build_ubatch_prefill_metadata_balanced( # ---- compress_plans: group's GLOBAL per-request (compressor all-gathers # the group to full order). Built from global cu / context_lens slices. ---- if self._unique_compress_ratios_overlap: - gcu = var["cu_seqlens_q"].np # GLOBAL (not overwritten for balanced) + gcu = var["cu_seqlens_q"].np # GLOBAL (not overwritten for request-boundary split) ext = (gcu[rs0 + 1 : rs1 + 1] - gcu[rs0:rs1]).astype(np.int32) ctx = np.asarray(var["context_lens"].np[rs0:rs1], dtype=np.int32) plan_bufs = self._get_ubatch_compress_plan_buffers(ubatch_idx) @@ -1882,8 +1882,8 @@ def _build_ubatch_prefill_metadata_balanced( ub.compress_plans = {} # ---- reindex the group to 1/pcp (proven path) ---- - # positions: group's GLOBAL positions (forward_vars stay global for - # balanced). _apply_pcp_reindex pads group_total to pcp + strides — + # positions: group's GLOBAL positions (forward_vars stay global for the + # request-boundary split). _apply_pcp_reindex pads group_total to pcp + strides — # matching run_model's per-group pcp_round_robin_split. group_positions = var["positions"].gpu[gts:gte] self._apply_pcp_reindex(ub, group_positions, group_bs, group_total) @@ -1896,7 +1896,7 @@ def _build_ubatch_prefill_metadata_balanced( # Clone GPU tensors that are slices/views into shared CpuGpuBuffers, so a # later ubatch (or fwd) reusing the same buffer can't overwrite this - # ubatch's data (mirrors the non-balanced path's clones). + # ubatch's data (mirrors the token-split path's clones). # n_committed_csa_per_seq is a view of src's shared buffer (the [rs0:rs1] # .contiguous() slice above stays a view when already contiguous). if ub.n_committed_csa_per_seq is not None: diff --git a/docs/context_parallel_guide.md b/docs/context_parallel_guide.md new file mode 100644 index 0000000000..cfb256b415 --- /dev/null +++ b/docs/context_parallel_guide.md @@ -0,0 +1,161 @@ +# Prefill Context Parallel (PCP) Guide + +For long-context serving, prefill is bottlenecked on the **sequence dimension**: +the DSA indexer scores every query against all history, and that cost grows with +sequence length and is replicated across Tensor-Parallel (TP) ranks. Plain TP +shards weights/heads/experts, not tokens, so it cannot reduce this cost. + +**Prefill Context Parallel (PCP)** is an independent parallelism dimension that +splits the **prefill token sequence** across the PCP process group, so each GPU +processes only `1/pcp` of the tokens during prefill. This cuts the per-GPU +prefill work (and the indexer's sequence-length cost) to `1/pcp`, lowering TTFT +and raising long-prefill throughput. Decode is left unchanged. PCP composes with +TP and Expert Parallelism (EP): the total world size is `world = tp × pcp`. + +``` + pcp = 2, prefill tokens: 0 1 2 3 4 5 + GPU (pcp rank 0): 0 2 4 ← each GPU processes 1/pcp of the tokens + GPU (pcp rank 1): 1 3 5 + Full KV is kept on every rank; the 1/pcp outputs are all-gathered back to the + full sequence before the LM head. Decode runs as usual (no split). +``` + +> **Model support.** PCP currently supports **DeepSeek-V4** only. Support for +> more models will be added incrementally. + +## When to use PCP + +- **Best fit**: long-context / large-prompt prefill on DeepSeek-V4, where prefill + TTFT dominates. +- **Combine with**: `--enable-tbo` (prefill) to overlap the MoE communication PCP + introduces (see [Overlapping communication with TBO](#overlapping-communication-with-tbo)). +- **Requires**: `world = tp × pcp` GPUs, e.g. `-tp 4 -pcp 2` on 8 GPUs. +- **Little benefit**: decode-heavy or short-prompt workloads — use TP/EP as usual. + +## Quick Reference + +| Flag / Variable | Default | Purpose | +|-----------------|---------|---------| +| `-pcp N` / `--prefill-context-parallel-size N` | `1` | Enable PCP with size `N` (`world = tp × pcp`) | +| `ATOM_PCP_MOE_MERGE` | `1` | Whether to shard MoE across the PCP ranks too | +| `--enable-tbo [prefill\|all]` | off | Overlap compute with PCP communication. With PCP, only prefill TBO is supported | + +| Goal (8 GPUs) | Command | +|------|---------| +| Long-context prefill | `-tp 4 -pcp 2` | +| Long-context prefill + overlap | `-tp 4 -pcp 2 --enable-tbo` | +| Disable PCP (baseline) | `-tp 8 -pcp 1` | + +## CLI usage + +```bash +-pcp N # or --prefill-context-parallel-size N; world = tp × pcp +--enable-tbo # prefill-only TBO overlap (prefill only supported with PCP) +``` + +```bash +ATOM_PCP_MOE_MERGE=1 # default: shard MoE across PCP ranks (gather/scatter) +ATOM_PCP_MOE_MERGE=0 # run MoE per-rank on its 1/pcp shard, no extra MoE comm +``` + +`ATOM_PCP_MOE_MERGE` only has an effect when PCP is enabled (`pcp > 1`): + +| Value | MoE behaviour | When to use | +|---|---|---| +| `1` (default, recommended) | PCP is folded into the MoE tensor/expert sharding, so MoE weights are also sharded across PCP ranks. Lowers per-GPU MoE weight memory, at the cost of one hidden gather/scatter per MoE layer (which TBO overlaps). | Most deployments | +| `0` | Each GPU runs MoE independently on its `1/pcp` shard with no extra MoE communication; MoE weights are replicated across PCP ranks. | Avoid extra MoE comm and have memory headroom for replicated MoE weights | + +## Launching server + +### DeepSeek-V4: TP4 + PCP2 (8 GPUs) + +```bash +python -m atom.entrypoints.openai_server \ + --model deepseek-ai/DeepSeek-V4 \ + -tp 4 -pcp 2 \ + --kv_cache_dtype fp8 +``` + +### DeepSeek-V4: TP4 + PCP2 + prefill TBO overlap (8 GPUs) + +```bash +python -m atom.entrypoints.openai_server \ + --model deepseek-ai/DeepSeek-V4 \ + -tp 4 -pcp 2 \ + --enable-tbo \ + --kv_cache_dtype fp8 +``` + +Tips: +- `-tp 8 -pcp 1` (or omitting `-pcp`) disables PCP and serves as the baseline. +- `--enable-tbo` overlaps the MoE communication introduced by + `ATOM_PCP_MOE_MERGE=1`. It only helps in that mode: with `ATOM_PCP_MOE_MERGE=0` + there is no MoE communication to overlap, so TBO is auto-disabled (a warning is + logged). +- Under PCP, TBO uses a **request-boundary split** (each micro-batch is a whole + subset of requests) — the non-default TBO split mode — instead of the + token-midpoint split used without PCP. `ATOM_TBO_PREFILL_TOKEN_SPLIT` + therefore has no effect when PCP is enabled. +- `--torch-profiler-dir ./log` can be added to collect traces for performance + analysis. + +## Performance baseline + +Benchmark against a running server with a long input length (PCP targets prefill): + +```bash +python -m atom.benchmarks.benchmark_serving \ + --model=deepseek-ai/DeepSeek-V4 --backend=vllm --base-url=http://localhost:7777 \ + --dataset-name=random \ + --random-input-len=32768 --random-output-len=512 \ + --num-prompts=128 --max-concurrency=64 \ + --request-rate=inf --ignore-eos +``` + +Compare `-tp 4 -pcp 2` against the `-tp 8` baseline and watch **Mean TTFT** and +output throughput; the gap widens as `--random-input-len` grows. + +> PCP was introduced in [ROCm/ATOM#1220](https://github.com/ROCm/ATOM/pull/1220), +> which reported, on 8×MI308 for `-tp 4 -pcp 2` vs `-tp 8`, a **35–43%** Mean-TTFT +> reduction and up to **~49%** higher throughput on long prefill. Actual gains +> depend on model, sequence length, and hardware. + +## Constraints & Compatibility + +| Constraint | Notes | +|-----------|-------| +| Models | DeepSeek-V4 only (more coming) | +| World size | `tp × pcp ≤ 8`; multi-node not yet validated | +| PCP + DP-attention | Not supported (raises at startup) | +| PCP + TBO decode (`--enable-tbo all`) | Not supported (raises at startup); use `--enable-tbo` prefill-only | +| `ATOM_PCP_MOE_MERGE=0` + `--enable-tbo` | TBO auto-disabled (warning logged) | + +PCP + TBO **prefill** is supported. Decode is unchanged by PCP in all +configurations. + +## How it works + +1. At the start of the prefill forward, the token sequence is split round-robin + across the PCP ranks (token `i` → rank `i % pcp`), padded so the count divides + evenly. +2. Each rank runs attention / indexer / compressor on its `1/pcp` token shard. + The full KV is kept on every rank (all-gathered as needed), so the attention + kernels are unchanged. +3. MoE either runs on the local `1/pcp` shard (`ATOM_PCP_MOE_MERGE=0`) or gathers + to the full sequence and scatters back (`=1`, default). +4. After the final layer, the `1/pcp` hidden states are all-gathered back to the + full sequence, the original token order is restored, and the LM head runs. +5. Decode is untouched: every rank keeps the full KV and runs normally, so PCP + adds no decode-time cost. + +## Source Files + +| File | Description | +|------|-------------| +| `atom/model_engine/arg_utils.py` | `--prefill-context-parallel-size`, `--enable-tbo` CLI | +| `atom/utils/envs.py` | `ATOM_PCP_MOE_MERGE` | +| `atom/distributed/pcp_utils.py` | PCP communication and helper primitives | +| `atom/models/deepseek_v4.py` | DeepSeek-V4 PCP forward path and MoE handling | +| `atom/model_ops/attentions/deepseek_v4_attn.py` | PCP attention metadata (incl. PCP + TBO prefill) | +| `atom/model_engine/model_runner.py` | PCP token split and PCP + TBO grouping | +| `atom/model_engine/llm_engine.py` | PCP / TBO / DP-attention validation | diff --git a/docs/index.rst b/docs/index.rst index 4938b4bf19..42cdc77fd8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -25,6 +25,7 @@ ATOM Documentation model_ops_guide scheduling_kv_cache_guide distributed_guide + context_parallel_guide compilation_cudagraph_guide serving_benchmarking_guide From 42870d4600b08d980cdf525171b4370f5701271f Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 10 Jul 2026 02:40:19 +0000 Subject: [PATCH 10/17] Limit PCP+TBO only when TP=1 --- atom/model_engine/llm_engine.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index c86c7b0269..7adfeb5c93 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -88,6 +88,26 @@ def __init__(self, model, tokenizer=None, **kwargs): ) 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 " From 5b98657efec42f86089cdaa4fcb96b9e997fc67d Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 10 Jul 2026 02:57:56 +0000 Subject: [PATCH 11/17] Update docs --- docs/context_parallel_guide.md | 39 ++++++++++++++++++++++++++++------ 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/docs/context_parallel_guide.md b/docs/context_parallel_guide.md index cfb256b415..65179fd026 100644 --- a/docs/context_parallel_guide.md +++ b/docs/context_parallel_guide.md @@ -29,6 +29,8 @@ TP and Expert Parallelism (EP): the total world size is `world = tp × pcp`. TTFT dominates. - **Combine with**: `--enable-tbo` (prefill) to overlap the MoE communication PCP introduces (see [Overlapping communication with TBO](#overlapping-communication-with-tbo)). + TBO is only usable with `ATOM_PCP_MOE_MERGE=1` **and `-tp 1`** (`tp > 1` hangs — + see [Constraints & Compatibility](#constraints--compatibility)). - **Requires**: `world = tp × pcp` GPUs, e.g. `-tp 4 -pcp 2` on 8 GPUs. - **Little benefit**: decode-heavy or short-prompt workloads — use TP/EP as usual. @@ -38,19 +40,22 @@ TP and Expert Parallelism (EP): the total world size is `world = tp × pcp`. |-----------------|---------|---------| | `-pcp N` / `--prefill-context-parallel-size N` | `1` | Enable PCP with size `N` (`world = tp × pcp`) | | `ATOM_PCP_MOE_MERGE` | `1` | Whether to shard MoE across the PCP ranks too | -| `--enable-tbo [prefill\|all]` | off | Overlap compute with PCP communication. With PCP, only prefill TBO is supported | +| `--enable-tbo [prefill\|all]` | off | Overlap compute with PCP communication. With PCP, only prefill TBO is supported, and only when `ATOM_PCP_MOE_MERGE=1` and `-tp 1` | | Goal (8 GPUs) | Command | |------|---------| | Long-context prefill | `-tp 4 -pcp 2` | -| Long-context prefill + overlap | `-tp 4 -pcp 2 --enable-tbo` | +| Long-context prefill + overlap | `-tp 1 -pcp 8 --enable-tbo` | | Disable PCP (baseline) | `-tp 8 -pcp 1` | +> **TBO requires `-tp 1`.** Two-batch overlap is only supported with +> `ATOM_PCP_MOE_MERGE=1` **and TP=1** (all GPUs go to PCP, e.g. `-tp 1 -pcp 8`). + ## CLI usage ```bash -pcp N # or --prefill-context-parallel-size N; world = tp × pcp ---enable-tbo # prefill-only TBO overlap (prefill only supported with PCP) +--enable-tbo # prefill-only TBO overlap; requires ATOM_PCP_MOE_MERGE=1 and -tp 1 (prefill only supported with PCP) ``` ```bash @@ -76,12 +81,16 @@ python -m atom.entrypoints.openai_server \ --kv_cache_dtype fp8 ``` -### DeepSeek-V4: TP4 + PCP2 + prefill TBO overlap (8 GPUs) +### DeepSeek-V4: TP1 + PCP8 + prefill TBO overlap (8 GPUs) + +TBO requires `ATOM_PCP_MOE_MERGE=1` (the default) **and `-tp 1`** — put every GPU +into PCP. With `tp > 1` the run hangs. ```bash +ATOM_PCP_MOE_MERGE=1 \ python -m atom.entrypoints.openai_server \ --model deepseek-ai/DeepSeek-V4 \ - -tp 4 -pcp 2 \ + -tp 1 -pcp 8 \ --enable-tbo \ --kv_cache_dtype fp8 ``` @@ -92,6 +101,8 @@ Tips: `ATOM_PCP_MOE_MERGE=1`. It only helps in that mode: with `ATOM_PCP_MOE_MERGE=0` there is no MoE communication to overlap, so TBO is auto-disabled (a warning is logged). +- `--enable-tbo` additionally requires `-tp 1`; with `tp > 1` the run **hangs** + (not yet supported). Give all GPUs to PCP instead, e.g. `-tp 1 -pcp 8`. - Under PCP, TBO uses a **request-boundary split** (each micro-batch is a whole subset of requests) — the non-default TBO split mode — instead of the token-midpoint split used without PCP. `ATOM_TBO_PREFILL_TOKEN_SPLIT` @@ -120,6 +131,19 @@ output throughput; the gap widens as `--random-input-len` grows. > reduction and up to **~49%** higher throughput on long prefill. Actual gains > depend on model, sequence length, and hardware. +### PCP + TBO (prefill overlap) + +TBO's benefit is **hardware-dependent**: + +- **MI308**: current testing shows PCP + TBO gives **almost no gain** over PCP + alone — the overlap does not meaningfully hide the MoE communication on this + hardware. Prefer plain PCP (`-tp 4 -pcp 2`) here. +- **MI355**: PCP + TBO **does** deliver a speedup. Enable it with + `ATOM_PCP_MOE_MERGE=1 -tp 1 -pcp 8 --enable-tbo`. + +Because TBO requires `-tp 1`, compare TBO configs against the matching `-tp 1` +PCP baseline (not the `-tp 4 -pcp 2` numbers above). + ## Constraints & Compatibility | Constraint | Notes | @@ -129,9 +153,10 @@ output throughput; the gap widens as `--random-input-len` grows. | PCP + DP-attention | Not supported (raises at startup) | | PCP + TBO decode (`--enable-tbo all`) | Not supported (raises at startup); use `--enable-tbo` prefill-only | | `ATOM_PCP_MOE_MERGE=0` + `--enable-tbo` | TBO auto-disabled (warning logged) | +| PCP + TBO with `tp > 1` | **Not supported — hangs.** TBO requires `-tp 1` (all GPUs in PCP, e.g. `-tp 1 -pcp 8`) | -PCP + TBO **prefill** is supported. Decode is unchanged by PCP in all -configurations. +PCP + TBO **prefill** is supported only with `ATOM_PCP_MOE_MERGE=1` **and +`-tp 1`**. Decode is unchanged by PCP in all configurations. ## How it works From 94a3b222ee620099eef0ecc4ee7a02fcc4deab20 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 10 Jul 2026 03:24:47 +0000 Subject: [PATCH 12/17] Update docs --- docs/context_parallel_guide.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/context_parallel_guide.md b/docs/context_parallel_guide.md index 65179fd026..4b8dc1ea22 100644 --- a/docs/context_parallel_guide.md +++ b/docs/context_parallel_guide.md @@ -29,8 +29,8 @@ TP and Expert Parallelism (EP): the total world size is `world = tp × pcp`. TTFT dominates. - **Combine with**: `--enable-tbo` (prefill) to overlap the MoE communication PCP introduces (see [Overlapping communication with TBO](#overlapping-communication-with-tbo)). - TBO is only usable with `ATOM_PCP_MOE_MERGE=1` **and `-tp 1`** (`tp > 1` hangs — - see [Constraints & Compatibility](#constraints--compatibility)). + TBO is only usable with `ATOM_PCP_MOE_MERGE=1` **and `-tp 1`** (see + [Constraints & Compatibility](#constraints--compatibility)). - **Requires**: `world = tp × pcp` GPUs, e.g. `-tp 4 -pcp 2` on 8 GPUs. - **Little benefit**: decode-heavy or short-prompt workloads — use TP/EP as usual. @@ -84,7 +84,7 @@ python -m atom.entrypoints.openai_server \ ### DeepSeek-V4: TP1 + PCP8 + prefill TBO overlap (8 GPUs) TBO requires `ATOM_PCP_MOE_MERGE=1` (the default) **and `-tp 1`** — put every GPU -into PCP. With `tp > 1` the run hangs. +into PCP. ```bash ATOM_PCP_MOE_MERGE=1 \ @@ -101,8 +101,9 @@ Tips: `ATOM_PCP_MOE_MERGE=1`. It only helps in that mode: with `ATOM_PCP_MOE_MERGE=0` there is no MoE communication to overlap, so TBO is auto-disabled (a warning is logged). -- `--enable-tbo` additionally requires `-tp 1`; with `tp > 1` the run **hangs** - (not yet supported). Give all GPUs to PCP instead, e.g. `-tp 1 -pcp 8`. +- `--enable-tbo` additionally requires `-tp 1`; with `tp > 1` it **may hang** under + long-sequence / high-concurrency workloads (not yet supported). Give all GPUs to + PCP instead, e.g. `-tp 1 -pcp 8`. - Under PCP, TBO uses a **request-boundary split** (each micro-batch is a whole subset of requests) — the non-default TBO split mode — instead of the token-midpoint split used without PCP. `ATOM_TBO_PREFILL_TOKEN_SPLIT` @@ -153,7 +154,7 @@ PCP baseline (not the `-tp 4 -pcp 2` numbers above). | PCP + DP-attention | Not supported (raises at startup) | | PCP + TBO decode (`--enable-tbo all`) | Not supported (raises at startup); use `--enable-tbo` prefill-only | | `ATOM_PCP_MOE_MERGE=0` + `--enable-tbo` | TBO auto-disabled (warning logged) | -| PCP + TBO with `tp > 1` | **Not supported — hangs.** TBO requires `-tp 1` (all GPUs in PCP, e.g. `-tp 1 -pcp 8`) | +| PCP + TBO with `tp > 1` | **Not supported.** May hang under long-sequence / high-concurrency workloads. TBO requires `-tp 1` (all GPUs in PCP, e.g. `-tp 1 -pcp 8`) | PCP + TBO **prefill** is supported only with `ATOM_PCP_MOE_MERGE=1` **and `-tp 1`**. Decode is unchanged by PCP in all configurations. From 461bb5e3a16439a02969e30df0fda86890069552 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 10 Jul 2026 06:36:28 +0000 Subject: [PATCH 13/17] Fix format --- atom/distributed/pcp_utils.py | 17 +++++++++------ atom/model_engine/model_runner.py | 8 +++++-- atom/model_ops/attentions/deepseek_v4_attn.py | 21 +++++++++++++------ atom/models/deepseek_v4.py | 4 ++++ atom/utils/envs.py | 2 +- 5 files changed, 37 insertions(+), 15 deletions(-) diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 7ad365ceac..4f601daa5f 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -29,11 +29,14 @@ class PcpBalGroup(NamedTuple): 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) + 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) + ) + from aiter.dist.parallel_state import ( get_pcp_group, @@ -170,7 +173,9 @@ def pcp_reduce_scatter( return get_pcp_group().reduce_scatter(input_.contiguous(), dim=0) -def pcp_all_reduce(input_: torch.Tensor, pcp_size: Optional[int] = None) -> torch.Tensor: +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. diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index ff1216c50d..530306083f 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1888,7 +1888,9 @@ def prepare_inputs(self, batch: ScheduledBatch, input_ids: torch.Tensor = None): pcp_size = self.config.prefill_context_parallel_size _pcp_tbo_balanced = ( - is_prefill and pcp_size > 1 and tbo_collective_active + is_prefill + and pcp_size > 1 + and tbo_collective_active and not batch.is_dummy_run and getattr(self, "_pcp_tbo_balanced_active", False) ) @@ -1910,7 +1912,9 @@ def prepare_inputs(self, batch: ScheduledBatch, input_ids: torch.Tensor = None): # request boundary whose cumulative token count is closest to target split_idx = int(np.searchsorted(cum, target, side="left")) + 1 split_idx = max(1, min(split_idx, num_prefill_reqs - 1)) - B = int(cum[split_idx - 1]) # global token count of group0 (reqs [0:split_idx]) + B = int( + cum[split_idx - 1] + ) # global token count of group0 (reqs [0:split_idx]) H0 = pcp_pad_len(B, pcp_size) H1 = pcp_pad_len(total_tok - B, pcp_size) l0 = H0 // pcp_size diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 7bf1beca90..8140530909 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1635,9 +1635,10 @@ def build_ubatch_prefill_metadata( # independent non-TBO PCP mini-batch. Slice the FULL (un-reindexed) # metadata to the group + call _apply_pcp_reindex on it (reuse the proven # reindex). Bypasses the token-split rebuild path entirely. - if getattr(self.model_runner, "_pcp_tbo_balanced_active", False) and getattr( - self.model_runner, "_pcp_bal_groups", None - ) is not None: + if ( + getattr(self.model_runner, "_pcp_tbo_balanced_active", False) + and getattr(self.model_runner, "_pcp_bal_groups", None) is not None + ): return self._build_ubatch_prefill_metadata_balanced( attn_metadata, ubatch_idx ) @@ -1826,7 +1827,9 @@ def _build_ubatch_prefill_metadata_balanced( if src.state_slot_mapping_cpu is not None: ub.state_slot_mapping_cpu = src.state_slot_mapping_cpu[rs0:rs1] if src.n_committed_csa_per_seq is not None: - ub.n_committed_csa_per_seq = src.n_committed_csa_per_seq[rs0:rs1].contiguous() + ub.n_committed_csa_per_seq = src.n_committed_csa_per_seq[ + rs0:rs1 + ].contiguous() if src.n_committed_csa_per_seq_cpu is not None: ub.n_committed_csa_per_seq_cpu = src.n_committed_csa_per_seq_cpu[rs0:rs1] if src.n_committed_hca_per_seq_cpu is not None: @@ -1867,7 +1870,9 @@ def _build_ubatch_prefill_metadata_balanced( # ---- compress_plans: group's GLOBAL per-request (compressor all-gathers # the group to full order). Built from global cu / context_lens slices. ---- if self._unique_compress_ratios_overlap: - gcu = var["cu_seqlens_q"].np # GLOBAL (not overwritten for request-boundary split) + gcu = var[ + "cu_seqlens_q" + ].np # GLOBAL (not overwritten for request-boundary split) ext = (gcu[rs0 + 1 : rs1 + 1] - gcu[rs0:rs1]).astype(np.int32) ctx = np.asarray(var["context_lens"].np[rs0:rs1], dtype=np.int32) plan_bufs = self._get_ubatch_compress_plan_buffers(ubatch_idx) @@ -1903,7 +1908,11 @@ def _build_ubatch_prefill_metadata_balanced( ub.n_committed_csa_per_seq = ub.n_committed_csa_per_seq.clone() if ub.indexer_meta is not None: im = ub.indexer_meta - for k in ("cu_committed_gpu", "batch_id_per_token_gpu", "n_committed_per_seq_gpu"): + for k in ( + "cu_committed_gpu", + "batch_id_per_token_gpu", + "n_committed_per_seq_gpu", + ): if im.get(k) is not None: im[k] = im[k].clone() return ub diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 34bacb3f51..272c30a44c 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1019,6 +1019,7 @@ def forward( tbo_yield_and_switch_from_compute_to_comm, tbo_switch_to_compute_sync, ) + _tbo = _tbo_active() if _tbo: tbo_yield_and_switch_from_compute_to_comm() @@ -1990,6 +1991,7 @@ def forward_impl( tbo_yield_and_switch_from_compute_to_comm, tbo_switch_to_compute_sync, ) + _tbo_attn = _tbo_active_attn() if _tbo_attn: tbo_yield_and_switch_from_compute_to_comm() @@ -2881,6 +2883,7 @@ def moe_pcp_merge_forward( tbo_yield_and_switch_from_compute_to_comm, tbo_switch_to_compute_sync, ) + _tbo = _tbo_active() # allgather: when TBO is active, yield to the partner ubatch so its # compute overlaps this collective on the comm stream (P2 overlap). @@ -3199,6 +3202,7 @@ def forward( full_padded_ids = None if use_pcp: from atom.utils.tbo.ubatching import tbo_active as _tbo_active + pcp_size = get_pcp_world_size() if _tbo_active(): # PCP+TBO prefill: the split was done upstream in run_model; diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 533899161a..394793b4fc 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -258,7 +258,7 @@ ), # --- PCP MoE comm mode --- # Fold the PCP (prefill-context-parallel) dim into the MoE tp/ep sharding. - # Only meaningful when prefill_context_parallel_size > 1; + # Only meaningful when prefill_context_parallel_size > 1; # Default "1": all-gather hidden 1/W -> full before MoE and slice # full -> 1/W after, so MoE sees the complete token set (MoE itself is # untouched / PCP-agnostic). Costs one extra hidden all-gather per layer. From eaad013af0e31e00338cfe2c0f1ab2c384e464ba Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 10 Jul 2026 06:38:09 +0000 Subject: [PATCH 14/17] Fix format --- atom/distributed/pcp_utils.py | 12 ++++++------ atom/model_engine/model_runner.py | 2 -- atom/model_ops/attentions/deepseek_v4_attn.py | 1 - 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 4f601daa5f..2c61431c1f 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -18,6 +18,12 @@ import torch +from aiter.dist.parallel_state import ( + get_pcp_group, + get_prefill_context_model_parallel_rank, + get_prefill_context_model_parallel_world_size, +) + class PcpBalGroup(NamedTuple): """One request group for PCP+TBO request-boundary split prefill. @@ -38,12 +44,6 @@ class PcpBalGroup(NamedTuple): ) -from aiter.dist.parallel_state import ( - get_pcp_group, - get_prefill_context_model_parallel_rank, - get_prefill_context_model_parallel_world_size, -) - logger = logging.getLogger("atom") diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 530306083f..4c90b95bce 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -46,7 +46,6 @@ ) from atom.distributed.pcp_utils import ( PcpBalGroup, - pcp_allgather_rankmajor, pcp_allgather_rerange, pcp_pad_len, pcp_round_robin_split, @@ -2088,7 +2087,6 @@ def run_model( and _pcp_bal_groups is not None ) if _pcp_tbo_balanced: - n_global = input_ids.shape[0] g_ids, g_pos = [], [] for grp in _pcp_bal_groups: seg_ids = input_ids[grp.tok_start : grp.tok_end] diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 8140530909..81fdc3b84f 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1799,7 +1799,6 @@ def _build_ubatch_prefill_metadata_balanced( grp = mr._pcp_bal_groups[ubatch_idx] # PcpBalGroup rs0, rs1 = grp.req_start, grp.req_stop gts, gte = grp.tok_start, grp.tok_end - H = grp.pad_total group_bs = rs1 - rs0 group_total = gte - gts # group's global token count (real, pre-pad) device = self.device From de6595b30bc7e7cd170e6953991f9b4570bf9f17 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 10 Jul 2026 07:49:29 +0000 Subject: [PATCH 15/17] Refactor model_runner --- atom/model_engine/model_runner.py | 165 ++++++++++++++++++------------ 1 file changed, 99 insertions(+), 66 deletions(-) diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 4c90b95bce..14a5f94352 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1894,48 +1894,11 @@ def prepare_inputs(self, batch: ScheduledBatch, input_ids: torch.Tensor = None): and getattr(self, "_pcp_tbo_balanced_active", False) ) if _pcp_tbo_balanced: - # split REQUESTS into two groups at a request boundary near the - # token midpoint. Each group is an independent "non-TBO PCP - # mini-batch": padded to a pcp multiple and round-robin striped - # as a whole, so every sequence stays intact in one group - # (root-fixes token-split R1/R2). forward_vars stay GLOBAL here; - # build_ubatch_prefill_metadata slices the FULL (un-reindexed) - # metadata per group and calls _apply_pcp_reindex on it. - num_prefill_reqs = batch.total_seqs_num_prefill - per_req = np.asarray( - num_scheduled_tokens[:num_prefill_reqs], dtype=np.int64 + # Request-boundary split for PCP+TBO prefill (see + # _build_pcp_balanced_slices). forward_vars stay GLOBAL here. + ubatch_slices, self._pcp_bal_groups = self._build_pcp_balanced_slices( + batch, num_scheduled_tokens, pcp_size ) - total_tok = int(per_req.sum()) - cum = np.cumsum(per_req) # cum[j] = sum of reqs [0..j] - target = total_tok // 2 - # request boundary whose cumulative token count is closest to target - split_idx = int(np.searchsorted(cum, target, side="left")) + 1 - split_idx = max(1, min(split_idx, num_prefill_reqs - 1)) - B = int( - cum[split_idx - 1] - ) # global token count of group0 (reqs [0:split_idx]) - H0 = pcp_pad_len(B, pcp_size) - H1 = pcp_pad_len(total_tok - B, pcp_size) - l0 = H0 // pcp_size - l1 = H1 // pcp_size - # token_slice is in the LOCAL concat space [g0_local | g1_local] that - # run_model produces (see per-group stripe there). - ubatch_slices = [ - UBatchSlice( - request_slice=slice(0, split_idx), - token_slice=slice(0, l0), - ), - UBatchSlice( - request_slice=slice(split_idx, num_prefill_reqs), - token_slice=slice(l0, l0 + l1), - ), - ] - # Per-group descriptors consumed by run_model (per-group stripe) and - # build_ubatch_prefill_metadata (slice+reindex). See PcpBalGroup. - self._pcp_bal_groups = [ - PcpBalGroup(0, split_idx, 0, B, H0), - PcpBalGroup(split_idx, num_prefill_reqs, B, total_tok, H1), - ] else: ubatch_slices = self._maybe_create_tbo_slices( batch, @@ -2024,6 +1987,95 @@ def prepare_model(self, batch: ScheduledBatch): needs_independent_noise, ) + def _build_pcp_balanced_slices( + self, + batch: ScheduledBatch, + num_scheduled_tokens: np.ndarray, + pcp_size: int, + ) -> "tuple[list[UBatchSlice], list[PcpBalGroup]]": + """Build request-boundary-split ubatch slices for PCP+TBO prefill. + + Split REQUESTS into two groups at a request boundary near the token + midpoint. Each group is an independent "non-TBO PCP mini-batch": padded + to a pcp multiple and round-robin striped as a whole, so every sequence + stays intact in one group (root-fixes token-split R1/R2). forward_vars + stay GLOBAL here; build_ubatch_prefill_metadata slices the FULL + (un-reindexed) metadata per group and calls _apply_pcp_reindex on it. + + Returns (ubatch_slices, groups): token_slice is in the LOCAL concat + space [g0_local | g1_local] that run_model produces (see + _apply_pcp_balanced_stripe); groups are the PcpBalGroup descriptors + consumed by run_model (per-group stripe) and + build_ubatch_prefill_metadata (slice + reindex). + """ + num_prefill_reqs = batch.total_seqs_num_prefill + per_req = np.asarray(num_scheduled_tokens[:num_prefill_reqs], dtype=np.int64) + total_tok = int(per_req.sum()) + cum = np.cumsum(per_req) # cum[j] = sum of reqs [0..j] + target = total_tok // 2 + # request boundary whose cumulative token count is closest to target + split_idx = int(np.searchsorted(cum, target, side="left")) + 1 + split_idx = max(1, min(split_idx, num_prefill_reqs - 1)) + # global token count of group0 (reqs [0:split_idx]) + b0 = int(cum[split_idx - 1]) + h0 = pcp_pad_len(b0, pcp_size) + h1 = pcp_pad_len(total_tok - b0, pcp_size) + l0 = h0 // pcp_size + l1 = h1 // pcp_size + ubatch_slices = [ + UBatchSlice( + request_slice=slice(0, split_idx), + token_slice=slice(0, l0), + ), + UBatchSlice( + request_slice=slice(split_idx, num_prefill_reqs), + token_slice=slice(l0, l0 + l1), + ), + ] + groups = [ + PcpBalGroup(0, split_idx, 0, b0, h0), + PcpBalGroup(split_idx, num_prefill_reqs, b0, total_tok, h1), + ] + return ubatch_slices, groups + + def _apply_pcp_balanced_stripe( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + groups: "list[PcpBalGroup]", + pcp_size: int, + forward_context, + ) -> "tuple[torch.Tensor, torch.Tensor]": + """PCP+TBO prefill per-group round-robin stripe, before UBatchWrapper. + + Each request group is padded to a pcp multiple and round-robin striped + as a WHOLE (so sequences stay intact per group), then the two groups' + 1/pcp shards are concatenated into [g0_local | g1_local]. token_slice + (built in prepare_inputs) indexes into this concat. Returns the striped + (input_ids, positions). + """ + g_ids, g_pos = [], [] + for grp in groups: + seg_ids = input_ids[grp.tok_start : grp.tok_end] + seg_pos = positions[grp.tok_start : grp.tok_end] + pad = grp.pad_total - (grp.tok_end - grp.tok_start) + if pad > 0: + seg_ids = torch.cat([seg_ids, seg_ids.new_zeros(pad)]) + seg_pos = torch.cat([seg_pos, seg_pos.new_zeros(pad)]) + g_ids.append(pcp_round_robin_split(seg_ids, pcp_size)) + g_pos.append(pcp_round_robin_split(seg_pos, pcp_size)) + input_ids = torch.cat(g_ids) + positions = torch.cat(g_pos) + # context.positions = local per-group concat so _make_ubatch_context + # slices each ubatch's forward positions correctly. + forward_context.context.positions = positions + # Hash MoE: local per-group-concat ids. Each ForCausalLM.forward + # allgathers its ubatch's slice (g_i local, H_i/pcp) across pcp ranks → + # H_i ids, matching moe_pcp_merge_forward's per-ubatch hidden allgather. + if envs.ATOM_PCP_MOE_MERGE: + forward_context.context.input_ids = input_ids + return input_ids, positions + def _restore_pcp_balanced_output( self, mo: torch.Tensor, @@ -2071,11 +2123,9 @@ def run_model( # internally short-circuits for prefill / cudagraph. forward_mode.assert_shape_contract(input_ids, forward_context.attn_metadata) - # PCP+TBO prefill: PER-GROUP stripe here, before UBatchWrapper. Each - # request group is padded to a pcp multiple and round-robin striped as a - # WHOLE (so sequences stay intact per group), then the two groups' 1/pcp - # shards are concatenated. token_slice (built in prepare_inputs) indexes - # into this [g0_local | g1_local] concat. + # PCP+TBO prefill: per-group round-robin stripe before UBatchWrapper (see + # _apply_pcp_balanced_stripe). _pcp_tbo_balanced also gates the per-group + # output restore further below. _pcp_size = self.config.prefill_context_parallel_size _pcp_bal_groups = getattr(self, "_pcp_bal_groups", None) _pcp_tbo_balanced = ( @@ -2087,26 +2137,9 @@ def run_model( and _pcp_bal_groups is not None ) if _pcp_tbo_balanced: - g_ids, g_pos = [], [] - for grp in _pcp_bal_groups: - seg_ids = input_ids[grp.tok_start : grp.tok_end] - seg_pos = positions[grp.tok_start : grp.tok_end] - pad = grp.pad_total - (grp.tok_end - grp.tok_start) - if pad > 0: - seg_ids = torch.cat([seg_ids, seg_ids.new_zeros(pad)]) - seg_pos = torch.cat([seg_pos, seg_pos.new_zeros(pad)]) - g_ids.append(pcp_round_robin_split(seg_ids, _pcp_size)) - g_pos.append(pcp_round_robin_split(seg_pos, _pcp_size)) - input_ids = torch.cat(g_ids) - positions = torch.cat(g_pos) - # context.positions = local per-group concat so _make_ubatch_context - # slices each ubatch's forward positions correctly. - forward_context.context.positions = positions - # Hash MoE: local per-group-concat ids. Each ForCausalLM.forward - # allgathers its ubatch's slice (g_i local, H_i/pcp) across pcp ranks → - # H_i ids, matching moe_pcp_merge_forward's per-ubatch hidden allgather. - if envs.ATOM_PCP_MOE_MERGE: - forward_context.context.input_ids = input_ids + input_ids, positions = self._apply_pcp_balanced_stripe( + input_ids, positions, _pcp_bal_groups, _pcp_size, forward_context + ) if not forward_mode.use_cudagraph: # prefill, or decode forced eager (enforce_eager / DP peer From 272505e063915803607291e18fce7d0771c6c789 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Fri, 10 Jul 2026 09:11:28 +0000 Subject: [PATCH 16/17] Update docs --- docs/context_parallel_guide.md | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/docs/context_parallel_guide.md b/docs/context_parallel_guide.md index 4b8dc1ea22..f3fd678e95 100644 --- a/docs/context_parallel_guide.md +++ b/docs/context_parallel_guide.md @@ -32,7 +32,11 @@ TP and Expert Parallelism (EP): the total world size is `world = tp × pcp`. TBO is only usable with `ATOM_PCP_MOE_MERGE=1` **and `-tp 1`** (see [Constraints & Compatibility](#constraints--compatibility)). - **Requires**: `world = tp × pcp` GPUs, e.g. `-tp 4 -pcp 2` on 8 GPUs. -- **Little benefit**: decode-heavy or short-prompt workloads — use TP/EP as usual. +- **Little benefit / avoid**: decode-heavy or short-prompt workloads. PCP only + shards **prefill** tokens to lower TTFT on long sequences; decode is left + unsharded and runs redundantly across the PCP ranks. For long-decode + (large `output_len`) workloads PCP can **hurt TPOT** — do not enable it there; + use TP/EP as usual. ## Quick Reference @@ -41,6 +45,8 @@ TP and Expert Parallelism (EP): the total world size is `world = tp × pcp`. | `-pcp N` / `--prefill-context-parallel-size N` | `1` | Enable PCP with size `N` (`world = tp × pcp`) | | `ATOM_PCP_MOE_MERGE` | `1` | Whether to shard MoE across the PCP ranks too | | `--enable-tbo [prefill\|all]` | off | Overlap compute with PCP communication. With PCP, only prefill TBO is supported, and only when `ATOM_PCP_MOE_MERGE=1` and `-tp 1` | +| `--no-enable_chunked_prefill` | chunked on | Disable chunked prefill. Recommended with PCP (see [Tuning for long sequences](#tuning-for-long-sequences)) | +| `--max-num-batched-tokens N` | `16384` | Max tokens scheduled per step. Raise (e.g. `131072`) for long-sequence PCP so a full long prompt is prefilled in one step | | Goal (8 GPUs) | Command | |------|---------| @@ -111,6 +117,26 @@ Tips: - `--torch-profiler-dir ./log` can be added to collect traces for performance analysis. +## Tuning for long sequences + +PCP targets long-prefill TTFT, but the default scheduler settings work against it. When serving long prompts: + +- **Disable chunked prefill** (`--no-enable_chunked_prefill`), or enlarge the chunk (`--attn_prefill_chunk_size`). It is on by default and splits a long prompt across steps, so PCP gets fewer tokens to shard per step and TBO's request-boundary split (needs ≥2 whole sequences per step) falls back to the non-overlapped path. +- **Raise `--max-num-batched-tokens`** to `131072` (≥ `input_len`; ≥ `2 × input_len` for TBO's balanced split). The default `16384` forces a long prompt across multiple steps instead of prefilling it in one. +- **Avoid PCP for long-decode workloads.** PCP shards only prefill; decode runs unsharded. A large `output_len` is decode-dominated, where PCP adds no benefit and can raise TPOT — use plain TP/EP. + +Example (long-context prefill, chunked off, larger batch budget): + +```bash +ATOM_PCP_MOE_MERGE=1 \ +python -m atom.entrypoints.openai_server \ + --model deepseek-ai/DeepSeek-V4 \ + -tp 1 -pcp 8 --enable-tbo \ + --no-enable_chunked_prefill \ + --max-num-batched-tokens 131072 \ + --kv_cache_dtype fp8 +``` + ## Performance baseline Benchmark against a running server with a long input length (PCP targets prefill): From ce9e8a49bc299b5107b28a5cf8e3f21f32dc4da5 Mon Sep 17 00:00:00 2001 From: Wang Yiting Date: Mon, 13 Jul 2026 06:47:23 +0000 Subject: [PATCH 17/17] Fix with PR1423 --- atom/model_ops/attentions/deepseek_v4_attn.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 4251c9ed81..069b4e8e76 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1933,6 +1933,11 @@ def _build_ubatch_prefill_metadata_balanced( ub.n_committed_csa_per_seq_cpu = src.n_committed_csa_per_seq_cpu[rs0:rs1] if src.n_committed_hca_per_seq_cpu is not None: ub.n_committed_hca_per_seq_cpu = src.n_committed_hca_per_seq_cpu[rs0:rs1] + # paged-SWA block tables (added by #1423): per-request [bs, MB], required + # by swa_write in prefill. split_attn_metadata does not carry this DSV4 + # field, so slice it to the group's requests explicitly (else None -> crash). + if src.swa_block_tables is not None: + ub.swa_block_tables = src.swa_block_tables[rs0:rs1].contiguous() # ---- per-token DSV4 fields sliced by the GLOBAL token range [gts,gte) ---- owned = torch.arange(gts, gte, device=device)