feat(normal): add optional diagnostic probes for DeepEP normal mode

[Marsssqqq]

Summary

This PR adds six opt-in cumulative probes to DeepEP's internode Normal dispatch and combine kernels: three full-kernel Notify timers, two Dispatch completion probes, and one Combine logical completion probe. Together they expose rank-arrival and data-path progress evidence without changing existing Normal return values or adding a collective or inter-rank synchronization.

The probes are disabled when their optional tensor arguments are absent.

Motivation

Normal mode is the throughput-oriented path used by MoE training and inference prefill. Its dispatch and combine kernels move larger token sets through inter-node RDMA and intra-node NVLink stages while interacting with upstream computation and scheduling.

An end-to-end slowdown can therefore have different causes: a rank may enter the communication stage late, the RDMA path may slow globally, a subset of gateways or rank pairs may degrade, or progress may stall near the final receive stage. Host-side timing exposes the symptom but cannot identify the device-side stage. This change records evidence at existing synchronization and queue-progress boundaries without adding a new collective to the data path.

Implementation

Normal API and integration contract

The public Normal APIs gain two optional trailing arguments:

• normal_dispatch_stats: an opaque bundle of ten cumulative int64 CUDA tensors;
• normal_combine_stats: an opaque bundle of five cumulative int64 CUDA tensors.

The tensors are supplied by the caller and passed through the Python, pybind, C++, and CUDA layers. For example, a caller using the companion DeepXTrace PR can obtain and pass the bundles unchanged:

from deepxtrace import diagnose as ds

buffer = deep_ep.Buffer(group, num_nvl_bytes, num_rdma_bytes)
diagnose = ds.Diagnose.create_from_env(
    group=group,
    enable_ll_diagnose=False,
    enable_normal_diagnose=True,
    snapshot_stream=buffer.get_comm_stream(),
)

dispatch_stats = combine_stats = None
if diagnose is not None:
    dispatch_stats = diagnose.get_normal_dispatch_stats_tensor()
    combine_stats = diagnose.get_normal_combine_stats_tensor()
    diagnose.start()

buffer.dispatch(..., normal_dispatch_stats=dispatch_stats)
buffer.combine(..., normal_combine_stats=combine_stats)

if diagnose is not None:
    diagnose.stop()

The Python layer checks only bundle arity. The native layer validates CUDA placement, int64 dtype, expected shapes, and all-or-none enablement for each Notify pair or Completion triplet.

The per-launch Notify timer scratch is allocated and managed by the DeepEP runtime. It is not part of either caller-visible bundle.

Notify full-kernel duration

The implementation times three existing coordination kernels:

• non-cached Dispatch Notify;
• cached Dispatch Notify;
• cached Combine Notify.

At kernel entry, thread 0 of every block publishes its %globaltimer value to the per-launch timer state. The complemented timestamp and atomicMax select the earliest block start. At the epilogue, every block increments the completion counter; the last block records the interval from the earliest block start to the last block completion and increments the cumulative launch count. The timer state is reset on the same runtime stream before the launch.

This produces exactly one full-grid duration sample per rank and launch. It intentionally includes the time spent by early blocks waiting for the rest of the coordination grid, which makes cross-rank arrival skew visible to the consumer.

The Notify control flow converges at a common timing epilogue so every launched block contributes to the full-grid completion count. Probe-disabled launches preserve the existing communication and metadata work.

Dispatch completion probes

The Dispatch kernel assigns warps to kRDMASender, kRDMASenderCoordinator, kRDMAAndNVLForwarder, kForwarderCoordinator, and kNVLReceivers roles. The two probes target the receive-side data-moving roles:

• dispatch_rdma_recv runs in a kRDMAAndNVLForwarder warp. Each such warp owns one destination NVLink rank, and each active lane tracks one source RDMA rank. After metadata arrives, the lane waits until the acquire-loaded RDMA tail reaches its expected token count and records elapsed clock64() cycles, one sample, and the expected token count. Only the warp targeting destination NVLink rank 0 records, avoiding eight duplicate observations; the resulting source index is a gateway proxy rather than precise source-GPU timing.
• dispatch_final runs in a kNVLReceivers warp. Each receiver warp owns one source NVLink rank, while lanes map source RDMA ranks. A lane records only after the last expected token for that logical source completes its TMA store to the output. Combining the lane's source RDMA rank with the warp's source NVLink rank produces the logical global source rank; gathering destination-local vectors yields a dense logical source-by-destination matrix.

Combine logical completion probe

The Combine kernel assigns warps to kNVLSender, kNVLAndRDMAForwarder, kRDMAReceiver, and kCoordinator roles. combine_logical_recv runs only in the kCoordinator warp on the non-forwarder SM. Each active coordinator lane owns one source RDMA rank and reconstructs the eight logical NVLink-source dependencies within that source node from the routing bitmap. The lane derives each logical source's last expected RDMA head and observes the receive tail in the coordinator's existing progress loop. Once the tail passes that head, it records elapsed clock64() cycles, one sample, and the logical source token count.

Counter contract

Probe	Accumulated values	Observation point
notify_dispatch	full-grid duration in ns, launch count	non-cached dispatch Notify
cached_notify_dispatch	full-grid duration in ns, launch count	cached dispatch Notify
cached_notify_combine	full-grid duration in ns, launch count	cached combine Notify
dispatch_final	cost cycles, sample count, token count by logical source rank	final NVLink receive/TMA completion
dispatch_rdma_recv	cost cycles, sample count, token count by source gateway proxy rank	RDMA gateway receive completion
combine_logical_recv	cost cycles, sample count, token count by reconstructed logical source rank	logical source receive completion

Notify duration is accumulated in nanoseconds from %globaltimer. Completion cost is accumulated in device cycles from clock64(). Each Completion probe also records sample count and token count, allowing a consumer to normalize cost / sample_count while retaining traffic volume as context.

The counters are monotonic and are not reset by DeepEP. Snapshot ordering, window differencing, distributed aggregation, diagnosis, logging, and artifacts belong to the external diagnostics provider.

CUDA 13 build support

setup.py also discovers CUDA 13 CCCL headers when they are installed under CUDA_HOME/include/cccl.

Validation

The end-to-end test used two 8-GPU NVIDIA H20-3e nodes (EP16), an asynchronous one-second interval with a one-second collection warm-up, 4,096 tokens, hidden size 4,096, top-k 8, 128 experts, 512 workload warm-up iterations, and 512 measured iterations per phase. The Normal APIs were exercised after initializing NVSHMEM with low_latency_mode=True and explicit NIC-PE mapping across four bonded HCAs. Dispatch alternated between 256 non-cached and 256 cached calls; cached Combine ran for all 512 iterations. Both node launchers and all RDMA pressure processes exited successfully.

Phase	Representative window	Probe response and diagnosis
Baseline	3	All six probes ok; dense matrices showed only the expected intra-node/inter-node block structure
Rank-9 pre-dispatch compute delay	8	notify_dispatch and cached_notify_dispatch isolated rank 9; downstream Completion probes remained normal; two pre_stage_arrival_skew incidents
Rank-9 pre-combine compute delay	16	cached_notify_combine isolated rank 9; Completion probes remained normal; pre_stage_arrival_skew
All-HCA RDMA pressure	26	Dispatch final, Dispatch RDMA receive, and Combine logical receive all increased and produced global_data_path_slowdown
Single-HCA RDMA pressure	58	Dispatch final and Combine formed the same localized scope and produced localized_data_path_slowdown for token-owner ranks 0/1/8/9

Scope and compatibility

• Both cached and non-cached internode dispatch paths receive the Dispatch bundle; cached combine receives the Combine bundle.
• Intranode and low-latency execution paths are unchanged.
• The two arguments are trailing, optional, and default to None; existing dispatch/combine return values are unchanged.

Companion PR

Configuration, lifecycle, snapshotting, semantic diagnosis, logging, NPZ artifacts, and heatmap rendering are implemented in antgroup/DeepXTrace#17.

[Copilot]

Pull request overview

Adds optional DeepXTrace instrumentation to DeepEP’s normal-mode internode dispatch/combine path, wiring stage-aware cumulative counters from Python → C++ → CUDA kernels while keeping DeepXTrace as an optional extra dependency.

Changes:

• Introduces optional DeepXTrace enablement/handshake + schema validation in Buffer, and plumbs normal-mode stats tensors into dispatch/combine runtime calls.
• Adds CUDA-side probe implementations for notify full-kernel timing and dispatch/combine completion counters, plus corresponding API surface changes in C++ headers/bindings.
• Updates packaging to add a deepxtrace extra and improves CUDA 13 build support by optionally including CUDA_HOME/include/cccl.

Reviewed changes

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

Show a summary per file

File	Description
setup.py	Adds CUDA 13 cccl include path detection and declares deepxtrace as an optional extra.
deep_ep/buffer.py	Adds DeepXTrace optional initialization, rank-wide agreement, normal stats tensor plumbing, and exposes sync collection API.
csrc/kernels/internode.cu	Implements notify full-kernel timers and dispatch/combine completion probes; extends kernel entrypoints with optional stats pointers.
csrc/kernels/api.cuh	Extends kernel API declarations to accept optional normal-mode stats pointers.
csrc/deep_ep.hpp	Extends Buffer::internode_dispatch/compute signatures with optional stats tensors (API surface).
csrc/deep_ep.cpp	Validates stats tensor shapes/dtypes and plumbs optional pointers through to CUDA kernels.

Suppressed comments (2)

deep_ep/buffer.py:951

• Similar to low_latency_dispatch: low_latency_combine auto-populates combine_wait_recv_cost_stats from self.diagnose.get_stats_ll_stats_tensor() even though Diagnose is instantiated with enable_ll_diagnose=False. This can unintentionally couple low-latency APIs to DeepXTrace or cause runtime errors. Keep LL stats strictly caller-provided.

        if (combine_wait_recv_cost_stats is None and
                self.diagnose is not None):
            combine_wait_recv_cost_stats = \
                self.diagnose.get_stats_ll_stats_tensor()[1]

deep_ep/buffer.py:264

• Comment capitalization is inconsistent with the rest of the PR description: "DeepXtrace" vs "DeepXTrace".

        # End DeepXtrace

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

新的默认值会破坏受支持的无 NVSHMEM 节点内构建；异步诊断线程也缺少生命周期管理，可能在销毁后继续运行。

v4f

The PR adds optional DeepXTrace instrumentation to DeepEP's normal internode dispatch/combine paths: three Notify full-kernel probe families and three Completion probe families, wired through Buffer (lazy import, schema validation, rank-wide agreement), the C++/pybind boundary (null-pointer-gated optional tensor args with strict shape/type/contiguity validation), and the kernels (last-block timer accumulation with per-launch reset for CUDA-graph stability; per-source-lane completion probes reusing existing queue state). It also flips the Buffer.low_latency_mode default to True and adds a 'deepxtrace' package extra plus CUDA-13 CCCL include detection. The kernel-side instrumentation is carefully done and internally consistent: I verified the notify timer math (inverted-atomicMax earliest-start, last-block accumulation, per-launch cudaMemsetAsync reset on the same stream, all blocks reach begin/end in both notify_dispatch and the restructured cached_notify, %globaltimer at 1ns on sm_70+ matching the *_ns naming), the C++/Python boundary validation (int64/contiguous/num_ranks-sized checks, enablement triples consistent with kernel indexing and the rdma-recv (cost&&sample)||token condition), and the per-call channel-tail semantics that make the completion comparisons valid. No test changes accompany the new host-side logic. Main concerns: (1) a contradiction between the 'LL diagnosis disabled' claim and low_latency_dispatch/combine now auto-binding LL stats tensors from Diagnose whenever it exists; (2) the low_latency_mode default flip and enable_deepxtrace=True default (rank-0 RuntimeWarning for all existing users without DeepXTrace, MPI-only groups included) are compatibility/UX risks bundled into a feat PR; (3) the RDMA-receive completion probe attributes a whole source node's aggregated channel to a single 'source gateway rank' using the local NVL rank as a proxy, which should be documented/verified for topologies beyond the validated H20 setup; (4) a debug-build trap (EP_DEVICE_ASSERT) in the combine reconstruction path. Overall: sound design, merge after addressing the LL-path contradiction and confirming the default-value changes are intentional.

v5

本 PR 为 DeepEP normal 模式的 internode dispatch/combine 路径增加了可选的 DeepXTrace 诊断埋点，实现完整且与描述一致：变更覆盖恰好 6 个文件（deep_ep/buffer.py、csrc/deep_ep.cpp、csrc/deep_ep.hpp、csrc/kernels/api.cuh、csrc/kernels/internode.cu、setup.py）。Buffer 集成实现了懒加载 deepxtrace.diagnose、15 字段有序 NORMAL_STATS_SCHEMA 校验（要求 deepxtrace>=0.2.0,<0.3.0）、EP 组内全员一致性协商（任一 rank 未就绪则全组禁用并由 rank 0 发出一次告警，全禁用组保持静默）、异步/同步两种采集模式一致性检查以及 torch.distributed 进程组要求；MPI-only 构造保持 DeepXTrace 关闭，LL 诊断显式禁用。C++ 边界对全部 18 个新增可选统计张量做了 int64/连续性/形状与配对启用校验，禁用时向内核传空指针，Python 调用点与 C++ 签名参数数量完全对齐（dispatch 32 个、combine 21 个）。内核侧：Notify 全网格计时通过 %globaltimer + 取反 atomicMax（首块发布最早时间戳）+ 末块累加实现，计时状态在通信流上以 cudaMemsetAsync 复位、指针对 CUDA Graph 稳定，cached_notify 的提前 return 被重构为条件块以保证所有 block 到达计时终点；dispatch 的 RDMA 接收探针位于 RDMAAndNVLForwarder 并限定 dst_nvl_rank==0 单车道记录，最终完成探针位于 NVLReceivers 且在 tma_store_wait 之后记录；combine 的逻辑源完成探针在接收 SM 协调者中由路由位图与 RDMA head 重建依赖并观测既有队列 tail，全部复用现有队列状态、未引入新通信协议。setup.py 的 CUDA 13 include/cccl 探测实现正确且对旧 CUDA 版本无副作用，deepxtrace 以可选 extras（deepxtrace>=0.2.0,<0.3.0）声明而非强依赖，版本约束与 buffer.py 的 schema 校验错误信息一致。总体质量良好，可合入；评论中列出的问题均为非阻塞项：pybind 绑定未注册带默认值的 py::arg（与『运行时 API 兼容』的说法有出入）、low_latency_mode 默认值翻转属行为变更需显著提示、_load_deepxtrace 异常捕获范围偏窄、RDMA 接收探针代码三处重复、rdma_recv token_count 启用组合校验较宽松、Notify 用 ns 而 Completion 用 clock64 周期两种单位、位图整字读取的对齐前提未加断言、diagnose_normal_sync 错配调用缺乏防护说明。

Files reviewed: 6
Issues found: 🔴 3 critical | 🟡 8 warning | 🔵 7 suggestion
Inline comments posted: 18

[wangfakang]

Would it be better to encapsulate it as a diagnostic type using a struct?

[wangfakang]

same issue

[wangfakang]

same issue

[wangfakang]

why not atomicMin?

[wangfakang]

This synchronization can be removed; it appears to be meaningless.