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

[Jiarwang77]

Summary

Add a fused Triton kernel for real packed MX9 quantization (convert_to_mx9 / convert_from_mx9), complementing the existing fake-quant path in alto.modifiers.quantization.mx.

Unlike the fake-quant path (which outputs dequantized bf16/fp32 for training), this kernel produces a compact three-part packed representation:

• q — [n_blocks, 16] int8: per-element quantized values, symmetric clamp ±127.
• max_exp — [n_blocks] uint8: per-block shared exponent in E8M0 format (true exponent + 127 bias).
• prime — [n_blocks, 1] uint8: per-block prime bitmap, 1 bit per pair of 2 elements, byte-packed.

This achieves 9 bits/element (8b value + 0.5b amortized exponent + 0.5b prime), enabling real weight compression for inference/storage.

Key design choices

• Symmetric clamp ±127 (not ±128): fits signed int8 without overflow; intentionally diverges from mx.py's 255-clamp at the rare demoted-±128 boundary (~0.1% of elements).
• Scale exponent clamped to -126 (minimum fp32 normal): eliminates FTZ dependence and 0/0 for degenerate blocks, following the mxfp4 kernel pattern.
• Exponent extracted from native dtype bits (not pre-cast fp32): ensures bit-exact alignment with mx.py's _t_exponent.

Test coverage (70 tests, all passing)

• Layer 1: Storage format assertions (dtype/shape/range)
• Layer 1b: Per-component intermediate verification (q/max_exp/prime individually vs torch recomputation)
• Layer 2: Round-trip bit-exact vs clamp-127 reference (6 shapes × 2 dtypes + axis + blocks_per_program sweep)
• Layer 3: Boundary & properties (zeros, Inf, block independence, idempotency, non-contiguous input, padding, ±128 clamp edge, prime bitmap symmetry, subnormal/tiny blocks, dual-reference cross-validation, outlier/heavy-tailed data, SNR quality floor)
• Layer 4: Invalid input rejection (wrong dtype, fp16, block_size≠16, bad out_dtype)

Test plan

• pytest tests/unittest/mx9_mx6/test_mx9_quantization.py — 70 passed (ROCm/HIP, MI300X)
• Verified SNR: ~46 dB (clean) / ~42 dB (outlier) — well above conservative floors (25/15 dB)
• Bit-exact round-trip agreement with two independent PyTorch references (handwritten + mx.py-based)
• Future: register as @triton_op + register_fake for torch.compile integration
• Future: wire into patcher.py / inference pipeline for end-to-end weight compression

[Copilot]

Pull request overview

Adds real packed MX9 quantization support via a fused Triton pack/unpack kernel, enabling a compact 3-part representation (q / max_exp / prime) intended for inference/storage compression, along with an extensive unit test suite validating bit-exact round-trips against a clamp-127 reference.

Changes:

• Introduces convert_to_mx9 / convert_from_mx9 Triton implementation for real packed MX9 (int8 q + E8M0-like exponent byte + prime bitmap).
• Adds comprehensive MX9 pack/unpack tests (format, intermediates, round-trip, edge cases, invalid inputs).
• Adds alto.kernels.mx package initializer and documentation.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
alto/kernels/mx/mx9_quantization.py	New Triton MX9 pack/unpack kernel + Python wrappers for producing/consuming the packed 3-tensor representation.
tests/unittest/mx9_mx6/test_mx9_quantization.py	New end-to-end and component-level tests validating storage format and bit-exact round-trip behavior.
alto/kernels/mx/__init__.py	Declares the MX “real kernel” package and documents its relationship to the fake-quant reference path.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[hann-wang]

This is misleading. We don't have real mx gemm now.

[hann-wang]

is it possible to use triton.language.div_rn?

[Jiarwang77]

Yes, I found "PyTorch eager uses the equivalent of Triton's div_rn" in torch inductor. It can be more accurate.