[DSV4] perf(tbo): overlap pure-TP all_reduce on TBO + Delay TP - #1681
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
ZhangLirong-amd
changed the title
Contributor
There was a problem hiding this comment.
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_reducepath that runs TP all-reduce on the TBO comm stream (opt-out viaATOM_TBO_TP_AR_MODE=inline). - Route DeepSeek-V4 attention (
wo_b) and MoEcombine_outputsTP 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.
Comment on lines
+324
to
+326
| global _TBO_TP_UBATCH_COMMS | ||
| if _TBO_TP_UBATCH_COMMS is not None: | ||
| return _TBO_TP_UBATCH_COMMS[ubatch_id] |
Comment on lines
+609
to
+612
| # 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() |
Comment on lines
+384
to
+390
| # 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" | ||
| ), |
Contributor
There was a problem hiding this comment.
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()
Comment on lines
+324
to
+347
| 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] |
Contributor
There was a problem hiding this comment.
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()
Comment on lines
1812
to
1816
| # 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
reviewed
Jul 23, 2026
| 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) |
Collaborator
There was a problem hiding this comment.
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
Collaborator
There was a problem hiding this comment.
i would like something like from ato,.model_ops.communication_op import tensor_model_parallel_all_reduce
| _TP_AR_ORDER_EVENT: "torch.cuda.Event | None" = None | ||
|
|
||
|
|
||
| def tbo_get_tp_ar_order_event() -> "torch.cuda.Event | None": |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (2)
atom/utils/tbo/ubatching.py:319
- tbo_get_ubatch_tp_comm caches _TBO_TP_UBATCH_COMMS based on the then-current _NUM_UBATCHES, but _NUM_UBATCHES can be reassigned by make_tbo_contexts(). If TBO is ever run with a different num_micro_batches later (or ubatch_id is out of range), this will index the wrong communicator or raise IndexError in the hot path. Add a bounds check and rebuild the cached list when its length doesn’t match _NUM_UBATCHES.
global _TBO_TP_UBATCH_COMMS
if _TBO_TP_UBATCH_COMMS is not None:
return _TBO_TP_UBATCH_COMMS[ubatch_id]
atom/models/deepseek_v4.py:1815
- This file is
@support_torch_compile-decorated(DeepseekV4Model), and the new TBO routing logic is being implemented directly inside the model definition. That increases the risk of Dynamo/cudagraph regressions for all compiled runs of DSv4. Consider moving TBO-specific routing behind a stable op boundary (e.g., the existing custom-op layer) so the model code stays unchanged and tracing remains stable.
# 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)
Comment on lines
+927
to
+930
| 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, tbo_aware=getattr(self, "tbo_aware", False) | ||
| ) |
Comment on lines
+34
to
+37
| if tbo_aware: | ||
| return torch.ops.aiter.tbo_all_reduce(x) | ||
| from aiter.dist.communication_op import tensor_model_parallel_all_reduce as ar | ||
|
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (4)
atom/utils/tbo/ubatching.py:311
- _TBO_TP_UBATCH_COMMS is lazily initialized but not protected against concurrent first-use by multiple TBO ubatch threads. Because communicator construction performs collectives, a same-rank race here can deadlock (different threads on different ranks may enter the warmup collectives in different orders). Add a module-level lock to serialize 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:320
- tbo_get_ubatch_tp_comm() initializes per-ubatch PyNcclCommunicator instances without synchronization. Under TBO, both ubatch threads can hit this path on first use and concurrently run the communicator warmup collectives, which is a realistic hang scenario. Serialize initialization with _TBO_TP_UBATCH_COMMS_LOCK (and add a bounds assert for ubatch_id).
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/model_ops/communication_op.py:36
- tensor_model_parallel_all_reduce(..., tbo_aware=True) calls torch.ops.aiter.tbo_all_reduce(), but this op only exists after atom.model_ops.module_dispatch_ops has been imported (it performs the registration). As written, any call site that sets tbo_aware=True without having imported module_dispatch_ops first will raise AttributeError at runtime. Ensure registration happens before invoking the op.
if tbo_aware:
return torch.ops.aiter.tbo_all_reduce(x)
from aiter.dist.communication_op import tensor_model_parallel_all_reduce as ar
atom/models/deepseek_v4.py:1756
- This file contains
@support_torch_compileusage (see deepseek_v4.py around the model forward), and the repo documents a convention to avoid modifying@support_torch_compile-decoratedmodel files because it can break/perturb Dynamo tracing. The new TBO routing logic is currently implemented by changing deepseek_v4.py (new helper + wiring tbo_aware into RowParallelLinear and MoE). Consider moving the TBO-aware routing to a call site outside this model file (e.g., ModelRunner/run_model or another non-decorated integration layer), or get explicit maintainer sign-off that touching this file is safe for torch.compile capture.
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
Comment on lines
+378
to
+383
| # 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"), |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (4)
atom/utils/tbo/ubatching.py:320
- tbo_get_ubatch_tp_comm lazily initializes a global communicator list without any thread synchronization. Under TBO there are multiple Python threads per rank; two threads can concurrently see _TBO_TP_UBATCH_COMMS=None and both construct PyNcclCommunicator instances (each construction runs a collective warmup), which can hang or leave ranks with mismatched communicator state. Add a lock + double-check around initialization, and validate ubatch_id bounds before indexing.
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/model_ops/module_dispatch_ops.py:165
- The comment claims the default mode is "inline", but envs.ATOM_TBO_TP_AR_MODE defaults to "overlap" in envs.py. This mismatch is confusing for reviewers/users and makes it unclear what behavior is the default at runtime. Update the comment to match the actual default (or change the env default if inline is intended).
# 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)
atom/model_ops/communication_op.py:36
- tensor_model_parallel_all_reduce(tbo_aware=True) calls torch.ops.aiter.tbo_all_reduce unconditionally. That op only exists after module_dispatch_ops has been imported for its side-effect registration; otherwise this will raise AttributeError at runtime. Make the wrapper robust by ensuring the custom op is registered (or by falling back to the plain aiter all_reduce) when tbo_aware=True.
if tbo_aware:
return torch.ops.aiter.tbo_all_reduce(x)
from aiter.dist.communication_op import tensor_model_parallel_all_reduce as ar
atom/models/deepseek_v4.py:1760
- This PR changes atom/models/deepseek_v4.py, but the file contains
@support_torch_compile(DeepseekV4Model). Project guidance in CLAUDE.md explicitly warns against modifying@support_torch_compile-decoratedmodel files because it can break Dynamo tracing even with --enforce-eager; instrumentation should be done at call sites (e.g., ModelRunner/EagleProposer) instead. Consider moving the TBO-aware all_reduce routing/gating out of this model file (e.g., into model_ops or the runner) so DSv4 stays trace-stable.
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.
ZhangLirong-amd
force-pushed
the
feat/tp-tbo-ar-overlap
branch
from
July 27, 2026 12:43
6ddfe32 to
8b389a3
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (3)
atom/utils/tbo/ubatching.py:337
- tbo_get_ubatch_tp_comm() is called from two ubatch threads; the current lazy init can race so both threads construct PyNcclCommunicator objects concurrently. Since PyNcclCommunicator construction performs a warmup all_reduce (collective), concurrent init can lead to mismatched collective ordering across ranks or duplicate communicator creation, risking hangs.
from aiter.dist.device_communicators.communicator_pynccl import (
PyNcclCommunicator,
)
from aiter.dist.parallel_state import get_tp_group
atom/model_ops/module_dispatch_ops.py:163
- The comment says the default mode is "inline", but envs.ATOM_TBO_TP_AR_MODE defaults to "overlap" (atom/utils/envs.py). This is misleading when debugging hangs/perf differences.
# Default "inline": correct Plan-A baseline (no overlap, never hangs). Only
# move the AR onto the comm stream when explicitly opted into "overlap".
atom/models/deepseek_v4.py:1757
- This PR changes atom/models/deepseek_v4.py, but the file contains a
@support_torch_compile-decoratedmodel (DeepseekV4Model). The repo’s guidance (CLAUDE.md:44) explicitly forbids modifying@support_torch_compilemodel files because it can break Dynamo tracing; the TBO routing should be instrumented at call sites instead (e.g., ModelRunner/ops wrappers).
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
Comment on lines
+293
to
+295
| if self.cpu_group is not None: | ||
| try: | ||
| torch.distributed.all_reduce( |
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.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (2)
atom/model_engine/engine_core.py:136
self.schedulercan beNonein disaggregated-decode mode (config.disagg_is_decode=True), but this block unconditionally callsself.scheduler.set_prefill_delayer(...). That will raiseAttributeError: 'NoneType' object has no attribute 'set_prefill_delayer'when TBO is enabled in disagg decode workers.
if (
config.enable_tbo
and envs.ATOM_ENABLE_PREFILL_DELAYER
and config.parallel_config.data_parallel_size <= 1
):
atom/model_engine/prefill_delayer.py:297
- If
dp_size > 1butcpu_groupis accidentallyNone, this will silently skip the cross-rankall_reduce, making each rank decide FIRE/HOLD based on local state only. That can break DP phase alignment and lead to mismatched collectives (hang) later. It’s safer to assert the invariant (cpu_group must be set when dp_size>1) rather than quietly degrading.
if self.cpu_group is not None:
try:
torch.distributed.all_reduce(
self._reduce_buf,
op=torch.distributed.ReduceOp.SUM,
Comment on lines
+131
to
+135
| # 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 |
Comment on lines
+384
to
+389
| # 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"), |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (4)
atom/utils/tbo/ubatching.py:323
- _TBO_TP_UBATCH_COMMS is lazily initialised but can be reached concurrently by both TBO ubatch threads. Because PyNcclCommunicator construction runs a collective warmup, a race here can create mismatched collective enqueue order (hang) or double initialisation. Add a module-level lock for the one-time initialisation.
# 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:333
- tbo_get_ubatch_tp_comm() performs a one-time communicator build that includes collectives, but it is not protected against concurrent calls from the two ubatch threads. This can interleave communicator construction and deadlock. Guard the initialisation path with the lock (double-checked) so only one thread performs the warmup collectives.
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/models/deepseek_v4.py:1768
- deepseek_v4.py is a
@support_torch_compile-decoratedmodel file (see DeepseekV4Model at ~line 3580). Changes here can easily introduce Dynamo/compile/cudagraph regressions. If possible, keep TBO all_reduce routing decisions outside the compiled model file (e.g., in runner/dispatch layers), or add an explicit torch.compile regression test for the DSv4 TBO path to ensure this stays traceable.
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
)
atom/model_ops/module_dispatch_ops.py:165
- The comment above the env-var gate contradicts the actual default: envs.ATOM_TBO_TP_AR_MODE defaults to "overlap" (see atom/utils/envs.py), but this comment says the default is "inline" and implies overlap is opt-in. This is misleading when reasoning about rollout/safety. Update the comment to match the real default (or change the env default if inline is intended).
# 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)
Comment on lines
+131
to
+136
| # 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 | ||
| ): |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (7)
atom/utils/tbo/ubatching.py:324
- _TBO_TP_UBATCH_COMMS is initialized lazily but there is no per-process lock. In TBO, two ubatch threads can hit tbo_get_ubatch_tp_comm concurrently on first use, racing communicator construction (which runs a warmup collective) and potentially deadlocking/hanging.
# 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:333
- tbo_get_ubatch_tp_comm performs communicator construction without synchronization. Since TBO runs two threads per rank, both threads can enter the initialization path and construct different numbers/orders of PyNcclCommunicator instances, which can hang because communicator creation executes collectives.
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/model_ops/communication_op.py:56
- tensor_model_parallel_all_reduce() can call torch.ops.aiter.tbo_all_reduce without ensuring the custom op has been registered. If a model enables TBO but does not import atom.model_ops.module_dispatch_ops (which registers the op via side effects), this will raise at runtime.
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
atom/model_engine/engine_core.py:139
- This block unconditionally calls self.scheduler.set_prefill_delayer(), but self.scheduler is None when config.disagg_is_decode is True. If enable_tbo is set in disagg-decode mode, this will raise AttributeError during EngineCore initialization.
# 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(
atom/model_engine/prefill_delayer.py:295
- With dp_size > 1 but cpu_group=None, the cross-rank all_reduce is silently skipped and the subsequent global decision logic uses only local counters, breaking cross-rank alignment. This should fail fast (or at least warn) because dp_size>1 requires a valid process group.
if self.cpu_group is not None:
try:
torch.distributed.all_reduce(
atom/model_ops/module_dispatch_ops.py:165
- The comment says the default mode is "inline" and overlap is only used when explicitly opted in, but envs.py sets ATOM_TBO_TP_AR_MODE default to "overlap". This mismatch is likely to confuse debugging / rollout expectations.
# 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)
atom/models/deepseek_v4.py:2149
- This file contains
@support_torch_compile(see later in the file) and project guidance discourages modifying such model files because small changes can invalidate Dynamo tracing assumptions. Consider moving TBO routing commentary/logic out of the model file (e.g., keep behavior entirely in atom/model_ops) or add an explicit torch.compile coverage check for the DSv4 forward path with TBO enabled.
"""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.
valarLip
approved these changes
Jul 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
For each 16k prefill, time 567ms -> 495ms, about 15% improve for each prefill
Technical Details
Test Plan
Test Result
Submission Checklist