[NPU] Pin torch/torch_npu=2.9.0, triton-ascend=3.2.2 and cann=9.1.0 for NPU

[zheliuyu]

Summary

Adapt three Ascend kernels to the new stack (torch/torch_npu 2.9.0, triton-ascend 3.2.2, CANN 9.1 / BiShengIR 1.2.0). triton itself stays 3.2.0; the failures come from compiler/runtime default changes in triton-ascend 3.2.2, not from test logic.

Each change is scoped to the kernel that actually hits a 3.2.2 constraint. GPU paths are untouched. The same sources still run on the previous stack (torch/torch_npu 2.7.1 + triton-ascend 3.2.1).

Details

Stack delta that matters:

	previous (green)	new (broken)
torch / torch_npu	2.7.1	2.9.0
triton	3.2.0	3.2.0
triton-ascend	3.2.1	3.2.2
CANN / BiShengIR	9.0 / older	9.1 / 1.2.0

3.2.2 defaults multibuffer=True and cannot lower some mix cube+vector programs to static UB shapes. Fixes:

1. Fused MoE — fp32 _fused_up_proj_swiglu_kernel failed compile with ConvertLinalgRToBinary / Failed to obtain op buffer shape size which should be static. Cause: gather + dual tl.dot + SwiGLU in one kernel. Fix: cube kernel writes pre_act only; new vector kernel _swiglu_from_pre_act_kernel does silu(gate) * up. Routing / down-proj / backward unchanged.

2. mHC — auto-multi-buffer races fp32 mix cube+vector results and hangs compiling mix + tl.atomic_add. Fix: pass multibuffer=False only on those launches (mm_norm_fwd/bwd, pre_bwd, post_res_bwd, coeffs_bwd_assemble). Vector-only forward kernels keep the 3.2.1 default.

3. Fused neighborhood attention — consecutive forwards on the same stream hang until QK+AV retire (test_fused_neighborhood_attention_deterministic). Fix: torch.npu.current_stream().synchronize() after the AV launch (stream-local, not a device-wide barrier).

Testing Done

• Hardware Type: Ascend NPU (torch 2.9.0, torch_npu 2.9.0, triton-ascend 3.2.2, CANN 9.1)
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

Targeted pytest on 3.2.2:

Coming soon.