fix: enable TBO compute/comm overlap in Deepseek V4 and fix mooncake gather hang for PD prefill

[ZhangLirong-amd]

Summary

Two fixes for DeepSeek-V4-Pro PD-disaggregation prefill with --enable-tbo:

1. Enable real compute/comm overlap.

2. Fix mooncake gather hang under high-concurrency TBO prefill.

Trace results

DeepSeek-V4-Pro, TP=8, --enable-dp-attention --enable-tbo, c=256, ISL=8192:

	TBO NCCL overlap with compute
before	3.8 %
after	91.7 %

Test plan

• PD prefill bench (c=256, ISL=8192, OSL=1024) runs to completion without hang
• Trace shows 91.7 % of TBO NCCL overlapping with compute on the new node
• CI: existing model accuracy tests

[Copilot]

Pull request overview

This PR targets DeepSeek-V4-Pro PD-disaggregation prefill stability and performance under TBO by (1) moving specific DP/TP collectives off the main compute stream to improve true compute/comm overlap, and (2) preventing a Mooncake RDMA race by synchronizing after staging-buffer gathers.

Changes:

• Added a per-device auxiliary CUDA stream helper and routed DP input_ids all-gather + MoE combine_outputs TP all-reduce through it to avoid NCCL interleaving on the compute lane during TBO ping-pong.
• Added an explicit CUDA stream synchronization after the staging-buffer gather in Mooncake block+slot transfers to prevent GPUDirect RDMA reads racing in-flight gathers.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
atom/models/deepseek_v4.py	Adds _run_on_tbo_comm_stream and uses it for DP id all-gather + TP all-reduce to improve TBO overlap behavior.
atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py	Synchronizes after _gather_slot(...) to avoid RDMA reading an incompletely-populated staging buffer.

Comments suppressed due to low confidence (1)

atom/models/deepseek_v4.py:54

• _run_on_tbo_comm_stream introduces torch.cuda.stream(...) / wait_stream(...) control flow into the eager Python graph. DeepseekV4Model is @support_torch_compile, and when MoE dual-stream is disabled (common when shared experts are fused), combine_outputs() executes in the compiled forward path. Dynamo/Inductor generally cannot trace stream-context switching reliably, which can break compilation or produce unexpected graph breaks.

Consider explicitly marking this helper as non-compilable so the rest of the model can still compile, while the stream manipulation runs eagerly at runtime.

    tensor_model_parallel_all_reduce,
)
from aiter.dist.parallel_state import (
    get_tensor_model_parallel_world_size,
)
from aiter.jit.utils.torch_guard import torch_compile_guard
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.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
from atom.config import (

---

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

[Copilot]

Pull request overview

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