feat: mx9: add fused MX9 quantization kernel with pack/unpack

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``).
+"""

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)