From 1294df63098c3528369376f98788f76ab5be3d46 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Wed, 1 Jul 2026 20:47:29 +0000
Subject: [PATCH 01/10] Add TensorQuantizer rotate-back mode

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/config.py         |  9 ++-
 .../nn/modules/tensor_quantizer.py            | 28 +++++--
 .../quantization/test_tensor_quant_cpu.py     | 78 ++++++++++++++++++-
 3 files changed, 106 insertions(+), 9 deletions(-)

diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py
index f4a935239a8..8ad3299466e 100644
--- a/modelopt/torch/quantization/config.py
+++ b/modelopt/torch/quantization/config.py
@@ -287,6 +287,7 @@ class RotateConfig(ModeloptBaseConfig):
     """
 
     enable: bool = False
+    mode: Literal["rotate", "rotate_back"] = "rotate"
     rotate_fp32: bool = False
     block_size: int | None = None
 
@@ -344,7 +345,7 @@ def _validate_effective_bits(cls, v: float | None) -> float | None:
     def validate_config(cls, values):
         """Validate quantizer config."""
 
-        def _validate_recursive(value):
+        def _validate_recursive(value, field_name=None):
             """Recursively validate config structure."""
             if value is None:
                 return
@@ -353,14 +354,16 @@ def _validate_recursive(value):
                 for item in value:
                     _validate_recursive(item)
             elif isinstance(value, dict):
+                if field_name == "rotate":
+                    return
                 if len(value) == 1 and "enable" in value and value["enable"] is True:
                     raise ValueError(
                         "Invalid quantizer config: Cannot specify only {'enable': True}. "
                         "Additional parameters are required when enabling quantization."
                     )
                 # Recurse into nested dicts
-                for v in value.values():
-                    _validate_recursive(v)
+                for k, v in value.items():
+                    _validate_recursive(v, k)
 
         _validate_recursive(values)
         return values
diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
index c50804dd2a9..7a47e769aea 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -619,6 +619,22 @@ def rotate_block_size(self):
             return self._rotate.get("block_size", None)
         return None
 
+    @property
+    def rotate_back_is_enabled(self):
+        """Check if inverse rotation should be applied after quantization."""
+        if isinstance(self._rotate, RotateConfig):
+            return self._rotate.enable and self._rotate.mode == "rotate_back"
+        if isinstance(self._rotate, dict) and self.rotate_is_enabled:
+            return self._rotate.get("mode", "rotate") == "rotate_back"
+        return False
+
+    def _rotate_inputs(self, inputs):
+        return normalized_hadamard_transform(
+            inputs,
+            rotate_fp32=self.rotate_is_fp32,
+            block_size=self.rotate_block_size,
+        )
+
     def disable_calib(self):
         """Disable calibration."""
         self._if_calib = False
@@ -1092,11 +1108,7 @@ def forward(self, inputs):
 
         # Rotating the input
         if self.rotate_is_enabled:
-            inputs = normalized_hadamard_transform(
-                inputs,
-                rotate_fp32=self.rotate_is_fp32,
-                block_size=self.rotate_block_size,
-            )
+            inputs = self._rotate_inputs(inputs)
 
         if self._disabled:
             # if quantizer is disabled, we still need to track the input dtype for saving the model
@@ -1149,6 +1161,8 @@ def forward(self, inputs):
                 with same_device_as(inputs):
                     outputs = self._fake_quantize(inputs)
             elif not self._dequantize:
+                if self.rotate_back_is_enabled:
+                    raise ValueError("rotate_back mode is only supported with fake_quant=True.")
                 outputs = self._real_quantize(inputs)
             else:
                 raise ValueError(
@@ -1159,6 +1173,9 @@ def forward(self, inputs):
         if self.is_static_block_quant:
             outputs = self._reset_to_original_shape(outputs)
 
+        if self.rotate_back_is_enabled and isinstance(outputs, torch.Tensor):
+            outputs = self._rotate_inputs(outputs)
+
         return outputs
 
     def _short_amax(self, fmt=".2e"):
@@ -1209,6 +1226,7 @@ def extra_repr(self):
             else ""
         )
         s += " rotated" if self.rotate_is_enabled else ""
+        s += " (rotate_back)" if self.rotate_back_is_enabled else ""
         s += " (fp32)" if self.rotate_is_fp32 else ""
         if self.rotate_block_size is not None:
             s += f" (block={self.rotate_block_size})"
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 218acd76f5d..b242aa3d0d2 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -22,8 +22,13 @@
 from _test_utils.torch.quantization.tensor_quant_common import FakeTensorQuantTester
 
 import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
-from modelopt.torch.quantization.nn import TensorQuantizer
+from modelopt.torch.quantization.nn import (
+    TensorQuantizer,
+    register_quant_backend,
+    unregister_quant_backend,
+)
 
 
 class TestFakeTensorQuantCPU(FakeTensorQuantTester):
@@ -56,6 +61,18 @@ def test_from_to_dict(self, verbose):
         quant_attr_cfg_2 = QuantizerAttributeConfig(**quant_attr_cfg_1.dict())
         assert quant_attr_cfg_1 == quant_attr_cfg_2
 
+    def test_rotate_mode_serialization(self):
+        quant_attr_cfg = QuantizerAttributeConfig(
+            rotate={"enable": True, "mode": "rotate_back", "rotate_fp32": True, "block_size": 8}
+        )
+
+        assert quant_attr_cfg.model_dump(exclude_unset=True)["rotate"] == {
+            "enable": True,
+            "mode": "rotate_back",
+            "rotate_fp32": True,
+            "block_size": 8,
+        }
+
     def test_num_bits(self):
         """Test num_bits for both integer and tuple cases."""
 
@@ -92,6 +109,65 @@ def test_num_bits(self):
             QuantizerAttributeConfig(enable=True, num_bits=(-1, 2))
 
 
+def _run_rotated_backend(monkeypatch, rotate):
+    calls = []
+
+    def rotate_fn(inputs, rotate_fp32=False, block_size=None):
+        calls.append((rotate_fp32, block_size))
+        return inputs + 10
+
+    def backend(inputs, _tq):
+        return inputs * 2
+
+    monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn)
+    register_quant_backend("test_rotate_mode_backend", backend)
+    try:
+        quantizer = TensorQuantizer(
+            QuantizerAttributeConfig(rotate=rotate, backend="test_rotate_mode_backend")
+        )
+        inputs = torch.tensor([[1.0, 2.0]])
+        return quantizer(inputs), inputs, calls, quantizer
+    finally:
+        unregister_quant_backend("test_rotate_mode_backend")
+
+
+def test_tensor_quantizer_rotate_mode_preserves_default_path(monkeypatch):
+    outputs, inputs, calls, quantizer = _run_rotated_backend(monkeypatch, rotate={"enable": True})
+
+    assert not quantizer.rotate_back_is_enabled
+    assert torch.equal(outputs, (inputs + 10) * 2)
+    assert calls == [(False, None)]
+
+
+def test_tensor_quantizer_rotate_mode_can_rotate_back(monkeypatch):
+    outputs, inputs, calls, quantizer = _run_rotated_backend(
+        monkeypatch,
+        rotate={"enable": True, "mode": "rotate_back", "rotate_fp32": True, "block_size": 8},
+    )
+
+    assert quantizer.rotate_back_is_enabled
+    assert torch.equal(outputs, ((inputs + 10) * 2) + 10)
+    assert calls == [(True, 8), (True, 8)]
+
+
+def test_tensor_quantizer_rotate_back_rejects_real_quant(monkeypatch):
+    monkeypatch.setattr(
+        tensor_quantizer_module,
+        "normalized_hadamard_transform",
+        lambda inputs, rotate_fp32=False, block_size=None: inputs,
+    )
+    quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits=8,
+            fake_quant=False,
+            rotate={"enable": True, "mode": "rotate_back"},
+        )
+    )
+
+    with pytest.raises(ValueError, match="rotate_back mode is only supported with fake_quant=True"):
+        quantizer(torch.tensor([[1.0, 2.0]]))
+
+
 WINT4INT8_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},

From 28deef13592e09116cf7214b3e0266a3a32f822f Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Wed, 1 Jul 2026 23:33:22 +0000
Subject: [PATCH 02/10] Add Hadamard roundtrip GPU test

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 tests/gpu/torch/quantization/test_hadamard.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/gpu/torch/quantization/test_hadamard.py b/tests/gpu/torch/quantization/test_hadamard.py
index 1173eba2097..7034feda75c 100644
--- a/tests/gpu/torch/quantization/test_hadamard.py
+++ b/tests/gpu/torch/quantization/test_hadamard.py
@@ -48,6 +48,9 @@ def test_hadamard_transform(dim):
     xxt_h = x_h @ x_h.T
     # The numerical error can be large, especially for 16-bit floats.
     assert torch.allclose(xxt_h, xxt, atol=0.05)
+    x_roundtrip = normalized_hadamard_transform(x_h)
+    assert torch.allclose(x_roundtrip, x, rtol=1e-5, atol=1e-6)
+
     x_h_fp32 = normalized_hadamard_transform(x, rotate_fp32=True)
     xxt_h_fp32 = x_h_fp32 @ x_h_fp32.T
     assert torch.allclose(xxt_h_fp32, xxt, atol=0.05)
@@ -66,6 +69,8 @@ def test_hadamard_transform_block(dim, block_size):
     # Use rtol instead of atol: float32 accumulated error scales with value magnitude,
     # which grows with dim. 1e-3 relative tolerance is appropriate for float32 block RHT.
     assert torch.allclose(xxt_h, xxt, rtol=1e-3, atol=1e-6)
+    x_roundtrip = normalized_hadamard_transform(x_h, block_size=block_size)
+    assert torch.allclose(x_roundtrip, x, rtol=1e-5, atol=1e-6)
 
 
 @pytest.mark.parametrize(

From b2f5fc837ff341e7dd9e11a797986fcbf5434051 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Thu, 2 Jul 2026 00:23:35 +0000
Subject: [PATCH 03/10] Address rotate-back review feedback

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 CHANGELOG.rst                                              | 1 +
 modelopt/torch/quantization/nn/modules/tensor_quantizer.py | 5 +++--
 tests/unit/torch/quantization/test_tensor_quant_cpu.py     | 5 ++++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 32d5a3388e2..d1791fbc800 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,6 +19,7 @@ Changelog
 
 **New Features**
 
+- Add ``rotate.mode`` to torch quantizer configs. The default ``"rotate"`` keeps the existing rotate-before-quantize behavior; ``"rotate_back"`` enables fake-quant rotate → quantize → rotate-back for TensorQuantizer.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.
diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
index 7a47e769aea..755dce7a45c 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -1106,6 +1106,9 @@ def forward(self, inputs):
         if self.pre_quant_scale is not None:
             inputs = inputs * self.pre_quant_scale
 
+        if self.rotate_back_is_enabled and self._if_quant and not self.fake_quant:
+            raise ValueError("rotate_back mode is only supported with fake_quant=True.")
+
         # Rotating the input
         if self.rotate_is_enabled:
             inputs = self._rotate_inputs(inputs)
@@ -1161,8 +1164,6 @@ def forward(self, inputs):
                 with same_device_as(inputs):
                     outputs = self._fake_quantize(inputs)
             elif not self._dequantize:
-                if self.rotate_back_is_enabled:
-                    raise ValueError("rotate_back mode is only supported with fake_quant=True.")
                 outputs = self._real_quantize(inputs)
             else:
                 raise ValueError(
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index b242aa3d0d2..139ca232b1b 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -151,10 +151,13 @@ def test_tensor_quantizer_rotate_mode_can_rotate_back(monkeypatch):
 
 
 def test_tensor_quantizer_rotate_back_rejects_real_quant(monkeypatch):
+    def fail_if_rotated(inputs, rotate_fp32=False, block_size=None):
+        raise AssertionError("rotate_back with fake_quant=False should fail before rotation")
+
     monkeypatch.setattr(
         tensor_quantizer_module,
         "normalized_hadamard_transform",
-        lambda inputs, rotate_fp32=False, block_size=None: inputs,
+        fail_if_rotated,
     )
     quantizer = TensorQuantizer(
         QuantizerAttributeConfig(

From de15abf861255c1311f5e35107add2ec4a808a05 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Fri, 3 Jul 2026 17:41:35 +0000
Subject: [PATCH 04/10] Refine TensorQuantizer rotate-back handling

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 .../torch/export/plugins/vllm_fakequant_hf.py |  22 +--
 modelopt/torch/quantization/model_quant.py    |   6 +-
 modelopt/torch/quantization/nn/functional.py  |  60 ++-----
 .../quantization/nn/modules/quant_module.py   |  42 +++--
 .../nn/modules/tensor_quantizer.py            |  18 ++
 .../torch/quantization/plugins/huggingface.py |  32 +---
 modelopt/torch/quantization/plugins/vllm.py   |  19 +-
 .../plugins/test_fused_experts.py             |  45 +++++
 .../quantization/test_tensor_quant_cpu.py     | 169 +++++++++++++++++-
 9 files changed, 303 insertions(+), 110 deletions(-)

diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py
index 8883964daac..acb1968e070 100644
--- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py
+++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py
@@ -27,7 +27,6 @@
 import torch.nn as nn
 
 import modelopt.torch.opt as mto
-from modelopt.torch.quantization.config import RotateConfig
 from modelopt.torch.quantization.conversion import quantizer_state
 from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection
 from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer
@@ -137,15 +136,6 @@ def _check_all_weight_quantizers_disabled(model: nn.Module) -> None:
                     )
 
 
-def disable_rotate(quantizer: TensorQuantizer):
-    """Return a disabled copy of the quantizer's ``_rotate`` field, preserving its type."""
-    if isinstance(quantizer._rotate, RotateConfig):
-        return RotateConfig(enable=False)
-    if isinstance(quantizer._rotate, dict):  # backward compat: old checkpoints stored a dict
-        return dict(quantizer._rotate, enable=False)
-    return False
-
-
 def _fakequant_fused_experts_weights(
     module: nn.Module,
     module_name: str,
@@ -638,16 +628,12 @@ def export_hf_vllm_fq_checkpoint(
                     if isinstance(quantizer, SequentialQuantizer):
                         quantizer.disable()
                         for sub in quantizer:
-                            orig_rotate = sub._rotate
-                            if sub.rotate_is_enabled:
-                                sub._rotate = disable_rotate(sub)
-                            wqs_to_restore.append((sub, orig_rotate))
+                            wqs_to_restore.append((sub, sub._rotate))
+                            sub.disable_rotate()
                     elif isinstance(quantizer, TensorQuantizer):
                         quantizer.disable()
-                        orig_rotate = quantizer._rotate
-                        if quantizer.rotate_is_enabled:
-                            quantizer._rotate = disable_rotate(quantizer)
-                        wqs_to_restore.append((quantizer, orig_rotate))
+                        wqs_to_restore.append((quantizer, quantizer._rotate))
+                        quantizer.disable_rotate()
 
         quantizer_state_dict = get_quantizer_state_dict(model)
         for key in list(quantizer_state_dict):
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 7dbdd36d04e..1adda0511a6 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -619,7 +619,11 @@ def print_quant_summary(model: nn.Module, output_dir: str | None = None):
 
 
 def fold_weight(model: nn.Module, keep_attrs: bool = False):
-    """Fold weight quantizer for fast evaluation."""
+    """Fold weight quantizer for fast evaluation.
+
+    Any weight-quantizer rotation is folded into the weights and disabled so subsequent
+    forwards do not re-rotate the already-folded weights.
+    """
     for name, module in model.named_modules():
         if isinstance(module, QuantModule):
             module.fold_weight(keep_attrs)
diff --git a/modelopt/torch/quantization/nn/functional.py b/modelopt/torch/quantization/nn/functional.py
index 94cbcf74fda..533e7ebf68d 100644
--- a/modelopt/torch/quantization/nn/functional.py
+++ b/modelopt/torch/quantization/nn/functional.py
@@ -73,26 +73,6 @@ def backward(ctx, grad_output):
 clip = ClipFunction.apply
 
 
-class FastHadamardTransform(Function):
-    """The fast Hadamard transform.
-
-    This only works for inputs.shape[-1] == power of 2.
-    """
-
-    @staticmethod
-    def forward(ctx, inputs):
-        """Hadamard forward."""
-        assert utils.is_pow2(inputs.shape[-1]), (
-            "Fast hadamard only works for inputs.shape[-1] == power of 2."
-        )
-        return fast_hadamard_transform.hadamard_transform(inputs)  # type: ignore[name-defined]
-
-    @staticmethod
-    def backward(ctx, grad_outputs):
-        """Hadamard backward."""
-        return fast_hadamard_transform.hadamard_transform(grad_outputs)  # type: ignore[name-defined]
-
-
 def _largest_pow2_divisor(n: int) -> int:
     """Return the largest power of 2 that divides n."""
     return n & (-n)
@@ -130,35 +110,29 @@ def normalized_hadamard_transform(inputs, rotate_fp32=False, block_size=None):
     if rotate_fp32:
         inputs = inputs.to(torch.float32)
 
-    if block_size is None and utils.is_pow2(dim):
-        # Full-dimension FHT (original behavior)
-        outputs = FastHadamardTransform.apply(inputs) / torch.sqrt(
-            torch.tensor(dim, dtype=torch.float32)
-        )
-    else:
-        # Block-granular RHT
-        if block_size is None:
+    # Full-dimension FHT is just block-granular RHT with block_size == dim.
+    if block_size is None:
+        if utils.is_pow2(dim):
+            block_size = dim
+        else:
             block_size = _largest_pow2_divisor(dim)
             if block_size < 2:
                 raise RuntimeError(
                     f"Block RHT: dimension {dim} has no power-of-2 divisor >= 2. "
                     "Set rotate.block_size explicitly (e.g. 128) or use a dimension divisible by a power of 2."
                 )
-        if not utils.is_pow2(block_size):
-            raise ValueError(f"Block RHT: block_size must be power of 2, got {block_size}.")
-        if dim % block_size != 0:
-            raise RuntimeError(
-                f"Block RHT: inputs.shape[-1]={dim} is not divisible by block_size={block_size}. "
-                f"Use a block_size that divides {dim} (e.g. {_largest_pow2_divisor(dim)})."
-            )
-        n_blocks = dim // block_size
-        # Reshape to (..., n_blocks, block_size)
-        flat = inputs.reshape(-1, dim)
-        blocks = flat.reshape(-1, n_blocks, block_size)
-        # Apply FHT per block (last dim)
-        rotated = FastHadamardTransform.apply(blocks) / torch.sqrt(
-            torch.tensor(block_size, dtype=torch.float32)
+    if not utils.is_pow2(block_size):
+        raise ValueError(f"Block RHT: block_size must be power of 2, got {block_size}.")
+    if dim % block_size != 0:
+        raise RuntimeError(
+            f"Block RHT: inputs.shape[-1]={dim} is not divisible by block_size={block_size}. "
+            f"Use a block_size that divides {dim} (e.g. {_largest_pow2_divisor(dim)})."
         )
-        outputs = rotated.reshape(inputs.shape)
+
+    # hadamard_transform is autograd-aware and fuses the normalization scale into the kernel.
+    rotated = fast_hadamard_transform.hadamard_transform(  # type: ignore[name-defined]
+        inputs.reshape(-1, dim // block_size, block_size).contiguous(), scale=block_size**-0.5
+    )
+    outputs = rotated.reshape(inputs.shape)
 
     return outputs.to(dtype) if rotate_fp32 else outputs
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 419c6f4924f..137fa465d15 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -17,6 +17,7 @@
 
 import contextlib
 import warnings
+from collections.abc import Iterable
 from typing import Any
 
 import torch
@@ -127,8 +128,36 @@ def iter_weights_for_calibration(self):
             weight_quantizer = getattr(self, quantizer_attr_names(weight_name).weight_quantizer)
             yield getattr(self, weight_name), weight_quantizer
 
+    @staticmethod
+    @torch.no_grad()
+    def _fold_weight_quantizer(
+        quantizer: TensorQuantizer,
+        weights: Iterable[torch.Tensor],
+        keep_attrs: bool = False,
+    ):
+        """Fold ``quantizer`` into each weight view in place, then disable and clean it once.
+
+        ``weights`` is an iterable so a single quantizer shared across views (e.g. one
+        per-tensor quantizer over all experts of a fused MoE weight) can be folded view by
+        view while disabling and dropping its calibration attrs exactly once.
+        """
+        for weight in weights:
+            weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
+        quantizer.disable()
+        quantizer.disable_rotate()
+        if not keep_attrs:
+            for attr_name in ("_pre_quant_scale", "_amax"):
+                if hasattr(quantizer, attr_name):
+                    delattr(quantizer, attr_name)
+
     def fold_weight(self, keep_attrs: bool = False):
-        """Fold the weight for faster eval."""
+        """Bake each fake-quant weight quantizer into its weight for faster eval.
+
+        Every fake-quant weight quantizer is folded regardless of its enabled state. Both its
+        quantization and its input rotation are baked into the stored weight and then disabled,
+        so subsequent forwards neither re-quantize nor re-rotate. Calibration buffers
+        (``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``.
+        """
         # Handle all attributes that end with _weight_quantizer
         for name in dir(self):
             attr = getattr(self, name)
@@ -144,16 +173,7 @@ def fold_weight(self, keep_attrs: bool = False):
                     f"{name} doesn't have a corresponding {weight_name} in {self.__class__.__name__}"
                 )
                 weight = getattr(self, weight_name)
-                weight.data.copy_(attr(weight.float()).to(weight.dtype))
-                attr.disable()
-                if not keep_attrs:
-                    _attrs = [
-                        "_pre_quant_scale",
-                        "_amax",
-                    ]
-                    for attr_name in _attrs:
-                        if hasattr(attr, attr_name):
-                            delattr(attr, attr_name)
+                self._fold_weight_quantizer(attr, (weight,), keep_attrs)
 
 
 QuantModuleRegistry = _DMRegistryCls("Quant", QuantModule)
diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
index 755dce7a45c..20cdcfc3c77 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -628,6 +628,19 @@ def rotate_back_is_enabled(self):
             return self._rotate.get("mode", "rotate") == "rotate_back"
         return False
 
+    def disable_rotate(self):
+        """Disable rotation while preserving the ``_rotate`` field's type.
+
+        Idempotent. Used after folding a weight quantizer so the baked-in rotation is not
+        re-applied on subsequent forwards.
+        """
+        if isinstance(self._rotate, RotateConfig):
+            self._rotate = self._rotate.model_copy(update={"enable": False})
+        elif isinstance(self._rotate, dict):  # backward compat: old checkpoints stored a dict
+            self._rotate = dict(self._rotate, enable=False)
+        else:
+            self._rotate = False
+
     def _rotate_inputs(self, inputs):
         return normalized_hadamard_transform(
             inputs,
@@ -1118,6 +1131,10 @@ def forward(self, inputs):
             # TODO: This is a temporary solution and needs to be removed once megatron supports
             # non-homogeneous layers
             self._input_dtype = inputs.dtype if hasattr(inputs, "dtype") else None
+            # Even when quantization is disabled, honor rotate_back so a rotate/rotate_back
+            # pair stays a no-op roundtrip instead of leaving the tensor rotated.
+            if self.rotate_back_is_enabled:
+                inputs = self._rotate_inputs(inputs)
             return inputs
 
         if (
@@ -1530,6 +1547,7 @@ class SequentialQuantizer(nn.Sequential):
     _delegated_methods = [
         "reset_amax",
         "disable",
+        "disable_rotate",
         "enable",
         "load_calib_amax",
         "load_calib_bias",
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 6779c3f9ade..c86c90eaa59 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1091,31 +1091,17 @@ def iter_weights_for_calibration(self):
                 yield weight[idx], q
 
     def fold_weight(self, keep_attrs: bool = False):
-        """Fold per-expert weight quantizers into the fused 3-D weights.
+        """Bake each per-expert weight quantizer into its slice of the fused 3-D weight.
 
-        The base ``fold_weight`` only handles singular ``*_weight_quantizer``
-        attributes. Fused experts use ``nn.ModuleList`` of per-expert quantizers
-        (``<first_proj>_weight_quantizers``, ``down_proj_weight_quantizers``),
-        which would otherwise be skipped, leaving ``_amax`` on every quantizer.
+        The base ``fold_weight`` only handles singular ``*_weight_quantizer`` attributes and
+        would skip the ``nn.ModuleList`` of per-expert quantizers used here. The per-expert
+        ``(weight_slice, quantizer)`` pairs are the same ones :meth:`iter_weights_for_calibration`
+        yields, so we reuse it; each fake-quant quantizer's quantization and rotation are folded
+        in and disabled, and calibration buffers are dropped unless ``keep_attrs``.
         """
-        for weight_name, quantizers_name in (
-            (self._first_proj_attr, self._first_proj_weight_quantizers_attr),
-            ("down_proj", "down_proj_weight_quantizers"),
-        ):
-            weight = getattr(self, weight_name, None)
-            quantizers = getattr(self, quantizers_name, None)
-            if weight is None or quantizers is None:
-                continue
-            for idx, q in enumerate(quantizers):
-                if not (isinstance(q, TensorQuantizer) and q.fake_quant):
-                    continue
-                slice_ = weight.data[idx]
-                slice_.copy_(q(slice_.float()).to(weight.dtype))
-                q.disable()
-                if not keep_attrs:
-                    for attr_name in ("_pre_quant_scale", "_amax"):
-                        if hasattr(q, attr_name):
-                            delattr(q, attr_name)
+        for weight_slice, q in self.iter_weights_for_calibration():
+            if isinstance(q, TensorQuantizer) and q.fake_quant:
+                self._fold_weight_quantizer(q, (weight_slice,), keep_attrs)
 
 
 class _QuantNonGatedFusedExperts(_QuantFusedExperts):
diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py
index 95ca3240b73..aa6141dd48b 100644
--- a/modelopt/torch/quantization/plugins/vllm.py
+++ b/modelopt/torch/quantization/plugins/vllm.py
@@ -462,20 +462,13 @@ def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor):
     @torch.no_grad()
     def fold_weight(self, keep_attrs: bool = False):
         # the MoE weights can be super large, it consumes too much memory, so we need to fold the weight one by one
-        for i in range(self.w13_weight.shape[0]):
-            self.w13_weight[i].copy_(
-                self.w13_weight_quantizer(self.w13_weight[i].float().contiguous()).to(
-                    self.w13_weight.dtype
-                )
-            )
-        self.w13_weight_quantizer.disable()
-        for i in range(self.w2_weight.shape[0]):
-            self.w2_weight[i].copy_(
-                self.w2_weight_quantizer(self.w2_weight[i].float().contiguous()).to(
-                    self.w2_weight.dtype
-                )
+        for weight, quantizer in (
+            (self.w13_weight, self.w13_weight_quantizer),
+            (self.w2_weight, self.w2_weight_quantizer),
+        ):
+            self._fold_weight_quantizer(
+                quantizer, (weight[i] for i in range(weight.shape[0])), keep_attrs
             )
-        self.w2_weight_quantizer.disable()
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 1829777e87f..9fa836bb620 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -25,8 +25,10 @@
 from _test_utils.torch.quantization.tied_modules import tie_fused_experts_3d_params
 
 import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
 from modelopt.torch.export.moe_utils import _export_fused_experts
 from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name
 from modelopt.torch.quantization.model_calib import local_hessian_calibrate
 from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer
@@ -336,6 +338,49 @@ def test_expert_index_recovery(self):
             assert recovered_idx == idx, f"Expected {idx}, got {recovered_idx}"
         self._cleanup_registry(expert_type)
 
+    def _make_rotated_fused_experts(self, monkeypatch):
+        monkeypatch.setattr(
+            tensor_quantizer_module,
+            "normalized_hadamard_transform",
+            lambda inputs, rotate_fp32=False, block_size=None: inputs,
+        )
+        model = _TinyMoEModel()
+        expert_type = type(model.moe.experts)
+        self._cleanup_registry(expert_type)
+        register_fused_experts_on_the_fly(model)
+        converted = QuantModuleRegistry.convert(model.moe.experts)
+        expert_quantizers = list(converted.gate_up_proj_weight_quantizers) + list(
+            converted.down_proj_weight_quantizers
+        )
+        for q in expert_quantizers:
+            q.set_from_attribute_config(
+                QuantizerAttributeConfig(num_bits=8, rotate={"enable": True})
+            )
+            q.amax = torch.tensor(6.0)
+        return converted, expert_quantizers, expert_type
+
+    def test_fold_weight_disables_per_expert_quantizers_and_rotation(self, monkeypatch):
+        converted, expert_quantizers, expert_type = self._make_rotated_fused_experts(monkeypatch)
+        try:
+            converted.fold_weight()
+            for q in expert_quantizers:
+                assert not q.is_enabled
+                assert not q.rotate_is_enabled
+                assert not hasattr(q, "_amax")
+        finally:
+            self._cleanup_registry(expert_type)
+
+    def test_fold_weight_keep_attrs_keeps_amax_disables_rotation(self, monkeypatch):
+        converted, expert_quantizers, expert_type = self._make_rotated_fused_experts(monkeypatch)
+        try:
+            converted.fold_weight(keep_attrs=True)
+            for q in expert_quantizers:
+                assert not q.is_enabled
+                assert not q.rotate_is_enabled
+                assert hasattr(q, "_amax")
+        finally:
+            self._cleanup_registry(expert_type)
+
 
 # ---------------------------------------------------------------------------
 # Tests for export
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 139ca232b1b..a15fff12e58 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -23,8 +23,10 @@
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
-from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization import QuantModuleRegistry
+from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig
 from modelopt.torch.quantization.nn import (
+    SequentialQuantizer,
     TensorQuantizer,
     register_quant_backend,
     unregister_quant_backend,
@@ -171,6 +173,171 @@ def fail_if_rotated(inputs, rotate_fp32=False, block_size=None):
         quantizer(torch.tensor([[1.0, 2.0]]))
 
 
+def test_tensor_quantizer_rotate_back_roundtrips_when_disabled(monkeypatch):
+    calls = []
+
+    def rotate_fn(inputs, rotate_fp32=False, block_size=None):
+        calls.append((rotate_fp32, block_size))
+        return inputs + 10
+
+    monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn)
+    quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back"}, enable=False)
+    )
+    inputs = torch.tensor([[1.0, 2.0]])
+
+    outputs = quantizer(inputs)
+
+    assert quantizer.rotate_back_is_enabled
+    # A disabled quantizer must still round-trip: forward rotate + rotate_back = no-op (+10 twice).
+    assert torch.equal(outputs, inputs + 20)
+    assert len(calls) == 2
+
+
+def test_tensor_quantizer_rotate_only_applies_once_when_disabled(monkeypatch):
+    calls = []
+
+    def rotate_fn(inputs, rotate_fp32=False, block_size=None):
+        calls.append((rotate_fp32, block_size))
+        return inputs + 10
+
+    monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn)
+    quantizer = TensorQuantizer(QuantizerAttributeConfig(rotate={"enable": True}, enable=False))
+    inputs = torch.tensor([[1.0, 2.0]])
+
+    outputs = quantizer(inputs)
+
+    assert not quantizer.rotate_back_is_enabled
+    assert torch.equal(outputs, inputs + 10)
+    assert len(calls) == 1
+
+
+def test_disable_rotate_preserves_type():
+    # RotateConfig: enable off, other fields retained.
+    quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back", "block_size": 8})
+    )
+    assert isinstance(quantizer._rotate, RotateConfig)
+    quantizer.disable_rotate()
+    assert isinstance(quantizer._rotate, RotateConfig)
+    assert quantizer._rotate.enable is False
+    assert quantizer._rotate.mode == "rotate_back"
+    assert quantizer._rotate.block_size == 8
+    assert not quantizer.rotate_is_enabled
+    quantizer.disable_rotate()  # idempotent
+    assert quantizer._rotate.enable is False
+
+    # Raw dict (old checkpoints).
+    quantizer._rotate = {"enable": True, "mode": "rotate", "block_size": 4}
+    quantizer.disable_rotate()
+    assert quantizer._rotate == {"enable": False, "mode": "rotate", "block_size": 4}
+
+    # Bool.
+    quantizer._rotate = True
+    quantizer.disable_rotate()
+    assert quantizer._rotate is False
+
+
+def test_sequential_quantizer_disable_rotate_delegates():
+    q0 = TensorQuantizer(QuantizerAttributeConfig(rotate={"enable": True}))
+    q1 = TensorQuantizer(QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back"}))
+    seq = SequentialQuantizer(q0, q1)
+
+    seq.disable_rotate()
+
+    assert not q0.rotate_is_enabled
+    assert not q1.rotate_is_enabled
+
+
+def _make_rotated_qlinear(monkeypatch, calls, rotate, backend_name):
+    def rotate_fn(inputs, rotate_fp32=False, block_size=None):
+        calls.append((rotate_fp32, block_size))
+        return inputs + 10
+
+    def backend(inputs, _tq):
+        return inputs * 2
+
+    monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn)
+    register_quant_backend(backend_name, backend)
+    qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3))
+    qlinear.input_quantizer.disable()
+    qlinear.output_quantizer.disable()
+    qlinear.weight_quantizer.set_from_attribute_config(
+        QuantizerAttributeConfig(rotate=rotate, backend=backend_name)
+    )
+    return qlinear
+
+
+def test_fold_weight_disables_rotation_no_double_rotate(monkeypatch):
+    calls = []
+    backend_name = "test_fold_rotate_backend"
+    qlinear = _make_rotated_qlinear(monkeypatch, calls, {"enable": True}, backend_name)
+    try:
+        qlinear.weight_quantizer.amax = torch.tensor(1.0)
+        weight0 = qlinear.weight.detach().clone()
+        x = torch.randn(2, 4)
+        out_before = qlinear(x)
+
+        qlinear.fold_weight()
+
+        # Rotation (+10) then backend (*2) is baked into the stored weight.
+        assert torch.allclose(qlinear.weight, (weight0 + 10) * 2)
+        assert not qlinear.weight_quantizer.is_enabled
+        assert not qlinear.weight_quantizer.rotate_is_enabled
+        assert not hasattr(qlinear.weight_quantizer, "_amax")
+
+        calls_after_fold = len(calls)
+        out_after = qlinear(x)
+        assert torch.allclose(out_after, out_before)  # fails pre-fix (weight re-rotated)
+        assert len(calls) == calls_after_fold  # no re-rotation on forward
+
+        # Second fold is a no-op: quantizer is disabled.
+        weight_after = qlinear.weight.detach().clone()
+        qlinear.fold_weight()
+        assert torch.allclose(qlinear.weight, weight_after)
+    finally:
+        unregister_quant_backend(backend_name)
+
+
+def test_fold_weight_keep_attrs_keeps_amax_disables_rotation(monkeypatch):
+    calls = []
+    backend_name = "test_fold_rotate_backend_keep"
+    qlinear = _make_rotated_qlinear(monkeypatch, calls, {"enable": True}, backend_name)
+    try:
+        qlinear.weight_quantizer.amax = torch.tensor(1.0)
+
+        qlinear.fold_weight(keep_attrs=True)
+
+        assert hasattr(qlinear.weight_quantizer, "_amax")
+        assert not qlinear.weight_quantizer.is_enabled
+        assert not qlinear.weight_quantizer.rotate_is_enabled
+    finally:
+        unregister_quant_backend(backend_name)
+
+
+def test_fold_weight_rotate_back_no_double_rotate(monkeypatch):
+    calls = []
+    backend_name = "test_fold_rotate_back_backend"
+    qlinear = _make_rotated_qlinear(
+        monkeypatch, calls, {"enable": True, "mode": "rotate_back"}, backend_name
+    )
+    try:
+        x = torch.randn(2, 4)
+        out_before = qlinear(x)
+
+        qlinear.fold_weight()
+
+        assert not qlinear.weight_quantizer.rotate_is_enabled
+        assert not qlinear.weight_quantizer.is_enabled
+
+        calls_after_fold = len(calls)
+        out_after = qlinear(x)
+        assert torch.allclose(out_after, out_before)  # fails pre-fix (weight re-rotated)
+        assert len(calls) == calls_after_fold  # no rotate calls added on forward
+    finally:
+        unregister_quant_backend(backend_name)
+
+
 WINT4INT8_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},

From e0c70b3c8cc0777068a9732c28cbb607f7b4bc29 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Fri, 3 Jul 2026 18:39:04 +0000
Subject: [PATCH 05/10] Address tensor quantizer rotation review comments

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/config.py         |  27 +++-
 .../quantization/nn/modules/quant_module.py   |   8 +-
 .../nn/modules/tensor_quantizer.py            |  17 ++-
 tests/unit/torch/quantization/test_print.py   |  20 +++
 .../quantization/test_tensor_quant_cpu.py     | 126 ++++++++----------
 5 files changed, 110 insertions(+), 88 deletions(-)

diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py
index 8ad3299466e..44ae162c474 100644
--- a/modelopt/torch/quantization/config.py
+++ b/modelopt/torch/quantization/config.py
@@ -286,10 +286,29 @@ class RotateConfig(ModeloptBaseConfig):
     for transform details.
     """
 
-    enable: bool = False
-    mode: Literal["rotate", "rotate_back"] = "rotate"
-    rotate_fp32: bool = False
-    block_size: int | None = None
+    enable: bool = ModeloptField(
+        default=False,
+        title="Enable input rotation.",
+        description="If True, applies a normalized Hadamard transform before quantization.",
+    )
+    mode: Literal["rotate", "rotate_back"] = ModeloptField(
+        default="rotate",
+        title="Rotation mode.",
+        description=(
+            "Use 'rotate' for input rotation only, or 'rotate_back' to apply the transform "
+            "again after fake quantization."
+        ),
+    )
+    rotate_fp32: bool = ModeloptField(
+        default=False,
+        title="Run rotation in float32.",
+        description="If True, computes the rotation in float32 before casting back to the input dtype.",
+    )
+    block_size: int | None = ModeloptField(
+        default=None,
+        title="Rotation block size.",
+        description="Positive block size for block-wise rotation, or None to rotate the full input.",
+    )
 
     @field_validator("block_size", mode="before")
     @classmethod
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 137fa465d15..9c9aee478a8 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -153,10 +153,10 @@ def _fold_weight_quantizer(
     def fold_weight(self, keep_attrs: bool = False):
         """Bake each fake-quant weight quantizer into its weight for faster eval.
 
-        Every fake-quant weight quantizer is folded regardless of its enabled state. Both its
-        quantization and its input rotation are baked into the stored weight and then disabled,
-        so subsequent forwards neither re-quantize nor re-rotate. Calibration buffers
-        (``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``.
+        Every fake-quant weight quantizer is folded regardless of its enabled state. The folded
+        transform is baked into the stored weight and then disabled, so subsequent forwards use
+        the stored weight directly. Calibration buffers (``_pre_quant_scale``, ``_amax``) are
+        dropped unless ``keep_attrs``.
         """
         # Handle all attributes that end with _weight_quantizer
         for name in dir(self):
diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
index 20cdcfc3c77..7ecb95f9547 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -1220,6 +1220,14 @@ def _short_tensor(self, tensor: torch.Tensor, fmt=".2e"):
             return f"{tensor.item():{fmt}}"
         return f"[{tensor.min().item():{fmt}}, {tensor.max().item():{fmt}}]({tensor.numel()})"
 
+    def _rotation_extra_repr(self):
+        s = " rotated" if self.rotate_is_enabled else ""
+        s += " (rotate_back)" if self.rotate_back_is_enabled else ""
+        s += " (fp32)" if self.rotate_is_fp32 else ""
+        if self.rotate_block_size is not None:
+            s += f" (block={self.rotate_block_size})"
+        return s
+
     def extra_repr(self):
         """Set the extra information about this module."""
         if self._disabled:
@@ -1229,7 +1237,8 @@ def extra_repr(self):
                 if self.pre_quant_scale is not None
                 else ""
             )
-            return "disabled"
+            s += self._rotation_extra_repr()
+            return s
         s = f"{'unsigned ' if self._unsigned else ''}{self._num_bits} bit"
         s += " narrow" if (self._narrow_range) else ""
         s += " fake" if (self._fake_quant) else ""
@@ -1243,11 +1252,7 @@ def extra_repr(self):
             if self.pre_quant_scale is not None
             else ""
         )
-        s += " rotated" if self.rotate_is_enabled else ""
-        s += " (rotate_back)" if self.rotate_back_is_enabled else ""
-        s += " (fp32)" if self.rotate_is_fp32 else ""
-        if self.rotate_block_size is not None:
-            s += f" (block={self.rotate_block_size})"
+        s += self._rotation_extra_repr()
         s += (
             f" calibrator={self._calibrator.__class__.__name__}"
             if (self._calibrator is not None)
diff --git a/tests/unit/torch/quantization/test_print.py b/tests/unit/torch/quantization/test_print.py
index 56a351e799e..d6e3e0d3181 100644
--- a/tests/unit/torch/quantization/test_print.py
+++ b/tests/unit/torch/quantization/test_print.py
@@ -20,6 +20,7 @@
 
 from modelopt.torch.quantization import calib, tensor_quant
 from modelopt.torch.quantization import nn as qnn
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
 
 
@@ -32,6 +33,25 @@ def test_print_tensor_quantizer(self):
         test_quantizer = TensorQuantizer()
         print(test_quantizer)
 
+    def test_disabled_tensor_quantizer_repr_shows_enabled_state(self):
+        test_quantizer = TensorQuantizer(
+            QuantizerAttributeConfig(
+                enable=False,
+                rotate={
+                    "enable": True,
+                    "mode": "rotate_back",
+                    "rotate_fp32": True,
+                    "block_size": 8,
+                },
+            )
+        )
+        test_quantizer.pre_quant_scale = torch.tensor([1.0, 2.0])
+
+        assert test_quantizer.extra_repr() == (
+            "disabled pre_quant_scale=[1.00e+00, 2.00e+00](2)"
+            " rotated (rotate_back) (fp32) (block=8)"
+        )
+
     def test_print_module(self):
         class _TestModule(nn.Module):
             def __init__(self):
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index a15fff12e58..72ebb1e0b98 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -133,23 +133,29 @@ def backend(inputs, _tq):
         unregister_quant_backend("test_rotate_mode_backend")
 
 
-def test_tensor_quantizer_rotate_mode_preserves_default_path(monkeypatch):
-    outputs, inputs, calls, quantizer = _run_rotated_backend(monkeypatch, rotate={"enable": True})
-
-    assert not quantizer.rotate_back_is_enabled
-    assert torch.equal(outputs, (inputs + 10) * 2)
-    assert calls == [(False, None)]
-
-
-def test_tensor_quantizer_rotate_mode_can_rotate_back(monkeypatch):
+@pytest.mark.parametrize(
+    ("rotate", "rotate_back_enabled", "expected_calls", "expected_fn"),
+    [
+        ({"enable": True}, False, [(False, None)], lambda inputs: (inputs + 10) * 2),
+        (
+            {"enable": True, "mode": "rotate_back", "rotate_fp32": True, "block_size": 8},
+            True,
+            [(True, 8), (True, 8)],
+            lambda inputs: ((inputs + 10) * 2) + 10,
+        ),
+    ],
+)
+def test_tensor_quantizer_rotate_modes(
+    monkeypatch, rotate, rotate_back_enabled, expected_calls, expected_fn
+):
     outputs, inputs, calls, quantizer = _run_rotated_backend(
         monkeypatch,
-        rotate={"enable": True, "mode": "rotate_back", "rotate_fp32": True, "block_size": 8},
+        rotate=rotate,
     )
 
-    assert quantizer.rotate_back_is_enabled
-    assert torch.equal(outputs, ((inputs + 10) * 2) + 10)
-    assert calls == [(True, 8), (True, 8)]
+    assert quantizer.rotate_back_is_enabled is rotate_back_enabled
+    assert torch.equal(outputs, expected_fn(inputs))
+    assert calls == expected_calls
 
 
 def test_tensor_quantizer_rotate_back_rejects_real_quant(monkeypatch):
@@ -173,28 +179,16 @@ def fail_if_rotated(inputs, rotate_fp32=False, block_size=None):
         quantizer(torch.tensor([[1.0, 2.0]]))
 
 
-def test_tensor_quantizer_rotate_back_roundtrips_when_disabled(monkeypatch):
-    calls = []
-
-    def rotate_fn(inputs, rotate_fp32=False, block_size=None):
-        calls.append((rotate_fp32, block_size))
-        return inputs + 10
-
-    monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn)
-    quantizer = TensorQuantizer(
-        QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back"}, enable=False)
-    )
-    inputs = torch.tensor([[1.0, 2.0]])
-
-    outputs = quantizer(inputs)
-
-    assert quantizer.rotate_back_is_enabled
-    # A disabled quantizer must still round-trip: forward rotate + rotate_back = no-op (+10 twice).
-    assert torch.equal(outputs, inputs + 20)
-    assert len(calls) == 2
-
-
-def test_tensor_quantizer_rotate_only_applies_once_when_disabled(monkeypatch):
+@pytest.mark.parametrize(
+    ("rotate", "rotate_back_enabled", "expected_call_count", "expected_fn"),
+    [
+        ({"enable": True}, False, 1, lambda inputs: inputs + 10),
+        ({"enable": True, "mode": "rotate_back"}, True, 2, lambda inputs: inputs + 20),
+    ],
+)
+def test_tensor_quantizer_disabled_rotate_modes_roundtrip(
+    monkeypatch, rotate, rotate_back_enabled, expected_call_count, expected_fn
+):
     calls = []
 
     def rotate_fn(inputs, rotate_fp32=False, block_size=None):
@@ -202,14 +196,14 @@ def rotate_fn(inputs, rotate_fp32=False, block_size=None):
         return inputs + 10
 
     monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn)
-    quantizer = TensorQuantizer(QuantizerAttributeConfig(rotate={"enable": True}, enable=False))
+    quantizer = TensorQuantizer(QuantizerAttributeConfig(rotate=rotate, enable=False))
     inputs = torch.tensor([[1.0, 2.0]])
 
     outputs = quantizer(inputs)
 
-    assert not quantizer.rotate_back_is_enabled
-    assert torch.equal(outputs, inputs + 10)
-    assert len(calls) == 1
+    assert quantizer.rotate_back_is_enabled is rotate_back_enabled
+    assert torch.equal(outputs, expected_fn(inputs))
+    assert len(calls) == expected_call_count
 
 
 def test_disable_rotate_preserves_type():
@@ -249,7 +243,7 @@ def test_sequential_quantizer_disable_rotate_delegates():
     assert not q1.rotate_is_enabled
 
 
-def _make_rotated_qlinear(monkeypatch, calls, rotate, backend_name):
+def _make_qlinear_with_backend(monkeypatch, calls, backend_name, rotate=False):
     def rotate_fn(inputs, rotate_fp32=False, block_size=None):
         calls.append((rotate_fp32, block_size))
         return inputs + 10
@@ -268,10 +262,19 @@ def backend(inputs, _tq):
     return qlinear
 
 
-def test_fold_weight_disables_rotation_no_double_rotate(monkeypatch):
+@pytest.mark.parametrize(
+    ("rotate", "expected_weight_fn"),
+    [
+        ({"enable": True}, lambda weight: (weight + 10) * 2),
+        ({"enable": True, "mode": "rotate_back"}, lambda weight: ((weight + 10) * 2) + 10),
+    ],
+)
+def test_fold_weight_disables_quantizer_without_extra_transform(
+    monkeypatch, rotate, expected_weight_fn
+):
     calls = []
-    backend_name = "test_fold_rotate_backend"
-    qlinear = _make_rotated_qlinear(monkeypatch, calls, {"enable": True}, backend_name)
+    backend_name = "test_fold_backend"
+    qlinear = _make_qlinear_with_backend(monkeypatch, calls, backend_name, rotate=rotate)
     try:
         qlinear.weight_quantizer.amax = torch.tensor(1.0)
         weight0 = qlinear.weight.detach().clone()
@@ -280,16 +283,15 @@ def test_fold_weight_disables_rotation_no_double_rotate(monkeypatch):
 
         qlinear.fold_weight()
 
-        # Rotation (+10) then backend (*2) is baked into the stored weight.
-        assert torch.allclose(qlinear.weight, (weight0 + 10) * 2)
+        assert torch.allclose(qlinear.weight, expected_weight_fn(weight0))
         assert not qlinear.weight_quantizer.is_enabled
         assert not qlinear.weight_quantizer.rotate_is_enabled
         assert not hasattr(qlinear.weight_quantizer, "_amax")
 
         calls_after_fold = len(calls)
         out_after = qlinear(x)
-        assert torch.allclose(out_after, out_before)  # fails pre-fix (weight re-rotated)
-        assert len(calls) == calls_after_fold  # no re-rotation on forward
+        assert torch.allclose(out_after, out_before)
+        assert len(calls) == calls_after_fold
 
         # Second fold is a no-op: quantizer is disabled.
         weight_after = qlinear.weight.detach().clone()
@@ -299,10 +301,10 @@ def test_fold_weight_disables_rotation_no_double_rotate(monkeypatch):
         unregister_quant_backend(backend_name)
 
 
-def test_fold_weight_keep_attrs_keeps_amax_disables_rotation(monkeypatch):
+def test_fold_weight_keep_attrs_keeps_amax(monkeypatch):
     calls = []
-    backend_name = "test_fold_rotate_backend_keep"
-    qlinear = _make_rotated_qlinear(monkeypatch, calls, {"enable": True}, backend_name)
+    backend_name = "test_fold_backend_keep"
+    qlinear = _make_qlinear_with_backend(monkeypatch, calls, backend_name)
     try:
         qlinear.weight_quantizer.amax = torch.tensor(1.0)
 
@@ -310,30 +312,6 @@ def test_fold_weight_keep_attrs_keeps_amax_disables_rotation(monkeypatch):
 
         assert hasattr(qlinear.weight_quantizer, "_amax")
         assert not qlinear.weight_quantizer.is_enabled
-        assert not qlinear.weight_quantizer.rotate_is_enabled
-    finally:
-        unregister_quant_backend(backend_name)
-
-
-def test_fold_weight_rotate_back_no_double_rotate(monkeypatch):
-    calls = []
-    backend_name = "test_fold_rotate_back_backend"
-    qlinear = _make_rotated_qlinear(
-        monkeypatch, calls, {"enable": True, "mode": "rotate_back"}, backend_name
-    )
-    try:
-        x = torch.randn(2, 4)
-        out_before = qlinear(x)
-
-        qlinear.fold_weight()
-
-        assert not qlinear.weight_quantizer.rotate_is_enabled
-        assert not qlinear.weight_quantizer.is_enabled
-
-        calls_after_fold = len(calls)
-        out_after = qlinear(x)
-        assert torch.allclose(out_after, out_before)  # fails pre-fix (weight re-rotated)
-        assert len(calls) == calls_after_fold  # no rotate calls added on forward
     finally:
         unregister_quant_backend(backend_name)
 

From b82e5e874169e704b02a89217d22c2fd8eac41fb Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Fri, 3 Jul 2026 21:35:09 +0000
Subject: [PATCH 06/10] Add W4A4 rotate eval config workaround

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 examples/llm_eval/quantization_utils.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/examples/llm_eval/quantization_utils.py b/examples/llm_eval/quantization_utils.py
index 80117bd1627..d7add825ca1 100644
--- a/examples/llm_eval/quantization_utils.py
+++ b/examples/llm_eval/quantization_utils.py
@@ -49,6 +49,29 @@
         ],
         "algorithm": "max",
     },
+    "NVFP4_W4A4_ROTATE": {
+        "quant_cfg": [
+            *mtq.config._base_disable_all,
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {**mtq.config._nvfp4_cfg, "rotate": {"enable": True}},
+                "enable": True,
+            },
+            {
+                "quantizer_name": "*input_quantizer",
+                "cfg": {**mtq.config._nvfp4_cfg, "rotate": {"enable": True}},
+                "enable": True,
+            },
+            *mtq.config._default_disabled_quantizer_cfg,
+            # TODO: The embedding token/input path creates this rotate issue; replace with a better fix.
+            {
+                "quantizer_name": "*embed_tokens*quantizer",
+                "cfg": {"rotate": False},
+                "enable": False,
+            },
+        ],
+        "algorithm": "max",
+    },
 }
 
 

From 7ec1afc63f49248e91c2c3c560f0bf1dbd23b758 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Sat, 4 Jul 2026 02:26:38 +0000
Subject: [PATCH 07/10] Clear rotation when disabling tensor quantizers

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 .../nn/modules/tensor_quantizer.py            |  3 ++
 .../quantization/test_quant_embedding.py      | 36 +++++++++++++++++++
 .../quantization/test_tensor_quant_cpu.py     | 16 +++++++++
 3 files changed, 55 insertions(+)

diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
index 7ecb95f9547..c7649f9383d 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -266,6 +266,9 @@ def _block_sizes_setter(val):
             )
             setattr(self, _tq_attribute_name, _setter(val))
 
+        if isinstance(attribute_cfg, dict) and attribute_cfg == {"enable": False}:
+            self.disable_rotate()
+
         if self.is_mx_format:
             self._pass_through_bwd = True
 
diff --git a/tests/unit/torch/quantization/test_quant_embedding.py b/tests/unit/torch/quantization/test_quant_embedding.py
index 5a3d2535dcf..d7053d1e09f 100644
--- a/tests/unit/torch/quantization/test_quant_embedding.py
+++ b/tests/unit/torch/quantization/test_quant_embedding.py
@@ -108,6 +108,42 @@ def test_wildcard_config_keeps_input_quantizer_disabled(self):
         # Forward still works — input_quantizer is disabled and never applied.
         qemb(torch.randint(0, VOCAB_SIZE, (4, 6)))
 
+    def test_disable_update_clears_hard_disabled_input_quantizer_rotate_state(self):
+        qemb = _make_quant_embedding()
+        set_quantizer_attributes_partial(
+            qemb,
+            "*input_quantizer",
+            {"enable": True, "num_bits": 4, "rotate": {"enable": True, "mode": "rotate_back"}},
+        )
+        assert not qemb.input_quantizer.is_enabled
+        assert qemb.input_quantizer.num_bits == 4
+        assert qemb.input_quantizer.rotate_is_enabled
+        assert qemb.input_quantizer.extra_repr() == "disabled rotated (rotate_back)"
+
+        set_quantizer_attributes_partial(qemb, "*input_quantizer", {"enable": False})
+
+        assert not qemb.input_quantizer.is_enabled
+        assert qemb.input_quantizer.num_bits == 4
+        assert not qemb.input_quantizer.rotate_is_enabled
+        assert qemb.input_quantizer.extra_repr() == "disabled"
+
+    def test_disable_update_clears_weight_quantizer_rotate_state(self):
+        qemb = _make_quant_embedding()
+        set_quantizer_attributes_partial(
+            qemb,
+            "*weight_quantizer",
+            {"enable": True, "num_bits": 4, "rotate": {"enable": True, "mode": "rotate_back"}},
+        )
+        assert qemb.weight_quantizer.is_enabled
+        assert qemb.weight_quantizer.rotate_is_enabled
+        assert "rotated (rotate_back)" in qemb.weight_quantizer.extra_repr()
+
+        set_quantizer_attributes_partial(qemb, "*weight_quantizer", {"enable": False})
+
+        assert not qemb.weight_quantizer.is_enabled
+        assert not qemb.weight_quantizer.rotate_is_enabled
+        assert "rotated" not in qemb.weight_quantizer.extra_repr()
+
 
 # Export-path tests for QuantEmbedding live in tests/gpu/torch/export/test_export_embedding.py
 # because _export_quantized_weight bottoms out in torch.cuda.empty_cache(), which raises on
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 72ebb1e0b98..ba352ec2162 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -206,6 +206,22 @@ def rotate_fn(inputs, rotate_fp32=False, block_size=None):
     assert len(calls) == expected_call_count
 
 
+def test_disable_only_update_clears_regular_quantizer_rotate_state():
+    quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back", "block_size": 8})
+    )
+    assert quantizer.rotate_is_enabled
+    assert quantizer.rotate_back_is_enabled
+
+    quantizer.set_from_attribute_config({"enable": False})
+
+    assert not quantizer.is_enabled
+    assert not quantizer.rotate_is_enabled
+    assert isinstance(quantizer._rotate, RotateConfig)
+    assert quantizer._rotate.mode == "rotate_back"
+    assert quantizer._rotate.block_size == 8
+
+
 def test_disable_rotate_preserves_type():
     # RotateConfig: enable off, other fields retained.
     quantizer = TensorQuantizer(

From d4cda04f42dc795329dbbe0c8a14b4e070dcafa2 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Sat, 4 Jul 2026 02:29:09 +0000
Subject: [PATCH 08/10] Remove temporary W4A4 rotate eval config

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 examples/llm_eval/quantization_utils.py | 23 -----------------------
 1 file changed, 23 deletions(-)

diff --git a/examples/llm_eval/quantization_utils.py b/examples/llm_eval/quantization_utils.py
index d7add825ca1..80117bd1627 100644
--- a/examples/llm_eval/quantization_utils.py
+++ b/examples/llm_eval/quantization_utils.py
@@ -49,29 +49,6 @@
         ],
         "algorithm": "max",
     },
-    "NVFP4_W4A4_ROTATE": {
-        "quant_cfg": [
-            *mtq.config._base_disable_all,
-            {
-                "quantizer_name": "*weight_quantizer",
-                "cfg": {**mtq.config._nvfp4_cfg, "rotate": {"enable": True}},
-                "enable": True,
-            },
-            {
-                "quantizer_name": "*input_quantizer",
-                "cfg": {**mtq.config._nvfp4_cfg, "rotate": {"enable": True}},
-                "enable": True,
-            },
-            *mtq.config._default_disabled_quantizer_cfg,
-            # TODO: The embedding token/input path creates this rotate issue; replace with a better fix.
-            {
-                "quantizer_name": "*embed_tokens*quantizer",
-                "cfg": {"rotate": False},
-                "enable": False,
-            },
-        ],
-        "algorithm": "max",
-    },
 }
 
 

From 5de968325b2cacdf6612f6cd6556e9967299eecc Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Mon, 6 Jul 2026 17:51:11 +0000
Subject: [PATCH 09/10] Move rotate.mode changelog entry under features

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 CHANGELOG.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index d1791fbc800..48aed5eb8b9 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,7 +19,6 @@ Changelog
 
 **New Features**
 
-- Add ``rotate.mode`` to torch quantizer configs. The default ``"rotate"`` keeps the existing rotate-before-quantize behavior; ``"rotate_back"`` enables fake-quant rotate → quantize → rotate-back for TensorQuantizer.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.
@@ -43,6 +42,7 @@ Changelog
 - Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet.
 - Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline.
 - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface/<model>/auto_quantize/``.
+- Add ``rotate.mode`` to torch quantizer configs. The default ``"rotate"`` keeps the existing rotate-before-quantize behavior; ``"rotate_back"`` enables fake-quant rotate → quantize → rotate-back for TensorQuantizer.
 
 **Bug Fixes**
 

From a600718ea7ccb0f9dc8fd3a81d4ce655b5c25559 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Mon, 6 Jul 2026 17:59:00 +0000
Subject: [PATCH 10/10] Test rotation-enabled quantizer backward

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 .../nn/modules/quant_embedding.py             |  3 +++
 tests/gpu/torch/quantization/test_hadamard.py | 23 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/modelopt/torch/quantization/nn/modules/quant_embedding.py b/modelopt/torch/quantization/nn/modules/quant_embedding.py
index 5004c11b3c1..690afd8a334 100644
--- a/modelopt/torch/quantization/nn/modules/quant_embedding.py
+++ b/modelopt/torch/quantization/nn/modules/quant_embedding.py
@@ -40,6 +40,9 @@ class _QuantEmbedding(QuantModule):
     table (weight) and the lookup output (an activation feeding downstream layers)
     are quantizable.
 
+    TODO: Remove the example-side ``*embed_tokens*quantizer`` rotation workaround
+    once rotation configs stop targeting embedding token/input paths.
+
     Quantizer roles:
         - ``weight_quantizer``: quantizes the embedding table (``self.weight``).
         - ``input_quantizer``: permanently disabled placeholder
diff --git a/tests/gpu/torch/quantization/test_hadamard.py b/tests/gpu/torch/quantization/test_hadamard.py
index 7034feda75c..d426a3e2ffa 100644
--- a/tests/gpu/torch/quantization/test_hadamard.py
+++ b/tests/gpu/torch/quantization/test_hadamard.py
@@ -30,11 +30,13 @@
 from _test_utils.torch.quantization.models import SDPAAttention
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.conversion import (
     set_quantizer_by_cfg,
     set_quantizer_by_cfg_context,
 )
 from modelopt.torch.quantization.nn.functional import normalized_hadamard_transform
+from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
 
 
 @pytest.mark.parametrize(
@@ -109,3 +111,24 @@ def test_kv_rotate(rotate_fp32):
     assert not torch.allclose(output_ref, output_test1, atol=0.05)
 
     mtq.unregister(SDPAAttention)
+
+
+@pytest.mark.parametrize(
+    "mode",
+    ["rotate", "rotate_back"],
+)
+def test_rotate_backward(mode):
+    """Autograd backward should flow through a rotation-enabled TensorQuantizer."""
+    x = torch.randn(4, 64, device="cuda", requires_grad=True)
+    quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(num_bits=8, axis=None, rotate={"enable": True, "mode": mode}),
+        amax=x.detach().abs().amax(),
+    ).cuda()
+
+    out = quantizer(x)
+    out.sum().backward()
+
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+    assert torch.isfinite(x.grad).all()
+    assert not torch.all(x.grad == 0)