linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

optimize group_norm for ASCEND_NPU (#1154)

Browse Source 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
This commit is contained in:
sunyi0505
2026-04-13 14:48:32 +00:00
committed by GitHub
parent b7025a7027
commit 76a0821780
1 changed files with 409 additions and 372 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/ops/backends/_ascend/ops/group_norm.py
+409 -372
View File
@@ -4,444 +4,471 @@ import triton.language as tl
from triton.language.math import rsqrt
from liger_kernel.ops.backends._ascend.ub_manager import compute_default_tiling_strategy
from liger_kernel.ops.utils import ensure_contiguous
from liger_kernel.ops.utils import get_npu_core_count
# -----------------------------------------------------------------------------
# Kernels (2D row/col tiling + persistent programs)
# -----------------------------------------------------------------------------
MAX_FUSED_SIZE = 16384
MAX_FUSED_FORWARD_SIZE = 4096
GRID_OVERSUB_FACTOR = 8
REDUCE_BATCH_CHUNK = 32
MAX_GRID_SIZE = 65535
def _select_reduce_block_size(hidden_size: int) -> int:
if hidden_size >= 2048:
return 1024
if hidden_size >= 1024:
return 512
if hidden_size >= 512:
return 256
return 128
def _group_norm_forward_stats_block_size(hidden_size: int, element_size: int) -> int:
vv_alignment = 32
required = vv_alignment // element_size
block_h = _select_reduce_block_size(hidden_size)
if hidden_size <= MAX_FUSED_SIZE:
block_h = max(block_h, triton.next_power_of_2(hidden_size))
block_h = min(block_h, 1024)
block_h = max(block_h, required)
return block_h
def _group_norm_forward_stats_batch_block_size(hidden_size: int) -> int:
block_h = min(1024, triton.next_power_of_2(hidden_size))
return max(1, min(8, 4096 // block_h))
def _group_norm_forward_affine_spatial_block_size(hidden_size_per_channel: int) -> int:
return min(256, max(32, triton.next_power_of_2(hidden_size_per_channel)))
def _group_norm_forward_affine_channel_block_size(channels_per_group: int, block_h: int) -> int:
max_block_ch = min(8, triton.next_power_of_2(channels_per_group))
target_block_ch = max(1, 1024 // block_h)
block_ch = 1
while block_ch * 2 <= max_block_ch and block_ch * 2 <= target_block_ch:
block_ch *= 2
return block_ch
def _group_norm_forward_affine_batch_block_size(block_h: int, block_ch: int) -> int:
tile_elems = block_h * block_ch
return max(1, min(8, 2048 // tile_elems))
def _group_norm_forward_launch_config(num_tiles: int, batch_size: int, block_batch: int) -> tuple[int, int]:
batch_blocks = triton.cdiv(batch_size, block_batch)
num_cores = get_npu_core_count()
if num_tiles * batch_blocks >= num_cores * GRID_OVERSUB_FACTOR:
return block_batch, batch_blocks
target_batch_blocks = min(batch_size, max(batch_blocks, triton.cdiv(num_cores * GRID_OVERSUB_FACTOR, num_tiles)))
block_batch = max(1, triton.cdiv(batch_size, target_batch_blocks))
return block_batch, triton.cdiv(batch_size, block_batch)
def _group_norm_backward_spatial_block_size(hidden_size_per_channel: int) -> int:
return min(1024, triton.next_power_of_2(hidden_size_per_channel))
def _group_norm_backward_batch_block_size(hidden_size_per_channel: int) -> int:
block_h = _group_norm_backward_spatial_block_size(hidden_size_per_channel)
return max(1, min(8, 4096 // block_h))
@triton.jit
def _group_norm_forward_kernel(
Y_ptr, # pointer to output, shape (B, G, hidden_size)
Y_row_stride, # stride of each batch row in Y
Y_col_stride, # stride of each group row in Y
X_ptr, # pointer to input, shape (B, G, hidden_size)
X_row_stride, # stride of each batch row in X
X_col_stride, # stride of each group row in X
Mean_ptr, # pointer to mean output, shape (B, G)
Mean_row_stride, # stride of each batch row in Mean
Mean_col_stride, # stride of each group row in Mean
RSTD_ptr, # pointer to rstd output, shape (B, G)
RSTD_row_stride, # stride of each batch row in RSTD
RSTD_col_stride, # stride of each group row in RSTD
W_ptr, # pointer to affine scale weights, shape (C)
B_ptr, # pointer to affine bias weights, shape (C)
n_rows, # total logical rows = B * G
hidden_size,
channels_per_group,
def _group_norm_forward_single_task_kernel(
Y_ptr,
X_ptr,
Mean_ptr,
RSTD_ptr,
W_ptr,
B_ptr,
num_groups,
SINGLE_CHANNEL_TILE: tl.constexpr,
total_tasks,
hidden_size: tl.constexpr,
channels_per_group: tl.constexpr,
eps,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_H: tl.constexpr,
BLOCK_CH: tl.constexpr,
):
hidden_size_per_channel: tl.constexpr = hidden_size // channels_per_group
pid = tl.program_id(0)
num_progs = tl.num_programs(0)
grid_stride = num_progs
grid_m = tl.cdiv(n_rows, BLOCK_SIZE_M)
num_col_blocks = tl.cdiv(hidden_size, BLOCK_SIZE_N)
hidden_size_per_channel = hidden_size // channels_per_group
hidden_size_inv = 1.0 / hidden_size
row_offsets = tl.arange(0, BLOCK_SIZE_M)
col_offsets_base = tl.arange(0, BLOCK_SIZE_N)
ch_offsets = tl.arange(0, BLOCK_CH)
h_offsets = tl.arange(0, BLOCK_H)
ch_mask = ch_offsets < channels_per_group
h_mask = h_offsets < hidden_size_per_channel
mask = ch_mask[:, None] & h_mask[None, :]
# Persistent-program loop over row tiles.
for block_m in tl.range(pid, grid_m, num_progs):
row_idx = block_m * BLOCK_SIZE_M + row_offsets
row_mask = row_idx < n_rows
batch_idx = row_idx // num_groups
group_idx = row_idx % num_groups
for task_iter in range(0, total_tasks, grid_stride):
pid_task = task_iter + pid
row_sum = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
row_square_sum = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
pid_group = pid_task % num_groups
task_base = pid_task * hidden_size
ptrs = task_base + ch_offsets[:, None] * hidden_size_per_channel + h_offsets[None, :]
# Pass 1: accumulate E[x] and E[x^2] for each row tile.
for cb in range(num_col_blocks):
col_offsets = cb * BLOCK_SIZE_N + col_offsets_base
col_mask = col_offsets < hidden_size
mask = row_mask[:, None] & col_mask[None, :]
X_ptrs = (
X_ptr + batch_idx[:, None] * X_row_stride + group_idx[:, None] * X_col_stride + col_offsets[None, :]
)
X_block = tl.load(X_ptrs, mask=mask, other=0.0).to(tl.float32)
row_sum += tl.sum(X_block, axis=1)
row_square_sum += tl.sum(X_block * X_block, axis=1)
X_vals = tl.load(X_ptr + ptrs, mask=mask, other=0.0).to(tl.float32)
sum_acc = tl.sum(tl.sum(X_vals, axis=1), axis=0)
sum_sq_acc = tl.sum(tl.sum(X_vals * X_vals, axis=1), axis=0)
mean = row_sum * hidden_size_inv
var = row_square_sum * hidden_size_inv - mean * mean
rstd = rsqrt(tl.maximum(var, 0.0) + eps)
mean = sum_acc / hidden_size
variance = tl.maximum(sum_sq_acc / hidden_size - mean * mean, 0.0)
rstd = rsqrt(variance + eps)
mean_ptrs = Mean_ptr + batch_idx * Mean_row_stride + group_idx * Mean_col_stride
rstd_ptrs = RSTD_ptr + batch_idx * RSTD_row_stride + group_idx * RSTD_col_stride
tl.store(mean_ptrs, mean, mask=row_mask)
tl.store(rstd_ptrs, rstd, mask=row_mask)
group_ch_start = pid_group * channels_per_group
W_vals = tl.load(W_ptr + group_ch_start + ch_offsets, mask=ch_mask, other=1.0).to(tl.float32)
B_vals = tl.load(B_ptr + group_ch_start + ch_offsets, mask=ch_mask, other=0.0).to(tl.float32)
Y_vals = (X_vals - mean) * rstd
Y_vals = Y_vals * W_vals[:, None] + B_vals[:, None]
tl.store(Y_ptr + ptrs, Y_vals.to(Y_ptr.dtype.element_ty), mask=mask)
# Pass 2: normalize + affine transform.
# SINGLE_CHANNEL_TILE indicates the current column tile maps to one channel,
# so W/B can be loaded once per row and broadcast to the tile.
for cb in range(num_col_blocks):
col_offsets = cb * BLOCK_SIZE_N + col_offsets_base
col_mask = col_offsets < hidden_size
mask = row_mask[:, None] & col_mask[None, :]
X_ptrs = (
X_ptr + batch_idx[:, None] * X_row_stride + group_idx[:, None] * X_col_stride + col_offsets[None, :]
)
X_block = tl.load(X_ptrs, mask=mask, other=0.0).to(tl.float32)
if SINGLE_CHANNEL_TILE:
local_channel = (cb * BLOCK_SIZE_N) // hidden_size_per_channel
global_channel = group_idx * channels_per_group + local_channel
W_block = tl.load(W_ptr + global_channel, mask=row_mask, other=0.0).to(tl.float32)[:, None]
B_block = tl.load(B_ptr + global_channel, mask=row_mask, other=0.0).to(tl.float32)[:, None]
else:
local_channel = col_offsets // hidden_size_per_channel
global_channel = group_idx[:, None] * channels_per_group + local_channel[None, :]
W_block = tl.load(W_ptr + global_channel, mask=mask, other=0.0).to(tl.float32)
B_block = tl.load(B_ptr + global_channel, mask=mask, other=0.0).to(tl.float32)
Y_block = (X_block - mean[:, None]) * rstd[:, None] * W_block + B_block
Y_ptrs = (
Y_ptr + batch_idx[:, None] * Y_row_stride + group_idx[:, None] * Y_col_stride + col_offsets[None, :]
)
tl.store(Y_ptrs, Y_block, mask=mask)
tl.store(Mean_ptr + pid_task, mean)
tl.store(RSTD_ptr + pid_task, rstd)
@triton.jit
def _group_norm_backward_kernel(
X_ptr, # pointer to input, shape (B, G, hidden_size)
X_row_stride, # stride of each batch row in X
X_col_stride, # stride of each group row in X
W_ptr, # pointer to affine scale weights, shape (C)
Mean_ptr, # pointer to saved group mean, shape (B, G)
Mean_row_stride, # stride of each batch row in Mean
Mean_col_stride, # stride of each group row in Mean
RSTD_ptr, # pointer to saved reciprocal std, shape (B, G)
DX_ptr, # pointer to input gradients, shape (B, G, hidden_size)
DW_scratch_ptr, # pointer to scratch buffer for dW partial sums, shape (grid, C)
DW_scratch_stride, # row stride for DW_scratch
DB_scratch_ptr, # pointer to scratch buffer for dB partial sums, shape (grid, C)
DB_scratch_stride, # row stride for DB_scratch
DY_ptr, # pointer to upstream gradients, shape (B, G, hidden_size)
DY_row_stride, # stride of each batch row in DY
DY_col_stride, # stride of each group row in DY
n_rows, # total logical rows = B * G
hidden_size,
channels_per_group,
def _group_norm_forward_stats_kernel(
X_ptr,
Mean_ptr,
RSTD_ptr,
batch_size,
num_groups,
SINGLE_CHANNEL_TILE: tl.constexpr,
COMPUTE_PARAM_GRAD: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
hidden_size: tl.constexpr,
eps,
BLOCK_SIZE_H: tl.constexpr,
BLOCK_BATCH: tl.constexpr,
):
pid_group = tl.program_id(0)
pid_batch_block = tl.program_id(1)
batch_offsets = pid_batch_block * BLOCK_BATCH + tl.arange(0, BLOCK_BATCH)
batch_mask = batch_offsets < batch_size
task_offsets = batch_offsets * num_groups + pid_group
inv_hidden_size = 1.0 / hidden_size
sum_acc = tl.zeros([BLOCK_BATCH], dtype=tl.float32)
sum_sq_acc = tl.zeros([BLOCK_BATCH], dtype=tl.float32)
h_offsets = tl.arange(0, BLOCK_SIZE_H)
for start in tl.range(0, hidden_size, BLOCK_SIZE_H):
idx = start + h_offsets
h_mask = idx < hidden_size
mask = batch_mask[:, None] & h_mask[None, :]
ptrs = task_offsets[:, None] * hidden_size + idx[None, :]
X_vals = tl.load(X_ptr + ptrs, mask=mask, other=0.0).to(tl.float32)
sum_acc += tl.sum(X_vals, axis=1)
sum_sq_acc += tl.sum(X_vals * X_vals, axis=1)
mean = sum_acc * inv_hidden_size
variance = tl.maximum(sum_sq_acc * inv_hidden_size - mean * mean, 0.0)
rstd = rsqrt(variance + eps)
tl.store(Mean_ptr + task_offsets, mean, mask=batch_mask)
tl.store(RSTD_ptr + task_offsets, rstd, mask=batch_mask)
@triton.jit
def _group_norm_forward_affine_kernel(
Y_ptr,
X_ptr,
Mean_ptr,
RSTD_ptr,
W_ptr,
B_ptr,
batch_size,
num_groups,
hidden_size: tl.constexpr,
channels_per_group: tl.constexpr,
channel_blocks,
BLOCK_SIZE_H: tl.constexpr,
BLOCK_SIZE_CH: tl.constexpr,
BLOCK_BATCH: tl.constexpr,
):
hidden_size_per_channel: tl.constexpr = hidden_size // channels_per_group
pid_tile = tl.program_id(0)
pid_batch_block = tl.program_id(1)
pid_group = pid_tile // channel_blocks
pid_channel_block = pid_tile - pid_group * channel_blocks
batch_offsets = pid_batch_block * BLOCK_BATCH + tl.arange(0, BLOCK_BATCH)
batch_mask = batch_offsets < batch_size
task_offsets = batch_offsets * num_groups + pid_group
mean = tl.load(Mean_ptr + task_offsets, mask=batch_mask, other=0.0).to(tl.float32)
rstd = tl.load(RSTD_ptr + task_offsets, mask=batch_mask, other=0.0).to(tl.float32)
ch_offsets = pid_channel_block * BLOCK_SIZE_CH + tl.arange(0, BLOCK_SIZE_CH)
ch_mask = ch_offsets < channels_per_group
group_ch_offsets = pid_group * channels_per_group + ch_offsets
W_vals = tl.load(W_ptr + group_ch_offsets, mask=ch_mask, other=1.0).to(tl.float32)
B_vals = tl.load(B_ptr + group_ch_offsets, mask=ch_mask, other=0.0).to(tl.float32)
h_offsets = tl.arange(0, BLOCK_SIZE_H)
for start in tl.range(0, hidden_size_per_channel, BLOCK_SIZE_H):
idx = start + h_offsets
h_mask = idx < hidden_size_per_channel
mask = batch_mask[:, None, None] & ch_mask[None, :, None] & h_mask[None, None, :]
ptrs = (
task_offsets[:, None, None] * hidden_size
+ ch_offsets[None, :, None] * hidden_size_per_channel
+ idx[None, None, :]
)
X_vals = tl.load(X_ptr + ptrs, mask=mask, other=0.0).to(tl.float32)
Y_vals = (X_vals - mean[:, None, None]) * rstd[:, None, None]
Y_vals = Y_vals * W_vals[None, :, None] + B_vals[None, :, None]
tl.store(Y_ptr + ptrs, Y_vals.to(Y_ptr.dtype.element_ty), mask=mask)
@triton.jit
def _group_norm_backward_dx_dwdb_kernel(
X_ptr,
W_ptr,
Mean_ptr,
RSTD_ptr,
DX_ptr,
DW_partial_ptr,
DB_partial_ptr,
DY_ptr,
batch_size,
num_groups,
hidden_size: tl.constexpr,
channels_per_group: tl.constexpr,
BLOCK_H: tl.constexpr,
BLOCK_BATCH: tl.constexpr,
):
hidden_size_per_channel: tl.constexpr = hidden_size // channels_per_group
pid_group = tl.program_id(0)
pid_batch_block = tl.program_id(1)
batch_offsets = pid_batch_block * BLOCK_BATCH + tl.arange(0, BLOCK_BATCH)
batch_mask = batch_offsets < batch_size
task_offsets = batch_offsets * num_groups + pid_group
h_offsets = tl.arange(0, BLOCK_H)
mean = tl.load(Mean_ptr + task_offsets, mask=batch_mask, other=0.0).to(tl.float32)
rstd = tl.load(RSTD_ptr + task_offsets, mask=batch_mask, other=0.0).to(tl.float32)
inv_hidden_size = 1.0 / hidden_size
group_ch_start = pid_group * channels_per_group
c1 = tl.zeros([BLOCK_BATCH], dtype=tl.float32)
c2 = tl.zeros([BLOCK_BATCH], dtype=tl.float32)
for local_ch in range(0, channels_per_group):
channel_base = local_ch * hidden_size_per_channel
W = tl.load(W_ptr + group_ch_start + local_ch).to(tl.float32)
dW_acc = 0.0
dB_acc = 0.0
for start in range(0, hidden_size_per_channel, BLOCK_H):
idx = start + h_offsets
h_mask = idx < hidden_size_per_channel
mask = batch_mask[:, None] & h_mask[None, :]
ptrs = task_offsets[:, None] * hidden_size + channel_base + idx[None, :]
X_vals = tl.load(X_ptr + ptrs, mask=mask, other=0.0).to(tl.float32)
DY_vals = tl.load(DY_ptr + ptrs, mask=mask, other=0.0).to(tl.float32)
x_hat = (X_vals - mean[:, None]) * rstd[:, None]
wdy = W * DY_vals
dW_acc += tl.sum(tl.sum(tl.where(mask, DY_vals * x_hat, 0.0), axis=1), axis=0)
dB_acc += tl.sum(tl.sum(tl.where(mask, DY_vals, 0.0), axis=1), axis=0)
c1 += tl.sum(tl.where(mask, x_hat * wdy, 0.0), axis=1)
c2 += tl.sum(tl.where(mask, wdy, 0.0), axis=1)
partial_offset = pid_batch_block * (num_groups * channels_per_group) + group_ch_start + local_ch
tl.store(DW_partial_ptr + partial_offset, dW_acc)
tl.store(DB_partial_ptr + partial_offset, dB_acc)
c1 *= inv_hidden_size
c2 *= inv_hidden_size
for local_ch in range(0, channels_per_group):
channel_base = local_ch * hidden_size_per_channel
W = tl.load(W_ptr + group_ch_start + local_ch).to(tl.float32)
for start in range(0, hidden_size_per_channel, BLOCK_H):
idx = start + h_offsets
h_mask = idx < hidden_size_per_channel
mask = batch_mask[:, None] & h_mask[None, :]
ptrs = task_offsets[:, None] * hidden_size + channel_base + idx[None, :]
X_vals = tl.load(X_ptr + ptrs, mask=mask, other=0.0).to(tl.float32)
DY_vals = tl.load(DY_ptr + ptrs, mask=mask, other=0.0).to(tl.float32)
x_hat = (X_vals - mean[:, None]) * rstd[:, None]
DX_vals = (W * DY_vals - (x_hat * c1[:, None] + c2[:, None])) * rstd[:, None]
tl.store(DX_ptr + ptrs, DX_vals.to(DX_ptr.dtype.element_ty), mask=mask)
@triton.jit
def _group_norm_reduce_param_grads_kernel(
DW_partial_ptr,
DB_partial_ptr,
DW_ptr,
DB_ptr,
num_partial_rows,
num_channels,
BLOCK_SIZE_CH: tl.constexpr,
BLOCK_SIZE_ROWS: tl.constexpr,
):
pid = tl.program_id(0)
num_progs = tl.num_programs(0)
ch_offsets = pid * BLOCK_SIZE_CH + tl.arange(0, BLOCK_SIZE_CH)
ch_mask = ch_offsets < num_channels
row_offsets = tl.arange(0, BLOCK_SIZE_ROWS)
grid_m = tl.cdiv(n_rows, BLOCK_SIZE_M)
num_col_blocks = tl.cdiv(hidden_size, BLOCK_SIZE_N)
hidden_size_per_channel = hidden_size // channels_per_group
N_inv = 1.0 / hidden_size
row_offsets = tl.arange(0, BLOCK_SIZE_M)
col_offsets_base = tl.arange(0, BLOCK_SIZE_N)
dW_acc = tl.zeros([BLOCK_SIZE_CH], dtype=tl.float32)
dB_acc = tl.zeros([BLOCK_SIZE_CH], dtype=tl.float32)
if COMPUTE_PARAM_GRAD:
DW_scratch_base = DW_scratch_ptr + pid * DW_scratch_stride
DB_scratch_base = DB_scratch_ptr + pid * DB_scratch_stride
for row_start in range(0, num_partial_rows, BLOCK_SIZE_ROWS):
rows = row_start + row_offsets
row_mask = rows < num_partial_rows
mask = row_mask[:, None] & ch_mask[None, :]
ptrs = rows[:, None] * num_channels + ch_offsets[None, :]
dW_vals = tl.load(DW_partial_ptr + ptrs, mask=mask, other=0.0)
dB_vals = tl.load(DB_partial_ptr + ptrs, mask=mask, other=0.0)
dW_acc += tl.sum(dW_vals, axis=0)
dB_acc += tl.sum(dB_vals, axis=0)
# Persistent-program loop over row tiles.
for block_m in tl.range(pid, grid_m, num_progs):
row_idx = block_m * BLOCK_SIZE_M + row_offsets
row_mask = row_idx < n_rows
batch_idx = row_idx // num_groups
group_idx = row_idx % num_groups
mean = tl.load(
Mean_ptr + batch_idx * Mean_row_stride + group_idx * Mean_col_stride,
mask=row_mask,
other=0.0,
).to(tl.float32)
rstd = tl.load(
RSTD_ptr + batch_idx * Mean_row_stride + group_idx * Mean_col_stride,
mask=row_mask,
other=0.0,
).to(tl.float32)
sum_x_hat_wdy = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
sum_wdy = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
# Pass 1: compute row-wise reduction terms (c1, c2).
for cb in range(num_col_blocks):
col_offsets = cb * BLOCK_SIZE_N + col_offsets_base
col_mask = col_offsets < hidden_size
mask = row_mask[:, None] & col_mask[None, :]
X_ptrs = (
X_ptr + batch_idx[:, None] * X_row_stride + group_idx[:, None] * X_col_stride + col_offsets[None, :]
)
DY_ptrs = (
DY_ptr + batch_idx[:, None] * DY_row_stride + group_idx[:, None] * DY_col_stride + col_offsets[None, :]
)
X_block = tl.load(X_ptrs, mask=mask, other=0.0).to(tl.float32)
DY_block = tl.load(DY_ptrs, mask=mask, other=0.0).to(tl.float32)
if SINGLE_CHANNEL_TILE:
local_channel = (cb * BLOCK_SIZE_N) // hidden_size_per_channel
global_channel = group_idx * channels_per_group + local_channel
W_block = tl.load(W_ptr + global_channel, mask=row_mask, other=0.0).to(tl.float32)[:, None]
else:
local_channel = col_offsets // hidden_size_per_channel
global_channel = group_idx[:, None] * channels_per_group + local_channel[None, :]
W_block = tl.load(W_ptr + global_channel, mask=mask, other=0.0).to(tl.float32)
x_hat = (X_block - mean[:, None]) * rstd[:, None]
wdy = W_block * DY_block
sum_x_hat_wdy += tl.sum(tl.where(mask, x_hat * wdy, 0.0), axis=1)
sum_wdy += tl.sum(tl.where(mask, wdy, 0.0), axis=1)
c1 = sum_x_hat_wdy * N_inv
c2 = sum_wdy * N_inv
# Pass 2: compute DX and optionally accumulate DW/DB.
# COMPUTE_PARAM_GRAD=False is used to skip expensive atomics in cases
# where host-side dense reduction is faster/more stable.
for cb in range(num_col_blocks):
col_offsets = cb * BLOCK_SIZE_N + col_offsets_base
col_mask = col_offsets < hidden_size
mask = row_mask[:, None] & col_mask[None, :]
X_ptrs = (
X_ptr + batch_idx[:, None] * X_row_stride + group_idx[:, None] * X_col_stride + col_offsets[None, :]
)
DY_ptrs = (
DY_ptr + batch_idx[:, None] * DY_row_stride + group_idx[:, None] * DY_col_stride + col_offsets[None, :]
)
X_block = tl.load(X_ptrs, mask=mask, other=0.0).to(tl.float32)
DY_block = tl.load(DY_ptrs, mask=mask, other=0.0).to(tl.float32)
if SINGLE_CHANNEL_TILE:
local_channel = (cb * BLOCK_SIZE_N) // hidden_size_per_channel
global_channel = group_idx * channels_per_group + local_channel
W_block = tl.load(W_ptr + global_channel, mask=row_mask, other=0.0).to(tl.float32)[:, None]
else:
local_channel = col_offsets // hidden_size_per_channel
global_channel = group_idx[:, None] * channels_per_group + local_channel[None, :]
W_block = tl.load(W_ptr + global_channel, mask=mask, other=0.0).to(tl.float32)
x_hat = (X_block - mean[:, None]) * rstd[:, None]
wdy = W_block * DY_block
DX_block = (wdy - (x_hat * c1[:, None] + c2[:, None])) * rstd[:, None]
DX_ptrs = (
DX_ptr + batch_idx[:, None] * X_row_stride + group_idx[:, None] * X_col_stride + col_offsets[None, :]
)
tl.store(DX_ptrs, DX_block.to(X_ptr.dtype.element_ty), mask=mask)
if COMPUTE_PARAM_GRAD:
if SINGLE_CHANNEL_TILE:
dW_partial = tl.sum(tl.where(mask, DY_block * x_hat, 0.0), axis=1)
dB_partial = tl.sum(tl.where(mask, DY_block, 0.0), axis=1)
tl.atomic_add(DW_scratch_base + global_channel, dW_partial, mask=row_mask)
tl.atomic_add(DB_scratch_base + global_channel, dB_partial, mask=row_mask)
# -----------------------------------------------------------------------------
# Helper: call compute_default_tiling_strategy
# -----------------------------------------------------------------------------
def get_optimal_block_size(n_rows, dtype_size, BLOCK_SIZE_N, is_backward: bool = False):
# Backward keeps larger live-state than forward in this kernel.
multiplier = 7.0 if is_backward else 6.0
# Use fp32-size as conservative UB estimate for tiling.
dtype_size = max(dtype_size, 4)
tile_shapes = compute_default_tiling_strategy(
safety_margin=0.9,
dtype_size=dtype_size,
memory_multiplier=multiplier,
shapes=((n_rows, BLOCK_SIZE_N),),
tiling_dims=(0,),
)
if tile_shapes and len(tile_shapes) > 0:
return tile_shapes[0][0]
return triton.next_power_of_2(min(128, n_rows))
tl.store(DW_ptr + ch_offsets, dW_acc, mask=ch_mask)
tl.store(DB_ptr + ch_offsets, dB_acc, mask=ch_mask)
def group_norm_forward(X, num_channels, num_groups, W, B, eps):
shape = X.shape
batch_size = shape[0]
channels_per_group = num_channels // num_groups
# Reshape X so that the mean / std are computed across each group
X = X.view(batch_size, num_groups, -1).contiguous()
X = X.view(batch_size, num_groups, -1)
hidden_size = X.shape[-1]
hidden_size_per_channel = hidden_size // channels_per_group
n_rows = batch_size * num_groups
Y = torch.empty_like(X)
total_tasks = batch_size * num_groups
Mean = torch.empty(total_tasks, dtype=torch.float32, device=X.device)
RSTD = torch.empty(total_tasks, dtype=torch.float32, device=X.device)
BLOCK_SIZE_N = min(128, triton.next_power_of_2(hidden_size))
BLOCK_SIZE_M = get_optimal_block_size(n_rows, X.element_size(), BLOCK_SIZE_N)
if hidden_size <= MAX_FUSED_FORWARD_SIZE:
block_h = _group_norm_backward_spatial_block_size(hidden_size_per_channel)
block_ch = triton.next_power_of_2(channels_per_group)
num_cores = get_npu_core_count()
grid_size = min(total_tasks, num_cores * 2, MAX_GRID_SIZE)
_group_norm_forward_single_task_kernel[(grid_size,)](
Y,
X,
Mean,
RSTD,
W,
B,
num_groups,
total_tasks,
hidden_size,
channels_per_group,
eps,
BLOCK_H=block_h,
BLOCK_CH=block_ch,
)
return Y.view(*shape), X.view(*shape), Mean.view(batch_size, num_groups), RSTD.view(batch_size, num_groups)
# Fast path condition: each column tile must lie entirely inside one channel
# segment of length `hidden_size_per_channel`.
#
# Layout of a row:
# | channel0 | channel1 | channel2 | ...
# |----Hc----|----Hc----|
# Hc = hidden_size_per_channel
#
# The kernel processes tiles of shape (BLOCK_SIZE_M, BLOCK_SIZE_N).
# Channel boundaries exist only along the column dimension, because
# each row corresponds to a different (batch, group).
#
# Therefore only BLOCK_SIZE_N matters for whether a tile crosses
# channel boundaries; BLOCK_SIZE_M does not affect channel mapping.
#
# If BLOCK_SIZE_N divides Hc and is <= Hc, each column tile belongs
# to exactly one channel. In that case W/B can be loaded once and
# broadcast across the tile (fast path).
#
# Otherwise a tile may span multiple channels, requiring per-element
# channel index computation and parameter loads (slow path).
single_channel_tile = BLOCK_SIZE_N <= hidden_size_per_channel and hidden_size_per_channel % BLOCK_SIZE_N == 0
stats_block_h = _group_norm_forward_stats_block_size(hidden_size, X.element_size())
stats_block_batch = _group_norm_forward_stats_batch_block_size(hidden_size)
stats_block_batch, stats_grid_batch = _group_norm_forward_launch_config(num_groups, batch_size, stats_block_batch)
_group_norm_forward_stats_kernel[(num_groups, stats_grid_batch)](
X,
Mean,
RSTD,
batch_size,
num_groups,
hidden_size,
eps,
BLOCK_SIZE_H=stats_block_h,
BLOCK_BATCH=stats_block_batch,
)
num_cores = get_npu_core_count()
grid = min(num_cores, triton.cdiv(n_rows, BLOCK_SIZE_M))
Y = torch.empty((batch_size, num_groups, hidden_size), dtype=X.dtype, device=X.device)
Mean = torch.empty((batch_size, num_groups), dtype=X.dtype, device=X.device)
RSTD = torch.empty((batch_size, num_groups), dtype=X.dtype, device=X.device)
_group_norm_forward_kernel[(grid,)](
affine_block_h = _group_norm_forward_affine_spatial_block_size(hidden_size_per_channel)
affine_block_ch = _group_norm_forward_affine_channel_block_size(channels_per_group, affine_block_h)
channel_blocks = triton.cdiv(channels_per_group, affine_block_ch)
affine_block_batch = _group_norm_forward_affine_batch_block_size(affine_block_h, affine_block_ch)
affine_block_batch, affine_grid_batch = _group_norm_forward_launch_config(
num_groups * channel_blocks,
batch_size,
affine_block_batch,
)
_group_norm_forward_affine_kernel[(num_groups * channel_blocks, affine_grid_batch)](
Y,
Y.stride(0),
Y.stride(1),
X,
X.stride(0),
X.stride(1),
Mean,
Mean.stride(0),
Mean.stride(1),
RSTD,
RSTD.stride(0),
RSTD.stride(1),
W,
B,
n_rows,
batch_size,
num_groups,
hidden_size,
channels_per_group,
num_groups,
SINGLE_CHANNEL_TILE=single_channel_tile,
eps=eps,
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE_N,
channel_blocks,
BLOCK_SIZE_H=affine_block_h,
BLOCK_SIZE_CH=affine_block_ch,
BLOCK_BATCH=affine_block_batch,
)
return Y.view(*shape), X.view(*shape), Mean, RSTD
return Y.view(*shape), X.view(*shape), Mean.view(batch_size, num_groups), RSTD.view(batch_size, num_groups)
def group_norm_backward(dY, X, W, B, Mean, RSTD, num_channels, num_groups):
def group_norm_backward(dY, X, W, Mean, RSTD, num_channels, num_groups, bias_dtype):
shape = dY.shape
batch_size = shape[0]
channels_per_group = num_channels // num_groups
X_grouped = X.view(batch_size, num_groups, -1)
dY_grouped = dY.view(batch_size, num_groups, -1)
hidden_size = dY_grouped.shape[-1]
dY = dY.view(batch_size, num_groups, -1)
X = X.view(batch_size, num_groups, -1)
hidden_size = dY.shape[-1]
hidden_size_per_channel = hidden_size // channels_per_group
n_rows = batch_size * num_groups
BLOCK_SIZE_N = min(128, triton.next_power_of_2(hidden_size))
BLOCK_SIZE_M = get_optimal_block_size(
n_rows,
X.element_size(),
BLOCK_SIZE_N,
is_backward=True,
)
DX = torch.empty_like(dY)
block_h = _group_norm_backward_spatial_block_size(hidden_size_per_channel)
block_batch = _group_norm_backward_batch_block_size(hidden_size_per_channel)
num_partial_rows = triton.cdiv(batch_size, block_batch)
# Same condition as forward:
# if true, each BLOCK_SIZE_N tile maps cleanly to one channel segment.
single_channel_tile = BLOCK_SIZE_N <= hidden_size_per_channel and hidden_size_per_channel % BLOCK_SIZE_N == 0
DW_partial = torch.empty((num_partial_rows, num_channels), dtype=torch.float32, device=W.device)
DB_partial = torch.empty((num_partial_rows, num_channels), dtype=torch.float32, device=W.device)
DW_accum = torch.empty(num_channels, dtype=torch.float32, device=W.device)
DB_accum = torch.empty(num_channels, dtype=torch.float32, device=W.device)
num_cores = get_npu_core_count()
grid = min(num_cores, triton.cdiv(n_rows, BLOCK_SIZE_M))
# For non-single-channel tiles, per-element atomic updates are costly.
# In that case, kernel computes DX only and DW/DB are reduced on host side.
compute_param_grad = single_channel_tile
DX = torch.empty((batch_size, num_groups, hidden_size), dtype=X.dtype, device=X.device)
if compute_param_grad:
DW_scratch = torch.zeros((grid, num_channels), dtype=torch.float32, device=W.device)
DB_scratch = torch.zeros((grid, num_channels), dtype=torch.float32, device=W.device)
else:
# Not used when COMPUTE_PARAM_GRAD=False.
# Intentionally set to None to enforce fail-fast behavior if accidentally accessed.
DW_scratch = None
DB_scratch = None
_group_norm_backward_kernel[(grid,)](
X_grouped,
X_grouped.stride(0),
X_grouped.stride(1),
_group_norm_backward_dx_dwdb_kernel[(num_groups, num_partial_rows)](
X,
W,
Mean,
Mean.stride(0),
Mean.stride(1),
RSTD,
DX,
DW_scratch,
0 if not compute_param_grad else DW_scratch.stride(0),
DB_scratch,
0 if not compute_param_grad else DB_scratch.stride(0),
dY_grouped,
dY_grouped.stride(0),
dY_grouped.stride(1),
n_rows,
DW_partial,
DB_partial,
dY,
batch_size,
num_groups,
hidden_size,
channels_per_group,
num_groups,
SINGLE_CHANNEL_TILE=single_channel_tile,
COMPUTE_PARAM_GRAD=compute_param_grad,
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE_N,
BLOCK_H=block_h,
BLOCK_BATCH=block_batch,
)
# Precision note:
# - In-kernel atomic_add on floating-point values is order-dependent under parallel
# scheduling (non-associative summation), which can introduce run-to-run numerical
# differences in DW/DB for contention-heavy shapes.
# - Host-side dense reduction provides a more stable accumulation pattern for these
# difficult layouts.
if compute_param_grad:
DW = DW_scratch.sum(dim=0).to(W.dtype)
DB = DB_scratch.sum(dim=0).to(W.dtype)
else:
# Fallback path to avoid severe atomic contention when SINGLE_CHANNEL_TILE=False.
# Layout: [B, G, hidden_size] -> [B, G, C_per_G, hidden_per_channel]
X4 = X_grouped.reshape(batch_size, num_groups, channels_per_group, hidden_size_per_channel).to(torch.float32)
dY4 = dY_grouped.reshape(batch_size, num_groups, channels_per_group, hidden_size_per_channel).to(torch.float32)
mean4 = Mean.reshape(batch_size, num_groups, 1, 1).to(torch.float32)
rstd4 = RSTD.reshape(batch_size, num_groups, 1, 1).to(torch.float32)
reduce_block_ch = min(256, triton.next_power_of_2(num_channels))
_group_norm_reduce_param_grads_kernel[(triton.cdiv(num_channels, reduce_block_ch),)](
DW_partial,
DB_partial,
DW_accum,
DB_accum,
num_partial_rows,
num_channels,
BLOCK_SIZE_CH=reduce_block_ch,
BLOCK_SIZE_ROWS=REDUCE_BATCH_CHUNK,
)
x_hat4 = (X4 - mean4) * rstd4
DW = (dY4 * x_hat4).sum(dim=(0, 3)).reshape(-1).to(W.dtype)
DB = dY4.sum(dim=(0, 3)).reshape(-1).to(W.dtype)
return DX.view(*shape), DW, DB
return DX.view(*shape), DW_accum.to(W.dtype), DB_accum.to(bias_dtype)
class LigerGroupNormFunction(torch.autograd.Function):
"""
Group Normalization autograd function for Ascend NPU.
Forward computes, for each sample/group:
y = (x - mean) * rstd * weight + bias
where:
mean = E[x], rstd = 1 / sqrt(Var[x] + eps)
The kernel uses row/column tiling with persistent programs. Backward computes
input gradients in Triton and computes parameter gradients either via Triton
atomics (fast path) or host-side dense reduction (fallback path), depending
on the tile/channel layout.
"""
@staticmethod
@ensure_contiguous
def forward(
@@ -463,12 +490,22 @@ class LigerGroupNormFunction(torch.autograd.Function):
)
ctx.num_channels = num_channels
ctx.num_groups = num_groups
ctx.save_for_backward(X, affine_scaling_weight, affine_shifting_bias, Mean, RSTD)
ctx.bias_dtype = affine_shifting_bias.dtype
ctx.save_for_backward(X, affine_scaling_weight, Mean, RSTD)
return Y
@staticmethod
@ensure_contiguous
def backward(ctx, dY):
X, W, B, Mean, RSTD = ctx.saved_tensors
DX, DW, DB = group_norm_backward(dY, X, W, B, Mean, RSTD, ctx.num_channels, ctx.num_groups)
X, W, Mean, RSTD = ctx.saved_tensors
DX, DW, DB = group_norm_backward(
dY,
X,
W,
Mean,
RSTD,
ctx.num_channels,
ctx.num_groups,
ctx.bias_dtype,
)
return DX, DW, DB, None, None, None