Expand npu device for KernelConfig

[zheliuyu]

What does this PR do?

Expand npu device for KernelConfig.

test

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

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**

[vasqu]

cc @MekkCyber

[MekkCyber]

Thank you @zheliuyu ! lgtm

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