[DSV4] perf(tbo): overlap pure-TP all_reduce on TBO + Delay TP by ZhangLirong-amd · Pull Request #1681 · ROCm/ATOM

ZhangLirong-amd · 2026-07-23T08:27:19Z

Motivation

AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1 python -m atom.entrypoints.openai_server   --model=/mnt/m2m_nobackup/lirzhang/DeepSeek-V4-Pro/   --tensor-parallel-size=8  --max-num-batched-tokens=16384   --kv-cache-dtype=fp8   --gpu-memory-utilization=0.85   --enable-tbo

For each 16k prefill, time 567ms -> 495ms, about 15% improve for each prefill

Technical Details

Test Plan

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

github-actions · 2026-07-23T08:27:59Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1681 --add-label <label>

Copilot

Pull request overview

This PR improves DeepSeek-V4 pure-TP + TBO prefill latency by overlapping the TP all-reduce with the partner ubatch’s compute via a new TBO-aware tbo_all_reduce custom op, and enables PrefillDelayer in TP-only + TBO mode to better coalesce prefills into effective ubatches.

Changes:

Add a TBO-aware torch.ops.aiter.tbo_all_reduce path that runs TP all-reduce on the TBO comm stream (opt-out via ATOM_TBO_TP_AR_MODE=inline).
Route DeepSeek-V4 attention (wo_b) and MoE combine_outputs TP reductions through the TBO-aware custom op in the pure TP+TBO (no DP) case.
Enable PrefillDelayer in TP-only + TBO mode (single-rank behavior: no cross-rank all_reduce), and add a benchmark scenario entry for TBO.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`atom/utils/tbo/ubatching.py`	Adds per-ubatch TP communicators and introduces an (currently unused) event intended for TP AR ordering.
`atom/utils/envs.py`	Updates PrefillDelayer docs; adds `ATOM_TBO_TP_AR_MODE` and a debug env var entry.
`atom/models/deepseek_v4.py`	Enables TBO-aware TP all-reduce routing for `wo_b` and MoE `combine_outputs` in the pure TP+TBO case.
`atom/model_ops/module_dispatch_ops.py`	Registers `torch.ops.aiter.tbo_all_reduce` custom op implementing overlap behavior under TBO.
`atom/model_ops/linear.py`	Adds `tbo_aware` flag to RowParallelLinear to route its TP reduction through the new custom op.
`atom/model_engine/prefill_delayer.py`	Adds TP-only (cpu_group=None) behavior and guards the DP all_reduce accordingly.
`atom/model_engine/engine_core.py`	Enables PrefillDelayer automatically in TP-only + TBO mode when allowed by env.
`.github/benchmark/models.json`	Adds a TBO benchmark variant configuration.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+    global _TBO_TP_UBATCH_COMMS
+    if _TBO_TP_UBATCH_COMMS is not None:
+        return _TBO_TP_UBATCH_COMMS[ubatch_id]


+    # 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()


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


Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

atom/utils/tbo/ubatching.py:610

make_tbo_contexts() allocates a fresh _TP_AR_ORDER_EVENT each step, but _TP_AR_ORDER_EVENT/tbo_get_tp_ar_order_event() are not referenced anywhere in the repo. This adds per-step overhead and is misleading; either wire the event into the pure-TP all_reduce ordering logic or stop creating it until it’s actually used.

    global _NUM_UBATCHES, _CURRENT_CONTEXTS, _TP_AR_ORDER_EVENT
    assert num_micro_batches > 1

    _NUM_UBATCHES = num_micro_batches
    # Grow the global context list if needed
    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()

+    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]


Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (3)

atom/utils/tbo/ubatching.py:318

_TBO_TP_UBATCH_COMMS is lazily initialized from the TBO worker threads, but there is no synchronization protecting the first-time construction. Because PyNcclCommunicator construction runs a warmup collective (per the comment below), two ubatch threads racing here can interleave communicator creation differently across ranks and hang. Add a module-level lock dedicated to guarding communicator initialization.

# Per-ubatch independent pynccl TP communicators for the pure-TP all_reduce
# overlap (see `tbo_all_reduce`).
_TBO_TP_UBATCH_COMMS: "list | None" = None

atom/utils/tbo/ubatching.py:328

tbo_get_ubatch_tp_comm() lazily builds PyNcclCommunicator instances without any mutual exclusion. Under TBO, both ubatch threads can hit this on the first forward, racing the constructor's warmup all_reduce and potentially deadlocking. Guard the initialization with the shared lock and double-check after acquiring it.

    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 (

atom/utils/tbo/ubatching.py:610

_TP_AR_ORDER_EVENT is created fresh each step and exported via tbo_get_tp_ar_order_event(), but it is not used anywhere (the new TBO all_reduce path in module_dispatch_ops doesn’t consult it). This dead state + docstring (“serializes the pure-TP all_reduce across ubatches”) is misleading and risks future callers relying on an ordering guarantee that doesn’t exist. Either wire the event into the overlap all_reduce path (e.g., record/wait between ubatches) or remove the event + getter entirely.

    # 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()

+        # 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


valarLip · 2026-07-23T09:27:32Z

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


let’s have a ATOM AR layer to deal with, so we don't have to add these "if..else.." for every model/op with AR

i would like something like from ato,.model_ops.communication_op import tensor_model_parallel_all_reduce

valarLip · 2026-07-23T09:30:38Z

+_TP_AR_ORDER_EVENT: "torch.cuda.Event | None" = None
+
+
+def tbo_get_tp_ar_order_event() -> "torch.cuda.Event | None":


never used?

[DSV4] perf(tbo): overlap pure-TP all_reduce

80a87eb

Copilot AI review requested due to automatic review settings July 23, 2026 08:27

Copilot started reviewing on behalf of ZhangLirong-amd July 23, 2026 08:28 View session

style: black format

569b216

ZhangLirong-amd changed the title ~~[DSV4] perf(tbo): overlap pure-TP all_reduce on TBO~~ [DSV4] perf(tbo): overlap pure-TP all_reduce on TBO + Delay TP Jul 23, 2026

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings July 23, 2026 08:34

docs: restore delayer comment

dede7d2

Copilot started reviewing on behalf of ZhangLirong-amd July 23, 2026 08:35 View session

chore: remove unused ATOM_TBO_DEBUG_CHECK

e65c612

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings July 23, 2026 08:52

Copilot started reviewing on behalf of ZhangLirong-amd July 23, 2026 08:53 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

ZhangLirong-amd requested a review from yitingw1 July 23, 2026 09:03

valarLip reviewed Jul 23, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

[DSV4] perf(tbo): overlap pure-TP all_reduce on TBO + Delay TP#1681
ZhangLirong-amd wants to merge 4 commits into
mainfrom
feat/tp-tbo-ar-overlap

ZhangLirong-amd commented Jul 23, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

valarLip Jul 23, 2026

Uh oh!

valarLip Jul 23, 2026

Uh oh!

ZhangLirong-amd Jul 23, 2026

Uh oh!

valarLip Jul 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		_TP_AR_ORDER_EVENT: "torch.cuda.Event \| None" = None


		def tbo_get_tp_ar_order_event() -> "torch.cuda.Event \| None":

Uh oh!

Conversation

ZhangLirong-amd commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Jul 23, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

valarLip Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

valarLip Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

ZhangLirong-amd Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

valarLip Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ZhangLirong-amd commented Jul 23, 2026 •

edited

Loading