feat(normal): support DeepEP normal-mode diagnosis

[Marsssqqq]

Summary

This PR adds end-to-end diagnosis for DeepEP's throughput-oriented Normal dispatch and combine path. It allocates the cumulative probe counters, collects stream-ordered snapshots, constructs independent diagnosis windows, localizes six Notify and Completion probes, correlates their evidence into conservative incident classes, and exports structured logs, NPZ artifacts, and heatmaps.

Motivation

Normal mode and low-latency mode serve different workloads. Normal mode is used for MoE training and inference prefill, where larger token sets traverse a hierarchical RDMA/NVLink path and communication overlaps with upstream computation and scheduling. A slow operation can therefore mean either that a rank arrived late or that progress slowed after arrival.

The existing low-latency receive-wait matrix cannot express all of these stages. Normal diagnosis needs separate arrival and completion evidence, topology-aware localization, and conservative cross-probe correlation.

Runtime integration

Diagnosis is disabled by default and configured by DeepXTrace:

export DEEPEP_DIAGNOSE_ENABLE=1
# 1: periodic asynchronous collection; 0: caller-driven synchronous collection.
export DEEPEP_DIAGNOSE_ASYNC=1

Every EP rank calls Diagnose.create_from_env(). The factory verifies that enablement and collection mode are consistent across the group. When disabled, it returns None.

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 and not diagnose.enable_async:
    diagnose.diagnose_normal_sync()

if diagnose is not None:
    diagnose.stop()

The Dispatch getter returns a ten-tensor opaque bundle; the Combine getter returns a five-tensor opaque bundle. Applications pass both bundles without unpacking or reordering them. Notify timer scratch remains private to DeepEP.

Counter layout and normalization

All Normal counters are views into one contiguous CUDA int64 allocation:

Counter group	Local shape	Meaning
three Notify duration/count pairs	3 x 2	cumulative full-grid duration in ns and launch count
Dispatch final cost/sample/token	3 x EP	cumulative values by logical source rank
Dispatch RDMA receive cost/sample/token	3 x EP	cumulative values by source gateway proxy rank
Combine logical receive cost/sample/token	3 x EP	cumulative values by logical source rank

After gathering one local vector from every destination rank, the Completion groups are transposed from [destination, source] into source-oriented matrices. Notify metrics use duration_sum_ns / count; Completion metrics use cost_sum_cycles / sample_count, while tokens / sample_count remains available as load context.

Ordered snapshot and window collection

Both schedules operate on the same monotonic counter block:

• synchronous collection enqueues one D2D copy on the supplied counter-writer/DeepEP communication stream, waits for the staging event, and gathers the stable GPU snapshot through the original EP process group;
• asynchronous collection enqueues the same D2D copy, makes a private CUDA stream wait for its event, copies the stable staging buffer to pinned CPU memory, and gathers the CPU snapshots through a matching Gloo group. The D2H transfer therefore does not read counters while later communication kernels update them.

Passing the producer stream is the ordered path used by the companion integration. If asynchronous collection is explicitly used without it, DeepXTrace emits a warning and falls back to an unordered best-effort D2H snapshot.

For non-world or non-contiguous EP subgroups, DeepXTrace resolves the ordered global ranks, creates the matching Gloo group for asynchronous collection, and uses the correct global root. Rank 0 subtracts consecutive cumulative snapshots to form the current window. At the end of warm-up, a new cumulative boundary is captured and all earlier samples are discarded.

Normal probes

Probe	Evidence	Diagnostic role
notify_dispatch	per-rank full-kernel duration/count	non-cached dispatch arrival skew
cached_notify_dispatch	per-rank full-kernel duration/count	cached dispatch arrival skew
cached_notify_combine	per-rank full-kernel duration/count	cached combine arrival skew
dispatch_final	dense source × destination cost/sample/token matrices	final dispatch completion
dispatch_rdma_recv	gateway-oriented cost/sample/token matrices	RDMA gateway receive completion
combine_logical_recv	dense source × destination cost/sample/token matrices	logical source receive completion

Semantic diagnosis

Independent probe analysis

Each Notify probe compares valid per-rank average duration with the current peer median. A ratio at or below 0.25 is raw late-arrival evidence: a rank entering a synchronizing Notify kernel later typically spends less time waiting inside that kernel. The three Notify paths are not merged.

Completion probes first construct a same-window topology reference:

• dense Dispatch-final and Combine matrices use the median of each source-node/destination-node block, so intra-node and inter-node paths are compared with their own peers;
• the sparse RDMA receive matrix uses the median of its valid gateway edges.

The diagnoser applies log(actual / reference), removes the window-wide, source-row, and destination-column effects in order, and then tests the remaining point/edge residual. Dense matrices additionally compare node blocks with other blocks of the same intra-node or inter-node class. This produces source rows, destination columns, points, node blocks, source/destination gateways, and gateway edges without reporting the same shared slowdown repeatedly at every cell.

Global slowdown is intentionally simpler than per-cell historical comparison: after three accepted normal windows establish a baseline, one topology-aggregated scalar per probe is compared with the median of up to 16 previous normal windows. An anomalous window is not learned into that history. This baseline bootstrap is separate from the collection warm-up boundary described above.

Decision	Default
Notify late-arrival ratio	duration / peer median <= 0.25
Source row, destination column, gateway axis, or dense block slowdown	>= 1.25x
Point or gateway-edge residual slowdown	>= 2.0x
Robust within-window outlier score for row/column/point/edge	>= 3.5
Global slowdown versus recent normal history	>= 1.3x
Global-history bootstrap / horizon	3 / 16 accepted normal windows

The three user-facing ratios remain configurable through DEEPEP_DIAGNOSE_* environment variables; the history and robust-score constants are implementation policy.

Cross-probe correlation

Correlation is conservative and does not erase the six independent probe results:

Raw evidence	Incident result
notify_dispatch or cached_notify_dispatch anomaly, with both dispatch_final and dispatch_rdma_recv normal	pre_stage_arrival_skew for that Notify probe
cached_notify_combine anomaly, with combine_logical_recv normal	pre_stage_arrival_skew for that Notify probe
Notify anomaly with anomalous or insufficient downstream Completion evidence	retain raw anomaly; do not force a pre-stage incident
Dispatch-final and Combine both show global slowdown	global_data_path_slowdown
Dispatch-final and Combine have intersecting local scopes	localized_data_path_slowdown
Completion anomaly without supporting cross-probe evidence	unclassified

RDMA receive evidence is attached to a global or localized data-path incident when its scope supports the same conclusion; it is not treated as a fixed HCA-to-rank mapping.

Logging, artifacts, and heatmaps

• Log level 0 emits the compact per-window summary.
• DEEPEP_DIAGNOSE_LOG_DETAILS=1 additionally emits all six per-probe details and the full source-oriented vectors/matrices. Every Normal summary, detail, and values line carries the same InstanceID, EPSize, and WindowIndex context; each detail is immediately followed by its values.
• Optional NPZ export stores one compressed, versioned, self-describing Normal window and supports bounded retention without changing the runtime probe contract.
• tools/deepxtrace_heatmap.py --all <window.npz> renders one Notify chart and three Completion heatmaps.

Validation

EP16 setup

The functional run used two 8-GPU NVIDIA H20-3e nodes (EP16), PyTorch 2.11.0+cu130, an asynchronous one-second interval with a one-second collection warm-up, detail logging level 1, 4,096 tokens, hidden size 4,096, top-k 8, 128 experts, and 24 communication SMs. The Normal APIs were exercised after initializing NVSHMEM with low_latency_mode=True and explicit NIC-PE mapping across four bonded HCAs. The workload ran 512 warm-up iterations followed by 512 measured iterations per phase. 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.

Five-phase result

The window column identifies the NPZ artifact used for the heatmaps; numeric signatures summarize all 512 measured iterations in that phase.

Phase	Representative heatmap window	Phase-level signature	Diagnosis
Baseline	3	All probes ok; no fixed source-row or destination-column bias	none
Rank-9 pre-dispatch compute delay	8	rank 9 Notify ratios: 0.036 non-cached and 0.006 cached; all Completion medians remained within 0.2% of baseline	two pre_stage_arrival_skew incidents for rank 9
Rank-9 pre-combine compute delay	16	rank 9 cached-Combine Notify ratio 0.006; all Completion medians remained within 0.3% of baseline	pre_stage_arrival_skew for rank 9
All-HCA RDMA pressure	26	median Dispatch-final / RDMA-receive / Combine completion cost rose to 1.371x / 1.447x / 1.514x baseline	global_data_path_slowdown
Single-HCA RDMA pressure	58	localized Dispatch/Combine scope agreed on token-owner ranks 0/1/8/9	localized_data_path_slowdown

The current tools/deepxtrace_heatmap.py --all implementation generated four images from each representative NPZ window: one three-panel Notify-duration chart plus Dispatch-final, Dispatch-RDMA-receive, and Combine-logical-receive cost-per-sample heatmaps.

Heatmap evidence

Baseline — window 3

Rank-9 pre-dispatch compute delay — window 8

Rank-9 pre-combine compute delay — window 16

All-HCA RDMA pressure — window 26

Single-HCA RDMA pressure — window 58

Each heatmap uses an independent automatic color scale; cross-window comparisons should use the annotated cell values rather than colors alone.

Compatibility

• Existing low-latency diagnosis APIs remain available.
• The NPZ artifact schema has its own version because artifacts are persistent files; the in-process bundle interface does not expose a schema version.
• The package version changes from 0.1.0 to 0.2.0.

Companion DeepEP PR

The device-side Normal Notify, Dispatch, and Combine probes, together with the optional Dispatch/Combine tensor-bundle arguments, are implemented in deepseek-ai/DeepEP#720.

[Copilot]

Pull request overview

Adds DeepEP “normal mode” semantic diagnosis to DeepXTrace, including a stable cumulative-counter stats contract, windowing/correlation logic, optional structured NPZ artifacts, and updated heatmap tooling to visualize Normal windows (while keeping LL diagnosis intact).

Changes:

• Introduces normal-mode semantic diagnoser + structured normal-window artifact format/load/export.
• Extends Diagnose to collect/gather normal-mode cumulative counters (sync or async), correlate evidence, and optionally export window artifacts.
• Updates tools/deepxtrace_heatmap.py + docs/tests to render Normal .npz artifacts and enforce stable default layout.

Reviewed changes

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

Show a summary per file

File	Description
tools/README.md	Documents Normal artifact rendering, stable layout presets, and new CLI flags.
tools/deepxtrace_heatmap.py	Adds Normal artifact rendering, stable layouts, log-scale completion heatmaps, and Notify panel rendering.
tests/test_normal_tools.py	Adds tests for artifact round-trip and heatmap CLI behavior (mocked plotting stack).
tests/test_normal_diagnose.py	Adds unit tests for normal-mode notify/completion diagnosis and cross-probe correlation.
tests/test_diagnose.py	Adds tests for subgroup root handling, normal windowing, and async snapshot ordering.
src/deepxtrace/normal_diagnose.py	Implements normal-mode preprocessing, per-probe diagnosis, and incident correlation.
src/deepxtrace/normal_artifact.py	Implements schema-versioned .npz artifact export + metric loading/validation.
src/deepxtrace/diagnose.py	Adds normal-mode stats schema, snapshot/gather/windowing, logging/export, and async loop updates.
setup.py	Bumps package version to 0.2.0.
README.md	Documents normal-mode integration contract, logging, and artifact/heatmap workflow.
.gitignore	Ignores generated visualization artifacts under tools/.

Suppressed comments (1)

tools/deepxtrace_heatmap.py:100

• calculate_nonzero_scale_min() exits when there are no positive values. Combined with LogNorm + masking, an all-zero matrix can be rendered safely; exiting here prevents generating heatmaps for sparse/empty probes.

def calculate_nonzero_scale_min(matrix):
    """Use the smallest finite positive value as the matrix scale minimum."""
    positive_values = matrix[np.isfinite(matrix) & (matrix > 0)]
    if positive_values.size == 0:
        raise SystemExit(
            "Error: Heatmap data must contain at least one positive value.")

---

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