@triton.jit
def load_scales(a_scale_ptr, b_scale_ptr, ROW_WISE_SCALING: tl.constexpr):
    if ROW_WISE_SCALING:
        return a_scale_ptr, b_scale_ptr
    else:
        a_scale = tl.load(a_scale_ptr)
        b_scale = tl.load(b_scale_ptr)
        return a_scale, b_scale
