deepseek-ai/DeepEP
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00
Code Issues 308 Projects Packages Wiki Activity

[GH-ISSUE #595] [Feature Request] Support Expert-Choice Routing via Native Expert-Centric API ([num_experts, capacity]) #366

New issue
Open
opened 2026-10-26 09:59:12 +00:00 by bsakaguchi · 1 comment
bsakaguchi commented 2026-10-26 09:59:12 +00:00
Owner
Copy link

Originally created by @niyunsheng on GitHub (Mar 30, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/595

Background

Expert-Choice routing (Zhou et al., Mixture-of-Experts with Expert Choice Routing, NeurIPS 2022) inverts standard Token-Choice routing: instead of tokens selecting experts, each expert selects its top-C preferred tokens.

This paradigm is increasingly adopted in large-scale models, such as EC-DiT (EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing, ICLR 2025). In configurations with 64 or 128 experts, a popular token can be selected by well over 32 experts.

Current Limitation & Performance Bottleneck

DeepEP currently enforces a Token-centric design where topk_idx expects a shape of [num_tokens, num_topk]. Mapping Expert-Choice outputs back to this Token-centric format introduces two critical issues:

  1. Assertion Failure: A single token selected by >32 experts triggers EP_DEVICE_ASSERT(num_topk <= 32) (e.g., intranode.cu#L408).
  2. Padding Overhead (Structural Inefficiency): To map Expert-Choice to [num_tokens, num_topk], num_topk must equal the maximum number of times any single token is selected. Because token selection distribution is highly skewed in EC routing, the resulting topk_idx tensor will be extremely sparse, filled with padding (e.g., -1). Processing this padded tensor through the current dense kernels wastes registers, shared memory, and memory bandwidth.

Proposed Solutions

Solution 1: Native Expert-Centric API (Preferred)
Introduce a new set of dispatch/combine APIs that natively accept Expert-centric indices.

  • Input Shape: token_indices of shape [num_experts, expert_capacity].

  • Advantages: Directly aligns with the output format of Expert-Choice models, avoiding expensive inverse mapping.

    • Completely eliminates padding overhead and structural sparsity.
    • Bypasses the num_topk <= 32 constraint naturally, as the iteration is over capacity rather than a variable number of experts per token.

  • Expected API Behavior:

    # Expert-Choice directly produces selected token indices
token_indices = torch.randint(0, num_tokens, (num_experts, capacity), dtype=torch.int32, device="cuda")
expert_weights = torch.randn(num_experts, capacity, device="cuda")

# New dispatch interface
buffer.dispatch_expert_choice(x, token_indices=token_indices, weights=expert_weights, ...)

Solution 2: Relax num_topk Limit with Sparse Branching (Alternative / Fallback)
If adding a new API surface is temporarily out of scope, the existing [num_tokens, num_topk] API can be modified to tolerate EC routing.

  • Implementation: Remove or increase the EP_DEVICE_ASSERT(num_topk <= 32). To mitigate the padding overhead described above, the CUDA kernels (intranode.cu / internode.cu) must be updated with dynamic branching to early-exit (break) warp execution when encountering padding indices (-1), minimizing invalid memory accesses.
  • Drawback: This remains suboptimal due to the unavoidable pre-processing mapping overhead and register pressure from large num_topk loops.
Originally created by @niyunsheng on GitHub (Mar 30, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/595 ## Background Expert-Choice routing (Zhou et al., *Mixture-of-Experts with Expert Choice Routing*, NeurIPS 2022) inverts standard Token-Choice routing: instead of tokens selecting experts, each expert selects its top-C preferred tokens. This paradigm is increasingly adopted in large-scale models, such as EC-DiT (*EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing*, ICLR 2025). In configurations with 64 or 128 experts, a popular token can be selected by well over 32 experts. ## Current Limitation & Performance Bottleneck DeepEP currently enforces a Token-centric design where `topk_idx` expects a shape of `[num_tokens, num_topk]`. Mapping Expert-Choice outputs back to this Token-centric format introduces two critical issues: 1. **Assertion Failure:** A single token selected by >32 experts triggers `EP_DEVICE_ASSERT(num_topk <= 32)` (e.g., `intranode.cu#L408`). 2. **Padding Overhead (Structural Inefficiency):** To map Expert-Choice to `[num_tokens, num_topk]`, `num_topk` must equal the maximum number of times *any* single token is selected. Because token selection distribution is highly skewed in EC routing, the resulting `topk_idx` tensor will be extremely sparse, filled with padding (e.g., `-1`). Processing this padded tensor through the current dense kernels wastes registers, shared memory, and memory bandwidth. ## Proposed Solutions **Solution 1: Native Expert-Centric API (Preferred)** Introduce a new set of dispatch/combine APIs that natively accept Expert-centric indices. * **Input Shape:** `token_indices` of shape `[num_experts, expert_capacity]`. * **Advantages:** Directly aligns with the output format of Expert-Choice models, avoiding expensive inverse mapping. * Completely eliminates padding overhead and structural sparsity. * Bypasses the `num_topk <= 32` constraint naturally, as the iteration is over capacity rather than a variable number of experts per token. * **Expected API Behavior:** ```python # Expert-Choice directly produces selected token indices token_indices = torch.randint(0, num_tokens, (num_experts, capacity), dtype=torch.int32, device="cuda") expert_weights = torch.randn(num_experts, capacity, device="cuda") # New dispatch interface buffer.dispatch_expert_choice(x, token_indices=token_indices, weights=expert_weights, ...) ``` **Solution 2: Relax num_topk Limit with Sparse Branching (Alternative / Fallback)** If adding a new API surface is temporarily out of scope, the existing `[num_tokens, num_topk]` API can be modified to tolerate EC routing. * **Implementation**: Remove or increase the `EP_DEVICE_ASSERT(num_topk <= 32)`. To mitigate the padding overhead described above, the CUDA kernels (`intranode.cu` / `internode.cu`) must be updated with dynamic branching to early-exit (break) warp execution when encountering padding indices (-1), minimizing invalid memory accesses. * **Drawback**: This remains suboptimal due to the unavoidable pre-processing mapping overhead and register pressure from large `num_topk` loops.
bsakaguchi commented 2026-10-26 09:59:13 +00:00
Author
Owner
Copy link

@0z5a commented on GitHub (Sep 19, 2026):

I’d like to work on this.
@niyunsheng

For a first PR, I’d prefer to keep the scope around the routing contract and metadata path rather than immediately modifying every transport kernel.

<!-- gh-comment-id:5740755986 --> @0z5a commented on GitHub (Sep 19, 2026): I’d like to work on this. @niyunsheng For a first PR, I’d prefer to keep the scope around the routing contract and metadata path rather than immediately modifying every transport kernel.
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
hybrid-ep
chore/remove-clang-format-yapf
v2.5
nvDev
try-fix-cu128
mori-ep
antgroup-opt
compat-old-torch
add-linter
eager-rdma
v1.2.1
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
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#366
No description provided.