Skip to content

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

Open
Open
[RFC] Roadmap: NPU support for flash-linear-attention#942
Labels
ascend-npuAscend NPU (triton_ascend) relatedenhancementNew feature or request

Description

@zheliuyu
zheliuyu
opened on Jun 9, 2026

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:

We have validated this approach for NPU adaptation work:

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:

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:

Reference：

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:

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

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:

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.

Activity

  1. zheliuyu
    added
    enhancementNew feature or request
    on Jun 9, 2026

  2. sunyi0505 commented on Jul 15, 2026

    @sunyi0505
    sunyi0505
    on Jul 15, 2026
    Contributor

    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

  3. sunyi0505 commented on Jul 15, 2026

    @sunyi0505
    sunyi0505
    on Jul 15, 2026
    Contributor

    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

  4. zheliuyu commented on Jul 18, 2026

    @zheliuyu
    zheliuyu
    on Jul 18, 2026
    CollaboratorAuthor

    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.

  5. sunyi0505 commented on Jul 20, 2026

    @sunyi0505
    sunyi0505
    on Jul 20, 2026
    Contributor

    ascend-npu-profiling skills

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

  6. zhiyuan1i commented on Jul 25, 2026

    @zhiyuan1i
    zhiyuan1i
    on Jul 25, 2026
    Collaborator

    @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.

  7. sunyi0505 commented on Jul 26, 2026

    @sunyi0505
    sunyi0505
    on Jul 26, 2026
    Contributor

    @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.

  8. zheliuyu commented on Aug 25, 2026

    @zheliuyu
    zheliuyu
    on Aug 25, 2026
    CollaboratorAuthor

    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)
    • 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.

    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.

    Graph Capture (Prefill and Training)

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

    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.

    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

  9. zhudada0120 commented on Sep 21, 2026

    @zhudada0120
    zhudada0120
    on Sep 21, 2026
    Contributor

    Ascend Adaptation Plan for fla/ops Operators

    The goal is to adapt fla operators to multiple triton-ascend versions. The work will proceed layer by layer: ops first, then layers, then models. This plan is the first step. It covers the fla/ops operators (the kernel level); fla/layers and fla/models will be covered in follow-up plans.

    Work within the ops layer is split into two phases. Phase 1 completes the adaptation under triton-ascend 3.2.2 and is this list. Phase 2 performs full-version testing on top of Phase 1 and will be detailed after Phase 1 completes.

    This list contains 23 items. They come from a batch run of all tests under tests/ops with triton-ascend 3.2.2: the 23 test files that had failing cases (snapshot from 2026-09-18).

    Acceptance Criteria

    Item Requirement
    CANN 9.1.0
    torch 2.9.0
    torch_npu 2.9.0.post6
    triton-ascend 3.2.2
    Hardware Ascend A2, A3, A5

    Note: torch_npu must be the .post6 release; plain 2.9.0 without .post is not sufficient.

    With the version stack above, run pytest tests/ops/test_<op>.py on each of the A2/A3/A5 hardware platforms. All test cases must pass (0 failed, 0 error). Do not modify any assertion or tolerance in the test cases.

    TODO

    ----- List of Pending Community Tasks -----

    • Based Linear Attention operator (tests/ops/test_based.py)
    • Comba state-space attention operator (tests/ops/test_comba.py)
    • Gated Delta Product operator (chunk tests) (tests/ops/test_delta_product.py)
    • DeltaFormer operator (tests/ops/test_deltaformer.py)
    • DPLR Delta Rule operator (tests/ops/test_dplr_delta.py)
    • HGRN gated recurrent network operator (tests/ops/test_hgrn.py)
    • IPLR Delta Rule operator (tests/ops/test_iplr_delta.py)
    • Log Linear Attention operator (tests/ops/test_log_linear_attn.py)
    • Mesa-Net operator (tests/ops/test_mesa.py)
    • NSA (Native Sparse Attention) operator (tests/ops/test_nsa.py)

    Notes

    During adaptation, reusable lessons that can help other operators will be recorded in .agents/skills/fla-ascend-performance/ (cases.md, TRAPS.md).

  10. zhudada0120 commented on Sep 21, 2026

    @zhudada0120
    zhudada0120
    on Sep 21, 2026
    Contributor

    cc @zheliuyu @TecJesh

  11. zheliuyu commented on Sep 21, 2026

    @zheliuyu
    zheliuyu
    on Sep 21, 2026
    CollaboratorAuthor

    @zhudada0120 Welcome your team’s contribution. I remember you were involved in building triton-ascend itself, and looking forward to more performance optimizations from you.

    Two quick notes:

    • The AttnRes, Precond Gated Delta Rule, and Simple GLA kernels have already been migrated or are about to have PRs opened. The remaining kernels are still pending migration and are available to pick up anytime.
    • I’d recommend prioritizing the Delta Rule and Precond KDA first, then moving on to the rest.
  12. zheliuyu
    added
    ascend-npuAscend NPU (triton_ascend) related
    on Sep 21, 2026

  13. zhudada0120 commented on Sep 21, 2026

    @zhudada0120
    zhudada0120
    on Sep 21, 2026
    Contributor

    Hi @zheliuyu, thanks for the suggestion. We will prioritize Delta Rule and Precond KDA kernels and will not duplicate claims for AttnRes, Precond Gated Delta Rule, and Simple GLA kernels. Feel free to share any further suggestions anytime.

  14. zheliuyu commented on Oct 10, 2026

    @zheliuyu
    zheliuyu
    on Oct 10, 2026
    CollaboratorAuthor

    Just a reminder:

    1. Please follow the contribution principles added recently in #1347 and the related PRs.
    2. A new triton-ascend v3.6.0 release is expected around the end of October, and that code is already on main branch. If you can, please verify newly submitted PRs against this version ahead of time. Within about a week of the release, I will update the Ascend CI dependency and image in FLA to match. 🤝

    cc @qianjinqiu @Ginray @hazelduan @FightingZhen @OsirisDuan @sunyi0505 @zhudada0120 @TecJesh

  15. zheliuyu commented on Oct 26, 2026

    @zheliuyu
    zheliuyu
    on Oct 26, 2026
    CollaboratorAuthor

    Notes, 2026.10.26:

    1. FLA has been upgraded to triton-ascend v3.6.0. Some kernels show performance regressions, and we're actively working on them.

    2. Some good news: triton-ascend is considering adding FLA to their daily ci pipeline. That would mean we no longer have to manually upgrade FLA after each triton-ascend release, which should also be an effective boost to the software health of FLA's NPU side. This also explains why there have been many new triton-ascend kernel PRs recently.

    3. Validation for 950 has started. However, we haven't found any usable 950 CI resources yet. The kernel code hasn't been open-sourced yet either — I'll keep you posted as soon as there's progress. Stay tuned.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    ascend-npuAscend NPU (triton_ascend) relatedenhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions