From 05949dfb284fcdcb7c41bf2bca5bf8101890e84e Mon Sep 17 00:00:00 2001 From: kuri780 Date: Mon, 19 Oct 2026 11:39:34 +0000 Subject: [PATCH 1/7] quant: add SwiGLU forward/backward VMI kernels Port the ASC SwiGLU forward and backward kernels to the VMI (PTO) backend following the PTOAS #1616 migration rules: keep the outer T.Kernel/T.Persistent/T.copy/tiling/SF layout unchanged and rewrite only the T.SimdVF bodies with logical T.vmi operations. Co-Authored-By: Claude Code --- tile_kernels/quant/swiglu_backward_vmi.py | 546 ++++++++++++++++++++++ tile_kernels/quant/swiglu_forward_vmi.py | 462 ++++++++++++++++++ 2 files changed, 1008 insertions(+) create mode 100644 tile_kernels/quant/swiglu_backward_vmi.py create mode 100644 tile_kernels/quant/swiglu_forward_vmi.py diff --git a/tile_kernels/quant/swiglu_backward_vmi.py b/tile_kernels/quant/swiglu_backward_vmi.py new file mode 100644 index 00000000..c9f9140e --- /dev/null +++ b/tile_kernels/quant/swiglu_backward_vmi.py @@ -0,0 +1,546 @@ +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 tilelang.transform import PassConfigKey + +from tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape +from tile_kernels.utils import align + +# Match the ASC source of truth: bare S.vdiv is precise unless fast math is +# explicitly enabled. PTO VMI's hardware vdiv is the explicit 'ftz_true' path. +PTO_FAST_MATH = {PassConfigKey.TL_ENABLE_FAST_MATH: False} + + +@tilelang.jit(target='pto', pass_configs=PTO_FAST_MATH) +def get_swiglu_backward_kernel_vmi( + hidden: int, + num_experts: int, + alignment: int, + with_weight: bool, + with_routed_scaling: bool, + use_clamp: bool, + x_dtype: T.dtype, + act_x_grad_dtype: T.dtype, + out_config: CastOutputConfig, + do_recompute: bool, + num_vec_cores: int, +): + num_stages = 2 + + with_quant = out_config.with_sf + out_dtype = T.bfloat16 if with_quant else out_config.dtype + + if with_quant: + assert x_dtype == T.float8_e4m3fn + assert out_config.dtype == T.float8_e4m3fn + assert out_config.sf_block == (1, 32) + assert not out_config.use_packed_ue8m0 or out_config.round_sf + else: + assert x_dtype in (T.bfloat16, T.float32) + assert out_dtype in (T.bfloat16, T.float32) + assert hidden % 128 == 0 + + # Keep each SF column at the 32-byte DMA granularity. + block_m = (16 if out_config.use_packed_ue8m0 else 8) if with_quant else 1 + if num_experts: + assert alignment % block_m == 0, f'Ascend swiglu backward requires alignment % {block_m} == 0, got {alignment}' + + vec_size = 64 + block_k = min(hidden, 4096) + while hidden % block_k != 0: + block_k -= vec_size + assert block_k >= 128 and hidden % block_k == 0 + + pack_factor = get_packed_ue8m0_pack_factor() if with_quant and out_config.use_packed_ue8m0 else 1 + num_hidden_blocks = hidden // block_k + num_groups = hidden // 32 if with_quant else 0 + num_block_groups = block_k // 32 if with_quant else 0 + num_stored_groups = num_groups // pack_factor + num_stored_block_groups = num_block_groups // pack_factor + # Pad each hidden block to one full vector store. + sf_block_stride = align(num_stored_block_groups, vec_size // pack_factor) if with_quant else 0 + sf_half_stride = sf_block_stride * num_hidden_blocks + sf_row_stride = sf_half_stride * 2 + num_scale_vectors = (num_block_groups + vec_size - 1) // vec_size if with_quant else 0 + sf_storage_dtype = T.uint16 if out_config.use_packed_ue8m0 else T.float32 + stored_groups_per_vector = vec_size * pack_factor // block_m + # Pad transpose UB to avoid unsupported partial-vector PAT_VL masks. + aligned_num_stored_block_groups = align(num_stored_block_groups, stored_groups_per_vector) + + num_expanded_tokens = T.dynamic('num_expanded_tokens') + out_sf_stride = T.dynamic('out_sf_stride') + x_sf_shape = (num_expanded_tokens, num_groups * 2) if with_quant else (1, 1) + out_sf_shape = get_sf_shape((num_expanded_tokens, hidden * 2), out_config) if with_quant else (1, 1) + x_grad_fp8_shape = (num_expanded_tokens, hidden * 2) if with_quant else (1, 1) + + @T.macro + def vdiv_0ulp(a, b): + """Correct a / b to 0 ULP (same algorithm as the ASC source).""" + full = V.create_mask(vec_size, size=vec_size) + z = V.vdiv(a, b, full) + y = V.vneg(b) + r = V.alloc_var(T.float32, size=vec_size) + r = V.vmula(a, z, y, full) + z_bits = V.vinterpret_cast(z, 'int32') + z_pre = V.vinterpret_cast(V.vadds(z_bits, T.int32(-1), full), 'float32') + z_next = V.vinterpret_cast(V.vadds(z_bits, T.int32(1), full), 'float32') + r_pre = V.alloc_var(T.float32, size=vec_size) + r_pre = V.vmula(a, z_pre, y, full) + r_next = V.alloc_var(T.float32, size=vec_size) + r_next = V.vmula(a, z_next, y, full) + r = V.vabs(r) + r_pre = V.vabs(r_pre) + r_next = V.vabs(r_next) + take_pre = V.vcmp(r_pre, r, full, 'le') + best_r = V.vsel(take_pre, r_pre, r) + best_z = V.vsel(take_pre, z_pre, z) + take_next = V.vcmp(r_next, best_r, full, 'lt') + corrected = V.vsel(take_next, z_next, best_z) + is_zero = V.vcmps(z, T.float32(0.0), full, 'eq') + return V.vsel(is_zero, z, corrected) + + @T.macro + def sigmoid_vdiv_0ulp(a, b): + """Correct a / b for sigmoid (same algorithm as the ASC source). + + PTO vmi.vsel only accepts vector-register operands, so the predicate + xor is materialized as int32 vectors. + """ + full = V.create_mask(vec_size, size=vec_size) + one_i32 = V.vbrc(T.int32(1), size=vec_size) + zero_i32 = V.vbrc(T.int32(0), size=vec_size) + z = V.vdiv(a, b, full) + neg_b = V.vneg(b) + + residual = V.alloc_var(T.float32, size=vec_size) + residual = V.vmula(a, z, neg_b, full) + + residual_bits = V.vinterpret_cast(residual, 'int32') + sign_mask = V.vshrs(residual_bits, T.int32(31), full) + use_bits_minus = V.vcmps(residual_bits, T.int32(0), full, 'lt') + delta = V.vor(sign_mask, one_i32, full) + adjacent = V.vinterpret_cast(V.vadd(V.vinterpret_cast(z, 'int32'), delta, full), 'float32') + + adjacent_residual = V.alloc_var(T.float32, size=vec_size) + adjacent_residual = V.vmula(a, adjacent, neg_b, full) + residual_sum = V.vadd(residual, adjacent_residual, full) + take_bits_plus = V.vcmps(residual_sum, T.float32(0.0), full, 'gt') + is_nonzero = V.vcmps(z, T.float32(0.0), full, 'ne') + + take_plus_i = V.vsel(take_bits_plus, one_i32, zero_i32) + use_minus_i = V.vsel(use_bits_minus, one_i32, zero_i32) + nonzero_i = V.vsel(is_nonzero, one_i32, zero_i32) + take_adjacent_i = V.vand(V.vxor(take_plus_i, use_minus_i, full), nonzero_i, full) + take_adjacent = V.vcmps(take_adjacent_i, T.int32(0), full, 'ne') + return V.vsel(take_adjacent, adjacent, z) + + @T.macro + def get_scale_and_inverse(amax): + # Compute one vector of output scales and inverse scales from per-group amax. + full = V.create_mask(vec_size, size=vec_size) + clamped_amax = V.vmaxs(amax, T.float32(out_config.clamp_min_value), full) + if out_config.round_sf: + exponent_mask = V.vbrc(T.uint32(0x7F800000), size=vec_size) + exponent_carry = V.vbrc(T.uint32(0x001FFFFF), size=vec_size) + inverse_bias = V.vbrc(T.uint32(0x83000000), size=vec_size) + scale_bias = V.vbrc(T.uint32(0x04000000), size=vec_size) + exponent = V.vand( + V.vadd(V.vinterpret_cast(clamped_amax, 'uint32'), exponent_carry, full), + exponent_mask, + full, + ) + scale = V.vinterpret_cast(V.vsub(exponent, scale_bias, full), 'float32') + inverse = V.vinterpret_cast(V.vsub(inverse_bias, exponent, full), 'float32') + return scale, inverse + quant_max = V.vbrc(T.float32(448.0), size=vec_size) + return ( + V.vdiv(clamped_amax, quant_max, full, precision='exact'), + V.vdiv(quant_max, clamped_amax, full, precision='exact'), + ) + + @T.macro + def init_transpose_indices(indices): + # Precompute gather indices shared by every token block when converting + # dense token-major scales to the output layout. + with T.SimdVF(): + if out_config.use_packed_ue8m0: + lane = V.vci(T.int16(0), size=128) + row = V.vand(V.vinterpret_cast(lane, 'uint16'), V.vbrc(T.uint16(block_m - 1), size=128)) + group = V.vinterpret_cast(V.vshr(lane, V.vbrc(T.int16(block_m.bit_length() - 1), size=128)), 'uint16') + source = V.vadd(V.vmul(row, V.vbrc(T.uint16(sf_row_stride), size=128)), group) + V.vstore(source, indices[0]) + else: + lane = V.vci(T.int32(0), size=64) + row = V.vand(V.vinterpret_cast(lane, 'uint32'), V.vbrc(T.uint32(block_m - 1), size=64)) + group = V.vinterpret_cast(V.vshr(lane, V.vbrc(T.int32(block_m.bit_length() - 1), size=64)), 'uint32') + source = V.vadd(V.vmul(row, V.vbrc(T.uint32(sf_row_stride), size=64)), group) + V.vstore(source, indices[0]) + + @T.macro + def transpose_sf(src, dst, indices, source_col): + # Transpose one hidden block of dense scales into the stored layout. + with T.SimdVF(): + if out_config.use_packed_ue8m0: + vector_lanes = 128 + else: + vector_lanes = 64 + 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_stored_block_groups + stored_groups_per_vector - 1) // stored_groups_per_vector, explicit=True): + group_base = output_vector * stored_groups_per_vector + values = V.vgather(src[0, source_col + group_base], source_indices, full_mask) + V.vstore(values, dst[group_base, 0]) + + @T.prim_func + def swiglu_backward_kernel_vmi( + x: T.Tensor[(num_expanded_tokens, hidden * 2), x_dtype], + x_sf: T.Tensor[x_sf_shape, T.float32], + act_x_grad: T.Tensor[(num_expanded_tokens, hidden), act_x_grad_dtype], + topk_weights: T.Tensor[(num_expanded_tokens,), T.float32], + routed_scaling_factor: T.float32, + psum_num_tokens_per_expert: T.Tensor[(num_experts,), T.int32], + out: T.Tensor[(num_expanded_tokens, hidden), out_dtype], + x_grad_fp8: T.Tensor[x_grad_fp8_shape, out_config.dtype], + x_grad_fp8_sf: T.StridedTensor[out_sf_shape, (out_sf_stride, 1), out_config.sf_dtype], + x_grad: T.Tensor[(num_expanded_tokens, hidden * 2), out_dtype], + weight_grad: T.Tensor[(num_expanded_tokens,), T.float32], + clamp_value: T.float32, + ): + with T.Kernel(num_vec_cores) as core_id: + x_ub = T.alloc_shared((block_k,), x_dtype) + y_ub = T.alloc_shared((block_k,), x_dtype) + grad_ub = T.alloc_shared((block_k,), act_x_grad_dtype) + x_grad_ub = T.alloc_shared((block_k,), out_dtype) + y_grad_ub = T.alloc_shared((block_k,), out_dtype) + if with_quant: + x_sf_ub = T.alloc_shared((sf_block_stride * pack_factor,), T.float32) + y_sf_ub = T.alloc_shared((sf_block_stride * pack_factor,), T.float32) + x_grad_f32_ub = T.alloc_shared((block_k,), T.float32) + y_grad_f32_ub = T.alloc_shared((block_k,), T.float32) + x_grad_fp8_ub = T.alloc_shared((block_k,), T.float8_e4m3fn) + y_grad_fp8_ub = T.alloc_shared((block_k,), T.float8_e4m3fn) + # Reused for compact amax and inverse scales (one value per 32 channels). + x_quant_scratch_ub = T.alloc_shared((sf_block_stride * pack_factor * 4,), T.float32) + y_quant_scratch_ub = T.alloc_shared((sf_block_stride * pack_factor * 4,), T.float32) + # Packed SF is written as bytes and transposed through a uint16 view. + sf_dense_ub = T.alloc_shared((block_m, sf_row_stride), sf_storage_dtype) + if out_config.use_packed_ue8m0: + sf_dense_u8 = T.Tensor((block_m, sf_row_stride * 2), T.uint8, sf_dense_ub.data) + if out_config.use_tma_aligned_col_major_sf: + x_sf_out_ub = T.alloc_shared((aligned_num_stored_block_groups, block_m), sf_storage_dtype) + y_sf_out_ub = T.alloc_shared((aligned_num_stored_block_groups, block_m), sf_storage_dtype) + transpose_indices_ub = T.alloc_shared((vec_size * pack_factor,), T.uint16 if out_config.use_packed_ue8m0 else T.uint32) + + if with_weight: + weight_acc_ub = T.alloc_shared((block_m, vec_size), T.float32) + weight_out_ub = T.alloc_shared((vec_size,), T.float32) + if do_recompute: + act_out_ub = T.alloc_shared((block_k,), out_dtype) + if num_experts: + psum_ub = T.alloc_shared((num_experts,), T.int32) + + versions = { + x_ub: num_stages, + y_ub: num_stages, + grad_ub: num_stages, + x_grad_ub: num_stages, + y_grad_ub: num_stages, + } + if with_quant: + versions.update( + { + x_sf_ub: num_stages, + y_sf_ub: num_stages, + # Vector-local scratch does not need MTE overlap versions. + x_grad_f32_ub: 1, + y_grad_f32_ub: 1, + x_grad_fp8_ub: num_stages, + y_grad_fp8_ub: num_stages, + x_quant_scratch_ub: 1, + y_quant_scratch_ub: 1, + sf_dense_ub: 1, + } + ) + if with_quant and out_config.use_tma_aligned_col_major_sf: + versions[x_sf_out_ub] = 1 + versions[y_sf_out_ub] = 1 + versions[transpose_indices_ub] = 1 + if with_weight: + versions[weight_acc_ub] = 1 + versions[weight_out_ub] = num_stages + if do_recompute: + versions[act_out_ub] = num_stages + if num_experts: + versions[psum_ub] = 1 + T.annotate_buffer_versions(versions) + + if with_quant and out_config.use_tma_aligned_col_major_sf: + init_transpose_indices(transpose_indices_ub) + if num_experts: + T.copy(psum_num_tokens_per_expert, psum_ub[:num_experts]) + expert_id = T.alloc_var(T.int32, init=0) + valid_token_end = T.alloc_var(T.int32, init=0) + + num_token_blocks = num_expanded_tokens // block_m + blocks_per_core = num_token_blocks // num_vec_cores + num_long_cores = num_token_blocks % num_vec_cores + # Assign contiguous block ranges per core to spread concurrent GM traffic across L2 home paths. + core_block_start = core_id * blocks_per_core + T.min(core_id, num_long_cores) + + for scheduler_pid in T.Persistent( + [num_token_blocks], + num_vec_cores, + core_id, + group_size=1, + num_stages=num_stages, + ): + wave = scheduler_pid // num_vec_cores + token_block = core_block_start + wave + token_base = token_block * block_m + if num_experts: + while expert_id < num_experts and token_base >= align(psum_ub[expert_id], alignment): + expert_id += 1 + if expert_id < num_experts: + valid_token_end = psum_ub[expert_id] + else: + valid_token_end = 0 + + for hidden_block in range(num_hidden_blocks): + col_offset = hidden_block * block_k + sf_offset = hidden_block * num_block_groups + sf_block_offset = hidden_block * sf_block_stride + + for row in T.serial(block_m): + token_id = token_base + row + valid_token_f32 = T.alloc_var(T.float32, init=1.0) + if num_experts: + # PTO codegen does not support scalar bool -> float32 casts. + valid_token_f32 = T.float32(0.0) + if expert_id < num_experts and token_id < valid_token_end: + valid_token_f32 = T.float32(1.0) + + effective_weight = T.alloc_var(T.float32, init=valid_token_f32) + if with_weight: + effective_weight = topk_weights[token_id] * valid_token_f32 + if with_routed_scaling: + effective_weight = effective_weight * routed_scaling_factor + + T.copy(x[token_id, col_offset : col_offset + block_k], x_ub) + T.copy(x[token_id, hidden + col_offset : hidden + col_offset + block_k], y_ub) + if with_quant: + T.copy(x_sf[token_id, sf_offset : sf_offset + num_block_groups], x_sf_ub[:num_block_groups]) + T.copy(x_sf[token_id, num_groups + sf_offset : num_groups + sf_offset + num_block_groups], y_sf_ub[:num_block_groups]) + T.copy(act_x_grad[token_id, col_offset : col_offset + block_k], grad_ub) + + with T.SimdVF(): + full = V.create_mask(vec_size, size=vec_size) + zeros = V.vbrc(T.float32(0.0), size=vec_size) + ones = V.vbrc(T.float32(1.0), size=vec_size) + if with_quant: + high_half_mask = V.vcmps(V.vci(T.int32(0), size=vec_size), T.int32(31), full, 'gt') + low_half_mask = V.create_mask(32, size=vec_size) + if with_weight: + weight_acc = V.alloc_var(T.float32, size=vec_size) + # hidden_block is a runtime TIR loop variable; use an + # if statement so PTO does not see a vector select. + if hidden_block == 0: + weight_acc = zeros + else: + weight_acc = V.vload(weight_acc_ub[row, 0], size=vec_size) + + for vector in T.serial(block_k // vec_size): + col = vector * vec_size + if with_quant: + x_scale = V.vbrc(V.vload(x_sf_ub[vector * 2], size=2, group=2, stride=1), size=vec_size, group=2) + y_scale = V.vbrc(V.vload(y_sf_ub[vector * 2], size=2, group=2, stride=1), size=vec_size, group=2) + x_raw = V.vcvt(V.vload(x_ub[col], size=vec_size), 'float32') + y_raw = V.vcvt(V.vload(y_ub[col], size=vec_size), 'float32') + x_value = V.vmul(x_raw, x_scale, full) + y_value = V.vmul(y_raw, y_scale, full) + elif x_dtype == T.bfloat16: + x_value = V.vcvt(V.vload(x_ub[col], size=vec_size), 'float32') + y_value = V.vcvt(V.vload(y_ub[col], size=vec_size), 'float32') + else: + x_value = V.vload(x_ub[col], size=vec_size) + y_value = V.vload(y_ub[col], size=vec_size) + if act_x_grad_dtype == T.bfloat16: + grad = V.vcvt(V.vload(grad_ub[col], size=vec_size), 'float32') + else: + grad = V.vload(grad_ub[col], size=vec_size) + + if use_clamp: + is_clamped_x = V.vcmps(x_value, T.float32(clamp_value), full, 'gt') + is_clamped_y = V.vcmps(V.vabs(y_value), T.float32(clamp_value), full, 'gt') + x_compute_value = V.vmin(x_value, V.vbrc(T.float32(clamp_value), size=vec_size)) + y_compute_value = V.vmin( + V.vmax(y_value, V.vbrc(T.float32(-clamp_value), size=vec_size)), + V.vbrc(T.float32(clamp_value), size=vec_size), + ) + else: + x_compute_value = x_value + y_compute_value = y_value + + sigmoid_denominator = V.vadds(V.vexpdif(zeros, x_compute_value, full), T.float32(1.0), full) + sigmoid = sigmoid_vdiv_0ulp(ones, sigmoid_denominator) + if with_weight or do_recompute: + # Keep a separate 0-ULP division for the required FP32 rounding. + act_value = V.vmul(vdiv_0ulp(x_compute_value, sigmoid_denominator), y_compute_value, full) + if with_weight: + weight_contribution = V.vmul(grad, act_value, full) + weight_value = V.vmuls(weight_contribution, valid_token_f32, full) if num_experts else weight_contribution + weight_acc = V.vadd(weight_acc, weight_value, full) + grad_ws = V.vmul(V.vmuls(grad, effective_weight, full), sigmoid, full) + inner = V.vadd(ones, V.vmul(x_compute_value, V.vsub(ones, sigmoid, full), full), full) + x_grad_unclamped = V.vmul(V.vmul(grad_ws, y_compute_value, full), inner, full) + y_grad_unclamped = V.vmul(grad_ws, x_compute_value, full) + x_grad_value = V.vsel(is_clamped_x, zeros, x_grad_unclamped) if use_clamp else x_grad_unclamped + y_grad_value = V.vsel(is_clamped_y, zeros, y_grad_unclamped) if use_clamp else y_grad_unclamped + if not with_quant: + if out_dtype == T.bfloat16: + V.vstore(V.vcvt(x_grad_value, 'bfloat16'), x_grad_ub[col]) + V.vstore(V.vcvt(y_grad_value, 'bfloat16'), y_grad_ub[col]) + else: + V.vstore(x_grad_value, x_grad_ub[col]) + V.vstore(y_grad_value, y_grad_ub[col]) + else: + V.vstore(x_grad_value, x_grad_f32_ub[col]) + V.vstore(y_grad_value, y_grad_f32_ub[col]) + x_abs = V.vabs(x_grad_value) + y_abs = V.vabs(y_grad_value) + V.vstore(V.vcmax(x_abs, full, group=2), x_quant_scratch_ub[vector * 2]) + V.vstore(V.vcmax(y_abs, full, group=2), y_quant_scratch_ub[vector * 2]) + if do_recompute: + act_out_value = V.vmuls(act_value, effective_weight, full) + if out_dtype == T.bfloat16: + V.vstore(V.vcvt(act_out_value, 'bfloat16'), act_out_ub[col]) + else: + V.vstore(act_out_value, act_out_ub[col]) + + if with_weight: + if hidden_block + 1 == num_hidden_blocks: + one_mask = V.create_mask(1, size=1) + reduced_weight_sum = V.vcadd(weight_acc, full, reassoc=True) + weight_sum = ( + V.vmuls(reduced_weight_sum, T.float32(routed_scaling_factor), one_mask) + if with_routed_scaling + else reduced_weight_sum + ) + V.vstore(weight_sum, weight_out_ub[0], group=1, stride=1) + else: + V.vstore(weight_acc, weight_acc_ub[row, 0]) + + if with_quant: + S.mem_bar('VST_VLD') + if num_experts: + valid_token_vec = V.vbrc(valid_token_f32, size=vec_size) + # Preserve unread amax while overwriting the scratch with inverses. + for scale_vector_index in range(num_scale_vectors): + scale_vector = num_scale_vectors - 1 - scale_vector_index + scale_offset = scale_vector * vec_size + x_scale_value, x_inv = get_scale_and_inverse(V.vload(x_quant_scratch_ub[scale_offset], size=vec_size)) + y_scale_value, y_inv = get_scale_and_inverse(V.vload(y_quant_scratch_ub[scale_offset], size=vec_size)) + V.vstore(x_inv, x_quant_scratch_ub[scale_offset]) + V.vstore(y_inv, y_quant_scratch_ub[scale_offset]) + x_scale_out = V.vmul(x_scale_value, valid_token_vec, full) if num_experts else x_scale_value + y_scale_out = V.vmul(y_scale_value, valid_token_vec, full) if num_experts else y_scale_value + if out_config.use_packed_ue8m0: + x_exp = V.vshr(V.vinterpret_cast(x_scale_out, 'uint32'), V.vbrc(T.uint32(23), size=vec_size)) + y_exp = V.vshr(V.vinterpret_cast(y_scale_out, 'uint32'), V.vbrc(T.uint32(23), size=vec_size)) + V.vstore(V.vcvt(x_exp, 'uint8'), sf_dense_u8[row, sf_block_offset * 2 + scale_offset]) + V.vstore( + V.vcvt(y_exp, 'uint8'), + sf_dense_u8[row, sf_half_stride * 2 + sf_block_offset * 2 + scale_offset], + ) + else: + V.vstore(x_scale_out, sf_dense_ub[row, sf_block_offset + scale_offset]) + V.vstore(y_scale_out, sf_dense_ub[row, sf_half_stride + sf_block_offset + scale_offset]) + + S.mem_bar('VST_VLD') + if num_experts: + valid_token_mask = V.vcmps(valid_token_vec, T.float32(0.0), full, 'gt') + for vector in T.serial(block_k // vec_size): + col = vector * vec_size + x_inv_vec = V.vbrc(V.vload(x_quant_scratch_ub[vector * 2], size=2, group=2, stride=1), size=vec_size, group=2) + y_inv_vec = V.vbrc(V.vload(y_quant_scratch_ub[vector * 2], size=2, group=2, stride=1), size=vec_size, group=2) + x_grad_quant_value = V.vload(x_grad_f32_ub[col], size=vec_size) + y_grad_quant_value = V.vload(y_grad_f32_ub[col], size=vec_size) + x_scaled_value = V.vmul(x_grad_quant_value, x_inv_vec, full) + y_scaled_value = V.vmul(y_grad_quant_value, y_inv_vec, full) + V.vstore(V.vcvt(x_grad_quant_value, 'bfloat16'), x_grad_ub[col]) + V.vstore(V.vcvt(y_grad_quant_value, 'bfloat16'), y_grad_ub[col]) + x_scaled = V.vsel(valid_token_mask, x_scaled_value, zeros) if num_experts else x_scaled_value + y_scaled = V.vsel(valid_token_mask, y_scaled_value, zeros) if num_experts else y_scaled_value + V.vstore(V.vcvt(x_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), x_grad_fp8_ub[col]) + V.vstore(V.vcvt(y_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), y_grad_fp8_ub[col]) + + T.copy(x_grad_ub, x_grad[token_id, col_offset : col_offset + block_k]) + T.copy(y_grad_ub, x_grad[token_id, hidden + col_offset : hidden + col_offset + block_k]) + if with_quant: + T.copy(x_grad_fp8_ub, x_grad_fp8[token_id, col_offset : col_offset + block_k]) + T.copy(y_grad_fp8_ub, x_grad_fp8[token_id, hidden + col_offset : hidden + col_offset + block_k]) + if do_recompute: + T.copy(act_out_ub, out[token_id, col_offset : col_offset + block_k]) + if with_weight and hidden_block + 1 == num_hidden_blocks: + T.copy(weight_out_ub[:1], weight_grad[token_id : token_id + 1]) + + if with_quant: + if out_config.use_packed_ue8m0: + # PTO's DMA requires matching source/destination element widths. + # Write packed SF through uint8 views instead of ASC's uint16 view. + if out_config.use_tma_aligned_col_major_sf: + output_sf = T.Tensor((num_stored_groups * 2, num_expanded_tokens * 2), T.uint8, x_grad_fp8_sf.data) + else: + output_sf = T.Tensor((num_expanded_tokens, num_stored_groups * 4), T.uint8, x_grad_fp8_sf.data) + else: + output_sf = x_grad_fp8_sf + + for output_block in range(num_hidden_blocks): + sf_source_offset = output_block * sf_block_stride + sf_output_offset = output_block * num_stored_block_groups + if out_config.use_tma_aligned_col_major_sf: + if out_config.use_packed_ue8m0: + x_sf_out_u8 = T.Tensor((aligned_num_stored_block_groups, block_m * 2), T.uint8, x_sf_out_ub.data) + y_sf_out_u8 = T.Tensor((aligned_num_stored_block_groups, block_m * 2), T.uint8, y_sf_out_ub.data) + transpose_sf(sf_dense_ub, x_sf_out_ub, transpose_indices_ub, sf_source_offset) + T.copy( + x_sf_out_u8[0:num_stored_block_groups, 0 : block_m * 2], + output_sf[sf_output_offset, token_base * 2], + ) + transpose_sf(sf_dense_ub, y_sf_out_ub, transpose_indices_ub, sf_half_stride + sf_source_offset) + T.copy( + y_sf_out_u8[0:num_stored_block_groups, 0 : block_m * 2], + output_sf[num_stored_groups + sf_output_offset, token_base * 2], + ) + else: + transpose_sf(sf_dense_ub, x_sf_out_ub, transpose_indices_ub, sf_source_offset) + T.copy(x_sf_out_ub[0:num_stored_block_groups, 0:block_m], output_sf[sf_output_offset, token_base]) + transpose_sf(sf_dense_ub, y_sf_out_ub, transpose_indices_ub, sf_half_stride + sf_source_offset) + T.copy(y_sf_out_ub[0:num_stored_block_groups, 0:block_m], output_sf[num_stored_groups + sf_output_offset, token_base]) + else: + if out_config.use_packed_ue8m0: + T.copy( + sf_dense_u8[0:block_m, sf_source_offset * 2 : (sf_source_offset + num_stored_block_groups) * 2], + output_sf[token_base, sf_output_offset * 2], + ) + T.copy( + sf_dense_u8[ + 0:block_m, + (sf_half_stride + sf_source_offset) * 2 : (sf_half_stride + sf_source_offset + num_stored_block_groups) * 2, + ], + output_sf[token_base, (num_stored_groups + sf_output_offset) * 2], + ) + else: + T.copy( + sf_dense_ub[0:block_m, sf_source_offset : sf_source_offset + num_stored_block_groups], + output_sf[token_base, sf_output_offset], + ) + T.copy( + sf_dense_ub[ + 0:block_m, + sf_half_stride + sf_source_offset : sf_half_stride + sf_source_offset + num_stored_block_groups, + ], + output_sf[token_base, num_stored_groups + sf_output_offset], + ) + + return swiglu_backward_kernel_vmi diff --git a/tile_kernels/quant/swiglu_forward_vmi.py b/tile_kernels/quant/swiglu_forward_vmi.py new file mode 100644 index 00000000..2e49b9b0 --- /dev/null +++ b/tile_kernels/quant/swiglu_forward_vmi.py @@ -0,0 +1,462 @@ +import math + +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 tilelang.transform import PassConfigKey + +from tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape +from tile_kernels.utils import align + +# Match the ASC source of truth: bare S.vdiv is precise unless fast math is +# explicitly enabled. PTO VMI's hardware vdiv is the explicit 'ftz_true' path. +PTO_FAST_MATH = {PassConfigKey.TL_ENABLE_FAST_MATH: False} + + +@tilelang.jit(target='pto', pass_configs=PTO_FAST_MATH) +def get_swiglu_forward_kernel_vmi( + hidden: int, + num_experts: int, + alignment: int, + with_weight: bool, + with_routed_scaling: bool, + clamp_value, + count_clamp: bool, + in_dtype: T.dtype, + out_config: CastOutputConfig, + num_vec_cores: int, +): + has_clamp = clamp_value is not None + has_psum = num_experts > 0 + + assert not count_clamp or has_clamp + assert not with_routed_scaling or with_weight + assert hidden % 128 == 0, f'Ascend swiglu requires hidden % 128 == 0, got {hidden}' + assert not out_config.use_tma_aligned_col_major_sf or out_config.use_packed_ue8m0 + if out_config.with_sf: + assert out_config.sf_block == (1, 32), 'Ascend swiglu only supports sf_block=(1, 32)' + + token_group = 32 if out_config.with_sf and out_config.use_tma_aligned_col_major_sf else 1 + if has_psum: + assert alignment % token_group == 0, f'Ascend swiglu col-major path requires alignment % {token_group} == 0, got {alignment}' + use_compact_groups = has_psum and token_group > 1 + # Fast certificate is only enabled for bf16 -> e4m3 with power-of-two sf. + # It checks that fast and exact division give the same e4m3/sf result; + # unsafe tokens fall back to exact division. + use_fast_certificate = in_dtype == T.bfloat16 and out_config.with_sf and out_config.round_sf and out_config.use_tma_aligned_col_major_sf + num_stages = 3 + # Multiplication by a weight requires a larger certificate error bound. + certificate_error = 5 if with_weight else 2 + vec_size = 64 + num_vecs = hidden // vec_size + hidden_aligned = align(hidden, 256) + + num_expanded_tokens = T.dynamic('num_expanded_tokens') + sf_stride = T.dynamic('sf_stride') + sf_shape = get_sf_shape((num_expanded_tokens, hidden), out_config) if out_config.with_sf else (1, 1) + num_groups = hidden // 32 + pack_factor = get_packed_ue8m0_pack_factor() + + @T.macro + def vdiv_0ulp(a, b): + """Correct a / b to 0 ULP (same algorithm as the ASC backward source).""" + full = V.create_mask(vec_size, size=vec_size) + z = V.vdiv(a, b, full) + y = V.vneg(b) + r = V.alloc_var(T.float32, size=vec_size) + r = V.vmula(a, z, y, full) + z_bits = V.vinterpret_cast(z, 'int32') + z_pre = V.vinterpret_cast(V.vadds(z_bits, T.int32(-1), full), 'float32') + z_next = V.vinterpret_cast(V.vadds(z_bits, T.int32(1), full), 'float32') + r_pre = V.alloc_var(T.float32, size=vec_size) + r_pre = V.vmula(a, z_pre, y, full) + r_next = V.alloc_var(T.float32, size=vec_size) + r_next = V.vmula(a, z_next, y, full) + r = V.vabs(r) + r_pre = V.vabs(r_pre) + r_next = V.vabs(r_next) + take_pre = V.vcmp(r_pre, r, full, 'le') + best_r = V.vsel(take_pre, r_pre, r) + best_z = V.vsel(take_pre, z_pre, z) + take_next = V.vcmp(r_next, best_r, full, 'lt') + corrected = V.vsel(take_next, z_next, best_z) + is_zero = V.vcmps(z, T.float32(0.0), full, 'eq') + return V.vsel(is_zero, z, corrected) + + @T.macro + def get_scale_and_inverse(amax): + full = V.create_mask(vec_size, size=vec_size) + clamped_amax = V.vmaxs(amax, T.float32(out_config.clamp_min_value), full) + scale_bits = V.vinterpret_cast(V.vmuls(clamped_amax, T.float32(1.0 / 448.0), full), 'uint32') + exponent = V.vadds( + V.vshrs(V.vsub(scale_bits, V.vbrc(T.uint32(1), size=vec_size), full), T.uint32(23), full), + T.uint32(1), + full, + ) + inverse = V.vinterpret_cast( + V.vshls(V.vsub(V.vbrc(T.uint32(254), size=vec_size), exponent, full), T.uint32(23), full), + 'float32', + ) + if out_config.use_packed_ue8m0: + return exponent, inverse + return V.vinterpret_cast(V.vshls(exponent, T.uint32(23), full), 'float32'), inverse + + @T.macro + def init_transpose_indices(indices): + # Precompute gather indices shared by every token group when converting the + # dense token-major scale buffer to the packed col-major output layout. + with T.SimdVF(): + lane = V.vci(T.int16(0), size=128) + output_row = V.vand(V.vinterpret_cast(lane, 'uint16'), V.vbrc(T.uint16(token_group - 1), size=128)) + group_in_vector = V.vinterpret_cast( + V.vshr(lane, V.vbrc(T.int16(int(math.log2(token_group))), size=128)), + 'uint16', + ) + source_indices = V.vadd( + V.vmul(output_row, V.vbrc(T.uint16(align(num_groups, 64) // pack_factor), size=128)), + group_in_vector, + ) + V.vstore(source_indices, indices[0]) + + @T.macro + def transpose_packed_scales(src, dst, indices): + # Transpose one token group's packed UE8M0 scales. + with T.SimdVF(): + groups_per_vector = 128 // token_group + num_output_values = num_groups // pack_factor * token_group + num_full_vectors = num_output_values // 128 + num_tail_values = num_output_values % 128 + source_indices = V.vload(indices[0], size=128) + full_mask = V.create_mask(128, size=128) + src_u16 = T.Tensor((token_group, align(num_groups, 64) // pack_factor), T.uint16, src.data) + + if num_full_vectors > 0: + for output_vector in T.unroll(num_full_vectors, explicit=True): + group_base = output_vector * groups_per_vector + values = V.vgather(src_u16[0, group_base], source_indices, full_mask) + V.vstore(values, dst[group_base, 0]) + + if num_tail_values > 0: + tail_group_base = num_full_vectors * groups_per_vector + tail_mask = V.create_mask(num_tail_values, size=128) + values = V.vgather(src_u16[0, tail_group_base], source_indices, tail_mask) + V.vstore(values, dst[tail_group_base, 0], tail_mask) + + @T.macro + def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, with_certificate, pacc_buf, count_clamp, acc=None): + # Compute SiLU(x_l) * x_r, optionally collecting clamp counters, + # quantization amax values, and the fast-division safety certificate. + full = V.create_mask(vec_size, size=vec_size) + zeros = V.vbrc(T.float32(0.0), size=vec_size) + if count_clamp: + if acc is None: + acc = V.alloc_local((3,), V.vreg(vec_size, T.float32)) + for i in T.unroll(3, explicit=True): + acc[i] = zeros + if with_certificate: + certificate_min = V.alloc_var(T.uint32, size=vec_size) + certificate_max = V.alloc_var(T.uint32, size=vec_size) + certificate_min = V.vbrc(T.uint32(0xFFFFFFFF), size=vec_size) + certificate_max = V.vbrc(T.uint32(0), size=vec_size) + else: + certificate_min = None + certificate_max = None + if out_config.with_sf: + amax_acc = V.alloc_local((4,), V.vreg(vec_size, T.float32)) + for i in T.unroll(4, explicit=True): + amax_acc[i] = zeros + + for tile_id in T.serial(T.ceildiv(num_vecs, 4)): + for sub in T.unroll(4, explicit=True): + vector_id = tile_id * 4 + sub + if vector_id < num_vecs: + col = vector_id * vec_size + if in_dtype == T.bfloat16: + val_l = V.vcvt(V.vload(xl_ub[col], size=vec_size), 'float32') + val_r = V.vcvt(V.vload(xr_ub[col], size=vec_size), 'float32') + else: + val_l = V.vload(xl_ub[col], size=vec_size) + val_r = V.vload(xr_ub[col], size=vec_size) + if has_clamp: + silu_l = V.vmin(val_l, V.vbrc(T.float32(clamp_value), size=vec_size)) + mul_r = V.vmax( + V.vmin(val_r, V.vbrc(T.float32(clamp_value), size=vec_size)), + V.vbrc(T.float32(-clamp_value), size=vec_size), + ) + else: + silu_l = val_l + mul_r = val_r + denominator = V.vadds(V.vexpdif(zeros, silu_l, full), T.float32(1.0), full) + if precision == 'ftz_true': + quotient = V.vdiv(silu_l, denominator, full) + else: + quotient = vdiv_0ulp(silu_l, denominator) + val = V.vmul(quotient, mul_r, full) + out_val = V.vmuls(val, weight, full) if with_weight else val + + if out_config.with_sf: + V.vstore(out_val, values_ub[col]) + elif out_config.dtype == T.bfloat16: + V.vstore(V.vcvt(out_val, 'bfloat16'), out_ub[col]) + else: + V.vstore(out_val, out_ub[col]) + + if with_certificate: + # Keep the low 19 mantissa bits; values too close to either end of this window use the exact fallback. + metric = V.vand(V.vinterpret_cast(out_val, 'uint32'), V.vbrc(T.uint32(0x7FFFF), size=vec_size)) + certificate_min = V.vmin(certificate_min, metric) + certificate_max = V.vmax(certificate_max, metric) + if count_clamp: + clamped_mask_0 = V.vcmps(val_l, T.float32(clamp_value), full, 'gt') + clamped_mask_1 = V.vcmps(val_r, T.float32(clamp_value), full, 'gt') + clamped_mask_2 = V.vcmps(val_r, T.float32(-clamp_value), full, 'lt') + # Merge-mode accumulation: add 1 on active lanes, keep the + # previous value elsewhere (PTO vadds only supports zeroing). + one = V.vbrc(T.float32(1.0), size=vec_size) + acc[0] = V.vadd(acc[0], V.vsel(clamped_mask_0, one, zeros), full) + acc[1] = V.vadd(acc[1], V.vsel(clamped_mask_1, one, zeros), full) + acc[2] = V.vadd(acc[2], V.vsel(clamped_mask_2, one, zeros), full) + if out_config.with_sf: + amax_acc[sub] = V.vabs(out_val) + if out_config.with_sf: + # 256-element tile amax: two 32-channel groups per 64-lane vector. + for sub in T.unroll(4, explicit=True): + amax = V.vcmax(amax_acc[sub], full, group=2) + V.vstore(amax, amax_ub[tile_id * 8 + sub * 2]) + if count_clamp: + for i in T.unroll(3, explicit=True): + V.vstore(V.vadd(V.vload(pacc_buf[i, 0], size=vec_size), acc[i], full), pacc_buf[i, 0]) + return certificate_min, certificate_max + + @T.macro + def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offset): + # Get inverse scale and quantize FP32 SwiGLU result to FP8. + with T.SimdVF(): + full = V.create_mask(vec_size, size=vec_size) + for group_batch in T.serial(T.ceildiv(num_groups, 64)): + group = group_batch * 64 + amax = V.vload(amax_ub[group], size=vec_size) + if out_config.round_sf: + scales_or_exponents, inverses = get_scale_and_inverse(amax) + if out_config.use_packed_ue8m0: + V.vstore(V.vcvt(scales_or_exponents, 'uint8'), sf_dense_ub[token_offset, group]) + else: + V.vstore(scales_or_exponents, sf_dense_ub[token_offset, group]) + V.vstore(inverses, inverse_ub[group]) + else: + clamped_amax = V.vmaxs(amax, T.float32(out_config.clamp_min_value), full) + quant_max = V.vbrc(T.float32(448.0), size=vec_size) + V.vstore(V.vdiv(clamped_amax, quant_max, full, precision='exact'), sf_dense_ub[token_offset, group]) + V.vstore(V.vdiv(quant_max, clamped_amax, full, precision='exact'), inverse_ub[group]) + + S.mem_bar('VST_VLD') + group_lane = V.vshrs(V.vci(T.int32(0), size=vec_size), T.int32(5), full) + for group_batch in T.unroll(T.ceildiv(num_groups, 64), explicit=True): + group_base = group_batch * 64 + vector_base = group_batch * 32 + inverses = V.vload(inverse_ub[group_base], size=vec_size) + for vector in T.serial(T.min(32, num_vecs - vector_base)): + col = (vector_base + vector) * vec_size + inverse_indices = V.vadds(group_lane, T.int32(vector * 2), full) + expanded_inverse = V.vselr(inverses, inverse_indices) + quantized = V.vcvt( + V.vmul(V.vload(values_ub[col], size=vec_size), expanded_inverse), + 'float8_e4m3fn', + rounding='R', + saturate='SAT', + ) + V.vstore(quantized, out_ub[col]) + + @T.macro + def load_weight(topk_weights, token_id, routed_scaling_factor): + weight = T.alloc_var(T.float32, init=1.0) + if with_weight: + weight = topk_weights[token_id] + if with_routed_scaling: + weight = weight * routed_scaling_factor + return weight + + @T.macro + def is_valid_token(token_id, expert_id, psum_ub): + valid = T.alloc_var(T.bool, init=token_id < num_expanded_tokens) + if has_psum: + valid = valid and (expert_id < num_experts) and (token_id < psum_ub[expert_id]) + return valid + + @T.macro + def store_sf_row(sf_dense_ub, out_sf, row, token_id): + if out_config.with_sf and not out_config.use_tma_aligned_col_major_sf: + T.copy(sf_dense_ub[row, 0:num_groups], out_sf[token_id, 0]) + + @T.macro + def load_inputs(x_gm, xl, xr, topk_w, rsf, tid): + weight = load_weight(topk_w, tid, rsf) + T.copy(x_gm[tid, 0], xl) + T.copy(x_gm[tid, hidden], xr) + return weight + + @T.macro + def finish_token(vals, out, amax, inv, sf, out_gm, out_sf_gm, off, tid, acc_buf, count_clamp): + if out_config.with_sf: + quantize_f32(vals, amax, inv, sf, out, off) + T.copy(out, out_gm[tid, 0]) + store_sf_row(sf, out_sf_gm, off, tid) + if count_clamp: + acc_buf[3] += hidden + + @T.prim_func + def swiglu_forward_kernel( + x: T.Tensor[(num_expanded_tokens, hidden * 2), in_dtype], + out: T.Tensor[(num_expanded_tokens, hidden), out_config.dtype], + out_sf: T.StridedTensor[sf_shape, (sf_stride, 1), out_config.sf_dtype], + psum_num_tokens_per_expert: T.Tensor[num_experts, T.int32], + topk_weights: T.Tensor[(num_expanded_tokens,), T.float32], + routed_scaling_factor: T.float32, + clamped_count: T.Tensor[(4,), T.int64], + ): + with T.Kernel(num_vec_cores) as core_id: + if count_clamp: + acc_ub = T.alloc_shared((4,), T.int32) + pacc_ub = T.alloc_shared((3, vec_size), T.float32) # per lane accumulator pacc_ub. use vcadd at the last moment + + for i in T.serial(4): + acc_ub[i] = 0 + with T.SimdVF(): + T.clear(pacc_ub) + else: + acc_ub = None + pacc_ub = None + + if has_psum: + psum_ub = T.alloc_shared((num_experts,), T.int32) + T.copy(psum_num_tokens_per_expert, psum_ub) + expert_id = T.alloc_var(T.int32, init=0) + if use_compact_groups: + group_psum_ub = T.alloc_shared((num_experts,), T.int32) + group_psum_ub[0] = T.ceildiv(psum_ub[0], token_group) + for expert in T.serial(1, num_experts): + expert_start = (psum_ub[expert - 1] + alignment - 1) // alignment * alignment + expert_groups = T.ceildiv(psum_ub[expert] - expert_start, token_group) + group_psum_ub[expert] = group_psum_ub[expert - 1] + expert_groups + + xl_ub = T.alloc_shared((hidden_aligned,), in_dtype) + xr_ub = T.alloc_shared((hidden_aligned,), in_dtype) + out_ub = T.alloc_shared((hidden_aligned,), out_config.dtype) + values_ub = T.alloc_shared((hidden_aligned,), T.float32) + if use_fast_certificate: + xl_fb_ub = T.alloc_shared((hidden_aligned,), in_dtype) + xr_fb_ub = T.alloc_shared((hidden_aligned,), in_dtype) + out_fb_ub = T.alloc_shared((hidden_aligned,), out_config.dtype) + values_fb_ub = T.alloc_shared((hidden_aligned,), T.float32) + if out_config.with_sf: + amax_ub = T.alloc_shared((align(num_groups, 64),), T.float32) + if use_fast_certificate: + certificate_ub = T.alloc_shared((token_group, 2), T.uint32) + inverse_ub = T.alloc_shared((align(num_groups, 64),), T.float32) + sf_dense_ub = T.alloc_shared( + (token_group, align(num_groups, 64)), + T.uint8 if out_config.use_packed_ue8m0 else T.float32, + ) + if out_config.use_tma_aligned_col_major_sf: + sf_ub = T.alloc_shared((num_groups // pack_factor, token_group), T.uint16) + transpose_indices_ub = T.alloc_shared((128,), T.uint16) + else: + amax_ub = None + inverse_ub = None + sf_dense_ub = None + T.annotate_buffer_versions({xl_ub: num_stages, xr_ub: num_stages, out_ub: num_stages}) + + if out_config.with_sf and out_config.use_tma_aligned_col_major_sf: + init_transpose_indices(transpose_indices_ub) + + for pid in T.Persistent( + [T.ceildiv(num_expanded_tokens, token_group)], + num_vec_cores, + core_id, + group_size=1, + num_stages=num_stages, + ): + token_base = T.alloc_var(T.int32, init=pid * token_group) + if use_compact_groups: + while expert_id < num_experts and pid >= group_psum_ub[expert_id]: + expert_id += 1 + expert_group_base = T.alloc_var(T.int32, init=0) + expert_token_base = T.alloc_var(T.int32, init=0) + if expert_id > 0: + expert_group_base = group_psum_ub[expert_id - 1] + expert_token_base = (psum_ub[expert_id - 1] + alignment - 1) // alignment * alignment + token_base = expert_token_base + (pid - expert_group_base) * token_group + elif has_psum: + while expert_id < num_experts and token_base >= (psum_ub[expert_id] + alignment - 1) // alignment * alignment: + expert_id += 1 + if has_psum: + valid_group = token_base < num_expanded_tokens and expert_id < num_experts and token_base < psum_ub[expert_id] + else: + valid_group = token_base < num_expanded_tokens + + if valid_group: + for token_offset in T.serial(token_group): + token_id = T.alloc_var(T.int32, init=token_base + token_offset) + valid = is_valid_token(token_id, expert_id if has_psum else 0, psum_ub if has_psum else None) + if valid: + weight = load_inputs(x, xl_ub, xr_ub, topk_weights, routed_scaling_factor, token_id) + with T.SimdVF(): + full = V.create_mask(vec_size, size=vec_size) + vdiv_precision = 'ftz_true' if use_fast_certificate else 'exact' + cert_min, cert_max = compute_swiglu( + xl_ub, xr_ub, values_ub, out_ub, weight, vdiv_precision, amax_ub, use_fast_certificate, pacc_ub, count_clamp + ) + if use_fast_certificate: + V.vstore(V.vcmin(cert_min, full), certificate_ub[token_offset, 0], group=1, stride=1) + V.vstore(V.vcmax(cert_max, full), certificate_ub[token_offset, 1], group=1, stride=1) + finish_token( + values_ub, out_ub, amax_ub, inverse_ub, sf_dense_ub, out, out_sf, token_offset, token_id, acc_ub, count_clamp + ) + + if use_fast_certificate: # fall back to exact vdiv for unsafe tokens + for fallback_offset in T.serial(token_group): + token_id = T.alloc_var(T.int32, init=token_base + fallback_offset) + valid = is_valid_token(token_id, expert_id if has_psum else 0, psum_ub if has_psum else None) + unsafe = T.alloc_var( + T.bool, + init=( + certificate_ub[fallback_offset, 0] <= certificate_error + or certificate_ub[fallback_offset, 1] >= 0x80000 - certificate_error + ), + ) + if valid and unsafe: + weight = load_inputs(x, xl_fb_ub, xr_fb_ub, topk_weights, routed_scaling_factor, token_id) + with T.SimdVF(): + # Clamp counts were already accumulated in the main path; fallback must set count_clamp to False + _, _ = compute_swiglu(xl_fb_ub, xr_fb_ub, values_fb_ub, out_fb_ub, weight, 'exact', amax_ub, False, None, False) + finish_token( + values_fb_ub, out_fb_ub, amax_ub, inverse_ub, sf_dense_ub, out, out_sf, fallback_offset, token_id, None, False + ) + + if out_config.with_sf and out_config.use_tma_aligned_col_major_sf: + if valid_group: + transpose_packed_scales(sf_dense_ub, sf_ub, transpose_indices_ub) + # PTO's DMA requires matching source/destination element widths. + # Copy through uint8 views instead of the ASC uint16 view. + output_sf = T.Tensor((sf_shape[0], sf_shape[1]), T.uint8, out_sf.data) + sf_ub_u8 = T.Tensor((num_groups // pack_factor, token_group * 2), T.uint8, sf_ub.data) + valid_token_end = psum_ub[expert_id] if has_psum else num_expanded_tokens + valid_tokens = T.alloc_var(T.int32, init=T.min(token_group, valid_token_end - token_base)) + T.copy( + sf_ub_u8[:, 0 : valid_tokens * 2], + output_sf[:, token_base * 2 : token_base * 2 + valid_tokens * 2], + ) + + if count_clamp: + with T.SimdVF(): + full = V.create_mask(vec_size, size=vec_size) + for i in T.unroll(3, explicit=True): + reduced = V.vcadd(V.vload(pacc_ub[i, 0], size=vec_size), full, reassoc=True) + V.vstore(V.vcvt(reduced, 'int32'), acc_ub[i], group=1, stride=1) + + with T.SimtVF(threads=1): + for i in T.serial(4): + T.atomic_add(clamped_count[i], T.int64(acc_ub[i])) + + return swiglu_forward_kernel From 806181363cb6bdf2a9ddace58f34519844b83b24 Mon Sep 17 00:00:00 2001 From: kuri780 Date: Mon, 19 Oct 2026 12:09:31 +0000 Subject: [PATCH 2/7] quant: clarify exact vdiv path and guard packed SF - Forward: assert packed UE8M0 requires round_sf, matching the backward kernel and the public config builder. - Make the vdiv precision dispatch explicit; raise on unknown modes instead of silently treating every non-ftz_true string as exact. - Rename the 0-ULP corrections to vdiv_exact_0ulp / sigmoid_vdiv_exact_0ulp and document that they start from the hardware quotient and intentionally do not use the slower built-in V.vdiv(precision="exact") expansion. Co-Authored-By: Claude Code --- tile_kernels/quant/swiglu_backward_vmi.py | 29 ++++++++++++++++------- tile_kernels/quant/swiglu_forward_vmi.py | 16 ++++++++++--- 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/tile_kernels/quant/swiglu_backward_vmi.py b/tile_kernels/quant/swiglu_backward_vmi.py index c9f9140e..e57fe4e7 100644 --- a/tile_kernels/quant/swiglu_backward_vmi.py +++ b/tile_kernels/quant/swiglu_backward_vmi.py @@ -75,8 +75,15 @@ def get_swiglu_backward_kernel_vmi( x_grad_fp8_shape = (num_expanded_tokens, hidden * 2) if with_quant else (1, 1) @T.macro - def vdiv_0ulp(a, b): - """Correct a / b to 0 ULP (same algorithm as the ASC source).""" + def vdiv_exact_0ulp(a, b): + """Correct a / b to 0 ULP (the ASC exact path). + + Starts from the hardware quotient and searches the two neighboring + representables, matching the ASC ``vdiv_0ulp`` algorithm. This + intentionally does not call ``V.vdiv(precision='exact')``: that built-in + expansion is a different (slower) software algorithm and would be + redundant here. + """ full = V.create_mask(vec_size, size=vec_size) z = V.vdiv(a, b, full) y = V.vneg(b) @@ -101,11 +108,15 @@ def vdiv_0ulp(a, b): return V.vsel(is_zero, z, corrected) @T.macro - def sigmoid_vdiv_0ulp(a, b): - """Correct a / b for sigmoid (same algorithm as the ASC source). - - PTO vmi.vsel only accepts vector-register operands, so the predicate - xor is materialized as int32 vectors. + def sigmoid_vdiv_exact_0ulp(a, b): + """Correct a / b for sigmoid to 0 ULP (the ASC sigmoid algorithm). + + Starts from the hardware quotient and checks the adjacent representable, + matching the ASC ``sigmoid_vdiv_0ulp`` algorithm. This intentionally + does not call ``V.vdiv(precision='exact')``; the ASC source uses this + cheaper specialized correction for the sigmoid. PTO ``vmi.vsel`` only + accepts vector-register operands, so the predicate xor is materialized + as int32 vectors. """ full = V.create_mask(vec_size, size=vec_size) one_i32 = V.vbrc(T.int32(1), size=vec_size) @@ -382,10 +393,10 @@ def swiglu_backward_kernel_vmi( y_compute_value = y_value sigmoid_denominator = V.vadds(V.vexpdif(zeros, x_compute_value, full), T.float32(1.0), full) - sigmoid = sigmoid_vdiv_0ulp(ones, sigmoid_denominator) + sigmoid = sigmoid_vdiv_exact_0ulp(ones, sigmoid_denominator) if with_weight or do_recompute: # Keep a separate 0-ULP division for the required FP32 rounding. - act_value = V.vmul(vdiv_0ulp(x_compute_value, sigmoid_denominator), y_compute_value, full) + act_value = V.vmul(vdiv_exact_0ulp(x_compute_value, sigmoid_denominator), y_compute_value, full) if with_weight: weight_contribution = V.vmul(grad, act_value, full) weight_value = V.vmuls(weight_contribution, valid_token_f32, full) if num_experts else weight_contribution diff --git a/tile_kernels/quant/swiglu_forward_vmi.py b/tile_kernels/quant/swiglu_forward_vmi.py index 2e49b9b0..4cd91c6f 100644 --- a/tile_kernels/quant/swiglu_forward_vmi.py +++ b/tile_kernels/quant/swiglu_forward_vmi.py @@ -34,6 +34,7 @@ def get_swiglu_forward_kernel_vmi( assert not with_routed_scaling or with_weight assert hidden % 128 == 0, f'Ascend swiglu requires hidden % 128 == 0, got {hidden}' assert not out_config.use_tma_aligned_col_major_sf or out_config.use_packed_ue8m0 + assert not out_config.use_packed_ue8m0 or out_config.round_sf if out_config.with_sf: assert out_config.sf_block == (1, 32), 'Ascend swiglu only supports sf_block=(1, 32)' @@ -59,8 +60,15 @@ def get_swiglu_forward_kernel_vmi( pack_factor = get_packed_ue8m0_pack_factor() @T.macro - def vdiv_0ulp(a, b): - """Correct a / b to 0 ULP (same algorithm as the ASC backward source).""" + def vdiv_exact_0ulp(a, b): + """Correct a / b to 0 ULP (the ASC exact path). + + Starts from the hardware quotient and searches the two neighboring + representables, matching the ASC ``vdiv_0ulp`` algorithm. This + intentionally does not call ``V.vdiv(precision='exact')``: that built-in + expansion is a different (slower) software algorithm and would be + redundant here. + """ full = V.create_mask(vec_size, size=vec_size) z = V.vdiv(a, b, full) y = V.vneg(b) @@ -190,8 +198,10 @@ def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, denominator = V.vadds(V.vexpdif(zeros, silu_l, full), T.float32(1.0), full) if precision == 'ftz_true': quotient = V.vdiv(silu_l, denominator, full) + elif precision == 'exact': + quotient = vdiv_exact_0ulp(silu_l, denominator) else: - quotient = vdiv_0ulp(silu_l, denominator) + raise ValueError(f"unsupported vdiv precision {precision!r}; expected 'ftz_true' or 'exact'") val = V.vmul(quotient, mul_r, full) out_val = V.vmuls(val, weight, full) if with_weight else val From 6eb553a56c50823263c99b244fcddcdb39c943b6 Mon Sep 17 00:00:00 2001 From: kuri780 Date: Tue, 20 Oct 2026 03:42:25 +0000 Subject: [PATCH 3/7] quant: align forward exact vdiv with ASC and address review - forward exact vdiv uses the 'lt' tie-break of ASC's built-in vdiv_0ulp_ftz_true; backward keeps 'le' as in swiglu_backward_asc.py - spell out NOSAT for packed UE8M0 exponent bytes and FP8 output casts, matching ASC's truncation/NaN semantics for amax=inf and non-finite values - remove dead V.alloc_var locals, unused masks, reverse scratch traversal, the redundant packed/round_sf assert, and the no-op fast-math config - shrink the quant scratch to the compact layout and use scalar vmins/vmaxs - add vmi_swiglu_plugin.py and ASC-vs-VMI coverage/adversarial scripts Co-Authored-By: Claude Code --- review_adversarial_ab.py | 146 ++++++++++++++++++++ review_coverage_ab.py | 156 ++++++++++++++++++++++ tile_kernels/quant/swiglu_backward_vmi.py | 86 ++++++------ tile_kernels/quant/swiglu_forward_vmi.py | 54 ++++---- vmi_swiglu_plugin.py | 25 ++++ 5 files changed, 404 insertions(+), 63 deletions(-) create mode 100644 review_adversarial_ab.py create mode 100644 review_coverage_ab.py create mode 100644 vmi_swiglu_plugin.py diff --git a/review_adversarial_ab.py b/review_adversarial_ab.py new file mode 100644 index 00000000..dbec926e --- /dev/null +++ b/review_adversarial_ab.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""Adversarial ASC-vs-VMI checks for the PTOAS #1616 review risks. + +Targets the cases that random data does not reach: + * forward exact division with subnormal operands (the 'le'/'lt' tie case) + * forward packed UE8M0 with amax = inf (ASC truncation vs PTODSL SAT) + * backward exact divisions with subnormal operands + +Run inside the A5 container after sourcing activate_vmi_migration.sh. +""" + +import torch +import torch_npu + +import tile_kernels +import tile_kernels.quant.swiglu_backward_asc as ba +import tile_kernels.quant.swiglu_forward_asc as fa +from tile_kernels.quant.swiglu_backward_vmi import get_swiglu_backward_kernel_vmi +from tile_kernels.quant.swiglu_forward_vmi import get_swiglu_forward_kernel_vmi + +H = 2048 +T = 1024 +of, ob = fa.get_swiglu_forward_kernel_asc, ba.get_swiglu_backward_kernel_asc + + +def cpu(t): + t = t.detach().cpu() + return t.view(torch.uint8) if t.dtype == torch.float8_e4m3fn else t + + +def bits_equal(a, b): + if isinstance(a, tuple): + return len(a) == len(b) and all(bits_equal(x, y) for x, y in zip(a, b)) + if isinstance(a, torch.Tensor): + x, y = cpu(a), cpu(b) + if x.dtype.is_floating_point: + return torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + return torch.equal(x, y) + return a == b + + +def run_f(vmi, x, **kw): + fn = tile_kernels.quant.swiglu_forward_and_per_token_cast if kw.get('fmt') == 'e4m3' else tile_kernels.quant.swiglu_forward + fa.get_swiglu_forward_kernel_asc = get_swiglu_forward_kernel_vmi if vmi else of + try: + return fn(x, **kw) + finally: + fa.get_swiglu_forward_kernel_asc = of + + +def run_b(vmi, x, g, **kw): + ba.get_swiglu_backward_kernel_asc = get_swiglu_backward_kernel_vmi if vmi else ob + try: + return tile_kernels.quant.swiglu_backward(x, g, **kw) + finally: + ba.get_swiglu_backward_kernel_asc = ob + + +def subnormal_bits(lo, hi, count, negative): + """Evenly sample float32 subnormal bit patterns.""" + step = max(1, (hi - lo + 1) // count) + idx = torch.arange(lo, hi + 1, step, dtype=torch.int64) + if idx.numel() < count: + idx = torch.cat((idx, idx[-1].repeat(count - idx.numel()))) + idx = idx[:count] + if negative: + idx = idx - (1 << 32) + return idx.to(torch.int32).view(torch.float32) + + +def check(tag, a, b): + ok = bits_equal(a, b) + print(f'{tag}:', 'OK' if ok else 'MISMATCH', flush=True) + return ok + + +def forward_subnormal(dtype, fmt): + tile_kernels.set_token_alignment(128) + x = torch.ones((T, 2 * H), dtype=dtype, device='npu') + if dtype == torch.float32: + # All negative float32 subnormals, evenly sampled. + vals = subnormal_bits(0x80000001, 0x807FFFFF, T * H, negative=True).view(T, H).to('npu') + else: + # All negative bf16 subnormals, tiled. + bits = torch.arange(0x8001, 0x8080, dtype=torch.int64).to(torch.int16) + vals = bits.view(torch.bfloat16).to('npu') + vals = vals.repeat((T * H + vals.numel() - 1) // vals.numel())[: T * H].view(T, H) + x[:, :H] = vals + x[:, H:] = 1.0 + return check(f'fwd_{fmt}_subnormal', run_f(False, x, fmt=fmt), run_f(True, x, fmt=fmt)) + + +def forward_nonfinite(): + tile_kernels.set_token_alignment(256) + x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') + rows = [ + (3e38, 3e38), # +inf result -> amax = inf -> exponent 256 + (3e38, -3e38), # -inf result + (float('nan'), 1.0), # NaN result + (-100.0, 100.0), # -0 result + (float('inf'), 1.0), # inf input + (float('-inf'), 1.0), # -inf input + ] + for i, (left, right) in enumerate(rows): + x[i, :H] = left + x[i, H:] = right + kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) + return check('fwd_e4m3_nonfinite', run_f(False, x, **kw), run_f(True, x, **kw)) + + +def forward_sf_only(): + # Focused check for the vcmax(group=...) scale mapping: compare only SF. + tile_kernels.set_token_alignment(256) + torch.manual_seed(20261020) + x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') + kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) + a = run_f(False, x, **kw) + b = run_f(True, x, **kw) + return check('fwd_e4m3_sf_only', a[1], b[1]) + + +def backward_subnormal(dtype, fmt): + tile_kernels.set_token_alignment(128) + x = torch.ones((T, 2 * H), dtype=dtype, device='npu') + if dtype == torch.float32: + vals = subnormal_bits(0x80000001, 0x807FFFFF, T * H, negative=True).view(T, H).to('npu') + else: + bits = torch.arange(0x8001, 0x8080, dtype=torch.int64).to(torch.int16) + vals = bits.view(torch.bfloat16).to('npu') + vals = vals.repeat((T * H + vals.numel() - 1) // vals.numel())[: T * H].view(T, H) + x[:, :H] = vals + x[:, H:] = 1.0 + g = torch.randn((T, H), dtype=dtype, device='npu') + return check(f'bwd_{fmt}_subnormal', run_b(False, x, g, fmt=fmt), run_b(True, x, g, fmt=fmt)) + + +if __name__ == '__main__': + results = [ + forward_subnormal(torch.float32, 'fp32'), + forward_subnormal(torch.bfloat16, 'bf16'), + forward_nonfinite(), + forward_sf_only(), + backward_subnormal(torch.float32, 'fp32'), + backward_subnormal(torch.bfloat16, 'bf16'), + ] + print('ALL_OK' if all(results) else 'HAS_MISMATCH') diff --git a/review_coverage_ab.py b/review_coverage_ab.py new file mode 100644 index 00000000..ad048166 --- /dev/null +++ b/review_coverage_ab.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python3 +"""Extra ASC-vs-VMI coverage requested in review.""" + +import torch +import torch_npu +import tile_kernels +import tile_kernels.quant.swiglu_forward_asc as fa +import tile_kernels.quant.swiglu_backward_asc as ba +from tile_kernels.quant.swiglu_forward_vmi import get_swiglu_forward_kernel_vmi +from tile_kernels.quant.swiglu_backward_vmi import get_swiglu_backward_kernel_vmi + +of, ob = fa.get_swiglu_forward_kernel_asc, ba.get_swiglu_backward_kernel_asc +H = 2048 +PSUM = torch.tensor([600, 1600], dtype=torch.int32, device='npu') + + +def cpu(t): + t = t.detach().cpu() + return t.view(torch.uint8) if t.dtype == torch.float8_e4m3fn else t + + +def cmp_tree(a, b, rows): + if isinstance(a, tuple): + return all(cmp_tree(x, y, rows) for x, y in zip(a, b)) + if isinstance(a, torch.Tensor): + return torch.equal(cpu(a)[rows], cpu(b)[rows]) + return a == b + + +def layout(use_psum): + if not use_psum: + return None, torch.ones(1024, dtype=torch.bool) + mask = torch.zeros(1792, dtype=torch.bool) + mask[0:600] = True + mask[768:1600] = True + return PSUM, mask + + +def run_f(vmi, x, fn, **kw): + fa.get_swiglu_forward_kernel_asc = get_swiglu_forward_kernel_vmi if vmi else of + try: + return fn(x=x, **kw) + finally: + fa.get_swiglu_forward_kernel_asc = of + + +def run_b(vmi, x, g, fn, **kw): + ba.get_swiglu_backward_kernel_asc = get_swiglu_backward_kernel_vmi if vmi else ob + try: + return fn(x, g, **kw) + finally: + ba.get_swiglu_backward_kernel_asc = ob + + +def sf_kwargs(fmt, packed, round_sf, tma): + if fmt != 'e4m3': + return {} + return dict( + fmt=fmt, + num_per_channels=32, + use_tma_aligned_col_major_sf=tma, + round_sf=round_sf, + use_packed_ue8m0=packed, + ) + + +def fwd(tag, dtype, use_psum, fmt, weights=False, rsf=None, clamp=None, packed=False, round_sf=False, tma=False): + tile_kernels.set_token_alignment(256 if use_psum else 128) + psum, mask = layout(use_psum) + torch.manual_seed(101) + T = 1792 if use_psum else 1024 + x = torch.randn((T, 2 * H), device='npu', dtype=dtype) + w = torch.rand((T,), device='npu', dtype=torch.float32) + if use_psum: + m = mask.to('npu') + x[~m] = 0 + w[~m] = 0 + kw = sf_kwargs(fmt, packed, round_sf, tma) + if fmt != 'e4m3': + kw['fmt'] = fmt + if weights: + kw.update(topk_weights=w, routed_scaling_factor=rsf) + if clamp is not None: + kw.update(clamp_value=clamp) + if fmt == 'e4m3': + fn = tile_kernels.quant.swiglu_forward_and_per_token_cast + else: + fn = tile_kernels.quant.swiglu_forward + a = run_f(False, x, fn, psum_num_tokens_per_expert=psum, **kw) + b = run_f(True, x, fn, psum_num_tokens_per_expert=psum, **kw) + torch.npu.synchronize() + print(f'{tag}:', 'OK' if cmp_tree(a, b, mask) else 'MISMATCH') + + +def bwd(tag, fmt, use_psum, dtype, weights=False, rsf=None, clamp=None, recompute=False, packed=False, round_sf=False, tma=False): + tile_kernels.set_token_alignment(256 if (use_psum or fmt == 'e4m3') else 128) + psum, mask = layout(use_psum) + torch.manual_seed(102) + T = 1792 if use_psum else 1024 + g_dtype = torch.bfloat16 if fmt == 'e4m3' else dtype + x_raw = torch.randn((T, 2 * H), device='npu', dtype=torch.bfloat16) + g = torch.randn((T, H), device='npu', dtype=g_dtype) + w = torch.rand((T,), device='npu', dtype=torch.float32) + if use_psum: + m = mask.to('npu') + x_raw[~m] = 0 + g[~m] = 0 + w[~m] = 0 + kw = sf_kwargs(fmt, packed, round_sf, tma) + if fmt != 'e4m3': + kw['fmt'] = fmt + if weights: + kw.update(topk_weights=w, routed_scaling_factor=rsf) + if clamp is not None: + kw.update(clamp_value=clamp) + kw.update(do_recompute=recompute) + if fmt == 'e4m3': + x = tile_kernels.quant.per_token_cast(x_raw, 'e4m3', num_per_channels=32) + fn = tile_kernels.quant.swiglu_backward_and_per_token_cast + else: + x = x_raw.to(dtype) + fn = tile_kernels.quant.swiglu_backward + a = run_b(False, x, g, fn, psum_num_tokens_per_expert=psum, **kw) + b = run_b(True, x, g, fn, psum_num_tokens_per_expert=psum, **kw) + torch.npu.synchronize() + print(f'{tag}:', 'OK' if cmp_tree(a, b, mask) else 'MISMATCH') + + +fwd('fwd_fp32_in_bf16_out_no_psum', torch.float32, False, 'bf16') +fwd('fwd_fp32_in_fp32_out_no_psum_clamp', torch.float32, False, 'fp32', clamp=0.5) +fwd('fwd_fp32_in_bf16_out_weights_rsf', torch.float32, False, 'bf16', weights=True, rsf=2.5) +fwd('fwd_fp32_in_e4m3_packed_no_tma', torch.float32, False, 'e4m3', packed=True, round_sf=True, tma=False) +fwd('fwd_e4m3_packed_no_tma_psum_weights', torch.bfloat16, True, 'e4m3', weights=True, rsf=2.5, packed=True, round_sf=True, tma=False) +fwd('fwd_e4m3_dense_exact_no_psum', torch.bfloat16, False, 'e4m3', round_sf=False, tma=False) +fwd('fwd_e4m3_tma_weights_clamp', torch.bfloat16, True, 'e4m3', weights=True, rsf=2.5, clamp=0.5, packed=True, round_sf=True, tma=True) + +bwd('bwd_bf16_no_psum_no_weights', 'bf16', False, torch.bfloat16) +bwd('bwd_bf16_psum_weights_no_rsf', 'bf16', True, torch.bfloat16, weights=True, rsf=None) +bwd('bwd_bf16_psum_weights_recompute_clamp', 'bf16', True, torch.bfloat16, weights=True, rsf=2.5, clamp=0.5, recompute=True) +bwd('bwd_fp32_no_psum_weights_recompute', 'fp32', False, torch.float32, weights=True, rsf=2.5, recompute=True) +bwd( + 'bwd_e4m3_no_psum_tma_weights_recompute', + 'e4m3', + False, + torch.bfloat16, + weights=True, + rsf=2.5, + recompute=True, + packed=True, + round_sf=True, + tma=True, +) +bwd('bwd_e4m3_psum_packed_no_tma_weights', 'e4m3', True, torch.bfloat16, weights=True, rsf=2.5, packed=True, round_sf=True, tma=False) +bwd('bwd_e4m3_psum_dense_exact_recompute', 'e4m3', True, torch.bfloat16, weights=True, rsf=2.5, recompute=True, round_sf=False, tma=False) +bwd('bwd_e4m3_psum_tma_no_weights_recompute', 'e4m3', True, torch.bfloat16, recompute=True, packed=True, round_sf=True, tma=True) +print('REVIEW_COVERAGE_DONE') diff --git a/tile_kernels/quant/swiglu_backward_vmi.py b/tile_kernels/quant/swiglu_backward_vmi.py index e57fe4e7..2558a24a 100644 --- a/tile_kernels/quant/swiglu_backward_vmi.py +++ b/tile_kernels/quant/swiglu_backward_vmi.py @@ -2,17 +2,12 @@ from tilelang.ascend import language as T from tilelang.ascend.language import simd as S from tilelang.ascend.language import vmi as V -from tilelang.transform import PassConfigKey from tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape from tile_kernels.utils import align -# Match the ASC source of truth: bare S.vdiv is precise unless fast math is -# explicitly enabled. PTO VMI's hardware vdiv is the explicit 'ftz_true' path. -PTO_FAST_MATH = {PassConfigKey.TL_ENABLE_FAST_MATH: False} - -@tilelang.jit(target='pto', pass_configs=PTO_FAST_MATH) +@tilelang.jit(target='pto') def get_swiglu_backward_kernel_vmi( hidden: int, num_experts: int, @@ -65,7 +60,7 @@ def get_swiglu_backward_kernel_vmi( num_scale_vectors = (num_block_groups + vec_size - 1) // vec_size if with_quant else 0 sf_storage_dtype = T.uint16 if out_config.use_packed_ue8m0 else T.float32 stored_groups_per_vector = vec_size * pack_factor // block_m - # Pad transpose UB to avoid unsupported partial-vector PAT_VL masks. + # Pad the transpose UB so every VMI gather/store uses a full vector width. aligned_num_stored_block_groups = align(num_stored_block_groups, stored_groups_per_vector) num_expanded_tokens = T.dynamic('num_expanded_tokens') @@ -76,25 +71,25 @@ def get_swiglu_backward_kernel_vmi( @T.macro def vdiv_exact_0ulp(a, b): - """Correct a / b to 0 ULP (the ASC exact path). + """Correct a / b to 0 ULP with ASC backward ``vdiv_0ulp`` semantics. + + ASC backward uses the ``le`` tie-break, unlike both the built-in + ``V.vdiv(precision='exact')`` expansion and the forward kernel, which + use ``lt``. Keep ``le`` for bitwise parity with + swiglu_backward_asc.py; do not unify the two helpers. - Starts from the hardware quotient and searches the two neighboring - representables, matching the ASC ``vdiv_0ulp`` algorithm. This - intentionally does not call ``V.vdiv(precision='exact')``: that built-in - expansion is a different (slower) software algorithm and would be - redundant here. + The built-in expansion and this version perform the same neighbor + search; the omitted denominator prescaling is a power of two and does + not change residual ordering in this operand domain. """ full = V.create_mask(vec_size, size=vec_size) z = V.vdiv(a, b, full) y = V.vneg(b) - r = V.alloc_var(T.float32, size=vec_size) r = V.vmula(a, z, y, full) z_bits = V.vinterpret_cast(z, 'int32') z_pre = V.vinterpret_cast(V.vadds(z_bits, T.int32(-1), full), 'float32') z_next = V.vinterpret_cast(V.vadds(z_bits, T.int32(1), full), 'float32') - r_pre = V.alloc_var(T.float32, size=vec_size) r_pre = V.vmula(a, z_pre, y, full) - r_next = V.alloc_var(T.float32, size=vec_size) r_next = V.vmula(a, z_next, y, full) r = V.vabs(r) r_pre = V.vabs(r_pre) @@ -124,7 +119,6 @@ def sigmoid_vdiv_exact_0ulp(a, b): z = V.vdiv(a, b, full) neg_b = V.vneg(b) - residual = V.alloc_var(T.float32, size=vec_size) residual = V.vmula(a, z, neg_b, full) residual_bits = V.vinterpret_cast(residual, 'int32') @@ -133,7 +127,6 @@ def sigmoid_vdiv_exact_0ulp(a, b): delta = V.vor(sign_mask, one_i32, full) adjacent = V.vinterpret_cast(V.vadd(V.vinterpret_cast(z, 'int32'), delta, full), 'float32') - adjacent_residual = V.alloc_var(T.float32, size=vec_size) adjacent_residual = V.vmula(a, adjacent, neg_b, full) residual_sum = V.vadd(residual, adjacent_residual, full) take_bits_plus = V.vcmps(residual_sum, T.float32(0.0), full, 'gt') @@ -232,9 +225,11 @@ def swiglu_backward_kernel_vmi( x_grad_fp8_ub = T.alloc_shared((block_k,), T.float8_e4m3fn) y_grad_fp8_ub = T.alloc_shared((block_k,), T.float8_e4m3fn) # Reused for compact amax and inverse scales (one value per 32 channels). - x_quant_scratch_ub = T.alloc_shared((sf_block_stride * pack_factor * 4,), T.float32) - y_quant_scratch_ub = T.alloc_shared((sf_block_stride * pack_factor * 4,), T.float32) - # Packed SF is written as bytes and transposed through a uint16 view. + # VMI overwrites in place, so no fourfold expansion is needed. + x_quant_scratch_ub = T.alloc_shared((sf_block_stride * pack_factor,), T.float32) + y_quant_scratch_ub = T.alloc_shared((sf_block_stride * pack_factor,), T.float32) + # Packed SF is stored as uint16 words: quantization writes bytes + # through sf_dense_u8, and the transpose reads the same buffer as uint16. sf_dense_ub = T.alloc_shared((block_m, sf_row_stride), sf_storage_dtype) if out_config.use_packed_ue8m0: sf_dense_u8 = T.Tensor((block_m, sf_row_stride * 2), T.uint8, sf_dense_ub.data) @@ -348,9 +343,6 @@ def swiglu_backward_kernel_vmi( full = V.create_mask(vec_size, size=vec_size) zeros = V.vbrc(T.float32(0.0), size=vec_size) ones = V.vbrc(T.float32(1.0), size=vec_size) - if with_quant: - high_half_mask = V.vcmps(V.vci(T.int32(0), size=vec_size), T.int32(31), full, 'gt') - low_half_mask = V.create_mask(32, size=vec_size) if with_weight: weight_acc = V.alloc_var(T.float32, size=vec_size) # hidden_block is a runtime TIR loop variable; use an @@ -360,6 +352,11 @@ def swiglu_backward_kernel_vmi( else: weight_acc = V.vload(weight_acc_ub[row, 0], size=vec_size) + # VMI exposes separate widen/narrow conversions. PTO + # codegen already supports fused T.simd dist='UNPK/PK' + # forms, but #1616 forbids using ASC physical layouts + # here; a VMI widen/narrow load/store is the upstream + # feature request described in the PR. for vector in T.serial(block_k // vec_size): col = vector * vec_size if with_quant: @@ -383,11 +380,8 @@ def swiglu_backward_kernel_vmi( if use_clamp: is_clamped_x = V.vcmps(x_value, T.float32(clamp_value), full, 'gt') is_clamped_y = V.vcmps(V.vabs(y_value), T.float32(clamp_value), full, 'gt') - x_compute_value = V.vmin(x_value, V.vbrc(T.float32(clamp_value), size=vec_size)) - y_compute_value = V.vmin( - V.vmax(y_value, V.vbrc(T.float32(-clamp_value), size=vec_size)), - V.vbrc(T.float32(clamp_value), size=vec_size), - ) + x_compute_value = V.vmins(x_value, T.float32(clamp_value), full) + y_compute_value = V.vmins(V.vmaxs(y_value, T.float32(-clamp_value), full), T.float32(clamp_value), full) else: x_compute_value = x_value y_compute_value = y_value @@ -430,6 +424,10 @@ def swiglu_backward_kernel_vmi( if with_weight: if hidden_block + 1 == num_hidden_blocks: + # ASC S.vcadd is a pairwise reduction. + # PTODSL only accepts reassoc=True here, so + # weight_grad can differ in the last mantissa + # bits; it is validated at the repo's 1e-4. one_mask = V.create_mask(1, size=1) reduced_weight_sum = V.vcadd(weight_acc, full, reassoc=True) weight_sum = ( @@ -445,9 +443,9 @@ def swiglu_backward_kernel_vmi( S.mem_bar('VST_VLD') if num_experts: valid_token_vec = V.vbrc(valid_token_f32, size=vec_size) - # Preserve unread amax while overwriting the scratch with inverses. - for scale_vector_index in range(num_scale_vectors): - scale_vector = num_scale_vectors - 1 - scale_vector_index + # Each iteration reads its amax before overwriting + # the same slot with the inverse, so order is free. + for scale_vector in range(num_scale_vectors): scale_offset = scale_vector * vec_size x_scale_value, x_inv = get_scale_and_inverse(V.vload(x_quant_scratch_ub[scale_offset], size=vec_size)) y_scale_value, y_inv = get_scale_and_inverse(V.vload(y_quant_scratch_ub[scale_offset], size=vec_size)) @@ -456,11 +454,19 @@ def swiglu_backward_kernel_vmi( x_scale_out = V.vmul(x_scale_value, valid_token_vec, full) if num_experts else x_scale_value y_scale_out = V.vmul(y_scale_value, valid_token_vec, full) if num_experts else y_scale_value if out_config.use_packed_ue8m0: - x_exp = V.vshr(V.vinterpret_cast(x_scale_out, 'uint32'), V.vbrc(T.uint32(23), size=vec_size)) - y_exp = V.vshr(V.vinterpret_cast(y_scale_out, 'uint32'), V.vbrc(T.uint32(23), size=vec_size)) - V.vstore(V.vcvt(x_exp, 'uint8'), sf_dense_u8[row, sf_block_offset * 2 + scale_offset]) + x_exp = V.vshrs(V.vinterpret_cast(x_scale_out, 'uint32'), T.uint32(23), full) + y_exp = V.vshrs(V.vinterpret_cast(y_scale_out, 'uint32'), T.uint32(23), full) + # ASC reinterprets the shifted exponent bits + # as uint8 (truncation). Spell out NOSAT + # because PTODSL otherwise defaults to SAT for + # this narrowing cast, which differs for + # exponents >= 256 (amax = inf). + V.vstore( + V.vcvt(x_exp, 'uint8', saturate='NOSAT'), + sf_dense_u8[row, sf_block_offset * 2 + scale_offset], + ) V.vstore( - V.vcvt(y_exp, 'uint8'), + V.vcvt(y_exp, 'uint8', saturate='NOSAT'), sf_dense_u8[row, sf_half_stride * 2 + sf_block_offset * 2 + scale_offset], ) else: @@ -482,8 +488,12 @@ def swiglu_backward_kernel_vmi( V.vstore(V.vcvt(y_grad_quant_value, 'bfloat16'), y_grad_ub[col]) x_scaled = V.vsel(valid_token_mask, x_scaled_value, zeros) if num_experts else x_scaled_value y_scaled = V.vsel(valid_token_mask, y_scaled_value, zeros) if num_experts else y_scaled_value - V.vstore(V.vcvt(x_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), x_grad_fp8_ub[col]) - V.vstore(V.vcvt(y_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), y_grad_fp8_ub[col]) + # ASC's FP8 conversion keeps non-finite values as + # NaN; PTODSL's SAT would clamp them. NOSAT matches + # ASC for non-finite inputs and for finite results + # at the 448 boundary. + V.vstore(V.vcvt(x_scaled, 'float8_e4m3fn', rounding='R', saturate='NOSAT'), x_grad_fp8_ub[col]) + V.vstore(V.vcvt(y_scaled, 'float8_e4m3fn', rounding='R', saturate='NOSAT'), y_grad_fp8_ub[col]) T.copy(x_grad_ub, x_grad[token_id, col_offset : col_offset + block_k]) T.copy(y_grad_ub, x_grad[token_id, hidden + col_offset : hidden + col_offset + block_k]) diff --git a/tile_kernels/quant/swiglu_forward_vmi.py b/tile_kernels/quant/swiglu_forward_vmi.py index 4cd91c6f..b6d8b418 100644 --- a/tile_kernels/quant/swiglu_forward_vmi.py +++ b/tile_kernels/quant/swiglu_forward_vmi.py @@ -4,17 +4,12 @@ from tilelang.ascend import language as T from tilelang.ascend.language import simd as S from tilelang.ascend.language import vmi as V -from tilelang.transform import PassConfigKey from tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape from tile_kernels.utils import align -# Match the ASC source of truth: bare S.vdiv is precise unless fast math is -# explicitly enabled. PTO VMI's hardware vdiv is the explicit 'ftz_true' path. -PTO_FAST_MATH = {PassConfigKey.TL_ENABLE_FAST_MATH: False} - -@tilelang.jit(target='pto', pass_configs=PTO_FAST_MATH) +@tilelang.jit(target='pto') def get_swiglu_forward_kernel_vmi( hidden: int, num_experts: int, @@ -34,7 +29,6 @@ def get_swiglu_forward_kernel_vmi( assert not with_routed_scaling or with_weight assert hidden % 128 == 0, f'Ascend swiglu requires hidden % 128 == 0, got {hidden}' assert not out_config.use_tma_aligned_col_major_sf or out_config.use_packed_ue8m0 - assert not out_config.use_packed_ue8m0 or out_config.round_sf if out_config.with_sf: assert out_config.sf_block == (1, 32), 'Ascend swiglu only supports sf_block=(1, 32)' @@ -61,30 +55,32 @@ def get_swiglu_forward_kernel_vmi( @T.macro def vdiv_exact_0ulp(a, b): - """Correct a / b to 0 ULP (the ASC exact path). - - Starts from the hardware quotient and searches the two neighboring - representables, matching the ASC ``vdiv_0ulp`` algorithm. This - intentionally does not call ``V.vdiv(precision='exact')``: that built-in - expansion is a different (slower) software algorithm and would be - redundant here. + """Correct a / b to 0 ULP with ASC forward exact-division semantics. + + This is the same neighbor search as CANN ``vdiv_0ulp_ftz_true`` (the + expansion behind ``V.vdiv(precision='exact')``) without its denominator + prescaling. In this kernel b = 1 + exp(-silu_l) >= 1, so the prescale is + a power of two <= 1 and preserves residual ordering, while + |z * b| ~= |a| cannot overflow the FMA. The built-in inf/nan fallback is + also unnecessary: non-finite z makes every residual comparison false and + is returned unchanged, while z == 0 is handled by the signed-zero check. + + The ``lt`` tie-break matches the built-in exact path used by ASC + forward. ASC backward intentionally uses ``le``; do not unify them. """ full = V.create_mask(vec_size, size=vec_size) z = V.vdiv(a, b, full) y = V.vneg(b) - r = V.alloc_var(T.float32, size=vec_size) r = V.vmula(a, z, y, full) z_bits = V.vinterpret_cast(z, 'int32') z_pre = V.vinterpret_cast(V.vadds(z_bits, T.int32(-1), full), 'float32') z_next = V.vinterpret_cast(V.vadds(z_bits, T.int32(1), full), 'float32') - r_pre = V.alloc_var(T.float32, size=vec_size) r_pre = V.vmula(a, z_pre, y, full) - r_next = V.alloc_var(T.float32, size=vec_size) r_next = V.vmula(a, z_next, y, full) r = V.vabs(r) r_pre = V.vabs(r_pre) r_next = V.vabs(r_next) - take_pre = V.vcmp(r_pre, r, full, 'le') + take_pre = V.vcmp(r_pre, r, full, 'lt') best_r = V.vsel(take_pre, r_pre, r) best_z = V.vsel(take_pre, z_pre, z) take_next = V.vcmp(r_next, best_r, full, 'lt') @@ -163,6 +159,8 @@ def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, for i in T.unroll(3, explicit=True): acc[i] = zeros if with_certificate: + # Loop-carried across tile_id: alloc_var is what lets the eager + # builder rebind these accumulators inside the loop. certificate_min = V.alloc_var(T.uint32, size=vec_size) certificate_max = V.alloc_var(T.uint32, size=vec_size) certificate_min = V.vbrc(T.uint32(0xFFFFFFFF), size=vec_size) @@ -187,11 +185,8 @@ def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, val_l = V.vload(xl_ub[col], size=vec_size) val_r = V.vload(xr_ub[col], size=vec_size) if has_clamp: - silu_l = V.vmin(val_l, V.vbrc(T.float32(clamp_value), size=vec_size)) - mul_r = V.vmax( - V.vmin(val_r, V.vbrc(T.float32(clamp_value), size=vec_size)), - V.vbrc(T.float32(-clamp_value), size=vec_size), - ) + silu_l = V.vmins(val_l, T.float32(clamp_value), full) + mul_r = V.vmaxs(V.vmins(val_r, T.float32(clamp_value), full), T.float32(-clamp_value), full) else: silu_l = val_l mul_r = val_r @@ -250,7 +245,11 @@ def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offs if out_config.round_sf: scales_or_exponents, inverses = get_scale_and_inverse(amax) if out_config.use_packed_ue8m0: - V.vstore(V.vcvt(scales_or_exponents, 'uint8'), sf_dense_ub[token_offset, group]) + # ASC reinterprets the shifted exponent bits as uint8 + # (truncation). Spell out NOSAT because PTODSL would + # otherwise default to SAT for this narrowing cast, + # which differs for exponents >= 256 (amax = inf). + V.vstore(V.vcvt(scales_or_exponents, 'uint8', saturate='NOSAT'), sf_dense_ub[token_offset, group]) else: V.vstore(scales_or_exponents, sf_dense_ub[token_offset, group]) V.vstore(inverses, inverse_ub[group]) @@ -270,11 +269,14 @@ def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offs col = (vector_base + vector) * vec_size inverse_indices = V.vadds(group_lane, T.int32(vector * 2), full) expanded_inverse = V.vselr(inverses, inverse_indices) + # ASC's FP8 conversion keeps non-finite values as NaN; + # PTODSL's SAT would clamp them. NOSAT matches ASC for both + # non-finite values and finite results at the 448 boundary. quantized = V.vcvt( V.vmul(V.vload(values_ub[col], size=vec_size), expanded_inverse), 'float8_e4m3fn', rounding='R', - saturate='SAT', + saturate='NOSAT', ) V.vstore(quantized, out_ub[col]) @@ -462,6 +464,8 @@ def swiglu_forward_kernel( with T.SimdVF(): full = V.create_mask(vec_size, size=vec_size) for i in T.unroll(3, explicit=True): + # These are exact integer counts (< 2^24), so the + # reassociation allowed by PTODSL cannot change them. reduced = V.vcadd(V.vload(pacc_ub[i, 0], size=vec_size), full, reassoc=True) V.vstore(V.vcvt(reduced, 'int32'), acc_ub[i], group=1, stride=1) diff --git a/vmi_swiglu_plugin.py b/vmi_swiglu_plugin.py new file mode 100644 index 00000000..0786e269 --- /dev/null +++ b/vmi_swiglu_plugin.py @@ -0,0 +1,25 @@ +"""Run the existing SwiGLU tests against the VMI kernels. + +Usage (from the repository root): + TK_USE_SWIGLU_VMI=1 python -m pytest tests/quant/test_swiglu_forward.py \ + tests/quant/test_swiglu_backward.py -q -p vmi_swiglu_plugin + +The dispatchers import the ASC getter lazily inside the public API, so +replacing the module attribute before the tests run is enough to exercise the +VMI implementation without modifying the repository. +""" + +import os + + +def pytest_configure(config): + if os.environ.get('TK_USE_SWIGLU_VMI') != '1': + return + + import tile_kernels.quant.swiglu_forward_asc as forward_asc + import tile_kernels.quant.swiglu_backward_asc as backward_asc + from tile_kernels.quant.swiglu_forward_vmi import get_swiglu_forward_kernel_vmi + from tile_kernels.quant.swiglu_backward_vmi import get_swiglu_backward_kernel_vmi + + forward_asc.get_swiglu_forward_kernel_asc = get_swiglu_forward_kernel_vmi + backward_asc.get_swiglu_backward_kernel_asc = get_swiglu_backward_kernel_vmi From 7d6330a86614f043684100d009d18b0ad2685b04 Mon Sep 17 00:00:00 2001 From: kuri780 Date: Tue, 20 Oct 2026 05:27:28 +0000 Subject: [PATCH 4/7] quant: report saturation mismatch and re-measure benchmarks - keep FP8 output casts on saturate='SAT' (faithful ASC translation); the minimal repro shows ASC RS_ENABLE currently behaves like NOSAT for f32->e4m3/bf16 while VMI SAT really saturates (toolchain bug, issue draft) - add repro_fp8_saturate.py and extend adversarial coverage with backward e4m3 non-finite cases; document that the subnormal sweep does not cover the exact-midpoint tie branch - 5-run median benchmarks: forward geomean 1.023 (22/24 >= 0.90), backward 0.904 (8/24 >= 0.90); a concurrent MASBench container on the same NPU is the dominant noise source Co-Authored-By: Claude Code --- repro_fp8_saturate.py | 98 +++++++++++++++++++++++ review_adversarial_ab.py | 85 ++++++++++++++++++-- tile_kernels/quant/swiglu_backward_vmi.py | 17 ++-- tile_kernels/quant/swiglu_forward_vmi.py | 16 +++- 4 files changed, 201 insertions(+), 15 deletions(-) create mode 100644 repro_fp8_saturate.py diff --git a/repro_fp8_saturate.py b/repro_fp8_saturate.py new file mode 100644 index 00000000..c958bd78 --- /dev/null +++ b/repro_fp8_saturate.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Minimal f32 -> e4m3 / bf16 saturation repro for ASC S.vcvt vs VMI V.vcvt. + +Only one conversion is performed per kernel; no SwiGLU scale or mask is involved. +Run inside the A5 container after sourcing activate_vmi_migration.sh. +""" + +import torch +import torch_npu + +import tilelang +import tilelang.ascend.language as T + +N = 64 +VALUES = [ + float('inf'), + float('-inf'), + float('nan'), + -0.0, + 448.0, + 448.00003, + 464.0, + 480.0, + 1e30, +] +VARIANTS = [ + ('asc_sat', False, True), + ('asc_nosat', False, False), + ('vmi_sat', True, True), + ('vmi_nosat', True, False), +] +TARGETS = [ + ('e4m3', 'float8_e4m3fn', T.float8_e4m3fn), + ('bf16', 'bfloat16', T.bfloat16), +] + + +def make_program(vmi: bool, sat: bool, target: str, target_tir): + @T.prim_func + def main(x: T.Buffer((N,), 'float32'), y: T.Buffer((N,), target)): + with T.Kernel(1) as _: + x_ub = T.alloc_shared((N,), 'float32') + y_ub = T.alloc_shared((N,), target) + T.copy(x, x_ub) + with T.SimdVF(): + full = T.vmi.create_mask(N, size=N) if vmi else T.simd.pset(32) + if vmi: + value = T.vmi.vload(x_ub[0], size=N) + converted = T.vmi.vcvt( + value, + target, + rounding='R', + saturate='SAT' if sat else 'NOSAT', + ) + T.vmi.vstore(converted, y_ub[0], full) + else: + value = T.simd.vld(x_ub[0]) + converted = T.simd.vcvt(value, target_tir, sat=sat) + dist = 'PK4_B32' if target == 'float8_e4m3fn' else 'PK_B32' + T.simd.vsts(y_ub[0], converted, full, dist=dist) + T.copy(y_ub, y) + + return main + + +def main(): + x_cpu = torch.tensor([VALUES[i % len(VALUES)] for i in range(N)], dtype=torch.float32) + x = x_cpu.to('npu') + for target, target_param, target_tir in TARGETS: + results = {} + for name, vmi, sat in VARIANTS: + program = make_program(vmi, sat, target_param, target_tir) + kernel = tilelang.compile(program, target='pto' if vmi else 'asc', out_idx=-1) + y = kernel(x) + torch.npu.synchronize() + if target_param == 'bfloat16': + bits = (y.cpu().view(torch.int16).to(torch.int32) & 0xFFFF).tolist() + else: + bits = y.cpu().view(torch.uint8).to(torch.int32).tolist() + results[name] = bits[: len(VALUES)] + source = kernel.get_kernel_source() + path = f'/tmp/repro_{target}_saturate_{name}.txt' + with open(path, 'w') as f: + f.write(source) + print(f'--- {target} {name} source: {path}') + for line in source.splitlines(): + if any(tok in line for tok in ('vcvt', 'saturate', 'RS_ENABLE', 'RS_DISABLE')): + print(' ', line.strip()) + + print(f'\n=== {target} bit patterns (hex) ===') + print('value ' + ' '.join(f'{name:>10}' for name, _, _ in VARIANTS)) + for i, value in enumerate(VALUES): + print(f'{value!r:<22}' + ' '.join(f'{results[name][i]:>#10x}' for name, _, _ in VARIANTS)) + print() + + +if __name__ == '__main__': + main() diff --git a/review_adversarial_ab.py b/review_adversarial_ab.py index dbec926e..bb5d327d 100644 --- a/review_adversarial_ab.py +++ b/review_adversarial_ab.py @@ -2,9 +2,12 @@ """Adversarial ASC-vs-VMI checks for the PTOAS #1616 review risks. Targets the cases that random data does not reach: - * forward exact division with subnormal operands (the 'le'/'lt' tie case) - * forward packed UE8M0 with amax = inf (ASC truncation vs PTODSL SAT) - * backward exact divisions with subnormal operands + * forward/backward exact divisions with subnormal operands (finite-domain + check; the subnormal sweep does NOT exercise the 'le'/'lt' tie branch: + b = 2 there, a/2 is exact, and r = 0) + * forward packed UE8M0 with amax = inf and non-finite/overflow outputs + (documents the known ASC RS_ENABLE vs VMI SAT toolchain difference) + * backward e4m3 non-finite inputs with packed/TMA and packed non-TMA SF Run inside the A5 container after sourcing activate_vmi_migration.sh. """ @@ -49,9 +52,10 @@ def run_f(vmi, x, **kw): def run_b(vmi, x, g, **kw): + fn = tile_kernels.quant.swiglu_backward_and_per_token_cast if kw.get('fmt') == 'e4m3' else tile_kernels.quant.swiglu_backward ba.get_swiglu_backward_kernel_asc = get_swiglu_backward_kernel_vmi if vmi else ob try: - return tile_kernels.quant.swiglu_backward(x, g, **kw) + return fn(x, g, **kw) finally: ba.get_swiglu_backward_kernel_asc = ob @@ -75,6 +79,8 @@ def check(tag, a, b): def forward_subnormal(dtype, fmt): + # For tiny |x|, b = 1 + exp(-x) = 2 exactly and a/2 is exact, so r = 0. + # This sweep checks finite subnormal handling, not the le/lt tie branch. tile_kernels.set_token_alignment(128) x = torch.ones((T, 2 * H), dtype=dtype, device='npu') if dtype == torch.float32: @@ -94,6 +100,9 @@ def forward_nonfinite(): tile_kernels.set_token_alignment(256) x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') rows = [ + (20.0, 448.0 / 20.0), # finite control, ~448 + (20.0, 448.00003 / 20.0), # finite control, just above 448 + (20.0, 464.0 / 20.0), # finite control, RNE overflow threshold (3e38, 3e38), # +inf result -> amax = inf -> exponent 256 (3e38, -3e38), # -inf result (float('nan'), 1.0), # NaN result @@ -105,7 +114,26 @@ def forward_nonfinite(): x[i, :H] = left x[i, H:] = right kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) - return check('fwd_e4m3_nonfinite', run_f(False, x, **kw), run_f(True, x, **kw)) + asc_out, asc_sf = run_f(False, x, **kw) + vmi_out, vmi_sf = run_f(True, x, **kw) + asc_bytes = asc_out.view(torch.uint8).cpu() + vmi_bytes = vmi_out.view(torch.uint8).cpu() + # Known toolchain difference: ASC's RS_ENABLE does not saturate f32->e4m3 + # (non-finite/overflow -> NaN 0x7f), while VMI SAT does. See + # repro_fp8_saturate.py and deliverables/ISSUE_fp8_saturate.md. + asc_nan = asc_bytes == 0x7F + finite_ok = bool(torch.equal(asc_bytes[~asc_nan], vmi_bytes[~asc_nan])) + sat_ok = bool((vmi_bytes[asc_nan] != asc_bytes[asc_nan]).all()) + sf_ok = bool(torch.equal(asc_sf.cpu(), vmi_sf.cpu())) + ok = finite_ok and sat_ok and sf_ok + print( + 'fwd_e4m3_nonfinite:', + f'OK (finite rows match; known ASC RS_ENABLE difference on {int(asc_nan.sum())} lanes)' + if ok + else f'MISMATCH finite={finite_ok} sat={sat_ok} sf={sf_ok}', + flush=True, + ) + return ok def forward_sf_only(): @@ -134,6 +162,51 @@ def backward_subnormal(dtype, fmt): return check(f'bwd_{fmt}_subnormal', run_b(False, x, g, fmt=fmt), run_b(True, x, g, fmt=fmt)) +def backward_nonfinite(tma): + tile_kernels.set_token_alignment(256) + torch.manual_seed(20261021) + x_f32 = torch.randn((T, 2 * H), dtype=torch.float32, device='npu') + x_f32[0, :H] = float('nan') + x_f32[0, H:] = 1.0 + x_f32[1, :H] = -0.0 + x_f32[1, H:] = 1.0 + x_f32[2, :H] = -1e-40 # underflows to -0 in e4m3 + x_f32[2, H:] = 1.0 + x_data = x_f32.to(torch.float8_e4m3fn) + x_sf = torch.ones((T, 2 * H // 32), dtype=torch.float32, device='npu') + g = torch.randn((T, H), dtype=torch.float32, device='npu') + g[0, :] = float('inf') + g[1, :] = float('-inf') + g[2, :] = float('nan') + g[3, :] = -0.0 + g[4, :] = -1e-40 + kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=tma, round_sf=True, use_packed_ue8m0=True) + a = run_b(False, (x_data, x_sf), g, **kw) + b = run_b(True, (x_data, x_sf), g, **kw) + a_fp8, a_sf = a[1] + b_fp8, b_sf = b[1] + asc_bytes = a_fp8.view(torch.uint8).cpu() + vmi_bytes = b_fp8.view(torch.uint8).cpu() + # Same known ASC RS_ENABLE vs VMI SAT difference as the forward path. + asc_nan = asc_bytes == 0x7F + finite_ok = bool(torch.equal(asc_bytes[~asc_nan], vmi_bytes[~asc_nan])) + sat_ok = bool((vmi_bytes[asc_nan] != asc_bytes[asc_nan]).all()) + sf_ok = bits_equal(a_sf, b_sf) + # x_grad is bf16. The same ASC RS_ENABLE quirk also affects f32->bf16 + # narrowing (inf/NaN), so report it without making it a pass criterion. + xgrad_ok = bits_equal(a[0], b[0]) + ok = finite_ok and sat_ok and sf_ok + tag = 'tma' if tma else 'packed' + print( + f'bwd_e4m3_nonfinite_{tag}:', + f'OK (finite x_grad_fp8/SF match; known ASC RS_ENABLE difference on {int(asc_nan.sum())} lanes; x_grad bitwise={xgrad_ok})' + if ok + else f'MISMATCH finite={finite_ok} sat={sat_ok} sf={sf_ok} xgrad={xgrad_ok}', + flush=True, + ) + return ok + + if __name__ == '__main__': results = [ forward_subnormal(torch.float32, 'fp32'), @@ -142,5 +215,7 @@ def backward_subnormal(dtype, fmt): forward_sf_only(), backward_subnormal(torch.float32, 'fp32'), backward_subnormal(torch.bfloat16, 'bf16'), + backward_nonfinite(True), + backward_nonfinite(False), ] print('ALL_OK' if all(results) else 'HAS_MISMATCH') diff --git a/tile_kernels/quant/swiglu_backward_vmi.py b/tile_kernels/quant/swiglu_backward_vmi.py index 2558a24a..ce6ae3ed 100644 --- a/tile_kernels/quant/swiglu_backward_vmi.py +++ b/tile_kernels/quant/swiglu_backward_vmi.py @@ -488,12 +488,17 @@ def swiglu_backward_kernel_vmi( V.vstore(V.vcvt(y_grad_quant_value, 'bfloat16'), y_grad_ub[col]) x_scaled = V.vsel(valid_token_mask, x_scaled_value, zeros) if num_experts else x_scaled_value y_scaled = V.vsel(valid_token_mask, y_scaled_value, zeros) if num_experts else y_scaled_value - # ASC's FP8 conversion keeps non-finite values as - # NaN; PTODSL's SAT would clamp them. NOSAT matches - # ASC for non-finite inputs and for finite results - # at the 448 boundary. - V.vstore(V.vcvt(x_scaled, 'float8_e4m3fn', rounding='R', saturate='NOSAT'), x_grad_fp8_ub[col]) - V.vstore(V.vcvt(y_scaled, 'float8_e4m3fn', rounding='R', saturate='NOSAT'), y_grad_fp8_ub[col]) + # Faithful translation of ASC S.vcvt(sat=True). + # Toolchain caveat: on this build ASC's RS_ENABLE + # behaves like NOSAT for narrowing conversions + # (non-finite -> NaN), while VMI SAT really + # saturates. This is a toolchain inconsistency, + # not a semantic choice; see repro_fp8_saturate.py + # and the upstream issue draft. Do not "fix" it by + # switching to NOSAT. Finite values are unaffected: + # value <= 448*(1+2^-23) << 464, the RNE threshold. + V.vstore(V.vcvt(x_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), x_grad_fp8_ub[col]) + V.vstore(V.vcvt(y_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), y_grad_fp8_ub[col]) T.copy(x_grad_ub, x_grad[token_id, col_offset : col_offset + block_k]) T.copy(y_grad_ub, x_grad[token_id, hidden + col_offset : hidden + col_offset + block_k]) diff --git a/tile_kernels/quant/swiglu_forward_vmi.py b/tile_kernels/quant/swiglu_forward_vmi.py index b6d8b418..3cd5a720 100644 --- a/tile_kernels/quant/swiglu_forward_vmi.py +++ b/tile_kernels/quant/swiglu_forward_vmi.py @@ -67,6 +67,9 @@ def vdiv_exact_0ulp(a, b): The ``lt`` tie-break matches the built-in exact path used by ASC forward. ASC backward intentionally uses ``le``; do not unify them. + No kernel-level input in this port has been shown to reach an exact + midpoint tie (b is a transcendental result); the subnormal sweep in + review_adversarial_ab.py does not claim to cover that branch. """ full = V.create_mask(vec_size, size=vec_size) z = V.vdiv(a, b, full) @@ -269,14 +272,19 @@ def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offs col = (vector_base + vector) * vec_size inverse_indices = V.vadds(group_lane, T.int32(vector * 2), full) expanded_inverse = V.vselr(inverses, inverse_indices) - # ASC's FP8 conversion keeps non-finite values as NaN; - # PTODSL's SAT would clamp them. NOSAT matches ASC for both - # non-finite values and finite results at the 448 boundary. + # Faithful translation of ASC S.vcvt(sat=True). Toolchain + # caveat: on this build ASC's RS_ENABLE behaves like NOSAT + # for narrowing conversions (non-finite -> NaN), while VMI + # SAT really saturates. This is a toolchain inconsistency, + # not a semantic choice; see repro_fp8_saturate.py and the + # upstream issue draft. Do not "fix" it by switching to + # NOSAT. Finite values are unaffected: value <= + # 448*(1+2^-23) << 464, the RNE overflow threshold. quantized = V.vcvt( V.vmul(V.vload(values_ub[col], size=vec_size), expanded_inverse), 'float8_e4m3fn', rounding='R', - saturate='NOSAT', + saturate='SAT', ) V.vstore(quantized, out_ub[col]) From 27593c1e11fc1500858fb51af15a367edb38770b Mon Sep 17 00:00:00 2001 From: kuri780 Date: Tue, 20 Oct 2026 07:45:49 +0000 Subject: [PATCH 5/7] quant: use NOSAT for f32->e4m3/bf16 narrowing casts - FP8 and bf16 output casts explicitly use saturate='NOSAT': ASC's RS_ENABLE is a no-op (bug A) and VMI SAT maps NaN to 0x00 (bug B); NOSAT matches ASC today and preserves NaN. Add explicit NOSAT to the fp->int clamp-count cast and separate the int-narrow comment from the fp-narrow decision. - extend repro_fp8_saturate.py to f32->bf16 and add bf16 non-finite adversarial cases (forward output, backward x_grad/y_grad/act_out); all adversarial checks are now full bitwise equality, with no "assert a difference exists" logic. - 5-run medians on a quiet device: forward 1.035 (20/24 >= 0.90), backward 0.904 (8/24), combined 0.967 (28/48); per-config CV median 0.2%, max 1.6% forward / 4.4% backward. Co-Authored-By: Claude Code --- repro_fp8_saturate.py | 6 ++ review_adversarial_ab.py | 102 ++++++++++++---------- tile_kernels/quant/swiglu_backward_vmi.py | 49 ++++++----- tile_kernels/quant/swiglu_forward_vmi.py | 38 ++++---- 4 files changed, 112 insertions(+), 83 deletions(-) diff --git a/repro_fp8_saturate.py b/repro_fp8_saturate.py index c958bd78..a68e29eb 100644 --- a/repro_fp8_saturate.py +++ b/repro_fp8_saturate.py @@ -2,6 +2,8 @@ """Minimal f32 -> e4m3 / bf16 saturation repro for ASC S.vcvt vs VMI V.vcvt. Only one conversion is performed per kernel; no SwiGLU scale or mask is involved. +Upstream issue: (bug A: ASC RS_ENABLE is a no-op; +bug B: VMI SAT maps NaN to 0x00). Run inside the A5 container after sourcing activate_vmi_migration.sh. """ @@ -22,6 +24,10 @@ 464.0, 480.0, 1e30, + 3.4e38, + -3.4e38, + 1e-40, + -1e-40, ] VARIANTS = [ ('asc_sat', False, True), diff --git a/review_adversarial_ab.py b/review_adversarial_ab.py index bb5d327d..e8d2b955 100644 --- a/review_adversarial_ab.py +++ b/review_adversarial_ab.py @@ -5,9 +5,9 @@ * forward/backward exact divisions with subnormal operands (finite-domain check; the subnormal sweep does NOT exercise the 'le'/'lt' tie branch: b = 2 there, a/2 is exact, and r = 0) - * forward packed UE8M0 with amax = inf and non-finite/overflow outputs - (documents the known ASC RS_ENABLE vs VMI SAT toolchain difference) - * backward e4m3 non-finite inputs with packed/TMA and packed non-TMA SF + * forward e4m3/bf16 non-finite outputs, including amax = inf + * backward e4m3 (packed TMA / packed non-TMA) and bf16 non-finite inputs, + comparing x_grad, x_grad_fp8, SF and act_out bitwise Run inside the A5 container after sourcing activate_vmi_migration.sh. """ @@ -114,26 +114,25 @@ def forward_nonfinite(): x[i, :H] = left x[i, H:] = right kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) - asc_out, asc_sf = run_f(False, x, **kw) - vmi_out, vmi_sf = run_f(True, x, **kw) - asc_bytes = asc_out.view(torch.uint8).cpu() - vmi_bytes = vmi_out.view(torch.uint8).cpu() - # Known toolchain difference: ASC's RS_ENABLE does not saturate f32->e4m3 - # (non-finite/overflow -> NaN 0x7f), while VMI SAT does. See - # repro_fp8_saturate.py and deliverables/ISSUE_fp8_saturate.md. - asc_nan = asc_bytes == 0x7F - finite_ok = bool(torch.equal(asc_bytes[~asc_nan], vmi_bytes[~asc_nan])) - sat_ok = bool((vmi_bytes[asc_nan] != asc_bytes[asc_nan]).all()) - sf_ok = bool(torch.equal(asc_sf.cpu(), vmi_sf.cpu())) - ok = finite_ok and sat_ok and sf_ok - print( - 'fwd_e4m3_nonfinite:', - f'OK (finite rows match; known ASC RS_ENABLE difference on {int(asc_nan.sum())} lanes)' - if ok - else f'MISMATCH finite={finite_ok} sat={sat_ok} sf={sf_ok}', - flush=True, - ) - return ok + return check('fwd_e4m3_nonfinite', run_f(False, x, **kw), run_f(True, x, **kw)) + + +def forward_nonfinite_bf16(): + tile_kernels.set_token_alignment(128) + x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') + rows = [ + (float('inf'), 1.0), + (float('-inf'), 1.0), + (float('nan'), 1.0), + (-0.0, 1.0), + (-1e-40, 1.0), # bf16 subnormal + (3e38, 1.0), # bf16 finite max vicinity + (2.0, 3.0), # finite control + ] + for i, (left, right) in enumerate(rows): + x[i, :H] = left + x[i, H:] = right + return check('fwd_bf16_nonfinite', run_f(False, x, fmt='bf16'), run_f(True, x, fmt='bf16')) def forward_sf_only(): @@ -180,31 +179,38 @@ def backward_nonfinite(tma): g[2, :] = float('nan') g[3, :] = -0.0 g[4, :] = -1e-40 - kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=tma, round_sf=True, use_packed_ue8m0=True) - a = run_b(False, (x_data, x_sf), g, **kw) - b = run_b(True, (x_data, x_sf), g, **kw) - a_fp8, a_sf = a[1] - b_fp8, b_sf = b[1] - asc_bytes = a_fp8.view(torch.uint8).cpu() - vmi_bytes = b_fp8.view(torch.uint8).cpu() - # Same known ASC RS_ENABLE vs VMI SAT difference as the forward path. - asc_nan = asc_bytes == 0x7F - finite_ok = bool(torch.equal(asc_bytes[~asc_nan], vmi_bytes[~asc_nan])) - sat_ok = bool((vmi_bytes[asc_nan] != asc_bytes[asc_nan]).all()) - sf_ok = bits_equal(a_sf, b_sf) - # x_grad is bf16. The same ASC RS_ENABLE quirk also affects f32->bf16 - # narrowing (inf/NaN), so report it without making it a pass criterion. - xgrad_ok = bits_equal(a[0], b[0]) - ok = finite_ok and sat_ok and sf_ok - tag = 'tma' if tma else 'packed' - print( - f'bwd_e4m3_nonfinite_{tag}:', - f'OK (finite x_grad_fp8/SF match; known ASC RS_ENABLE difference on {int(asc_nan.sum())} lanes; x_grad bitwise={xgrad_ok})' - if ok - else f'MISMATCH finite={finite_ok} sat={sat_ok} sf={sf_ok} xgrad={xgrad_ok}', - flush=True, + kw = dict( + fmt='e4m3', + num_per_channels=32, + use_tma_aligned_col_major_sf=tma, + round_sf=True, + use_packed_ue8m0=True, + do_recompute=True, ) - return ok + tag = 'tma' if tma else 'packed' + return check(f'bwd_e4m3_nonfinite_{tag}', run_b(False, (x_data, x_sf), g, **kw), run_b(True, (x_data, x_sf), g, **kw)) + + +def backward_nonfinite_bf16(): + tile_kernels.set_token_alignment(128) + torch.manual_seed(20261022) + x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') + x[0, :H] = float('inf') + x[1, :H] = float('-inf') + x[2, :H] = float('nan') + x[3, :H] = -0.0 + x[4, :H] = -1e-40 # bf16 subnormal + x[5, :H] = 3e38 + for row in range(6): + x[row, H:] = 1.0 + g = torch.randn((T, H), dtype=torch.bfloat16, device='npu') + g[0, :] = float('inf') + g[1, :] = float('-inf') + g[2, :] = float('nan') + g[3, :] = -0.0 + g[4, :] = -1e-40 + g[5, :] = 3e38 + return check('bwd_bf16_nonfinite', run_b(False, x, g, fmt='bf16', do_recompute=True), run_b(True, x, g, fmt='bf16', do_recompute=True)) if __name__ == '__main__': @@ -212,10 +218,12 @@ def backward_nonfinite(tma): forward_subnormal(torch.float32, 'fp32'), forward_subnormal(torch.bfloat16, 'bf16'), forward_nonfinite(), + forward_nonfinite_bf16(), forward_sf_only(), backward_subnormal(torch.float32, 'fp32'), backward_subnormal(torch.bfloat16, 'bf16'), backward_nonfinite(True), backward_nonfinite(False), + backward_nonfinite_bf16(), ] print('ALL_OK' if all(results) else 'HAS_MISMATCH') diff --git a/tile_kernels/quant/swiglu_backward_vmi.py b/tile_kernels/quant/swiglu_backward_vmi.py index ce6ae3ed..7eb8d02e 100644 --- a/tile_kernels/quant/swiglu_backward_vmi.py +++ b/tile_kernels/quant/swiglu_backward_vmi.py @@ -403,8 +403,11 @@ def swiglu_backward_kernel_vmi( y_grad_value = V.vsel(is_clamped_y, zeros, y_grad_unclamped) if use_clamp else y_grad_unclamped if not with_quant: if out_dtype == T.bfloat16: - V.vstore(V.vcvt(x_grad_value, 'bfloat16'), x_grad_ub[col]) - V.vstore(V.vcvt(y_grad_value, 'bfloat16'), y_grad_ub[col]) + # Explicit NOSAT: same bug A/B rationale as + # the FP8 output; do not rely on PTODSL's + # implicit SAT default for narrowing. + V.vstore(V.vcvt(x_grad_value, 'bfloat16', saturate='NOSAT'), x_grad_ub[col]) + V.vstore(V.vcvt(y_grad_value, 'bfloat16', saturate='NOSAT'), y_grad_ub[col]) else: V.vstore(x_grad_value, x_grad_ub[col]) V.vstore(y_grad_value, y_grad_ub[col]) @@ -418,7 +421,8 @@ def swiglu_backward_kernel_vmi( if do_recompute: act_out_value = V.vmuls(act_value, effective_weight, full) if out_dtype == T.bfloat16: - V.vstore(V.vcvt(act_out_value, 'bfloat16'), act_out_ub[col]) + # Explicit NOSAT, same bug A/B rationale. + V.vstore(V.vcvt(act_out_value, 'bfloat16', saturate='NOSAT'), act_out_ub[col]) else: V.vstore(act_out_value, act_out_ub[col]) @@ -456,11 +460,13 @@ def swiglu_backward_kernel_vmi( if out_config.use_packed_ue8m0: x_exp = V.vshrs(V.vinterpret_cast(x_scale_out, 'uint32'), T.uint32(23), full) y_exp = V.vshrs(V.vinterpret_cast(y_scale_out, 'uint32'), T.uint32(23), full) - # ASC reinterprets the shifted exponent bits - # as uint8 (truncation). Spell out NOSAT - # because PTODSL otherwise defaults to SAT for - # this narrowing cast, which differs for - # exponents >= 256 (amax = inf). + # int-narrow (uint32 -> uint8): NOSAT is + # exactly ASC's reinterpret-truncation; + # unrelated to the FP8 output saturate + # decision below. Spell it out because PTODSL + # otherwise defaults to SAT for narrowing, + # which differs for exponents >= 256 + # (amax = inf). V.vstore( V.vcvt(x_exp, 'uint8', saturate='NOSAT'), sf_dense_u8[row, sf_block_offset * 2 + scale_offset], @@ -484,21 +490,22 @@ def swiglu_backward_kernel_vmi( y_grad_quant_value = V.vload(y_grad_f32_ub[col], size=vec_size) x_scaled_value = V.vmul(x_grad_quant_value, x_inv_vec, full) y_scaled_value = V.vmul(y_grad_quant_value, y_inv_vec, full) - V.vstore(V.vcvt(x_grad_quant_value, 'bfloat16'), x_grad_ub[col]) - V.vstore(V.vcvt(y_grad_quant_value, 'bfloat16'), y_grad_ub[col]) + # Explicit NOSAT, same bug A/B rationale. + V.vstore(V.vcvt(x_grad_quant_value, 'bfloat16', saturate='NOSAT'), x_grad_ub[col]) + V.vstore(V.vcvt(y_grad_quant_value, 'bfloat16', saturate='NOSAT'), y_grad_ub[col]) x_scaled = V.vsel(valid_token_mask, x_scaled_value, zeros) if num_experts else x_scaled_value y_scaled = V.vsel(valid_token_mask, y_scaled_value, zeros) if num_experts else y_scaled_value - # Faithful translation of ASC S.vcvt(sat=True). - # Toolchain caveat: on this build ASC's RS_ENABLE - # behaves like NOSAT for narrowing conversions - # (non-finite -> NaN), while VMI SAT really - # saturates. This is a toolchain inconsistency, - # not a semantic choice; see repro_fp8_saturate.py - # and the upstream issue draft. Do not "fix" it by - # switching to NOSAT. Finite values are unaffected: - # value <= 448*(1+2^-23) << 464, the RNE threshold. - V.vstore(V.vcvt(x_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), x_grad_fp8_ub[col]) - V.vstore(V.vcvt(y_scaled, 'float8_e4m3fn', rounding='R', saturate='SAT'), y_grad_fp8_ub[col]) + # FP8 output saturate: explicit NOSAT. ASC's + # RS_ENABLE is a no-op for f32->e4m3 here (bug A), + # while VMI SAT maps NaN to 0x00 (bug B). NOSAT is + # the only mode that both reproduces ASC today and + # preserves NaN; switch back to SAT after bug A is + # fixed upstream. Both bugs: + # (draft ready; to be filed), repro_fp8_saturate.py. + # Finite values are + # unaffected: value <= 448*(1+2^-23) << 464. + V.vstore(V.vcvt(x_scaled, 'float8_e4m3fn', rounding='R', saturate='NOSAT'), x_grad_fp8_ub[col]) + V.vstore(V.vcvt(y_scaled, 'float8_e4m3fn', rounding='R', saturate='NOSAT'), y_grad_fp8_ub[col]) T.copy(x_grad_ub, x_grad[token_id, col_offset : col_offset + block_k]) T.copy(y_grad_ub, x_grad[token_id, hidden + col_offset : hidden + col_offset + block_k]) diff --git a/tile_kernels/quant/swiglu_forward_vmi.py b/tile_kernels/quant/swiglu_forward_vmi.py index 3cd5a720..5c1c62ba 100644 --- a/tile_kernels/quant/swiglu_forward_vmi.py +++ b/tile_kernels/quant/swiglu_forward_vmi.py @@ -206,7 +206,11 @@ def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, if out_config.with_sf: V.vstore(out_val, values_ub[col]) elif out_config.dtype == T.bfloat16: - V.vstore(V.vcvt(out_val, 'bfloat16'), out_ub[col]) + # Explicit NOSAT for the same bug A/B reason as the FP8 + # output: ASC RS_ENABLE is a no-op and VMI SAT would + # turn NaN into 0 / clamp inf. PTOAS accepts NOSAT for + # this narrowing (verified in repro_fp8_saturate.py). + V.vstore(V.vcvt(out_val, 'bfloat16', saturate='NOSAT'), out_ub[col]) else: V.vstore(out_val, out_ub[col]) @@ -248,10 +252,11 @@ def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offs if out_config.round_sf: scales_or_exponents, inverses = get_scale_and_inverse(amax) if out_config.use_packed_ue8m0: - # ASC reinterprets the shifted exponent bits as uint8 - # (truncation). Spell out NOSAT because PTODSL would - # otherwise default to SAT for this narrowing cast, - # which differs for exponents >= 256 (amax = inf). + # int-narrow (uint32 -> uint8): NOSAT is exactly ASC's + # reinterpret-truncation, unrelated to the FP8 output + # saturate decision below. Spell it out because PTODSL + # otherwise defaults to SAT for narrowing, which differs + # for exponents >= 256 (amax = inf). V.vstore(V.vcvt(scales_or_exponents, 'uint8', saturate='NOSAT'), sf_dense_ub[token_offset, group]) else: V.vstore(scales_or_exponents, sf_dense_ub[token_offset, group]) @@ -272,19 +277,19 @@ def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offs col = (vector_base + vector) * vec_size inverse_indices = V.vadds(group_lane, T.int32(vector * 2), full) expanded_inverse = V.vselr(inverses, inverse_indices) - # Faithful translation of ASC S.vcvt(sat=True). Toolchain - # caveat: on this build ASC's RS_ENABLE behaves like NOSAT - # for narrowing conversions (non-finite -> NaN), while VMI - # SAT really saturates. This is a toolchain inconsistency, - # not a semantic choice; see repro_fp8_saturate.py and the - # upstream issue draft. Do not "fix" it by switching to - # NOSAT. Finite values are unaffected: value <= - # 448*(1+2^-23) << 464, the RNE overflow threshold. + # FP8 output saturate: explicit NOSAT. ASC's RS_ENABLE is a + # no-op for f32->e4m3 here (bug A: sat=True and sat=False + # produce the same bytes), while VMI SAT maps NaN to 0x00 + # (bug B). NOSAT is the only mode that both reproduces ASC + # today and preserves NaN; switch back to SAT after bug A is + # fixed upstream. Both bugs: (draft ready; + # to be filed), repro_fp8_saturate.py. Finite values are unaffected: + # value <= 448*(1+2^-23) << 464 (e4m3 RNE overflow threshold). quantized = V.vcvt( V.vmul(V.vload(values_ub[col], size=vec_size), expanded_inverse), 'float8_e4m3fn', rounding='R', - saturate='SAT', + saturate='NOSAT', ) V.vstore(quantized, out_ub[col]) @@ -475,7 +480,10 @@ def swiglu_forward_kernel( # These are exact integer counts (< 2^24), so the # reassociation allowed by PTODSL cannot change them. reduced = V.vcadd(V.vload(pacc_ub[i, 0], size=vec_size), full, reassoc=True) - V.vstore(V.vcvt(reduced, 'int32'), acc_ub[i], group=1, stride=1) + # fp->int narrow: the count is < 2^24, so SAT/NOSAT are + # identical. Spell out NOSAT instead of relying on + # PTODSL's implicit SAT default. + V.vstore(V.vcvt(reduced, 'int32', saturate='NOSAT'), acc_ub[i], group=1, stride=1) with T.SimtVF(threads=1): for i in T.serial(4): From 10b5bc50ae27501d26182f27f0b81c8179626ff5 Mon Sep 17 00:00:00 2001 From: kuri780 Date: Tue, 20 Oct 2026 08:55:49 +0000 Subject: [PATCH 6/7] chore: keep only the two SwiGLU VMI kernels in the PR Move the reproduction/coverage scripts out of the repository. They remain available locally under deliverables/repro_scripts/ but are not part of the mainline change. Co-Authored-By: Claude Code --- repro_fp8_saturate.py | 104 ------------------ review_adversarial_ab.py | 229 --------------------------------------- review_coverage_ab.py | 156 -------------------------- vmi_swiglu_plugin.py | 25 ----- 4 files changed, 514 deletions(-) delete mode 100644 repro_fp8_saturate.py delete mode 100644 review_adversarial_ab.py delete mode 100644 review_coverage_ab.py delete mode 100644 vmi_swiglu_plugin.py diff --git a/repro_fp8_saturate.py b/repro_fp8_saturate.py deleted file mode 100644 index a68e29eb..00000000 --- a/repro_fp8_saturate.py +++ /dev/null @@ -1,104 +0,0 @@ -#!/usr/bin/env python3 -"""Minimal f32 -> e4m3 / bf16 saturation repro for ASC S.vcvt vs VMI V.vcvt. - -Only one conversion is performed per kernel; no SwiGLU scale or mask is involved. -Upstream issue: (bug A: ASC RS_ENABLE is a no-op; -bug B: VMI SAT maps NaN to 0x00). -Run inside the A5 container after sourcing activate_vmi_migration.sh. -""" - -import torch -import torch_npu - -import tilelang -import tilelang.ascend.language as T - -N = 64 -VALUES = [ - float('inf'), - float('-inf'), - float('nan'), - -0.0, - 448.0, - 448.00003, - 464.0, - 480.0, - 1e30, - 3.4e38, - -3.4e38, - 1e-40, - -1e-40, -] -VARIANTS = [ - ('asc_sat', False, True), - ('asc_nosat', False, False), - ('vmi_sat', True, True), - ('vmi_nosat', True, False), -] -TARGETS = [ - ('e4m3', 'float8_e4m3fn', T.float8_e4m3fn), - ('bf16', 'bfloat16', T.bfloat16), -] - - -def make_program(vmi: bool, sat: bool, target: str, target_tir): - @T.prim_func - def main(x: T.Buffer((N,), 'float32'), y: T.Buffer((N,), target)): - with T.Kernel(1) as _: - x_ub = T.alloc_shared((N,), 'float32') - y_ub = T.alloc_shared((N,), target) - T.copy(x, x_ub) - with T.SimdVF(): - full = T.vmi.create_mask(N, size=N) if vmi else T.simd.pset(32) - if vmi: - value = T.vmi.vload(x_ub[0], size=N) - converted = T.vmi.vcvt( - value, - target, - rounding='R', - saturate='SAT' if sat else 'NOSAT', - ) - T.vmi.vstore(converted, y_ub[0], full) - else: - value = T.simd.vld(x_ub[0]) - converted = T.simd.vcvt(value, target_tir, sat=sat) - dist = 'PK4_B32' if target == 'float8_e4m3fn' else 'PK_B32' - T.simd.vsts(y_ub[0], converted, full, dist=dist) - T.copy(y_ub, y) - - return main - - -def main(): - x_cpu = torch.tensor([VALUES[i % len(VALUES)] for i in range(N)], dtype=torch.float32) - x = x_cpu.to('npu') - for target, target_param, target_tir in TARGETS: - results = {} - for name, vmi, sat in VARIANTS: - program = make_program(vmi, sat, target_param, target_tir) - kernel = tilelang.compile(program, target='pto' if vmi else 'asc', out_idx=-1) - y = kernel(x) - torch.npu.synchronize() - if target_param == 'bfloat16': - bits = (y.cpu().view(torch.int16).to(torch.int32) & 0xFFFF).tolist() - else: - bits = y.cpu().view(torch.uint8).to(torch.int32).tolist() - results[name] = bits[: len(VALUES)] - source = kernel.get_kernel_source() - path = f'/tmp/repro_{target}_saturate_{name}.txt' - with open(path, 'w') as f: - f.write(source) - print(f'--- {target} {name} source: {path}') - for line in source.splitlines(): - if any(tok in line for tok in ('vcvt', 'saturate', 'RS_ENABLE', 'RS_DISABLE')): - print(' ', line.strip()) - - print(f'\n=== {target} bit patterns (hex) ===') - print('value ' + ' '.join(f'{name:>10}' for name, _, _ in VARIANTS)) - for i, value in enumerate(VALUES): - print(f'{value!r:<22}' + ' '.join(f'{results[name][i]:>#10x}' for name, _, _ in VARIANTS)) - print() - - -if __name__ == '__main__': - main() diff --git a/review_adversarial_ab.py b/review_adversarial_ab.py deleted file mode 100644 index e8d2b955..00000000 --- a/review_adversarial_ab.py +++ /dev/null @@ -1,229 +0,0 @@ -#!/usr/bin/env python3 -"""Adversarial ASC-vs-VMI checks for the PTOAS #1616 review risks. - -Targets the cases that random data does not reach: - * forward/backward exact divisions with subnormal operands (finite-domain - check; the subnormal sweep does NOT exercise the 'le'/'lt' tie branch: - b = 2 there, a/2 is exact, and r = 0) - * forward e4m3/bf16 non-finite outputs, including amax = inf - * backward e4m3 (packed TMA / packed non-TMA) and bf16 non-finite inputs, - comparing x_grad, x_grad_fp8, SF and act_out bitwise - -Run inside the A5 container after sourcing activate_vmi_migration.sh. -""" - -import torch -import torch_npu - -import tile_kernels -import tile_kernels.quant.swiglu_backward_asc as ba -import tile_kernels.quant.swiglu_forward_asc as fa -from tile_kernels.quant.swiglu_backward_vmi import get_swiglu_backward_kernel_vmi -from tile_kernels.quant.swiglu_forward_vmi import get_swiglu_forward_kernel_vmi - -H = 2048 -T = 1024 -of, ob = fa.get_swiglu_forward_kernel_asc, ba.get_swiglu_backward_kernel_asc - - -def cpu(t): - t = t.detach().cpu() - return t.view(torch.uint8) if t.dtype == torch.float8_e4m3fn else t - - -def bits_equal(a, b): - if isinstance(a, tuple): - return len(a) == len(b) and all(bits_equal(x, y) for x, y in zip(a, b)) - if isinstance(a, torch.Tensor): - x, y = cpu(a), cpu(b) - if x.dtype.is_floating_point: - return torch.equal(x.view(torch.uint8), y.view(torch.uint8)) - return torch.equal(x, y) - return a == b - - -def run_f(vmi, x, **kw): - fn = tile_kernels.quant.swiglu_forward_and_per_token_cast if kw.get('fmt') == 'e4m3' else tile_kernels.quant.swiglu_forward - fa.get_swiglu_forward_kernel_asc = get_swiglu_forward_kernel_vmi if vmi else of - try: - return fn(x, **kw) - finally: - fa.get_swiglu_forward_kernel_asc = of - - -def run_b(vmi, x, g, **kw): - fn = tile_kernels.quant.swiglu_backward_and_per_token_cast if kw.get('fmt') == 'e4m3' else tile_kernels.quant.swiglu_backward - ba.get_swiglu_backward_kernel_asc = get_swiglu_backward_kernel_vmi if vmi else ob - try: - return fn(x, g, **kw) - finally: - ba.get_swiglu_backward_kernel_asc = ob - - -def subnormal_bits(lo, hi, count, negative): - """Evenly sample float32 subnormal bit patterns.""" - step = max(1, (hi - lo + 1) // count) - idx = torch.arange(lo, hi + 1, step, dtype=torch.int64) - if idx.numel() < count: - idx = torch.cat((idx, idx[-1].repeat(count - idx.numel()))) - idx = idx[:count] - if negative: - idx = idx - (1 << 32) - return idx.to(torch.int32).view(torch.float32) - - -def check(tag, a, b): - ok = bits_equal(a, b) - print(f'{tag}:', 'OK' if ok else 'MISMATCH', flush=True) - return ok - - -def forward_subnormal(dtype, fmt): - # For tiny |x|, b = 1 + exp(-x) = 2 exactly and a/2 is exact, so r = 0. - # This sweep checks finite subnormal handling, not the le/lt tie branch. - tile_kernels.set_token_alignment(128) - x = torch.ones((T, 2 * H), dtype=dtype, device='npu') - if dtype == torch.float32: - # All negative float32 subnormals, evenly sampled. - vals = subnormal_bits(0x80000001, 0x807FFFFF, T * H, negative=True).view(T, H).to('npu') - else: - # All negative bf16 subnormals, tiled. - bits = torch.arange(0x8001, 0x8080, dtype=torch.int64).to(torch.int16) - vals = bits.view(torch.bfloat16).to('npu') - vals = vals.repeat((T * H + vals.numel() - 1) // vals.numel())[: T * H].view(T, H) - x[:, :H] = vals - x[:, H:] = 1.0 - return check(f'fwd_{fmt}_subnormal', run_f(False, x, fmt=fmt), run_f(True, x, fmt=fmt)) - - -def forward_nonfinite(): - tile_kernels.set_token_alignment(256) - x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') - rows = [ - (20.0, 448.0 / 20.0), # finite control, ~448 - (20.0, 448.00003 / 20.0), # finite control, just above 448 - (20.0, 464.0 / 20.0), # finite control, RNE overflow threshold - (3e38, 3e38), # +inf result -> amax = inf -> exponent 256 - (3e38, -3e38), # -inf result - (float('nan'), 1.0), # NaN result - (-100.0, 100.0), # -0 result - (float('inf'), 1.0), # inf input - (float('-inf'), 1.0), # -inf input - ] - for i, (left, right) in enumerate(rows): - x[i, :H] = left - x[i, H:] = right - kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) - return check('fwd_e4m3_nonfinite', run_f(False, x, **kw), run_f(True, x, **kw)) - - -def forward_nonfinite_bf16(): - tile_kernels.set_token_alignment(128) - x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') - rows = [ - (float('inf'), 1.0), - (float('-inf'), 1.0), - (float('nan'), 1.0), - (-0.0, 1.0), - (-1e-40, 1.0), # bf16 subnormal - (3e38, 1.0), # bf16 finite max vicinity - (2.0, 3.0), # finite control - ] - for i, (left, right) in enumerate(rows): - x[i, :H] = left - x[i, H:] = right - return check('fwd_bf16_nonfinite', run_f(False, x, fmt='bf16'), run_f(True, x, fmt='bf16')) - - -def forward_sf_only(): - # Focused check for the vcmax(group=...) scale mapping: compare only SF. - tile_kernels.set_token_alignment(256) - torch.manual_seed(20261020) - x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') - kw = dict(fmt='e4m3', num_per_channels=32, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) - a = run_f(False, x, **kw) - b = run_f(True, x, **kw) - return check('fwd_e4m3_sf_only', a[1], b[1]) - - -def backward_subnormal(dtype, fmt): - tile_kernels.set_token_alignment(128) - x = torch.ones((T, 2 * H), dtype=dtype, device='npu') - if dtype == torch.float32: - vals = subnormal_bits(0x80000001, 0x807FFFFF, T * H, negative=True).view(T, H).to('npu') - else: - bits = torch.arange(0x8001, 0x8080, dtype=torch.int64).to(torch.int16) - vals = bits.view(torch.bfloat16).to('npu') - vals = vals.repeat((T * H + vals.numel() - 1) // vals.numel())[: T * H].view(T, H) - x[:, :H] = vals - x[:, H:] = 1.0 - g = torch.randn((T, H), dtype=dtype, device='npu') - return check(f'bwd_{fmt}_subnormal', run_b(False, x, g, fmt=fmt), run_b(True, x, g, fmt=fmt)) - - -def backward_nonfinite(tma): - tile_kernels.set_token_alignment(256) - torch.manual_seed(20261021) - x_f32 = torch.randn((T, 2 * H), dtype=torch.float32, device='npu') - x_f32[0, :H] = float('nan') - x_f32[0, H:] = 1.0 - x_f32[1, :H] = -0.0 - x_f32[1, H:] = 1.0 - x_f32[2, :H] = -1e-40 # underflows to -0 in e4m3 - x_f32[2, H:] = 1.0 - x_data = x_f32.to(torch.float8_e4m3fn) - x_sf = torch.ones((T, 2 * H // 32), dtype=torch.float32, device='npu') - g = torch.randn((T, H), dtype=torch.float32, device='npu') - g[0, :] = float('inf') - g[1, :] = float('-inf') - g[2, :] = float('nan') - g[3, :] = -0.0 - g[4, :] = -1e-40 - kw = dict( - fmt='e4m3', - num_per_channels=32, - use_tma_aligned_col_major_sf=tma, - round_sf=True, - use_packed_ue8m0=True, - do_recompute=True, - ) - tag = 'tma' if tma else 'packed' - return check(f'bwd_e4m3_nonfinite_{tag}', run_b(False, (x_data, x_sf), g, **kw), run_b(True, (x_data, x_sf), g, **kw)) - - -def backward_nonfinite_bf16(): - tile_kernels.set_token_alignment(128) - torch.manual_seed(20261022) - x = torch.randn((T, 2 * H), dtype=torch.bfloat16, device='npu') - x[0, :H] = float('inf') - x[1, :H] = float('-inf') - x[2, :H] = float('nan') - x[3, :H] = -0.0 - x[4, :H] = -1e-40 # bf16 subnormal - x[5, :H] = 3e38 - for row in range(6): - x[row, H:] = 1.0 - g = torch.randn((T, H), dtype=torch.bfloat16, device='npu') - g[0, :] = float('inf') - g[1, :] = float('-inf') - g[2, :] = float('nan') - g[3, :] = -0.0 - g[4, :] = -1e-40 - g[5, :] = 3e38 - return check('bwd_bf16_nonfinite', run_b(False, x, g, fmt='bf16', do_recompute=True), run_b(True, x, g, fmt='bf16', do_recompute=True)) - - -if __name__ == '__main__': - results = [ - forward_subnormal(torch.float32, 'fp32'), - forward_subnormal(torch.bfloat16, 'bf16'), - forward_nonfinite(), - forward_nonfinite_bf16(), - forward_sf_only(), - backward_subnormal(torch.float32, 'fp32'), - backward_subnormal(torch.bfloat16, 'bf16'), - backward_nonfinite(True), - backward_nonfinite(False), - backward_nonfinite_bf16(), - ] - print('ALL_OK' if all(results) else 'HAS_MISMATCH') diff --git a/review_coverage_ab.py b/review_coverage_ab.py deleted file mode 100644 index ad048166..00000000 --- a/review_coverage_ab.py +++ /dev/null @@ -1,156 +0,0 @@ -#!/usr/bin/env python3 -"""Extra ASC-vs-VMI coverage requested in review.""" - -import torch -import torch_npu -import tile_kernels -import tile_kernels.quant.swiglu_forward_asc as fa -import tile_kernels.quant.swiglu_backward_asc as ba -from tile_kernels.quant.swiglu_forward_vmi import get_swiglu_forward_kernel_vmi -from tile_kernels.quant.swiglu_backward_vmi import get_swiglu_backward_kernel_vmi - -of, ob = fa.get_swiglu_forward_kernel_asc, ba.get_swiglu_backward_kernel_asc -H = 2048 -PSUM = torch.tensor([600, 1600], dtype=torch.int32, device='npu') - - -def cpu(t): - t = t.detach().cpu() - return t.view(torch.uint8) if t.dtype == torch.float8_e4m3fn else t - - -def cmp_tree(a, b, rows): - if isinstance(a, tuple): - return all(cmp_tree(x, y, rows) for x, y in zip(a, b)) - if isinstance(a, torch.Tensor): - return torch.equal(cpu(a)[rows], cpu(b)[rows]) - return a == b - - -def layout(use_psum): - if not use_psum: - return None, torch.ones(1024, dtype=torch.bool) - mask = torch.zeros(1792, dtype=torch.bool) - mask[0:600] = True - mask[768:1600] = True - return PSUM, mask - - -def run_f(vmi, x, fn, **kw): - fa.get_swiglu_forward_kernel_asc = get_swiglu_forward_kernel_vmi if vmi else of - try: - return fn(x=x, **kw) - finally: - fa.get_swiglu_forward_kernel_asc = of - - -def run_b(vmi, x, g, fn, **kw): - ba.get_swiglu_backward_kernel_asc = get_swiglu_backward_kernel_vmi if vmi else ob - try: - return fn(x, g, **kw) - finally: - ba.get_swiglu_backward_kernel_asc = ob - - -def sf_kwargs(fmt, packed, round_sf, tma): - if fmt != 'e4m3': - return {} - return dict( - fmt=fmt, - num_per_channels=32, - use_tma_aligned_col_major_sf=tma, - round_sf=round_sf, - use_packed_ue8m0=packed, - ) - - -def fwd(tag, dtype, use_psum, fmt, weights=False, rsf=None, clamp=None, packed=False, round_sf=False, tma=False): - tile_kernels.set_token_alignment(256 if use_psum else 128) - psum, mask = layout(use_psum) - torch.manual_seed(101) - T = 1792 if use_psum else 1024 - x = torch.randn((T, 2 * H), device='npu', dtype=dtype) - w = torch.rand((T,), device='npu', dtype=torch.float32) - if use_psum: - m = mask.to('npu') - x[~m] = 0 - w[~m] = 0 - kw = sf_kwargs(fmt, packed, round_sf, tma) - if fmt != 'e4m3': - kw['fmt'] = fmt - if weights: - kw.update(topk_weights=w, routed_scaling_factor=rsf) - if clamp is not None: - kw.update(clamp_value=clamp) - if fmt == 'e4m3': - fn = tile_kernels.quant.swiglu_forward_and_per_token_cast - else: - fn = tile_kernels.quant.swiglu_forward - a = run_f(False, x, fn, psum_num_tokens_per_expert=psum, **kw) - b = run_f(True, x, fn, psum_num_tokens_per_expert=psum, **kw) - torch.npu.synchronize() - print(f'{tag}:', 'OK' if cmp_tree(a, b, mask) else 'MISMATCH') - - -def bwd(tag, fmt, use_psum, dtype, weights=False, rsf=None, clamp=None, recompute=False, packed=False, round_sf=False, tma=False): - tile_kernels.set_token_alignment(256 if (use_psum or fmt == 'e4m3') else 128) - psum, mask = layout(use_psum) - torch.manual_seed(102) - T = 1792 if use_psum else 1024 - g_dtype = torch.bfloat16 if fmt == 'e4m3' else dtype - x_raw = torch.randn((T, 2 * H), device='npu', dtype=torch.bfloat16) - g = torch.randn((T, H), device='npu', dtype=g_dtype) - w = torch.rand((T,), device='npu', dtype=torch.float32) - if use_psum: - m = mask.to('npu') - x_raw[~m] = 0 - g[~m] = 0 - w[~m] = 0 - kw = sf_kwargs(fmt, packed, round_sf, tma) - if fmt != 'e4m3': - kw['fmt'] = fmt - if weights: - kw.update(topk_weights=w, routed_scaling_factor=rsf) - if clamp is not None: - kw.update(clamp_value=clamp) - kw.update(do_recompute=recompute) - if fmt == 'e4m3': - x = tile_kernels.quant.per_token_cast(x_raw, 'e4m3', num_per_channels=32) - fn = tile_kernels.quant.swiglu_backward_and_per_token_cast - else: - x = x_raw.to(dtype) - fn = tile_kernels.quant.swiglu_backward - a = run_b(False, x, g, fn, psum_num_tokens_per_expert=psum, **kw) - b = run_b(True, x, g, fn, psum_num_tokens_per_expert=psum, **kw) - torch.npu.synchronize() - print(f'{tag}:', 'OK' if cmp_tree(a, b, mask) else 'MISMATCH') - - -fwd('fwd_fp32_in_bf16_out_no_psum', torch.float32, False, 'bf16') -fwd('fwd_fp32_in_fp32_out_no_psum_clamp', torch.float32, False, 'fp32', clamp=0.5) -fwd('fwd_fp32_in_bf16_out_weights_rsf', torch.float32, False, 'bf16', weights=True, rsf=2.5) -fwd('fwd_fp32_in_e4m3_packed_no_tma', torch.float32, False, 'e4m3', packed=True, round_sf=True, tma=False) -fwd('fwd_e4m3_packed_no_tma_psum_weights', torch.bfloat16, True, 'e4m3', weights=True, rsf=2.5, packed=True, round_sf=True, tma=False) -fwd('fwd_e4m3_dense_exact_no_psum', torch.bfloat16, False, 'e4m3', round_sf=False, tma=False) -fwd('fwd_e4m3_tma_weights_clamp', torch.bfloat16, True, 'e4m3', weights=True, rsf=2.5, clamp=0.5, packed=True, round_sf=True, tma=True) - -bwd('bwd_bf16_no_psum_no_weights', 'bf16', False, torch.bfloat16) -bwd('bwd_bf16_psum_weights_no_rsf', 'bf16', True, torch.bfloat16, weights=True, rsf=None) -bwd('bwd_bf16_psum_weights_recompute_clamp', 'bf16', True, torch.bfloat16, weights=True, rsf=2.5, clamp=0.5, recompute=True) -bwd('bwd_fp32_no_psum_weights_recompute', 'fp32', False, torch.float32, weights=True, rsf=2.5, recompute=True) -bwd( - 'bwd_e4m3_no_psum_tma_weights_recompute', - 'e4m3', - False, - torch.bfloat16, - weights=True, - rsf=2.5, - recompute=True, - packed=True, - round_sf=True, - tma=True, -) -bwd('bwd_e4m3_psum_packed_no_tma_weights', 'e4m3', True, torch.bfloat16, weights=True, rsf=2.5, packed=True, round_sf=True, tma=False) -bwd('bwd_e4m3_psum_dense_exact_recompute', 'e4m3', True, torch.bfloat16, weights=True, rsf=2.5, recompute=True, round_sf=False, tma=False) -bwd('bwd_e4m3_psum_tma_no_weights_recompute', 'e4m3', True, torch.bfloat16, recompute=True, packed=True, round_sf=True, tma=True) -print('REVIEW_COVERAGE_DONE') diff --git a/vmi_swiglu_plugin.py b/vmi_swiglu_plugin.py deleted file mode 100644 index 0786e269..00000000 --- a/vmi_swiglu_plugin.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Run the existing SwiGLU tests against the VMI kernels. - -Usage (from the repository root): - TK_USE_SWIGLU_VMI=1 python -m pytest tests/quant/test_swiglu_forward.py \ - tests/quant/test_swiglu_backward.py -q -p vmi_swiglu_plugin - -The dispatchers import the ASC getter lazily inside the public API, so -replacing the module attribute before the tests run is enough to exercise the -VMI implementation without modifying the repository. -""" - -import os - - -def pytest_configure(config): - if os.environ.get('TK_USE_SWIGLU_VMI') != '1': - return - - import tile_kernels.quant.swiglu_forward_asc as forward_asc - import tile_kernels.quant.swiglu_backward_asc as backward_asc - from tile_kernels.quant.swiglu_forward_vmi import get_swiglu_forward_kernel_vmi - from tile_kernels.quant.swiglu_backward_vmi import get_swiglu_backward_kernel_vmi - - forward_asc.get_swiglu_forward_kernel_asc = get_swiglu_forward_kernel_vmi - backward_asc.get_swiglu_backward_kernel_asc = get_swiglu_backward_kernel_vmi From 12fc4e1fa174f1f924efbded0eb66fdc7034deff Mon Sep 17 00:00:00 2001 From: kuri780 Date: Thu, 22 Oct 2026 10:47:20 +0000 Subject: [PATCH 7/7] quant: address SwiGLU VMI review comments - Drop redundant full masks on VMI ops that support unmasked forms; keep masks required by the API (vadds/vmuls/vmaxs/vmins/vshrs/vshls, vexpdif, vmula, vcmp/vsel, vcadd/vcmax, vgather and exact vdiv). - Forward amax: use the widest vector whose 4-vector tile stays full. 2048 -> 256-lane vcmax(group=8); 3584 -> 128-lane vcmax(group=4); clamp-count and non-quantized paths keep 64-lane vcmax(group=2). vcmax group is the output count, so 256-lane/group=4 would produce 64-channel groups, which is wrong for sf_block=(1,32). - ptoas 0.66 does not lower 256-lane full reductions (group=1); use grouped certificate/clamp reductions plus a scalar combine. - Document that tilelang PR #14 (codex/unify-asc-mte-pointer-lowering) does not apply to the pinned 5f105e99 snapshot; keep the uint8-view workaround for packed SF copies. Validated: level-0 48/48, review_adversarial_ab.py ALL_OK, hidden=384 64-lane fallback passes. Co-Authored-By: Claude Code --- tile_kernels/quant/swiglu_backward_vmi.py | 51 ++++--- tile_kernels/quant/swiglu_forward_vmi.py | 175 ++++++++++++++++------ 2 files changed, 156 insertions(+), 70 deletions(-) diff --git a/tile_kernels/quant/swiglu_backward_vmi.py b/tile_kernels/quant/swiglu_backward_vmi.py index 7eb8d02e..8eef93b9 100644 --- a/tile_kernels/quant/swiglu_backward_vmi.py +++ b/tile_kernels/quant/swiglu_backward_vmi.py @@ -83,7 +83,7 @@ def vdiv_exact_0ulp(a, b): not change residual ordering in this operand domain. """ full = V.create_mask(vec_size, size=vec_size) - z = V.vdiv(a, b, full) + z = V.vdiv(a, b) y = V.vneg(b) r = V.vmula(a, z, y, full) z_bits = V.vinterpret_cast(z, 'int32') @@ -116,7 +116,7 @@ def sigmoid_vdiv_exact_0ulp(a, b): full = V.create_mask(vec_size, size=vec_size) one_i32 = V.vbrc(T.int32(1), size=vec_size) zero_i32 = V.vbrc(T.int32(0), size=vec_size) - z = V.vdiv(a, b, full) + z = V.vdiv(a, b) neg_b = V.vneg(b) residual = V.vmula(a, z, neg_b, full) @@ -124,18 +124,18 @@ def sigmoid_vdiv_exact_0ulp(a, b): residual_bits = V.vinterpret_cast(residual, 'int32') sign_mask = V.vshrs(residual_bits, T.int32(31), full) use_bits_minus = V.vcmps(residual_bits, T.int32(0), full, 'lt') - delta = V.vor(sign_mask, one_i32, full) - adjacent = V.vinterpret_cast(V.vadd(V.vinterpret_cast(z, 'int32'), delta, full), 'float32') + delta = V.vor(sign_mask, one_i32) + adjacent = V.vinterpret_cast(V.vadd(V.vinterpret_cast(z, 'int32'), delta), 'float32') adjacent_residual = V.vmula(a, adjacent, neg_b, full) - residual_sum = V.vadd(residual, adjacent_residual, full) + residual_sum = V.vadd(residual, adjacent_residual) take_bits_plus = V.vcmps(residual_sum, T.float32(0.0), full, 'gt') is_nonzero = V.vcmps(z, T.float32(0.0), full, 'ne') take_plus_i = V.vsel(take_bits_plus, one_i32, zero_i32) use_minus_i = V.vsel(use_bits_minus, one_i32, zero_i32) nonzero_i = V.vsel(is_nonzero, one_i32, zero_i32) - take_adjacent_i = V.vand(V.vxor(take_plus_i, use_minus_i, full), nonzero_i, full) + take_adjacent_i = V.vand(V.vxor(take_plus_i, use_minus_i), nonzero_i) take_adjacent = V.vcmps(take_adjacent_i, T.int32(0), full, 'ne') return V.vsel(take_adjacent, adjacent, z) @@ -149,13 +149,9 @@ def get_scale_and_inverse(amax): exponent_carry = V.vbrc(T.uint32(0x001FFFFF), size=vec_size) inverse_bias = V.vbrc(T.uint32(0x83000000), size=vec_size) scale_bias = V.vbrc(T.uint32(0x04000000), size=vec_size) - exponent = V.vand( - V.vadd(V.vinterpret_cast(clamped_amax, 'uint32'), exponent_carry, full), - exponent_mask, - full, - ) - scale = V.vinterpret_cast(V.vsub(exponent, scale_bias, full), 'float32') - inverse = V.vinterpret_cast(V.vsub(inverse_bias, exponent, full), 'float32') + exponent = V.vand(V.vadd(V.vinterpret_cast(clamped_amax, 'uint32'), exponent_carry), exponent_mask) + scale = V.vinterpret_cast(V.vsub(exponent, scale_bias), 'float32') + inverse = V.vinterpret_cast(V.vsub(inverse_bias, exponent), 'float32') return scale, inverse quant_max = V.vbrc(T.float32(448.0), size=vec_size) return ( @@ -364,8 +360,8 @@ def swiglu_backward_kernel_vmi( y_scale = V.vbrc(V.vload(y_sf_ub[vector * 2], size=2, group=2, stride=1), size=vec_size, group=2) x_raw = V.vcvt(V.vload(x_ub[col], size=vec_size), 'float32') y_raw = V.vcvt(V.vload(y_ub[col], size=vec_size), 'float32') - x_value = V.vmul(x_raw, x_scale, full) - y_value = V.vmul(y_raw, y_scale, full) + x_value = V.vmul(x_raw, x_scale) + y_value = V.vmul(y_raw, y_scale) elif x_dtype == T.bfloat16: x_value = V.vcvt(V.vload(x_ub[col], size=vec_size), 'float32') y_value = V.vcvt(V.vload(y_ub[col], size=vec_size), 'float32') @@ -390,15 +386,15 @@ def swiglu_backward_kernel_vmi( sigmoid = sigmoid_vdiv_exact_0ulp(ones, sigmoid_denominator) if with_weight or do_recompute: # Keep a separate 0-ULP division for the required FP32 rounding. - act_value = V.vmul(vdiv_exact_0ulp(x_compute_value, sigmoid_denominator), y_compute_value, full) + act_value = V.vmul(vdiv_exact_0ulp(x_compute_value, sigmoid_denominator), y_compute_value) if with_weight: - weight_contribution = V.vmul(grad, act_value, full) + weight_contribution = V.vmul(grad, act_value) weight_value = V.vmuls(weight_contribution, valid_token_f32, full) if num_experts else weight_contribution - weight_acc = V.vadd(weight_acc, weight_value, full) - grad_ws = V.vmul(V.vmuls(grad, effective_weight, full), sigmoid, full) - inner = V.vadd(ones, V.vmul(x_compute_value, V.vsub(ones, sigmoid, full), full), full) - x_grad_unclamped = V.vmul(V.vmul(grad_ws, y_compute_value, full), inner, full) - y_grad_unclamped = V.vmul(grad_ws, x_compute_value, full) + weight_acc = V.vadd(weight_acc, weight_value) + grad_ws = V.vmul(V.vmuls(grad, effective_weight, full), sigmoid) + inner = V.vadd(ones, V.vmul(x_compute_value, V.vsub(ones, sigmoid))) + x_grad_unclamped = V.vmul(V.vmul(grad_ws, y_compute_value), inner) + y_grad_unclamped = V.vmul(grad_ws, x_compute_value) x_grad_value = V.vsel(is_clamped_x, zeros, x_grad_unclamped) if use_clamp else x_grad_unclamped y_grad_value = V.vsel(is_clamped_y, zeros, y_grad_unclamped) if use_clamp else y_grad_unclamped if not with_quant: @@ -455,8 +451,8 @@ def swiglu_backward_kernel_vmi( y_scale_value, y_inv = get_scale_and_inverse(V.vload(y_quant_scratch_ub[scale_offset], size=vec_size)) V.vstore(x_inv, x_quant_scratch_ub[scale_offset]) V.vstore(y_inv, y_quant_scratch_ub[scale_offset]) - x_scale_out = V.vmul(x_scale_value, valid_token_vec, full) if num_experts else x_scale_value - y_scale_out = V.vmul(y_scale_value, valid_token_vec, full) if num_experts else y_scale_value + x_scale_out = V.vmul(x_scale_value, valid_token_vec) if num_experts else x_scale_value + y_scale_out = V.vmul(y_scale_value, valid_token_vec) if num_experts else y_scale_value if out_config.use_packed_ue8m0: x_exp = V.vshrs(V.vinterpret_cast(x_scale_out, 'uint32'), T.uint32(23), full) y_exp = V.vshrs(V.vinterpret_cast(y_scale_out, 'uint32'), T.uint32(23), full) @@ -488,8 +484,8 @@ def swiglu_backward_kernel_vmi( y_inv_vec = V.vbrc(V.vload(y_quant_scratch_ub[vector * 2], size=2, group=2, stride=1), size=vec_size, group=2) x_grad_quant_value = V.vload(x_grad_f32_ub[col], size=vec_size) y_grad_quant_value = V.vload(y_grad_f32_ub[col], size=vec_size) - x_scaled_value = V.vmul(x_grad_quant_value, x_inv_vec, full) - y_scaled_value = V.vmul(y_grad_quant_value, y_inv_vec, full) + x_scaled_value = V.vmul(x_grad_quant_value, x_inv_vec) + y_scaled_value = V.vmul(y_grad_quant_value, y_inv_vec) # Explicit NOSAT, same bug A/B rationale. V.vstore(V.vcvt(x_grad_quant_value, 'bfloat16', saturate='NOSAT'), x_grad_ub[col]) V.vstore(V.vcvt(y_grad_quant_value, 'bfloat16', saturate='NOSAT'), y_grad_ub[col]) @@ -521,6 +517,9 @@ def swiglu_backward_kernel_vmi( if out_config.use_packed_ue8m0: # PTO's DMA requires matching source/destination element widths. # Write packed SF through uint8 views instead of ASC's uint16 view. + # tilelang PR #14 (codex/unify-asc-mte-pointer-lowering) removes + # this restriction, but it targets a newer pto-dev and does not + # apply to our pinned 5f105e99 snapshot; keep the byte views. if out_config.use_tma_aligned_col_major_sf: output_sf = T.Tensor((num_stored_groups * 2, num_expanded_tokens * 2), T.uint8, x_grad_fp8_sf.data) else: diff --git a/tile_kernels/quant/swiglu_forward_vmi.py b/tile_kernels/quant/swiglu_forward_vmi.py index 5c1c62ba..e4726711 100644 --- a/tile_kernels/quant/swiglu_forward_vmi.py +++ b/tile_kernels/quant/swiglu_forward_vmi.py @@ -43,7 +43,27 @@ def get_swiglu_forward_kernel_vmi( num_stages = 3 # Multiplication by a weight requires a larger certificate error bound. certificate_error = 5 if with_weight else 2 + # The amax path uses the widest logical vector whose outer tile (4 vectors) is + # always full. A partially-valid tile puts the amax accumulator behind a + # dynamic `vector_id < num_vecs` branch, and ptoas then sees two VMI layouts + # for the same 256/128-lane value on the scf.if result (contiguous vs + # deinterleaved). Non-quantized outputs keep the original 64-lane form, and + # clamp counting stays at 64 lanes too: its per-lane counters plus the + # certificate spill the register file at 256/128 lanes (measured 2.3x slower). vec_size = 64 + if out_config.with_sf and not count_clamp: + for candidate in (256, 128): + if hidden % candidate == 0 and (hidden // candidate) % 4 == 0: + vec_size = candidate + break + groups_per_vector = vec_size // 32 + # A full reduction over one physical part (64 lanes) is the largest form + # ptoas lowers. Wider vectors use a grouped reduction plus a scalar combine. + cert_groups = vec_size // 64 + # Scale conversion and the 64-value amax rows stay one vector wide. + scale_batch = 64 + # 64 amax values (one per 32 channels) cover 2048 channels per batch. + vectors_per_batch = scale_batch * 32 // vec_size num_vecs = hidden // vec_size hidden_aligned = align(hidden, 256) @@ -72,7 +92,7 @@ def vdiv_exact_0ulp(a, b): review_adversarial_ab.py does not claim to cover that branch. """ full = V.create_mask(vec_size, size=vec_size) - z = V.vdiv(a, b, full) + z = V.vdiv(a, b) y = V.vneg(b) r = V.vmula(a, z, y, full) z_bits = V.vinterpret_cast(z, 'int32') @@ -93,16 +113,16 @@ def vdiv_exact_0ulp(a, b): @T.macro def get_scale_and_inverse(amax): - full = V.create_mask(vec_size, size=vec_size) + full = V.create_mask(scale_batch, size=scale_batch) clamped_amax = V.vmaxs(amax, T.float32(out_config.clamp_min_value), full) scale_bits = V.vinterpret_cast(V.vmuls(clamped_amax, T.float32(1.0 / 448.0), full), 'uint32') exponent = V.vadds( - V.vshrs(V.vsub(scale_bits, V.vbrc(T.uint32(1), size=vec_size), full), T.uint32(23), full), + V.vshrs(V.vsub(scale_bits, V.vbrc(T.uint32(1), size=scale_batch)), T.uint32(23), full), T.uint32(1), full, ) inverse = V.vinterpret_cast( - V.vshls(V.vsub(V.vbrc(T.uint32(254), size=vec_size), exponent, full), T.uint32(23), full), + V.vshls(V.vsub(V.vbrc(T.uint32(254), size=scale_batch), exponent), T.uint32(23), full), 'float32', ) if out_config.use_packed_ue8m0: @@ -171,7 +191,9 @@ def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, else: certificate_min = None certificate_max = None - if out_config.with_sf: + if out_config.with_sf and vec_size != 64: + # Batched reduction wins on the wide forms; the 64-lane form is + # cheaper to reduce in place and avoids four pinned registers. amax_acc = V.alloc_local((4,), V.vreg(vec_size, T.float32)) for i in T.unroll(4, explicit=True): amax_acc[i] = zeros @@ -195,12 +217,12 @@ def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, mul_r = val_r denominator = V.vadds(V.vexpdif(zeros, silu_l, full), T.float32(1.0), full) if precision == 'ftz_true': - quotient = V.vdiv(silu_l, denominator, full) + quotient = V.vdiv(silu_l, denominator) elif precision == 'exact': quotient = vdiv_exact_0ulp(silu_l, denominator) else: raise ValueError(f"unsupported vdiv precision {precision!r}; expected 'ftz_true' or 'exact'") - val = V.vmul(quotient, mul_r, full) + val = V.vmul(quotient, mul_r) out_val = V.vmuls(val, weight, full) if with_weight else val if out_config.with_sf: @@ -226,19 +248,24 @@ def compute_swiglu(xl_ub, xr_ub, values_ub, out_ub, weight, precision, amax_ub, # Merge-mode accumulation: add 1 on active lanes, keep the # previous value elsewhere (PTO vadds only supports zeroing). one = V.vbrc(T.float32(1.0), size=vec_size) - acc[0] = V.vadd(acc[0], V.vsel(clamped_mask_0, one, zeros), full) - acc[1] = V.vadd(acc[1], V.vsel(clamped_mask_1, one, zeros), full) - acc[2] = V.vadd(acc[2], V.vsel(clamped_mask_2, one, zeros), full) + acc[0] = V.vadd(acc[0], V.vsel(clamped_mask_0, one, zeros)) + acc[1] = V.vadd(acc[1], V.vsel(clamped_mask_1, one, zeros)) + acc[2] = V.vadd(acc[2], V.vsel(clamped_mask_2, one, zeros)) if out_config.with_sf: - amax_acc[sub] = V.vabs(out_val) - if out_config.with_sf: - # 256-element tile amax: two 32-channel groups per 64-lane vector. + if vec_size == 64: + amax = V.vcmax(V.vabs(out_val), full, group=groups_per_vector) + V.vstore(amax, amax_ub[tile_id * 4 * groups_per_vector + sub * groups_per_vector]) + else: + amax_acc[sub] = V.vabs(out_val) + if out_config.with_sf and vec_size != 64: + # One 32-channel scale group per 32 lanes: a 256-lane vector + # yields eight groups in one vcmax, a 128-lane vector four. for sub in T.unroll(4, explicit=True): - amax = V.vcmax(amax_acc[sub], full, group=2) - V.vstore(amax, amax_ub[tile_id * 8 + sub * 2]) + amax = V.vcmax(amax_acc[sub], full, group=groups_per_vector) + V.vstore(amax, amax_ub[tile_id * 4 * groups_per_vector + sub * groups_per_vector]) if count_clamp: for i in T.unroll(3, explicit=True): - V.vstore(V.vadd(V.vload(pacc_buf[i, 0], size=vec_size), acc[i], full), pacc_buf[i, 0]) + V.vstore(V.vadd(V.vload(pacc_buf[i, 0], size=vec_size), acc[i]), pacc_buf[i, 0]) return certificate_min, certificate_max @T.macro @@ -246,9 +273,10 @@ def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offs # Get inverse scale and quantize FP32 SwiGLU result to FP8. with T.SimdVF(): full = V.create_mask(vec_size, size=vec_size) - for group_batch in T.serial(T.ceildiv(num_groups, 64)): - group = group_batch * 64 - amax = V.vload(amax_ub[group], size=vec_size) + full_scale = V.create_mask(scale_batch, size=scale_batch) + for group_batch in T.serial(T.ceildiv(num_groups, scale_batch)): + group = group_batch * scale_batch + amax = V.vload(amax_ub[group], size=scale_batch) if out_config.round_sf: scales_or_exponents, inverses = get_scale_and_inverse(amax) if out_config.use_packed_ue8m0: @@ -262,21 +290,25 @@ def quantize_f32(values_ub, amax_ub, inverse_ub, sf_dense_ub, out_ub, token_offs V.vstore(scales_or_exponents, sf_dense_ub[token_offset, group]) V.vstore(inverses, inverse_ub[group]) else: - clamped_amax = V.vmaxs(amax, T.float32(out_config.clamp_min_value), full) - quant_max = V.vbrc(T.float32(448.0), size=vec_size) - V.vstore(V.vdiv(clamped_amax, quant_max, full, precision='exact'), sf_dense_ub[token_offset, group]) - V.vstore(V.vdiv(quant_max, clamped_amax, full, precision='exact'), inverse_ub[group]) + clamped_amax = V.vmaxs(amax, T.float32(out_config.clamp_min_value), full_scale) + quant_max = V.vbrc(T.float32(448.0), size=scale_batch) + V.vstore(V.vdiv(clamped_amax, quant_max, full_scale, precision='exact'), sf_dense_ub[token_offset, group]) + V.vstore(V.vdiv(quant_max, clamped_amax, full_scale, precision='exact'), inverse_ub[group]) S.mem_bar('VST_VLD') - group_lane = V.vshrs(V.vci(T.int32(0), size=vec_size), T.int32(5), full) - for group_batch in T.unroll(T.ceildiv(num_groups, 64), explicit=True): - group_base = group_batch * 64 - vector_base = group_batch * 32 - inverses = V.vload(inverse_ub[group_base], size=vec_size) - for vector in T.serial(T.min(32, num_vecs - vector_base)): + for group_batch in T.unroll(T.ceildiv(num_groups, scale_batch), explicit=True): + group_base = group_batch * scale_batch + vector_base = group_batch * vectors_per_batch + for vector in T.serial(T.min(vectors_per_batch, num_vecs - vector_base)): col = (vector_base + vector) * vec_size - inverse_indices = V.vadds(group_lane, T.int32(vector * 2), full) - expanded_inverse = V.vselr(inverses, inverse_indices) + # Broadcast each 32-channel inverse to its 32 lanes. + inverse_slots = V.vload( + inverse_ub[group_base + vector * groups_per_vector], + size=groups_per_vector, + group=groups_per_vector, + stride=1, + ) + expanded_inverse = V.vbrc(inverse_slots, size=vec_size, group=groups_per_vector) # FP8 output saturate: explicit NOSAT. ASC's RS_ENABLE is a # no-op for f32->e4m3 here (bug A: sat=True and sat=False # produce the same bytes), while VMI SAT maps NaN to 0x00 @@ -344,6 +376,8 @@ def swiglu_forward_kernel( if count_clamp: acc_ub = T.alloc_shared((4,), T.int32) pacc_ub = T.alloc_shared((3, vec_size), T.float32) # per lane accumulator pacc_ub. use vcadd at the last moment + if cert_groups > 1: + pacc_partial_ub = T.alloc_shared((3, cert_groups), T.int32) for i in T.serial(4): acc_ub[i] = 0 @@ -377,7 +411,8 @@ def swiglu_forward_kernel( if out_config.with_sf: amax_ub = T.alloc_shared((align(num_groups, 64),), T.float32) if use_fast_certificate: - certificate_ub = T.alloc_shared((token_group, 2), T.uint32) + # cert_groups partial minima followed by cert_groups partial maxima. + certificate_ub = T.alloc_shared((token_group, 2 * cert_groups), T.uint32) inverse_ub = T.alloc_shared((align(num_groups, 64),), T.float32) sf_dense_ub = T.alloc_shared( (token_group, align(num_groups, 64)), @@ -433,8 +468,25 @@ def swiglu_forward_kernel( xl_ub, xr_ub, values_ub, out_ub, weight, vdiv_precision, amax_ub, use_fast_certificate, pacc_ub, count_clamp ) if use_fast_certificate: - V.vstore(V.vcmin(cert_min, full), certificate_ub[token_offset, 0], group=1, stride=1) - V.vstore(V.vcmax(cert_max, full), certificate_ub[token_offset, 1], group=1, stride=1) + if cert_groups == 1: + V.vstore(V.vcmin(cert_min, full), certificate_ub[token_offset, 0], group=1, stride=1) + V.vstore(V.vcmax(cert_max, full), certificate_ub[token_offset, 1], group=1, stride=1) + else: + # A 256-lane full reduction is not lowered by + # ptoas; keep one partial per physical part and + # combine the scalars when checking the token. + V.vstore( + V.vcmin(cert_min, full, group=cert_groups), + certificate_ub[token_offset, 0], + group=cert_groups, + stride=1, + ) + V.vstore( + V.vcmax(cert_max, full, group=cert_groups), + certificate_ub[token_offset, cert_groups], + group=cert_groups, + stride=1, + ) finish_token( values_ub, out_ub, amax_ub, inverse_ub, sf_dense_ub, out, out_sf, token_offset, token_id, acc_ub, count_clamp ) @@ -443,13 +495,28 @@ def swiglu_forward_kernel( for fallback_offset in T.serial(token_group): token_id = T.alloc_var(T.int32, init=token_base + fallback_offset) valid = is_valid_token(token_id, expert_id if has_psum else 0, psum_ub if has_psum else None) - unsafe = T.alloc_var( - T.bool, - init=( + if cert_groups == 1: + unsafe_init = ( certificate_ub[fallback_offset, 0] <= certificate_error or certificate_ub[fallback_offset, 1] >= 0x80000 - certificate_error - ), - ) + ) + elif cert_groups == 2: + cert_min_all = T.min(certificate_ub[fallback_offset, 0], certificate_ub[fallback_offset, 1]) + cert_max_all = T.max(certificate_ub[fallback_offset, 2], certificate_ub[fallback_offset, 3]) + unsafe_init = cert_min_all <= certificate_error or cert_max_all >= 0x80000 - certificate_error + else: + # cert_groups == 4: four partial minima, then four + # partial maxima, stored by the main path. + cert_min_all = T.min( + T.min(certificate_ub[fallback_offset, 0], certificate_ub[fallback_offset, 1]), + T.min(certificate_ub[fallback_offset, 2], certificate_ub[fallback_offset, 3]), + ) + cert_max_all = T.max( + T.max(certificate_ub[fallback_offset, 4], certificate_ub[fallback_offset, 5]), + T.max(certificate_ub[fallback_offset, 6], certificate_ub[fallback_offset, 7]), + ) + unsafe_init = cert_min_all <= certificate_error or cert_max_all >= 0x80000 - certificate_error + unsafe = T.alloc_var(T.bool, init=unsafe_init) if valid and unsafe: weight = load_inputs(x, xl_fb_ub, xr_fb_ub, topk_weights, routed_scaling_factor, token_id) with T.SimdVF(): @@ -464,6 +531,9 @@ def swiglu_forward_kernel( transpose_packed_scales(sf_dense_ub, sf_ub, transpose_indices_ub) # PTO's DMA requires matching source/destination element widths. # Copy through uint8 views instead of the ASC uint16 view. + # tilelang PR #14 (codex/unify-asc-mte-pointer-lowering) removes + # this restriction, but it targets a newer pto-dev and does not + # apply to our pinned 5f105e99 snapshot; keep the byte views. output_sf = T.Tensor((sf_shape[0], sf_shape[1]), T.uint8, out_sf.data) sf_ub_u8 = T.Tensor((num_groups // pack_factor, token_group * 2), T.uint8, sf_ub.data) valid_token_end = psum_ub[expert_id] if has_psum else num_expanded_tokens @@ -479,13 +549,30 @@ def swiglu_forward_kernel( for i in T.unroll(3, explicit=True): # These are exact integer counts (< 2^24), so the # reassociation allowed by PTODSL cannot change them. - reduced = V.vcadd(V.vload(pacc_ub[i, 0], size=vec_size), full, reassoc=True) - # fp->int narrow: the count is < 2^24, so SAT/NOSAT are - # identical. Spell out NOSAT instead of relying on - # PTODSL's implicit SAT default. - V.vstore(V.vcvt(reduced, 'int32', saturate='NOSAT'), acc_ub[i], group=1, stride=1) + if cert_groups == 1: + reduced = V.vcadd(V.vload(pacc_ub[i, 0], size=vec_size), full, reassoc=True) + # fp->int narrow: the count is < 2^24, so SAT/NOSAT are + # identical. Spell out NOSAT instead of relying on + # PTODSL's implicit SAT default. + V.vstore(V.vcvt(reduced, 'int32', saturate='NOSAT'), acc_ub[i], group=1, stride=1) + else: + # ptoas does not lower a full 256-lane reduction; keep + # one partial per physical part and sum them scalar-side. + reduced = V.vcadd(V.vload(pacc_ub[i, 0], size=vec_size), full, group=cert_groups, reassoc=True) + V.vstore( + V.vcvt(reduced, 'int32', saturate='NOSAT'), + pacc_partial_ub[i, 0], + group=cert_groups, + stride=1, + ) with T.SimtVF(threads=1): + if cert_groups == 2: + for i in T.serial(3): + acc_ub[i] = pacc_partial_ub[i, 0] + pacc_partial_ub[i, 1] + elif cert_groups == 4: + for i in T.serial(3): + acc_ub[i] = pacc_partial_ub[i, 0] + pacc_partial_ub[i, 1] + pacc_partial_ub[i, 2] + pacc_partial_ub[i, 3] for i in T.serial(4): T.atomic_add(clamped_count[i], T.int64(acc_ub[i]))