[Ops] Add Triton-Ascend compatibility for GDN-2

[hazelduan]

Summary

This PR enables the GDN-2 chunk training/prefill path on Ascend NPU through a dedicated Triton-Ascend backend while preserving the existing implementation for non-NPU platforms.

The change:

• adds a GDN-2 backend for the two hardware-sensitive leaf stages, chunk_gdn2_fwd_intra and chunk_gdn2_bwd_wy_dqkg_fused;
• uses the existing backend registry and verifiers to select the Ascend implementation only for NPU tensors with supported dtypes and chunk_size=64;
• keeps the public API, autograd orchestration, state recurrence, output computation, fused-recurrent path, and original non-NPU function bodies unchanged;
• covers dense and packed variable-length forward/backward execution, NPU dispatch routing, native numerical references, and an opt-in 32K stress mode in the existing tests/ops/test_gdn2.py file;
• registers chunk_gdn2 in the unified operator benchmark runner; and
• adds tests/ops/test_gdn2.py to the Ascend A2 CI workflow.

Design and implementation

The backend is registered under fla.ops.gdn2.backends and is imported lazily by @dispatch('gdn2'). On Ascend, the verifier checks the device, dtype, and fixed 64-token chunk contract before selecting the specialized implementation. On CUDA, ROCm, and other platforms, the Ascend backend is unavailable and dispatch executes the original function body as the semantic fallback.

For the forward intra-chunk stage, each 64-token chunk is decomposed into four 16-token sub-blocks. The Ascend kernels build the causal Aqk and lower-triangular Akk terms, solve the diagonal blocks, and combine the inter-block terms while keeping the live working set within Ascend UB constraints. Host launch helpers split large grids and preserve packed-sequence boundaries for long and variable-length inputs. The backward path separates value/write gradients, the shared key/state gradient structure, GDN-2 erase/gate gradients, and the final dA transformation. It reuses the existing Triton-Ascend KDA key-side kernel where the mathematics is shared, while retaining dedicated GDN-2 kernels for its channel-wise gate terms. Hot time-strided inputs use a contiguous time layout where needed, and long-sequence/task offsets use 64-bit address arithmetic.

The small changes in the Triton-Ascend KDA gate/backward and shared cumsum launchers only remove CUDA-style num_warps arguments that are not valid NPU tuning parameters. They do not change those kernels' public interfaces, mathematics, tiling, or output contracts.

Test plan

35 passed, 62 warnings in 62.78s

The 35 nodes include native-reference checks, forward and backward gradients, dense and packed varlen execution, initial/final states, gate-in-kernel modes, GVA, non-power-of-two key dimensions, short-convolution layer coverage, and explicit assertions that the NPU forward and backward calls route through the Triton-Ascend backend.

Benchmark / NCU (kernel changes only)

chunk_gdn2 is registered in benchmarks/ops/registry.py for forward and forward+backward timing.
A current-code smoke benchmark was run on one Ascend910_9382 device with the runner's bfloat16 inputs, chunk_size=64, and use_qk_l2norm_in_kernel=True:

FLA_BENCH_OP_WARMUP_ITERS=1 \
FLA_BENCH_WARMUP_MS=5 \
FLA_BENCH_REP_MS=20 \
python -m benchmarks.ops.run \
  --op chunk_gdn2 \
  --base "" \
  --custom-shapes '{"npu_baseline":{"B":1,"T":64,"H":1,"D":32},"training_smoke":{"B":8,"T":1024,"H":8,"D":64}}' \
  --modes fwd fwdbwd

Median latency:

Shape	Forward	Forward + backward
B1 T64 H1 K32 V32 bf16	0.852 ms	2.389 ms
B8 T1024 H8 K64 V64 bf16	7.668 ms	16.158 ms

There is no numeric upstream-before result because this PR establishes the supported Ascend backend rather than optimizing an existing supported NPU implementation. These short-window measurements verify input generation, warmup, forward, backward, synchronization, and result reporting; they are not a speedup claim. No NCU or formal Ascend profiler comparison is included, and 32K throughput and peak-memory benchmarking remain follow-up performance work. Conclusion: neutral for performance claims; functional Ascend enablement.

Breaking changes

None. There are no public API, checkpoint, tolerance, or intended numerical-behavior changes. The specialized Ascend chunk path requires the existing GDN-2 chunk_size=64 contract, and fused_recurrent_gdn2 continues to use its existing implementation.

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.

[zheliuyu]

Just wondering — is the experimental code in test_gdn2 still not cleaned up? I wouldn't expect the test file to require this many changes.

[hazelduan]

Thanks for pointing this out. Initial diff mixed temporary NPU validation code into the canonical GDN2 UT and was broader than necessary. The current change in tests/ops/test_gdn2.py only replaces the CUDA-only skip condition with the repository-level accelerator guard covering NVIDIA, AMD, and Ascend NPU.

[sunyi0505]

nit: UT only tests small shapes. It is recommended to add the benchmark to ascend‑a2‑benchmark‑ci.yml to guard functionality and performance under large‑shape scenarios.

[zhiyuan1i]

Style issues:

1. P1: chunk_intra.py:394,398 — b_Ai22_c = b_Ai22 + 0.0 and b_Ai33_c are dead code (never used).

2. P1: chunk_bwd.py:163-165,291 and chunk_intra.py:392-402 — multiple x_c = x + 0.0 copies with no comment. If these are CANN compiler workarounds, each needs a one-line comment saying what it avoids. _c/_c2/_c3 naming is also uninformative.

3. P1: chunk_bwd.py:54 — _launch_dA_finalize is called only once (line 523). Inline the 25-line launch branch into chunk_gdn2_bwd_wy_dqkg_fused_npu.