Adding ScatterMoE kernel support for Granite models. 
Author: shawntan
Pull request: yes

# What does this PR do?


Adds ScatterMoE kernel support for Granite MoE models. 
Started in #40365 but has significantly deviated in approach, so starting a new pull request. 


## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@MekkCyber  already started to provide some comments in #40365. 


Comments:
So far I've simply been copying from `torch-ext` to `build/torch-universal` to "build" the kernel. 

Should I follow the example here: https://github.com/huggingface/kernels-community/tree/main/trimul_gpumode

It's another triton based model with `universal` "build". 

Should I be targetting `main` or the latest release?

Some changes have been made to GraniteMoE in the latest `main` that affects how the kernel is used: https://github.com/huggingface/transformers/compare/e6a8e7debed69d4cf5d493158f3f347931c24c9a...7938e91faabb051f3a001cd39c173d4697c2d81c#diff-49f545f0c27ad25a565417d120d925f6bfcde532d6e0ea5539dcec182ad978aaL373

I haven't been able to test out the current `main` branch successfully, there seems to be many breaking changes, so I need some help in modifying the current PR:
1. Should I target the `main` branch or the latest release?
2. Should I be modifying GraniteMoE to the way it is supposed to be?
3. Changing `GraniteMoeMoE` back and regenerating the `modeling_` files of the downstream classes would fix the issue, but what other models with MoEs are affected and how was the change decided?

@MekkCyber will need some help here. Which branch should I target for the PR: `main` or the latest release? 

`main` has issues with the hf API URLs. The current state of the PR targets the latest release.

Hey @shawntan, let's target main! I think you did some faulty rebase, let's only keep the relevant files in the PR. 
I'm not sure what you mean by modifying GraniteMoE but we shouldn't touch the modeling files, the kernel should be applied using the decorator solely.

I've reverted the PR back to the original one against `main`

My issue was this line: 

`main`: https://github.com/huggingface/transformers/blob/main/src/transformers/models/granitemoe/modeling_granitemoe.py#L235

`v4.57.1`: https://github.com/huggingface/transformers/blob/v4.57.1/src/transformers/models/granitemoe/modeling_granitemoe.py#L320

`router_logits` has been removed. I cannot figure out why or how it got removed, but I think it breaks the Granite models downstream of that class. It's also hard for me to test this since pulling from the hub is also broken in the latest `main`.

> router_logits has been removed

Yes it was removed in the MoE refactor for vllm compatibility with transformers but this shouldn't break anything in the transformers implementation

> It's also hard for me to test this since pulling from the hub is also broken in the latest main.

What do you mean ?

> > router_logits has been removed
> 
> Yes it was removed in the MoE refactor for vllm compatibility with transformers but this shouldn't break anything in the transformers implementation

I see. I will need to change the layer definition in the kernel since it will only produce one output, instead of the tuple it returns right now. I am also confused as to how the auxiliary loss will work if the MoE doesn't return the `router_logits`. I will look into it.

Update:
Looking at it, the current changes break the `output_router_logits=True` path. 

The attribute is checked for during the forward call:
https://github.com/huggingface/transformers/blob/ac815417784320cb17ed7d6b077dd19cb6854fa1/src/transformers/models/granitemoe/modeling_granitemoe.py#L725 

But it is never passed to the model:
https://github.com/huggingface/transformers/blob/ac815417784320cb17ed7d6b077dd19cb6854fa1/src/transformers/models/granitemoe/modeling_granitemoe.py#L696-L704

There is also no pathway for passing the router logits, or recomputing them in `GraniteMoeModel`:
https://github.com/huggingface/transformers/blob/ac815417784320cb17ed7d6b077dd19cb6854fa1/src/transformers/models/granitemoe/modeling_granitemoe.py#L492-L551
 
What's the ideal way that the HF team is thinking of allowing for this while still maintaining vllm compatibility? I can make the necessary changes.

> 
> > It's also hard for me to test this since pulling from the hub is also broken in the latest main.
> 
> What do you mean ?

Hmm, sorry, I can't seem to reproduce the same issue I saw before with the jinja files. 

Still, there are some incompatibilities from the Mamba end of things, right now with both repositories on `main`, the error being thrown is this: 
```
  File "/proj/checkpoints/shawntan/transformers/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py", line 47, in <module>
    from mamba_ssm.ops.triton.selective_state_update import selective_state_update
  File "/u/shawntan/mamba/mamba_ssm/__init__.py", line 6, in <module>
    from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
  File "/u/shawntan/mamba/mamba_ssm/models/mixer_seq_simple.py", line 20, in <module>
    from mamba_ssm.utils.generation import GenerationMixin
  File "/u/shawntan/mamba/mamba_ssm/utils/generation.py", line 14, in <module>
    from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer
ImportError: cannot import name 'GreedySearchDecoderOnlyOutput' from 'transformers.generation' (/proj/checkpoints/shawntan/transformers/src/transformers/generation/__init__.py)
```

TL;DR:

1. If I change the kernel to suit the current state of `GraniteMoeMoE`, it would need 1 output, but the problem would be the calculating of the aux loss would be broken.
2. I don't believe that is the intended behaviour, so my question is: How should the `GraniteMoeMoE` class be modified to best support the vLLM compatibility while not breaking previous functionality (aux loss)?

@MekkCyber 

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

run-slow: granitemoe, granitemoehybrid, granitemoeshared

@ArthurZucker that effort is patching for Qwen2MoE, and is working. However there are other issues that prevent OOB working:

There's a PR and an issue ongoing:
 - https://github.com/huggingface/kernels-community/pull/82
 - https://github.com/huggingface/kernels-community/issues/76
 
 The SiLU kernel issue affects more than the Granite models, and could potentially break every model using GeLU via the SiLU activation. Once `.split` is called before SiLU, `use_kernels=True` will almost certainly break.
 
The PRs are useful and work for ScatterMoE, but the SiLU issue prevents `use_kernels=True` from working out of the box. Not sure if that is a requirement for this PR to be merged.



Okay! We merged https://huggingface.co/docs/transformers/main/en/experts_interface which should also make this simpler I think

This helps if the model is not using `use_kernels=True`, but the SiLU kernel issue still remains.

The issue with activation functions should be fixed now: https://huggingface.co/kernels-community/activation/blob/main/build/torch210-cxx11-cu126-x86_64-linux/layers.py#L20