[Perf] Optimize Ascend prepare_wy_repr_bwd with larger kv tiles, bf16 mmad, and fused a2/dg

[sunyi0505]

Summary

Speed up the Ascend (triton_ascend) prepare_wy_repr_bwd_npu pipeline. Public API is unchanged; GPU path is untouched.
1. Larger K/V tiles for the kv stage (BK=BV up to 256)

• Split tiling: kv reuses K/V slabs in-place, so a 2.25× UB model can compile BK=BV=256. finalize_k / a2 keep a conservative 4.5× slab (max=128) because BK=256 overflows 192KB UB.
• Load A once per task instead of reloading it every K/V tile; hoist block-pointer bases; keep K/V/dw/du in input dtype and downcast only for reductions.
  2. Cube bf16 mmad with fp32 dA accumulation
• Match CUDA: tl.dot(..., acc=b_dA) in input dtype, dA stays fp32. Copy A before downcast so lhs clobber is safe (b_A_c = b_A.to(b_dw.dtype) + 0.0).
• On 910, fp32 64×64×256 mmad was scalar-setup bound (~2 µs/gemm). bf16 mmad halves Cube scalar wall time on the D=256 hotspot.
  3. Fuse finalize A2 + dg; drop GM scratch
• prepare_wy_repr_bwd_finalize_a2_dg_npu computes A2 = (k k^T) * beta in UB, then dg += row(dA*A2) - col(dA*A2). Removes a2_scr / col_acc_scr and the separate BC-tiled dg kernel.
• da_gate is full-BT Vector (drop _DG_BLK nested GM). dA_* scratch is empty_like (kernels fully overwrite).

Test plan

hardware: Atlas 800T A3(X86)
FLA_NPU_XDIST=1 ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 pytest --exitfirst -n 8 --dist load tests/modules tests/ops/utils tests/ops/test_gdn_kernels.py tests/ops/test_gdn.py tests/ops/test_kda.py tests/ops/test_attnres.py tests/ops/test_solve_tril.py tests/ops/test_gla.py

Benchmark / NCU (kernel changes only)

Breaking changes

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.

If you ticked the "minor" box above

Standalone minor PRs are normally not accepted (see No busywork PRs).
Justify here why yours is worth a maintainer's review time — minor PRs without a justification may be closed without review:

[sunyi0505]

@zhiyuan1i This pr is ready for review. Could you please review this PR? Thanks!

[zheliuyu]

LGTM. Optional: the zeros_like → empty_like switch is safe today because every (chunk, head) task unconditionally full-writes its tile and all consumers read via boundary_check — this invariant is what keeps it from becoming another #1070. Would you consider stating it in a one-line comment at the dA_* allocations, or simply keeping zeros_like? I lean toward the comment.