[RFC] Roadmap: NPU support for flash-linear-attention

[zheliuyu]

Thanks to the flash-linear-attention community for merging our native NPU support PR.

This roadmap outlines upcoming plans for native NPU support. Feedback and discussion are welcome.

Native NPU Support in flash-linear-attention

Native support is built on top of flash-linear-attention's dispatch system:

• [Backend] Introduce dispatch system #741

We have validated this approach for NPU adaptation work:

• [Ops] Feat: Add triton-ascend backend support for FLA module kernels #927

The following diagram illustrates the flow for adapting kernels under fla/modules/:

User / model code (unchanged)
        │
        ▼
fla/modules/<module>.py              # Public API + @dispatch('modules') hooks only
        │
        ▼
fla/ops/backends                     # Runtime backend selection (BackendRegistry)
        │
        ▼
fla/modules/backends/triton_ascend/  # NPU Triton kernels + launchers
        ├── __init__.py              # TritonAscendBackend registration
        ├── rotary.py
        ├── fused_cross_entropy.py
        ├── fused_linear_cross_entropy.py
        ├── activations.py
        ├── fused_kl_div.py
        ├── layernorm.py
        ├── grpo.py
        └── causal_conv1d.py

Dependency Installation

Note that flash-linear-attention already supports other devices such as XPU and AMD:

• [Docs] Add Ascend NPU installation section to README #938
• [Refactor] Decouple torch/triton backend from base deps #937

NPU CI Setup

Future native adaptations will introduce a substantial amount of triton-ascend code, making NPU CI especially important. We plan to use the CI tooling provided by ascend-gha-runners:

• [CI] Add ascend-a2-ci workflow for Atlas A2 NPU module tests #944
• [Ops] Expand triton-ascend backends(fla/modules, fla/ops/utils) and upgrade CANN to 9.0.0 #973
• [CI] Parallelize Ascend A2 CI tests with pytest-xdist across 8 NPUs #987

Reference：

• https://github.com/ascend-gha-runners/docs/blob/main/docs/user-manual-gha-en.md
• https://github.com/ascend-gha-runners/docs/blob/main/docs/user-manual-gha-zh.md

This is our top priority at the current stage, and work is still in progress.

Adaptation Order

At a high level, adaptation work is organized as follows:

• fla/modules/
  • [Ops] Feat: Add triton-ascend backend support for FLA module kernels #927
  • [Conv] Chunk Ascend causal_conv1d launches for 65535 grid limit #992
  • [Conv] Harden triton-ascend causal_conv1d silu/add against Ascend launch limits #997
  • [Perf] Optimize Ascend NPU activation kernels #995
  • [Fix] remove redundant decorators, add PowGlu kernel and pin the version of torch_npu to 2.7.1.post4 for npu. #982
  • [Fix] fix benchamrk issues and add ub management for npu #1001
  • [Ops] Add triton-ascend backend for l2norm and fused_norm_gate kernels #1002
  • [Fix] [Test] Remove unused GeLU/SquaredReLU NPU kernels and add Ascend UB Manager tests #1006
• fla/ops/
• fla/layers/
• fla/models/

For the remaining work, we propose the following order within fla/ops/:

• fla/ops/based/native.py, tests/ops/test_based.py
• fla/ops/based/parallel.py, tests/ops/test_based.py
• fla/ops/based/fused_chunk, tests/ops/test_based.py
• fla/ops/utils/cumsum.py, tests/ops/utils/test_cumsum.py
  • [Ops] Expand triton-ascend backends(fla/modules, fla/ops/utils) and upgrade CANN to 9.0.0 #973
• fla/ops/utils/index.py, tests/ops/utils/test_index.py
  • [Ops] Expand triton-ascend backends(fla/modules, fla/ops/utils) and upgrade CANN to 9.0.0 #973
• fla/ops/utils/cache.py, tests/ops/test_cache.py
  • [Ops] Expand triton-ascend backends(fla/modules, fla/ops/utils) and upgrade CANN to 9.0.0 #973
• fla/ops/common/*
• fla/ops/gated_delta_rule/*
  • [Ops] Add triton-ascend backend for GDN kernels #1011
• fla/ops/gated_delta_product/*
• fla/ops/gdn2/*
• fla/ops/kda/*
• Other kernels

Skills

Because triton-ascend differs from upstream triton in programming model and performance characteristics, migrating kernels to NPU requires additional tuning time. We plan to capture reusable migration skills to improve efficiency over time.

Hardware Coverage

For current migration work, we recommend validating on the following hardware:

• Atlas 800T A2 (32G, x86)
• Atlas 800T A3 (64G, x86)

We expect to publish validation results on additional hardware platforms in future updates.

Integration with Training Frameworks

Once native NPU adaptation of flash-linear-attention reaches a sufficient level of maturity, we need to demonstrate real-world usability by running end-to-end SFT or RL workflows in production training frameworks. We plan to validate integration through the following projects:

• verl-project/verl
• hiyouga/LlamaFactory
• modelscope/ms-swift

tilelang-ascend Support

flash-linear-attention also provides a tilelang DSL backend. Ascend offers tilelang-ascend, and native support may be feasible in the future. Whether we start this effort will depend on community feedback. For now, the primary focus remains on triton-ascend.

[sunyi0505]

FLA Module Benchmark — NVIDIA H100

Benchmark results for Flash Linear Attention (FLA) fused operators on NVIDIA H100.

Activation Function

Details

B	T	D	sigmoid_fwd	sigmoid_fwdbwd	logsigmoid_fwd	logsigmoid_fwdbwd	swish_fwd	swish_fwdbwd	gelu_fwd	gelu_fwdbwd	sqrelu_fwd	sqrelu_fwdbwd	swiglu_fwd	swiglu_fwdbwd	powglu_fwd	powglu_fwdbwd
4	512	1024	0.0268	0.99475	0.2602	1.62685	0.06557	1.84797	0.00922	1.67446	0.00819	1.67334	0.07795	1.98016	0.0825	1.99043
4	512	2048	0.06563	1.85744	0.07334	1.88675	0.07174	1.87174	0.01219	1.7853	0.01171	1.79216	0.08154	2.00474	0.08454	2.00989
4	512	4096	0.06443	1.86726	0.09402	1.82891	0.06637	1.85565	0.04909	1.92317	0.0543	2.2439	0.08378	2.42134	0.09648	2.45182
4	1024	1024	0.07155	2.25019	0.06274	2.34496	0.06669	2.30595	0.01859	2.25872	0.01341	0.42922	0.07709	1.97427	0.08416	1.99242
4	1024	2048	0.0635	1.83888	0.07005	1.8683	0.05462	1.84413	0.04907	1.76931	0.04214	1.77914	0.07758	1.95789	0.0833	0.796
4	1024	4096	0.07893	0.61494	0.07014	0.6156	0.06179	0.58582	0.04926	1.77482	0.04304	1.78264	0.08502	1.90653	0.0845	1.99546
4	2048	1024	0.06365	1.83747	0.06826	1.8545	0.05661	1.84789	0.04882	1.76816	0.04704	1.0869	0.0767	0.68878	0.08406	0.70427
4	2048	2048	0.06365	0.6008	0.07046	1.8279	0.0664	1.83651	0.05126	1.77827	0.05104	1.77728	0.08509	1.95734	0.08064	1.94682
4	2048	4096	0.07347	1.92882	0.07648	1.8271	0.0633	0.8745	0.07546	2.16331	0.06701	2.20901	0.09114	1.94534	0.12539	2.42261
4	4096	1024	0.07149	2.25546	0.07533	2.24875	0.06864	2.25056	0.05446	2.21352	0.05194	1.77802	0.08386	1.97088	0.09053	1.91782
4	4096	2048	0.06973	0.62707	0.07453	0.65262	0.05888	0.57914	0.07514	0.48662	0.06589	1.78643	0.08474	2.11826	0.1088	2.4399
4	4096	4096	0.10678	2.31178	0.13594	2.28171	0.10278	2.2848	0.14515	2.25582	0.12595	2.2576	0.14746	2.37179	0.22342	2.45984
4	8192	1024	0.07677	2.30826	0.0831	2.3384	0.07774	2.31456	0.07645	2.26522	0.06656	2.21232	0.09088	2.2296	0.12092	2.46267
4	8192	2048	0.10541	2.28365	0.13712	2.32051	0.10245	2.17403	0.14515	2.02147	0.12522	0.64955	0.14773	0.82878	0.2096	0.79136
4	8192	4096	0.20701	1.8864	0.26448	1.95248	0.19728	1.9388	0.28179	2.23227	0.24288	2.24202	0.28378	2.4752	0.43888	2.56938

Causal Conv1D

Details

T	D	causal_conv1d_fwd	causal_conv1d_fwdbwd
256	256	0.029664	1.857376
512	256	0.031808	2.254656
1024	256	0.0352	2.847104
2048	256	0.032656	2.446720
4096	256	0.033632	1.964000
8192	256	0.033056	1.740832
16384	256	0.032928	1.519552
32768	256	0.035904	0.372032
65536	256	0.03872	0.430848
256	512	0.033152	0.427520
512	512	0.032992	0.410272
1024	512	0.034048	1.992416
2048	512	0.03232	1.976416
4096	512	0.032928	0.901984
8192	512	0.03312	0.374112
16384	512	0.034272	1.753968
32768	512	0.039488	0.883520
65536	512	0.068096	0.920160
256	1024	0.032448	1.970880
512	1024	0.032448	2.741440
1024	1024	0.032544	2.938272
2048	1024	0.03328	2.445104
4096	1024	0.032704	1.765168
8192	1024	0.035456	2.015520
16384	1024	0.039776	1.132096
32768	1024	0.068224	1.668800
65536	1024	0.126336	2.310624
256	2048	0.032832	2.382704
512	2048	0.033632	2.716448
1024	2048	0.032448	2.946272
2048	2048	0.033696	1.952864
4096	2048	0.03484	1.931264
8192	2048	0.039376	1.932576
16384	2048	0.06848	1.973504
32768	2048	0.126176	1.596848
65536	2048	0.240688	1.753920
256	4096	0.0344	2.224704
512	4096	0.033408	2.858880
1024	4096	0.032992	2.971712
2048	4096	0.035072	1.929536
4096	4096	0.040576	0.946368
8192	4096	0.069472	1.896704
16384	4096	0.12608	2.036560
32768	4096	0.239488	1.772224
65536	4096	0.468192	3.441536

Cross Entropy

Details

T	naive	fused	fused_linear	naive_bwd	fused_bwd	fused_linear_bwd
128	0.0859808	0.812064	28.4412	3.37282	3.94562	30.5018
256	1.62803	1.61277	37.9333	5.40365	5.03837	39.8904
512	3.34845	2.96659	57.7802	9.73034	9.05408	59.4665
1024	6.94774	6.00122	100.518	19.0548	17.2757	101.83
2048	13.9516	11.8873	186.783	36.3447	36.1541	188.172
4096	26.4369	23.8772	358.459	74.3272	69.2011	359.768
8192	53.5101	54.3401	716.053	150.659	147.641	717.562
16384	121.335	99.8885	1424.24	321.297	274.294	1428.72

L2 Norm

Details

B	T	H	D	naive	compiled	fused	naive_bwd	compiled_bwd	fused_bwd
16	256	1	2048	0.04496	0.020704	0.014304	0.438176	0.487808	1.93664
16	512	1	2048	0.073088	0.047328	0.054394	1.79851	2.14442	2.46931
16	1024	1	2048	0.130304	0.063232	0.060576	1.78323	2.10822	2.46794
16	2048	1	2048	0.249088	0.11624	0.096608	1.91341	2.09219	2.3296
16	4096	1	2048	0.476832	0.226208	0.184832	3.74554	2.15514	2.45278
16	8192	1	2048	0.93192	0.43552	0.361632	7.40138	2.21082	2.52781
16	16384	1	2048	1.8295	0.859232	0.715232	14.6922	2.86326	2.79222
16	256	2	1024	0.02752	0.051552	0.05616	1.79045	2.29054	2.46346
16	512	2	1024	0.040288	0.049568	0.053088	1.79659	2.12904	2.46333
16	1024	2	1024	0.076384	0.053216	0.055072	1.80456	2.18682	2.45914
16	2048	2	1024	0.135744	0.075008	0.060704	1.80136	2.13229	2.48422
16	4096	2	1024	0.252416	0.142592	0.095648	1.92654	2.0291	2.39066
16	8192	2	1024	0.482448	0.276944	0.18456	3.76051	2.01555	2.28502
16	16384	2	1024	0.938432	0.5464	0.359584	7.4201	1.98243	2.23882
16	256	4	512	0.020192	0.050944	0.070688	1.77131	2.16435	2.61709
16	512	4	512	0.029184	0.049408	0.068544	1.77182	2.19477	2.596
16	1024	4	512	0.042144	0.05168	0.069824	1.81308	2.18803	2.61322
16	2048	4	512	0.078816	0.069632	0.069088	1.80765	2.09366	2.58806
16	4096	4	512	0.134816	0.132384	0.075232	1.78573	2.11677	2.52182
16	8192	4	512	0.258112	0.256224	0.096064	1.93341	2.0232	2.58061
16	16384	4	512	0.494496	0.505216	0.183488	3.77677	1.9633	2.42195
16	256	8	256	0.018464	0.050144	0.068832	1.80541	2.12728	2.60464
16	512	8	256	0.02064	0.049296	0.06736	1.79205	2.13066	2.584
16	1024	8	256	0.029952	0.09632	0.067616	1.79835	2.20838	2.60131
16	2048	8	256	0.045952	0.06672	0.067872	1.79507	2.0837	2.59216
16	4096	8	256	0.08352	0.125472	0.068128	1.78642	2.16614	2.56386
16	8192	8	256	0.146544	0.243904	0.074384	1.79038	2.09869	2.5688
16	16384	8	256	0.274416	0.480864	0.095808	1.96506	2.09784	2.55946
16	256	16	128	0.013984	0.050752	0.068448	1.78382	2.17901	2.55653
16	512	16	128	0.018304	0.050752	0.068896	1.74541	2.12547	2.56
16	1024	16	128	0.020512	0.050336	0.06896	1.75773	2.16211	2.39672
16	2048	16	128	0.032896	0.064896	0.068	1.76622	2.18272	2.5516
16	4096	16	128	0.04864	0.123072	0.069856	1.76208	2.10962	2.5608
16	8192	16	128	0.090816	0.240096	0.068752	1.76266	2.1185	2.57362
16	16384	16	128	0.16432	0.47224	0.074336	1.78822	2.13208	2.5853

LayerNorm / GroupNorm

Details

T	naive_ln	fused_ln	naive_gn	fused_gn	naive_ln_bwd	fused_ln_bwd	naive_gn_bwd	fused_gn_bwd	naive_ln_gate	fused_ln_gate	naive_rms_gate	fused_rms_gate	naive_ln_gate_bwd	fused_ln_gate_bwd	naive_rms_gate_bwd	fused_rms_gate_bwd
128	0.020032	0.00896	0.026016	0.009216	0.092192	1.72099	0.275936	1.53941	0.023712	0.092256	0.020608	0.083680	1.479696	3.176880	1.718512	3.044304
256	0.01696	0.01248	0.041344	0.01152	0.156096	1.72762	0.292752	2.6752	0.036128	0.091344	0.032448	0.081632	1.446944	3.194944	1.736816	3.056528
512	0.026112	0.081664	0.072832	0.090208	1.08317	2.90694	1.54915	3.12269	0.056800	0.092192	0.050208	0.081584	1.474624	3.198064	1.731360	3.049728
1024	0.045504	0.088032	0.136736	0.094656	1.15766	2.97059	1.59597	3.12755	0.108864	0.093216	0.094176	0.085104	1.476448	3.204160	1.709872	3.055520
2048	0.0832	0.095184	0.255872	0.102048	1.14986	2.75011	1.70464	3.12934	0.199216	0.098272	0.174560	0.090192	1.457232	3.203040	1.719728	3.066992
4096	0.156816	0.09648	0.493696	0.101792	1.1575	2.99526	2.94032	3.02125	0.383072	0.139520	0.330976	0.138096	1.483424	3.182624	1.828816	3.050880
8192	0.305664	0.185984	0.968681	0.186432	1.30707	2.9793	5.83619	3.13933	0.749792	0.269824	0.642848	0.267680	2.575712	3.097952	2.443632	3.170048
16384	0.602144	0.36176	1.92502	0.362368	1.88546	3.00792	11.6732	2.98059	1.479328	0.527456	1.268160	0.523744	5.100896	3.935104	4.853264	3.651840

Token Shift

Details

T	naive_token_shift	fused_token_shift	naive_token_shift_bwd	fused_token_shift_bwd
128	0.03024	0.031328	1.08864	2.19381
256	0.047712	0.031072	1.27306	1.81235
512	0.0984	0.033504	1.28248	1.98714
1024	0.18096	0.0536	1.256	1.88304
2048	0.346032	0.0984	1.42048	2.20317
4096	0.674208	0.1904	1.78438	2.03648
8192	1.33181	0.379296	3.53974	2.06998
16384	2.6537	0.74552	7.05104	2.52397
32768	5.2953	1.47898	14.0695	5.03008

GDN

Details

B	T	H	D	OP	fwd	fwdbwd
1	8192	96	128	chunk_gdn	1.923	7.501
2	16384	16	128	chunk_gdn	1.337	5.657
4	2048	16	128	chunk_gdn	1.121	6.055
4	4096	64	128	chunk_gdn	2.424	9.363
8	1024	8	64	chunk_gdn	0.86	6.093
8	2048	32	256	chunk_gdn	2.778	12.489

KDA

Details

B	T	H	D	OP	fwd	fwdbwd
1	8192	96	128	chunk_kda	3.112	13.966
2	16384	16	128	chunk_kda	2.167	9.423
4	2048	16	128	chunk_kda	1.011	7.943
4	4096	64	128	chunk_kda	3.971	18.098
8	1024	8	64	chunk_kda	0.975	7.124
8	2048	32	256	chunk_kda	4.495	20.487

Fused_attnres

Details

L	B	T	H	D	OP	fwd	fwdbwd
8	1	8192	1	2048	fused_attnres	0.270	3.809
8	1	32768	1	2048	fused_attnres	0.414	3.917
10	1	8192	1	4096	fused_attnres	0.324	4.048
10	1	32768	1	4096	fused_attnres	1.006	8.152
64	1	8192	1	8192	fused_attnres	3.011	26.266

[sunyi0505]

FLA Module Benchmark — Atlas 800T A3(X86)

Benchmark results for Flash Linear Attention (FLA) fused operators on Atlas 800T A3(X86).

Activation Functions

Details

B	T	D	sigmoid_fwd	sigmoid_fwdbwd	logsigmoid_fwd	logsigmoid_fwdbwd	swish_fwd	swish_fwdbwd	gelu_fwd	gelu_fwdbwd	sqrelu_fwd	sqrelu_fwdbwd	swiglu_fwd	swiglu_fwdbwd	powglu_fwd	powglu_fwdbwd
4	512	1024	0.04262	0.34546	0.06788	0.38144	0.03734	0.3437	0.0146	0.26542	0.0087	0.2467	0.06771	0.38542	0.08213	0.38735
4	512	2048	0.07342	0.68075	0.16145	0.75625	0.08368	0.68411	0.02548	0.51746	0.01466	0.48424	0.14554	0.75994	0.15749	0.75956
4	512	4096	0.13366	1.3569	0.27779	1.50844	0.15608	1.35632	0.0514	1.02254	0.03954	0.95746	0.31205	1.50162	0.27836	1.51169
4	1024	1024	0.07327	0.68132	0.16139	0.75787	0.07034	0.68427	0.02538	0.51712	0.0145	0.48278	0.13368	0.75628	0.13456	0.76216
4	1024	2048	0.16652	1.35464	0.33102	1.50124	0.16655	1.35788	0.05156	1.0229	0.03958	0.95732	0.26604	1.50964	0.26674	1.50372
4	1024	4096	0.28976	2.69033	0.62228	3.0071	0.26614	2.70124	0.1007	2.03009	0.0801	1.90118	0.57814	2.99388	0.57882	2.99582
4	2048	1024	0.1665	1.35275	0.33104	1.50126	0.15694	1.35044	0.05178	1.02326	0.03958	0.9574	0.31208	1.50171	0.2904	1.51148
4	2048	2048	0.33163	2.70262	0.64148	3.00754	0.32183	2.70114	0.10074	2.03033	0.08	1.90104	0.57813	2.97918	0.57881	2.99571
4	2048	4096	0.62239	5.31059	1.2816	5.90949	0.64219	5.32494	0.19654	3.99411	0.15887	3.77564	1.24328	5.97637	1.10773	5.92047
4	4096	1024	0.31198	2.7038	0.52953	2.99048	0.33162	2.68934	0.10055	2.02967	0.08006	1.9001	0.5543	2.98675	0.53148	2.97306
4	4096	2048	0.64222	5.26967	1.059	5.93241	0.64234	5.29	0.19684	3.98427	0.16114	3.756	1.05942	5.85039	1.28265	5.92026
4	4096	4096	1.3223	10.5798	2.11273	11.8947	1.24338	10.9353	0.38288	7.95438	0.29078	7.57626	2.48128	11.837	2.48132	12.0004
4	8192	1024	0.5781	5.30909	1.10685	5.89651	0.62256	5.2915	0.19708	3.98373	0.16114	3.75632	1.15489	5.89332	1.32187	5.90375
4	8192	2048	1.05976	10.6136	2.11245	11.8215	1.28296	10.5832	0.38284	7.9549	0.29078	7.576	2.30818	11.8316	2.30889	11.9131
4	8192	4096	1.243	17.6545	2.11154	18.9639	1.06068	17.6121	0.7514	15.9368	0.50376	15.1496	2.65125	18.8846	2.6536	19.5055

Causal Conv1D

Details

T	D	causal_conv1d_fwd	causal_conv1d_fwdbwd
256	256	0.16208	0.780900
512	256	0.15376	1.450160
1024	256	0.15796	2.179740
2048	256	0.34107	4.617440
4096	256	0.67093	8.171141
8192	256	1.3247	14.921970
16384	256	2.60389	30.752060
32768	256	5.18669	62.436619
65536	256	9.54013	120.371002
256	512	0.15436	1.412950
512	512	0.15378	2.223760
1024	512	0.30574	4.272150
2048	512	0.76724	8.591800
4096	512	1.39462	16.069780
8192	512	2.70656	31.237860
16384	512	5.25713	61.816380
32768	512	10.3468	120.237381
65536	512	18.6829	333.311615
256	1024	0.2285	2.544930
512	1024	0.53066	4.951860
1024	1024	1.04092	9.471340
2048	1024	2.47644	19.822439
4096	1024	4.76655	39.303108
8192	1024	9.38647	81.764618
16384	1024	18.2665	181.484787
32768	1024	36.6472	413.333374
65536	1024	39.745	766.867432
256	2048	0.46026	5.780900
512	2048	1.07362	11.486020
1024	2048	2.36876	23.041180
2048	2048	5.26004	49.941761
4096	2048	10.19	104.298477
8192	2048	19.1096	201.118179
16384	2048	37.7728	402.170837
32768	2048	39.5	733.571716
65536	2048	79.4147	1475.955322
256	4096	1.22638	11.554461
512	4096	2.33048	22.809330
1024	4096	4.60796	182.262344
2048	4096	10.5026	106.817017
4096	4096	20.2351	247.983246
8192	4096	38.4435	427.983398
16384	4096	39.8535	724.088928
32768	4096	79.1002	1451.673096
65536	4096	158.56	2901.156006

Cross Entropy

Details

T	naive	fused	fused_linear	naive_bwd	fused_bwd	fused_linear_bwd
128	2.40062	3.1423	73.5167	10.21	12.6815	75.7418
256	4.41135	5.88796	78.4045	16.9374	21.9013	80.8865
512	8.30542	11.2386	96.3361	28.7631	34.5546	99.2921
1024	16.2331	17.6384	140.369	54.2092	56.9704	143.202
2048	32.3341	35.019	233.915	128.834	136.116	236.796
4096	72.0124	67.6623	403.963	254.08	263.977	407.314
8192	128.718	141.754	754.499	547.5	571.965	757.811
16384	258.086	281.594	1676.76	1153.31	1222.89	1672.74

L2 Norm

Details

B	T	H	D	naive	compiled	fused	naive_bwd	compiled_bwd	fused_bwd
16	256	1	2048	0.07344	0.03532	0.28981	0.48182	0.1394	0.64614
16	512	1	2048	0.13372	0.07976	0.52926	0.78596	0.2597	1.27957
16	1024	1	2048	0.24422	0.16178	1.15482	1.77686	0.48177	2.52666
16	2048	1	2048	0.68368	0.29407	2.3081	4.46698	1.04487	4.98398
16	4096	1	2048	1.35733	0.5123	4.84245	9.00919	2.07137	9.86212
16	8192	1	2048	2.53641	0.935	9.62394	17.5997	4.06399	19.8899
16	16384	1	2048	4.85434	1.80442	19.3486	34.7248	8.08501	40.5396
16	256	2	1024	0.04646	0.03478	0.31172	0.32664	0.24338	0.6217
16	512	2	1024	0.0806	0.0479	0.65833	0.48312	0.23784	1.25412
16	1024	2	1024	0.1355	0.10618	1.15473	0.79089	0.32609	2.50258
16	2048	2	1024	0.24674	0.21186	2.4759	1.75276	0.61772	4.94248
16	4096	2	1024	0.6809	0.40258	4.8109	4.41423	1.2635	9.91559
16	8192	2	1024	1.34914	0.72776	9.89258	9.0306	2.51052	20.4831
16	16384	2	1024	2.50226	1.36338	19.7654	17.6452	4.90233	40.2144
16	256	4	512	0.03984	0.0345	0.3113	0.24854	0.14638	0.61786
16	512	4	512	0.0496	0.03608	0.63966	0.32612	0.12209	1.22278
16	1024	4	512	0.083	0.05058	1.10507	0.48072	0.17808	2.42134
16	2048	4	512	0.14217	0.10832	2.3081	0.79112	0.33198	4.6537
16	4096	4	512	0.25616	0.21424	4.8671	1.75495	0.61176	9.7154
16	8192	4	512	0.68879	0.40286	9.61788	4.39225	1.28874	19.0802
16	16384	4	512	1.34624	0.74069	18.6306	8.98464	2.5678	38.7345
16	256	8	256	0.0348	0.034	0.26456	0.20562	0.15253	0.60928
16	512	8	256	0.03442	0.03348	0.52756	0.2461	0.15056	1.20265
16	1024	8	256	0.04864	0.03164	1.24162	0.3282	0.12106	2.42387
16	2048	8	256	0.0859	0.0533	2.30802	0.48738	0.18343	4.87296
16	4096	8	256	0.14514	0.11022	4.8665	0.79522	0.33919	9.54088
16	8192	8	256	0.25046	0.21704	9.61652	1.74114	0.63764	18.9053
16	16384	8	256	0.68876	0.41428	19.7355	4.37208	1.35154	38.855
16	256	16	128	0.02086	0.03014	0.28936	0.18846	0.14997	0.60886
16	512	16	128	0.03344	0.02957	0.52753	0.20812	0.14487	1.18589
16	1024	16	128	0.03528	0.03056	1.31286	0.25323	0.14246	2.42608
16	2048	16	128	0.05076	0.03186	2.47032	0.3457	0.12648	4.8863
16	4096	16	128	0.09341	0.06364	4.80924	0.50018	0.1918	9.17559
16	8192	16	128	0.14394	0.11285	9.71326	0.79964	0.36876	18.8228
16	16384	16	128	0.25334	0.22073	19.1032	1.74704	0.70486	40.2414

LayerNorm / GroupNorm

Details

T	naive_ln	fused_ln	naive_gn	fused_gn	naive_ln_bwd	fused_ln_bwd	naive_gn_bwd	fused_gn_bwd	naive_ln_gate	fused_ln_gate	naive_rms_gate	fused_rms_gate	naive_ln_gate_bwd	fused_ln_gate_bwd	naive_rms_gate_bwd	fused_rms_gate_bwd
128	0.02548	0.156820	0.15702	0.528790	0.12303	0.34714	0.483180	0.763160	0.04574	0.03222	0.13050	0.02442	0.202020	0.24884	0.460300	0.22422
256	0.04578	0.290060	0.30080	1.313340	0.17508	0.44084	0.896380	1.287790	0.07244	0.04364	0.18788	0.03846	0.308150	0.23024	0.660230	0.20588
512	0.07638	0.622650	0.59014	2.626420	0.28892	0.72859	1.749040	2.831280	0.13846	0.07678	0.28744	0.06580	0.510480	0.35112	1.070320	0.31952
1024	0.13322	1.057170	1.16774	4.866820	0.51099	1.50578	3.435040	5.476120	0.25585	0.15308	0.47206	0.14010	0.897500	0.60070	2.258260	0.52202
2048	0.25390	2.112620	2.32184	9.174120	0.92653	2.83149	7.239680	10.840790	0.49891	0.30520	1.12814	0.27910	1.850680	1.23974	5.260710	1.07054
4096	0.49178	4.871150	4.65556	18.718460	1.95364	6.43468	14.555190	22.245289	1.09557	0.56401	2.70498	0.49798	3.986800	2.48848	11.858080	2.13882
8192	0.96017	9.687070	9.29855	39.363430	3.77312	12.77004	29.113800	46.030029	2.14504	1.08081	5.39023	0.89248	7.900110	4.90880	23.595928	4.17668
16384	1.90668	19.797819	18.60710	73.682938	7.48866	26.19804	58.519741	88.377617	4.24612	2.11378	10.71250	1.72040	15.729731	9.74247	46.837112	8.28558

Token Shift

Details

T	naive_token_shift	fused_token_shift	naive_token_shift_bwd	fused_token_shift_bwd
128	0.08065	0.04271	0.1881	0.17992
256	0.11752	0.07067	0.29015	0.18865
512	0.1923	0.15702	0.46784	0.36819
1024	0.36658	0.26693	0.90892	0.70716
2048	1.03242	0.55584	2.44902	1.52018
4096	2.11166	1.11011	5.25737	3.15448
8192	4.13322	0.99889	10.4277	3.37673
16384	8.27564	2.08444	20.8441	7.05924
32768	16.15	3.76722	41.5001	13.3548

GDN

Details

B	T	H	D	OP	fwd	fwdbwd
1	8192	96	128	chunk_gdn	1004.271	3672.884
2	16384	16	128	chunk_gdn	658.528	2470.066
4	2048	16	128	chunk_gdn	166.759	616.442
4	4096	64	128	chunk_gdn	1338.315	5067.84
8	1024	8	64	chunk_gdn	53.776	174.102
8	2048	32	256	chunk_gdn	1179.903	5462.45

KDA

Details

B	T	H	D	OP	fwd	fwdbwd
1	8192	96	128	chunk_kda	140.444	3526.229
2	16384	16	128	chunk_kda	116.530	2541.330
4	2048	16	128	chunk_kda	29.313	639.462
4	4096	64	128	chunk_kda	190.247	4685.728
8	1024	8	64	chunk_kda	9.511	281.453
8	2048	32	256	chunk_kda	180.317	2970.449

Fused_attnres

Details

L	B	T	H	D	OP	fwd	fwdbwd
8	1	8192	1	2048	fused_attnres	62.996	332.270
8	1	32768	1	2048	fused_attnres	253.667	1329.242
10	1	8192	1	4096	fused_attnres	157.186	826.988
10	1	32768	1	4096	fused_attnres	636.500	3314.368
64	1	8192	1	8192	fused_attnres	2000.703	10538.138

[zheliuyu]

News

[2026.07.18] Since flash-linear-attention natively provides triton-ascend backend support for the conv1d and gated_delta_rule kernels, ms-swift has officially started using flash-linear-attention on NPUs, maintaining the same user experience as on cuda: modelscope/ms-swift#9746

At the same time, this also marks that the work related to flash-linear-attention + triton-ascend has officially entered the high-performance optimization phase. We'll share more news soon — stay with us.

[sunyi0505]

ascend-npu-profiling skills

https://github.com/sunyi0505/flash-linear-attention/tree/sync/ascend-npu-profiling-skills/.agents/skills/ascend-npu-profiling

[zhiyuan1i]

@sunyi0505 Thanks for sharing the ascend-npu-profiling skill — this looks really useful. Would you be interested in contributing it to the repo's .agents/skills/ directory (alongside fla-kda, fla-dispatch-backends, etc.) so more people (and agents) can pick it up? Happy to help turn it into a PR if you'd like.

[sunyi0505]

@sunyi0505 Thanks for sharing the ascend-npu-profiling skill — this looks really useful. Would you be interested in contributing it to the repo's .agents/skills/ directory (alongside fla-kda, fla-dispatch-backends, etc.) so more people (and agents) can pick it up? Happy to help turn it into a PR if you'd like.

Thank you for acknowledging the practical value of this skill component. As this component is still in the iterative development phase, multiple granular details remain to be optimized and supplemented. I am prepared to submit it for repository merge work immediately after all relevant refinements are fully accomplished.

[zheliuyu]

2026.09-2026.12 Roadmap

This block is the September–December 2026 plan: close the remaining production gaps on the kernels training and serving already call, keep the companion stack (CANN / torch_npu / triton-ascend) from breaking us, make short-sequence training and prefill graph-capturable, make Transformers actually bind native FLA on NPU (Qwen3.5 / Next / Olmo Hybrid), deepen ms-swift so real packing / SFT / RL stays on that path, show end-to-end evidence in other trainers, and land a native tilelang-ascend backend. It is not a second pass over every unported research kernel.

Feedback and discussion are welcome.

Remaining Adaptation

Priority is production demand. fla/layers/ and fla/models/ stay backend-agnostic — September–December “layers / models” work is e2e tests, cache / generation, and training-framework integration, not new NPU files.

Paths that training and decode already call

• fused_recurrent GDN (decode / short prefill)
  • [GDN] feat(triton_ascend): add fused_recurrent GDN fwd kernel and dispatch #1078
  • Shared fla/ops/common/fused_recurrent.py is still a CUDA-style kernel with NPU compile pins on backward.
• fused_recurrent KDA backward (fwd only today)
• fla/ops/gated_delta_product/ dedicated chunk_deltaproduct_h / chunk_deltaproduct_o
  • Fwd/bwd of the product state do not go through chunk_delta_h.
  • Reuse GDN WY / chunk_gated_delta_rule_bwd where GDP already does
• Verify — do not blindly rewrite — ops that already inherit the common NPU chunk path:
  • fla/ops/gdn2/ (chunk_delta_h)
  • fla/ops/delta_rule/ (chunk_delta_h + chunk_o + kkt)
• fla/modules/token_shift.py NPU backend (RWKV / token-mix layers; still generic Triton, already comments on the 65535 grid cap)

Demand-gated (not the default September–December queue)

• fla/ops/based/ (native / parallel / fused_chunk) — port if a training stack asks
• fla/ops/rwkv6/, fla/ops/rwkv7/ native backends — token_shift + fused_recurrent cover a large part of the gap first
• precond_gated_delta_rule / precond_kda
• Other kernels (retention, gsa, nsa, path_attn, mesa_net, …) — community request only

High-performance Optimization

June–August already entered this phase. Frozen pytest (fwd and bwd, under NaN poisoning) stays the gate; do not loosen tolerances.

• Finish in-flight perf PRs
  • [Perf] Rewrite Ascend local_cumsum #1061 (local_cumsum)
  • [Ops] Use block matrix inversion to speed up solve_tril for the Ascend NPU backend #1145 (solve_tril block inversion)
• Re-profile GDN / KDA backward on the current companion stack before picking the next kernel — those ops still dominate fwdbwd
• GLA chunk ([Ops] Add triton-ascend backend for chunk_gla #1121) has not had an NPU perf pass
• causal_conv1d leftovers recorded in the skill case file: fwd window + extract_slice; host ReduceSum of per-chunk dw / db
• chunk_delta_h bwd dhu is still AIC-scalar dominated after tiling work
• Drop num_warps / num_stages from any shared kernel that NPU still launches (common fused_recurrent, GDP, leftover utils)
• Publish an A2 (and later A3 / A5) number table for GDN / KDA / GLA / conv / attnres

Compiler and Companion-stack Compatibility

triton-ascend 3.2.2 changes defaults under us (AutoBlockify on, auto-multi-buffer / num_stages=2, extra.cann vs extra.ascend). This is a standing September–December track, not a one-off hotfix.

Supported stack at the start of September–December: CANN 9.1.0, torch==2.9.0, torch_npu==2.9.0.post6, triton-ascend==3.2.2. Plain torch_npu==2.9.0 (no .post) is not sufficient. See INSTALL.md.

• Land the 3.2.2 / CANN 9.1.0 pin
  • [Fix] Pin NPU kernels for torch_npu=2.9.0.post6, triton-ascend=3.2.2 and cann=9.1.0 #1161
• When a later compiler lands, A/B one launch at a time and drop the pin if the default is safe
• Inductor-generated kernels (gelu / sqrelu, GRPO old_logps) stay “compat shims + skip compile”, not FLA math rewrites. New inductor numeric bugs go back to torch_npu

Graph Capture (Prefill and Training)

Backend-neutral RFC. CUDA Graphs and Ascend ACL Graph / torch.npu.NPUGraph share the same contract:

• [RFC] Graph capture support for prefill and training: GDN/KDA chunk and convolution paths #1155

Today’s varlen GDN/KDA + conv path is capture-illegal on NPU for the same reasons as CUDA: host prepare_chunk_indices, data-dependent grids, .item() / D2H, host launch chunking over NT.

• Step 1: capture-safe utils (prepare_chunk_indices / offsets) under NPUGraph — static [NT_max, 2] + sentinel, device-side searchsorted, padded cu_seqlens
• Step 2: graph-path launch sequence uses the existing 1D core-grid (task_num / num_core) instead of host-splitting over NT (a Python loop is frozen at capture time)
• Step 3: GDN + KDA chunk fwd/bwd and causal_conv1d fwd/bwd, bounded by (T_max, N_max)
• Capture / replay tests next to the CUDA ones. Do not rely on graph_task_update
• Decode (fused_recurrent, conv update-with-cache) stays out of scope for this RFC — serving stacks already own that path
• Context parallelism (fla/ops/cp/) is a later wave; [RFC] Graph capture support for prefill and training: GDN/KDA chunk and convolution paths #1155 already records what CP needs so the single-rank contract does not block it. HCCL warmup is the NPU analogue of the NCCL side-stream warmup

Decode and Inference

Chunk kernels cover training. Serving and short-prefill need the recurrent / cache path to be first-class on NPU.

• GDN fused_recurrent fwd ([GDN] feat(triton_ascend): add fused_recurrent GDN fwd kernel and dispatch #1078) + bwd if training-from-recurrent is required
• KDA fused_recurrent bwd
• causal_conv1d cache / update / compute_dh0 — already have NPU kernels; keep them correct on short chunks (see [Fix] Correct Ascend causal_conv1d short-chunk cache backward #1146, [Fix] Correct the final-state gradient for chunks shorter than W - 1 #1141, [Fix] Honor the real layouts of y and dy in compute_dh0_kernel #1139)
• Layer / model generation + FLACache e2e on NPU for GatedDeltaNet, KDA, GLA
• Document the intended serving path (ms-swift / vLLM-Ascend / native generate) so we do not grow a second decode stack

NPU CI

A2 correctness and benchmark workflows already exist. September–December extends coverage, still using ascend-gha-runners.

• Add the ops that inherit the common NPU path but are not in the A2 suite yet: tests/ops/test_simple_gla.py, tests/ops/test_delta.py / test_delta_rule, tests/ops/test_gdn2.py (names as in-tree)
• A3 / A5 job once a runner is available

Skills

• Keep .agents/skills/fla-ascend-performance/ (cases.md, TRAPS.md) current when a generalizable trap is found (compile-flag pins, lhs clobber, constexpr DMA split, int64 pointer math). Do not bundle doc dumps into perf PRs

Hardware Coverage

• Publish a one-page A2 validation matrix on the current companion stack (ops in A2 CI + GDN/KDA/GLA model smoke)
• Atlas 800T A3: core counts, UB, and the 3.2.2 A2/A3 compile split differ. Need at least a manual validation pass, then an A3 CI job if a runner is allocated
• Atlas A5: core counts, UB, and the compile path differ from both A2 and A3 (3.2.2 uses compile_on_910_95 / 950 SIMT — not A2 with a different core count). Need at least a manual validation pass, then an A5 CI job if a runner is allocated
  • SIMT unstructured access may disagree with the current constexpr DMA split / make_block_ptr assumptions; if compile or numerics fail, A/B compile_mode="simd" first
  • UB capacity differs (A2 paths hard-code 192KB in places, e.g. _UB_BYTES = 196608). Size tiles from get_soc_spec("UB_SIZE"); do not copy the A2 constant
  • Re-A/B the A2 multibuffer / AutoBlockify pins per launch on A5; do not reuse the A2 table as-is
• Do not silently assume A2 tiles / num_aicore / num_vectorcore on A3 / A5

We expect to publish validation results on additional hardware platforms in future updates.

FLA and Transformers

• Upstream Transformers: is_flash_linear_attention_available() also accepts is_torch_npu_available()
• Upstream Transformers: Qwen3.5 / Next / Olmo Hybrid conv fallback tries FLA causal_conv1d / causal_conv1d_update when Dao-AILab causal_conv1d is absent
• FLA: export recurrent_gated_delta_rule as an alias of fused_recurrent_gated_delta_rule so Transformers’ decode resolve succeeds (pairs with [GDN] feat(triton_ascend): add fused_recurrent GDN fwd kernel and dispatch #1078)
• Keep the public signatures Transformers already binds stable: chunk_gated_delta_rule + packing cu_seqlens, causal_conv1d / causal_conv1d_update
• NPU smoke: load Transformers Qwen3.5 and assert the bound chunk / recurrent / conv callables are FLA Triton-Ascend, not the torch fallback

Integration with Training Frameworks

• Deepen ms-swift past the 4B LoRA smoke: full SFT, larger Qwen3.5 / MoE, longer max_length / packing, with a public recipe
• At least one RL / GRPO path on GatedDeltaNet (and KDA if the framework exposes it)
• Keep winning the dispatch fight after MindSpeed repatch(): no stale mcore-bridge cache. The CUDA-only Transformers availability check is owned by the section above; do not re-patch it only inside ms-swift long-term
• Re-enable FLA FusedRMSNormGated on the NPU packing path now that torch.cuda.current_device() is gone; if anything else still blocks it, fix it in FLA
• Packing configs #9746 did not validate (TP / CP > 1)
• verl-project/verl: one documented NPU recipe, or an explicit “blocked on X” note
• hiyouga/LlamaFactory: same
• Graph capture ([RFC] Graph capture support for prefill and training: GDN/KDA chunk and convolution paths #1155) is the FLA-side work that makes short-sequence training in those frameworks cheap
• Keep GRPO on the eager path on NPU until inductor’s old_logps compile is safe (already skipped in TritonAscendBackend)

tilelang-ascend Support

September–December will land native tilelang-ascend support. triton-ascend stays the default production path; tilelang is a second backend on the same @dispatch / BaseBackend registry already used by the in-tree NVIDIA tilelang backends (common, KDA, DPLR, RWKV6).

• Register a tilelang_ascend backend (IS_NPU + tilelang-ascend, lazy import). Do not reuse the NVIDIA tilelang backend’s has_usable_nvcc() gate
• First wave: ops that already have a CUDA tilelang backend and an NPU triton_ascend path, so the two NPU DSLs can be A/B’d
  • fla/ops/common/ tilelang chunk / parallel-attn
  • fla/ops/kda/
  • GDN-family chunk (no CUDA tilelang in-tree yet; start from the existing triton_ascend GDN surface)
• Same correctness gate as triton-ascend (fwd and bwd, NaN poisoning). No Torch fallback on a claimed path
• Pin a validated tilelang-ascend version in INSTALL.md / extras once the first wave passes A2
• A2 CI slice for the first-wave ops (not the full model matrix)
• Document when to pick tilelang-ascend vs triton-ascend (perf, compile time, coverage). September–December does not switch the default away from triton-ascend