From 2de0111de532a447c72e4f8ca9e0f53f6b59d80e Mon Sep 17 00:00:00 2001 From: jiarwang Date: Mon, 29 Jun 2026 09:43:59 +0000 Subject: [PATCH 1/7] feat: mx9: add fused MX9 quantization kernel with pack/unpack --- alto/kernels/mx/__init__.py | 9 + alto/kernels/mx/mx9_quantization.py | 484 +++++++++++++++ .../unittest/mx9_mx6/test_mx9_quantization.py | 567 ++++++++++++++++++ 3 files changed, 1060 insertions(+) create mode 100644 alto/kernels/mx/__init__.py create mode 100644 alto/kernels/mx/mx9_quantization.py create mode 100644 tests/unittest/mx9_mx6/test_mx9_quantization.py diff --git a/alto/kernels/mx/__init__.py b/alto/kernels/mx/__init__.py new file mode 100644 index 00000000..f04649f0 --- /dev/null +++ b/alto/kernels/mx/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""MX real (fused Triton) kernels. + +The fake-quant reference / emulation lives in +``alto.modifiers.quantization.mx``; this package holds the fused GPU kernel +implementation of the same MX QDQ. +""" diff --git a/alto/kernels/mx/mx9_quantization.py b/alto/kernels/mx/mx9_quantization.py new file mode 100644 index 00000000..d95f3e9b --- /dev/null +++ b/alto/kernels/mx/mx9_quantization.py @@ -0,0 +1,484 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Packed MX9 quantization Triton device kernels (called by grid kernels). + +Unlike the fake-quant path (``alto/kernels/mx/quantize_triton.py`` which outputs +dequantized bf16), this module performs *real packed* quantization: each MX9 block +is compressed into a three-part tuple stored separately: + - ``q`` : per-element quantized integer (int8, symmetric clamp +/-127) + - ``max_exp`` : per-block shared exponent, uint8 E8M0 (floor(log2(amax)) + 127 bias) + - ``prime`` : per-block prime bitmap (1 bit per pair of 2 elements, indicates + whether that pair gets a 1-exponent demotion) + +The exponent extraction / scale / prime math is aligned with +``alto/modifiers/quantization/mx.py`` (exponent extracted from native dtype bits += floor; QDQ computed in fp32). The **only intentional divergence** is the +quantized value clamp: + + True 8-bit variant: quantized integers are symmetrically clamped to +/-127, + stored as **int8**, yielding exactly 9 bits/element (8b value + 0.5b amortized + max_exp + 0.5b prime). In contrast, mx.py uses quant_max=255 for demoted pairs + and can preserve +/-128. The two differ only for elements that are both demoted + AND whose quantized value reaches +/-128. + +Therefore the acceptance reference is **not** ``mx9_fake_quantize`` but its +clamp-127 variant (see tests): +``unpack(pack(x)) == mx9_clamp127_ref(x)`` bit-exact. + +Constraints: + - Blocks are formed along the last dimension only (host transposes the quant + axis to dim -1 then calls .contiguous()). + - ``block_size`` is fixed at 16: each block has 8 pairs which pack into + exactly 1 byte of prime bitmap. + +Key design decisions (lessons learned / non-obvious trade-offs -- understand +these before modifying): + 1. q uses int8 + clamp +/-127, not int16: demoted pairs use half-scale so + the quantized value can reach +/-128 (quant_max=255), which overflows signed + int8. We previously used int16 to preserve +/-128 bit-exactness, but that + yields ~17 bits/element (larger than bf16), defeating the purpose of packing. + Choosing clamp +/-127 achieves true 9 bits/element, at the cost of a 1-LSB + divergence from mx.py on the rare "demoted AND reaching +/-128" elements + (measured ~0.1%). + 2. Scale exponent is clamped to the minimum normal exponent -126 (following + mxfp4 / PyTorch #125557), rather than a post-hoc scale==0 guard: the latter + implicitly depends on GPU FTZ flushing subnormal scales to 0, which once + caused divergence of ~1e-41 for deep-subnormal blocks vs a no-FTZ reference. + Clamping keeps scale always normal, FTZ-independent, and degenerate blocks + deterministically output 0. + 3. Exponents are extracted from native dtype bits (fp32/bf16 >>... - 127); + do NOT pre-cast to fp32 on host. + 4. max_exp is stored as uint8 E8M0 (true exponent + 127), not int32: true + exponents for any finite input fall in [-127, 127], +127 -> [0, 254] which + fits uint8 without overflow; this achieves 1 byte/block and 9 bits/element. + +Known limitations: + - NaN: packed integers cannot represent NaN; blocks containing NaN will have + garbage q values (only the fake-quant path supports NaN pass-through). + - The acceptance reference is the custom clamp-127 variant (Quark-divergent), + not Quark/mx.py itself. +""" + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +BLOCK_SIZE = 16 +QUANT_BIT = 8 +PRIME_GROUP = 2 +BLOCKS_PER_PROG_DEFAULT = 64 + +# torch dtype -> triton dtype (used by unpack output cast). +_TORCH_TO_TL = { + torch.float32: tl.float32, + torch.bfloat16: tl.bfloat16, +} + + +@triton.jit +def _sanitize(x): + """Replace NaN / +/-Inf with 0; used only for the amax / exponent statistics + copy, does not touch the QDQ data path. + + (Faithfully replicates Quark/reference: exponent statistics use sanitized + values to prevent Inf from polluting the block scale; the actual round still + operates on the original x, letting Inf be clamped and NaN pass through.) + """ + x = tl.where(x != x, 0.0, x) + x = tl.where(x == float("inf"), 0.0, x) + x = tl.where(x == float("-inf"), 0.0, x) + return x + + +@triton.jit +def _floor_exp(x): + """floor(log2(|x|)): extract exponent field directly from native float bits. + + Branches by dtype, aligned with ``mx.py``'s ``_exponent_frexp_no_exception``: + - fp32 : (bits>>23)&0xFF - 127 + - bf16 : (bits>>7) &0xFF - 127 + The ``& mask`` also clears sign-bit extension from arithmetic right shift, + so negative values are safe. Always returns int32 to avoid int16 mixing. + """ + 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), pure tl + implementation to avoid libdevice dependency.""" + 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_mx9_exp( + x, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, +): + """Compute per-block shared exponent + per-element shared_exp + per-pair + prime bits. + + Input ``x``: this program's tile ``[BLOCKS_PER_PROG, BLOCK_SIZE]`` in + **native dtype** (each row is one MX9 block). Exponent extraction must be + done on native dtype, so do NOT pre-cast to fp32 on host. + + Returns: + - ``shared_exp`` : [BPP, BLOCK_SIZE] int32, per-element shared exponent + (demoted pairs get -1) + - ``max_exp`` : [BPP, 1] int32, per-block max exponent + - ``pair`` : [BPP, N_PAIRS] int32(0/1), per-pair prime bit + (1 = that pair gets 1-exponent demotion) + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + + # Sanitized copy used only for exponent statistics (NaN/Inf -> 0). + clean = _sanitize(x) + + # Block shared exponent: per-row (block) amax -> floor extract exponent. + # tl.max on some backends promotes the reduce result to fp32; cast back to + # native dtype to ensure _floor_exp uses the same branch as per-element + # t_exp, consistent with the mxfp4 kernel (see mxfp_quantization.py:72). + amax = tl.max(tl.abs(clean), axis=1, keep_dims=True).to(clean.dtype) # [BPP, 1] native dtype + max_exp = _floor_exp(amax) # [BPP, 1] int32 + + # Per-element exponent + demote flag: whether at least 1 octave below block max. + t_exp = _floor_exp(clean) # [BPP, BLOCK_SIZE] int32 + demote = (max_exp - t_exp) >= 1 # [BPP, BLOCK_SIZE] bool (max_exp broadcasts) + + # Prime: adjacent PRIME_GROUP(=2) elements form a pair; both must be demoted + # for the pair to be demoted. + d = demote.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP) + pair_keep = tl.sum(d, axis=2, keep_dims=True) == PRIME_GROUP # [BPP, N_PAIRS, 1] bool + + # Broadcast back to per-element to get per-element shared_exp. + pair_b = tl.broadcast_to(pair_keep, (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) # [BPP, BLOCK_SIZE] int32 + + # Pair bits (not broadcast) are kept for prime bitmap packing. + pair = pair_keep.reshape(BLOCKS_PER_PROG, N_PAIRS).to(tl.int32) + return shared_exp, max_exp.reshape(BLOCKS_PER_PROG), pair + + +@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 + pack: original x -> (q:int8 values, prime:uint8 bitmap). + + True 8-bit variant: quantized integers are symmetrically clamped to +/-127, + stored as int8 (achieving 9 bits/element). This intentionally diverges from + mx.py (which uses quant_max=255 for demoted pairs, preserving +/-128) at + the +/-128 boundary -- this variant serves as its own reference (see the + clamp-127 reference in tests). ``prime`` packs every 8 pair bits into one + byte. + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + N_PRIME_BYTES: tl.constexpr = N_PAIRS // 8 # block 16 -> 1; block 32 -> 2 + + # Per-element scale = 2^(shared_exp - quant_bit + 2). + # Following mxfp4 (mxfp_quantization.py:81-89 / PyTorch #125557): clamp the + # scale exponent to fp32 minimum normal exponent -126, ensuring scale is + # always normal (never subnormal or zero), eliminating 0/0 and GPU FTZ + # dependence from the source (only affects blocks with amax < 2^-120, which + # real data never reaches). + scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) + scale = tl.exp2(scale_exp.to(tl.float32)) # [BPP, BLOCK_SIZE], always normal (!=0) + + # Q part of QDQ operates on the *original* x (NaN will produce garbage, see below). + xf = x.to(tl.float32) + q = _round_half_even(xf / scale) + # True 8-bit: symmetric clamp to +/-127. + # - Normal data: demoted elements have |q|<128 mathematically; only rounding + # edge cases can produce 128, which is clamped to 127. + # - +/-Inf input: Inf/scale=Inf, clamped to +/-127 (does not pollute block + # scale since _sanitize already cleared it). + # - Intentional divergence from mx.py: mx.py uses quant_max=255 for demoted + # pairs, preserving +/-128. + q = tl.minimum(tl.maximum(q, -127.0), 127.0) + + # NaN limitation: packed integers cannot represent NaN; blocks containing NaN + # will have garbage q values (round-trip tests must exclude NaN inputs; only + # the fake-quant path supports NaN pass-through). + q_int = q.to(tl.int8) + + # Prime packing: N_PAIRS pair bits -> every 8 bits packed into 1 byte. + # weights = [1,2,4,...,128], generated via exp2 to avoid 1< reconstructed values. + + Rebuilds scale = 2^((max_exp - pair) - quant_bit + 2), then y = q * scale. + Note: max_exp has already been unbiased from E8M0 uint8 at the kernel entry; + it arrives here as int32. + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + N_PRIME_BYTES: tl.constexpr = N_PAIRS // 8 + + # Unpack prime bitmap: extract 8 bits from each byte -> one flag per pair. + weights = tl.exp2(tl.arange(0, 8).to(tl.float32)).to(tl.int32) # [8] + pbytes = prime.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PRIME_BYTES, 1) + bits = (pbytes // weights[None, None, :]) % 2 # [BPP, NB, 8] + pair = bits.reshape(BLOCKS_PER_PROG, N_PAIRS) # [BPP, N_PAIRS] + + # Broadcast back to per-element, rebuild per-element shared_exp. + 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 # [BPP, BLOCK_SIZE] (max_exp broadcasts) + + # Same scale exponent clamping as pack (-126 minimum) to ensure consistency. + scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) + scale = tl.exp2(scale_exp.to(tl.float32)) + y = q.to(tl.float32) * scale + return y.to(out_dtype) + + +# ============================================================================ +# Grid kernel (entry point): 1D grid, each program handles BLOCKS_PER_PROG +# complete 16-element blocks. Inputs/outputs are addressed as flattened +# [n_blocks, BLOCK_SIZE] views; the three-part tuple layout: +# q : [n_blocks, BLOCK_SIZE] (per-element, int8, clamp +/-127) +# max_exp : [n_blocks] (per-block, uint8 E8M0, true exp + 127) +# prime : [n_blocks, N_PRIME_BYTES] (per-block, N_PRIME_BYTES bytes, uint8 bitmap) +# ============================================================================ + + +@triton.jit +def _convert_to_mx9_kernel( + x_ptr, + q_ptr, + e_ptr, + p_ptr, + n_blocks, + BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, +): + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + + pid = tl.program_id(0) + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) # block indices for this program [BPP] + col = tl.arange(0, BLOCK_SIZE) + offs = blk[:, None] * BLOCK_SIZE + col[None, :] # input/output q offsets [BPP, BLOCK] + mask = blk[:, None] < n_blocks # tail block-level mask + blk_mask = blk < n_blocks + + tl.static_assert( + x_ptr.type.element_ty == tl.float32 or + x_ptr.type.element_ty == tl.bfloat16, + "x must be fp32 / bf16", + ) + x = tl.load(x_ptr + offs, mask=mask, other=0.0) # native dtype tile + + shared_exp, max_exp, pair = _calculate_mx9_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, + ) + + tl.store(q_ptr + offs, q_int, mask=mask) + # max_exp is the true exponent ([BPP] int32); store as E8M0 uint8: +127 bias. + e_store = (max_exp + 127).to(tl.uint8) + tl.store(e_ptr + blk, e_store, mask=blk_mask) + + pcol = tl.arange(0, N_PRIME_BYTES) + offs_p = blk[:, None] * N_PRIME_BYTES + pcol[None, :] + tl.store(p_ptr + offs_p, prime, mask=blk_mask[:, None]) + + +@triton.jit +def _convert_from_mx9_kernel( + q_ptr, + e_ptr, + p_ptr, + y_ptr, + n_blocks, + BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + + pid = tl.program_id(0) + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) + col = tl.arange(0, BLOCK_SIZE) + offs = blk[:, None] * BLOCK_SIZE + col[None, :] + mask = blk[:, None] < n_blocks + blk_mask = blk < n_blocks + + q = tl.load(q_ptr + offs, mask=mask, other=0) # int8 + # max_exp is stored as E8M0 uint8 (with +127 bias); subtract back to true exponent. + max_exp_u = tl.load(e_ptr + blk, mask=blk_mask, other=0) # [BPP] uint8 + max_exp = max_exp_u.to(tl.int32) - 127 + + pcol = tl.arange(0, N_PRIME_BYTES) + offs_p = blk[:, None] * N_PRIME_BYTES + pcol[None, :] + prime = tl.load(p_ptr + offs_p, mask=blk_mask[:, None], other=0) # [BPP, NB] uint8 + + y = _unpack_mx9( + q, max_exp[:, None], prime, OUT_DTYPE, + BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT, + ) + tl.store(y_ptr + offs, y, mask=mask) + + +# ============================================================================ +# Host wrappers (not yet registered as @triton_op / register_fake; using bare +# launch for ease of round-trip verification). +# ============================================================================ + + +def convert_to_mx9( + data_hp: torch.Tensor, + block_size: int = BLOCK_SIZE, + axis: int = -1, + blocks_per_program: int = BLOCKS_PER_PROG_DEFAULT, +): + """High-precision tensor -> packed MX9 three-part tuple (q:int8, max_exp:uint8 + E8M0, prime:uint8). + + Blocks are formed along the ``axis`` dimension (axis is transposed to the last + dim internally). The returned three-part tuple is a flattened view: + q : [n_blocks, block_size] int8 (clamp +/-127) + max_exp : [n_blocks] uint8 (E8M0, true exponent + 127) + prime : [n_blocks, n_prime_bytes] uint8 + Shape reconstruction info (original shape / axis / padding) must be passed + back by the caller in convert_from_mx9. + """ + assert data_hp.dtype in (torch.float32, torch.bfloat16), \ + f"mx9_quantization only supports fp32 / bf16, got {data_hp.dtype}" + assert block_size == 16, f"block_size only supports 16, got {block_size}" + + # Transpose quant axis to last dim, matching mxfp4 host processing. + data_hp = data_hp.transpose(axis, -1) + last = data_hp.shape[-1] + x2d = data_hp.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) + x_blocks = x2d.reshape(n_blocks, block_size).contiguous() + + n_prime_bytes = (block_size // 2) // 8 + q = torch.empty((n_blocks, block_size), dtype=torch.int8, device=data_hp.device) + max_exp = torch.empty((n_blocks,), dtype=torch.uint8, device=data_hp.device) + prime = torch.empty((n_blocks, n_prime_bytes), dtype=torch.uint8, device=data_hp.device) + + grid = (triton.cdiv(n_blocks, blocks_per_program),) + _convert_to_mx9_kernel[grid]( + x_blocks, q, max_exp, prime, n_blocks, + BLOCK_SIZE=block_size, + BLOCKS_PER_PROG=blocks_per_program, + PRIME_GROUP=PRIME_GROUP, + QUANT_BIT=QUANT_BIT, + ) + return q, max_exp, prime + + +def convert_from_mx9( + q: torch.Tensor, + max_exp: torch.Tensor, + prime: 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 three-part tuple -> reconstructed high-precision tensor. + + out_shape / axis must match the original convert_to_mx9 call, used to + reverse the transpose and padding. Returns shape = out_shape, dtype = out_dtype. + """ + assert out_dtype in _TORCH_TO_TL + assert block_size == 16, f"block_size only supports 16, got {block_size}" + + # Three-part consistency check: q/max_exp/prime must share the same n_blocks + # and shapes must match the storage spec; otherwise the kernel addressing / + # reshape would read misaligned data or raise unhelpful dimension errors. + n_blocks = q.shape[0] + n_prime_bytes = (block_size // 2) // 8 + assert q.shape == (n_blocks, block_size), \ + f"q shape should be ({n_blocks}, {block_size}), got {tuple(q.shape)}" + assert max_exp.shape == (n_blocks,), \ + f"max_exp shape should be ({n_blocks},), got {tuple(max_exp.shape)}" + assert prime.shape == (n_blocks, n_prime_bytes), \ + f"prime shape should be ({n_blocks}, {n_prime_bytes}), got {tuple(prime.shape)}" + + y_blocks = torch.empty((n_blocks, block_size), dtype=out_dtype, device=q.device) + + grid = (triton.cdiv(n_blocks, blocks_per_program),) + _convert_from_mx9_kernel[grid]( + q.contiguous(), max_exp.contiguous(), prime.contiguous(), y_blocks, n_blocks, + 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], + ) + + # Reverse: remove padding, reshape back to post-transpose shape, then + # transpose back to original axis. + # out_shape is the original (pre-transpose) shape; axis indicates which + # dimension was the quant axis. + transposed_shape = list(out_shape) + transposed_shape[axis], transposed_shape[-1] = transposed_shape[-1], transposed_shape[axis] + last = transposed_shape[-1] + rows = 1 + for s in transposed_shape[:-1]: + rows *= s + pad = (block_size - last % block_size) % block_size + padded_cols = last + pad + # out_shape/axis must be consistent with convert_to: if the inferred block + # count doesn't match the three-part tuple, intercept here with a readable + # error rather than letting the reshape below raise a cryptic dimension error. + assert rows * padded_cols == n_blocks * block_size, ( + f"out_shape/axis/block_size inconsistent with three-part tuple: " + f"inferred rows*padded_cols={rows * padded_cols}, " + f"but tuple has n_blocks*block_size={n_blocks * block_size}" + ) + y2d = y_blocks.reshape(rows, padded_cols) + y2d = y2d[:, :last] + return y2d.reshape(transposed_shape).transpose(axis, -1).contiguous() 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..ae9568f3 --- /dev/null +++ b/tests/unittest/mx9_mx6/test_mx9_quantization.py @@ -0,0 +1,567 @@ +# 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. + +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 + +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 convert_to_mx9, convert_from_mx9 +from alto.kernels.fp4.testing_utils import calc_snr + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + +# --------------------------------------------------------------------------- # +# 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) + + +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): + q, max_exp, prime = convert_to_mx9(x, block_size=block_size, axis=axis) + return convert_from_mx9(q, max_exp, prime, 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() + q, max_exp, prime = convert_to_mx9(x, block_size=16) + n_blocks = (4 * 64) // 16 + assert q.dtype == torch.int8 + assert max_exp.dtype == torch.uint8 + assert prime.dtype == torch.uint8 + assert q.shape == (n_blocks, 16) + assert max_exp.shape == (n_blocks,) + assert prime.shape == (n_blocks, 1) # block_size=16 -> 8 pairs -> 1 byte + + +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() + _, max_exp, _ = convert_to_mx9(x, block_size=16) + 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() + q, _, _ = convert_to_mx9(x, block_size=16) + 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() + q, max_exp, prime = convert_to_mx9(x, block_size=BLOCK_SIZE) + rq, re, rp, _ = _torch_intermediates(x, block_size=BLOCK_SIZE) + assert torch.equal(q, rq), f"q mismatch: {(q != rq).sum().item()} elements" + assert torch.equal(max_exp, re), f"max_exp mismatch" + assert torch.equal(prime, rp), f"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] + q, max_exp, prime = convert_to_mx9(x) + rq, re, rp, _ = _torch_intermediates(x) + assert torch.equal(q, rq) + assert torch.equal(max_exp, re) + assert torch.equal(prime, 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) + q, me, pr = convert_to_mx9(x, block_size=BLOCK_SIZE, + blocks_per_program=blocks_per_program) + out = convert_from_mx9(q, me, pr, 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. + + Host-side data_hp.contiguous() handles the re-layout; the kernel always + receives contiguous memory. This test verifies that path is not skipped. + """ + 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_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 + q, max_exp, prime = 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 prime[0, 0].item() == 0x0F, f"got {hex(prime[0, 0].item())}" + out = convert_from_mx9(q, max_exp, prime, 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 + + +def test_rejects_unknown_out_dtype(): + x = _rand((4, 16), torch.float32).cuda() + q, me, pr = convert_to_mx9(x) + with pytest.raises(AssertionError): + convert_from_mx9(q, me, pr, torch.int32, x.shape) From 95b11148635542d5754c80a6c049401186bdff0f Mon Sep 17 00:00:00 2001 From: jiarwang Date: Tue, 30 Jun 2026 11:02:56 +0000 Subject: [PATCH 2/7] feat: mx9: wire packed kernel into patcher and fix convert_from_mx9 - patcher: format=="mx9" now dispatches to convert_to/from_mx9 (packed Triton kernel) instead of mx9_fake_quantize (pure PyTorch). Validated: LLaMA-3.2-1B wikitext loss 2.1141 (== fake-quant 2.1140). - mx9_quantization: use tl.div_rn for IEEE round-nearest division (fixes potential 2-ULP fast-div rounding at .5 boundaries); add dtype asserts on q/max_exp/prime in convert_from_mx9. - kernels/mx/__init__: update docstring. --- alto/kernels/mx/__init__.py | 8 ++++---- alto/kernels/mx/mx9_quantization.py | 13 ++++++++++--- alto/models/patcher.py | 14 ++++++-------- 3 files changed, 20 insertions(+), 15 deletions(-) diff --git a/alto/kernels/mx/__init__.py b/alto/kernels/mx/__init__.py index f04649f0..506cad3e 100644 --- a/alto/kernels/mx/__init__.py +++ b/alto/kernels/mx/__init__.py @@ -1,9 +1,9 @@ # Copyright (c) 2026 Advanced Micro Devices, Inc. # # SPDX-License-Identifier: MIT -"""MX real (fused Triton) kernels. +"""MX packed-quantization Triton kernels (currently MX9 only). -The fake-quant reference / emulation lives in -``alto.modifiers.quantization.mx``; this package holds the fused GPU kernel -implementation of the same MX QDQ. +Fake-quant reference / emulation lives in ``alto.modifiers.quantization.mx``. +This package provides GPU pack/unpack for MX block formats; it does not yet +include MX quantized GEMM or fused quant+matmul (see ``mxfp4`` / ``mxfp8``). """ diff --git a/alto/kernels/mx/mx9_quantization.py b/alto/kernels/mx/mx9_quantization.py index d95f3e9b..a4f7f8b3 100644 --- a/alto/kernels/mx/mx9_quantization.py +++ b/alto/kernels/mx/mx9_quantization.py @@ -84,7 +84,8 @@ def _sanitize(x): (Faithfully replicates Quark/reference: exponent statistics use sanitized values to prevent Inf from polluting the block scale; the actual round still - operates on the original x, letting Inf be clamped and NaN pass through.) + operates on the original x, letting Inf be clamped. NaN is not representable + in the packed path (q is int8), so NaN inputs yield undefined q values.) """ x = tl.where(x != x, 0.0, x) x = tl.where(x == float("inf"), 0.0, x) @@ -206,8 +207,10 @@ def _pack_mx9( scale = tl.exp2(scale_exp.to(tl.float32)) # [BPP, BLOCK_SIZE], always normal (!=0) # Q part of QDQ operates on the *original* x (NaN will produce garbage, see below). + # div_rn (IEEE RN) matches PyTorch eager x/scale; plain / lowers to fast div + # (max 2 ULP error) which can flip round at .5 boundaries. xf = x.to(tl.float32) - q = _round_half_even(xf / scale) + q = _round_half_even(tl.div_rn(xf, scale)) # True 8-bit: symmetric clamp to +/-127. # - Normal data: demoted elements have |q|<128 mathematically; only rounding # edge cases can produce 128, which is clamped to 127. @@ -432,7 +435,11 @@ def convert_from_mx9( out_shape / axis must match the original convert_to_mx9 call, used to reverse the transpose and padding. Returns shape = out_shape, dtype = out_dtype. """ - assert out_dtype in _TORCH_TO_TL + assert out_dtype in _TORCH_TO_TL, \ + f"out_dtype must be one of {tuple(_TORCH_TO_TL)}, got {out_dtype}" + assert q.dtype == torch.int8, f"q dtype must be int8, got {q.dtype}" + assert max_exp.dtype == torch.uint8, f"max_exp dtype must be uint8, got {max_exp.dtype}" + assert prime.dtype == torch.uint8, f"prime dtype must be uint8, got {prime.dtype}" assert block_size == 16, f"block_size only supports 16, got {block_size}" # Three-part consistency check: q/max_exp/prime must share the same n_blocks diff --git a/alto/models/patcher.py b/alto/models/patcher.py index 8c751354..186c7363 100644 --- a/alto/models/patcher.py +++ b/alto/models/patcher.py @@ -58,16 +58,14 @@ 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, + 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), + q, max_exp, prime = convert_to_mx9(x) + return convert_from_mx9( + q, max_exp, prime, x.dtype, x.shape, ) if getattr(args, "format", None) == "mx6": from alto.modifiers.quantization.mx import ( From d3ce92b6197d84b0aa1d3da743c299b7104409a8 Mon Sep 17 00:00:00 2001 From: jiarwang Date: Tue, 14 Jul 2026 08:16:46 +0000 Subject: [PATCH 3/7] feat: mx6: add packed quantization kernel --- alto/kernels/mx/mx6_quantization.py | 405 ++++++++++ .../unittest/mx9_mx6/test_mx6_quantization.py | 763 ++++++++++++++++++ 2 files changed, 1168 insertions(+) create mode 100644 alto/kernels/mx/mx6_quantization.py create mode 100644 tests/unittest/mx9_mx6/test_mx6_quantization.py diff --git a/alto/kernels/mx/mx6_quantization.py b/alto/kernels/mx/mx6_quantization.py new file mode 100644 index 00000000..6aa6612c --- /dev/null +++ b/alto/kernels/mx/mx6_quantization.py @@ -0,0 +1,405 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Packed MX6 quantization Triton device kernels (called by grid kernels). + +MX6 shares the *exact same block algorithm* as MX9 (exponent extraction, prime +demotion, power-of-two scale); the only numeric difference is the element +integer width ``quant_bit`` (MX9 = 8, MX6 = 5), exactly mirroring Quark's +``fake_quantize_mx6_mx9`` (see ``alto/modifiers/quantization/mx.py``). Everything +in this module is therefore a straightforward specialization of +``mx9_quantization.py``; the two intentional divergences (to be implemented in +later stages) are: + + 1. Element width: ``QUANT_BIT = 5`` -> quantized integers are symmetrically + clamped to +/-15 (=2^(quant_bit-1)-1), instead of MX9's +/-127. + 2. **Python-export-compatible packing** (achieves a true 6 bits/element): MX9 + stores ``q`` as int8 (one byte per element), dense at its 8-bit width. At 5 + bits, one byte per element would waste 3 bits and defeat the "6" in MX6. + To match the Python export path, each 16-element block is laid out as: + - max_exp : 1 byte (E8M0, true exponent + 127) + - prime : 1 byte (same bitmap as ``idx2`` in the Python packer) + - sign : 2 bytes (16 sign bits, LSB = element 0) + - mantissa : 8 bytes (two ``q & 0xF`` nibbles per byte) + -> 12 bytes/block = 6 bits/element. The 4-bit mantissa nibbles are the low + bits of the signed quantized integer, not ``abs(q)``. + +The packed bytes are laid out per block as ``[max_exp, prime, sign bitmap ..., +mantissa nibbles ...]``. For the mantissa segment, low nibble = even element and +high nibble = odd element. + +Development is incremental. This file currently implements: + - Stage 1: stateless device helpers ``_sanitize`` / ``_floor_exp`` / + ``_round_half_even`` (identical to MX9). + - Stage 2: ``_calculate_mx6_exp`` (per-block shared exponent + per-element + shared_exp + per-pair prime bits; width-independent, identical to MX9). + - Stage 3: ``_pack_mx6`` / ``_unpack_mx6`` (Python-export QDQ pack + inverse). + - Stage 4: ``_convert_to_mx6_kernel`` / ``_convert_from_mx6_kernel`` grid entries + (stride-addressed load/store, ragged-tail masking, E8M0 bias). + +Still to come: the ``convert_to_mx6`` / ``convert_from_mx6`` host wrappers (Stage 5). + +See ``mx9_quantization.py`` for the full rationale behind the shared design +decisions (scale exponent clamp to -126 / FTZ independence, exponent extraction +from native dtype bits, E8M0 max_exp, stride addressing, ragged-tail masking). +""" + +import triton +import triton.language as tl + +BLOCK_SIZE = 16 +QUANT_BIT = 5 +PRIME_GROUP = 2 +BLOCKS_PER_PROG_DEFAULT = 64 + +# Python export-compatible packing (Stage 3+): per block, mantissa nibbles +# occupy block_size/2 bytes and the sign bitmap occupies block_size/8 bytes. +MANTISSA_BIT = QUANT_BIT - 1 # 4-bit low mantissa bits (q & 0xF) + + +@triton.jit +def _sanitize(x): + """Replace NaN / +/-Inf with 0; used only for the amax / exponent statistics + copy, does not touch the QDQ data path (identical to MX9).""" + x = tl.where(x != x, 0.0, x) + x = tl.where(x == float("inf"), 0.0, x) + x = tl.where(x == float("-inf"), 0.0, x) + return x + + +@triton.jit +def _floor_exp(x): + """floor(log2(|x|)): extract exponent field directly from native float bits + (fp32: (bits>>23)&0xFF - 127; bf16: (bits>>7)&0xFF - 127). Identical to MX9.""" + 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), pure tl. Identical to MX9.""" + 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_mx6_exp( + x, + BLOCKS_PER_PROG: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PRIME_GROUP: tl.constexpr, +): + """Compute per-block shared exponent + per-element shared_exp + per-pair + prime bits. + + The exponent / prime math is *width-independent* (does not involve + ``quant_bit``), so this is byte-for-byte the MX9 ``_calculate_mx9_exp``. + + Input ``x``: this program's tile ``[BLOCKS_PER_PROG, BLOCK_SIZE]`` in native + dtype (each row is one MX6 block). Exponent extraction must be done on native + dtype, so do NOT pre-cast to fp32 on host. + + Returns: + - ``shared_exp`` : [BPP, BLOCK_SIZE] int32, per-element shared exponent + (demoted pairs get max_exp - 1) + - ``max_exp`` : [BPP] int32, per-block max exponent + - ``pair`` : [BPP, N_PAIRS] int32(0/1), per-pair prime bit + (1 = that pair gets 1-exponent demotion) + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + + # Sanitized copy used only for exponent statistics (NaN/Inf -> 0). + clean = _sanitize(x) + + # Block shared exponent: per-row (block) amax -> floor extract exponent. + # Cast the reduce result back to native dtype so _floor_exp uses the same + # branch as the per-element t_exp (some backends promote tl.max to fp32). + amax = tl.max(tl.abs(clean), axis=1, keep_dims=True).to(clean.dtype) # [BPP, 1] native + max_exp = _floor_exp(amax) # [BPP, 1] int32 + + # Per-element exponent + demote flag (at least 1 octave below block max). + t_exp = _floor_exp(clean) # [BPP, BLOCK_SIZE] int32 + demote = (max_exp - t_exp) >= 1 # [BPP, BLOCK_SIZE] bool (broadcast) + + # Prime: adjacent PRIME_GROUP(=2) elements form a pair; both must be demoted + # for the pair to be demoted. + d = demote.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP) + pair_keep = tl.sum(d, axis=2, keep_dims=True) == PRIME_GROUP # [BPP, N_PAIRS, 1] bool + + # Broadcast back to per-element to get per-element shared_exp. + pair_b = tl.broadcast_to(pair_keep, (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) # [BPP, BLOCK_SIZE] int32 + + # Pair bits (not broadcast) are kept for prime bitmap packing (Stage 3). + pair = pair_keep.reshape(BLOCKS_PER_PROG, N_PAIRS).to(tl.int32) + return shared_exp, max_exp.reshape(BLOCKS_PER_PROG), pair + + +@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, +): + """Quantize + Python-export-compatible 5-bit pack. + + Returns three uint8 tiles used by the final single packed block layout: + - ``sign_bytes`` : [BPP, BLOCK_SIZE//8] sign bitmap (LSB = element 0) + - ``mantissa_bytes`` : [BPP, BLOCK_SIZE//2] 4-bit ``q & 0xF`` nibbles + (low nibble = even element, high nibble = odd element) + - ``prime`` : [BPP, N_PAIRS//8] prime bitmap (identical to MX9) + + Quantized integers are clamped to +/-15, matching the BFPQuantizer path. The + mantissa is the low 4 bits of the signed integer, matching the Python export + packer; it is intentionally not ``abs(q)``. + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + N_PRIME_BYTES: tl.constexpr = N_PAIRS // 8 + N_MANTISSA_BYTES: tl.constexpr = BLOCK_SIZE // 2 # 2 nibbles per byte + N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 # 8 sign bits per byte + Q_HI: tl.constexpr = (1 << (QUANT_BIT - 1)) - 1 # 15 for quant_bit=5 + + # Same scale as MX9 (clamp scale exponent to -126 for FTZ independence). + scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) + scale = tl.exp2(scale_exp.to(tl.float32)) # [BPP, BLOCK_SIZE], always normal + + xf = x.to(tl.float32) + q = _round_half_even(tl.div_rn(xf, scale)) + q = tl.minimum(tl.maximum(q, -(Q_HI * 1.0)), Q_HI * 1.0) # clamp +/-15 + + qi = q.to(tl.int32) + mantissa = qi & 0xF # [BPP, BLOCK_SIZE], low 4 bits + sign = (qi < 0).to(tl.int32) # [BPP, BLOCK_SIZE], 0/1 + + # Sign bitmap: 8 sign bits per byte via weights [1, 2, ..., 128]. + weights = tl.exp2(tl.arange(0, 8).to(tl.float32)).to(tl.int32) # [8] + sign_g = sign.reshape(BLOCKS_PER_PROG, N_SIGN_BYTES, 8) + sign_bytes = tl.sum(sign_g * weights[None, None, :], axis=2) # [BPP, N_SIGN_BYTES] + + # Mantissa: pack two 4-bit low-nibble values per byte via weights [1, 16]. + nib_w = tl.exp2((4 * tl.arange(0, 2)).to(tl.float32)).to(tl.int32) # [2] = [1, 16] + mantissa_g = mantissa.reshape(BLOCKS_PER_PROG, N_MANTISSA_BYTES, 2) + mantissa_bytes = tl.sum(mantissa_g * nib_w[None, None, :], axis=2) # [BPP, N_MANTISSA_BYTES] + + # Prime bitmap (identical to MX9): 8 pair bits -> 1 byte, LSB = pair0. + pb = pair.reshape(BLOCKS_PER_PROG, N_PRIME_BYTES, 8) # [BPP, NB, 8] + prime = tl.sum(pb * weights[None, None, :], axis=2) # [BPP, NB] + return sign_bytes.to(tl.uint8), mantissa_bytes.to(tl.uint8), prime.to(tl.uint8) + + +@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, +): + """Dequantize Python-export-compatible tiles + shared exponent. + + Inverse of ``_pack_mx6``: extract 4-bit ``q & 0xF`` mantissa nibbles and + sign bits, rebuild the signed integer, then apply + scale = 2^((max_exp - pair) - quant_bit + 2). ``max_exp`` arrives already + unbiased (int32, [BPP, 1]); ``prime`` is the per-pair bitmap. + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + N_PRIME_BYTES: tl.constexpr = N_PAIRS // 8 + N_MANTISSA_BYTES: tl.constexpr = BLOCK_SIZE // 2 + N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 + + # Unpack sign bitmap (same scheme as prime / Python uint16 little-endian view). + weights = tl.exp2(tl.arange(0, 8).to(tl.float32)).to(tl.int32) # [8] + sb = sign_bytes.to(tl.int32).reshape(BLOCKS_PER_PROG, N_SIGN_BYTES, 1) + sign = (sb // weights[None, None, :]) % 2 # [BPP, N_SIGN_BYTES, 8] + sign = sign.reshape(BLOCKS_PER_PROG, BLOCK_SIZE) + + # Unpack mantissa nibbles: (byte // [1,16]) % 16 -> [even, odd] element. + nib_w = tl.exp2((4 * tl.arange(0, 2)).to(tl.float32)).to(tl.int32) # [1, 16] + mb = mantissa_bytes.to(tl.int32).reshape(BLOCKS_PER_PROG, N_MANTISSA_BYTES, 1) + mantissa = (mb // nib_w[None, None, :]) % 16 # [BPP, N_MANTISSA_BYTES, 2] + mantissa = mantissa.reshape(BLOCKS_PER_PROG, BLOCK_SIZE) + + # Rebuild signed q from sign + low 4 bits. sign=1, mantissa=13 -> -3. + q = tl.where(sign == 1, mantissa - 16, mantissa) # [BPP, BLOCK_SIZE] int32 + + # Unpack prime bitmap -> pair -> per-element shared_exp (identical to MX9). + pbytes = prime.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PRIME_BYTES, 1) + bits = (pbytes // weights[None, None, :]) % 2 # [BPP, NB, 8] + pair = bits.reshape(BLOCKS_PER_PROG, N_PAIRS) # [BPP, N_PAIRS] + + 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 # [BPP, BLOCK_SIZE] + + scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) + scale = tl.exp2(scale_exp.to(tl.float32)) + y = q.to(tl.float32) * scale + return y.to(out_dtype) + + +# ============================================================================ +# Grid kernel (entry point): 1D grid, each program handles BLOCKS_PER_PROG +# 16-element blocks. High-precision side addressed by stride (no forced copy). +# Python export-compatible single packed tensor layout: +# packed : [n_blocks, N_PACKED_BYTES] per block uint8 bytes +# [0] max_exp E8M0 (true exp + 127) +# [1 : 1 + N_PRIME_BYTES] prime / idx2 bitmap +# next N_SIGN_BYTES sign bitmap +# last N_MANTISSA_BYTES 4-bit q low nibbles +# For BLOCK_SIZE=16 this is [max_exp, prime, sign0, sign1, mantissa0..7]. +# ============================================================================ + + +@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, +): + """Input ``x`` addressed by stride as 2D ``[rows, last]`` (``last`` = unpadded + quant-axis length), so the host does not force a ``.contiguous()`` copy. Each + block maps back to (row, block-within-row); columns beyond ``last`` (the + ragged tail) are masked to 0. Output is a freshly allocated contiguous packed + tensor using the same per-block byte order as the Python export path. + """ + 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 + N_PACKED_BYTES: tl.constexpr = 1 + N_PRIME_BYTES + N_SIGN_BYTES + N_MANTISSA_BYTES + PRIME_OFFSET: tl.constexpr = 1 + SIGN_OFFSET: tl.constexpr = PRIME_OFFSET + N_PRIME_BYTES + MANTISSA_OFFSET: tl.constexpr = SIGN_OFFSET + N_SIGN_BYTES + + pid = tl.program_id(0) + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) # global block indices [BPP] + blk_mask = blk < n_blocks + + # Map each block back to its source (row, block-within-row) coordinate. + row = blk // blocks_per_row # [BPP] + brow = blk % blocks_per_row # [BPP] + col = tl.arange(0, BLOCK_SIZE) # [BLOCK] + col_g = brow[:, None] * BLOCK_SIZE + col[None, :] # column in the row [BPP, BLOCK] + + in_range = col_g < last + load_mask = blk_mask[:, None] & in_range + in_offs = row[:, None] * stride_row + col_g * stride_col + + tl.static_assert( + x_ptr.type.element_ty == tl.float32 or + x_ptr.type.element_ty == tl.bfloat16, + "x must be fp32 / bf16", + ) + x = tl.load(x_ptr + in_offs, mask=load_mask, other=0.0) # native dtype tile + + shared_exp, max_exp, pair = _calculate_mx6_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, + ) + + # Byte 0: max_exp true exponent ([BPP] int32) -> E8M0 uint8: +127 bias. + e_store = (max_exp + 127).to(tl.uint8) + tl.store(packed_ptr + blk * N_PACKED_BYTES, e_store, mask=blk_mask) + + # Byte 1..: prime / idx2 bitmap. + 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]) + + # Then sign bitmap bytes, matching quantized_weight_sign.view(torch.uint8). + 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]) + + # Last segment: q low 4-bit mantissa nibbles. + 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_packed_blk, + stride_packed_col, + BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROG: tl.constexpr, + PRIME_GROUP: tl.constexpr, + QUANT_BIT: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + """Packed input addressed by stride; output ``y`` is freshly allocated + contiguous ``[n_blocks, BLOCK_SIZE]`` with flat offsets. + + The packed row is split as ``[max_exp, prime/idx2, sign, mantissa]`` before + dequantizing. + """ + 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 + + 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] + + # max_exp stored as E8M0 uint8 (+127 bias); subtract back to true exponent. + max_exp_u = tl.load( + packed_ptr + blk * stride_packed_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_packed_blk + (PRIME_OFFSET + pcol[None, :]) * stride_packed_col + prime = tl.load(packed_ptr + offs_p, mask=mask, other=0) # [BPP, NB] uint8 + + signcol = tl.arange(0, N_SIGN_BYTES) + sign_offs = blk[:, None] * stride_packed_blk + (SIGN_OFFSET + signcol[None, :]) * stride_packed_col + sign_bytes = tl.load(packed_ptr + sign_offs, mask=mask, other=0) # [BPP, N_SIGN_BYTES] + + mantcol = tl.arange(0, N_MANTISSA_BYTES) + mant_offs = blk[:, None] * stride_packed_blk + (MANTISSA_OFFSET + mantcol[None, :]) * stride_packed_col + mantissa_bytes = tl.load(packed_ptr + mant_offs, mask=mask, other=0) # [BPP, N_MANTISSA_BYTES] + + 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, :] # contiguous y offsets + tl.store(y_ptr + out_offs, y, mask=mask) + 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..b74e4da1 --- /dev/null +++ b/tests/unittest/mx9_mx6/test_mx6_quantization.py @@ -0,0 +1,763 @@ +# 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: ``_sanitize`` / ``_floor_exp`` / ``_round_half_even`` + - 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 with +transpose / padding / shape restore -- is added later, reusing the MX9 host +test suite.) +""" + +import pytest +import torch +import triton +import triton.language as tl + +from alto.kernels.mx.mx6_quantization import ( + BLOCK_SIZE, + PRIME_GROUP, + QUANT_BIT, + BLOCKS_PER_PROG_DEFAULT, + _sanitize, + _floor_exp, + _round_half_even, + _calculate_mx6_exp, + _pack_mx6, + _unpack_mx6, + _convert_to_mx6_kernel, + _convert_from_mx6_kernel, +) +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_sanitize_kernel(x_ptr, o_ptr, N: tl.constexpr): + i = tl.arange(0, N) + tl.store(o_ptr + i, _sanitize(tl.load(x_ptr + i))) + + +@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_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: _sanitize +# --------------------------------------------------------------------------- # +def test_sanitize_replaces_nonfinite_with_zero(): + x = torch.tensor( + [1.0, float("nan"), float("inf"), float("-inf"), -2.0, 0.0, 3.5, float("nan")], + dtype=torch.float32, device="cuda", + ) + out = torch.empty_like(x) + _probe_sanitize_kernel[(1,)](x, out, N=x.numel()) + ref = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# 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 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, + ) + 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, + ) + 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], + ) + 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 + + +# --------------------------------------------------------------------------- # +# Layer: differential test against production mx.py (fake-quant reference) +# The tests above compare against a hand-written torch reimplementation. This +# layer cross-checks against the *production* Quark-aligned code in +# alto/modifiers/quantization/mx.py, in two ways: +# (1) a SECOND independent reference that reuses mx.py's own primitives +# (_reshape_to_blocks / _t_exponent / SHARED_PRIME_BIT_GROUP), only +# swapping the value clamp to +/-15 -- it must agree with our hand ref +# AND with the kernel, so a shared bug in one reference cannot hide; +# (2) a divergence check vs the raw mx6_fake_quantize (clamp 31 for demoted +# pairs): the two may only differ at the rare "demoted AND rounds to +# +/-16" edge, so the mismatch fraction must be tiny. +# --------------------------------------------------------------------------- # +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) + + +@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 + n_blocks = rows * cols // BLOCK_SIZE + 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"])) From db210c08b3f557438c3dc93e82f31702ba245ae3 Mon Sep 17 00:00:00 2001 From: jiarwang Date: Tue, 14 Jul 2026 08:16:46 +0000 Subject: [PATCH 4/7] perf: mx9: consume strided quantization inputs --- alto/kernels/mx/mx9_quantization.py | 102 +++++++++++++----- .../unittest/mx9_mx6/test_mx9_quantization.py | 5 +- 2 files changed, 78 insertions(+), 29 deletions(-) diff --git a/alto/kernels/mx/mx9_quantization.py b/alto/kernels/mx/mx9_quantization.py index a4f7f8b3..51084dd2 100644 --- a/alto/kernels/mx/mx9_quantization.py +++ b/alto/kernels/mx/mx9_quantization.py @@ -28,7 +28,9 @@ Constraints: - Blocks are formed along the last dimension only (host transposes the quant - axis to dim -1 then calls .contiguous()). + axis to dim -1). The tensor is passed to the kernel **by stride** (no forced + ``.contiguous()`` copy); the kernel addresses elements via row/column strides + and masks the ragged tail block instead of host-side ``F.pad``. - ``block_size`` is fixed at 16: each block has 8 pairs which pack into exactly 1 byte of prime bitmap. @@ -61,7 +63,6 @@ """ import torch -import torch.nn.functional as F import triton import triton.language as tl @@ -274,8 +275,8 @@ def _unpack_mx9( # ============================================================================ # Grid kernel (entry point): 1D grid, each program handles BLOCKS_PER_PROG -# complete 16-element blocks. Inputs/outputs are addressed as flattened -# [n_blocks, BLOCK_SIZE] views; the three-part tuple layout: +# 16-element blocks. The high-precision side is addressed by stride (no forced +# copy); the packed three-part tuple layout: # q : [n_blocks, BLOCK_SIZE] (per-element, int8, clamp +/-127) # max_exp : [n_blocks] (per-block, uint8 E8M0, true exp + 127) # prime : [n_blocks, N_PRIME_BYTES] (per-block, N_PRIME_BYTES bytes, uint8 bitmap) @@ -289,26 +290,49 @@ def _convert_to_mx9_kernel( e_ptr, p_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, ): + """Input ``x`` is addressed **by stride** as a 2D ``[rows, last]`` tensor + (``last`` = unpadded quant-axis length), so the host does not force a + ``.contiguous()`` copy. Each of the ``BLOCKS_PER_PROG`` blocks maps back to + ``(row, block-within-row)``; columns beyond ``last`` (the ragged tail block) + are masked to 0, replacing host-side ``F.pad``. + + Outputs (q / max_exp / prime) are freshly allocated and contiguous, so they + are addressed with the classic flat ``[n_blocks, ...]`` offsets. + """ N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 pid = tl.program_id(0) - blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) # block indices for this program [BPP] - col = tl.arange(0, BLOCK_SIZE) - offs = blk[:, None] * BLOCK_SIZE + col[None, :] # input/output q offsets [BPP, BLOCK] - mask = blk[:, None] < n_blocks # tail block-level mask + blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) # global block indices [BPP] blk_mask = blk < n_blocks + # Map each block back to its source (row, block-within-row) coordinate. + row = blk // blocks_per_row # [BPP] + brow = blk % blocks_per_row # [BPP] + col = tl.arange(0, BLOCK_SIZE) # [BLOCK] + col_g = brow[:, None] * BLOCK_SIZE + col[None, :] # column in the (unpadded) row [BPP, BLOCK] + + # Strided gather: address = row*stride_row + col*stride_col. Columns past the + # unpadded width (last) are the padding tail -> masked to 0 (same statistics + # as the previous F.pad-with-zeros behaviour). + in_range = col_g < last + load_mask = blk_mask[:, None] & in_range + in_offs = row[:, None] * stride_row + col_g * stride_col + tl.static_assert( x_ptr.type.element_ty == tl.float32 or x_ptr.type.element_ty == tl.bfloat16, "x must be fp32 / bf16", ) - x = tl.load(x_ptr + offs, mask=mask, other=0.0) # native dtype tile + x = tl.load(x_ptr + in_offs, mask=load_mask, other=0.0) # native dtype tile shared_exp, max_exp, pair = _calculate_mx9_exp( x, BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP @@ -318,7 +342,8 @@ def _convert_to_mx9_kernel( BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT, ) - tl.store(q_ptr + offs, q_int, mask=mask) + out_offs = blk[:, None] * BLOCK_SIZE + col[None, :] # contiguous q offsets [BPP, BLOCK] + tl.store(q_ptr + out_offs, q_int, mask=blk_mask[:, None]) # max_exp is the true exponent ([BPP] int32); store as E8M0 uint8: +127 bias. e_store = (max_exp + 127).to(tl.uint8) tl.store(e_ptr + blk, e_store, mask=blk_mask) @@ -335,35 +360,45 @@ def _convert_from_mx9_kernel( p_ptr, y_ptr, n_blocks, + stride_q_blk, + stride_q_col, + stride_e, + stride_p_blk, + stride_p_col, BLOCK_SIZE: tl.constexpr, BLOCKS_PER_PROG: tl.constexpr, PRIME_GROUP: tl.constexpr, QUANT_BIT: tl.constexpr, OUT_DTYPE: tl.constexpr, ): + """Packed inputs (q / max_exp / prime) are addressed **by stride** so the host + does not force a ``.contiguous()`` copy. The output ``y`` is freshly allocated + contiguous ``[n_blocks, BLOCK_SIZE]`` and uses classic flat offsets. + """ N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 pid = tl.program_id(0) blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) col = tl.arange(0, BLOCK_SIZE) - offs = blk[:, None] * BLOCK_SIZE + col[None, :] - mask = blk[:, None] < n_blocks blk_mask = blk < n_blocks + mask = blk_mask[:, None] - q = tl.load(q_ptr + offs, mask=mask, other=0) # int8 + q_offs = blk[:, None] * stride_q_blk + col[None, :] * stride_q_col + q = tl.load(q_ptr + q_offs, mask=mask, other=0) # int8 # max_exp is stored as E8M0 uint8 (with +127 bias); subtract back to true exponent. - max_exp_u = tl.load(e_ptr + blk, mask=blk_mask, other=0) # [BPP] uint8 + max_exp_u = tl.load(e_ptr + blk * stride_e, mask=blk_mask, other=0) # [BPP] uint8 max_exp = max_exp_u.to(tl.int32) - 127 pcol = tl.arange(0, N_PRIME_BYTES) - offs_p = blk[:, None] * N_PRIME_BYTES + pcol[None, :] - prime = tl.load(p_ptr + offs_p, mask=blk_mask[:, None], other=0) # [BPP, NB] uint8 + offs_p = blk[:, None] * stride_p_blk + pcol[None, :] * stride_p_col + prime = tl.load(p_ptr + offs_p, mask=mask, other=0) # [BPP, NB] uint8 y = _unpack_mx9( q, max_exp[:, None], prime, OUT_DTYPE, BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT, ) - tl.store(y_ptr + offs, y, mask=mask) + out_offs = blk[:, None] * BLOCK_SIZE + col[None, :] # contiguous y offsets + tl.store(y_ptr + out_offs, y, mask=mask) # ============================================================================ @@ -394,24 +429,26 @@ def convert_to_mx9( assert block_size == 16, f"block_size only supports 16, got {block_size}" # Transpose quant axis to last dim, matching mxfp4 host processing. + # No forced .contiguous(): reshape returns a (possibly non-contiguous) view + # when possible, and the kernel gathers by stride. Padding of the ragged tail + # block is handled in-kernel via column masking instead of F.pad. data_hp = data_hp.transpose(axis, -1) last = data_hp.shape[-1] - x2d = data_hp.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) - x_blocks = x2d.reshape(n_blocks, block_size).contiguous() + x2d = data_hp.reshape(-1, last) + rows = x2d.shape[0] + blocks_per_row = triton.cdiv(last, block_size) + n_blocks = rows * blocks_per_row n_prime_bytes = (block_size // 2) // 8 q = torch.empty((n_blocks, block_size), dtype=torch.int8, device=data_hp.device) max_exp = torch.empty((n_blocks,), dtype=torch.uint8, device=data_hp.device) prime = torch.empty((n_blocks, n_prime_bytes), dtype=torch.uint8, device=data_hp.device) + stride_row, stride_col = x2d.stride() grid = (triton.cdiv(n_blocks, blocks_per_program),) _convert_to_mx9_kernel[grid]( - x_blocks, q, max_exp, prime, n_blocks, + x2d, q, max_exp, prime, n_blocks, last, blocks_per_row, + stride_row, stride_col, BLOCK_SIZE=block_size, BLOCKS_PER_PROG=blocks_per_program, PRIME_GROUP=PRIME_GROUP, @@ -456,9 +493,17 @@ def convert_from_mx9( y_blocks = torch.empty((n_blocks, block_size), dtype=out_dtype, device=q.device) + # No forced .contiguous(): pass the packed inputs by stride so already-laid-out + # (or non-contiguous view) tensors are consumed without an extra copy. + stride_q_blk, stride_q_col = q.stride() + stride_e = max_exp.stride(0) + stride_p_blk, stride_p_col = prime.stride() grid = (triton.cdiv(n_blocks, blocks_per_program),) _convert_from_mx9_kernel[grid]( - q.contiguous(), max_exp.contiguous(), prime.contiguous(), y_blocks, n_blocks, + q, max_exp, prime, y_blocks, n_blocks, + stride_q_blk, stride_q_col, + stride_e, + stride_p_blk, stride_p_col, BLOCK_SIZE=block_size, BLOCKS_PER_PROG=blocks_per_program, PRIME_GROUP=PRIME_GROUP, @@ -486,6 +531,9 @@ def convert_from_mx9( f"inferred rows*padded_cols={rows * padded_cols}, " f"but tuple has n_blocks*block_size={n_blocks * block_size}" ) + # Aligned with mxfp4 (convert_from_mxfp4): return the transposed stride view + # without a forced .contiguous() copy. Only the ragged-tail slice reshape may + # trigger an internal copy; the common (no-pad / axis=-1) path stays a view. y2d = y_blocks.reshape(rows, padded_cols) y2d = y2d[:, :last] - return y2d.reshape(transposed_shape).transpose(axis, -1).contiguous() + return y2d.reshape(transposed_shape).transpose(axis, -1) diff --git a/tests/unittest/mx9_mx6/test_mx9_quantization.py b/tests/unittest/mx9_mx6/test_mx9_quantization.py index ae9568f3..ae3479e0 100644 --- a/tests/unittest/mx9_mx6/test_mx9_quantization.py +++ b/tests/unittest/mx9_mx6/test_mx9_quantization.py @@ -362,8 +362,9 @@ def test_non_contiguous_input(dtype): """Non-contiguous memory input (transposed stride) produces the same result as the equivalent contiguous tensor. - Host-side data_hp.contiguous() handles the re-layout; the kernel always - receives contiguous memory. This test verifies that path is not skipped. + 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 From c5ce9157b85c49c289cad30d053b7d7c650420ff Mon Sep 17 00:00:00 2001 From: jiarwang Date: Fri, 17 Jul 2026 08:54:08 +0000 Subject: [PATCH 5/7] feat: mx6: add convert_to_mx6 / convert_from_mx6 host wrappers --- alto/kernels/mx/mx6_quantization.py | 144 +++++++++++++- .../unittest/mx9_mx6/test_mx6_quantization.py | 183 ++++++++++++++++-- 2 files changed, 307 insertions(+), 20 deletions(-) diff --git a/alto/kernels/mx/mx6_quantization.py b/alto/kernels/mx/mx6_quantization.py index 6aa6612c..97c0a15f 100644 --- a/alto/kernels/mx/mx6_quantization.py +++ b/alto/kernels/mx/mx6_quantization.py @@ -36,14 +36,15 @@ - Stage 3: ``_pack_mx6`` / ``_unpack_mx6`` (Python-export QDQ pack + inverse). - Stage 4: ``_convert_to_mx6_kernel`` / ``_convert_from_mx6_kernel`` grid entries (stride-addressed load/store, ragged-tail masking, E8M0 bias). - -Still to come: the ``convert_to_mx6`` / ``convert_from_mx6`` host wrappers (Stage 5). + - Stage 5: ``convert_to_mx6`` / ``convert_from_mx6`` host wrappers (axis + transpose, padding-aware shape restore -- mirrors ``mx9_quantization.py``). See ``mx9_quantization.py`` for the full rationale behind the shared design decisions (scale exponent clamp to -126 / FTZ independence, exponent extraction from native dtype bits, E8M0 max_exp, stride addressing, ragged-tail masking). """ +import torch import triton import triton.language as tl @@ -52,9 +53,11 @@ PRIME_GROUP = 2 BLOCKS_PER_PROG_DEFAULT = 64 -# Python export-compatible packing (Stage 3+): per block, mantissa nibbles -# occupy block_size/2 bytes and the sign bitmap occupies block_size/8 bytes. -MANTISSA_BIT = QUANT_BIT - 1 # 4-bit low mantissa bits (q & 0xF) +# torch dtype -> triton dtype (used by unpack output cast, Stage 5 host wrapper). +_TORCH_TO_TL = { + torch.float32: tl.float32, + torch.bfloat16: tl.bfloat16, +} @triton.jit @@ -403,3 +406,134 @@ def _convert_from_mx6_kernel( out_offs = blk[:, None] * BLOCK_SIZE + col[None, :] # contiguous y offsets tl.store(y_ptr + out_offs, y, mask=mask) + +# ============================================================================ +# Host wrappers (bare launch, mirroring mx9_quantization.py). Unlike MX9's +# three-part tuple (q / max_exp / prime), MX6 packs everything into a SINGLE +# uint8 tensor per the Python-export-compatible byte layout described above, +# so these wrappers take/return one tensor instead of three. +# ============================================================================ + + +def _mx6_packed_bytes(block_size: int) -> int: + """Bytes per block for the packed layout [max_exp | prime | sign | mantissa].""" + n_prime_bytes = (block_size // 2) // 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 tensor (Python-export-compatible byte + layout: ``[max_exp(1B) | prime(1B) | sign(2B) | mantissa(8B)]`` per block). + + Blocks are formed along the ``axis`` dimension (axis is transposed to the last + dim internally). The returned tensor is a flattened view: + packed : [n_blocks, n_packed_bytes] uint8 + Shape reconstruction info (original shape / axis / padding) must be passed + back by the caller in convert_from_mx6. + """ + assert data_hp.dtype in (torch.float32, torch.bfloat16), \ + f"mx6_quantization only supports fp32 / bf16, got {data_hp.dtype}" + assert block_size == 16, f"block_size only supports 16, got {block_size}" + + # Transpose quant axis to last dim, matching mx9 host processing. No forced + # .contiguous(): reshape returns a (possibly non-contiguous) view when + # possible, and the kernel gathers by stride. The ragged tail block is + # handled in-kernel via column masking instead of host-side F.pad. + data_hp = data_hp.transpose(axis, -1) + last = data_hp.shape[-1] + x2d = data_hp.reshape(-1, last) + rows = x2d.shape[0] + blocks_per_row = triton.cdiv(last, block_size) + n_blocks = rows * blocks_per_row + + n_packed_bytes = _mx6_packed_bytes(block_size) + 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),) + _convert_to_mx6_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, + ) + return packed + + +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 / axis must match the original convert_to_mx6 call, used to + reverse the transpose and padding. Returns shape = out_shape, dtype = out_dtype. + """ + 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 block_size == 16, f"block_size only supports 16, got {block_size}" + + # Layout consistency check: packed must carry a whole number of blocks with + # the expected per-block byte count; otherwise the kernel addressing would + # read misaligned bytes or raise an unhelpful dimension error downstream. + n_blocks = packed.shape[0] + n_packed_bytes = _mx6_packed_bytes(block_size) + assert packed.shape == (n_blocks, n_packed_bytes), \ + f"packed shape should be ({n_blocks}, {n_packed_bytes}), got {tuple(packed.shape)}" + + y_blocks = torch.empty((n_blocks, block_size), dtype=out_dtype, device=packed.device) + + # No forced .contiguous(): pass the packed input by stride so an already + # laid-out (or non-contiguous view) tensor is consumed without an extra copy. + stride_blk, stride_col = packed.stride() + grid = (triton.cdiv(n_blocks, blocks_per_program),) + _convert_from_mx6_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], + ) + + # Reverse: remove padding, reshape back to post-transpose shape, then + # transpose back to original axis. out_shape is the original (pre-transpose) + # shape; axis indicates which dimension was the quant axis. + transposed_shape = list(out_shape) + transposed_shape[axis], transposed_shape[-1] = transposed_shape[-1], transposed_shape[axis] + last = transposed_shape[-1] + rows = 1 + for s in transposed_shape[:-1]: + rows *= s + pad = (block_size - last % block_size) % block_size + padded_cols = last + pad + # out_shape/axis must be consistent with convert_to: if the inferred block + # count doesn't match the packed tensor, intercept here with a readable + # error rather than letting the reshape below raise a cryptic dimension error. + 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 tuple has n_blocks*block_size={n_blocks * block_size}" + ) + # Aligned with mx9 (convert_from_mx9): return the transposed stride view + # without a forced .contiguous() copy. Only the ragged-tail slice reshape may + # trigger an internal copy; the common (no-pad / axis=-1) path stays a view. + y2d = y_blocks.reshape(rows, padded_cols) + y2d = y2d[:, :last] + return y2d.reshape(transposed_shape).transpose(axis, -1) + diff --git a/tests/unittest/mx9_mx6/test_mx6_quantization.py b/tests/unittest/mx9_mx6/test_mx6_quantization.py index b74e4da1..fa0348c1 100644 --- a/tests/unittest/mx9_mx6/test_mx6_quantization.py +++ b/tests/unittest/mx9_mx6/test_mx6_quantization.py @@ -23,10 +23,9 @@ - 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 with -transpose / padding / shape restore -- is added later, reusing the MX9 host -test suite.) + - 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 @@ -47,6 +46,8 @@ _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 @@ -673,17 +674,10 @@ def test_kernel_zeros_and_block_independence(): # --------------------------------------------------------------------------- # -# Layer: differential test against production mx.py (fake-quant reference) -# The tests above compare against a hand-written torch reimplementation. This -# layer cross-checks against the *production* Quark-aligned code in -# alto/modifiers/quantization/mx.py, in two ways: -# (1) a SECOND independent reference that reuses mx.py's own primitives -# (_reshape_to_blocks / _t_exponent / SHARED_PRIME_BIT_GROUP), only -# swapping the value clamp to +/-15 -- it must agree with our hand ref -# AND with the kernel, so a shared bug in one reference cannot hide; -# (2) a divergence check vs the raw mx6_fake_quantize (clamp 31 for demoted -# pairs): the two may only differ at the rare "demoted AND rounds to -# +/-16" edge, so the mismatch fraction must be tiny. +# 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. @@ -721,6 +715,165 @@ def _mx6_clamp15_ref_via_mxpy(x, block_size=BLOCK_SIZE, quant_bit=QUANT_BIT): 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_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 + + +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): From abf9559e473c7d1697f3407d468e4186f1a1acd7 Mon Sep 17 00:00:00 2001 From: jiarwang Date: Fri, 17 Jul 2026 08:54:08 +0000 Subject: [PATCH 6/7] refactor: mx: pack mx9 into a single uint8 tensor and dispatch mx6/mx9 to the packed kernels --- alto/kernels/mx/mx9_quantization.py | 184 ++++++++++-------- alto/models/patcher.py | 30 ++- examples/llama3.2_1b_mx6.sh | 65 +++++++ examples/llama3.2_1b_mx9.sh | 5 +- tests/unittest/mx9_mx6/test_mx6_dispatch.py | 31 ++- .../mx9_mx6/test_mx6_wa_integration.py | 46 +++-- tests/unittest/mx9_mx6/test_mx9_dispatch.py | 31 ++- .../unittest/mx9_mx6/test_mx9_quantization.py | 72 ++++--- .../mx9_mx6/test_mx9_wa_integration.py | 47 +++-- 9 files changed, 350 insertions(+), 161 deletions(-) create mode 100755 examples/llama3.2_1b_mx6.sh diff --git a/alto/kernels/mx/mx9_quantization.py b/alto/kernels/mx/mx9_quantization.py index 51084dd2..a623e811 100644 --- a/alto/kernels/mx/mx9_quantization.py +++ b/alto/kernels/mx/mx9_quantization.py @@ -4,12 +4,18 @@ """Packed MX9 quantization Triton device kernels (called by grid kernels). Unlike the fake-quant path (``alto/kernels/mx/quantize_triton.py`` which outputs -dequantized bf16), this module performs *real packed* quantization: each MX9 block -is compressed into a three-part tuple stored separately: - - ``q`` : per-element quantized integer (int8, symmetric clamp +/-127) - - ``max_exp`` : per-block shared exponent, uint8 E8M0 (floor(log2(amax)) + 127 bias) - - ``prime`` : per-block prime bitmap (1 bit per pair of 2 elements, indicates - whether that pair gets a 1-exponent demotion) +dequantized bf16), this module performs *real packed* quantization: each MX9 +block is compressed into a **single packed uint8 tensor**, with each block's +bytes laid out as ``[max_exp(1B) | prime(1B) | q(16B)]`` (mirroring MX6's +prefix convention in ``mx6_quantization.py``): + - ``max_exp`` : 1 byte, per-block shared exponent, uint8 E8M0 + (floor(log2(amax)) + 127 bias) + - ``prime`` : 1 byte, per-block prime bitmap (1 bit per pair of 2 elements, + indicates whether that pair gets a 1-exponent demotion) + - ``q`` : 16 bytes, per-element quantized integer (symmetric clamp + +/-127), stored as the bit-identical uint8 reinterpretation of + int8 -- reconstruct via ``packed[..., 2:].view(torch.int8)``. +-> 18 bytes/block = 9 bits/element (8b value + 0.5b amortized max_exp + 0.5b prime). The exponent extraction / scale / prime math is aligned with ``alto/modifiers/quantization/mx.py`` (exponent extracted from native dtype bits @@ -54,6 +60,12 @@ 4. max_exp is stored as uint8 E8M0 (true exponent + 127), not int32: true exponents for any finite input fall in [-127, 127], +127 -> [0, 254] which fits uint8 without overflow; this achieves 1 byte/block and 9 bits/element. + 5. All three components live in one uint8 tensor (not three separate tensors): + ``q`` is computed as int8 in registers (see ``_pack_mx9`` / ``_unpack_mx9``, + unchanged) but stored/loaded via the packed uint8 tensor using + ``.to(tl.uint8/tl.int8, bitcast=True)`` -- a pure bit reinterpretation, not a + value cast (which would incorrectly clamp/wrap negative values). This mirrors + ``mx6_quantization.py``'s single-tensor, Python-export-compatible layout. Known limitations: - NaN: packed integers cannot represent NaN; blocks containing NaN will have @@ -276,19 +288,19 @@ def _unpack_mx9( # ============================================================================ # Grid kernel (entry point): 1D grid, each program handles BLOCKS_PER_PROG # 16-element blocks. The high-precision side is addressed by stride (no forced -# copy); the packed three-part tuple layout: -# q : [n_blocks, BLOCK_SIZE] (per-element, int8, clamp +/-127) -# max_exp : [n_blocks] (per-block, uint8 E8M0, true exp + 127) -# prime : [n_blocks, N_PRIME_BYTES] (per-block, N_PRIME_BYTES bytes, uint8 bitmap) +# copy). Single packed tensor layout (mirrors mx6_quantization.py): +# packed : [n_blocks, N_PACKED_BYTES] per block uint8 bytes +# [0] max_exp E8M0 (true exp + 127) +# [1 : 1 + N_PRIME_BYTES] prime bitmap +# [1 + N_PRIME_BYTES : ] q values (int8, bit-reinterpreted uint8) +# For BLOCK_SIZE=16 this is [max_exp, prime, q0..q15] = 18 bytes/block. # ============================================================================ @triton.jit def _convert_to_mx9_kernel( x_ptr, - q_ptr, - e_ptr, - p_ptr, + packed_ptr, n_blocks, last, blocks_per_row, @@ -305,10 +317,14 @@ def _convert_to_mx9_kernel( ``(row, block-within-row)``; columns beyond ``last`` (the ragged tail block) are masked to 0, replacing host-side ``F.pad``. - Outputs (q / max_exp / prime) are freshly allocated and contiguous, so they - are addressed with the classic flat ``[n_blocks, ...]`` offsets. + Output is a freshly allocated contiguous packed tensor; ``q`` (int8 in + registers) is bit-reinterpreted to uint8 before storing (a bitcast, not a + value cast, so negative values are preserved exactly). """ N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + N_PACKED_BYTES: tl.constexpr = 1 + N_PRIME_BYTES + BLOCK_SIZE + PRIME_OFFSET: tl.constexpr = 1 + Q_OFFSET: tl.constexpr = PRIME_OFFSET + N_PRIME_BYTES pid = tl.program_id(0) blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) # global block indices [BPP] @@ -342,40 +358,44 @@ def _convert_to_mx9_kernel( BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP, QUANT_BIT, ) - out_offs = blk[:, None] * BLOCK_SIZE + col[None, :] # contiguous q offsets [BPP, BLOCK] - tl.store(q_ptr + out_offs, q_int, mask=blk_mask[:, None]) - # max_exp is the true exponent ([BPP] int32); store as E8M0 uint8: +127 bias. + # Byte 0: max_exp true exponent ([BPP] int32) -> E8M0 uint8: +127 bias. e_store = (max_exp + 127).to(tl.uint8) - tl.store(e_ptr + blk, e_store, mask=blk_mask) + tl.store(packed_ptr + blk * N_PACKED_BYTES, e_store, mask=blk_mask) + # Byte 1..: prime bitmap. pcol = tl.arange(0, N_PRIME_BYTES) - offs_p = blk[:, None] * N_PRIME_BYTES + pcol[None, :] - tl.store(p_ptr + offs_p, prime, mask=blk_mask[:, None]) + offs_p = blk[:, None] * N_PACKED_BYTES + (PRIME_OFFSET + pcol[None, :]) + tl.store(packed_ptr + offs_p, prime, mask=blk_mask[:, None]) + + # Last BLOCK_SIZE bytes: q values, bit-reinterpreted int8 -> uint8. + 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( - q_ptr, - e_ptr, - p_ptr, + packed_ptr, y_ptr, n_blocks, - stride_q_blk, - stride_q_col, - stride_e, - stride_p_blk, - stride_p_col, + 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, ): - """Packed inputs (q / max_exp / prime) are addressed **by stride** so the host - does not force a ``.contiguous()`` copy. The output ``y`` is freshly allocated - contiguous ``[n_blocks, BLOCK_SIZE]`` and uses classic flat offsets. + """Packed input is addressed **by stride** so the host does not force a + ``.contiguous()`` copy. The output ``y`` is freshly allocated contiguous + ``[n_blocks, BLOCK_SIZE]`` and uses classic flat offsets. + + The packed row is split as ``[max_exp, prime, q]``; ``q`` bytes are loaded + as uint8 then bit-reinterpreted back to int8 (a bitcast, not a value cast) + before dequantizing. """ N_PRIME_BYTES: tl.constexpr = (BLOCK_SIZE // PRIME_GROUP) // 8 + PRIME_OFFSET: tl.constexpr = 1 + Q_OFFSET: tl.constexpr = PRIME_OFFSET + N_PRIME_BYTES pid = tl.program_id(0) blk = pid * BLOCKS_PER_PROG + tl.arange(0, BLOCKS_PER_PROG) @@ -383,15 +403,17 @@ def _convert_from_mx9_kernel( blk_mask = blk < n_blocks mask = blk_mask[:, None] - q_offs = blk[:, None] * stride_q_blk + col[None, :] * stride_q_col - q = tl.load(q_ptr + q_offs, mask=mask, other=0) # int8 - # max_exp is stored as E8M0 uint8 (with +127 bias); subtract back to true exponent. - max_exp_u = tl.load(e_ptr + blk * stride_e, mask=blk_mask, other=0) # [BPP] uint8 + # max_exp stored as E8M0 uint8 (+127 bias); subtract back to true exponent. + max_exp_u = tl.load(packed_ptr + blk * stride_blk, mask=blk_mask, other=0) # [BPP] uint8 max_exp = max_exp_u.to(tl.int32) - 127 pcol = tl.arange(0, N_PRIME_BYTES) - offs_p = blk[:, None] * stride_p_blk + pcol[None, :] * stride_p_col - prime = tl.load(p_ptr + offs_p, mask=mask, other=0) # [BPP, NB] uint8 + offs_p = blk[:, None] * stride_blk + (PRIME_OFFSET + pcol[None, :]) * stride_col + prime = tl.load(packed_ptr + offs_p, mask=mask, other=0) # [BPP, NB] uint8 + + q_offs = blk[:, None] * stride_blk + (Q_OFFSET + col[None, :]) * stride_col + q_u8 = tl.load(packed_ptr + q_offs, mask=mask, other=0) # uint8 + q = q_u8.to(tl.int8, bitcast=True) # bit-reinterpret -> int8 y = _unpack_mx9( q, max_exp[:, None], prime, OUT_DTYPE, @@ -403,26 +425,32 @@ def _convert_from_mx9_kernel( # ============================================================================ # Host wrappers (not yet registered as @triton_op / register_fake; using bare -# launch for ease of round-trip verification). +# launch for ease of round-trip verification). Single packed uint8 tensor +# in/out (mirrors mx6_quantization.py), not a three-tensor tuple. # ============================================================================ +def _mx9_packed_bytes(block_size: int) -> int: + """Bytes per block for the packed layout [max_exp(1B) | prime | q(block_size)B].""" + n_prime_bytes = (block_size // 2) // 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, -): - """High-precision tensor -> packed MX9 three-part tuple (q:int8, max_exp:uint8 - E8M0, prime:uint8). +) -> torch.Tensor: + """High-precision tensor -> packed MX9 tensor (single uint8 tensor, byte + layout ``[max_exp(1B) | prime(1B) | q(block_size bytes)]`` per block). Blocks are formed along the ``axis`` dimension (axis is transposed to the last - dim internally). The returned three-part tuple is a flattened view: - q : [n_blocks, block_size] int8 (clamp +/-127) - max_exp : [n_blocks] uint8 (E8M0, true exponent + 127) - prime : [n_blocks, n_prime_bytes] uint8 + dim internally). The returned tensor is a flattened view: + packed : [n_blocks, n_packed_bytes] uint8 Shape reconstruction info (original shape / axis / padding) must be passed - back by the caller in convert_from_mx9. + back by the caller in convert_from_mx9. ``q`` is recoverable bit-exact via + ``packed[..., 2:].view(torch.int8)`` (int8, clamp +/-127). """ assert data_hp.dtype in (torch.float32, torch.bfloat16), \ f"mx9_quantization only supports fp32 / bf16, got {data_hp.dtype}" @@ -439,71 +467,57 @@ def convert_to_mx9( blocks_per_row = triton.cdiv(last, block_size) n_blocks = rows * blocks_per_row - n_prime_bytes = (block_size // 2) // 8 - q = torch.empty((n_blocks, block_size), dtype=torch.int8, device=data_hp.device) - max_exp = torch.empty((n_blocks,), dtype=torch.uint8, device=data_hp.device) - prime = torch.empty((n_blocks, n_prime_bytes), dtype=torch.uint8, device=data_hp.device) + n_packed_bytes = _mx9_packed_bytes(block_size) + 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),) _convert_to_mx9_kernel[grid]( - x2d, q, max_exp, prime, n_blocks, last, blocks_per_row, + 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, ) - return q, max_exp, prime + return packed def convert_from_mx9( - q: torch.Tensor, - max_exp: torch.Tensor, - prime: torch.Tensor, + 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 three-part tuple -> reconstructed high-precision tensor. + """Packed MX9 tensor -> reconstructed high-precision tensor. out_shape / axis must match the original convert_to_mx9 call, used to reverse the transpose and padding. Returns shape = out_shape, dtype = out_dtype. """ assert out_dtype in _TORCH_TO_TL, \ f"out_dtype must be one of {tuple(_TORCH_TO_TL)}, got {out_dtype}" - assert q.dtype == torch.int8, f"q dtype must be int8, got {q.dtype}" - assert max_exp.dtype == torch.uint8, f"max_exp dtype must be uint8, got {max_exp.dtype}" - assert prime.dtype == torch.uint8, f"prime dtype must be uint8, got {prime.dtype}" + assert packed.dtype == torch.uint8, f"packed dtype must be uint8, got {packed.dtype}" assert block_size == 16, f"block_size only supports 16, got {block_size}" - # Three-part consistency check: q/max_exp/prime must share the same n_blocks - # and shapes must match the storage spec; otherwise the kernel addressing / - # reshape would read misaligned data or raise unhelpful dimension errors. - n_blocks = q.shape[0] - n_prime_bytes = (block_size // 2) // 8 - assert q.shape == (n_blocks, block_size), \ - f"q shape should be ({n_blocks}, {block_size}), got {tuple(q.shape)}" - assert max_exp.shape == (n_blocks,), \ - f"max_exp shape should be ({n_blocks},), got {tuple(max_exp.shape)}" - assert prime.shape == (n_blocks, n_prime_bytes), \ - f"prime shape should be ({n_blocks}, {n_prime_bytes}), got {tuple(prime.shape)}" - - y_blocks = torch.empty((n_blocks, block_size), dtype=out_dtype, device=q.device) - - # No forced .contiguous(): pass the packed inputs by stride so already-laid-out - # (or non-contiguous view) tensors are consumed without an extra copy. - stride_q_blk, stride_q_col = q.stride() - stride_e = max_exp.stride(0) - stride_p_blk, stride_p_col = prime.stride() + # Layout consistency check: packed must carry a whole number of blocks with + # the expected per-block byte count; otherwise the kernel addressing would + # read misaligned bytes or raise an unhelpful dimension error downstream. + n_blocks = packed.shape[0] + n_packed_bytes = _mx9_packed_bytes(block_size) + assert packed.shape == (n_blocks, n_packed_bytes), \ + f"packed shape should be ({n_blocks}, {n_packed_bytes}), got {tuple(packed.shape)}" + + y_blocks = torch.empty((n_blocks, block_size), dtype=out_dtype, device=packed.device) + + # No forced .contiguous(): pass the packed input by stride so an already + # laid-out (or non-contiguous view) tensor is consumed without an extra copy. + stride_blk, stride_col = packed.stride() grid = (triton.cdiv(n_blocks, blocks_per_program),) _convert_from_mx9_kernel[grid]( - q, max_exp, prime, y_blocks, n_blocks, - stride_q_blk, stride_q_col, - stride_e, - stride_p_blk, stride_p_col, + packed, y_blocks, n_blocks, + stride_blk, stride_col, BLOCK_SIZE=block_size, BLOCKS_PER_PROG=blocks_per_program, PRIME_GROUP=PRIME_GROUP, @@ -524,10 +538,10 @@ def convert_from_mx9( pad = (block_size - last % block_size) % block_size padded_cols = last + pad # out_shape/axis must be consistent with convert_to: if the inferred block - # count doesn't match the three-part tuple, intercept here with a readable + # count doesn't match the packed tensor, intercept here with a readable # error rather than letting the reshape below raise a cryptic dimension error. assert rows * padded_cols == n_blocks * block_size, ( - f"out_shape/axis/block_size inconsistent with three-part tuple: " + f"out_shape/axis/block_size inconsistent with packed tensor: " f"inferred rows*padded_cols={rows * padded_cols}, " f"but tuple has n_blocks*block_size={n_blocks * block_size}" ) diff --git a/alto/models/patcher.py b/alto/models/patcher.py index 186c7363..e596a88d 100644 --- a/alto/models/patcher.py +++ b/alto/models/patcher.py @@ -58,27 +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": + # 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, ) - q, max_exp, prime = convert_to_mx9(x) + packed = convert_to_mx9(x) return convert_from_mx9( - q, max_exp, prime, x.dtype, x.shape, + 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_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 index ae3479e0..820410fe 100644 --- a/tests/unittest/mx9_mx6/test_mx9_quantization.py +++ b/tests/unittest/mx9_mx6/test_mx9_quantization.py @@ -7,6 +7,14 @@ 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. @@ -28,6 +36,12 @@ 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 # --------------------------------------------------------------------------- # @@ -167,8 +181,8 @@ def _rand_with_outliers(shape, dtype, seed=0, p=0.005, spike=100.0): def _roundtrip(x, block_size=BLOCK_SIZE, axis=-1): - q, max_exp, prime = convert_to_mx9(x, block_size=block_size, axis=axis) - return convert_from_mx9(q, max_exp, prime, x.dtype, x.shape, + 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) @@ -178,28 +192,26 @@ def _roundtrip(x, block_size=BLOCK_SIZE, axis=-1): def test_output_dtypes_and_shapes(): x = _rand((4, 64), torch.float32).cuda() - q, max_exp, prime = convert_to_mx9(x, block_size=16) + packed = convert_to_mx9(x, block_size=16) n_blocks = (4 * 64) // 16 - assert q.dtype == torch.int8 - assert max_exp.dtype == torch.uint8 - assert prime.dtype == torch.uint8 - assert q.shape == (n_blocks, 16) - assert max_exp.shape == (n_blocks,) - assert prime.shape == (n_blocks, 1) # block_size=16 -> 8 pairs -> 1 byte + 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() - _, max_exp, _ = convert_to_mx9(x, block_size=16) + 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() - q, _, _ = convert_to_mx9(x, block_size=16) + 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 @@ -255,11 +267,12 @@ def _torch_intermediates(x, block_size=BLOCK_SIZE, quant_bit=8): def test_intermediates_match_torch_reference(shape, dtype): """Each component (q, max_exp, prime) matches torch recomputation bit-exact.""" x = _rand(shape, dtype).cuda() - q, max_exp, prime = convert_to_mx9(x, block_size=BLOCK_SIZE) + 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(max_exp, re), f"max_exp mismatch" - assert torch.equal(prime, rp), f"prime mismatch" + assert torch.equal(packed[:, 0], re), f"max_exp mismatch" + assert torch.equal(packed[:, _PRIME_OFFSET:_Q_OFFSET], rp), f"prime mismatch" def test_intermediates_constructed_demote(): @@ -272,11 +285,11 @@ def test_intermediates_constructed_demote(): x[0, 0] = 4.0 _, _, _, pair = _torch_intermediates(x) assert pair[0].tolist() == [0, 1, 1, 1, 1, 1, 1, 1] - q, max_exp, prime = convert_to_mx9(x) + packed = convert_to_mx9(x) rq, re, rp, _ = _torch_intermediates(x) - assert torch.equal(q, rq) - assert torch.equal(max_exp, re) - assert torch.equal(prime, rp) + 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) # --------------------------------------------------------------------------- # @@ -312,9 +325,9 @@ def test_roundtrip_axis(axis): 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) - q, me, pr = convert_to_mx9(x, block_size=BLOCK_SIZE, - blocks_per_program=blocks_per_program) - out = convert_from_mx9(q, me, pr, x.dtype, x.shape, + 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) @@ -415,10 +428,10 @@ def test_prime_bitmap_round_trips(): 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 - q, max_exp, prime = convert_to_mx9(x, block_size=BLOCK_SIZE) + 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 prime[0, 0].item() == 0x0F, f"got {hex(prime[0, 0].item())}" - out = convert_from_mx9(q, max_exp, prime, x.dtype, x.shape) + 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) @@ -563,6 +576,13 @@ def test_rejects_block_size_not_16(): def test_rejects_unknown_out_dtype(): x = _rand((4, 16), torch.float32).cuda() - q, me, pr = convert_to_mx9(x) + 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(q, me, pr, torch.int32, x.shape) + convert_from_mx9(packed, torch.float32, x.shape) 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]) From 087a8763b6a5d8e911b36701871ee8f039592753 Mon Sep 17 00:00:00 2001 From: jiarwang Date: Thu, 23 Jul 2026 07:40:36 +0000 Subject: [PATCH 7/7] refactor: mx: centralize shared MX6/MX9 helpers in _mx_common Move the width-/packing-independent Triton device helpers (_sanitize, _floor_exp, _round_half_even, _calculate_mx_exp, _scale_from_shared_exp) and the shared host pieces (_TORCH_TO_TL dtype map, _unpad_and_transpose shape restore) into alto/kernels/mx/_mx_common.py. MX6 and MX9 now import and reuse them instead of keeping byte-for-byte copies, so both formats provably share the same exponent/prime/scale math and diverge only in quant_bit and byte packing. Add MX9 device-helper probe tests mirroring MX6. --- alto/kernels/mx/_mx_common.py | 175 ++++++++++ alto/kernels/mx/mx6_quantization.py | 195 ++--------- alto/kernels/mx/mx9_quantization.py | 302 ++++-------------- .../unittest/mx9_mx6/test_mx6_quantization.py | 8 +- .../unittest/mx9_mx6/test_mx9_quantization.py | 179 ++++++++++- 5 files changed, 453 insertions(+), 406 deletions(-) create mode 100644 alto/kernels/mx/_mx_common.py diff --git a/alto/kernels/mx/_mx_common.py b/alto/kernels/mx/_mx_common.py new file mode 100644 index 00000000..28f2c342 --- /dev/null +++ b/alto/kernels/mx/_mx_common.py @@ -0,0 +1,175 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Shared MX packed-quantization Triton device helpers. + +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; the per-format ``_pack_*`` / ``_unpack_*`` and grid kernels stay +in their respective modules. + + - ``_sanitize`` : NaN / +/-Inf -> 0 for exponent statistics only. + - ``_floor_exp`` : floor(log2|x|) from native float bits. + - ``_round_half_even`` : round-to-nearest-ties-to-even (matches torch). + - ``_calculate_mx_exp`` : per-block shared exponent + per-element + shared_exp + per-pair prime bits. + - ``_scale_from_shared_exp`` : 2^(shared_exp - quant_bit + 2), clamped so the + scale is always normal (FTZ-independent). +""" + +import math + +import torch +import triton +import triton.language as tl + +# Shared dtype map used by both mx6 and mx9 host wrappers. +_TORCH_TO_TL = { + torch.float32: tl.float32, + torch.bfloat16: tl.bfloat16, +} + + +@triton.jit +def _sanitize(x): + """Replace NaN / +/-Inf with 0; used only for the amax / exponent statistics + copy, does not touch the QDQ data path. + + (Faithfully replicates Quark/reference: exponent statistics use sanitized + values to prevent Inf from polluting the block scale; the actual round still + operates on the original x, letting Inf be clamped. NaN is not representable + in the packed path, so NaN inputs yield undefined quantized values.) + """ + x = tl.where(x != x, 0.0, x) + x = tl.where(x == float("inf"), 0.0, x) + x = tl.where(x == float("-inf"), 0.0, x) + return x + + +@triton.jit +def _floor_exp(x): + """floor(log2(|x|)): extract exponent field directly from native float bits. + + Branches by dtype, aligned with ``mx.py``'s ``_exponent_frexp_no_exception``: + - fp32 : (bits>>23)&0xFF - 127 + - bf16 : (bits>>7) &0xFF - 127 + The ``& mask`` also clears sign-bit extension from arithmetic right shift, + so negative values are safe. Always returns int32 to avoid int16 mixing. + """ + 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), pure tl + implementation to avoid libdevice dependency.""" + 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, +): + """Compute per-block shared exponent + per-element shared_exp + per-pair + prime bits. Width-independent (does not involve ``quant_bit``), so MX6 and + MX9 share this verbatim. + + Input ``x``: this program's tile ``[BLOCKS_PER_PROG, BLOCK_SIZE]`` in native + dtype (each row is one MX block). Exponent extraction must be done on native + dtype, so do NOT pre-cast to fp32 on host. + + Returns: + - ``shared_exp`` : [BPP, BLOCK_SIZE] int32, per-element shared exponent + (demoted pairs get max_exp - 1) + - ``max_exp`` : [BPP] int32, per-block max exponent + - ``pair`` : [BPP, N_PAIRS] int32(0/1), per-pair prime bit + (1 = that pair gets 1-exponent demotion) + """ + N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP + + # Sanitized copy used only for exponent statistics (NaN/Inf -> 0). + clean = _sanitize(x) + + # Block shared exponent: per-row (block) amax -> floor extract exponent. + # tl.max on some backends promotes the reduce result to fp32; cast back to + # native dtype so _floor_exp takes the same branch as the per-element t_exp + # (consistent with the mxfp4 kernel). + amax = tl.max(tl.abs(clean), axis=1, keep_dims=True).to(clean.dtype) # [BPP, 1] native dtype + max_exp = _floor_exp(amax) # [BPP, 1] int32 + + # Per-element exponent + demote flag: whether at least 1 octave below block max. + t_exp = _floor_exp(clean) # [BPP, BLOCK_SIZE] int32 + demote = (max_exp - t_exp) >= 1 # [BPP, BLOCK_SIZE] bool (max_exp broadcasts) + + # Prime: adjacent PRIME_GROUP(=2) elements form a pair; both must be demoted + # for the pair to be demoted. + d = demote.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP) + pair_keep = tl.sum(d, axis=2, keep_dims=True) == PRIME_GROUP # [BPP, N_PAIRS, 1] bool + + # Broadcast back to per-element to get per-element shared_exp. + pair_b = tl.broadcast_to(pair_keep, (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) # [BPP, BLOCK_SIZE] int32 + + # Pair bits (not broadcast) are kept for prime bitmap packing. + pair = pair_keep.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): + """Per-element power-of-two scale = 2^(shared_exp - quant_bit + 2). + + Following mxfp4 (mxfp_quantization.py / PyTorch #125557): clamp the scale + exponent to the fp32 minimum normal exponent -126, ensuring the scale is + always normal (never subnormal or zero). This eliminates 0/0 and GPU FTZ + dependence at the source (only affects blocks with amax < 2^-120, which real + data never reaches). Shared by MX6 / MX9 pack and unpack so the scale -- and + hence the -126 clamp behaviour -- is provably identical across both. + """ + scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) + return tl.exp2(scale_exp.to(tl.float32)) # always normal (!=0) + + +def _unpad_and_transpose( + y_blocks: torch.Tensor, + out_shape, + axis: int, + block_size: int, + n_blocks: int, +) -> torch.Tensor: + """Reverse the transpose + padding applied by convert_to_mx* host wrappers. + + y_blocks : [n_blocks, block_size] freshly dequantized output from the kernel. + out_shape : original (pre-transpose) tensor shape passed by the caller. + axis : quant axis that was transposed to dim -1 before quantization. + """ + transposed_shape = list(out_shape) + transposed_shape[axis], transposed_shape[-1] = transposed_shape[-1], transposed_shape[axis] + last = transposed_shape[-1] + rows = math.prod(transposed_shape[:-1]) + pad = (block_size - last % block_size) % block_size + padded_cols = last + pad + 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 tuple has n_blocks*block_size={n_blocks * block_size}" + ) + 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 index 97c0a15f..62aa09b1 100644 --- a/alto/kernels/mx/mx6_quantization.py +++ b/alto/kernels/mx/mx6_quantization.py @@ -3,149 +3,46 @@ # SPDX-License-Identifier: MIT """Packed MX6 quantization Triton device kernels (called by grid kernels). -MX6 shares the *exact same block algorithm* as MX9 (exponent extraction, prime -demotion, power-of-two scale); the only numeric difference is the element -integer width ``quant_bit`` (MX9 = 8, MX6 = 5), exactly mirroring Quark's -``fake_quantize_mx6_mx9`` (see ``alto/modifiers/quantization/mx.py``). Everything -in this module is therefore a straightforward specialization of -``mx9_quantization.py``; the two intentional divergences (to be implemented in -later stages) are: - - 1. Element width: ``QUANT_BIT = 5`` -> quantized integers are symmetrically - clamped to +/-15 (=2^(quant_bit-1)-1), instead of MX9's +/-127. - 2. **Python-export-compatible packing** (achieves a true 6 bits/element): MX9 - stores ``q`` as int8 (one byte per element), dense at its 8-bit width. At 5 - bits, one byte per element would waste 3 bits and defeat the "6" in MX6. - To match the Python export path, each 16-element block is laid out as: +MX6 uses the same block algorithm as MX9 (exponent extraction, prime demotion, +power-of-two scale); it differs only in the element integer width +(``QUANT_BIT`` = 5 vs 8) and the byte packing, mirroring Quark's +``fake_quantize_mx6_mx9`` (see ``alto/modifiers/quantization/mx.py``). The +width-independent math lives in ``_mx_common``; ``mx9_quantization.py`` documents +the shared design decisions (scale clamp to -126, native-dtype exponent +extraction, E8M0 max_exp, stride addressing, ragged-tail masking). + +The two MX6-specific divergences are: + + 1. Element width: ``QUANT_BIT = 5``, so quantized integers are clamped to + +/-15 (= 2^(quant_bit-1)-1) instead of MX9's +/-127. + 2. Packing at a true 6 bits/element (MX9 stores q as a full int8 byte). At 5 + bits, one byte per element would waste 3 bits, so each 16-element block is + laid out to match the Python export path: - max_exp : 1 byte (E8M0, true exponent + 127) - prime : 1 byte (same bitmap as ``idx2`` in the Python packer) - sign : 2 bytes (16 sign bits, LSB = element 0) - - mantissa : 8 bytes (two ``q & 0xF`` nibbles per byte) - -> 12 bytes/block = 6 bits/element. The 4-bit mantissa nibbles are the low - bits of the signed quantized integer, not ``abs(q)``. - -The packed bytes are laid out per block as ``[max_exp, prime, sign bitmap ..., -mantissa nibbles ...]``. For the mantissa segment, low nibble = even element and -high nibble = odd element. - -Development is incremental. This file currently implements: - - Stage 1: stateless device helpers ``_sanitize`` / ``_floor_exp`` / - ``_round_half_even`` (identical to MX9). - - Stage 2: ``_calculate_mx6_exp`` (per-block shared exponent + per-element - shared_exp + per-pair prime bits; width-independent, identical to MX9). - - Stage 3: ``_pack_mx6`` / ``_unpack_mx6`` (Python-export QDQ pack + inverse). - - Stage 4: ``_convert_to_mx6_kernel`` / ``_convert_from_mx6_kernel`` grid entries - (stride-addressed load/store, ragged-tail masking, E8M0 bias). - - Stage 5: ``convert_to_mx6`` / ``convert_from_mx6`` host wrappers (axis - transpose, padding-aware shape restore -- mirrors ``mx9_quantization.py``). - -See ``mx9_quantization.py`` for the full rationale behind the shared design -decisions (scale exponent clamp to -126 / FTZ independence, exponent extraction -from native dtype bits, E8M0 max_exp, stride addressing, ragged-tail masking). + - mantissa : 8 bytes (two ``q & 0xF`` nibbles per byte, low = even element) + = 12 bytes/block = 6 bits/element. The 4-bit mantissa nibbles are the low + bits of the signed integer, not ``abs(q)``. """ import torch import triton import triton.language as tl +from ._mx_common import ( + _round_half_even, + _calculate_mx_exp, + _scale_from_shared_exp, + _TORCH_TO_TL, + _unpad_and_transpose, +) + BLOCK_SIZE = 16 QUANT_BIT = 5 PRIME_GROUP = 2 BLOCKS_PER_PROG_DEFAULT = 64 -# torch dtype -> triton dtype (used by unpack output cast, Stage 5 host wrapper). -_TORCH_TO_TL = { - torch.float32: tl.float32, - torch.bfloat16: tl.bfloat16, -} - - -@triton.jit -def _sanitize(x): - """Replace NaN / +/-Inf with 0; used only for the amax / exponent statistics - copy, does not touch the QDQ data path (identical to MX9).""" - x = tl.where(x != x, 0.0, x) - x = tl.where(x == float("inf"), 0.0, x) - x = tl.where(x == float("-inf"), 0.0, x) - return x - - -@triton.jit -def _floor_exp(x): - """floor(log2(|x|)): extract exponent field directly from native float bits - (fp32: (bits>>23)&0xFF - 127; bf16: (bits>>7)&0xFF - 127). Identical to MX9.""" - 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), pure tl. Identical to MX9.""" - 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_mx6_exp( - x, - BLOCKS_PER_PROG: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - PRIME_GROUP: tl.constexpr, -): - """Compute per-block shared exponent + per-element shared_exp + per-pair - prime bits. - - The exponent / prime math is *width-independent* (does not involve - ``quant_bit``), so this is byte-for-byte the MX9 ``_calculate_mx9_exp``. - - Input ``x``: this program's tile ``[BLOCKS_PER_PROG, BLOCK_SIZE]`` in native - dtype (each row is one MX6 block). Exponent extraction must be done on native - dtype, so do NOT pre-cast to fp32 on host. - - Returns: - - ``shared_exp`` : [BPP, BLOCK_SIZE] int32, per-element shared exponent - (demoted pairs get max_exp - 1) - - ``max_exp`` : [BPP] int32, per-block max exponent - - ``pair`` : [BPP, N_PAIRS] int32(0/1), per-pair prime bit - (1 = that pair gets 1-exponent demotion) - """ - N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP - - # Sanitized copy used only for exponent statistics (NaN/Inf -> 0). - clean = _sanitize(x) - - # Block shared exponent: per-row (block) amax -> floor extract exponent. - # Cast the reduce result back to native dtype so _floor_exp uses the same - # branch as the per-element t_exp (some backends promote tl.max to fp32). - amax = tl.max(tl.abs(clean), axis=1, keep_dims=True).to(clean.dtype) # [BPP, 1] native - max_exp = _floor_exp(amax) # [BPP, 1] int32 - - # Per-element exponent + demote flag (at least 1 octave below block max). - t_exp = _floor_exp(clean) # [BPP, BLOCK_SIZE] int32 - demote = (max_exp - t_exp) >= 1 # [BPP, BLOCK_SIZE] bool (broadcast) - - # Prime: adjacent PRIME_GROUP(=2) elements form a pair; both must be demoted - # for the pair to be demoted. - d = demote.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP) - pair_keep = tl.sum(d, axis=2, keep_dims=True) == PRIME_GROUP # [BPP, N_PAIRS, 1] bool - - # Broadcast back to per-element to get per-element shared_exp. - pair_b = tl.broadcast_to(pair_keep, (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) # [BPP, BLOCK_SIZE] int32 - - # Pair bits (not broadcast) are kept for prime bitmap packing (Stage 3). - pair = pair_keep.reshape(BLOCKS_PER_PROG, N_PAIRS).to(tl.int32) - return shared_exp, max_exp.reshape(BLOCKS_PER_PROG), pair - @triton.jit def _pack_mx6( @@ -175,13 +72,12 @@ def _pack_mx6( N_SIGN_BYTES: tl.constexpr = BLOCK_SIZE // 8 # 8 sign bits per byte Q_HI: tl.constexpr = (1 << (QUANT_BIT - 1)) - 1 # 15 for quant_bit=5 - # Same scale as MX9 (clamp scale exponent to -126 for FTZ independence). - scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) - scale = tl.exp2(scale_exp.to(tl.float32)) # [BPP, BLOCK_SIZE], always normal + # Same power-of-two scale as MX9 (shared helper; -126 FTZ-independence clamp). + scale = _scale_from_shared_exp(shared_exp, QUANT_BIT) # [BPP, BLOCK_SIZE], always normal xf = x.to(tl.float32) q = _round_half_even(tl.div_rn(xf, scale)) - q = tl.minimum(tl.maximum(q, -(Q_HI * 1.0)), Q_HI * 1.0) # clamp +/-15 + q = tl.minimum(tl.maximum(q, -(Q_HI * 1.0)), Q_HI * 1.0) qi = q.to(tl.int32) mantissa = qi & 0xF # [BPP, BLOCK_SIZE], low 4 bits @@ -251,8 +147,7 @@ def _unpack_mx6( pair_b = pair_b.reshape(BLOCKS_PER_PROG, BLOCK_SIZE) shared_exp = max_exp - pair_b # [BPP, BLOCK_SIZE] - scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) - scale = tl.exp2(scale_exp.to(tl.float32)) + scale = _scale_from_shared_exp(shared_exp, QUANT_BIT) y = q.to(tl.float32) * scale return y.to(out_dtype) @@ -319,7 +214,7 @@ def _convert_to_mx6_kernel( ) x = tl.load(x_ptr + in_offs, mask=load_mask, other=0.0) # native dtype tile - shared_exp, max_exp, pair = _calculate_mx6_exp( + shared_exp, max_exp, pair = _calculate_mx_exp( x, BLOCKS_PER_PROG, BLOCK_SIZE, PRIME_GROUP ) sign_bytes, mantissa_bytes, prime = _pack_mx6( @@ -417,7 +312,7 @@ def _convert_from_mx6_kernel( def _mx6_packed_bytes(block_size: int) -> int: """Bytes per block for the packed layout [max_exp | prime | sign | mantissa].""" - n_prime_bytes = (block_size // 2) // 8 + 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 @@ -511,29 +406,5 @@ def convert_from_mx6( OUT_DTYPE=_TORCH_TO_TL[out_dtype], ) - # Reverse: remove padding, reshape back to post-transpose shape, then - # transpose back to original axis. out_shape is the original (pre-transpose) - # shape; axis indicates which dimension was the quant axis. - transposed_shape = list(out_shape) - transposed_shape[axis], transposed_shape[-1] = transposed_shape[-1], transposed_shape[axis] - last = transposed_shape[-1] - rows = 1 - for s in transposed_shape[:-1]: - rows *= s - pad = (block_size - last % block_size) % block_size - padded_cols = last + pad - # out_shape/axis must be consistent with convert_to: if the inferred block - # count doesn't match the packed tensor, intercept here with a readable - # error rather than letting the reshape below raise a cryptic dimension error. - 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 tuple has n_blocks*block_size={n_blocks * block_size}" - ) - # Aligned with mx9 (convert_from_mx9): return the transposed stride view - # without a forced .contiguous() copy. Only the ragged-tail slice reshape may - # trigger an internal copy; the common (no-pad / axis=-1) path stays a view. - y2d = y_blocks.reshape(rows, padded_cols) - y2d = y2d[:, :last] - return y2d.reshape(transposed_shape).transpose(axis, -1) + return _unpad_and_transpose(y_blocks, out_shape, axis, block_size, n_blocks) diff --git a/alto/kernels/mx/mx9_quantization.py b/alto/kernels/mx/mx9_quantization.py index a623e811..fef94dec 100644 --- a/alto/kernels/mx/mx9_quantization.py +++ b/alto/kernels/mx/mx9_quantization.py @@ -3,190 +3,61 @@ # SPDX-License-Identifier: MIT """Packed MX9 quantization Triton device kernels (called by grid kernels). -Unlike the fake-quant path (``alto/kernels/mx/quantize_triton.py`` which outputs -dequantized bf16), this module performs *real packed* quantization: each MX9 -block is compressed into a **single packed uint8 tensor**, with each block's -bytes laid out as ``[max_exp(1B) | prime(1B) | q(16B)]`` (mirroring MX6's -prefix convention in ``mx6_quantization.py``): - - ``max_exp`` : 1 byte, per-block shared exponent, uint8 E8M0 - (floor(log2(amax)) + 127 bias) - - ``prime`` : 1 byte, per-block prime bitmap (1 bit per pair of 2 elements, - indicates whether that pair gets a 1-exponent demotion) - - ``q`` : 16 bytes, per-element quantized integer (symmetric clamp - +/-127), stored as the bit-identical uint8 reinterpretation of - int8 -- reconstruct via ``packed[..., 2:].view(torch.int8)``. --> 18 bytes/block = 9 bits/element (8b value + 0.5b amortized max_exp + 0.5b prime). - -The exponent extraction / scale / prime math is aligned with -``alto/modifiers/quantization/mx.py`` (exponent extracted from native dtype bits -= floor; QDQ computed in fp32). The **only intentional divergence** is the -quantized value clamp: - - True 8-bit variant: quantized integers are symmetrically clamped to +/-127, - stored as **int8**, yielding exactly 9 bits/element (8b value + 0.5b amortized - max_exp + 0.5b prime). In contrast, mx.py uses quant_max=255 for demoted pairs - and can preserve +/-128. The two differ only for elements that are both demoted - AND whose quantized value reaches +/-128. - -Therefore the acceptance reference is **not** ``mx9_fake_quantize`` but its -clamp-127 variant (see tests): -``unpack(pack(x)) == mx9_clamp127_ref(x)`` bit-exact. - -Constraints: - - Blocks are formed along the last dimension only (host transposes the quant - axis to dim -1). The tensor is passed to the kernel **by stride** (no forced - ``.contiguous()`` copy); the kernel addresses elements via row/column strides - and masks the ragged tail block instead of host-side ``F.pad``. - - ``block_size`` is fixed at 16: each block has 8 pairs which pack into - exactly 1 byte of prime bitmap. - -Key design decisions (lessons learned / non-obvious trade-offs -- understand -these before modifying): - 1. q uses int8 + clamp +/-127, not int16: demoted pairs use half-scale so - the quantized value can reach +/-128 (quant_max=255), which overflows signed - int8. We previously used int16 to preserve +/-128 bit-exactness, but that - yields ~17 bits/element (larger than bf16), defeating the purpose of packing. - Choosing clamp +/-127 achieves true 9 bits/element, at the cost of a 1-LSB - divergence from mx.py on the rare "demoted AND reaching +/-128" elements - (measured ~0.1%). - 2. Scale exponent is clamped to the minimum normal exponent -126 (following - mxfp4 / PyTorch #125557), rather than a post-hoc scale==0 guard: the latter - implicitly depends on GPU FTZ flushing subnormal scales to 0, which once - caused divergence of ~1e-41 for deep-subnormal blocks vs a no-FTZ reference. - Clamping keeps scale always normal, FTZ-independent, and degenerate blocks - deterministically output 0. - 3. Exponents are extracted from native dtype bits (fp32/bf16 >>... - 127); - do NOT pre-cast to fp32 on host. - 4. max_exp is stored as uint8 E8M0 (true exponent + 127), not int32: true - exponents for any finite input fall in [-127, 127], +127 -> [0, 254] which - fits uint8 without overflow; this achieves 1 byte/block and 9 bits/element. - 5. All three components live in one uint8 tensor (not three separate tensors): - ``q`` is computed as int8 in registers (see ``_pack_mx9`` / ``_unpack_mx9``, - unchanged) but stored/loaded via the packed uint8 tensor using - ``.to(tl.uint8/tl.int8, bitcast=True)`` -- a pure bit reinterpretation, not a - value cast (which would incorrectly clamp/wrap negative values). This mirrors - ``mx6_quantization.py``'s single-tensor, Python-export-compatible layout. - -Known limitations: - - NaN: packed integers cannot represent NaN; blocks containing NaN will have - garbage q values (only the fake-quant path supports NaN pass-through). - - The acceptance reference is the custom clamp-127 variant (Quark-divergent), - not Quark/mx.py itself. +Unlike the fake-quant path (``alto/kernels/mx/quantize_triton.py``, which outputs +dequantized bf16), this module produces real packed quantization: each MX9 block +is compressed into a single uint8 tensor laid out as +``[max_exp(1B) | prime(1B) | q(16B)]`` (mirroring MX6's prefix convention): + - ``max_exp`` : per-block shared exponent, uint8 E8M0 (floor(log2(amax)) + 127) + - ``prime`` : per-block prime bitmap, 1 bit per pair of 2 elements (whether + that pair gets a 1-exponent demotion) + - ``q`` : per-element quantized integer, clamped +/-127, stored as the + bit-identical uint8 view of int8 -- reconstruct via + ``packed[..., 2:].view(torch.int8)``. +This is 18 bytes/block = 9 bits/element (8b value + 0.5b max_exp + 0.5b prime). + +The exponent / scale / prime math follows ``alto/modifiers/quantization/mx.py`` +(exponent from native dtype bits = floor; QDQ in fp32); the shared, width- +independent pieces live in ``_mx_common``. This module (and MX6) reference it for +the -126 scale clamp, native-dtype exponent extraction, and E8M0 max_exp. + +Design decisions worth knowing before modifying: + - q uses int8 + clamp +/-127, not int16. Demoted pairs use half-scale, so a + quantized value can reach +/-128 (mx.py's quant_max=255), which overflows + signed int8. Preserving +/-128 would need int16 (~17 bits/element, larger + than bf16 and defeating the point of packing). Clamping to +/-127 keeps a + true 9 bits/element at the cost of a 1-LSB divergence from mx.py on the rare + "demoted and reaching +/-128" elements (~0.1%). The acceptance reference is + therefore the clamp-127 variant, not ``mx9_fake_quantize``: + ``unpack(pack(x)) == mx9_clamp127_ref(x)`` bit-exact (see tests). + - q lives in the same uint8 tensor as max_exp/prime via a bitcast + (``.to(tl.uint8/tl.int8, bitcast=True)``), a bit reinterpretation rather than + a value cast (a value cast would wrap negatives). + - block_size is fixed at 16 (8 pairs pack into exactly 1 prime byte). Blocks + form along the last dim; the host transposes the quant axis to dim -1 and + passes the tensor by stride (no forced ``.contiguous()``), with the ragged + tail masked in-kernel instead of ``F.pad``. + +NaN limitation: packed integers cannot represent NaN, so blocks containing NaN +yield garbage q values; only the fake-quant path passes NaN through. """ import torch import triton import triton.language as tl +from ._mx_common import ( + _round_half_even, + _calculate_mx_exp, + _scale_from_shared_exp, + _TORCH_TO_TL, + _unpad_and_transpose, +) + BLOCK_SIZE = 16 QUANT_BIT = 8 PRIME_GROUP = 2 BLOCKS_PER_PROG_DEFAULT = 64 -# torch dtype -> triton dtype (used by unpack output cast). -_TORCH_TO_TL = { - torch.float32: tl.float32, - torch.bfloat16: tl.bfloat16, -} - - -@triton.jit -def _sanitize(x): - """Replace NaN / +/-Inf with 0; used only for the amax / exponent statistics - copy, does not touch the QDQ data path. - - (Faithfully replicates Quark/reference: exponent statistics use sanitized - values to prevent Inf from polluting the block scale; the actual round still - operates on the original x, letting Inf be clamped. NaN is not representable - in the packed path (q is int8), so NaN inputs yield undefined q values.) - """ - x = tl.where(x != x, 0.0, x) - x = tl.where(x == float("inf"), 0.0, x) - x = tl.where(x == float("-inf"), 0.0, x) - return x - - -@triton.jit -def _floor_exp(x): - """floor(log2(|x|)): extract exponent field directly from native float bits. - - Branches by dtype, aligned with ``mx.py``'s ``_exponent_frexp_no_exception``: - - fp32 : (bits>>23)&0xFF - 127 - - bf16 : (bits>>7) &0xFF - 127 - The ``& mask`` also clears sign-bit extension from arithmetic right shift, - so negative values are safe. Always returns int32 to avoid int16 mixing. - """ - 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), pure tl - implementation to avoid libdevice dependency.""" - 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_mx9_exp( - x, - BLOCKS_PER_PROG: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - PRIME_GROUP: tl.constexpr, -): - """Compute per-block shared exponent + per-element shared_exp + per-pair - prime bits. - - Input ``x``: this program's tile ``[BLOCKS_PER_PROG, BLOCK_SIZE]`` in - **native dtype** (each row is one MX9 block). Exponent extraction must be - done on native dtype, so do NOT pre-cast to fp32 on host. - - Returns: - - ``shared_exp`` : [BPP, BLOCK_SIZE] int32, per-element shared exponent - (demoted pairs get -1) - - ``max_exp`` : [BPP, 1] int32, per-block max exponent - - ``pair`` : [BPP, N_PAIRS] int32(0/1), per-pair prime bit - (1 = that pair gets 1-exponent demotion) - """ - N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP - - # Sanitized copy used only for exponent statistics (NaN/Inf -> 0). - clean = _sanitize(x) - - # Block shared exponent: per-row (block) amax -> floor extract exponent. - # tl.max on some backends promotes the reduce result to fp32; cast back to - # native dtype to ensure _floor_exp uses the same branch as per-element - # t_exp, consistent with the mxfp4 kernel (see mxfp_quantization.py:72). - amax = tl.max(tl.abs(clean), axis=1, keep_dims=True).to(clean.dtype) # [BPP, 1] native dtype - max_exp = _floor_exp(amax) # [BPP, 1] int32 - - # Per-element exponent + demote flag: whether at least 1 octave below block max. - t_exp = _floor_exp(clean) # [BPP, BLOCK_SIZE] int32 - demote = (max_exp - t_exp) >= 1 # [BPP, BLOCK_SIZE] bool (max_exp broadcasts) - - # Prime: adjacent PRIME_GROUP(=2) elements form a pair; both must be demoted - # for the pair to be demoted. - d = demote.to(tl.int32).reshape(BLOCKS_PER_PROG, N_PAIRS, PRIME_GROUP) - pair_keep = tl.sum(d, axis=2, keep_dims=True) == PRIME_GROUP # [BPP, N_PAIRS, 1] bool - - # Broadcast back to per-element to get per-element shared_exp. - pair_b = tl.broadcast_to(pair_keep, (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) # [BPP, BLOCK_SIZE] int32 - - # Pair bits (not broadcast) are kept for prime bitmap packing. - pair = pair_keep.reshape(BLOCKS_PER_PROG, N_PAIRS).to(tl.int32) - return shared_exp, max_exp.reshape(BLOCKS_PER_PROG), pair - @triton.jit def _pack_mx9( @@ -198,48 +69,25 @@ def _pack_mx9( PRIME_GROUP: tl.constexpr, QUANT_BIT: tl.constexpr, ): - """Quantize + pack: original x -> (q:int8 values, prime:uint8 bitmap). - - True 8-bit variant: quantized integers are symmetrically clamped to +/-127, - stored as int8 (achieving 9 bits/element). This intentionally diverges from - mx.py (which uses quant_max=255 for demoted pairs, preserving +/-128) at - the +/-128 boundary -- this variant serves as its own reference (see the - clamp-127 reference in tests). ``prime`` packs every 8 pair bits into one - byte. + """Quantize + pack: x -> (q:int8 clamped +/-127, prime:uint8 bitmap). + + ``prime`` packs every 8 pair bits into one byte. """ N_PAIRS: tl.constexpr = BLOCK_SIZE // PRIME_GROUP N_PRIME_BYTES: tl.constexpr = N_PAIRS // 8 # block 16 -> 1; block 32 -> 2 - # Per-element scale = 2^(shared_exp - quant_bit + 2). - # Following mxfp4 (mxfp_quantization.py:81-89 / PyTorch #125557): clamp the - # scale exponent to fp32 minimum normal exponent -126, ensuring scale is - # always normal (never subnormal or zero), eliminating 0/0 and GPU FTZ - # dependence from the source (only affects blocks with amax < 2^-120, which - # real data never reaches). - scale_exp = tl.maximum(shared_exp - QUANT_BIT + 2, -126) - scale = tl.exp2(scale_exp.to(tl.float32)) # [BPP, BLOCK_SIZE], always normal (!=0) - - # Q part of QDQ operates on the *original* x (NaN will produce garbage, see below). - # div_rn (IEEE RN) matches PyTorch eager x/scale; plain / lowers to fast div - # (max 2 ULP error) which can flip round at .5 boundaries. + scale = _scale_from_shared_exp(shared_exp, QUANT_BIT) # [BPP, BLOCK_SIZE], always normal (!=0) + + # div_rn (IEEE round-nearest) matches PyTorch eager x/scale; plain / lowers to + # fast div (up to 2 ULP) which can flip rounding at .5 boundaries. xf = x.to(tl.float32) q = _round_half_even(tl.div_rn(xf, scale)) - # True 8-bit: symmetric clamp to +/-127. - # - Normal data: demoted elements have |q|<128 mathematically; only rounding - # edge cases can produce 128, which is clamped to 127. - # - +/-Inf input: Inf/scale=Inf, clamped to +/-127 (does not pollute block - # scale since _sanitize already cleared it). - # - Intentional divergence from mx.py: mx.py uses quant_max=255 for demoted - # pairs, preserving +/-128. + # Clamp to +/-127 (see module docstring); also catches Inf/scale=Inf. q = tl.minimum(tl.maximum(q, -127.0), 127.0) + q_int = q.to(tl.int8) # NaN inputs yield garbage here; excluded from round-trip tests - # NaN limitation: packed integers cannot represent NaN; blocks containing NaN - # will have garbage q values (round-trip tests must exclude NaN inputs; only - # the fake-quant path supports NaN pass-through). - q_int = q.to(tl.int8) - - # Prime packing: N_PAIRS pair bits -> every 8 bits packed into 1 byte. - # weights = [1,2,4,...,128], generated via exp2 to avoid 1< 1 byte per 8, weights [1,2,...,128] via exp2 + # (avoids the 1< int: """Bytes per block for the packed layout [max_exp(1B) | prime | q(block_size)B].""" - n_prime_bytes = (block_size // 2) // 8 + n_prime_bytes = (block_size // PRIME_GROUP) // 8 return 1 + n_prime_bytes + block_size @@ -525,29 +372,4 @@ def convert_from_mx9( OUT_DTYPE=_TORCH_TO_TL[out_dtype], ) - # Reverse: remove padding, reshape back to post-transpose shape, then - # transpose back to original axis. - # out_shape is the original (pre-transpose) shape; axis indicates which - # dimension was the quant axis. - transposed_shape = list(out_shape) - transposed_shape[axis], transposed_shape[-1] = transposed_shape[-1], transposed_shape[axis] - last = transposed_shape[-1] - rows = 1 - for s in transposed_shape[:-1]: - rows *= s - pad = (block_size - last % block_size) % block_size - padded_cols = last + pad - # out_shape/axis must be consistent with convert_to: if the inferred block - # count doesn't match the packed tensor, intercept here with a readable - # error rather than letting the reshape below raise a cryptic dimension error. - 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 tuple has n_blocks*block_size={n_blocks * block_size}" - ) - # Aligned with mxfp4 (convert_from_mxfp4): return the transposed stride view - # without a forced .contiguous() copy. Only the ragged-tail slice reshape may - # trigger an internal copy; the common (no-pad / axis=-1) path stays a view. - y2d = y_blocks.reshape(rows, padded_cols) - y2d = y2d[:, :last] - return y2d.reshape(transposed_shape).transpose(axis, -1) + return _unpad_and_transpose(y_blocks, out_shape, axis, block_size, n_blocks) diff --git a/tests/unittest/mx9_mx6/test_mx6_quantization.py b/tests/unittest/mx9_mx6/test_mx6_quantization.py index fa0348c1..9b8a19b7 100644 --- a/tests/unittest/mx9_mx6/test_mx6_quantization.py +++ b/tests/unittest/mx9_mx6/test_mx6_quantization.py @@ -33,15 +33,17 @@ import triton import triton.language as tl +from alto.kernels.mx._mx_common import ( + _sanitize, + _floor_exp, + _calculate_mx_exp as _calculate_mx6_exp, +) from alto.kernels.mx.mx6_quantization import ( BLOCK_SIZE, PRIME_GROUP, QUANT_BIT, BLOCKS_PER_PROG_DEFAULT, - _sanitize, - _floor_exp, _round_half_even, - _calculate_mx6_exp, _pack_mx6, _unpack_mx6, _convert_to_mx6_kernel, diff --git a/tests/unittest/mx9_mx6/test_mx9_quantization.py b/tests/unittest/mx9_mx6/test_mx9_quantization.py index 820410fe..2ad310b5 100644 --- a/tests/unittest/mx9_mx6/test_mx9_quantization.py +++ b/tests/unittest/mx9_mx6/test_mx9_quantization.py @@ -26,10 +26,24 @@ 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 convert_to_mx9, convert_from_mx9 +from alto.kernels.mx.mx9_quantization import ( + PRIME_GROUP, + QUANT_BIT, + BLOCKS_PER_PROG_DEFAULT, + convert_to_mx9, + convert_from_mx9, +) +from alto.kernels.mx._mx_common import ( + _sanitize, + _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( @@ -161,6 +175,169 @@ def _mx9_clamp127_ref_via_mxpy(x: torch.Tensor, block_size: int = BLOCK_SIZE, return out.transpose(axis, -1) +# --------------------------------------------------------------------------- # +# Probe kernels for device helpers (Stage 1 / Stage 2) +# --------------------------------------------------------------------------- # + +@triton.jit +def _probe_sanitize_kernel(x_ptr, o_ptr, N: tl.constexpr): + i = tl.arange(0, N) + tl.store(o_ptr + i, _sanitize(tl.load(x_ptr + i))) + + +@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: _sanitize +# --------------------------------------------------------------------------- # + +def test_sanitize_replaces_nonfinite_with_zero(): + x = torch.tensor( + [1.0, float("nan"), float("inf"), float("-inf"), -2.0, 0.0, 3.5, float("nan")], + dtype=torch.float32, device="cuda", + ) + out = torch.empty_like(x) + _probe_sanitize_kernel[(1,)](x, out, N=x.numel()) + ref = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- # +# 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)