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

README.md

@@ -137,9 +137,10 @@ pip3 install torch torchvision --index-url https://download.pytorch.org/whl/rocm
 
 #### Ascend NPU
 
-- `torch == 2.7.1`
-- `torch_npu == 2.7.1`
-- `triton-ascend == 3.2.1` Install from the Ascend PyPI mirror (not on default PyPI).
+- `torch == 2.9.0`
+- `torch_npu == 2.9.0`
+- `triton-ascend == 3.2.2` Install from the Ascend PyPI mirror (not on default PyPI).
+- `CANN == 9.1.0`
 
 ```bash
 pip install -e ".[dev]" --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple

setup.py

@@ -26,7 +26,7 @@ def get_default_dependencies():
             "torch>=2.6.0",
         ]
     elif platform == "npu":
-        return ["torch==2.7.1", "torch_npu==2.7.1", "triton-ascend==3.2.1"]
+        return ["torch==2.9.0", "torch_npu==2.9.0", "triton-ascend==3.2.2"]
 
 
 def get_optional_dependencies():

src/liger_kernel/ops/backends/_ascend/ops/fused_moe.py

@@ -1,7 +1,7 @@
 """
 Fused MoE expert computation via Triton grouped GEMM (Ascend backend).
 
-Forward: routing metadata (3 kernels) → fused gather+GEMM+SwiGLU → down-proj → token aggregation
+Forward: routing metadata (3 kernels) → gather+GEMM → SwiGLU epilogue → down-proj → token aggregation
 Backward: memory-efficient — recomputes dA' = dO@W2^T to avoid caching Y (TK×H bytes)
 """
 
@@ -28,6 +28,7 @@
 from .fused_moe_kernels import _moe_router_histogram_kernel
 from .fused_moe_kernels import _moe_router_prefix_sum_kernel
 from .fused_moe_kernels import _moe_router_scatter_kernel
+from .fused_moe_kernels import _swiglu_from_pre_act_kernel
 from .fused_moe_kernels import _token_gather_weighted_sum_kernel
 
 # Token-dimension tile size for M. Fixed (not autotuned) because tile_row_start,
@@ -193,7 +194,7 @@ def _token_scatter_sum(src, s_reverse_scatter_idx, T, K, H):
 class LigerFusedMoEFunction(torch.autograd.Function):
     """Fused grouped GEMM MoE forward + memory-efficient backward.
 
-    Forward: routing metadata → fused gather+GEMM+SwiGLU → down-proj → token aggregation
+    Forward: routing metadata → gather+GEMM → SwiGLU epilogue → down-proj → token aggregation
     Backward: avoids caching Y (TK×H) by recomputing dA' = dO@W2^T in backward
     """
 
@@ -240,7 +241,6 @@ def forward(ctx, x, gate_up_proj, down_proj, top_k_index, top_k_weights):
                 tile_row_start,
                 tile_expert,
                 pre_act,
-                post_act,
                 H_dim=H,
                 I_dim=intermediate_dim,
                 stride_x_T=x.stride(0),
@@ -250,12 +250,21 @@ def forward(ctx, x, gate_up_proj, down_proj, top_k_index, top_k_weights):
                 stride_w_K=gate_up_proj.stride(2),
                 stride_pre_TK=pre_act.stride(0),
                 stride_pre_N=pre_act.stride(1),
-                stride_post_TK=post_act.stride(0),
-                stride_post_N=post_act.stride(1),
                 BLOCK_M=BLOCK_M_TOKEN,
                 BLOCK_N=ASCEND_GEMM_BLOCK_N,
                 BLOCK_K=ASCEND_GEMM_BLOCK_K,
             )
+            _swiglu_from_pre_act_kernel[(min(TK, ASCEND_MAX_GRID_PROGRAMS),)](
+                pre_act,
+                post_act,
+                TK,
+                I_dim=intermediate_dim,
+                stride_pre_TK=pre_act.stride(0),
+                stride_pre_N=pre_act.stride(1),
+                stride_post_TK=post_act.stride(0),
+                stride_post_N=post_act.stride(1),
+                BLOCK_N=ASCEND_GEMM_BLOCK_N,
+            )
 
         Y = torch.empty(TK, H, dtype=x.dtype, device=x.device)
         if num_m_tiles > 0:

src/liger_kernel/ops/backends/_ascend/ops/fused_moe_kernels.py

@@ -213,7 +213,6 @@ def _fused_up_proj_swiglu_kernel(
     tile_row_start_ptr,  # (num_m_tiles,) int32 — row_start per M-tile
     tile_expert_ptr,  # (num_m_tiles,) int32 — expert index per M-tile
     pre_act_ptr,  # (TK, 2*I)  pre-SwiGLU activations [saved for backward]
-    post_act_ptr,  # (TK, I)    post-SwiGLU activations
     H_dim: tl.constexpr,
     I_dim: tl.constexpr,
     stride_x_T,
@@ -223,15 +222,15 @@ def _fused_up_proj_swiglu_kernel(
     stride_w_K: tl.constexpr,
     stride_pre_TK,
     stride_pre_N: tl.constexpr,
-    stride_post_TK,
-    stride_post_N: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
     """Grid: (num_m_tiles,). One CTA per M-tile; N-tiles iterated in-kernel.
 
-    Ascend: 1D grid avoids exceeding the ~32K launch limit at large T.
+    Cube-only: writes pre-SwiGLU [gate, up]. SwiGLU is a separate vector kernel
+    because triton-ascend 3.2.2 ConvertLinalgRToBinary cannot lower mix
+    cube+vector (gather + dual tl.dot + silu + GM stores) to static UB shapes.
     """
     pid_m = tl.program_id(0)
 
@@ -283,12 +282,44 @@ def _fused_up_proj_swiglu_kernel(
         tl.store(pre_gate_ptrs, acc_gate.to(pre_act_ptr.dtype.element_ty), mask=out_mask)
         tl.store(pre_up_ptrs, acc_up.to(pre_act_ptr.dtype.element_ty), mask=out_mask)
 
-        sig_gate = tl.sigmoid(acc_gate)
-        silu_gate = acc_gate * sig_gate
-        a_out = silu_gate * acc_up
 
-        post_ptrs = post_act_ptr + row_offs[:, None] * stride_post_TK + n_idx[None, :] * stride_post_N
-        tl.store(post_ptrs, a_out.to(post_act_ptr.dtype.element_ty), mask=out_mask)
+@triton.jit
+def _swiglu_from_pre_act_kernel(
+    pre_act_ptr,  # (TK, 2*I)
+    post_act_ptr,  # (TK, I)
+    TK,
+    I_dim: tl.constexpr,
+    stride_pre_TK,
+    stride_pre_N: tl.constexpr,
+    stride_post_TK,
+    stride_post_N: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Vector epilogue: post = silu(gate) * up. Grid-stride over TK rows."""
+    pid = tl.program_id(0)
+    n_prog = tl.num_programs(0)
+    n_offs = tl.arange(0, BLOCK_N)
+
+    for tk in tl.range(pid, TK, n_prog):
+        for n_start in tl.range(0, I_dim, BLOCK_N):
+            n_idx = n_start + n_offs
+            n_mask = n_idx < I_dim
+            gate = tl.load(
+                pre_act_ptr + tk * stride_pre_TK + n_idx * stride_pre_N,
+                mask=n_mask,
+                other=0.0,
+            ).to(tl.float32)
+            up = tl.load(
+                pre_act_ptr + tk * stride_pre_TK + (n_idx + I_dim) * stride_pre_N,
+                mask=n_mask,
+                other=0.0,
+            ).to(tl.float32)
+            out = gate * tl.sigmoid(gate) * up
+            tl.store(
+                post_act_ptr + tk * stride_post_TK + n_idx * stride_post_N,
+                out.to(post_act_ptr.dtype.element_ty),
+                mask=n_mask,
+            )
 
 
 @triton.jit

src/liger_kernel/ops/backends/_ascend/ops/fused_neighborhood_attention.py

@@ -2338,6 +2338,10 @@ def fused_neighborhood_attention_forward(
         block_n,
     )
 
+    # triton-ascend 3.2.2: consecutive forwards on the same stream hang unless
+    # QK+AV retire first (`test_fused_neighborhood_attention_deterministic`).
+    torch.npu.current_stream().synchronize()
+
     if return_lse:
         raise NotImplementedError("return_lse=True is not supported yet.")
 

src/liger_kernel/ops/backends/_ascend/ops/mhc.py

@@ -25,6 +25,11 @@
 from liger_kernel.ops.utils import ensure_contiguous
 from liger_kernel.ops.utils import get_npu_core_count
 
+# Only on mix cube+vector and tl.atomic_add launches. BiShengIR 1.2.0
+# (triton-ascend 3.2.2) auto-multi-buffer races fp32 mix results and hangs
+# compiling mix + atomic_add. Recognized on 3.2.1 as well (default was True).
+_NO_MULTIBUFFER = dict(multibuffer=False)
+
 # ---------------------------------------------------------------------------
 # UB-aware block size helpers  (via unified compute_default_tiling_strategy)
 # ---------------------------------------------------------------------------
@@ -205,6 +210,7 @@ def mhc_mm_norm_fwd(
         BLOCK_N=block_n,
         BLOCK_K=block_k,
         BLOCK_M=block_m,
+        **_NO_MULTIBUFFER,
     )
     return out_mix, out_invr
 
@@ -365,6 +371,7 @@ def mhc_mm_norm_bwd(
         BLOCK_N=block_n,
         BLOCK_K=block_k,
         BLOCK_M=block_m,
+        **_NO_MULTIBUFFER,
     )
     if out_grad_phi.dtype != phi.dtype:
         out_grad_phi = out_grad_phi.to(phi.dtype)
@@ -999,6 +1006,7 @@ def mhc_pre_bwd(
         stride_ghh=out_grad_h.stride(1),
         BLOCK_N=block_n,
         BLOCK_C=block_c,
+        **_NO_MULTIBUFFER,
     )
     return out_grad_x, out_grad_h
 
@@ -1305,6 +1313,7 @@ def mhc_post_res_bwd(
         stride_ghrj=out_grad_hres.stride(2),
         BLOCK_N=block_n,
         BLOCK_C=block_c,
+        **_NO_MULTIBUFFER,
     )
     return out_grad_x, out_grad_f, out_grad_hpost, out_grad_hres
 
@@ -1458,6 +1467,7 @@ def mhc_coeffs_bwd_assemble(
         stride_grn=grad_res_flat.stride(0),
         BLOCK_HC=BLOCK_HC,
         BLOCK_RES=BLOCK_RES,
+        **_NO_MULTIBUFFER,
     )
     return (
         grad_mix,