deepseek-ai/DeepEP

Fork 0

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

[PR #745] [MERGED] [HybridEP] Fix aarch64 inter-node JIT, CUDA graph capture, and test EP topology #750

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/745
Author: @QiaoK
Created: 9/1/2026
Status: ✅ Merged
Merged: 10/16/2026
Merged by: @jershi425

Base: hybrid-ep ← Head: fix/hybrid-ep-aarch64-jit-cudagraph-testtopo

📝 Commits (6)

4f03978 [HybridEP] Use the correct multiarch NIXL lib dir in the inter-node JIT
8747c2a [HybridEP] Let the caller size the dispatch output so it can be captured
abb09bd [HybridEP] Derive the test EP topology from the buffer, not the launcher
7869925 [HybridEP] Cover capturing dispatch() in the graph test
f186b2d [HybridEP] Derive the ragged test EP topology from the buffer too
660cfa8 [HybridEP] Print the derived EP topology in the tests

📊 Changes

11 files changed (+152 additions, -39 deletions)

View changed files

📝 csrc/hybrid_ep/executor/executor.cu (+7 -2)
📝 csrc/hybrid_ep/executor/executor.cuh (+4 -0)
📝 csrc/hybrid_ep/hybrid_ep.cu (+4 -1)
📝 csrc/hybrid_ep/hybrid_ep.cuh (+2 -1)
📝 csrc/hybrid_ep/jit/compiler.cu (+11 -1)
📝 csrc/hybrid_ep/pybind_hybrid_ep.cu (+2 -1)
📝 deep_ep/hybrid_ep_buffer.py (+12 -0)
📝 setup.py (+12 -7)
📝 tests/test_graphed_hybrid_ep.py (+70 -9)
📝 tests/test_hybrid_ep.py (+14 -9)
📝 tests/test_ragged_hybrid_ep.py (+14 -8)

📄 Description

Three independent bugs found while bringing HybridEP up on 2x GB200 nodes with the NIXL backend. None depends on the others, so they are separate commits and can be taken independently.

Rebased onto 15f4acf after #744 (ragged dispatch) merged. All commits applied without conflict, and everything below was re-verified on that base.

Two of them make HybridEP unusable on aarch64 / under CUDA graphs. The third makes a passing configuration look like a data-corruption bug, which is how I found the other two.

1. The inter-node JIT cannot link libnixl on aarch64

setup.py picks the multiarch suffix from platform.machine():

nixl_lib_suffix = "lib/aarch64-linux-gnu" if machine == "aarch64" else "lib/x86_64-linux-gnu"

but csrc/hybrid_ep/jit/compiler.cu hardcodes the x86_64 one. On aarch64 the JIT emits a -L pointing at a directory that does not exist, so every inter-node NIXL compile fails at link time and the NIXL cross-node path is unreachable on GB200 / Grace-Hopper.

Observed on baseline, with NIXL_HOME set correctly:

RuntimeError: Failed to compile the code, compile command: /usr/local/cuda/bin/nvcc -std=c++17 ...
  -L/lustre/.../nixl/install/lib/x86_64-linux-gnu -lnixl -lnixl_build -lnixl_common ...

Note the prefix is right and only the suffix is wrong.

Fix: select the same directory the extension links against. Since that suffix now appears in two places and only covers Debian-style layouts, both honour a NIXL_LIB_DIR override, so other distributions and future architectures do not need another patch and the two selections cannot drift apart.

2. `dispatch()` cannot be CUDA-graph captured

HybridEPBuffer.dispatch() synchronized the stream to make the pinned num_dispatched_tokens_tensor host-readable:

if num_dispatched_tokens is None:
    torch.cuda.current_stream().synchronize()

Synchronizing is illegal while a graph is being captured, so capturing a dispatch without an explicit token count fails:

torch.AcceleratorError: CUDA error: operation not permitted when stream is capturing
Search for `cudaErrorStreamCaptureUnsupported' ...

This is not a hypothetical derived from reading the code: I hit it in TensorRT-LLM's bench_moe_comm.py, which graph-captures the MoE layer, and it is what currently prevents that integration from using HybridEP at all.

Deleting the synchronize is not a fix, and an earlier revision of this PR did exactly that. Thanks to the review for catching it. The underlying problem is that the no-permute dispatch has a data-dependent output shape, and the executor obtains that shape with a host read:

// csrc/hybrid_ep/executor/executor.cu, dispatch_postprocess
int num_dispatched_tokens = args.num_dispatched_tokens_tensor.value().item<int>();
args.local_expert_output_token = torch::empty({num_dispatched_tokens, config.hidden_dim}, ...);
auto res_sz = static_cast<size_t>(num_dispatched_tokens) * config.hidden_dim * sizeof_token_data_type;
CUDA_CHECK(cudaMemcpyAsync(..., res_sz, cudaMemcpyDeviceToDevice, args.stream));

That read happens on the host at capture time, and both the allocation and the D2D copy sizes are frozen into the graph. The kernel that produces the real count is only recorded, never run, and with handle=None metadata_preprocessing allocates a fresh (uninitialized) pinned tensor for it. So removing the synchronize just replaces an illegal wait with a read of memory nothing wrote. Measured on baseline, capturing with an explicit count silently produced a 0-row output — the fresh pinned page read as zero.

The other half of the problem is that num_dispatched_tokens never reached C++ at all. runtime.dispatch() took only hidden, probs, scaling_factor, handle, with_probs, so passing a count merely skipped the synchronize and the value itself was discarded, leaving the executor reading the pinned tensor regardless.

Fix: plumb num_dispatched_tokens through to the executor and honour it, so a caller that needs static shapes can supply an upper bound and get a graph whose allocation and copies are sized by that bound. When no count is supplied under capture, raise an error that says so instead of capturing a shape that cannot be known. Eager behaviour is unchanged: with no count the executor still reads the tensor back, after the same synchronize as before.

tests/test_graphed_hybrid_ep.py could not catch any of this, because it captures dispatch_with_permute(..., num_permuted_tokens=...) and therefore always supplies a count, so the no-permute path was never captured. A case covering it is included, checking both halves of the contract: capture without a count is refused with an actionable message, and a supplied count decides the output shape. The count used deliberately differs from the one the warmup leaves in the pinned buffer, so an implementation that reads the tensor back instead of honouring it fails the test.

3. The tests derive EP topology from the launcher instead of the buffer

The tests only read the detected NVLink-domain width when USE_MNNVL was set, and otherwise assumed one domain per physical node holding --num-processes ranks. On MNNVL hardware a domain can span several nodes, so the buffer forms a wider domain than the launcher implies.

The dispatched prob vector is indexed by rank-within-domain, so when the two disagree the reference slices the wrong expert columns and dispatch appears to return corrupt probabilities:

AssertionError: Dispatch probs mismatch (sparse routing, with_probs=True):
32552/178688 elements (18.22%), 65104 bytes differ, shape=[22336, 8]

Hidden states still compare clean because they are indexed by global rank, which makes this look like a transport bug specific to probs.

Diagnosis. On 2x GB200 nodes sharing an MNNVL clique the buffer reports num_of_hybrid_ep_ranks_per_nvlink_domain=8 and num_of_nodes=1, so dispatched_probs is 64 wide (8 local experts x 8 ranks). The tests assumed 4 ranks/node and sliced at (rank % 4) * 8, so ranks 0-3 were correct by coincidence and ranks 4-7 compared against the columns belonging to ranks 0-3. Re-slicing the same run at the buffer-derived offset (rank % 8) * 8 gives 0 mismatches on all 8 ranks, i.e. the transported data was correct throughout.

Fix: take the topology from the buffer unconditionally. USE_MNNVL is still honoured by the C++ allocator as a switch to disable fabric allocation (USE_MNNVL=0); it just no longer selects how the reference is built. The now-unused Python-side USE_MNNVL constant is removed from the tests.

tests/test_ragged_hybrid_ep.py, added by #744 after this PR was opened, derives its topology the same way and so inherits the same bug. It is fixed here too, in its own commit. Unfixed it reports Ragged dispatch probs mismatch ... 18.43%, the same signature as above.

A/B verification

Same test, same nodes, same container; the only difference is the source tree, selected with PYTHONPATH so both stay built and nothing is rebuilt between the A and B runs.

#	Test	Config	Baseline `58397ea`	This PR
1	`test_hybrid_ep.py`	NIXL cross-node	FAIL — `Failed to compile the code`, `-L .../lib/x86_64-linux-gnu`	PASS (`rc=0`)
2	graph capture, no count	fabric	FAIL — opaque `cudaErrorStreamCaptureInvalidated`	PASS — refused, error names `num_dispatched_tokens`
3	graph capture, explicit count	fabric	FAIL — silently produced 0 rows instead of the requested count	PASS — output shaped by the supplied count
4	`test_hybrid_ep.py`	fabric	FAIL — `Dispatch probs mismatch ... 18.22%`	PASS (`rc=0`)
5	`test_graphed_hybrid_ep.py` (with the new case)	fabric	FAIL — capture error is not actionable	PASS (`rc=0`)
6	`test_ragged_hybrid_ep.py`	fabric	FAIL — `Ragged dispatch probs mismatch ... 18.43%`	PASS (`rc=0`)

Row 6 was run with the same library and only the test file swapped, so it isolates the test-side topology bug. Its baseline run ends rc=1 on one node and rc=124 on the other, the deadlock described below.

Rows 2 and 3 are the two halves of the added test case; run standalone against both trees they give the contrast above, and row 5 is the same check as committed.

On provenance: bugs 1 and 3 were both observed with the stock tests/test_hybrid_ep.py, unmodified. Bug 2 had no native test that could reach it, for the reason given in its section, and the failure it restates was first seen in TensorRT-LLM.

Two notes on how the baseline failures terminate, so the logs are not surprising: in row 4 rank 1 raises the assertion and the remaining ranks then block in the following collective, and in row 1 all ranks raise the compile error together.

Performance

test_hybrid_ep.py run back to back on both trees, same allocation and nodes, both rc=0 (measured on the pre-rebase base 58397ea; the delta this PR adds is unchanged by the rebase). The baseline was given the test-only topology fix so it reaches the benchmark section, so the library is the only difference:

32 metrics (us, lower is better) | mean +0.02%  best -0.08%  worst +0.34%

Dispatch, combine, permute/unpermute and the fused paths are flat across BF16 and FP8. The worst case (FP8 dispatch with probs, 383.0 -> 384.3 us) is well inside run-to-run noise; the min/max spread within a single run is around 7% for the same kernels.

No kernel code changes here. The only hot-path addition is one host-side integer comparison in dispatch_postprocess choosing between the caller's count and reading the tensor back, and with no count supplied the eager path runs exactly the sequence it did before.

How to reproduce

Environment used: 2x GB200 nodes, 4 GPUs each (EP=8), nvcr.io/nvidia/pytorch:26.07-py3, DeepEP built with HYBRID_EP_MULTINODE=1 USE_NIXL=1, UCX and NIXL installed, an etcd reachable from both nodes, and the defaults from tests/test_hybrid_ep.py (hidden 7168, 8 local experts, top-k 8, 4096 tokens/rank).

Common environment:

export NIXL_HOME=/path/to/nixl UCX_HOME=/path/to/ucx CUDA_HOME=/usr/local/cuda
export LD_LIBRARY_PATH=$NIXL_HOME/lib/aarch64-linux-gnu:$NIXL_HOME/lib:$UCX_HOME/lib:$LD_LIBRARY_PATH
export NIXL_PLUGIN_DIR=$NIXL_HOME/lib/aarch64-linux-gnu/plugins
export NIXL_ETCD_ENDPOINTS=http://$HEAD:2379
export MASTER_ADDR=$HEAD MASTER_PORT=31501 WORLD_SIZE=2 RANK=$SLURM_NODEID

fabric configuration (default auto-detect; on GB200 both nodes join one MNNVL domain) — reproduces bug 3:

python tests/test_hybrid_ep.py --only-bf16

NIXL cross-node configuration — reproduces bug 1. Fabric allocation has to be disabled and the domain pinned to one physical node, otherwise the two nodes form a single NVLink domain and the inter-node path is never taken:

export USE_MNNVL=0
export NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=4
python tests/test_hybrid_ep.py --only-bf16

CUDA graph capture — reproduces bug 2. This is now covered by tests/test_graphed_hybrid_ep.py, so on this branch it is just:

python tests/test_graphed_hybrid_ep.py

To see the baseline behaviour, copy that file onto the baseline tree and run it there, or use this standalone version:

import torch, torch.distributed as dist
import deep_ep
from utils import init_dist

def main(local_rank, num_local_ranks, _args=None):
    rank, world, group = init_dist(local_rank, num_local_ranks)
    num_experts = 8 * world
    buffer = deep_ep.HybridEPBuffer(group=group, hidden_dim=7168,
                                    max_num_of_tokens_per_rank=4096,
                                    num_local_experts=8, use_fp8=False)
    hidden = torch.randn(4096, 7168, device="cuda", dtype=torch.bfloat16)
    scaling_factor = torch.randn(4096, 56, device="cuda", dtype=torch.float32)
    probs = torch.zeros(4096, num_experts, device="cuda", dtype=torch.float32)
    routing_map = torch.zeros(4096, num_experts, device="cuda", dtype=torch.bool)
    for i in range(4096):
        sel = torch.randperm(num_experts, device="cuda")[:8]
        routing_map[i, sel] = True
        probs[i, sel] = 1.0

    # Warm up outside capture. This also leaves the warmup's own count in the
    # pinned buffer, which the explicit count below is checked against.
    warm, _, _, _ = buffer.dispatch(hidden=hidden, scaling_factor=scaling_factor,
                                    routing_map=routing_map, probs=probs)
    warm_count = warm.shape[0]
    torch.cuda.synchronize(); dist.barrier()

    # No count: the shape cannot be known, so this must be refused clearly.
    try:
        with torch.cuda.graph(torch.cuda.CUDAGraph()):
            buffer.dispatch(hidden=hidden, scaling_factor=scaling_factor,
                            routing_map=routing_map, probs=probs)
        print("case 1: capture unexpectedly succeeded")
    except RuntimeError as exc:
        ok = "num_dispatched_tokens" in str(exc)
        print(f"case 1: {'clear' if ok else 'UNCLEAR'} error: {exc}")
    torch.cuda.synchronize()

    # Explicit count: it must decide the output shape. Staying below the warmup
    # count keeps the sized copies in bounds while differing from the pinned value.
    target = max(1, warm_count - 128)
    with torch.cuda.graph(torch.cuda.CUDAGraph()):
        out, _, _, _ = buffer.dispatch(hidden=hidden, scaling_factor=scaling_factor,
                                       routing_map=routing_map, probs=probs,
                                       num_dispatched_tokens=target)
    print(f"case 2: {out.shape[0]} rows, expected {target} (warmup was {warm_count})")
    dist.barrier(); dist.destroy_process_group()

if __name__ == "__main__":
    torch.multiprocessing.spawn(main, args=(4, None), nprocs=4)

On baseline this reports an unclear cudaErrorStreamCaptureInvalidated for case 1 and 0 rows for case 2; on this branch case 1 names num_dispatched_tokens and case 2 returns exactly target rows.

_{🔄 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/745 **Author:** [@QiaoK](https://github.com/QiaoK) **Created:** 9/1/2026 **Status:** ✅ Merged **Merged:** 10/16/2026 **Merged by:** [@jershi425](https://github.com/jershi425) **Base:** `hybrid-ep` ← **Head:** `fix/hybrid-ep-aarch64-jit-cudagraph-testtopo` --- ### 📝 Commits (6) - [`4f03978`](https://github.com/deepseek-ai/DeepEP/commit/4f039785b98f7c7fc029d7bf5f96bce08cd17f08) [HybridEP] Use the correct multiarch NIXL lib dir in the inter-node JIT - [`8747c2a`](https://github.com/deepseek-ai/DeepEP/commit/8747c2a63551d09556709d1950fbd465e51e13f8) [HybridEP] Let the caller size the dispatch output so it can be captured - [`abb09bd`](https://github.com/deepseek-ai/DeepEP/commit/abb09bdf586e6590e619516a957389ac4a296ff2) [HybridEP] Derive the test EP topology from the buffer, not the launcher - [`7869925`](https://github.com/deepseek-ai/DeepEP/commit/78699256f8ebbd902aa80fa41ece49264c8c2a13) [HybridEP] Cover capturing dispatch() in the graph test - [`f186b2d`](https://github.com/deepseek-ai/DeepEP/commit/f186b2d68cb859f30af9bc8c2116d8c68c50e913) [HybridEP] Derive the ragged test EP topology from the buffer too - [`660cfa8`](https://github.com/deepseek-ai/DeepEP/commit/660cfa8ece60c314d13f610503cf42704f33167a) [HybridEP] Print the derived EP topology in the tests ### 📊 Changes **11 files changed** (+152 additions, -39 deletions) <details> <summary>View changed files</summary> 📝 `csrc/hybrid_ep/executor/executor.cu` (+7 -2) 📝 `csrc/hybrid_ep/executor/executor.cuh` (+4 -0) 📝 `csrc/hybrid_ep/hybrid_ep.cu` (+4 -1) 📝 `csrc/hybrid_ep/hybrid_ep.cuh` (+2 -1) 📝 `csrc/hybrid_ep/jit/compiler.cu` (+11 -1) 📝 `csrc/hybrid_ep/pybind_hybrid_ep.cu` (+2 -1) 📝 `deep_ep/hybrid_ep_buffer.py` (+12 -0) 📝 `setup.py` (+12 -7) 📝 `tests/test_graphed_hybrid_ep.py` (+70 -9) 📝 `tests/test_hybrid_ep.py` (+14 -9) 📝 `tests/test_ragged_hybrid_ep.py` (+14 -8) </details> ### 📄 Description Three independent bugs found while bringing HybridEP up on 2x GB200 nodes with the NIXL backend. None depends on the others, so they are separate commits and can be taken independently. Rebased onto `15f4acf` after #744 (ragged dispatch) merged. All commits applied without conflict, and everything below was re-verified on that base. Two of them make HybridEP unusable on aarch64 / under CUDA graphs. The third makes a passing configuration look like a data-corruption bug, which is how I found the other two. --- ## 1. The inter-node JIT cannot link libnixl on aarch64 `setup.py` picks the multiarch suffix from `platform.machine()`: ```python nixl_lib_suffix = "lib/aarch64-linux-gnu" if machine == "aarch64" else "lib/x86_64-linux-gnu" ``` but `csrc/hybrid_ep/jit/compiler.cu` hardcodes the x86_64 one. On aarch64 the JIT emits a `-L` pointing at a directory that does not exist, so **every** inter-node NIXL compile fails at link time and the NIXL cross-node path is unreachable on GB200 / Grace-Hopper. Observed on baseline, with `NIXL_HOME` set correctly: ``` RuntimeError: Failed to compile the code, compile command: /usr/local/cuda/bin/nvcc -std=c++17 ... -L/lustre/.../nixl/install/lib/x86_64-linux-gnu -lnixl -lnixl_build -lnixl_common ... ``` Note the prefix is right and only the suffix is wrong. **Fix:** select the same directory the extension links against. Since that suffix now appears in two places and only covers Debian-style layouts, both honour a `NIXL_LIB_DIR` override, so other distributions and future architectures do not need another patch and the two selections cannot drift apart. ## 2. `dispatch()` cannot be CUDA-graph captured `HybridEPBuffer.dispatch()` synchronized the stream to make the pinned `num_dispatched_tokens_tensor` host-readable: ```python if num_dispatched_tokens is None: torch.cuda.current_stream().synchronize() ``` Synchronizing is illegal while a graph is being captured, so capturing a dispatch without an explicit token count fails: ``` torch.AcceleratorError: CUDA error: operation not permitted when stream is capturing Search for `cudaErrorStreamCaptureUnsupported' ... ``` This is not a hypothetical derived from reading the code: I hit it in TensorRT-LLM's `bench_moe_comm.py`, which graph-captures the MoE layer, and it is what currently prevents that integration from using HybridEP at all. **Deleting the synchronize is not a fix, and an earlier revision of this PR did exactly that.** Thanks to the review for catching it. The underlying problem is that the no-permute dispatch has a *data-dependent output shape*, and the executor obtains that shape with a host read: ```cpp // csrc/hybrid_ep/executor/executor.cu, dispatch_postprocess int num_dispatched_tokens = args.num_dispatched_tokens_tensor.value().item<int>(); args.local_expert_output_token = torch::empty({num_dispatched_tokens, config.hidden_dim}, ...); auto res_sz = static_cast<size_t>(num_dispatched_tokens) * config.hidden_dim * sizeof_token_data_type; CUDA_CHECK(cudaMemcpyAsync(..., res_sz, cudaMemcpyDeviceToDevice, args.stream)); ``` That read happens on the host *at capture time*, and both the allocation and the D2D copy sizes are frozen into the graph. The kernel that produces the real count is only recorded, never run, and with `handle=None` `metadata_preprocessing` allocates a fresh (uninitialized) pinned tensor for it. So removing the synchronize just replaces an illegal wait with a read of memory nothing wrote. Measured on baseline, capturing with an explicit count silently produced a **0-row output** — the fresh pinned page read as zero. The other half of the problem is that `num_dispatched_tokens` never reached C++ at all. `runtime.dispatch()` took only `hidden, probs, scaling_factor, handle, with_probs`, so passing a count merely skipped the synchronize and the value itself was discarded, leaving the executor reading the pinned tensor regardless. **Fix:** plumb `num_dispatched_tokens` through to the executor and honour it, so a caller that needs static shapes can supply an upper bound and get a graph whose allocation and copies are sized by that bound. When no count is supplied under capture, raise an error that says so instead of capturing a shape that cannot be known. Eager behaviour is unchanged: with no count the executor still reads the tensor back, after the same synchronize as before. `tests/test_graphed_hybrid_ep.py` could not catch any of this, because it captures `dispatch_with_permute(..., num_permuted_tokens=...)` and therefore always supplies a count, so the no-permute path was never captured. A case covering it is included, checking both halves of the contract: capture without a count is refused with an actionable message, and a supplied count decides the output shape. The count used deliberately differs from the one the warmup leaves in the pinned buffer, so an implementation that reads the tensor back instead of honouring it fails the test. ## 3. The tests derive EP topology from the launcher instead of the buffer The tests only read the detected NVLink-domain width when `USE_MNNVL` was set, and otherwise assumed one domain per physical node holding `--num-processes` ranks. On MNNVL hardware a domain can span several nodes, so the buffer forms a wider domain than the launcher implies. The dispatched prob vector is indexed by rank-within-domain, so when the two disagree the reference slices the wrong expert columns and dispatch appears to return corrupt probabilities: ``` AssertionError: Dispatch probs mismatch (sparse routing, with_probs=True): 32552/178688 elements (18.22%), 65104 bytes differ, shape=[22336, 8] ``` Hidden states still compare clean because they are indexed by global rank, which makes this look like a transport bug specific to probs. **Diagnosis.** On 2x GB200 nodes sharing an MNNVL clique the buffer reports `num_of_hybrid_ep_ranks_per_nvlink_domain=8` and `num_of_nodes=1`, so `dispatched_probs` is 64 wide (8 local experts x 8 ranks). The tests assumed 4 ranks/node and sliced at `(rank % 4) * 8`, so ranks 0-3 were correct by coincidence and ranks 4-7 compared against the columns belonging to ranks 0-3. Re-slicing the same run at the buffer-derived offset `(rank % 8) * 8` gives **0 mismatches on all 8 ranks**, i.e. the transported data was correct throughout. **Fix:** take the topology from the buffer unconditionally. `USE_MNNVL` is still honoured by the C++ allocator as a switch to disable fabric allocation (`USE_MNNVL=0`); it just no longer selects how the reference is built. The now-unused Python-side `USE_MNNVL` constant is removed from the tests. `tests/test_ragged_hybrid_ep.py`, added by #744 after this PR was opened, derives its topology the same way and so inherits the same bug. It is fixed here too, in its own commit. Unfixed it reports `Ragged dispatch probs mismatch ... 18.43%`, the same signature as above. --- ## A/B verification Same test, same nodes, same container; the only difference is the source tree, selected with `PYTHONPATH` so both stay built and nothing is rebuilt between the A and B runs. | # | Test | Config | Baseline `58397ea` | This PR | |---|---|---|---|---| | 1 | `test_hybrid_ep.py` | NIXL cross-node | FAIL — `Failed to compile the code`, `-L .../lib/x86_64-linux-gnu` | PASS (`rc=0`) | | 2 | graph capture, no count | fabric | FAIL — opaque `cudaErrorStreamCaptureInvalidated` | PASS — refused, error names `num_dispatched_tokens` | | 3 | graph capture, explicit count | fabric | FAIL — silently produced **0 rows** instead of the requested count | PASS — output shaped by the supplied count | | 4 | `test_hybrid_ep.py` | fabric | FAIL — `Dispatch probs mismatch ... 18.22%` | PASS (`rc=0`) | | 5 | `test_graphed_hybrid_ep.py` (with the new case) | fabric | FAIL — capture error is not actionable | PASS (`rc=0`) | | 6 | `test_ragged_hybrid_ep.py` | fabric | FAIL — `Ragged dispatch probs mismatch ... 18.43%` | PASS (`rc=0`) | Row 6 was run with the same library and only the test file swapped, so it isolates the test-side topology bug. Its baseline run ends `rc=1` on one node and `rc=124` on the other, the deadlock described below. Rows 2 and 3 are the two halves of the added test case; run standalone against both trees they give the contrast above, and row 5 is the same check as committed. On provenance: bugs 1 and 3 were both observed with the stock `tests/test_hybrid_ep.py`, unmodified. Bug 2 had no native test that could reach it, for the reason given in its section, and the failure it restates was first seen in TensorRT-LLM. Two notes on how the baseline failures terminate, so the logs are not surprising: in row 4 rank 1 raises the assertion and the remaining ranks then block in the following collective, and in row 1 all ranks raise the compile error together. ## Performance `test_hybrid_ep.py` run back to back on both trees, same allocation and nodes, both `rc=0` (measured on the pre-rebase base `58397ea`; the delta this PR adds is unchanged by the rebase). The baseline was given the test-only topology fix so it reaches the benchmark section, so the library is the only difference: ``` 32 metrics (us, lower is better) | mean +0.02% best -0.08% worst +0.34% ``` Dispatch, combine, permute/unpermute and the fused paths are flat across BF16 and FP8. The worst case (FP8 dispatch with probs, 383.0 -> 384.3 us) is well inside run-to-run noise; the min/max spread within a single run is around 7% for the same kernels. No kernel code changes here. The only hot-path addition is one host-side integer comparison in `dispatch_postprocess` choosing between the caller's count and reading the tensor back, and with no count supplied the eager path runs exactly the sequence it did before. ## How to reproduce Environment used: 2x GB200 nodes, 4 GPUs each (EP=8), `nvcr.io/nvidia/pytorch:26.07-py3`, DeepEP built with `HYBRID_EP_MULTINODE=1 USE_NIXL=1`, UCX and NIXL installed, an etcd reachable from both nodes, and the defaults from `tests/test_hybrid_ep.py` (hidden 7168, 8 local experts, top-k 8, 4096 tokens/rank). Common environment: ```bash export NIXL_HOME=/path/to/nixl UCX_HOME=/path/to/ucx CUDA_HOME=/usr/local/cuda export LD_LIBRARY_PATH=$NIXL_HOME/lib/aarch64-linux-gnu:$NIXL_HOME/lib:$UCX_HOME/lib:$LD_LIBRARY_PATH export NIXL_PLUGIN_DIR=$NIXL_HOME/lib/aarch64-linux-gnu/plugins export NIXL_ETCD_ENDPOINTS=http://$HEAD:2379 export MASTER_ADDR=$HEAD MASTER_PORT=31501 WORLD_SIZE=2 RANK=$SLURM_NODEID ``` **fabric configuration** (default auto-detect; on GB200 both nodes join one MNNVL domain) — reproduces bug 3: ```bash python tests/test_hybrid_ep.py --only-bf16 ``` **NIXL cross-node configuration** — reproduces bug 1. Fabric allocation has to be disabled and the domain pinned to one physical node, otherwise the two nodes form a single NVLink domain and the inter-node path is never taken: ```bash export USE_MNNVL=0 export NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=4 python tests/test_hybrid_ep.py --only-bf16 ``` **CUDA graph capture** — reproduces bug 2. This is now covered by `tests/test_graphed_hybrid_ep.py`, so on this branch it is just: ```bash python tests/test_graphed_hybrid_ep.py ``` To see the baseline behaviour, copy that file onto the baseline tree and run it there, or use this standalone version: ```python import torch, torch.distributed as dist import deep_ep from utils import init_dist def main(local_rank, num_local_ranks, _args=None): rank, world, group = init_dist(local_rank, num_local_ranks) num_experts = 8 * world buffer = deep_ep.HybridEPBuffer(group=group, hidden_dim=7168, max_num_of_tokens_per_rank=4096, num_local_experts=8, use_fp8=False) hidden = torch.randn(4096, 7168, device="cuda", dtype=torch.bfloat16) scaling_factor = torch.randn(4096, 56, device="cuda", dtype=torch.float32) probs = torch.zeros(4096, num_experts, device="cuda", dtype=torch.float32) routing_map = torch.zeros(4096, num_experts, device="cuda", dtype=torch.bool) for i in range(4096): sel = torch.randperm(num_experts, device="cuda")[:8] routing_map[i, sel] = True probs[i, sel] = 1.0 # Warm up outside capture. This also leaves the warmup's own count in the # pinned buffer, which the explicit count below is checked against. warm, _, _, _ = buffer.dispatch(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs) warm_count = warm.shape[0] torch.cuda.synchronize(); dist.barrier() # No count: the shape cannot be known, so this must be refused clearly. try: with torch.cuda.graph(torch.cuda.CUDAGraph()): buffer.dispatch(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs) print("case 1: capture unexpectedly succeeded") except RuntimeError as exc: ok = "num_dispatched_tokens" in str(exc) print(f"case 1: {'clear' if ok else 'UNCLEAR'} error: {exc}") torch.cuda.synchronize() # Explicit count: it must decide the output shape. Staying below the warmup # count keeps the sized copies in bounds while differing from the pinned value. target = max(1, warm_count - 128) with torch.cuda.graph(torch.cuda.CUDAGraph()): out, _, _, _ = buffer.dispatch(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs, num_dispatched_tokens=target) print(f"case 2: {out.shape[0]} rows, expected {target} (warmup was {warm_count})") dist.barrier(); dist.destroy_process_group() if __name__ == "__main__": torch.multiprocessing.spawn(main, args=(4, None), nprocs=4) ``` On baseline this reports an unclear `cudaErrorStreamCaptureInvalidated` for case 1 and `0 rows` for case 2; on this branch case 1 names `num_dispatched_tokens` and case 2 returns exactly `target` rows. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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

closed this issue
added the
pull-request
label

bsakaguchi referenced this issue

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

[PR #750] [MERGED] [HybridEP] Add DMA-BUF fallback for GPU RDMA memory registration #754

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepEP#750

No description provided.

Rows
Columns