From 1090531f8283e164bbc2da7a65a69de67989e6be Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Thu, 23 Jul 2026 08:22:21 +0000 Subject: [PATCH 01/10] [DSV4] perf(tbo): overlap pure-TP all_reduce --- .github/benchmark/models.json | 8 ++++ atom/model_engine/engine_core.py | 23 ++++++++++++ atom/model_engine/prefill_delayer.py | 35 +++++++++++------- atom/model_ops/linear.py | 10 ++++- atom/model_ops/module_dispatch_ops.py | 53 +++++++++++++++++++++++++++ atom/models/deepseek_v4.py | 42 ++++++++++++++++++++- atom/utils/envs.py | 33 ++++++++++++++--- atom/utils/tbo/ubatching.py | 50 ++++++++++++++++++++++++- 8 files changed, 231 insertions(+), 23 deletions(-) diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index 8cd79f1a4b..f13a53a8bf 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -35,6 +35,14 @@ "conc_min": 4, "conc_max": 256 }, + { + "label": "TBO", + "suffix": "-tbo", + "extra_args": "--enable-tbo", + "env_vars": "GPU_MAX_HW_QUEUES=5\nATOM_NUMA_BIND=1", + "conc_min": 64, + "conc_max": 256 + }, { "label": "DPA", "suffix": "-dpa", diff --git a/atom/model_engine/engine_core.py b/atom/model_engine/engine_core.py index 6d62eb2983..8448eca81c 100644 --- a/atom/model_engine/engine_core.py +++ b/atom/model_engine/engine_core.py @@ -128,6 +128,29 @@ def __init__(self, config: Config, input_address: str, output_address: str): if not config.disagg_is_decode: self.scheduler = Scheduler(config) + # Enable delayer prefill when TP TBO on + if ( + config.enable_tbo + and envs.ATOM_ENABLE_PREFILL_DELAYER + and config.parallel_config.data_parallel_size <= 1 + ): + from atom.model_engine.prefill_delayer import PrefillDelayer + + self.scheduler.set_prefill_delayer( + PrefillDelayer( + dp_size=1, + cpu_group=None, + max_num_batched_tokens=config.max_num_batched_tokens, + target_fill=envs.ATOM_PREFILL_DELAYER_TARGET_FILL, + ttft_max_ticks=envs.ATOM_PREFILL_DELAYER_TTFT_MAX_TICKS, + partial_max_ticks=envs.ATOM_PREFILL_DELAYER_PARTIAL_MAX_TICKS, + stall_ticks=envs.ATOM_PREFILL_DELAYER_STALL_TICKS, + kv_high_watermark=envs.ATOM_PREFILL_DELAYER_KV_HIGH_WATERMARK, + token_usage_low_watermark=envs.ATOM_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK, + max_queue_ms=envs.ATOM_PREFILL_DELAYER_MAX_QUEUE_MS, + ) + ) + self.kv_transfer_enabled = bool(config.kv_transfer_config) if self.kv_transfer_enabled: self.kv_aggregator = KVOutputAggregator( diff --git a/atom/model_engine/prefill_delayer.py b/atom/model_engine/prefill_delayer.py index f723e375c5..11640b593f 100644 --- a/atom/model_engine/prefill_delayer.py +++ b/atom/model_engine/prefill_delayer.py @@ -14,6 +14,14 @@ While it holds, decode keeps running (nothing is wasted); TTFT is bounded so a held request never starves. +Single-rank / TP-only mode +-------------------------- +With ``cpu_group=None`` (``dp_size=1``) the delayer runs the same coalescer +locally and skips the cross-rank ``all_reduce`` — a single scheduler drives all +TP workers, so there is no cross-rank phase to align. Only the batching value +remains (fill / stall / ttft / kv bounds decide FIRE/HOLD from this rank's own +counts); the ``n_prefillable < dp_size`` alignment gate is vacuous. + It is NOT about mixing prefill+decode in one forward. It DOES preserve cross-DP phase alignment: it only releases when every rank is prefill-ready (so all ranks enter prefill together and the MoE collective stays aligned), except when a @@ -282,19 +290,20 @@ def should_allow_prefill( self._reduce_buf[4] = 1 if kv_low else 0 self._reduce_buf[5] = 1 if has_partial else 0 self._reduce_buf[6] = 1 if queue_hot else 0 - try: - torch.distributed.all_reduce( - self._reduce_buf, - op=torch.distributed.ReduceOp.SUM, - group=self.cpu_group, - ) - except RuntimeError: - logger.warning( - "PrefillDelayer all_reduce failed (peer down?); admitting prefill " - "and deferring to the DP-state shutdown barrier." - ) - self._reset() - return True + if self.cpu_group is not None: + try: + torch.distributed.all_reduce( + self._reduce_buf, + op=torch.distributed.ReduceOp.SUM, + group=self.cpu_group, + ) + except RuntimeError: + logger.warning( + "PrefillDelayer all_reduce failed (peer down?); admitting " + "prefill and deferring to the DP-state shutdown barrier." + ) + self._reset() + return True # One host<-device readback for all 7 slots (a single .tolist() beats # seven .item() boundary crossings on this per-tick hot path). diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 50a15f3010..0286d998d5 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -924,7 +924,10 @@ def forward( if self.bias is not None: y += self.bias if self.tp_dim == 1 and self.tp_size > 1 and self.reduce_results: - y = get_tp_group().all_reduce(y, ca_fp8_quant=False) + if getattr(self, "tbo_aware", False): + y = torch.ops.aiter.tbo_all_reduce(y) + else: + y = get_tp_group().all_reduce(y, ca_fp8_quant=False) return y @@ -1818,6 +1821,7 @@ def __init__( bias: bool = False, quant_config: Optional[QuantizationConfig] = None, reduce_results: bool = True, + tbo_aware: bool = False, source_quant_dtype: torch.dtype = None, prefix: str = "", **kwargs, @@ -1833,6 +1837,10 @@ def __init__( source_quant_dtype=source_quant_dtype, prefix=prefix, ) + # When True, the built-in TP reduce goes through the TBO-aware custom op + # (overlaps with the partner ubatch under TBO). Default off so all other + # RowParallelLinear layers keep the plain all_reduce unchanged. + self.tbo_aware = tbo_aware def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): param_data = param.data diff --git a/atom/model_ops/module_dispatch_ops.py b/atom/model_ops/module_dispatch_ops.py index cb7a2ae17a..77be66e8fa 100644 --- a/atom/model_ops/module_dispatch_ops.py +++ b/atom/model_ops/module_dispatch_ops.py @@ -149,3 +149,56 @@ def _indexer_score_topk_fake( fake_impl=_indexer_score_topk_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) + + +def tbo_all_reduce(x: torch.Tensor) -> torch.Tensor: + from aiter.dist.communication_op import tensor_model_parallel_all_reduce + + from atom.utils.tbo.ubatching import tbo_active + + if not tbo_active(): + return tensor_model_parallel_all_reduce(x) + + # Default "inline": correct Plan-A baseline (no overlap, never hangs). Only + # move the AR onto the comm stream when explicitly opted into "overlap". + if envs.ATOM_TBO_TP_AR_MODE != "overlap": + return tensor_model_parallel_all_reduce(x) + + from atom.utils.tbo.ubatching import ( + tbo_current_ubatch_id, + tbo_get_comm_stream, + tbo_get_ubatch_tp_comm, + tbo_switch_to_compute_sync, + tbo_yield_and_switch_from_compute_to_comm, + ) + + ubatch_id = tbo_current_ubatch_id() + + ub_comm = tbo_get_ubatch_tp_comm(ubatch_id) + if ub_comm is None: + # world_size == 1: nothing to reduce. + return x + + # Hand the CPU baton to the partner ubatch and move onto the comm stream, + # so this AR overlaps the partner's compute. + tbo_yield_and_switch_from_compute_to_comm() + comm_stream = tbo_get_comm_stream() + # out-of-place all_reduce on this ubatch's dedicated communicator, launched + # on the comm stream (current stream after the switch above). + x = ub_comm.all_reduce(x, stream=comm_stream) + + tbo_switch_to_compute_sync() + return x + + +def _tbo_all_reduce_fake(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + +direct_register_custom_op( + op_name="tbo_all_reduce", + op_func=tbo_all_reduce, + mutates_args=(), + fake_impl=_tbo_all_reduce_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 703f6f3b25..fcb95ab9ff 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1749,6 +1749,24 @@ def _score_topk_decode_ragged( # --------------------------------------------------------------------------- +def _tbo_aware_tp_reduce(tp_size: int) -> bool: + """True iff a pure-TP all_reduce should route through the TBO-aware custom op. + + Only the pure TP+TBO case (tp>1, TBO on, no DP) benefits: the ubatch's AR + then overlaps the partner ubatch's compute. Non-TBO keeps the plain + all_reduce (no custom-op / Dynamo-barrier indirection), and TBO+DP is left + untouched (that path overlaps via DP gather/scatter, not this pure-TP AR). + Shared by attention (wo_b) and MoE (combine_outputs) so both gate identically. + """ + if tp_size <= 1: + return False + cfg = get_current_atom_config() + return ( + getattr(cfg, "enable_tbo", False) + and cfg.parallel_config.data_parallel_size <= 1 + ) + + class DeepseekV4Attention(nn.Module): """Hybrid attention: MQA + grouped output LoRA + sliding window + attn_sink. @@ -1791,6 +1809,10 @@ def __init__( args.o_groups % tp_size == 0 ), f"o_groups={args.o_groups} not divisible by tp={tp_size}" self.tp_size = tp_size + # Route wo_b's TP all_reduce through the TBO-aware custom op only for the + # pure TP+TBO case (see _tbo_aware_tp_reduce). Non-TBO / TBO+DP keep the + # original all_reduce untouched. + self.tbo_aware = _tbo_aware_tp_reduce(tp_size) self.n_local_heads = args.n_heads // tp_size self.q_lora_rank = args.q_lora_rank self.o_lora_rank = args.o_lora_rank @@ -1854,6 +1876,7 @@ def __init__( self.dim, bias=False, quant_config=qc, + tbo_aware=self.tbo_aware, prefix=f"{p}.wo_b", ) self.softmax_scale = self.head_dim**-0.5 @@ -2141,7 +2164,11 @@ def _wo_a_grouped_lora(self, o: torch.Tensor, prefix: str = "") -> torch.Tensor: return torch.einsum("sgd,grd->sgr", o, wo_a) def _attn_post(self, o: torch.Tensor) -> torch.Tensor: - """Grouped output LoRA + wo_b (graphable, num_tokens-shaped).""" + """Grouped output LoRA + wo_b (graphable, num_tokens-shaped). + + wo_b is built with tbo_aware=True, so its built-in TP all_reduce is + TBO-aware (overlaps the partner ubatch's compute under TBO). + """ o = self._wo_a_grouped_lora(o, prefix=f"{self.layer_name}.wo_a") return self.wo_b(o.flatten(1)) @@ -2773,6 +2800,9 @@ def __init__( self.routed_scaling_factor = args.route_scale self.swiglu_limit = args.swiglu_limit self.tp_size = get_tensor_model_parallel_world_size() + # Pure TP+TBO: route combine_outputs' TP all_reduce through the TBO-aware + # custom op (overlap with partner ubatch). Non-TBO / TBO+DP unchanged. + self.tbo_aware = _tbo_aware_tp_reduce(self.tp_size) self.alt_stream = alt_stream qc = args.quant_config @@ -3003,7 +3033,15 @@ def combine_outputs( shared = shared * (1.0 / get_pcp_world_size()) routed = routed + shared if self.tp_size > 1: - routed = tensor_model_parallel_all_reduce(routed) + if self.tbo_aware: + # Pure TP+TBO: TBO-aware AR (custom op) overlaps this ubatch's TP + # all_reduce with the partner ubatch's compute. The custom op is a + # Dynamo barrier so no tbo_active() branch lands in the (possibly + # compiled) call site — see module_dispatch_ops.tbo_all_reduce. + routed = torch.ops.aiter.tbo_all_reduce(routed) + else: + # Non-TBO / TBO+DP: original plain TP all_reduce, untouched. + routed = tensor_model_parallel_all_reduce(routed) return routed def single_stream_moe_forward( diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 60044e5d70..a5bb94e167 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -308,12 +308,18 @@ # They are configured via --dspark-config (JSON dict) and carried in # config.dspark (see atom/config.py DSparkConfig). See # recipes/DeepSeek-V4-DSpark.md. - # --- PrefillDelayer (cross-DP prefill alignment) --- - # Master switch; default on. Set "0" to disable construction. - # The delayer is a prefill COALESCER: it holds back prefill admission under - # DP-attention until the accumulated prefill fills a worthwhile forward, so - # fragmented short-input prefills / small partial tail chunks batch into one - # forward instead of firing many tiny ones. + # --- PrefillDelayer (prefill coalescer) --- + # Master switch / kill switch; default on. Set "0" to disable construction. + # The delayer is a prefill COALESCER: it holds back prefill admission until + # the accumulated prefill fills a worthwhile forward, so fragmented + # short-input prefills / small partial tail chunks batch into one forward + # instead of firing many tiny ones. + # Enabled in two cases: + # - DP (data_parallel_size > 1): always on with this switch — also enforces + # cross-rank prefill alignment (reduces over dp_group). + # - TP-only / single-rank with TBO on: rides on TBO, since a fuller forward + # is exactly what TBO needs to split into two useful ubatches. Runs in + # single-rank mode (no all_reduce). This switch still force-disables it. "ATOM_ENABLE_PREFILL_DELAYER": lambda: ( os.getenv("ATOM_ENABLE_PREFILL_DELAYER", "1") == "1" ), @@ -381,6 +387,21 @@ # 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: in split_attn_metadata, cross-check the host-computed per-ubatch + # max_seqlen_q/k against the device tensors and fall back to the device value + # (with a warning) on mismatch. Off by default — turn on to diagnose stale + # attach_tbo_cpu_lens snapshots. Adds two .item() D2H syncs per ubatch. + "ATOM_TBO_DEBUG_CHECK": lambda: ( + os.getenv("ATOM_TBO_DEBUG_CHECK", "0") == "1" + ), + # Pure-TP TBO all_reduce overlap mode (see module_dispatch_ops.tbo_all_reduce): + # "overlap" (default): move the AR onto the comm stream so it overlaps the + # partner ubatch's compute. Per-ubatch pynccl comms keep the + # cross-rank enqueue order consistent (hang-free). + # "inline": run the AR on the current stream, no overlap. Kill switch — + # the Plan-A baseline; fall back here if an overlap-path hang + # ever resurfaces on some shape/topology. + "ATOM_TBO_TP_AR_MODE": lambda: os.getenv("ATOM_TBO_TP_AR_MODE", "overlap"), # --- 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/ubatching.py b/atom/utils/tbo/ubatching.py index 08e5637500..953a9e0c62 100644 --- a/atom/utils/tbo/ubatching.py +++ b/atom/utils/tbo/ubatching.py @@ -318,6 +318,49 @@ def sync_dp_metadata( _CURRENT_CONTEXTS: list["TBOContext | None"] = [] _NUM_UBATCHES: int = 2 +_TP_AR_ORDER_EVENT: "torch.cuda.Event | None" = None + + +def tbo_get_tp_ar_order_event() -> "torch.cuda.Event | None": + """Shared event that serializes the pure-TP all_reduce across ubatches.""" + return _TP_AR_ORDER_EVENT + + +# Per-ubatch independent pynccl TP communicators for the pure-TP all_reduce +# overlap (see `tbo_all_reduce`). +_TBO_TP_UBATCH_COMMS: "list | None" = None + + +def tbo_get_ubatch_tp_comm(ubatch_id: int): + """Return this ubatch's dedicated pynccl TP communicator, building the pair + lazily on first use. Returns None if TP world size == 1 (no AR needed).""" + global _TBO_TP_UBATCH_COMMS + if _TBO_TP_UBATCH_COMMS is not None: + return _TBO_TP_UBATCH_COMMS[ubatch_id] + + from aiter.dist.device_communicators.communicator_pynccl import ( + PyNcclCommunicator, + ) + from aiter.dist.parallel_state import get_tp_group + + tp = get_tp_group() + if tp.world_size == 1: + _TBO_TP_UBATCH_COMMS = [None] * _NUM_UBATCHES + return None + + # One independent communicator per ubatch. Construction runs a warmup + # all_reduce (a collective), so every rank must build the same number in the + # same order — guaranteed here because all ranks run TBO in lockstep and hit + # this on the same forward. Bound to the TP cpu_group (non-NCCL backend), + # exactly like the shared pynccl_comm. + comms = [] + for _ in range(_NUM_UBATCHES): + comms.append( + PyNcclCommunicator(group=tp.cpu_group, device=tp.device) + ) + _TBO_TP_UBATCH_COMMS = comms + return _TBO_TP_UBATCH_COMMS[ubatch_id] + class TBOContext: """Context manager for micro-batch dual-thread overlap. @@ -568,7 +611,7 @@ def make_tbo_contexts( Threading events are arranged in a ring so that each context's ``cpu_signal_event`` is the *next* context's ``cpu_wait_event``. """ - global _NUM_UBATCHES, _CURRENT_CONTEXTS + global _NUM_UBATCHES, _CURRENT_CONTEXTS, _TP_AR_ORDER_EVENT assert num_micro_batches > 1 _NUM_UBATCHES = num_micro_batches @@ -576,6 +619,11 @@ def make_tbo_contexts( while len(_CURRENT_CONTEXTS) < num_micro_batches: _CURRENT_CONTEXTS.append(None) + # Fresh per-step event that serializes the pure-TP all_reduce across + # ubatches (see `_TP_AR_ORDER_EVENT` note above). A new event each step + # means the first ubatch's AR never waits on a prior step's recording. + _TP_AR_ORDER_EVENT = torch.cuda.Event() + cpu_events = [threading.Event() for _ in range(num_micro_batches)] gpu_comm_done_events = [torch.Event() for _ in range(num_micro_batches)] gpu_compute_done_events = [torch.Event() for _ in range(num_micro_batches)] From 88bc1be502940652886ffa56df32b0531e7f2729 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Thu, 23 Jul 2026 08:33:08 +0000 Subject: [PATCH 02/10] style: black format --- atom/utils/envs.py | 4 +--- atom/utils/tbo/ubatching.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index a5bb94e167..f9c6f078c2 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -391,9 +391,7 @@ # max_seqlen_q/k against the device tensors and fall back to the device value # (with a warning) on mismatch. Off by default — turn on to diagnose stale # attach_tbo_cpu_lens snapshots. Adds two .item() D2H syncs per ubatch. - "ATOM_TBO_DEBUG_CHECK": lambda: ( - os.getenv("ATOM_TBO_DEBUG_CHECK", "0") == "1" - ), + "ATOM_TBO_DEBUG_CHECK": lambda: (os.getenv("ATOM_TBO_DEBUG_CHECK", "0") == "1"), # Pure-TP TBO all_reduce overlap mode (see module_dispatch_ops.tbo_all_reduce): # "overlap" (default): move the AR onto the comm stream so it overlaps the # partner ubatch's compute. Per-ubatch pynccl comms keep the diff --git a/atom/utils/tbo/ubatching.py b/atom/utils/tbo/ubatching.py index 953a9e0c62..5378a7493c 100644 --- a/atom/utils/tbo/ubatching.py +++ b/atom/utils/tbo/ubatching.py @@ -355,9 +355,7 @@ def tbo_get_ubatch_tp_comm(ubatch_id: int): # exactly like the shared pynccl_comm. comms = [] for _ in range(_NUM_UBATCHES): - comms.append( - PyNcclCommunicator(group=tp.cpu_group, device=tp.device) - ) + comms.append(PyNcclCommunicator(group=tp.cpu_group, device=tp.device)) _TBO_TP_UBATCH_COMMS = comms return _TBO_TP_UBATCH_COMMS[ubatch_id] From e334c2330a6d70063ad790fa90b9d8dc62256806 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Thu, 23 Jul 2026 08:35:19 +0000 Subject: [PATCH 03/10] docs: restore delayer comment --- atom/utils/envs.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index f9c6f078c2..1b792cf5c7 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -308,18 +308,12 @@ # They are configured via --dspark-config (JSON dict) and carried in # config.dspark (see atom/config.py DSparkConfig). See # recipes/DeepSeek-V4-DSpark.md. - # --- PrefillDelayer (prefill coalescer) --- - # Master switch / kill switch; default on. Set "0" to disable construction. - # The delayer is a prefill COALESCER: it holds back prefill admission until - # the accumulated prefill fills a worthwhile forward, so fragmented - # short-input prefills / small partial tail chunks batch into one forward - # instead of firing many tiny ones. - # Enabled in two cases: - # - DP (data_parallel_size > 1): always on with this switch — also enforces - # cross-rank prefill alignment (reduces over dp_group). - # - TP-only / single-rank with TBO on: rides on TBO, since a fuller forward - # is exactly what TBO needs to split into two useful ubatches. Runs in - # single-rank mode (no all_reduce). This switch still force-disables it. + # --- PrefillDelayer (cross-DP prefill alignment) --- + # Master switch; default on. Set "0" to disable construction. + # The delayer is a prefill COALESCER: it holds back prefill admission under + # DP-attention until the accumulated prefill fills a worthwhile forward, so + # fragmented short-input prefills / small partial tail chunks batch into one + # forward instead of firing many tiny ones. "ATOM_ENABLE_PREFILL_DELAYER": lambda: ( os.getenv("ATOM_ENABLE_PREFILL_DELAYER", "1") == "1" ), From dd35f08cf505338f220d6fa3ea1672bc30d752fd Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Thu, 23 Jul 2026 08:37:28 +0000 Subject: [PATCH 04/10] chore: remove unused ATOM_TBO_DEBUG_CHECK --- atom/utils/envs.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 1b792cf5c7..01cc78de03 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -381,11 +381,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: in split_attn_metadata, cross-check the host-computed per-ubatch - # max_seqlen_q/k against the device tensors and fall back to the device value - # (with a warning) on mismatch. Off by default — turn on to diagnose stale - # attach_tbo_cpu_lens snapshots. Adds two .item() D2H syncs per ubatch. - "ATOM_TBO_DEBUG_CHECK": lambda: (os.getenv("ATOM_TBO_DEBUG_CHECK", "0") == "1"), # Pure-TP TBO all_reduce overlap mode (see module_dispatch_ops.tbo_all_reduce): # "overlap" (default): move the AR onto the comm stream so it overlaps the # partner ubatch's compute. Per-ubatch pynccl comms keep the From 7493c39d6a878f6233de81be7dc8e73e54313067 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Mon, 27 Jul 2026 05:46:22 +0000 Subject: [PATCH 05/10] [DSV4] refactor(tbo): unify TP AR into communication_op layer --- atom/model_ops/communication_op.py | 38 ++++++++++++++++++++++++++++++ atom/model_ops/linear.py | 8 +++---- atom/models/deepseek_v4.py | 17 ++++++------- atom/utils/envs.py | 4 +--- atom/utils/tbo/ubatching.py | 15 +----------- 5 files changed, 51 insertions(+), 31 deletions(-) create mode 100644 atom/model_ops/communication_op.py diff --git a/atom/model_ops/communication_op.py b/atom/model_ops/communication_op.py new file mode 100644 index 0000000000..adcfc3b078 --- /dev/null +++ b/atom/model_ops/communication_op.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""ATOM tensor-parallel all_reduce layer. + +A thin wrapper over the aiter TP all_reduce that centralises the TBO-aware +routing, so call sites (model AR points such as attention output-proj and MoE +combine) don't each hand-roll an ``if tbo_aware: tbo_all_reduce else: plain`` +branch. Pass ``tbo_aware=True`` at a site that wants its all_reduce to overlap +the partner ubatch's compute under Two-Batch Overlap; leave it False (default) +for the plain reduce. + +``tbo_aware`` is a per-site opt-in rather than an internal ``tbo_active()`` +probe on purpose: TBO+DP overlaps via the DP gather/scatter path, not this +pure-TP reduce, so only sites explicitly on the pure-TP+TBO path should route +here. The ``tbo_all_reduce`` custom op still no-ops back to a plain reduce when +TBO is inactive or ``ATOM_TBO_TP_AR_MODE != overlap``. +""" + +import torch + + +def tensor_model_parallel_all_reduce( + x: torch.Tensor, tbo_aware: bool = False +) -> torch.Tensor: + """TP all_reduce; routes through the TBO-aware custom op when tbo_aware. + + The aiter import is lazy (inside the function): this module is pulled in + very early via atom.model_ops.__init__ -> ... -> linear, before aiter.dist + is fully initialised, so a top-level ``import aiter.dist.communication_op`` + resolves wrong / circularly. Importing at call time (same pattern as + module_dispatch_ops.tbo_all_reduce) sidesteps that. + """ + if tbo_aware: + return torch.ops.aiter.tbo_all_reduce(x) + from aiter.dist.communication_op import tensor_model_parallel_all_reduce as ar + + return ar(x) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 0286d998d5..05649fd1f5 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -23,6 +23,7 @@ from aiter.tuned_gemm import tgemm from aiter.utility import fp4_utils from atom.config import QuantizationConfig, get_current_atom_config +from atom.model_ops.communication_op import tensor_model_parallel_all_reduce from atom.quant_spec import LayerQuantConfig, should_skip_online_quant from atom.model_ops.utils import ( atom_parameter, @@ -924,10 +925,9 @@ def forward( if self.bias is not None: y += self.bias if self.tp_dim == 1 and self.tp_size > 1 and self.reduce_results: - if getattr(self, "tbo_aware", False): - y = torch.ops.aiter.tbo_all_reduce(y) - else: - y = get_tp_group().all_reduce(y, ca_fp8_quant=False) + y = tensor_model_parallel_all_reduce( + y, tbo_aware=getattr(self, "tbo_aware", False) + ) return y diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index fcb95ab9ff..5e63a41062 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -35,7 +35,7 @@ rope_rotate_activation, ) from aiter import silu_and_mul as aiter_silu_and_mul -from aiter.dist.communication_op import ( +from atom.model_ops.communication_op import ( tensor_model_parallel_all_reduce, ) from aiter.dist.parallel_state import ( @@ -3033,15 +3033,12 @@ def combine_outputs( shared = shared * (1.0 / get_pcp_world_size()) routed = routed + shared if self.tp_size > 1: - if self.tbo_aware: - # Pure TP+TBO: TBO-aware AR (custom op) overlaps this ubatch's TP - # all_reduce with the partner ubatch's compute. The custom op is a - # Dynamo barrier so no tbo_active() branch lands in the (possibly - # compiled) call site — see module_dispatch_ops.tbo_all_reduce. - routed = torch.ops.aiter.tbo_all_reduce(routed) - else: - # Non-TBO / TBO+DP: original plain TP all_reduce, untouched. - routed = tensor_model_parallel_all_reduce(routed) + # ATOM AR layer routes through the TBO-aware custom op when + # self.tbo_aware (pure TP+TBO); plain all_reduce otherwise + # (non-TBO / TBO+DP), untouched. + routed = tensor_model_parallel_all_reduce( + routed, tbo_aware=self.tbo_aware + ) return routed def single_stream_moe_forward( diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 01cc78de03..8680bed4e9 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -385,9 +385,7 @@ # "overlap" (default): move the AR onto the comm stream so it overlaps the # partner ubatch's compute. Per-ubatch pynccl comms keep the # cross-rank enqueue order consistent (hang-free). - # "inline": run the AR on the current stream, no overlap. Kill switch — - # the Plan-A baseline; fall back here if an overlap-path hang - # ever resurfaces on some shape/topology. + # "inline": run the AR on the current stream, no overlap. Plan-A baseline. "ATOM_TBO_TP_AR_MODE": lambda: os.getenv("ATOM_TBO_TP_AR_MODE", "overlap"), # --- NUMA binding --- # Master switch: pin each GPU worker to its GPU-local NUMA node's CPU cores diff --git a/atom/utils/tbo/ubatching.py b/atom/utils/tbo/ubatching.py index 5378a7493c..90458cfa64 100644 --- a/atom/utils/tbo/ubatching.py +++ b/atom/utils/tbo/ubatching.py @@ -318,14 +318,6 @@ def sync_dp_metadata( _CURRENT_CONTEXTS: list["TBOContext | None"] = [] _NUM_UBATCHES: int = 2 -_TP_AR_ORDER_EVENT: "torch.cuda.Event | None" = None - - -def tbo_get_tp_ar_order_event() -> "torch.cuda.Event | None": - """Shared event that serializes the pure-TP all_reduce across ubatches.""" - return _TP_AR_ORDER_EVENT - - # Per-ubatch independent pynccl TP communicators for the pure-TP all_reduce # overlap (see `tbo_all_reduce`). _TBO_TP_UBATCH_COMMS: "list | None" = None @@ -609,7 +601,7 @@ def make_tbo_contexts( Threading events are arranged in a ring so that each context's ``cpu_signal_event`` is the *next* context's ``cpu_wait_event``. """ - global _NUM_UBATCHES, _CURRENT_CONTEXTS, _TP_AR_ORDER_EVENT + global _NUM_UBATCHES, _CURRENT_CONTEXTS assert num_micro_batches > 1 _NUM_UBATCHES = num_micro_batches @@ -617,11 +609,6 @@ def make_tbo_contexts( while len(_CURRENT_CONTEXTS) < num_micro_batches: _CURRENT_CONTEXTS.append(None) - # Fresh per-step event that serializes the pure-TP all_reduce across - # ubatches (see `_TP_AR_ORDER_EVENT` note above). A new event each step - # means the first ubatch's AR never waits on a prior step's recording. - _TP_AR_ORDER_EVENT = torch.cuda.Event() - cpu_events = [threading.Event() for _ in range(num_micro_batches)] gpu_comm_done_events = [torch.Event() for _ in range(num_micro_batches)] gpu_compute_done_events = [torch.Event() for _ in range(num_micro_batches)] From 43e75bc0aafa96f0b6cff215221598c8cde1b964 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Mon, 27 Jul 2026 05:52:29 +0000 Subject: [PATCH 06/10] style: black format --- atom/models/deepseek_v4.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 5e63a41062..af4f4f1f0a 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -3036,9 +3036,7 @@ def combine_outputs( # ATOM AR layer routes through the TBO-aware custom op when # self.tbo_aware (pure TP+TBO); plain all_reduce otherwise # (non-TBO / TBO+DP), untouched. - routed = tensor_model_parallel_all_reduce( - routed, tbo_aware=self.tbo_aware - ) + routed = tensor_model_parallel_all_reduce(routed, tbo_aware=self.tbo_aware) return routed def single_stream_moe_forward( From 8b389a3a6add9b69322a6747ce294ad3e49fce96 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Mon, 27 Jul 2026 06:10:34 +0000 Subject: [PATCH 07/10] style: sort imports, Optional->|None (ruff) --- atom/model_ops/linear.py | 13 +++++++------ atom/models/deepseek_v4.py | 29 +++++++++++++++-------------- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 05649fd1f5..04ab616c76 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -11,9 +11,9 @@ dtypes, gemm_a4w4, gemm_a8w8, + gemm_a8w8_blockscale, gemm_a8w8_blockscale_bpreshuffle, gemm_a8w8_bpreshuffle, - gemm_a8w8_blockscale, get_hip_quant, ) @@ -22,22 +22,23 @@ from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.tuned_gemm import tgemm from aiter.utility import fp4_utils +from torch import nn + from atom.config import QuantizationConfig, get_current_atom_config from atom.model_ops.communication_op import tensor_model_parallel_all_reduce -from atom.quant_spec import LayerQuantConfig, should_skip_online_quant from atom.model_ops.utils import ( atom_parameter, normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale, shuffle_weights, ) -from atom.utils import envs -from atom.utils.decorators import mark_trace +from atom.quant_spec import LayerQuantConfig, should_skip_online_quant from atom.quantization.quark.utils import ( dequant_weight_online, quant_weight_online, ) -from torch import nn +from atom.utils import envs +from atom.utils.decorators import mark_trace logger = logging.getLogger("atom") @@ -1819,7 +1820,7 @@ def __init__( input_size: int, output_size: int, bias: bool = False, - quant_config: Optional[QuantizationConfig] = None, + quant_config: QuantizationConfig | None = None, reduce_results: bool = True, tbo_aware: bool = False, source_quant_dtype: torch.dtype = None, diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index af4f4f1f0a..44345c4137 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -35,23 +35,10 @@ rope_rotate_activation, ) from aiter import silu_and_mul as aiter_silu_and_mul -from atom.model_ops.communication_op import ( - tensor_model_parallel_all_reduce, -) from aiter.dist.parallel_state import ( get_tensor_model_parallel_world_size, ) from aiter.jit.utils.chip_info import get_gfx -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 ( @@ -59,6 +46,8 @@ ) from aiter.ops.triton.gemm.batched.batched_gemm_bf16 import batched_gemm_bf16 from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits +from torch import nn + from atom.config import ( Config, LayerQuantConfig, @@ -66,6 +55,15 @@ QuantType, get_current_atom_config, ) +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.model_loader.loader import WeightsMapper # Side-effect import: registers `torch.ops.aiter.maybe_dual_stream_forward` @@ -74,6 +72,9 @@ # code as opaque; Indexer.forward_batched dispatches via the latter to hide # its dynamic-shape internals from Dynamo / fake-tensor mode. from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401 +from atom.model_ops.communication_op import ( + tensor_model_parallel_all_reduce, +) from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import RMSNorm, rmsnorm2d_fwd_ from atom.model_ops.linear import ( @@ -107,9 +108,9 @@ update_compressor_states, ) from atom.utils import envs, mark_spliting_op +from atom.utils.custom_register import direct_register_custom_op from atom.utils.decorators import mark_trace, support_torch_compile from atom.utils.forward_context import AttnState, get_forward_context -from torch import nn logger = logging.getLogger(__name__) From 564412ad7cb9ddef73827fdf47fc89d361436628 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Mon, 27 Jul 2026 13:26:30 +0000 Subject: [PATCH 08/10] test(prefill_delayer): give dp>1 delayers a non-None cpu_group MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cross-DP all_reduce only runs when cpu_group is not None (the TP-only dp=1 path skips it). Multi-rank tests mock all_reduce to inject sibling contributions, so that branch must execute — but make_delayer defaulted cpu_group=None even for dp_size=2, so the reduce was skipped, the sibling contribution never applied, n_prefillable stayed 1 < dp_size, and the alignment gate held instead of firing (test_aligned_fires_when_full). Set a sentinel cpu_group when dp_size>1 (all_reduce is mocked in call(), so the value is only used for the is-not-None check; real dp>1 always has a cpu_group). Production code unchanged. --- tests/test_prefill_delayer.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_prefill_delayer.py b/tests/test_prefill_delayer.py index d0cb142c60..39e4d605bf 100644 --- a/tests/test_prefill_delayer.py +++ b/tests/test_prefill_delayer.py @@ -50,6 +50,12 @@ def make_delayer(**kw): token_usage_low_watermark=None, ) defaults.update(kw) + # The cross-DP all_reduce only runs when cpu_group is not None (dp=1/TP-only + # skips it — there is no sibling to align with). Multi-rank tests mock + # all_reduce to inject sibling contributions, so they need a non-None group + # for that branch to execute. Real dp>1 always has a cpu_group; mirror that. + if defaults["dp_size"] > 1 and defaults["cpu_group"] is None: + defaults["cpu_group"] = object() # sentinel; all_reduce is mocked in call() return PrefillDelayer(**defaults) From 6bee85aef1da6eef784a7f03c74954073881a6e0 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Mon, 27 Jul 2026 13:41:00 +0000 Subject: [PATCH 09/10] style(test): use dict literal in make_delayer (ruff C408) --- tests/test_prefill_delayer.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/test_prefill_delayer.py b/tests/test_prefill_delayer.py index 39e4d605bf..525bfed30b 100644 --- a/tests/test_prefill_delayer.py +++ b/tests/test_prefill_delayer.py @@ -38,17 +38,17 @@ def f(tensor, op=None, group=None): def make_delayer(**kw): - defaults = dict( - dp_size=1, - cpu_group=None, - max_num_batched_tokens=MAX_BATCHED, - target_fill=0.7, - ttft_max_ticks=30, - partial_max_ticks=8, - stall_ticks=3, - kv_high_watermark=0.9, - token_usage_low_watermark=None, - ) + defaults = { + "dp_size": 1, + "cpu_group": None, + "max_num_batched_tokens": MAX_BATCHED, + "target_fill": 0.7, + "ttft_max_ticks": 30, + "partial_max_ticks": 8, + "stall_ticks": 3, + "kv_high_watermark": 0.9, + "token_usage_low_watermark": None, + } defaults.update(kw) # The cross-DP all_reduce only runs when cpu_group is not None (dp=1/TP-only # skips it — there is no sibling to align with). Multi-rank tests mock From 69c3926d2b3fbb9278a2306f3720bbdc95f979a7 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Tue, 28 Jul 2026 03:29:36 +0000 Subject: [PATCH 10/10] refactor(tbo): drop tbo_aware param, decide in AR layer --- atom/model_ops/communication_op.py | 48 +++++++++++++++++++++--------- atom/model_ops/linear.py | 9 +----- atom/models/deepseek_v4.py | 39 +++++------------------- 3 files changed, 42 insertions(+), 54 deletions(-) diff --git a/atom/model_ops/communication_op.py b/atom/model_ops/communication_op.py index adcfc3b078..641a2de577 100644 --- a/atom/model_ops/communication_op.py +++ b/atom/model_ops/communication_op.py @@ -6,24 +6,42 @@ A thin wrapper over the aiter TP all_reduce that centralises the TBO-aware routing, so call sites (model AR points such as attention output-proj and MoE combine) don't each hand-roll an ``if tbo_aware: tbo_all_reduce else: plain`` -branch. Pass ``tbo_aware=True`` at a site that wants its all_reduce to overlap -the partner ubatch's compute under Two-Batch Overlap; leave it False (default) -for the plain reduce. - -``tbo_aware`` is a per-site opt-in rather than an internal ``tbo_active()`` -probe on purpose: TBO+DP overlaps via the DP gather/scatter path, not this -pure-TP reduce, so only sites explicitly on the pure-TP+TBO path should route -here. The ``tbo_all_reduce`` custom op still no-ops back to a plain reduce when -TBO is inactive or ``ATOM_TBO_TP_AR_MODE != overlap``. +branch — nor even thread a ``tbo_aware`` flag through. Whether an all_reduce +should overlap the partner ubatch's compute is decided here, once, from the +runtime config via ``_tbo_aware_tp_reduce``. + +Only the pure TP+TBO case routes through the overlap op: TBO+DP overlaps via the +DP gather/scatter path (not this pure-TP reduce), and non-TBO wants the plain +reduce with no custom-op / Dynamo-barrier indirection. The ``tbo_all_reduce`` +custom op itself still no-ops back to a plain reduce when TBO is inactive or +``ATOM_TBO_TP_AR_MODE != overlap``. """ import torch -def tensor_model_parallel_all_reduce( - x: torch.Tensor, tbo_aware: bool = False -) -> torch.Tensor: - """TP all_reduce; routes through the TBO-aware custom op when tbo_aware. +def _tbo_aware_tp_reduce(tp_size: int) -> bool: + """True iff a pure-TP all_reduce should route through the TBO-aware custom op. + + Only the pure TP+TBO case (tp>1, TBO on, no DP) benefits: the ubatch's AR + then overlaps the partner ubatch's compute. Non-TBO keeps the plain + all_reduce (no custom-op / Dynamo-barrier indirection), and TBO+DP is left + untouched (that path overlaps via DP gather/scatter, not this pure-TP AR). + """ + if tp_size <= 1: + return False + from atom.config import get_current_atom_config + + cfg = get_current_atom_config() + return ( + getattr(cfg, "enable_tbo", False) + and cfg.parallel_config.data_parallel_size <= 1 + ) + + +def tensor_model_parallel_all_reduce(x: torch.Tensor) -> torch.Tensor: + """TP all_reduce; routes through the TBO-aware custom op on the pure-TP+TBO + path (decided internally via ``_tbo_aware_tp_reduce``), else a plain reduce. The aiter import is lazy (inside the function): this module is pulled in very early via atom.model_ops.__init__ -> ... -> linear, before aiter.dist @@ -31,7 +49,9 @@ def tensor_model_parallel_all_reduce( resolves wrong / circularly. Importing at call time (same pattern as module_dispatch_ops.tbo_all_reduce) sidesteps that. """ - if tbo_aware: + from aiter.dist.parallel_state import get_tp_group + + if _tbo_aware_tp_reduce(get_tp_group().world_size): return torch.ops.aiter.tbo_all_reduce(x) from aiter.dist.communication_op import tensor_model_parallel_all_reduce as ar diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 04ab616c76..639b82d03b 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -926,9 +926,7 @@ def forward( if self.bias is not None: y += self.bias if self.tp_dim == 1 and self.tp_size > 1 and self.reduce_results: - y = tensor_model_parallel_all_reduce( - y, tbo_aware=getattr(self, "tbo_aware", False) - ) + y = tensor_model_parallel_all_reduce(y) return y @@ -1822,7 +1820,6 @@ def __init__( bias: bool = False, quant_config: QuantizationConfig | None = None, reduce_results: bool = True, - tbo_aware: bool = False, source_quant_dtype: torch.dtype = None, prefix: str = "", **kwargs, @@ -1838,10 +1835,6 @@ def __init__( source_quant_dtype=source_quant_dtype, prefix=prefix, ) - # When True, the built-in TP reduce goes through the TBO-aware custom op - # (overlaps with the partner ubatch under TBO). Default off so all other - # RowParallelLinear layers keep the plain all_reduce unchanged. - self.tbo_aware = tbo_aware def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): param_data = param.data diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 44345c4137..0627dc0a22 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1750,24 +1750,6 @@ def _score_topk_decode_ragged( # --------------------------------------------------------------------------- -def _tbo_aware_tp_reduce(tp_size: int) -> bool: - """True iff a pure-TP all_reduce should route through the TBO-aware custom op. - - Only the pure TP+TBO case (tp>1, TBO on, no DP) benefits: the ubatch's AR - then overlaps the partner ubatch's compute. Non-TBO keeps the plain - all_reduce (no custom-op / Dynamo-barrier indirection), and TBO+DP is left - untouched (that path overlaps via DP gather/scatter, not this pure-TP AR). - Shared by attention (wo_b) and MoE (combine_outputs) so both gate identically. - """ - if tp_size <= 1: - return False - cfg = get_current_atom_config() - return ( - getattr(cfg, "enable_tbo", False) - and cfg.parallel_config.data_parallel_size <= 1 - ) - - class DeepseekV4Attention(nn.Module): """Hybrid attention: MQA + grouped output LoRA + sliding window + attn_sink. @@ -1810,10 +1792,6 @@ def __init__( args.o_groups % tp_size == 0 ), f"o_groups={args.o_groups} not divisible by tp={tp_size}" self.tp_size = tp_size - # Route wo_b's TP all_reduce through the TBO-aware custom op only for the - # pure TP+TBO case (see _tbo_aware_tp_reduce). Non-TBO / TBO+DP keep the - # original all_reduce untouched. - self.tbo_aware = _tbo_aware_tp_reduce(tp_size) self.n_local_heads = args.n_heads // tp_size self.q_lora_rank = args.q_lora_rank self.o_lora_rank = args.o_lora_rank @@ -1877,7 +1855,6 @@ def __init__( self.dim, bias=False, quant_config=qc, - tbo_aware=self.tbo_aware, prefix=f"{p}.wo_b", ) self.softmax_scale = self.head_dim**-0.5 @@ -2167,8 +2144,9 @@ def _wo_a_grouped_lora(self, o: torch.Tensor, prefix: str = "") -> torch.Tensor: def _attn_post(self, o: torch.Tensor) -> torch.Tensor: """Grouped output LoRA + wo_b (graphable, num_tokens-shaped). - wo_b is built with tbo_aware=True, so its built-in TP all_reduce is - TBO-aware (overlaps the partner ubatch's compute under TBO). + wo_b's RowParallelLinear TP all_reduce goes through the ATOM AR layer, + which routes it through the TBO-aware custom op on the pure-TP+TBO path + (overlaps the partner ubatch's compute) and a plain reduce otherwise. """ o = self._wo_a_grouped_lora(o, prefix=f"{self.layer_name}.wo_a") return self.wo_b(o.flatten(1)) @@ -2801,9 +2779,6 @@ def __init__( self.routed_scaling_factor = args.route_scale self.swiglu_limit = args.swiglu_limit self.tp_size = get_tensor_model_parallel_world_size() - # Pure TP+TBO: route combine_outputs' TP all_reduce through the TBO-aware - # custom op (overlap with partner ubatch). Non-TBO / TBO+DP unchanged. - self.tbo_aware = _tbo_aware_tp_reduce(self.tp_size) self.alt_stream = alt_stream qc = args.quant_config @@ -3034,10 +3009,10 @@ def combine_outputs( shared = shared * (1.0 / get_pcp_world_size()) routed = routed + shared if self.tp_size > 1: - # ATOM AR layer routes through the TBO-aware custom op when - # self.tbo_aware (pure TP+TBO); plain all_reduce otherwise - # (non-TBO / TBO+DP), untouched. - routed = tensor_model_parallel_all_reduce(routed, tbo_aware=self.tbo_aware) + # ATOM AR layer decides internally (via _tbo_aware_tp_reduce) whether + # to route through the TBO-aware custom op (pure TP+TBO) or a plain + # all_reduce (non-TBO / TBO+DP). + routed = tensor_model_parallel_all_reduce(routed) return routed def single_stream_moe_forward(