Support Gated MLP

[Insideyyy]

Hello! FlashDMoE is a very great piece of work!
As Gated MLP is widely used as FFN in many LLMs (like DeepSeekV3, Qwen3, Llama), is there a plan to support it?

Gated MLP based on swiGLU:

$$ \begin{split}Swish(x) = x \otimes Sigmoid(x) \\ swiGLU(src,W_1,W_2) = (src \cdot W_1) \otimes Swish(src \cdot W_2) \\ FFN(src,W_1,W_2,V) = swiGLU(src,W_1,W_2) \cdot V\end{split} $$

[osayamenja]

@Insideyyy we welcome contributions :) Feel free to work on this and commit a PR. The current implementation should be a sufficiently good enough starter to implement this. Some changes (#5 ) are incoming to improve CTA GEMM performance, so keep that in mind if you seek to work on this.

I'm oversubscribed currently but I'll keep this in the back of my mind, I think there some nice optimizations to do here.

1. To do tile-based fusion, $src\cdot W_1$ and $src\cdot W_2$ would have to be computed by the same CTA.
2. $Sigmoid(...)$ is elementwise, so it can be applied directly on the register output of $src\cdot W_2$.
3. Key concern here is maintaining the output of $src \cdot W_1)$ and $src \cdot W_2$ in fast memory (preferably registers).

[Insideyyy]

@osayamenja I'd be glad to work on this :)
It may require some time, though.

[zhangshaofeng-zsf]

@osayamenja I'd be glad to work on this :) It may require some time, though.

I am looking forward to your work results, which will be very useful to me

[zhangshaofeng-zsf]

@osayamenja Will the v0.1.0 version support swiGLU?

[osayamenja]

@zhangshaofeng-zsf yes, we've implemented it already, just testing now.