Difference between mom/ and mom_native/

[a1600012888]

Hi authors,

Great work! Really enjoy it. The idea of using varlen for MoE is quite cool!

Just curious the difference between mom/ and mom_native/

Best,
Tianyuan

[weigao266]

We optimized the kernel implementation for MoM before, mom_native is the version before optimization. We recommend to try MoM in the FLA repo, as the latest implementation was already upstreamed to FLA (https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/mom.py).