[Kernels] Add new kernels to models: CE, Swiglu, Geglu(, Linear)

[vasqu]

As per title, branched out from #46039 to propogate the actual kernels so we can use them on models. Since the dep tree is a bit more complicated and we should sync sooner than later, we move this to a separate PR

Currently to be applied

The core changes are currently in

• Llama
  • Swiglu MLP
  • CE
• Gemma2
  • Geglu MLP
  • CE

Potential other kernel

We could also add kernels around torch native nn.Linear applying tiled forwards/backwards similar to their fused couterparts (e.g. Tiled Swiglu MLP). That would only be enabled during trainin making a simple wrapper/patcher like

@use_kernel_forward_from_hub("Linear")
class Linear(nn.Linear):
    pass

nn.Linear = Linear

Benchmark

Using Llama with this script

[github-actions]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46575&sha=648cc9

[github-actions]

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

run-slow: gemma2, llama

[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.