Repository navigation
Conversation
Take over PR PTO-ISA#20, remove operator-local soft division, restore the ASC shell and document validation and remaining performance gaps. Co-authored-by: kuri780 <kuri780@users.noreply.github.com>
Author
Forward：ASC → VMI 完整源码 diff对应 PR head --- a/tile_kernels/quant/swiglu_forward_asc.py
+++ b/tile_kernels/quant/swiglu_forward_vmi.py
@@ -3,13 +3,14 @@
import tilelang
from tilelang.ascend import language as T
from tilelang.ascend.language import simd as S
+from tilelang.ascend.language import vmi as V
from tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape
from tile_kernels.utils import align
-@tilelang.jit
-def get_swiglu_forward_kernel_asc(
+@tilelang.jit(target='pto')
+def get_swiglu_forward_kernel_vmi(
hidden: int,
num_experts: int,
alignment: int,
@@ -54,24 +55,38 @@
@T.macro
def get_scale_and_inverse(amax):
- clamped_amax = S.vmaxs(amax, out_config.clamp_min_value)
- scale_bits = T.reinterpret(S.vmuls(clamped_amax, 1.0 / 448.0), 'uint32x64')
- exponent = S.vadds(S.vshrs(S.vsub(scale_bits, S.vdup(1, T.uint32)), 23), 1)
- inverse = T.reinterpret(S.vshls(S.vsub(S.vdup(254, T.uint32), exponent), 23), 'float32x64')
+ 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 T.reinterpret(S.vshls(exponent, 23), 'float32x64'), 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 = S.vci(0, T.int16)
- output_row = S.vand(T.reinterpret(lane, 'uint16x128'), S.vdup(token_group - 1, T.uint16))
- group_in_vector = T.reinterpret(S.vshrs(lane, int(math.log2(token_group))), 'uint16x128')
- source_indices = S.vadd(S.vmuls(output_row, align(num_groups, 64) // pack_factor), group_in_vector)
- S.vsts(indices[0], source_indices, dist='NORM_B16', extent=128)
+ 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):
@@ -81,134 +96,166 @@
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 = S.vld(indices[0])
- src_u16 = T.Tensor((token_group, align(num_groups, 64) // pack_factor), T.uint16, src.data)
+ 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
- values = S.vgather2(src_u16[0, group_base], source_indices)
- S.vsts(dst[group_base, 0], values, dist='NORM_B16', extent=128)
+ 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 = S.pset(16, f'PAT_VL{num_tail_values}')
- values = S.vgather2(src_u16[0, tail_group_base], source_indices, tail_mask)
- S.vsts(dst[tail_group_base, 0], values, tail_mask, dist='NORM_B16', extent=num_tail_values)
+ 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.
- zeros = S.vdup(0.0, T.float32)
+ # 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 = S.alloc_local((3,), T.float32)
+ 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:
- certificate_min = S.alloc_var(T.uint32)
- certificate_max = S.alloc_var(T.uint32)
- certificate_min = S.vdup(0xFFFFFFFF, T.uint32)
- certificate_max = S.vdup(0, T.uint32)
+ # 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
- if out_config.with_sf:
- amax_store_mask = S.pset(32, 'PAT_VL8')
- amax_acc = S.alloc_local((4,), T.float32)
- for i in T.unroll(4, explicit=True):
- amax_acc[i] = S.vdup(0.0, T.float32)
-
- 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
+ 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 = S.vcvt(S.vld(xl_ub[col], dist='UNPK_B16'), T.float32, part=0)
- val_r = S.vcvt(S.vld(xr_ub[col], dist='UNPK_B16'), T.float32, part=0)
+ 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 = S.vld(xl_ub[col])
- val_r = S.vld(xr_ub[col])
+ 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 = S.vmins(val_l, clamp_value)
- mul_r = S.vmaxs(S.vmins(val_r, clamp_value), -clamp_value)
+ 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 = S.vadds(S.vexpdif(zeros, silu_l), 1.0)
- quotient = S.vdiv(silu_l, denominator, precision=precision)
- val = S.vmul(quotient, mul_r)
- out_val = S.vmuls(val, weight) if with_weight else val
+ 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:
- S.vsts(values_ub[col], out_val)
+ V.vstore(out_val, values_ub[col])
elif out_config.dtype == T.bfloat16:
- S.vsts(out_ub[col], S.vcvt(out_val, T.bfloat16), dist='PK_B32')
+ V.vstore(V.vcvt(out_val, 'bfloat16', saturate='NOSAT'), out_ub[col])
else:
- S.vsts(out_ub[col], out_val)
+ 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 = S.vand(T.reinterpret(out_val, 'uint32x64'), S.vdup(0x7FFFF, T.uint32))
- certificate_min = S.vmin(certificate_min, metric)
- certificate_max = S.vmax(certificate_max, metric)
+ 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 = S.vcmps(val_l, clamp_value, op='gt')
- clamped_mask_1 = S.vcmps(val_r, clamp_value, op='gt')
- clamped_mask_2 = S.vcmps(val_r, -clamp_value, op='lt')
- acc[0] = S.vadds(acc[0], 1.0, clamped_mask_0, mode='MODE_MERGING')
- acc[1] = S.vadds(acc[1], 1.0, clamped_mask_1, mode='MODE_MERGING')
- acc[2] = S.vadds(acc[2], 1.0, clamped_mask_2, mode='MODE_MERGING')
+ 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_acc[sub] = S.vabs(out_val)
- if out_config.with_sf:
- # 256-element tile amax reduction
- low01, high01 = S.vdintlv(amax_acc[0], amax_acc[1])
- low23, high23 = S.vdintlv(amax_acc[2], amax_acc[3])
- compressed01 = S.vmax(low01, high01)
- compressed23 = S.vmax(low23, high23)
- low, high = S.vdintlv(compressed01, compressed23)
- amax = S.vcgmax(S.vmax(low, high))
- S.vsts(amax_ub[tile_id * 8], amax, amax_store_mask, dist='NORM_B32', extent=8)
+ 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):
- S.vsts(pacc_buf[i, 0], S.vadd(S.vld(pacc_buf[i, 0]), acc[i]))
- return certificate_min, certificate_max
+ 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 = S.vld(amax_ub[group])
+ 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:
- S.vsts(sf_dense_ub[token_offset, group], T.reinterpret(scales_or_exponents, 'uint8x256'), dist='PK4_B32')
+ # 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:
- S.vsts(sf_dense_ub[token_offset, group], scales_or_exponents)
- S.vsts(inverse_ub[group], inverses)
+ V.vstore(scales_or_exponents, sf_dense_ub[token_offset, group])
+ V.vstore(inverses, inverse_ub[group])
else:
- clamped_amax = S.vmaxs(amax, out_config.clamp_min_value)
- quant_max = S.vdup(448.0, T.float32)
- S.vsts(sf_dense_ub[token_offset, group], S.vdiv(clamped_amax, quant_max))
- S.vsts(inverse_ub[group], S.vdiv(quant_max, clamped_amax))
+ 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')
- group_lane = S.vshrs(S.vci(0, T.int32), 5)
for group_batch in T.unroll(T.ceildiv(num_groups, 64), explicit=True):
group_base = group_batch * 64
vector_base = group_batch * 32
- inverses = S.vld(inverse_ub[group_base])
for vector in T.serial(T.min(32, num_vecs - vector_base)):
col = (vector_base + vector) * vec_size
- inverse_indices = S.vadds(group_lane, vector * 2)
- expanded_inverse = S.vselr(inverses, inverse_indices)
- quantized = S.vcvt(S.vmul(S.vld(values_ub[col]), expanded_inverse), T.float8_e4m3fn)
- S.vsts(out_ub[col], quantized, dist='PK4_B32')
+ # 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):
@@ -349,8 +396,8 @@
xl_ub, xr_ub, values_ub, out_ub, weight, vdiv_precision, amax_ub, use_fast_certificate, pacc_ub, count_clamp
)
if use_fast_certificate:
- S.vsts(certificate_ub[token_offset, 0], S.vcmin(cert_min), dist='ONEPT_B32')
- S.vsts(certificate_ub[token_offset, 1], S.vcmax(cert_max), dist='ONEPT_B32')
+ 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
)
@@ -386,7 +433,7 @@
if count_clamp:
with T.SimdVF():
for i in T.unroll(3, explicit=True):
- S.vsts(acc_ub[i], S.vcvt(S.vcadd(S.vld(pacc_ub[i, 0])), T.int32), dist='ONEPT_B32')
+ 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): |
Author
Backward：ASC → VMI 完整源码 diff对应 PR head --- a/tile_kernels/quant/swiglu_backward_asc.py
+++ b/tile_kernels/quant/swiglu_backward_vmi.py
@@ -1,13 +1,14 @@
import tilelang
from tilelang.ascend import language as T
from tilelang.ascend.language import simd as S
+from tilelang.ascend.language import vmi as V
from tile_kernels.quant.common import CastOutputConfig, get_packed_ue8m0_pack_factor, get_sf_shape
from tile_kernels.utils import align
-@tilelang.jit
-def get_swiglu_backward_kernel_asc(
+@tilelang.jit(target='pto')
+def get_swiglu_backward_kernel_vmi(
hidden: int,
num_experts: int,
alignment: int,
@@ -68,84 +69,28 @@
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 when the hardware quotient is finite; preserve signed zero."""
- z = S.vdiv(a, b, precision='ftz_true')
- y = S.vneg(b)
- r = S.alloc_var(T.float32)
- r = a
- S.vmula(r, z, y)
- z_pre = T.reinterpret(S.vadds(T.reinterpret(z, 'int32x64'), -1), 'float32x64')
- z_next = T.reinterpret(S.vadds(T.reinterpret(z, 'int32x64'), 1), 'float32x64')
- r_pre = S.alloc_var(T.float32)
- r_pre = a
- S.vmula(r_pre, z_pre, y)
- r_next = S.alloc_var(T.float32)
- r_next = a
- S.vmula(r_next, z_next, y)
- r = S.vabs(r)
- r_pre = S.vabs(r_pre)
- r_next = S.vabs(r_next)
- take_pre = S.vcmp(r_pre, r, op='le')
- best_r = S.vsel(r_pre, r, take_pre)
- best_z = S.vsel(z_pre, z, take_pre)
- take_next = S.vcmp(r_next, best_r, op='lt')
- corrected = S.vsel(z_next, best_z, take_next)
- is_zero = S.vcmps(z, 0.0, op='eq')
- return S.vsel(z, corrected, is_zero)
-
- @T.macro
- def sigmoid_vdiv_0ulp(a, b):
- """Correct a / b for sigmoid."""
- z = S.vdiv(a, b, precision='ftz_true')
- neg_b = S.vneg(b)
-
- residual = S.alloc_var(T.float32)
- residual = a
- S.vmula(residual, z, neg_b)
-
- residual_bits = T.reinterpret(residual, 'int32x64')
- sign_mask = S.vshrs(residual_bits, 31)
- use_bits_minus = S.vcmps(residual_bits, 0, op='lt')
- delta = S.vor(sign_mask, S.vdup(1, T.int32))
- adjacent = T.reinterpret(S.vadd(T.reinterpret(z, 'int32x64'), delta), 'float32x64')
-
- adjacent_residual = S.alloc_var(T.float32)
- adjacent_residual = a
- S.vmula(adjacent_residual, adjacent, neg_b)
- residual_sum = S.vadd(residual, adjacent_residual)
- take_bits_plus = S.vcmps(residual_sum, 0.0, op='gt')
- is_nonzero = S.vcmps(z, 0.0, op='ne')
- take_adjacent = S.pxor(take_bits_plus, use_bits_minus, is_nonzero)
- return S.vsel(adjacent, z, take_adjacent)
+
@T.macro
def get_scale_and_inverse(amax):
# Compute one vector of output scales and inverse scales from per-group amax.
- clamped_amax = S.vmaxs(amax, out_config.clamp_min_value)
+ 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 = S.vdup(0x7F800000, T.uint32)
- exponent_carry = S.vdup(0x001FFFFF, T.uint32)
- inverse_bias = S.vdup(0x83000000, T.uint32)
- scale_bias = S.vdup(0x04000000, T.uint32)
- exponent = S.vand(S.vadd(T.reinterpret(clamped_amax, 'uint32x64'), exponent_carry), exponent_mask)
- scale = T.reinterpret(S.vsub(exponent, scale_bias), 'float32x64')
- inverse = T.reinterpret(S.vsub(inverse_bias, exponent), 'float32x64')
+ 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 = S.vdup(448.0, T.float32)
- return S.vdiv(clamped_amax, quant_max), S.vdiv(quant_max, clamped_amax)
-
- @T.macro
- def store_fourfold_expanded(dst, offset, value):
- # Repeat each compact FP32 value four times for subsequent E2B_B32 loads.
- value_2_lo, value_2_hi = S.vintlv(value, value)
- value_4_0, value_4_1 = S.vintlv(value_2_lo, value_2_lo)
- value_4_2, value_4_3 = S.vintlv(value_2_hi, value_2_hi)
- S.vsts(dst[offset], value_4_0)
- S.vsts(dst[offset + vec_size], value_4_1)
- S.vsts(dst[offset + vec_size * 2], value_4_2)
- S.vsts(dst[offset + vec_size * 3], value_4_3)
+ 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):
@@ -153,31 +98,35 @@
# dense token-major scales to the output layout.
with T.SimdVF():
if out_config.use_packed_ue8m0:
- lane = S.vci(0, T.int16)
- row = S.vand(T.reinterpret(lane, 'uint16x128'), S.vdup(block_m - 1, T.uint16))
- group = T.reinterpret(S.vshrs(lane, block_m.bit_length() - 1), 'uint16x128')
- source = S.vadd(S.vmuls(row, sf_row_stride), group)
- S.vsts(indices[0], source, dist='NORM_B16')
+ 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 = S.vci(0, T.int32)
- row = S.vand(T.reinterpret(lane, 'uint32x64'), S.vdup(block_m - 1, T.uint32))
- group = T.reinterpret(S.vshrs(lane, block_m.bit_length() - 1), 'uint32x64')
- source = S.vadd(S.vmuls(row, sf_row_stride), group)
- S.vsts(indices[0], source, dist='NORM_B32')
+ 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():
- source_indices = S.vld(indices[0])
- store_dist = 'NORM_B16' if out_config.use_packed_ue8m0 else 'NORM_B32'
+ 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 = S.vgather2(src[0, source_col + group_base], source_indices)
- S.vsts(dst[group_base, 0], values, dist=store_dist)
+ 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_asc(
+ 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],
@@ -315,147 +264,172 @@
T.copy(act_x_grad[token_id, col_offset : col_offset + block_k], grad_ub)
with T.SimdVF():
- if with_quant:
- for scale_vector in range(num_scale_vectors):
- scale_offset = scale_vector * vec_size
- expanded_offset = scale_offset * 4
- x_scale_compact = S.vld(x_sf_ub[scale_offset])
- y_scale_compact = S.vld(y_sf_ub[scale_offset])
- store_fourfold_expanded(x_quant_scratch_ub, expanded_offset, x_scale_compact)
- store_fourfold_expanded(y_quant_scratch_ub, expanded_offset, y_scale_compact)
- S.mem_bar('VST_VLD')
- ones = S.vdup(1.0, T.float32)
- zeros = S.vdup(0.0, T.float32)
- if with_quant:
- high_half_mask = S.vcmps(S.vci(T.int32(0), T.int32), 31, op='gt')
- low_half_mask = S.pset(32, 'PAT_VL32')
+ 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 = S.alloc_var(T.float32)
- weight_acc = zeros if hidden_block == 0 else S.vld(weight_acc_ub[row, 0])
-
+ 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 = S.vld(x_quant_scratch_ub[vector * 8], dist='E2B_B32')
- y_scale = S.vld(y_quant_scratch_ub[vector * 8], dist='E2B_B32')
- x_raw = S.vcvt(S.vld(x_ub[col], dist='UNPK4_B8'), T.float32)
- y_raw = S.vcvt(S.vld(y_ub[col], dist='UNPK4_B8'), T.float32)
- x_value = S.vmul(x_raw, x_scale)
- y_value = S.vmul(y_raw, y_scale)
+ 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 = S.vcvt(S.vld(x_ub[col], dist='UNPK_B16'), T.float32, part=0)
- y_value = S.vcvt(S.vld(y_ub[col], dist='UNPK_B16'), T.float32, part=0)
- else:
- x_value = S.vld(x_ub[col])
- y_value = S.vld(y_ub[col])
+ 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 = S.vcvt(S.vld(grad_ub[col], dist='UNPK_B16'), T.float32, part=0)
- else:
- grad = S.vld(grad_ub[col])
+ 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 = S.vcmps(x_value, clamp_value, op='gt')
- is_clamped_y = S.vcmps(S.vabs(y_value), clamp_value, op='gt')
- x_compute_value = S.vmins(x_value, clamp_value)
- y_compute_value = S.vmins(S.vmaxs(y_value, -clamp_value), clamp_value)
+ 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 = S.vadds(S.vexpdif(zeros, x_compute_value), 1.0)
- sigmoid = sigmoid_vdiv_0ulp(ones, sigmoid_denominator)
+ 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 exact division for the required FP32 rounding.
- act_value = S.vmul(vdiv_0ulp(x_compute_value, sigmoid_denominator), y_compute_value)
+ # 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 = S.vmul(grad, act_value)
- weight_value = S.vmuls(weight_contribution, valid_token_f32) if num_experts else weight_contribution
- weight_acc = S.vadd(weight_acc, weight_value)
- grad_ws = S.vmul(S.vmuls(grad, effective_weight), sigmoid)
- inner = S.vadd(ones, S.vmul(x_compute_value, S.vsub(ones, sigmoid)))
- x_grad_unclamped = S.vmul(S.vmul(grad_ws, y_compute_value), inner)
- y_grad_unclamped = S.vmul(grad_ws, x_compute_value)
- x_grad_value = S.vsel(zeros, x_grad_unclamped, is_clamped_x) if use_clamp else x_grad_unclamped
- y_grad_value = S.vsel(zeros, y_grad_unclamped, is_clamped_y) if use_clamp else y_grad_unclamped
+ 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:
- S.vsts(x_grad_ub[col], S.vcvt(x_grad_value, T.bfloat16), dist='PK_B32')
- S.vsts(y_grad_ub[col], S.vcvt(y_grad_value, T.bfloat16), dist='PK_B32')
+ # 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:
- S.vsts(x_grad_ub[col], x_grad_value)
- S.vsts(y_grad_ub[col], y_grad_value)
- else:
- S.vsts(x_grad_f32_ub[col], x_grad_value)
- S.vsts(y_grad_f32_ub[col], y_grad_value)
- x_abs = S.vabs(x_grad_value)
- y_abs = S.vabs(y_grad_value)
- S.vsts(x_quant_scratch_ub[vector * 2], S.vcmax(x_abs, low_half_mask), dist='ONEPT_B32')
- S.vsts(x_quant_scratch_ub[vector * 2 + 1], S.vcmax(x_abs, high_half_mask), dist='ONEPT_B32')
- S.vsts(y_quant_scratch_ub[vector * 2], S.vcmax(y_abs, low_half_mask), dist='ONEPT_B32')
- S.vsts(y_quant_scratch_ub[vector * 2 + 1], S.vcmax(y_abs, high_half_mask), dist='ONEPT_B32')
+ 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 = S.vmuls(act_value, effective_weight)
+ act_out_value = V.vmuls(act_value, effective_weight, full)
if out_dtype == T.bfloat16:
- S.vsts(act_out_ub[col], S.vcvt(act_out_value, T.bfloat16), dist='PK_B32')
+ # Explicit NOSAT, same bug A/B rationale.
+ V.vstore(V.vcvt(act_out_value, 'bfloat16', saturate='NOSAT'), act_out_ub[col])
else:
- S.vsts(act_out_ub[col], act_out_value)
+ V.vstore(act_out_value, act_out_ub[col])
if with_weight:
- one_mask = S.pset(32, 'PAT_VL1')
if hidden_block + 1 == num_hidden_blocks:
- reduced_weight_sum = S.vcadd(weight_acc)
- weight_sum = S.vmuls(reduced_weight_sum, routed_scaling_factor) if with_routed_scaling else reduced_weight_sum
- S.vsts(weight_out_ub[0], weight_sum, one_mask, dist='ONEPT_B32')
- else:
- S.vsts(weight_acc_ub[row, 0], weight_acc)
+ # 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 = S.vdup(valid_token_f32, T.float32)
- # 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
+ 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
- inverse_offset = scale_offset * 4
- x_scale_value, x_inv = get_scale_and_inverse(S.vld(x_quant_scratch_ub[scale_offset]))
- y_scale_value, y_inv = get_scale_and_inverse(S.vld(y_quant_scratch_ub[scale_offset]))
- store_fourfold_expanded(x_quant_scratch_ub, inverse_offset, x_inv)
- store_fourfold_expanded(y_quant_scratch_ub, inverse_offset, y_inv)
- x_scale_out = S.vmul(x_scale_value, valid_token_vec) if num_experts else x_scale_value
- y_scale_out = S.vmul(y_scale_value, valid_token_vec) if num_experts else y_scale_value
+ 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:
- S.vsts(
+ 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],
- T.reinterpret(S.vshrs(T.reinterpret(x_scale_out, 'uint32x64'), 23), 'uint8x256'),
- dist='PK4_B32',
)
- S.vsts(
+ V.vstore(
+ V.vcvt(y_exp, 'uint8', saturate='NOSAT'),
sf_dense_u8[row, sf_half_stride * 2 + sf_block_offset * 2 + scale_offset],
- T.reinterpret(S.vshrs(T.reinterpret(y_scale_out, 'uint32x64'), 23), 'uint8x256'),
- dist='PK4_B32',
)
else:
- S.vsts(sf_dense_ub[row, sf_block_offset + scale_offset], x_scale_out)
- S.vsts(sf_dense_ub[row, sf_half_stride + sf_block_offset + scale_offset], y_scale_out)
+ 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 = S.vcmps(valid_token_vec, 0.0, op='gt')
+ 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 = S.vld(x_quant_scratch_ub[vector * 8], dist='E2B_B32')
- y_inv_vec = S.vld(y_quant_scratch_ub[vector * 8], dist='E2B_B32')
- x_grad_quant_value = S.vld(x_grad_f32_ub[col])
- y_grad_quant_value = S.vld(y_grad_f32_ub[col])
- x_scaled_value = S.vmul(x_grad_quant_value, x_inv_vec)
- y_scaled_value = S.vmul(y_grad_quant_value, y_inv_vec)
- S.vsts(x_grad_ub[col], S.vcvt(x_grad_quant_value, T.bfloat16), dist='PK_B32')
- S.vsts(y_grad_ub[col], S.vcvt(y_grad_quant_value, T.bfloat16), dist='PK_B32')
- x_scaled = S.vsel(x_scaled_value, zeros, valid_token_mask) if num_experts else x_scaled_value
- y_scaled = S.vsel(y_scaled_value, zeros, valid_token_mask) if num_experts else y_scaled_value
- S.vsts(x_grad_fp8_ub[col], S.vcvt(x_scaled, T.float8_e4m3fn), dist='PK4_B32')
- S.vsts(y_grad_fp8_ub[col], S.vcvt(y_scaled, T.float8_e4m3fn), dist='PK4_B32')
+ 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])
@@ -497,4 +471,4 @@
output_sf[token_base, num_stored_groups + sf_output_offset],
)
- return swiglu_backward_kernel_asc
+ return swiglu_backward_kernel_vmi |
Author
ASC → VMI：两个公开入口的完整 diff对应 PR head 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_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(), |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Port SwiGLU forward and backward to VMI while preserving the ASC scheduling shell, DMA copies, buffer geometry, scalar predicates, and synchronization. The Ascend public API now selects the VMI factories. This takes over #20, retaining provenance from kuri780's implementation.
Operator-local software division and sigmoid correction routines are removed. Exact division delegates to
V.vdiv(..., precision='exact'). Optional full masks are omitted where supported. Forward grouped amax usesgroup=8for 256 lanes to retain SF32 semantics;group=4would incorrectly combine64 channels.This is a draft pending final correctness/revision validation. Performance below the original0.95 target is accepted by the task owner and is documented rather than used as a merge gate. Historical candidatev5 passed24 forward and24 backward core configurations with three paired timing rounds. Extended forward26 and backward26 cases pass; finite-boundary forward7 cases also pass. Remaining checks are final submitted-source correctness and toolchain dependency verification. The final source will be revalidated before merge.
Historical median ASC/VMI ratios: forward BF16/FP32/FP8=1.033/1.024/0.972; backward=0.903/0.992/0.877. These medians hide individual gaps, particularly forwardFP8 clamp. Detailed historical per-case measurements are listed below. The submitted source includes the partial-width fix validated by the extended matrix. Clamp optimization remains experimental and is not included; historical numbers must not be attributed to a different source revision.
Toolchain: TileLang pto-dev
4eb83fdeb2fb317a0485052f51d0b77f4238d8b5, GitCode cann/pto-as master04c785491e844ec6d642e444a4e3ee84d1a60496, CANN9.2.0. PTOAS also requires a local Python atomic signedness patch for unchanged ASC-style clamp-count atomics. This dependency is not unmodified upstream master and must be resolved/reproduced before merge.Original review disposition: optional masks omitted where supported; ASC DMA expressions restored; 256-lane SF32 reduction uses group=8. TileLang PR #14 remains open; the restored copies compile on the pinned pto-dev. ASC reference files remain unchanged. Investigation scripts, logs and compiler patches are retained locally rather than added to the operator diff.
Final submitted-source validation is running on both NPUs: repository torch-reference tests ondevice1, and a fresh48-case paired ASC/VMI three-round sweep ondevice0. Unmerged clamp performance experiments are excluded from this PR.
The authenticated author liuzidi can publish this fork PR but currently has no upstream push/merge permission (GitHub repository permissions API). A maintainer will need to execute the eventual merge after the correctness and dependency checks complete.
Historical v5 per-case performance (ASC/VMI; not final-source results)
Candidate v5 full level-0 performance
ASC/VMI ratio; three alternating-order msprof rounds. All24 forward and24 backward configs passed; this remains a historical candidate, not final acceptance.