From 539552670bb1e13e611895cb3b40b61509183a8a Mon Sep 17 00:00:00 2001 From: learning-chip Date: Wed, 30 Sep 2026 10:09:27 +0200 Subject: [PATCH] Port four quant kernels to PTO/VMI Rewrite the per-token, per-block, per-channel, and cast-back SimdVF implementations with logical VMI while preserving their ASC schedules and host contracts. --- tile_kernels/quant/cast_back_asc.py | 135 +++++++------ tile_kernels/quant/per_block_cast_asc.py | 192 +++++++++---------- tile_kernels/quant/per_channel_cast_asc.py | 167 +++++++--------- tile_kernels/quant/per_token_cast_asc.py | 213 ++++++++------------- 4 files changed, 312 insertions(+), 395 deletions(-) diff --git a/tile_kernels/quant/cast_back_asc.py b/tile_kernels/quant/cast_back_asc.py index 0ae3c0c7..c6ee5e92 100644 --- a/tile_kernels/quant/cast_back_asc.py +++ b/tile_kernels/quant/cast_back_asc.py @@ -2,14 +2,14 @@ import tilelang from tilelang.ascend import language as T -from tilelang.ascend.language import simd as S +from tilelang.ascend.language import vmi as V from tile_kernels.config import get_max_ub_per_vector_core from tile_kernels.quant.common import CastInputConfig, get_packed_ue8m0_pack_factor, get_sf_shape from tile_kernels.utils import align, ceil_div -@tilelang.jit() +@tilelang.jit(target='pto', pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) def get_cast_back_kernel_asc( hidden: int, in_config: CastInputConfig, @@ -59,6 +59,7 @@ def get_cast_back_kernel_asc( sf_ub_shape = (num_sf_cols, num_sf_rows) if keep_unpacked_col_major else (num_sf_rows, num_sf_cols) if not in_config.use_packed_ue8m0: sf_ub_shape = (sf_ub_shape[0], align(sf_ub_shape[1], 32 // in_config.sf_dtype.bytes)) + num_sf_rows_padded = sf_ub_shape[1] @T.macro def load_scale_tile(x_sf_raw, sf_raw_ub, sf_ub, token_base, channel_base): @@ -74,107 +75,101 @@ def load_scale_tile(x_sf_raw, sf_raw_ub, sf_ub, token_base, channel_base): if use_e2b_scale: assert num_sf_rows * num_sf_cols % 128 == 0 with T.SimdVF(): - lane_indices = T.reinterpret(S.vci(0, T.int16), 'uint16x128') - byte_mask = S.vdup(0x00FF, T.uint16) + mask = V.create_mask(128, size=128) + lanes = V.vinterpret_cast(V.vci(T.int16(0), size=128), 'uint16') + byte_mask = V.vbrc(T.uint16(0xFF), size=128) for output_vector in T.serial(num_sf_rows * num_sf_cols // 128): - output_base = output_vector * 128 - output_indices = S.vadds(lane_indices, output_base) - output_rows = T.reinterpret(S.vshrs(output_indices, int(math.log2(num_sf_cols))), 'uint16x128') - output_cols = S.vand(output_indices, S.vdup(num_sf_cols - 1, T.uint16)) - packed_channel = S.vshrs(output_cols, 1) - byte_indices = S.vand(output_cols, S.vdup(1, T.uint16)) + base = output_vector * 128 + idx = V.vadds(lanes, base, mask) + rows = V.vshrs(idx, int(math.log2(num_sf_cols)), mask) + cols = V.vand(idx, V.vbrc(T.uint16(num_sf_cols - 1), size=128)) + packed_channel = V.vshrs(cols, 1, mask) + byte_index = V.vand(cols, V.vbrc(T.uint16(1), size=128)) if in_config.use_tma_aligned_col_major_sf: - source_rows, source_cols = packed_channel, output_rows + src_rows, src_cols = packed_channel, rows else: - source_rows, source_cols = output_rows, packed_channel - source_indices = S.vadd(S.vmuls(source_rows, sf_raw_cols), source_cols) - packed = S.vgather2(sf_raw_ub[0, 0], source_indices) - shifts = T.reinterpret(S.vmuls(byte_indices, 8), 'int16x128') - exponents = S.vand(S.vshr(packed, shifts), byte_mask) - scales = T.reinterpret(S.vshls(exponents, 7), 'bfloat16x128') - S.vsts(sf_ub[output_base // num_sf_cols, output_base % num_sf_cols], scales, dist='NORM_B16') + src_rows, src_cols = rows, packed_channel + src_idx = V.vadd(V.vmuls(src_rows, sf_raw_cols, mask), src_cols) + raw = V.vgather(sf_raw_ub[0, 0], src_idx, mask) + low_exp = V.vand(raw, byte_mask) + exp = V.vsel(V.vcmp(byte_index, V.vbrc(T.uint16(0), size=128), mask, 'eq'), low_exp, V.vshrs(raw, 8, mask)) + V.vstore(V.vinterpret_cast(V.vshls(exp, 7, mask), 'bfloat16'), sf_ub[base // num_sf_cols, base % num_sf_cols]) elif num_per_channels == 1: with T.SimdVF(): - sf_exp_mask = S.vdup(0x7F800000, T.uint32) + mask = V.create_mask(64, size=64) + byte_mask = V.vbrc(T.uint16(0xFF), size=64) for group in T.serial(num_sf_rows): packed_index = group + token_base % pack_factor - shift = 23 - packed_index % pack_factor * 8 + shift = packed_index % pack_factor * 8 for chunk in T.serial(num_sf_cols // 64): col = chunk * 64 - packed = S.vld(sf_raw_ub[packed_index // pack_factor, col], dist='US_B16') - packed_u32 = T.reinterpret(packed, 'uint32x64') - scales = T.reinterpret(S.vand(S.vshls(packed_u32, shift), sf_exp_mask), 'float32x64') - S.vsts(sf_ub[group, col], scales) + raw = V.vload(sf_raw_ub[packed_index // pack_factor, col], size=64) + exp = V.vand(V.vshrs(raw, shift, mask), byte_mask) + V.vstore(V.vcvt(V.vinterpret_cast(V.vshls(exp, 7, mask), 'bfloat16'), 'float32'), sf_ub[group, col]) else: T.copy(source, sf_ub) @T.macro - def load_scale(sf_ub, group, col, mask_low_f32, sf_shift, sf_exp_mask): + def load_scale(sf_ub, group, col, size, sf_shift=None, sf_exp_mask=None): if use_lazy_packed_scale: - packed = ( - S.vld(sf_ub[col // 64, group], dist='BRC_B16') - if in_config.use_tma_aligned_col_major_sf - else S.vld(sf_ub[group, col // 64], dist='BRC_B16') - ) - packed_u32 = T.reinterpret(packed, 'uint32x64') - return T.reinterpret(S.vand(S.vshl(packed_u32, sf_shift), sf_exp_mask), 'float32x64') + if in_config.use_tma_aligned_col_major_sf: + packed_v = V.vload(sf_ub[col // 64, group], size=size * 2, stride=sf_raw_cols, dist_mode='brc', group=size // 64) + else: + packed_v = V.vload(sf_ub[group, col // 64], size=size * 2, stride=1, dist_mode='brc', group=size // 64) + packed_u = V.vinterpret_cast(packed_v, 'uint32') + return V.vinterpret_cast(V.vand(V.vshl(packed_u, sf_shift), sf_exp_mask), 'float32') if num_per_channels == 1: - return S.vld(sf_ub[group, col]) + return V.vload(sf_ub[group, col], size=size) if keep_unpacked_col_major: - scale_low = S.vld(sf_ub[col // 32, group], dist='BRC_B32') - scale_high = S.vld(sf_ub[col // 32 + 1, group], dist='BRC_B32') - else: - scale_low = S.vld(sf_ub[group, col // 32], dist='BRC_B32') - scale_high = S.vld(sf_ub[group, col // 32 + 1], dist='BRC_B32') - return S.vsel(scale_low, scale_high, mask_low_f32) + return V.vload(sf_ub[col // 32, group], size=size, stride=num_sf_rows_padded, dist_mode='brc', group=size // 32) + return V.vload(sf_ub[group, col // 32], size=size, stride=1, dist_mode='brc', group=size // 32) @T.macro def store_out(out_ub, row, col, out_f32): if out_dtype == T.float32: - S.vsts(out_ub[row, col], out_f32) + V.vstore(out_f32, out_ub[row, col]) else: - out_bf16 = S.vcvt(out_f32, T.bfloat16) - S.vsts(out_ub[row, col], out_bf16, dist='PK_B32') + V.vstore(V.vcvt(out_f32, 'bfloat16'), out_ub[row, col]) + + @T.macro + def lazy_scale_args(size): + return ( + V.vbrc(T.uint32(0x7F800000), size=size) if use_lazy_packed_scale else None, + V.vsel( + V.create_mask(32, size=size, group=size // 64), + V.vbrc(T.uint32(23), size=size), + V.vbrc(T.uint32(15), size=size), + ) + if use_lazy_packed_scale and num_per_channels == 32 + else None, + ) @T.macro def transform_tile(x_ub, sf_ub, out_ub): with T.SimdVF(): - mask_low_f32 = S.pset(32, 'PAT_VL32') if num_per_channels == 32 else None - sf_shift = S.vsel(S.vdup(23, T.int32), S.vdup(15, T.int32), mask_low_f32) if use_lazy_packed_scale and num_per_channels == 32 else None - sf_exp_mask = S.vdup(0x7F800000, T.uint32) if use_lazy_packed_scale else None if use_e2b_scale: for row in T.serial(block_m): for chunk in T.serial(block_k // 256): col = chunk * 256 - scale_e2b = S.vld(sf_ub[row, chunk * 8], dist='E2B_B16') - scale_low, scale_high = S.vintlv(scale_e2b, scale_e2b) - x_low = S.vcvt(S.vld(x_ub[row, col], dist='UNPK4_B8'), T.bfloat16) - x_high = S.vcvt(S.vld(x_ub[row, col + 128], dist='UNPK4_B8'), T.bfloat16) - S.vsts(out_ub[row, col], S.vmul(x_low, scale_low), dist='NORM_B16') - S.vsts(out_ub[row, col + 128], S.vmul(x_high, scale_high), dist='NORM_B16') - elif in_config.dtype == T.float4_e2m1fn: - zero_bf16 = S.vdup(0.0, T.bfloat16) - for group in T.serial(num_sf_rows): - row_base = group * num_per_tokens - for chunk in T.serial(block_k // 128): - col = chunk * 128 - scale_low = load_scale(sf_ub, group, col, mask_low_f32, sf_shift, sf_exp_mask) - scale_high = load_scale(sf_ub, group, col + 64, mask_low_f32, sf_shift, sf_exp_mask) - for row in T.serial(num_per_tokens): - x_raw = S.vld(x_ub[row_base + row, col], dist='UNPK4_B8') - x_bf16 = S.vcvt(x_raw, T.bfloat16) - x_uint_low, x_uint_high = S.vintlv(zero_bf16, x_bf16) - store_out(out_ub, row_base + row, col, S.vmul(T.reinterpret(x_uint_low, 'float32x64'), scale_low)) - store_out(out_ub, row_base + row, col + 64, S.vmul(T.reinterpret(x_uint_high, 'float32x64'), scale_high)) + scale256 = V.vload(sf_ub[row, chunk * 8], size=256, stride=1, dist_mode='brc', group=8) + V.vstore(V.vmul(V.vcvt(V.vload(x_ub[row, col], size=256), 'bfloat16'), scale256), out_ub[row, col]) else: + scale_size = 128 if in_config.dtype == T.float4_e2m1fn else 64 + zero_bf16 = V.vbrc(T.bfloat16(0.0), size=scale_size) + sf_exp_mask, sf_shift = lazy_scale_args(scale_size) for group in T.serial(num_sf_rows): row_base = group * num_per_tokens - for chunk in T.serial(block_k // 64): - col = chunk * 64 - scale = load_scale(sf_ub, group, col, mask_low_f32, sf_shift, sf_exp_mask) + for chunk in T.serial(block_k // scale_size): + col = chunk * scale_size + scale = load_scale(sf_ub, group, col, scale_size, sf_shift, sf_exp_mask) for row in T.serial(num_per_tokens): - x_raw = S.vld(x_ub[row_base + row, col], dist='UNPK4_B8') - store_out(out_ub, row_base + row, col, S.vmul(S.vcvt(x_raw, T.float32), scale)) + raw = V.vload(x_ub[row_base + row, col], size=scale_size) + x_f32 = ( + V.vzip(zero_bf16, V.vcvt(raw, 'bfloat16'), 'float32') + if in_config.dtype == T.float4_e2m1fn + else V.vcvt(raw, 'float32') + ) + store_out(out_ub, row_base + row, col, V.vmul(x_f32, scale)) @T.prim_func def cast_back_kernel_ascend( diff --git a/tile_kernels/quant/per_block_cast_asc.py b/tile_kernels/quant/per_block_cast_asc.py index ae1c4f9c..5be79c99 100644 --- a/tile_kernels/quant/per_block_cast_asc.py +++ b/tile_kernels/quant/per_block_cast_asc.py @@ -6,10 +6,10 @@ # Re-bind below the star import: tile_kernels.quant.common exports its own `T`. from tilelang.ascend import language as T -from tilelang.ascend.language import simd as S +from tilelang.ascend.language import vmi as V -@tilelang.jit() +@tilelang.jit(target='pto', pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) def get_per_block_cast_kernel_asc( hidden: int, in_config: CastInputConfig, @@ -42,132 +42,104 @@ def get_per_block_cast_kernel_asc( sf_stride = T.dynamic('sf_stride') @T.macro - def f32_pair_to_bf16(x0, x1, round_odd=False): # round_odd=False means round_to_zero here - # in this kernel, we ensure that round_odd=False is used only when the low 16bit of fp32 inputs are all zero. - low, high = S.vdintlv(T.reinterpret(x0, 'uint16x128'), T.reinterpret(x1, 'uint16x128')) - if round_odd: - high = S.vor(high, S.vmins(low, 1)) - return T.reinterpret(high, 'bfloat16x128') + def f32_to_bf16_round_odd(x): + low, high = V.vunzip(x, 'uint16') + return V.vinterpret_cast(V.vor(high, V.vmin(low, V.vbrc(T.uint16(1), size=128))), 'bfloat16') @T.macro - def get_scale_and_inv(amax, fp8_max, inv_fp8_max): + def get_scale_and_inv(amax, fp8_max, inv_fp8_max, mask, lanes): if not out_config.round_sf: - scale = S.vdiv(amax, fp8_max) - return scale, S.vdiv(fp8_max, amax) - - scale_raw = S.vmul(amax, inv_fp8_max) - bits = T.reinterpret(scale_raw, 'uint32x64') - exp = S.vadds(S.vshrs(S.vsub(bits, S.vdup(1, T.uint32)), 23), 1) - scale_bits = S.vshls(exp, 23) - scale = T.reinterpret(scale_bits, 'float32x64') - inv_exp = S.vsub(S.vdup(254, T.uint32), exp) - inv_bits = S.vshls(inv_exp, 23) - return scale, T.reinterpret(inv_bits, 'float32x64') + return V.vdiv(amax, fp8_max, mask, precision='exact'), V.vdiv(fp8_max, amax, mask, precision='exact') + raw = V.vmul(amax, inv_fp8_max) + bits = V.vinterpret_cast(raw, 'uint32') + exp = V.vadds(V.vshrs(V.vsub(bits, V.vbrc(T.uint32(1), size=lanes)), 23, mask), T.uint32(1), mask) + scale = V.vinterpret_cast(V.vshls(exp, 23, mask), 'float32') + inv = V.vinterpret_cast(V.vshls(V.vsub(V.vbrc(T.uint32(254), size=lanes), exp), 23, mask), 'float32') + return scale, inv @T.macro - def pack_two_block_amax(amax_vec, m_low, m_high): - a0 = S.vcmax(amax_vec, m_low) - a1 = S.vcmax(amax_vec, m_high) - return S.vsel(S.vdupv(a0), S.vdupv(a1), m_low) - - @T.macro - def store_scale_pair(sf_ub, sf_row, sf_j, scale01): - if out_config.use_packed_ue8m0: - exponents = S.vshrs(T.reinterpret(scale01, 'uint32x64'), 23) - exponent_high = S.vdupv(exponents, pos='POS_HIGHEST') - packed_exponents = S.vor(exponents, S.vshls(exponent_high, 8)) - if out_config.use_tma_aligned_col_major_sf: - S.vsts(sf_ub[sf_j // pack_factor, sf_row], T.reinterpret(packed_exponents, 'uint16x128'), dist='ONEPT_B16') - else: - S.vsts(sf_ub[sf_row, sf_j // pack_factor], T.reinterpret(packed_exponents, 'uint16x128'), dist='ONEPT_B16') + def store_scale_four(sf_ub, sf_f32_ub, sf_row, sf_j, scale4): + mask4 = V.create_mask(4, size=4) + if out_config.use_packed_ue8m0 and out_config.use_tma_aligned_col_major_sf: + V.vstore(scale4, sf_f32_ub[0, sf_j], group=4, stride=1) + elif out_config.use_packed_ue8m0: + exponent_u8 = V.vcvt(V.vshrs(V.vinterpret_cast(scale4, 'uint32'), 23, mask4), 'uint8') + sf_u8 = T.view(sf_ub, (sf_ub.shape[0], sf_ub.shape[1] * pack_factor), T.uint8) + V.vstore(exponent_u8, sf_u8[sf_row, sf_j]) elif out_config.use_tma_aligned_col_major_sf: - S.vsts(sf_ub[sf_j, sf_row], scale01, dist='ONEPT_B32') - S.vsts(sf_ub[sf_j + 1, sf_row], S.vdupv(scale01, pos='POS_HIGHEST'), dist='ONEPT_B32') + V.vstore(scale4, sf_ub[sf_j, sf_row], group=4, stride=token_group) else: - S.vsts(sf_ub[sf_row, sf_j], scale01, dist='ONEPT_B32') - S.vsts(sf_ub[sf_row, sf_j + 1], S.vdupv(scale01, pos='POS_HIGHEST'), dist='ONEPT_B32') + V.vstore(scale4, sf_ub[sf_row, sf_j], group=4, stride=1) @T.macro - def load_scale_pair(sf_ub, sf_j, m_low): - scale0 = S.vld(sf_ub[0, sf_j], dist='BRC_B32') - scale1 = S.vld(sf_ub[0, sf_j + 1], dist='BRC_B32') - return S.vsel(scale0, scale1, m_low) + def pack_tma_sf(sf_f32_ub, sf_ub, sf_row): + mask1 = V.create_mask(1, size=1) + for sf_j in T.serial(num_stored_sf_cols_per_block): + s0 = V.vload(sf_f32_ub[0, sf_j * pack_factor], size=1) + s1 = V.vload(sf_f32_ub[0, sf_j * pack_factor + 1], size=1) + e0 = V.vcvt(V.vshrs(V.vinterpret_cast(s0, 'uint32'), 23, mask1), 'uint16') + e1 = V.vcvt(V.vshrs(V.vinterpret_cast(s1, 'uint32'), 23, mask1), 'uint16') + V.vstore(V.vor(e0, V.vshls(e1, 8, mask1)), sf_ub[sf_j, sf_row]) @T.macro - def quant_block(x_ub, sf_ub, out_ub, sf_row): + def quant_block(x_ub, sf_ub, sf_f32_ub, out_ub, sf_row): with T.SimdVF(): - m_low = S.pset(32, 'PAT_VL32') - m_high = S.vcmps(S.vci(0, T.int32), 31, op='gt') - one_reg = S.vdup(1.0, T.float32) - eps = S.vdup(out_config.clamp_min_value, T.float32) - fp8_max = S.vdup(quant_max, T.float32) - inv_fp8_max = S.vdup(1.0 / quant_max, T.float32) - zero_bf16 = S.vdup(0.0, T.bfloat16) - abs_mask = T.reinterpret(S.vdup(0x7FFF, T.uint16), 'bfloat16x128') + mask4 = V.create_mask(4, size=4) + mask64 = V.create_mask(64, size=64) + mask128 = V.create_mask(128, size=128) + one128 = V.vbrc(T.float32(1.0), size=128) + eps128 = V.vbrc(T.float32(out_config.clamp_min_value), size=128) + fp8_max4 = V.vbrc(T.float32(quant_max), size=4) + inv_fp8_max4 = V.vbrc(T.float32(1.0 / quant_max), size=4) + u254 = V.vbrc(T.uint32(254), size=128) for chunk in T.serial(block_k // 128): col = chunk * 128 - sf_col = chunk * 4 # 4 = 128 // num_per_channels - inverse = S.alloc_local((2,), T.float32) + sf_col = chunk * 4 if cast_only: - for half in T.unroll(2, explicit=True): - scale = load_scale_pair(sf_ub, sf_col + half * 2, m_low) - if out_config.round_sf: - exponent = S.vshrs(T.reinterpret(scale, 'uint32x64'), 23) - inverse_bits = S.vshls(S.vsub(S.vdup(254, T.uint32), exponent), 23) - inverse[half] = T.reinterpret(inverse_bits, 'float32x64') - else: - inverse[half] = S.vdiv(one_reg, scale) + scale128 = V.vload(sf_ub[0, sf_col], size=128, group=4, stride=1, dist_mode='brc') + if out_config.round_sf: + inverse_vec = V.vinterpret_cast( + V.vshls(V.vsub(u254, V.vshrs(V.vinterpret_cast(scale128, 'uint32'), 23, mask128)), 23, mask128), + 'float32', + ) + else: + inverse_vec = V.vdiv(one128, scale128, mask128, precision='exact') else: - amax = S.alloc_local((2,), T.float32) if in_config.dtype == T.bfloat16: - amax_bf16 = S.alloc_var(T.bfloat16) - amax_bf16 = S.vdup(0.0, T.bfloat16) + amax_bf16 = V.vbrc(T.bfloat16(0.0), size=128) for row in T.serial(num_per_tokens): - amax_bf16 = S.vmax(amax_bf16, S.vand(S.vld(x_ub[row, col]), abs_mask)) - amax_low, amax_high = S.vintlv(zero_bf16, amax_bf16) - amax[0] = T.reinterpret(amax_low, 'float32x64') - amax[1] = T.reinterpret(amax_high, 'float32x64') + amax_bf16 = V.vmax(amax_bf16, V.vabs(V.vload(x_ub[row, col], size=128))) + amax128 = V.vcvt(amax_bf16, 'float32') else: - for half in T.unroll(2, explicit=True): - amax[half] = S.vdup(0.0, T.float32) + amax128 = V.vbrc(T.float32(0.0), size=128) for row in T.serial(num_per_tokens): - for half in T.unroll(2, explicit=True): - amax[half] = S.vmax(amax[half], S.vabs(S.vld(x_ub[row, col + half * 64]))) - for half in T.unroll(2, explicit=True): - amax[half] = S.vmax(amax[half], eps) - scale, inverse[half] = get_scale_and_inv(pack_two_block_amax(amax[half], m_low, m_high), fp8_max, inv_fp8_max) - store_scale_pair(sf_ub, sf_row, sf_col + half * 2, scale) + amax128 = V.vmax(amax128, V.vabs(V.vload(x_ub[row, col], size=128))) + amax4 = V.vcmax(V.vmax(amax128, eps128), mask128, group=4) + scale4, inverse4 = get_scale_and_inv(amax4, fp8_max4, inv_fp8_max4, mask4, 4) + inverse_vec = V.vbrc(inverse4, size=128, group=4) + store_scale_four(sf_ub, sf_f32_ub, sf_row, sf_col, scale4) if in_config.dtype == T.bfloat16 and out_config.round_sf: - sf_inv = f32_pair_to_bf16(inverse[0], inverse[1]) + sf_inv = V.vcvt(inverse_vec, 'bfloat16', rounding='Z') if not sf_only: for row in T.serial(num_per_tokens): if in_config.dtype == T.bfloat16 and out_config.round_sf: - quantized_bf16 = S.vmul(S.vld(x_ub[row, col]), sf_inv) + quantized_bf16 = V.vmul(V.vload(x_ub[row, col], size=128), sf_inv) if out_config.dtype == T.float8_e4m3fn: - interleaved_low, interleaved_high = S.vintlv(zero_bf16, quantized_bf16) - quantized_low = T.reinterpret(interleaved_low, 'float32x64') - quantized_high = T.reinterpret(interleaved_high, 'float32x64') + quantized = V.vcvt(quantized_bf16, 'float32') else: - if in_config.dtype == T.bfloat16: - interleaved_low, interleaved_high = S.vintlv(zero_bf16, S.vld(x_ub[row, col])) - input_low = T.reinterpret(interleaved_low, 'float32x64') - input_high = T.reinterpret(interleaved_high, 'float32x64') - else: - input_low = S.vld(x_ub[row, col]) - input_high = S.vld(x_ub[row, col + 64]) - quantized_low = S.vmul(input_low, inverse[0]) - quantized_high = S.vmul(input_high, inverse[1]) + values = V.vload(x_ub[row, col], size=128) + quantized = V.vmul(V.vcvt(values, 'float32') if in_config.dtype == T.bfloat16 else values, inverse_vec) if out_config.dtype == T.float4_e2m1fn: - quantized_bf16 = f32_pair_to_bf16(quantized_low, quantized_high, round_odd=True) - + quantized_bf16 = f32_to_bf16_round_odd(quantized) if out_config.dtype == T.float4_e2m1fn: - S.vsts(out_ub[row, col], S.vcvt(quantized_bf16, T.float4_e2m1fn), dist='PK4_B32') + V.vstore(V.vcvt(quantized_bf16, 'float4_e2m1fn', rounding='R'), out_ub[row, col], mask64) else: - S.vsts(out_ub[row, col], S.vcvt(quantized_low, T.float8_e4m3fn), dist='PK4_B32') - S.vsts(out_ub[row, col + 64], S.vcvt(quantized_high, T.float8_e4m3fn), dist='PK4_B32') + V.vstore(V.vcvt(quantized, 'float8_e4m3fn', rounding='R', saturate='SAT'), out_ub[row, col]) + if not cast_only and out_config.use_packed_ue8m0 and out_config.use_tma_aligned_col_major_sf: + pack_tma_sf(sf_f32_ub, sf_ub, sf_row) @T.prim_func def per_block_cast_kernel( @@ -182,6 +154,11 @@ def per_block_cast_kernel( sf_ub = T.alloc_shared((sf_cols, token_group), sf_ub_dtype) else: sf_ub = T.alloc_shared((token_group, sf_cols), sf_ub_dtype) + sf_f32_ub = ( + T.alloc_shared((1, num_sf_cols_per_block), T.float32) + if out_config.use_packed_ue8m0 and out_config.use_tma_aligned_col_major_sf + else None + ) out_ub = T.alloc_shared((block_m, block_k), out_config.dtype) versions = {x_ub: num_stages, sf_ub: num_stages} @@ -204,18 +181,27 @@ def per_block_cast_kernel( T.copy(x[token_block * block_m, pid_y * block_k], x_ub[0:block_m, 0:block_k]) if cast_only: T.copy(out_sf[token_block, sf_k_base], sf_ub[0:1, 0:num_sf_cols_per_block]) - quant_block(x_ub, sf_ub, out_ub, token_offset) + quant_block(x_ub, sf_ub, sf_f32_ub, out_ub, token_offset) if not sf_only: T.copy(out_ub[0:block_m, 0:block_k], out[token_block * block_m, pid_y * block_k]) if not cast_only: if out_config.use_packed_ue8m0: - output_sf = T.Tensor((sf_shape[0], sf_shape[1] // pack_factor), T.uint16, out_sf.data) - else: - output_sf = out_sf - if out_config.use_tma_aligned_col_major_sf: - T.copy(sf_ub[0:num_stored_sf_cols_per_block, 0:token_group], output_sf[stored_sf_k_base, sf_m_base]) + if out_config.use_tma_aligned_col_major_sf: + sf_out_ub = T.view(sf_ub, (sf_cols, token_group * pack_factor), T.uint8) + T.copy( + sf_out_ub[0:num_stored_sf_cols_per_block, 0 : token_group * pack_factor], + out_sf[stored_sf_k_base, sf_m_base * pack_factor], + ) + else: + sf_out_ub = T.view(sf_ub, (token_group, sf_cols * pack_factor), T.uint8) + T.copy( + sf_out_ub[0:token_group, 0:num_sf_cols_per_block], + out_sf[sf_m_base, stored_sf_k_base * pack_factor], + ) + elif out_config.use_tma_aligned_col_major_sf: + T.copy(sf_ub[0:num_stored_sf_cols_per_block, 0:token_group], out_sf[stored_sf_k_base, sf_m_base]) else: - T.copy(sf_ub[0:token_group, 0:num_stored_sf_cols_per_block], output_sf[sf_m_base, stored_sf_k_base]) + T.copy(sf_ub[0:token_group, 0:num_stored_sf_cols_per_block], out_sf[sf_m_base, stored_sf_k_base]) return per_block_cast_kernel diff --git a/tile_kernels/quant/per_channel_cast_asc.py b/tile_kernels/quant/per_channel_cast_asc.py index f03cf89b..22f98169 100644 --- a/tile_kernels/quant/per_channel_cast_asc.py +++ b/tile_kernels/quant/per_channel_cast_asc.py @@ -6,11 +6,11 @@ # Re-bind below the star import: tile_kernels.quant.common exports its own `T`. from tilelang.ascend import language as T -from tilelang.ascend.language import simd as S +from tilelang.ascend.language import vmi as V from tile_kernels.quant.common import get_packed_ue8m0_pack_factor -@tilelang.jit +@tilelang.jit(target='pto', pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) def get_per_channel_cast_kernel_asc( hidden: int, in_config: CastInputConfig, @@ -82,67 +82,61 @@ def get_per_channel_cast_kernel_asc( dequant_ub_dtype = T.bfloat16 if use_bf16_compute else T.float32 @T.macro - def compute_sf_and_inv(amax): - clamped_amax = S.vmax(amax, S.vdup(out_config.clamp_min_value, T.float32)) + def compute_sf_and_inv(amax, clamp_v, qmax_v, inv_q, u1, u254, mask): + clamped_amax = V.vmax(amax, clamp_v) if not out_config.round_sf: - sf = S.vdiv(clamped_amax, S.vdup(quant_max, T.float32)) - sf_inv = S.vdiv(S.vdup(quant_max, T.float32), clamped_amax) - return sf, sf_inv - raw_sf = S.vmul(clamped_amax, S.vdup(1.0 / quant_max, T.float32)) - raw_sf_bits = T.reinterpret(raw_sf, 'uint32x64') - sf_exponent = S.vadds(S.vshrs(S.vsub(raw_sf_bits, S.vdup(1, T.uint32)), 23), 1) - sf_inv_exponent = S.vsub(S.vdup(254, T.uint32), sf_exponent) - sf_inv = T.reinterpret(S.vshls(sf_inv_exponent, 23), 'float32x64') + return V.vdiv(clamped_amax, qmax_v, mask, precision='exact'), V.vdiv(qmax_v, clamped_amax, mask, precision='exact') + raw_sf_bits = V.vinterpret_cast(V.vmul(clamped_amax, inv_q), 'uint32') + sf_exponent = V.vadds(V.vshrs(V.vsub(raw_sf_bits, u1), 23, mask), T.uint32(1), mask) + sf_inv = V.vinterpret_cast(V.vshls(V.vsub(u254, sf_exponent), 23, mask), 'float32') if out_config.use_packed_ue8m0: - return sf_exponent, sf_inv - sf = T.reinterpret(S.vshls(sf_exponent, 23), 'float32x64') - return sf, sf_inv + return V.vcvt(sf_exponent, 'uint8', saturate='SAT'), sf_inv + return V.vinterpret_cast(V.vshls(sf_exponent, 23, mask), 'float32'), sf_inv @T.macro def quantize_block(x_ub, out_ub, sf_out_ub, sf_in_ub, dequant_ub, sf_row_base=0): with T.SimdVF(): - # Prepare and build gather offsets per input sf layout - sf_dist = 'PK4_B32' if out_config.use_packed_ue8m0 else 'NORM_B32' - amax = S.alloc_local((num_vecs,), T.float32) - sf_inv = S.alloc_local((num_vecs,), T.float32) + f32_lanes = 128 if use_bf16_compute else 64 + mask64 = V.create_mask(64, size=64) + mask128 = V.create_mask(128, size=128) + f32_mask = mask128 if use_bf16_compute else mask64 + f32 = V.vreg(f32_lanes, T.float32) + amax = V.alloc_local((num_vecs // compute_vec_factor,), f32) + sf_inv = V.alloc_local((num_vecs // compute_vec_factor,), f32) + clamp_v = V.vbrc(T.float32(out_config.clamp_min_value), size=f32_lanes) + qmax_v = V.vbrc(T.float32(quant_max), size=f32_lanes) + inv_q = V.vbrc(T.float32(1.0 / quant_max), size=f32_lanes) + u1 = V.vbrc(T.uint32(1), size=f32_lanes) + u254 = V.vbrc(T.uint32(254), size=f32_lanes) if use_bf16_compute: - zero_bf16 = S.vdup(0.0, T.bfloat16) - bf16_abs_mask = T.reinterpret(S.vdup(0x7FFF, T.uint16), 'bfloat16x128') + bf16_abs_mask = V.vbrc(T.uint16(0x7FFF), size=128) if in_config.with_sf: - sf_lane = T.reinterpret(S.vshrs(S.vci(0, T.int16), 5), 'uint16x128') # 0x32 1x32 2x32 3x32 + sf_lane = V.vshrs(V.vci(T.uint16(0), size=128), 5, mask128) if in_config.use_tma_aligned_col_major_sf: - # Col-major: int16 entries viewed as byte pairs, flat (k-block, token-pair) order. sf_flat_ub = T.view(sf_in_ub, (block_k // sf_in_block_k * block_m * pack_factor,), T.uint8) - sf_offsets = S.vadd( - S.vmuls(S.vshrs(sf_lane, 1), block_m * pack_factor), - S.vand(sf_lane, S.vdup(1, T.uint16)), - ) # 0x32 1x32 (2 * block_m)x32 (2 * block_m + 1)x32 - sf_mask = S.pset(8) + sf_offsets = V.vadd( + V.vmuls(V.vshrs(sf_lane, 1, mask128), block_m * pack_factor, mask128), + V.vand(sf_lane, V.vbrc(T.uint16(1), size=128)), + ) sf_col_stride = block_m * pack_factor sf_row_stride = pack_factor else: sf_flat_ub = T.view(sf_in_ub, (sf_in_rows * sf_in_cols,), T.uint8) sf_offsets = sf_lane - sf_mask = None sf_col_stride = 1 sf_row_stride = sf_in_cols else: - f32_abs_mask = T.reinterpret(S.vdup(0x7FFFFFFF, T.uint32), 'float32x64') - sf_lane = T.reinterpret(S.vshrs(S.vci(0, T.int32), 5), 'uint32x64') # 0x32 1x32 + sf_idx = V.vshrs(V.vci(T.uint32(0), size=64), 5, mask64) for sf_row in T.serial(sf_rows_per_group): row_base = sf_row * num_per_tokens - # Dequantize and accumulate amax for chunk in T.serial(block_k // quant_chunk_size): chunk_base = chunk * quant_chunk_size - if use_bf16_compute: - amax_bf16 = S.alloc_local((num_vecs // compute_vec_factor,), T.bfloat16) for vec in T.unroll(num_vecs // compute_vec_factor, explicit=True): vec_col = vec * 64 * compute_vec_factor if use_bf16_compute: - # bf16 path - amax_bf16[vec] = S.vdup(0.0, T.bfloat16) + acc_bf16 = V.vbrc(T.bfloat16(0.0), size=128) if in_config.with_sf: if in_config.use_tma_aligned_col_major_sf: sf_col = chunk * (quant_chunk_size // 64) + vec * 2 @@ -150,79 +144,66 @@ def quantize_block(x_ub, out_ub, sf_out_ub, sf_in_ub, dequant_ub, sf_row_base=0) sf_col = chunk * (quant_chunk_size // 32) + vec * 4 for row in T.serial(num_per_tokens): if in_config.with_sf: - raw_f32_0 = S.vcvt(S.vld(x_ub[row_base + row, chunk_base + vec_col], dist='UNPK4_B8'), T.float32) - raw_f32_1 = S.vcvt(S.vld(x_ub[row_base + row, chunk_base + vec_col + 64], dist='UNPK4_B8'), T.float32) - _, raw_bf16 = S.vdintlv(T.reinterpret(raw_f32_0, 'bfloat16x128'), T.reinterpret(raw_f32_1, 'bfloat16x128')) - sf_exponents = S.vgather2( + raw_bf16 = V.vcvt(V.vload(x_ub[row_base + row, chunk_base + vec_col], size=128), 'bfloat16') + exp_u16 = V.vgather( sf_flat_ub[(row_base + row) * sf_row_stride + sf_col * sf_col_stride], sf_offsets, - mask=sf_mask, + mask128, ) - sf_bits = S.vshls(sf_exponents, 7) - values_bf16 = S.vmul(raw_bf16, T.reinterpret(sf_bits, 'bfloat16x128')) - S.vsts(dequant_ub[row, vec_col], values_bf16, dist='NORM_B16') + values_bf16 = V.vmul(raw_bf16, V.vinterpret_cast(V.vshls(exp_u16, 7, mask128), 'bfloat16')) + V.vstore(values_bf16, dequant_ub[row, vec_col]) else: - values_bf16 = S.vld(x_ub[row_base + row, chunk_base + vec_col]) - amax_bf16[vec] = S.vmax(amax_bf16[vec], S.vand(values_bf16, bf16_abs_mask)) + values_bf16 = V.vload(x_ub[row_base + row, chunk_base + vec_col], size=128) + acc_bf16 = V.vmax( + acc_bf16, + V.vinterpret_cast(V.vand(V.vinterpret_cast(values_bf16, 'uint16'), bf16_abs_mask), 'bfloat16'), + ) + amax[vec] = V.vcvt(acc_bf16, 'float32') else: - # f32 path - amax[vec] = S.vdup(0.0, T.float32) + acc_f32 = V.vbrc(T.float32(0.0), size=64) if in_config.use_tma_aligned_col_major_sf: - sf_offsets = S.vadds(S.vmuls(sf_lane, block_m), (chunk * 8 + vec * 2) * block_m) + sf_offsets = V.vadds(V.vmuls(sf_idx, block_m, mask64), (chunk * 8 + vec * 2) * block_m, mask64) else: - sf_offsets = S.vadds(sf_lane, vec * 2) + sf_offsets = V.vadds(sf_idx, vec * 2, mask64) for row in T.serial(num_per_tokens): + raw_f32 = V.vcvt(V.vload(x_ub[row_base + row, chunk_base + vec_col], size=64), 'float32') if in_config.use_tma_aligned_col_major_sf: - input_scale = S.vgather2(sf_in_ub[0, 0], S.vadds(sf_offsets, row_base + row)) - else: - compact_sf = S.vld(sf_in_ub[row_base + row, chunk * 8]) - input_scale = S.vselr(compact_sf, sf_offsets) - dequant_f32 = S.vmul( - S.vcvt(S.vld(x_ub[row_base + row, chunk_base + vec_col], dist='UNPK4_B8'), T.float32), - input_scale, - ) - S.vsts(dequant_ub[row, vec_col], dequant_f32) - amax[vec] = S.vmax(amax[vec], S.vand(dequant_f32, f32_abs_mask)) - if use_bf16_compute: - # Merge the 128-lane bf16 amax into two f32 half-vectors - amax_lo, amax_hi = S.vintlv(zero_bf16, amax_bf16[vec]) - amax[vec * 2] = T.reinterpret(amax_lo, 'float32x64') - amax[vec * 2 + 1] = T.reinterpret(amax_hi, 'float32x64') - - # Mem bar on dequant_ub - S.mem_bar('VST_VLD') - for pair in T.unroll(num_vecs // 2, explicit=True): - # Write sf_out_ub - for half in T.unroll(2, explicit=True): - vec = pair * 2 + half - sf, sf_inv[vec] = compute_sf_and_inv(amax[vec]) - S.vsts(sf_out_ub[sf_row_base + sf_row, chunk_base + vec * 64], sf, dist=sf_dist) - - # Calculate and write out_ub - for half in T.unroll(2, explicit=True): - vec = pair * 2 + half - vec_col = vec * 64 - for row in T.serial(num_per_tokens): - if use_bf16_compute: - value_ub, value_row, value_col = (dequant_ub, row, 0) if in_config.with_sf else (x_ub, row_base + row, chunk_base) - dequant_f32 = S.vcvt(S.vld(value_ub[value_row, value_col + vec_col], dist='UNPK_B16'), T.float32) + input_scale = V.vgather(sf_in_ub[0, 0], V.vadds(sf_offsets, row_base + row, mask64), mask64) else: - dequant_f32 = S.vld(dequant_ub[row, vec_col]) - scaled = S.vmul(dequant_f32, sf_inv[vec]) - S.vsts( - out_ub[row_base + row, chunk_base + vec_col], - S.vcvt(scaled, T.float8_e4m3fn), - dist='PK4_B32', - ) + input_scale = V.vgather(sf_in_ub[row_base + row, 0], sf_offsets, mask64) + dequant_f32 = V.vmul(raw_f32, input_scale) + V.vstore(dequant_f32, dequant_ub[row, vec_col]) + acc_f32 = V.vmax(acc_f32, V.vabs(dequant_f32)) + amax[vec] = acc_f32 + T.simd.mem_bar('VST_VLD') + for vec in T.unroll(num_vecs // compute_vec_factor, explicit=True): + sf, sf_inv[vec] = compute_sf_and_inv(amax[vec], clamp_v, qmax_v, inv_q, u1, u254, f32_mask) + V.vstore(sf, sf_out_ub[sf_row_base + sf_row, chunk_base + vec * 64 * compute_vec_factor]) + vec_col = vec * 64 * compute_vec_factor + for row in T.serial(num_per_tokens): + if use_bf16_compute: + value_ub, value_row, value_col = (dequant_ub, row, 0) if in_config.with_sf else (x_ub, row_base + row, chunk_base) + dequant_f32 = V.vcvt(V.vload(value_ub[value_row, value_col + vec_col], size=128), 'float32') + else: + dequant_f32 = V.vload(dequant_ub[row, vec_col], size=64) + V.vstore( + V.vcvt(V.vmul(dequant_f32, sf_inv[vec]), 'float8_e4m3fn', rounding='R', saturate='SAT'), + out_ub[row_base + row, chunk_base + vec_col], + ) @T.macro def pack_sf_rows(unpacked_sf_ub, packed_sf_ub, unpacked_row, packed_row): with T.SimdVF(): for col in T.serial(0, block_k, vector_width): - packed_0, packed_1 = S.vintlv(S.vld(unpacked_sf_ub[unpacked_row, col]), S.vld(unpacked_sf_ub[unpacked_row + 1, col])) - S.vsts(packed_sf_ub[packed_row, col * 2], packed_0, dist='NORM_B8') + pack_mask = V.create_mask(vector_width, size=vector_width) + packed_0, packed_1 = V.vintlv( + V.vload(unpacked_sf_ub[unpacked_row, col], size=vector_width), + V.vload(unpacked_sf_ub[unpacked_row + 1, col], size=vector_width), + pack_mask, + ) + V.vstore(packed_0, packed_sf_ub[packed_row, col * 2]) if block_k - col >= vector_width: - S.vsts(packed_sf_ub[packed_row, col * 2 + vector_width], packed_1, dist='NORM_B8') + V.vstore(packed_1, packed_sf_ub[packed_row, col * 2 + vector_width]) @T.prim_func def per_channel_cast_kernel( diff --git a/tile_kernels/quant/per_token_cast_asc.py b/tile_kernels/quant/per_token_cast_asc.py index b736cf3d..f680a0b7 100644 --- a/tile_kernels/quant/per_token_cast_asc.py +++ b/tile_kernels/quant/per_token_cast_asc.py @@ -2,13 +2,13 @@ import tilelang from tilelang.ascend import language as T -from tilelang.ascend.language import simd as S +from tilelang.ascend.language import vmi as V from tile_kernels.quant.common import CastInputConfig, CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape from tile_kernels.utils import align -@tilelang.jit +@tilelang.jit(target='pto') def get_per_token_cast_kernel_asc( hidden: int, token_stride: int, @@ -57,210 +57,165 @@ def get_per_token_cast_kernel_asc( x_sf_shape = get_sf_shape((num_tokens, hidden), in_config) if in_config.with_sf else (1, 1) @T.macro - def compute_scale(amax): - clamped = S.vmaxs(amax, out_config.clamp_min_value) + def compute_scale(amax, lanes): + mask = V.create_mask(lanes, size=lanes) + clamped = V.vmax(amax, V.vbrc(T.float32(out_config.clamp_min_value), size=lanes)) if not out_config.round_sf: - max_value = S.vdup(quant_max, T.float32) - return S.vdiv(clamped, max_value), S.vdiv(max_value, clamped) - bits = T.reinterpret(S.vmuls(clamped, inv_quant_max), 'uint32x64') - exponent = S.vadds(S.vshrs(S.vsub(bits, S.vdup(1, T.uint32)), 23), 1) - inverse = T.reinterpret(S.vshls(S.vsub(S.vdup(254, T.uint32), exponent), 23), 'float32x64') + max_value = V.vbrc(T.float32(quant_max), size=lanes) + return V.vdiv(clamped, max_value, mask, precision='exact'), V.vdiv(max_value, clamped, mask, precision='exact') + one = V.vbrc(T.uint32(1), size=lanes) + shift = V.vbrc(T.uint32(23), size=lanes) + bits = V.vinterpret_cast(V.vmul(clamped, V.vbrc(T.float32(inv_quant_max), size=lanes)), 'uint32') + exponent = V.vadd(V.vshr(V.vsub(bits, one), shift), one) + inverse = V.vinterpret_cast(V.vshl(V.vsub(V.vbrc(T.uint32(254), size=lanes), exponent), shift), 'float32') if out_config.use_packed_ue8m0: return exponent, inverse - return T.reinterpret(S.vshls(exponent, 23), 'float32x64'), inverse + return V.vinterpret_cast(V.vshl(exponent, shift), 'float32'), inverse @T.macro - def load_packed_input_scale_pair(scales_ub, group, sf_shift, sf_exp_mask): - packed = S.vld(scales_ub[group // pack_factor], dist='BRC_B16') - packed_u32 = T.reinterpret(packed, 'uint32x64') - scale_bits = S.vand(S.vshl(packed_u32, sf_shift), sf_exp_mask) - return T.reinterpret(scale_bits, 'float32x64') + def f32_to_bf16_round_odd(x): + low, high = V.vunzip(x, 'uint16') + return V.vinterpret_cast(V.vor(high, V.vmin(low, V.vbrc(T.uint16(1), size=128))), 'bfloat16') + + @T.macro + def load_packed_input_scale_bf16(scales_ub, group, sf_shift, sf_exp_mask): + packed = V.vload(scales_ub[group // pack_factor], size=256, dist_mode='brc', group=2, stride=1) + scale_bits = V.vand(V.vshl(V.vinterpret_cast(packed, 'uint32'), sf_shift), sf_exp_mask) + scale_u16 = V.vcvt(V.vshr(scale_bits, V.vbrc(T.uint32(16), size=128)), 'uint16', saturate='NOSAT') + return V.vinterpret_cast(scale_u16, 'bfloat16') @T.macro def init_transpose_indices(indices): with T.SimdVF(): if out_config.use_packed_ue8m0: - lane = S.vci(0, T.int16) - output_row = S.vand(T.reinterpret(lane, 'uint16x128'), S.vdup(block_m - 1, T.uint16)) - group_in_vector = T.reinterpret(S.vshrs(lane, log2_block_m), 'uint16x128') - source_indices = S.vadd(S.vmuls(output_row, sf_stride // pack_factor), group_in_vector) - S.vsts(indices[0], source_indices, dist='NORM_B16') + lane = V.vci(T.int16(0), size=128) + output_row = V.vand(V.vinterpret_cast(lane, 'uint16'), V.vbrc(T.uint16(block_m - 1), size=128)) + group_in_vector = V.vinterpret_cast(V.vshr(lane, V.vbrc(T.int16(log2_block_m), size=128)), 'uint16') + V.vstore(V.vadd(V.vmul(output_row, V.vbrc(T.uint16(sf_stride // pack_factor), size=128)), group_in_vector), indices[0]) else: - lane = S.vci(0, T.int32) - output_row = S.vand(T.reinterpret(lane, 'uint32x64'), S.vdup(block_m - 1, T.uint32)) - group_in_vector = T.reinterpret(S.vshrs(lane, log2_block_m), 'uint32x64') - source_indices = S.vadd(S.vmuls(output_row, sf_stride), group_in_vector) - S.vsts(indices[0], source_indices, dist='NORM_B32') + lane = V.vci(T.int32(0), size=64) + output_row = V.vand(V.vinterpret_cast(lane, 'uint32'), V.vbrc(T.uint32(block_m - 1), size=64)) + group_in_vector = V.vinterpret_cast(V.vshr(lane, V.vbrc(T.int32(log2_block_m), size=64)), 'uint32') + V.vstore(V.vadd(V.vmul(output_row, V.vbrc(T.uint32(sf_stride), size=64)), group_in_vector), indices[0]) @T.macro - def transpose_vectors(src, dst, indices, vector_lanes, num_output_groups, store_dist): + def transpose_vectors(src, dst, indices, vector_lanes, num_output_groups): with T.SimdVF(): groups_per_vector = vector_lanes // block_m num_output_values = num_output_groups * block_m num_full_vectors = num_output_values // vector_lanes num_tail_values = num_output_values % vector_lanes - source_indices = S.vld(indices[0]) - + source_indices = V.vload(indices[0], size=vector_lanes) + full_mask = V.create_mask(vector_lanes, size=vector_lanes) for output_vector in T.unroll(num_full_vectors, explicit=True): full_group_base = output_vector * groups_per_vector - values = S.vgather2(src[0, full_group_base], source_indices) - S.vsts(dst[full_group_base, 0], values, dist=store_dist) - + V.vstore(V.vgather(src[0, full_group_base], source_indices, full_mask), dst[full_group_base, 0]) if num_tail_values > 0: tail_group_base = num_full_vectors * groups_per_vector - tail_mask = S.pset(2048 // vector_lanes, f'PAT_VL{num_tail_values}') - values = S.vgather2(src[0, tail_group_base], source_indices, tail_mask) - S.vsts(dst[tail_group_base, 0], values, tail_mask, dist=store_dist, extent=num_tail_values) + tail_mask = V.create_mask(num_tail_values, size=vector_lanes) + V.vstore(V.vgather(src[0, tail_group_base], source_indices, tail_mask), dst[tail_group_base, 0], tail_mask) @T.macro def transpose_output_sf(src, dst, indices): if out_config.use_packed_ue8m0: src_u16 = T.Tensor((block_m, sf_stride // pack_factor), T.uint16, src.data) dst_u16 = T.Tensor((num_groups // pack_factor, block_m), T.uint16, dst.data) - transpose_vectors(src_u16, dst_u16, indices, 128, num_groups // pack_factor, 'NORM_B16') + transpose_vectors(src_u16, dst_u16, indices, 128, num_groups // pack_factor) else: - transpose_vectors(src, dst, indices, 64, num_groups, 'NORM_B32') + transpose_vectors(src, dst, indices, 64, num_groups) @T.macro def dequantize_input(x_ub, scales_ub, out_ub): with T.SimdVF(): - mask_low = S.pset(32, 'PAT_VL32') if in_config.use_packed_ue8m0: - sf_shift = S.vsel(S.vdup(23, T.int32), S.vdup(15, T.int32), mask_low) - sf_exp_mask = S.vdup(0x7F800000, T.uint32) + mask_low = V.create_mask(32, size=128) + sf_shift = V.vsel(mask_low, V.vbrc(T.uint32(23), size=128), V.vbrc(T.uint32(15), size=128)) + sf_exp_mask = V.vbrc(T.uint32(0x7F800000), size=128) for pair in T.serial(block_k // 128): col = pair * 128 group = pair * 4 if in_config.use_packed_ue8m0: - scale0 = load_packed_input_scale_pair(scales_ub, group, sf_shift, sf_exp_mask) - scale1 = load_packed_input_scale_pair(scales_ub, group + 2, sf_shift, sf_exp_mask) + scale = load_packed_input_scale_bf16(scales_ub, group, sf_shift, sf_exp_mask) + if value_dtype != T.bfloat16: + scale = V.vcvt(scale, 'float32') else: - scale0 = S.vsel(S.vld(scales_ub[group], dist='BRC_B32'), S.vld(scales_ub[group + 1], dist='BRC_B32'), mask_low) - scale1 = S.vsel(S.vld(scales_ub[group + 2], dist='BRC_B32'), S.vld(scales_ub[group + 3], dist='BRC_B32'), mask_low) - if value_dtype == T.bfloat16: - _, scale_bf16 = S.vdintlv(T.reinterpret(scale0, 'uint16x128'), T.reinterpret(scale1, 'uint16x128')) - scale_bf16 = T.reinterpret(scale_bf16, 'bfloat16x128') + scale = V.vload(scales_ub[group], size=128, dist_mode='brc', group=4, stride=1) + if value_dtype == T.bfloat16: + scale = V.vcvt(scale, 'bfloat16') for row in T.serial(block_m): + raw = V.vload(x_ub[row, col], size=128) if value_dtype == T.bfloat16: - if in_config.dtype == T.float8_e4m3fn: - x0 = S.vcvt(S.vld(x_ub[row, col], dist='UNPK4_B8'), T.float32) - x1 = S.vcvt(S.vld(x_ub[row, col + 64], dist='UNPK4_B8'), T.float32) - _, x_bf16 = S.vdintlv(T.reinterpret(x0, 'uint16x128'), T.reinterpret(x1, 'uint16x128')) - x_bf16 = T.reinterpret(x_bf16, 'bfloat16x128') - else: - x_bf16 = S.vcvt(S.vld(x_ub[row, col], dist='UNPK4_B8'), T.bfloat16) - S.vsts(out_ub[row, col], S.vmul(x_bf16, scale_bf16), dist='NORM_B16') + values = V.vcvt(V.vcvt(raw, 'float32'), 'bfloat16') if in_config.dtype == T.float8_e4m3fn else V.vcvt(raw, 'bfloat16') else: - if in_config.dtype == T.float8_e4m3fn: - x0 = S.vcvt(S.vld(x_ub[row, col], dist='UNPK4_B8'), T.float32) - x1 = S.vcvt(S.vld(x_ub[row, col + 64], dist='UNPK4_B8'), T.float32) - else: - packed = S.vld(x_ub[row, col], dist='UNPK4_B8') - x_bf16 = S.vcvt(packed, T.bfloat16) - even = S.vcvt(x_bf16, T.float32, part=0) - odd = S.vcvt(x_bf16, T.float32, part=1) - x0, x1 = S.vintlv(even, odd) - S.vsts(out_ub[row, col], S.vmul(x0, scale0)) - S.vsts(out_ub[row, col + 64], S.vmul(x1, scale1)) + values = V.vcvt(raw, 'float32') if in_config.dtype == T.float8_e4m3fn else V.vcvt(V.vcvt(raw, 'bfloat16'), 'float32') + V.vstore(V.vmul(values, scale), out_ub[row, col]) @T.macro def quantize_bf16(x_ub, amax_ub, sf_inv_ub, sf_dense_ub, out_ub): with T.SimdVF(): - mask_vl8 = S.pset(32, 'PAT_VL8') - abs_mask = S.vdup(0x7FFF, T.uint16) - zero = S.vdup(0.0, T.bfloat16) + mask256 = V.create_mask(256, size=256) + abs_mask = V.vbrc(T.uint16(0x7FFF), size=256) for row in T.serial(block_m): - for tile in T.serial(block_k // 256): - x0, x1 = S.vld2(x_ub[row, tile * 256], dist='DINTLV_B16') - abs_x0 = S.vand(T.reinterpret(x0, 'uint16x128'), abs_mask) - abs_x1 = S.vand(T.reinterpret(x1, 'uint16x128'), abs_mask) - maxima = S.vcgmax(S.vmax(abs_x0, abs_x1)) - dense, _ = S.vintlv(zero, T.reinterpret(maxima, 'bfloat16x128')) - S.vsts(amax_ub[row, tile * 8], T.reinterpret(dense, 'float32x64'), mask_vl8, dist='NORM_B32', extent=8) - S.mem_bar('VST_VLD') + for tile in T.unroll(block_k // 256, explicit=True): + col = tile * 256 + amax_u = V.vcmax(V.vand(V.vinterpret_cast(V.vload(x_ub[row, col], size=256), 'uint16'), abs_mask), mask256, group=8) + V.vstore(V.vcvt(V.vinterpret_cast(amax_u, 'bfloat16'), 'float32'), amax_ub[row, tile * 8], group=8, stride=1) + T.simd.mem_bar('VST_VLD') for row in T.serial(block_m): - scale, inverse = compute_scale(S.vld(amax_ub[row, 0])) - S.vsts(sf_dense_ub[row, 0], T.reinterpret(scale, 'uint8x256'), dist='PK4_B32') - S.vsts(sf_inv_ub[row, 0], S.vcvt(inverse, T.bfloat16), dist='PK_B32') + scale, inverse = compute_scale(V.vload(amax_ub[row, 0], size=64), 64) + V.vstore(V.vcvt(scale, 'uint8'), sf_dense_ub[row, 0]) + V.vstore(V.vcvt(inverse, 'bfloat16'), sf_inv_ub[row, 0]) if not sf_only: - S.mem_bar('VST_VLD') + T.simd.mem_bar('VST_VLD') for row in T.serial(block_m): for tile in T.serial(block_k // 256): col = tile * 256 - x0, x1 = S.vld2(x_ub[row, col], dist='DINTLV_B16') - inverse = S.vld(sf_inv_ub[row, tile * 8], dist='E2B_B16') - quantized0 = S.vmul(x0, inverse) - quantized1 = S.vmul(x1, inverse) + inverse = V.vload(sf_inv_ub[row, tile * 8], size=256, dist_mode='brc', group=8, stride=1) + values = V.vmul(V.vload(x_ub[row, col], size=256), inverse) if out_config.dtype == T.float4_e2m1fn: - c0, c1 = S.vintlv(quantized0, quantized1) - S.vsts(out_ub[row, col], S.vcvt(c0, T.float4_e2m1fn), dist='PK4_B32') - S.vsts(out_ub[row, col + 128], S.vcvt(c1, T.float4_e2m1fn), dist='PK4_B32') + V.vstore(V.vcvt(values, 'float4_e2m1fn', rounding='R'), out_ub[row, col]) else: - quantized_fp8_0 = S.vcvt(S.vcvt(quantized0, T.float32, part=0), T.float8_e4m3fn, part=0) - quantized_fp8_2 = S.vcvt(S.vcvt(quantized0, T.float32, part=1), T.float8_e4m3fn, part=2) - quantized_fp8_1 = S.vcvt(S.vcvt(quantized1, T.float32, part=0), T.float8_e4m3fn, part=1) - quantized_fp8_3 = S.vcvt(S.vcvt(quantized1, T.float32, part=1), T.float8_e4m3fn, part=3) - merged = S.vor( - S.vor(T.reinterpret(quantized_fp8_0, 'uint8x256'), T.reinterpret(quantized_fp8_2, 'uint8x256')), - S.vor(T.reinterpret(quantized_fp8_1, 'uint8x256'), T.reinterpret(quantized_fp8_3, 'uint8x256')), - ) - S.vsts(out_ub[row, col], T.reinterpret(merged, 'float8_e4m3fnx256'), dist='NORM_B8') + V.vstore(V.vcvt(V.vcvt(values, 'float32'), 'float8_e4m3fn', rounding='R', saturate='SAT'), out_ub[row, col]) @T.macro def load_f32_pair(x_ub, row, col): if in_config.dtype == T.bfloat16: - return S.vcvt(S.vld(x_ub[row, col], dist='UNPK_B16'), T.float32), S.vcvt(S.vld(x_ub[row, col + 64], dist='UNPK_B16'), T.float32) - return S.vld(x_ub[row, col]), S.vld(x_ub[row, col + 64]) + return V.vcvt(V.vload(x_ub[row, col], size=128), 'float32') + return V.vload(x_ub[row, col], size=128) @T.macro def quantize_f32(x_ub, amax_ub, sf_inv_ub, sf_dense_ub, out_ub): with T.SimdVF(): - mask_low = S.pset(32, 'PAT_VL32') + mask64 = V.create_mask(64, size=64) + mask128 = V.create_mask(128, size=128) if cast_only: - one = S.vdup(1.0, T.float32) + one = V.vbrc(T.float32(1.0), size=64) for row in T.serial(block_m): - S.vsts(sf_inv_ub[row, 0], S.vdiv(one, S.vld(sf_dense_ub[row, 0]))) + V.vstore(V.vdiv(one, V.vload(sf_dense_ub[row, 0], size=64), mask64, precision='exact'), sf_inv_ub[row, 0]) else: - mask_all = S.pset(32, 'PAT_ALL') - mask_high = S.pnot(mask_low, mask_all) for row in T.serial(block_m): for pair in T.serial(block_k // 128): col = pair * 128 - x0, x1 = load_f32_pair(x_ub, row, col) - group = pair * 4 - abs_x0, abs_x1 = S.vabs(x0), S.vabs(x1) - S.vsts(amax_ub[row, group], S.vcmax(abs_x0, mask_low), dist='ONEPT_B32') - S.vsts(amax_ub[row, group + 1], S.vcmax(abs_x0, mask_high), dist='ONEPT_B32') - S.vsts(amax_ub[row, group + 2], S.vcmax(abs_x1, mask_low), dist='ONEPT_B32') - S.vsts(amax_ub[row, group + 3], S.vcmax(abs_x1, mask_high), dist='ONEPT_B32') - S.mem_bar('VST_VLD') + V.vstore(V.vcmax(V.vabs(load_f32_pair(x_ub, row, col)), mask128, group=4), amax_ub[row, pair * 4]) + T.simd.mem_bar('VST_VLD') for row in T.serial(block_m): - scale, inverse = compute_scale(S.vld(amax_ub[row, 0])) - S.vsts(sf_inv_ub[row, 0], inverse) + scale, inverse = compute_scale(V.vload(amax_ub[row, 0], size=64), 64) if out_config.use_packed_ue8m0: - S.vsts(sf_dense_ub[row, 0], T.reinterpret(scale, 'uint8x256'), dist='PK4_B32') + V.vstore(V.vcvt(scale, 'uint8'), sf_dense_ub[row, 0]) else: - S.vsts(sf_dense_ub[row, 0], scale) + V.vstore(scale, sf_dense_ub[row, 0]) + V.vstore(inverse, sf_inv_ub[row, 0]) if not sf_only: - S.mem_bar('VST_VLD') + T.simd.mem_bar('VST_VLD') for row in T.serial(block_m): for pair in T.serial(block_k // 128): col = pair * 128 - group = pair * 4 - sf_inv0 = S.vld(sf_inv_ub[row, group], dist='BRC_B32') - sf_inv1 = S.vld(sf_inv_ub[row, group + 1], dist='BRC_B32') - sf_inv2 = S.vld(sf_inv_ub[row, group + 2], dist='BRC_B32') - sf_inv3 = S.vld(sf_inv_ub[row, group + 3], dist='BRC_B32') - x0, x1 = load_f32_pair(x_ub, row, col) - quantized0 = S.vmul(x0, S.vsel(sf_inv0, sf_inv1, mask_low)) - quantized1 = S.vmul(x1, S.vsel(sf_inv2, sf_inv3, mask_low)) + inverse = V.vload(sf_inv_ub[row, pair * 4], size=128, dist_mode='brc', group=4, stride=1) + quantized = V.vmul(load_f32_pair(x_ub, row, col), inverse) if out_config.dtype == T.float4_e2m1fn: - low, high = S.vdintlv(T.reinterpret(quantized0, 'uint16x128'), T.reinterpret(quantized1, 'uint16x128')) - packed = T.reinterpret(S.vor(high, S.vmins(low, 1)), 'bfloat16x128') - S.vsts(out_ub[row, col], S.vcvt(packed, T.float4_e2m1fn), dist='PK4_B32') + V.vstore(V.vcvt(f32_to_bf16_round_odd(quantized), 'float4_e2m1fn', rounding='R'), out_ub[row, col]) else: - S.vsts(out_ub[row, col], S.vcvt(quantized0, T.float8_e4m3fn), dist='PK4_B32') - S.vsts(out_ub[row, col + 64], S.vcvt(quantized1, T.float8_e4m3fn), dist='PK4_B32') + V.vstore(V.vcvt(quantized, 'float8_e4m3fn', rounding='R', saturate='SAT'), out_ub[row, col]) @T.prim_func def per_token_cast_kernel(