[RFC, NPU] NPU benchmark - as a baseline for follow-up tasks.

[zheliuyu]

Motivation

Show the current state of NPU benchmarks. It also sets the direction for the next phase of performance optimization tasks.

Evaluation criteria: The speedup ratio of each kernel compared to huggingface/torch should be greater than 1.

Env

Test machine environment

• platform: https://www.autodl.com/
• NPU: Atlas 900 A2 PoD(64G)
• HDK=25.2.0
• CANN=8.5.0
• CPU: 24 vCPU Kunpeng-920
• OS: ubuntu22.04

Software dependencies

• python=3.10.8
• Liger-Kernel=0.7.0, Commit ID: 781083b
• torch=2.6.0
• torch_npu=2.6.0
• torchvision==0.21.0
• triton-ascend=3.2.0
• transformers=5.2.0

[zheliuyu]

Test results

Evaluation criteria: The speedup ratio of each kernel compared to huggingface/torch should be greater than 1.

kernel name	speed forward	speed backward	speed full	memory
rms_norm	❌	❌	❌	✅
swiglu	✅	❌	❌	✅
rope	✅	✅	✅	✅
cross_entropy	❌	✅	❌	✅
softmax	❌	❌	❌	❌
layer_norm	❌	❌	❌	✅
group_norm	❌	❌	❌	✅
geglu	✅	✅	✅	✅
fused_linear_grpo_loss_sequence	❌	✅	❌	❌
fused_linear_grpo_loss_token	❌	✅	❌	❌
dyt_beta=False	❌	❌	❌	✅
dyt_beta=True	❌	❌	❌	✅
jsd	✅	✅	✅	✅
kl_div	✅	❌	❌	✅
llama4_rope	✅	✅	✅	✅
poly_norm	✅	✅	✅	✅
qwen2vl_mrope	✅	❌	✅	✅
tvd	✅	✅	✅	✅

1. High priority

1.1 rms_norm

forward: partial pass ❌

backward: partial pass ❌

full: partial pass ❌

memory: ✔

1.2 swiglu

forward: ✔

backward: ❌

full: ❌

memory: ✔

1.3 rope

Note: The transformers version is 4.57.3. Related issue: #1155

forward: ✔

backward: ✔

full: ✔

memory: ✔

1.4 cross_entropy

forward: ❌

backward: ✔

full: ❌

no_grad_forward: ❌

memory: ✔

2. Others

2.1 softmax

forward: ❌

backward: ❌

full: ❌

memory: ❌

2.2 layer_norm

forward: ❌

backward: ❌

full: ❌

memory:✔

2.3 group_norm

forward: ❌

backward: ❌

full: ❌

memory: ✔

2.4 geglu

forward: ✔

backward: ✔

full: ✔

memory: ✔

2.5 fused_linear_grpo_loss_sequence

forward: ❌

backward: ✔

full: ❌

memory: ❌

2.6 fused_linear_grpo_loss_token

forward: ❌

backward: ✔

full: ❌

memory: ❌

2.7 dyt_beta=False

forward: ❌

backward: ❌

full: ❌

memory: ✔

2.8 dyt_beta=True

forward: ❌

backward: ❌

full: ❌

memory: ✔

2.9 jsd

forward: ✔

backward: ✔

full: ✔

memory: ✔

2.10 kl_div

forward: ✔

backward: ❌

full: ❌

memory: ✔

2.11 llama4_rope

forward: ✔

backward: ✔

full: ✔

memory: ✔

2.12 poly_norm

forward: ✔

backward: ✔

full: ✔

memory: ✔

2.13 qwen2vl_mrope

forward: ✔

backward: ❌

full: ✔

memory: ✔

2.14 tvd

forward: ✔

backward: ✔

full: ✔

memory: ✔