[GH-ISSUE #48] How to apply Moe MLP to recv_x returned by dispatch when num_expert_per_rank > 1 #43

New issue

Closed

opened 2026-10-26 09:53:10 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-26 09:53:10 +00:00

Owner

Originally created by @CyCle1024 on GitHub (Mar 5, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/48

In the case of normal kernels and num_expert_per_rank > 1, for example, token A in rank 0 is dispatched to rank 1 for experts M and N.
A is dispatched only once into recv_x of rank 1 with topk indices. How can I use group gemm based MLP to feed forward token A in recv_x?
If I duplicate A into two token seperately for experts M and N, it's more complicated in combine process. Because combine function would use the origin handle returned by dispatch function, and shape[0] of the new hidden states computed with duplicating A and applying MLP is different from shape[0] of the recv_x returned by dispatch function.

Originally created by @CyCle1024 on GitHub (Mar 5, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/48 In the case of normal kernels and `num_expert_per_rank > 1`, for example, token A in rank 0 is dispatched to rank 1 for experts M and N. A is dispatched only once into recv_x of rank 1 with topk indices. How can I use group gemm based MLP to feed forward token A in recv_x? If I duplicate A into two token seperately for experts M and N, it's more complicated in combine process. Because combine function would use the origin handle returned by dispatch function, and shape[0] of the new hidden states computed with duplicating A and applying MLP is different from shape[0] of the recv_x returned by dispatch function.

bsakaguchi closed this issue

2026-10-26 09:53:10 +00:00

bsakaguchi commented

2026-10-26 09:53:11 +00:00

Author

Owner

@LyricZhao commented on GitHub (Mar 6, 2025):

You can refer to our design the Chrome tracing in https://github.com/deepseek-ai/profile-data.

How can I use group gemm based MLP to feed forward token A in recv_x?

If you use DeepGEMM, you must expand one tokens to several experts.

If I duplicate A into two token seperately for experts M and N, it's more complicated in combine process.

So there is a reduction kernel before the combine all2all.

@LyricZhao commented on GitHub (Mar 6, 2025): You can refer to our design the Chrome tracing in https://github.com/deepseek-ai/profile-data. > How can I use group gemm based MLP to feed forward token A in recv_x? If you use DeepGEMM, you must expand one tokens to several experts. > If I duplicate A into two token seperately for experts M and N, it's more complicated in combine process. So there is a reduction kernel before the combine all2all.

bsakaguchi commented

2026-10-26 09:53:11 +00:00

Author

Owner

@LyricZhao commented on GitHub (Mar 6, 2025):

We follow the current design, because if we expand at the end of dispatch, for some models with more than one experts per GPU selected for a single token, it will cost much more VRAM. This is a tradeoff between VRAM and performance (expand/reduce kernel is not very slow).

@LyricZhao commented on GitHub (Mar 6, 2025): We follow the current design, because if we expand at the end of dispatch, for some models with more than one experts per GPU selected for a single token, it will cost much more VRAM. This is a tradeoff between VRAM and performance (expand/reduce kernel is not very slow).

bsakaguchi referenced this issue

2026-10-26 09:56:30 +00:00

[GH-ISSUE #270] A timeout error occurred when running test_internode.py on four devices. #200