@triton.jit
def apply_scaling(accumulator, a_scale, b_scale, ROW_WISE_SCALING: tl.
    constexpr, offs_cm, offs_cn, M, N, stride_a_scale_m, stride_b_scale_n):
    if ROW_WISE_SCALING:
        a_scales = tl.load(a_scale + offs_cm * stride_a_scale_m, mask=
            offs_cm < M, other=0.0)
        b_scales = tl.load(b_scale + offs_cn * stride_b_scale_n, mask=
            offs_cn < N, other=0.0)
        acc_scale = a_scales[:, None] * b_scales[None, :]
    else:
        acc_scale = a_scale * b_scale
    return accumulator * acc_scale
