diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 79ccf37585..2c61431c1f 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -13,7 +13,8 @@ `layers/utils/cp_utils.py:cp_all_gather_rerange_output`). """ -from typing import Optional +import logging +from typing import NamedTuple, Optional import torch @@ -24,6 +25,28 @@ ) +class PcpBalGroup(NamedTuple): + """One request group for PCP+TBO request-boundary split prefill. + + A prefill batch is split into request groups at request boundaries (never + inside a sequence); each group is processed as an independent non-TBO PCP + mini-batch (padded to a pcp multiple, round-robin striped, reindexed on its + own). Consumed by ModelRunner.run_model (per-group stripe / restore) and the + attn builder's `_build_ubatch_prefill_metadata_balanced` (slice + reindex). + """ + + req_start: int # first request index of this group (inclusive) + req_stop: int # last request index of this group (exclusive) + tok_start: int # global token offset of the group's first token + tok_end: int # global token offset past the group's last REAL token + pad_total: ( + int # tok count padded to a pcp multiple = pcp_pad_len(tok_end-tok_start, pcp) + ) + + +logger = logging.getLogger("atom") + + def get_pcp_world_size() -> int: return get_prefill_context_model_parallel_world_size() @@ -39,22 +62,27 @@ def pcp_is_enabled() -> bool: def pcp_pad_len( total_tokens: int, pcp_size: Optional[int] = None, + multiple: int = 1, ) -> int: - """Padded token count so the global sequence is divisible by pcp_size. + """Padded token count so the global sequence is divisible by pcp_size * multiple. Round-robin split requires the global token count to be divisible by pcp_size - (see SGLang `can_dsa_cp_split` assert / HIP `apply_cp_reindex`). Returns the - padded length (>= total_tokens); callers pad per-token tensors to this - length with dummy tokens (KV length 0) before splitting. + (see SGLang `can_dsa_cp_split` assert / HIP `apply_cp_reindex`). `multiple` is + an extra factor applied on top of pcp_size when the sequence must additionally + be evenly divisible by some multiplier. Returns the padded length + (>= total_tokens); callers pad per-token tensors to this length with dummy + tokens (KV length 0) before splitting. + """ if pcp_size is None: pcp_size = get_pcp_world_size() - if pcp_size <= 1: + divisor = pcp_size * max(multiple, 1) + if divisor <= 1: return total_tokens - rem = total_tokens % pcp_size + rem = total_tokens % divisor if rem == 0: return total_tokens - return total_tokens + (pcp_size - rem) + return total_tokens + (divisor - rem) def pcp_round_robin_split( @@ -111,6 +139,55 @@ def pcp_allgather_rerange( return out +# ==== MoE-path PCP collectives (rank-major gather + reduce_scatter) ==== +# Rank-major all_gather + reduce_scatter are a mutually-inverse pair: +# - gather (1/W -> full): all_gather(dim=0) concats rank-major, so rank r's +# 1/W stripe lands at rows [r*L:(r+1)*L]. MoE is per-token so the rank-major +# (not global) order is fine. +# - reduce_scatter (full partial-sum -> 1/W): sums the pcp-half across ranks +# AND scatters dim0 back so rank r receives the summed chunk r == its own +# original stripe tokens. No rerange/slice needed. + + +def pcp_allgather_rankmajor( + input_: torch.Tensor, pcp_size: Optional[int] = None +) -> torch.Tensor: + """Gather this rank's 1/W stripe shard into the full rank-major sequence + via a plain all_gather (dim=0). Inverse of pcp_reduce_scatter.""" + if pcp_size is None: + pcp_size = get_pcp_world_size() + if pcp_size <= 1: + return input_ + return get_pcp_group().all_gather(input_.contiguous(), dim=0) + + +def pcp_reduce_scatter( + input_: torch.Tensor, pcp_size: Optional[int] = None +) -> torch.Tensor: + """Sum the pcp-half across ranks and scatter dim0 back to this rank's 1/W + stripe via a plain reduce_scatter (dim=0). Inverse of pcp_allgather_rankmajor.""" + if pcp_size is None: + pcp_size = get_pcp_world_size() + if pcp_size <= 1: + return input_ + return get_pcp_group().reduce_scatter(input_.contiguous(), dim=0) + + +def pcp_all_reduce( + input_: torch.Tensor, pcp_size: Optional[int] = None +) -> torch.Tensor: + """All-reduce (sum) over the PCP group, no token reshaping. DECODE path: + tokens are pcp-redundant (every rank holds the same full batch), so just sum + the pcp-half of the intermediate that combine_outputs' tp all_reduce missed. + Uses aiter's compile-safe custom-op all_reduce. + """ + if pcp_size is None: + pcp_size = get_pcp_world_size() + if pcp_size <= 1: + return input_ + return get_pcp_group().all_reduce(input_) + + def pcp_round_robin_query_indices( n_global_q: int, pcp_size: Optional[int] = None, pcp_rank: Optional[int] = None ) -> torch.Tensor: diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index a20b019dc2..7adfeb5c93 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -70,20 +70,58 @@ def __init__(self, model, tokenizer=None, **kwargs): "them (use -tp N -pcp M without DP-attention, or -dp N " "--enable-dp-attention without -pcp)." ) - # PCP and TBO (two-batch overlap) are not yet compatible: TBO's - # UBatchWrapper calls ForCausalLM.forward once per micro-batch with a - # sub-slice of tokens, and PCP stripe-splits at that entry — so each - # ubatch would be split independently (double-split), giving each rank - # ~1/(2*pcp) tokens and a corrupted all-gather restore. + # PCP + TBO prefill: supported via coordinated splitting. + # PCP + TBO decode: not yet supported (pcp_all_reduce semantics under + # per-request ubatch split are unverified). if config.prefill_context_parallel_size > 1 and config.enable_tbo: - raise ValueError( - "prefill_context_parallel_size > 1 (-pcp) combined with " - "--enable-tbo is not supported yet (may be supported in a " - "future release): TBO calls ForCausalLM.forward per micro-batch " - "with a token sub-slice, so PCP would stripe-split each ubatch " - "independently (double-split) and corrupt the output. For now, " - "disable one of them." - ) + # TBO overlaps compute with the attn<->MoE PCP collectives, which + # only exist in MoE merge mode (ATOM_PCP_MOE_MERGE=1). With + # ATOM_PCP_MOE_MERGE=0, MoE runs on each rank's 1/W token shard + # with NO extra comm between attn and MoE, so TBO has nothing + # to overlap and only adds ubatch-splitting overhead. Force it off. + if not envs.ATOM_PCP_MOE_MERGE: + logger.warning( + "Disabling TBO because ATOM_PCP_MOE_MERGE=0: in this " + "situation it runs MoE on each rank's 1/W token shard with " + "no extra attn<->MoE communication, so TBO has nothing to " + "overlap and only adds overhead." + ) + config.enable_tbo = False + config.enable_tbo_decode = False + elif config.tensor_parallel_size > 1: + # Cross-communicator (PCP x TP) RCCL deadlock: + # under TBO the PCP collectives (comm_stream) run concurrently + # and UNORDERED with the TP-group all_reduces (compute_stream, + # from attention wo_b / MoE RowParallelLinear). On the TPxPCP + # rank grid with large collectives this forms a cross-rank + # circular wait -> hang (reproduced MI355 TP4PCP2/TP2PCP4 merge + # +TBO, 64k/c32). Serialized (non-TBO single stream) is fine, and + # TP=1 has no TP communicator so it cannot form the cycle. Only + # TP=1 + PCP keeps TBO; TP>1 falls back to non-TBO. + logger.warning( + "Disabling TBO: prefill_context_parallel_size > 1 (-pcp) " + "with tensor_parallel_size > 1 (-tp) hangs under TBO due to " + "a PCP<->TP cross-communicator RCCL deadlock (concurrent " + "unordered collectives on comm/compute streams; see " + "PCP_TBO.md 14.4). TBO with PCP is only supported at -tp 1 " + "(e.g. -tp 1 -pcp 8)." + ) + config.enable_tbo = False + config.enable_tbo_decode = False + elif config.enable_tbo_decode: + raise ValueError( + "prefill_context_parallel_size > 1 (-pcp) combined with " + "--enable-tbo all (decode TBO) is not supported yet. " + "Use --enable-tbo (prefill only) with -pcp." + ) + # Under PCP, TBO prefill uses a request-boundary split (the + # non-default TBO split mode; never token-midpoint split), so + # ATOM_TBO_PREFILL_TOKEN_SPLIT is ignored. + if config.enable_tbo and envs.ATOM_TBO_PREFILL_TOKEN_SPLIT: + logger.warning( + "ATOM_TBO_PREFILL_TOKEN_SPLIT is ignored under PCP: TBO " + "prefill uses request-boundary balanced grouping." + ) self.rquest_ids = set() self.io_processor = InputOutputProcessor( config, self.tokenizer, config.kv_cache_block_size diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index e7768679b8..5cd7ed6b7f 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -39,11 +39,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 ( + PcpBalGroup, + pcp_allgather_rerange, + pcp_pad_len, + pcp_round_robin_split, +) from atom.utils.forward_context import ( Context, DPMetadata, @@ -1834,6 +1841,31 @@ def _preprocess( self.config, batch, is_prefill, num_scheduled_tokens ) + # 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 (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 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 + # 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 + n_prefill = batch.total_tokens_num_prefill + # 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. # dp_uniform_decode=True mirrors the DP-disabled case in the @@ -1954,14 +1986,29 @@ 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_balanced = ( + is_prefill + and pcp_size > 1 + and tbo_collective_active + and not batch.is_dummy_run + and getattr(self, "_pcp_tbo_balanced_active", False) + ) + if _pcp_tbo_balanced: + # 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 + ) + 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, @@ -2041,6 +2088,117 @@ 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, + groups: "list[PcpBalGroup]", + pcp_size: int, + ) -> torch.Tensor: + """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 + 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,6 +2224,24 @@ 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 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 = ( + _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_balanced: + 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 # prefill / bs above the largest captured graph). @@ -2115,6 +2291,25 @@ def run_model( model_output = self.model( input_ids, positions, inputs_embeds=inputs_embeds ) + # 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 + # 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 + model_output = ( + self._restore_pcp_balanced_output( + _h, _pcp_bal_groups, _pcp_size + ), + _aux, + ) + else: + 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 f37da9b235..069b4e8e76 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1567,7 +1567,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 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. + _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 @@ -1613,6 +1618,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 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 owned_q = pcp_round_robin_query_indices(padded_total, pcp_size).to(device) @@ -1703,9 +1711,30 @@ 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 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). + - Token-split TBO (default, §11): uses `ub_slice` / `padded_bs`. + """ from atom.utils.tbo.ubatch_splitting import split_attn_metadata + # 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. + 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 @@ -1847,6 +1876,151 @@ 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 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 + 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 + 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] + # 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) + 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 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) + 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 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) + + # 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 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: + 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, diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index ad2a5c57f2..e6d4887c0f 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -150,6 +150,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 = ( + envs.ATOM_PCP_MOE_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 2efb436f09..4379fc60db 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -41,6 +41,16 @@ from aiter.dist.parallel_state import ( get_tensor_model_parallel_world_size, ) +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 ( @@ -1064,7 +1074,18 @@ 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 @@ -2242,6 +2263,15 @@ def _attn_core( 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`). @@ -2250,6 +2280,8 @@ def _attn_core( # 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 @@ -2473,10 +2505,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__( @@ -2585,9 +2619,13 @@ 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 + # — the latter requires registration regardless of + # dual-stream, so register whenever either consumer is active. + _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 @@ -2702,6 +2740,18 @@ def combine_outputs( all-reduce across TP ranks. """ if shared is not None: + # 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 + # 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) @@ -2792,6 +2842,9 @@ def __init__( indexer_stream=indexer_stream, ) self.ffn = MoE(layer_id, args, prefix=f"{prefix}.ffn", alt_stream=alt_stream) + self._moe_merge_enabled = get_pcp_world_size() > 1 and bool( + envs.ATOM_PCP_MOE_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 @@ -2987,9 +3040,12 @@ def forward( self.norm_eps, ) x = hc_state.x_prev - x = self.ffn( - x - ) # [num_tokens, dim] (input_ids read from forward_context for hash MoE) + 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 @@ -3060,6 +3116,80 @@ def forward( return self.get_logits(norm(x)) # [bs, vocab] +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: + 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 + + +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. @@ -3073,6 +3203,37 @@ def _pcp_active() -> bool: return fc.context.is_prefill and not fc.context.is_dummy_run +def _moe_pcp_merge_active() -> bool: + """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 + 1/W shard with no extra comm. + """ + if not _pcp_active(): + return False + return bool(envs.ATOM_PCP_MOE_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, 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(envs.ATOM_PCP_MOE_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. @@ -3161,7 +3322,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 ) @@ -3294,15 +3455,46 @@ 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). + 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. + 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." + ) + 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) + if _tbo_active(): + # 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: + 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) + # 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 @@ -3317,15 +3509,34 @@ 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: + 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 bbe542b7ee..aeeabf3201 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -273,6 +273,14 @@ "ATOM_TBO_PREFILL_MIN_TOKENS": lambda: int( os.getenv("ATOM_TBO_PREFILL_MIN_TOKENS", "8192") ), + # --- 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", # --- 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 b21a7eca36..d3dbc9dd22 100644 --- a/atom/utils/tbo/ubatch_wrapper.py +++ b/atom/utils/tbo/ubatch_wrapper.py @@ -542,6 +542,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( diff --git a/docs/context_parallel_guide.md b/docs/context_parallel_guide.md new file mode 100644 index 0000000000..f3fd678e95 --- /dev/null +++ b/docs/context_parallel_guide.md @@ -0,0 +1,213 @@ +# 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)). + 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 / 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 + +| 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, 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 | +|------|---------| +| Long-context prefill | `-tp 4 -pcp 2` | +| 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; requires ATOM_PCP_MOE_MERGE=1 and -tp 1 (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: TP1 + PCP8 + prefill TBO overlap (8 GPUs) + +TBO requires `ATOM_PCP_MOE_MERGE=1` (the default) **and `-tp 1`** — put every GPU +into PCP. + +```bash +ATOM_PCP_MOE_MERGE=1 \ +python -m atom.entrypoints.openai_server \ + --model deepseek-ai/DeepSeek-V4 \ + -tp 1 -pcp 8 \ + --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). +- `--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` + therefore has no effect when PCP is enabled. +- `--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): + +```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. + +### 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 | +|-----------|-------| +| 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 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. + +## 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