diff --git a/tile_kernels/quant/swiglu_backward_kernel.py b/tile_kernels/quant/swiglu_backward_kernel.py index cc65f9e..f2edc00 100644 --- a/tile_kernels/quant/swiglu_backward_kernel.py +++ b/tile_kernels/quant/swiglu_backward_kernel.py @@ -107,9 +107,9 @@ def swiglu_backward_and_per_token_cast_impl( # Get kernel implement if is_ascend(): - from tile_kernels.quant.swiglu_backward_asc import get_swiglu_backward_kernel_asc + from tile_kernels.quant.swiglu_backward_vmi import get_swiglu_backward_kernel_vmi - kernel = get_swiglu_backward_kernel_asc( + kernel = get_swiglu_backward_kernel_vmi( hidden, num_experts, get_token_alignment(), diff --git a/tile_kernels/quant/swiglu_backward_vmi.py b/tile_kernels/quant/swiglu_backward_vmi.py new file mode 100644 index 0000000..b885af1 --- /dev/null +++ b/tile_kernels/quant/swiglu_backward_vmi.py @@ -0,0 +1,474 @@ +import tilelang +from tilelang.ascend import language as T +from tilelang.ascend.language import simd as S +from tilelang.ascend.language import vmi as V + +from tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape +from tile_kernels.utils import align + + +@tilelang.jit(target='pto') +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 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), 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 ( + 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 expanded input scales, amax, and inverse scales. + 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: + valid_token_f32 = T.cast(expert_id < num_experts and token_id < valid_token_end, T.float32) + + 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_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) + + # 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: + 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) + 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') + 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.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 + + sigmoid_denominator = V.vadds(V.vexpdif(zeros, x_compute_value, full), T.float32(1.0), full) + sigmoid = V.vdiv(ones, sigmoid_denominator, full, precision='exact') + if with_weight or do_recompute: + # Keep a separate 0-ULP division for the required FP32 rounding. + act_value = V.vmul(V.vdiv(x_compute_value, sigmoid_denominator, full, precision='exact'), y_compute_value) + if with_weight: + 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) + 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: + if out_dtype == T.bfloat16: + # 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]) + 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: + # 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]) + + 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 = ( + 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) + # 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)) + 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) 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) + # 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], + ) + V.vstore( + V.vcvt(y_exp, 'uint8', saturate='NOSAT'), + 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) + 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]) + 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 + # 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: upstream + # (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]) + 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: + if out_config.use_tma_aligned_col_major_sf: + output_sf = T.Tensor((num_stored_groups * 2, num_expanded_tokens), T.uint16, x_grad_fp8_sf.data) + else: + output_sf = T.Tensor((num_expanded_tokens, num_stored_groups * 2), T.uint16, 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: + 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: + 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_kernel.py b/tile_kernels/quant/swiglu_forward_kernel.py index a602be8..9285d17 100644 --- a/tile_kernels/quant/swiglu_forward_kernel.py +++ b/tile_kernels/quant/swiglu_forward_kernel.py @@ -46,9 +46,9 @@ def swiglu_forward_and_per_token_cast_impl( # Get kernel implement if is_ascend(): - from tile_kernels.quant.swiglu_forward_asc import get_swiglu_forward_kernel_asc + from tile_kernels.quant.swiglu_forward_vmi import get_swiglu_forward_kernel_vmi - kernel = get_swiglu_forward_kernel_asc( + kernel = get_swiglu_forward_kernel_vmi( hidden, num_experts, get_token_alignment(), diff --git a/tile_kernels/quant/swiglu_forward_vmi.py b/tile_kernels/quant/swiglu_forward_vmi.py new file mode 100644 index 0000000..bd96b8b --- /dev/null +++ b/tile_kernels/quant/swiglu_forward_vmi.py @@ -0,0 +1,442 @@ +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 tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape +from tile_kernels.utils import align + + +@tilelang.jit(target='pto') +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 get_scale_and_inverse(amax): + full = V.create_mask(64, size=64) + 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=64)), T.uint32(23), full), + T.uint32(1), + full, + ) + inverse = V.vinterpret_cast( + V.vshls(V.vsub(V.vbrc(T.uint32(254), size=64), exponent), 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) + # Read byte-owned SF storage without introducing a typed pointer alias. + byte_indices = V.vmul(source_indices, V.vbrc(T.uint16(2), size=128)) + + if num_full_vectors > 0: + for output_vector in T.unroll(num_full_vectors, explicit=True): + group_base = output_vector * groups_per_vector + low = V.vgather(src[0, group_base * 2], byte_indices, full_mask) + high = V.vgather(src[0, group_base * 2 + 1], byte_indices, full_mask) + values = V.vor(low, V.vshls(high, T.uint16(8), 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) + low = V.vgather(src[0, tail_group_base * 2], byte_indices, tail_mask) + high = V.vgather(src[0, tail_group_base * 2 + 1], byte_indices, tail_mask) + values = V.vor(low, V.vshls(high, T.uint16(8), 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. + # Logical compute width is local to SIMD; ASC buffers and scheduler + # retain their original 64-lane geometry and certificate storage. + logical_size = (256 if hidden % 256 == 0 else 128) if out_config.with_sf and not count_clamp else 64 + logical_vectors = hidden // logical_size + compute_unroll = 1 if logical_size > 64 else 4 + full = V.create_mask(logical_size, size=logical_size) + zeros = V.vbrc(T.float32(0.0), size=logical_size) + if count_clamp: + if acc is None: + acc = V.alloc_local((3,), V.vreg(logical_size, T.float32)) + 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=logical_size) + certificate_max = V.alloc_var(T.uint32, size=logical_size) + certificate_min = V.vbrc(T.uint32(0xFFFFFFFF), size=logical_size) + certificate_max = V.vbrc(T.uint32(0), size=logical_size) + else: + certificate_min = None + certificate_max = None + for tile_id in T.serial(T.ceildiv(logical_vectors, compute_unroll)): + for sub in T.unroll(compute_unroll, explicit=True): + vector_id = tile_id * compute_unroll + sub + if vector_id < logical_vectors: + col = vector_id * logical_size + if in_dtype == T.bfloat16: + val_l = V.vcvt(V.vload(xl_ub[col], size=logical_size), 'float32') + val_r = V.vcvt(V.vload(xr_ub[col], size=logical_size), 'float32') + else: + val_l = V.vload(xl_ub[col], size=logical_size) + val_r = V.vload(xr_ub[col], size=logical_size) + if has_clamp: + 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 + denominator = V.vadds(V.vexpdif(zeros, silu_l, full), T.float32(1.0), full) + if precision == 'ftz_true': + quotient = V.vdiv(silu_l, denominator) + elif precision == 'exact': + quotient = V.vdiv(silu_l, denominator, full, precision='exact') + else: + raise ValueError(f"unsupported vdiv precision {precision!r}; expected 'ftz_true' or 'exact'") + val = V.vmul(quotient, mul_r) + 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', saturate='NOSAT'), 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=logical_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=logical_size) + 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 = V.vcmax(V.vabs(out_val), full, group=logical_size // 32) + V.vstore(amax, amax_ub[vector_id * (logical_size // 32)]) + if count_clamp: + for i in T.unroll(3, explicit=True): + V.vstore(V.vadd(V.vload(pacc_buf[i, 0], size=logical_size), acc[i]), pacc_buf[i, 0]) + if with_certificate and logical_size != 64: + # The public helper contract remains two 64-lane values, allowing + # ASC's outer certificate buffers and scalar fallback to stay exact. + certificate_min_out = V.vbrc(V.vcmin(certificate_min, full), size=64, group=1) + certificate_max_out = V.vbrc(V.vcmax(certificate_max, full), size=64, group=1) + else: + certificate_min_out = certificate_min + certificate_max_out = certificate_max + return certificate_min_out, certificate_max_out + + @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) + full_scale = V.create_mask(64, size=64) + for group_batch in T.serial(T.ceildiv(num_groups, 64)): + group = group_batch * 64 + amax = V.vload(amax_ub[group], size=64) + if out_config.round_sf: + scales_or_exponents, inverses = get_scale_and_inverse(amax) + if out_config.use_packed_ue8m0: + # 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]) + V.vstore(inverses, inverse_ub[group]) + else: + clamped_amax = V.vmaxs(amax, T.float32(out_config.clamp_min_value), full_scale) + quant_max = V.vbrc(T.float32(448.0), size=64) + 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') + for group_batch in T.unroll(T.ceildiv(num_groups, 64), explicit=True): + group_base = group_batch * 64 + vector_base = group_batch * 32 + for vector in T.serial(T.min(32, num_vecs - vector_base)): + col = (vector_base + vector) * vec_size + # Broadcast each 32-channel inverse to its 32 lanes. + inverse_slots = V.vload( + inverse_ub[group_base + vector * 2], + size=2, + group=2, + stride=1, + ) + expanded_inverse = V.vbrc(inverse_slots, size=vec_size, group=2) + quantized = V.vcvt( + V.vmul(V.vload(values_ub[col], size=vec_size), expanded_inverse), + 'float8_e4m3fn', + rounding='R', + saturate='NOSAT', + ) + 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(): + 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, V.create_mask(vec_size, size=vec_size)), certificate_ub[token_offset, 0]) + V.vstore(V.vcmax(cert_max, V.create_mask(vec_size, size=vec_size)), certificate_ub[token_offset, 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) + output_sf = T.Tensor((sf_shape[0], sf_shape[1] // pack_factor), T.uint16, out_sf.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[:, 0:valid_tokens], output_sf[:, token_base : token_base + valid_tokens]) + + if count_clamp: + with T.SimdVF(): + for i in T.unroll(3, explicit=True): + V.vstore(V.vcvt(V.vcadd(V.vload(pacc_ub[i, 0], size=vec_size), V.create_mask(vec_size, size=vec_size), reassoc=True), 'int32', saturate='NOSAT'), acc_ub[i]) + + 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