Skip to content

[NPU]: remove num_stages and num_warps from NPU Triton kernels #1111

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Tcc0403 merged 3 commits into from
Mar 2, 2026
Merged

[NPU]: remove num_stages and num_warps from NPU Triton kernels #1111

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
63a5997
[NPU]: remove num_stages and num_warps from NPU Triton kernels
noemotiovon Feb 27, 2026
06c79af
Remove duplicate comments
noemotiovon Mar 2, 2026
f9151e1
Remove duplicate comments
noemotiovon Mar 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 2 additions & 6 deletions src/liger_kernel/ops/backends/_ascend/ops/embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ def embedding_forward_kernel(
embedding_dim: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
pid = tl.program_id(0)
num_progs = tl.num_programs(0)
Expand All @@ -25,7 +24,7 @@ def embedding_forward_kernel(
grid_n = tl.cdiv(embedding_dim, BLOCK_SIZE_N)
total_2d_blocks = grid_m * grid_n

for block_idx in tl.range(pid, total_2d_blocks, num_progs, num_stages=NUM_STAGES):
for block_idx in tl.range(pid, total_2d_blocks, num_progs):
block_m = block_idx // grid_n
block_n = block_idx % grid_n

Expand Down Expand Up @@ -66,7 +65,6 @@ def embedding_backward_kernel(
embedding_dim: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
pid = tl.program_id(0)
num_progs = tl.num_programs(0)
Expand All @@ -75,7 +73,7 @@ def embedding_backward_kernel(
grid_n = tl.cdiv(embedding_dim, BLOCK_SIZE_N)
total_2d_blocks = grid_m * grid_n

for block_idx in tl.range(pid, total_2d_blocks, num_progs, num_stages=NUM_STAGES):
for block_idx in tl.range(pid, total_2d_blocks, num_progs):
block_m = block_idx // grid_n
block_n = block_idx % grid_n

Expand Down Expand Up @@ -164,7 +162,6 @@ def embedding_forward(embeddings, indices):
embedding_dim=embedding_dim,
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE_N,
NUM_STAGES=3,
)

return output.view(*ori_shape, -1)
Expand All @@ -191,7 +188,6 @@ def embedding_backward(embeddings, indices, grad_output):
embedding_dim=embedding_dim,
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE_N,
NUM_STAGES=3,
)

return grad_weight
Expand Down
16 changes: 4 additions & 12 deletions src/liger_kernel/ops/backends/_ascend/ops/fused_add_rms_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,6 @@ def _fused_add_rms_norm_forward_kernel_no_tiling(
X_DTYPE: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
"""
NPU-optimized fused_add_rms_norm forward kernel for small n_cols (< 2048).
Expand Down Expand Up @@ -80,7 +79,7 @@ def _fused_add_rms_norm_forward_kernel_no_tiling(
W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0.0)

# Grid-stride loop over row blocks
for i in tl.range(num_iterations, num_stages=NUM_STAGES):
for i in tl.range(num_iterations):
row_idx = i * grid_stride + pid * BLOCK_SIZE_M + row_offsets
row_mask = row_idx < n_rows
block_mask = row_mask[:, None] & col_mask[None, :]
Expand Down Expand Up @@ -157,7 +156,6 @@ def _fused_add_rms_norm_forward_kernel_npu(
casting_mode: tl.constexpr,
X_DTYPE: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
"""
NPU-optimized fused_add_rms_norm forward kernel.
Expand All @@ -181,7 +179,7 @@ def _fused_add_rms_norm_forward_kernel_npu(

offsets = tl.arange(0, BLOCK_SIZE)
# Grid-stride loop over rows
for row_idx in tl.range(pid, n_rows, num_progs, num_stages=NUM_STAGES):
for row_idx in tl.range(pid, n_rows, num_progs):
Y_row_ptr = Y_ptr + row_idx * Y_row_stride
S_row_ptr = S_ptr + row_idx * S_row_stride
X_row_ptr = X_ptr + row_idx * X_row_stride
Expand Down Expand Up @@ -279,7 +277,6 @@ def _fused_add_rms_norm_backward_kernel_no_tiling(
casting_mode: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
NUM_STAGES: tl.constexpr,
has_dS_out: tl.constexpr,
):
"""
Expand Down Expand Up @@ -308,7 +305,7 @@ def _fused_add_rms_norm_backward_kernel_no_tiling(
W_offset = W_row + offset

# Grid-stride loop over row blocks
for i in tl.range(num_iterations, num_stages=NUM_STAGES):
for i in tl.range(num_iterations):
row_idx = i * grid_stride + pid * BLOCK_SIZE_M + row_offsets
row_mask = row_idx < n_rows
block_mask = row_mask[:, None] & col_mask[None, :]
Expand Down Expand Up @@ -402,7 +399,6 @@ def _fused_add_rms_norm_backward_kernel_npu(
offset,
casting_mode: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
NUM_STAGES: tl.constexpr,
has_dS_out: tl.constexpr,
):
"""
Expand All @@ -419,7 +415,7 @@ def _fused_add_rms_norm_backward_kernel_npu(
offsets = tl.arange(0, BLOCK_SIZE)

# Grid-stride loop over rows
for row_idx in tl.range(pid, n_rows, num_progs, num_stages=NUM_STAGES):
for row_idx in tl.range(pid, n_rows, num_progs):
# Base pointers for this row
dY_row_ptr = dY_ptr + row_idx * dY_row_stride
dX_row_ptr = dX_ptr + row_idx * dX_row_stride
Expand Down Expand Up @@ -625,7 +621,6 @@ def fused_add_rms_norm_forward(X, R, W, eps, offset, casting_mode):
X_DTYPE,
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE,
NUM_STAGES=2,
)
else:
# Use tiled kernel for large n_cols
Expand All @@ -648,7 +643,6 @@ def fused_add_rms_norm_forward(X, R, W, eps, offset, casting_mode):
casting_mode,
X_DTYPE,
BLOCK_SIZE=BLOCK_SIZE,
NUM_STAGES=2,
)

return Y.view(*shape), S.view(*shape), RSTD, casting_mode
Expand Down Expand Up @@ -703,7 +697,6 @@ def fused_add_rms_norm_backward(dY, dS_out, S, W, RSTD, offset, casting_mode, in
casting_mode,
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE,
NUM_STAGES=2,
has_dS_out=dS_out is not None,
)
else:
Expand All @@ -728,7 +721,6 @@ def fused_add_rms_norm_backward(dY, dS_out, S, W, RSTD, offset, casting_mode, in
offset,
casting_mode,
BLOCK_SIZE=BLOCK_SIZE,
NUM_STAGES=2,
has_dS_out=dS_out is not None,
)

Expand Down
16 changes: 6 additions & 10 deletions src/liger_kernel/ops/backends/_ascend/ops/geglu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@


@triton.jit
def _geglu_forward_kernel_flat(a_ptr, b_ptr, c_ptr, total_elements, BLOCK_SIZE: tl.constexpr, NUM_STAGES: tl.constexpr):
def _geglu_forward_kernel_flat(a_ptr, b_ptr, c_ptr, total_elements, BLOCK_SIZE: tl.constexpr):
"""
High-performance GEGLU forward kernel using flatten 1D approach.

Expand All @@ -27,7 +27,7 @@ def _geglu_forward_kernel_flat(a_ptr, b_ptr, c_ptr, total_elements, BLOCK_SIZE:
sqrt_2_over_pi = 0.7978845608028654 # sqrt(2 / pi)
gelu_coeff = 0.044715

for idx in tl.range(start_idx, total_elements, stride, num_stages=NUM_STAGES):
for idx in tl.range(start_idx, total_elements, stride):
offsets = idx + tl.arange(0, BLOCK_SIZE)
mask = offsets < total_elements

Expand All @@ -45,9 +45,7 @@ def _geglu_forward_kernel_flat(a_ptr, b_ptr, c_ptr, total_elements, BLOCK_SIZE:


@triton.jit
def _geglu_backward_kernel_flat(
dc_ptr, a_ptr, b_ptr, da_ptr, db_ptr, total_elements, BLOCK_SIZE: tl.constexpr, NUM_STAGES: tl.constexpr
):
def _geglu_backward_kernel_flat(dc_ptr, a_ptr, b_ptr, da_ptr, db_ptr, total_elements, BLOCK_SIZE: tl.constexpr):
"""
High-performance GEGLU backward kernel using flatten 1D approach.

Expand All @@ -62,7 +60,7 @@ def _geglu_backward_kernel_flat(
sqrt_2_over_pi = 0.7978845608028654 # sqrt(2 / pi)
gelu_coeff = 0.044715

for idx in tl.range(start_idx, total_elements, stride, num_stages=NUM_STAGES):
for idx in tl.range(start_idx, total_elements, stride):
offsets = idx + tl.arange(0, BLOCK_SIZE)
mask = offsets < total_elements

Expand Down Expand Up @@ -143,7 +141,7 @@ def geglu_forward(a, b):
num_cores = get_npu_core_count()
grid_size = min(num_cores, (total_elements + block_size - 1) // block_size)

_geglu_forward_kernel_flat[(grid_size,)](a, b, c, total_elements, BLOCK_SIZE=block_size, NUM_STAGES=3, num_warps=4)
_geglu_forward_kernel_flat[(grid_size,)](a, b, c, total_elements, BLOCK_SIZE=block_size)
return c


Expand All @@ -167,9 +165,7 @@ def geglu_backward(a, b, dc):
num_cores = get_npu_core_count()
grid_size = min(num_cores, (total_elements + block_size - 1) // block_size)

_geglu_backward_kernel_flat[(grid_size,)](
dc, a, b, grad_a, grad_b, total_elements, BLOCK_SIZE=block_size, NUM_STAGES=3, num_warps=4
)
_geglu_backward_kernel_flat[(grid_size,)](dc, a, b, grad_a, grad_b, total_elements, BLOCK_SIZE=block_size)
return grad_a, grad_b


Expand Down
5 changes: 1 addition & 4 deletions src/liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ def _triton_qwen2vl_mrope_npu(
mrope_section_h: tl.constexpr,
BLOCK_Q: tl.constexpr,
BLOCK_K: tl.constexpr,
NUM_STAGES: tl.constexpr,
BACKWARD_PASS: tl.constexpr = False,
):
program_id = tl.program_id(0)
Expand All @@ -34,7 +33,7 @@ def _triton_qwen2vl_mrope_npu(
start_row = program_id * rows_per_program
actual_rows = tl.minimum(rows_per_program, total_rows - start_row)

for row_offset in tl.range(0, actual_rows, num_stages=NUM_STAGES):
for row_offset in tl.range(0, actual_rows):
pid = start_row + row_offset

t_end = mrope_section_t
Expand Down Expand Up @@ -196,7 +195,6 @@ def qwen2vl_mrope_forward(q, k, cos, sin, mrope_section):
mrope_section[1],
BLOCK_Q,
BLOCK_K,
NUM_STAGES=3,
BACKWARD_PASS=False,
)
return q.transpose(1, 2), k.transpose(1, 2), cos, sin
Expand Down Expand Up @@ -241,7 +239,6 @@ def qwen2vl_mrope_backward(dq, dk, cos, sin, mrope_section):
mrope_section[1],
BLOCK_Q,
BLOCK_K,
NUM_STAGES=3,
BACKWARD_PASS=True,
)
return dq.transpose(1, 2), dk.transpose(1, 2)
Expand Down
16 changes: 4 additions & 12 deletions src/liger_kernel/ops/backends/_ascend/ops/rms_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,6 @@ def _rms_norm_forward_kernel_no_tiling(
elementwise_affine: tl.constexpr,
X_DTYPE: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
"""
NPU-optimized rms_norm forward kernel for small n_cols (< 2048).
Expand All @@ -68,7 +67,7 @@ def _rms_norm_forward_kernel_no_tiling(
mask = col_offsets < n_cols

# Grid-stride loop over rows
for row_idx in tl.range(pid, n_rows, num_progs, num_stages=NUM_STAGES):
for row_idx in tl.range(pid, n_rows, num_progs):
Y_row_ptr = Y_ptr + row_idx * Y_row_stride
X_row_ptr = X_ptr + row_idx * X_row_stride
RSTD_row_ptr = RSTD_ptr + row_idx * RSTD_row_stride
Expand Down Expand Up @@ -146,7 +145,6 @@ def _rms_norm_forward_kernel_tiled(
elementwise_affine: tl.constexpr,
X_DTYPE: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
"""
NPU-optimized rms_norm forward kernel for large n_cols (>= 2048).
Expand All @@ -170,7 +168,7 @@ def _rms_norm_forward_kernel_tiled(

offsets = tl.arange(0, BLOCK_SIZE)
# Grid-stride loop over rows
for row_idx in tl.range(pid, n_rows, num_progs, num_stages=NUM_STAGES):
for row_idx in tl.range(pid, n_rows, num_progs):
Y_row_ptr = Y_ptr + row_idx * Y_row_stride
X_row_ptr = X_ptr + row_idx * X_row_stride
RSTD_row_ptr = RSTD_ptr + row_idx * RSTD_row_stride
Expand Down Expand Up @@ -266,7 +264,6 @@ def _rms_norm_backward_kernel_no_tiling(
casting_mode: tl.constexpr,
elementwise_affine: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
"""
NPU-optimized rms_norm backward kernel for small n_cols (< 2048).
Expand All @@ -282,7 +279,7 @@ def _rms_norm_backward_kernel_no_tiling(
mask = col_offsets < n_cols

# Grid-stride loop over rows
for row_idx in tl.range(pid, n_rows, num_progs, num_stages=NUM_STAGES):
for row_idx in tl.range(pid, n_rows, num_progs):
# Base pointers for this row
dY_row_ptr = dY_ptr + row_idx * dY_row_stride
dX_row_ptr = dX_ptr + row_idx * dX_row_stride
Expand Down Expand Up @@ -373,7 +370,6 @@ def _rms_norm_backward_kernel_tiled(
casting_mode: tl.constexpr,
elementwise_affine: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
"""
NPU-optimized rms_norm backward kernel for large n_cols (>= 2048).
Expand All @@ -389,7 +385,7 @@ def _rms_norm_backward_kernel_tiled(
offsets = tl.arange(0, BLOCK_SIZE)

# Grid-stride loop over rows
for row_idx in tl.range(pid, n_rows, num_progs, num_stages=NUM_STAGES):
for row_idx in tl.range(pid, n_rows, num_progs):
# Base pointers for this row
dY_row_ptr = dY_ptr + row_idx * dY_row_stride
dX_row_ptr = dX_ptr + row_idx * dX_row_stride
Expand Down Expand Up @@ -604,7 +600,6 @@ def rms_norm_forward(X, W, eps, offset, casting_mode):
elementwise_affine,
X_DTYPE,
BLOCK_SIZE=BLOCK_SIZE,
NUM_STAGES=2,
)
else:
# Use tiled kernel for large n_cols
Expand All @@ -624,7 +619,6 @@ def rms_norm_forward(X, W, eps, offset, casting_mode):
elementwise_affine,
X_DTYPE,
BLOCK_SIZE=BLOCK_SIZE,
NUM_STAGES=2,
)

return Y.view(*shape), X, RSTD, casting_mode
Expand Down Expand Up @@ -678,7 +672,6 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, in_place):
casting_mode,
elementwise_affine,
BLOCK_SIZE=BLOCK_SIZE,
NUM_STAGES=2,
)
else:
# Use tiled kernel for large n_cols
Expand All @@ -701,7 +694,6 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, in_place):
casting_mode,
elementwise_affine,
BLOCK_SIZE=BLOCK_SIZE,
NUM_STAGES=2,
)

dX = dX.view(*shape)
Expand Down
5 changes: 1 addition & 4 deletions src/liger_kernel/ops/backends/_ascend/ops/rope.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ def _triton_rope_npu(
hd: tl.constexpr,
BLOCK_Q: tl.constexpr,
BLOCK_K: tl.constexpr,
NUM_STAGES: tl.constexpr,
BACKWARD_PASS: tl.constexpr = False,
):
program_id = tl.program_id(0)
Expand All @@ -34,7 +33,7 @@ def _triton_rope_npu(
start_row = program_id * rows_per_program
actual_rows = tl.minimum(rows_per_program, total_rows - start_row)

for row_offset in tl.range(0, actual_rows, num_stages=NUM_STAGES):
for row_offset in tl.range(0, actual_rows):
pid = start_row + row_offset

row_idx = pid % sl
Expand Down Expand Up @@ -186,7 +185,6 @@ def rope_forward(q, k, cos, sin):
head_dim,
BLOCK_Q,
BLOCK_K,
NUM_STAGES=3,
BACKWARD_PASS=False,
)
return q.transpose(1, 2), k.transpose(1, 2), cos, sin
Expand Down Expand Up @@ -232,7 +230,6 @@ def rope_backward(dq, dk, cos, sin):
head_dim,
BLOCK_Q,
BLOCK_K,
NUM_STAGES=3,
BACKWARD_PASS=True,
)
return dq.transpose(1, 2), dk.transpose(1, 2)
Expand Down
Loading
Loading