diff --git a/alto/kernels/mx/__init__.py b/alto/kernels/mx/__init__.py new file mode 100644 index 00000000..cca0f020 --- /dev/null +++ b/alto/kernels/mx/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""MX packed-quantization Triton kernels for the MX6 and MX9 block formats. + +Fake-quant reference / emulation lives in ``alto.modifiers.quantization.mx``. +This package provides GPU pack/unpack only. +""" diff --git a/alto/kernels/mx/_mx_common.py b/alto/kernels/mx/_mx_common.py new file mode 100644 index 00000000..394de0a2 --- /dev/null +++ b/alto/kernels/mx/_mx_common.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Shared MX packed-quantization helpers, device and host. + +MX6 and MX9 share the same block algorithm -- exponent extraction, prime/pair +demotion, and the power-of-two scale (with its -126 clamp) -- and differ only in +element integer width (``quant_bit`` = 8 for MX9, 5 for MX6) and byte packing. +The width-/packing-independent pieces are centralised here so both formats use +identical math and an identical host path; only the per-format ``_pack_*`` / +``_unpack_*``, the grid kernels, and the byte layout stay in their own modules. +""" + +import math + +import torch +import triton +import triton.language as tl + +_TORCH_TO_TL = { + torch.float32: tl.float32, + torch.bfloat16: tl.bfloat16, +} + + +@triton.jit +def _floor_exp(x): + """Extract the unbiased exponent field directly from native float bits. + """ + if x.type.element_ty == tl.float32: + bits = x.to(tl.int32, bitcast=True) + return (((bits >> 23) & 0xFF) - 127).to(tl.int32) + elif x.type.element_ty == tl.bfloat16: + bits = x.to(tl.int16, bitcast=True) + return (((bits >> 7) & 0xFF) - 127).to(tl.int32) + else: + tl.static_assert(False, "x must be fp32 / bf16") + + +@triton.jit +def _round_half_even(y): + """Round-to-nearest-ties-to-even (matches torch.round).""" + rounded = tl.floor(y + 0.5) + is_tie = (y - tl.floor(y)) == 0.5 + is_odd = (rounded - 2.0 * tl.floor(rounded * 0.5)) == 1.0 + return tl.where(is_tie & is_odd, rounded - 1.0, rounded) + + +@triton.jit +def _calculate_mx_exp( + x, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, +): + """Per-block shared exponent, per-element shared_exp, and per-pair prime bits. + + x is [BPP, BLOCK_SIZE] in native dtype; do NOT pre-cast to fp32 on the host + since exponent extraction must match the per-element branch in _floor_exp. + + Returns (shared_exp [BPP, BLOCK_SIZE], max_exp [BPP], pair [BPP, N_PAIRS]), + all int32. pair[i,j]=1 means that pair gets a 1-exponent demotion. + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + + # Sanitize the exponent statistics only. Quantization keeps the original x, + # so Inf is clamped rather than replaced by 0. + clean = tl.where(x != x, 0.0, x) + clean = tl.where(tl.abs(clean) == float("inf"), 0.0, clean) + + # tl.max on some backends promotes the reduce to fp32; cast back to native + # dtype so _floor_exp takes the same branch as the per-element call below. + amax = tl.max(tl.abs(clean), axis=1, keep_dims=True).to(clean.dtype) + max_exp = _floor_exp(amax) + + t_exp = _floor_exp(clean) + demote = (max_exp - t_exp) >= 1 + + demote_groups = demote.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP) + pair_demote = tl.sum(demote_groups, axis=2, keep_dims=True) == PRIME_GROUP + + pair_b = tl.broadcast_to(pair_demote, (BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP)) + pair_b = pair_b.reshape(BLOCKS_PER_PROG, BLOCK_SIZE) + shared_exp = max_exp - pair_b.to(tl.int32) + + pair = pair_demote.reshape(BLOCKS_PER_PROG, N_PAIRS).to(tl.int32) + return shared_exp, max_exp.reshape(BLOCKS_PER_PROG), pair + + +@triton.jit +def _scale_from_shared_exp(shared_exp, QUANT_BIT: tl.constexpr): + """2^(shared_exp - quant_bit + 2), scale exponent clamped to -126. + + Clamping to the minimum normal fp32 exponent keeps the scale nonzero and + avoids backend-dependent subnormal/FTZ behavior. The exact activation + threshold depends on ``shared_exp`` and ``QUANT_BIT``. + """ + scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) + return tl.exp2(scale_exp.to(tl.float32)) + + +@triton.jit +def _quantize_clamped(x, shared_exp, QUANT_BIT: tl.constexpr, Q_MAX: tl.constexpr): + """Round-half-even to [-Q_MAX, Q_MAX] and return fp32 q for format packing.""" + scale = _scale_from_shared_exp(shared_exp, QUANT_BIT) + q = _round_half_even(tl.div_rn(x.to(tl.float32), scale)) + return tl.minimum(tl.maximum(q, -(Q_MAX * 1.0)), Q_MAX * 1.0) + + +@triton.jit +def _pack_bits8(bits, BLOCKS_PER_PROG: tl.constexpr, N_BYTES: tl.constexpr): + """Pack a [BPP, N_BYTES*8] 0/1 tensor into [BPP, N_BYTES] uint8, LSB first.""" + weights = (1 << tl.arange(0, 8)).to(tl.int32) + g = bits.reshape(BLOCKS_PER_PROG, N_BYTES, 8) + return tl.sum(g * weights[None, None, :], axis=2).to(tl.uint8) + + +@triton.jit +def _unpack_bits8(packed, BLOCKS_PER_PROG: tl.constexpr, N_BYTES: tl.constexpr): + """Inverse of ``_pack_bits8``: [BPP, N_BYTES] uint8 -> [BPP, N_BYTES*8] 0/1 int32.""" + shifts = tl.arange(0, 8).to(tl.int32) + p = packed.to(tl.int32).reshape(BLOCKS_PER_PROG, N_BYTES, 1) # nonnegative, so >> is safe + bits = (p >> shifts[None, None, :]) & 1 + return bits.reshape(BLOCKS_PER_PROG, N_BYTES * 8) + + +@triton.jit +def _dequantize_decoded_q( + q, + pair, + max_exp, + out_dtype: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, +): + """Reconstruct values from an already-decoded integer ``q`` and per-pair + ``pair`` demotion bitmap: rebuild per-element shared_exp, then + y = q * 2^(shared_exp - quant_bit + 2). + + ``q`` may be int8 (MX9, straight from storage) or int32 (MX6, rebuilt from + sign + mantissa nibble); both are cast to fp32 here before scaling. + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + pair_b = tl.broadcast_to(pair[:, :, None], (BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP)) + pair_b = pair_b.reshape(BLOCKS_PER_PROG, BLOCK_SIZE) + shared_exp = max_exp - pair_b + scale = _scale_from_shared_exp(shared_exp, QUANT_BIT) + return (q.to(tl.float32) * scale).to(out_dtype) + + +def _convert_to_mx_host( + data_hp: torch.Tensor, + *, + kernel, + n_packed_bytes: int, + fmt: str, + block_size: int, + prime_group: int, + quant_bit: int, + axis: int, + blocks_per_program: int, +) -> torch.Tensor: + """Shared host path behind ``convert_to_mx6`` / ``convert_to_mx9``. + + Transposes the quant axis to the last dim and launches ``kernel`` over the + resulting blocks. No ``.contiguous()`` is forced: ``reshape`` keeps a view + where it can and the kernel gathers by stride, with the ragged tail block + masked in-kernel instead of padded here. + + ``n_packed_bytes`` sizes the output rows and is also handed to the kernel, + which static-asserts it against its own segment offsets. + """ + assert block_size == 16, f"block_size only supports 16, got {block_size}" + assert (isinstance(blocks_per_program, int) and blocks_per_program > 0 and + (blocks_per_program & (blocks_per_program - 1)) + == 0), f"blocks_per_program must be a positive power of two, got {blocks_per_program}" + assert data_hp.dtype in (torch.float32, torch.bfloat16), \ + f"{fmt} only supports fp32 / bf16, got {data_hp.dtype}" + + data_hp = data_hp.transpose(axis, -1) + last = data_hp.shape[-1] + assert last > 0, f"{fmt} requires a non-empty quantization axis" + x2d = data_hp.reshape(-1, last) + rows = x2d.shape[0] + blocks_per_row = triton.cdiv(last, block_size) + n_blocks = rows * blocks_per_row + + packed = torch.empty((n_blocks, n_packed_bytes), dtype=torch.uint8, device=data_hp.device) + + stride_row, stride_col = x2d.stride() + grid = (triton.cdiv(n_blocks, blocks_per_program),) + kernel[grid]( + x2d, + packed, + n_blocks, + last, + blocks_per_row, + stride_row, + stride_col, + BLOCK_SIZE=block_size, + BLOCKS_PER_PROG=blocks_per_program, + PRIME_GROUP=prime_group, + QUANT_BIT=quant_bit, + N_PACKED_BYTES=n_packed_bytes, + ) + return packed + + +def _convert_from_mx_host( + packed: torch.Tensor, + *, + kernel, + n_packed_bytes: int, + fmt: str, + out_dtype: torch.dtype, + out_shape, + block_size: int, + prime_group: int, + quant_bit: int, + axis: int, + blocks_per_program: int, +) -> torch.Tensor: + """Shared host path behind ``convert_from_mx6`` / ``convert_from_mx9``. + + ``packed`` is passed by stride, so an existing non-contiguous view is consumed + without a copy. The packed bytes contain no shape or axis metadata: + ``out_shape`` / ``axis`` must exactly match the convert_to_* call that + produced them. Some mismatches have the same block count and cannot be + detected, so violating this requirement may silently reorder values. + """ + assert block_size == 16, f"block_size only supports 16, got {block_size}" + assert (isinstance(blocks_per_program, int) and blocks_per_program > 0 and + (blocks_per_program & (blocks_per_program - 1)) + == 0), f"blocks_per_program must be a positive power of two, got {blocks_per_program}" + assert out_dtype in _TORCH_TO_TL, \ + f"out_dtype must be one of {tuple(_TORCH_TO_TL)}, got {out_dtype}" + assert packed.dtype == torch.uint8, f"packed dtype must be uint8, got {packed.dtype}" + assert packed.ndim == 2 and packed.shape[1] == n_packed_bytes, \ + f"{fmt} packed tensor must be [n_blocks, {n_packed_bytes}], got {tuple(packed.shape)}" + + n_blocks = packed.shape[0] + transposed_shape = list(out_shape) + assert transposed_shape, f"{fmt} out_shape must have at least one dimension" + assert -len(transposed_shape) <= axis < len(transposed_shape), \ + f"axis {axis} is out of range for out_shape={tuple(out_shape)}" + transposed_shape[axis], transposed_shape[-1] = transposed_shape[-1], transposed_shape[axis] + last = transposed_shape[-1] + assert last > 0, f"{fmt} requires a non-empty quantization axis" + rows = math.prod(transposed_shape[:-1]) + padded_cols = triton.cdiv(last, block_size) * block_size + assert rows * padded_cols == n_blocks * block_size, ( + f"out_shape/axis/block_size inconsistent with packed tensor: " + f"inferred rows*padded_cols={rows * padded_cols}, " + f"but packed holds n_blocks*block_size={n_blocks * block_size}") + + y_blocks = torch.empty((n_blocks, block_size), dtype=out_dtype, device=packed.device) + + stride_blk, stride_col = packed.stride() + grid = (triton.cdiv(n_blocks, blocks_per_program),) + kernel[grid]( + packed, + y_blocks, + n_blocks, + stride_blk, + stride_col, + BLOCK_SIZE=block_size, + BLOCKS_PER_PROG=blocks_per_program, + PRIME_GROUP=prime_group, + QUANT_BIT=quant_bit, + OUT_DTYPE=_TORCH_TO_TL[out_dtype], + N_PACKED_BYTES=n_packed_bytes, + ) + y2d = y_blocks.reshape(rows, padded_cols) + y2d = y2d[:, :last] + return y2d.reshape(transposed_shape).transpose(axis, -1) diff --git a/alto/kernels/mx/mx6_quantization.py b/alto/kernels/mx/mx6_quantization.py new file mode 100644 index 00000000..da3d8bfb --- /dev/null +++ b/alto/kernels/mx/mx6_quantization.py @@ -0,0 +1,319 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Packed MX6 quantization: each 16-element block becomes 12 uint8 bytes. + +Per-block layout ``[max_exp(1B) | prime(1B) | sign(2B) | mantissa(8B)]`` += 6 bits/element, byte-compatible with Quark's MX6 export packer (not vendored +here; its ``idx2`` bitmap is ``prime`` below): + - ``max_exp`` : E8M0 (true exponent + 127) + - ``prime`` : 1 bit per pair of 2 elements, set when the pair takes a + 1-exponent demotion + - ``sign`` : 16 sign bits, LSB = element 0 + - ``mantissa`` : two ``q & 0xF`` nibbles per byte, low nibble = even element + +The mantissa nibbles are the low bits of the signed integer, not ``abs(q)``. + +MX6 runs the same block algorithm as MX9 and differs only in element width +(``QUANT_BIT`` = 5, so integers clamp to +/-15) and this byte packing; the shared +math lives in ``_mx_common`` and follows the MX6/MX9 fake-quant port in +``alto/modifiers/quantization/mx.py``. +""" + +import torch +import triton +import triton.language as tl + +from ._mx_common import ( + _calculate_mx_exp, + _quantize_clamped, + _pack_bits8, + _unpack_bits8, + _dequantize_decoded_q, + _convert_to_mx_host, + _convert_from_mx_host, +) + +BLOCK_SIZE = 16 +QUANT_BIT = 5 +PRIME_GROUP = 2 +BLOCKS_PER_PROG_DEFAULT = 64 + + +@triton.jit +def _pack_mx6( + x, + shared_exp, + pair, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, +): + + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + N_MANTISSA_BYTES: tl.constexpr = BLOCK_SIZE // 2 + N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 + Q_HI: tl.constexpr = (1 << (QUANT_BIT - 1)) - 1 + + q = _quantize_clamped(x, shared_exp, QUANT_BIT, Q_HI) + + qi = q.to(tl.int32) + mantissa = qi & 0xF + sign = (qi < 0).to(tl.int32) + + sign_bytes = _pack_bits8(sign, BLOCKS_PER_PROG, N_SIGN_BYTES) + + nib_w = (1 << (4 * tl.arange(0, 2))).to(tl.int32) + mantissa_g = mantissa.reshape(BLOCKS_PER_PROG, N_MANTISSA_BYTES, 2) + mantissa_bytes = tl.sum(mantissa_g * nib_w[None, None, :], axis=2) + + prime = _pack_bits8(pair, BLOCKS_PER_PROG, N_PRIME_BYTES) + return sign_bytes, mantissa_bytes.to(tl.uint8), prime + + +@triton.jit +def _unpack_mx6( + sign_bytes, + mantissa_bytes, + max_exp, + prime, + out_dtype: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, +): + + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + N_MANTISSA_BYTES: tl.constexpr = BLOCK_SIZE // 2 + N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 + + sign = _unpack_bits8(sign_bytes, BLOCKS_PER_PROG, N_SIGN_BYTES) + + nib_shift = (4 * tl.arange(0, 2)).to(tl.int32) + mb = mantissa_bytes.to(tl.int32).reshape(BLOCKS_PER_PROG, N_MANTISSA_BYTES, 1) + mantissa = (mb >> nib_shift[None, None, :]) & 0xF + mantissa = mantissa.reshape(BLOCKS_PER_PROG, BLOCK_SIZE) + + # Two's complement on 4 bits: sign=1 with mantissa=13 means q=-3. + q = tl.where(sign == 1, mantissa - 16, mantissa) + + pair = _unpack_bits8(prime, BLOCKS_PER_PROG, N_PRIME_BYTES) + return _dequantize_decoded_q( + q, + pair, + max_exp, + out_dtype, + BLOCKS_PER_PROG, + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + ) + + +# Grid kernels + + +@triton.jit +def _convert_to_mx6_kernel( + x_ptr, + packed_ptr, + n_blocks, + last, + blocks_per_row, + stride_row, + stride_col, + BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, + N_PACKED_BYTES: tl.constexpr, +): + + tl.static_assert(BLOCK_SIZE % PRIME_GROUP == 0, "BLOCK_SIZE must be divisible by PRIME_GROUP") + tl.static_assert((BLOCK_SIZE // PRIME_GROUP) % 8 == 0, "MX6 prime bitmap requires pair count divisible by 8") + tl.static_assert(BLOCK_SIZE % 8 == 0, "MX6 sign bitmap requires BLOCK_SIZE divisible by 8") + + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + N_MANTISSA_BYTES: tl.constexpr = BLOCK_SIZE // 2 + N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 + PRIME_OFFSET: tl.constexpr = 1 + SIGN_OFFSET: tl.constexpr = PRIME_OFFSET + N_PRIME_BYTES + MANTISSA_OFFSET: tl.constexpr = SIGN_OFFSET + N_SIGN_BYTES + + # The row size the host allocated must match the segments written below, + # otherwise a layout change on one side silently writes past the row. + tl.static_assert(N_PACKED_BYTES == MANTISSA_OFFSET + N_MANTISSA_BYTES, + "N_PACKED_BYTES does not match the MX6 segment layout") + + pid = tl.program_id(0) + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) + blk_mask = blk < n_blocks + + row = blk // blocks_per_row + brow = blk % blocks_per_row + col = tl.arange(0, BLOCK_SIZE) + col_g = brow[:, None] * BLOCK_SIZE + col[None, :] + + in_range = col_g < last + load_mask = blk_mask[:, None] & in_range + in_offs = row[:, None] * stride_row + col_g * stride_col + x = tl.load(x_ptr + in_offs, mask=load_mask, other=0.0) + + shared_exp, max_exp, pair = _calculate_mx_exp(x, BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP) + sign_bytes, mantissa_bytes, prime = _pack_mx6( + x, + shared_exp, + pair, + BLOCKS_PER_PROG, + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + ) + + # E8M0 bias: store the true exponent + 127. + e_store = (max_exp + 127).to(tl.uint8) + tl.store(packed_ptr + blk * N_PACKED_BYTES, e_store, mask=blk_mask) + + pcol = tl.arange(0, N_PRIME_BYTES) + offs_p = blk[:, None] * N_PACKED_BYTES + (PRIME_OFFSET + pcol[None, :]) + tl.store(packed_ptr + offs_p, prime, mask=blk_mask[:, None]) + + signcol = tl.arange(0, N_SIGN_BYTES) + sign_offs = blk[:, None] * N_PACKED_BYTES + (SIGN_OFFSET + signcol[None, :]) + tl.store(packed_ptr + sign_offs, sign_bytes, mask=blk_mask[:, None]) + + mantcol = tl.arange(0, N_MANTISSA_BYTES) + mant_offs = blk[:, None] * N_PACKED_BYTES + (MANTISSA_OFFSET + mantcol[None, :]) + tl.store(packed_ptr + mant_offs, mantissa_bytes, mask=blk_mask[:, None]) + + +@triton.jit +def _convert_from_mx6_kernel( + packed_ptr, + y_ptr, + n_blocks, + stride_blk, + stride_col, + BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, + OUT_DTYPE: tl.constexpr, + N_PACKED_BYTES: tl.constexpr, +): + + tl.static_assert(BLOCK_SIZE % PRIME_GROUP == 0, "BLOCK_SIZE must be divisible by PRIME_GROUP") + tl.static_assert((BLOCK_SIZE // PRIME_GROUP) % 8 == 0, "MX6 prime bitmap requires pair count divisible by 8") + tl.static_assert(BLOCK_SIZE % 8 == 0, "MX6 sign bitmap requires BLOCK_SIZE divisible by 8") + + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + N_MANTISSA_BYTES: tl.constexpr = BLOCK_SIZE // 2 + N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 + PRIME_OFFSET: tl.constexpr = 1 + SIGN_OFFSET: tl.constexpr = PRIME_OFFSET + N_PRIME_BYTES + MANTISSA_OFFSET: tl.constexpr = SIGN_OFFSET + N_SIGN_BYTES + tl.static_assert(N_PACKED_BYTES == MANTISSA_OFFSET + N_MANTISSA_BYTES, + "N_PACKED_BYTES does not match the MX6 segment layout") + + pid = tl.program_id(0) + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) + col = tl.arange(0, BLOCK_SIZE) + blk_mask = blk < n_blocks + mask = blk_mask[:, None] + + # Undo the E8M0 bias to recover the true exponent. + max_exp_u = tl.load(packed_ptr + blk * stride_blk, mask=blk_mask, other=0) + max_exp = max_exp_u.to(tl.int32) - 127 + + pcol = tl.arange(0, N_PRIME_BYTES) + offs_p = blk[:, None] * stride_blk + (PRIME_OFFSET + pcol[None, :]) * stride_col + prime = tl.load(packed_ptr + offs_p, mask=mask, other=0) + + signcol = tl.arange(0, N_SIGN_BYTES) + sign_offs = blk[:, None] * stride_blk + (SIGN_OFFSET + signcol[None, :]) * stride_col + sign_bytes = tl.load(packed_ptr + sign_offs, mask=mask, other=0) + + mantcol = tl.arange(0, N_MANTISSA_BYTES) + mant_offs = blk[:, None] * stride_blk + (MANTISSA_OFFSET + mantcol[None, :]) * stride_col + mantissa_bytes = tl.load(packed_ptr + mant_offs, mask=mask, other=0) + + y = _unpack_mx6( + sign_bytes, + mantissa_bytes, + max_exp[:, None], + prime, + OUT_DTYPE, + BLOCKS_PER_PROG, + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + ) + out_offs = blk[:, None] * BLOCK_SIZE + col[None, :] + tl.store(y_ptr + out_offs, y, mask=mask) + + +# Host wrappers + + +def _mx6_packed_bytes(block_size: int) -> int: + """Return bytes per MX6 packed block.""" + n_prime_bytes = (block_size // PRIME_GROUP) // 8 + n_sign_bytes = block_size // 8 + n_mantissa_bytes = block_size // 2 + return 1 + n_prime_bytes + n_sign_bytes + n_mantissa_bytes + + +def convert_to_mx6( + data_hp: torch.Tensor, + block_size: int = BLOCK_SIZE, + axis: int = -1, + blocks_per_program: int = BLOCKS_PER_PROG_DEFAULT, +) -> torch.Tensor: + """High-precision tensor -> packed MX6 ``[n_blocks, n_packed_bytes]`` uint8. + + Blocks are formed along ``axis``, which is transposed to the last dim + internally. The original shape / axis are not stored, so the caller must pass + them back to convert_from_mx6. + """ + return _convert_to_mx_host( + data_hp, + kernel=_convert_to_mx6_kernel, + n_packed_bytes=_mx6_packed_bytes(block_size), + fmt="mx6_quantization", + block_size=block_size, + prime_group=PRIME_GROUP, + quant_bit=QUANT_BIT, + axis=axis, + blocks_per_program=blocks_per_program, + ) + + +def convert_from_mx6( + packed: torch.Tensor, + out_dtype: torch.dtype, + out_shape, + block_size: int = BLOCK_SIZE, + axis: int = -1, + blocks_per_program: int = BLOCKS_PER_PROG_DEFAULT, +) -> torch.Tensor: + """Packed MX6 tensor -> reconstructed high-precision tensor. + + ``out_shape`` and ``axis`` must exactly match the original convert_to_mx6 + call. They are not stored in ``packed``; a mismatch may silently reorder + values when it happens to produce the same block count. + """ + return _convert_from_mx_host( + packed, + kernel=_convert_from_mx6_kernel, + n_packed_bytes=_mx6_packed_bytes(block_size), + fmt="mx6_quantization", + out_dtype=out_dtype, + out_shape=out_shape, + block_size=block_size, + prime_group=PRIME_GROUP, + quant_bit=QUANT_BIT, + axis=axis, + blocks_per_program=blocks_per_program, + ) diff --git a/alto/kernels/mx/mx9_quantization.py b/alto/kernels/mx/mx9_quantization.py new file mode 100644 index 00000000..9d5d5b15 --- /dev/null +++ b/alto/kernels/mx/mx9_quantization.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Packed MX9 quantization: each 16-element block becomes 18 uint8 bytes. + +Per-block layout ``[max_exp(1B) | prime(1B) | q(16B)]`` = 9 bits/element: + - ``max_exp`` : shared exponent, E8M0 (floor(log2(amax)) + 127) + - ``prime`` : 1 bit per pair of 2 elements, set when the pair takes a + 1-exponent demotion + - ``q`` : quantized integer clamped +/-127, stored as the bit-identical + uint8 view of int8 (recover with ``.view(torch.int8)``) + +Exponent / scale / prime math follows ``alto/modifiers/quantization/mx.py``; the +parts shared with MX6 live in ``_mx_common``. + +q is clamped to +/-127 instead of widened to int16: demoted pairs use half-scale, +so a value can reach +/-128 (mx.py's quant_max=255) and overflow int8, but int16 +would cost ~17 bits/element -- more than bf16, defeating the packing. The clamp +diverges from mx.py by 1 LSB on rare elements that are both demoted and reach ++/-128, so this module's reference is the clamp-127 variant of ``mx9_fake_quantize``. + +block_size is fixed at 16 so 8 pairs pack into exactly one prime byte. Packed +integers cannot represent NaN, so q is undefined for blocks containing NaN. +""" + +import torch +import triton +import triton.language as tl + +from ._mx_common import ( + _calculate_mx_exp, + _quantize_clamped, + _pack_bits8, + _unpack_bits8, + _dequantize_decoded_q, + _convert_to_mx_host, + _convert_from_mx_host, +) + +BLOCK_SIZE = 16 +QUANT_BIT = 8 +PRIME_GROUP = 2 +BLOCKS_PER_PROG_DEFAULT = 64 + + +@triton.jit +def _pack_mx9( + x, + shared_exp, + pair, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, +): + """Quantize x to int8 q and pack pair bits into prime bytes.""" + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + Q_HI: tl.constexpr = (1 << (QUANT_BIT - 1)) - 1 + + q = _quantize_clamped(x, shared_exp, QUANT_BIT, Q_HI) + q_int = q.to(tl.int8) + + prime = _pack_bits8(pair, BLOCKS_PER_PROG, N_PRIME_BYTES) + return q_int, prime + + +@triton.jit +def _unpack_mx9( + q, + max_exp, + prime, + out_dtype: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, +): + + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + + pair = _unpack_bits8(prime, BLOCKS_PER_PROG, N_PRIME_BYTES) + return _dequantize_decoded_q( + q, + pair, + max_exp, + out_dtype, + BLOCKS_PER_PROG, + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + ) + + +# Grid kernels + + +@triton.jit +def _convert_to_mx9_kernel( + x_ptr, + packed_ptr, + n_blocks, + last, + blocks_per_row, + stride_row, + stride_col, + BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, + N_PACKED_BYTES: tl.constexpr, +): + + tl.static_assert(BLOCK_SIZE % PRIME_GROUP == 0, "BLOCK_SIZE must be divisible by PRIME_GROUP") + tl.static_assert((BLOCK_SIZE // PRIME_GROUP) % 8 == 0, "MX9 prime bitmap requires pair count divisible by 8") + + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + PRIME_OFFSET: tl.constexpr = 1 + Q_OFFSET: tl.constexpr = PRIME_OFFSET + N_PRIME_BYTES + + # The row size the host allocated must match the segments written below, + # otherwise a layout change on one side silently writes past the row. + tl.static_assert(N_PACKED_BYTES == Q_OFFSET + BLOCK_SIZE, "N_PACKED_BYTES does not match the MX9 segment layout") + + pid = tl.program_id(0) + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) + blk_mask = blk < n_blocks + + row = blk // blocks_per_row + brow = blk % blocks_per_row + col = tl.arange(0, BLOCK_SIZE) + col_g = brow[:, None] * BLOCK_SIZE + col[None, :] + + in_range = col_g < last + load_mask = blk_mask[:, None] & in_range + in_offs = row[:, None] * stride_row + col_g * stride_col + x = tl.load(x_ptr + in_offs, mask=load_mask, other=0.0) + + shared_exp, max_exp, pair = _calculate_mx_exp(x, BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP) + q_int, prime = _pack_mx9( + x, + shared_exp, + pair, + BLOCKS_PER_PROG, + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + ) + + # E8M0 bias: store the true exponent + 127. + e_store = (max_exp + 127).to(tl.uint8) + tl.store(packed_ptr + blk * N_PACKED_BYTES, e_store, mask=blk_mask) + + pcol = tl.arange(0, N_PRIME_BYTES) + offs_p = blk[:, None] * N_PACKED_BYTES + (PRIME_OFFSET + pcol[None, :]) + tl.store(packed_ptr + offs_p, prime, mask=blk_mask[:, None]) + + # Preserve the signed int8 byte representation in uint8 storage. + q_offs = blk[:, None] * N_PACKED_BYTES + (Q_OFFSET + col[None, :]) + tl.store(packed_ptr + q_offs, q_int.to(tl.uint8, bitcast=True), mask=blk_mask[:, None]) + + +@triton.jit +def _convert_from_mx9_kernel( + packed_ptr, + y_ptr, + n_blocks, + stride_blk, + stride_col, + BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, + OUT_DTYPE: tl.constexpr, + N_PACKED_BYTES: tl.constexpr, +): + + tl.static_assert(BLOCK_SIZE % PRIME_GROUP == 0, "BLOCK_SIZE must be divisible by PRIME_GROUP") + tl.static_assert((BLOCK_SIZE // PRIME_GROUP) % 8 == 0, "MX9 prime bitmap requires pair count divisible by 8") + + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + PRIME_OFFSET: tl.constexpr = 1 + Q_OFFSET: tl.constexpr = PRIME_OFFSET + N_PRIME_BYTES + tl.static_assert(N_PACKED_BYTES == Q_OFFSET + BLOCK_SIZE, "N_PACKED_BYTES does not match the MX9 segment layout") + + pid = tl.program_id(0) + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) + col = tl.arange(0, BLOCK_SIZE) + blk_mask = blk < n_blocks + mask = blk_mask[:, None] + + # Undo the E8M0 bias to recover the true exponent. + max_exp_u = tl.load(packed_ptr + blk * stride_blk, mask=blk_mask, other=0) + max_exp = max_exp_u.to(tl.int32) - 127 + + pcol = tl.arange(0, N_PRIME_BYTES) + offs_p = blk[:, None] * stride_blk + (PRIME_OFFSET + pcol[None, :]) * stride_col + prime = tl.load(packed_ptr + offs_p, mask=mask, other=0) + + q_offs = blk[:, None] * stride_blk + (Q_OFFSET + col[None, :]) * stride_col + q_u8 = tl.load(packed_ptr + q_offs, mask=mask, other=0) + q = q_u8.to(tl.int8, bitcast=True) + + y = _unpack_mx9( + q, + max_exp[:, None], + prime, + OUT_DTYPE, + BLOCKS_PER_PROG, + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + ) + out_offs = blk[:, None] * BLOCK_SIZE + col[None, :] + tl.store(y_ptr + out_offs, y, mask=mask) + + +# Host wrappers + + +def _mx9_packed_bytes(block_size: int) -> int: + """Return bytes per MX9 packed block.""" + n_prime_bytes = (block_size // PRIME_GROUP) // 8 + return 1 + n_prime_bytes + block_size + + +def convert_to_mx9( + data_hp: torch.Tensor, + block_size: int = BLOCK_SIZE, + axis: int = -1, + blocks_per_program: int = BLOCKS_PER_PROG_DEFAULT, +) -> torch.Tensor: + """High-precision tensor -> packed MX9 ``[n_blocks, n_packed_bytes]`` uint8. + + Blocks are formed along ``axis``, which is transposed to the last dim + internally. The original shape / axis are not stored, so the caller must pass + them back to convert_from_mx9. + """ + return _convert_to_mx_host( + data_hp, + kernel=_convert_to_mx9_kernel, + n_packed_bytes=_mx9_packed_bytes(block_size), + fmt="mx9_quantization", + block_size=block_size, + prime_group=PRIME_GROUP, + quant_bit=QUANT_BIT, + axis=axis, + blocks_per_program=blocks_per_program, + ) + + +def convert_from_mx9( + packed: torch.Tensor, + out_dtype: torch.dtype, + out_shape, + block_size: int = BLOCK_SIZE, + axis: int = -1, + blocks_per_program: int = BLOCKS_PER_PROG_DEFAULT, +) -> torch.Tensor: + """Packed MX9 tensor -> reconstructed high-precision tensor. + + ``out_shape`` and ``axis`` must exactly match the original convert_to_mx9 + call. They are not stored in ``packed``; a mismatch may silently reorder + values when it happens to produce the same block count. + """ + return _convert_from_mx_host( + packed, + kernel=_convert_from_mx9_kernel, + n_packed_bytes=_mx9_packed_bytes(block_size), + fmt="mx9_quantization", + out_dtype=out_dtype, + out_shape=out_shape, + block_size=block_size, + prime_group=PRIME_GROUP, + quant_bit=QUANT_BIT, + axis=axis, + blocks_per_program=blocks_per_program, + ) diff --git a/alto/models/patcher.py b/alto/models/patcher.py index 8c751354..e596a88d 100644 --- a/alto/models/patcher.py +++ b/alto/models/patcher.py @@ -58,29 +58,37 @@ class FakeQuantizeFunction(torch.autograd.Function): @staticmethod def forward(ctx, x, scale, zero_point, args, g_idx, global_scale): if getattr(args, "format", None) == "mx9": - from alto.modifiers.quantization.mx import ( - BLOCK_SIZE, - MX9_QUANT_BIT, - mx9_fake_quantize, + # Dispatches to the REAL packed Triton kernel (quantize -> + # bit-packed bytes -> dequantize), not the torch fake-quant + # emulation -- this is the intentional, validated design + # (commit 95b1114: LLaMA-3.2-1B wikitext loss 2.1141 == the + # fake-quant path's 2.1140). Requires a CUDA tensor (Triton + # kernel launch); mx6 below mirrors this same method. + from alto.kernels.mx.mx9_quantization import ( + convert_to_mx9, + convert_from_mx9, ) - return mx9_fake_quantize( - x, - block_size=(args.group_size or BLOCK_SIZE), - quant_bit=(args.num_bits or MX9_QUANT_BIT), + packed = convert_to_mx9(x) + return convert_from_mx9( + packed, x.dtype, x.shape, ) if getattr(args, "format", None) == "mx6": - from alto.modifiers.quantization.mx import ( + # Same method as "mx9" above: dispatches to the REAL packed + # Triton kernel (quantize -> bit-packed bytes -> dequantize), + # not the mx6_fake_quantize emulation. Requires a CUDA tensor + # (Triton kernel launch). quant_bit is not a parameter here: + # the packed kernel's 5-bit width is a fixed module constant + # (QUANT_BIT), not configurable. + from alto.kernels.mx.mx6_quantization import ( BLOCK_SIZE, - MX6_QUANT_BIT, - mx6_fake_quantize, + convert_to_mx6, + convert_from_mx6, ) - return mx6_fake_quantize( - x, - block_size=(args.group_size or BLOCK_SIZE), - quant_bit=(args.num_bits or MX6_QUANT_BIT), - ) + packed = convert_to_mx6(x, block_size=(args.group_size or BLOCK_SIZE)) + return convert_from_mx6(packed, x.dtype, x.shape, + block_size=(args.group_size or BLOCK_SIZE)) return original_fake_quantize(x, scale, zero_point, args, g_idx, global_scale) @staticmethod diff --git a/examples/llama3.2_1b_mx6.sh b/examples/llama3.2_1b_mx6.sh new file mode 100755 index 00000000..9ee69f39 --- /dev/null +++ b/examples/llama3.2_1b_mx6.sh @@ -0,0 +1,65 @@ +#!/usr/bin/bash +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +# MX6 W5A5 dynamic validation on Llama-3.2-1B +# Weight and input activations are quantized dynamically through +# alto/models/llama3/configs/mx6_wa_recipe.yaml. format=="mx6" dispatches to the +# packed Triton kernel (convert_to_mx6 / convert_from_mx6) in alto/models/patcher.py, +# not the mx6_fake_quantize emulation. +# +# Usage (MODEL_PATH is required, point it at your local Llama-3.2-1B dir): +# MODEL_PATH=/path/to/Llama-3.2-1B bash examples/llama3.2_1b_mx6.sh +# MODEL_PATH=/path/to/Llama-3.2-1B VALIDATOR_STEPS=100 bash examples/llama3.2_1b_mx6.sh +# MODEL_PATH=/path/to/Llama-3.2-1B VALIDATOR_STEPS=-1 bash examples/llama3.2_1b_mx6.sh # full validation set once +# MODEL_PATH=/path/to/Llama-3.2-1B CONFIG=llama3_1b bash examples/llama3.2_1b_mx6.sh # BF16 baseline +rm -rf outputs/ +set -ex + +NGPU=${NGPU:-"1"} +export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} +export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-${CUDA_VISIBLE_DEVICES}} +export LOG_RANK=${LOG_RANK:-0} +TRAIN_FILE=${TRAIN_FILE:-"alto.train"} +MODULE=${MODULE:-"llama3"} +CONFIG=${CONFIG:-"llama3_1b_mx6_wa"} +COMM_MODE=${COMM_MODE:-""} + +MODEL_PATH=${MODEL_PATH:-""} +if [ -z "${MODEL_PATH}" ]; then + echo "ERROR: MODEL_PATH must be set to your local Llama-3.2-1B directory, e.g." >&2 + echo " MODEL_PATH=/path/to/Llama-3.2-1B bash $0" >&2 + exit 1 +fi +VALIDATOR_STEPS=${VALIDATOR_STEPS:-"10"} +CHECKPOINT_FOLDER=${CHECKPOINT_FOLDER:-"./outputs/ckpt_${CONFIG}_$(date +%Y%m%d_%H%M%S)"} + +TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE:-"http://localhost:29510"} + +if [ -n "$COMM_MODE" ]; then + echo "Running with comm_mode=${COMM_MODE}" + NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} \ + --module ${MODULE} \ + --config ${CONFIG} \ + --hf_assets_path "${MODEL_PATH}" \ + --checkpoint.initial_load_path "${MODEL_PATH}" \ + --checkpoint.folder "${CHECKPOINT_FOLDER}" \ + --validator.steps "${VALIDATOR_STEPS}" \ + "$@" \ + --comm.mode=${COMM_MODE} \ + --training.steps 1 +else + PYTORCH_ALLOC_CONF="expandable_segments:True" \ + TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ + torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ + --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ + -m ${TRAIN_FILE} \ + --module ${MODULE} \ + --config ${CONFIG} \ + --hf_assets_path "${MODEL_PATH}" \ + --checkpoint.initial_load_path "${MODEL_PATH}" \ + --checkpoint.folder "${CHECKPOINT_FOLDER}" \ + --validator.steps "${VALIDATOR_STEPS}" \ + "$@" +fi diff --git a/examples/llama3.2_1b_mx9.sh b/examples/llama3.2_1b_mx9.sh index 204794a6..ef127d3a 100755 --- a/examples/llama3.2_1b_mx9.sh +++ b/examples/llama3.2_1b_mx9.sh @@ -3,9 +3,10 @@ # # SPDX-License-Identifier: MIT -# MX9 W8A8 dynamic fake-quant validation on Llama-3.2-1B +# MX9 W8A8 dynamic validation on Llama-3.2-1B # Weight and input activations are quantized dynamically through -# alto/models/llama3/configs/mx9_wa_recipe.yaml +# alto/models/llama3/configs/mx9_wa_recipe.yaml. format=="mx9" dispatches to the +# packed Triton kernel (convert_to_mx9 / convert_from_mx9) in alto/models/patcher.py. # # Usage (MODEL_PATH is required, point it at your local Llama-3.2-1B dir): # MODEL_PATH=/path/to/Llama-3.2-1B bash examples/llama3.2_1b_mx9.sh diff --git a/tests/unittest/mx9_mx6/test_mx6_dispatch.py b/tests/unittest/mx9_mx6/test_mx6_dispatch.py index 870f776a..a1ff7b46 100644 --- a/tests/unittest/mx9_mx6/test_mx6_dispatch.py +++ b/tests/unittest/mx9_mx6/test_mx6_dispatch.py @@ -3,11 +3,16 @@ # SPDX-License-Identifier: MIT """Tests for the MX6 dispatch wiring in ``ModelPatcher.patch_fake_quantize``. -The MX6 kernel itself is covered by ``test_mx6_quantize.py``. This file tests the -*wiring* one layer up: after ``patch_fake_quantize()`` replaces +The MX6 kernel itself is covered by ``test_mx6_quantize.py`` (fake-quant +emulation) and ``test_mx6_quantization.py`` (real packed kernel). This file +tests the *wiring* one layer up: after ``patch_fake_quantize()`` replaces ``compressed_tensors...forward.fake_quantize``, a ``QuantizationArgs`` carrying -``format == "mx6"`` must route to ``mx6_fake_quantize``, while plain int8 args -(no ``format``) must fall through to the original implementation untouched. +``format == "mx6"`` must route to the REAL packed Triton kernel +(``convert_to_mx6`` / ``convert_from_mx6``), NOT the ``mx6_fake_quantize`` +emulation -- this mirrors "mx9"'s method exactly (see +``alto/models/patcher.py`` and commit ``95b1114`` for mx9's validated +precedent). Plain int8 args (no ``format``) must fall through to the original +implementation untouched. Also checks that ``format_registry.inject_format_field()`` makes the ``format`` field survive pydantic validation (otherwise the recipe value is silently @@ -68,21 +73,31 @@ def test_format_field_defaults_none_for_plain_int8(): # --------------------------------------------------------------------------- # # dispatch routing # --------------------------------------------------------------------------- # +@pytest.mark.skipif(not torch.cuda.is_available(), reason="mx6 dispatch launches the real Triton kernel") def test_mx6_args_dispatch_to_mx6_kernel(): - """format == "mx6" must route fake_quantize to mx6_fake_quantize bit-exact.""" + """format == "mx6" must route fake_quantize to the REAL packed kernel + (convert_to_mx6 -> convert_from_mx6 round trip) bit-exact. Requires CUDA + since the packed kernel launches a Triton kernel; NOT compared against + mx6_fake_quantize, which intentionally diverges at the rare demoted + +/-16 boundary (see test_mx6_quantization.py:: + test_divergence_vs_mxpy_31clamp_is_tiny).""" + from alto.kernels.mx.mx6_quantization import convert_to_mx6, convert_from_mx6 + torch.manual_seed(6) - x = torch.randn(3, 40, dtype=torch.float32) + x = torch.randn(3, 40, dtype=torch.float32).cuda() args = _mx6_args() patched_out = forward_module.fake_quantize( x=x, - scale=torch.ones(1), + scale=torch.ones(1, device="cuda"), zero_point=None, args=args, g_idx=None, global_scale=None, ) - expected = mx6_fake_quantize(x, block_size=BLOCK_SIZE) + + packed = convert_to_mx6(x) + expected = convert_from_mx6(packed, x.dtype, x.shape) assert torch.equal(patched_out, expected) diff --git a/tests/unittest/mx9_mx6/test_mx6_quantization.py b/tests/unittest/mx9_mx6/test_mx6_quantization.py new file mode 100644 index 00000000..fb427b97 --- /dev/null +++ b/tests/unittest/mx9_mx6/test_mx6_quantization.py @@ -0,0 +1,993 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Unit tests for packed MX6 quantization, built up stage by stage. + +Testing strategy +---------------- +Stages 1-3 are ``@triton.jit`` *device* helpers (callable only from within a +kernel), so each is exercised through a tiny single-program "probe" kernel that +loads inputs, calls the helper, and stores the result back. Stage 4 exposes real +grid kernels, launched directly via minimal inline launchers. In every case the +GPU result is compared **bit-exact** (``torch.equal``) against an independent +pure-PyTorch recomputation of the same math, over a matrix of shapes / dtypes / +block counts, plus constructed edge cases (zeros, NaN/Inf, prime patterns, +sign bitmaps, the +/-15 clamp boundary, block independence). + +Coverage by stage +----------------- + - Stage 1: ``_floor_exp`` / ``_round_half_even`` / + ``_pack_bits8`` / ``_unpack_bits8`` + - Stage 2: ``_calculate_mx6_exp`` (shared_exp / max_exp / prime pair bits) + - Stage 3: ``_pack_mx6`` (Python-export-compatible encoding) and the pack->unpack + round-trip vs the clamp-15 reference + - Stage 4: ``_convert_to_mx6_kernel`` / ``_convert_from_mx6_kernel`` (storage + format, single packed tensor layout, E8M0 bias, blocks-per-program invariance, + end-to-end round-trip) + - Stage 5: ``convert_to_mx6`` / ``convert_from_mx6`` host wrappers (axis + transpose, padding-aware shape restore, non-contiguous input, invalid-input + rejection), mirroring the MX9 host test suite. +""" + +import pytest +import torch +import triton +import triton.language as tl + +from alto.kernels.mx._mx_common import ( + _floor_exp, + _round_half_even, + _pack_bits8, + _unpack_bits8, + _calculate_mx_exp as _calculate_mx6_exp, +) +from alto.kernels.mx.mx6_quantization import ( + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + BLOCKS_PER_PROG_DEFAULT, + _pack_mx6, + _unpack_mx6, + _convert_to_mx6_kernel, + _convert_from_mx6_kernel, + convert_to_mx6, + convert_from_mx6, +) +from alto.modifiers.quantization import mx as _mxref +from alto.modifiers.quantization.mx import mx6_fake_quantize + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +# --------------------------------------------------------------------------- # +# Pure-PyTorch references (independent reimplementation) +# --------------------------------------------------------------------------- # +def _rand(shape, dtype, seed=6): + torch.manual_seed(seed) + return (torch.randn(*shape, dtype=dtype) * 5.0) + + +def _bit_exponent(t: torch.Tensor) -> torch.Tensor: + """Extract exponent field from native float bits (approx floor(log2|t|)), + branching by dtype. Must run on native dtype (no pre-cast to float).""" + if t.dtype == torch.float32: + return (((t.view(torch.int32) >> 23) & 0xFF) - 127).long() + if t.dtype == torch.bfloat16: + return (((t.view(torch.int16) >> 7) & 0xFF) - 127).long() + raise ValueError(f"unsupported dtype {t.dtype}") + + +def _torch_calc_ref(xb: torch.Tensor, block_size: int, prime_group: int): + """Reference for _calculate_mx6_exp. xb: [n_blocks, block_size] native dtype. + Returns (shared_exp[int64], max_exp[int64], pair[int64]).""" + n_blocks = xb.shape[0] + clean = torch.nan_to_num(xb, nan=0.0, posinf=0.0, neginf=0.0) + amax = clean.abs().amax(dim=1, keepdim=True) # [nb, 1] native dtype + max_exp = _bit_exponent(amax) # [nb, 1] + t_exp = _bit_exponent(clean) # [nb, bs] + demote = (max_exp - t_exp) >= 1 + + n_pairs = block_size // prime_group + d = demote.reshape(n_blocks, n_pairs, prime_group) + pair = (d.sum(dim=-1) == prime_group) # [nb, n_pairs] bool + pair_b = pair[:, :, None].expand(n_blocks, n_pairs, prime_group).reshape(n_blocks, block_size) + shared_exp = max_exp - pair_b.long() + return shared_exp, max_exp.reshape(n_blocks), pair.long() + + +# --------------------------------------------------------------------------- # +# Probe kernels (test-only harnesses that invoke the device helpers) +# --------------------------------------------------------------------------- # +@triton.jit +def _probe_floor_exp_kernel(x_ptr, o_ptr, N: tl.constexpr): + i = tl.arange(0, N) + tl.store(o_ptr + i, _floor_exp(tl.load(x_ptr + i))) + + +@triton.jit +def _probe_round_kernel(y_ptr, o_ptr, N: tl.constexpr): + i = tl.arange(0, N) + tl.store(o_ptr + i, _round_half_even(tl.load(y_ptr + i))) + + +@triton.jit +def _probe_bits8_kernel( + bits_ptr, + packed_ptr, + unpacked_ptr, + BPP: tl.constexpr, + N_BYTES: tl.constexpr, +): + N_BITS: tl.constexpr = N_BYTES * 8 + rows = tl.arange(0, BPP) + bit_cols = tl.arange(0, N_BITS) + bit_offs = rows[:, None] * N_BITS + bit_cols[None, :] + bits = tl.load(bits_ptr + bit_offs) + + packed = _pack_bits8(bits, BPP, N_BYTES) + byte_cols = tl.arange(0, N_BYTES) + byte_offs = rows[:, None] * N_BYTES + byte_cols[None, :] + tl.store(packed_ptr + byte_offs, packed) + + unpacked = _unpack_bits8(packed, BPP, N_BYTES) + tl.store(unpacked_ptr + bit_offs, unpacked) + + +@triton.jit +def _probe_calc_kernel( + x_ptr, se_ptr, me_ptr, pr_ptr, + BPP: tl.constexpr, BLOCK_SIZE: tl.constexpr, PRIME_GROUP: tl.constexpr, +): + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + rows = tl.arange(0, BPP) + cols = tl.arange(0, BLOCK_SIZE) + offs = rows[:, None] * BLOCK_SIZE + cols[None, :] + x = tl.load(x_ptr + offs) + shared_exp, max_exp, pair = _calculate_mx6_exp(x, BPP, BLOCK_SIZE, PRIME_GROUP) + tl.store(se_ptr + offs, shared_exp) + tl.store(me_ptr + rows, max_exp) + pcols = tl.arange(0, N_PAIRS) + tl.store(pr_ptr + rows[:, None] * N_PAIRS + pcols[None, :], pair) + + +# --------------------------------------------------------------------------- # +# Stage 1: _floor_exp +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_floor_exp_matches_bit_extraction(dtype): + torch.manual_seed(6) + # Finite non-zero values only (_floor_exp does not sanitize NaN/Inf/0). + x = (torch.randn(64, dtype=dtype, device="cuda") * 37.0) + x = torch.where(x.abs() < 1e-3, torch.full_like(x, 1e-3), x) + out = torch.empty(64, dtype=torch.int32, device="cuda") + _probe_floor_exp_kernel[(1,)](x, out, N=64) + ref = _bit_exponent(x).to(torch.int32) + assert torch.equal(out, ref) + + +def test_floor_exp_known_powers_of_two(): + x = torch.tensor([1.0, 2.0, 3.9, 4.0, 0.5, 0.25, 130.0, -8.0], + dtype=torch.float32, device="cuda") + out = torch.empty(8, dtype=torch.int32, device="cuda") + _probe_floor_exp_kernel[(1,)](x, out, N=8) + # floor(log2|x|): 1->0, 2->1, 3.9->1, 4->2, 0.5->-1, 0.25->-2, 130->7, -8->3 + ref = torch.tensor([0, 1, 1, 2, -1, -2, 7, 3], dtype=torch.int32, device="cuda") + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# Stage 1: _round_half_even +# --------------------------------------------------------------------------- # +def test_round_half_even_ties_and_regular(): + y = torch.tensor( + [-2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5, + 0.4, 0.6, -0.4, -0.6, 2.4, 2.6, -2.4, -2.6], + dtype=torch.float32, device="cuda", + ) + out = torch.empty_like(y) + _probe_round_kernel[(1,)](y, out, N=y.numel()) + ref = torch.round(y) # torch.round is round-half-to-even + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# Stage 1: _pack_bits8 / _unpack_bits8 +# --------------------------------------------------------------------------- # +def test_bits8_pack_order_and_roundtrip(): + bits = torch.tensor( + [ + [0] * 8 + [1] * 8, + [1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1], + ], + dtype=torch.int32, + device="cuda", + ) + packed = torch.empty((2, 2), dtype=torch.uint8, device="cuda") + unpacked = torch.empty_like(bits) + + _probe_bits8_kernel[(1,)](bits, packed, unpacked, BPP=2, N_BYTES=2) + + expected_packed = torch.tensor( + [[0x00, 0xFF], [0x55, 0xAA]], + dtype=torch.uint8, + device="cuda", + ) + assert torch.equal(packed, expected_packed) + assert torch.equal(unpacked, bits) + + +# --------------------------------------------------------------------------- # +# Stage 2: _calculate_mx6_exp +# --------------------------------------------------------------------------- # +def _run_calc(xb, block_size=BLOCK_SIZE, prime_group=PRIME_GROUP): + n_blocks = xb.shape[0] + n_pairs = block_size // prime_group + se = torch.empty((n_blocks, block_size), dtype=torch.int32, device="cuda") + me = torch.empty((n_blocks,), dtype=torch.int32, device="cuda") + pr = torch.empty((n_blocks, n_pairs), dtype=torch.int32, device="cuda") + _probe_calc_kernel[(1,)](xb, se, me, pr, BPP=n_blocks, + BLOCK_SIZE=block_size, PRIME_GROUP=prime_group) + return se, me, pr + + +@pytest.mark.parametrize("n_blocks", [8, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_calculate_exp_matches_torch(n_blocks, dtype): + torch.manual_seed(6) + xb = torch.randn(n_blocks, BLOCK_SIZE, dtype=dtype, device="cuda") * 5.0 + se, me, pr = _run_calc(xb) + rse, rme, rpr = _torch_calc_ref(xb, BLOCK_SIZE, PRIME_GROUP) + assert torch.equal(se.long(), rse), "shared_exp mismatch" + assert torch.equal(me.long(), rme), "max_exp mismatch" + assert torch.equal(pr.long(), rpr), "prime pair mismatch" + + +def test_calculate_exp_constructed_prime_pattern(): + """Deterministic prime bitmap: block = [8]*8 + [1]*8. + amax=8 (max_exp=3); first 8 elems=8 (t_exp=3, not demoted); + last 8 elems=1 (t_exp=0, demoted). Pairs 0-3 -> 0, pairs 4-7 -> 1.""" + xb = torch.empty(1, BLOCK_SIZE, dtype=torch.float32, device="cuda") + xb[0, :8] = 8.0 + xb[0, 8:] = 1.0 + se, me, pr = _run_calc(xb) + assert me[0].item() == 3 + assert pr[0].tolist() == [0, 0, 0, 0, 1, 1, 1, 1] + # demoted pairs use max_exp - 1 = 2; non-demoted stay at 3. + assert se[0].tolist() == [3] * 8 + [2] * 8 + + +def test_calculate_exp_all_zero_block(): + """All-zero block: amax=0 -> _floor_exp(0)=-127; nothing demoted.""" + xb = torch.zeros(2, BLOCK_SIZE, dtype=torch.float32, device="cuda") + se, me, pr = _run_calc(xb) + assert torch.all(me == -127) + assert torch.all(pr == 0) + assert torch.all(se == -127) + + +def test_calculate_exp_nan_inf_do_not_pollute(): + """NaN/Inf are sanitized out of the exponent statistics: a block of finite + values plus one Inf must yield the same max_exp as without the Inf.""" + xb = torch.full((1, BLOCK_SIZE), 3.0, dtype=torch.float32, device="cuda") + xb[0, 0] = float("inf") + xb[0, 1] = float("nan") + se, me, pr = _run_calc(xb) + # Finite amax = 3 -> max_exp = 1 (Inf/NaN sanitized to 0, ignored). + assert me[0].item() == 1 + rse, rme, rpr = _torch_calc_ref(xb, BLOCK_SIZE, PRIME_GROUP) + assert torch.equal(me.long(), rme) + assert torch.equal(pr.long(), rpr) + + +# --------------------------------------------------------------------------- # +# Stage 3 references: Python-export-compatible pack encoding + clamp-15 QDQ +# --------------------------------------------------------------------------- # +def _torch_q_and_shared_exp(xb, block_size, prime_group, quant_bit): + """Shared helper: per-element clamp-15 quantized int q and shared_exp.""" + n_blocks = xb.shape[0] + clean = torch.nan_to_num(xb, nan=0.0, posinf=0.0, neginf=0.0) + amax = clean.abs().amax(dim=1, keepdim=True) + max_exp = _bit_exponent(amax) + t_exp = _bit_exponent(clean) + demote = (max_exp - t_exp) >= 1 + n_pairs = block_size // prime_group + d = demote.reshape(n_blocks, n_pairs, prime_group) + pair = (d.sum(dim=-1) == prime_group) + pair_b = pair[:, :, None].expand(n_blocks, n_pairs, prime_group).reshape(n_blocks, block_size) + shared_exp = max_exp - pair_b.long() + q_hi = (1 << (quant_bit - 1)) - 1 + scale_exp = torch.clamp(shared_exp - quant_bit + 2, min=-126) + scale = torch.exp2(scale_exp.float()) + q = torch.round(xb.float() / scale).clamp(-q_hi, q_hi) + return q, scale, pair.long() + + +def _torch_pack_ref(xb, block_size, prime_group, quant_bit): + """Reference Python export byte encoding. + + Returns ``(sign_bytes, mantissa_bytes, prime_bytes, max_exp_biased)``. The + full block layout is ``[max_exp, prime, sign, mantissa]``. + """ + n_blocks = xb.shape[0] + q, _, pair = _torch_q_and_shared_exp(xb, block_size, prime_group, quant_bit) + qi = q.to(torch.int64) + mantissa = qi & 0xF + sign = (qi < 0).to(torch.int64) + + n_sign_bytes = block_size // 8 + w8 = (2 ** torch.arange(8, device=xb.device)).to(torch.int64) + sign_bytes = (sign.reshape(n_blocks, n_sign_bytes, 8) * w8).sum(-1) + + n_mantissa_bytes = block_size // 2 + nib_w = torch.tensor([1, 16], device=xb.device, dtype=torch.int64) + mantissa_bytes = (mantissa.reshape(n_blocks, n_mantissa_bytes, 2) * nib_w).sum(-1) + + n_pairs = block_size // prime_group + n_prime_bytes = n_pairs // 8 + prime_bytes = (pair.reshape(n_blocks, n_prime_bytes, 8) * w8).sum(-1) + + clean = torch.nan_to_num(xb, nan=0.0, posinf=0.0, neginf=0.0) + max_exp_biased = _bit_exponent(clean.abs().amax(1, keepdim=True)).reshape(n_blocks) + 127 + return sign_bytes, mantissa_bytes, prime_bytes, max_exp_biased + + +def _mx6_clamp15_blockref(xb, block_size, prime_group, quant_bit): + """Clamp-15 QDQ reference at block granularity: xb[nb, bs] -> y[nb, bs] fp32.""" + q, scale, _ = _torch_q_and_shared_exp(xb, block_size, prime_group, quant_bit) + return q * scale + + +_TL_DTYPE = {torch.float32: tl.float32, torch.bfloat16: tl.bfloat16} + + +@triton.jit +def _probe_pack_kernel( + x_ptr, sign_ptr, mantissa_ptr, prime_ptr, + BPP: tl.constexpr, BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, QUANT_BIT: tl.constexpr, +): + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 + N_MANTISSA_BYTES: tl.constexpr = BLOCK_SIZE // 2 + N_PRIME_BYTES: tl.constexpr = N_PAIRS // 8 + rows = tl.arange(0, BPP) + cols = tl.arange(0, BLOCK_SIZE) + x = tl.load(x_ptr + rows[:, None] * BLOCK_SIZE + cols[None, :]) + shared_exp, max_exp, pair = _calculate_mx6_exp(x, BPP, BLOCK_SIZE, PRIME_GROUP) + sign_b, mantissa_b, prime_b = _pack_mx6(x, shared_exp, pair, BPP, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT) + scol = tl.arange(0, N_SIGN_BYTES) + tl.store(sign_ptr + rows[:, None] * N_SIGN_BYTES + scol[None, :], sign_b) + mcol = tl.arange(0, N_MANTISSA_BYTES) + tl.store(mantissa_ptr + rows[:, None] * N_MANTISSA_BYTES + mcol[None, :], mantissa_b) + pcol = tl.arange(0, N_PRIME_BYTES) + tl.store(prime_ptr + rows[:, None] * N_PRIME_BYTES + pcol[None, :], prime_b) + + +@triton.jit +def _probe_roundtrip_kernel( + x_ptr, y_ptr, + BPP: tl.constexpr, BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, QUANT_BIT: tl.constexpr, OUT_DTYPE: tl.constexpr, +): + rows = tl.arange(0, BPP) + cols = tl.arange(0, BLOCK_SIZE) + offs = rows[:, None] * BLOCK_SIZE + cols[None, :] + x = tl.load(x_ptr + offs) + shared_exp, max_exp, pair = _calculate_mx6_exp(x, BPP, BLOCK_SIZE, PRIME_GROUP) + sign_b, mantissa_b, prime_b = _pack_mx6(x, shared_exp, pair, BPP, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT) + y = _unpack_mx6(sign_b, mantissa_b, max_exp[:, None], prime_b, OUT_DTYPE, + BPP, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT) + tl.store(y_ptr + offs, y) + + +def _run_pack(xb): + n_blocks = xb.shape[0] + sign = torch.empty((n_blocks, BLOCK_SIZE // 8), dtype=torch.uint8, device="cuda") + mantissa = torch.empty((n_blocks, BLOCK_SIZE // 2), dtype=torch.uint8, device="cuda") + prime = torch.empty((n_blocks, (BLOCK_SIZE // PRIME_GROUP) // 8), dtype=torch.uint8, device="cuda") + _probe_pack_kernel[(1,)](xb, sign, mantissa, prime, BPP=n_blocks, BLOCK_SIZE=BLOCK_SIZE, + PRIME_GROUP=PRIME_GROUP, QUANT_BIT=QUANT_BIT) + return sign, mantissa, prime + + +def _run_roundtrip(xb): + n_blocks = xb.shape[0] + y = torch.empty((n_blocks, BLOCK_SIZE), dtype=xb.dtype, device="cuda") + _probe_roundtrip_kernel[(1,)](xb, y, BPP=n_blocks, BLOCK_SIZE=BLOCK_SIZE, + PRIME_GROUP=PRIME_GROUP, QUANT_BIT=QUANT_BIT, + OUT_DTYPE=_TL_DTYPE[xb.dtype]) + return y + + +# --------------------------------------------------------------------------- # +# Stage 3: _pack_mx6 encoding (sign bitmap / low-4-bit mantissa / prime) +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("n_blocks", [8, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_pack_encoding_matches_torch(n_blocks, dtype): + torch.manual_seed(6) + xb = torch.randn(n_blocks, BLOCK_SIZE, dtype=dtype, device="cuda") * 5.0 + sign, mantissa, prime = _run_pack(xb) + rsign, rmantissa, rprime, _ = _torch_pack_ref(xb, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT) + assert torch.equal(sign.long(), rsign), "sign bitmap mismatch" + assert torch.equal(mantissa.long(), rmantissa), "mantissa nibble mismatch" + assert torch.equal(prime.long(), rprime), "prime bitmap mismatch" + + +def test_pack_mantissa_fits_4_bits(): + """Every mantissa nibble must be <= 15.""" + torch.manual_seed(1) + xb = torch.randn(16, BLOCK_SIZE, dtype=torch.float32, device="cuda") * 50.0 + _, mantissa, _ = _run_pack(xb) + lo = mantissa & 0x0F + hi = (mantissa >> 4) & 0x0F + assert lo.max().item() <= 15 and hi.max().item() <= 15 + + +def test_pack_mantissa_uses_twos_complement_low4(): + """q=-3 is encoded as sign=1 plus low4=0xD, matching the Python packer.""" + xb = torch.zeros(1, BLOCK_SIZE, dtype=torch.float32, device="cuda") + xb[0, 0] = 8.0 # max_exp=3, scale=1 for pair0 + xb[0, 1] = -3.0 + sign, mantissa, _ = _run_pack(xb) + assert sign[0, 0].item() & 0b10 + assert mantissa[0, 0].item() == 0xD8 + + +def test_pack_sign_bitmap_constructed(): + """Alternating signs -> sign byte 0b10101010 = 0xAA (odd elements negative).""" + xb = torch.ones(1, BLOCK_SIZE, dtype=torch.float32, device="cuda") + xb[0, 1::2] = -1.0 + sign, _, _ = _run_pack(xb) + # elements 0..7 -> byte0: bits 1,3,5,7 set = 0xAA; elements 8..15 -> byte1 = 0xAA + assert sign[0, 0].item() == 0xAA + assert sign[0, 1].item() == 0xAA + + +# --------------------------------------------------------------------------- # +# Stage 3: pack -> unpack round-trip (bit-exact vs clamp-15 reference) +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("n_blocks", [8, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_roundtrip_matches_clamp15_ref(n_blocks, dtype): + torch.manual_seed(6) + xb = torch.randn(n_blocks, BLOCK_SIZE, dtype=dtype, device="cuda") * 5.0 + y = _run_roundtrip(xb) + ref = _mx6_clamp15_blockref(xb, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT).to(dtype) + assert y.dtype == dtype + assert torch.equal(y, ref), f"max diff={(y.float() - ref.float()).abs().max().item()}" + + +def test_roundtrip_zeros_stay_zero(): + xb = torch.zeros(4, BLOCK_SIZE, dtype=torch.float32, device="cuda") + y = _run_roundtrip(xb) + assert torch.equal(y, xb) + + +def test_roundtrip_sign_preserved(): + torch.manual_seed(3) + xb = torch.randn(8, BLOCK_SIZE, dtype=torch.float32, device="cuda") * 5.0 + y = _run_roundtrip(xb) + ref = _mx6_clamp15_blockref(xb, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT) + # nonzero reconstructed values keep the reference sign + nz = ref != 0 + assert torch.equal(torch.sign(y[nz]), torch.sign(ref[nz])) + + +def test_roundtrip_idempotent(): + """pack->unpack output already lies on the MX6 grid; a second pass is identical.""" + torch.manual_seed(0) + xb = torch.randn(8, BLOCK_SIZE, dtype=torch.float32, device="cuda") * 5.0 + once = _run_roundtrip(xb) + twice = _run_roundtrip(once) + assert torch.equal(once, twice) + + +def test_roundtrip_demote_boundary_clamp15(): + """Demoted element rounding to +/-16 must be clamped to +/-15. + + max_exp=3 (amax in [8,16)); demote scale = 2^(3-4)=0.5. An element at 7.6 in a + demotable pair -> round(7.6/0.5)=round(15.2)=15 stays; push to reach 16: + value 7.9 -> round(15.8)=16 -> clamp 15 -> reconstruct 15*0.5=7.5. + """ + xb = torch.zeros(1, BLOCK_SIZE, dtype=torch.float32, device="cuda") + xb[0, 0] = 9.0 # amax -> max_exp=3, non-demote anchor + xb[0, 1] = 8.0 # pair0 neighbor (non-demote) + xb[0, 2] = 7.9 # demote pair1 + xb[0, 3] = 0.5 # demote pair1 neighbor -> both demotable + y = _run_roundtrip(xb) + ref = _mx6_clamp15_blockref(xb, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT) + assert torch.equal(y, ref) + assert y[0, 2].item() == pytest.approx(7.5) + + +# --------------------------------------------------------------------------- # +# Stage 4: grid kernels (_convert_to/from_mx6_kernel) +# Minimal inline launchers drive the kernels directly on a contiguous 2D +# tensor whose last dim is a multiple of block_size (no transpose / padding -- +# those belong to the Stage 5 host wrappers). This isolates the kernel wiring: +# stride-addressed load/store, single packed tensor layout, E8M0 bias, masking. +# --------------------------------------------------------------------------- # +_N_PRIME_BYTES = (BLOCK_SIZE // PRIME_GROUP) // 8 # 1 +_N_SIGN_BYTES = BLOCK_SIZE // 8 # 2 +_N_MANTISSA_BYTES = BLOCK_SIZE // 2 # 8 +_N_PACKED_BYTES = 1 + _N_PRIME_BYTES + _N_SIGN_BYTES + _N_MANTISSA_BYTES # 12 +_PRIME_OFFSET = 1 +_SIGN_OFFSET = _PRIME_OFFSET + _N_PRIME_BYTES +_MANTISSA_OFFSET = _SIGN_OFFSET + _N_SIGN_BYTES + + +def _launch_to(x2d, bpp=BLOCKS_PER_PROG_DEFAULT): + rows, cols = x2d.shape + assert cols % BLOCK_SIZE == 0 + blocks_per_row = cols // BLOCK_SIZE + n_blocks = rows * blocks_per_row + packed = torch.empty((n_blocks, _N_PACKED_BYTES), dtype=torch.uint8, device="cuda") + sr, sc = x2d.stride() + grid = (triton.cdiv(n_blocks, bpp),) + _convert_to_mx6_kernel[grid]( + x2d, packed, n_blocks, cols, blocks_per_row, sr, sc, + BLOCK_SIZE=BLOCK_SIZE, BLOCKS_PER_PROG=bpp, + PRIME_GROUP=PRIME_GROUP, QUANT_BIT=QUANT_BIT, + N_PACKED_BYTES=_N_PACKED_BYTES, + ) + return packed, n_blocks + + +def _launch_to_ragged(x2d, bpp=BLOCKS_PER_PROG_DEFAULT): + """Like ``_launch_to`` but does NOT require ``cols`` to be a multiple of + ``BLOCK_SIZE``. Drives the in-kernel ragged-tail path: ``blocks_per_row`` is + ``cdiv`` and ``last`` = the unpadded column count, so columns >= last are + masked to 0 by ``_convert_to_mx6_kernel`` (equivalent to the host-side + zero-padding the Stage 5 wrapper will do explicitly).""" + rows, cols = x2d.shape + blocks_per_row = triton.cdiv(cols, BLOCK_SIZE) + n_blocks = rows * blocks_per_row + packed = torch.empty((n_blocks, _N_PACKED_BYTES), dtype=torch.uint8, device="cuda") + sr, sc = x2d.stride() + grid = (triton.cdiv(n_blocks, bpp),) + _convert_to_mx6_kernel[grid]( + x2d, packed, n_blocks, cols, blocks_per_row, sr, sc, + BLOCK_SIZE=BLOCK_SIZE, BLOCKS_PER_PROG=bpp, + PRIME_GROUP=PRIME_GROUP, QUANT_BIT=QUANT_BIT, + N_PACKED_BYTES=_N_PACKED_BYTES, + ) + return packed, n_blocks + + +def _launch_from(packed, n_blocks, out_dtype, bpp=BLOCKS_PER_PROG_DEFAULT): + y = torch.empty((n_blocks, BLOCK_SIZE), dtype=out_dtype, device="cuda") + spb, spc = packed.stride() + grid = (triton.cdiv(n_blocks, bpp),) + _convert_from_mx6_kernel[grid]( + packed, y, n_blocks, spb, spc, + BLOCK_SIZE=BLOCK_SIZE, BLOCKS_PER_PROG=bpp, + PRIME_GROUP=PRIME_GROUP, QUANT_BIT=QUANT_BIT, + OUT_DTYPE=_TL_DTYPE[out_dtype], + N_PACKED_BYTES=_N_PACKED_BYTES, + ) + return y + + +def test_kernel_output_dtypes_and_shapes(): + x = _rand((4, 64), torch.float32).cuda() + packed, n_blocks = _launch_to(x) + assert n_blocks == (4 * 64) // BLOCK_SIZE + assert packed.dtype == torch.uint8 and packed.shape == (n_blocks, _N_PACKED_BYTES) + + +def test_kernel_max_exp_biased_range(): + x = _rand((16, 64), torch.float32).cuda() + packed, _ = _launch_to(x) + assert packed[:, 0].min().item() >= 1 + assert packed[:, 0].max().item() <= 254 + + +def test_kernel_packed_bytes_match_pack_ref(): + """packed row must match Python export layout: [max_exp, prime, sign, mantissa].""" + for dtype in (torch.float32, torch.bfloat16): + x = _rand((8, 96), dtype).cuda() + packed, n_blocks = _launch_to(x) + xb = x.reshape(n_blocks, BLOCK_SIZE) + rsign, rmantissa, rprime, rme = _torch_pack_ref(xb, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT) + assert torch.equal(packed[:, 0].long(), rme), f"{dtype} max_exp" + assert torch.equal(packed[:, _PRIME_OFFSET:_SIGN_OFFSET].long(), rprime), f"{dtype} prime" + assert torch.equal(packed[:, _SIGN_OFFSET:_MANTISSA_OFFSET].long(), rsign), f"{dtype} sign" + assert torch.equal(packed[:, _MANTISSA_OFFSET:].long(), rmantissa), f"{dtype} mantissa" + + +@pytest.mark.parametrize("shape", [(4, 64), (8, 16), (16, 128), (32, 512)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_kernel_roundtrip_matches_clamp15_ref(shape, dtype): + x = _rand(shape, dtype).cuda() + packed, n_blocks = _launch_to(x) + y = _launch_from(packed, n_blocks, dtype) + ref = _mx6_clamp15_blockref(x.reshape(n_blocks, BLOCK_SIZE), + BLOCK_SIZE, PRIME_GROUP, QUANT_BIT).to(dtype) + assert y.shape == (n_blocks, BLOCK_SIZE) + assert torch.equal(y, ref), f"max diff={(y.float() - ref.float()).abs().max().item()}" + + +# --------------------------------------------------------------------------- # +# Stage 4: ragged tail (cols not a multiple of BLOCK_SIZE) +# The launcher forces cols % BLOCK_SIZE == 0 everywhere else, so the kernel's +# `in_range = col_g < last` tail masking + brow/row remapping never runs. These +# cases exercise it directly and check equivalence to explicit zero-padding. +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("shape", [(4, 70), (3, 17), (8, 100), (1, 1), (5, 33)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_kernel_ragged_tail_equals_zero_padding(shape, dtype): + rows, cols = shape + x = _rand(shape, dtype).cuda() + packed, n_blocks = _launch_to_ragged(x) + y = _launch_from(packed, n_blocks, dtype) + + # In-kernel tail masking (other=0) is equivalent to zero-padding the quant + # axis up to a multiple of BLOCK_SIZE, then block-quantizing. + pad = (BLOCK_SIZE - cols % BLOCK_SIZE) % BLOCK_SIZE + xpad = torch.nn.functional.pad(x, (0, pad)) + xb = xpad.reshape(n_blocks, BLOCK_SIZE) + ref = _mx6_clamp15_blockref(xb, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT).to(dtype) + + assert n_blocks == rows * triton.cdiv(cols, BLOCK_SIZE) + assert y.shape == (n_blocks, BLOCK_SIZE) + assert torch.equal(y, ref), f"max diff={(y.float() - ref.float()).abs().max().item()}" + + +def test_kernel_ragged_tail_block_is_partly_padded(): + """A row whose width leaves a 1-real-element tail block: that block's amax / + q must be computed from the single real value only (rest masked to 0).""" + # cols=17 -> 2 blocks/row; block1 holds real element 16 then 15 masked zeros. + x = torch.zeros(1, 17, dtype=torch.float32, device="cuda") + x[0, :16] = 1.0 + x[0, 16] = 6.0 # sole real element of the tail block + packed, n_blocks = _launch_to_ragged(x) + y = _launch_from(packed, n_blocks, torch.float32) + assert n_blocks == 2 + # Tail block: amax=6 -> max_exp=2 -> E8M0 = 129; element 0 reconstructs to 6, + # the 15 padded slots stay 0. + assert packed[1, 0].item() == 2 + 127 + assert y[1, 0].item() == pytest.approx(6.0) + assert torch.equal(y[1, 1:], torch.zeros(15, device="cuda")) + + +# --------------------------------------------------------------------------- # +# Stage 4: strided / non-contiguous high-precision input +# The module's whole point is "addressed by stride, no forced .contiguous()". +# Every other test passes a contiguous tensor; these pass non-contiguous views +# (non-unit stride_col, and a transposed stride_row=1 layout) and assert the +# packed bytes + dequant are bit-identical to the contiguous equivalent. +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("layout", ["column_slice", "transpose"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_kernel_strided_input_matches_contiguous(layout, dtype): + rows, cols = 8, 64 + if layout == "column_slice": + # Allocate 2x wide, take every other column -> stride (2*cols, 2). + big = _rand((rows, cols * 2), dtype).cuda() + x_strided = big[:, ::2] + else: + # Transpose a contiguous (cols, rows) tensor -> shape (rows, cols), + # stride (1, rows): stride_row=1, non-unit stride_col. + base = _rand((cols, rows), dtype).cuda() + x_strided = base.t() + assert not x_strided.is_contiguous() + assert x_strided.shape == (rows, cols) + + packed_s, nb_s = _launch_to(x_strided) + y_s = _launch_from(packed_s, nb_s, dtype) + + x_contig = x_strided.contiguous() + packed_c, nb_c = _launch_to(x_contig) + y_c = _launch_from(packed_c, nb_c, dtype) + + assert torch.equal(packed_s, packed_c), "packed bytes differ between strided and contiguous" + assert torch.equal(y_s, y_c), "dequant differs" + + +@pytest.mark.parametrize("bpp", [1, 16, 64, 256]) +def test_kernel_blocks_per_program_invariant(bpp): + x = _rand((32, 256), torch.float32).cuda() + packed0, nb = _launch_to(x, bpp=BLOCKS_PER_PROG_DEFAULT) + y0 = _launch_from(packed0, nb, torch.float32, bpp=BLOCKS_PER_PROG_DEFAULT) + packed1, nb1 = _launch_to(x, bpp=bpp) + y1 = _launch_from(packed1, nb1, torch.float32, bpp=bpp) + assert torch.equal(packed0, packed1) + assert torch.equal(y0, y1) + + +def test_kernel_zeros_and_block_independence(): + big = torch.full((1, BLOCK_SIZE), 100.0, device="cuda") + small = torch.full((1, BLOCK_SIZE), 0.01, device="cuda") + joint_x = torch.cat([big, small], dim=0) + packed_j, nbj = _launch_to(joint_x) + yj = _launch_from(packed_j, nbj, torch.float32) + packed_a, nba = _launch_to(small) + ya = _launch_from(packed_a, nba, torch.float32) + assert torch.equal(yj[1:2], ya) # small block unaffected by the large one + + +# --------------------------------------------------------------------------- # +# Stage 5 reference: full-tensor clamp-15 QDQ reusing mx.py primitives +# Unlike _mx6_clamp15_blockref (block-granularity, no transpose/padding), this +# operates on the full tensor at axis=-1 (transpose + pad handled internally +# by mx.py's _reshape_to_blocks), matching what the Stage 5 host wrappers do. +# --------------------------------------------------------------------------- # +def _mx6_clamp15_ref_via_mxpy(x, block_size=BLOCK_SIZE, quant_bit=QUANT_BIT): + """Second independent clamp-15 reference reusing mx.py production primitives. + + Same math as mx6_fake_quantize but (a) value clamp fixed to +/-15 (matching + the pack path) and (b) scale exponent clamped to -126 (matching the kernel / + mxfp4 FTZ-independence). Operates on a full tensor (axis=-1).""" + axis = -1 + input_dtype = x.dtype + input_shape = list(x.shape) + input_shape[-1], input_shape[axis] = input_shape[axis], input_shape[-1] + + block_x = _mxref._reshape_to_blocks(x.detach(), block_size, axis) + block_x = torch.nan_to_num(block_x, nan=0.0, posinf=0.0, neginf=0.0) + amax, _ = torch.max(torch.abs(block_x), dim=-1, keepdim=True) + max_exp = _mxref._t_exponent(amax) + + inp = _mxref._reshape_to_blocks(x, block_size, axis) + t_exp = _mxref._t_exponent(inp) + demote = (max_exp - t_exp) >= 1 + + n = _mxref.SHARED_PRIME_BIT_GROUP + fs = demote.shape + demote = demote.reshape(*fs[:-1], fs[-1] // n, n) + demote = torch.sum(demote, -1, keepdim=True) == n + demote = demote.repeat(*([1] * (demote.dim() - 1)), n).reshape(fs) + + shared_exp = max_exp - demote.long() + q_hi = (1 << (quant_bit - 1)) - 1 + scale_exp = torch.clamp(shared_exp - quant_bit + 2, min=-126) + scale = torch.pow(2.0, scale_exp) + out = torch.round(inp / scale).clamp(-q_hi, q_hi) * scale + out = out.reshape(out.size(0), -1) + out = out[:, : input_shape[-1]].reshape(input_shape).to(input_dtype) + return out.transpose(axis, -1) + + +# --------------------------------------------------------------------------- # +# Stage 5: host wrappers (convert_to_mx6 / convert_from_mx6) +# The Stage 4 tests above launch the grid kernels directly on an already-2D, +# block-aligned tensor. These tests instead go through the public host +# wrappers, exercising what Stage 4 deliberately skips: axis transpose, +# arbitrary tensor rank, padding-aware shape restore, and non-contiguous +# input arising naturally from the transpose (rather than constructed via +# slicing/``.t()`` as in the Stage 4 strided-input tests). +# --------------------------------------------------------------------------- # +def _host_roundtrip(x, block_size=BLOCK_SIZE, axis=-1): + packed = convert_to_mx6(x, block_size=block_size, axis=axis) + return convert_from_mx6(packed, x.dtype, x.shape, block_size=block_size, axis=axis) + + +def test_host_output_dtype_and_shape(): + x = _rand((4, 64), torch.float32).cuda() + packed = convert_to_mx6(x, block_size=BLOCK_SIZE) + n_blocks = (4 * 64) // BLOCK_SIZE + assert packed.dtype == torch.uint8 + assert packed.shape == (n_blocks, _N_PACKED_BYTES) + + +def test_host_max_exp_biased_range(): + # E8M0 stores true exponent + 127; finite fp32 true exponent in [-126, 127], + # biased to [1, 254]; should never be 0 or 255. + x = _rand((16, 64), torch.float32).cuda() + packed = convert_to_mx6(x, block_size=BLOCK_SIZE) + assert packed[:, 0].min().item() >= 1 + assert packed[:, 0].max().item() <= 254 + + +@pytest.mark.parametrize("shape", [(4, 64), (8, 16), (2, 3, 32), (3, 40), (1, 4096)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_host_roundtrip_matches_clamp15_ref(shape, dtype): + x = _rand(shape, dtype).cuda() + ref = _mx6_clamp15_ref_via_mxpy(x, block_size=BLOCK_SIZE) + out = _host_roundtrip(x) + assert out.shape == x.shape + assert out.dtype == dtype + assert torch.equal(out, ref), \ + f"max diff={(out.float() - ref.float()).abs().max().item()}" + + +@pytest.mark.parametrize("axis", [0, 1, -1]) +def test_host_roundtrip_axis(axis): + x = _rand((32, 64), torch.float32).cuda() + out = _host_roundtrip(x, block_size=BLOCK_SIZE, axis=axis) + assert out.shape == x.shape + assert out.dtype == x.dtype + ref = _mx6_clamp15_ref_via_mxpy( + x.transpose(axis, -1).contiguous(), block_size=BLOCK_SIZE + ).transpose(axis, -1).contiguous() + assert torch.equal(out, ref), \ + f"axis={axis} max diff={(out - ref).abs().max().item()}" + + +@pytest.mark.parametrize("blocks_per_program", [1, 16, 64, 256]) +def test_host_blocks_per_program_does_not_change_numerics(blocks_per_program): + x = _rand((32, 512), torch.float32).cuda() + ref = _host_roundtrip(x) + packed = convert_to_mx6(x, block_size=BLOCK_SIZE, blocks_per_program=blocks_per_program) + out = convert_from_mx6(packed, x.dtype, x.shape, block_size=BLOCK_SIZE, + blocks_per_program=blocks_per_program) + assert torch.equal(out, ref) + + +def test_host_zeros_stay_zero(): + x = torch.zeros((4, 64), dtype=torch.float32).cuda() + out = _host_roundtrip(x) + assert torch.equal(out, x) + + +def test_host_block_independence(): + big = torch.full((1, BLOCK_SIZE), 100.0).cuda() + small = torch.full((1, BLOCK_SIZE), 0.01).cuda() + joint = _host_roundtrip(torch.cat([big, small], dim=0)) + alone = _host_roundtrip(small) + assert torch.equal(joint[1:2], alone) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_host_non_contiguous_input(dtype): + """Non-contiguous memory input (transposed stride) produces the same result + as the equivalent contiguous tensor: the host's ``transpose(axis, -1)`` plus + the kernel's stride-addressed load means no ``.contiguous()`` copy is forced.""" + x = _rand((64, 32), dtype).cuda() + x_t = x.T # shape (32, 64), non-contiguous + assert not x_t.is_contiguous() + out_t = _host_roundtrip(x_t, axis=0) + ref_t = _mx6_clamp15_ref_via_mxpy(x, block_size=BLOCK_SIZE).T.contiguous() + assert torch.equal(out_t, ref_t), \ + f"non-contiguous max diff={(out_t.float() - ref_t.float()).abs().max().item()}" + + +def test_host_non_contiguous_packed_input(): + x = _rand((4, 64), torch.float32).cuda() + packed = convert_to_mx6(x) + storage = torch.empty( + (packed.shape[0], packed.shape[1] * 2), + dtype=packed.dtype, + device=packed.device, + ) + packed_view = storage[:, ::2] + packed_view.copy_(packed) + assert not packed_view.is_contiguous() + + expected = convert_from_mx6(packed, x.dtype, x.shape) + actual = convert_from_mx6(packed_view, x.dtype, x.shape) + assert torch.equal(actual, expected) + + +def test_host_padding_non_divisible_last_dim(): + x = _rand((3, 40), torch.float32).cuda() # 40 not divisible by 16 + out = _host_roundtrip(x, block_size=BLOCK_SIZE) + ref = _mx6_clamp15_ref_via_mxpy(x, block_size=BLOCK_SIZE) + assert out.shape == x.shape + assert torch.equal(out, ref) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_host_divergence_vs_mxpy_31clamp_is_tiny(dtype): + """Host-level counterpart of test_divergence_vs_mxpy_31clamp_is_tiny: the + +/-15 clamp only diverges from raw mx6_fake_quantize at the rare "demoted + AND rounds to +/-16" edge.""" + x = _rand((8, 256), dtype).cuda() + out = _host_roundtrip(x) + mxpy = mx6_fake_quantize(x, block_size=BLOCK_SIZE) + n_div = (out.float() != mxpy.float()).sum().item() + assert n_div < x.numel() * 0.05, \ + f"divergence {n_div}/{x.numel()} too large, likely a bug rather than the +/-16 edge" + + +# --------------------------------------------------------------------------- # +# Stage 5: invalid input rejection +# --------------------------------------------------------------------------- # +def test_host_rejects_non_float_dtype(): + x = torch.randint(0, 10, (4, 16)).cuda() + with pytest.raises(AssertionError): + convert_to_mx6(x) + + +def test_host_rejects_float16(): + x = _rand((4, 16), torch.float16).cuda() + with pytest.raises(AssertionError): + convert_to_mx6(x) + + +def test_host_rejects_block_size_not_16(): + x = _rand((4, 64), torch.float32).cuda() + with pytest.raises(AssertionError): + convert_to_mx6(x, block_size=24) # not 16 + with pytest.raises(AssertionError): + convert_to_mx6(x, block_size=32) # even though multiple of 16, only 16 is supported + + +@pytest.mark.parametrize("blocks_per_program", [0, 3, -1]) +def test_host_rejects_invalid_blocks_per_program(blocks_per_program): + x = _rand((4, 16), torch.float32).cuda() + packed = convert_to_mx6(x) + with pytest.raises(AssertionError, match="positive power of two"): + convert_to_mx6(x, blocks_per_program=blocks_per_program) + with pytest.raises(AssertionError, match="positive power of two"): + convert_from_mx6( + packed, + x.dtype, + x.shape, + blocks_per_program=blocks_per_program, + ) + + +def test_host_rejects_empty_quantization_axis(): + x = torch.empty((2, 0), dtype=torch.float32, device="cuda") + with pytest.raises(AssertionError, match="non-empty quantization axis"): + convert_to_mx6(x) + + packed = torch.empty((0, _N_PACKED_BYTES), dtype=torch.uint8, device="cuda") + with pytest.raises(AssertionError, match="non-empty quantization axis"): + convert_from_mx6(packed, torch.float32, x.shape) + + +def test_host_rejects_unknown_out_dtype(): + x = _rand((4, 16), torch.float32).cuda() + packed = convert_to_mx6(x) + with pytest.raises(AssertionError): + convert_from_mx6(packed, torch.int32, x.shape) + + +def test_host_rejects_bad_packed_dtype(): + x = _rand((4, 16), torch.float32).cuda() + packed = convert_to_mx6(x).to(torch.int8) # wrong dtype (should be uint8) + with pytest.raises(AssertionError): + convert_from_mx6(packed, torch.float32, x.shape) + + +def test_host_rejects_out_shape_inconsistent_with_packed(): + x = _rand((4, 64), torch.float32).cuda() + packed = convert_to_mx6(x) + with pytest.raises(AssertionError): + convert_from_mx6(packed, torch.float32, (3, 64)) # wrong row count for n_blocks + + +@pytest.mark.parametrize("shape", [(4, 64), (8, 96), (16, 256)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_two_independent_refs_agree_and_match_kernel(shape, dtype): + x = _rand(shape, dtype).cuda() + rows, cols = shape + n_blocks = rows * cols // BLOCK_SIZE + ref_hand = _mx6_clamp15_blockref( + x.reshape(n_blocks, BLOCK_SIZE), BLOCK_SIZE, PRIME_GROUP, QUANT_BIT + ).to(dtype).reshape(rows, cols) + ref_mxpy = _mx6_clamp15_ref_via_mxpy(x) + # (1) two independent references must agree bit-exact. + assert torch.equal(ref_hand, ref_mxpy), "hand ref vs mx.py-based ref diverge" + # (2) kernel round-trip must match the mx.py-based reference. + packed, nb = _launch_to(x) + y = _launch_from(packed, nb, dtype).reshape(rows, cols) + assert torch.equal(y, ref_mxpy), "kernel vs mx.py-based ref diverge" + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_divergence_vs_mxpy_31clamp_is_tiny(dtype): + """Our +/-15 clamp diverges from raw mx6_fake_quantize (demoted quant_max=31) + only where a demoted element rounds to +/-16; that fraction must be tiny.""" + x = _rand((32, 256), dtype).cuda() + rows, cols = x.shape + packed, nb = _launch_to(x) + y = _launch_from(packed, nb, dtype).reshape(rows, cols) + mxpy = mx6_fake_quantize(x, block_size=BLOCK_SIZE) + n_div = (y.float() != mxpy.float()).sum().item() + assert n_div < x.numel() * 0.05, ( + f"divergence {n_div}/{x.numel()} too large, likely a bug rather than " + f"the +/-16 demote edge" + ) + # Where they differ, mx.py must have the larger magnitude (it kept +/-16..31). + diff = y.float() != mxpy.float() + assert torch.all(mxpy.float()[diff].abs() >= y.float()[diff].abs()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v", "-p", "no:cacheprovider"])) diff --git a/tests/unittest/mx9_mx6/test_mx6_wa_integration.py b/tests/unittest/mx9_mx6/test_mx6_wa_integration.py index a9b4b9d0..d916d564 100644 --- a/tests/unittest/mx9_mx6/test_mx6_wa_integration.py +++ b/tests/unittest/mx9_mx6/test_mx6_wa_integration.py @@ -8,6 +8,12 @@ recipe yaml -> QuantizationModifier -> Linear quantization_scheme -> post_step dynamic-weight skip -> wrapped Linear forward -> MX6 W+A QDQ. + +MX6 W+A QDQ dispatches to the REAL packed Triton kernel (``convert_to_mx6`` / +``convert_from_mx6``), not the ``mx6_fake_quantize`` emulation -- see +``alto/models/patcher.py`` (mirrors mx9's method, commit ``95b1114``). This +test therefore requires CUDA and counts calls on both functions to prove the +real kernel path fires and the emulation does not. """ import importlib.util @@ -15,6 +21,7 @@ import sys import types +import pytest import torch import yaml @@ -129,9 +136,10 @@ def _load_mx6_modifier_from_recipe(monkeypatch): return QuantizationModifier(**mod_args), quantize_mod +@pytest.mark.skipif(not torch.cuda.is_available(), reason="mx6 dispatch launches the real Triton kernel") def test_mx6_wa_recipe_toy_linear_lifecycle(monkeypatch): modifier, quantize_mod = _load_mx6_modifier_from_recipe(monkeypatch) - model = torch.nn.Sequential(torch.nn.Linear(16, 16, bias=False)) + model = torch.nn.Sequential(torch.nn.Linear(16, 16, bias=False)).cuda() linear = model[0] modifier.initialize([model]) @@ -144,16 +152,31 @@ def test_mx6_wa_recipe_toy_linear_lifecycle(monkeypatch): assert not hasattr(linear, "weight_observer") assert not hasattr(linear, "input_observer") - calls = [] - original_mx6 = quantize_mod.mx6_fake_quantize + # format=="mx6" dispatches to the REAL packed kernel (see patcher.py), not + # the mx6_fake_quantize emulation. Count calls on BOTH to prove which path + # actually fires: emulation must stay untouched, the real kernel must be + # hit exactly twice (weight + input activation). + import alto.kernels.mx.mx6_quantization as mx6_kernel_mod + + fake_calls = [] + original_mx6_fake = quantize_mod.mx6_fake_quantize + + def counted_mx6_fake(input_tensor, *args, **kwargs): + fake_calls.append(tuple(input_tensor.shape)) + return original_mx6_fake(input_tensor, *args, **kwargs) + + monkeypatch.setattr(quantize_mod, "mx6_fake_quantize", counted_mx6_fake) + + packed_calls = [] + original_convert_to_mx6 = mx6_kernel_mod.convert_to_mx6 - def counted_mx6(input_tensor, *args, **kwargs): - calls.append(tuple(input_tensor.shape)) - return original_mx6(input_tensor, *args, **kwargs) + def counted_convert_to_mx6(input_tensor, *args, **kwargs): + packed_calls.append(tuple(input_tensor.shape)) + return original_convert_to_mx6(input_tensor, *args, **kwargs) - monkeypatch.setattr(quantize_mod, "mx6_fake_quantize", counted_mx6) + monkeypatch.setattr(mx6_kernel_mod, "convert_to_mx6", counted_convert_to_mx6) - x = torch.randn(2, 16) + x = torch.randn(2, 16).cuda() modifier.pre_step([model]) modifier.post_step([model]) # must skip static weight baking for MX6 dynamic weight assert not hasattr(linear, "weight_observer") @@ -161,8 +184,9 @@ def counted_mx6(input_tensor, *args, **kwargs): out = model(x) assert out.shape == (2, 16) - assert len(calls) == 2 - assert (2, 16) in calls # input activation QDQ - assert (16, 16) in calls # weight QDQ + assert len(fake_calls) == 0 # emulation must NOT be used + assert len(packed_calls) == 2 # real packed kernel used for weight + activation + assert (2, 16) in packed_calls # input activation QDQ + assert (16, 16) in packed_calls # weight QDQ modifier.finalize([model]) diff --git a/tests/unittest/mx9_mx6/test_mx9_dispatch.py b/tests/unittest/mx9_mx6/test_mx9_dispatch.py index c22745a2..7e835f56 100644 --- a/tests/unittest/mx9_mx6/test_mx9_dispatch.py +++ b/tests/unittest/mx9_mx6/test_mx9_dispatch.py @@ -3,11 +3,16 @@ # SPDX-License-Identifier: MIT """Tests for the MX9 dispatch wiring in ``ModelPatcher.patch_fake_quantize``. -The MX9 kernel itself is covered by ``test_mx9_quantize.py``. This file tests the -*wiring* one layer up: after ``patch_fake_quantize()`` replaces +The MX9 kernel itself is covered by ``test_mx9_quantize.py`` (fake-quant +emulation) and ``test_mx9_quantization.py`` (real packed kernel). This file +tests the *wiring* one layer up: after ``patch_fake_quantize()`` replaces ``compressed_tensors...forward.fake_quantize``, a ``QuantizationArgs`` carrying -``format == "mx9"`` must route to ``mx9_fake_quantize``, while plain int8 args -(no ``format``) must fall through to the original implementation untouched. +``format == "mx9"`` must route to the REAL packed Triton kernel +(``convert_to_mx9`` / ``convert_from_mx9``), NOT the ``mx9_fake_quantize`` +emulation -- this is the intentional, validated design (see commit +``95b1114``: "feat: mx9: wire packed kernel into patcher and fix +convert_from_mx9", validated end-to-end on LLaMA-3.2-1B). Plain int8 args (no +``format``) must fall through to the original implementation untouched. Also checks that ``format_registry.inject_format_field()`` makes the ``format`` field survive pydantic validation (otherwise the recipe value is silently @@ -68,21 +73,31 @@ def test_format_field_defaults_none_for_plain_int8(): # --------------------------------------------------------------------------- # # dispatch routing # --------------------------------------------------------------------------- # +@pytest.mark.skipif(not torch.cuda.is_available(), reason="mx9 dispatch launches the real Triton kernel") def test_mx9_args_dispatch_to_mx9_kernel(): - """format == "mx9" must route fake_quantize to mx9_fake_quantize bit-exact.""" + """format == "mx9" must route fake_quantize to the REAL packed kernel + (convert_to_mx9 -> convert_from_mx9 round trip) bit-exact. Requires CUDA + since the packed kernel launches a Triton kernel; NOT compared against + mx9_fake_quantize, which intentionally diverges at the rare demoted + +/-128 boundary (see test_mx9_quantization.py:: + test_divergence_vs_mxpy_255clamp_is_tiny).""" + from alto.kernels.mx.mx9_quantization import convert_to_mx9, convert_from_mx9 + torch.manual_seed(0) - x = torch.randn(3, 40, dtype=torch.float32) + x = torch.randn(3, 40, dtype=torch.float32).cuda() args = _mx9_args() patched_out = forward_module.fake_quantize( x=x, - scale=torch.ones(1), + scale=torch.ones(1, device="cuda"), zero_point=None, args=args, g_idx=None, global_scale=None, ) - expected = mx9_fake_quantize(x, block_size=BLOCK_SIZE) + + packed = convert_to_mx9(x) + expected = convert_from_mx9(packed, x.dtype, x.shape) assert torch.equal(patched_out, expected) diff --git a/tests/unittest/mx9_mx6/test_mx9_quantization.py b/tests/unittest/mx9_mx6/test_mx9_quantization.py new file mode 100644 index 00000000..93dccac5 --- /dev/null +++ b/tests/unittest/mx9_mx6/test_mx9_quantization.py @@ -0,0 +1,790 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Tests for packed MX9 quantization (convert_to_mx9 / convert_from_mx9). + +Acceptance criterion: unpack(pack(x)) == mx9_clamp127_ref(x) bit-exact. +clamp127_ref is the clamp-127 variant of mx9_fake_quantize (demoted +/-128 +truncated to +/-127), serving as the true reference for the pack path. + +``convert_to_mx9``/``convert_from_mx9`` exchange a SINGLE packed uint8 tensor +(byte layout ``[max_exp(1B) | prime(1B) | q(16B)]`` per block; ``q`` is int8 +values bit-reinterpreted as uint8), not a three-tensor tuple -- mirroring +``mx6_quantization.py``. Most tests below go through the ``_roundtrip`` helper +and never see the packed layout directly; only the tests that inspect +individual components (storage-format / intermediates / prime-bitmap / +blocks-per-program / invalid-dtype tests) slice into ``packed`` explicitly. + +Tests are organized in four layers: + 1. Storage format assertions: dtype / shape correctness. + 2. Round-trip bit-exact: vs clamp127_ref, covering shape/dtype/axis. + 3. Boundary & properties: zeros, Inf, block independence, idempotency, + padding, +/-128 edge. + 4. Invalid input rejection. +""" + +import pytest +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from alto.modifiers.quantization import mx as _mxref +from alto.modifiers.quantization.mx import mx9_fake_quantize, BLOCK_SIZE +from alto.kernels.mx.mx9_quantization import ( + PRIME_GROUP, + convert_to_mx9, + convert_from_mx9, +) +from alto.kernels.mx._mx_common import ( + _floor_exp, + _round_half_even, + _calculate_mx_exp as _calculate_mx9_exp, +) +from alto.kernels.fp4.testing_utils import calc_snr + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + +# Packed byte layout (single tensor): [max_exp(1B) | prime(N_PRIME_BYTES) | q(BLOCK_SIZE)]. +_N_PRIME_BYTES = (BLOCK_SIZE // 2) // 8 # 1 +_PRIME_OFFSET = 1 +_Q_OFFSET = _PRIME_OFFSET + _N_PRIME_BYTES # 2 +_N_PACKED_BYTES = _Q_OFFSET + BLOCK_SIZE # 18 + +# --------------------------------------------------------------------------- # +# Reference implementation: clamp-127 variant of mx9_fake_quantize +# --------------------------------------------------------------------------- # + +def _bit_exponent(t: torch.Tensor) -> torch.Tensor: + """Extract exponent field from float bits with bias removal (approx + floor(log2|t|)), branching by dtype. + + Aligned with the kernel's ``_floor_exp`` / mx.py's + ``_exponent_frexp_no_exception``, but independently inlined here (does not + call mx.py) to keep the two reference implementations independent. + Must operate on *native dtype* (fp16 bias=15, bf16/fp32 bias=127); callers + must NOT pre-cast to .float(). NaN/Inf are zeroed before extraction to + prevent pollution. ``& mask`` clears sign-extension bits from arithmetic + right shift, so negative values are safe. Returns long. + """ + t = torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0) + if t.dtype == torch.float32: + return (((t.view(torch.int32) >> 23) & 0xFF) - 127).long() + if t.dtype == torch.bfloat16: + return (((t.view(torch.int16) >> 7) & 0xFF) - 127).long() + if t.dtype == torch.float16: + return (((t.view(torch.int16) >> 10) & 0x1F) - 15).long() + raise ValueError(f"unsupported dtype {t.dtype}") + + +def _mx9_clamp127_ref(x: torch.Tensor, block_size: int = BLOCK_SIZE) -> torch.Tensor: + """Clamp-127 variant of mx9_fake_quantize, serving as the pack path reference. + + Only differs from mx9_fake_quantize in that demoted elements' quant_max is + also truncated to 127 (matching the pack path's int8 clamp ceiling). The two + diverge only for the extremely rare "demoted AND round reaches +/-128" edge. + + Exponent extraction uses bit-field extraction (``_bit_exponent``, on native + dtype), aligned with kernel / mx.py -- an earlier version using + ``floor(log2())`` had fp32 rounding errors near octave boundaries, causing + occasional 1-octave divergence in small-magnitude regions; now eliminated. + """ + orig_dtype = x.dtype + orig_shape = x.shape + last = orig_shape[-1] + + # Keep native dtype for blocking (exponent extraction requires native dtype). + pad = (block_size - last % block_size) % block_size + x2d = x.reshape(-1, last) + if pad: + x2d = F.pad(x2d, (0, pad)) + rows, cols = x2d.shape + x3d = x2d.reshape(rows, cols // block_size, block_size) # native dtype + + # Block shared exponent: amax(native dtype) -> bit extract exponent. + clean = torch.nan_to_num(x3d, nan=0.0, posinf=0.0, neginf=0.0) + amax = clean.abs().amax(dim=-1, keepdim=True) # [rows, nblk, 1] native dtype + max_exp = _bit_exponent(amax) # long + + # Per-element exponent + demote flag. + t_exp = _bit_exponent(x3d) # long + demote_elem = (max_exp - t_exp) >= 1 + + d = demote_elem.reshape(rows, cols // block_size, block_size // 2, 2) + pair = (d.sum(dim=-1, keepdim=True) == 2) + pair = pair.expand_as(d).reshape(rows, cols // block_size, block_size) + + shared_exp = max_exp - pair.long() + # Sync with kernel: scale exponent clamped to minimum normal -126. + scale_exp = torch.clamp(shared_exp - 8 + 2, min=-126) # quant_bit=8 + scale = torch.exp2(scale_exp.float()) + + q = torch.round(x3d.float() / scale).clamp(-127, 127) # clamp-127 variant + out3d = q * scale + + out2d = out3d.reshape(rows, cols) + if pad: + out2d = out2d[:, :last] + return out2d.reshape(orig_shape).to(orig_dtype) + + +def _mx9_clamp127_ref_via_mxpy(x: torch.Tensor, block_size: int = BLOCK_SIZE, + quant_bit: int = 8) -> torch.Tensor: + """Second independent reference: reuses mx.py's real primitives (_t_exponent / + _reshape_to_blocks / SHARED_PRIME_BIT_GROUP), only clamping the quantized + value to +/-127. + + Computes the same result as `_mx9_clamp127_ref` above (pure handwritten + reimplementation), but **reuses production code mx.py's building blocks**. + Cross-checking two independent references prevents "kernel and one reference + share the same bug" from going undetected. + """ + axis = -1 + input_dtype = x.dtype + input_shape = list(x.shape) + input_shape[-1], input_shape[axis] = input_shape[axis], input_shape[-1] + + block_x = _mxref._reshape_to_blocks(x.detach(), block_size, axis) + block_x = torch.nan_to_num(block_x, nan=0.0, posinf=0.0, neginf=0.0) + amax, _ = torch.max(torch.abs(block_x), dim=-1, keepdim=True) + max_exp = _mxref._t_exponent(amax) + + inp = _mxref._reshape_to_blocks(x, block_size, axis) + t_exp = _mxref._t_exponent(inp) + demote = (max_exp - t_exp) >= 1 + + n = _mxref.SHARED_PRIME_BIT_GROUP + fs = demote.shape + demote = demote.reshape(*fs[:-1], fs[-1] // n, n) + demote = torch.sum(demote, -1, keepdim=True) == n + demote = demote.repeat(*([1] * (demote.dim() - 1)), n).reshape(fs) + + shared_exp = max_exp - demote.long() + # Sync with kernel: scale exponent clamped to minimum normal -126 (mxfp4 style). + scale_exp = torch.clamp(shared_exp - quant_bit + 2, min=-126) + scale = torch.pow(2.0, scale_exp) + out = torch.round(inp / scale) + out = torch.clamp(out, -127, 127) * scale # clamp-127 variant + out = out.reshape(out.size(0), -1) + out = out[:, : input_shape[-1]].reshape(input_shape).to(input_dtype) + return out.transpose(axis, -1) + + +# --------------------------------------------------------------------------- # +# Probe kernels for device helpers (Stage 1 / Stage 2) +# --------------------------------------------------------------------------- # + +@triton.jit +def _probe_floor_exp_kernel(x_ptr, o_ptr, N: tl.constexpr): + i = tl.arange(0, N) + tl.store(o_ptr + i, _floor_exp(tl.load(x_ptr + i))) + + +@triton.jit +def _probe_round_kernel(y_ptr, o_ptr, N: tl.constexpr): + i = tl.arange(0, N) + tl.store(o_ptr + i, _round_half_even(tl.load(y_ptr + i))) + + +@triton.jit +def _probe_calc_kernel( + x_ptr, se_ptr, me_ptr, pr_ptr, + BPP: tl.constexpr, BLOCK_SIZE: tl.constexpr, PRIME_GROUP: tl.constexpr, +): + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + rows = tl.arange(0, BPP) + cols = tl.arange(0, BLOCK_SIZE) + offs = rows[:, None] * BLOCK_SIZE + cols[None, :] + x = tl.load(x_ptr + offs) + shared_exp, max_exp, pair = _calculate_mx9_exp(x, BPP, BLOCK_SIZE, PRIME_GROUP) + tl.store(se_ptr + offs, shared_exp) + tl.store(me_ptr + rows, max_exp) + pcols = tl.arange(0, N_PAIRS) + tl.store(pr_ptr + rows[:, None] * N_PAIRS + pcols[None, :], pair) + + +def _run_calc(xb, block_size=BLOCK_SIZE, prime_group=PRIME_GROUP): + n_blocks = xb.shape[0] + n_pairs = block_size // prime_group + se = torch.empty((n_blocks, block_size), dtype=torch.int32, device="cuda") + me = torch.empty((n_blocks,), dtype=torch.int32, device="cuda") + pr = torch.empty((n_blocks, n_pairs), dtype=torch.int32, device="cuda") + _probe_calc_kernel[(1,)](xb, se, me, pr, BPP=n_blocks, + BLOCK_SIZE=block_size, PRIME_GROUP=prime_group) + return se, me, pr + + +# --------------------------------------------------------------------------- # +# Stage 1: _floor_exp +# --------------------------------------------------------------------------- # + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_floor_exp_matches_bit_extraction(dtype): + torch.manual_seed(6) + x = (torch.randn(64, dtype=dtype, device="cuda") * 37.0) + x = torch.where(x.abs() < 1e-3, torch.full_like(x, 1e-3), x) + out = torch.empty(64, dtype=torch.int32, device="cuda") + _probe_floor_exp_kernel[(1,)](x, out, N=64) + ref = _bit_exponent(x).to(torch.int32) + assert torch.equal(out, ref) + + +def test_floor_exp_known_powers_of_two(): + x = torch.tensor([1.0, 2.0, 3.9, 4.0, 0.5, 0.25, 130.0, -8.0], + dtype=torch.float32, device="cuda") + out = torch.empty(8, dtype=torch.int32, device="cuda") + _probe_floor_exp_kernel[(1,)](x, out, N=8) + ref = torch.tensor([0, 1, 1, 2, -1, -2, 7, 3], dtype=torch.int32, device="cuda") + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# Stage 1: _round_half_even +# --------------------------------------------------------------------------- # + +def test_round_half_even_ties_and_regular(): + y = torch.tensor( + [-2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5, + 0.4, 0.6, -0.4, -0.6, 2.4, 2.6, -2.4, -2.6], + dtype=torch.float32, device="cuda", + ) + out = torch.empty_like(y) + _probe_round_kernel[(1,)](y, out, N=y.numel()) + ref = torch.round(y) + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# Stage 2: _calculate_mx9_exp +# --------------------------------------------------------------------------- # + +@pytest.mark.parametrize("n_blocks", [8, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_calculate_exp_matches_torch(n_blocks, dtype): + torch.manual_seed(6) + xb = torch.randn(n_blocks, BLOCK_SIZE, dtype=dtype, device="cuda") * 5.0 + se, me, pr = _run_calc(xb) + + # Reference: recompute in pure PyTorch. + clean = torch.nan_to_num(xb, nan=0.0, posinf=0.0, neginf=0.0) + amax = clean.abs().amax(dim=1, keepdim=True) + rme = _bit_exponent(amax).reshape(n_blocks) + t_exp = _bit_exponent(clean) + demote = (rme[:, None] - t_exp) >= 1 + n_pairs = BLOCK_SIZE // PRIME_GROUP + d = demote.reshape(n_blocks, n_pairs, PRIME_GROUP) + rpr = (d.sum(-1) == PRIME_GROUP).long() + pair_b = rpr.reshape(n_blocks, n_pairs, 1).expand(n_blocks, n_pairs, PRIME_GROUP).reshape(n_blocks, BLOCK_SIZE) + rse = rme[:, None] - pair_b + + assert torch.equal(me.long(), rme), "max_exp mismatch" + assert torch.equal(se.long(), rse), "shared_exp mismatch" + assert torch.equal(pr.long(), rpr), "prime pair mismatch" + + +def test_calculate_exp_constructed_prime_pattern(): + """amax=8 (max_exp=3); first 8 elements at 8 (not demoted), last 8 at 1 + (demoted). Pairs 0-3 -> prime=0; pairs 4-7 -> prime=1.""" + xb = torch.empty(1, BLOCK_SIZE, dtype=torch.float32, device="cuda") + xb[0, :8] = 8.0 + xb[0, 8:] = 1.0 + se, me, pr = _run_calc(xb) + assert me[0].item() == 3 + assert pr[0].tolist() == [0, 0, 0, 0, 1, 1, 1, 1] + assert se[0].tolist() == [3] * 8 + [2] * 8 + + +def test_calculate_exp_all_zero_block(): + """All-zero block: amax=0 -> _floor_exp(0)=-127; nothing demoted.""" + xb = torch.zeros(2, BLOCK_SIZE, dtype=torch.float32, device="cuda") + se, me, pr = _run_calc(xb) + assert torch.all(me == -127) + assert torch.all(pr == 0) + assert torch.all(se == -127) + + +def test_calculate_exp_nan_inf_do_not_pollute(): + """Inf/NaN sanitized out of exponent statistics: max_exp driven by finite values.""" + xb = torch.full((1, BLOCK_SIZE), 3.0, dtype=torch.float32, device="cuda") + xb[0, 0] = float("inf") + xb[0, 1] = float("nan") + se, me, pr = _run_calc(xb) + assert me[0].item() == 1 # finite amax=3 -> max_exp=1 + + +def _rand(shape, dtype, seed=0): + torch.manual_seed(seed) + return torch.randn(*shape, dtype=dtype) + + +def _rand_with_outliers(shape, dtype, seed=0, p=0.005, spike=100.0): + """Random data + sparse outlier spikes (ported from mxfp4 tests' prepare_data). + + ~p fraction of elements get a spike* addition, creating heavy-tailed blocks + that stress demote/prime/max_exp: one huge value in a block pulls the shared + scale up, squashing the remaining small values. Pure randn never hits this. + """ + torch.manual_seed(seed) + x = torch.randn(*shape, dtype=dtype) + mask = torch.bernoulli(torch.full_like(x, p)) + x = x + spike * torch.randn_like(x) * mask + return x + + +def _roundtrip(x, block_size=BLOCK_SIZE, axis=-1): + packed = convert_to_mx9(x, block_size=block_size, axis=axis) + return convert_from_mx9(packed, x.dtype, x.shape, + block_size=block_size, axis=axis) + + +# --------------------------------------------------------------------------- # +# Layer 1: Storage format assertions +# --------------------------------------------------------------------------- # + +def test_output_dtypes_and_shapes(): + x = _rand((4, 64), torch.float32).cuda() + packed = convert_to_mx9(x, block_size=16) + n_blocks = (4 * 64) // 16 + assert packed.dtype == torch.uint8 + assert packed.shape == (n_blocks, _N_PACKED_BYTES) # 1 + 1 + 16 = 18 + + +def test_max_exp_biased_range(): + # E8M0 stores true exponent + 127; finite fp32 true exponent in [-126, 127], + # biased to [1, 254]; should never be 0 or 255. + x = _rand((16, 64), torch.float32).cuda() + packed = convert_to_mx9(x, block_size=16) + max_exp = packed[:, 0] + assert max_exp.min().item() >= 1 + assert max_exp.max().item() <= 254 + + +def test_q_within_int8_range(): + x = _rand((8, 128), torch.float32).cuda() + packed = convert_to_mx9(x, block_size=16) + q = packed[:, _Q_OFFSET:].view(torch.int8) # bit-reinterpret uint8 -> int8 + assert q.min().item() >= -127 + assert q.max().item() <= 127 + + +# --------------------------------------------------------------------------- # +# Layer 1b: Per-component intermediate value verification +# (absorbed from _mx9_intermediate_check.py scaffold) +# Round-trip only verifies "combined result is correct" which can mask +# individual component errors; here we compare convert_to_mx9's q / max_exp / +# prime individually against torch-recomputed intermediates. +# --------------------------------------------------------------------------- # + +def _torch_intermediates(x, block_size=BLOCK_SIZE, quant_bit=8): + """Pure PyTorch recomputation of the three-part tuple (q_int8, max_exp_u8, + prime_u8, pair).""" + last = x.shape[-1] + x2d = x.contiguous().reshape(-1, last) + pad = (block_size - last % block_size) % block_size + if pad: + x2d = F.pad(x2d, (0, pad)) + rows, cols = x2d.shape + n_blocks = rows * (cols // block_size) + xb = x2d.reshape(n_blocks, block_size) + + clean = torch.nan_to_num(xb, nan=0.0, posinf=0.0, neginf=0.0) + amax = clean.abs().amax(dim=1, keepdim=True) + max_exp = _bit_exponent(amax) + t_exp = _bit_exponent(xb) + demote = (max_exp - t_exp) >= 1 + + n_pairs = block_size // 2 + d = demote.reshape(n_blocks, n_pairs, 2) + pair = (d.sum(-1) == 2).to(torch.int64) + pair_b = pair.reshape(n_blocks, n_pairs, 1).expand(n_blocks, n_pairs, 2).reshape(n_blocks, block_size) + shared_exp = max_exp - pair_b + + scale_exp = torch.clamp(shared_exp - quant_bit + 2, min=-126) + scale = torch.pow(2.0, scale_exp.to(torch.float32)) + q = torch.round(xb.to(torch.float32) / scale) + q = torch.clamp(q, -127, 127).to(torch.int8) + + max_exp_u8 = (max_exp.reshape(n_blocks) + 127).to(torch.uint8) + + n_prime_bytes = n_pairs // 8 + weights = (2 ** torch.arange(8, device=x.device)).to(torch.int64) + pb = pair.reshape(n_blocks, n_prime_bytes, 8) + prime_u8 = (pb * weights.view(1, 1, 8)).sum(-1).to(torch.uint8) + return q, max_exp_u8, prime_u8, pair + + +@pytest.mark.parametrize("shape", [(4, 64), (8, 16), (3, 40), (2, 3, 32)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_intermediates_match_torch_reference(shape, dtype): + """Each component (q, max_exp, prime) matches torch recomputation bit-exact.""" + x = _rand(shape, dtype).cuda() + packed = convert_to_mx9(x, block_size=BLOCK_SIZE) + rq, re, rp, _ = _torch_intermediates(x, block_size=BLOCK_SIZE) + q = packed[:, _Q_OFFSET:].view(torch.int8) + assert torch.equal(q, rq), f"q mismatch: {(q != rq).sum().item()} elements" + assert torch.equal(packed[:, 0], re), "max_exp mismatch" + assert torch.equal(packed[:, _PRIME_OFFSET:_Q_OFFSET], rp), "prime mismatch" + + +def test_intermediates_constructed_demote(): + """Constructed case: verify deterministic prime bitmap behavior. + + blk = [4, 1,1,...,1]: 4 has exp=2 (max); 1 has exp=0 (demoted). + pair0=(4,1) only one demoted -> prime=0; pair1..7=(1,1) both demoted -> prime=1. + """ + x = torch.ones((1, BLOCK_SIZE), dtype=torch.float32).cuda() + x[0, 0] = 4.0 + _, _, _, pair = _torch_intermediates(x) + assert pair[0].tolist() == [0, 1, 1, 1, 1, 1, 1, 1] + packed = convert_to_mx9(x) + rq, re, rp, _ = _torch_intermediates(x) + assert torch.equal(packed[:, _Q_OFFSET:].view(torch.int8), rq) + assert torch.equal(packed[:, 0], re) + assert torch.equal(packed[:, _PRIME_OFFSET:_Q_OFFSET], rp) + + +# --------------------------------------------------------------------------- # +# Layer 2: Round-trip bit-exact vs clamp127_ref +# --------------------------------------------------------------------------- # + +@pytest.mark.parametrize("shape", [(4, 64), (8, 16), (2, 3, 32), (3, 40), (1, 4096), (128, 4095)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("block_size", [16]) +def test_roundtrip_matches_clamp127_ref(shape, dtype, block_size): + x = _rand(shape, dtype).cuda() + ref = _mx9_clamp127_ref(x, block_size=block_size) + out = _roundtrip(x, block_size=block_size) + assert out.shape == x.shape + assert out.dtype == dtype + assert torch.equal(out, ref), \ + f"max diff={( out.float() - ref.float()).abs().max().item()}" + + +@pytest.mark.parametrize("axis", [0, 1, -1]) +def test_roundtrip_axis(axis): + x = _rand((32, 64), torch.float32).cuda() + out = _roundtrip(x, block_size=16, axis=axis) + assert out.shape == x.shape + assert out.dtype == x.dtype + ref = _mx9_clamp127_ref(x.transpose(axis, -1).contiguous(), + block_size=16).transpose(axis, -1).contiguous() + assert torch.equal(out, ref), \ + f"axis={axis} max diff={(out - ref).abs().max().item()}" + + +@pytest.mark.parametrize("blocks_per_program", [1, 16, 64, 256]) +def test_blocks_per_program_does_not_change_numerics(blocks_per_program): + x = _rand((32, 512), torch.float32).cuda() + ref = _roundtrip(x, block_size=BLOCK_SIZE) + packed = convert_to_mx9(x, block_size=BLOCK_SIZE, + blocks_per_program=blocks_per_program) + out = convert_from_mx9(packed, x.dtype, x.shape, + block_size=BLOCK_SIZE, + blocks_per_program=blocks_per_program) + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# Layer 3: Boundary & properties +# --------------------------------------------------------------------------- # + +def test_zeros_stay_zero(): + x = torch.zeros((4, 64), dtype=torch.float32).cuda() + out = _roundtrip(x) + assert torch.equal(out, x) + + +def test_inf_clamped_to_finite(): + x = torch.full((1, BLOCK_SIZE), 4.0).cuda() + x[0, 0] = float("inf") + x[0, 1] = float("-inf") + out = _roundtrip(x) + assert not torch.isinf(out).any() + ref = _mx9_clamp127_ref(x) + assert torch.equal(out[0, 2:], ref[0, 2:]) + + +def test_block_independence(): + big = torch.full((1, BLOCK_SIZE), 100.0).cuda() + small = torch.full((1, BLOCK_SIZE), 0.01).cuda() + joint = _roundtrip(torch.cat([big, small], dim=0)) + alone = _roundtrip(small) + assert torch.equal(joint[1:2], alone) + + +def test_qdq_idempotent(): + # Output of pack->unpack already lies on the MX9 grid; a second round-trip + # should produce identical results. + x = _rand((8, 64), torch.float32).cuda() + once = _roundtrip(x) + twice = _roundtrip(once) + assert torch.equal(once, twice) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_non_contiguous_input(dtype): + """Non-contiguous memory input (transposed stride) produces the same result + as the equivalent contiguous tensor. + + The kernel gathers the high-precision input by stride (no host-side + .contiguous() copy), so non-contiguous layouts are consumed directly. This + test verifies that path is correct. + """ + x = _rand((64, 32), dtype).cuda() + x_t = x.T # shape (32, 64), non-contiguous + assert not x_t.is_contiguous() + out_t = _roundtrip(x_t, axis=0) + ref_t = _mx9_clamp127_ref(x, block_size=BLOCK_SIZE).T.contiguous() + assert torch.equal(out_t, ref_t), \ + f"non-contiguous max diff={(out_t.float()-ref_t.float()).abs().max().item()}" + + +def test_non_contiguous_packed_input(): + x = _rand((4, 64), torch.float32).cuda() + packed = convert_to_mx9(x) + storage = torch.empty( + (packed.shape[0], packed.shape[1] * 2), + dtype=packed.dtype, + device=packed.device, + ) + packed_view = storage[:, ::2] + packed_view.copy_(packed) + assert not packed_view.is_contiguous() + + expected = convert_from_mx9(packed, x.dtype, x.shape) + actual = convert_from_mx9(packed_view, x.dtype, x.shape) + assert torch.equal(actual, expected) + + +def test_padding_non_divisible_last_dim(): + x = _rand((3, 40), torch.float32).cuda() # 40 not divisible by 16 + out = _roundtrip(x, block_size=16) + ref = _mx9_clamp127_ref(x, block_size=16) + assert out.shape == x.shape + assert torch.equal(out, ref) + + +def test_demote_boundary_clamp127(): + """Construct a demoted element that rounds exactly to +/-128, verify it is + clamped to +/-127. + + Setup: max_exp=7 (amax in [128,256)), demote scale = 2^(7-7)=1. + x=127.6 -> round(127.6/1)=128 -> clamp127 truncates to 127 -> reconstructs 127*1=127. + """ + x = torch.zeros((1, BLOCK_SIZE), dtype=torch.float32).cuda() + x[0, 0] = 130.0 # amax, max_exp=7, non-demote + x[0, 1] = 128.0 # pair[0] neighbor (non-demote, ensures amax unchanged) + x[0, 2] = 127.6 # demote pair[1,2] + x[0, 3] = 1.0 # demote pair[1,2] neighbor, both demotable + out = _roundtrip(x) + ref = _mx9_clamp127_ref(x) + assert torch.equal(out, ref) + assert out[0, 2].item() == pytest.approx(127.0) + + +def test_prime_bitmap_round_trips(): + """Verify prime bitmap pack/unpack symmetry. + + Construct a block: first half (pairs 0-3) all demotable, second half + (pairs 4-7) not demotable. Expected prime byte: low 4 bits = 0xF, + high 4 bits = 0x0, i.e. 0x0F. + Demotable condition: max_exp - t_exp >= 1, i.e. element is at least 1 + octave below amax. + amax=8 (max_exp=3), first 8 elements=1.0 (t_exp=0, diff 3 >= 1, demotable); + last 8 elements=8.0 (t_exp=3, diff 0, not demotable). + """ + x = torch.zeros((1, BLOCK_SIZE), dtype=torch.float32).cuda() + x[0, :8] = 1.0 # t_exp=0, max_exp=3, diff 3 -> demotable + x[0, 8:] = 8.0 # t_exp=3, max_exp=3, diff 0 -> not demotable, amax + packed = convert_to_mx9(x, block_size=BLOCK_SIZE) + # pairs 0-3 (idx 0-7) all demote -> bits 0-3 = 1; pairs 4-7 not demote -> bits 4-7 = 0 + assert packed[0, _PRIME_OFFSET].item() == 0x0F, f"got {hex(packed[0, _PRIME_OFFSET].item())}" + out = convert_from_mx9(packed, x.dtype, x.shape) + ref = _mx9_clamp127_ref(x) + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# Layer 3b: Subnormal / tiny-value degenerate blocks +# (absorbed from _mx9_edge_check.py scaffold) +# Numerical kernels are most likely to break on 0/0, FTZ, and scale underflow; +# random data never reaches these cases. +# --------------------------------------------------------------------------- # + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_all_zero_blocks(dtype): + x = torch.zeros((4, 64), dtype=dtype).cuda() + out = _roundtrip(x) + assert torch.equal(out, _mx9_clamp127_ref(x)) + + +def test_one_zero_block_among_normal(): + x = _rand((4, 64), torch.float32).cuda() + x[1, 16:32] = 0.0 # row 1, second 16-block is all zeros + out = _roundtrip(x) + assert torch.equal(out, _mx9_clamp127_ref(x)) + + +@pytest.mark.parametrize("scale_pow", [-120, -135, -140]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_tiny_and_subnormal_blocks(scale_pow, dtype): + # -120: near scale clamp boundary (-126) normal small; -135/-140: subnormal. + # Verify scale clamped to minimum normal produces no 0/0, no FTZ dependence, + # and matches reference bit-exact. + x = (_rand((4, 64), dtype) * (2.0 ** scale_pow)).cuda() + out = _roundtrip(x) + ref = _mx9_clamp127_ref(x) + assert torch.equal(out, ref) + assert not torch.isnan(out).any() + + +# --------------------------------------------------------------------------- # +# Layer 3c: Dual independent reference cross-validation +# (absorbed from _mx9_roundtrip_check.py scaffold) +# Asserts "two independent references agree with each other" AND "kernel +# agrees with both", preventing shared bugs from going undetected. +# --------------------------------------------------------------------------- # + +@pytest.mark.parametrize("shape", [(4, 64), (3, 40), (2, 3, 32)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_two_independent_refs_agree_and_match_kernel(shape, dtype): + x = _rand(shape, dtype).cuda() + ref_handwritten = _mx9_clamp127_ref(x) # pure handwritten reimplementation + ref_via_mxpy = _mx9_clamp127_ref_via_mxpy(x) # reuses mx.py production primitives + # (1) Two independent references must agree (otherwise one has a bug). + assert torch.equal(ref_handwritten, ref_via_mxpy) + # (2) Kernel round-trip must match both references. + out = _roundtrip(x) + assert torch.equal(out, ref_via_mxpy) + + +def test_divergence_vs_mxpy_255clamp_is_tiny(): + """Verify that the "intentional divergence vs mx.py original (255-clamp)" is + tiny, confirming it only occurs at the rare +/-128 boundary. + + Divergence should be far less than total element count (only elements that + are both demoted AND have quantized value reaching +/-128). + """ + x = _rand((8, 256), torch.float32).cuda() + out = _roundtrip(x) + mxpy = mx9_fake_quantize(x) + n_div = (out.float() != mxpy.float()).sum().item() + assert n_div < x.numel() * 0.01, \ + f"divergence {n_div}/{x.numel()} too large, likely a bug rather than +/-128 edge" + + +# --------------------------------------------------------------------------- # +# Layer 3d: Outlier / heavy-tailed data + quality floor (inspired by mxfp4 tests) +# Pure randn cannot reach "one huge value in block pulls scale up, squashing +# remaining small values"; outlier spikes specifically stress demote/prime/ +# max_exp. SNR floor is a human-readable sanity check ("is quantization +# overall usable?") complementing bit-exact assertions: a completely broken +# kernel would yield near-0 or negative dB. +# --------------------------------------------------------------------------- # + +@pytest.mark.parametrize("shape", [(8, 256), (128, 512), (2, 3, 64)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_roundtrip_matches_ref_with_outliers(shape, dtype): + """Round-trip with heavy-tailed/outlier data still matches reference bit-exact.""" + x = _rand_with_outliers(shape, dtype).cuda() + ref = _mx9_clamp127_ref(x) + out = _roundtrip(x) + assert out.shape == x.shape + assert out.dtype == dtype + assert torch.equal(out, ref), \ + f"max diff={(out.float() - ref.float()).abs().max().item()}" + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_roundtrip_snr_floor_clean(dtype): + """Clean random data round-trip SNR should be well above a conservative floor. + + Floor is set conservatively to catch total failures only, not as a tight bound + (8-bit/block quantization on clean data typically yields 40+ dB). + """ + x = _rand((64, 512), dtype).cuda() + out = _roundtrip(x) + snr = calc_snr(x, out) + assert snr > 25.0, f"SNR={snr:.1f} dB too low, quantization likely broken" + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_roundtrip_snr_floor_outliers(dtype): + """With outlier spikes, SNR will drop (small values in outlier blocks get + squashed), but should still exceed a more conservative floor.""" + x = _rand_with_outliers((64, 512), dtype).cuda() + out = _roundtrip(x) + snr = calc_snr(x, out) + assert snr > 15.0, f"SNR={snr:.1f} dB too low, quantization likely broken" + + +# --------------------------------------------------------------------------- # +# Layer 4: Invalid input rejection +# --------------------------------------------------------------------------- # + +def test_rejects_non_float_dtype(): + x = torch.randint(0, 10, (4, 16)).cuda() + with pytest.raises(AssertionError): + convert_to_mx9(x) + + +def test_rejects_float16(): + x = _rand((4, 16), torch.float16).cuda() + with pytest.raises(AssertionError): + convert_to_mx9(x) + + +def test_rejects_block_size_not_16(): + x = _rand((4, 64), torch.float32).cuda() + with pytest.raises(AssertionError): + convert_to_mx9(x, block_size=24) # not 16 + with pytest.raises(AssertionError): + convert_to_mx9(x, block_size=32) # even though multiple of 16, only 16 is supported + + +@pytest.mark.parametrize("blocks_per_program", [0, 3, -1]) +def test_rejects_invalid_blocks_per_program(blocks_per_program): + x = _rand((4, 16), torch.float32).cuda() + packed = convert_to_mx9(x) + with pytest.raises(AssertionError, match="positive power of two"): + convert_to_mx9(x, blocks_per_program=blocks_per_program) + with pytest.raises(AssertionError, match="positive power of two"): + convert_from_mx9( + packed, + x.dtype, + x.shape, + blocks_per_program=blocks_per_program, + ) + + +def test_rejects_empty_quantization_axis(): + x = torch.empty((2, 0), dtype=torch.float32, device="cuda") + with pytest.raises(AssertionError, match="non-empty quantization axis"): + convert_to_mx9(x) + + packed = torch.empty((0, _N_PACKED_BYTES), dtype=torch.uint8, device="cuda") + with pytest.raises(AssertionError, match="non-empty quantization axis"): + convert_from_mx9(packed, torch.float32, x.shape) + + +def test_rejects_unknown_out_dtype(): + x = _rand((4, 16), torch.float32).cuda() + packed = convert_to_mx9(x) + with pytest.raises(AssertionError): + convert_from_mx9(packed, torch.int32, x.shape) + + +def test_rejects_bad_packed_dtype(): + x = _rand((4, 16), torch.float32).cuda() + packed = convert_to_mx9(x).to(torch.int8) # wrong dtype (should be uint8) + with pytest.raises(AssertionError): + convert_from_mx9(packed, torch.float32, x.shape) + + +def test_rejects_out_shape_inconsistent_with_packed(): + x = _rand((4, 64), torch.float32).cuda() + packed = convert_to_mx9(x) + with pytest.raises(AssertionError, match="inconsistent with packed tensor"): + convert_from_mx9(packed, torch.float32, (3, 64)) diff --git a/tests/unittest/mx9_mx6/test_mx9_wa_integration.py b/tests/unittest/mx9_mx6/test_mx9_wa_integration.py index f0901809..87dbfb89 100644 --- a/tests/unittest/mx9_mx6/test_mx9_wa_integration.py +++ b/tests/unittest/mx9_mx6/test_mx9_wa_integration.py @@ -8,6 +8,13 @@ recipe yaml -> QuantizationModifier -> Linear quantization_scheme -> post_step dynamic-weight skip -> wrapped Linear forward -> MX9 W+A QDQ. + +MX9 W+A QDQ dispatches to the REAL packed Triton kernel (``convert_to_mx9`` / +``convert_from_mx9``), not the ``mx9_fake_quantize`` emulation -- see +``alto/models/patcher.py`` and commit ``95b1114`` (validated end-to-end on +LLaMA-3.2-1B: wikitext loss 2.1141 vs the fake-quant path's 2.1140). This test +therefore requires CUDA and counts calls on both functions to prove the real +kernel path fires and the emulation does not. """ import importlib.util @@ -15,6 +22,7 @@ import sys import types +import pytest import torch import yaml @@ -127,9 +135,10 @@ def _load_mx9_modifier_from_recipe(monkeypatch): return QuantizationModifier(**mod_args), quantize_mod +@pytest.mark.skipif(not torch.cuda.is_available(), reason="mx9 dispatch launches the real Triton kernel") def test_mx9_wa_recipe_toy_linear_lifecycle(monkeypatch): modifier, quantize_mod = _load_mx9_modifier_from_recipe(monkeypatch) - model = torch.nn.Sequential(torch.nn.Linear(16, 16, bias=False)) + model = torch.nn.Sequential(torch.nn.Linear(16, 16, bias=False)).cuda() linear = model[0] modifier.initialize([model]) @@ -142,16 +151,31 @@ def test_mx9_wa_recipe_toy_linear_lifecycle(monkeypatch): assert not hasattr(linear, "weight_observer") assert not hasattr(linear, "input_observer") - calls = [] - original_mx9 = quantize_mod.mx9_fake_quantize + # format=="mx9" dispatches to the REAL packed kernel (see patcher.py), not + # the mx9_fake_quantize emulation. Count calls on BOTH to prove which path + # actually fires: emulation must stay untouched, the real kernel must be + # hit exactly twice (weight + input activation). + import alto.kernels.mx.mx9_quantization as mx9_kernel_mod + + fake_calls = [] + original_mx9_fake = quantize_mod.mx9_fake_quantize + + def counted_mx9_fake(input_tensor, *args, **kwargs): + fake_calls.append(tuple(input_tensor.shape)) + return original_mx9_fake(input_tensor, *args, **kwargs) + + monkeypatch.setattr(quantize_mod, "mx9_fake_quantize", counted_mx9_fake) + + packed_calls = [] + original_convert_to_mx9 = mx9_kernel_mod.convert_to_mx9 - def counted_mx9(input_tensor, *args, **kwargs): - calls.append(tuple(input_tensor.shape)) - return original_mx9(input_tensor, *args, **kwargs) + def counted_convert_to_mx9(input_tensor, *args, **kwargs): + packed_calls.append(tuple(input_tensor.shape)) + return original_convert_to_mx9(input_tensor, *args, **kwargs) - monkeypatch.setattr(quantize_mod, "mx9_fake_quantize", counted_mx9) + monkeypatch.setattr(mx9_kernel_mod, "convert_to_mx9", counted_convert_to_mx9) - x = torch.randn(2, 16) + x = torch.randn(2, 16).cuda() modifier.pre_step([model]) modifier.post_step([model]) # must skip static weight baking for MX9 dynamic weight assert not hasattr(linear, "weight_observer") @@ -159,8 +183,9 @@ def counted_mx9(input_tensor, *args, **kwargs): out = model(x) assert out.shape == (2, 16) - assert len(calls) == 2 - assert (2, 16) in calls # input activation QDQ - assert (16, 16) in calls # weight QDQ + assert len(fake_calls) == 0 # emulation must NOT be used + assert len(packed_calls) == 2 # real packed kernel used for weight + activation + assert (2, 16) in packed_calls # input activation QDQ + assert (16, 16) in packed_calls # weight QDQ modifier.finalize([model])