From 05bd76c413b88f5bce17f09a53f3c0946094618d Mon Sep 17 00:00:00 2001
From: SID <99672439+SID-6921@users.noreply.github.com>
Date: Sat, 26 Sep 2026 19:09:56 -0400
Subject: [PATCH 1/8] fix(onnx): trace Transpose between DequantizeLinear and
 its consumer in qdq_to_dq

qdq_to_dq assumed DequantizeLinear feeds Conv/Gemm/MatMul directly, with only
Cast tolerated in between. A weight that reaches MatMul transposed -- as in
MaxViT's attention blocks -- put Transpose in front of _convert_weight, whose
axis_map has no entry for it:

    Unsupported op_type for real weight quantization: Transpose

Tracing through the Transpose is not enough on its own. axis_map is expressed in
the consumer's layout, while the scale indexes the weight as stored, so the axis
has to be mapped back through the Transpose's permutation. Skipping the Transpose
the way Cast is skipped raises "Scale shape (96,) does not match weight shape
(96, 32) along axis 1" on non-square weights, and on square weights (out == in,
common for attention projections) it silently quantizes along the wrong axis --
907 of 1024 elements wrong in a 32x32 case.

_get_successive_consumers now walks any Cast/Transpose chain and returns the
composed permutation; _convert_weight applies it to the axis, asserting the
permutation matches the weight's rank.

Verified on the reported model: quantized maxvit_tiny_rw_256 exported with
dq_only=True raises on main and exports here, and the resulting DQ-only graph is
bit-identical to the QDQ graph under onnxruntime (max abs diff 0.0 across 457
quantizers).

Fixes #2442

Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  1 +
 modelopt/onnx/quantization/qdq_utils.py       | 37 ++++++++---
 .../unit/onnx/quantization/test_qdq_utils.py  | 62 +++++++++++++++++++
 3 files changed, 93 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 4f621df5dea..f0508daa690 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -94,6 +94,7 @@ Changelog
 - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected.
 - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths.
 - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand.
+- Fix ``qdq_to_dq`` failing with ``Unsupported op_type for real weight quantization: Transpose`` when a quantized weight reaches its consumer through a Transpose, as in MaxViT's attention blocks. The Transpose is now traced through and its permutation applied to the quantization axis, so the scale keeps indexing the stored weight: tracing through without remapping the axis would raise on non-square weights and silently quantize square ones along the wrong axis.
 
 0.47.0 (2026-09-23)
 ^^^^^^^^^^^^^^^^^^^
diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py
index d0010571035..2975f0c7717 100644
--- a/modelopt/onnx/quantization/qdq_utils.py
+++ b/modelopt/onnx/quantization/qdq_utils.py
@@ -579,11 +579,11 @@ def _get_scale_and_zp(
 
 def _get_successive_consumers(
     node: onnx.NodeProto, tensor_consumers: dict[str, list[onnx.NodeProto]]
-) -> tuple[onnx.NodeProto, onnx.NodeProto]:
+) -> tuple[onnx.NodeProto, onnx.NodeProto, list[int] | None]:
     """Get the DequantizeLinear node and its consumer node for a given QuantizeLinear node.
 
     This function validates and retrieves the next two nodes in the quantization chain:
-    QuantizeLinear -> DequantizeLinear -> Operation
+    QuantizeLinear -> DequantizeLinear -> [Cast | Transpose]* -> Operation
 
     Args:
         node: The QuantizeLinear node to find consumers for
@@ -593,6 +593,8 @@ def _get_successive_consumers(
         Tuple containing:
             - dq_node: The DequantizeLinear node that consumes the QuantizeLinear output
             - quantized_node: The operation node that consumes the DequantizeLinear output
+            - perm: Permutation mapping the operation's axes back onto the weight's own axes,
+              or None when no Transpose separates them
     """
     dq_node = tensor_consumers.get(node.output[0], [None])[0]
     if not dq_node or dq_node.op_type != "DequantizeLinear":
@@ -601,13 +603,24 @@ def _get_successive_consumers(
     quantized_node = tensor_consumers.get(dq_node.output[0], [None])[0]
     if not quantized_node:
         raise ValueError(f"No consumer found for {dq_node.name}")
-    if quantized_node.op_type == "Cast":
+
+    perm = None
+    while quantized_node.op_type in ("Cast", "Transpose"):
+        if quantized_node.op_type == "Transpose":
+            node_perm = next(
+                (list(attr.ints) for attr in quantized_node.attribute if attr.name == "perm"), None
+            )
+            assert node_perm is not None, f"Permutation not found for {quantized_node.name}"
+            # Composing back to front keeps perm relative to the weight across a Transpose chain.
+            perm = node_perm if perm is None else [perm[i] for i in node_perm]
         next_node = tensor_consumers.get(quantized_node.output[0], [None])[0]
         if not next_node:
-            raise ValueError(f"No consumer found after Cast for {quantized_node.name}")
+            raise ValueError(
+                f"No consumer found after {quantized_node.op_type} for {quantized_node.name}"
+            )
         quantized_node = next_node
 
-    return dq_node, quantized_node
+    return dq_node, quantized_node, perm
 
 
 def _convert_weight(
@@ -615,6 +628,7 @@ def _convert_weight(
     scale: onnx.TensorProto,
     zp: onnx.TensorProto,
     quantized_node: onnx.NodeProto,
+    perm: list[int] | None = None,
 ) -> np.ndarray:
     """Convert a weight tensor to INT8/FP8 format based on scale and zero point.
 
@@ -623,6 +637,7 @@ def _convert_weight(
         scale: The scale tensor for quantization
         zp: The zero point tensor for quantization
         quantized_node: The operation node that will use the converted weight
+        perm: Permutation of a Transpose between the weight and quantized_node, if any
 
     Returns:
         The converted weight tensor as a numpy array
@@ -662,6 +677,14 @@ def _convert_weight(
 
     axis = axis_map[op_type]
 
+    if perm is not None:
+        # axis_map is in the consumer's layout; a Transpose in between means that is not
+        # the weight's own layout. Map it back, since the scale indexes the stored weight.
+        assert len(perm) == len(weight_shape), (
+            f"Transpose perm {perm} does not match weight rank {len(weight_shape)}"
+        )
+        axis = perm[axis]
+
     if scale_array.shape and scale_array.shape[0] != weight_shape[axis]:
         raise ValueError(
             f"Scale shape {scale_array.shape} does not match weight shape {weight_shape} along axis {axis}"
@@ -754,10 +777,10 @@ def qdq_to_dq(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
             scale, zp = _get_scale_and_zp(node, initializers, tensor_producers)
 
             # Validate Q->DQ->Op pattern and get consumers
-            dq_node, quantized_node = _get_successive_consumers(node, tensor_consumers)
+            dq_node, quantized_node, perm = _get_successive_consumers(node, tensor_consumers)
 
             # Convert weight
-            scaled = _convert_weight(weight_array, scale, zp, quantized_node)
+            scaled = _convert_weight(weight_array, scale, zp, quantized_node, perm)
 
             # Create and update new weight tensor
             if zp.data_type == onnx_dtype_map["Float8"]:
diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py
index 95944e1f7ab..dadd2e5a547 100644
--- a/tests/unit/onnx/quantization/test_qdq_utils.py
+++ b/tests/unit/onnx/quantization/test_qdq_utils.py
@@ -1162,6 +1162,68 @@ def test_constant_node_scale_path_still_patched(self):
         assert (scale_arr > 0).all()
 
 
+def create_test_model_with_dq_transpose_matmul(out_features: int, in_features: int):
+    """Create weight[out, in] -> Q -> DQ -> Transpose -> MatMul, per-channel on the out axis.
+
+    This is how a torch ``Linear`` exports when the weight reaches MatMul transposed, as in
+    MaxViT's attention blocks.
+    """
+    rng = np.random.RandomState(7)
+    weight = rng.randn(out_features, in_features).astype(np.float32)
+    scale = (np.abs(weight).max(axis=1) / 127.0).astype(np.float32)
+    zero_point = np.zeros(out_features, dtype=np.int8)
+
+    nodes = [
+        helper.make_node(
+            "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q", axis=0
+        ),
+        helper.make_node(
+            "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", axis=0
+        ),
+        helper.make_node("Transpose", ["w_dq"], ["w_t"], name="w_t", perm=[1, 0]),
+        helper.make_node("MatMul", ["input", "w_t"], ["output"], name="matmul"),
+    ]
+    graph = helper.make_graph(
+        nodes=nodes,
+        name="dq_transpose_matmul",
+        inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [3, in_features])],
+        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [3, out_features])],
+        initializer=[
+            numpy_helper.from_array(weight, "weight"),
+            numpy_helper.from_array(scale, "w_scale"),
+            numpy_helper.from_array(zero_point, "w_zp"),
+        ],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)])
+    model.ir_version = 10
+    onnx.checker.check_model(model)
+    return model
+
+
+class TestQdqToDqTranspose:
+    """qdq_to_dq must trace a Transpose between DequantizeLinear and its consumer."""
+
+    # out == in is the case a naive fix gets wrong silently: the scale shape still matches
+    # the untransposed axis, so nothing raises and the weight is quantized along the wrong one.
+    @pytest.mark.parametrize(("out_features", "in_features"), [(96, 32), (32, 32)])
+    def test_dq_transpose_matmul_matches_qdq(self, out_features, in_features):
+        model = create_test_model_with_dq_transpose_matmul(out_features, in_features)
+        inputs = {"input": np.random.RandomState(11).randn(3, in_features).astype(np.float32)}
+
+        def run(m):
+            return ort.InferenceSession(
+                m.SerializeToString(), providers=["CPUExecutionProvider"]
+            ).run(None, inputs)[0]
+
+        expected = run(model)
+        converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
+
+        assert not [n for n in converted.graph.node if n.op_type == "QuantizeLinear"]
+        weight = next(t for t in converted.graph.initializer if t.name == "weight")
+        assert weight.data_type == TensorProto.INT8
+        assert np.array_equal(run(converted), expected)
+
+
 class TestQdqToDqValidation:
     """Regression tests for qdq_to_dq input validation."""
 

From f2349535dfbcc650b680de2b5519400d0776732b Mon Sep 17 00:00:00 2001
From: SID <99672439+SID-6921@users.noreply.github.com>
Date: Sat, 26 Sep 2026 19:14:27 -0400
Subject: [PATCH 2/8] Trim the changelog entry to the user-visible effect

AGENTS.md keeps root-cause analysis and implementation detail out of
CHANGELOG.rst; those belong in the PR description.

Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com>
---
 CHANGELOG.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index f0508daa690..2a3b12ea6b9 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -94,7 +94,7 @@ Changelog
 - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected.
 - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths.
 - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand.
-- Fix ``qdq_to_dq`` failing with ``Unsupported op_type for real weight quantization: Transpose`` when a quantized weight reaches its consumer through a Transpose, as in MaxViT's attention blocks. The Transpose is now traced through and its permutation applied to the quantization axis, so the scale keeps indexing the stored weight: tracing through without remapping the axis would raise on non-square weights and silently quantize square ones along the wrong axis.
+- Fix ONNX export of quantized models failing with ``Unsupported op_type for real weight quantization: Transpose`` when a weight reaches its consumer transposed, as in MaxViT's attention blocks.
 
 0.47.0 (2026-09-23)
 ^^^^^^^^^^^^^^^^^^^

From bdd4c0d19652195a9988ab3b8bc26aff246a4f00 Mon Sep 17 00:00:00 2001
From: SID <99672439+SID-6921@users.noreply.github.com>
Date: Sat, 26 Sep 2026 19:48:25 -0400
Subject: [PATCH 3/8] Default a Transpose's missing perm to reversing the axes

ONNX makes Transpose's `perm` optional and defaults it to reversing the input
axes, so asserting the attribute exists rejected a spec-valid graph. Derive the
reverse permutation from the weight's rank instead.

onnxruntime cannot execute DQ -> Transpose(no perm) -> MatMul at all (its QDQ
optimizer raises "vector too long"), so the regression test compares the
converted weight against the equivalent explicit permutation rather than
asserting runtime equivalence.

Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com>
---
 modelopt/onnx/quantization/qdq_utils.py       | 12 ++++++---
 .../unit/onnx/quantization/test_qdq_utils.py  | 25 ++++++++++++++++---
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py
index 2975f0c7717..d4ffa99997f 100644
--- a/modelopt/onnx/quantization/qdq_utils.py
+++ b/modelopt/onnx/quantization/qdq_utils.py
@@ -578,7 +578,7 @@ def _get_scale_and_zp(
 
 
 def _get_successive_consumers(
-    node: onnx.NodeProto, tensor_consumers: dict[str, list[onnx.NodeProto]]
+    node: onnx.NodeProto, tensor_consumers: dict[str, list[onnx.NodeProto]], weight_rank: int
 ) -> tuple[onnx.NodeProto, onnx.NodeProto, list[int] | None]:
     """Get the DequantizeLinear node and its consumer node for a given QuantizeLinear node.
 
@@ -588,6 +588,7 @@ def _get_successive_consumers(
     Args:
         node: The QuantizeLinear node to find consumers for
         tensor_consumers: Dictionary mapping tensor names to their consumer nodes
+        weight_rank: Rank of the weight being converted, used when a Transpose omits `perm`
 
     Returns:
         Tuple containing:
@@ -608,9 +609,10 @@ def _get_successive_consumers(
     while quantized_node.op_type in ("Cast", "Transpose"):
         if quantized_node.op_type == "Transpose":
             node_perm = next(
-                (list(attr.ints) for attr in quantized_node.attribute if attr.name == "perm"), None
+                (list(attr.ints) for attr in quantized_node.attribute if attr.name == "perm"),
+                # ONNX defaults a missing `perm` to reversing the axes.
+                list(reversed(range(weight_rank))),
             )
-            assert node_perm is not None, f"Permutation not found for {quantized_node.name}"
             # Composing back to front keeps perm relative to the weight across a Transpose chain.
             perm = node_perm if perm is None else [perm[i] for i in node_perm]
         next_node = tensor_consumers.get(quantized_node.output[0], [None])[0]
@@ -777,7 +779,9 @@ def qdq_to_dq(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
             scale, zp = _get_scale_and_zp(node, initializers, tensor_producers)
 
             # Validate Q->DQ->Op pattern and get consumers
-            dq_node, quantized_node, perm = _get_successive_consumers(node, tensor_consumers)
+            dq_node, quantized_node, perm = _get_successive_consumers(
+                node, tensor_consumers, weight_array.ndim
+            )
 
             # Convert weight
             scaled = _convert_weight(weight_array, scale, zp, quantized_node, perm)
diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py
index dadd2e5a547..19e78138697 100644
--- a/tests/unit/onnx/quantization/test_qdq_utils.py
+++ b/tests/unit/onnx/quantization/test_qdq_utils.py
@@ -1162,11 +1162,14 @@ def test_constant_node_scale_path_still_patched(self):
         assert (scale_arr > 0).all()
 
 
-def create_test_model_with_dq_transpose_matmul(out_features: int, in_features: int):
+def create_test_model_with_dq_transpose_matmul(
+    out_features: int, in_features: int, perm: list[int] | None = [1, 0]
+):
     """Create weight[out, in] -> Q -> DQ -> Transpose -> MatMul, per-channel on the out axis.
 
     This is how a torch ``Linear`` exports when the weight reaches MatMul transposed, as in
-    MaxViT's attention blocks.
+    MaxViT's attention blocks. ``perm=None`` omits the attribute, which ONNX defines as
+    reversing the axes.
     """
     rng = np.random.RandomState(7)
     weight = rng.randn(out_features, in_features).astype(np.float32)
@@ -1180,7 +1183,9 @@ def create_test_model_with_dq_transpose_matmul(out_features: int, in_features: i
         helper.make_node(
             "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", axis=0
         ),
-        helper.make_node("Transpose", ["w_dq"], ["w_t"], name="w_t", perm=[1, 0]),
+        helper.make_node(
+            "Transpose", ["w_dq"], ["w_t"], name="w_t", **({} if perm is None else {"perm": perm})
+        ),
         helper.make_node("MatMul", ["input", "w_t"], ["output"], name="matmul"),
     ]
     graph = helper.make_graph(
@@ -1223,6 +1228,20 @@ def run(m):
         assert weight.data_type == TensorProto.INT8
         assert np.array_equal(run(converted), expected)
 
+    def test_dq_transpose_without_perm_reverses_axes(self):
+        """ONNX lets Transpose omit `perm`, which reverses the axes."""
+
+        def converted_weight(perm):
+            model = create_test_model_with_dq_transpose_matmul(96, 32, perm=perm)
+            converted = qdq_to_dq(model)
+            return numpy_helper.to_array(
+                next(t for t in converted.graph.initializer if t.name == "weight")
+            )
+
+        # onnxruntime cannot execute DQ -> Transpose(no perm) -> MatMul, so compare the
+        # converted weight against the equivalent explicit permutation instead.
+        assert np.array_equal(converted_weight(None), converted_weight([1, 0]))
+
 
 class TestQdqToDqValidation:
     """Regression tests for qdq_to_dq input validation."""

From 46770f395faf04d2a4186561ce697c085cb1bfc0 Mon Sep 17 00:00:00 2001
From: SID <99672439+SID-6921@users.noreply.github.com>
Date: Sat, 26 Sep 2026 20:05:00 -0400
Subject: [PATCH 4/8] Quantize along the axis the source Q/DQ pair declares

_convert_weight derived the conversion axis from the consumer op, but the
DequantizeLinear left in the graph dequantizes along the axis it declares. When
the two disagree the weight is quantized along the wrong axis, and for a square
weight the scale length still matches, so validation passes and dq_only export
silently changes the model's output.

Use the declared axis when the Q/DQ pair has one, keeping the consumer mapping
as the fallback for pairs that do not. This also fixes the same mismatch without
a Transpose, which reproduces on main: a square weight with axis=0 feeding
MatMul converts along axis 1 today.

Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  2 +-
 modelopt/onnx/quantization/qdq_utils.py       | 31 ++++++----
 .../unit/onnx/quantization/test_qdq_utils.py  | 57 +++++++++++++++----
 3 files changed, 68 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2a3b12ea6b9..0f630ce1d82 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -94,7 +94,7 @@ Changelog
 - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected.
 - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths.
 - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand.
-- Fix ONNX export of quantized models failing with ``Unsupported op_type for real weight quantization: Transpose`` when a weight reaches its consumer transposed, as in MaxViT's attention blocks.
+- Fix ONNX export of quantized models failing with ``Unsupported op_type for real weight quantization: Transpose`` when a weight reaches its consumer transposed, as in MaxViT's attention blocks. The same export now also quantizes each weight along the axis its Q/DQ pair declares, correcting silently wrong ``dq_only`` output when that axis differed from the consumer's default.
 
 0.47.0 (2026-09-23)
 ^^^^^^^^^^^^^^^^^^^
diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py
index d4ffa99997f..49935db0e56 100644
--- a/modelopt/onnx/quantization/qdq_utils.py
+++ b/modelopt/onnx/quantization/qdq_utils.py
@@ -630,7 +630,9 @@ def _convert_weight(
     scale: onnx.TensorProto,
     zp: onnx.TensorProto,
     quantized_node: onnx.NodeProto,
+    *,
     perm: list[int] | None = None,
+    dq_axis: int | None = None,
 ) -> np.ndarray:
     """Convert a weight tensor to INT8/FP8 format based on scale and zero point.
 
@@ -640,6 +642,7 @@ def _convert_weight(
         zp: The zero point tensor for quantization
         quantized_node: The operation node that will use the converted weight
         perm: Permutation of a Transpose between the weight and quantized_node, if any
+        dq_axis: The axis declared by the source Q/DQ pair, if it declares one
 
     Returns:
         The converted weight tensor as a numpy array
@@ -677,15 +680,20 @@ def _convert_weight(
     if op_type not in axis_map:
         raise ValueError(f"Unsupported op_type for real weight quantization: {op_type}")
 
-    axis = axis_map[op_type]
-
-    if perm is not None:
-        # axis_map is in the consumer's layout; a Transpose in between means that is not
-        # the weight's own layout. Map it back, since the scale indexes the stored weight.
-        assert len(perm) == len(weight_shape), (
-            f"Transpose perm {perm} does not match weight rank {len(weight_shape)}"
-        )
-        axis = perm[axis]
+    if dq_axis is not None:
+        # The retained DequantizeLinear dequantizes along the axis it declares, so the weight
+        # has to be quantized along that same one. It already indexes the stored weight, so no
+        # Transpose remapping applies.
+        axis = dq_axis % len(weight_shape)
+    else:
+        axis = axis_map[op_type]
+        if perm is not None:
+            # axis_map is in the consumer's layout; a Transpose in between means that is not
+            # the weight's own layout. Map it back, since the scale indexes the stored weight.
+            assert len(perm) == len(weight_shape), (
+                f"Transpose perm {perm} does not match weight rank {len(weight_shape)}"
+            )
+            axis = perm[axis]
 
     if scale_array.shape and scale_array.shape[0] != weight_shape[axis]:
         raise ValueError(
@@ -784,7 +792,10 @@ def qdq_to_dq(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
             )
 
             # Convert weight
-            scaled = _convert_weight(weight_array, scale, zp, quantized_node, perm)
+            dq_axis = next((attr.i for attr in dq_node.attribute if attr.name == "axis"), None)
+            scaled = _convert_weight(
+                weight_array, scale, zp, quantized_node, perm=perm, dq_axis=dq_axis
+            )
 
             # Create and update new weight tensor
             if zp.data_type == onnx_dtype_map["Float8"]:
diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py
index 19e78138697..82cc0029238 100644
--- a/tests/unit/onnx/quantization/test_qdq_utils.py
+++ b/tests/unit/onnx/quantization/test_qdq_utils.py
@@ -1163,7 +1163,11 @@ def test_constant_node_scale_path_still_patched(self):
 
 
 def create_test_model_with_dq_transpose_matmul(
-    out_features: int, in_features: int, perm: list[int] | None = [1, 0]
+    out_features: int,
+    in_features: int,
+    perm: list[int] | None = [1, 0],
+    q_axis: int = 0,
+    transpose: bool = True,
 ):
     """Create weight[out, in] -> Q -> DQ -> Transpose -> MatMul, per-channel on the out axis.
 
@@ -1173,26 +1177,38 @@ def create_test_model_with_dq_transpose_matmul(
     """
     rng = np.random.RandomState(7)
     weight = rng.randn(out_features, in_features).astype(np.float32)
-    scale = (np.abs(weight).max(axis=1) / 127.0).astype(np.float32)
-    zero_point = np.zeros(out_features, dtype=np.int8)
+    reduced = tuple(axis for axis in range(2) if axis != q_axis)
+    scale = (np.abs(weight).max(axis=reduced) / 127.0).astype(np.float32)
+    zero_point = np.zeros(weight.shape[q_axis], dtype=np.int8)
 
     nodes = [
         helper.make_node(
-            "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q", axis=0
+            "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q", axis=q_axis
         ),
         helper.make_node(
-            "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", axis=0
+            "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", axis=q_axis
         ),
-        helper.make_node(
-            "Transpose", ["w_dq"], ["w_t"], name="w_t", **({} if perm is None else {"perm": perm})
-        ),
-        helper.make_node("MatMul", ["input", "w_t"], ["output"], name="matmul"),
     ]
+    matmul_rhs = "w_dq"
+    if transpose:
+        nodes.append(
+            helper.make_node(
+                "Transpose",
+                ["w_dq"],
+                ["w_t"],
+                name="w_t",
+                **({} if perm is None else {"perm": perm}),
+            )
+        )
+        matmul_rhs = "w_t"
+    nodes.append(helper.make_node("MatMul", ["input", matmul_rhs], ["output"], name="matmul"))
+
+    rows, cols = (in_features, out_features) if transpose else (out_features, in_features)
     graph = helper.make_graph(
         nodes=nodes,
         name="dq_transpose_matmul",
-        inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [3, in_features])],
-        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [3, out_features])],
+        inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [3, rows])],
+        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [3, cols])],
         initializer=[
             numpy_helper.from_array(weight, "weight"),
             numpy_helper.from_array(scale, "w_scale"),
@@ -1242,6 +1258,25 @@ def converted_weight(perm):
         # converted weight against the equivalent explicit permutation instead.
         assert np.array_equal(converted_weight(None), converted_weight([1, 0]))
 
+    # Square weights are the dangerous shape: the scale length matches either axis, so an
+    # axis mismatch passes validation and silently changes the model's output.
+    @pytest.mark.parametrize(("q_axis", "transpose"), [(0, False), (1, True)])
+    def test_conversion_follows_the_source_qdq_axis(self, q_axis, transpose):
+        """The retained DequantizeLinear dequantizes along the axis it declares."""
+        model = create_test_model_with_dq_transpose_matmul(
+            32, 32, q_axis=q_axis, transpose=transpose
+        )
+        inputs = {"input": np.random.RandomState(11).randn(3, 32).astype(np.float32)}
+
+        def run(m):
+            return ort.InferenceSession(
+                m.SerializeToString(), providers=["CPUExecutionProvider"]
+            ).run(None, inputs)[0]
+
+        expected = run(model)
+        converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
+        assert np.array_equal(run(converted), expected)
+
 
 class TestQdqToDqValidation:
     """Regression tests for qdq_to_dq input validation."""

From b9c43c82f6826c8b95eb472a4df66dfbd5e7011c Mon Sep 17 00:00:00 2001
From: SID <99672439+SID-6921@users.noreply.github.com>
Date: Sat, 26 Sep 2026 20:18:10 -0400
Subject: [PATCH 5/8] Honour ONNX's default axis when the Q/DQ pair omits it

A per-axis DequantizeLinear that omits `axis` dequantizes along axis 1, but the
fallback mapped the consumer's axis through the Transpose permutation instead.
For a square weight behind Transpose -> MatMul that picks axis 0, so the
conversion disagreed with the DQ left in the graph and silently changed the
model's output.

Per-axis conversion now takes the axis from the Q/DQ pair in every case,
declared or defaulted. That axis already indexes the stored weight whatever the
consumer's layout is, so the permutation is no longer needed to pick it and the
plumbing that carried it is gone; Transpose is still traced to reach the real
consumer. The consumer mapping remains for non-per-axis scales.

Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com>
---
 modelopt/onnx/quantization/qdq_utils.py       | 45 +++++--------------
 .../unit/onnx/quantization/test_qdq_utils.py  | 25 ++++++-----
 2 files changed, 23 insertions(+), 47 deletions(-)

diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py
index 49935db0e56..cd34e52c597 100644
--- a/modelopt/onnx/quantization/qdq_utils.py
+++ b/modelopt/onnx/quantization/qdq_utils.py
@@ -578,8 +578,8 @@ def _get_scale_and_zp(
 
 
 def _get_successive_consumers(
-    node: onnx.NodeProto, tensor_consumers: dict[str, list[onnx.NodeProto]], weight_rank: int
-) -> tuple[onnx.NodeProto, onnx.NodeProto, list[int] | None]:
+    node: onnx.NodeProto, tensor_consumers: dict[str, list[onnx.NodeProto]]
+) -> tuple[onnx.NodeProto, onnx.NodeProto]:
     """Get the DequantizeLinear node and its consumer node for a given QuantizeLinear node.
 
     This function validates and retrieves the next two nodes in the quantization chain:
@@ -588,14 +588,11 @@ def _get_successive_consumers(
     Args:
         node: The QuantizeLinear node to find consumers for
         tensor_consumers: Dictionary mapping tensor names to their consumer nodes
-        weight_rank: Rank of the weight being converted, used when a Transpose omits `perm`
 
     Returns:
         Tuple containing:
             - dq_node: The DequantizeLinear node that consumes the QuantizeLinear output
             - quantized_node: The operation node that consumes the DequantizeLinear output
-            - perm: Permutation mapping the operation's axes back onto the weight's own axes,
-              or None when no Transpose separates them
     """
     dq_node = tensor_consumers.get(node.output[0], [None])[0]
     if not dq_node or dq_node.op_type != "DequantizeLinear":
@@ -605,16 +602,7 @@ def _get_successive_consumers(
     if not quantized_node:
         raise ValueError(f"No consumer found for {dq_node.name}")
 
-    perm = None
     while quantized_node.op_type in ("Cast", "Transpose"):
-        if quantized_node.op_type == "Transpose":
-            node_perm = next(
-                (list(attr.ints) for attr in quantized_node.attribute if attr.name == "perm"),
-                # ONNX defaults a missing `perm` to reversing the axes.
-                list(reversed(range(weight_rank))),
-            )
-            # Composing back to front keeps perm relative to the weight across a Transpose chain.
-            perm = node_perm if perm is None else [perm[i] for i in node_perm]
         next_node = tensor_consumers.get(quantized_node.output[0], [None])[0]
         if not next_node:
             raise ValueError(
@@ -622,7 +610,7 @@ def _get_successive_consumers(
             )
         quantized_node = next_node
 
-    return dq_node, quantized_node, perm
+    return dq_node, quantized_node
 
 
 def _convert_weight(
@@ -631,7 +619,6 @@ def _convert_weight(
     zp: onnx.TensorProto,
     quantized_node: onnx.NodeProto,
     *,
-    perm: list[int] | None = None,
     dq_axis: int | None = None,
 ) -> np.ndarray:
     """Convert a weight tensor to INT8/FP8 format based on scale and zero point.
@@ -641,7 +628,6 @@ def _convert_weight(
         scale: The scale tensor for quantization
         zp: The zero point tensor for quantization
         quantized_node: The operation node that will use the converted weight
-        perm: Permutation of a Transpose between the weight and quantized_node, if any
         dq_axis: The axis declared by the source Q/DQ pair, if it declares one
 
     Returns:
@@ -680,20 +666,13 @@ def _convert_weight(
     if op_type not in axis_map:
         raise ValueError(f"Unsupported op_type for real weight quantization: {op_type}")
 
-    if dq_axis is not None:
-        # The retained DequantizeLinear dequantizes along the axis it declares, so the weight
-        # has to be quantized along that same one. It already indexes the stored weight, so no
-        # Transpose remapping applies.
-        axis = dq_axis % len(weight_shape)
+    if scale_array.ndim == 1:
+        # Per-axis: the DequantizeLinear left in the graph dequantizes along the axis it
+        # declares -- ONNX defaults that to 1 -- and that axis already indexes the stored
+        # weight, whatever the consumer's layout is. Quantize along the same one.
+        axis = (1 if dq_axis is None else dq_axis) % len(weight_shape)
     else:
         axis = axis_map[op_type]
-        if perm is not None:
-            # axis_map is in the consumer's layout; a Transpose in between means that is not
-            # the weight's own layout. Map it back, since the scale indexes the stored weight.
-            assert len(perm) == len(weight_shape), (
-                f"Transpose perm {perm} does not match weight rank {len(weight_shape)}"
-            )
-            axis = perm[axis]
 
     if scale_array.shape and scale_array.shape[0] != weight_shape[axis]:
         raise ValueError(
@@ -787,15 +766,11 @@ def qdq_to_dq(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
             scale, zp = _get_scale_and_zp(node, initializers, tensor_producers)
 
             # Validate Q->DQ->Op pattern and get consumers
-            dq_node, quantized_node, perm = _get_successive_consumers(
-                node, tensor_consumers, weight_array.ndim
-            )
+            dq_node, quantized_node = _get_successive_consumers(node, tensor_consumers)
 
             # Convert weight
             dq_axis = next((attr.i for attr in dq_node.attribute if attr.name == "axis"), None)
-            scaled = _convert_weight(
-                weight_array, scale, zp, quantized_node, perm=perm, dq_axis=dq_axis
-            )
+            scaled = _convert_weight(weight_array, scale, zp, quantized_node, dq_axis=dq_axis)
 
             # Create and update new weight tensor
             if zp.data_type == onnx_dtype_map["Float8"]:
diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py
index 82cc0029238..097a779a873 100644
--- a/tests/unit/onnx/quantization/test_qdq_utils.py
+++ b/tests/unit/onnx/quantization/test_qdq_utils.py
@@ -1244,19 +1244,20 @@ def run(m):
         assert weight.data_type == TensorProto.INT8
         assert np.array_equal(run(converted), expected)
 
-    def test_dq_transpose_without_perm_reverses_axes(self):
-        """ONNX lets Transpose omit `perm`, which reverses the axes."""
-
-        def converted_weight(perm):
-            model = create_test_model_with_dq_transpose_matmul(96, 32, perm=perm)
-            converted = qdq_to_dq(model)
-            return numpy_helper.to_array(
-                next(t for t in converted.graph.initializer if t.name == "weight")
-            )
+    def test_dq_transpose_without_perm_is_traced(self):
+        """ONNX lets Transpose omit `perm`; the MatMul behind it must still be found."""
+        model = create_test_model_with_dq_transpose_matmul(96, 32, perm=None)
+        initializer = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer}
+
+        converted = qdq_to_dq(model)
 
-        # onnxruntime cannot execute DQ -> Transpose(no perm) -> MatMul, so compare the
-        # converted weight against the equivalent explicit permutation instead.
-        assert np.array_equal(converted_weight(None), converted_weight([1, 0]))
+        # onnxruntime cannot execute DQ -> Transpose(no perm) -> MatMul, so check the
+        # converted weight directly. The Q/DQ pair declares axis 0.
+        expected = np.clip(
+            np.round(initializer["weight"] / initializer["w_scale"][:, None]), -128, 127
+        ).astype(np.int8)
+        actual = next(t for t in converted.graph.initializer if t.name == "weight")
+        assert np.array_equal(numpy_helper.to_array(actual), expected)
 
     # Square weights are the dangerous shape: the scale length matches either axis, so an
     # axis mismatch passes validation and silently changes the model's output.

From 3ce246b8eebddbc003946d3032d1e67ec05a0e22 Mon Sep 17 00:00:00 2001
From: SID <99672439+SID-6921@users.noreply.github.com>
Date: Sat, 26 Sep 2026 20:20:58 -0400
Subject: [PATCH 6/8] Cover an omitted Q/DQ axis in the conversion-axis test

The fixture always declared axis, so the ONNX-default path had no regression
case. Parametrize it, and make the fixture's optional arguments keyword-only.

Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com>
---
 .../unit/onnx/quantization/test_qdq_utils.py  | 23 ++++++++++++-------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py
index 097a779a873..6e5a7cc50ba 100644
--- a/tests/unit/onnx/quantization/test_qdq_utils.py
+++ b/tests/unit/onnx/quantization/test_qdq_utils.py
@@ -1165,9 +1165,11 @@ def test_constant_node_scale_path_still_patched(self):
 def create_test_model_with_dq_transpose_matmul(
     out_features: int,
     in_features: int,
-    perm: list[int] | None = [1, 0],
+    *,
+    perm: list[int] | None = (1, 0),
     q_axis: int = 0,
     transpose: bool = True,
+    declare_axis: bool = True,
 ):
     """Create weight[out, in] -> Q -> DQ -> Transpose -> MatMul, per-channel on the out axis.
 
@@ -1181,12 +1183,14 @@ def create_test_model_with_dq_transpose_matmul(
     scale = (np.abs(weight).max(axis=reduced) / 127.0).astype(np.float32)
     zero_point = np.zeros(weight.shape[q_axis], dtype=np.int8)
 
+    # ONNX defaults an omitted `axis` to 1, so a caller omitting it must pass q_axis=1.
+    axis_attr = {"axis": q_axis} if declare_axis else {}
     nodes = [
         helper.make_node(
-            "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q", axis=q_axis
+            "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q", **axis_attr
         ),
         helper.make_node(
-            "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", axis=q_axis
+            "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", **axis_attr
         ),
     ]
     matmul_rhs = "w_dq"
@@ -1197,7 +1201,7 @@ def create_test_model_with_dq_transpose_matmul(
                 ["w_dq"],
                 ["w_t"],
                 name="w_t",
-                **({} if perm is None else {"perm": perm}),
+                **({} if perm is None else {"perm": list(perm)}),
             )
         )
         matmul_rhs = "w_t"
@@ -1261,11 +1265,14 @@ def test_dq_transpose_without_perm_is_traced(self):
 
     # Square weights are the dangerous shape: the scale length matches either axis, so an
     # axis mismatch passes validation and silently changes the model's output.
-    @pytest.mark.parametrize(("q_axis", "transpose"), [(0, False), (1, True)])
-    def test_conversion_follows_the_source_qdq_axis(self, q_axis, transpose):
-        """The retained DequantizeLinear dequantizes along the axis it declares."""
+    @pytest.mark.parametrize(
+        ("q_axis", "transpose", "declare_axis"),
+        [(0, False, True), (1, True, True), (1, True, False)],
+    )
+    def test_conversion_follows_the_source_qdq_axis(self, q_axis, transpose, declare_axis):
+        """The retained DequantizeLinear dequantizes along its axis, declared or defaulted."""
         model = create_test_model_with_dq_transpose_matmul(
-            32, 32, q_axis=q_axis, transpose=transpose
+            32, 32, q_axis=q_axis, transpose=transpose, declare_axis=declare_axis
         )
         inputs = {"input": np.random.RandomState(11).randn(3, 32).astype(np.float32)}
 

From 4ba0613b4b66d19f0f67c5d1f78a135ab462d89a Mon Sep 17 00:00:00 2001
From: SID <99672439+SID-6921@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:51:39 -0400
Subject: [PATCH 7/8] Keep a one-element scale on the consumer axis, and reject
 disagreeing Q/DQ axes

Review caught a regression: gating the per-axis branch on `ndim == 1` also
matched a per-tensor scale spelled as a one-element 1-D tensor. A weight whose
consumer axis has length 1 -- a single-output-channel Conv, a Gemm with one
output -- converted on main and started raising here, because the branch picked
the DQ's default axis 1 and the shape check then rejected it. Gate on the
element count instead, so a single scale keeps using the consumer axis; the axis
is irrelevant for a scale that is broadcast anyway.

Also, since the conversion follows the DequantizeLinear and drops the
QuantizeLinear, a per-axis pair that disagrees on the axis would silently change
the output. Compare the two effective axes and raise instead. Nodes that omit
the attribute default to 1, so a pair that declares it on only one side is not
a disagreement.

Bounds-check the axis rather than letting the modulo wrap an out-of-range one,
and fold the two near-identical before/after comparisons into one
parametrization.

Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com>
---
 modelopt/onnx/quantization/qdq_utils.py       | 31 +++++--
 .../unit/onnx/quantization/test_qdq_utils.py  | 86 +++++++++++++++----
 2 files changed, 91 insertions(+), 26 deletions(-)

diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py
index cd34e52c597..5366e605cca 100644
--- a/modelopt/onnx/quantization/qdq_utils.py
+++ b/modelopt/onnx/quantization/qdq_utils.py
@@ -577,6 +577,11 @@ def _get_scale_and_zp(
     return scale, zp
 
 
+def _quantization_axis(node: onnx.NodeProto) -> int:
+    """Returns a Q/DQ node's quantization axis; ONNX defaults the attribute to 1."""
+    return next((attr.i for attr in node.attribute if attr.name == "axis"), 1)
+
+
 def _get_successive_consumers(
     node: onnx.NodeProto, tensor_consumers: dict[str, list[onnx.NodeProto]]
 ) -> tuple[onnx.NodeProto, onnx.NodeProto]:
@@ -666,12 +671,17 @@ def _convert_weight(
     if op_type not in axis_map:
         raise ValueError(f"Unsupported op_type for real weight quantization: {op_type}")
 
-    if scale_array.ndim == 1:
-        # Per-axis: the DequantizeLinear left in the graph dequantizes along the axis it
-        # declares -- ONNX defaults that to 1 -- and that axis already indexes the stored
-        # weight, whatever the consumer's layout is. Quantize along the same one.
-        axis = (1 if dq_axis is None else dq_axis) % len(weight_shape)
+    if scale_array.size > 1:
+        # The DequantizeLinear left in the graph dequantizes along the axis it declares,
+        # and that axis indexes the stored weight whatever the consumer's layout is.
+        axis = 1 if dq_axis is None else dq_axis
+        if not -len(weight_shape) <= axis < len(weight_shape):
+            raise ValueError(
+                f"Quantization axis {axis} is out of range for weight shape {weight_shape}"
+            )
+        axis %= len(weight_shape)
     else:
+        # A single scale is broadcast over the whole weight, so only the check below cares.
         axis = axis_map[op_type]
 
     if scale_array.shape and scale_array.shape[0] != weight_shape[axis]:
@@ -768,8 +778,15 @@ def qdq_to_dq(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
             # Validate Q->DQ->Op pattern and get consumers
             dq_node, quantized_node = _get_successive_consumers(node, tensor_consumers)
 
-            # Convert weight
-            dq_axis = next((attr.i for attr in dq_node.attribute if attr.name == "axis"), None)
+            # Convert weight. The conversion follows the DQ and drops the Q, so a per-axis
+            # pair that disagrees on the axis would silently change the output.
+            dq_axis, q_axis = _quantization_axis(dq_node), _quantization_axis(node)
+            if int(np.prod(scale.dims)) > 1 and (
+                q_axis % weight_array.ndim != dq_axis % weight_array.ndim
+            ):
+                raise ValueError(
+                    f"QuantizeLinear axis {q_axis} disagrees with DequantizeLinear axis {dq_axis}"
+                )
             scaled = _convert_weight(weight_array, scale, zp, quantized_node, dq_axis=dq_axis)
 
             # Create and update new weight tensor
diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py
index 6e5a7cc50ba..fb59785bc84 100644
--- a/tests/unit/onnx/quantization/test_qdq_utils.py
+++ b/tests/unit/onnx/quantization/test_qdq_utils.py
@@ -1166,7 +1166,7 @@ def create_test_model_with_dq_transpose_matmul(
     out_features: int,
     in_features: int,
     *,
-    perm: list[int] | None = (1, 0),
+    perm: tuple[int, ...] | None = (1, 0),
     q_axis: int = 0,
     transpose: bool = True,
     declare_axis: bool = True,
@@ -1226,14 +1226,33 @@ def create_test_model_with_dq_transpose_matmul(
 
 
 class TestQdqToDqTranspose:
-    """qdq_to_dq must trace a Transpose between DequantizeLinear and its consumer."""
+    """qdq_to_dq must trace a Transpose and follow the Q/DQ pair's own quantization axis."""
 
-    # out == in is the case a naive fix gets wrong silently: the scale shape still matches
-    # the untransposed axis, so nothing raises and the weight is quantized along the wrong one.
-    @pytest.mark.parametrize(("out_features", "in_features"), [(96, 32), (32, 32)])
-    def test_dq_transpose_matmul_matches_qdq(self, out_features, in_features):
-        model = create_test_model_with_dq_transpose_matmul(out_features, in_features)
-        inputs = {"input": np.random.RandomState(11).randn(3, in_features).astype(np.float32)}
+    # Square weights are the dangerous shape: the scale length matches either axis, so a
+    # wrong axis passes validation and silently changes the model's output.
+    @pytest.mark.parametrize(
+        ("out_features", "in_features", "q_axis", "transpose", "declare_axis"),
+        [
+            (96, 32, 0, True, True),
+            (32, 32, 0, True, True),
+            (32, 32, 0, False, True),
+            (32, 32, 1, True, True),
+            (32, 32, 1, True, False),
+        ],
+    )
+    def test_conversion_matches_qdq(
+        self, out_features, in_features, q_axis, transpose, declare_axis
+    ):
+        """The DequantizeLinear left in the graph dequantizes along its own axis."""
+        model = create_test_model_with_dq_transpose_matmul(
+            out_features,
+            in_features,
+            q_axis=q_axis,
+            transpose=transpose,
+            declare_axis=declare_axis,
+        )
+        rows = in_features if transpose else out_features
+        inputs = {"input": np.random.RandomState(11).randn(3, rows).astype(np.float32)}
 
         def run(m):
             return ort.InferenceSession(
@@ -1263,18 +1282,38 @@ def test_dq_transpose_without_perm_is_traced(self):
         actual = next(t for t in converted.graph.initializer if t.name == "weight")
         assert np.array_equal(numpy_helper.to_array(actual), expected)
 
-    # Square weights are the dangerous shape: the scale length matches either axis, so an
-    # axis mismatch passes validation and silently changes the model's output.
-    @pytest.mark.parametrize(
-        ("q_axis", "transpose", "declare_axis"),
-        [(0, False, True), (1, True, True), (1, True, False)],
-    )
-    def test_conversion_follows_the_source_qdq_axis(self, q_axis, transpose, declare_axis):
-        """The retained DequantizeLinear dequantizes along its axis, declared or defaulted."""
-        model = create_test_model_with_dq_transpose_matmul(
-            32, 32, q_axis=q_axis, transpose=transpose, declare_axis=declare_axis
+    def test_one_element_scale_is_not_per_axis(self):
+        """A one-element scale is broadcast over the whole weight, not indexed per axis."""
+        # A single-output-channel Conv: reading the (1,) scale as per-axis would pick the
+        # DQ's default axis 1 and reject a graph that converts fine.
+        weight = np.random.RandomState(3).randn(1, 3, 3, 3).astype(np.float32)
+        graph = helper.make_graph(
+            nodes=[
+                helper.make_node(
+                    "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q"
+                ),
+                helper.make_node(
+                    "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq"
+                ),
+                helper.make_node(
+                    "Conv", ["input", "w_dq"], ["output"], name="conv", kernel_shape=[3, 3]
+                ),
+            ],
+            name="one_element_scale",
+            inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 8, 8])],
+            outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 1, 6, 6])],
+            initializer=[
+                numpy_helper.from_array(weight, "weight"),
+                numpy_helper.from_array(
+                    np.array([np.abs(weight).max() / 127.0], np.float32), "w_scale"
+                ),
+                numpy_helper.from_array(np.array([0], np.int8), "w_zp"),
+            ],
         )
-        inputs = {"input": np.random.RandomState(11).randn(3, 32).astype(np.float32)}
+        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)])
+        model.ir_version = 10
+        onnx.checker.check_model(model)
+        inputs = {"input": np.random.RandomState(5).randn(1, 3, 8, 8).astype(np.float32)}
 
         def run(m):
             return ort.InferenceSession(
@@ -1285,6 +1324,15 @@ def run(m):
         converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
         assert np.array_equal(run(converted), expected)
 
+    def test_disagreeing_qdq_axes_raise(self):
+        """Following the DQ while the Q declares another axis would change the output."""
+        model = create_test_model_with_dq_transpose_matmul(32, 32, q_axis=0, transpose=False)
+        q_node = next(node for node in model.graph.node if node.name == "w_q")
+        next(attr for attr in q_node.attribute if attr.name == "axis").i = 1
+
+        with pytest.raises(RuntimeError, match="disagrees with DequantizeLinear axis"):
+            qdq_to_dq(model)
+
 
 class TestQdqToDqValidation:
     """Regression tests for qdq_to_dq input validation."""

From 28002ff973570c625efd3edc4e611c574873a60a Mon Sep 17 00:00:00 2001
From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com>
Date: Fri, 2 Oct 2026 11:33:59 -0400
Subject: [PATCH 8/8] Let a single scale broadcast instead of borrowing the
 consumer's axis

A one-element scale was routed to the consumer's axis_map entry, which only
converts when that dimension happens to be 1. A single-output projection is the
counterexample: weight (1, 32) with a (1,) scale picked MatMul axis 1 and failed
the length check against 32. It fails the same way without a Transpose, so the
gate was wrong rather than the tracing.

A single scale covers the whole weight. It broadcasts, and no axis applies to
it, so leave reshape_dims all ones and run the length check only for a per-axis
scale.

That leaves nothing reading axis_map's values -- the per-axis branch takes its
axis from the Q/DQ pair -- so it becomes the set of supported consumers it was
really being used as, and the Gemm transB lookup that fed it goes with it.

One consequence worth flagging: a scalar scale now converts as well, which is
the subject of #2613.

Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com>
---
 modelopt/onnx/quantization/qdq_utils.py       | 43 ++++++-------------
 .../unit/onnx/quantization/test_qdq_utils.py  | 31 +++++++++++--
 2 files changed, 42 insertions(+), 32 deletions(-)

diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py
index 5366e605cca..aa3f92830ea 100644
--- a/modelopt/onnx/quantization/qdq_utils.py
+++ b/modelopt/onnx/quantization/qdq_utils.py
@@ -577,6 +577,10 @@ def _get_scale_and_zp(
     return scale, zp
 
 
+# Consumers whose quantized weights qdq_to_dq knows how to convert.
+_REAL_QUANT_WEIGHT_CONSUMERS = {"Conv", "ConvTranspose", "Gemm", "MatMul"}
+
+
 def _quantization_axis(node: onnx.NodeProto) -> int:
     """Returns a Q/DQ node's quantization axis; ONNX defaults the attribute to 1."""
     return next((attr.i for attr in node.attribute if attr.name == "axis"), 1)
@@ -645,7 +649,6 @@ def _convert_weight(
         - INT8 weights are clipped to [-128, 127]
         - FP8 weights use float8e4m3fn format
     """
-    # Per-op quantization axis mapping (must match ORT config)
     weight_shape = weight_array.shape
     op_type = quantized_node.op_type
 
@@ -653,44 +656,26 @@ def _convert_weight(
     scale_array = onnx.numpy_helper.to_array(scale)
     zp_array = onnx.numpy_helper.to_array(zp)
 
-    # Dynamically determine transB for Gemm
-    trans_b = 0
-    if op_type == "Gemm":
-        for attr in quantized_node.attribute:
-            if attr.name == "transB":
-                trans_b = attr.i
-                break
-
-    axis_map = {
-        "Conv": 0,
-        "ConvTranspose": 1,
-        "Gemm": 0 if trans_b else 1,
-        "MatMul": 1,
-    }
-
-    if op_type not in axis_map:
+    if op_type not in _REAL_QUANT_WEIGHT_CONSUMERS:
         raise ValueError(f"Unsupported op_type for real weight quantization: {op_type}")
 
+    # A single scale covers the whole weight, so it broadcasts and no axis applies to it.
+    reshape_dims = [1] * len(weight_shape)
     if scale_array.size > 1:
-        # The DequantizeLinear left in the graph dequantizes along the axis it declares,
-        # and that axis indexes the stored weight whatever the consumer's layout is.
+        # Per-axis: the DequantizeLinear left in the graph dequantizes along the axis it
+        # declares, and that axis indexes the stored weight whatever the consumer's layout is.
         axis = 1 if dq_axis is None else dq_axis
         if not -len(weight_shape) <= axis < len(weight_shape):
             raise ValueError(
                 f"Quantization axis {axis} is out of range for weight shape {weight_shape}"
             )
         axis %= len(weight_shape)
-    else:
-        # A single scale is broadcast over the whole weight, so only the check below cares.
-        axis = axis_map[op_type]
-
-    if scale_array.shape and scale_array.shape[0] != weight_shape[axis]:
-        raise ValueError(
-            f"Scale shape {scale_array.shape} does not match weight shape {weight_shape} along axis {axis}"
-        )
+        if scale_array.shape[0] != weight_shape[axis]:
+            raise ValueError(
+                f"Scale shape {scale_array.shape} does not match weight shape {weight_shape} along axis {axis}"
+            )
+        reshape_dims[axis] = scale_array.shape[0]
 
-    reshape_dims = [1] * len(weight_shape)
-    reshape_dims[axis] = scale_array.shape[0]
     scale_array = scale_array.reshape(*reshape_dims)
     zp_array = zp_array.reshape(*reshape_dims)
 
diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py
index fb59785bc84..522b2e870b8 100644
--- a/tests/unit/onnx/quantization/test_qdq_utils.py
+++ b/tests/unit/onnx/quantization/test_qdq_utils.py
@@ -1170,6 +1170,7 @@ def create_test_model_with_dq_transpose_matmul(
     q_axis: int = 0,
     transpose: bool = True,
     declare_axis: bool = True,
+    single_scale: bool = False,
 ):
     """Create weight[out, in] -> Q -> DQ -> Transpose -> MatMul, per-channel on the out axis.
 
@@ -1179,9 +1180,14 @@ def create_test_model_with_dq_transpose_matmul(
     """
     rng = np.random.RandomState(7)
     weight = rng.randn(out_features, in_features).astype(np.float32)
-    reduced = tuple(axis for axis in range(2) if axis != q_axis)
-    scale = (np.abs(weight).max(axis=reduced) / 127.0).astype(np.float32)
-    zero_point = np.zeros(weight.shape[q_axis], dtype=np.int8)
+    if single_scale:
+        # One element: covers the whole weight and broadcasts, whatever the axis says.
+        scale = np.array([np.abs(weight).max() / 127.0], dtype=np.float32)
+        zero_point = np.zeros(1, dtype=np.int8)
+    else:
+        reduced = tuple(axis for axis in range(2) if axis != q_axis)
+        scale = (np.abs(weight).max(axis=reduced) / 127.0).astype(np.float32)
+        zero_point = np.zeros(weight.shape[q_axis], dtype=np.int8)
 
     # ONNX defaults an omitted `axis` to 1, so a caller omitting it must pass q_axis=1.
     axis_attr = {"axis": q_axis} if declare_axis else {}
@@ -1324,6 +1330,25 @@ def run(m):
         converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
         assert np.array_equal(run(converted), expected)
 
+    @pytest.mark.parametrize(("out_features", "in_features"), [(1, 32), (32, 1)])
+    def test_one_element_scale_broadcasts_through_transpose(self, out_features, in_features):
+        """A single scale covers the whole weight, so no axis of it has to be length 1."""
+        model = create_test_model_with_dq_transpose_matmul(
+            out_features, in_features, single_scale=True
+        )
+        inputs = {"input": np.random.RandomState(11).randn(3, in_features).astype(np.float32)}
+
+        def run(m):
+            return ort.InferenceSession(
+                m.SerializeToString(), providers=["CPUExecutionProvider"]
+            ).run(None, inputs)[0]
+
+        expected = run(model)
+        converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
+
+        assert not [n for n in converted.graph.node if n.op_type == "QuantizeLinear"]
+        assert np.array_equal(run(converted), expected)
+
     def test_disagreeing_qdq_axes_raise(self):
         """Following the DQ while the Q declares another axis would change the output."""
         model = create_test_model_with_dq_transpose_matmul(32, 32, q_axis=0, transpose=False)