Add RMSNorm kernels for npu

[zheliuyu]

What does this PR do?

Add support for kernels-ext-npu/RMSNorm acceleration on npu.

Basework

• Discussed: How to use other acceleration apis of npu? #39105
• Supported kernels:
  • [Merged through #155] Add support for npu kernelize kernels#146
  • Add support for NPU kernelize/layers kernels#155

Test

Prepare the env

git clone https://github.com/huggingface/kernels
pip install -e kernels

git clone https://github.com/zheliuyu/transformers-add-kernels
pip install -e transformers-add-kernels

Test script

• Set logging.basicConfig(level=logging.DEBUG) to observe if "kernels-ext-npu/RMSNorm" takes effect.
• Set use_kernels=True
• Set use_kernel_forward_from_hub("RMSNorm") on modeling_qwen3.py

from transformers import AutoModelForCausalLM, AutoTokenizer
import logging


logging.basicConfig(level=logging.DEBUG)

model_name = "Qwen/Qwen3-0.6B"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    use_kernels=True,
    device_map="auto"
)

# prepare the model input
prompt = "Give me a short introduction to large language model."
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True # Switches between thinking and non-thinking modes. Default is True.
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# conduct text completion
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32768
)
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 

# parsing thinking content
try:
    # rindex finding 151668 (</think>)
    index = len(output_ids) - output_ids[::-1].index(151668)
except ValueError:
    index = 0

thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

print("thinking content:", thinking_content)
print("content:", content)

Output

INFO:root:Using layer `RMSNorm` from repo `kernels-ext-npu/RMSNorm` (revision: main), layer `RMSNorm`
DEBUG:root:kernelize mode: Mode.INFERENCE, repo mode: Mode.FALLBACK


=========repeat x N=============

INFO:root:Using layer `RMSNorm` from repo `kernels-ext-npu/RMSNorm` (revision: main), layer `RMSNorm`
DEBUG:root:kernelize mode: Mode.INFERENCE, repo mode: Mode.FALLBACK


thinking content: <think>
Okay, the user wants a short introduction to a large language model. Let me start by recalling what I know. Large language models are AI systems that can understand and generate human language. They're used in various fields like translation, writing, and customer service. I should mention their capabilities and applications.

Wait, should I include specific examples? Maybe mention something like translating from Spanish to English or helping with customer service. Also, emphasize their adaptability to different languages and topics. Oh, and maybe touch on their training data sources to add depth. Let me check if I'm covering all key points without being too technical. Keep it concise but informative. Alright, that should do it.
</think>
content: A large language model (LLM) is an advanced AI system capable of understanding and generating human language, enabling tasks such as translation, text generation, and customer service. These models are trained on vast datasets to learn patterns and nuances in language, allowing them to interact with users in natural and meaningful ways. They are used across industries for tasks ranging from content creation to complex problem-solving.

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you write any new necessary tests?

@ArthurZucker Ready for review.

[github-actions]

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

run-slow: qwen3, qwen3_moe, qwen3_next, qwen3_omni_moe, qwen3_vl, qwen3_vl_moe

[zheliuyu]

Supplement test by LLaMA-Factory

env

cann=8.2.RC1
torch=2.6.0
torch_npu=2.6.0
device_nums=1

config

### model
model_name_or_path: Qwen/Qwen3-0.6B
trust_remote_code: true

### method
stage: sft
do_train: true
finetuning_type: lora
lora_rank: 8
lora_target: all

### dataset
dataset: identity,alpaca_en_demo
template: llama3
cutoff_len: 2048
max_samples: 1000
overwrite_cache: true
preprocessing_num_workers: 16
dataloader_num_workers: 4

### output
output_dir: saves/Qwen/sft
logging_steps: 10
save_steps: 500
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]

### train
per_device_train_batch_size: 1
gradient_accumulation_steps: 8
learning_rate: 1.0e-4
num_train_epochs: 1.0
lr_scheduler_type: cosine
warmup_ratio: 0.1
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null

result

Using LLaMA-Factory, we fine-tuned Qwen3-0.6B.

llamafactory-cli train custom.yaml

No RMSNorm kernels

{'loss': 1.6764, 'grad_norm': 1.1894797086715698, 'learning_rate': 6.428571428571429e-05, 'epoch': 0.07}
{'loss': 1.6344, 'grad_norm': 1.0937652587890625, 'learning_rate': 9.95928267727387e-05, 'epoch': 0.15}
{'loss': 1.638, 'grad_norm': 1.061297059059143, 'learning_rate': 9.637512255510475e-05, 'epoch': 0.22}
{'loss': 1.5791, 'grad_norm': 1.171653151512146, 'learning_rate': 9.014848666026138e-05, 'epoch': 0.29}
{'loss': 1.619, 'grad_norm': 1.0558280944824219, 'learning_rate': 8.131691849466153e-05, 'epoch': 0.37}
{'loss': 1.4807, 'grad_norm': 1.0464788675308228, 'learning_rate': 7.045343185856701e-05, 'epoch': 0.44}
{'loss': 1.4885, 'grad_norm': 1.1576353311538696, 'learning_rate': 5.826287640548425e-05, 'epoch': 0.51}
{'loss': 1.3543, 'grad_norm': 0.846390426158905, 'learning_rate': 4.5536205268140294e-05, 'epoch': 0.59}
{'loss': 1.5345, 'grad_norm': 1.0654412508010864, 'learning_rate': 3.309915607957487e-05, 'epoch': 0.66}
{'loss': 1.4327, 'grad_norm': 1.3742737770080566, 'learning_rate': 2.175867509055033e-05, 'epoch': 0.73}
{'loss': 1.5616, 'grad_norm': 0.8401137590408325, 'learning_rate': 1.2250560517925746e-05, 'epoch': 0.81}
{'loss': 1.5832, 'grad_norm': 1.0662145614624023, 'learning_rate': 5.191722151947226e-06, 'epoch': 0.88}
{'loss': 1.4545, 'grad_norm': 0.7635636329650879, 'learning_rate': 1.040154736207194e-06, 'epoch': 0.95}
{'train_runtime': 631.4651, 'train_samples_per_second': 1.726, 'train_steps_per_second': 0.217, 'train_loss': 1.5288616653776517, 'epoch': 1.0}

***** train metrics *****
    "epoch": 1.0,
    "total_flos": 547575576920064.0,
    "train_loss": 1.5288616653776517,
    "train_runtime": 631.4651,
    "train_samples_per_second": 1.726,
    "train_steps_per_second": 0.217

Add RMSNorm kernels

{'loss': 1.6767, 'grad_norm': 1.1918253898620605, 'learning_rate': 6.428571428571429e-05, 'epoch': 0.07}
{'loss': 1.6352, 'grad_norm': 1.0958017110824585, 'learning_rate': 9.95928267727387e-05, 'epoch': 0.15}
{'loss': 1.6379, 'grad_norm': 1.058767318725586, 'learning_rate': 9.637512255510475e-05, 'epoch': 0.22}
{'loss': 1.5788, 'grad_norm': 1.165596842765808, 'learning_rate': 9.014848666026138e-05, 'epoch': 0.29}
{'loss': 1.6186, 'grad_norm': 1.061798095703125, 'learning_rate': 8.131691849466153e-05, 'epoch': 0.37}
{'loss': 1.48, 'grad_norm': 1.0418390035629272, 'learning_rate': 7.045343185856701e-05, 'epoch': 0.44}
{'loss': 1.4884, 'grad_norm': 1.1575311422348022, 'learning_rate': 5.826287640548425e-05, 'epoch': 0.51}
{'loss': 1.3542, 'grad_norm': 0.8492579460144043, 'learning_rate': 4.5536205268140294e-05, 'epoch': 0.59}
{'loss': 1.5346, 'grad_norm': 1.066164493560791, 'learning_rate': 3.309915607957487e-05, 'epoch': 0.66}
{'loss': 1.4338, 'grad_norm': 1.3807560205459595, 'learning_rate': 2.175867509055033e-05, 'epoch': 0.73}
{'loss': 1.5608, 'grad_norm': 0.838101863861084, 'learning_rate': 1.2250560517925746e-05, 'epoch': 0.81}
{'loss': 1.5819, 'grad_norm': 1.069312572479248, 'learning_rate': 5.191722151947226e-06, 'epoch': 0.88}
{'loss': 1.4535, 'grad_norm': 0.7640863060951233, 'learning_rate': 1.040154736207194e-06, 'epoch': 0.95}
{'train_runtime': 558.9375, 'train_samples_per_second': 1.95, 'train_steps_per_second': 0.245, 'train_loss': 1.5287004526514207, 'epoch': 1.0}

***** train metrics *****
    "epoch": 1.0,
    "total_flos": 547575576920064.0,
    "train_loss": 1.5287004526514207,
    "train_runtime": 558.9375,
    "train_samples_per_second": 1.95,
    "train_steps_per_second": 0.245

(631.4651 - 558.9375) / 631.4651 ≈ 11.5%

The results show an approximate 11% speedup from adding the RMSNorm kernels.

[zheliuyu]

@ArthurZucker I've added a supplementary experiment. Could you please review this PR.

[Rocketknight1]

cc @ivarflakstad

[MekkCyber]

Thanks for this integration @zheliuyu ! though i'm not really sure it's useful or not since it's basically just calling : torch_npu.npu_rms_norm(x, self.weight, epsilon=self.variance_epsilon)[0] ?

[zheliuyu]

@MekkCyber Yes, the torch_npu.npu_rms_norm() API accelerates the sft process. I've added supplementary experiments above to validate its effectiveness

torch_npu.npu_rms_norm accelerates computation by fusing the operators below into a single kernel upon call, reducing data transfer time. Although the code appears quite simple.

[MekkCyber]

Thanks for the explanation @zheliuyu, but doesn't torch dispatch automatically to this implementation when having a npu device

[zheliuyu]

@MekkCyber Thanks for your review.

torch_npu.npu_rms_norm() is not enabled by default. it is provided by the torch_npu.

If it's not explicitly called, the transformers library will use the npu logic from the standard torch.

[zheliuyu]

@MekkCyber Hey, I was wondering if you'd consider merging this PR. Alternatively, if you have any suggestions, please do let me know.

[MekkCyber]

Hey @zheliuyu, my only concern is that we are using big machinery to just call torch_npu, also we cannot use kernels from the hub from many sources because it's code executing on the user device and we need to ensure that the code is trustworthy. If you can share some benchmarks on npu devices comparing the latency with and without kernels, we can move the kernel to kernels-community and host it there in case it leads to big improvements

[zheliuyu]

@MekkCyber I believe you're right. Let me check if we can directly reuse the kernels from the kernels-community project.
Thank you for the suggestion.