[Fix] Use unit KDA gate scale when A_log is omitted by taking-lying-flat · Pull Request #1294 · fla-org/flash-linear-attention

taking-lying-flat · 2026-09-27T00:14:36Z

Summary

With use_gate_in_kernel=True and no A_log, unbounded recurrent and fused-gate KDA paths used exp(1) as the decay scale, while chunk KDA rejected the call. Use a unit scale consistently across chunk, recurrent, fused-gate, and Triton-Ascend paths, equivalent to passing an explicit zero A_log.

The PyTorch reference and existing parameterized tests cover this optional-argument contract. Explicit A_log, bounded-gate formulas, parameter shapes, and checkpoint layouts retain their existing behavior. This is a correctness fix for callers that omit A_log; it does not claim a training-throughput improvement.

Test plan

107 passed, 8 deselected on an RTX A1000 Laptop GPU (4 GiB), PyTorch 2.12.1+cu130, and Triton 3.7.1. Backend dispatch is disabled to exercise the changed Triton kernels directly. The original tolerances and NaN-memory-poisoning fixture are retained. All pre-commit hooks and git diff --check passed.

FLA_CI_ENV=0 FLA_DISABLE_BACKEND_DISPATCH=1 python -m pytest \
  tests/ops/test_kda.py::test_chunk \
  tests/ops/test_kda.py::test_fused_recurrent_gate_in_kernel \
  tests/ops/test_kda.py::test_gate \
  -k 'not (test_chunk and (D100 or D128 or T2048 or HV8))' \
  -q --tb=short -o log_cli=false
pre-commit run --all-files
git diff --check

The selection includes the omitted-A_log regressions, explicit-parameter controls, bounded and unbounded gates, FP32/FP16/BF16, output and final-state comparisons, and supported gradients. Large chunk cases are excluded to fit local GPU memory. Two explicit-A_log BF16 chunk cases emitted dA warnings under the existing warning=True checks (relative RMS errors 0.006700 and 0.003616); no tolerances were changed. Full-model, context-parallel, and local NPU tests were not run.

The prior Ascend CI run at f85ebbbc stopped on three activation stride failures (sigmoid, swiglu, and powglu), before completing KDA coverage. This branch now includes the upstream fix from #1374. The refreshed Ascend run remains the NPU validation gate.

Benchmark / NCU (kernel changes only)

Previously recorded same-machine measurements compare CUDA fix 039ba807 with baseline 0259f80c on an RTX A1000 Laptop GPU, PyTorch 2.13.0+cu130, and Triton 3.7.1. The CUDA formulas are unchanged by the subsequent test consolidation and mainline merges.

Workload: B=1, T=73, H=1, HV=2, K=V=64, FP32, explicit zero A_log, nonzero initial state, final-state output enabled, and varlen boundaries [0, 3, 73]. Median latency uses triton.testing.do_bench(fn, warmup=100, rep=300, quantiles=[0.5, 0.2, 0.8]).

Path	Before (µs)	After (µs)
Chunk, dense	52.224	52.224
Chunk, varlen	53.248	53.248
Recurrent, dense	76.800	77.152
Recurrent, varlen	73.728	74.752
Gate forward + backward	9.216	8.192

Performance is neutral within the measurement spread. Nsight Compute --set full could not collect counters because the driver returned ERR_NVGPUCTRPERM. The table records the original CUDA-fix validation; current-branch correctness is checked separately above.

Breaking changes

No parameter shapes or checkpoint layouts change. Previously rejected chunk calls now work. Unbounded recurrent and fused-gate calls that omit A_log change from the unintended exp(1) scale to unit scale. Pass an explicit ones tensor as A_log to reproduce the former recurrent scaling.

Checklist

I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.

…pgrade CANN to 9.0.0 (fla-org#973) * refactor activations kernels * fix npu token_shift * init ops/base * skip npu compile on ops/utils/index * add ops/utils/cumsum kernel for npu * fix causal_conv1d &cross_entropy for triton-ascend=3.2.1 * Upgrade Ascend NPU stack to CANN 9.0.0 / PyTorch 2.7.1 * Upgrade NPU CI to CANN 9.0.0 / torch 2.7.1 * Expand NPU CI test coverage * remove torch.npu.synchronize in ops/based * remove fla/ops/based/backends/triton-ascend/kernels

…u to 2.7.1.post4 for npu. (fla-org#982) Co-authored-by squash-merge of fla-org#982. - Remove redundant @torch.compiler.disable from dispatched activations (centralized in @dispatch). - Add PowGLU fwd/fwdbwd/linear NPU kernels + dispatch registration (closes fla-org#973 gap). - Pin torch_npu to 2.7.1.post4 for CANN 9.0.0 / triton-ascend 3.2.1 compatibility.

…#983) * [GDN] Re-enable Triton backend for Triton >= 3.7.1 on Hopper The gated chunk_bwd_dqkwg precision issue (fla-org#640) is fixed in Triton 3.7.1, so narrow the RuntimeError guard to the affected range [3.4.0, 3.7.1) and stop skipping the Triton baseline in CI. - fla/utils: add TRITON_ABOVE_3_7_1 flag - chunk_bwd_dqkwg: only raise on Hopper + 3.4.0 <= Triton < 3.7.1 - CI: always run the dispatch-disabled Triton baseline (continue-on-error); if it fails, post a comment on the PR instead of blocking CI * [CI] Make Triton baseline a hard gate; comment on dispatch-coverage failure - Triton baseline (dispatch disabled) now fails the job when it breaks (it is the correctness reference). - Dispatch coverage (backend) is continue-on-error: when it fails, post a PR comment with the last 40 lines of the failure log instead of blocking CI. * [CI] Fix startup_failure: grant issues:write via caller, not reusable job The reusable workflow's test-ops job declared `issues: write`, which is a permission escalation that the caller (nvidia-h100.yml grants only issues: none) rejected, causing the H100 workflow to fail at startup. Move the grant to the caller: drop the job-level permissions block in the reusable workflow (it inherits the caller's token) and add `issues: write` to nvidia-h100.yml's top-level permissions so the PR-comment step can post.

…fla-org#986) OJ-style kernel debugging on self-hosted runners: - oj.yml: workflow_dispatch that reads tests/oj/env.yml (runner + pytest_args) and runs ONLY tests/oj/ on the selected backend (npu-a2 / h100 / b580), reusing each backend's existing setup. - check-no-oj.yml: PR gate that fails if tests/oj/ is still present, so the throwaway debug directory can never reach main. Usage (after this lands on main): gh workflow run oj.yml --ref <branch-with-tests/oj>

…g#985) * [Test] Add per-kernel torch baselines + tests for GDN kernels Add tests/ops/test_gdn_kernels.py: each GDN Triton kernel is tested in isolation against a pure-PyTorch baseline implementing the same math, so that every kernel can be validated independently (kernel-by-kernel). Covered kernels (forward + a subset of backward): - recompute_w_u_fwd - chunk_gated_delta_rule_fwd_kkt_solve - gdn_gate_fwd / gdn_gate_chunk_cumsum - chunk_fwd_o - chunk_gated_delta_rule_fwd_h - chunk_bwd_dv_local - gdn_gate_bwd Also add test_gdn_full_bwd: validates the full GDN backward pipeline against torch.autograd on the naive reference forward (the same ground-truth style as tests/ops/test_gdn.py), which catches kernel bugs that per-kernel baselines might miss. Three complex backward kernels (chunk_bwd_dqkwg, prepare_wy_repr_bwd, chunk_gated_delta_rule_bwd_dhu) are covered by test_gdn_full_bwd but their standalone per-kernel baselines are left as TODO (they compute intermediate/partial gradients that need a carefully-derived baseline to match the kernel contract). * [Test] Enable per-kernel backward tests for GDN Rewrite the three remaining GDN backward-kernel baselines so each kernel's partial-gradient contract is matched directly, and drop their skips: - chunk_bwd_dqkwg: autodiff of the local output + chunk-end state sub-loss; account for the kernel's RCP_LN2 gate-gradient convention. - prepare_wy_repr_bwd: autodiff through the true (I + L)^{-1} so the WY-inverse dependence is covered; L2-normalize k to keep the inverse well-conditioned. - chunk_gated_delta_rule_bwd_dhu: evaluate the kernel's reverse state-gradient recurrence directly (autodiff on the stored forward states cannot reproduce the running reverse accumulator). All 41 tests in tests/ops/test_gdn_kernels.py pass on H100 (Triton 3.7.1).

fla-org#988) Triton miscompiles prepare_wy_repr_bwd_kernel with num_warps=4 on Hopper (sm_90) for BT=64, producing incorrect dk/dbeta (and sometimes an illegal memory access). num_warps=2 is correct, and num_warps=4 works fine on Blackwell (GB200), so the bug is Hopper-specific. Restrict the Hopper autotune configs to num_warps=2 until the upstream compiler issue is resolved. Fixes fla-org#984.

…la-org#987) Squash-merge of fla-org#987. - Switch A2 NPU CI runner from linux-aarch64-a2-1 to linux-aarch64-a2-8 and expose ASCEND_RT_VISIBLE_DEVICES=0..7. - Run module + utils op tests with pytest -n 8 --dist loadfile (FLA_NPU_XDIST=1). - conftest.py maps each xdist worker (gwN) to a distinct NPU via ASCEND_RT_VISIBLE_DEVICES before torch_npu init. - Add pytest-xdist to test deps; lint-exclude tests/conftest.py (early-exec import ordering).

…#989) The GLA chunk backward autotune config pruner was a no-op: it recovered K/V from positional named_args, but the kernels are launched by keyword, so K/V resolved to None and no configs were pruned. Read K/V directly from the launch kwargs instead. Also drop configs whose tile exactly matches the tensor dim when a smaller tile is available: then the dim loop has a single iteration and Triton's software-pipelined make_block_ptr prefetches one tile past the end, causing an illegal memory access even with boundary_check. Fixes the H100 CI failure on tests/ops/test_gla.py::test_chunk_state_v_first[B2-T1024-H4-D128-torch.float16].

Add GVA (HV > H) support to the TileLang chunk_bwd_dqkwg kernel and verifier, default-enable the TileLang backend on Hopper with Triton >= 3.4, and use 2 warps for small head dims (< 64) so the kernel compiles for D=32.

Run each op test file in its own pytest process (single concurrency) so memory is released between files, avoiding the ~15GB runner OOM when the full op suite runs in one process.

… import cost (fla-org#993)

…la-org#992)

…d limit (fla-org#997)

Co-authored-by: tianjiayun <tianjiayun@bytedance.com>

…1001)

…d NPU (fla-org#1002) * [Ops] Add triton-ascend l2norm backend for Ascend NPU * [Ops] Add triton-ascend fused_norm_gate backend for Ascend NPU * [Ops] Refactor triton-ascend l2norm and fused_norm_gate for UB manager

…d UB Manager tests (fla-org#1006)

* [GDN] Add FlashQLA backend dispatch * style: fix autopep8 formatting * add dtype verifier for flashqla backend * test(gdn): add FlashQLA dtype coverage, tighten state_v_first equivalence checks to exact equality * enable auto cp by default in flashqla backend * [Utils] Wrap over-length lines in _device.py * [README] Add FlashQLA backend news entry * [GDN] Consolidate FlashQLA tests and unify flash_qla/flash_kda naming - Rename backend modules to flash_qla.py / flash_kda.py; align backend_type with package_name - Verifier: read use_gate_in_kernel from kwargs (arg order matches the op), reject deprecated transpose_state_layout - Make the verifier rejection test hardware-independent via monkeypatch - Fold dtype/GVA cases into a single parametrized test_flash_qla_chunk --------- Co-authored-by: Yu Zhang <yzhang.cs@outlook.com>

…la-org#1008)

…-org#1004) chunk_rwkv7 delegates to chunk_dplr_delta_rule but only forwarded a subset of the backend's kwargs. cp_context (and disable_recompute) exist on the backend but were absent from the wrapper, and because chunk_rwkv7 calls the backend with an explicit arg list (no **kwargs splat), any cp_context=... passed by the caller was silently swallowed into the wrapper's own **kwargs and dropped. Why it matters: with cp_context dropped, a Context-Parallel run computes each rank from zero state instead of receiving the previous rank's recurrent state - wrong output on non-first ranks, with no error raised. The dplr backend already has full CP support; only the thin wrapper was hiding it. Changes: - fla/ops/rwkv7/chunk.py: add cp_context (FLACPContext) and disable_recompute to chunk_rwkv7 and forward both to the backend. Docstring notes cp_context is mutually exclusive with initial_state/output_final_state (mirrors the backend contract). - tests/context_parallel/test_cp_dplr.py: add an `op=chunk_dplr_delta_rule` injection point to the worker + spawn helper and call it at the CP forward site. Default preserves existing behavior (no backend change). - tests/context_parallel/test_cp_rwkv7.py: new thin test reusing the dplr CP harness via the op injection, swapping in chunk_rwkv7 (r=q, w=gk) so cp_context is exercised through the wrapper. Adapter aligns scale=None since chunk_rwkv7 defaults scale=1.0 while the dplr_recurrence reference uses 1/sqrt(K). Test benches (H800, 2x GPU, bf16, T=1024 H=4 D=64): - test_cp_rwkv7.py::test_cp2_wrapper_sequence_cut PASSED; output and all of dq/dk/dv/da/db/dgk match dplr_recurrence at ratio ~0.002-0.004 (threshold 8e-3). - test_cp_dplr.py backend control: all 4 CP2 scenarios still PASS (SequenceCut/BoundaryAligned/ManyShortSequences/SafeGate) - confirms the harness op edit is a no-op for the backend path. - CP4 scenarios skipped (need >=4 GPUs). The wrapper-path error profile is byte-identical to the backend's CP2_SequenceCut, confirming chunk_rwkv7(cp_context=) reaches the same CP code path as chunk_dplr_delta_rule(cp_context=). Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Zhiyuan Li <uniartisan2017@gmail.com>

* [Ops] Add Gluon backend for AttnRes Opt-in Gluon backend for fused AttnRes via `backend='gluon'`, wired through the fla.ops.backends dispatch registry (never auto-selected; the Triton path is unchanged). Adds a generic `--backend` flag to the ops benchmark runner and parametrizes the parity test over the backend. * [Ops] Select AttnRes Gluon backend via FLA_ATTNRES_GLUON env var Switch the Gluon backend from an explicit `backend='gluon'` kwarg to the standard env-gated dispatch used by the other FLA backends (e.g. FlashQLA): `FLA_ATTNRES_GLUON=1` opt-in, off by default, with a runtime verifier (CUDA / SM80+ / D*itemsize >= 128) that otherwise falls back to Triton. Drops the `backend` argument from `fused_attnres`; the benchmark `--backend` flag now toggles the dispatch env var via a new OpConfig.backend_env map. * [Fix] Zero AttnRes Gluon dq at L=1 With a single residual source the depth softmax is constant, so the logit gradient vanishes analytically. Branch `_dv_tile` on the `L` constexpr to set ds to exactly zero at L=1 instead of relying on two fp32 reductions cancelling bitwise, which left a spurious ~2e-3 dq (naive reference is exactly 0). * [Fix] Convert lambda assignments to def in AttnRes Gluon smem prune (autopep8) * [Ops] Trim AttnRes Gluon module docstring to a high-level summary * [Ops] Polish AttnRes Gluon docstrings and comments * [Test] Keep test_attnres backend-agnostic; select Gluon via FLA_ATTNRES_GLUON * [Ops] Rewrite AttnRes Gluon bwd with mbarrier double buffering Stage the V tiles through an mbarrier-driven ring in the backward, gate the backend on gluon availability (`package_name = triton.experimental.gluon`, so Triton < 3.5 falls back cleanly), and normalise loop indices to `i_l`. Full parity (15/15, fwd + bwd) reconfirmed on GB200 / Triton 3.6.0. * [Ops] Revert AttnRes Gluon bwd to wait_group double buffering The mbarrier per-slot pipeline is indistinguishable from the simpler commit_group / wait_group scheme in benchmarks (this op is bandwidth-bound; large shapes saturate HBM and small-shape deltas sit inside measurement noise), so drop the extra complexity. Keep the block-first constexpr order and the `package_name` availability gate. Parity 15/15 reconfirmed on GB200.

Co-authored-by: Yu Zhang <yzhang.cs@outlook.com>

)

…ed (fla-org#1014) Models whose layers unconditionally call into flash-attn (BitAttention, full Attention, MLA, Samba attention layers, DeltaFormer, MoBA, NSA sliding window, YOCO swa) currently fail with ImportError in environments where flash-attn can't be installed (e.g. newer PyTorch images without matching flash-attn wheels). Skip them instead, mirroring the existing module-level importorskip pattern used for mamba_ssm: - module-level pytest.importorskip for fully flash-attn-dependent suites (bitnet, transformer, mla, samba, deltaformer, moba, nsa model tests, and the deltaformer op test) - per-configuration skips for partially dependent suites (YOCO swa rows, NSA op rows with window_size > 0) Co-authored-by: lizhiyuan <lizhiyuan@moonshot.ai>

## Summary This PR adds a triton-ascend backend for Gated Delta Net (GDN), following FLA's existing sub-operator dispatch pattern (same as `cumsum`, `l2norm`, and other NPU backends). On NPU, `chunk_gated_delta_rule` and its sub-ops are routed through `@dispatch`-decorated public APIs to NPU-specific Triton implementations. Results: - GDN forward and backward run on Ascend NPU via triton-ascend kernels. - `tests/ops/test_gdn_kernels.py` and `tests/ops/test_gdn.py` pass on NPU CI. --- ### Motivation GDN is a core FLA operator, previously available only on CUDA Triton. Ascend NPU needs a separate backend because: - No Triton autotune / `num_warps` - **UB limits and `ASCEND_MAX_GRID_DIM` grid caps** - Compiler-specific constraints (e.g. HV==1 dot paths, identity-gate workaround) - Numerical stability requirements (fp32 accumulation, etc.) This ports the GDN pipeline to NPU while keeping the CUDA default unchanged (`FLA_DISABLE_BACKEND_DISPATCH=1` preserves existing behavior). --- ## File Changes ### Added — NPU kernels (`fla/ops/**/backends/triton_ascend/`) | File | Role | Why | |------|------|-----| | `fla/ops/gated_delta_rule/backends/triton_ascend/__init__.py` | Register `TritonAscendGDNBackend` | Route GDN-specific ops on `IS_NPU` | | `fla/ops/gated_delta_rule/backends/triton_ascend/gate.py` | `gdn_gate_fwd/chunk_cumsum/bwd` NPU kernels | Gate path cannot use CUDA autotune on NPU | | `fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py` | `recompute_w_u_fwd`, `prepare_wy_repr_bwd` NPU kernels | WY bwd split into stages for UB; fwd uses static tiles | | `fla/ops/gated_delta_rule/backends/triton_ascend/chunk_fwd.py` | `chunk_gated_delta_rule_fwd_intra_npu` | NPU always uses unfused `kkt + solve_tril` (fused kernel exceeds UB) | | `fla/ops/common/backends/triton_ascend/__init__.py` | Register `TritonAscendCommonBackend` | Route shared chunk ops in the GDN pipeline | | `fla/ops/common/backends/triton_ascend/chunk_scaled_dot_kkt.py` | `chunk_scaled_dot_kkt_fwd_npu` | BC-subblock KKT lowers UB peak; grid chunking for large shapes | | `fla/ops/common/backends/triton_ascend/chunk_delta_h.py` | `fwd_h` / `bwd_dhu` NPU kernels | Recurrence with `V_OFFSET`/`NH_OFFSET` grid chunking | | `fla/ops/common/backends/triton_ascend/chunk_o.py` | `chunk_fwd_o`, `chunk_bwd_dv_local`, `chunk_bwd_dqkwg` | Largest GDN submodule; HV==1 full-BT paths and identity-gate host workaround | | `fla/ops/common/backends/triton_ascend/gate.py` | `fused_beta_sigmoid_fwd/bwd` NPU kernels | Shared beta-sigmoid on the GDN gate path | | `fla/ops/utils/backends/triton_ascend/solve_tril.py` | `solve_tril_npu` | Triangular solve without CUDA TMA; required by intra-chunk pipeline | New files use the **same basename** as upstream defaults (`chunk_o.py`, `wy_fast.py`, etc.), matching existing `triton_ascend` conventions. Host wrappers use the `*_npu` suffix; internal Triton kernels use `*_kernel_npu` where applicable. --- ### Modified — Dispatch wiring (default public API layer) Minimal changes so the NPU backend is selected at runtime: | File | Change | Why | |------|--------|-----| | `fla/ops/gated_delta_rule/backends/__init__.py` | Register `TritonAscendGDNBackend` (priority 0) alongside `FlashQLABackend` | Required for NPU routing | | `fla/ops/common/backends/__init__.py` | Register `TritonAscendCommonBackend` | Shared chunk ops routable | | `fla/ops/utils/backends/__init__.py` | Register `TritonAscendUtilsBackend` | `solve_tril` / cumsum routable | | `fla/ops/gated_delta_rule/chunk_fwd.py` | `@dispatch('gated_delta_rule')` on `chunk_gated_delta_rule_fwd_intra` | NPU can override intra path | | `fla/ops/gated_delta_rule/gate.py` | `@dispatch` on gate APIs | Gate ops routable | | `fla/ops/gated_delta_rule/wy_fast.py` | `@dispatch` on `recompute_w_u_fwd`, `prepare_wy_repr_bwd` | WY ops routable | | `fla/ops/common/chunk_delta_h.py` | `@dispatch('common')` on `fwd_h`, `bwd_dhu` | Core recurrence ops routable | | `fla/ops/common/chunk_o.py` | `@dispatch` on `chunk_fwd_o`, `chunk_bwd_dv_local`, `chunk_bwd_dqkwg` | Output / bwd-local ops routable | | `fla/ops/common/gate.py` | `@dispatch` on `fused_beta_sigmoid_*` | Beta-sigmoid routable | | `fla/ops/common/chunk_scaled_dot_kkt.py` | `@dispatch('common')` on `chunk_scaled_dot_kkt_fwd` | Shared op; avoids per-op `*_npu` imports | | `fla/ops/utils/solve_tril.py` | `@dispatch('utils')` on `solve_tril` | Same; benefits delta_rule, path_attn, etc. on NPU | **`chunk_fwd_npu` routing:** `fla/ops/gated_delta_rule/backends/triton_ascend/chunk_fwd.py` now calls public APIs (`chunk_scaled_dot_kkt_fwd`, `solve_tril`, `recompute_w_u_fwd`) instead of importing `*_npu` helpers directly. NPU routing stays in the dispatch layer, aligned with the CUDA default path. --- ### Modified — Tests & CI | File | Change | Why | |------|--------|-----| | `tests/ops/test_gdn.py` | NPU-aware GQA skip (`IS_NPU`); merge with upstream FlashQLA | NPU has no TileLang backend; GQA cases should not be skipped on NPU | | `.github/workflows/ascend-a2-ci.yml` | Add `tests/ops/test_gdn.py` to CI | End-to-end GDN coverage on Ascend A2, not only per-kernel tests | --- ### Modified — Bundled fix | File | Change | Why | |------|--------|-----| | `fla/modules/backends/triton_ascend/l2norm.py` | `rstd.reshape(-1)` + pre-launch assert | Pre-existing NPU l2norm bug found during GDN CI; small, low-risk fix | --- ## Architecture ``` chunk_gated_delta_rule() [@dispatch gated_delta_rule] └─ sub-ops each dispatch separately: chunk_gated_delta_rule_fwd_intra → TritonAscendGDNBackend ├─ chunk_scaled_dot_kkt_fwd → TritonAscendCommonBackend ├─ solve_tril → TritonAscendUtilsBackend └─ recompute_w_u_fwd → TritonAscendGDNBackend chunk_gated_delta_rule_fwd_h → TritonAscendCommonBackend chunk_fwd_o / chunk_bwd_* → TritonAscendCommonBackend gdn_gate_* / prepare_wy_repr_bwd → TritonAscendGDNBackend ``` - CUDA: `FLA_DISABLE_BACKEND_DISPATCH=1` keeps the original Triton defaults. - NPU: `TritonAscend*Backend` (priority 0) on `IS_NPU`. --- ## Key NPU adaptations (vs CUDA default) | Adaptation | Reason | |------------|--------| | `ascend_ub_manager` static tiles | No autotune on NPU | | Grid-axis chunking (`*_OFFSET` launch) | `ASCEND_MAX_GRID_DIM` hard limit | | Always unfused `kkt + solve_tril` on NPU | Fused kernel exceeds UB | | Multi-stage `prepare_wy_repr_bwd` | Monolithic bwd kernel exceeds UB | | `chunk_o` HV==1 full-BT path | BC-split `tl.dot` issues on Ascend | | Identity gate when `USE_G=False` | Host `torch.zeros`; compute stays in Triton | | `allow_tf32=False` and dot operand reordering | NPU precision and compiler stability | ## test ### Atlas 800T A3(x86) ```bash pytest tests/ops/test_gdn.py tests/ops_gdn_kernels.py ``` <img width="862" height="33" alt="image" src="https://github.com/user-attachments/assets/673b97bb-ba09-473a-8fa9-253fe4837cf9" /> ### Atlas 800T A2(arm) NPU CI: https://github.com/fla-org/flash-linear-attention/actions/runs/29007393672/job/86082214593?pr=1011 Co-authored-by: Zhiyuan Li <uniartisan2017@gmail.com>

… support (fla-org#1019) ## Summary Adds **faceted visualization** to the FLA benchmark workflow and aligns module/ops benchmark output and device handling, making results easier to read and compare across commits and hardware backends (including Ascend NPU). ### 1. New faceted visualization module (`benchmarks/visualize.py`) - Supports two input types: - **Triton `perf_report` CSV** (module benchmark output) - **`benchmark_results.json`** (cross-commit comparison from `scripts/run_benchmark_compare.py`) - For multi-dimensional shape sweeps (B/T/H/D/L), automatically picks the x-axis and splits remaining dimensions into **facet subplots** - Separates Forward / Backward providers and generates `*_faceted.png` plus an HTML gallery (`results.html`) - Comparison mode produces a **percent-change bar chart** and a **base vs head latency chart** - CLI entry point: `python -m benchmarks.visualize perf-report|comparison` ### 2. New shared module benchmark runner (`benchmarks/modules/runner.py`) - Introduces `run_module_benchmark()`: runs Triton `perf_report` and optionally auto-generates faceted plots - Auto-infers `save_dir` (e.g. `benchmark_layernorm.py` → `./layernorm_benchmark`) - Auto-detects `plot_name` and calls `visualize_perf_report_csv` after the benchmark finishes ### 3. Updated 6 module benchmark scripts These scripts now use `run_module_benchmark()` instead of calling `benchmark.run()` directly, and auto-generate visualizations after each run: - `benchmark_activations.py` - `benchmark_conv.py` - `benchmark_cross_entropy.py` - `benchmark_l2norm.py` - `benchmark_layernorm.py` - `benchmark_tokenshift.py` ### 4. Multi-device support for ops benchmarks (`benchmarks/ops/run.py`) - Replaces hardcoded `'cuda'` with `fla.utils.device_name` - Adds `_device_synchronize()` for consistent device sync during warmup and timing - Extends `_get_machine_info()` for NPU: records `device_platform`, NPU device name, and memory - Adds `_format_machine_line()` for consistent machine info output (CUDA / NPU / PyTorch version) ### 5. Ops verify updates (`benchmarks/ops/verify.py`) - Profiler selects `ProfilerActivity` based on device type (CUDA uses `self_cuda_time_total`, otherwise CPU) - Input generation and synchronization aligned with `run.py` for non-CUDA devices ### 6. Cross-commit comparison script enhancement (`scripts/run_benchmark_compare.py`) - Adds `--plot` flag to auto-generate visualization plots after comparison completes ### 7. Dependency update (`pyproject.toml`) - Adds `pandas` to the `benchmark` optional extra (required for CSV parsing in visualization) --- ## Usage Examples ```bash # Module benchmark (auto-generates plots) python benchmarks/modules/benchmark_activations.py ``` <img width="959" height="815" alt="image" src="https://github.com/user-attachments/assets/f7d3fddb-b2c2-439b-845d-916775eee1b4" /> ```shell # Manually visualize an existing CSV python -m benchmarks.visualize perf-report benchmarks/modules/activation_benchmark/activation_performance.csv # Cross-commit comparison + visualization python scripts/run_benchmark_compare.py --base HEAD --output results.json --plot plots/ ```

* [Refactor] Centralize dispatch and operation backend registration * [Refactor] Simplify backend registration and dependency coverage * [Fix] Keep public imports off deprecated backend entry point * [Refactor] Align shared backend kernels with FLA conventions * [Refactor] Align TileLang attention with kernel conventions * [Test] Strengthen backend dispatch contract coverage

…1348) Preserve existing normalization imports while grouping implementations and Ascend backends under fla.modules.norm.

…scale

Move convolution backend ownership under fla.modules.conv while preserving public imports and dispatch behavior.

…ent paths run on Ascend NPU (fla-org#1307) * [Ascend] Add the simple GLA NPU backend for the Retention paths Retention's chunk, fused_chunk, parallel and fused_recurrent paths all wrap simple GLA. On the pinned Ascend stack (CANN 9.1.0 / torch_npu 2.9.0 / triton-ascend 3.2.2) the unpatched tree cannot complete tests/ops/test_retention.py: fused_chunk_bwd_kernel faults with an unaligned vector UB access (AICore 507015) and parallel_simple_gla_fwd_kernel crashes bishengir-compile, after which the poisoned device aborts the run. Add a TritonAscendSimpleGLABackend that takes over the simple_gla entry points on NPU. fused_chunk and parallel are realized through the chunk decomposition and output_attentions=True is materialized by a dedicated NPU kernel behind @input_guard. fused_recurrent uses an NPU copy of the shared kernels with pinned launch parameters, and chunk_bwd_dv gains a common dispatch entry with an offset-aware kernel; both split launches above the Ascend per-launch block budget. The shared kernels are untouched and FLA_DISABLE_BACKEND_DISPATCH=1 keeps the previous implementations as fallback. Verified on an Ascend 910_9382 card (CANN 9.1.0, torch_npu 2.9.0.post6, triton-ascend 3.2.2): - tests/ops/test_retention.py: 29 passed (unpatched: aborts) - tests/ops/test_simple_gla.py: 50 passed / 4 skipped (unpatched: 13 failed) - tests/ops/test_gla.py: 34 passed, the same 2 pre-existing failures as unpatched - tests/ops/test_gsa.py: the failure-set delta is the known flaky test_fused_recurrent cases * [Ascend] Fix simple GLA backend kernels Preserve the K tile across consecutive dot operations and remove unsupported launch parameters from Ascend autotune configs. --------- Co-authored-by: quipiter <qsj18912269550@163.com> Co-authored-by: Zhiyuan Li <uniartisan2017@gmail.com>

* [Refactor] Flatten normalization module layout * [Refactor] Keep existing normalization imports stable * [Refactor] Limit norm compatibility to public imports * [Refactor] Align normalization formatting with contribution guidelines * [Refactor] Normalize remaining normalization call formatting * [Refactor] Align Ascend normalization adapter signatures

* [Refactor] Separate causal and long convolution layouts * [Refactor] Remove redundant convolution forwarding package * [Refactor] Use causal_conv1d and ops consistently * [Refactor] Import convolution entries through their package * [Style] Leave autograd contexts unannotated * [Refactor] Expose convolution kernels through the package * [Refactor] Keep convolution backend implementations in sibling files * [Test] Patch the convolution module that owns backend selection * [Refactor] Align convolution context parallel naming with existing CP modules * [Refactor] Group convolution implementations in the backend package * [Refactor] Keep convolution kernels out of package exports * [Test] Import convolution kernels from their implementation module

…nts (fla-org#1223)

…1345) * [Refactor] Centralize dispatch and operation backend registration * [Refactor] Simplify backend registration and dependency coverage * [Fix] Keep public imports off deprecated backend entry point * [Refactor] Align shared backend kernels with FLA conventions * [Refactor] Align TileLang attention with kernel conventions * [Refactor] Group normalization implementations and backends * [Refactor] Keep convolution backends with their operation * [Refactor] Give remaining modules their own backends * [Docs] Simplify repository skills and contribution guidance * [Docs] Clarify documentation and test conventions [skip test]

Preserve legacy 0.5.2 module imports through centralized aliases and announce 0.6.1 deprecation. Remove forwarding files and align BitLinear kernels and autograd with module conventions.

Use named Triton JIT forwarding functions so Inductor-generated callers can resolve softplus and softplus2. Preserve the existing arithmetic and platform selection.

Extract quantized normalization kernels and autograd into norm, align kernel conventions, and use portable Triton rounding. Preserve BitLinear entry points and existing numerical tolerances.

Keep GRPO out of the eager compatibility import loop and remove the blanket legacy-import warning. Drop obsolete warning assertions while retaining public import checks.

* [Fix] Honor strides in vector activations * [Test] Integrate vector strides into activation coverage --------- Co-authored-by: Xucheng Zhou <209601943+adenzhou1350@users.noreply.github.com> Co-authored-by: Yu Zhang <yzhang.cs@outlook.com>

…1295) Accumulate FP32 bias-gradient partials in the gate backward kernels before reducing them in the wrappers, while storing input gradients directly in their original dtype.

Accept the dispatcher arguments in the KDA TileLang adapter and verifier, and keep graph capture on Triton. Cover eager routing and graph rejection with focused backend tests.

)

Skip unused weight-gradient computation and activation recomputation when SwiGLU linear inputs are frozen. Cover partial gradients and optional bias in the existing activation tests. Validated 81 affected tests on H200, 144 operator parity cases, six AMP GatedMLP parity cases, and scoped lint checks.

…1303)

* [GDN] Add configurable output gate activation * [Test] Add GatedDeltaNet output gate coverage * [CI] Add smoke tests (fla-org#1226) * [GDN] Consolidate output gate tests and verify numerical behavior * [GDN] Align output gate tests with contribution conventions * [GDN] Apply contribution style to output gate integration * [GDN] Remove redundant output gate activation state * [GDN] Keep output gate tests focused on configuration wiring * [GDN] Compact output gate validation and documentation * [GDN] Drop unrelated model test changes --------- Co-authored-by: Zhiyuan Li <uniartisan2017@gmail.com> Co-authored-by: Yu Zhang <yzhang.cs@outlook.com>

…s in test selection (fla-org#1374)

…scale

zhiyuan1i · 2026-10-15T14:55:56Z

Hi @taking-lying-flat — sorry, this PR was auto-closed on Oct 15 (UTC) by a maintainer-side incident (a mistaken force-push rewrote main's history), and we are unable to reopen it.

Could you open a new PR from the same branch (taking-lying-flat:fix/kda-default-gate-scale → fla-org:main) and reference #1294 for the previous discussion? Your commits are unchanged, so it takes only a minute. Feel free to assign me (@zhiyuan1i) as the reviewer and I will prioritize it.

Sorry again, and thank you for contributing to FLA!

zheliuyu and others added 30 commits June 26, 2026 10:55

[TileLang] Support GVA in chunk_bwd_dqkwg (fla-org#975)

6afd197

Add GVA (HV > H) support to the TileLang chunk_bwd_dqkwg kernel and verifier, default-enable the TileLang backend on Hopper with Triton >= 3.4, and use 2 warps for small head dims (< 64) so the kernel compiles for D=32.

[CI] Run op tests one file per process to avoid runner OOM (fla-org#991)

30d2a52

Run each op test file in its own pytest process (single concurrency) so memory is released between files, avoiding the ~15GB runner OOM when the full op suite runs in one process.

[Fix] Make T a runtime arg in path_attn cumprod_householder_bwd

a074cc4

[Fix] Probe TileLang availability via find_spec to avoid per-dispatch…

a6cfa08

… import cost (fla-org#993)

[Conv][Fix] Chunk Ascend causal_conv1d launches for 65535 grid limit (f…

e0661d0

…la-org#992)

[Perf] optimize activation for npu (fla-org#995)

e1355bb

[Perf] Cache find_spec lookups in backend dispatch (fla-org#996)

2e6aee2

[Conv][Fix] Chunk triton-ascend causal_conv1d silu/add for Ascend gri…

c50c127

…d limit (fla-org#997)

Restrict Blackwell gated delta bwd autotune (fla-org#1000)

ac6c648

Co-authored-by: tianjiayun <tianjiayun@bytedance.com>

[Fix] Add Ascend UB management and fix NPU benchmark issues (fla-org#…

a954753

…1001)

[Fix] drop deprecated tl.extra.cuda.libdevice prefix (fla-org#1005)

adb8f38

[Fix] [Test] Remove unused GeLU/SquaredReLU NPU kernels and add Ascen…

6bd9069

…d UB Manager tests (fla-org#1006)

[Fix] avoid SymInt stride reads in activations under torch.compile (f…

bce39a4

…la-org#1008)

[Fix] Correct Mamba-3 decay parameterization (fla-org#1012)

9b04758

Co-authored-by: Yu Zhang <yzhang.cs@outlook.com>

Add get_max_length to FLA cache layer for transformers 5.x (fla-org#1009

953e435

)

yzhangcs and others added 26 commits October 9, 2026 09:10

[Refactor] Group normalization implementations and backends (fla-org#…

d21ece4

…1348) Preserve existing normalization imports while grouping implementations and Ascend backends under fla.modules.norm.

Merge remote-tracking branch 'origin/main' into fix/kda-default-gate-…

abe3b52

…scale

[Fix] Cover omitted KDA gate scale

05cf28a

[Refactor] Keep convolution backends with their operation (fla-org#1349)

1b79e83

Move convolution backend ownership under fla.modules.conv while preserving public imports and dispatch behavior.

[Refactor] Register module backend implementations at their entry poi…

b3d7bc9

…nts (fla-org#1223)

[Refactor] Consolidate legacy module imports (fla-org#1360)

a788006

Preserve legacy 0.5.2 module imports through centralized aliases and announce 0.6.1 deprecation. Remove forwarding files and align BitLinear kernels and autograd with module conventions.

[Fix] Preserve softplus names in compiled Triton kernels (fla-org#1353)

51e80c5

Use named Triton JIT forwarding functions so Inductor-generated callers can resolve softplus and softplus2. Preserve the existing arithmetic and platform selection.

[Fix] Accept local loss tensors in LinearLossParallel (fla-org#1362)

b46f774

[Refactor] Move quantized normalization into norm (fla-org#1366)

c86f6d9

Extract quantized normalization kernels and autograd into norm, align kernel conventions, and use portable Triton rounding. Preserve BitLinear entry points and existing numerical tolerances.

[Fix] Avoid eager GRPO imports and spurious warnings (fla-org#1364)

f3bec72

Keep GRPO out of the eager compatibility import loop and remove the blanket legacy-import warning. Drop obsolete warning assertions while retaining public import checks.

[Fix] Reduce GDN/KDA gate bias gradients before downcasting (fla-org#…

3a59c2a

…1295) Accumulate FP32 bias-gradient partials in the gate backward kernels before reducing them in the wrappers, while storing input gradients directly in their original dtype.

[KDA] Restore TileLang dispatch for the fused WY backward (fla-org#1361)

f282893

Accept the dispatcher arguments in the KDA TileLang adapter and verifier, and keep graph capture on Triton. Cover eager routing and graph rejection with focused backend tests.

[Fix] Accept only the modes LightNetAttention can dispatch (fla-org#1305

99f864d

)

[Fix] Forward the configured epsilon in the bit-linear norm (fla-org#…

4a412ed

…1303)

[Fix] Resolve KDA gate scale test conflict with main

6820ab5

[Test] Consolidate KDA default gate scale coverage

f85ebbb

[Fix] Honor 1D strides in Ascend activations and follow module aliase…

1fe52b5

…s in test selection (fla-org#1374)

Merge remote-tracking branch 'origin/main' into fix/kda-default-gate-…

ccea674

…scale

zhiyuan1i closed this Oct 15, 2026

zhiyuan1i force-pushed the main branch from f68fabe to 81ada54 Compare October 15, 2026 13:00

yzhangcs mentioned this pull request Oct 16, 2026

[Fix] Use unit KDA gate scale when A_log is omitted #1442

Draft

4 of 5 tasks

Navigation Menu

Repository navigation

[Fix] Use unit KDA gate scale when A_log is omitted - #1294

[Fix] Use unit KDA gate scale when A_log is omitted#1294
taking-lying-flat wants to merge 2231 commits into
fla-org:mainfrom
taking-lying-flat:fix/kda-default-gate-scale

taking-lying-flat commented Sep 27, 2026 •

edited

Loading

Uh oh!

zhiyuan1i commented Oct 15, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

19 participants

Conversation

taking-lying-flat commented Sep 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Benchmark / NCU (kernel changes only)

Breaking changes

Checklist

Uh oh!

zhiyuan1i commented Oct 15, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

19 participants

taking-lying-flat commented Sep 27, 2026 •

edited

Loading