flash-attn2: add flash-attn2-npu kernel

[zheliuyu]

Introduction

Following the design of huggingface/transformers and huggingface/kernels, I've created the flash-attn2-npu project to supersede the current NPU flash attention implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/npu_flash_attention.py

Dependencies

• huggingface/kernels. Add support for NPU kernelize/layers kernels#155
• huggingface/kernels-builder. The huggingface/kernels-builder requires no additional NPU adaptations and can be used directly. Becasue, most NPU acceleration methods are implemented in torch_npu, which is a pure python project. Of course, I'll keep a close watch on any future NPU-specific adaptations needed for kernels-builder.

Follow CONTRIBUTING.md

• Create a new project flash-attn2-npu.
  • Should I put it in the flash-attn2 folder?
• Add flake.nix.
  • set doGetKernelCheck = false to avoid torch_npu ImportError when no GPUs. Instead, I will provide the validation results for the LoRA algorithm within the transformers.
• Add build.toml.
• Test by test_flash_attn2_npu.py.
• Test by transformers.

Test by transformers + LLaMA-Factory

LLaMA-Factory: https://github.com/hiyouga/LLaMA-Factory

Config

Qwen3_lora_sft.yaml, copied from examples/train_lora/llama3_lora_sft.yaml

### 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/
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

Run

llamafactory-cli train examples/train_lora/Qwen3_lora_sft.yaml

AutoModel.from_pretrained(attn_implementation="flash_attention_2")

Built in implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/npu_flash_attention.py

"epoch": 1.0
"total_flos": 547575576920064.0
"train_loss": 1.5288773348731717
"train_runtime": 666.4749

AutoModel.from_pretrained(attn_implementation="kernels-ext-npu/flash-attn2-npu")

Built from kernels-community: this pr.

The build outputs are temporarily stored for use by transformers: https://huggingface.co/kernels-ext-npu/flash-attn2-npu/tree/main

"epoch": 1.0
"total_flos": 547575576920064.0
"train_loss": 1.5290487386884482
"train_runtime": 670.0417

[zheliuyu]

Ready for review. @danieldk @drbh @MekkCyber

[MekkCyber]

Hey @zheliuyu, Thanks so much for your contribution! It looks really great. Right now, we can’t merge it into kernels-community because we’re trying to keep the number of maintained kernels small and manageable—we usually only include the most widely used ones or those needed in our frameworks.

But don’t let that stop you! You can absolutely host your kernel on your Hub organization and use it from there. The whole idea of the kernels project is to give the community a space to share their work and let others benefit from it.

[zheliuyu]

Hey @zheliuyu, Thanks so much for your contribution! It looks really great. Right now, we can’t merge it into kernels-community because we’re trying to keep the number of maintained kernels small and manageable—we usually only include the most widely used ones or those needed in our frameworks.

But don’t let that stop you! You can absolutely host your kernel on your Hub organization and use it from there. The whole idea of the kernels project is to give the community a space to share their work and let others benefit from it.

While this code wasn't merged, your suggestions were convincing. Thanks very much. I'll now focus on reusing work from kernels-community.