:rotating_light: [`Kernels`] Sync to latest version
Author: vasqu
Pull request: yes

## :rotating_light: Breaking change
Breaking by design as kernels no longer accepts model type repos (e.g. see my newly builts kernels for the tests for exactly that reason). This is more of a breaking change on kernels than on us but we cannot avoid this!

## General initial overview
- Torch compile issues
    - RMS Norm does not work whatsoever, producing wrong outputs
    - Causal loss can encounter weird IMA and cublas errors; since it is a fused op in itself, intentional disable (graph break) added to not lose the benefits of the chunked implementation
- Upstream fixes :white_check_mark: 
    - Sync repos
      - https://github.com/huggingface/kernels-community/pull/948
      - https://github.com/huggingface/kernels-community/pull/907
    - Enable functions with torch compile
      - https://github.com/huggingface/kernels/pull/607
- Future PR
  - Actually use the new kernels and propogate through the code base
      - See #46575

## Additional targeted issues fixed with this PR
- Remove unnecessary logic needed for previous versions (try catch logics especially)
- Fixes security vulnerabilities
    - `register_kernel_replacements_and_fusions` disregarded the context manager bypassing any safety mechanisms on lower kernel versions
        - Essentially allows any kernel to be loaded via kernel config 
    - We now mostly relay a lot of this logic to the original kernels package as well
-  Allows the inclusion of metadata (version/revision/trust_remote_code)
    - For our existing kernels, this acts as BC as we fallback to version 1 which is equivalent to what we currently use
    - For new kernels, it allows user more control --> version, remote code etc
        - This is especially handled within the kernel config as a new entry style
        - Prev: Simple str of the layer name with the attached class layer `dummy-kernel:DummyLayer`
        - Now: Optional tuple structure to pass the necessary data `(dummy-kernel:DummyLayer, {"version": 0, "trust_remote_code": True})`   
- Fixed various side effects during testing
    - A lot of patching was not working as expected --> isolated tests passed but in succession failed
- Fix kernelize pattern not inheriting the default mapping

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46039). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Yep, will do :saluting_face: thinking about memory benchmarks mostly since this is the key behind the new loss kernel from liger

Hi @vasqu, thanks a lot for taking this on! 🙏

We're currently hitting the incompatibility between the latest `transformers` and the latest `kernels`, and this PR looks like exactly the fix we need, so we've been following it closely.

I completely understand it's still a Draft and depends on a few upstream pieces landing first (the `kernels-community` sync PRs #906 / #907 / #928, `kernels#607`, and the follow-up version bump). I just wanted to gently check in on the current status — is there anything still blocking it, and is there an estimated path to getting it out of Draft?

More importantly, **is there anything I can help with to move this forward?** A few things I'd be happy to take on:

- Run and provide **benchmarks** (including the memory benchmarks for the new liger loss kernel) that @ArthurZucker requested for the docs — note that I can only test on **XPU**.
- **Test the compatibility** of this branch against `kernels>=0.15.0` on my setup and report back any issues.
- **Collaborate on** the upstream `kernels-community` and `kernels` PRs that this one depends on.
- Validate the remaining TODOs (e.g. propagating the decorator to other models / models that don't want the loss overwrite) on real models.

Just let me know what would be most useful and I'll jump on it. Thanks again for the great work! 🚀

Hey @jiqing-feng :wave: thanks for checking in!

I was out for the past 2 days mostly so it got a bit delayed. The main concern is having a proper liger kernels version that has the proper flags. I.e. it needs some thorough testing to match the expected losses.

I think, I will split this PR a bit to ease the merge of this PR
1. Only make syncing possible with existing kernels
2. Move to new liger kernels to be used (swiglu etc)
    - I feel like we should add another feature here to make sure the exchange can only happen on certain conditions, e.g. no bias on linear projections (mixture of kernels and transformers side)
    - Verify the losses across kernels and compile vs eager (no compile)
3. Then we can enable and propogate this to all the models

What I need the most help is verifying this across different stacks (torch versions, compile, etc). The development moved fast. 
 

Note that the kernels need to manually patched based on the PRs I opened on the kernels community. They are not uploaded atm


Hey @vasqu 👋 I ran a cross-stack verification on **Intel XPU (Data Center GPU Max 1550)** against this branch. TL;DR: everything works on XPU, including the `torch.compile` + kernels combo. 🎉

## Environment

- `torch 2.12.0+xpu` (release)
- `kernels 0.15.2` (matches the PR's `>=0.15.2,<0.16` pin)
- `transformers` @ this PR branch (`5.8.0.dev0`)
- device: `xpu`, dtype: `bfloat16`

## Coverage note

The new `SwiGLUMLP` (`LigerSwiGLUMLP`) and `ForCausalLMLoss` (`LigerForCausalLMLoss`) kernels are **`cuda`-only in the mapping**, so on XPU they fall back to the eager / Python implementations — I could not exercise the liger kernels themselves here. What this run actually validates on XPU is:
- the **framework-side changes** (fused-loss fallback when `logits=None`, the `logits is None` skip, mode switching, `torch.compile`), and
- the existing **`RMSNorm` + rotary XPU kernels** (`kernels-community/rmsnorm`, `kernels-community/rotary`), which load and run fine.

## Results

| Scenario | loss | diff vs eager |
|---|---|---|
| eager reference (`use_kernels=False`) | 7.75389 | — |
| `ForCausalLMLoss`: `hidden_states` path vs `logits` path | 7.75389 | **0.00e+00** |
| `use_kernels=True`, inference | 7.75361 | 2.79e-04 |
| `use_kernels=True`, training (`logits is None` ✓) | 7.75362 | 2.65e-04 |
| `torch.compile`, no kernels | 7.75394 | 5.44e-05 |
| **`torch.compile` + `use_kernels=True`** | 7.75368 | 2.09e-04 |

## Notes

- The new fused-loss fallback (`logits=None` → compute from `hidden_states` + `lm_head_weight`) is **numerically identical** to the old logits path (diff = 0). Note this exercises the transformers Python fallback, not the liger CUDA kernel.
- The RMSNorm + rotary XPU kernels load and run fine, and the **`torch.compile` + kernels** combo works without the IMA/cublas issues you saw on CUDA — losses stay within bf16 noise (~1e-4). The fused-loss path correctly skips logit materialization (`logits is None`) in both train and inference.
- I couldn't exercise the `cuda`-only `LigerForCausalLMLoss` / `LigerSwiGLUMLP` or the memory benchmark on XPU. Happy to help validate those (and re-run on CUDA) once the patched liger kernels are uploaded — or to add XPU entries to the mapping if you'd like the liger kernels covered there too.

## Repro script

```python
"""Verification for transformers PR #46039 (SwiGLU + fused ForCausalLMLoss kernels) on Intel XPU.

Run inside an XPU container:
    HF_HOME=/path/to/cache python verify_pr46039.py
"""

import torch
import transformers
from transformers import AutoConfig, AutoModelForCausalLM
from transformers.loss.loss_utils import ForCausalLMLoss

DEVICE = "xpu" if torch.xpu.is_available() else "cpu"
DTYPE = torch.bfloat16


def build():
    cfg = AutoConfig.for_model(
        "llama", vocab_size=2048, hidden_size=512, intermediate_size=1376,
        num_hidden_layers=4, num_attention_heads=8, num_key_value_heads=4,
        max_position_embeddings=512, tie_word_embeddings=False,
    )
    torch.manual_seed(0)
    return AutoModelForCausalLM.from_config(cfg, dtype=DTYPE).to(DEVICE).eval()


def make_kernel_model(ref, train=False):
    m = build()
    m.load_state_dict(ref.state_dict())
    m.train() if train else m.eval()
    m.set_use_kernels(True)
    return m


def main():
    import kernels
    print(f"transformers {transformers.__version__} | torch {torch.__version__} | "
          f"kernels {kernels.__version__} | device {DEVICE} | dtype {DTYPE}")

    torch.manual_seed(1)
    input_ids = torch.randint(0, 2048, (2, 64), device=DEVICE)
    labels = input_ids.clone()
    ref = build()
    results = []

    # 0) eager reference
    with torch.no_grad():
        out = ref(input_ids=input_ids, labels=labels)
    ref_loss = out.loss.item()
    results.append(("eager reference (use_kernels=False)", ref_loss, None))

    # 1) new fused-loss fallback: logits=None -> compute from hidden_states + lm_head_weight
    with torch.no_grad():
        hidden = ref.model(input_ids=input_ids).last_hidden_state
    common = dict(labels=labels, vocab_size=ref.config.vocab_size)
    l_logits = ForCausalLMLoss(logits=ref.lm_head(hidden), **common).item()
    l_hidden = ForCausalLMLoss(logits=None, hidden_states=hidden,
                               lm_head_weight=ref.lm_head.weight, **common).item()
    results.append(("ForCausalLMLoss: hidden_states path vs logits path", l_hidden, abs(l_hidden - l_logits)))

    # 2) kernels, inference
    with torch.no_grad():
        out = make_kernel_model(ref)(input_ids=input_ids, labels=labels)
    results.append(("use_kernels=True, inference", out.loss.item(), abs(out.loss.item() - ref_loss)))

    # 3) kernels, training (fused loss path, logits should be None)
    out = make_kernel_model(ref, train=True)(input_ids=input_ids, labels=labels)
    tag = "use_kernels=True, training" + ("  [logits is None ok]" if out.logits is None else "  [logits NOT None FAIL]")
    results.append((tag, out.loss.item(), abs(out.loss.item() - ref_loss)))

    # 4) torch.compile, no kernels
    mc = build(); mc.load_state_dict(ref.state_dict()); mc.eval()
    with torch.no_grad():
        out = torch.compile(mc)(input_ids=input_ids, labels=labels)
    results.append(("torch.compile, no kernels", out.loss.item(), abs(out.loss.item() - ref_loss)))

    # 5) torch.compile + kernels
    with torch.no_grad():
        out = torch.compile(make_kernel_model(ref))(input_ids=input_ids, labels=labels)
    results.append(("torch.compile + use_kernels=True", out.loss.item(), abs(out.loss.item() - ref_loss)))

    print(f"\n{'scenario':<52}{'loss':>12}{'diff vs eager':>16}")
    print("-" * 80)
    for name, loss, diff in results:
        diff_s = "-" if diff is None else f"{diff:.2e}"
        print(f"{name:<52}{loss:>12.5f}{diff_s:>16}")


if __name__ == "__main__":
    main()
```


@jiqing-feng I think they could potentially work on XPU as well, just need to adjust the mapping to add them to XPU as well. I expect this to work OOB as it's triton and not somethig device specific. But that's already nice news!! :partying_face: 

I also wanted to check the tiled swiglu version, especially across different seq lens.

Hey @jiqing-feng I updated the PR description, please see the newest infos I added there :hugs: 

I decided to split this a bit
1. Sync for latest kernels
2. Propogate to actual usage on models

This should fasten the sync which is the most important before we diverge too much. I guess what remains are potential other calls via `get_kernel` or anything that currently fails. I expect that I need to recheck FA

XPU verification (liger kernels)

Hey @vasqu 👋 — followed your suggestion and wired the liger `SwiGLUMLP` + `ForCausalLMLoss`
into the **XPU** mapping, then ran a cross-stack verification on **Intel Data Center GPU Max 1550**.
TL;DR: the liger Triton kernels **do run on XPU** and match eager in eager/inference/training.
The only thing that fails is `torch.compile` **with the liger kernels enabled** — the Intel Triton
backend hits `PassManager::run failed`. I narrowed this down: it fails even for SwiGLU-only
inference (no labels, so no fused CE in the graph), so it's a compile-path / Intel-Triton-backend
issue, not specific to the fused loss.

## Environment

- `torch 2.12.0+xpu`
- `kernels 0.15.2` (matches the PR pin `>=0.15.2,<0.16`)
- `transformers` @ this PR branch (`5.8.0.dev0`)
- device `xpu`, dtype `bfloat16`

## What I changed

**1. Mapping** (`integrations/hub_kernels.py`): mirrored the existing `cuda` entries to `xpu`
for both `SwiGLUMLP` (`LigerSwiGLUMLP`) and the `ForCausalLMLoss` func
(`LigerForCausalLMLoss`, `Mode.TRAINING | Mode.TORCH_COMPILE`):

```diff
@@ SwiGLUMLP mapping @@
             layer_name="LigerSwiGLUMLP",
             version=1,
         ),
+        "xpu": LayerRepository(
+            repo_id="kernels-community/liger-kernels",
+            layer_name="LigerSwiGLUMLP",
+            version=1,
+        ),
@@ ForCausalLMLoss func mapping @@
                 repo_id="kernels-community/liger-kernels", func_name="LigerForCausalLMLoss", version=1
             ),
         },
+        "xpu": {
+            Mode.TRAINING | Mode.TORCH_COMPILE: FuncRepository(
+                repo_id="kernels-community/liger-kernels", func_name="LigerForCausalLMLoss", version=1
+            ),
+        },
```

**2. Kernel patches** (local, against the `kernels-community/liger-kernels` `torch-xpu` build) —
these mirror your upstream liger sync in
https://github.com/huggingface/kernels-community/pull/948 (open; supersedes the now-closed
https://github.com/huggingface/kernels-community/pull/906) and rely on the flag support added in
https://github.com/huggingface/kernels/pull/607 (merged). Since the huggingface/kernels-community#948
build isn't uploaded to the Hub yet, I applied the equivalent changes directly to the cached build
so the prebuilt kernels pass validation. Once huggingface/kernels-community#948 lands and the build
is re-uploaded, the mapping above should work as-is:
- `LigerSwiGLUMLP`: made stateless (dropped the `__init__`/members) so it passes
  `kernels._validate_layer` (must not override `nn.Module.__init__`).
- `LigerForCausalLMLoss`: aligned the signature to the transformers `ForCausalLMLoss` reference
  (8 args + `**kwargs`) so the func validation (`different number of arguments`) passes.
- Added `can_torch_compile = True` / `has_backward = True` flags to the loss func (the feature
  enabled by huggingface/kernels#607).

## Results

| Scenario | loss / val | note |
|---|---|---|
| eager reference (no kernels) | 7.75389 | — |
| `ForCausalLMLoss` hidden_states path vs logits path | 7.75389 | diff **0.00e+00** |
| `use_kernels` inference (liger SwiGLU) | 7.75361 | liger SwiGLU engaged ✓ |
| `use_kernels` training (liger SwiGLU + fused CE) | 7.75362 | `logits is None` ✓, liger engaged ✓ |
| `torch.compile`, no kernels | 7.75394 | ✓ |
| `torch.compile` + `use_kernels` | **FAIL** | Intel Triton `PassManager::run failed` |
| liger SwiGLU seq_len=16 (max\|Δlogits\| vs eager) | 0.02734 | engaged ✓ |
| liger SwiGLU seq_len=128 | 0.02734 | engaged ✓ |
| liger SwiGLU seq_len=512 | 0.03906 | engaged ✓ |

## Notes

- **liger SwiGLU works on XPU** out of the box (Triton) once the validation patch is applied —
  it engages across all seq lengths I tried (16/128/512) and stays within bf16 noise vs eager.
  Confirmed via `m.model.layers[0].mlp.forward.__module__` pointing at the liger hub module.
- **Fused liger CE loss works in eager training** (`logits is None`, computed from
  `hidden_states` + `lm_head_weight`) and matches eager loss within bf16 noise (~3e-4).
- **`torch.compile` + liger kernels fails on XPU**: the Intel Triton backend errors with
  `RuntimeError: PassManager::run failed` (`triton/backends/intel/compiler.py:401 make_llir`).
  I narrowed it down — it fails even with **SwiGLU-only inference** (no labels → no fused CE in the
  graph), so this is a compile-path / Intel-Triton-backend issue, not specific to the fused loss.
  (My local patch only covers the validation bits in `layers.py`; I didn't replicate the
  `@torch.compiler.disable` you added around the fused CE in
  huggingface/kernels-community#948, and the un-touched `fused_linear_cross_entropy.py` still has a
  `.item()` — `total_n_non_ignore = target_mask.sum().item()` — that graph-breaks regardless. But
  since SwiGLU-only already fails, the loss isn't the blocker here.)
  Eager works in all modes, so the framework-side handling looks fine.
- `torch.compile` with no kernels is clean (7.75394). Each scenario is run in its own subprocess in
  the script — running kernelized + compile scenarios in one process leaks `kernelize` state into
  later compiles and produced spurious NaNs.
- I used the **non-tiled** `LigerSwiGLUMLP`. `LigerTiledSwiGLUMLP` also fails `_validate_layer`
  (it carries `_mlp_forward` / `num_shards` members) and would need the same stateless patch before
  it can be exercised on XPU.

## Repro script

`verify_pr46039.py` (subprocess-isolated, prints the table above):

```python
"""Verification for transformers PR #46039 + liger kernels on Intel XPU.

Adds liger `SwiGLUMLP` / `ForCausalLMLoss` to the XPU mapping (per @vasqu) and
checks that they engage on XPU and stay numerically close to eager across
inference / training / torch.compile.

Each scenario runs in its own subprocess so that global `kernelize` state from
one scenario cannot leak into a later `torch.compile` (which otherwise produces
NaNs / spurious recompiles).

Run inside an XPU container:
    HF_HOME=/path/to/cache python verify_pr46039.py
"""

import os
import sys
import subprocess

import torch
from transformers import AutoConfig, AutoModelForCausalLM
from transformers.loss.loss_utils import ForCausalLMLoss

DEVICE = "xpu" if torch.xpu.is_available() else "cpu"
DTYPE = torch.bfloat16
SEED = 0


def build():
    cfg = AutoConfig.for_model(
        "llama", vocab_size=2048, hidden_size=512, intermediate_size=1376,
        num_hidden_layers=4, num_attention_heads=8, num_key_value_heads=4,
        max_position_embeddings=512, tie_word_embeddings=False,
    )
    torch.manual_seed(SEED)
    return AutoModelForCausalLM.from_config(cfg, dtype=DTYPE).to(DEVICE).eval()


def data(seq_len=64):
    torch.manual_seed(1)
    ids = torch.randint(0, 2048, (2, seq_len), device=DEVICE)
    return ids, ids.clone()


def mlp_is_liger(m):
    return "liger" in getattr(m.model.layers[0].mlp.forward, "__module__", "").lower()


def sc_eager():
    m = build(); ids, lab = data()
    with torch.no_grad():
        print(f"{m(input_ids=ids, labels=lab).loss.item():.5f} | reference")


def sc_fallback():
    m = build(); ids, lab = data()
    with torch.no_grad():
        hidden = m.model(input_ids=ids).last_hidden_state
    kw = dict(labels=lab, vocab_size=m.config.vocab_size)
    a = ForCausalLMLoss(logits=m.lm_head(hidden), **kw).item()
    b = ForCausalLMLoss(logits=None, hidden_states=hidden, lm_head_weight=m.lm_head.weight, **kw).item()
    print(f"{b:.5f} | logits-path vs hidden-path diff {abs(a-b):.2e}")


def sc_inference():
    m = build(); ids, lab = data()
    m.set_use_kernels(True)
    with torch.no_grad():
        loss = m(input_ids=ids, labels=lab).loss.item()
    print(f"{loss:.5f} | SwiGLU=liger:{mlp_is_liger(m)}")


def sc_training():
    m = build(); ids, lab = data()
    m.train(); m.set_use_kernels(True)
    out = m(input_ids=ids, labels=lab)
    print(f"{out.loss.item():.5f} | SwiGLU=liger:{mlp_is_liger(m)} logitsNone:{out.logits is None}")


def sc_compile():
    m = build(); ids, lab = data()
    with torch.no_grad():
        loss = torch.compile(m)(input_ids=ids, labels=lab).loss.item()
    print(f"{loss:.5f} | no kernels")


def sc_compile_kernels():
    m = build(); ids, lab = data()
    m.set_use_kernels(True)
    with torch.no_grad():
        loss = torch.compile(m)(input_ids=ids, labels=lab).loss.item()
    print(f"{loss:.5f} | SwiGLU=liger:{mlp_is_liger(m)}")


def sc_seqlen():
    sl = int(os.environ["SEQ_LEN"])
    ref = build(); ids, _ = data(sl)
    km = build(); km.load_state_dict(ref.state_dict()); km.set_use_kernels(True)
    with torch.no_grad():
        rl = ref(input_ids=ids).logits
        kl = km(input_ids=ids).logits
    md = float((kl.float() - rl.float()).abs().max())
    print(f"{md:.5f} | max|dlogits| SwiGLU=liger:{mlp_is_liger(km)}")


SCENARIOS = {
    "eager": sc_eager,
    "fallback": sc_fallback,
    "inference": sc_inference,
    "training": sc_training,
    "compile": sc_compile,
    "compile_kernels": sc_compile_kernels,
    "seqlen": sc_seqlen,
}


def main():
    import transformers, kernels
    print(f"transformers {transformers.__version__} | torch {torch.__version__} | "
          f"kernels {kernels.__version__} | device {DEVICE} | dtype {DTYPE}\n")

    plan = [
        ("eager reference (no kernels)", "eager", {}),
        ("ForCausalLMLoss hidden_states path", "fallback", {}),
        ("use_kernels inference (liger SwiGLU)", "inference", {}),
        ("use_kernels training (liger SwiGLU + fused CE)", "training", {}),
        ("torch.compile, no kernels", "compile", {}),
        ("torch.compile + use_kernels", "compile_kernels", {}),
        ("liger SwiGLU seq_len=16", "seqlen", {"SEQ_LEN": "16"}),
        ("liger SwiGLU seq_len=128", "seqlen", {"SEQ_LEN": "128"}),
        ("liger SwiGLU seq_len=512", "seqlen", {"SEQ_LEN": "512"}),
    ]
    results = []
    for label, key, env in plan:
        e = dict(os.environ, SCENARIO=key, **env)
        p = subprocess.run([sys.executable, __file__], capture_output=True, text=True, env=e)
        out = (p.stdout or "").strip().splitlines()
        if p.returncode == 0 and out and "|" in out[-1]:
            val, note = (x.strip() for x in out[-1].split("|", 1))
        else:
            tail = (p.stderr or "").strip().splitlines()
            err = next((l for l in reversed(tail) if "Error" in l or "Exception" in l), "")
            val, note = "FAIL", err[:70]
        results.append((label, val, note))

    w = max(len(r[0]) for r in results)
    print(f"{'scenario':<{w}}  {'loss/val':>10}  note")
    print("-" * (w + 50))
    for label, val, note in results:
        print(f"{label:<{w}}  {val:>10}  {note}")


if __name__ == "__main__":
    key = os.environ.get("SCENARIO")
    if key:
        SCENARIOS[key]()
    else:
        main()
```


cc @michaelbenayoun if you could check the fusion portion. It works on kernels 0.13 but fails on 0.14+ - seems like the kernels were built on an older version which is not compatible

Just that feature remaining then we should be good to merge 🤗 expecting sometime early next week

Failing CI is related to something different opening another PR to fix (see #46695)

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46039&sha=9e4a2f

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

run-slow: llama4, mra, rwkv, sam3_video, yoso, mxfp4