[IlyasMoutawwakil]

Depends on #49615.

What this adds

A "deepep" expert-parallel implementation. It runs DeepEP's dispatch and combine, loaded from the Hub (kernels-community/deep-ep), as a new EPImplementation in ALL_EP_IMPLEMENTATIONS.

• Expanded layout. Each rank receives one row per (token, local expert) slot, grouped by local expert, with that slot's routing weight. The experts only see their own routes and apply the weights; the combine only sums.
• Eager vs captured. Eagerly, DeepEP reads the received counts back to the host, so the grouped GEMM gets exactly the received rows (the prefill / large-batch path). Under CUDA graph capture or torch.compile (torch.cuda.is_current_stream_capturing() / torch.compiler.is_compiling()), it runs without the host sync. The rows are sized for the worst case (every rank sending its whole capacity with all of its top-k experts on this rank), so they cannot overflow. Rows past the received total are zeroed and routed to the sentinel expert, which the experts skip.
• Custom ops. Dispatch, combine and the combine's backward (a cached dispatch) are torch.library custom ops with fake implementations and registered autograd, so the implementation trains, captures and compiles (fullgraph=True at the experts level). The DeepEP handle can't go through an op, so the ops pass its tensors and rebuild it.
• Training. Gradients reach the tokens and the routing weights. The weight gradients go back through combine(..., topk_weights=slot_grads), and the combine's backward dispatches along the forward's layout.
• Buffers. There is one DeepEP buffer per EP group, shared by every layer (like Mega MoE's symmetric buffers). It is built lazily and rebuilt larger, as a collective, when a batch outgrows it, and it is never allocated during capture. Warm up eagerly at the capture batch size, as usual; otherwise the capture raises an error that says so. transformers.integrations.deepep.destroy_deepep_buffers() frees the buffers.
• TP. With tp_size > 1, the ranks of a TP group each dispatch a disjoint slice of the batch they share, like AllToAllEP. Both implementations share the slicing and gathering helpers.
• integrations/deepep.py mirrors integrations/deepgemm.py. is_deepep_loadable(raise_error=False) (lru-cached) checks for CUDA, SM ≥ 90, a CUDA toolkit ≥ 13.1 (found by DeepGEMM's _get_cuda_home / _get_nvcc_version, which are reused), exactly one NCCL ≥ 2.32.3 loaded by PyTorch (it reads the runtime version through ncclGetVersion, since torch.cuda.nccl.version() is the build-time one), and whether the Hub kernel loads. load_deepep_kernel() raises the ImportError that names the missing requirement. ep_implementation="deepep" fails at load with that message.

Scope: intranode only (hybrid mode is left at DeepEP's default and is untested across nodes). There is no low-latency mode; it is gone from DeepEP 2.5. FP8 dispatch (use_fp8_dispatch, half the dispatch bytes for FP8 recipes) is a follow-up.

Usage

DistributedConfig(tp_size=1, fsdp_size=8, ep_size=8, ep_implementation="deepep")

Requirements: the kernels-community/deep-ep Hub kernel (version 1, from kernels-community#1231, not merged yet; until then, the staged build works through LOCAL_KERNELS=kernels-community/deep-ep=<snapshot of kernels-staging/deep-ep@pr-1231>). You also need the NCCL that PyTorch loads at ≥ 2.32.3 (e.g. nvidia-nccl-cu13==2.32.3), a CUDA toolkit ≥ 13.1 at runtime (nvcc, through CUDA_HOME, since DeepEP JIT-compiles its kernels), SM90+, and a hidden size that's a multiple of 256 (DeepEP's combine; anything else gets a clear ValueError).

"auto"

"auto" now prefers DeepEP wherever it used to pick all-to-all. Whether DeepEP is usable is decided once per process: an NCCL EP group and is_deepep_loadable(). If it isn't, "auto" falls back to all-to-all with a single logger.info_once.

regime	DeepEP usable	DeepEP not usable
shared batch (ep_size == tp_size), inference	all_reduce	all_reduce
training, or ranks holding different batches (tp_size < ep_size), eager	deepep (host sync)	all_to_all
any token dispatch under CUDA graph capture / torch.compile	deepep, sync-free	all_to_all (not capturable)

No token threshold: in DeepEP's expanded layout, each rank only runs its own expert routes, so it never loses to the all-to-all by more than run-to-run noise (below ~1024 tokens/rank every eager implementation is host-bound) and wins from ~4096 tokens/rank. A cutoff tuned on one EP size and GPU would not carry over to others.

Benchmarks

Hardware and software: B200 node (NVLink), torch 2.14.0+cu130, NCCL 2.32.3, DeepEP 2.5.0 (staged kernels-staging/deep-ep@337d4ae, nvcc 13.2), experts grouped_mm in bf16.

Setup:

• Layer shapes: 30b is Qwen3-30B-A3B (E=128, H=2048, I=768, top-8) and h4096 is E=128, H=4096, I=1536, top-8.
• Each config runs four distinct experts layers back to back, so the weights are cold-ish, in one warmed process. The time is the median of 5 windows of back-to-back steps, per layer, on the slowest rank.
• "tok" is the batch each rank holds. With TP = EP, that batch is shared and each rank dispatches 1/TP of it.
• CUDA graph rows capture after an eager warm-up. all_to_all is not capturable (n/c) because it reads split sizes back to the host.
• EP=8 was not run: two of the eight GPUs were serving other jobs.

Bold marks the fastest implementation in each cell.

inference, eager (ms per layer; all_to_all / all_reduce / deepep)

shape	EP	TP	1 tok	256 tok	1024 tok	4096 tok	16384 tok
30b	2	1	1.44 / – / 1.44	1.26 / – / 0.94	1.44 / – / 1.52	2.96 / – / 2.36	9.75 / – / 6.40
30b	2	2	3.39 / 0.52 / 1.59	3.12 / 1.38 / 1.58	3.28 / 0.90 / 1.56	3.46 / 1.33 / 1.59	5.35 / 4.26 / 3.40
30b	4	1	2.14 / – / 1.48	2.14 / – / 1.43	2.27 / – / 1.59	3.17 / – / 2.35	10.14 / – / 6.62
30b	4	4	3.58 / 0.88 / 2.59	3.28 / 0.91 / 2.37	3.45 / 0.95 / 2.59	3.57 / 1.27 / 2.66	3.68 / 3.96 / 2.71
h4096	2	1	2.07 / – / 1.50	2.05 / – / 1.25	2.44 / – / 1.64	5.60 / – / 3.92	19.46 / – / 13.70
h4096	2	2	2.19 / 0.83 / 2.54	2.06 / 0.89 / 2.04	2.23 / 1.18 / 2.57	3.37 / 2.68 / 2.98	10.90 / 9.15 / 7.29
h4096	4	1	2.11 / – / 1.46	2.05 / – / 1.52	2.36 / – / 1.87	5.51 / – / 4.15	21.05 / – / 14.12
h4096	4	4	3.49 / 0.87 / 2.63	3.25 / 0.88 / 2.37	3.38 / 0.94 / 2.53	3.49 / 2.37 / 2.59	5.80 / 7.99 / 4.09

training fwd+bwd, eager (ms per layer; all_to_all / all_reduce / deepep)

shape	EP	TP	1 tok	256 tok	1024 tok	4096 tok	16384 tok
30b	2	1	3.39 / – / 4.50	3.50 / – / 3.64	3.88 / – / 4.70	6.38 / – / 5.74	22.51 / – / 15.86
30b	2	2	8.08 / 1.77 / 6.72	7.48 / 1.93 / 6.08	8.10 / 2.62 / 6.53	8.23 / 3.81 / 5.47	11.63 / 11.69 / 9.82
30b	4	1	4.82 / – / 4.71	4.94 / – / 4.71	5.13 / – / 4.79	6.58 / – / 5.70	22.13 / – / 16.19
30b	4	4	8.20 / 2.81 / 7.77	7.66 / 2.89 / 7.14	8.24 / 3.19 / 7.83	8.46 / 3.44 / 7.92	8.36 / 10.72 / 7.90
h4096	2	1	5.16 / – / 5.05	5.31 / – / 4.63	6.43 / – / 5.47	13.74 / – / 11.51	45.80 / – / 36.27
h4096	2	2	5.88 / 2.91 / 7.59	5.81 / 3.44 / 6.19	6.04 / 4.38 / 7.92	10.29 / 8.40 / 8.30	25.57 / 26.02 / 19.64
h4096	4	1	4.67 / – / 4.59	4.67 / – / 4.73	5.53 / – / 5.26	12.82 / – / 10.49	48.21 / – / 36.31
h4096	4	4	7.61 / 2.53 / 7.79	7.19 / 2.87 / 7.24	7.72 / 3.08 / 7.60	7.93 / 6.86 / 7.81	13.35 / 22.13 / 10.78

decode, CUDA graph (ms per layer; all_to_all / all_reduce / deepep)

shape	EP	TP	1 tok	256 tok
30b	2	1	n/c / – / 0.17	n/c / – / 0.46
30b	2	2	n/c / 0.12 / 0.17	n/c / 0.32 / 0.41
30b	4	1	n/c / – / 0.17	n/c / – / 0.58
30b	4	4	n/c / 0.12 / 0.17	n/c / 0.27 / 0.35
h4096	2	1	n/c / – / 0.24	n/c / – / 0.93
h4096	2	2	n/c / 0.17 / 0.22	n/c / 0.74 / 0.84
h4096	4	1	n/c / – / 0.24	n/c / – / 1.05
h4096	4	4	n/c / 0.16 / 0.21	n/c / 0.50 / 0.60

Takeaway.

• DeepEP is 1.2-1.5x faster than all-to-all from 4096 tokens per rank, for both prefill and training, at EP=2 and EP=4.
• With a shared batch it also beats the all-reduce at 16384 tokens (2.7 vs 4.0 ms and 4.1 vs 8.0 ms per layer at EP=4).
• Below about 1024 tokens per rank, every eager implementation is host-bound, and the deepep/all_to_all gap is within run-to-run noise (±20%).
• Under CUDA graphs, DeepEP's sync-free path is the only token dispatch that captures: 0.17-0.24 ms per layer at 1 token, against 0.12-0.17 ms for the all-reduce on a shared batch.

Tests

• tests/expert_parallel/test_expert_parallel.py, after merging the updated stack: 12 passed on 2x B200 with the staged kernel. The new tests are:
  • test_is_deepep_loadable: mocked environments, each missing requirement returns False, and with raise_error=True raises its message.
  • test_auto_dispatches_with_deepep_where_it_runs: NCCL and loadable picks deepep, otherwise all_to_all; "deepep" raises the missing requirement.
  • test_deepep_matches_one_process (2 GPUs): forward and input/routing-weight gradients match a single process, on a shared batch and on per-rank batches. It also checks that "auto" trains with DeepEP.
  • test_deepep_captures_and_compiles (2 GPUs): a CUDA graph capture after an eager warm-up replays like eager, and torch.compile(fullgraph=True) matches eager.
  • The GPU tests skip cleanly without 2 GPUs or without a loadable DeepEP.
• Tiny Qwen3-MoE end to end (from_pretrained with distributed_config, bf16, forward and backward on the same inputs), deepep against all_to_all and a single process:
  • Configs (tp, fsdp, ep): 2 GPUs (1,2,2), (2,1,2); 4 GPUs (1,4,4), (2,2,4), (4,1,4), (1,4,2), (2,2,2). Each ran with equal batches and with uneven per-batch lengths (different batch sizes across ranks, uneven TP slices).
  • Forward logits match all_to_all exactly and match a single process (exactly with TP=1; within TP's bf16 noise otherwise).
  • Gradients of all 25 parameters match all_to_all to ≤ 5e-5 relative.
• Before the stack update, the full FSDP model captured in a CUDA graph with deepep and replayed exactly. A model-level torch.compile(fullgraph=True) is blocked by FSDP's hooks, not by this implementation; the experts-level fullgraph compile passes.
• ruff check / ruff format, ty check on touched files, utils/check_copies.py, utils/check_docstrings.py.

🤖 Generated with Claude Code

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Pages changed in this PR:

• en/expert_parallelism

[vasqu]

Are we sure the branch is in sync? Seems like a lot of other changes are included that shouldnt be in this prs diff imo

But I see this PR as less prio overall: 1-5 are core and need merge, this is a nice perf optim to get after

[IlyasMoutawwakil]

yes it was behind, merged the updated stack so the diff is only deepep now