@triton.jit
def triton_normalization_backward_kernel_2(grad_weight_partial_ptr,
    grad_bias_partial_ptr, grad_weight_ptr, grad_bias_ptr, m, n_cols,
    has_bias: tl.constexpr, accumulate_grad: tl.constexpr, block_size_m: tl
    .constexpr, block_size_n: tl.constexpr):
    pid = tl.program_id(0)
    cols = pid * block_size_n + tl.arange(0, block_size_n)
    grad_weight_partial_sum = tl.zeros((block_size_m, block_size_n), dtype=
        tl.float32)
    if has_bias:
        grad_bias_partial_sum = tl.zeros((block_size_m, block_size_n),
            dtype=tl.float32)
    col_mask = cols < n_cols
    for i in range(0, m, block_size_m):
        rows = i + tl.arange(0, block_size_m)
        mask = (rows[:, None] < m) & (cols[None, :] < n_cols)
        offsets = rows[:, None] * n_cols + cols[None, :]
        grad_weight_partial_sum += tl.load(grad_weight_partial_ptr +
            offsets, mask=mask, other=0.0)
        if has_bias:
            grad_bias_partial_sum += tl.load(grad_bias_partial_ptr +
                offsets, mask=mask, other=0.0)
    grad_weight = tl.sum(grad_weight_partial_sum, axis=0)
    if accumulate_grad:
        grad_weight = tl.load(grad_weight_ptr + cols, mask=col_mask
            ) + grad_weight
    tl.store(grad_weight_ptr + cols, grad_weight, mask=col_mask)
    if has_bias:
        grad_bias = tl.sum(grad_bias_partial_sum, axis=0)
        if accumulate_grad:
            grad_bias = tl.load(grad_bias_ptr + cols, mask=col_mask
                ) + grad_bias
        tl.store(grad_bias_ptr + cols, grad_bias, mask=col_mask)
