From b7dbdce8bf615fb3b454bc73c427a9e111a70278 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Sat, 15 Aug 2026 07:44:29 -0700
Subject: [PATCH 01/10] WIP: add opaque quantized weight export state

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 493 +++++++++++++++++-
 .../nn/modules/tensor_quantizer.py            |   7 +-
 .../export/test_quantized_weight_state.py     | 171 ++++++
 3 files changed, 665 insertions(+), 6 deletions(-)
 create mode 100644 tests/unit/torch/export/test_quantized_weight_state.py

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index c86af3aa9f5..49b9e8b2d56 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -16,8 +16,10 @@
 """Utils for quantization including scaling factors adjustments."""
 
 import logging
-from collections import defaultdict
-from collections.abc import Generator
+import re
+from collections import OrderedDict, defaultdict
+from collections.abc import Generator, Iterable, Mapping, Sequence
+from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
 from warnings import warn
@@ -777,6 +779,11 @@ def process_layer_quant_config(layer_config_dict):
     # If we have more than one quantization format, infer MIXED_PRECISION
     if len(quantization_formats) > 1:
         per_layer_config["quant_algo"] = "MIXED_PRECISION"
+        per_layer_config["exclude_modules"] = sorted(
+            _prefix_wildcard_summarize_exclude_modules(
+                exclude_modules, per_layer_config["quantized_layers"].keys()
+            )
+        )
     elif len(quantization_formats) == 1 and quantization_config is not None:
         per_layer_config.update(quantization_config)
         per_layer_config["exclude_modules"] = sorted(
@@ -938,6 +945,488 @@ def to_quantized_weight(
     raise NotImplementedError(f"quantization format {quantization} not supported")
 
 
+_FUNCTIONAL_EXPORT_FORMATS = {
+    QUANTIZATION_FP8,
+    QUANTIZATION_FP8_PB_WO,
+    QUANTIZATION_FP8_PC_PT,
+    QUANTIZATION_MXFP8,
+    QUANTIZATION_NVFP4,
+    QUANTIZATION_W4A16_NVFP4,
+}
+_NVFP4_EXPORT_FORMATS = {QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4}
+
+
+@dataclass(frozen=True)
+class _ExportStateTensor:
+    name: str
+    value: torch.Tensor
+    axes: tuple[int, ...] = ()
+    block_sizes: tuple[int, ...] = ()
+
+
+@dataclass(frozen=True)
+class QuantizedWeightExportState:
+    """Opaque state required to export one logical quantized weight."""
+
+    quantization_format: str
+    block_size: int
+    weight_shape: tuple[int, ...]
+    tensors: tuple[_ExportStateTensor, ...]
+    packing_permutation: tuple[int, ...]
+    static_nvfp4: bool = False
+    four_over_six: bool = False
+
+
+def _same_storage(left: object, right: object) -> bool:
+    if left is right:
+        return True
+    if not isinstance(left, torch.Tensor) or not isinstance(right, torch.Tensor):
+        return False
+    if left.device.type == "meta" or right.device.type == "meta":
+        return False
+    return (
+        left.device == right.device
+        and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
+        and left.storage_offset() == right.storage_offset()
+    )
+
+
+def _resolve_weight_quantizer(
+    module: nn.Module, weight_name: str
+) -> tuple[torch.Tensor, TensorQuantizer | SequentialQuantizer]:
+    weight = getattr(module, weight_name)
+    iter_weights = getattr(module, "iter_weights_for_calibration", None)
+    if iter_weights is not None:
+        for candidate, quantizer in iter_weights():
+            if _same_storage(candidate, weight):
+                return candidate, quantizer
+
+    quantizer = representative_weight_quantizer(module, weight_name)
+    if quantizer is None and weight_name.startswith("weight"):
+        quantizer = representative_weight_quantizer(module)
+    if quantizer is None:
+        raise RuntimeError(f"Missing weight quantizer for {weight_name!r}")
+    return weight, quantizer
+
+
+def _packing_permutation(weight: torch.Tensor, quantized_view: torch.Tensor) -> tuple[int, ...]:
+    ndim = weight.ndim
+    identity = tuple(range(ndim))
+    if tuple(quantized_view.shape) == tuple(weight.shape):
+        return identity
+    if ndim >= 2:
+        transposed = (*weight.shape[:-2], weight.shape[-1], weight.shape[-2])
+        if tuple(quantized_view.shape) == transposed:
+            return (*range(ndim - 2), ndim - 1, ndim - 2)
+    raise NotImplementedError(
+        f"Unsupported quantized weight view {tuple(quantized_view.shape)} for {tuple(weight.shape)}"
+    )
+
+
+def _state_tensor(
+    name: str,
+    value: torch.Tensor,
+    packed_shape: tuple[int, ...],
+    *,
+    block_sizes: tuple[int, ...] | None = None,
+) -> _ExportStateTensor:
+    value = value.detach().clone()
+    if value.numel() == 1:
+        return _ExportStateTensor(name, value.reshape(()))
+    if value.ndim > len(packed_shape):
+        raise NotImplementedError(
+            f"Cannot relate {name} shape {tuple(value.shape)} to weight shape {packed_shape}"
+        )
+
+    axes = tuple(range(value.ndim))
+    if block_sizes is None:
+        inferred = []
+        for axis, size in enumerate(value.shape):
+            if packed_shape[axis] % size:
+                raise NotImplementedError(
+                    f"Cannot relate {name} shape {tuple(value.shape)} to weight shape {packed_shape}"
+                )
+            inferred.append(packed_shape[axis] // size)
+        block_sizes = tuple(inferred)
+    if len(block_sizes) != value.ndim:
+        raise ValueError(f"Invalid block layout for {name}: {block_sizes}")
+    return _ExportStateTensor(name, value, axes, block_sizes)
+
+
+def _input_quantizer(module: nn.Module, weight_name: str):
+    quantizer = getattr(module, quantizer_attr_names(weight_name).input_quantizer, None)
+    if quantizer is None:
+        quantizer = getattr(module, "input_quantizer", None)
+    return quantizer
+
+
+def capture_quantized_weight_export_state(
+    module: nn.Module,
+    weight_name: str = "weight",
+) -> QuantizedWeightExportState:
+    """Capture detached export state without mutating the quantized module."""
+    weight = getattr(module, weight_name)
+    if isinstance(weight, QTensorWrapper):
+        raise NotImplementedError("Functional export requires an uncompressed source weight")
+
+    quantized_view, weight_quantizer = _resolve_weight_quantizer(module, weight_name)
+    quantization_format = get_quantization_format(module)
+    if quantization_format not in _FUNCTIONAL_EXPORT_FORMATS:
+        raise NotImplementedError(f"Functional export does not support {quantization_format!r}")
+    if isinstance(weight_quantizer, SequentialQuantizer):
+        weight_quantizer = weight_quantizer[0]
+    if not weight_quantizer.is_enabled:
+        raise RuntimeError(f"Weight quantizer for {weight_name!r} is disabled")
+
+    permutation = _packing_permutation(weight, quantized_view)
+    packed_shape = tuple(weight.shape[axis] for axis in permutation)
+    block_config = getattr(weight_quantizer, "block_sizes", None) or {}
+    block_size = int(block_config.get(-1, 0)) if isinstance(block_config, dict) else 0
+    tensors = []
+    static_nvfp4 = (
+        quantization_format in _NVFP4_EXPORT_FORMATS
+        and NVFP4QTensor._is_static_quantizer(weight_quantizer)
+    )
+
+    if static_nvfp4:
+        if not block_size or packed_shape[-1] % block_size:
+            raise RuntimeError(f"Invalid static NVFP4 block size for {weight_name!r}")
+        per_block_amax = weight_quantizer.amax
+        global_amax = NVFP4QTensor._get_static_global_amax(weight_quantizer)
+        if per_block_amax is None or global_amax is None:
+            raise RuntimeError(f"Missing calibrated static NVFP4 state for {weight_name!r}")
+        block_shape = (*packed_shape[:-1], packed_shape[-1] // block_size)
+        tensors.append(
+            _state_tensor(
+                "weight_block_amax",
+                per_block_amax.reshape(block_shape),
+                packed_shape,
+                block_sizes=(1,) * (len(packed_shape) - 1) + (block_size,),
+            )
+        )
+        tensors.append(_state_tensor("weight_global_amax", global_amax, packed_shape))
+    elif quantization_format in _NVFP4_EXPORT_FORMATS:
+        weight_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
+        tensors.append(_state_tensor("weight_scale_2", weight_scale_2, packed_shape))
+    elif quantization_format == QUANTIZATION_MXFP8:
+        weight_scale = MXFP8QTensor.get_weights_scaling_factor_from_quantizer(
+            quantized_view, weight_quantizer
+        )
+        tensors.append(
+            _state_tensor(
+                "weight_scale",
+                weight_scale,
+                packed_shape,
+                block_sizes=(1,) * (len(packed_shape) - 1) + (block_size,),
+            )
+        )
+    else:
+        weight_scale = get_scaling_factor(weight_quantizer)
+        if weight_scale is None:
+            raise RuntimeError(f"Missing calibrated weight scale for {weight_name!r}")
+        tensors.append(_state_tensor("weight_scale", weight_scale, packed_shape))
+
+    input_quantizer = _input_quantizer(module, weight_name)
+    if input_quantizer is not None and input_quantizer.is_enabled:
+        if input_quantizer.export_amax() is None:
+            raise RuntimeError(f"Missing calibrated input scale for {weight_name!r}")
+        input_scale = (
+            NVFP4QTensor.get_activation_scaling_factor(input_quantizer)
+            if quantization_format == QUANTIZATION_NVFP4
+            else get_scaling_factor(input_quantizer)
+        )
+        tensors.append(_state_tensor("input_scale", input_scale, packed_shape))
+
+    return QuantizedWeightExportState(
+        quantization_format=quantization_format,
+        block_size=block_size,
+        weight_shape=tuple(weight.shape),
+        tensors=tuple(tensors),
+        packing_permutation=permutation,
+        static_nvfp4=static_nvfp4,
+        four_over_six=bool(block_config.get("four_over_six", False)),
+    )
+
+
+def _normalize_weight_dim(weight_dim: int, ndim: int) -> int:
+    if not -ndim <= weight_dim < ndim:
+        raise IndexError(f"Weight dimension {weight_dim} is invalid for a rank-{ndim} weight")
+    return weight_dim % ndim
+
+
+def _merge_state_tensors(
+    records: Sequence[_ExportStateTensor], packed_dim: int
+) -> _ExportStateTensor:
+    reference = records[0]
+    if any(
+        record.name != reference.name
+        or record.axes != reference.axes
+        or record.block_sizes != reference.block_sizes
+        for record in records[1:]
+    ):
+        raise ValueError("Quantized weight shards have incompatible export state")
+    if packed_dim in reference.axes:
+        tensor_dim = reference.axes.index(packed_dim)
+        value = torch.cat([record.value for record in records], dim=tensor_dim)
+    else:
+        if len({tuple(record.value.shape) for record in records}) != 1:
+            raise ValueError("Replicated export tensors have incompatible shapes")
+        value = torch.stack([record.value for record in records]).amax(dim=0)
+    return _ExportStateTensor(reference.name, value, reference.axes, reference.block_sizes)
+
+
+def merge_quantized_weight_export_states(
+    states: Sequence[QuantizedWeightExportState],
+    weight_dim: int,
+) -> QuantizedWeightExportState:
+    """Merge state for logical weight shards concatenated along ``weight_dim``."""
+    if not states:
+        raise ValueError("At least one quantized weight state is required")
+    reference = states[0]
+    ndim = len(reference.weight_shape)
+    weight_dim = _normalize_weight_dim(weight_dim, ndim)
+    if any(
+        state.quantization_format != reference.quantization_format
+        or state.block_size != reference.block_size
+        or state.packing_permutation != reference.packing_permutation
+        or state.static_nvfp4 != reference.static_nvfp4
+        or state.four_over_six != reference.four_over_six
+        or len(state.tensors) != len(reference.tensors)
+        or len(state.weight_shape) != ndim
+        or any(
+            size != reference.weight_shape[axis]
+            for axis, size in enumerate(state.weight_shape)
+            if axis != weight_dim
+        )
+        for state in states[1:]
+    ):
+        raise ValueError("Quantized weight shards are incompatible")
+
+    shape = list(reference.weight_shape)
+    shape[weight_dim] = sum(state.weight_shape[weight_dim] for state in states)
+    packed_dim = reference.packing_permutation.index(weight_dim)
+    tensors = tuple(
+        _merge_state_tensors([state.tensors[index] for state in states], packed_dim)
+        for index in range(len(reference.tensors))
+    )
+    return QuantizedWeightExportState(
+        reference.quantization_format,
+        reference.block_size,
+        tuple(shape),
+        tensors,
+        reference.packing_permutation,
+        reference.static_nvfp4,
+        reference.four_over_six,
+    )
+
+
+def _select_state_tensor(
+    record: _ExportStateTensor,
+    packed_dim: int,
+    indices: torch.Tensor,
+) -> _ExportStateTensor:
+    if packed_dim not in record.axes:
+        return record
+    tensor_dim = record.axes.index(packed_dim)
+    block_size = record.block_sizes[tensor_dim]
+    if block_size == 1:
+        selected = indices
+    else:
+        selected = torch.div(indices, block_size, rounding_mode="floor")
+        unique, counts = selected.unique_consecutive(return_counts=True)
+        if not torch.all(counts == block_size):
+            raise ValueError("Weight selection must preserve complete quantization blocks")
+        selected = unique
+    value = record.value.index_select(tensor_dim, selected.to(record.value.device))
+    return _ExportStateTensor(record.name, value, record.axes, record.block_sizes)
+
+
+def select_quantized_weight_export_state(
+    state: QuantizedWeightExportState,
+    weight_dim: int,
+    indices: Iterable[int] | torch.Tensor,
+) -> QuantizedWeightExportState:
+    """Select logical weight indices and their corresponding opaque export state."""
+    ndim = len(state.weight_shape)
+    weight_dim = _normalize_weight_dim(weight_dim, ndim)
+    indices = torch.as_tensor(list(indices) if not isinstance(indices, torch.Tensor) else indices)
+    indices = indices.to(dtype=torch.long, device="cpu").reshape(-1)
+    if indices.numel() == 0:
+        raise ValueError("Weight selection cannot be empty")
+    if indices.min() < 0 or indices.max() >= state.weight_shape[weight_dim]:
+        raise IndexError("Weight selection is out of range")
+
+    shape = list(state.weight_shape)
+    shape[weight_dim] = indices.numel()
+    packed_dim = state.packing_permutation.index(weight_dim)
+    return QuantizedWeightExportState(
+        state.quantization_format,
+        state.block_size,
+        tuple(shape),
+        tuple(_select_state_tensor(record, packed_dim, indices) for record in state.tensors),
+        state.packing_permutation,
+        state.static_nvfp4,
+        state.four_over_six,
+    )
+
+
+def permute_quantized_weight_export_state(
+    state: QuantizedWeightExportState,
+    dims: Sequence[int],
+) -> QuantizedWeightExportState:
+    """Permute logical weight dimensions while retaining quantizer packing axes."""
+    dims = tuple(dims)
+    ndim = len(state.weight_shape)
+    if sorted(dims) != list(range(ndim)):
+        raise ValueError(f"Invalid permutation {dims} for a rank-{ndim} weight")
+    inverse = {old_dim: new_dim for new_dim, old_dim in enumerate(dims)}
+    return QuantizedWeightExportState(
+        state.quantization_format,
+        state.block_size,
+        tuple(state.weight_shape[dim] for dim in dims),
+        state.tensors,
+        tuple(inverse[dim] for dim in state.packing_permutation),
+        state.static_nvfp4,
+        state.four_over_six,
+    )
+
+
+def _restore_packing_permutation(tensor: torch.Tensor, permutation: tuple[int, ...]):
+    if permutation == tuple(range(len(permutation))):
+        return tensor
+    return tensor.permute(tuple(permutation.index(dim) for dim in range(len(permutation))))
+
+
+def _restore_state_tensor(
+    record: _ExportStateTensor,
+    permutation: tuple[int, ...],
+) -> torch.Tensor:
+    if not record.axes:
+        return record.value
+    logical_axes = tuple(permutation[axis] for axis in record.axes)
+    order = tuple(sorted(range(len(logical_axes)), key=logical_axes.__getitem__))
+    if order == tuple(range(len(order))):
+        return record.value
+    return record.value.permute(order)
+
+
+def export_quantized_weight_tensors(
+    weight: torch.Tensor,
+    state: QuantizedWeightExportState,
+    dtype: torch.dtype,
+    weight_name: str = "weight",
+) -> OrderedDict[str, torch.Tensor]:
+    """Pack a logical weight into canonical ModelOpt checkpoint tensors."""
+    if tuple(weight.shape) != state.weight_shape:
+        raise ValueError(
+            f"Weight shape {tuple(weight.shape)} does not match export state {state.weight_shape}"
+        )
+    packed_weight = weight.permute(state.packing_permutation)
+    records = {record.name: record for record in state.tensors}
+
+    weight_scale_2 = None
+    if state.static_nvfp4:
+        block_amax = records["weight_block_amax"].value
+        global_amax = records["weight_global_amax"].value
+        quantizer = SimpleNamespace(
+            block_sizes={
+                -1: state.block_size,
+                "scale_bits": (4, 3),
+                "four_over_six": state.four_over_six,
+            },
+            _amax=block_amax,
+            _global_amax=global_amax,
+            global_amax=global_amax,
+        )
+        weight_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(quantizer)
+        weight_scale = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(
+            quantizer, packed_weight, weight_scale_2
+        )[0]
+    elif state.quantization_format in _NVFP4_EXPORT_FORMATS:
+        weight_scale_2 = records["weight_scale_2"].value
+        weight_scale = NVFP4QTensor.get_weights_scaling_factor(
+            packed_weight,
+            state.block_size,
+            weights_scaling_factor_2=weight_scale_2.to(packed_weight.device),
+        )[0]
+    else:
+        weight_scale = records["weight_scale"].value
+
+    quantized_weight = to_quantized_weight(
+        packed_weight.to(dtype),
+        weight_scale,
+        state.quantization_format,
+        weight_scale_2,
+        state.block_size,
+    )
+    attrs = quantizer_attr_names(weight_name)
+    output = OrderedDict(
+        ((weight_name, _restore_packing_permutation(quantized_weight, state.packing_permutation)),)
+    )
+
+    weight_scale_record = _state_tensor(
+        "weight_scale",
+        weight_scale,
+        tuple(packed_weight.shape),
+        block_sizes=(1,) * (packed_weight.ndim - 1) + (state.block_size,)
+        if state.quantization_format in _NVFP4_EXPORT_FORMATS
+        else None,
+    )
+    output[attrs.weight_scale] = _restore_state_tensor(
+        weight_scale_record, state.packing_permutation
+    )
+    if weight_scale_2 is not None:
+        scale_2_record = records.get("weight_scale_2") or records["weight_global_amax"]
+        output[attrs.weight_scale_2] = _restore_state_tensor(
+            _ExportStateTensor(
+                scale_2_record.name,
+                weight_scale_2.squeeze(),
+                scale_2_record.axes,
+                scale_2_record.block_sizes,
+            ),
+            state.packing_permutation,
+        )
+    if "input_scale" in records:
+        output[attrs.input_scale] = _restore_state_tensor(
+            records["input_scale"], state.packing_permutation
+        ).squeeze()
+    return output
+
+
+def _quantized_layer_name(weight_name: str) -> str:
+    name = weight_name.removesuffix(".weight")
+    return re.sub(r"(\.experts)\.\d+(?=\.)", r"\1", name)
+
+
+def build_hf_quantization_config(
+    named_states: Mapping[str, QuantizedWeightExportState | None]
+    | Iterable[tuple[str, QuantizedWeightExportState | None]],
+) -> dict[str, Any]:
+    """Build canonical ModelOpt HF configuration from named opaque states."""
+    layers: dict[str, tuple[str, int] | None] = {}
+    for weight_name, state in dict(named_states).items():
+        layer_name = _quantized_layer_name(weight_name)
+        value = None if state is None else (state.quantization_format, state.block_size)
+        previous = layers.setdefault(layer_name, value)
+        if previous != value:
+            raise ValueError(f"Inconsistent quantization state for {layer_name}")
+
+    layer_config = {}
+    for layer_name, value in layers.items():
+        layer_config[f"{layer_name}.quantization"] = (
+            QUANTIZATION_NONE if value is None else value[0]
+        )
+        layer_config[f"{layer_name}.awq_block_size"] = 0 if value is None else value[1]
+    config = {
+        "producer": {"name": "modelopt", "version": __version__},
+        "quantization": process_layer_quant_config(layer_config),
+    }
+    config["quantization"].setdefault("kv_cache_quant_algo", QUANTIZATION_NONE)
+    from .convert_hf_config import convert_hf_quant_config_format
+
+    return convert_hf_quant_config_format(config)
+
+
 def from_quantized_weight(
     weight: torch.Tensor,
     weights_scaling_factor: torch.Tensor,
diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
index 18b97ac2774..e3e870e3c63 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -1092,10 +1092,9 @@ def export_amax(self) -> torch.Tensor | None:
         if self.amax is None:
             return None
 
-        if not hasattr(self, "_amax_shape_for_export"):
-            amax = self.amax
-        else:
-            amax = self.amax.reshape(self._amax_shape_for_export)
+        amax = self.amax.detach().clone()
+        if hasattr(self, "_amax_shape_for_export"):
+            amax = amax.reshape(self._amax_shape_for_export)
         amax[amax == 0] = self.maxbound
         amax = torch.nan_to_num(amax, nan=self.maxbound)
         clamp_min, clamp_max = torch.finfo(amax.dtype).tiny, torch.finfo(amax.dtype).max
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
new file mode 100644
index 00000000000..bb1999b113f
--- /dev/null
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+
+import pytest
+import torch
+import torch.nn as nn
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.export.quant_utils import (
+    build_hf_quantization_config,
+    capture_quantized_weight_export_state,
+    export_quantized_weight_tensors,
+    merge_quantized_weight_export_states,
+    permute_quantized_weight_export_state,
+    select_quantized_weight_export_state,
+)
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer
+
+
+def _fp8_linear() -> nn.Linear:
+    module = nn.Linear(4, 4, bias=False)
+    with torch.no_grad():
+        module.weight.copy_(torch.arange(16, dtype=torch.float32).reshape(4, 4) / 8 - 1)
+    return mtq.quantize(module, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m(torch.ones(2, 4)))
+
+
+def _static_w4a16_linear(
+    weight: torch.Tensor,
+    per_block_amax: torch.Tensor,
+    global_amax: torch.Tensor,
+) -> nn.Linear:
+    module = nn.Linear(weight.shape[1], weight.shape[0], bias=False)
+    module.weight.data.copy_(weight)
+    cfg = QuantizerAttributeConfig(
+        num_bits=(2, 1),
+        block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)},
+    )
+    quantizer = NVFP4StaticQuantizer(quant_attribute_cfg=cfg)
+    quantizer.amax = per_block_amax.clone()
+    quantizer.global_amax = global_amax.clone()
+    module.weight_quantizer = quantizer
+    module.input_quantizer = TensorQuantizer()
+    module.input_quantizer.disable()
+    return module
+
+
+def test_capture_does_not_modify_zero_amax():
+    module = _fp8_linear()
+    module.weight_quantizer._amax.zero_()
+    module.input_quantizer._amax.zero_()
+    weight_amax = module.weight_quantizer._amax.clone()
+    input_amax = module.input_quantizer._amax.clone()
+
+    capture_quantized_weight_export_state(module)
+
+    torch.testing.assert_close(module.weight_quantizer._amax, weight_amax)
+    torch.testing.assert_close(module.input_quantizer._amax, input_amax)
+
+
+def test_functional_fp8_export_is_repeatable_and_supports_permutation():
+    module = _fp8_linear()
+    original_weight = module.weight
+    original_buffers = {name: value.clone() for name, value in module.named_buffers()}
+    state = capture_quantized_weight_export_state(module)
+
+    base = export_quantized_weight_tensors(module.weight, state, torch.float32)
+    repeated = export_quantized_weight_tensors(module.weight, state, torch.float32)
+    transposed = export_quantized_weight_tensors(
+        module.weight.T,
+        permute_quantized_weight_export_state(state, (1, 0)),
+        torch.float32,
+    )
+
+    assert module.weight is original_weight
+    assert set(dict(module.named_buffers())) == set(original_buffers)
+    for name, value in module.named_buffers():
+        torch.testing.assert_close(value, original_buffers[name])
+    for name in base:
+        torch.testing.assert_close(base[name], repeated[name])
+    torch.testing.assert_close(transposed["weight"], base["weight"].T)
+
+
+def test_static_nvfp4_merge_recomputes_scales_from_merged_amax():
+    left_weight = torch.arange(32, dtype=torch.float32).reshape(2, 16) / 32
+    right_weight = torch.arange(32, 64, dtype=torch.float32).reshape(2, 16) / 16
+    left_amax = left_weight.abs().amax(dim=1, keepdim=True)
+    right_amax = right_weight.abs().amax(dim=1, keepdim=True)
+    left = _static_w4a16_linear(left_weight, left_amax, torch.tensor(1.0))
+    right = _static_w4a16_linear(right_weight, right_amax, torch.tensor(4.0))
+
+    state = merge_quantized_weight_export_states(
+        (
+            capture_quantized_weight_export_state(left),
+            capture_quantized_weight_export_state(right),
+        ),
+        0,
+    )
+    weight = torch.cat((left_weight, right_weight), dim=0)
+    actual = export_quantized_weight_tensors(weight, state, torch.float32)
+
+    reference = _static_w4a16_linear(
+        weight,
+        torch.cat((left_amax, right_amax), dim=0),
+        torch.tensor(4.0),
+    )
+    expected = export_quantized_weight_tensors(
+        reference.weight,
+        capture_quantized_weight_export_state(reference),
+        torch.float32,
+    )
+    for name in actual:
+        torch.testing.assert_close(actual[name], expected[name], rtol=0, atol=0)
+
+
+def test_static_nvfp4_noncontiguous_selection_tracks_block_amax():
+    weight = torch.arange(64, dtype=torch.float32).reshape(4, 16) / 32
+    per_block_amax = weight.abs().amax(dim=1, keepdim=True)
+    module = _static_w4a16_linear(weight, per_block_amax, torch.tensor(2.0))
+    indices = torch.tensor([0, 2])
+
+    state = select_quantized_weight_export_state(
+        capture_quantized_weight_export_state(module), 0, indices
+    )
+    actual = export_quantized_weight_tensors(weight.index_select(0, indices), state, torch.float32)
+
+    reference = _static_w4a16_linear(
+        weight.index_select(0, indices),
+        per_block_amax.index_select(0, indices),
+        torch.tensor(2.0),
+    )
+    expected = export_quantized_weight_tensors(
+        reference.weight,
+        capture_quantized_weight_export_state(reference),
+        torch.float32,
+    )
+    for name in actual:
+        torch.testing.assert_close(actual[name], expected[name], rtol=0, atol=0)
+
+
+def test_state_rejects_invalid_dimensions():
+    state = capture_quantized_weight_export_state(_fp8_linear())
+    with pytest.raises(IndexError, match="invalid"):
+        merge_quantized_weight_export_states((state, state), 2)
+    with pytest.raises(IndexError, match="invalid"):
+        select_quantized_weight_export_state(state, -3, (0,))
+
+
+def test_mixed_config_groups_moe_experts_by_projection_family():
+    fp8 = capture_quantized_weight_export_state(_fp8_linear())
+    weight = torch.ones(4, 16)
+    w4a16 = capture_quantized_weight_export_state(
+        _static_w4a16_linear(weight, torch.ones(4, 1), torch.tensor(1.0))
+    )
+    config = build_hf_quantization_config(
+        {
+            "model.layers.0.mlp.experts.0.gate_proj.weight": fp8,
+            "model.layers.0.mlp.experts.1.gate_proj.weight": fp8,
+            "model.layers.0.mlp.experts.0.down_proj.weight": w4a16,
+            "model.layers.0.mlp.experts.1.down_proj.weight": w4a16,
+            "lm_head.weight": None,
+        }
+    )
+
+    assert set(config["quantized_layers"]) == {
+        "model.layers.0.mlp.experts.gate_proj",
+        "model.layers.0.mlp.experts.down_proj",
+    }
+    assert "lm_head" in config["ignore"]

From d66e5c6eebeeeda9d4c5493f78f899dd33a8a26d Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Sat, 15 Aug 2026 07:47:26 -0700
Subject: [PATCH 02/10] WIP: use functional state in HF weight export

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    | 23 ++++++++
 tests/gpu/torch/export/test_export.py         |  1 +
 .../torch/export/test_export_weight_gpu.py    | 52 ++++++++++++++++++-
 3 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 77429b1cfaf..945fe346168 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -82,6 +82,7 @@
 from .model_config import (
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
+    QUANTIZATION_FP8_PB_WO,
     QUANTIZATION_FP8_PC_PT,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
@@ -100,6 +101,8 @@
     revert_weight_conversion_quant_aware,
 )
 from .quant_utils import (
+    capture_quantized_weight_export_state,
+    export_quantized_weight_tensors,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -607,6 +610,26 @@ def _export_quantized_weight(
         sub_module, quantizer_attrs.output_quantizer, None
     )
 
+    if not isinstance(weight, QTensorWrapper) and quantization_format in {
+        QUANTIZATION_FP8,
+        QUANTIZATION_FP8_PB_WO,
+        QUANTIZATION_FP8_PC_PT,
+        QUANTIZATION_MXFP8,
+        QUANTIZATION_NVFP4,
+        QUANTIZATION_W4A16_NVFP4,
+    }:
+        state = capture_quantized_weight_export_state(sub_module, weight_name)
+        exported = export_quantized_weight_tensors(weight, state, dtype, weight_name)
+        setattr(
+            sub_module,
+            weight_name,
+            nn.Parameter(exported.pop(weight_name), requires_grad=False),
+        )
+        for name, value in exported.items():
+            sub_module.register_buffer(name, value)
+        torch.cuda.empty_cache()
+        return
+
     # Already real-quantized weights (``mtq.compress`` / ``hf_ptq --low_memory_mode``) hold packed
     # nibbles -- half the logical last dim -- so per-block scales cannot be recomputed from them.
     # Use the scale the quantizer captured at compression time instead.
diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py
index 55137a64639..45a7e3e02a7 100644
--- a/tests/gpu/torch/export/test_export.py
+++ b/tests/gpu/torch/export/test_export.py
@@ -106,6 +106,7 @@ def test_get_quantization_format(config, expected):
             {
                 "quant_algo": "MIXED_PRECISION",
                 "kv_cache_quant_algo": None,
+                "exclude_modules": ["layer8"],
                 "quantized_layers": {
                     "layer1": {"quant_algo": "NVFP4", "group_size": 16},
                     "layer3": {
diff --git a/tests/gpu/torch/export/test_export_weight_gpu.py b/tests/gpu/torch/export/test_export_weight_gpu.py
index 9db2b51114b..dfca4a74bf4 100644
--- a/tests/gpu/torch/export/test_export_weight_gpu.py
+++ b/tests/gpu/torch/export/test_export_weight_gpu.py
@@ -16,6 +16,7 @@
 import copy
 import math
 
+import pytest
 import torch
 import torch.nn as nn
 from _test_utils.torch.export.utils import ToyModel, partial_w4a8_config
@@ -23,7 +24,17 @@
 from torch.nn import init
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.quant_utils import postprocess_state_dict
+from modelopt.torch.export.quant_utils import (
+    capture_quantized_weight_export_state,
+    export_quantized_weight_tensors,
+    get_activation_scaling_factor,
+    get_quantization_format,
+    get_weight_block_size,
+    get_weight_scaling_factor,
+    get_weight_scaling_factor_2,
+    postprocess_state_dict,
+    to_quantized_weight,
+)
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
 from modelopt.torch.quantization.nn.modules.quant_module import QuantModule, QuantModuleRegistry
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
@@ -125,6 +136,45 @@ def test_export_per_block_quantized_weight():
     assert not hasattr(model.linears[2], quantizer_attrs.output_scale)
 
 
+@pytest.mark.parametrize("quant_cfg", [mtq.NVFP4_DEFAULT_CFG, mtq.W4A16_NVFP4_CFG])
+def test_functional_nvfp4_export_matches_existing_helpers_without_mutation(quant_cfg):
+    in_features = 256
+    torch.manual_seed(0)
+    module = nn.Linear(in_features, in_features, bias=False, device="cuda", dtype=torch.bfloat16)
+    calib_input = torch.randn(2, 4, in_features, device="cuda", dtype=torch.bfloat16)
+    module = mtq.quantize(module, copy.deepcopy(quant_cfg), lambda model: model(calib_input))
+    original_weight = module.weight.detach().clone()
+    original_buffers = {name: value.detach().clone() for name, value in module.named_buffers()}
+
+    state = capture_quantized_weight_export_state(module)
+    actual = export_quantized_weight_tensors(module.weight, state, torch.float16)
+
+    quantization_format = get_quantization_format(module)
+    weight_scale = get_weight_scaling_factor(module)
+    weight_scale_2 = get_weight_scaling_factor_2(module)
+    expected = {
+        "weight": to_quantized_weight(
+            module.weight.to(torch.float16),
+            weight_scale,
+            quantization_format,
+            weight_scale_2,
+            get_weight_block_size(module),
+        ),
+        "weight_scale": weight_scale,
+        "weight_scale_2": weight_scale_2.squeeze(),
+    }
+    if module.input_quantizer.is_enabled:
+        expected["input_scale"] = get_activation_scaling_factor(module).squeeze()
+
+    assert actual.keys() == expected.keys()
+    for name, value in actual.items():
+        torch.testing.assert_close(value, expected[name], rtol=0, atol=0)
+    torch.testing.assert_close(module.weight, original_weight)
+    assert set(dict(module.named_buffers())) == set(original_buffers)
+    for name, value in module.named_buffers():
+        torch.testing.assert_close(value, original_buffers[name])
+
+
 def test_export_compressed_nvfp4_weight():
     """``mtq.compress`` (used by ``hf_ptq --low_memory_mode``) leaves the weight as packed NVFP4
     nibbles, so per-block scales cannot be recomputed from it. The export must reuse the scales

From 898c0c07cc21e35a1ee54e7d86be36697675df6f Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Sat, 15 Aug 2026 08:10:31 -0700
Subject: [PATCH 03/10] WIP: tighten quantized export state contract

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 296 ++++++++----------
 .../export/test_quantized_weight_state.py     |  46 +++
 2 files changed, 178 insertions(+), 164 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 49b9e8b2d56..5214c230836 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -42,7 +42,6 @@
     QTensorWrapper,
 )
 from modelopt.torch.quantization.utils import (
-    QuantizerAttrNames,
     quantizer_attr_names,
     representative_weight_quantizer,
     weight_attr_names,
@@ -484,116 +483,102 @@ def get_weight_block_size(module: nn.Module, weight_name: str = "weight") -> int
     return 0
 
 
-def get_quantization_format(module) -> str | None:
-    """Gets the quantization string.
-
-    Gets the quantization string by iterating through the module and its children.
-    The first non-None quantization string is returned.
-    """
-
-    def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames):
-        # Singular form first, plural ModuleList fallback (fused-experts).
-        # Strip the "_weight_quantizer" suffix to recover the weight attr name.
-        weight_attr = quantizer_attr_names.weight_quantizer
-        weight_name = weight_attr[: -len("_weight_quantizer")].rstrip("_") or "weight"
-        weight_quantizer = representative_weight_quantizer(layer, weight_name)
-        input_quantizer = getattr(layer, quantizer_attr_names.input_quantizer, None)
-
-        if weight_quantizer is None or not weight_quantizer.is_enabled:
-            return QUANTIZATION_NONE
-
-        # Handle SequentialQuantizer
-        if isinstance(weight_quantizer, SequentialQuantizer):
-            assert (
-                len(weight_quantizer) == 2
-                and weight_quantizer[0].num_bits == 4
-                and weight_quantizer[1].num_bits == (4, 3)
-            ), "Unsupported SequentialQuantizer configuration"
-            assert (
-                weight_quantizer[0].block_sizes
-                and len(weight_quantizer[0].block_sizes) > 0
-                and weight_quantizer[0].block_sizes[-1] > 0
-            ), "Invalid block_sizes for SequentialQuantizer"
-
-            return QUANTIZATION_W4A8_AWQ
-
-        # Handle individual num_bits cases
-        if weight_quantizer.num_bits == 4:
-            assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, (
-                "Invalid block_sizes for INT4 quantizer"
-            )
-            return QUANTIZATION_INT4_AWQ
+def _get_quantization_from_quantizers(
+    layer: nn.Module,
+    weight_quantizer: TensorQuantizer | SequentialQuantizer | None,
+    input_quantizer: TensorQuantizer | SequentialQuantizer | None,
+) -> str:
+    if weight_quantizer is None or not weight_quantizer.is_enabled:
+        return QUANTIZATION_NONE
 
-        if weight_quantizer.num_bits == 8:
-            if input_quantizer is not None and input_quantizer.is_enabled:
-                return QUANTIZATION_INT8_SQ
-            else:
-                return QUANTIZATION_INT8_WO
+    if isinstance(weight_quantizer, SequentialQuantizer):
+        assert (
+            len(weight_quantizer) == 2
+            and weight_quantizer[0].num_bits == 4
+            and weight_quantizer[1].num_bits == (4, 3)
+        ), "Unsupported SequentialQuantizer configuration"
+        assert (
+            weight_quantizer[0].block_sizes
+            and len(weight_quantizer[0].block_sizes) > 0
+            and weight_quantizer[0].block_sizes[-1] > 0
+        ), "Invalid block_sizes for SequentialQuantizer"
+        return QUANTIZATION_W4A8_AWQ
+
+    if weight_quantizer.num_bits == 4:
+        assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, (
+            "Invalid block_sizes for INT4 quantizer"
+        )
+        return QUANTIZATION_INT4_AWQ
 
-        if weight_quantizer.num_bits == (4, 3):
-            if weight_quantizer.block_sizes:
-                assert weight_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer"
-                # Check if this is MXFP8 (dynamic block quantization with scale_bits (8, 0))
-                block_sizes = getattr(weight_quantizer, "block_sizes")
-                if (
-                    isinstance(block_sizes, dict)
-                    and block_sizes.get("type", "static") == "dynamic"
-                    and block_sizes.get("scale_bits") == (8, 0)
-                ):
-                    return QUANTIZATION_MXFP8
-                if weight_quantizer.fake_quant:
-                    return QUANTIZATION_FP8_PB_WO
-                else:
-                    return QUANTIZATION_FP8_PB_REAL
-            if weight_quantizer.axis == 0:
-                return QUANTIZATION_FP8_PC_PT
-            return QUANTIZATION_FP8
+    if weight_quantizer.num_bits == 8:
+        if input_quantizer is not None and input_quantizer.is_enabled:
+            return QUANTIZATION_INT8_SQ
+        return QUANTIZATION_INT8_WO
 
-        if weight_quantizer.num_bits == (2, 1):
-            # FP4 formats are all block quantization
+    if weight_quantizer.num_bits == (4, 3):
+        if weight_quantizer.block_sizes:
+            assert weight_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer"
             block_sizes = getattr(weight_quantizer, "block_sizes")
-            scale_bits = block_sizes.get("scale_bits")
-
-            if input_quantizer is not None and hasattr(weight_quantizer, "svdquant_lora_a"):
-                return QUANTIZATION_NVFP4_SVDQUANT
-            if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"):
-                return QUANTIZATION_NVFP4_AWQ
-            if getattr(layer, "fused_with_prequant", False):
-                return QUANTIZATION_NVFP4_AWQ
-            if input_quantizer is None or not input_quantizer.is_enabled:
-                if scale_bits == (4, 3):
-                    return QUANTIZATION_W4A16_NVFP4
-            assert input_quantizer is not None, (
-                f"input_quantizer is None for {quantizer_attr_names}"
-            )
-            if (
-                block_sizes.get("type", "static") == "dynamic"
-                and scale_bits == (8, 0)
-                and input_quantizer.is_enabled
-                and input_quantizer.num_bits == (4, 3)
-                and input_quantizer.block_sizes is None
-            ):
-                return QUANTIZATION_W4A8_MXFP4_FP8
             if (
-                block_sizes.get("type", "static") == "dynamic"
-                and scale_bits == (4, 3)
-                and input_quantizer.is_enabled
-                and input_quantizer.num_bits == (4, 3)
-                and input_quantizer.block_sizes is None
+                isinstance(block_sizes, dict)
+                and block_sizes.get("type", "static") == "dynamic"
+                and block_sizes.get("scale_bits") == (8, 0)
             ):
-                return QUANTIZATION_W4A8_NVFP4_FP8
+                return QUANTIZATION_MXFP8
+            if weight_quantizer.fake_quant:
+                return QUANTIZATION_FP8_PB_WO
+            return QUANTIZATION_FP8_PB_REAL
+        if weight_quantizer.axis == 0:
+            return QUANTIZATION_FP8_PC_PT
+        return QUANTIZATION_FP8
+
+    if weight_quantizer.num_bits == (2, 1):
+        block_sizes = getattr(weight_quantizer, "block_sizes")
+        scale_bits = block_sizes.get("scale_bits")
+
+        if input_quantizer is not None and hasattr(weight_quantizer, "svdquant_lora_a"):
+            return QUANTIZATION_NVFP4_SVDQUANT
+        if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"):
+            return QUANTIZATION_NVFP4_AWQ
+        if getattr(layer, "fused_with_prequant", False):
+            return QUANTIZATION_NVFP4_AWQ
+        if input_quantizer is None or not input_quantizer.is_enabled:
             if scale_bits == (4, 3):
-                return QUANTIZATION_NVFP4
-            elif scale_bits == (8, 0):
-                return QUANTIZATION_MXFP4
+                return QUANTIZATION_W4A16_NVFP4
+        assert input_quantizer is not None, "input_quantizer is required for weight-activation FP4"
+        if (
+            block_sizes.get("type", "static") == "dynamic"
+            and scale_bits == (8, 0)
+            and input_quantizer.is_enabled
+            and input_quantizer.num_bits == (4, 3)
+            and input_quantizer.block_sizes is None
+        ):
+            return QUANTIZATION_W4A8_MXFP4_FP8
+        if (
+            block_sizes.get("type", "static") == "dynamic"
+            and scale_bits == (4, 3)
+            and input_quantizer.is_enabled
+            and input_quantizer.num_bits == (4, 3)
+            and input_quantizer.block_sizes is None
+        ):
+            return QUANTIZATION_W4A8_NVFP4_FP8
+        if scale_bits == (4, 3):
+            return QUANTIZATION_NVFP4
+        if scale_bits == (8, 0):
+            return QUANTIZATION_MXFP4
+
+    raise NotImplementedError(f"Unsupported quantizer with num_bits: {weight_quantizer.num_bits}")
 
-        # Raise error for unsupported num_bits
-        raise NotImplementedError(
-            f"Unsupported quantizer with num_bits: {weight_quantizer.num_bits}"
-        )
 
+def get_quantization_format(module) -> str | None:
+    """Return the first enabled weight quantization format in a module tree."""
     for weight_name in weight_attr_names(module):
-        quantization = _get_quantization_from_layer(module, quantizer_attr_names(weight_name))
+        attr_names = quantizer_attr_names(weight_name)
+        quantization = _get_quantization_from_quantizers(
+            module,
+            representative_weight_quantizer(module, weight_name),
+            getattr(module, attr_names.input_quantizer, None),
+        )
         if quantization != QUANTIZATION_NONE:
             return quantization
 
@@ -947,9 +932,6 @@ def to_quantized_weight(
 
 _FUNCTIONAL_EXPORT_FORMATS = {
     QUANTIZATION_FP8,
-    QUANTIZATION_FP8_PB_WO,
-    QUANTIZATION_FP8_PC_PT,
-    QUANTIZATION_MXFP8,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A16_NVFP4,
 }
@@ -965,7 +947,7 @@ class _ExportStateTensor:
 
 
 @dataclass(frozen=True)
-class QuantizedWeightExportState:
+class _QuantizedWeightExportState:
     """Opaque state required to export one logical quantized weight."""
 
     quantization_format: str
@@ -1012,44 +994,40 @@ def _resolve_weight_quantizer(
 def _packing_permutation(weight: torch.Tensor, quantized_view: torch.Tensor) -> tuple[int, ...]:
     ndim = weight.ndim
     identity = tuple(range(ndim))
-    if tuple(quantized_view.shape) == tuple(weight.shape):
-        return identity
     if ndim >= 2:
-        transposed = (*weight.shape[:-2], weight.shape[-1], weight.shape[-2])
-        if tuple(quantized_view.shape) == transposed:
+        transposed = weight.transpose(-1, -2)
+        if (
+            tuple(quantized_view.shape) == tuple(transposed.shape)
+            and quantized_view.stride() == transposed.stride()
+        ):
             return (*range(ndim - 2), ndim - 1, ndim - 2)
+    if (
+        tuple(quantized_view.shape) == tuple(weight.shape)
+        and quantized_view.stride() == weight.stride()
+    ):
+        return identity
     raise NotImplementedError(
-        f"Unsupported quantized weight view {tuple(quantized_view.shape)} for {tuple(weight.shape)}"
+        "Unsupported quantized weight view "
+        f"shape/stride={tuple(quantized_view.shape)}/{quantized_view.stride()} for "
+        f"shape/stride={tuple(weight.shape)}/{weight.stride()}"
     )
 
 
 def _state_tensor(
     name: str,
     value: torch.Tensor,
-    packed_shape: tuple[int, ...],
     *,
-    block_sizes: tuple[int, ...] | None = None,
+    axes: tuple[int, ...] = (),
+    block_sizes: tuple[int, ...] = (),
 ) -> _ExportStateTensor:
     value = value.detach().clone()
     if value.numel() == 1:
         return _ExportStateTensor(name, value.reshape(()))
-    if value.ndim > len(packed_shape):
-        raise NotImplementedError(
-            f"Cannot relate {name} shape {tuple(value.shape)} to weight shape {packed_shape}"
+    if len(axes) != value.ndim or len(block_sizes) != value.ndim:
+        raise ValueError(
+            f"Explicit axis and block metadata is required for non-scalar {name}: "
+            f"shape={tuple(value.shape)}, axes={axes}, block_sizes={block_sizes}"
         )
-
-    axes = tuple(range(value.ndim))
-    if block_sizes is None:
-        inferred = []
-        for axis, size in enumerate(value.shape):
-            if packed_shape[axis] % size:
-                raise NotImplementedError(
-                    f"Cannot relate {name} shape {tuple(value.shape)} to weight shape {packed_shape}"
-                )
-            inferred.append(packed_shape[axis] // size)
-        block_sizes = tuple(inferred)
-    if len(block_sizes) != value.ndim:
-        raise ValueError(f"Invalid block layout for {name}: {block_sizes}")
     return _ExportStateTensor(name, value, axes, block_sizes)
 
 
@@ -1063,14 +1041,17 @@ def _input_quantizer(module: nn.Module, weight_name: str):
 def capture_quantized_weight_export_state(
     module: nn.Module,
     weight_name: str = "weight",
-) -> QuantizedWeightExportState:
+) -> _QuantizedWeightExportState:
     """Capture detached export state without mutating the quantized module."""
     weight = getattr(module, weight_name)
     if isinstance(weight, QTensorWrapper):
         raise NotImplementedError("Functional export requires an uncompressed source weight")
 
     quantized_view, weight_quantizer = _resolve_weight_quantizer(module, weight_name)
-    quantization_format = get_quantization_format(module)
+    input_quantizer = _input_quantizer(module, weight_name)
+    quantization_format = _get_quantization_from_quantizers(
+        module, weight_quantizer, input_quantizer
+    )
     if quantization_format not in _FUNCTIONAL_EXPORT_FORMATS:
         raise NotImplementedError(f"Functional export does not support {quantization_format!r}")
     if isinstance(weight_quantizer, SequentialQuantizer):
@@ -1100,33 +1081,20 @@ def capture_quantized_weight_export_state(
             _state_tensor(
                 "weight_block_amax",
                 per_block_amax.reshape(block_shape),
-                packed_shape,
+                axes=tuple(range(len(packed_shape))),
                 block_sizes=(1,) * (len(packed_shape) - 1) + (block_size,),
             )
         )
-        tensors.append(_state_tensor("weight_global_amax", global_amax, packed_shape))
+        tensors.append(_state_tensor("weight_global_amax", global_amax))
     elif quantization_format in _NVFP4_EXPORT_FORMATS:
         weight_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
-        tensors.append(_state_tensor("weight_scale_2", weight_scale_2, packed_shape))
-    elif quantization_format == QUANTIZATION_MXFP8:
-        weight_scale = MXFP8QTensor.get_weights_scaling_factor_from_quantizer(
-            quantized_view, weight_quantizer
-        )
-        tensors.append(
-            _state_tensor(
-                "weight_scale",
-                weight_scale,
-                packed_shape,
-                block_sizes=(1,) * (len(packed_shape) - 1) + (block_size,),
-            )
-        )
+        tensors.append(_state_tensor("weight_scale_2", weight_scale_2))
     else:
         weight_scale = get_scaling_factor(weight_quantizer)
         if weight_scale is None:
             raise RuntimeError(f"Missing calibrated weight scale for {weight_name!r}")
-        tensors.append(_state_tensor("weight_scale", weight_scale, packed_shape))
+        tensors.append(_state_tensor("weight_scale", weight_scale))
 
-    input_quantizer = _input_quantizer(module, weight_name)
     if input_quantizer is not None and input_quantizer.is_enabled:
         if input_quantizer.export_amax() is None:
             raise RuntimeError(f"Missing calibrated input scale for {weight_name!r}")
@@ -1135,9 +1103,9 @@ def capture_quantized_weight_export_state(
             if quantization_format == QUANTIZATION_NVFP4
             else get_scaling_factor(input_quantizer)
         )
-        tensors.append(_state_tensor("input_scale", input_scale, packed_shape))
+        tensors.append(_state_tensor("input_scale", input_scale))
 
-    return QuantizedWeightExportState(
+    return _QuantizedWeightExportState(
         quantization_format=quantization_format,
         block_size=block_size,
         weight_shape=tuple(weight.shape),
@@ -1176,9 +1144,9 @@ def _merge_state_tensors(
 
 
 def merge_quantized_weight_export_states(
-    states: Sequence[QuantizedWeightExportState],
+    states: Sequence[_QuantizedWeightExportState],
     weight_dim: int,
-) -> QuantizedWeightExportState:
+) -> _QuantizedWeightExportState:
     """Merge state for logical weight shards concatenated along ``weight_dim``."""
     if not states:
         raise ValueError("At least one quantized weight state is required")
@@ -1209,7 +1177,7 @@ def merge_quantized_weight_export_states(
         _merge_state_tensors([state.tensors[index] for state in states], packed_dim)
         for index in range(len(reference.tensors))
     )
-    return QuantizedWeightExportState(
+    return _QuantizedWeightExportState(
         reference.quantization_format,
         reference.block_size,
         tuple(shape),
@@ -1242,11 +1210,11 @@ def _select_state_tensor(
 
 
 def select_quantized_weight_export_state(
-    state: QuantizedWeightExportState,
+    state: _QuantizedWeightExportState,
     weight_dim: int,
     indices: Iterable[int] | torch.Tensor,
-) -> QuantizedWeightExportState:
-    """Select logical weight indices and their corresponding opaque export state."""
+) -> _QuantizedWeightExportState:
+    """Select logical indices that preserve complete quantization blocks."""
     ndim = len(state.weight_shape)
     weight_dim = _normalize_weight_dim(weight_dim, ndim)
     indices = torch.as_tensor(list(indices) if not isinstance(indices, torch.Tensor) else indices)
@@ -1259,7 +1227,7 @@ def select_quantized_weight_export_state(
     shape = list(state.weight_shape)
     shape[weight_dim] = indices.numel()
     packed_dim = state.packing_permutation.index(weight_dim)
-    return QuantizedWeightExportState(
+    return _QuantizedWeightExportState(
         state.quantization_format,
         state.block_size,
         tuple(shape),
@@ -1271,16 +1239,16 @@ def select_quantized_weight_export_state(
 
 
 def permute_quantized_weight_export_state(
-    state: QuantizedWeightExportState,
+    state: _QuantizedWeightExportState,
     dims: Sequence[int],
-) -> QuantizedWeightExportState:
+) -> _QuantizedWeightExportState:
     """Permute logical weight dimensions while retaining quantizer packing axes."""
     dims = tuple(dims)
     ndim = len(state.weight_shape)
     if sorted(dims) != list(range(ndim)):
         raise ValueError(f"Invalid permutation {dims} for a rank-{ndim} weight")
     inverse = {old_dim: new_dim for new_dim, old_dim in enumerate(dims)}
-    return QuantizedWeightExportState(
+    return _QuantizedWeightExportState(
         state.quantization_format,
         state.block_size,
         tuple(state.weight_shape[dim] for dim in dims),
@@ -1312,7 +1280,7 @@ def _restore_state_tensor(
 
 def export_quantized_weight_tensors(
     weight: torch.Tensor,
-    state: QuantizedWeightExportState,
+    state: _QuantizedWeightExportState,
     dtype: torch.dtype,
     weight_name: str = "weight",
 ) -> OrderedDict[str, torch.Tensor]:
@@ -1367,10 +1335,10 @@ def export_quantized_weight_tensors(
     weight_scale_record = _state_tensor(
         "weight_scale",
         weight_scale,
-        tuple(packed_weight.shape),
+        axes=tuple(range(packed_weight.ndim)),
         block_sizes=(1,) * (packed_weight.ndim - 1) + (state.block_size,)
         if state.quantization_format in _NVFP4_EXPORT_FORMATS
-        else None,
+        else (1,) * packed_weight.ndim,
     )
     output[attrs.weight_scale] = _restore_state_tensor(
         weight_scale_record, state.packing_permutation
@@ -1399,8 +1367,8 @@ def _quantized_layer_name(weight_name: str) -> str:
 
 
 def build_hf_quantization_config(
-    named_states: Mapping[str, QuantizedWeightExportState | None]
-    | Iterable[tuple[str, QuantizedWeightExportState | None]],
+    named_states: Mapping[str, _QuantizedWeightExportState | None]
+    | Iterable[tuple[str, _QuantizedWeightExportState | None]],
 ) -> dict[str, Any]:
     """Build canonical ModelOpt HF configuration from named opaque states."""
     layers: dict[str, tuple[str, int] | None] = {}
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index bb1999b113f..c15c6705467 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -47,6 +47,33 @@ def _static_w4a16_linear(
     return module
 
 
+class _SquareTransposedExperts(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gate_up_proj = nn.Parameter(torch.arange(512, dtype=torch.float32).reshape(2, 16, 16))
+        self.down_proj = nn.Parameter(torch.arange(512, dtype=torch.float32).reshape(2, 16, 16))
+
+        fp8_cfg = QuantizerAttributeConfig(num_bits=(4, 3))
+        self.gate_up_proj_weight_quantizer = TensorQuantizer(fp8_cfg)
+        self.gate_up_proj_weight_quantizer._amax = torch.tensor(2.0)
+        self.gate_up_proj_input_quantizer = TensorQuantizer(fp8_cfg)
+        self.gate_up_proj_input_quantizer._amax = torch.tensor(1.0)
+
+        nvfp4_cfg = QuantizerAttributeConfig(
+            num_bits=(2, 1),
+            block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)},
+        )
+        self.down_proj_weight_quantizer = NVFP4StaticQuantizer(quant_attribute_cfg=nvfp4_cfg)
+        self.down_proj_weight_quantizer.amax = torch.arange(32, dtype=torch.float32) + 1
+        self.down_proj_weight_quantizer.global_amax = torch.tensor(32.0)
+        self.down_proj_input_quantizer = TensorQuantizer()
+        self.down_proj_input_quantizer.disable()
+
+    def iter_weights_for_calibration(self):
+        for name in ("gate_up_proj", "down_proj"):
+            yield getattr(self, name).transpose(-1, -2), getattr(self, f"{name}_weight_quantizer")
+
+
 def test_capture_does_not_modify_zero_amax():
     module = _fp8_linear()
     module.weight_quantizer._amax.zero_()
@@ -148,6 +175,25 @@ def test_state_rejects_invalid_dimensions():
         select_quantized_weight_export_state(state, -3, (0,))
 
 
+def test_capture_uses_the_requested_weight_format_and_transposed_layout():
+    module = _SquareTransposedExperts()
+    gate_state = capture_quantized_weight_export_state(module, "gate_up_proj")
+    down_state = capture_quantized_weight_export_state(module, "down_proj")
+
+    config = build_hf_quantization_config(
+        {
+            "model.layers.0.mlp.gate_proj.weight": gate_state,
+            "model.layers.0.mlp.down_proj.weight": down_state,
+        }
+    )
+    assert config["quantized_layers"]["model.layers.0.mlp.gate_proj"]["quant_algo"] == "FP8"
+    assert config["quantized_layers"]["model.layers.0.mlp.down_proj"]["quant_algo"] == "W4A16_NVFP4"
+
+    with pytest.raises(ValueError, match="complete quantization blocks"):
+        select_quantized_weight_export_state(down_state, 1, (0,))
+    select_quantized_weight_export_state(down_state, 2, (0, 2))
+
+
 def test_mixed_config_groups_moe_experts_by_projection_family():
     fp8 = capture_quantized_weight_export_state(_fp8_linear())
     weight = torch.ones(4, 16)

From e12f71b11d218eb6b57973d2b13b37c297eaf79a Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Sat, 15 Aug 2026 08:25:08 -0700
Subject: [PATCH 04/10] WIP: preserve legacy export fallbacks

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 48 ++++++++++++-------
 modelopt/torch/export/unified_export_hf.py    | 12 ++---
 .../torch/export/test_export_weight_gpu.py    | 26 ++++++++++
 .../export/test_quantized_weight_state.py     | 39 +++++++++++++++
 4 files changed, 99 insertions(+), 26 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 5214c230836..1fa8f8d89d9 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -1019,9 +1019,10 @@ def _state_tensor(
     *,
     axes: tuple[int, ...] = (),
     block_sizes: tuple[int, ...] = (),
+    cpu: bool = True,
 ) -> _ExportStateTensor:
-    value = value.detach().clone()
-    if value.numel() == 1:
+    value = value.detach().cpu().clone() if cpu else value.detach()
+    if value.numel() == 1 and not axes and not block_sizes:
         return _ExportStateTensor(name, value.reshape(()))
     if len(axes) != value.ndim or len(block_sizes) != value.ndim:
         raise ValueError(
@@ -1200,11 +1201,16 @@ def _select_state_tensor(
     if block_size == 1:
         selected = indices
     else:
-        selected = torch.div(indices, block_size, rounding_mode="floor")
-        unique, counts = selected.unique_consecutive(return_counts=True)
-        if not torch.all(counts == block_size):
+        if indices.numel() % block_size:
+            raise ValueError("Weight selection must preserve complete quantization blocks")
+        blocks = indices.reshape(-1, block_size)
+        block_indices = torch.div(blocks, block_size, rounding_mode="floor")
+        selected = block_indices[:, 0]
+        expected = selected[:, None] * block_size + torch.arange(block_size)
+        if not torch.all(block_indices == selected[:, None]) or not torch.equal(
+            blocks.sort(dim=1).values, expected
+        ):
             raise ValueError("Weight selection must preserve complete quantization blocks")
-        selected = unique
     value = record.value.index_select(tensor_dim, selected.to(record.value.device))
     return _ExportStateTensor(record.name, value, record.axes, record.block_sizes)
 
@@ -1320,6 +1326,10 @@ def export_quantized_weight_tensors(
     else:
         weight_scale = records["weight_scale"].value
 
+    weight_scale = weight_scale.to(packed_weight.device)
+    if weight_scale_2 is not None:
+        weight_scale_2 = weight_scale_2.to(packed_weight.device)
+
     quantized_weight = to_quantized_weight(
         packed_weight.to(dtype),
         weight_scale,
@@ -1332,14 +1342,16 @@ def export_quantized_weight_tensors(
         ((weight_name, _restore_packing_permutation(quantized_weight, state.packing_permutation)),)
     )
 
-    weight_scale_record = _state_tensor(
-        "weight_scale",
-        weight_scale,
-        axes=tuple(range(packed_weight.ndim)),
-        block_sizes=(1,) * (packed_weight.ndim - 1) + (state.block_size,)
-        if state.quantization_format in _NVFP4_EXPORT_FORMATS
-        else (1,) * packed_weight.ndim,
-    )
+    if state.quantization_format in _NVFP4_EXPORT_FORMATS:
+        weight_scale_record = _state_tensor(
+            "weight_scale",
+            weight_scale,
+            axes=tuple(range(packed_weight.ndim)),
+            block_sizes=(1,) * (packed_weight.ndim - 1) + (state.block_size,),
+            cpu=False,
+        )
+    else:
+        weight_scale_record = _state_tensor("weight_scale", weight_scale, cpu=False)
     output[attrs.weight_scale] = _restore_state_tensor(
         weight_scale_record, state.packing_permutation
     )
@@ -1355,9 +1367,11 @@ def export_quantized_weight_tensors(
             state.packing_permutation,
         )
     if "input_scale" in records:
-        output[attrs.input_scale] = _restore_state_tensor(
-            records["input_scale"], state.packing_permutation
-        ).squeeze()
+        output[attrs.input_scale] = (
+            _restore_state_tensor(records["input_scale"], state.packing_permutation)
+            .squeeze()
+            .to(weight.device)
+        )
     return output
 
 
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 945fe346168..6bc8df91615 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -82,7 +82,6 @@
 from .model_config import (
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
-    QUANTIZATION_FP8_PB_WO,
     QUANTIZATION_FP8_PC_PT,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
@@ -610,14 +609,9 @@ def _export_quantized_weight(
         sub_module, quantizer_attrs.output_quantizer, None
     )
 
-    if not isinstance(weight, QTensorWrapper) and quantization_format in {
-        QUANTIZATION_FP8,
-        QUANTIZATION_FP8_PB_WO,
-        QUANTIZATION_FP8_PC_PT,
-        QUANTIZATION_MXFP8,
-        QUANTIZATION_NVFP4,
-        QUANTIZATION_W4A16_NVFP4,
-    }:
+    from .quant_utils import _FUNCTIONAL_EXPORT_FORMATS
+
+    if not isinstance(weight, QTensorWrapper) and quantization_format in _FUNCTIONAL_EXPORT_FORMATS:
         state = capture_quantized_weight_export_state(sub_module, weight_name)
         exported = export_quantized_weight_tensors(weight, state, dtype, weight_name)
         setattr(
diff --git a/tests/gpu/torch/export/test_export_weight_gpu.py b/tests/gpu/torch/export/test_export_weight_gpu.py
index dfca4a74bf4..86d7d591031 100644
--- a/tests/gpu/torch/export/test_export_weight_gpu.py
+++ b/tests/gpu/torch/export/test_export_weight_gpu.py
@@ -175,6 +175,32 @@ def test_functional_nvfp4_export_matches_existing_helpers_without_mutation(quant
         torch.testing.assert_close(value, original_buffers[name])
 
 
+@pytest.mark.parametrize(
+    "quant_cfg",
+    [
+        mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
+        mtq.FP8_PER_CHANNEL_PER_TOKEN_CFG,
+        mtq.MXFP8_DEFAULT_CFG,
+    ],
+)
+def test_unqualified_formats_keep_existing_export_path(monkeypatch, quant_cfg):
+    in_features = 256
+    module = nn.Linear(in_features, in_features, bias=False, device="cuda", dtype=torch.bfloat16)
+    calib_input = torch.randn(2, 4, in_features, device="cuda", dtype=torch.bfloat16)
+    module = mtq.quantize(module, copy.deepcopy(quant_cfg), lambda model: model(calib_input))
+
+    def fail_functional_capture(*args, **kwargs):
+        raise AssertionError("unqualified format entered the functional export path")
+
+    monkeypatch.setattr(
+        "modelopt.torch.export.unified_export_hf.capture_quantized_weight_export_state",
+        fail_functional_capture,
+    )
+    _export_quantized_weight(module, torch.float16)
+
+    assert hasattr(module, "weight_scale")
+
+
 def test_export_compressed_nvfp4_weight():
     """``mtq.compress`` (used by ``hf_ptq --low_memory_mode``) leaves the weight as packed NVFP4
     nibbles, so per-block scales cannot be recomputed from it. The export must reuse the scales
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index c15c6705467..39a252322f0 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import copy
+import pickle
 
 import pytest
 import torch
@@ -110,6 +111,15 @@ def test_functional_fp8_export_is_repeatable_and_supports_permutation():
     torch.testing.assert_close(transposed["weight"], base["weight"].T)
 
 
+def test_export_state_round_trips_through_object_transport():
+    module = _fp8_linear()
+    state = pickle.loads(pickle.dumps(capture_quantized_weight_export_state(module)))
+
+    exported = export_quantized_weight_tensors(module.weight, state, torch.float32)
+
+    assert exported["weight"].shape == module.weight.shape
+
+
 def test_static_nvfp4_merge_recomputes_scales_from_merged_amax():
     left_weight = torch.arange(32, dtype=torch.float32).reshape(2, 16) / 32
     right_weight = torch.arange(32, 64, dtype=torch.float32).reshape(2, 16) / 16
@@ -167,6 +177,35 @@ def test_static_nvfp4_noncontiguous_selection_tracks_block_amax():
         torch.testing.assert_close(actual[name], expected[name], rtol=0, atol=0)
 
 
+def test_static_nvfp4_single_block_preserves_scale_shape():
+    weight = torch.arange(16, dtype=torch.float32).reshape(1, 16)
+    module = _static_w4a16_linear(weight, weight.abs().amax().reshape(1, 1), torch.tensor(1.0))
+
+    exported = export_quantized_weight_tensors(
+        module.weight,
+        capture_quantized_weight_export_state(module),
+        torch.float32,
+    )
+
+    assert exported["weight_scale"].shape == (1, 1)
+
+
+def test_static_nvfp4_selection_rejects_duplicate_block_members():
+    weight = torch.arange(32, dtype=torch.float32).reshape(2, 16)
+    module = _static_w4a16_linear(
+        weight,
+        weight.abs().amax(dim=1, keepdim=True),
+        torch.tensor(1.0),
+    )
+
+    with pytest.raises(ValueError, match="complete quantization blocks"):
+        select_quantized_weight_export_state(
+            capture_quantized_weight_export_state(module),
+            1,
+            (0,) * 16,
+        )
+
+
 def test_state_rejects_invalid_dimensions():
     state = capture_quantized_weight_export_state(_fp8_linear())
     with pytest.raises(IndexError, match="invalid"):

From c092e841e84810880e56ce53509808eb58b9cbdc Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Sat, 15 Aug 2026 10:09:15 -0700
Subject: [PATCH 05/10] WIP: preserve real quant group size

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/convert_hf_config.py        |  5 +++++
 .../torch/export/test_quantized_weight_state.py   | 15 +++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py
index 45fa0c30f3b..be6346d555d 100644
--- a/modelopt/torch/export/convert_hf_config.py
+++ b/modelopt/torch/export/convert_hf_config.py
@@ -265,6 +265,11 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An
     if quant_algo_value:
         new_config["quant_algo"] = quant_algo_value
 
+    group_size = original_quantization_details.get("group_size")
+    if group_size is not None:
+        # ModelOpt consumers use this top-level value to size block scales.
+        new_config["group_size"] = group_size
+
     kv_cache_quant_algo = original_quantization_details.get("kv_cache_quant_algo")
     if kv_cache_quant_algo:
         if kv_cache_quant_algo == "FP8":
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index 39a252322f0..751a987a784 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -9,6 +9,7 @@
 import torch.nn as nn
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
 from modelopt.torch.export.quant_utils import (
     build_hf_quantization_config,
     capture_quantized_weight_export_state,
@@ -233,6 +234,20 @@ def test_capture_uses_the_requested_weight_format_and_transposed_layout():
     select_quantized_weight_export_state(down_state, 2, (0, 2))
 
 
+def test_converted_hf_config_preserves_nvfp4_group_size():
+    config = convert_hf_quant_config_format(
+        {
+            "quantization": {
+                "quant_algo": "W4A16_NVFP4",
+                "group_size": 32,
+            }
+        }
+    )
+
+    assert config["group_size"] == 32
+    assert config["config_groups"]["group_0"]["weights"]["group_size"] == 32
+
+
 def test_mixed_config_groups_moe_experts_by_projection_family():
     fp8 = capture_quantized_weight_export_state(_fp8_linear())
     weight = torch.ones(4, 16)

From 133a2f6372f655e750a7ac3fd14b3c195dc80741 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Sun, 16 Aug 2026 10:46:49 -0700
Subject: [PATCH 06/10] Complete functional noninteger weight export

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 215 ++++++++++++++----
 modelopt/torch/export/unified_export_hf.py    |  37 +--
 .../torch/export/test_export_weight_gpu.py    | 155 ++++++++++++-
 3 files changed, 350 insertions(+), 57 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 1fa8f8d89d9..16c30d4dc33 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -930,12 +930,16 @@ def to_quantized_weight(
     raise NotImplementedError(f"quantization format {quantization} not supported")
 
 
-_FUNCTIONAL_EXPORT_FORMATS = {
-    QUANTIZATION_FP8,
+_NVFP4_EXPORT_FORMATS = {
     QUANTIZATION_NVFP4,
+    QUANTIZATION_W4A8_NVFP4_FP8,
     QUANTIZATION_W4A16_NVFP4,
 }
-_NVFP4_EXPORT_FORMATS = {QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4}
+_MXFP4_EXPORT_FORMATS = {QUANTIZATION_MXFP4, QUANTIZATION_W4A8_MXFP4_FP8}
+
+
+class _UnsupportedQuantizedWeightExportFormat(NotImplementedError):
+    pass
 
 
 @dataclass(frozen=True)
@@ -1032,6 +1036,79 @@ def _state_tensor(
     return _ExportStateTensor(name, value, axes, block_sizes)
 
 
+def _axis_scale_state(
+    name: str,
+    value: torch.Tensor,
+    packed_shape: tuple[int, ...],
+    axis: int | Sequence[int],
+) -> _ExportStateTensor:
+    axes = (axis,) if isinstance(axis, int) else tuple(axis)
+    axes = tuple(_normalize_weight_dim(dim, len(packed_shape)) for dim in axes)
+    expected_shape = tuple(packed_shape[dim] for dim in axes)
+    if value.numel() != torch.Size(expected_shape).numel():
+        raise RuntimeError(
+            f"Scale {name!r} shape {tuple(value.shape)} does not match axes {axes} "
+            f"of packed weight shape {packed_shape}"
+        )
+    return _state_tensor(
+        name,
+        value.reshape(expected_shape),
+        axes=axes,
+        block_sizes=(1,) * len(axes),
+    )
+
+
+def _block_scale_state(
+    name: str,
+    value: torch.Tensor,
+    packed_shape: tuple[int, ...],
+    block_config: Mapping[int | str, Any],
+    *,
+    cpu: bool = True,
+) -> _ExportStateTensor:
+    ndim = len(packed_shape)
+    block_sizes = [1] * ndim
+    for dim, block_size in block_config.items():
+        if isinstance(dim, int) and block_size is not None:
+            block_sizes[_normalize_weight_dim(dim, ndim)] = int(block_size)
+    expected_shape = tuple(
+        (size + block_size - 1) // block_size
+        for size, block_size in zip(packed_shape, block_sizes)
+    )
+    if value.numel() != torch.Size(expected_shape).numel():
+        raise RuntimeError(
+            f"Scale {name!r} shape {tuple(value.shape)} does not match block sizes "
+            f"{tuple(block_sizes)} of packed weight shape {packed_shape}"
+        )
+
+    expanded_shape = []
+    expanded_axes = []
+    expanded_block_sizes = []
+    for axis, (size, block_size) in enumerate(zip(expected_shape, block_sizes)):
+        expanded_shape.append(size)
+        expanded_axes.append(axis)
+        expanded_block_sizes.append(block_size)
+        if block_size != 1:
+            expanded_shape.append(1)
+            expanded_axes.append(axis)
+            expanded_block_sizes.append(1)
+    if tuple(value.shape) == tuple(expanded_shape):
+        return _state_tensor(
+            name,
+            value,
+            axes=tuple(expanded_axes),
+            block_sizes=tuple(expanded_block_sizes),
+            cpu=cpu,
+        )
+    return _state_tensor(
+        name,
+        value.reshape(expected_shape),
+        axes=tuple(range(ndim)),
+        block_sizes=tuple(block_sizes),
+        cpu=cpu,
+    )
+
+
 def _input_quantizer(module: nn.Module, weight_name: str):
     quantizer = getattr(module, quantizer_attr_names(weight_name).input_quantizer, None)
     if quantizer is None:
@@ -1053,10 +1130,10 @@ def capture_quantized_weight_export_state(
     quantization_format = _get_quantization_from_quantizers(
         module, weight_quantizer, input_quantizer
     )
-    if quantization_format not in _FUNCTIONAL_EXPORT_FORMATS:
-        raise NotImplementedError(f"Functional export does not support {quantization_format!r}")
     if isinstance(weight_quantizer, SequentialQuantizer):
-        weight_quantizer = weight_quantizer[0]
+        raise _UnsupportedQuantizedWeightExportFormat(
+            f"Functional export does not support {quantization_format!r}"
+        )
     if not weight_quantizer.is_enabled:
         raise RuntimeError(f"Weight quantizer for {weight_name!r} is disabled")
 
@@ -1088,23 +1165,56 @@ def capture_quantized_weight_export_state(
         )
         tensors.append(_state_tensor("weight_global_amax", global_amax))
     elif quantization_format in _NVFP4_EXPORT_FORMATS:
-        weight_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
+        weight_scale_2 = (
+            weight_quantizer._amax.float() / 448.0
+            if quantization_format == QUANTIZATION_W4A8_NVFP4_FP8
+            else NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
+        )
         tensors.append(_state_tensor("weight_scale_2", weight_scale_2))
-    else:
+    elif quantization_format == QUANTIZATION_FP8:
         weight_scale = get_scaling_factor(weight_quantizer)
         if weight_scale is None:
             raise RuntimeError(f"Missing calibrated weight scale for {weight_name!r}")
         tensors.append(_state_tensor("weight_scale", weight_scale))
+    elif quantization_format == QUANTIZATION_FP8_PC_PT:
+        weight_scale = get_scaling_factor(weight_quantizer)
+        if weight_scale is None or weight_quantizer.axis is None:
+            raise RuntimeError(f"Missing calibrated per-channel scale for {weight_name!r}")
+        tensors.append(
+            _axis_scale_state(
+                "weight_scale",
+                weight_scale,
+                packed_shape,
+                weight_quantizer.axis,
+            )
+        )
+    elif quantization_format == QUANTIZATION_FP8_PB_WO:
+        weight_scale = get_scaling_factor(weight_quantizer)
+        if weight_scale is None:
+            raise RuntimeError(f"Missing calibrated block scale for {weight_name!r}")
+        tensors.append(
+            _block_scale_state("weight_scale", weight_scale, packed_shape, block_config)
+        )
+    elif quantization_format == QUANTIZATION_MXFP8:
+        cached_scale = getattr(weight_quantizer, "_scale", None)
+        if cached_scale is not None:
+            tensors.append(
+                _block_scale_state("weight_scale", cached_scale, packed_shape, block_config)
+            )
+    elif quantization_format not in _MXFP4_EXPORT_FORMATS:
+        raise _UnsupportedQuantizedWeightExportFormat(
+            f"Functional export does not support {quantization_format!r}"
+        )
 
     if input_quantizer is not None and input_quantizer.is_enabled:
-        if input_quantizer.export_amax() is None:
-            raise RuntimeError(f"Missing calibrated input scale for {weight_name!r}")
-        input_scale = (
-            NVFP4QTensor.get_activation_scaling_factor(input_quantizer)
-            if quantization_format == QUANTIZATION_NVFP4
-            else get_scaling_factor(input_quantizer)
-        )
-        tensors.append(_state_tensor("input_scale", input_scale))
+        input_amax = input_quantizer.export_amax()
+        if input_amax is not None:
+            input_scale = (
+                NVFP4QTensor.get_activation_scaling_factor(input_quantizer)
+                if quantization_format == QUANTIZATION_NVFP4
+                else get_scaling_factor(input_quantizer)
+            )
+            tensors.append(_state_tensor("input_scale", input_scale))
 
     return _QuantizedWeightExportState(
         quantization_format=quantization_format,
@@ -1201,20 +1311,25 @@ def _select_state_tensor(
     if block_size == 1:
         selected = indices
     else:
-        if indices.numel() % block_size:
-            raise ValueError("Weight selection must preserve complete quantization blocks")
-        blocks = indices.reshape(-1, block_size)
-        block_indices = torch.div(blocks, block_size, rounding_mode="floor")
-        selected = block_indices[:, 0]
-        expected = selected[:, None] * block_size + torch.arange(block_size)
-        if not torch.all(block_indices == selected[:, None]) or not torch.equal(
-            blocks.sort(dim=1).values, expected
-        ):
-            raise ValueError("Weight selection must preserve complete quantization blocks")
+        selected = _selected_block_indices(indices, block_size)
     value = record.value.index_select(tensor_dim, selected.to(record.value.device))
     return _ExportStateTensor(record.name, value, record.axes, record.block_sizes)
 
 
+def _selected_block_indices(indices: torch.Tensor, block_size: int) -> torch.Tensor:
+    if indices.numel() % block_size:
+        raise ValueError("Weight selection must preserve complete quantization blocks")
+    blocks = indices.reshape(-1, block_size)
+    block_indices = torch.div(blocks, block_size, rounding_mode="floor")
+    selected = block_indices[:, 0]
+    expected = selected[:, None] * block_size + torch.arange(block_size)
+    if not torch.all(block_indices == selected[:, None]) or not torch.equal(
+        blocks.sort(dim=1).values, expected
+    ):
+        raise ValueError("Weight selection must preserve complete quantization blocks")
+    return selected
+
+
 def select_quantized_weight_export_state(
     state: _QuantizedWeightExportState,
     weight_dim: int,
@@ -1233,6 +1348,8 @@ def select_quantized_weight_export_state(
     shape = list(state.weight_shape)
     shape[weight_dim] = indices.numel()
     packed_dim = state.packing_permutation.index(weight_dim)
+    if state.block_size > 1 and packed_dim == ndim - 1:
+        _selected_block_indices(indices, state.block_size)
     return _QuantizedWeightExportState(
         state.quantization_format,
         state.block_size,
@@ -1323,31 +1440,53 @@ def export_quantized_weight_tensors(
             state.block_size,
             weights_scaling_factor_2=weight_scale_2.to(packed_weight.device),
         )[0]
-    else:
+    elif "weight_scale" in records:
         weight_scale = records["weight_scale"].value
+    elif state.quantization_format == QUANTIZATION_MXFP8:
+        weight_scale = MXFP8QTensor.get_weights_scaling_factor(packed_weight)
+    elif state.quantization_format in _MXFP4_EXPORT_FORMATS:
+        quantized_weight, weight_scale = MXFP4QTensor.quantize(
+            packed_weight.to(dtype), block_size=state.block_size
+        )
+        quantized_weight = quantized_weight._quantized_data
+    else:
+        raise _UnsupportedQuantizedWeightExportFormat(
+            f"Functional export does not support {state.quantization_format!r}"
+        )
 
     weight_scale = weight_scale.to(packed_weight.device)
     if weight_scale_2 is not None:
         weight_scale_2 = weight_scale_2.to(packed_weight.device)
 
-    quantized_weight = to_quantized_weight(
-        packed_weight.to(dtype),
-        weight_scale,
-        state.quantization_format,
-        weight_scale_2,
-        state.block_size,
-    )
+    if state.quantization_format not in _MXFP4_EXPORT_FORMATS:
+        quantized_weight = to_quantized_weight(
+            packed_weight.to(dtype),
+            weight_scale,
+            state.quantization_format,
+            weight_scale_2,
+            state.block_size,
+        )
     attrs = quantizer_attr_names(weight_name)
     output = OrderedDict(
         ((weight_name, _restore_packing_permutation(quantized_weight, state.packing_permutation)),)
     )
 
-    if state.quantization_format in _NVFP4_EXPORT_FORMATS:
-        weight_scale_record = _state_tensor(
+    if "weight_scale" in records:
+        scale_state = records["weight_scale"]
+        weight_scale_record = _ExportStateTensor(
+            scale_state.name,
+            weight_scale,
+            scale_state.axes,
+            scale_state.block_sizes,
+        )
+    elif state.quantization_format in _NVFP4_EXPORT_FORMATS | {
+        QUANTIZATION_MXFP8,
+    } | _MXFP4_EXPORT_FORMATS:
+        weight_scale_record = _block_scale_state(
             "weight_scale",
             weight_scale,
-            axes=tuple(range(packed_weight.ndim)),
-            block_sizes=(1,) * (packed_weight.ndim - 1) + (state.block_size,),
+            tuple(packed_weight.shape),
+            {-1: state.block_size},
             cpu=False,
         )
     else:
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 6bc8df91615..d2e2ae8428d 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -100,6 +100,7 @@
     revert_weight_conversion_quant_aware,
 )
 from .quant_utils import (
+    _UnsupportedQuantizedWeightExportFormat,
     capture_quantized_weight_export_state,
     export_quantized_weight_tensors,
     fuse_prequant_layernorm,
@@ -609,20 +610,28 @@ def _export_quantized_weight(
         sub_module, quantizer_attrs.output_quantizer, None
     )
 
-    from .quant_utils import _FUNCTIONAL_EXPORT_FORMATS
-
-    if not isinstance(weight, QTensorWrapper) and quantization_format in _FUNCTIONAL_EXPORT_FORMATS:
-        state = capture_quantized_weight_export_state(sub_module, weight_name)
-        exported = export_quantized_weight_tensors(weight, state, dtype, weight_name)
-        setattr(
-            sub_module,
-            weight_name,
-            nn.Parameter(exported.pop(weight_name), requires_grad=False),
-        )
-        for name, value in exported.items():
-            sub_module.register_buffer(name, value)
-        torch.cuda.empty_cache()
-        return
+    if not isinstance(weight, QTensorWrapper):
+        try:
+            state = capture_quantized_weight_export_state(sub_module, weight_name)
+        except _UnsupportedQuantizedWeightExportFormat:
+            pass
+        else:
+            exported = export_quantized_weight_tensors(weight, state, dtype, weight_name)
+            setattr(
+                sub_module,
+                weight_name,
+                nn.Parameter(exported.pop(weight_name), requires_grad=False),
+            )
+            for name, value in exported.items():
+                sub_module.register_buffer(name, value)
+            if (
+                quantization_format == QUANTIZATION_MXFP8
+                and hasattr(weight_quantizer, "_scale")
+                and weight_quantizer._scale is not None
+            ):
+                del weight_quantizer._scale
+            torch.cuda.empty_cache()
+            return
 
     # Already real-quantized weights (``mtq.compress`` / ``hf_ptq --low_memory_mode``) hold packed
     # nibbles -- half the logical last dim -- so per-block scales cannot be recomputed from them.
diff --git a/tests/gpu/torch/export/test_export_weight_gpu.py b/tests/gpu/torch/export/test_export_weight_gpu.py
index 86d7d591031..eacbb6e02df 100644
--- a/tests/gpu/torch/export/test_export_weight_gpu.py
+++ b/tests/gpu/torch/export/test_export_weight_gpu.py
@@ -24,7 +24,9 @@
 from torch.nn import init
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export.model_config import QUANTIZATION_MXFP8
 from modelopt.torch.export.quant_utils import (
+    build_hf_quantization_config,
     capture_quantized_weight_export_state,
     export_quantized_weight_tensors,
     get_activation_scaling_factor,
@@ -32,7 +34,9 @@
     get_weight_block_size,
     get_weight_scaling_factor,
     get_weight_scaling_factor_2,
+    permute_quantized_weight_export_state,
     postprocess_state_dict,
+    select_quantized_weight_export_state,
     to_quantized_weight,
 )
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
@@ -181,24 +185,165 @@ def test_functional_nvfp4_export_matches_existing_helpers_without_mutation(quant
         mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
         mtq.FP8_PER_CHANNEL_PER_TOKEN_CFG,
         mtq.MXFP8_DEFAULT_CFG,
+        mtq.MXFP4_DEFAULT_CFG,
+        mtq.W4A8_MXFP4_FP8_CFG,
+        mtq.W4A8_NVFP4_FP8_CFG,
     ],
 )
-def test_unqualified_formats_keep_existing_export_path(monkeypatch, quant_cfg):
+def test_functional_export_matches_existing_noninteger_helpers(monkeypatch, quant_cfg):
     in_features = 256
+    torch.manual_seed(0)
     module = nn.Linear(in_features, in_features, bias=False, device="cuda", dtype=torch.bfloat16)
     calib_input = torch.randn(2, 4, in_features, device="cuda", dtype=torch.bfloat16)
     module = mtq.quantize(module, copy.deepcopy(quant_cfg), lambda model: model(calib_input))
+    original_weight = module.weight.detach().clone()
+    original_buffers = {name: value.detach().clone() for name, value in module.named_buffers()}
 
-    def fail_functional_capture(*args, **kwargs):
-        raise AssertionError("unqualified format entered the functional export path")
+    state = capture_quantized_weight_export_state(module)
+    actual = export_quantized_weight_tensors(module.weight, state, torch.float16)
+
+    quantization_format = get_quantization_format(module)
+    weight_scale = get_weight_scaling_factor(module)
+    weight_scale_2 = get_weight_scaling_factor_2(module)
+    expected = {
+        "weight": to_quantized_weight(
+            module.weight.to(torch.float16),
+            weight_scale,
+            quantization_format,
+            weight_scale_2,
+            get_weight_block_size(module),
+        ),
+        "weight_scale": weight_scale,
+    }
+    if weight_scale_2 is not None:
+        expected["weight_scale_2"] = weight_scale_2.squeeze()
+    if module.input_quantizer.is_enabled and module.input_quantizer.amax is not None:
+        expected["input_scale"] = get_activation_scaling_factor(module).squeeze()
+
+    assert actual.keys() == expected.keys()
+    for name, value in actual.items():
+        torch.testing.assert_close(value, expected[name], rtol=0, atol=0)
+    torch.testing.assert_close(module.weight, original_weight)
+    assert set(dict(module.named_buffers())) == set(original_buffers)
+    for name, value in module.named_buffers():
+        torch.testing.assert_close(value, original_buffers[name])
+
+    called = False
+
+    def record_functional_capture(*args, **kwargs):
+        nonlocal called
+        called = True
+        return capture_quantized_weight_export_state(*args, **kwargs)
 
     monkeypatch.setattr(
         "modelopt.torch.export.unified_export_hf.capture_quantized_weight_export_state",
-        fail_functional_capture,
+        record_functional_capture,
     )
     _export_quantized_weight(module, torch.float16)
 
-    assert hasattr(module, "weight_scale")
+    assert called
+    for name, value in actual.items():
+        torch.testing.assert_close(getattr(module, name), value, rtol=0, atol=0)
+    if quantization_format == QUANTIZATION_MXFP8:
+        assert hasattr(module, "weight_scale")
+        assert not hasattr(module.weight_quantizer, "_scale")
+
+
+def test_functional_mxfp8_preserves_cached_scale_during_selection():
+    features = 256
+    module = nn.Linear(features, features, bias=False, device="cuda", dtype=torch.bfloat16)
+    calib_input = torch.randn(2, 4, features, device="cuda", dtype=torch.bfloat16)
+    module = mtq.quantize(
+        module,
+        copy.deepcopy(mtq.MXFP8_DEFAULT_CFG),
+        lambda model: model(calib_input),
+    )
+    cached_scale = get_weight_scaling_factor(module).clone()
+    cached_scale[0].add_(1)
+    module.weight_quantizer._scale = cached_scale
+
+    indices = torch.arange(features // 2)
+    state = select_quantized_weight_export_state(
+        capture_quantized_weight_export_state(module),
+        0,
+        indices,
+    )
+    exported = export_quantized_weight_tensors(
+        module.weight.index_select(0, indices.to(module.weight.device)),
+        state,
+        torch.float16,
+    )
+
+    torch.testing.assert_close(exported["weight_scale"], cached_scale[: features // 2])
+
+
+def test_functional_block_scale_layout_follows_weight_permutation():
+    features = 256
+    module = nn.Linear(features, features, bias=False, device="cuda", dtype=torch.bfloat16)
+    calib_input = torch.randn(2, 4, features, device="cuda", dtype=torch.bfloat16)
+    module = mtq.quantize(
+        module,
+        copy.deepcopy(mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG),
+        lambda model: model(calib_input),
+    )
+    state = capture_quantized_weight_export_state(module)
+    base = export_quantized_weight_tensors(module.weight, state, torch.float16)
+    transposed = export_quantized_weight_tensors(
+        module.weight.T,
+        permute_quantized_weight_export_state(state, (1, 0)),
+        torch.float16,
+    )
+
+    torch.testing.assert_close(transposed["weight"], base["weight"].T)
+    torch.testing.assert_close(
+        transposed["weight_scale"],
+        base["weight_scale"].permute(2, 3, 0, 1),
+    )
+
+
+def test_weight_derived_mxfp4_state_rejects_partial_block_selection():
+    features = 256
+    module = nn.Linear(features, features, bias=False, device="cuda", dtype=torch.bfloat16)
+    calib_input = torch.randn(2, 4, features, device="cuda", dtype=torch.bfloat16)
+    module = mtq.quantize(
+        module,
+        copy.deepcopy(mtq.MXFP4_DEFAULT_CFG),
+        lambda model: model(calib_input),
+    )
+
+    with pytest.raises(ValueError, match="complete quantization blocks"):
+        select_quantized_weight_export_state(
+            capture_quantized_weight_export_state(module),
+            1,
+            (0,),
+        )
+
+
+def test_mixed_noninteger_states_build_one_canonical_config():
+    features = 256
+    states = {}
+    for name, quant_cfg in (
+        ("model.layers.0.self_attn.q_proj.weight", mtq.FP8_DEFAULT_CFG),
+        ("model.layers.0.self_attn.k_proj.weight", mtq.MXFP8_DEFAULT_CFG),
+        ("model.layers.0.mlp.down_proj.weight", mtq.NVFP4_DEFAULT_CFG),
+    ):
+        module = nn.Linear(features, features, bias=False, device="cuda", dtype=torch.bfloat16)
+        calib_input = torch.randn(2, 4, features, device="cuda", dtype=torch.bfloat16)
+        module = mtq.quantize(
+            module,
+            copy.deepcopy(quant_cfg),
+            lambda model: model(calib_input),
+        )
+        states[name] = capture_quantized_weight_export_state(module)
+
+    config = build_hf_quantization_config(states)
+
+    assert config["quant_algo"] == "MIXED_PRECISION"
+    assert set(config["quantized_layers"]) == {
+        "model.layers.0.self_attn.q_proj",
+        "model.layers.0.self_attn.k_proj",
+        "model.layers.0.mlp.down_proj",
+    }
 
 
 def test_export_compressed_nvfp4_weight():

From f0678d2cdf50bb32baf350ecd58415530348e059 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Sun, 16 Aug 2026 11:02:42 -0700
Subject: [PATCH 07/10] Fix functional export lint invariants

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/quant_utils.py       | 34 ++++++++++++----------
 modelopt/torch/export/unified_export_hf.py |  4 +--
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 16c30d4dc33..9780a3f38ac 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -487,7 +487,7 @@ def _get_quantization_from_quantizers(
     layer: nn.Module,
     weight_quantizer: TensorQuantizer | SequentialQuantizer | None,
     input_quantizer: TensorQuantizer | SequentialQuantizer | None,
-) -> str:
+) -> str | None:
     if weight_quantizer is None or not weight_quantizer.is_enabled:
         return QUANTIZATION_NONE
 
@@ -938,7 +938,7 @@ def to_quantized_weight(
 _MXFP4_EXPORT_FORMATS = {QUANTIZATION_MXFP4, QUANTIZATION_W4A8_MXFP4_FP8}
 
 
-class _UnsupportedQuantizedWeightExportFormat(NotImplementedError):
+class _UnsupportedQuantizedWeightExportFormatError(NotImplementedError):
     pass
 
 
@@ -1072,8 +1072,7 @@ def _block_scale_state(
         if isinstance(dim, int) and block_size is not None:
             block_sizes[_normalize_weight_dim(dim, ndim)] = int(block_size)
     expected_shape = tuple(
-        (size + block_size - 1) // block_size
-        for size, block_size in zip(packed_shape, block_sizes)
+        (size + block_size - 1) // block_size for size, block_size in zip(packed_shape, block_sizes)
     )
     if value.numel() != torch.Size(expected_shape).numel():
         raise RuntimeError(
@@ -1127,15 +1126,17 @@ def capture_quantized_weight_export_state(
 
     quantized_view, weight_quantizer = _resolve_weight_quantizer(module, weight_name)
     input_quantizer = _input_quantizer(module, weight_name)
+    if not weight_quantizer.is_enabled:
+        raise RuntimeError(f"Weight quantizer for {weight_name!r} is disabled")
     quantization_format = _get_quantization_from_quantizers(
         module, weight_quantizer, input_quantizer
     )
+    if quantization_format is None:
+        raise RuntimeError(f"Unable to resolve quantization format for {weight_name!r}")
     if isinstance(weight_quantizer, SequentialQuantizer):
-        raise _UnsupportedQuantizedWeightExportFormat(
+        raise _UnsupportedQuantizedWeightExportFormatError(
             f"Functional export does not support {quantization_format!r}"
         )
-    if not weight_quantizer.is_enabled:
-        raise RuntimeError(f"Weight quantizer for {weight_name!r} is disabled")
 
     permutation = _packing_permutation(weight, quantized_view)
     packed_shape = tuple(weight.shape[axis] for axis in permutation)
@@ -1192,9 +1193,7 @@ def capture_quantized_weight_export_state(
         weight_scale = get_scaling_factor(weight_quantizer)
         if weight_scale is None:
             raise RuntimeError(f"Missing calibrated block scale for {weight_name!r}")
-        tensors.append(
-            _block_scale_state("weight_scale", weight_scale, packed_shape, block_config)
-        )
+        tensors.append(_block_scale_state("weight_scale", weight_scale, packed_shape, block_config))
     elif quantization_format == QUANTIZATION_MXFP8:
         cached_scale = getattr(weight_quantizer, "_scale", None)
         if cached_scale is not None:
@@ -1202,7 +1201,7 @@ def capture_quantized_weight_export_state(
                 _block_scale_state("weight_scale", cached_scale, packed_shape, block_config)
             )
     elif quantization_format not in _MXFP4_EXPORT_FORMATS:
-        raise _UnsupportedQuantizedWeightExportFormat(
+        raise _UnsupportedQuantizedWeightExportFormatError(
             f"Functional export does not support {quantization_format!r}"
         )
 
@@ -1450,7 +1449,7 @@ def export_quantized_weight_tensors(
         )
         quantized_weight = quantized_weight._quantized_data
     else:
-        raise _UnsupportedQuantizedWeightExportFormat(
+        raise _UnsupportedQuantizedWeightExportFormatError(
             f"Functional export does not support {state.quantization_format!r}"
         )
 
@@ -1479,9 +1478,14 @@ def export_quantized_weight_tensors(
             scale_state.axes,
             scale_state.block_sizes,
         )
-    elif state.quantization_format in _NVFP4_EXPORT_FORMATS | {
-        QUANTIZATION_MXFP8,
-    } | _MXFP4_EXPORT_FORMATS:
+    elif (
+        state.quantization_format
+        in _NVFP4_EXPORT_FORMATS
+        | {
+            QUANTIZATION_MXFP8,
+        }
+        | _MXFP4_EXPORT_FORMATS
+    ):
         weight_scale_record = _block_scale_state(
             "weight_scale",
             weight_scale,
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index d2e2ae8428d..c7ea9629234 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -100,7 +100,7 @@
     revert_weight_conversion_quant_aware,
 )
 from .quant_utils import (
-    _UnsupportedQuantizedWeightExportFormat,
+    _UnsupportedQuantizedWeightExportFormatError,
     capture_quantized_weight_export_state,
     export_quantized_weight_tensors,
     fuse_prequant_layernorm,
@@ -613,7 +613,7 @@ def _export_quantized_weight(
     if not isinstance(weight, QTensorWrapper):
         try:
             state = capture_quantized_weight_export_state(sub_module, weight_name)
-        except _UnsupportedQuantizedWeightExportFormat:
+        except _UnsupportedQuantizedWeightExportFormatError:
             pass
         else:
             exported = export_quantized_weight_tensors(weight, state, dtype, weight_name)

From 192c88993d664870e346bb33d690c055e1adf978 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 26 Aug 2026 01:20:54 -0700
Subject: [PATCH 08/10] chore: add required test license header

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 .../unit/torch/export/test_quantized_weight_state.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index 751a987a784..9500e5a679e 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -1,5 +1,17 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import copy
 import pickle

From dcca192bd08d5125765130aede30021fde91b2a0 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Fri, 28 Aug 2026 06:46:18 -0700
Subject: [PATCH 09/10] docs(export): clarify legacy quantization fallback

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index c7ea9629234..47311642b27 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -614,6 +614,7 @@ def _export_quantized_weight(
         try:
             state = capture_quantized_weight_export_state(sub_module, weight_name)
         except _UnsupportedQuantizedWeightExportFormatError:
+            # AWQ and SmoothQuant remain on the legacy format-specific path below.
             pass
         else:
             exported = export_quantized_weight_tensors(weight, state, dtype, weight_name)

From 2615d4bf4d9b14880da5b4edbfbfcdf8220b138f Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Fri, 28 Aug 2026 22:37:55 -0700
Subject: [PATCH 10/10] fix: separate real quant specs from export state

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 152 +++++++++++++++---
 modelopt/torch/export/unified_export_hf.py    |  27 ----
 .../export/test_quantized_weight_state.py     |  81 ++++++++++
 3 files changed, 215 insertions(+), 45 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 9780a3f38ac..287d217c6f1 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -963,6 +963,32 @@ class _QuantizedWeightExportState:
     four_over_six: bool = False
 
 
+@dataclass(frozen=True)
+class _QuantizedWeightExportSpec:
+    """Stable deployment metadata for one logical quantized weight."""
+
+    quantization_format: str
+    block_size: int
+
+
+@dataclass(frozen=True)
+class _ExportStateTensorMetadata:
+    name: str
+    axes: tuple[int, ...]
+    block_sizes: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class _QuantizedWeightExportMetadata:
+    quantization_format: str
+    block_size: int
+    weight_shape: tuple[int, ...]
+    tensors: tuple[_ExportStateTensorMetadata, ...]
+    packing_permutation: tuple[int, ...]
+    static_nvfp4: bool
+    four_over_six: bool
+
+
 def _same_storage(left: object, right: object) -> bool:
     if left is right:
         return True
@@ -979,7 +1005,7 @@ def _same_storage(left: object, right: object) -> bool:
 
 def _resolve_weight_quantizer(
     module: nn.Module, weight_name: str
-) -> tuple[torch.Tensor, TensorQuantizer | SequentialQuantizer]:
+) -> tuple[torch.Tensor, TensorQuantizer | SequentialQuantizer] | None:
     weight = getattr(module, weight_name)
     iter_weights = getattr(module, "iter_weights_for_calibration", None)
     if iter_weights is not None:
@@ -991,7 +1017,7 @@ def _resolve_weight_quantizer(
     if quantizer is None and weight_name.startswith("weight"):
         quantizer = representative_weight_quantizer(module)
     if quantizer is None:
-        raise RuntimeError(f"Missing weight quantizer for {weight_name!r}")
+        return None
     return weight, quantizer
 
 
@@ -1025,7 +1051,7 @@ def _state_tensor(
     block_sizes: tuple[int, ...] = (),
     cpu: bool = True,
 ) -> _ExportStateTensor:
-    value = value.detach().cpu().clone() if cpu else value.detach()
+    value = value.detach().cpu().clone() if cpu else value.detach().clone()
     if value.numel() == 1 and not axes and not block_sizes:
         return _ExportStateTensor(name, value.reshape(()))
     if len(axes) != value.ndim or len(block_sizes) != value.ndim:
@@ -1041,6 +1067,8 @@ def _axis_scale_state(
     value: torch.Tensor,
     packed_shape: tuple[int, ...],
     axis: int | Sequence[int],
+    *,
+    cpu: bool = True,
 ) -> _ExportStateTensor:
     axes = (axis,) if isinstance(axis, int) else tuple(axis)
     axes = tuple(_normalize_weight_dim(dim, len(packed_shape)) for dim in axes)
@@ -1055,6 +1083,7 @@ def _axis_scale_state(
         value.reshape(expected_shape),
         axes=axes,
         block_sizes=(1,) * len(axes),
+        cpu=cpu,
     )
 
 
@@ -1115,19 +1144,30 @@ def _input_quantizer(module: nn.Module, weight_name: str):
     return quantizer
 
 
-def capture_quantized_weight_export_state(
+def _resolve_quantized_weight_export_inputs(
     module: nn.Module,
-    weight_name: str = "weight",
-) -> _QuantizedWeightExportState:
-    """Capture detached export state without mutating the quantized module."""
+    weight_name: str,
+) -> (
+    tuple[
+        torch.Tensor,
+        torch.Tensor,
+        TensorQuantizer,
+        TensorQuantizer | None,
+        str,
+    ]
+    | None
+):
     weight = getattr(module, weight_name)
     if isinstance(weight, QTensorWrapper):
         raise NotImplementedError("Functional export requires an uncompressed source weight")
 
-    quantized_view, weight_quantizer = _resolve_weight_quantizer(module, weight_name)
+    resolved = _resolve_weight_quantizer(module, weight_name)
+    if resolved is None:
+        return None
+    quantized_view, weight_quantizer = resolved
     input_quantizer = _input_quantizer(module, weight_name)
     if not weight_quantizer.is_enabled:
-        raise RuntimeError(f"Weight quantizer for {weight_name!r} is disabled")
+        return None
     quantization_format = _get_quantization_from_quantizers(
         module, weight_quantizer, input_quantizer
     )
@@ -1137,6 +1177,20 @@ def capture_quantized_weight_export_state(
         raise _UnsupportedQuantizedWeightExportFormatError(
             f"Functional export does not support {quantization_format!r}"
         )
+    return weight, quantized_view, weight_quantizer, input_quantizer, quantization_format
+
+
+def capture_quantized_weight_export_state(
+    module: nn.Module,
+    weight_name: str = "weight",
+    *,
+    cpu: bool = True,
+) -> _QuantizedWeightExportState | None:
+    """Capture detached state for one weight, or ``None`` when it is not quantized."""
+    resolved = _resolve_quantized_weight_export_inputs(module, weight_name)
+    if resolved is None:
+        return None
+    weight, quantized_view, weight_quantizer, input_quantizer, quantization_format = resolved
 
     permutation = _packing_permutation(weight, quantized_view)
     packed_shape = tuple(weight.shape[axis] for axis in permutation)
@@ -1162,21 +1216,22 @@ def capture_quantized_weight_export_state(
                 per_block_amax.reshape(block_shape),
                 axes=tuple(range(len(packed_shape))),
                 block_sizes=(1,) * (len(packed_shape) - 1) + (block_size,),
+                cpu=cpu,
             )
         )
-        tensors.append(_state_tensor("weight_global_amax", global_amax))
+        tensors.append(_state_tensor("weight_global_amax", global_amax, cpu=cpu))
     elif quantization_format in _NVFP4_EXPORT_FORMATS:
         weight_scale_2 = (
             weight_quantizer._amax.float() / 448.0
             if quantization_format == QUANTIZATION_W4A8_NVFP4_FP8
             else NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
         )
-        tensors.append(_state_tensor("weight_scale_2", weight_scale_2))
+        tensors.append(_state_tensor("weight_scale_2", weight_scale_2, cpu=cpu))
     elif quantization_format == QUANTIZATION_FP8:
         weight_scale = get_scaling_factor(weight_quantizer)
         if weight_scale is None:
             raise RuntimeError(f"Missing calibrated weight scale for {weight_name!r}")
-        tensors.append(_state_tensor("weight_scale", weight_scale))
+        tensors.append(_state_tensor("weight_scale", weight_scale, cpu=cpu))
     elif quantization_format == QUANTIZATION_FP8_PC_PT:
         weight_scale = get_scaling_factor(weight_quantizer)
         if weight_scale is None or weight_quantizer.axis is None:
@@ -1187,18 +1242,23 @@ def capture_quantized_weight_export_state(
                 weight_scale,
                 packed_shape,
                 weight_quantizer.axis,
+                cpu=cpu,
             )
         )
     elif quantization_format == QUANTIZATION_FP8_PB_WO:
         weight_scale = get_scaling_factor(weight_quantizer)
         if weight_scale is None:
             raise RuntimeError(f"Missing calibrated block scale for {weight_name!r}")
-        tensors.append(_block_scale_state("weight_scale", weight_scale, packed_shape, block_config))
+        tensors.append(
+            _block_scale_state("weight_scale", weight_scale, packed_shape, block_config, cpu=cpu)
+        )
     elif quantization_format == QUANTIZATION_MXFP8:
         cached_scale = getattr(weight_quantizer, "_scale", None)
         if cached_scale is not None:
             tensors.append(
-                _block_scale_state("weight_scale", cached_scale, packed_shape, block_config)
+                _block_scale_state(
+                    "weight_scale", cached_scale, packed_shape, block_config, cpu=cpu
+                )
             )
     elif quantization_format not in _MXFP4_EXPORT_FORMATS:
         raise _UnsupportedQuantizedWeightExportFormatError(
@@ -1213,7 +1273,7 @@ def capture_quantized_weight_export_state(
                 if quantization_format == QUANTIZATION_NVFP4
                 else get_scaling_factor(input_quantizer)
             )
-            tensors.append(_state_tensor("input_scale", input_scale))
+            tensors.append(_state_tensor("input_scale", input_scale, cpu=cpu))
 
     return _QuantizedWeightExportState(
         quantization_format=quantization_format,
@@ -1226,6 +1286,62 @@ def capture_quantized_weight_export_state(
     )
 
 
+def get_quantized_weight_export_spec(
+    module: nn.Module,
+    weight_name: str = "weight",
+) -> _QuantizedWeightExportSpec | None:
+    """Return stable deployment metadata without retaining quantizer tensors."""
+    resolved = _resolve_quantized_weight_export_inputs(module, weight_name)
+    if resolved is None:
+        return None
+    _, _, weight_quantizer, _, quantization_format = resolved
+    block_config = getattr(weight_quantizer, "block_sizes", None) or {}
+    block_size = int(block_config.get(-1, 0)) if isinstance(block_config, dict) else 0
+    return _QuantizedWeightExportSpec(quantization_format, block_size)
+
+
+def split_quantized_weight_export_state(
+    state: _QuantizedWeightExportState,
+) -> tuple[object, tuple[torch.Tensor, ...]]:
+    """Separate opaque state metadata from tensor values for typed transport."""
+    metadata = _QuantizedWeightExportMetadata(
+        state.quantization_format,
+        state.block_size,
+        state.weight_shape,
+        tuple(
+            _ExportStateTensorMetadata(record.name, record.axes, record.block_sizes)
+            for record in state.tensors
+        ),
+        state.packing_permutation,
+        state.static_nvfp4,
+        state.four_over_six,
+    )
+    return metadata, tuple(record.value for record in state.tensors)
+
+
+def restore_quantized_weight_export_state(
+    metadata: object,
+    tensors: Sequence[torch.Tensor],
+) -> _QuantizedWeightExportState:
+    """Restore opaque export state after its tensor values have been transported."""
+    if not isinstance(metadata, _QuantizedWeightExportMetadata):
+        raise TypeError(f"Invalid quantized weight export metadata: {type(metadata).__name__}")
+    if len(tensors) != len(metadata.tensors):
+        raise ValueError(f"Expected {len(metadata.tensors)} export tensors, got {len(tensors)}")
+    return _QuantizedWeightExportState(
+        metadata.quantization_format,
+        metadata.block_size,
+        metadata.weight_shape,
+        tuple(
+            _ExportStateTensor(spec.name, tensor, spec.axes, spec.block_sizes)
+            for spec, tensor in zip(metadata.tensors, tensors, strict=True)
+        ),
+        metadata.packing_permutation,
+        metadata.static_nvfp4,
+        metadata.four_over_six,
+    )
+
+
 def _normalize_weight_dim(weight_dim: int, ndim: int) -> int:
     if not -ndim <= weight_dim < ndim:
         raise IndexError(f"Weight dimension {weight_dim} is invalid for a rank-{ndim} weight")
@@ -1524,10 +1640,10 @@ def _quantized_layer_name(weight_name: str) -> str:
 
 
 def build_hf_quantization_config(
-    named_states: Mapping[str, _QuantizedWeightExportState | None]
-    | Iterable[tuple[str, _QuantizedWeightExportState | None]],
+    named_states: Mapping[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None]
+    | Iterable[tuple[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None]],
 ) -> dict[str, Any]:
-    """Build canonical ModelOpt HF configuration from named opaque states."""
+    """Build canonical ModelOpt HF configuration from named states or specs."""
     layers: dict[str, tuple[str, int] | None] = {}
     for weight_name, state in dict(named_states).items():
         layer_name = _quantized_layer_name(weight_name)
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 47311642b27..77429b1cfaf 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -100,9 +100,6 @@
     revert_weight_conversion_quant_aware,
 )
 from .quant_utils import (
-    _UnsupportedQuantizedWeightExportFormatError,
-    capture_quantized_weight_export_state,
-    export_quantized_weight_tensors,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -610,30 +607,6 @@ def _export_quantized_weight(
         sub_module, quantizer_attrs.output_quantizer, None
     )
 
-    if not isinstance(weight, QTensorWrapper):
-        try:
-            state = capture_quantized_weight_export_state(sub_module, weight_name)
-        except _UnsupportedQuantizedWeightExportFormatError:
-            # AWQ and SmoothQuant remain on the legacy format-specific path below.
-            pass
-        else:
-            exported = export_quantized_weight_tensors(weight, state, dtype, weight_name)
-            setattr(
-                sub_module,
-                weight_name,
-                nn.Parameter(exported.pop(weight_name), requires_grad=False),
-            )
-            for name, value in exported.items():
-                sub_module.register_buffer(name, value)
-            if (
-                quantization_format == QUANTIZATION_MXFP8
-                and hasattr(weight_quantizer, "_scale")
-                and weight_quantizer._scale is not None
-            ):
-                del weight_quantizer._scale
-            torch.cuda.empty_cache()
-            return
-
     # Already real-quantized weights (``mtq.compress`` / ``hf_ptq --low_memory_mode``) hold packed
     # nibbles -- half the logical last dim -- so per-block scales cannot be recomputed from them.
     # Use the scale the quantizer captured at compression time instead.
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index 9500e5a679e..98954e3bb3e 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -20,15 +20,19 @@
 import torch
 import torch.nn as nn
 
+import modelopt.torch.export.quant_utils as quant_utils
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
 from modelopt.torch.export.quant_utils import (
     build_hf_quantization_config,
     capture_quantized_weight_export_state,
     export_quantized_weight_tensors,
+    get_quantized_weight_export_spec,
     merge_quantized_weight_export_states,
     permute_quantized_weight_export_state,
+    restore_quantized_weight_export_state,
     select_quantized_weight_export_state,
+    split_quantized_weight_export_state,
 )
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer
@@ -88,6 +92,25 @@ def iter_weights_for_calibration(self):
             yield getattr(self, name).transpose(-1, -2), getattr(self, f"{name}_weight_quantizer")
 
 
+class _GroupedWeights(nn.Module):
+    """Minimal TEGroupedLinear-style numbered-weight layout."""
+
+    def __init__(self):
+        super().__init__()
+        self.weight0 = nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4))
+        self.weight1 = nn.Parameter(torch.arange(16, 32, dtype=torch.float32).reshape(4, 4))
+        cfg = QuantizerAttributeConfig(num_bits=(4, 3))
+        self.quantizers = nn.ModuleList([TensorQuantizer(cfg), TensorQuantizer(cfg)])
+        for index, quantizer in enumerate(self.quantizers, start=1):
+            quantizer._amax = torch.tensor(float(index))
+        self.input_quantizer = TensorQuantizer()
+        self.input_quantizer.disable()
+
+    def iter_weights_for_calibration(self):
+        yield self.weight0, self.quantizers[0]
+        yield self.weight1, self.quantizers[1]
+
+
 def test_capture_does_not_modify_zero_amax():
     module = _fp8_linear()
     module.weight_quantizer._amax.zero_()
@@ -133,6 +156,64 @@ def test_export_state_round_trips_through_object_transport():
     assert exported["weight"].shape == module.weight.shape
 
 
+def test_capture_resolves_numbered_grouped_weights_by_storage():
+    module = _GroupedWeights()
+
+    state0 = capture_quantized_weight_export_state(module, "weight0")
+    state1 = capture_quantized_weight_export_state(module, "weight1")
+
+    assert state0 is not None and state1 is not None
+    assert state0.quantization_format == state1.quantization_format == "fp8"
+    assert state0.tensors[0].value.item() != state1.tensors[0].value.item()
+
+
+def test_unquantized_weight_has_no_export_state_or_spec():
+    module = nn.Linear(4, 4, bias=False)
+
+    assert capture_quantized_weight_export_state(module) is None
+    assert get_quantized_weight_export_spec(module) is None
+
+
+def test_export_state_split_restore_preserves_output():
+    module = _fp8_linear()
+    state = capture_quantized_weight_export_state(module)
+    assert state is not None
+    metadata, tensors = split_quantized_weight_export_state(state)
+
+    restored = restore_quantized_weight_export_state(metadata, tensors)
+
+    expected = export_quantized_weight_tensors(module.weight, state, torch.float32)
+    actual = export_quantized_weight_tensors(module.weight, restored, torch.float32)
+    assert actual.keys() == expected.keys()
+    for name in actual:
+        torch.testing.assert_close(actual[name], expected[name], rtol=0, atol=0)
+
+
+def test_export_spec_builds_config_without_tensor_state():
+    spec = get_quantized_weight_export_spec(_fp8_linear())
+    assert spec is not None
+
+    config = build_hf_quantization_config({"model.layers.0.proj.weight": spec})
+
+    assert config["quant_algo"] == "FP8"
+
+
+def test_export_spec_does_not_materialize_static_scale_state(monkeypatch):
+    weight = torch.arange(32, dtype=torch.float32).reshape(2, 16) / 32
+    module = _static_w4a16_linear(weight, weight.abs().amax(dim=1, keepdim=True), torch.tensor(1.0))
+    monkeypatch.setattr(
+        quant_utils,
+        "_state_tensor",
+        lambda *args, **kwargs: pytest.fail("export spec materialized tensor state"),
+    )
+
+    spec = get_quantized_weight_export_spec(module)
+
+    assert spec is not None
+    assert spec.quantization_format == "w4a16_nvfp4"
+    assert spec.block_size == 16
+
+
 def test_static_nvfp4_merge_recomputes_scales_from_merged_amax():
     left_weight = torch.arange(32, dtype=torch.float32).reshape(2, 16) / 32
     right_weight = torch.arange(32, 64, dtype=torch.float32).reshape(2, 16) / 16