Repository: huggingface/transformers
Pull request: Expand npu device for KernelConfig

# What does this PR do?

Expand npu device for KernelConfig.

## test

```python
from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig
import logging
import time


# Set the level to `DEBUG` to see which kernels are being called.
logging.basicConfig(level=logging.DEBUG)

model_name = "Qwen/Qwen3-0.6B"

# load the tokenizer and the model
kernel_mapping = {
    "RMSNorm":
        "kernels-ext-npu/rmsnorm:rmsnorm",
}
kernel_config = KernelConfig(kernel_mapping)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    kernel_config = kernel_config,
)

# prepare the model input
prompt = "Output the first 20 digits of pi directly."
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False,
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# Print Runtime
start_time = time.time()
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32768
)
print("runtime: ", time.time()-start_time)

output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
content = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
print("content:", content)
```

## Output
```python
root@autodl-container:~# python test_rmsnorm.py 
A kernel_config was provided but use_kernels is False; setting use_kernels=True automatically. To suppress this warning, explicitly set use_kernels to True.
Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 311/311 [00:00<00:00, 635.96it/s, Materializing param=model.norm.weight]
Fetching 2 files: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 25890.77it/s]
Download complete: : 0.00B [00:00, ?B/s]                                                                                                                                                                                                                                                           | 0/2 [00:00<?, ?it/s]
DEBUG:root:kernelize mode: Mode.INFERENCE, repo mode: Mode.INFERENCE
INFO:root:Using layer `rmsnorm` from repo `kernels-ext-npu/rmsnorm` (revision: main), layer `rmsnorm`
runtime:  4.810910701751709
content: The first 20 digits of π (pi) are:

**3.14159265358979323846**
```

Changed files:
filename: src/transformers/utils/kernel_config.py
patch: @@ -57,8 +57,8 @@ def infer_device(model):
 def add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping):
     from kernels import LayerRepository
 
-    if device not in ["cuda", "rocm", "xpu"]:
-        raise ValueError(f"Only cuda, rocm, and xpu devices supported, got: {device}")
+    if device not in ["cuda", "rocm", "xpu", "npu"]:
+        raise ValueError(f"Only cuda, rocm, xpu and npu devices supported, got: {device}")
     repo_layer_name = repo_name.split(":")[1]
     repo_id = repo_name.split(":")[0]
     compatible_mapping[layer_name] = {
@@ -102,8 +102,8 @@ def sanitize_kernel_mapping(self, model):
         """
         Validates the kernel_mapping to ensure that:
         1. Each layer_name in the mapping is registered in the model (i.e., the model contains a module with a matching kernel_layer_name).
-        2. Each kernel value is either a string of the form 'org/repo:layer_name' or a dict mapping device types ("cuda", "rocm", "xpu") to such strings.
-        3. Each device key in a dict is one of "cuda", "rocm", or "xpu".
+        2. Each kernel value is either a string of the form 'org/repo:layer_name' or a dict mapping device types ("cuda", "rocm", "xpu", "npu") to such strings.
+        3. Each device key in a dict is one of "cuda", "rocm", "xpu", or "npu".
         4. Each repo_name is a valid repository and layer name in the format 'org/repo:layer_name' (i.e., a string containing both a slash and a colon).
 
         Args:
@@ -154,8 +154,8 @@ def sanitize_kernel_mapping(self, model):
 
             elif isinstance(kernel, dict):
                 for device, repo_name in kernel.items():
-                    if device not in ["cuda", "rocm", "xpu"]:
-                        raise ValueError(f"Only cuda, rocm, and xpu devices supported, got: {device}")
+                    if device not in ["cuda", "rocm", "xpu", "npu"]:
+                        raise ValueError(f"Only cuda, rocm, xpu and npu devices supported, got: {device}")
 
                     if not isinstance(repo_name, str) or "/" not in repo_name or ":" not in repo_name:
                         raise ValueError(