Repository navigation
[RFC] AscendC backend for FLA: Scope, roadmap, and performance comparison #1531
Description
Proposal
This RFC proposes adding a wheel-backed AscendC backend to the project (initially demonstrated in PR #1528) and outlines the necessary performance evaluation to validate this architectural change.
1. Adding the AscendC Backend
The proposed backend integrates through the existing registration and dispatch mechanism. It provides an opt-in path to use prebuilt operators (e.g., fla_npu) for specific workloads on Ascend hardware. (Refer to PR #1528 for the initial KDA integration).
2. Scope & Performance Evaluation Plan
Before committing fully to this backend, we must evaluate its performance benefits over the existing Triton-Ascend kernels. We are preparing a comprehensive benchmark suite covering:
- Specific Operators: KDA operators adapted in PR [KDA] Add wheel-backed AscendC backend #1528, as well as other representative core operators in the repository.
Next Steps
We will post the benchmark scripts and comparison results (AscendC vs. Triton-Ascend) in this issue once ready. We aim to use this data to determine the value of integrating AscendC, and to define the future scope and roadmap for this backend.
Rationale
Why this RFC is necessary
Introducing and maintaining a dual-backend approach (AscendC + Triton) is a fairly large undertaking. It increases the CI burden and long-term maintenance overhead. To make an informed decision about the roadmap, we need concrete performance data rather than assumptions.
Key Questions to Answer
- Does the AscendC backend provide a sufficient performance gain over Triton-Ascend to justify the maintenance cost?
- If integration is necessary, what is the preferred approach: wheel-based (prebuilt) integration or source-code integration?
This issue will serve as the central tracking point for these discussions.
Related PR: #1528
Activity
zheliuyu commentedon Oct 20, 2026
Looking forward to the AscendC backend delivering better kernel performance. I can’t wait to see the benchmark data.
chenjingxiang96 commentedon Oct 20, 2026
Data collected on: 2026-10-15
Raw results:
results_a2/,results_a5/(*.jsonl, oneRESULTJSON per line)1. Baseline & Environment
git archive origin/main=f623f2b02ddd7a06b115d8061cc805275761a8b6(2026-10-15 09:52). On A5 it is built into a wheel and installed; on A2 it runs from the same source tree with a manually embedded OPP (reason: see §6). Both sides use identical source codegithub.1485827954.workers.dev/fla-org/flash-linear-attentionmain =3995fd79a49210e1f1cb0b90b3bf3df9fef63707(2026-10-15 11:16), using its built-infla/ops/kda/backends/triton_ascendandfla/ops/gated_delta_rule/backends/triton_ascendbackends2. Benchmark Scenarios & Methodology
TP splitting is modeled as heads-per-rank = total heads / TP (B and T remain unchanged per rank):
Methodology:
torch.npu.synchronize(), median of 5 runs after 2 warmups, unit: ms.msopprof --aic-metrics=BasicInfo --replay-mode=application, sum ofTask Duration(us)fromOpBasicInfo*.csv; the profile covers 3 invocations (2 warmup + 1 timed) and values are normalized to a single invocation.BNSD(the actual head-major input used by models), fla-org uses its only supported[B,T,H,D].g/betasemantics aligned on both sides: KDA passes raw gate + raw beta logits (use_gate_in_kernel=True,use_beta_sigmoid_in_kernel=True,safe_gate=True,lower_bound=-5,use_qk_l2norm_in_kernel=True,disable_recompute=True); GDN passes the already-activated log-space decay (fla_npu's fused entry does not support in-kernel gate).3. Performance Results
Wall-clock median (ms, single op / single layer, B=1, T=8192, chunk=64, K=V=128). Bold = faster side for that row.
Notes:
SyncCommon.cpp:222crash).zhiyuan1i commentedon Oct 21, 2026
Thanks for the RFC and benchmark results. I would prefer a source-based JIT integration with the NPU kernel repository as a git submodule, rather than a prebuilt kernel wheel.
Why this approach:
safe_gate/lower_bound, dtype conversions, and backward recomputation. Reviewers and coding agents need to inspect the actual implementation, not just a binary interface.Distribution requirements:
flash-linear-attention[ascendc]should include the complete source snapshot for the pinned revision, including license and notice files.torch_npu, and compiler compatibility still need explicit validation.This provides a better balance than either a binary-only dependency or copying the entire implementation into FLA. A minimal KDA JIT proof of concept covering packaging, source-version reporting, and numerical tests would help validate the approach.
cc @zheliuyu
zheliuyu commentedon Oct 22, 2026
Findings
Looking at KDA alone: the shapes in this benchmark barely overlap with what
flash-linear-attentionalready tests.After collapsing the repeated backward rows, the KDA table lists 10
(B, T, H, D)shapes. On currentmain(cd61447d),tests/ops/test_kda.pycollects to 213 cases (pytest --collect-only) and 48 distinct(B, T, H, D)shapes (varlen cases counted as(1, total T, H, D), vLLM decode cases asT=1). The two sets share 1 shape, and that shape is hit by 2 cases.flash-linear-attentiontestsThe single shared shape is
(1, 8192, 96, 128), intest_chunk_use_beta_sigmoid_in_kernel, once withallow_neg_eigvalfalse and once true. The benchmark reports 6 switches:use_gate_in_kernel,use_beta_sigmoid_in_kernel,safe_gate,lower_bound=-5,use_qk_l2norm_in_kernel, anddisable_recompute=True; these 2 cases exercise onlyuse_beta_sigmoid_in_kernel(their q/k l2norm is applied outside the kernel, not viause_qk_l2norm_in_kernel). So the one shape that is covered is still not tested on the code path this benchmark runs.To be clear:
flash-linear-attention-npudoes show a solid performance edge overflash-linear-attention'striton-ascendpath on some of these shapes. But the near-zero test overlap says three things:flash-linear-attention's 48 shapes are mostly small, generic correctness cases, while the 10 benchmark shapes come from production Kimi/GLM workloads.flash-linear-attention-npu's numerical correctness is currently not verifiable insideflash-linear-attention's test suite — the shared test backend imported by its main KDA test file is not in the public repo either.flash-linear-attention-npu's correctness is unverified as far asflash-linear-attentionis concerned, let alone protected against regressions.On the data posted so far,
flash-linear-attentionis therefore not in a position to integrateflash-linear-attention-npu(the AscendC backend proposed in this issue).Where to go from here
If the two repos are judged to be moving in the same direction
flash-linear-attention-npushould first align withflash-linear-attention's existing KDA tests — at minimum the fulltest_chunk/test_fused_recurrentpaths (including their varlen variants) and production-adjacent large shapes such as(1, 8192, 96, 128)intest_chunk_use_beta_sigmoid_in_kernel. In parallel, the 10 production shapes above, run with the benchmark's 6-switch combination, should be added as numerical correctness tests (flash-linear-attentionhas none today, and none are publicly visible on theflash-linear-attention-npuside). On top of that,flash-linear-attention-npuneeds to demonstrate a stable, reproducible performance advantage overtriton-ascendon A2/A3 and 950 — i.e. become a drop-in replacement fortriton-ascendon these high-priority cases. That is the right time to integrate.If not
Shelve the integration for now. In the meantime, we suggest
flash-linear-attention-npukeep evolving as a standalone wheel and publish its own numerical test results so correctness can be checked; contributing the 10 shapes with the 6-switch combination as a test PR toflash-linear-attentionwould also be the most direct way to give the two test suites their first real intersection. The integration discussion can be reopened any time once that exists.