[zheliuyu]

Summary

Pin the Ascend NPU install to Triton-Ascend 3.6.0 and cap CV workspace slots to one. CANN stays 9.1.0. torch, torch_npu, and torchvision move together so the stack stays installable. Kernels that already had a Triton implementation stay on Triton-Ascend. No PyTorch stand-in was added.

Package	Before	After
torch	2.9.0	2.10.0
torch_npu	2.9.0	2.10.0.post7
torchvision	unpinned	0.25.0
triton-ascend	3.2.2	3.6.0
CANN	9.1.0	9.1.0

Details

Triton-Ascend 3.6.0 ships BiShengIR 2.0, which pipelines mix kernels with a 4-slot CV workspace by default. A2 and A3 have a 192KB unified buffer, so four slots overflow (ub overflow, requires 2883840 bits while 1572864 bits available, multi-buffer enabled). install_ascend_compile_kwargs() sets set_workspace_multibuffer=1 when NPUOptions has that field. Triton-Ascend 3.2.x does not have the field, so those launches stay unchanged. An explicit caller value is kept.

Two Triton lowerings also change on 3.6.0:

• A masked wide tl.cumsum no longer assigns unique MoE ranks. The permutation stops being invertible, and the next launch stores an out-of-range index (vector core 507035). The rank is now a scalar tl.atomic_add. GEMM and SwiGLU are unchanged. tl.sort is not used.
• tl.static_range unrolls into the same 192KB buffer. AttnRes backward used to cap that loop at 32 blocks, so N=40 dropped the tail. The loop is now tl.range(0, n_blocks). Forward stays tl.static_range because n_blocks is constexpr there.

Host-side adjustments still call the existing Triton kernels:

• RoPE forwards rotary_dim and rotates only that prefix. An explicit width is not inferred from a narrower cos/sin table (GPT-OSS).
• TVD uses an integer reduction constexpr. torch.compile re-emits a string constexpr as a bare name. A tie subgradient is 0, and an all-ignored batch does not divide.
• LUSPO materializes a (B, 1) scale as (B, L) so the kernel does not read off the row.

These tests are skipped on NPU. They are not replaced with PyTorch:

• MLP. On the 3.2.2 pin (fbdea0ff) the file passed only because pytest.importorskip("triton.tools.tensor_descriptor") skipped it. 3.2.2 does not ship that module. 3.6.0 does, the tests actually run, and the CUDA TMA kernel hits assert is_cuda. There was no NPU MLP kernel in that 3.2.2 run.
• Fused CE/TVD and fused linear score-centering have no Ascend backend. The shared files are unchanged.

mHC still disables auto-multi-buffer, and neighborhood attention still synchronizes between forwards. setup.py and the README list the new pins. docs/index.md is unchanged. Existing NPU environments need to reinstall.

Testing Done

• Hardware Type: Ascend NPU, CANN 9.1.0
• pytest test/transformers/
• pytest test/transformers/test_attn_res.py test/transformers/test_fused_moe.py after the final trim (52 passed, including N=40 and the routing invariants)
• ruff check on the touched files
• make test-convergence

[zheliuyu]

 pytest test/transformers/

[zheliuyu]

fb903ad

pytest test/transformers/

[Tcc0403]

LGTM

[zheliuyu]

@Tcc0403 Benchmarks for swiglu, rms_norm, rope, cross_entropy, and fused_linear_cross_entropy are done on CANN 9.1.0 (Ascend910_9382). No significant regression compared with triton-ascend 3.2.2. Results are below.

triton-ascend 3.2.2: #1426

Benchmark — CANN 9.1.0, Ascend910_9382

Token-length sweep: 1024 / 2048 / 4096 / 8192.

SwiGLU

Speed · forward	Speed · backward
Speed · full	Memory · full

RMSNorm

Speed · forward	Speed · backward
Speed · full	Memory · full

RoPE

Speed · forward	Speed · backward
Speed · full	Memory · full

Cross Entropy

Speed · forward	Speed · backward
Speed · full	Speed · no-grad forward
Memory · full

Fused Linear Cross Entropy

Speed · forward	Speed · backward
Speed · full	Memory · full