feat: petr: add fully-dynamic MX6/MX9 quantization recipes

alto/models/patcher.py

@@ -62,21 +62,16 @@ def forward(ctx, x, scale, zero_point, args, g_idx, global_scale):
                     # 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). num_bits / group_size are fixed by the
-                    # format rather than honoured, so they are validated here
-                    # instead of forwarded -- a recipe that disagrees with the
-                    # format would otherwise be silently ignored.
                     from alto.kernels.mx import MX_QUANT_BIT, convert_to_mx, convert_from_mx
 
                     assert args.group_size in (None, 16), \
                         f"{target_dtype} packed kernel only supports group_size 16, got {args.group_size}"
                     assert args.num_bits == MX_QUANT_BIT[target_dtype], \
                         (f"{target_dtype} packed kernel is fixed at num_bits "
                          f"{MX_QUANT_BIT[target_dtype]}, got {args.num_bits}")
-                    packed = convert_to_mx(x, target_dtype=target_dtype)
-                    return convert_from_mx(packed, target_dtype, x.dtype, x.shape)
+                    axis = getattr(args, "block_axis", -1)
+                    packed = convert_to_mx(x, target_dtype=target_dtype, axis=axis)
+                    return convert_from_mx(packed, target_dtype, x.dtype, x.shape, axis=axis)
                 return original_fake_quantize(x, scale, zero_point, args, g_idx, global_scale)
 
             @staticmethod

alto/models/petr/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# SPDX-License-Identifier: MIT

alto/models/petr/quantize.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# SPDX-License-Identifier: MIT
+"""Fully-dynamic MX6 / MX9 recipes.
+
+Model-agnostic: the part of a model to leave alone is named by pattern, not by
+an attribute name a particular architecture happens to use, so the same entry
+point works on a detector, an LLM, or anything else built out of ``nn.Conv2d``
+and ``nn.Linear``.
+
+Mixed precision comes from ``ignore``. A model already runs in some dtype, and
+the modules a recipe skips keep it, so "MX6 here, BF16 there" is one recipe with
+the BF16 half excluded.
+
+Scales come from the tensor being quantized on each forward pass, so no
+calibration data, observers or persisted scales are involved.
+
+The public surface is :func:`apply_mx_quantization`; building a recipe and
+driving its modifiers through ``convert -> initialize -> pre_step -> post_step
+-> finalize`` is internal.
+"""
+
+import re
+import torch.nn as nn
+from compressed_tensors.utils import match_named_modules
+from alto.config import Recipe
+
+_MX_FORMAT_BITS = {"mx6": 5, "mx9": 8}
+
+# Axis the MX blocks run along, per module type. A Conv2d is blocked along its
+# input channels;
+_MX_BLOCK_AXIS = {nn.Conv2d: 1, nn.Linear: -1}
+
+
+def _mx_quant_args(mx_format: str, block_axis: int) -> dict:
+    """Fully-dynamic MX quantization args.
+
+    ``num_bits`` is required by ``compressed_tensors`` for bookkeeping and
+    follows the format's fixed physical layout; the Triton codec still owns the
+    actual packed representation.
+    """
+    try:
+        num_bits = _MX_FORMAT_BITS[mx_format]
+    except KeyError as exc:
+        supported = ", ".join(sorted(_MX_FORMAT_BITS))
+        raise ValueError(
+            f"unsupported MX format {mx_format!r}; choose one of: {supported}"
+        ) from exc
+
+    return {
+        "num_bits": num_bits,
+        "type": "int",
+        "symmetric": True,
+        "strategy": "tensor",
+        "dynamic": True,
+        "format": mx_format,
+        "block_axis": block_axis,
+    }
+
+
+def _create_mx_recipe(mx_format: str, ignore_patterns=()) -> Recipe:
+    """Build an MX W+A recipe for every Conv2d and Linear.
+
+    Args:
+        mx_format: ``"mx6"`` or ``"mx9"``.
+        ignore_patterns: ``re:`` patterns naming the modules left out of the
+            recipe, so they keep the model's dtype.
+    """
+    config_groups = {
+        # Split per module type because MX blocks run along a different axis for
+        # each.
+        f"group_{cls.__name__.lower()}": {
+            "targets": [cls.__name__],
+            "weights": dict(_mx_quant_args(mx_format, axis)),
+            "input_activations": dict(_mx_quant_args(mx_format, axis)),
+        }
+        for cls, axis in _MX_BLOCK_AXIS.items()
+    }
+    # ignore sits on the modifier, not on a scheme: compressed_tensors applies it
+    # across every config group.
+    modifier = {"sequential": False, "config_groups": config_groups}
+    if ignore_patterns:
+        modifier["ignore"] = list(ignore_patterns)
+    return Recipe.from_dict(
+        {
+            "quantization_stage": {
+                "quantization_modifiers": {"QuantizationModifier": modifier}
+            }
+        }
+    )
+
+
+def _apply_recipe(model, recipe):
+    """Drive a recipe's modifiers through their lifecycle against ``model``."""
+    modifiers = recipe.modifiers
+    if not modifiers:
+        raise ValueError(f"recipe {recipe} produced no modifiers")
+
+    model_parts = [model]
+    for modifier in modifiers:
+        modifier.convert(model)
+    for modifier in modifiers:
+        modifier.initialize(model_parts)
+    for modifier in modifiers:
+        modifier.pre_step(model_parts)
+    for modifier in modifiers:
+        modifier.post_step(model_parts)
+    for modifier in modifiers:
+        modifier.finalize(model_parts)
+    return model
+
+
+def apply_mx_quantization(model, mx_format: str, ignore=()):
+    """Apply an MX6 or MX9 W+A recipe in place.
+
+    Every Conv2d and Linear is quantized unless ``ignore`` excludes it; excluded
+    modules keep the dtype the model already runs in.
+
+    Raises:
+        ValueError: if the recipe would reach nothing, or if ``ignore`` was given
+            but excludes nothing. Neither fails on its own -- a pattern that
+            matches no module quantizes the whole model quietly -- and both mean
+            the run will not measure what was asked for. Both are checked before
+            the model is touched, so a rejected call leaves it unmodified.
+    """
+    # A bare name is matched against the module name exactly, which would select
+    # only the container -- not a Conv2d or Linear, and so exclude nothing. Widen
+    # it to the subtree; re: patterns pass through.
+    patterns = [
+        p if p.startswith("re:") else f"re:{re.escape(p)}($|\\.)" for p in ignore
+    ]
+    recipe = _create_mx_recipe(mx_format, patterns)
+
+    # Match_named_modules is the matcher apply_quantization_config itself uses,
+    # so the scope is checked exactly, before the model is touched.
+    targets = [cls.__name__ for cls in _MX_BLOCK_AXIS]
+    total = [name for name, _ in match_named_modules(model, targets)]
+    kept = [name for name, _ in match_named_modules(model, targets, patterns)]
+    if not kept:
+        raise ValueError(
+            f"the {mx_format} recipe matched no {'/'.join(targets)} in "
+            f"{type(model).__name__}")
+    if ignore and len(kept) == len(total):
+        raise ValueError(
+            f"ignore={list(ignore)} excluded nothing: all {len(total)} module(s) "
+            f"would be quantized. A pattern is a module name, covering that "
+            f"module and everything under it, or 're:' plus a regex; names look "
+            f"like {total[0]!r}")
+
+    return _apply_recipe(model, recipe)

alto/modifiers/quantization/format_registry.py

@@ -3,38 +3,52 @@
 # SPDX-License-Identifier: MIT
 """Runtime patch that wires emulated formats into the standard quant path.
 
-Importing this module injects a real ``format`` field into
+Importing this module injects real ``format`` and ``block_axis`` fields into
 ``compressed_tensors.QuantizationArgs`` so recipe values like ``format: mx9``
-survive pydantic parsing and become readable via ``getattr(args, "format", None)``
-(by default unknown fields are silently dropped).
+survive pydantic parsing and become readable via
+``getattr(args, "format", None)`` (by default unknown fields are silently
+dropped).
+
+``block_axis`` is the tensor axis the packed formats group elements along, and
+defaults to the last one. A recipe only needs to set it where that is the wrong
+axis, such as MX on Conv2d weights, whose blocks run along the input channels
+rather than along the kernel width.
 
 The actual ``fake_quantize`` dispatch (``args.format == "mx9"`` -> mx9) lives in
 ``alto.models.patcher.ModelPatcher.patch_fake_quantize`` where the single wrap of
 ``compressed_tensors...forward.fake_quantize`` already happens.
 
 ``inject_format_field()`` is called at the top of this package's ``__init__`` (before
-``QuantizationModifier`` is imported) so the field exists before the modifier
+``QuantizationModifier`` is imported) so the fields exist before the modifier
 compiles its nested ``QuantizationScheme`` schema.
 """
 
 from typing import Optional
 
 _FORMAT_FIELD_INJECTED = False
 
+DEFAULT_BLOCK_AXIS = -1
+
 
 def inject_format_field() -> None:
-    """Add ``format: Optional[str] = None`` to ``QuantizationArgs`` (idempotent)."""
+    """Add the ALTO fields to ``QuantizationArgs`` (idempotent)."""
     global _FORMAT_FIELD_INJECTED
     if _FORMAT_FIELD_INJECTED:
         return
 
     from pydantic.fields import FieldInfo
     from compressed_tensors.quantization import QuantizationArgs, QuantizationConfig, QuantizationScheme
 
-    if "format" not in QuantizationArgs.model_fields:
-        QuantizationArgs.model_fields["format"] = FieldInfo(
-            annotation=Optional[str], default=None
-        )
+    added = False
+    for name, field in (
+        ("format", FieldInfo(annotation=Optional[str], default=None)),
+        ("block_axis", FieldInfo(annotation=int, default=DEFAULT_BLOCK_AXIS)),
+    ):
+        if name not in QuantizationArgs.model_fields:
+            QuantizationArgs.model_fields[name] = field
+            added = True
+
+    if added:
         QuantizationArgs.model_rebuild(force=True)
         # QuantizationArgs is nested inside these models. Rebuild them as well so
         # recipe dictionaries with weights/input_activations.format are accepted