make ScatterMoe more easily patchable for HF MoE modeling

[winglian]

This PR makes it easier to simply use something like the snippet below for a drop in replacement for all the MoEs based on qwen2 in transformers v5. The original ScatterMoEGatedMLP isn't quite usable since it relies on router instead of gate and input_linear, output_linear instead of gate_up_proj, down_proj in the v5 modeling.

register_kernel_mapping({
    "HFScatterMoEParallelExperts": {
        "cuda": {
            Mode.TRAINING: LayerRepository(
                repo_id="axolotl-ai-co/scattermoe",
                layer_name="HFScatterMoEGatedMLP",
            ),
            Mode.INFERENCE: LayerRepository(
                repo_id="axolotl-ai-co/scattermoe",
                layer_name="HFScatterMoEGatedMLP",
            ),
        },
    }
})
replace_kernel_forward_from_hub(Qwen2MoeSparseMoeBlock, "HFScatterMoEParallelExperts")

[MekkCyber]

cc @shawntan

[danieldk]

Hi, @shawntan, any chance you could review this PR?

[winglian]

Here's a quick SFT test of olmoe.

[shawntan]

Yeah this looks good and is a good idea, but have you tested it end-to-end with use_kernels=True?

I have an issue with another community kernel: #76

[winglian]

yes, this was tested with use_kernels=True

[shawntan]

Okay. I'm trying to test GraniteMoEHybrid with use_kernels=True to make sure these kernels work, but it seems the SiLU kernel does not work with non-contiguous tensors, and it seems both GraniteMoEHybrid and Qwen2MoE will be affected by this. See this comment: #76 (comment)

The decision seems to be to assert hidden_states.is_contiguous(). Which breaks both models, as far as I understand.

UPDATE: confirmed that it breaks Qwen2MoE as well.

@winglian Not fully sure how the error didn't show up when testing with use_kernels=True, but I'm interested to know how you get around it.

@MekkCyber FYI.