From a8968e2afb623692ea8a64adf9b47c324d6c3527 Mon Sep 17 00:00:00 2001 From: Vinayak Gokhale Date: Wed, 2 Sep 2026 12:01:07 -0500 Subject: [PATCH 1/3] [Triton] Deduplicate the MoE quant kernels against aiter/ops/triton/quant The MoE quant path had four kernels that reimplemented quant kernels living in aiter/ops/triton/quant. Fold them together. _mxfp4_quant_kernel: deleted. It already shared _mxfp4_quant_op with _dynamic_mxfp4_quant_kernel and was otherwise a strict subset of it (no int64 stride guard, no persistent loop, no cache modifier), so mxfp4_quant() is now a wrapper over dynamic_mxfp4_quant. dynamic_mxfp4_quant gains optional out params, mirroring dynamic_mxfp8_quant's `scale`, because the a4w4 GEMM needs row-major scales while dynamic_mxfp4_quant allocates a transposed view -- and its stride(0)==1 would flip the X_SCALES_TDM heuristic in moe_gemm_a4w4. Also drops an eager x.to(float32) that materialised the whole activation in fp32 before the launch; the kernel casts on load. _downcast_to_mxfp: keeps both scale schemes behind a new POW2_SCALE constexpr. False (default, unchanged) is amax/dtype_max with the exponent rounded per DEQUANT_SCALE_ROUNDING_MODE; True is the even_round scheme shared with _mxfp4_quant_op/_mxfp8_quant_op. They are not interchangeable -- they disagree on ~12% of fp4 block scales and ~0.3% of fp8 ones -- so the flag stays. With POW2_SCALE=True the output is now bit-identical to dynamic_mxfp4_quant (fp4) and dynamic_mxfp8_quant (e4m3). No perf difference between the two (0.98-1.09x across four shapes): the kernel is memory bound, so the log2/exp2 pair costs nothing and the name describes the policy, not a speed tier. _mxfp8_quant_op gains a defaulted LOG2_DTYPE_MAX so the fast path also covers e5m2, whose 15 its hardcoded 8 got wrong. _downcast_to_static_fp8: deleted, merged into _static_per_tensor_quant_fp8_i8_kernel behind FAST_CONVERT (reciprocal-multiply vs exact division). The two forms are not equivalent: at scale=448.0 a dense bf16 sweep puts them 1 fp8 ulp apart on 64 elements, so both survive. Here the name is literal -- the reciprocal form is 15-20% faster. Two side effects: - The merged kernel takes its output dtype from the pointer rather than hardcoding tl.float8e4nv, so gfx942 no longer converts fn->fnuz on store (the buffer is allocated fnuz there). This changes gfx942 output; it was double-rounding through the wrong exponent bias before. - The old MoE tile indexed with offs_m[None, :] / offs_n[:, None], putting the stride-1 axis first where Triton cannot prove stride_x_n == 1. Correcting that makes the MoE path 1.07-1.62x faster. Against the old row-per-program kernel: 0.94x worst case, 1.05-2.38x elsewhere, and it can now serve wide rows that did not fit in registers at all. The shared launcher views both tensors as 2D instead of trusting qx.stride(0), because per_tensor_quant_triton in aiter/ops/quant.py passes an N-D output alongside a 2-D input -- the old kernel got away with using x's row stride for the output. No test covered that. upcast_from_mxfp: moved to op_tests/triton_tests/utils/mxfp_ref.py. It is test scaffolding -- it builds bf16 references so the torch comparison sees the same values the kernel reads -- and had no callers under aiter/. Verified on gfx950: the exact-division and reciprocal paths are bit-identical to the kernels they replace across 9,431,040 elements x 307 scales, int8 included; mxfp4_quant is bit-identical across 6 shapes x 3 dtypes; the N-D output path matches for 3-D and 4-D, fp8 and int8, static and dynamic. op_tests/triton_tests/{moe,quant} plus test_activation.py: 6897 passed, 6396 skipped. --- .../_triton_kernels/moe/moe_op_gemm_a4w4.py | 74 ----- .../triton/_triton_kernels/moe/quant_moe.py | 276 ++++-------------- .../ops/triton/_triton_kernels/quant/quant.py | 68 ++++- aiter/ops/triton/moe/moe_op_gemm_a4w4.py | 46 +-- aiter/ops/triton/moe/quant_moe.py | 101 ++----- aiter/ops/triton/quant/quant.py | 96 ++++-- .../triton_tests/moe/test_moe_gemm_a16w4.py | 7 +- .../triton_tests/moe/test_moe_gemm_a4w4.py | 6 +- .../triton_tests/moe/test_moe_gemm_a8w4.py | 2 +- .../triton_tests/moe/test_moe_gemm_a8w8.py | 2 +- op_tests/triton_tests/utils/mxfp_ref.py | 236 +++++++++++++++ 11 files changed, 453 insertions(+), 461 deletions(-) create mode 100644 op_tests/triton_tests/utils/mxfp_ref.py diff --git a/aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a4w4.py b/aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a4w4.py index 370d6c3df1..fb110f858f 100644 --- a/aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a4w4.py +++ b/aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a4w4.py @@ -6,7 +6,6 @@ import triton.language as tl from aiter.ops.triton._triton_kernels.moe.activations import _swiglu -from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid @@ -127,79 +126,6 @@ def unswizzle_mx_scale_gfx1250( return scale_buffer_slice -_mxfp4_quant_kernel_repr = make_kernel_repr( - "_mxfp4_quant_kernel", - [ - "BLOCK_SIZE_M", - "BLOCK_SIZE_N", - "MXFP4_QUANT_BLOCK_SIZE", - "EVEN_M_N", - ], -) - - -@triton.jit(repr=_mxfp4_quant_kernel_repr) -def _mxfp4_quant_kernel( - x_ptr, - x_fp4_ptr, - bs_ptr, - stride_x_m, - stride_x_n, - stride_fp4_m, - stride_fp4_n, - stride_bs_m, - stride_bs_n, - M, - N, - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - MXFP4_QUANT_BLOCK_SIZE: tl.constexpr, - EVEN_M_N: tl.constexpr, -): - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - mask_m = offs_m < M - mask_n = offs_n < N - mask = mask_m[:, None] & mask_n[None, :] - x_offs = offs_m[:, None] * stride_x_m + offs_n[None, :] * stride_x_n - x = tl.load(x_ptr + x_offs, mask=mask, other=0).to(tl.float32) - - out_tensor, bs_e8m0 = _mxfp4_quant_op( - x, BLOCK_SIZE_N, BLOCK_SIZE_M, MXFP4_QUANT_BLOCK_SIZE - ) - - # Store quantized x blocks - out_offs_n = pid_n * BLOCK_SIZE_N // 2 + tl.arange(0, BLOCK_SIZE_N // 2) - out_offs = offs_m[:, None] * stride_fp4_m + out_offs_n[None, :] * stride_fp4_n - - if EVEN_M_N: - tl.store(x_fp4_ptr + out_offs, out_tensor) - else: - out_mask = (offs_m < M)[:, None] & (out_offs_n < (N // 2))[None, :] - tl.store(x_fp4_ptr + out_offs, out_tensor, mask=out_mask) - - # Store scale blocks - NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE - num_blocks_total = N // MXFP4_QUANT_BLOCK_SIZE - bs_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - bs_offs_n = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS) - - bs_offs = bs_offs_m[:, None] * stride_bs_m + bs_offs_n[None, :] * stride_bs_n - bs_mask = (bs_offs_m < M)[:, None] & (bs_offs_n < num_blocks_total)[None, :] - - if EVEN_M_N: - tl.store(bs_ptr + bs_offs, bs_e8m0) - else: - tl.store( - bs_ptr + bs_offs, - bs_e8m0, - mask=bs_mask, - ) - - _moe_gemm_a4w4_repr = make_kernel_repr( "_moe_gemm_a4w4", [ diff --git a/aiter/ops/triton/_triton_kernels/moe/quant_moe.py b/aiter/ops/triton/_triton_kernels/moe/quant_moe.py index 28e65e9599..d8aca59dbb 100644 --- a/aiter/ops/triton/_triton_kernels/moe/quant_moe.py +++ b/aiter/ops/triton/_triton_kernels/moe/quant_moe.py @@ -1,6 +1,10 @@ import triton import triton.language as tl +from aiter.ops.triton._triton_kernels.quant.quant import ( + _mxfp4_quant_op, + _mxfp8_quant_op, +) from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr @@ -12,62 +16,6 @@ def _compute_static_fp8_quant(tensor, scale): return tensor -_downcast_to_static_fp8_repr = make_kernel_repr( - "_downcast_to_static_fp8", - [ - "BLOCK_M", - "BLOCK_N", - ], -) - - -@triton.jit(repr=_downcast_to_static_fp8_repr) -def _downcast_to_static_fp8( - x_ptr, - stride_x_m, - stride_x_n, - y_ptr, - stride_y_m, - stride_y_n, - scale_ptr, - M, - N, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - - x_dtype: tl.constexpr = x_ptr.dtype.element_ty - tl.static_assert( - (x_dtype == tl.bfloat16) or (x_dtype == tl.float16) or (x_dtype == tl.float32), - f"{x_dtype=} must be bfloat16 or float16 or float32", - ) - - pid_m = tl.program_id(0).to(tl.int64) - pid_n = tl.program_id(1).to(tl.int64) - - start_m = pid_m * BLOCK_M - start_n = pid_n * BLOCK_N - - x_ptr += start_m * stride_x_m + start_n * stride_x_n - y_ptr += start_m * stride_y_m + start_n * stride_y_n - - offs_m = tl.arange(0, BLOCK_M)[None, :].to(tl.int64) - offs_n = tl.arange(0, BLOCK_N)[:, None].to(tl.int64) - - mask_m = start_m + offs_m < M - mask_n = start_n + offs_n < N - mask_xy = mask_m & mask_n - - offs_x = offs_m * stride_x_m + offs_n * stride_x_n - offs_y = offs_m * stride_y_m + offs_n * stride_y_n - - x = tl.load(x_ptr + offs_x, mask=mask_xy) - - y = _compute_static_fp8_quant(x, tl.load(scale_ptr)) - - tl.store(y_ptr + offs_y, y, mask=mask_xy) - - @triton.jit def _get_max_quant_val(dtype: tl.constexpr): if dtype == tl.uint8: @@ -86,7 +34,27 @@ def _compute_mx_quant_and_scale( valid_src_mask, mx_tensor_dtype: tl.constexpr, DEQUANT_SCALE_ROUNDING_MODE: tl.constexpr = 0, + POW2_SCALE: tl.constexpr = False, ): + """Quantize a [OUT_DIM, QUANT_DIM] fp32/bf16/fp16 tile to MXFP4 or MXFP8. + + Two scale-derivation schemes, selected by POW2_SCALE: + + * ``POW2_SCALE=False`` (default): ``scale = amax / dtype_max`` with the + exponent rounded per DEQUANT_SCALE_ROUNDING_MODE. Never saturates, but + leaves up to 2x of the dtype's range unused when amax sits just above a + power of two. + * ``POW2_SCALE=True``: the ``even_round`` scheme shared with + :func:`_mxfp4_quant_op` / :func:`_mxfp8_quant_op` -- amax is rounded to the + nearest power of two and the scale is taken relative to the largest power + of two the dtype holds (4 for e2m1, 256 for e4m3, 32768 for e5m2). Uses + the range better in the bulk at the cost of clipping the odd outlier. + + The two are NOT bit-compatible: they disagree on roughly 12% of fp4 block + scales and 0.3% of fp8 ones, so a tensor quantized with one must be + dequantized against the same one. DEQUANT_SCALE_ROUNDING_MODE is ignored + when POW2_SCALE is set. + """ is_fp8: tl.constexpr = ( mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5 ) @@ -94,6 +62,30 @@ def _compute_mx_quant_and_scale( BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1] BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // 32 + if POW2_SCALE: + # Padding lanes are zeroed rather than set to -1: zero is neutral for the + # group amax and is also what the tile stores for them either way. + masked = tl.where(valid_src_mask, src_tensor.to(tl.float32), 0.0) + if is_fp8: + LOG2_DTYPE_MAX: tl.constexpr = 8 if mx_tensor_dtype == tl.float8e4nv else 15 + grouped = tl.reshape( + masked, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32] + ) + scale_e8m0, quant_scale = _mxfp8_quant_op(grouped, 2, LOG2_DTYPE_MAX) + out_tensor = ( + (grouped * quant_scale) + .reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM]) + .to(mx_tensor_dtype) + ) + dequant_scale_exponent = scale_e8m0.reshape( + [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE] + ) + else: + out_tensor, dequant_scale_exponent = _mxfp4_quant_op( + masked, BLOCK_SIZE_QUANT_DIM, BLOCK_SIZE_OUT_DIM, 32 + ) + return out_tensor, dequant_scale_exponent + # Explicit cast to fp32 since most ops are not supported on bfloat16. We avoid needless conversions to and from bf16 f32_tensor = src_tensor.to(tl.float32) abs_tensor = tl.abs(f32_tensor) @@ -182,6 +174,7 @@ def _compute_mx_quant_and_scale( "BLOCK_SIZE_OUT_DIM", "BLOCK_SIZE_QUANT_DIM", "DEQUANT_SCALE_ROUNDING_MODE", + "POW2_SCALE", ], ) @@ -202,6 +195,7 @@ def _downcast_to_mxfp( BLOCK_SIZE_OUT_DIM: tl.constexpr, BLOCK_SIZE_QUANT_DIM: tl.constexpr, DEQUANT_SCALE_ROUNDING_MODE: tl.constexpr, + POW2_SCALE: tl.constexpr, ): tl.static_assert( @@ -276,177 +270,17 @@ def _downcast_to_mxfp( src_tensor = tl.load(src_ptr + src_tensor_offsets, mask=full_mask_src) out_tensor, scale_tensor = _compute_mx_quant_and_scale( - src_tensor, full_mask_src, mx_tensor_dtype, DEQUANT_SCALE_ROUNDING_MODE + src_tensor, + full_mask_src, + mx_tensor_dtype, + DEQUANT_SCALE_ROUNDING_MODE, + POW2_SCALE, ) tl.store(mx_scale_ptr + mx_scale_offsets, scale_tensor, mask=full_scale_mask) tl.store(mx_tensor_ptr + mx_tensor_offsets, out_tensor, mask=full_mask_mxt) -_upcast_from_mxfp_repr = make_kernel_repr( - "_upcast_from_mxfp", - [ - "BLOCK_SIZE_OUT_DIM", - "BLOCK_SIZE_QUANT_DIM", - ], -) - - -@triton.jit(repr=_upcast_from_mxfp_repr) -def _upcast_from_mxfp( - out_ptr, - stride_o_outer, - stride_o_quant: tl.constexpr, - mx_scale_ptr, - stride_scale_outer, - stride_scale_quant, - mx_tensor_ptr, - stride_tensor_outer, - stride_tensor_quant: tl.constexpr, - outer_dim, - quant_dim, - BLOCK_SIZE_OUT_DIM: tl.constexpr, - BLOCK_SIZE_QUANT_DIM: tl.constexpr, -): - - tl.static_assert( - stride_o_quant == 1, "the weight must be contiguous in the k dimension for mx" - ) - tl.static_assert( - BLOCK_SIZE_QUANT_DIM % 32 == 0, "BLOCK_SIZE_K must be a multiple of 32" - ) - # uint8 signifies two fp4 e2m1 values packed into a single byte - mx_tensor_dtype: tl.constexpr = mx_tensor_ptr.dtype.element_ty - dst_dtype: tl.constexpr = out_ptr.dtype.element_ty - tl.static_assert(dst_dtype == tl.float16 or dst_dtype == tl.bfloat16) - tl.static_assert( - mx_tensor_dtype == tl.uint8 - or ( - (mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5) - or mx_tensor_dtype == dst_dtype - ), - "mx_tensor_ptr must be uint8 or float8 or dst_dtype", - ) - tl.static_assert( - mx_scale_ptr.dtype.element_ty == tl.uint8, "mx_scale_ptr must be uint8" - ) - - # Determine if we are dealing with fp8 types. - is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8 - is_fp8: tl.constexpr = ( - mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5 - ) - K_DIVISOR: tl.constexpr = 2 if is_fp4 else 1 - BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 32 - BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // K_DIVISOR - - # Compute starting indices for the quantized (packed) dimension and the outer dimension. - outer_block = tl.program_id(0).to(tl.int64) - quant_block = tl.program_id(1).to(tl.int64) - - start_mxt_quant = quant_block * BLOCK_SIZE_QUANT_MX_TENSOR - start_out_quant = quant_block * BLOCK_SIZE_QUANT_DIM - start_mx_scale_quant = quant_block * BLOCK_SIZE_QUANT_MX_SCALE - start_out = outer_block * BLOCK_SIZE_OUT_DIM - - mx_tensor_ptr += ( - start_mxt_quant * stride_tensor_quant + start_out * stride_tensor_outer - ) - mx_scale_ptr += ( - start_mx_scale_quant * stride_scale_quant + start_out * stride_scale_outer - ) - out_ptr += start_out * stride_o_outer + start_out_quant * stride_o_quant - - # Compute offsets and masks. - offs_src_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :].to(tl.int64) - offs_out_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :].to(tl.int64) - offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None].to(tl.int64) - offs_scale = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :].to(tl.int64) - - mask_outer = start_out + offs_outer < outer_dim - mask_out_quant = start_out_quant + offs_out_quant < quant_dim - full_mask_out = mask_out_quant & mask_outer - - mask_src_quant = start_mxt_quant + offs_src_quant < tl.cdiv(quant_dim, K_DIVISOR) - full_mask_src = mask_src_quant & mask_outer - - mask_scale = start_mx_scale_quant + offs_scale < tl.cdiv(quant_dim, 32) - full_scale_mask = mask_scale & mask_outer - - tensor_offsets = ( - offs_src_quant * stride_tensor_quant + offs_outer * stride_tensor_outer - ) - scale_offsets = offs_scale * stride_scale_quant + offs_outer * stride_scale_outer - out_offsets = offs_out_quant * stride_o_quant + offs_outer * stride_o_outer - - # Load the packed tensor and scale. - tensor = tl.load(mx_tensor_ptr + tensor_offsets, mask=full_mask_src) - scale = tl.load(mx_scale_ptr + scale_offsets, mask=full_scale_mask) - - # Upcast the scale to the destination type. - if dst_dtype == tl.bfloat16: - dst_scale = (scale.to(tl.uint16) << 7).to(dst_dtype, bitcast=True) - else: - tl.static_assert(dst_dtype == tl.float16) - dst_scale = (scale.to(tl.uint32) << 23).to(tl.float32, bitcast=True) - dst_scale = dst_scale.to(tl.float16) - - # Now upcast the tensor. - if is_fp8: - dst_tensor = tensor.to(dst_dtype) - if tensor.dtype == tl.float8e5: - from_e_bits: tl.constexpr = 5 - from_m_bits: tl.constexpr = 2 - to_e_bits: tl.constexpr = 8 if dst_dtype == tl.bfloat16 else 5 - to_m_bits: tl.constexpr = 7 if dst_dtype == tl.bfloat16 else 10 - - # Preserve infs and nans. FIXME Fp8E5M2_to_Bf16 doesn't preserve them! - non_finite_mask_src: tl.constexpr = ((1 << from_e_bits) - 1) << from_m_bits - non_finite_mask_dst: tl.constexpr = ((1 << to_e_bits) - 1) << to_m_bits - dst_tensor = tl.where( - (tensor.to(tl.uint8, bitcast=True) & non_finite_mask_src) - == non_finite_mask_src, - (dst_tensor.to(tl.uint16, bitcast=True) | non_finite_mask_dst).to( - dst_dtype, bitcast=True - ), - dst_tensor, - ) - else: - assert is_fp4 - dst_bias: tl.constexpr = 127 if dst_dtype == tl.bfloat16 else 15 - dst_0p5: tl.constexpr = 16128 if dst_dtype == tl.bfloat16 else 0x3800 - dst_m_bits: tl.constexpr = 7 if dst_dtype == tl.bfloat16 else 10 - # e2m1 - em0 = tensor & 0x07 - em1 = tensor & 0x70 - x0 = (em0.to(tl.uint16) << (dst_m_bits - 1)) | ( - (tensor & 0x08).to(tl.uint16) << 12 - ) - x1 = (em1.to(tl.uint16) << (dst_m_bits - 5)) | ( - (tensor & 0x80).to(tl.uint16) << 8 - ) - # Three cases: - # 1) x is normal and non-zero: Correct bias - x0 = tl.where((em0 & 0x06) != 0, x0 + ((dst_bias - 1) << dst_m_bits), x0) - x1 = tl.where((em1 & 0x60) != 0, x1 + ((dst_bias - 1) << dst_m_bits), x1) - # 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in the dst type - x0 = tl.where(em0 == 0x01, dst_0p5 | (x0 & 0x8000), x0) - x1 = tl.where(em1 == 0x10, dst_0p5 | (x1 & 0x8000), x1) - # 3) x is zero, do nothing - dst_tensor = tl.interleave(x0, x1).to(dst_dtype, bitcast=True) - - # Reshape for proper broadcasting: the scale was stored with a 32-sized "inner" grouping. - dst_tensor = dst_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32]) - dst_scale = dst_scale.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1]) - scale = scale.reshape(dst_scale.shape) - - out_tensor = dst_tensor * dst_scale - # Correct any NaNs encoded via the scale. - out_tensor = tl.where(scale == 0xFF, float("nan"), out_tensor) - out_tensor = out_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM]) - tl.store(out_ptr + out_offsets, out_tensor, mask=full_mask_out) - - _smoothquant_fuse_quant_kernel_repr = make_kernel_repr( "_smoothquant_fuse_quant_kernel", [ diff --git a/aiter/ops/triton/_triton_kernels/quant/quant.py b/aiter/ops/triton/_triton_kernels/quant/quant.py index ec35bd32f1..8cb6b0728f 100644 --- a/aiter/ops/triton/_triton_kernels/quant/quant.py +++ b/aiter/ops/triton/_triton_kernels/quant/quant.py @@ -10,24 +10,59 @@ def _static_per_tensor_quant_fp8_i8_kernel( qx_ptr, x_in_ptr, scale_in_ptr, + rows: int, cols: int, - x_in_stride_r: int, - NUM_COL_POW2: tl.constexpr, + stride_x_m, + stride_x_n, + stride_q_m, + stride_q_n, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + FAST_CONVERT: tl.constexpr, ): - pid = tl.program_id(axis=0) - tl.assume(pid > 0) - tl.assume(x_in_stride_r > 0) + """Quantize x by a single tensor-wide scale into the dtype of ``qx_ptr`` + (fp8 e4m3/e5m2 or int8). - offs = pid * x_in_stride_r + tl.arange(0, NUM_COL_POW2) - mask = tl.arange(0, NUM_COL_POW2) < cols - x = tl.load(x_in_ptr + offs, mask=mask, cache_modifier=".cg") + FAST_CONVERT picks how the scale is applied: - scale = tl.load(scale_in_ptr) - scale_recip = 1 / scale + * ``True`` -- multiply by the reciprocal. One v_mul per element, the + reciprocal computed once per program. + * ``False`` -- divide. A correctly-rounded fp32 division per element, so it + matches an ``x / scale`` reference exactly. - qx = (x * scale_recip).to(qx_ptr.dtype.element_ty) + The two agree on nearly every input but not all of them: with scale=448.0 a + dense bf16 sweep puts them 1 fp8 ulp apart on ~0.2% of elements, because a + reciprocal that is half an fp32 ulp off can land on the far side of an fp8 + rounding boundary. + """ + # Fold the block origin into the base pointers in int64 so only the in-tile + # offsets, which always fit, stay 32-bit. + start_m = tl.program_id(axis=0).to(tl.int64) * BLOCK_M + start_n = tl.program_id(axis=1).to(tl.int64) * BLOCK_N + x_in_ptr += start_m * stride_x_m + start_n * stride_x_n + qx_ptr += start_m * stride_q_m + start_n * stride_q_n + + offs_m = tl.arange(0, BLOCK_M)[:, None] + offs_n = tl.arange(0, BLOCK_N)[None, :] + mask = (start_m + offs_m < rows) & (start_n + offs_n < cols) + + x = tl.load( + x_in_ptr + offs_m * stride_x_m + offs_n * stride_x_n, + mask=mask, + cache_modifier=".cg", + ) - tl.store(qx_ptr + offs, qx, mask=mask) + scale = tl.load(scale_in_ptr) + if FAST_CONVERT: + qx = x * (1 / scale) + else: + qx = x.to(tl.float32) / scale + + tl.store( + qx_ptr + offs_m * stride_q_m + offs_n * stride_q_n, + qx.to(qx_ptr.dtype.element_ty), + mask=mask, + ) @triton.jit @@ -377,19 +412,24 @@ def _dynamic_mxfp4_quant_kernel( @triton.jit -def _mxfp8_quant_op(x_grouped, QUANT_AXIS: tl.constexpr): +def _mxfp8_quant_op( + x_grouped, QUANT_AXIS: tl.constexpr, LOG2_DTYPE_MAX: tl.constexpr = 8 +): """Shared MXFP8 (1x32 e8m0) scale derivation. Given a fp32 tile where the QUANT_AXIS dim is sized QUANT_BLOCK_SIZE (=32), returns (scale_e8m0, quant_scale): the per-group uint8 e8m0 scale and the matching fp32 multiplicative scale. Both outputs keep QUANT_AXIS with size 1 so they broadcast against the input for in-place quantization. + + LOG2_DTYPE_MAX is log2 of the largest power of two the target dtype holds: + 8 for e4m3 (max 448 -> 256) and 15 for e5m2 (max 57344 -> 32768). """ amax = tl.max(tl.abs(x_grouped), axis=QUANT_AXIS, keep_dims=True) amax_i32 = amax.to(tl.int32, bitcast=True) amax_i32 = (amax_i32 + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000 amax_p2 = amax_i32.to(tl.float32, bitcast=True) - scale_unbiased = tl.log2(amax_p2).floor() - 8 + scale_unbiased = tl.log2(amax_p2).floor() - LOG2_DTYPE_MAX scale_unbiased = tl.clamp(scale_unbiased, min=-127, max=127) scale_e8m0 = (scale_unbiased.to(tl.int32) + 127).to(tl.uint8) quant_scale = tl.exp2(-scale_unbiased) diff --git a/aiter/ops/triton/moe/moe_op_gemm_a4w4.py b/aiter/ops/triton/moe/moe_op_gemm_a4w4.py index 22bd968b0c..f130f945e2 100644 --- a/aiter/ops/triton/moe/moe_op_gemm_a4w4.py +++ b/aiter/ops/triton/moe/moe_op_gemm_a4w4.py @@ -12,12 +12,10 @@ get_moe_a4w4_layouts_decode, get_moe_a4w4_layouts_prefill, ) -from aiter.ops.triton._triton_kernels.moe.moe_op_gemm_a4w4 import ( - _moe_gemm_a4w4, - _mxfp4_quant_kernel, -) +from aiter.ops.triton._triton_kernels.moe.moe_op_gemm_a4w4 import _moe_gemm_a4w4 from aiter.ops.triton.moe.moe_routing.routing import RoutingData from aiter.ops.triton.moe.reduce import reduce_grouped +from aiter.ops.triton.quant.quant import dynamic_mxfp4_quant from aiter.ops.triton.utils._triton.arch_info import get_arch from aiter.ops.triton.utils.gemm_config_utils import pick_gemm_num_stages from aiter.ops.triton.utils.moe_config_utils import get_moe_dispatch @@ -183,53 +181,29 @@ def get_kernel_config_gluon(m, n, k, routing_data): def mxfp4_quant( x: torch.Tensor, - block_size_m: int = 16, - block_size_n: int = 256, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Quantize a 2D tensor `x` of shape [M, K] (bf16/fp16/fp32) to MXFP4 (E2M1) format - quantized along the K dimension. + Quantize a 2D tensor `x` of shape [M, N] (bf16/fp16/fp32) to MXFP4 (E2M1) format + quantized along the N dimension. + + Thin wrapper over :func:`aiter.ops.triton.quant.quant.dynamic_mxfp4_quant`; it + exists only to hand that op the row-major scale buffer the a4w4 GEMM expects + (``dynamic_mxfp4_quant`` allocates a column-major one by default). Returns: - A packed MXFP4 tensor `x_fp4` of shape [M, N // 2] (stored as uint8), where each byte stores two 4-bit values. - A block-scale tensor `x_scale` of shape [M, N / 32], where each entry - corresponds to one MXFP4 quantization block of 32 elements along the K dimension. + corresponds to one MXFP4 quantization block of 32 elements along the N dimension. """ M, N = x.shape assert N % MXFP4_QUANT_BLOCK_SIZE == 0 - assert block_size_n % MXFP4_QUANT_BLOCK_SIZE == 0 - x_fp32 = x.to(torch.float32) x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device) x_scale = torch.empty( (M, N // MXFP4_QUANT_BLOCK_SIZE), dtype=torch.uint8, device=x.device ) - - grid = ( - triton.cdiv(M, block_size_m), - triton.cdiv(N, block_size_n), - ) - - _mxfp4_quant_kernel[grid]( - x_fp32, - x_fp4, - x_scale, - x_fp32.stride(0), - x_fp32.stride(1), - x_fp4.stride(0), - x_fp4.stride(1), - x_scale.stride(0), - x_scale.stride(1), - M, - N, - BLOCK_SIZE_M=block_size_m, - BLOCK_SIZE_N=block_size_n, - MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE, - EVEN_M_N=(M % block_size_m == 0) and (N % block_size_n == 0), - ) - - return x_fp4, x_scale + return dynamic_mxfp4_quant(x, x_fp4=x_fp4, blockscale_e8m0=x_scale) def moe_gemm_a4w4( diff --git a/aiter/ops/triton/moe/quant_moe.py b/aiter/ops/triton/moe/quant_moe.py index acafc103f8..9e5d71f608 100644 --- a/aiter/ops/triton/moe/quant_moe.py +++ b/aiter/ops/triton/moe/quant_moe.py @@ -5,11 +5,10 @@ from aiter.ops.triton._triton_kernels.moe.quant_moe import ( _downcast_to_mxfp, - _downcast_to_static_fp8, _smoothquant_fuse_quant_kernel, _smoothquant_fuse_quant_kernel_single_pass, - _upcast_from_mxfp, ) +from aiter.ops.triton.quant.quant import _static_per_tensor_quant_launch from aiter.ops.triton.utils._triton.arch_info import get_arch @@ -25,37 +24,20 @@ def downcast_to_static_fp8_3d(x: torch.Tensor, scale: torch.Tensor): def downcast_to_static_fp8(x: torch.Tensor, scale: torch.Tensor): + """Quantize ``x`` by a single tensor-wide ``scale`` to the arch's fp8 e4m3. + + Uses the exact-division form of the shared static per-tensor quant kernel + (``FAST_CONVERT=False``), which is what this path has always done; the + reciprocal form used by ``static_per_tensor_quant_fp8_i8`` differs on a + small fraction of inputs. + """ M, N = x.shape if get_arch() != "gfx942": dtype = torch.float8_e4m3fn else: dtype = torch.float8_e4m3fnuz - y = torch.empty((M, N), dtype=dtype, device="cuda") - - BLOCK_M = min(triton.next_power_of_2(M), 128) - if M <= 4096: - BLOCK_N = 32 - else: - BLOCK_N = 64 - grid_m = triton.cdiv(x.shape[0], BLOCK_M) - grid_n = triton.cdiv(x.shape[1], BLOCK_N) - - _downcast_to_static_fp8[(grid_m, grid_n)]( - x, - x.stride(0), - x.stride(1), - y, - y.stride(0), - y.stride(1), - scale, - M, - N, - BLOCK_M, - BLOCK_N, - num_warps=8, - ) - - return y + y = torch.empty((M, N), dtype=dtype, device=x.device) + return _static_per_tensor_quant_launch(y, x, scale, fast_convert=False) class DequantScaleRoundingMode(Enum): @@ -68,6 +50,7 @@ def downcast_to_mxfp( out_quant_type: torch.dtype, axis: int, DEQUANT_SCALE_ROUNDING_MODE: DequantScaleRoundingMode = DequantScaleRoundingMode.ROUND_UP, + pow2_scale: bool = False, ): """ Convert the src weights to mx format. The src weight is quantized along the axis dimension. @@ -77,6 +60,11 @@ def downcast_to_mxfp( If weight_quant_type is torch.float8_e4m3fn or torch.float8_e5m2, we output mxfp8 with the float8s are stored in their respective formats. + + ``pow2_scale`` selects the ``even_round`` scale scheme shared with + ``dynamic_mxfp4_quant`` / ``dynamic_mxfp8_quant`` instead of the default + ``amax / dtype_max`` one; see :func:`_compute_mx_quant_and_scale`. The two + disagree, so a tensor quantized one way must be dequantized the same way. """ ndim = src_tensor.ndim assert -ndim <= axis < ndim, f"Invalid axis {axis=}" @@ -120,6 +108,7 @@ def downcast_to_mxfp( BLOCK_OUT_DIM, BLOCK_QUANT_DIM, DEQUANT_SCALE_ROUNDING_MODE.value, + pow2_scale, num_warps=8, ) @@ -128,62 +117,6 @@ def downcast_to_mxfp( return out_quant_tensor, out_scale -def upcast_from_mxfp( - tensor: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, axis: int -): - """ - Upcasts an mxfp (packed) weight tensor back to float16 or bfloat16. - - The function assumes that the tensors were quantized along the given axis. - It permutes the tensor so that the quantized axis is last, reshapes to 2D, - launches the Triton upcast kernel, and then unpermutes back to the original order. - """ - ndim = tensor.ndim - assert -ndim <= axis < ndim, f"Invalid axis {axis=}" - axis = axis if axis >= 0 else axis + ndim - assert tensor.ndim == scale.ndim, ( - f"Weight and scale must have the same number of dimensions. " - f"Got {tensor.ndim=} and {scale.ndim=}" - ) - # dtype checks - assert tensor.dtype in { - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, - torch.float8_e4m3fnuz, - }, f"Invalid tensor dtype {tensor.dtype=}" - assert scale.dtype == torch.uint8, f"Invalid scale dtype {scale.dtype=}" - assert dtype in (torch.float16, torch.bfloat16), f"Invalid output dtype {dtype=}" - # upcast - logical_quant_dim = tensor.shape[axis] * (2 if tensor.dtype == torch.uint8 else 1) - tensor = tensor.transpose(axis, tensor.ndim - 1).contiguous() - scale = scale.transpose(axis, scale.ndim - 1).contiguous() - out = torch.empty( - (*tensor.shape[:-1], logical_quant_dim), dtype=dtype, device=tensor.device - ) - reshaped_out = out.view(-1, out.shape[-1]) - reshaped_tensor = tensor.view(-1, tensor.shape[-1]) - reshaped_scale = scale.view(-1, scale.shape[-1]) - BLOCK_OUT_DIM = 128 - BLOCK_QUANT_DIM = 32 - blocks_out_dim = triton.cdiv(reshaped_out.shape[0], BLOCK_OUT_DIM) - blocks_quant_dim = triton.cdiv(reshaped_out.shape[1], BLOCK_QUANT_DIM) - _upcast_from_mxfp[(blocks_out_dim, blocks_quant_dim)]( - reshaped_out, - *reshaped_out.stride(), - reshaped_scale, - *reshaped_scale.stride(), - reshaped_tensor, - *reshaped_tensor.stride(), - *reshaped_out.shape, - BLOCK_OUT_DIM, - BLOCK_QUANT_DIM, - num_warps=8, - ) - out = out.transpose(axis, scale.ndim - 1).contiguous() - return out - - def dequant_x_blockscale(x, x_scales, per_row_x_scale, group_shape): assert x_scales is not None group_shape_m, _, group_shape_k = group_shape diff --git a/aiter/ops/triton/quant/quant.py b/aiter/ops/triton/quant/quant.py index 534028ddf0..133926b0e3 100644 --- a/aiter/ops/triton/quant/quant.py +++ b/aiter/ops/triton/quant/quant.py @@ -42,6 +42,53 @@ _LOGGER = AiterTritonLogger() +def _static_per_tensor_quant_launch(qx, x_in, scale_in, fast_convert: bool): + """Shared launch for the static per-tensor quant kernel. + + Picks a tile shape from the row width: narrow rows get several rows per + program so the grid stays large enough, wide rows (the MoE expert-weight + case, where a whole row will not fit in registers) get split across + programs along the row instead. + """ + # Callers may hand in an output that still has the input's pre-flattened + # shape (per_tensor_quant_triton does exactly that: 2D x, N-D qx), so view + # both as 2D rather than trusting qx.stride(0). .view keeps it a view, so + # the kernel still writes into the caller's buffer. + x2d = x_in if x_in.ndim == 2 else x_in.view(-1, x_in.shape[-1]) + q2d = qx if qx.ndim == 2 else qx.view(-1, qx.shape[-1]) + assert x2d.shape == q2d.shape, f"{tuple(x2d.shape)=} != {tuple(q2d.shape)=}" + + rows, cols = x2d.shape + cols_pow2 = triton.next_power_of_2(cols) + if cols_pow2 >= 2048: + # Wide rows: one program per row segment, as many columns at a time as + # fit. Packing rows on top of this only shrinks the grid. + BLOCK_N = min(cols_pow2, 4096) + BLOCK_M = 1 + else: + # Narrow rows: a row per program leaves the grid too small and each + # program too short, so stack rows up to a ~2K-element tile. + BLOCK_N = min(cols_pow2, 512) + BLOCK_M = max(1, min(triton.next_power_of_2(rows), 2048 // BLOCK_N)) + grid = (triton.cdiv(rows, BLOCK_M), triton.cdiv(cols, BLOCK_N)) + _static_per_tensor_quant_fp8_i8_kernel[grid]( + q2d, + x2d, + scale_in, + rows, + cols, + x2d.stride(0), + x2d.stride(1), + q2d.stride(0), + q2d.stride(1), + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + FAST_CONVERT=fast_convert, + num_warps=4, + ) + return qx + + def static_per_tensor_quant_fp8_i8( qx: torch.Tensor, x_in: torch.Tensor, scale_in: torch.Tensor ): @@ -58,15 +105,7 @@ def static_per_tensor_quant_fp8_i8( """ _LOGGER.info(f"STAIC_PER_TENSOR_QUANT_FP8_I8: x={tuple(x_in.shape)}") assert scale_in.numel() == 1 # only single scale value - rows = x_in.shape[0] - cols = x_in.shape[1] - NUM_COL_POW2 = triton.next_power_of_2(cols) - grid = (rows,) - _static_per_tensor_quant_fp8_i8_kernel[grid]( - qx, x_in, scale_in, cols, x_in.stride(0), NUM_COL_POW2=NUM_COL_POW2 - ) - - return qx + return _static_per_tensor_quant_launch(qx, x_in, scale_in, fast_convert=True) def dynamic_per_tensor_quant_fp8_i8( @@ -88,8 +127,7 @@ def dynamic_per_tensor_quant_fp8_i8( rows = x_in.shape[0] cols = x_in.shape[1] NUM_COL_POW2 = triton.next_power_of_2(cols) - grid = (rows,) - _dynamic_per_tensor_quant_fp8_i8_kernel[grid]( + _dynamic_per_tensor_quant_fp8_i8_kernel[(rows,)]( x_in, scale_out, cols, @@ -102,9 +140,7 @@ def dynamic_per_tensor_quant_fp8_i8( ), ) - _static_per_tensor_quant_fp8_i8_kernel[grid]( - qx, x_in, scale_out, cols, x_in.stride(0), NUM_COL_POW2=NUM_COL_POW2 - ) + _static_per_tensor_quant_launch(qx, x_in, scale_out, fast_convert=True) return qx, scale_out @@ -150,7 +186,10 @@ def dynamic_per_token_quant_fp8_i8( def dynamic_mxfp4_quant( - x: torch.Tensor, scaling_mode: str = "even" + x: torch.Tensor, + scaling_mode: str = "even", + x_fp4: torch.Tensor | None = None, + blockscale_e8m0: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Quantize a tensor to MX FP4 format. @@ -160,6 +199,11 @@ def dynamic_mxfp4_quant( scaling_mode: The method to calculate MX block scaling. - "even" (default): `even_round` in `quark.torch.quantization.utils`. - etc. + x_fp4: Optional pre-allocated uint8 output of shape (M, N // 2). Allocated + here when omitted. + blockscale_e8m0: Optional pre-allocated uint8 scale output of shape + (M, N // 32). Allocated here when omitted, in a column-major layout; + callers that need row-major scales (the MoE a4w4 path) pass their own. Returns: A tuple of (x_fp4, blockscale_e8m0). """ @@ -171,12 +215,22 @@ def dynamic_mxfp4_quant( # This is fixed by spec for MXFP4. Do not tune this. MXFP4_QUANT_BLOCK_SIZE = 32 - x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device) - blockscale_e8m0 = torch.empty( - ((N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE, M), - dtype=torch.uint8, - device=x.device, - ).T + if x_fp4 is None: + x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device) + else: + assert x_fp4.shape == (M, N // 2) and x_fp4.dtype == torch.uint8 + n_scales = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE + if blockscale_e8m0 is None: + blockscale_e8m0 = torch.empty( + (n_scales, M), + dtype=torch.uint8, + device=x.device, + ).T + else: + assert ( + blockscale_e8m0.shape == (M, n_scales) + and blockscale_e8m0.dtype == torch.uint8 + ) # for large N values if M <= 32: diff --git a/op_tests/triton_tests/moe/test_moe_gemm_a16w4.py b/op_tests/triton_tests/moe/test_moe_gemm_a16w4.py index f789bb1474..fe9845a2b4 100644 --- a/op_tests/triton_tests/moe/test_moe_gemm_a16w4.py +++ b/op_tests/triton_tests/moe/test_moe_gemm_a16w4.py @@ -17,16 +17,13 @@ from aiter.ops.triton.moe.moe_routing.routing import routing # numerics utilities -from aiter.ops.triton.moe.quant_moe import ( - # downcast_to_static_fp8, - downcast_to_mxfp, - upcast_from_mxfp, -) +from aiter.ops.triton.moe.quant_moe import downcast_to_mxfp # target-specific utilities from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.shuffle import shuffle_scale_moe from aiter.ops.triton.utils.types import str_to_torch_dtype +from op_tests.triton_tests.utils.mxfp_ref import upcast_from_mxfp # --------------- # initialize data diff --git a/op_tests/triton_tests/moe/test_moe_gemm_a4w4.py b/op_tests/triton_tests/moe/test_moe_gemm_a4w4.py index f646b31051..ac3d0b2c7c 100644 --- a/op_tests/triton_tests/moe/test_moe_gemm_a4w4.py +++ b/op_tests/triton_tests/moe/test_moe_gemm_a4w4.py @@ -19,14 +19,12 @@ from aiter.ops.triton.moe.moe_routing.routing import routing # numerics utilities -from aiter.ops.triton.moe.quant_moe import ( - downcast_to_mxfp, - upcast_from_mxfp, -) +from aiter.ops.triton.moe.quant_moe import downcast_to_mxfp # target-specific utilities from aiter.ops.triton.utils._triton.arch_info import get_arch, is_fp4_avail from aiter.ops.triton.utils.shuffle import moe_weight_decode_view, shuffle_scale_moe +from op_tests.triton_tests.utils.mxfp_ref import upcast_from_mxfp def preshuffle_moe_weight(w: torch.Tensor) -> torch.Tensor: diff --git a/op_tests/triton_tests/moe/test_moe_gemm_a8w4.py b/op_tests/triton_tests/moe/test_moe_gemm_a8w4.py index c3fc8d5e07..5756cf5a51 100644 --- a/op_tests/triton_tests/moe/test_moe_gemm_a8w4.py +++ b/op_tests/triton_tests/moe/test_moe_gemm_a8w4.py @@ -21,12 +21,12 @@ from aiter.ops.triton.moe.quant_moe import ( downcast_to_mxfp, downcast_to_static_fp8, - upcast_from_mxfp, ) # target-specific utilities from aiter.ops.triton.utils._triton.arch_info import get_arch from aiter.ops.triton.utils.shuffle import moe_weight_decode_view, shuffle_scale_moe +from op_tests.triton_tests.utils.mxfp_ref import upcast_from_mxfp def preshuffle_moe_weight(w: torch.Tensor) -> torch.Tensor: diff --git a/op_tests/triton_tests/moe/test_moe_gemm_a8w8.py b/op_tests/triton_tests/moe/test_moe_gemm_a8w8.py index dfcc94e6c4..69616bc151 100644 --- a/op_tests/triton_tests/moe/test_moe_gemm_a8w8.py +++ b/op_tests/triton_tests/moe/test_moe_gemm_a8w8.py @@ -20,12 +20,12 @@ downcast_to_mxfp, downcast_to_static_fp8, downcast_to_static_fp8_3d, - upcast_from_mxfp, ) # target-specific utilities from aiter.ops.triton.utils._triton.arch_info import get_arch from aiter.ops.triton.utils.shuffle import shuffle_scale_moe +from op_tests.triton_tests.utils.mxfp_ref import upcast_from_mxfp # --------------- # initialize data diff --git a/op_tests/triton_tests/utils/mxfp_ref.py b/op_tests/triton_tests/utils/mxfp_ref.py new file mode 100644 index 0000000000..3f63a853a5 --- /dev/null +++ b/op_tests/triton_tests/utils/mxfp_ref.py @@ -0,0 +1,236 @@ +"""MX dequantization used to build bf16 references in the MoE GEMM tests. + +``downcast_to_mxfp`` throws away precision, so a bf16 torch reference fed the +original tensor would charge the kernel for quantization error that the kernel +did not cause. These tests instead round-trip the quantized tensor back to bf16 +and hand the reference exactly the values the kernel will read, which is what +lets them keep tight tolerances. + +This is test scaffolding, not an inference path: nothing under ``aiter/`` calls +it. It lived in ``aiter/ops/triton/moe/quant_moe.py`` until it was moved here. +""" + +import torch +import triton +import triton.language as tl + +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr + +_upcast_from_mxfp_repr = make_kernel_repr( + "_upcast_from_mxfp", + [ + "BLOCK_SIZE_OUT_DIM", + "BLOCK_SIZE_QUANT_DIM", + ], +) + + +@triton.jit(repr=_upcast_from_mxfp_repr) +def _upcast_from_mxfp( + out_ptr, + stride_o_outer, + stride_o_quant: tl.constexpr, + mx_scale_ptr, + stride_scale_outer, + stride_scale_quant, + mx_tensor_ptr, + stride_tensor_outer, + stride_tensor_quant: tl.constexpr, + outer_dim, + quant_dim, + BLOCK_SIZE_OUT_DIM: tl.constexpr, + BLOCK_SIZE_QUANT_DIM: tl.constexpr, +): + + tl.static_assert( + stride_o_quant == 1, "the weight must be contiguous in the k dimension for mx" + ) + tl.static_assert( + BLOCK_SIZE_QUANT_DIM % 32 == 0, "BLOCK_SIZE_K must be a multiple of 32" + ) + # uint8 signifies two fp4 e2m1 values packed into a single byte + mx_tensor_dtype: tl.constexpr = mx_tensor_ptr.dtype.element_ty + dst_dtype: tl.constexpr = out_ptr.dtype.element_ty + tl.static_assert(dst_dtype == tl.float16 or dst_dtype == tl.bfloat16) + tl.static_assert( + mx_tensor_dtype == tl.uint8 + or ( + (mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5) + or mx_tensor_dtype == dst_dtype + ), + "mx_tensor_ptr must be uint8 or float8 or dst_dtype", + ) + tl.static_assert( + mx_scale_ptr.dtype.element_ty == tl.uint8, "mx_scale_ptr must be uint8" + ) + + # Determine if we are dealing with fp8 types. + is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8 + is_fp8: tl.constexpr = ( + mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5 + ) + K_DIVISOR: tl.constexpr = 2 if is_fp4 else 1 + BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 32 + BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // K_DIVISOR + + # Compute starting indices for the quantized (packed) dimension and the outer dimension. + outer_block = tl.program_id(0).to(tl.int64) + quant_block = tl.program_id(1).to(tl.int64) + + start_mxt_quant = quant_block * BLOCK_SIZE_QUANT_MX_TENSOR + start_out_quant = quant_block * BLOCK_SIZE_QUANT_DIM + start_mx_scale_quant = quant_block * BLOCK_SIZE_QUANT_MX_SCALE + start_out = outer_block * BLOCK_SIZE_OUT_DIM + + mx_tensor_ptr += ( + start_mxt_quant * stride_tensor_quant + start_out * stride_tensor_outer + ) + mx_scale_ptr += ( + start_mx_scale_quant * stride_scale_quant + start_out * stride_scale_outer + ) + out_ptr += start_out * stride_o_outer + start_out_quant * stride_o_quant + + # Compute offsets and masks. + offs_src_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :].to(tl.int64) + offs_out_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :].to(tl.int64) + offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None].to(tl.int64) + offs_scale = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :].to(tl.int64) + + mask_outer = start_out + offs_outer < outer_dim + mask_out_quant = start_out_quant + offs_out_quant < quant_dim + full_mask_out = mask_out_quant & mask_outer + + mask_src_quant = start_mxt_quant + offs_src_quant < tl.cdiv(quant_dim, K_DIVISOR) + full_mask_src = mask_src_quant & mask_outer + + mask_scale = start_mx_scale_quant + offs_scale < tl.cdiv(quant_dim, 32) + full_scale_mask = mask_scale & mask_outer + + tensor_offsets = ( + offs_src_quant * stride_tensor_quant + offs_outer * stride_tensor_outer + ) + scale_offsets = offs_scale * stride_scale_quant + offs_outer * stride_scale_outer + out_offsets = offs_out_quant * stride_o_quant + offs_outer * stride_o_outer + + # Load the packed tensor and scale. + tensor = tl.load(mx_tensor_ptr + tensor_offsets, mask=full_mask_src) + scale = tl.load(mx_scale_ptr + scale_offsets, mask=full_scale_mask) + + # Upcast the scale to the destination type. + if dst_dtype == tl.bfloat16: + dst_scale = (scale.to(tl.uint16) << 7).to(dst_dtype, bitcast=True) + else: + tl.static_assert(dst_dtype == tl.float16) + dst_scale = (scale.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + dst_scale = dst_scale.to(tl.float16) + + # Now upcast the tensor. + if is_fp8: + dst_tensor = tensor.to(dst_dtype) + if tensor.dtype == tl.float8e5: + from_e_bits: tl.constexpr = 5 + from_m_bits: tl.constexpr = 2 + to_e_bits: tl.constexpr = 8 if dst_dtype == tl.bfloat16 else 5 + to_m_bits: tl.constexpr = 7 if dst_dtype == tl.bfloat16 else 10 + + # Preserve infs and nans. FIXME Fp8E5M2_to_Bf16 doesn't preserve them! + non_finite_mask_src: tl.constexpr = ((1 << from_e_bits) - 1) << from_m_bits + non_finite_mask_dst: tl.constexpr = ((1 << to_e_bits) - 1) << to_m_bits + dst_tensor = tl.where( + (tensor.to(tl.uint8, bitcast=True) & non_finite_mask_src) + == non_finite_mask_src, + (dst_tensor.to(tl.uint16, bitcast=True) | non_finite_mask_dst).to( + dst_dtype, bitcast=True + ), + dst_tensor, + ) + else: + assert is_fp4 + dst_bias: tl.constexpr = 127 if dst_dtype == tl.bfloat16 else 15 + dst_0p5: tl.constexpr = 16128 if dst_dtype == tl.bfloat16 else 0x3800 + dst_m_bits: tl.constexpr = 7 if dst_dtype == tl.bfloat16 else 10 + # e2m1 + em0 = tensor & 0x07 + em1 = tensor & 0x70 + x0 = (em0.to(tl.uint16) << (dst_m_bits - 1)) | ( + (tensor & 0x08).to(tl.uint16) << 12 + ) + x1 = (em1.to(tl.uint16) << (dst_m_bits - 5)) | ( + (tensor & 0x80).to(tl.uint16) << 8 + ) + # Three cases: + # 1) x is normal and non-zero: Correct bias + x0 = tl.where((em0 & 0x06) != 0, x0 + ((dst_bias - 1) << dst_m_bits), x0) + x1 = tl.where((em1 & 0x60) != 0, x1 + ((dst_bias - 1) << dst_m_bits), x1) + # 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in the dst type + x0 = tl.where(em0 == 0x01, dst_0p5 | (x0 & 0x8000), x0) + x1 = tl.where(em1 == 0x10, dst_0p5 | (x1 & 0x8000), x1) + # 3) x is zero, do nothing + dst_tensor = tl.interleave(x0, x1).to(dst_dtype, bitcast=True) + + # Reshape for proper broadcasting: the scale was stored with a 32-sized "inner" grouping. + dst_tensor = dst_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32]) + dst_scale = dst_scale.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1]) + scale = scale.reshape(dst_scale.shape) + + out_tensor = dst_tensor * dst_scale + # Correct any NaNs encoded via the scale. + out_tensor = tl.where(scale == 0xFF, float("nan"), out_tensor) + out_tensor = out_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM]) + tl.store(out_ptr + out_offsets, out_tensor, mask=full_mask_out) + + +def upcast_from_mxfp( + tensor: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, axis: int +): + """ + Upcasts an mxfp (packed) weight tensor back to float16 or bfloat16. + + The function assumes that the tensors were quantized along the given axis. + It permutes the tensor so that the quantized axis is last, reshapes to 2D, + launches the Triton upcast kernel, and then unpermutes back to the original order. + """ + ndim = tensor.ndim + assert -ndim <= axis < ndim, f"Invalid axis {axis=}" + axis = axis if axis >= 0 else axis + ndim + assert tensor.ndim == scale.ndim, ( + f"Weight and scale must have the same number of dimensions. " + f"Got {tensor.ndim=} and {scale.ndim=}" + ) + # dtype checks + assert tensor.dtype in { + torch.uint8, + torch.float8_e5m2, + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + }, f"Invalid tensor dtype {tensor.dtype=}" + assert scale.dtype == torch.uint8, f"Invalid scale dtype {scale.dtype=}" + assert dtype in (torch.float16, torch.bfloat16), f"Invalid output dtype {dtype=}" + # upcast + logical_quant_dim = tensor.shape[axis] * (2 if tensor.dtype == torch.uint8 else 1) + tensor = tensor.transpose(axis, tensor.ndim - 1).contiguous() + scale = scale.transpose(axis, scale.ndim - 1).contiguous() + out = torch.empty( + (*tensor.shape[:-1], logical_quant_dim), dtype=dtype, device=tensor.device + ) + reshaped_out = out.view(-1, out.shape[-1]) + reshaped_tensor = tensor.view(-1, tensor.shape[-1]) + reshaped_scale = scale.view(-1, scale.shape[-1]) + BLOCK_OUT_DIM = 128 + BLOCK_QUANT_DIM = 32 + blocks_out_dim = triton.cdiv(reshaped_out.shape[0], BLOCK_OUT_DIM) + blocks_quant_dim = triton.cdiv(reshaped_out.shape[1], BLOCK_QUANT_DIM) + _upcast_from_mxfp[(blocks_out_dim, blocks_quant_dim)]( + reshaped_out, + *reshaped_out.stride(), + reshaped_scale, + *reshaped_scale.stride(), + reshaped_tensor, + *reshaped_tensor.stride(), + *reshaped_out.shape, + BLOCK_OUT_DIM, + BLOCK_QUANT_DIM, + num_warps=8, + ) + out = out.transpose(axis, scale.ndim - 1).contiguous() + return out From f9a812d0c6db7a870bb8ab318939d0b0e745e1d2 Mon Sep 17 00:00:00 2001 From: Vinayak Gokhale Date: Thu, 3 Sep 2026 13:15:55 -0500 Subject: [PATCH 2/3] [Triton] Trim the comments added by the MoE quant dedup Comment-only follow-up to a8968e2af. The two constexpr flags each carried a long docstring explaining both of their branches; move each down to a short comment at the branch it actually governs (POW2_SCALE in _compute_mx_quant_and_scale, FAST_CONVERT in _static_per_tensor_quant_fp8_i8_kernel) and drop the rest. Same for the wrapper-level prose on mxfp4_quant, downcast_to_static_fp8, downcast_to_mxfp, _static_per_tensor_quant_launch and _mxfp8_quant_op's LOG2_DTYPE_MAX note. No functional change. Reverified on gfx950: both static-fp8 paths still bit-identical to the kernels they replaced over 9,431,040 elements x 307 scales, mxfp4_quant bit-identical, the N-D output path still correct for 3-D and 4-D across fp8/int8 and static/dynamic, and op_tests/triton_tests/{moe,quant} plus test_activation.py unchanged at 6897 passed / 6396 skipped. --- .../triton/_triton_kernels/moe/quant_moe.py | 23 ++++--------------- .../ops/triton/_triton_kernels/quant/quant.py | 21 +++-------------- aiter/ops/triton/moe/moe_op_gemm_a4w4.py | 7 ------ aiter/ops/triton/moe/quant_moe.py | 12 ---------- aiter/ops/triton/quant/quant.py | 20 ++++------------ 5 files changed, 11 insertions(+), 72 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/moe/quant_moe.py b/aiter/ops/triton/_triton_kernels/moe/quant_moe.py index d8aca59dbb..163e314509 100644 --- a/aiter/ops/triton/_triton_kernels/moe/quant_moe.py +++ b/aiter/ops/triton/_triton_kernels/moe/quant_moe.py @@ -36,25 +36,6 @@ def _compute_mx_quant_and_scale( DEQUANT_SCALE_ROUNDING_MODE: tl.constexpr = 0, POW2_SCALE: tl.constexpr = False, ): - """Quantize a [OUT_DIM, QUANT_DIM] fp32/bf16/fp16 tile to MXFP4 or MXFP8. - - Two scale-derivation schemes, selected by POW2_SCALE: - - * ``POW2_SCALE=False`` (default): ``scale = amax / dtype_max`` with the - exponent rounded per DEQUANT_SCALE_ROUNDING_MODE. Never saturates, but - leaves up to 2x of the dtype's range unused when amax sits just above a - power of two. - * ``POW2_SCALE=True``: the ``even_round`` scheme shared with - :func:`_mxfp4_quant_op` / :func:`_mxfp8_quant_op` -- amax is rounded to the - nearest power of two and the scale is taken relative to the largest power - of two the dtype holds (4 for e2m1, 256 for e4m3, 32768 for e5m2). Uses - the range better in the bulk at the cost of clipping the odd outlier. - - The two are NOT bit-compatible: they disagree on roughly 12% of fp4 block - scales and 0.3% of fp8 ones, so a tensor quantized with one must be - dequantized against the same one. DEQUANT_SCALE_ROUNDING_MODE is ignored - when POW2_SCALE is set. - """ is_fp8: tl.constexpr = ( mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5 ) @@ -62,6 +43,10 @@ def _compute_mx_quant_and_scale( BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1] BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // 32 + # POW2_SCALE picks the scale scheme: False (default) is amax / dtype_max with + # the exponent rounded per DEQUANT_SCALE_ROUNDING_MODE, True is the even_round + # scheme shared with _mxfp4_quant_op / _mxfp8_quant_op. The two are not + # bit-compatible, so a tensor must be dequantized against whichever made it. if POW2_SCALE: # Padding lanes are zeroed rather than set to -1: zero is neutral for the # group amax and is also what the tile stores for them either way. diff --git a/aiter/ops/triton/_triton_kernels/quant/quant.py b/aiter/ops/triton/_triton_kernels/quant/quant.py index 8cb6b0728f..2ad64a3509 100644 --- a/aiter/ops/triton/_triton_kernels/quant/quant.py +++ b/aiter/ops/triton/_triton_kernels/quant/quant.py @@ -20,21 +20,6 @@ def _static_per_tensor_quant_fp8_i8_kernel( BLOCK_N: tl.constexpr, FAST_CONVERT: tl.constexpr, ): - """Quantize x by a single tensor-wide scale into the dtype of ``qx_ptr`` - (fp8 e4m3/e5m2 or int8). - - FAST_CONVERT picks how the scale is applied: - - * ``True`` -- multiply by the reciprocal. One v_mul per element, the - reciprocal computed once per program. - * ``False`` -- divide. A correctly-rounded fp32 division per element, so it - matches an ``x / scale`` reference exactly. - - The two agree on nearly every input but not all of them: with scale=448.0 a - dense bf16 sweep puts them 1 fp8 ulp apart on ~0.2% of elements, because a - reciprocal that is half an fp32 ulp off can land on the far side of an fp8 - rounding boundary. - """ # Fold the block origin into the base pointers in int64 so only the in-tile # offsets, which always fit, stay 32-bit. start_m = tl.program_id(axis=0).to(tl.int64) * BLOCK_M @@ -53,6 +38,9 @@ def _static_per_tensor_quant_fp8_i8_kernel( ) scale = tl.load(scale_in_ptr) + # FAST_CONVERT multiplies by the reciprocal, one v_mul per element, instead of + # a correctly-rounded division. Not equivalent: a reciprocal half an fp32 ulp + # off can land on the far side of an fp8 rounding boundary. if FAST_CONVERT: qx = x * (1 / scale) else: @@ -421,9 +409,6 @@ def _mxfp8_quant_op( returns (scale_e8m0, quant_scale): the per-group uint8 e8m0 scale and the matching fp32 multiplicative scale. Both outputs keep QUANT_AXIS with size 1 so they broadcast against the input for in-place quantization. - - LOG2_DTYPE_MAX is log2 of the largest power of two the target dtype holds: - 8 for e4m3 (max 448 -> 256) and 15 for e5m2 (max 57344 -> 32768). """ amax = tl.max(tl.abs(x_grouped), axis=QUANT_AXIS, keep_dims=True) amax_i32 = amax.to(tl.int32, bitcast=True) diff --git a/aiter/ops/triton/moe/moe_op_gemm_a4w4.py b/aiter/ops/triton/moe/moe_op_gemm_a4w4.py index f130f945e2..af53aa8c1e 100644 --- a/aiter/ops/triton/moe/moe_op_gemm_a4w4.py +++ b/aiter/ops/triton/moe/moe_op_gemm_a4w4.py @@ -183,13 +183,6 @@ def mxfp4_quant( x: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Quantize a 2D tensor `x` of shape [M, N] (bf16/fp16/fp32) to MXFP4 (E2M1) format - quantized along the N dimension. - - Thin wrapper over :func:`aiter.ops.triton.quant.quant.dynamic_mxfp4_quant`; it - exists only to hand that op the row-major scale buffer the a4w4 GEMM expects - (``dynamic_mxfp4_quant`` allocates a column-major one by default). - Returns: - A packed MXFP4 tensor `x_fp4` of shape [M, N // 2] (stored as uint8), where each byte stores two 4-bit values. diff --git a/aiter/ops/triton/moe/quant_moe.py b/aiter/ops/triton/moe/quant_moe.py index 9e5d71f608..c696f1601d 100644 --- a/aiter/ops/triton/moe/quant_moe.py +++ b/aiter/ops/triton/moe/quant_moe.py @@ -24,13 +24,6 @@ def downcast_to_static_fp8_3d(x: torch.Tensor, scale: torch.Tensor): def downcast_to_static_fp8(x: torch.Tensor, scale: torch.Tensor): - """Quantize ``x`` by a single tensor-wide ``scale`` to the arch's fp8 e4m3. - - Uses the exact-division form of the shared static per-tensor quant kernel - (``FAST_CONVERT=False``), which is what this path has always done; the - reciprocal form used by ``static_per_tensor_quant_fp8_i8`` differs on a - small fraction of inputs. - """ M, N = x.shape if get_arch() != "gfx942": dtype = torch.float8_e4m3fn @@ -60,11 +53,6 @@ def downcast_to_mxfp( If weight_quant_type is torch.float8_e4m3fn or torch.float8_e5m2, we output mxfp8 with the float8s are stored in their respective formats. - - ``pow2_scale`` selects the ``even_round`` scale scheme shared with - ``dynamic_mxfp4_quant`` / ``dynamic_mxfp8_quant`` instead of the default - ``amax / dtype_max`` one; see :func:`_compute_mx_quant_and_scale`. The two - disagree, so a tensor quantized one way must be dequantized the same way. """ ndim = src_tensor.ndim assert -ndim <= axis < ndim, f"Invalid axis {axis=}" diff --git a/aiter/ops/triton/quant/quant.py b/aiter/ops/triton/quant/quant.py index 133926b0e3..5869689b40 100644 --- a/aiter/ops/triton/quant/quant.py +++ b/aiter/ops/triton/quant/quant.py @@ -43,17 +43,8 @@ def _static_per_tensor_quant_launch(qx, x_in, scale_in, fast_convert: bool): - """Shared launch for the static per-tensor quant kernel. - - Picks a tile shape from the row width: narrow rows get several rows per - program so the grid stays large enough, wide rows (the MoE expert-weight - case, where a whole row will not fit in registers) get split across - programs along the row instead. - """ - # Callers may hand in an output that still has the input's pre-flattened - # shape (per_tensor_quant_triton does exactly that: 2D x, N-D qx), so view - # both as 2D rather than trusting qx.stride(0). .view keeps it a view, so - # the kernel still writes into the caller's buffer. + # per_tensor_quant_triton hands in a 2D x with an N-D qx, so view both as 2D + # rather than trusting qx.stride(0); .view still writes the caller's buffer. x2d = x_in if x_in.ndim == 2 else x_in.view(-1, x_in.shape[-1]) q2d = qx if qx.ndim == 2 else qx.view(-1, qx.shape[-1]) assert x2d.shape == q2d.shape, f"{tuple(x2d.shape)=} != {tuple(q2d.shape)=}" @@ -199,11 +190,8 @@ def dynamic_mxfp4_quant( scaling_mode: The method to calculate MX block scaling. - "even" (default): `even_round` in `quark.torch.quantization.utils`. - etc. - x_fp4: Optional pre-allocated uint8 output of shape (M, N // 2). Allocated - here when omitted. - blockscale_e8m0: Optional pre-allocated uint8 scale output of shape - (M, N // 32). Allocated here when omitted, in a column-major layout; - callers that need row-major scales (the MoE a4w4 path) pass their own. + x_fp4, blockscale_e8m0: Optional pre-allocated uint8 outputs, shaped + (M, N // 2) and (M, N // 32); allocated column-major when omitted. Returns: A tuple of (x_fp4, blockscale_e8m0). """ From 7965462f9c32fe6cc5e9d77e3a7c40ec8b94a08b Mon Sep 17 00:00:00 2001 From: Vinayak Gokhale Date: Thu, 3 Sep 2026 15:55:51 -0500 Subject: [PATCH 3/3] [Triton] Fold the static per-tensor quant launcher into its wrapper The dedup left three names one underscore apart: the kernel (_static_per_tensor_quant_fp8_i8_kernel), the public wrapper (static_per_tensor_quant_fp8_i8) and a private launcher between them holding the tiling heuristic and the FAST_CONVERT choice. That middle tier is not a layer this tree has anywhere else -- kernels live in _triton_kernels/, host wrappers in the op module, and nothing sits in between -- so remove it. fast_convert becomes a keyword argument on static_per_tensor_quant_fp8_i8, defaulting to True so the existing positional callers (aiter/ops/quant.py) are unaffected, and the tiling and launch move into the wrapper body alongside every other op in that file. dynamic_per_tensor_quant_fp8_i8 and the MoE downcast_to_static_fp8 now both go through the public wrapper. It reads better as a public knob anyway: reciprocal-multiply vs exact division is a real accuracy/speed choice for a caller to make. Side effect: the dynamic path emits a second INFO log line, since it now calls the public wrapper rather than the kernel. That is gated behind AITER_TRITON_LOG_LEVEL (default WARNING) and is arguably accurate -- the dynamic op does perform a static quant internally. No behavior change otherwise. Reverified on gfx950: both scale-application paths still bit-identical to the kernels they replaced over 9,431,040 elements x 307 scales, int8 included; the N-D output path still correct for 3-D and 4-D across fp8/int8 and static/dynamic; op_tests/triton_tests/{moe,quant} plus test_activation.py at 6897 passed / 6396 skipped. --- aiter/ops/triton/moe/quant_moe.py | 4 +-- aiter/ops/triton/quant/quant.py | 44 ++++++++++++++++--------------- 2 files changed, 25 insertions(+), 23 deletions(-) diff --git a/aiter/ops/triton/moe/quant_moe.py b/aiter/ops/triton/moe/quant_moe.py index c696f1601d..192ded345c 100644 --- a/aiter/ops/triton/moe/quant_moe.py +++ b/aiter/ops/triton/moe/quant_moe.py @@ -8,7 +8,7 @@ _smoothquant_fuse_quant_kernel, _smoothquant_fuse_quant_kernel_single_pass, ) -from aiter.ops.triton.quant.quant import _static_per_tensor_quant_launch +from aiter.ops.triton.quant.quant import static_per_tensor_quant_fp8_i8 from aiter.ops.triton.utils._triton.arch_info import get_arch @@ -30,7 +30,7 @@ def downcast_to_static_fp8(x: torch.Tensor, scale: torch.Tensor): else: dtype = torch.float8_e4m3fnuz y = torch.empty((M, N), dtype=dtype, device=x.device) - return _static_per_tensor_quant_launch(y, x, scale, fast_convert=False) + return static_per_tensor_quant_fp8_i8(y, x, scale, fast_convert=False) class DequantScaleRoundingMode(Enum): diff --git a/aiter/ops/triton/quant/quant.py b/aiter/ops/triton/quant/quant.py index 5869689b40..67258a8569 100644 --- a/aiter/ops/triton/quant/quant.py +++ b/aiter/ops/triton/quant/quant.py @@ -42,7 +42,28 @@ _LOGGER = AiterTritonLogger() -def _static_per_tensor_quant_launch(qx, x_in, scale_in, fast_convert: bool): +def static_per_tensor_quant_fp8_i8( + qx: torch.Tensor, + x_in: torch.Tensor, + scale_in: torch.Tensor, + fast_convert: bool = True, +): + """ + Quantizes tensor using the provided scale to int8 or fp8 + + Parameters: + - qx: Output tensor of same shape as x_in. Must be fp8 or int8 dtype and allocated by the caller + - x_in: Input tensor of shape (M, N). + - scale_in: Input Scale tensor of shape (1,) and dtype fp32 + - fast_convert: multiply by the reciprocal of the scale instead of dividing + by it. Cheaper, and differs from the division on a small fraction of + inputs; see the kernel. + + Returns: + - qx: Quantized output values. + """ + _LOGGER.info(f"STAIC_PER_TENSOR_QUANT_FP8_I8: x={tuple(x_in.shape)}") + assert scale_in.numel() == 1 # only single scale value # per_tensor_quant_triton hands in a 2D x with an N-D qx, so view both as 2D # rather than trusting qx.stride(0); .view still writes the caller's buffer. x2d = x_in if x_in.ndim == 2 else x_in.view(-1, x_in.shape[-1]) @@ -80,25 +101,6 @@ def _static_per_tensor_quant_launch(qx, x_in, scale_in, fast_convert: bool): return qx -def static_per_tensor_quant_fp8_i8( - qx: torch.Tensor, x_in: torch.Tensor, scale_in: torch.Tensor -): - """ - Quantizes tensor using the provided scale to int8 or fp8 - - Parameters: - - qx: Output tensor of same shape as x_in. Must be fp8 or int8 dtype and allocated by the caller - - x_in: Input tensor of shape (M, N). - - scale_in: Input Scale tensor of shape (1,) and dtype fp32 - - Returns: - - qx: Quantized output values. - """ - _LOGGER.info(f"STAIC_PER_TENSOR_QUANT_FP8_I8: x={tuple(x_in.shape)}") - assert scale_in.numel() == 1 # only single scale value - return _static_per_tensor_quant_launch(qx, x_in, scale_in, fast_convert=True) - - def dynamic_per_tensor_quant_fp8_i8( qx: torch.Tensor, x_in: torch.Tensor, scale_out: torch.Tensor ): @@ -131,7 +133,7 @@ def dynamic_per_tensor_quant_fp8_i8( ), ) - _static_per_tensor_quant_launch(qx, x_in, scale_out, fast_convert=True) + static_per_tensor_quant_fp8_i8(qx, x_in, scale_out) return qx, scale_out