[Ops] Use block matrix inversion to speed up solve_tril for the Ascend NPU backend

[OsirisDuan]

Summary

Rewrites the Ascend NPU solve_tril kernels to align with the GPU mainline on signature, index width, and memory access form, while adopting an NPU-friendlier block-inversion algorithm. This is an in-backend implementation replacement for triton_ascend: no wrapper signature change, no registration change, no observable behavior change.

Changes (fla/ops/utils/backends/triton_ascend/solve_tril.py + __init__.py, +234/-197):

1. Kernel signature aligned with GPU mainline. The three kernels (solve_tril_16x16_kernel_npu / merge_16x16_to_32x32_inverse_kernel_npu / merge_16x16_to_64x64_inverse_kernel_npu) now share (A, Ai, cu_seqlens, chunk_indices, T, H, BT, USE_TMA, IS_VARLEN, DOT_PRECISION) with H/BT/USE_TMA/IS_VARLEN/DOT_PRECISION as tl.constexpr, @triton.heuristics for IS_VARLEN, and do_not_specialize=['T']. The previously merged NT_OFFSET/BH_OFFSET (and the grid-split mechanism they enabled) are removed.

2. Index width int32 -> int64 (correctness fix). program_id is cast to int64; base address (bos*H+i_h)*BT and row offset (i_t*BT+i)*H*BT stay int64 end-to-end; in the varlen branch i_n is cast to int32 and i_t/cu_seqlens to int64. The old int32 indexing could overflow for large B*T*H combos (e.g. B=16, T=131072, H=64, BT=64 -> base offset ~= 8.05e9 > 2^31).

3. Memory access: tl.make_block_ptr -> pointer arithmetic + mask. make_block_ptr only accepts int32 offsets, which conflicts with the int64 indices above (AssertionError: Block pointers only support 32 bit offsets/block_shape). Replaced with the GPU-mainline-style o_t = (i_t*BT + o_i).to(tl.int64) + p = base + o_t[:,None]*stride + ... + tl.load(..., mask=...).

4. [Algorithm change declaration] Diagonal block inversion: serial forward-substitution -> tri_inv_mch doubling method. For the four 16x16 diagonal blocks of the 64x64 kernel, replaces GPU mainline's serial row-by-row forward substitution (15-level vector dependency chain) with the tri_inv_mch doubling iteration (3-level Cube dependency chain, higher NPU parallelism). The two methods are mathematically equivalent under exact arithmetic (A strictly lower-triangular -> nilpotent -> finite convergent series, Sum_{k=0}^{15} B^k = (I - B)^{-1}); the only difference is fp rounding order. The varlen tail short-block path (remaining_rows < 16/32/48/64) keeps the serial fallback branch for boundary correctness.

5. Off-diagonal blocks: recursive 2x2 grouping. The 64x64 off-diagonal blocks use recursive 2x2 block inversion (reusing P00/P01/P10/P11 intermediates, 11 dots) instead of the GPU mainline expansion (9 dots). Mathematically equivalent; dependency tree is shallower. Total 16x16 dot count is identical (16).

6. Wrapper simplified. Removes _launch_solve_tril_kernel and the grid-split machinery; reverts to a single merge_fn[NT, B*H](...) launch with USE_TMA=False, DOT_PRECISION="ieee", and drops the triton-ascend extension launch args (num_stages/multibuffer/sync_solver).

7. Verifier hardened. solve_tril_verifier now validates IS_NPU, A.device.type == 'npu', A.dtype in {fp16, bf16, fp32}, and A.shape[-1] in {16, 32, 64} instead of unconditionally returning True, None.

Per AGENTS.md ("RFC first for precision/algorithm changes"), item 4 (diagonal block inversion algorithm change) is declared here explicitly. Verified: 8/8 fixed-shape cases report diff = 0.000000 (well within the 1e-4 tolerance).

Test plan

• Direct coverage: tests/ops/test_solve_tril.py
  • Hardware: Ascend 910B
  • Command: python -m pytest tests/ops/test_solve_tril.py -v -p no:cacheprovider
  • Result: 10 passed, 1 skipped (120s); diff/ratio all 0.000000
  • Coverage: BT=16/32/64 fixed-shape (B1-T63 / B2-T500/T1000 / B3-T1024 / B4-T2048); BT=16/32/64 varlen (incl. short-block cu_seqlens=[0,15], multi-segment [0,256,500,1000], [0,1,100,300,1200,2048])
  • Skipped: test_solve_tril_large_batch_offsets (Blackwell-only, expected on NPU)
• Dependent tests: python scripts/find_dependent_tests.py fla/ops/utils/backends/triton_ascend/solve_tril.py returns 47 files (test_gdn, test_gla, test_linear_attn, test_kda, ...). Not run in full here (too long); solve_tril is a utils op consumed by those kernels, recommended for CI coverage.

Benchmark / NCU (kernel changes only)

• Hardware: Ascend 910B
  <TBD>

Breaking changes

None. solve_tril_npu wrapper signature (A, cu_seqlens, chunk_indices, output_dtype) is unchanged; TritonAscendUtilsBackend.solve_tril registration in __init__.py is unchanged; output semantics ((I+A)^{-1}, shape/dtype) are unchanged. The algorithm change (item 4) shows no numerical difference within test tolerance.

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.
• This is not a minor/cosmetic-only PR (typo, formatting, style-only tweaks).
• Dependent tests pass locally or in CI; new behavior is covered by tests where applicable.
• Kernel changes include same-hardware before/after benchmark numbers (dense + varlen where applicable).

If you cannot tick the "not minor" box above

Standalone minor PRs are normally not accepted (see No busywork PRs).
If you believe yours is an exception, justify here why it is worth a maintainer's review time — PRs without a justification may be closed without review:

[zheliuyu]

The 32×32 / 64×64 kernels drop the + 0.0 lhs copies that the current NPU solve_tril already has. On Ascend, tl.dot(lhs, rhs) can overwrite lhs in UB; later reads (second lhs, rhs, store, or X = X + tl.dot(X, Y)) see the clobbered tile. This file is in the repo catalog for exactly that (b_Ai_22_c on 32×32; b_Ai_22_c / b_Ai_33_c{,2} / b_Ai_44_c{,2} / b_A_42_c / b_A_43_c on 64×64).

This site is the clear example: inner tl.dot(b_Ai_22, b_A_21) then tl.store(..., b_Ai_22) on L157. The 64×64 path is the same class of bug, denser (tl.dot(A, A), X = X + tl.dot(X, Y), lhs reused in b_Ai_31/b_Ai_32 and b_Ai_41/b_Ai_42).

Please restore copies before the first lhs tl.dot of each reused tile. A 0.0 assert_close on test_solve_tril does not show the compiler stopped clobbering; that failure is silent and layout-dependent.

[zheliuyu]

@OsirisDuan #1161 has landed (torch_npu==2.9.0.post6 / triton-ascend==3.2.2 / CANN 9.1.0). Could you rebase this onto latest main before continuing? The file overlap should be small, but the compiler defaults changed.