Llama Kernel integration

[MekkCyber]

What does this PR do?

PoC of Integrating kernels into the llama model. The kernels used are for attn, mlp, and rmsnorm

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[ArthurZucker]

Love the updates for rms and mlp!
Attention needs a bit more tinkering !
🤗 happy to hear your opinion as well!

[ArthurZucker]

to remove!
Using kernels for that layer should IMO still be done via the attn_implementation!
it should give a bit more freedom to use kernels for the rest, and use sdpa or flex or paged or whatever here!

Attn is something that can change very easily, which is why I thought ALL_ATTENTION_FUNCTIONS was a better entry point!

[MekkCyber]

Do you mean we pass something like attn_implementation="use_kenrel" in from_pretrained and then add a decorator with the name of the layer to use in the attention class ?

[ArthurZucker]

Yes something along those lines. Or you can just say attn_implementation="kernels-community/paged-attention" something like this. (with extra (kernel_config = "...")
I don't know the answer but the simpler the better!

[MekkCyber]

I changed the implementation so that we need to pass attn_implementation="use_kernel" and a kernel_config following the structure :
kernel_config={ "repo_id": "medmekk/triton-flash-attn", "kernel_name": "attention" }

[ArthurZucker]

this is super nice!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Missing some docs a but marvelous

[ArthurZucker]

I would allso support for example passing a repo _id directly no? We can get the device ourselves, this way its easy to just use 1 arg for simplicity!

[ArthurZucker]

We can create:

            "cuda": LayerRepository(
                repo_id="medmekk/triton-llama-mlp",
                layer_name="TritonLlamaMLP",
            )

base on the cls name for example

[ArthurZucker]

we can have a func wrapper in integration / hub kernels or something

[ArthurZucker]

And if repo is not supper we juste use eager + warning!

[ArthurZucker]

Flag to disable is coming1

[efsotr]

I believe the implementation should be compatible with torch.compile.
For most simple modules, such as RMSNorm and operations like silu_and_mul in MLP, torch.compile already provides sufficient optimization. Therefore, for these cases, we should avoid overriding the forward method when torch.compile is enabled.

Suggestions:
In the kernels-community kernels, consider leveraging torch.compile for optimization, for example by using PyTorch ops or Triton implementations.
For custom kernels that are neither PyTorch nor Triton based, compatibility with torch.compile should also be taken into account.
For details on supporting custom kernels, refer to the official documentation: https://pytorch.org/docs/stable/library.html

[MekkCyber]

Yeah that's right, once we have a stable api, we can think about how to make it compatible with torch.compile, but for now using kernels can be deactivated, by uninstalling the kernels library or just setting the DISABLE_KERNEL_MAPPING to 1 for example.
I'm also doing some benchmarks to see how much we gain or lose from using kernels instead of torch.compile, that would lay a good basis to see what we can do to improve the performance