Refactor quantization into a shared extension/llm/export/quant package (#21044) (#21044)

Summary:

Consolidates the quantization/packing logic that was previously
duplicated
across per-model directories into a single reusable package under
extension/llm/export/quant, and introduces a model-free checkpoint
loader.

Key changes:
- New extension/llm/export/quant package exposing a stable API:
- convert.py: Convert, to_exportable/to_default, fuse_along_output,
identity
- quantize.py: quantize_model/quantize_weight/quantize_stream +
dequantize
  - recipe.py: QuantConfig / QuantRecipe / QuantRule
- New extension/llm/export/load.py: model-free (iter_checkpoint) and
model-aware (load_checkpoint) readers that stream GGUF / safetensors
weights,
remap names to FQNs, and wrap quantized subclasses in portable
Exportable*
form. Backend-specific layout (CUDA coalesced int4, MLX gather buffers)
is a
  separate terminal pass.
- Extended int4.py / gguf.py to plug into the shared conversion path.
- Removed the old per-model examples/models/gemma4_31b/quant/ packers
  (pack.py, pack_cuda.py, pack_mlx.py) and moved packing into
  examples/models/gemma4_31b/cuda_packers.py, wiring gemma4_31b and
  qwen3_5_moe onto the shared package.
- Added unit tests for the new package and the gemma4_31b
packers/pipelines.

Reviewed By: Gasoonjia

Differential Revision: D112867629

Author: metascroy

.github/workflows/cuda.yml

@@ -211,7 +211,7 @@ jobs:
 
         # Run Gemma 4 31B tests (quant unit tests + pipeline integration tests)
         pip install gguf
-        python -m pytest examples/models/gemma4_31b/quant/tests/ examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts="
+        python -m pytest examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts="
 
   unittest-cuda-runtime:
     name: unittest-cuda-runtime

.github/workflows/mlx.yml

@@ -85,7 +85,6 @@ jobs:
           backends/mlx/test/test_serialization_dedup.py \
           backends/mlx/test/test_slot_recycling.py \
           backends/mlx/test/test_sample.py \
-          examples/models/gemma4_31b/quant/tests/test_pack_mlx.py \
           examples/models/gemma4_31b/tests/test_mlx_pipeline.py \
           -v
         echo "::endgroup::"

backends/cuda/coalesced_int4_tensor.py

@@ -46,14 +46,13 @@
 The coalesced [N, n_groups] layout is exactly what the W4A8 dp4a matvec kernel
 (``executorch_cuda::int4_plain_mm`` / ``int4_plain_mm.cuh``) reads row-for-row
 with qdata, so the exported decode graph carries no per-step transpose. The
-transpose (and the uint8 re-encoding) is owned by :meth:`from_int4_tensor` so it
-is baked into the serialized weight constant once at pack time.
+is owned by :meth:`from_exportable_int4_tensor` so
+it is baked into the serialized weight constant once at pack time.
 """
 
 from typing import List, Optional, Tuple
 
 import torch
-from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor
 from torchao.utils import TorchAOBaseTensor
 
 __all__ = [
@@ -113,10 +112,10 @@ def _unpack_nibble_qdata(qdata: torch.Tensor, N: int, K: int) -> torch.Tensor:
 class CudaCoalescedInt4Tensor(TorchAOBaseTensor):
     """INT4 weight, uint8 scale + per-256 fp16 step, uint8 zero + per-256 fp16 step.
 
-    ExecuTorch-internal; see the module docstring. Mirrors torchao
+        ExecuTorch-internal; see the module docstring. Mirrors torchao
     ``Int4Tensor``'s data/attribute layout (so the common tensor utilities and
-    serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose and
-    the uint8 re-encoding via :meth:`from_int4_tensor`.
+        serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose and
+        the uint8 re-encoding via :meth:`from_exportable_int4_tensor`.
     """
 
     tensor_data_names = [
@@ -181,31 +180,34 @@ def _quantization_type(self):
         return s
 
     @classmethod
-    def from_int4_tensor(cls, t: Int4Tensor) -> "CudaCoalescedInt4Tensor":
-        """Build a coalesced tensor from a torchao ``Int4Tensor``.
-
-        Owns the transpose AND the uint8 re-encoding: torchao stores
-        scale/zero_point as (n_groups, N) bf16. The CUDA decode kernel reads the
-        (N, n_groups) uint8 scale/zero *codes* plus per-256-super-block
-        (N, K/256) fp16 ``scale_step`` / ``zero_point_step`` (scale = scale_code *
-        scale_step[:, g//8], zero = zero_code * zero_point_step[:, g//8]). The
-        transpose + encode here is baked into the serialized weight constant so
-        the exported decode graph has no per-step transpose/clone.
+    def from_exportable_int4_tensor(
+        cls, t: "ExportableInt4Tensor"  # noqa: F821
+    ) -> "CudaCoalescedInt4Tensor":
+        """Build a coalesced tensor from an ``ExportableInt4Tensor``.
+
+        Owns the [n_groups, N] -> [N, n_groups] transpose and the uint8
+        re-encoding, baked into the serialized weight constant at pack time.
+        torchao stores scale/zero_point as (n_groups, N) bf16; the CUDA decode
+        kernel reads (N, n_groups) uint8 scale/zero *codes* plus per-256-super-
+        block (N, K/256) fp16 ``scale_step`` / ``zero_point_step`` (scale =
+        scale_code * scale_step[:, g//8], zero = zero_code * zero_point_step[:,
+        g//8]), so the exported decode graph has no per-step transpose/clone. An
+        ``ExportableInt4Tensor`` (e.g. a decoded GGUF Q4_K) stores ``group_size``
+        as a scalar rather than a ``block_size`` list and carries no activation
+        quantization, so ``block_size`` is reconstructed as ``[1, group_size]``.
         """
-        scale_codes, scale_step = _encode_uint8_per_super(t.scale, t.block_size[-1])
+        scale_codes, scale_step = _encode_uint8_per_super(t.scale, t.group_size)
         zero_codes, zero_point_step = _encode_uint8_per_super(
-            t.zero_point, t.block_size[-1]
+            t.zero_point, t.group_size
         )
         return cls(
             t.qdata,
             scale_codes,
             scale_step,
             zero_codes,
             zero_point_step,
-            t.block_size,
+            [1, t.group_size],
             t.shape,
-            t.act_pre_scale,
-            t.activation_dtype,
         )
 
     def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor:

backends/cuda/runtime/shims/tests/gen_plain_mm_test_vectors.py

@@ -23,8 +23,9 @@
   1. torch.manual_seed(case.seed) on CPU.
   2. Draw a random bf16 weight ``[N, K]`` then activation ``[M, K]`` (weight
      first, then activation: fixed order is part of the seed contract).
-  3. quantize_weight(..., bits=4, min_max, asymmetric) -> torchao Int4Tensor.
-  4. CudaCoalescedInt4Tensor.from_int4_tensor(...) -> qdata [N, K/2] uint8,
+  3. quantize_weight(..., bits=4, min_max, asymmetric) -> torchao Int4Tensor,
+     wrapped as the canonical ExportableInt4Tensor.
+  4. CudaCoalescedInt4Tensor.from_exportable_int4_tensor(...) -> qdata [N, K/2] uint8,
      scale codes [N, K/gs] uint8, scale_step [N, K/256] fp16, zero_point codes
      [N, K/gs] uint8, zero_point_step [N, K/256] fp16.
   5. expected = F.linear(A, tensor.dequantize(bf16)).
@@ -157,8 +158,9 @@ def _i8(t: torch.Tensor) -> List[int]:
 def build_int4(case: Case) -> Dict[str, tuple]:
     """Return {array_name: (ctype, [ints])} for one INT4 case (CPU only)."""
     from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor
-    from executorch.examples.models.gemma4_31b.quant.quantize import quantize_weight
-    from executorch.examples.models.gemma4_31b.quant.recipe import QuantConfig
+    from executorch.extension.llm.export.int4 import ExportableInt4Tensor
+    from executorch.extension.llm.export.quant.quantize import quantize_weight
+    from executorch.extension.llm.export.quant.recipe import QuantConfig
 
     torch.manual_seed(case.seed)
     # Weight first, then activation: fixed order is part of the seed contract.
@@ -167,7 +169,9 @@ def build_int4(case: Case) -> Dict[str, tuple]:
 
     config = QuantConfig(bits=4, group_size=case.gs, symmetric=False, method="min_max")
     int4 = quantize_weight(w, config)
-    c = CudaCoalescedInt4Tensor.from_int4_tensor(int4)
+    c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(
+        ExportableInt4Tensor.from_int4_tensor(int4)
+    )
 
     # bf16 dequant @ F.linear reference (kernel adds activation-quant noise).
     w_deq = c.dequantize(torch.bfloat16)

backends/cuda/tests/test_int4_dispatch.py

@@ -34,12 +34,13 @@
 import torch.nn.functional as F
 from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor
 from executorch.backends.cuda.quantize_op_dispatch.int4_dispatch import _dequant_matmul
-from executorch.examples.models.gemma4_31b.quant.pack_cuda import pack_linear_for_cuda
-from executorch.examples.models.gemma4_31b.quant.quantize import (
+from executorch.examples.models.gemma4_31b.cuda_packers import pack_linear_for_cuda
+from executorch.extension.llm.export.int4 import ExportableInt4Tensor
+from executorch.extension.llm.export.quant.quantize import (
     dequantize_weight,
     quantize_weight,
 )
-from executorch.examples.models.gemma4_31b.quant.recipe import QuantConfig
+from executorch.extension.llm.export.quant.recipe import QuantConfig
 
 
 def _require_cuda(tc: unittest.TestCase) -> None:
@@ -48,22 +49,26 @@ def _require_cuda(tc: unittest.TestCase) -> None:
 
 
 def _make_int4_linear(N, K, group_size=128, symmetric=False, bias=False):
-    """Build an nn.Linear with Int4Tensor weight and return (module, bf16_ref_weight).
+    """Build an nn.Linear with ExportableInt4Tensor weight + bf16 ref weight.
 
-    The bf16 reference is the original unquantized weight, so tests can
-    measure quantization error against the true value.
+    Mirrors production: weights are converted to ExportableInt4Tensor (the
+    canonical portable int4 form) before packing for CUDA. The bf16 reference
+    is the original unquantized weight, so tests can measure quantization
+    error against the true value.
     """
     w_bf16 = torch.randn(N, K, dtype=torch.bfloat16)
     config = QuantConfig(
         bits=4, group_size=group_size, symmetric=symmetric, method="min_max"
     )
-    int4_w = quantize_weight(w_bf16, config)
+    exportable_w = ExportableInt4Tensor.from_int4_tensor(
+        quantize_weight(w_bf16, config)
+    )
 
     # device="cuda" so the random init draws from the CUDA RNG to match the
     # same random weight as regular int4 dispatch and fit the same numerical
     # error tolerance.
     module = nn.Linear(K, N, bias=bias, dtype=torch.bfloat16, device="cuda")
-    pack_linear_for_cuda(module, {"weight": int4_w})
+    pack_linear_for_cuda(module, {"weight": exportable_w})
     module.cuda()
     return module, w_bf16.cuda()
 
@@ -184,9 +189,11 @@ def _check(self, out, ref, tol=0.15):
     def test_to_cuda(self):
         w_bf16 = torch.randn(256, 512, dtype=torch.bfloat16)
         config = QuantConfig(bits=4, group_size=128, symmetric=False, method="min_max")
-        int4_w = quantize_weight(w_bf16, config)
+        exportable_w = ExportableInt4Tensor.from_int4_tensor(
+            quantize_weight(w_bf16, config)
+        )
         module = nn.Linear(512, 256, bias=False)
-        pack_linear_for_cuda(module, {"weight": int4_w})
+        pack_linear_for_cuda(module, {"weight": exportable_w})
         module = module.to("cuda")
         x = torch.randn(1, 512, dtype=torch.bfloat16, device="cuda")
         self._check(module(x), F.linear(x, w_bf16.cuda()))
@@ -228,6 +235,12 @@ def _make_int4_tensor(N, K, group_size=128, symmetric=False):
     return quantize_weight(w, config), w
 
 
+def _make_exportable_int4_tensor(N, K, group_size=128, symmetric=False):
+    """Build an ``ExportableInt4Tensor`` (canonical portable int4) + bf16 ref."""
+    t, w = _make_int4_tensor(N, K, group_size=group_size, symmetric=symmetric)
+    return ExportableInt4Tensor.from_int4_tensor(t), w
+
+
 @contextlib.contextmanager
 def _record_int4_plain_mm():
     """Record calls to the decode custom op without needing a GPU.
@@ -278,8 +291,8 @@ def test_stock_int4tensor_does_not_route_to_int4_plain_mm(self):
 
     def test_coalesced_tensor_routes_to_int4_plain_mm(self):
         """CudaCoalescedInt4Tensor with M<=4 routes to the decode custom op."""
-        t, _ = _make_int4_tensor(16, 256, group_size=32)
-        c = CudaCoalescedInt4Tensor.from_int4_tensor(t)
+        t, _ = _make_exportable_int4_tensor(16, 256, group_size=32)
+        c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t)
         x = torch.randn(1, 256, dtype=torch.bfloat16)  # M=1 (decode regime)
         with _record_int4_plain_mm() as calls:
             out = F.linear(x, c)
@@ -288,8 +301,8 @@ def test_coalesced_tensor_routes_to_int4_plain_mm(self):
 
     def test_coalesced_tensor_prefill_uses_dequant(self):
         """M>4 uses inline dequant (no custom op) and is numerically correct."""
-        t, _ = _make_int4_tensor(16, 256, group_size=32)
-        c = CudaCoalescedInt4Tensor.from_int4_tensor(t)
+        t, _ = _make_exportable_int4_tensor(16, 256, group_size=32)
+        c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t)
         x = torch.randn(8, 256, dtype=torch.bfloat16)  # M=8 > 4 (prefill regime)
         with _record_int4_plain_mm() as calls:
             out = F.linear(x, c)
@@ -312,10 +325,10 @@ def test_square_shape_not_misrouted(self):
                 F.linear(x, t)
         self.assertEqual(calls, [])
 
-    def test_from_int4_tensor_transpose_correct(self):
-        """from_int4_tensor owns the (n_groups, N) -> (N, n_groups) transpose."""
-        t, _ = _make_int4_tensor(24, 256, group_size=64)
-        c = CudaCoalescedInt4Tensor.from_int4_tensor(t)
+    def test_from_exportable_int4_tensor_transpose_correct(self):
+        """from_exportable_int4_tensor owns the (n_groups, N) -> (N, n_groups) transpose."""
+        t, _ = _make_exportable_int4_tensor(24, 256, group_size=64)
+        c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t)
         n_groups = 256 // 64
         self.assertEqual(tuple(t.scale.shape), (n_groups, 24))  # torchao layout
         self.assertEqual(tuple(c.scale.shape), (24, n_groups))  # coalesced layout