[DSV4] perf(tbo): overlap pure-TP all_reduce on TBO + Delay TP

.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",

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(

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).

atom/model_ops/communication_op.py

@@ -0,0 +1,58 @@
+# 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 — 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 _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
+    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.
+    """
+    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
+
+    return ar(x)

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,21 +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.quant_spec import LayerQuantConfig, should_skip_online_quant
+from atom.model_ops.communication_op import tensor_model_parallel_all_reduce
 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")
 
@@ -924,7 +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 = get_tp_group().all_reduce(y, ca_fp8_quant=False)
+            y = tensor_model_parallel_all_reduce(y)
         return y
 
 
@@ -1816,7 +1818,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,
         source_quant_dtype: torch.dtype = None,
         prefix: str = "",

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,),
+)

atom/models/deepseek_v4.py

@@ -35,37 +35,35 @@
     rope_rotate_activation,
 )
 from aiter import silu_and_mul as aiter_silu_and_mul
-from aiter.dist.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 (
     fused_clamp_act_mul,
 )
 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,
     QuantizationConfig,
     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__)
 
@@ -2141,7 +2142,12 @@ 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'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))
 
@@ -3003,6 +3009,9 @@ def combine_outputs(
                 shared = shared * (1.0 / get_pcp_world_size())
             routed = routed + shared
         if self.tp_size > 1:
+            # 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
 

atom/utils/envs.py

@@ -381,6 +381,12 @@
     # 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",
+    # 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. 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
     # and preferred memory. Default off so baseline/pinned A/B stays clean.