How to use other acceleration apis of npu?
Author: zheliuyu
Pull request: no
Labels: name: Feature request

### Feature request

I noticed that transformers now support using flash attention directly in the npu by [```npu_flash_attention.py```](https://github.com/huggingface/transformers/pull/36696). There are many other acceleration apis that can be used in npu, such as shown in [doc](https://www.hiascend.com/document/detail/zh/Pytorch/700/ptmoddevg/trainingmigrguide/performance_tuning_0028.html).

 How can we use them directly in transformers? How to switch seamlessly between different devices?

### Motivation

Request to integrate other acceleration apis of npu in transformers. If this can be done, the ease of using transformers will be greatly improved in npu.

Comments:
track:
https://github.com/huggingface/transformers/pull/39238

PR reviewed!

> track: [#39238](https://github.com/huggingface/transformers/pull/39238)

#### already added
- torch_npu.npu_rms_norm()
- torch_npu.npu_swiglu()
- torch_npu.npu_rotary_mul()

#### next
- Refactoring the usage of npu fusion operator using `kernels`.

restart:
    by https://github.com/huggingface/kernels/pull/146

Add support for NPU kernelize/layers: 
https://github.com/huggingface/kernels/pull/155

🚀 

Related Work: 
transformers + kernels: https://github.com/huggingface/transformers/pull/42106

## 1 How to use ascend npu accelerated kernels in transformers? 

This issue has been elegantly resolved through npu's native support for `huggingface/kernels` <sup>[1][2]</sup>. Let's briefly recap.

### 1.1 Use  flash-attn on npu
- set `attn_implementation="flash_attention_2"` in `AutoModel.from_pretrained`,  it invokes the built-in logic in `transformers`, specifically the [npu_flash_attention.py](https://github.com/huggingface/transformers/pull/41542) file.
- (Optional) set `attn_implementation="kernels-ext-npu/flash-attn2-npu"` in `AutoModel.from_pretrained`, `transformers` will load flash attention logic through `huggingface/kernels` <sup>[3]</sup>, which is identical to the code of npu_flash-attention.py file.

### 1.2  Use other kernels on npu
- All components are accessed via `huggingface/kernels` and `huggingface/kernels-community`. For example, RMSNorm has now been enabled by `kernels-community/liger_kernels` <sup>[4]</sup>. 

## 2 *Discussion

I think there is only one question left to discuss now. Then we can close this issue.

Currently, npu acceleration kernels are primarily invoked through [torch_npu](https://pypi.org/project/torch-npu/), which is the officially recommended method. For example, `torch_npu.npu_rms_norm`, `torch_npu.npu_stiglu`, and `torch_npu.npu_rotary_rul`. I attempted to push them to `kernels-community` but was unsuccessful <sup>[5]</sup> due to its principle of "keep the number of maintained kernels small and manageable". I think this reason is convincing. Hence, I've created https://huggingface.co/kernels-ext-npu to host these npu accelerated kernels.

**My only remaining question is whether the transformers community would consider registering kernels from the kernels-ext-npu hub in [integrations/hub_kernels.py#L87](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/hub_kernels.py#L87) (only npu scenarios), even though they are not part of `kernels-community`.**

Eg:
``` diff
"npu": {
   Mode.INFERENCE: LayerRepository(
-        repo_id="kernels-community/liger_kernels",
-        layer_name="LigerRMSNorm",
   )
}

"npu": {
  Mode.INFERENCE: LayerRepository(
+       repo_id="kernels-ext-npu/rmsnorm",
+       layer_name="rmsnorm",
  )
}
```

We want demonstrate the performance optimization benefits of `torch_npu` kernels and submit a PR to enable them by default in transformers-npu. But before proceeding, looking forward to your suggestions😀. @ArthurZucker @MekkCyber @danieldk 

## 3 Reference
[1] Add support for npu kernelize. https://github.com/huggingface/kernels/pull/146
[2] Add support for NPU kernelize/layers. https://github.com/huggingface/kernels/pull/155
[3] Add conditional checks to _check_and_adjust_attn_implementation(). [huggingface/transformers#41542](https://github.com/huggingface/transformers/pull/41542)
[4] Add rmsnorm kernels support for npu. [huggingface/transformers#42106](https://github.com/huggingface/transformers/pull/42106)
[5] flash-attn2: add flash-attn2-npu kernel. https://github.com/huggingface/kernels-community/pull/64


## 4 Appendix

### 4.1 No kernel

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
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-8B"

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

# Prepare the model input
prompt = "Output the first 20 digits of pi."
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)

for _ in range(11):
    # 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)
```

### 4.2 Add rmsnorm
- Set `use_kernels=True`
- Add `kernels-ext-npu/rmsnorm` for integrations/hub_kernels.py https://github.com/huggingface/transformers/compare/main...zheliuyu:transformers-dev:main

<img width="742" height="396" alt="Image" src="https://github.com/user-attachments/assets/27a02481-fc71-40f0-aae4-5fe8afae827d" />

Oh right, transformers v5 + kernels already provides a method for this, so this issue can be closed.