[tarekziade]

use_kernels=True fails on MPS for every model with an RMSNorm: the kernel mapping points at kernels-community/mlx_rmsnorm, which the kernels API no longer serves (Repository Not Found for url: https://huggingface.co/api/kernels/kernels-community/mlx_rmsnorm/refs). The kernel lives at kernels-community/mlx-rmsnorm.

Reproducer (Apple M5, torch 2.15.0, kernels 0.18.2): fails on main, and with this change the kernelized RMSNorm matches eager to 9e-8.

import torch
from transformers import LlamaConfig, LlamaForCausalLM

config = LlamaConfig(vocab_size=128, hidden_size=64, intermediate_size=128, num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=4)
model = LlamaForCausalLM(config).to("mps").eval()
ids = torch.tensor([[1, 2, 3]], device="mps")
reference = model(ids).logits
model.set_use_kernels(True)
print((model(ids).logits - reference).abs().max())

[vasqu]

Thanks, there might be more since im not on mps to verify :/

[tarekziade]

Thanks, there might be more since im not on mps to verify :/

Yeah I wished we had MPS hardware to run test on CI...

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