feat: mx9: add fused MX9 quantization kernel with pack/unpack by Jiarwang77 · Pull Request #33 · AMD-AGI/ALTO

Jiarwang77 · 2026-06-29T11:05:42Z

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.

+    assert out_dtype in _TORCH_TO_TL
+    assert block_size == 16, f"block_size only supports 16, got {block_size}"
+


+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


hann-wang · 2026-06-30T07:44:25Z

+# Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# SPDX-License-Identifier: MIT
+"""MX real (fused Triton) kernels.


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

hann-wang · 2026-06-30T07:51:09Z

+
+    # 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)


is it possible to use triton.language.div_rn?

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

- 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.

…9 to the packed kernels

…imize them.

hann-wang · 2026-07-29T07:52:09Z

+    internally. The original shape / axis are not stored, so the caller must pass
+    them back to convert_from_mx6.
+    """
+    return _convert_to_mx_host(


Normally "host" refers to functions executed on CPU but here _convert_to_mx_host lauches a triton kernel that requires GPU. Better change this to "device" or "gpu".

blocks_per_program is a tunable parameter and should not be exposed to end-user. Keep it hidden for now would be fine.

hann-wang · 2026-07-29T07:52:23Z

+    call. They are not stored in ``packed``; a mismatch may silently reorder
+    values when it happens to produce the same block count.
+    """
+    return _convert_from_mx_host(


same as above.

hann-wang · 2026-07-29T07:55:34Z

+    """
+    return _convert_to_mx_host(
+        data_hp,
+        kernel=_convert_to_mx6_kernel,


I don't think we will have another equivalent kernel in the near future. Just launch this kernel inside your dispatch function. No need to pass it as an argument.

hann-wang · 2026-07-29T07:56:55Z

+    return _convert_to_mx_host(
+        data_hp,
+        kernel=_convert_to_mx6_kernel,
+        n_packed_bytes=_mx6_packed_bytes(block_size),


we have block_size passed as an argument. Why not calculating it inside the dispatch function?

On another thought, block_size, quant_bit, prime_group are all fixed parameters for MX6 (or MX9). Just signal whether you want MX6 or MX9 ("target_dtype=mx6/mx9") would be sufficient.

hann-wang · 2026-07-29T08:01:07Z

+        data_hp,
+        kernel=_convert_to_mx6_kernel,
+        n_packed_bytes=_mx6_packed_bytes(block_size),
+        fmt="mx6_quantization",


fmt="mx6_quantization" has no actual function (only used for error logging) here. Please read the AI-generated code before accepting it.

hann-wang · 2026-07-29T08:07:34Z

+        PRIME_GROUP,
+        QUANT_BIT,
+    )
+    out_offs = blk[:, None] * BLOCK_SIZE + col[None, :]


same as above. no strides for output tensors?

hann-wang · 2026-07-29T08:08:05Z

+
+    # E8M0 bias: store the true exponent + 127.
+    e_store = (max_exp + 127).to(tl.uint8)
+    tl.store(packed_ptr + blk * N_PACKED_BYTES, e_store, mask=blk_mask)


same as MX6. no strides for output tensors.

hann-wang · 2026-07-29T08:08:16Z

+        PRIME_GROUP,
+        QUANT_BIT,
+    )
+    out_offs = blk[:, None] * BLOCK_SIZE + col[None, :]


same as MX6. no strides for output tensors.

hann-wang · 2026-07-29T08:10:26Z

+                    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)


Is this path still working? original_fake_quantize

hann-wang · 2026-07-29T08:13:30Z

+export LOG_RANK=${LOG_RANK:-0}
+TRAIN_FILE=${TRAIN_FILE:-"alto.train"}
+MODULE=${MODULE:-"llama3"}
+CONFIG=${CONFIG:-"llama3_1b_mx6_wa"}


Please help us refactor all these example scripts.

Instead of copying examples/run.sh, re-use the script like this:

NGPU=1 MODULE=llama3 CONFIG=llama3_1b_mx6_wa ./examples/run.sh --hf_assets_path xxxx --validator.steps xxx

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

2de0111

hann-wang requested a review from Copilot June 30, 2026 07:30

Copilot started reviewing on behalf of hann-wang June 30, 2026 07:30 View session

Copilot AI reviewed Jun 30, 2026

View reviewed changes

hann-wang reviewed Jun 30, 2026

View reviewed changes

Jiarwang77 added 3 commits June 30, 2026 11:02

feat: mx6: add packed quantization kernel

d3ce92b

perf: mx9: consume strided quantization inputs

db210c0

Jiarwang77 force-pushed the jiarwang/mx9-fake-quant branch from 548d1e5 to db210c0 Compare July 14, 2026 08:18

Jiarwang77 added 2 commits July 17, 2026 08:54

feat: mx6: add convert_to_mx6 / convert_from_mx6 host wrappers

c5ce915

refactor: mx: pack mx9 into a single uint8 tensor and dispatch mx6/mx…

abf9559

…9 to the packed kernels

Jiarwang77 force-pushed the jiarwang/mx9-fake-quant branch 3 times, most recently from 207fb94 to 087a876 Compare July 23, 2026 08:39

refactor: mx: Extract the common parts of MX6 and MX9, review and opt…

a505eb5

…imize them.

Jiarwang77 force-pushed the jiarwang/mx9-fake-quant branch from 087a876 to a505eb5 Compare July 28, 2026 07:55

Jiarwang77 requested a review from hann-wang July 28, 2026 08:20

hann-wang reviewed Jul 29, 2026

View reviewed changes

		assert out_dtype in _TORCH_TO_TL
		assert block_size == 16, f"block_size only supports 16, got {block_size}"

Uh oh!

Conversation

Jiarwang77 commented Jun 29, 2026

Summary

Key design choices

Test coverage (70 tests, all passing)

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants