How to use the output of DeepEP as the input of DeepGemm

[hpz4311]

1. Suppose the number of experts on device 0 is 4, and the number of tokens assigned to each expert is 1, 2, 3, and 7, respectively.
   I think that after DeepEP, the activations received by expert 0 are (13, 7168).

2.In this case, calling m_grouped_gemm_fp8_fp8_bf16_nt_masked requires the activations to be arranged in the form of (4, 8, 7168), with masked_m=[1,2,3,7].

So, if we want to use the output of DeepEP as the input, do we need to manually rearrange the data from form 1 to form 2?
I have only run the DeepGemm demo and am not sure if I have misunderstood the output distribution of DeepEP .

[LyricZhao]

I think that after DeepEP, the activations received by expert 0 are (13, 7168).

Not exactly, it can be less than 13, as there can be some tokens selecting more than 1 expert.

So, if we want to use the output of DeepEP as the input, do we need to manually rearrange the data from form 1 to form 2?
I have only run the DeepGemm demo and am not sure if I have misunderstood the output distribution of DeepEP .

Yes, for outputs of the normal kernels of DeepEP, you have to compute a layout and expand the tokens into the format required by the contiguous kernel in DeepGEMM.

For low-latency kernels' output of DeepEP, it can be directly used by the masked kernel in DeepGEMM (but not useful for training, small batch inference only).

[LyricZhao]

We plan to open-source the small layout-related kernels later (as an optional utility C library), maybe next month. cc @xuzhean

[xutizhou]

Hi @LyricZhao, do you have a timeline for open-sourcing the small layout-related kernels?

[LyricZhao]

Sorry, we are still refactoring and optimizing the small and wgrad kernels (optimization took much more time as expected). Hard to tell a specific time, but at least before June.