diff --git a/README.md b/README.md index 491644892..0fe1e0c51 100644 --- a/README.md +++ b/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 diff --git a/setup.py b/setup.py index a680d1121..9d18fd306 100644 --- a/setup.py +++ b/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(): diff --git a/src/liger_kernel/ops/backends/_ascend/ops/fused_moe.py b/src/liger_kernel/ops/backends/_ascend/ops/fused_moe.py index 04740400f..e062b9749 100644 --- a/src/liger_kernel/ops/backends/_ascend/ops/fused_moe.py +++ b/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: diff --git a/src/liger_kernel/ops/backends/_ascend/ops/fused_moe_kernels.py b/src/liger_kernel/ops/backends/_ascend/ops/fused_moe_kernels.py index 37e4e4b4d..71bbea7f5 100644 --- a/src/liger_kernel/ops/backends/_ascend/ops/fused_moe_kernels.py +++ b/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 diff --git a/src/liger_kernel/ops/backends/_ascend/ops/fused_neighborhood_attention.py b/src/liger_kernel/ops/backends/_ascend/ops/fused_neighborhood_attention.py index 370363f5e..04621dc2a 100644 --- a/src/liger_kernel/ops/backends/_ascend/ops/fused_neighborhood_attention.py +++ b/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.") diff --git a/src/liger_kernel/ops/backends/_ascend/ops/mhc.py b/src/liger_kernel/ops/backends/_ascend/ops/mhc.py index 2a40a8136..7d49c2eb5 100644 --- a/src/liger_kernel/ops/backends/_ascend/ops/mhc.py +++ b/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,