deepseek-ai/DeepEP

Fork 0

mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00

[PR #720] feat(normal): add optional diagnostic probes for DeepEP normal mode #733

New issue

Open

opened 2026-10-27 21:05:43 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:05:43 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/720
Author: @Marsssqqq
Created: 8/10/2026
Status: 🔄 Open

Base: antgroup-opt ← Head: xym/antgroup-opt-normal-deepxtrace

📝 Commits (10+)

699d17b fix(build): locate CUDA 13 CCCL headers
5867e09 feat(deepxtrace): validate and bind normal stats contract
4365d77 feat(normal): add notify full-kernel duration probes
f46f3ff feat(normal): add dispatch final completion probes
b10f01d feat(normal): add dispatch RDMA receive completion probes
7ab8aa8 feat(normal): add combine logical receive completion probes
9572707 feat(deepxtrace): make normal diagnosis optional and configurable
fd9182a fix(buffer): preserve backward-compatible optional defaults
5caf04e fix(normal): harden optional DeepXTrace integration
77977ed fix(normal): isolate diagnostics from low-latency APIs

📊 Changes

7 files changed (+632 additions, -172 deletions)

View changed files

📝 csrc/deep_ep.cpp (+170 -7)
📝 csrc/deep_ep.hpp (+30 -2)
📝 csrc/kernels/api.cuh (+23 -2)
📝 csrc/kernels/internode.cu (+342 -136)
📝 deep_ep/buffer.py (+26 -15)
📝 setup.py (+4 -1)
📝 tests/test_internode.py (+37 -9)

📄 Description

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.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/720 **Author:** [@Marsssqqq](https://github.com/Marsssqqq) **Created:** 8/10/2026 **Status:** 🔄 Open **Base:** `antgroup-opt` ← **Head:** `xym/antgroup-opt-normal-deepxtrace` --- ### 📝 Commits (10+) - [`699d17b`](https://github.com/deepseek-ai/DeepEP/commit/699d17ba2767b119121b6a63f52aef6127607b63) fix(build): locate CUDA 13 CCCL headers - [`5867e09`](https://github.com/deepseek-ai/DeepEP/commit/5867e09bf50700cd8a09ffe51336009734cbafbe) feat(deepxtrace): validate and bind normal stats contract - [`4365d77`](https://github.com/deepseek-ai/DeepEP/commit/4365d7702371890351f1a27dcb01118434228858) feat(normal): add notify full-kernel duration probes - [`f46f3ff`](https://github.com/deepseek-ai/DeepEP/commit/f46f3ff223732cfa4c4192c241b3f06b6eaa9259) feat(normal): add dispatch final completion probes - [`b10f01d`](https://github.com/deepseek-ai/DeepEP/commit/b10f01d9adf57d716202ae49e4793d3c49104966) feat(normal): add dispatch RDMA receive completion probes - [`7ab8aa8`](https://github.com/deepseek-ai/DeepEP/commit/7ab8aa8d51a7c9957739e39d957469514dc17986) feat(normal): add combine logical receive completion probes - [`9572707`](https://github.com/deepseek-ai/DeepEP/commit/95727078476ecfaca923ab6b926e25ccf46dd4c9) feat(deepxtrace): make normal diagnosis optional and configurable - [`fd9182a`](https://github.com/deepseek-ai/DeepEP/commit/fd9182a8d4028664158b39853cdf7e1d5cdc3a6d) fix(buffer): preserve backward-compatible optional defaults - [`5caf04e`](https://github.com/deepseek-ai/DeepEP/commit/5caf04e9ae871ffdabf62f0dd3e087f9910b3ca0) fix(normal): harden optional DeepXTrace integration - [`77977ed`](https://github.com/deepseek-ai/DeepEP/commit/77977ed0d7cd81b268d422648731032575e8ed28) fix(normal): isolate diagnostics from low-latency APIs ### 📊 Changes **7 files changed** (+632 additions, -172 deletions) <details> <summary>View changed files</summary> 📝 `csrc/deep_ep.cpp` (+170 -7) 📝 `csrc/deep_ep.hpp` (+30 -2) 📝 `csrc/kernels/api.cuh` (+23 -2) 📝 `csrc/kernels/internode.cu` (+342 -136) 📝 `deep_ep/buffer.py` (+26 -15) 📝 `setup.py` (+4 -1) 📝 `tests/test_internode.py` (+37 -9) </details> ### 📄 Description ## 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: ```python 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](https://github.com/antgroup/DeepXTrace/pull/17). --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-27 21:05:43 +00:00

bsakaguchi referenced this issue

2026-10-27 21:05:48 +00:00

[PR #736] fix: return default_value from get_env<int> on unparsable env values #743

bsakaguchi referenced this issue

2026-10-27 21:05:50 +00:00

[PR #742] Parse EP_BUFFER_DEBUG as an integer on the Python side #747