diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index 8cd79f1a4..f13a53a8b 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 6d62eb298..8448eca81 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 f723e375c..11640b593 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/communication_op.py b/atom/model_ops/communication_op.py new file mode 100644 index 000000000..641a2de57 --- /dev/null +++ b/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) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 50a15f301..639b82d03 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,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 = "", diff --git a/atom/model_ops/module_dispatch_ops.py b/atom/model_ops/module_dispatch_ops.py index cb7a2ae17..77be66e8f 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 703f6f3b2..0627dc0a2 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 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 ( @@ -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__) @@ -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 diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 60044e5d7..8680bed4e 100644 --- a/atom/utils/envs.py +++ b/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. diff --git a/atom/utils/tbo/ubatching.py b/atom/utils/tbo/ubatching.py index 08e563750..90458cfa6 100644 --- a/atom/utils/tbo/ubatching.py +++ b/atom/utils/tbo/ubatching.py @@ -318,6 +318,39 @@ def sync_dp_metadata( _CURRENT_CONTEXTS: list["TBOContext | None"] = [] _NUM_UBATCHES: int = 2 +# 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. diff --git a/tests/test_prefill_delayer.py b/tests/test_prefill_delayer.py index d0cb142c6..525bfed30 100644 --- a/tests/test_prefill_delayer.py +++ b/tests/test_prefill_delayer.py @@ -38,18 +38,24 @@ 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 + # 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)