[PR #744] [MERGED] [HybridEP] Support unequal per-rank token counts (ragged dispatch) #749

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/744
Author: @abcdabcd987
Created: 9/1/2026
Status: ✅ Merged
Merged: 9/7/2026
Merged by: @jershi425

Base: hybrid-ep ← Head: ragged-tokens

📝 Commits (5)

efa9196 [HybridEP] Support unequal per-rank token counts (ragged dispatch)
0aa8753 Address review: inline handle pack/unpack, buffer-generation guard
2364852 Review: make buffer generation counter uint64_t
268ca93 Review: keep the ragged matrix end-padded only; dedicated dropped-token test
7c646d6 Review: remove dropped-token simulation from the ragged tests

📊 Changes

11 files changed (+845 additions, -52 deletions)

View changed files

📝 csrc/hybrid_ep/backend/hybrid_ep_backend.cuh (+45 -5)
📝 csrc/hybrid_ep/config.cuh (+10 -0)
📝 csrc/hybrid_ep/executor/executor.cu (+12 -2)
📝 csrc/hybrid_ep/executor/executor.cuh (+11 -0)
📝 csrc/hybrid_ep/hybrid_ep.cu (+34 -9)
📝 csrc/hybrid_ep/hybrid_ep.cuh (+13 -6)
📝 csrc/hybrid_ep/pybind_hybrid_ep.cu (+4 -0)
📝 deep_ep/hybrid_ep_buffer.py (+131 -9)
📝 docs/Hybrid-EP_Implementation.md (+23 -20)
📝 tests/test_hybrid_ep.py (+1 -1)
➕ tests/test_ragged_hybrid_ep.py (+561 -0)

📄 Description

Summary

HybridEP required every EP rank to dispatch the same number of tokens, forcing frameworks with variable-length (no-pad) batches to pad tokens up to the group max — pad rows that ride the dispatch payload and the expert GEMMs. A line-level audit showed the requirement is not fundamental: every payload path is already per-token routing-map-gated; only the routing-map allgather size, the metadata map strides, and the chunk-grid/flag arithmetic need a shared layout constant, not equal real counts.

This PR lifts the restriction:

New optional num_of_tokens_per_rank kwarg on dispatch / dispatch_with_permute: a group-uniform token-slot count S (identical on every rank, e.g. the group max of the local counts; rounded up to a multiple of 16 internally, which satisfies the TMA remainder-chunk constraint for every ranks_per_node). Each rank's hidden / probs / routing inputs keep their real N_i <= S rows. Omitted ⇒ defaults to hidden.size(0), i.e. the previous uniform behavior.
Only the routing tensor is padded internally to S rows, with routed-nowhere rows (all-False bool rows in sparse mode, all--1 sentinel rows in dense int16 top-k mode). Pad slots produce no NVLink/RDMA traffic, no permuted rows, and no expert GEMM rows — the only cost is ~(S-N_i) × K × 2 bytes of extra routing-map allgather payload (≈0.03% of dispatch traffic at typical shapes) and zero-length chunk flags.
The handle carries the real local count in a new trailing num_of_valid_tokens field (appended, so existing indices — including handle[10] overflow_flag — are unchanged; unpackers accept the old arity). Combine outputs are allocated at S (the combine kernel unconditionally writes every slot; un-routed rows come out as exact zeros) and returned narrowed to N_i rows.
A zero-token rank and a group-wide-empty step both work (slot count clamped to ≥16).

Bug fixes included

NIXL dispatch flag stride used the runtime chunk count as the per-sender stride while receivers poll at the template-MAX stride — wrong slot whenever ceil(N/chunk) != ceil(max/chunk) on 3+ nodes (pre-existing; now MAX-based, matching DOCA).
NIXL signal skipped for zero-token chunks while the dispatch receiver waits unconditionally per chunk — deadlock (pre-existing; now signals unconditionally, matching DOCA), plus a receiver-side end-of-round drain loop in combine so unobserved in-flight atomics can't race the residue fill when the chunk grid shrinks across calls (NIXL has no analog of DOCA's rdma_sync same-QP ordering).
NIXL prob restripe read num_of_tokens_per_rank rows of the user probs tensor — OOB when slots exceed real rows; now bounded by the tensor's real row count.
Replaying a cached handle after buffer growth (which reallocates buffers, resets flags, and re-registers RDMA strides) previously corrupted silently or deadlocked — now raises with a clear message in all four cached-handle paths.
dispatch() now asserts routing rows == hidden rows on the fresh path (mirroring dispatch_with_permute), so a mismatch can't be masked by the internal padding as silent token drops.

Note: the NIXL changes are code-reviewed but compile-untested here (no NIXL build environment); DOCA and single-domain MNNVL paths are unaffected by them.

API compatibility

Uniform-count callers with N % 16 == 0 (all existing tests / typical production configs) take a bit-identical path.
Uniform-count callers with N % 16 != 0 now get internal slot alignment + combine narrowing instead of a potential TMA failure; output shapes are unchanged (N rows).
Handle tuples grew 7→8 and 11→12 by appending num_of_valid_tokens. Positional indices of existing fields are unchanged; fixed-arity structural unpacks of the returned handle need updating (one site in tests/test_hybrid_ep.py was).

Validation

4×GB200 single node (MNNVL/IPC): existing tests/test_hybrid_ep.py (BF16+FP8, all six correctness modes) and tests/test_graphed_hybrid_ep.py (non-blocking + CUDA graphs) pass unchanged; new tests/test_ragged_hybrid_ep.py passes — bitwise vs a padded-reference TorchRef: sparse/dense routing × fused/non-fused permute × with/without probs × BF16/FP8, cached-handle replay bitwise (tokens, probs, scaling factors), dropped (zero-expert) tokens, zero-token ranks (incl. dense mode and FP8), and the N % 16 != 0 default path.
GB200 NVL72 cluster: same suites pass on 1 node × 4 GPUs and on 4 nodes × 16 GPUs inside a single rack segment (cross-OS-node MNNVL fabric, USE_MNNVL=1).
Adversarially reviewed (multi-agent, per-finding refutation): 11 confirmed findings, all addressed in this PR.

🤖 Generated with Claude Code

_{🔄 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/744 **Author:** [@abcdabcd987](https://github.com/abcdabcd987) **Created:** 9/1/2026 **Status:** ✅ Merged **Merged:** 9/7/2026 **Merged by:** [@jershi425](https://github.com/jershi425) **Base:** `hybrid-ep` ← **Head:** `ragged-tokens` --- ### 📝 Commits (5) - [`efa9196`](https://github.com/deepseek-ai/DeepEP/commit/efa9196f68769310f5a662d65e3a353b96d90209) [HybridEP] Support unequal per-rank token counts (ragged dispatch) - [`0aa8753`](https://github.com/deepseek-ai/DeepEP/commit/0aa8753abed27c41e6fe77db49e3c83f9c2df628) Address review: inline handle pack/unpack, buffer-generation guard - [`2364852`](https://github.com/deepseek-ai/DeepEP/commit/236485273d670df4a6d8d9a55528076f7a2fbead) Review: make buffer generation counter uint64_t - [`268ca93`](https://github.com/deepseek-ai/DeepEP/commit/268ca934a575b9b967af944b81feaf148489f343) Review: keep the ragged matrix end-padded only; dedicated dropped-token test - [`7c646d6`](https://github.com/deepseek-ai/DeepEP/commit/7c646d6fc2760e96ecedf47e46ae75caa827e21a) Review: remove dropped-token simulation from the ragged tests ### 📊 Changes **11 files changed** (+845 additions, -52 deletions) <details> <summary>View changed files</summary> 📝 `csrc/hybrid_ep/backend/hybrid_ep_backend.cuh` (+45 -5) 📝 `csrc/hybrid_ep/config.cuh` (+10 -0) 📝 `csrc/hybrid_ep/executor/executor.cu` (+12 -2) 📝 `csrc/hybrid_ep/executor/executor.cuh` (+11 -0) 📝 `csrc/hybrid_ep/hybrid_ep.cu` (+34 -9) 📝 `csrc/hybrid_ep/hybrid_ep.cuh` (+13 -6) 📝 `csrc/hybrid_ep/pybind_hybrid_ep.cu` (+4 -0) 📝 `deep_ep/hybrid_ep_buffer.py` (+131 -9) 📝 `docs/Hybrid-EP_Implementation.md` (+23 -20) 📝 `tests/test_hybrid_ep.py` (+1 -1) ➕ `tests/test_ragged_hybrid_ep.py` (+561 -0) </details> ### 📄 Description ## Summary HybridEP required every EP rank to dispatch the same number of tokens, forcing frameworks with variable-length (no-pad) batches to pad tokens up to the group max — pad rows that ride the dispatch payload and the expert GEMMs. A line-level audit showed the requirement is not fundamental: every payload path is already per-token routing-map-gated; only the routing-map allgather size, the metadata map strides, and the chunk-grid/flag arithmetic need a *shared layout constant*, not equal real counts. This PR lifts the restriction: - New optional `num_of_tokens_per_rank` kwarg on `dispatch` / `dispatch_with_permute`: a group-uniform token-slot count `S` (identical on every rank, e.g. the group max of the local counts; rounded up to a multiple of 16 internally, which satisfies the TMA remainder-chunk constraint for every `ranks_per_node`). Each rank's `hidden` / `probs` / routing inputs keep their real `N_i <= S` rows. Omitted ⇒ defaults to `hidden.size(0)`, i.e. the previous uniform behavior. - Only the **routing tensor** is padded internally to `S` rows, with routed-nowhere rows (all-`False` bool rows in sparse mode, all-`-1` sentinel rows in dense int16 top-k mode). Pad slots produce no NVLink/RDMA traffic, no permuted rows, and no expert GEMM rows — the only cost is ~`(S-N_i) × K × 2` bytes of extra routing-map allgather payload (≈0.03% of dispatch traffic at typical shapes) and zero-length chunk flags. - The handle carries the real local count in a new trailing `num_of_valid_tokens` field (appended, so existing indices — including `handle[10]` overflow_flag — are unchanged; unpackers accept the old arity). Combine outputs are allocated at `S` (the combine kernel unconditionally writes every slot; un-routed rows come out as exact zeros) and returned narrowed to `N_i` rows. - A zero-token rank and a group-wide-empty step both work (slot count clamped to ≥16). ## Bug fixes included - **NIXL dispatch flag stride** used the runtime chunk count as the per-sender stride while receivers poll at the template-MAX stride — wrong slot whenever `ceil(N/chunk) != ceil(max/chunk)` on 3+ nodes (pre-existing; now MAX-based, matching DOCA). - **NIXL signal skipped for zero-token chunks** while the dispatch receiver waits unconditionally per chunk — deadlock (pre-existing; now signals unconditionally, matching DOCA), plus a receiver-side end-of-round **drain loop** in combine so unobserved in-flight atomics can't race the residue fill when the chunk grid shrinks across calls (NIXL has no analog of DOCA's `rdma_sync` same-QP ordering). - **NIXL prob restripe** read `num_of_tokens_per_rank` rows of the user probs tensor — OOB when slots exceed real rows; now bounded by the tensor's real row count. - **Replaying a cached handle after buffer growth** (which reallocates buffers, resets flags, and re-registers RDMA strides) previously corrupted silently or deadlocked — now raises with a clear message in all four cached-handle paths. - `dispatch()` now asserts routing rows == hidden rows on the fresh path (mirroring `dispatch_with_permute`), so a mismatch can't be masked by the internal padding as silent token drops. Note: the NIXL changes are code-reviewed but compile-untested here (no NIXL build environment); DOCA and single-domain MNNVL paths are unaffected by them. ## API compatibility - Uniform-count callers with `N % 16 == 0` (all existing tests / typical production configs) take a bit-identical path. - Uniform-count callers with `N % 16 != 0` now get internal slot alignment + combine narrowing instead of a potential TMA failure; output shapes are unchanged (`N` rows). - Handle tuples grew `7→8` and `11→12` by appending `num_of_valid_tokens`. Positional indices of existing fields are unchanged; fixed-arity structural unpacks of the *returned* handle need updating (one site in `tests/test_hybrid_ep.py` was). ## Validation - **4×GB200 single node (MNNVL/IPC)**: existing `tests/test_hybrid_ep.py` (BF16+FP8, all six correctness modes) and `tests/test_graphed_hybrid_ep.py` (non-blocking + CUDA graphs) pass unchanged; new `tests/test_ragged_hybrid_ep.py` passes — bitwise vs a padded-reference `TorchRef`: sparse/dense routing × fused/non-fused permute × with/without probs × BF16/FP8, cached-handle replay bitwise (tokens, probs, scaling factors), dropped (zero-expert) tokens, zero-token ranks (incl. dense mode and FP8), and the `N % 16 != 0` default path. - **GB200 NVL72 cluster**: same suites pass on 1 node × 4 GPUs and on 4 nodes × 16 GPUs inside a single rack segment (cross-OS-node MNNVL fabric, `USE_MNNVL=1`). - Adversarially reviewed (multi-agent, per-finding refutation): 11 confirmed findings, all addressed in this PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <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:51 +00:00