diff --git a/alto/kernels/mx/__init__.py b/alto/kernels/mx/__init__.py new file mode 100644 index 0000000..506cad3 --- /dev/null +++ b/alto/kernels/mx/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""MX packed-quantization Triton kernels (currently MX9 only). + +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 new file mode 100644 index 0000000..a4f7f8b --- /dev/null +++ b/alto/kernels/mx/mx9_quantization.py @@ -0,0 +1,491 @@ +# 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. 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( + 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). + # 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(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. + 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, \ + 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 + # 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/alto/models/patcher.py b/alto/models/patcher.py index 8c75135..186c736 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 ( 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 0000000..ae9568f --- /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)