Repository navigation
[Ops] Add Triton-Ascend backend for gated Oja rule fused recurrent - #1372
Closed
zhudada0120 wants to merge 2216 commits into
Closed
zhudada0120 wants to merge 2216 commits into
zhudada0120 wants to merge 2216 commits into
Conversation
…g#964) The per-token base offset `(bos + i_t) * HQ*K` and the dq output offset `(i_v*all + bos + i_t) * HQ*K` overflow int32 once T*HQ*D approaches 2^31 (e.g. T=64K, HQ=128, D=256), triggering illegal memory accesses in the backward pass. Cast bos/eos to int64 at assignment across the selected (topk, fwd, bwd_dq) and compression (fwd, bwd_dq, bwd_dkv) kernels, while keeping sequence lengths int32 for block-pointer shapes, matching the convention already used in bwd_dkv and the chunk kernels (kda/gdn2/cp).
…autotune (fla-org#965) * [NSA] Add test cases * [NSA] Expand parallel_nsa_fwd autotune for large GQA groups At large G (HQ/H) the forward kernel was capped at num_warps<=4 with no num_stages search and G absent from the autotune key, leaving G=128 badly under-occupied (T=64K,HQ=128,D=256 fwd: ~99ms). Search num_warps up to 8 and num_stages in 1..3, and add G to the key so each group size specializes. Same shape: ~21ms (4.6x). Gate unchanged (47 passed). * [NSA] Expand parallel_nsa_bwd_dq autotune for large GQA groups Same occupancy fix as the forward: search num_warps up to 8 and num_stages in 1..3, and add G to the autotune key so each group size specializes. Backward latency on large shapes (T>=16K, G>=64) improves ~1.15-1.25x. Gate unchanged (47 passed). * [Attn] Use int64 addressing to avoid int32 overflow at large T The fwd/bwd_dq/bwd_dkv kernels addressed q/k/v/o with int32 (bos*HQ+...)*K offsets, and the bwd preprocess used an int32 program_id*V, both overflowing at large T*HQ*D (>= 2^31). Cast bos/eos to int64 and the preprocess program_id to int64, keeping sequence lengths int32 for block-pointer shapes (matching the chunk-kernel convention). * [PaTH] Use int64 addressing across remaining kernels for long sequences Follow-up to fla-org#769: the remaining parallel/intra/preprocess/transform kernels still used int32 bos/eos offsets that overflow at large T*HQ*D. Cast bos/eos to int64 at assignment, keeping sequence lengths int32 for block-pointer shapes. * [Log Linear Attention] Use int64 addressing to avoid int32 overflow at large T Cast bos/eos to int64 in the chunk kernels, keeping sequence lengths int32 for block-pointer shapes. * [Attn] Wall: use int64 addressing to avoid int32 overflow at large T Cast bos/eos (and the decode kernel's dense bases) to int64, keeping sequence lengths int32 for block-pointer shapes. * [Log Linear Attention] Wrap a >127 line to satisfy E501
…#966) Run one program per BT-query tile instead of per query: the [BT, S] block-index load and the count/scatter atomics become coalesced vectorized ops. Varlen finds each query's global block base with a vectorized in-kernel binary search over cu_seqlens into chunk_offsets, dropping the host-side token_indices array. Output contract is unchanged; tests/ops/utils/test_csr.py passes as-is.
* [CI] Pin nvcc 13.0 for PyTorch 2.12 H100 jobs PyTorch 2.12 dropped cu128 wheels (CUDA 13.0 is the new stable default), so the pytorch_2_12 conda env is now built on CUDA 13.0. The setup action still defaulted nvcc_toolkit_version to 12.8.1, so conda could not solve the env (cuda-nvcc 12.8.93 conflicts with the pre-installed cuda-nvcc-dev 13.0.88). Pin the 2.12 test and benchmark jobs to nvcc 13.0.1 to match the env. * [CI] Drop auto-installation of NVCC toolkit in CI workflows The H100 runners have manually maintained conda environments with pre-installed CUDA/PyTorch stacks. Let CI stop trying to install nvcc via conda, which currently conflicts with the env and blocks all H100 PyTorch 2.12 jobs. - Remove nvcc_toolkit_version input from setup-ci-env action - Drop conda install cuda-nvcc in setup-ci-env and benchmark workflow - Remove nvcc_toolkit_version passing in nvidia-h100.yml and reusable-ci-tests.yml - Remove unused pytorch_version/pytorch_cuda_version from reusable-ci-benchmarks.yml - Always uninstall old flash-linear-attention before installing the PR version --------- Co-authored-by: lizhiyuan <lizhiyuan@moonshot.ai>
…-org#970) Borrow five executable guardrails from BBuf/KDA-Pilot's diffusion kernel rules, sharpening the loop's principles into hard conditions: - symmetric numerics: ban one-sided allow_tf32 / accumulator-precision relaxation in the anti-reward-hacking table, with a matching TRAPS entry - no-go bar: a first candidate losing is not a no-go; require baseline + reasoned attempt + gate + bench + a named bound - profiler-explains-result as a promotion gate; unaccountable speedup is a silent-skip suspect - full final-claim stats block (median/mean/std/p10/p90/geomean + idle-clock evidence) cross-linked to fla-nvidia-performance - dispatch.md template for Phase 3 shape specialization
) * [Fix] Scale L2Wrap logits gradient by the upstream gradient L2Wrap.backward returned the logits gradient without multiplying it by the upstream grad_output, so a non-unit upstream gradient (gradient accumulation loss scaling, custom loss weighting, torch.autograd.grad(grad_outputs=k)) under-scaled the L2 penalty gradient on the logits. The loss.backward() path uses an implicit upstream gradient of 1 and is unaffected. The fused path already scales its full gradient, including the L2Warp contribution, by the upstream gradient. Signed-off-by: Ting Sun <suntcrick@gmail.com> * [Fix] Fold upstream grad into L2Wrap scatter; tidy types/docstring and expand tests Scatter the scalar-scaled maxx directly instead of allocating a second full [B, T, V] tensor for glogits * grad_output. Add type hints and an FLA-style docstring, and parametrize the equivalence test over grad_scale to cover the non-unit upstream gradient path. --------- Signed-off-by: Ting Sun <suntcrick@gmail.com> Co-authored-by: Yu Zhang <yzhang.cs@outlook.com>
- Correct every __init__ return annotation from `-> ClassName` to `-> None`, add return hints to forwards, and fix `dot_dim: int = None` -> `int | None`. - Fix LearnablePolySketchNonNegativeFeatureMap crashing on the default `sketch_size=None`: normalize it before building the sketch layers. - Cache the triu gather indices on the input device in flatten_diag_outer_product* to avoid rebuilding and host->device copies on every forward. - Minor forward cleanups (dedup relu/2x, use nn.init.eye_) and a typo fix. - Add tests/modules/test_feature_map.py covering every feature map.
…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.
* [Model] Add Cyclic Flow Attention (CyFA) * [Perf] Keep CyFA backward tile offsets local * [Docs] Add CyFA news emoji [skip test] * [Docs] Use wave emoji for CyFA news [skip test] * [Fix] Validate CyFA recurrent inputs and preserve eval autograd * [Fix] Respect padding in cached CyFA continuations --------- Co-authored-by: Zhiyuan Li <uniartisan2017@gmail.com>
…la-org#1270) Respect the UB manager's safe BV estimate instead of forcing a larger tile.
Handle padded and packed LightNet inputs through chunk and recurrent operator entries that reuse GLA, preserving per-sequence normalization and cached state. Restore forward state reduction and avoid materializing unused gate gradients. Verified with strict GB200 operator, layer, cache and generation checks, followed by focused numerical and performance validation of the gate corrections.
…1351) * [Fix] Match TileLang attention output gradient shape * [Fix] Avoid invalid TileLang attention pipeline lowering
* [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.
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
…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.
zhudada0120
force-pushed
the
oja-ascend-adapt
branch
2 times, most recently
from
October 12, 2026 02:10
b5ade4b to
ed095f6
Compare
6 of 16 tasks
Collaborator
|
Looks like some kernels aren't meeting the accuracy threshold required by the test cases. One moment — let me check. |
On NPU the mainline kernel fails to compile for K=V=128: the whole [BK, V] fp32 state tile plus its per-step temporaries exceeds the 192KiB UB budget (hivm-plan-memory reports 1851392 bits required vs 1572864 available). The Oja residual `k - h @ v` couples all V columns, so V cannot be split across programs; instead the state lives as sequential in-kernel V slabs whose width comes from the UB manager budget (blockdim64 family style, as in chunk_delta_h), at most two slabs for the supported V <= 128 envelope. The mainline stays platform-neutral: only a `@dispatch` decorator is added, and the NPU implementation registers via `@register`. Also register test_oja.py in the A2 CI and add an NPU routing test. Tests: pytest tests/ops/test_oja.py — 26 passed on Atlas A2 (CANN 9.1.0, triton-ascend 3.2.2).
zhudada0120
force-pushed
the
oja-ascend-adapt
branch
from
October 14, 2026 04:59
ed095f6 to
f3120dc
Compare
…hunk kernels Route chunk_oja_fwd_o, chunk_oja_bwd_dqk, and chunk_oja_bwd_dv_o through @dispatch('gated_oja_rule') to Triton-Ascend backend forks that avoid three miscompilations only present in the old triton-ascend 3.2.2 package: - chunk_oja_fwd_intra: with the per-sub-chunk loop (for i_j in range(0, i_i)) present at a runtime trip count, the compiler drops the following j-loop accumulation for the i_i == 0 programs, corrupting the first BC rows of affected chunks. The loop is folded into one full-chunk masked dot, which is mathematically identical: columns at or past the current sub-chunk contribute zero. - chunk_oja_bwd_kernel_dqk: the transpose-view dot tl.dot(tl.trans(b_dA), b_q) produces wrong dk_new at BT = 16 while dq on the same tile stays exact. Load the transposed dA tile through swapped indices instead of tl.trans. - chunk_oja_bwd_kernel_dv_o: a later-sub-chunk loop whose trip count can be zero corrupts those programs. Launch the kernel twice with a HAS_LATER constexpr so each instance either carries no loop at all (i_i == NC - 1) or a loop guaranteed at least one iteration (i_i < NC - 1). The mainline stays platform-neutral: only the @dispatch decorators are added, and the fixed kernels live in fla/ops/gated_oja_rule/backends/triton_ascend/. This is a temporary shim until the package is upgraded: none of the three bugs reproduce on triton-ascend 3.6.0, so this commit can be reverted once 3.6.0 ships. The fused recurrent backend from the previous commit is unaffected (its UB overflow is inherent to the 192KiB budget and still present in 3.6.0) and stays. Validated cold-cache: all 26 tests in tests/ops/test_oja.py pass on 3.2.2 with the routing active.
zhudada0120
force-pushed
the
oja-ascend-adapt
branch
from
October 14, 2026 10:09
ba4dbd4 to
18be770
Compare
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:oja-ascend-adapt → fla-org:main) and reference #1372 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! |
5 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
Two parts, one per commit. f3120dc adds a Triton-Ascend backend for the gated Oja rule fused recurrent — a permanent adaptation, since the UB overflow it works around is inherent to the 192KiB Ascend Unified Buffer and is still present with triton-ascend 3.6.0. 18be770 adds workarounds for chunk kernel miscompiles that only exist in the current triton-ascend 3.2.2 package — a temporary shim that can be reverted once 3.6.0 ships.
f3120dc: add a Triton-Ascend backend for gated Oja rule fused recurrent forward and dispatch it through @dispatch('gated_oja_rule') when running on NPU. The mainline kernel keeps the whole [BK, V] fp32 state in one tile, which exceeds the 192KiB Ascend UB budget at K = V = 128 (hivm-plan-memory reports 1851392 bits required vs 1572864 available) and fails to compile. Because the Oja residual k - h @ v couples all V columns, V cannot be split across programs; the NPU kernel instead keeps the state as sequential in-kernel V slabs (blockdim64 family style, as in chunk_delta_h), with the slab width sized by the UB manager budget and at most two slabs covering the supported V <= 128 envelope. The mainline stays platform-neutral: only the @dispatch decorator is added, and the NPU implementation registers via @register('gated_oja_rule'). Also register test_oja.py in the A2 CI, add an NPU routing test, and register fused_recurrent_oja in the benchmark op registry.
18be770: route chunk_oja_fwd_o, chunk_oja_bwd_dqk, and chunk_oja_bwd_dv_o through the same @dispatch('gated_oja_rule') entry point to Triton-Ascend backend forks that avoid three miscompilations only present in the current triton-ascend 3.2.2 package. A per-sub-chunk loop with a runtime trip count makes the compiler drop the following j-loop accumulation for the i_i == 0 programs; the loop is folded into one full-chunk masked dot, which is mathematically identical. The transpose-view dot tl.dot(tl.trans(b_dA), b_q) computes wrong dk_new at BT = 16 while dq on the same tile stays exact; the transposed tile is loaded through swapped indices instead. A later-sub-chunk loop whose trip count can be zero corrupts those programs; the kernel is launched twice with a HAS_LATER constexpr so each instance either has no loop at all or a loop guaranteed at least one iteration. The mainline again only gains the dispatch decorators; CUDA and other backends are untouched and fall back to the original kernels. With the routing active, the full tests/ops/test_oja.py suite (26 tests) passes on triton-ascend 3.2.2 from a cold Triton cache. None of the three bugs reproduce on 3.6.0, so 18be770 can be reverted when the package is upgraded.
Test plan
Benchmark / NCU
Breaking changes
None
Checklist
Large PR justification
This PR is to adapt oja on the NPU; it is just a single migration PR.