add scattermoe kernel for fast MoE training

[mayank31398]

No description provided.

[Rocketknight1]

cc @ArthurZucker

[ArthurZucker]

Hello @mayank31398 ! Nice pr, happy to add something like that, do you mind using kernels like what we do for GPT_OSS?!
This way we keep a slow path, compatible with all torch, all hardwares etc and don't have code changes for the core modeling, and just have the kernel on the hub!

WDYT? 🤗

[mayank31398]

@ArthurZucker scattermoe doesnt support bias for now, I will add this soon!
meanwhile supporting every model is hard since some models have expert weights as a moduleList instead of a 3D tensor :/

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: granitemoe

[shawntan]

Hello @mayank31398 ! Nice pr, happy to add something like that, do you mind using kernels like what we do for GPT_OSS?! This way we keep a slow path, compatible with all torch, all hardwares etc and don't have code changes for the core modeling, and just have the kernel on the hub!

WDYT? 🤗

Is there an existing triton kernel you could point to that I could follow?