Repository navigation
[Ops] Add Triton-Ascend backend for Parallax single-token decode - #1377
Closed
zhudada0120 wants to merge 2224 commits into
Closed
zhudada0120 wants to merge 2224 commits into
zhudada0120 wants to merge 2224 commits into
Conversation
…la-org#967) * Fix uninitialized rotary output on left-padded batched generation rotary_embedding_fwdbwd allocated its output with torch.empty_like, but the Triton kernel only stores rows whose rotary position o_cs is in range (m_t = (o_t>=0)&(o_t<T)&(o_cs>=0)&(o_cs<TR)). Under left-padding the per-sequence seqlen_offset is negative, so padding rows get o_cs < 0, are masked out of the store, and keep uninitialized memory. That garbage then leaks into attention on paths that do not unpad (e.g. torch SDPA): when it is NaN/Inf it propagates across the batch and RoPE models emit empty generations. (The flash-attn varlen path unpads padding rows before attention, which is why it only surfaces without flash-attn.) Fix: allocate the output with torch.zeros_like so the kernel-skipped rows are defined rather than uninitialized. Tests (tests/modules/test_rotary.py): - test_rotary_left_padding / test_rotary_partial: correctness on the negative-offset (left-padding) and partial-rotary paths, which were previously untested (existing tests cover only positive offsets); output is checked against rotary_embedding_ref. - test_rotary_left_padding_no_uninit_leak: regression guard for this bug -- seeds the allocator free list with NaN blocks the output buffer would reuse, then requires the result to be finite. Fails on the empty_like version, passes with the fix. * [Fix] Shorten rotary zeros_like comment to fit line-length budget * [Test] Cover fp32 in rotary uninitialized-leak regression test * [Test] Use float16 instead of bfloat16 in new rotary tests --------- Co-authored-by: Aviv Bick <abick@aws-cmh-slurm-1-vscode-01.cm.cluster> Co-authored-by: Yu Zhang <yzhang.cs@outlook.com>
…che path (fla-org#972) * [Perf] Fuse GDN qkv short convolution Use a single dense short-convolution call for GatedDeltaNet q/k/v projections in the no-cache path and add layer-level coverage and benchmarking. Co-authored-by: Cursor <cursoragent@cursor.com> * [Misc] Split GDN layer benchmark tooling Move the standalone GatedDeltaNet benchmark into its own PR so this branch only contains the fused qkv convolution runtime change and regression test. Co-authored-by: Cursor <cursoragent@cursor.com> * add assertion for gradients * extract condition to fn so we can mock * harden test * address comment --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Zhiyuan Li <uniartisan2017@gmail.com>
…fla-org#978) * [Ops] Add naive PyTorch reference for DeepSeek Sparse Attention (DSA) Add `fla/ops/dsa` with a pure-PyTorch DSA reference following the TileLang `examples/deepseek_v32` implementation: - `naive_dsa_indexer` runs the lightning indexer (MQA scoring with optional per-head weights and activation) and returns the causal top-k key indices. - `naive_dsa` attends over the selected keys only, with GQA, varlen (`cu_seqlens`), and precomputed-indices support. Also move the sparse-attention design-space README from `nsa/` to `dsa/` and add a high-level note on the granularity quality-speed trade-off. * [Ops] Add DSA paper and reference links to docstring * Update naive.py
…d coverage (fla-org#935) * [Test] Add chunk_size coverage for chunk kernels Cover chunk_size=16/32/64 for TTT chunk/fused/varlen paths, DPLR chunk/varlen paths, and RWKV7 chunk wrapper. Add contract tests for KDA and GDN2 unsupported chunk sizes, and extend chunk helper coverage for index, local cumsum, and Comba scalar cumsum. Validation: targeted pytest cases for modified TTT, DPLR, RWKV7, KDA, GDN2, Comba, index, and cumsum tests passed; pre-commit, compileall, header check, and git diff --check passed. * Use A.shape[-1] for wy_fast BT Avoid hardcoded BT=64 in delta, gated delta, Comba, and OJA WY helper functions so the chunk dimension follows the actual A tensor. Validation: - python -m py_compile fla/ops/delta_rule/wy_fast.py fla/ops/gated_delta_rule/wy_fast.py fla/ops/comba/wy_fast.py fla/ops/gated_oja_rule/wy_fast.py - git diff --check * [Ops] Propagate chunk_size chunk-size coverage - Plumb caller-provided chunk_size through delta, gated delta, Oja, and GLA-like chunk wrappers. - Use A.shape[-1] as the WY helper time chunk dimension and add fallback for non-64 GDN intra chunks. - Add focused dense forward/backward coverage for chunk_size=16/32/64 across affected kernels. * fix
…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.
Co-authored-by: tianjiayun <tianjiayun@bytedance.com>
…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>
…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
…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.
* [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)
zhudada0120
force-pushed
the
ops/parallax-ascend
branch
from
October 12, 2026 17:00
4a04fac to
36721ff
Compare
Contributor
Author
|
Hi @zheliuyu, CI has passed. PTAL thanks. |
Collaborator
|
Copy that. |
On NPU the mainline one-step kernel fails to compile for D=128 (BK=128): one KV block keeps two dtype tiles plus four fp32 product tiles live in UB, requiring 2118912 bits against the 192KiB budget of 1572864 bits (hivm-plan-memory "ub overflow"; the compiler's automatic multi-buffer bailout does not help). BS is a plain loop stride of the mainline kernel, so the backend relaunches the unchanged kernel with a UB-budgeted block length from the Ascend UB manager (compute_row_tile_block_size): BS stays 128 for BK <= 64 and drops to 64 for BK = 128, matching the footprint of the already-passing D=64 configuration. The mainline stays platform-neutral: only a @dispatch decorator is added, and the NPU implementation registers via @register. Also register parallel_parallax and parallax_decode_one_step in the benchmark registry (the training entry is constrained to D in [64, 128]; its D > 128 kernel still overflows UB), add parallel_parallax to the A2 benchmark CI, and register test_parallax.py in the A2 test CI. Tests: pytest tests/ops/test_parallax.py — 72 passed on Atlas A2 910B4 (CANN 9.1.0, triton-ascend 3.2.2).
zhudada0120
force-pushed
the
ops/parallax-ascend
branch
from
October 13, 2026 02:26
36721ff to
a9be267
Compare
6 of 16 tasks
sunyi0505
reviewed
Oct 15, 2026
| WINDOW_SIZE_LEFT=window_size_left, | ||
| USE_CACHE_START=cache_start is not None, | ||
| BS=_one_step_bs(BK), | ||
| num_warps=4, num_stages=2, |
Contributor
There was a problem hiding this comment.
nit: num_warps and num_stages do not take effect on the current chip. For code simplicity, it is recommended to delete them directly.
Collaborator
|
Hi @zhudada0120 — 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 (zhudada0120:ops/parallax-ascend → fla-org:main) and reference #1377 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! |
4 of 6 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
@dispatch('parallax')when running on NPU. The mainlineBS = BK = 128KV block exceeds the 192KiB UB budget, so the backend relaunches the unchanged kernel with a UB-budgetedBS(compute_row_tile_block_size).parallel_parallaxandparallax_decode_one_stepin the benchmark registry, addparallel_parallaxto the A2 benchmark CI, and registertests/ops/test_parallax.pyin the A2 test CI. The training entry is constrained to D ∈ [64, 128] since the D > 128 kernel still overflows UB on NPU.Test plan
pytest /home/zhudada/project/flash-linear-attention/tests/ops/test_parallax.py
Benchmark / NCU
python -m benchmarks.ops.run --op parallel_parallax --base '' --json benchmark_parallel_parallax.json
Breaking changes
None
Checklist