From dd5b59d62942a9025be8fb26958cb435ed2d4973 Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Sun, 26 Jul 2026 21:06:13 -0700
Subject: [PATCH 01/16] Add optional tooling to accumulate quantization scaling
 factors for inference.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 .../debug/features/log_fp8_tensor_stats.py    | 29 +-----
 transformer_engine/pytorch/module/_common.py  | 67 +++++++++++++-
 .../pytorch/module/grouped_linear.py          | 92 ++++++++++++++++++-
 .../pytorch/module/layernorm_linear.py        | 61 +++++++++++-
 .../pytorch/module/layernorm_mlp.py           | 78 +++++++++++++++-
 transformer_engine/pytorch/module/linear.py   | 71 +++++++++++++-
 transformer_engine/pytorch/tensor/utils.py    | 20 ++++
 7 files changed, 386 insertions(+), 32 deletions(-)

diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py
index 96f1b644cf..e05a90f0d8 100644
--- a/transformer_engine/debug/features/log_fp8_tensor_stats.py
+++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py
@@ -23,35 +23,12 @@
 )
 from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
 from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer
-
-try:
-    from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
-
-    _nvfp4_available = True
-except ImportError:
-    _nvfp4_available = False
-    NVFP4Quantizer = None
+from transformer_engine.pytorch.tensor.utils import get_quantization_recipe_name
 
 
 ALL_RECIPE_NAMES = ["fp8_delayed_scaling", "fp8_current_scaling", "mxfp8", "fp8_block_scaling"]
 
 
-def _get_recipe_name(quantizer: Optional[Quantizer]):
-    if quantizer is None:
-        return ""
-    if isinstance(quantizer, Float8Quantizer):
-        return "fp8_delayed_scaling"
-    if isinstance(quantizer, Float8CurrentScalingQuantizer):
-        return "fp8_current_scaling"
-    if isinstance(quantizer, MXFP8Quantizer):
-        return "mxfp8"
-    if isinstance(quantizer, Float8BlockQuantizer):
-        return "fp8_block_scaling"
-    if _nvfp4_available and isinstance(quantizer, NVFP4Quantizer):
-        return "nvfp4"
-    raise ValueError(f"Unsupported quantizer type: {type(quantizer)}")
-
-
 def _get_new_quantizer(recipe_name, fp8_dtype):
     if recipe_name == "fp8_block_scaling":
         return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True)
@@ -336,7 +313,9 @@ def inspect_tensor(
             )
             return
 
-        recipe_name = _get_recipe_name(quantizer)
+        recipe_name = get_quantization_recipe_name(quantizer)
+        if recipe_name == "nvfp4_rowwise":
+            recipe_name = "nvfp4"
 
         for stat in config["stats"]:
             self.check_if_stat_is_supported(
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 908e88f30d..cac60a74d0 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -6,7 +6,7 @@
 
 import dataclasses
 import queue
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 
@@ -14,9 +14,74 @@
 from ..constants import TE_DType
 from ..export import is_in_onnx_export_mode
 from ..tensor.hybrid_tensor import HybridQuantizer
+from ..tensor.utils import get_quantization_recipe_name
 from ..utils import get_default_init_method
 
 
+def _get_scale_buffer_info(
+    tensor_name: str,
+    tensor: Any,
+    quantizer: Any,
+) -> Optional[Tuple[str, Optional[torch.Tensor]]]:
+    """Get the calibration buffer name and value for a quantized tensor."""
+    recipe = get_quantization_recipe_name(quantizer)
+    if not recipe:
+        return None
+    if recipe == "fp8_current_scaling":
+        metadata_name = "scale_inv"
+        metadata = getattr(tensor, "_scale_inv", None)
+    else:
+        metadata_name = "amax_rowwise"
+        metadata = getattr(
+            tensor,
+            "_amax_rowwise",
+            getattr(quantizer, "amax", None),
+        )
+    buffer_name = f"{tensor_name}_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
+    return buffer_name, metadata
+
+
+def _update_scale_buffers(
+    scale_buffers: Dict[str, Optional[torch.Tensor]],
+    scale_updates: Dict[str, Optional[torch.Tensor]],
+    activation_scale_decay: float = 0.0,
+) -> None:
+    """Merge observed scaling factors into checkpoint buffers."""
+    for buffer_name, scale in scale_updates.items():
+        if scale is None:
+            continue
+        if buffer_name.startswith("input"):
+            observed_scale = scale.detach().float()
+            scale_buffer = scale_buffers.get(buffer_name)
+            if scale_buffer is not None and scale_buffer.shape != observed_scale.shape:
+                raise RuntimeError(
+                    "Quantized scaling-factor buffer shape changed from "
+                    f"{tuple(scale_buffer.shape)} to {tuple(observed_scale.shape)}"
+                )
+            if activation_scale_decay == 0.0:
+                # If not using scale decay, just buffer the current scaling.
+                scale_buffers[buffer_name] = observed_scale
+                continue
+            if scale_buffer is None:
+                # Initialize the rolling activation scaling factor.
+                # Requires CUDA graph warmup step.
+                scale_buffer = torch.zeros_like(observed_scale)
+                scale_buffers[buffer_name] = scale_buffer
+            # Track a decaying maximum so early-training activation
+            # outliers do not permanently determine the inference scale.
+            scale_buffer.mul_(activation_scale_decay)
+            torch.maximum(
+                scale_buffer,
+                observed_scale,
+                out=scale_buffer,
+            )
+        else:
+            # Keep a reference to the current weight metadata without
+            # allocating or copying a separate buffer.
+            # Requires CUDA graph warmup step.
+            scale_buffers[buffer_name] = scale.detach()
+
+
 def set_quantizer_amax_reduction_group(quantizer, amax_reduction_group) -> None:
     """Set the amax reduction group on a quantizer; no-op if it doesn't support it.
 
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 8b078944f1..221ea0d36c 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -4,7 +4,7 @@
 
 """GroupedLinear API"""
 
-from typing import Union, Optional, Callable, Tuple, List
+from typing import Union, Optional, Callable, Tuple, List, Dict
 from itertools import chain
 import os
 import warnings
@@ -31,7 +31,12 @@
     _clear_high_precision_init_val,
     _get_high_precision_init_val,
 )
-from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore
+from ._common import (
+    _get_scale_buffer_info,
+    _update_scale_buffers,
+    can_reconstruct_wgrad_input_from_original,
+    WeightGradStore,
+)
 from . import _split_quantization
 from ..quantization import FP8GlobalStateManager, QuantizerRole
 from ..utils import (
@@ -152,6 +157,29 @@ def is_module_grouped_tensor_path_supported(
     return False
 
 
+def _update_grouped_scale_buffers(
+    scale_buffers: Dict[str, Optional[torch.Tensor]],
+    input_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
+    weight_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
+    input_quantizer: Optional[Quantizer],
+    weight_quantizer: Optional[Quantizer],
+    activation_scale_decay: float,
+) -> None:
+    """Update GroupedLinear PTQ calibration buffers with per-GEMM metadata."""
+    scale_updates = {}
+    for index, tensor in enumerate(input_tensors):
+        scale_buffer = _get_scale_buffer_info(f"input_gemm{index}", tensor, input_quantizer)
+        if scale_buffer is not None:
+            scale_updates[scale_buffer[0]] = scale_buffer[1]
+    for index, tensor in enumerate(weight_tensors):
+        scale_buffer = _get_scale_buffer_info(
+            f"weight_gemm{index}", tensor, weight_quantizer
+        )
+        if scale_buffer is not None:
+            scale_updates[scale_buffer[0]] = scale_buffer[1]
+    _update_scale_buffers(scale_buffers, scale_updates, activation_scale_decay)
+
+
 class _GroupedLinear(torch.autograd.Function):
     """GroupedLinear semi-top level module
     Calls custom cuda extensions.
@@ -428,6 +456,8 @@ def _forward_grouped_tensor(
         save_original_input: bool,
         single_grouped_weight: bool,
         single_grouped_bias: bool,
+        scale_buffers: Optional[Dict[str, Optional[torch.Tensor]]],
+        quantized_scaling_factor_buffering_decay: float,
         weights: Tuple[torch.Tensor, ...],
         biases: Tuple[torch.Tensor, ...],
         out: Optional[torch.Tensor] = None,
@@ -536,6 +566,19 @@ def _forward_grouped_tensor(
             use_split_accumulator=use_split_accumulator,
         )
 
+        if scale_buffers is not None:
+            grouped_inputs = grouped_x.quantized_tensors
+            if grouped_inputs is None:
+                grouped_inputs = grouped_x.split_into_quantized_tensors()
+            _update_grouped_scale_buffers(
+                scale_buffers,
+                grouped_inputs,
+                weights_for_gemm,
+                input_quantizers[0],
+                weight_quantizers[0],
+                quantized_scaling_factor_buffering_decay,
+            )
+
         if is_grad_enabled:
             input_to_save = grouped_x
             if weight_requires_grad:
@@ -657,6 +700,8 @@ def forward(
             single_grouped_weight,
             single_grouped_bias,
             use_grouped_tensor,
+            scale_buffers,
+            quantized_scaling_factor_buffering_decay,
         ) = non_tensor_args
         recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
         backward_override = recipe.backward_override if recipe is not None else None
@@ -809,6 +854,10 @@ def forward(
                 save_original_input=save_original_input,
                 single_grouped_weight=single_grouped_weight,
                 single_grouped_bias=single_grouped_bias,
+                scale_buffers=scale_buffers,
+                quantized_scaling_factor_buffering_decay=(
+                    quantized_scaling_factor_buffering_decay
+                ),
                 weights=weights,
                 biases=biases,
                 out=out,
@@ -855,6 +904,16 @@ def forward(
         else:
             weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]
 
+        if scale_buffers is not None:
+            _update_grouped_scale_buffers(
+                scale_buffers,
+                inputmats,
+                weights_fp8,
+                input_quantizers[0],
+                weight_quantizers[0],
+                quantized_scaling_factor_buffering_decay,
+            )
+
         # Initialize biases
         bias_dtype = activation_dtype
         if fp8 and activation_dtype == torch.float32:
@@ -1590,6 +1649,10 @@ class GroupedLinear(TransformerEngineBaseModule):
                        and saving the original input tensor may reduce the memory usage.
                        Requires input quantizers that can safely reproduce their results from the
                        original input. Cannot work with FP8 DelayedScaling recipe.
+    buffer_quantized_scaling_factors : bool, default = False
+                       If set to ``True``, maintain nonpersistent input and weight quantization
+                       metadata buffers for inference checkpoint export. Buffers store metadata
+                       per grouped GEMM, using inverse scales directly for FP8 current scaling.
     single_grouped_weight : bool, default = False
                        If set to ``True``, grouped weights are stored as a single grouped parameter
                        instead of one parameter per GEMM.
@@ -1643,6 +1706,8 @@ def __init__(
         single_grouped_bias: bool = False,
         name: Optional[str] = None,
         use_grouped_tensor: Optional[bool] = None,
+        buffer_quantized_scaling_factors: bool = False,
+        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -1676,6 +1741,10 @@ def __init__(
                 f"use_grouped_tensor must be a bool or None, got {type(use_grouped_tensor)}."
             )
         self.use_grouped_tensor = use_grouped_tensor
+        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
+        self.quantized_scaling_factor_buffering_decay = (
+            quantized_scaling_factor_buffering_decay
+        )
         single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags(
             single_grouped_weight, single_grouped_bias
         )
@@ -2272,6 +2341,13 @@ def forward(
                     if cache_weight
                     else [None] * num_gemms
                 )
+            scale_buffers = None
+            if self.buffer_quantized_scaling_factors:
+                scale_buffers = {
+                    name: value
+                    for name, value in self._buffers.items()
+                    if name.endswith("_te_ptq_calibrated")
+                }
 
             non_tensor_args = (
                 self.apply_bias,
@@ -2298,6 +2374,8 @@ def forward(
                 self.single_grouped_weight,
                 use_grouped_bias,
                 self.use_grouped_tensor,
+                scale_buffers,
+                self.quantized_scaling_factor_buffering_decay,
             )
             out, new_workspaces = linear_fn(
                 *autograd_ctx,
@@ -2310,6 +2388,16 @@ def forward(
                 *bias_tensors,
             )
 
+            if scale_buffers is not None:
+                # Assign scaling-factor calibration buffers to the model.
+                # Materializing a new buffer requires a CUDA graph warmup step.
+                for name, value in scale_buffers.items():
+                    if value is not None:
+                        if name in self._buffers:
+                            setattr(self, name, value)
+                        else:
+                            self.register_buffer(name, value, persistent=False)
+
             if cache_weight:
                 for i, ws in enumerate(new_workspaces):
                     if ws is not None:
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 561e813348..e78d612f7b 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -17,7 +17,10 @@
 
 from transformer_engine.common.recipe import Recipe
 from transformer_engine.pytorch.torch_version import torch_version
-from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom
+from transformer_engine.pytorch.tensor.utils import (
+    clear_columnwise_cache,
+    is_custom,
+)
 from .base import (
     fill_userbuffers_buffer_for_all_gather,
     get_ub,
@@ -66,6 +69,8 @@
 from ..jit import no_torch_dynamo
 from ..graph import is_graph_capturing
 from ._common import (
+    _get_scale_buffer_info,
+    _update_scale_buffers,
     apply_normalization,
     noop_cat,
     set_quantizer_amax_reduction_group,
@@ -160,6 +165,8 @@ def forward(
             symmetric_ar_type,
             debug,
             is_fsdp2,
+            quantized_scaling_factor_buffering_decay,
+            scale_buffers,
         ) = non_tensor_args
         if fp8:
             backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override
@@ -383,6 +390,24 @@ def forward(
             if weight_quantizer is not None:
                 weight_quantizer.calibrate(weight)
 
+        if scale_buffers is not None:
+            scale_updates = {}
+            input_scale_buffer = _get_scale_buffer_info(
+                "input", ln_out_total, input_quantizer
+            )
+            if input_scale_buffer is not None:
+                scale_updates[input_scale_buffer[0]] = input_scale_buffer[1]
+            weight_scale_buffer = _get_scale_buffer_info(
+                "weight", weightmat, weight_quantizer
+            )
+            if weight_scale_buffer is not None:
+                scale_updates[weight_scale_buffer[0]] = weight_scale_buffer[1]
+            _update_scale_buffers(
+                scale_buffers,
+                scale_updates,
+                quantized_scaling_factor_buffering_decay,
+            )
+
         # Choose whether to use GEMM kernel with split accumulator
         use_split_accumulator = _2X_ACC_FPROP
         if fp8:
@@ -1301,6 +1326,15 @@ class LayerNormLinear(TransformerEngineBaseModule):
                    This can help in latency bound communication situations.
                    Requires PyTorch version 2.7.0 or higher. When set to ``None``, standard all-reduce
                    is used.
+    buffer_quantized_scaling_factors : bool, default = False
+                       If set to ``True``, maintain nonpersistent input and weight quantization
+                       metadata buffers for inference checkpoint export. Per-tensor buffers
+                       store raw global amaxes, except FP8 current scaling buffers, which store
+                       inverse scales directly.
+    quantized_scaling_factor_buffering_decay : float, default = 0.0
+                       Decay applied to buffered activation scaling factors before incorporating
+                       each new observation. Defaults to 0.0, in which case only the most recent
+                       scaling factor is buffered.
     """
 
     def __init__(
@@ -1333,6 +1367,8 @@ def __init__(
         delay_wgrad_compute: bool = False,
         symmetric_ar_type: Optional[str] = None,
         name: Optional[str] = None,
+        buffer_quantized_scaling_factors: bool = False,
+        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -1595,6 +1631,11 @@ def __init__(
                 if name in self.weight_names or name in self.bias_names:
                     param.skip_backward_post_hook = True
 
+        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
+        self.quantized_scaling_factor_buffering_decay = (
+            quantized_scaling_factor_buffering_decay
+        )
+
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""
         super().set_meta_tensor(fwd, recipe)
@@ -1765,6 +1806,14 @@ def forward(
                 self._fp8_workspaces.get(cache_name) if cache_name is not None else None
             )
 
+            scale_buffers = None
+            if self.buffer_quantized_scaling_factors:
+                scale_buffers = {
+                    name: value
+                    for name, value in self._buffers.items()
+                    if name.endswith("_te_ptq_calibrated")
+                }
+
             non_tensor_args = (
                 self.eps,
                 is_first_microbatch,
@@ -1805,6 +1854,8 @@ def forward(
                 self.symmetric_ar_type,
                 debug,
                 self.is_fsdp2,
+                self.quantized_scaling_factor_buffering_decay,
+                scale_buffers,
             )
             out, ln_out, new_weight_workspace = fwd_fn(
                 *autograd_ctx,
@@ -1817,6 +1868,14 @@ def forward(
                 non_tensor_args,
             )
 
+            if scale_buffers is not None:
+                for name, value in scale_buffers.items():
+                    if value is not None:
+                        if name in self._buffers:
+                            setattr(self, name, value)
+                        else:
+                            self.register_buffer(name, value, persistent=False)
+
             if new_weight_workspace is not None and cache_name is not None:
                 if isinstance(new_weight_workspace, torch.Tensor):
                     new_weight_workspace = new_weight_workspace.detach()
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 3ee0cda50c..8ca6fc1d0e 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -18,7 +18,10 @@
 
 from transformer_engine.common.recipe import Recipe
 from transformer_engine.pytorch.torch_version import torch_version
-from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom
+from transformer_engine.pytorch.tensor.utils import (
+    clear_columnwise_cache,
+    is_custom,
+)
 from .base import (
     fill_userbuffers_buffer_for_all_gather,
     _ub_communicators,
@@ -73,6 +76,8 @@
 from ..tensor.hybrid_tensor import HybridQuantizer
 from ..tensor.identity_tensor import IdentityQuantizer
 from ._common import (
+    _get_scale_buffer_info,
+    _update_scale_buffers,
     apply_normalization,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
@@ -244,6 +249,8 @@ def _forward(
             checkpoint,
             debug,
             is_fsdp2,
+            quantized_scaling_factor_buffering_decay,
+            scale_buffers,
             recompute_for_bwd,
         ) = non_tensor_args
         if fp8:
@@ -346,6 +353,10 @@ def _forward(
                 "checkpoint": checkpoint,
                 "debug": debug,
                 "is_fsdp2": is_fsdp2,
+                "quantized_scaling_factor_buffering_decay": (
+                    quantized_scaling_factor_buffering_decay
+                ),
+                "scale_buffers": scale_buffers,
                 "recompute_for_bwd": True,  # set this to true for recomputation phase
             }
         # Make sure input dimensions are compatible
@@ -564,6 +575,16 @@ def _forward(
             if fc1_weight_quantizer is not None:
                 fc1_weight_quantizer.calibrate(fc1_weight)
 
+        fc1_input_scale_buffer = None
+        fc1_weight_scale_buffer = None
+        if scale_buffers is not None:
+            fc1_input_scale_buffer = _get_scale_buffer_info(
+                "fc1_input", ln_out_total, fc1_input_quantizer
+            )
+            fc1_weight_scale_buffer = _get_scale_buffer_info(
+                "fc1_weight", fc1_weight_final, fc1_weight_quantizer
+            )
+
         # ------------------------------------------------------
         # FC1 GEMM
         # ------------------------------------------------------
@@ -678,6 +699,28 @@ def _forward(
                 if fc2_weight_quantizer is not None:
                     fc2_weight_quantizer.calibrate(fc2_weight)
 
+            if scale_buffers is not None:
+                scale_updates = {}
+                if fc1_input_scale_buffer is not None:
+                    scale_updates[fc1_input_scale_buffer[0]] = fc1_input_scale_buffer[1]
+                if fc1_weight_scale_buffer is not None:
+                    scale_updates[fc1_weight_scale_buffer[0]] = fc1_weight_scale_buffer[1]
+                fc2_input_scale_buffer = _get_scale_buffer_info(
+                    "fc2_input", act_out, fc2_input_quantizer
+                )
+                if fc2_input_scale_buffer is not None:
+                    scale_updates[fc2_input_scale_buffer[0]] = fc2_input_scale_buffer[1]
+                fc2_weight_scale_buffer = _get_scale_buffer_info(
+                    "fc2_weight", fc2_weight_final, fc2_weight_quantizer
+                )
+                if fc2_weight_scale_buffer is not None:
+                    scale_updates[fc2_weight_scale_buffer[0]] = fc2_weight_scale_buffer[1]
+                _update_scale_buffers(
+                    scale_buffers,
+                    scale_updates,
+                    quantized_scaling_factor_buffering_decay,
+                )
+
             # Configure Userbuffers reduce-scatter if needed
             ub_obj_fc2out = None
             reduce_scatter_out = None
@@ -1944,6 +1987,15 @@ class LayerNormMLP(TransformerEngineBaseModule):
                 whether to use selective activation checkpointing, where activations are not saved for bwd,
                 and instead are recomputed (skipping fc2, as it is not needed for backward). Trades compute
                 for memory. default is false, in which activations are saved in fwd. not supported for onnx forward
+    buffer_quantized_scaling_factors : bool, default = False
+                       If set to ``True``, maintain nonpersistent activation and weight quantization
+                       metadata buffers for both internal linear layers for inference checkpoint
+                       export. Per-tensor buffers store raw global amaxes, except FP8 current
+                       scaling buffers, which store inverse scales directly.
+    quantized_scaling_factor_buffering_decay : float, default = 0.0
+                       Decay applied to buffered activation scaling factors before incorporating
+                       each new observation. Defaults to 0.0, in which case only the most recent
+                       scaling factor is buffered.
     """
 
     def __init__(
@@ -1980,6 +2032,8 @@ def __init__(
         delay_wgrad_compute: bool = False,
         symmetric_ar_type: Optional[str] = None,
         checkpoint: bool = False,
+        buffer_quantized_scaling_factors: bool = False,
+        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -2001,6 +2055,10 @@ def __init__(
         self.zero_centered_gamma = zero_centered_gamma
         self.symmetric_ar_type = symmetric_ar_type
         self.checkpoint = checkpoint
+        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
+        self.quantized_scaling_factor_buffering_decay = (
+            quantized_scaling_factor_buffering_decay
+        )
 
         # GEMM-GELU fusion is currently only supported with split GEMM-AG overlap
         self.gemm_gelu_fusion = (
@@ -2381,6 +2439,14 @@ def forward(
                 self._fp8_workspaces.get(cache_name_fc2) if cache_name_fc2 is not None else None
             )
 
+            scale_buffers = None
+            if self.buffer_quantized_scaling_factors:
+                scale_buffers = {
+                    name: value
+                    for name, value in self._buffers.items()
+                    if name.endswith("_te_ptq_calibrated")
+                }
+
             non_tensor_args = (
                 self.eps,
                 is_first_microbatch,
@@ -2431,6 +2497,8 @@ def forward(
                 self.checkpoint,
                 debug,
                 self.is_fsdp2,
+                self.quantized_scaling_factor_buffering_decay,
+                scale_buffers,
             )
             out, ln_out, new_fc1_ws, new_fc2_ws = fwd_fn(
                 *autograd_ctx,
@@ -2446,6 +2514,14 @@ def forward(
                 non_tensor_args,
             )
 
+            if scale_buffers is not None:
+                for name, value in scale_buffers.items():
+                    if value is not None:
+                        if name in self._buffers:
+                            setattr(self, name, value)
+                        else:
+                            self.register_buffer(name, value, persistent=False)
+
             if new_fc1_ws is not None and cache_name_fc1 is not None:
                 if isinstance(new_fc1_ws, torch.Tensor):
                     new_fc1_ws = new_fc1_ws.detach()
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 94de69e975..24afbb9fb1 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -32,6 +32,8 @@
     _2X_ACC_WGRAD,
 )
 from ._common import (
+    _get_scale_buffer_info,
+    _update_scale_buffers,
     can_reconstruct_wgrad_input_from_original,
     noop_cat,
     set_quantizer_amax_reduction_group,
@@ -93,7 +95,10 @@
 )
 from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
-from ..tensor.utils import clear_columnwise_cache, is_custom
+from ..tensor.utils import (
+    clear_columnwise_cache,
+    is_custom,
+)
 from ..export import is_in_onnx_export_mode, assert_warmed_up
 from ..cpu_offload import (
     is_cpu_offload_enabled,
@@ -176,6 +181,10 @@ class LinearFwdArgs:
     fuse_wgrad_accumulation: bool
     wgrad_store: Optional[Any]
 
+    # Inference Scaling Factor Calibration Buffering
+    scale_buffers: Optional[Dict[str, Optional[torch.Tensor]]]
+    quantized_scaling_factor_buffering_decay: float
+
     # --- Misc ---
     cpu_offloading: bool
     is_grad_enabled: bool
@@ -370,6 +379,7 @@ def _linear_forward_impl(
     ctx_attrs)``. ``new_weight_workspace`` is the freshly produced FP8 weight
     workspace (returned alongside ``out`` so the caller can refresh its
     cache). The last two are ``None`` when gradients are disabled.
+    Scaling-factor checkpoint buffers are updated through ``args.scale_buffers``.
     """
 
     weight = args.weight
@@ -605,6 +615,25 @@ def _linear_forward_impl(
         if weight_quantizer is not None:
             weight_quantizer.calibrate(weight)
 
+    # Capture scaling metadata while it is still available.
+    if args.scale_buffers is not None:
+        scale_updates = {}
+        input_scale_buffer = _get_scale_buffer_info(
+            "input", inputmat_total, input_quantizer
+        )
+        if input_scale_buffer is not None:
+            scale_updates[input_scale_buffer[0]] = input_scale_buffer[1]
+        weight_scale_buffer = _get_scale_buffer_info(
+            "weight", weightmat, weight_quantizer
+        )
+        if weight_scale_buffer is not None:
+            scale_updates[weight_scale_buffer[0]] = weight_scale_buffer[1]
+        _update_scale_buffers(
+            args.scale_buffers,
+            scale_updates,
+            args.quantized_scaling_factor_buffering_decay,
+        )
+
     # Choose whether to use GEMM kernel with split accumulator
     use_split_accumulator = _2X_ACC_FPROP
     if fp8:
@@ -1990,6 +2019,18 @@ class Linear(TransformerEngineBaseModule):
                        and saving the original input tensor may reduce the memory usage.
                        Requires an input quantizer that can safely reproduce its result from the
                        original input. Cannot work with FP8 DelayedScaling recipe.
+    buffer_quantized_scaling_factors : bool, default = False
+                       If set to ``True``, maintain nonpersistent input and weight quantization
+                       metadata buffers for inference checkpoint export. Per-tensor buffers
+                       store raw global amaxes, except FP8 current scaling buffers, which store
+                       inverse scales directly.
+                       Each buffer is materialized only when its tensor uses a quantizer with a
+                       per-tensor scaling factor. Used to propagate scaling factors from training
+                       into inference.
+    quantized_scaling_factor_buffering_decay : float, default = 0.0
+                       Decay applied to buffered activation scaling factors before incorporating
+                       each new observation. Defaults to 0.0, in which case only the most recent
+                       scaling factor is buffered.
     """
 
     def __init__(
@@ -2019,6 +2060,8 @@ def __init__(
         symmetric_ar_type: Optional[str] = None,
         save_original_input: bool = False,
         name: Optional[str] = None,
+        buffer_quantized_scaling_factors: bool = False,
+        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -2246,6 +2289,9 @@ def __init__(
                 if name in self.weight_names or name in self.bias_names:
                     param.skip_backward_post_hook = True
 
+        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
+        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
+
     def get_quantizer_roles(
         self,
         *,
@@ -2475,7 +2521,13 @@ def forward(
                 bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None
             )
             wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None
-
+            scale_buffers = None
+            if self.buffer_quantized_scaling_factors:
+                scale_buffers = {
+                    name: value
+                    for name, value in self._buffers.items()
+                    if name.endswith("_te_ptq_calibrated")
+                }
             fwd_args = LinearFwdArgs(
                 # tensors
                 weight=weight_tensor,
@@ -2532,6 +2584,11 @@ def forward(
                 # weight-grad scheduling
                 fuse_wgrad_accumulation=self.fuse_wgrad_accumulation,
                 wgrad_store=wgrad_store,
+                # Inference Scaling Factor Calibration Buffering
+                scale_buffers=scale_buffers,
+                quantized_scaling_factor_buffering_decay=(
+                    self.quantized_scaling_factor_buffering_decay
+                ),
                 # misc
                 cpu_offloading=is_cpu_offload_enabled(),
                 is_grad_enabled=is_grad_enabled,
@@ -2555,6 +2612,16 @@ def forward(
                     weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled
                 )
 
+            if scale_buffers is not None:
+                # Assign the scaling factor calibration buffers to model.
+                # Requires CUDA graph warmup step.
+                for name, value in scale_buffers.items():
+                    if value is not None:
+                        if name in self._buffers:
+                            setattr(self, name, value)
+                        else:
+                            self.register_buffer(name, value, persistent=False)
+
             if new_weight_workspace is not None and cache_name is not None:
                 if isinstance(new_weight_workspace, torch.Tensor):
                     new_weight_workspace = new_weight_workspace.detach()
diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py
index cef45c0223..d7ba7b4982 100644
--- a/transformer_engine/pytorch/tensor/utils.py
+++ b/transformer_engine/pytorch/tensor/utils.py
@@ -27,6 +27,26 @@
 from ..constants import NVFP4_BLOCK_SCALING_SIZE, DType
 
 
+def get_quantization_recipe_name(quantizer: Optional[Quantizer]) -> str:
+    """Get a stable recipe name from a quantizer."""
+    quantizer = getattr(quantizer, "parent_quantizer", quantizer)
+    if quantizer is None:
+        return ""
+    if isinstance(quantizer, Float8Quantizer):
+        return "fp8_delayed_scaling"
+    if isinstance(quantizer, Float8CurrentScalingQuantizer):
+        return "fp8_current_scaling"
+    if isinstance(quantizer, MXFP8Quantizer):
+        return "mxfp8"
+    if isinstance(quantizer, Float8BlockQuantizer):
+        return "fp8_block_scaling"
+    if isinstance(quantizer, NVFP4Quantizer):
+        if quantizer.row_scaled_nvfp4:
+            return "nvfp4_rowwise"
+        return "nvfp4"
+    raise ValueError(f"Unsupported quantizer type: {type(quantizer)}")
+
+
 def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor):
     r"""Change a quantized tensor's data buffer while preserving values
 

From 9074933846544005f1a69127ae84f0a16f6eb90f Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Fri, 31 Jul 2026 11:56:03 -0700
Subject: [PATCH 02/16] Add tests.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 ...test_ptq_calibration_metadata_buffering.py | 104 ++++++++++++++++++
 transformer_engine/pytorch/module/_common.py  |  28 +++--
 2 files changed, 125 insertions(+), 7 deletions(-)
 create mode 100644 tests/pytorch/test_ptq_calibration_metadata_buffering.py

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
new file mode 100644
index 0000000000..94f1ba7664
--- /dev/null
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -0,0 +1,104 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from transformer_engine.pytorch.module import _common
+from transformer_engine.pytorch.module import grouped_linear
+
+
+@pytest.mark.parametrize(
+    ("recipe", "metadata_name", "expected_value"),
+    (
+        ("fp8_current_scaling", "scale_inv", 0.25),
+        ("fp8_delayed_scaling", "amax", 448.0),
+        ("nvfp4", "amax", 2688.0),
+        ("nvfp4_rowwise", "amax_rowwise", 1344.0),
+    ),
+)
+def test_scale_buffer_info_selects_recipe_metadata(
+    monkeypatch, recipe, metadata_name, expected_value
+):
+    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: recipe)
+    tensor = SimpleNamespace(
+        _scale_inv=torch.tensor([0.25], dtype=torch.float32),
+        _amax_rowwise=torch.tensor(
+            [2688.0 if recipe == "nvfp4" else 1344.0], dtype=torch.float32
+        ),
+    )
+    quantizer = SimpleNamespace(amax=torch.tensor([448.0], dtype=torch.float32))
+
+    buffer_name, value = _common._get_scale_buffer_info("input", tensor, quantizer)
+
+    assert buffer_name == f"input_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
+    torch.testing.assert_close(value, torch.tensor([expected_value]))
+
+
+@pytest.mark.parametrize("recipe", ("mxfp8", "fp8_block_scaling"))
+def test_scale_buffer_info_skips_non_global_scaling_recipes(monkeypatch, recipe):
+    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: recipe)
+    tensor = SimpleNamespace(_rowwise_scale_inv=torch.ones(2, 2))
+
+    assert _common._get_scale_buffer_info("input", tensor, object()) is None
+
+
+def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
+    monkeypatch.setattr(
+        _common, "get_quantization_recipe_name", lambda _: "fp8_current_scaling"
+    )
+    inputs = [
+        SimpleNamespace(_scale_inv=torch.tensor([0.25])),
+        SimpleNamespace(_scale_inv=torch.tensor([0.5])),
+    ]
+    weights = [
+        SimpleNamespace(_scale_inv=torch.tensor([0.75])),
+        SimpleNamespace(_scale_inv=torch.tensor([1.0])),
+    ]
+    scale_buffers = {}
+
+    grouped_linear._update_grouped_scale_buffers(
+        scale_buffers,
+        inputs,
+        weights,
+        object(),
+        object(),
+        activation_scale_decay=0.0,
+    )
+
+    assert set(scale_buffers) == {
+        "input_gemm0_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
+        "input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
+        "weight_gemm0_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
+        "weight_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
+    }
+    torch.testing.assert_close(
+        scale_buffers[
+            "input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+        ],
+        torch.tensor([0.5]),
+    )
+
+
+@pytest.mark.parametrize(
+    ("observed_scale", "expected_scale"),
+    (
+        # Decayed max is greater than the observed.
+        (1.0, 2.0),
+        # Decayed max is less than the observed.
+        (3.0, 3.0),
+    ),
+)
+def test_activation_scale_buffer_uses_decaying_maximum(observed_scale, expected_scale):
+    name = "input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+    scale_buffers = {name: torch.tensor([4.0])}
+
+    _common._update_scale_buffers(
+        scale_buffers,
+        {name: torch.tensor([observed_scale])},
+        activation_scale_decay=0.5,
+    )
+    torch.testing.assert_close(scale_buffers[name], torch.tensor([expected_scale]))
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index cac60a74d0..2293163650 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -27,16 +27,30 @@ def _get_scale_buffer_info(
     recipe = get_quantization_recipe_name(quantizer)
     if not recipe:
         return None
-    if recipe == "fp8_current_scaling":
+
+    if recipe == "fp8_delayed_scaling":
+        metadata_name = "amax"
+        metadata = getattr(quantizer, "amax", None)
+    elif recipe == "fp8_current_scaling":
         metadata_name = "scale_inv"
         metadata = getattr(tensor, "_scale_inv", None)
-    else:
+    elif recipe == "nvfp4":
+        metadata_name = "amax"
+        metadata = getattr(tensor, "_amax_rowwise", None)
+    elif recipe == "nvfp4_rowwise":
         metadata_name = "amax_rowwise"
-        metadata = getattr(
-            tensor,
-            "_amax_rowwise",
-            getattr(quantizer, "amax", None),
-        )
+        metadata = getattr(tensor, "_amax_rowwise", None)
+    elif recipe == "mxfp8":
+        # MXFP8 only exposes blockwise E8M0-encoded inverse scales, not a
+        # global FP32 scaling factor suitable for PTQ checkpoint export.
+        return None
+    elif recipe == "fp8_block_scaling":
+        # FP8 block scaling only exposes blockwise inverse scales, not a
+        # global FP32 scaling factor suitable for PTQ checkpoint export.
+        return None
+    else:
+        raise ValueError(f"Unsupported quantization recipe {recipe!r}")
+
     buffer_name = f"{tensor_name}_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
     return buffer_name, metadata
 

From 745658d9d5a62502608a0436565108995f7cf6fe Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 31 Jul 2026 19:08:25 +0000
Subject: [PATCH 03/16] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../test_ptq_calibration_metadata_buffering.py       | 12 +++---------
 transformer_engine/pytorch/module/grouped_linear.py  | 12 +++---------
 .../pytorch/module/layernorm_linear.py               | 12 +++---------
 transformer_engine/pytorch/module/layernorm_mlp.py   |  4 +---
 transformer_engine/pytorch/module/linear.py          |  8 ++------
 5 files changed, 12 insertions(+), 36 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index 94f1ba7664..0f0fa1d3f0 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -26,9 +26,7 @@ def test_scale_buffer_info_selects_recipe_metadata(
     monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: recipe)
     tensor = SimpleNamespace(
         _scale_inv=torch.tensor([0.25], dtype=torch.float32),
-        _amax_rowwise=torch.tensor(
-            [2688.0 if recipe == "nvfp4" else 1344.0], dtype=torch.float32
-        ),
+        _amax_rowwise=torch.tensor([2688.0 if recipe == "nvfp4" else 1344.0], dtype=torch.float32),
     )
     quantizer = SimpleNamespace(amax=torch.tensor([448.0], dtype=torch.float32))
 
@@ -47,9 +45,7 @@ def test_scale_buffer_info_skips_non_global_scaling_recipes(monkeypatch, recipe)
 
 
 def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
-    monkeypatch.setattr(
-        _common, "get_quantization_recipe_name", lambda _: "fp8_current_scaling"
-    )
+    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: "fp8_current_scaling")
     inputs = [
         SimpleNamespace(_scale_inv=torch.tensor([0.25])),
         SimpleNamespace(_scale_inv=torch.tensor([0.5])),
@@ -76,9 +72,7 @@ def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
         "weight_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
     }
     torch.testing.assert_close(
-        scale_buffers[
-            "input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
-        ],
+        scale_buffers["input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"],
         torch.tensor([0.5]),
     )
 
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 221ea0d36c..936e51e181 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -172,9 +172,7 @@ def _update_grouped_scale_buffers(
         if scale_buffer is not None:
             scale_updates[scale_buffer[0]] = scale_buffer[1]
     for index, tensor in enumerate(weight_tensors):
-        scale_buffer = _get_scale_buffer_info(
-            f"weight_gemm{index}", tensor, weight_quantizer
-        )
+        scale_buffer = _get_scale_buffer_info(f"weight_gemm{index}", tensor, weight_quantizer)
         if scale_buffer is not None:
             scale_updates[scale_buffer[0]] = scale_buffer[1]
     _update_scale_buffers(scale_buffers, scale_updates, activation_scale_decay)
@@ -855,9 +853,7 @@ def forward(
                 single_grouped_weight=single_grouped_weight,
                 single_grouped_bias=single_grouped_bias,
                 scale_buffers=scale_buffers,
-                quantized_scaling_factor_buffering_decay=(
-                    quantized_scaling_factor_buffering_decay
-                ),
+                quantized_scaling_factor_buffering_decay=(quantized_scaling_factor_buffering_decay),
                 weights=weights,
                 biases=biases,
                 out=out,
@@ -1742,9 +1738,7 @@ def __init__(
             )
         self.use_grouped_tensor = use_grouped_tensor
         self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
-        self.quantized_scaling_factor_buffering_decay = (
-            quantized_scaling_factor_buffering_decay
-        )
+        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
         single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags(
             single_grouped_weight, single_grouped_bias
         )
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index e78d612f7b..80e9a4bc45 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -392,14 +392,10 @@ def forward(
 
         if scale_buffers is not None:
             scale_updates = {}
-            input_scale_buffer = _get_scale_buffer_info(
-                "input", ln_out_total, input_quantizer
-            )
+            input_scale_buffer = _get_scale_buffer_info("input", ln_out_total, input_quantizer)
             if input_scale_buffer is not None:
                 scale_updates[input_scale_buffer[0]] = input_scale_buffer[1]
-            weight_scale_buffer = _get_scale_buffer_info(
-                "weight", weightmat, weight_quantizer
-            )
+            weight_scale_buffer = _get_scale_buffer_info("weight", weightmat, weight_quantizer)
             if weight_scale_buffer is not None:
                 scale_updates[weight_scale_buffer[0]] = weight_scale_buffer[1]
             _update_scale_buffers(
@@ -1632,9 +1628,7 @@ def __init__(
                     param.skip_backward_post_hook = True
 
         self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
-        self.quantized_scaling_factor_buffering_decay = (
-            quantized_scaling_factor_buffering_decay
-        )
+        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
 
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 8ca6fc1d0e..d9625d8759 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -2056,9 +2056,7 @@ def __init__(
         self.symmetric_ar_type = symmetric_ar_type
         self.checkpoint = checkpoint
         self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
-        self.quantized_scaling_factor_buffering_decay = (
-            quantized_scaling_factor_buffering_decay
-        )
+        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
 
         # GEMM-GELU fusion is currently only supported with split GEMM-AG overlap
         self.gemm_gelu_fusion = (
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 24afbb9fb1..d28d3dc5e0 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -618,14 +618,10 @@ def _linear_forward_impl(
     # Capture scaling metadata while it is still available.
     if args.scale_buffers is not None:
         scale_updates = {}
-        input_scale_buffer = _get_scale_buffer_info(
-            "input", inputmat_total, input_quantizer
-        )
+        input_scale_buffer = _get_scale_buffer_info("input", inputmat_total, input_quantizer)
         if input_scale_buffer is not None:
             scale_updates[input_scale_buffer[0]] = input_scale_buffer[1]
-        weight_scale_buffer = _get_scale_buffer_info(
-            "weight", weightmat, weight_quantizer
-        )
+        weight_scale_buffer = _get_scale_buffer_info("weight", weightmat, weight_quantizer)
         if weight_scale_buffer is not None:
             scale_updates[weight_scale_buffer[0]] = weight_scale_buffer[1]
         _update_scale_buffers(

From 25c68645444c0aec1a6c4ec7b18684f2933deb8e Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Fri, 31 Jul 2026 12:42:22 -0700
Subject: [PATCH 04/16] Fix minor bugs.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 ...test_ptq_calibration_metadata_buffering.py |  4 +--
 transformer_engine/pytorch/module/_common.py  | 10 +++----
 .../pytorch/module/grouped_linear.py          | 18 ++++++++++---
 .../pytorch/module/layernorm_linear.py        | 18 +++++++------
 .../pytorch/module/layernorm_mlp.py           | 26 ++++++++++++++-----
 transformer_engine/pytorch/module/linear.py   | 18 +++++++------
 transformer_engine/pytorch/tensor/utils.py    |  4 ++-
 7 files changed, 62 insertions(+), 36 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index 0f0fa1d3f0..da9b2c11a9 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
 
@@ -87,7 +87,7 @@ def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
     ),
 )
 def test_activation_scale_buffer_uses_decaying_maximum(observed_scale, expected_scale):
-    name = "input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+    name = "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
     scale_buffers = {name: torch.tensor([4.0])}
 
     _common._update_scale_buffers(
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 2293163650..f7abb57246 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -64,7 +64,7 @@ def _update_scale_buffers(
     for buffer_name, scale in scale_updates.items():
         if scale is None:
             continue
-        if buffer_name.startswith("input"):
+        if activation_scale_decay > 0.0:
             observed_scale = scale.detach().float()
             scale_buffer = scale_buffers.get(buffer_name)
             if scale_buffer is not None and scale_buffer.shape != observed_scale.shape:
@@ -72,10 +72,6 @@ def _update_scale_buffers(
                     "Quantized scaling-factor buffer shape changed from "
                     f"{tuple(scale_buffer.shape)} to {tuple(observed_scale.shape)}"
                 )
-            if activation_scale_decay == 0.0:
-                # If not using scale decay, just buffer the current scaling.
-                scale_buffers[buffer_name] = observed_scale
-                continue
             if scale_buffer is None:
                 # Initialize the rolling activation scaling factor.
                 # Requires CUDA graph warmup step.
@@ -90,8 +86,8 @@ def _update_scale_buffers(
                 out=scale_buffer,
             )
         else:
-            # Keep a reference to the current weight metadata without
-            # allocating or copying a separate buffer.
+            # Without scale history, keep a reference to the current metadata
+            # without allocating or copying a separate buffer.
             # Requires CUDA graph warmup step.
             scale_buffers[buffer_name] = scale.detach()
 
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 936e51e181..0974fdaca8 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -166,16 +166,26 @@ def _update_grouped_scale_buffers(
     activation_scale_decay: float,
 ) -> None:
     """Update GroupedLinear PTQ calibration buffers with per-GEMM metadata."""
-    scale_updates = {}
+    activation_scale_updates = {}
     for index, tensor in enumerate(input_tensors):
         scale_buffer = _get_scale_buffer_info(f"input_gemm{index}", tensor, input_quantizer)
         if scale_buffer is not None:
-            scale_updates[scale_buffer[0]] = scale_buffer[1]
+            activation_scale_updates[scale_buffer[0]] = scale_buffer[1]
+    weight_scale_updates = {}
     for index, tensor in enumerate(weight_tensors):
         scale_buffer = _get_scale_buffer_info(f"weight_gemm{index}", tensor, weight_quantizer)
         if scale_buffer is not None:
-            scale_updates[scale_buffer[0]] = scale_buffer[1]
-    _update_scale_buffers(scale_buffers, scale_updates, activation_scale_decay)
+            weight_scale_updates[scale_buffer[0]] = scale_buffer[1]
+    _update_scale_buffers(
+        scale_buffers,
+        activation_scale_updates,
+        activation_scale_decay,
+    )
+    _update_scale_buffers(
+        scale_buffers,
+        weight_scale_updates,
+        activation_scale_decay=0.0,
+    )
 
 
 class _GroupedLinear(torch.autograd.Function):
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 80e9a4bc45..2f3df6fe87 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -391,18 +391,20 @@ def forward(
                 weight_quantizer.calibrate(weight)
 
         if scale_buffers is not None:
-            scale_updates = {}
             input_scale_buffer = _get_scale_buffer_info("input", ln_out_total, input_quantizer)
             if input_scale_buffer is not None:
-                scale_updates[input_scale_buffer[0]] = input_scale_buffer[1]
+                _update_scale_buffers(
+                    scale_buffers,
+                    {input_scale_buffer[0]: input_scale_buffer[1]},
+                    quantized_scaling_factor_buffering_decay,
+                )
             weight_scale_buffer = _get_scale_buffer_info("weight", weightmat, weight_quantizer)
             if weight_scale_buffer is not None:
-                scale_updates[weight_scale_buffer[0]] = weight_scale_buffer[1]
-            _update_scale_buffers(
-                scale_buffers,
-                scale_updates,
-                quantized_scaling_factor_buffering_decay,
-            )
+                _update_scale_buffers(
+                    scale_buffers,
+                    {weight_scale_buffer[0]: weight_scale_buffer[1]},
+                    activation_scale_decay=0.0,
+                )
 
         # Choose whether to use GEMM kernel with split accumulator
         use_split_accumulator = _2X_ACC_FPROP
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index d9625d8759..c1964511fb 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -700,26 +700,40 @@ def _forward(
                     fc2_weight_quantizer.calibrate(fc2_weight)
 
             if scale_buffers is not None:
-                scale_updates = {}
+                activation_scale_updates = {}
+                weight_scale_updates = {}
                 if fc1_input_scale_buffer is not None:
-                    scale_updates[fc1_input_scale_buffer[0]] = fc1_input_scale_buffer[1]
+                    activation_scale_updates[fc1_input_scale_buffer[0]] = (
+                        fc1_input_scale_buffer[1]
+                    )
                 if fc1_weight_scale_buffer is not None:
-                    scale_updates[fc1_weight_scale_buffer[0]] = fc1_weight_scale_buffer[1]
+                    weight_scale_updates[fc1_weight_scale_buffer[0]] = (
+                        fc1_weight_scale_buffer[1]
+                    )
                 fc2_input_scale_buffer = _get_scale_buffer_info(
                     "fc2_input", act_out, fc2_input_quantizer
                 )
                 if fc2_input_scale_buffer is not None:
-                    scale_updates[fc2_input_scale_buffer[0]] = fc2_input_scale_buffer[1]
+                    activation_scale_updates[fc2_input_scale_buffer[0]] = (
+                        fc2_input_scale_buffer[1]
+                    )
                 fc2_weight_scale_buffer = _get_scale_buffer_info(
                     "fc2_weight", fc2_weight_final, fc2_weight_quantizer
                 )
                 if fc2_weight_scale_buffer is not None:
-                    scale_updates[fc2_weight_scale_buffer[0]] = fc2_weight_scale_buffer[1]
+                    weight_scale_updates[fc2_weight_scale_buffer[0]] = (
+                        fc2_weight_scale_buffer[1]
+                    )
                 _update_scale_buffers(
                     scale_buffers,
-                    scale_updates,
+                    activation_scale_updates,
                     quantized_scaling_factor_buffering_decay,
                 )
+                _update_scale_buffers(
+                    scale_buffers,
+                    weight_scale_updates,
+                    activation_scale_decay=0.0,
+                )
 
             # Configure Userbuffers reduce-scatter if needed
             ub_obj_fc2out = None
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index d28d3dc5e0..f56be159e5 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -617,18 +617,20 @@ def _linear_forward_impl(
 
     # Capture scaling metadata while it is still available.
     if args.scale_buffers is not None:
-        scale_updates = {}
         input_scale_buffer = _get_scale_buffer_info("input", inputmat_total, input_quantizer)
         if input_scale_buffer is not None:
-            scale_updates[input_scale_buffer[0]] = input_scale_buffer[1]
+            _update_scale_buffers(
+                args.scale_buffers,
+                {input_scale_buffer[0]: input_scale_buffer[1]},
+                args.quantized_scaling_factor_buffering_decay,
+            )
         weight_scale_buffer = _get_scale_buffer_info("weight", weightmat, weight_quantizer)
         if weight_scale_buffer is not None:
-            scale_updates[weight_scale_buffer[0]] = weight_scale_buffer[1]
-        _update_scale_buffers(
-            args.scale_buffers,
-            scale_updates,
-            args.quantized_scaling_factor_buffering_decay,
-        )
+            _update_scale_buffers(
+                args.scale_buffers,
+                {weight_scale_buffer[0]: weight_scale_buffer[1]},
+                activation_scale_decay=0.0,
+            )
 
     # Choose whether to use GEMM kernel with split accumulator
     use_split_accumulator = _2X_ACC_FPROP
diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py
index d7ba7b4982..f0b9898612 100644
--- a/transformer_engine/pytorch/tensor/utils.py
+++ b/transformer_engine/pytorch/tensor/utils.py
@@ -44,7 +44,9 @@ def get_quantization_recipe_name(quantizer: Optional[Quantizer]) -> str:
         if quantizer.row_scaled_nvfp4:
             return "nvfp4_rowwise"
         return "nvfp4"
-    raise ValueError(f"Unsupported quantizer type: {type(quantizer)}")
+    # Custom recipes may provide arbitrary Quantizer implementations without a
+    # stable recipe name or globally checkpointable scaling metadata.
+    return ""
 
 
 def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor):

From aa4292a34644644c3c5cd344560f36316264353c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 31 Jul 2026 19:46:32 +0000
Subject: [PATCH 05/16] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../pytorch/module/layernorm_mlp.py              | 16 ++++------------
 1 file changed, 4 insertions(+), 12 deletions(-)

diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index c1964511fb..0ca4e353dc 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -703,27 +703,19 @@ def _forward(
                 activation_scale_updates = {}
                 weight_scale_updates = {}
                 if fc1_input_scale_buffer is not None:
-                    activation_scale_updates[fc1_input_scale_buffer[0]] = (
-                        fc1_input_scale_buffer[1]
-                    )
+                    activation_scale_updates[fc1_input_scale_buffer[0]] = fc1_input_scale_buffer[1]
                 if fc1_weight_scale_buffer is not None:
-                    weight_scale_updates[fc1_weight_scale_buffer[0]] = (
-                        fc1_weight_scale_buffer[1]
-                    )
+                    weight_scale_updates[fc1_weight_scale_buffer[0]] = fc1_weight_scale_buffer[1]
                 fc2_input_scale_buffer = _get_scale_buffer_info(
                     "fc2_input", act_out, fc2_input_quantizer
                 )
                 if fc2_input_scale_buffer is not None:
-                    activation_scale_updates[fc2_input_scale_buffer[0]] = (
-                        fc2_input_scale_buffer[1]
-                    )
+                    activation_scale_updates[fc2_input_scale_buffer[0]] = fc2_input_scale_buffer[1]
                 fc2_weight_scale_buffer = _get_scale_buffer_info(
                     "fc2_weight", fc2_weight_final, fc2_weight_quantizer
                 )
                 if fc2_weight_scale_buffer is not None:
-                    weight_scale_updates[fc2_weight_scale_buffer[0]] = (
-                        fc2_weight_scale_buffer[1]
-                    )
+                    weight_scale_updates[fc2_weight_scale_buffer[0]] = fc2_weight_scale_buffer[1]
                 _update_scale_buffers(
                     scale_buffers,
                     activation_scale_updates,

From f53509bc2ced18d1ffef3c02fda94ae7bc80f34a Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Fri, 11 Sep 2026 08:42:51 -0700
Subject: [PATCH 06/16] Weed out NaN.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 ...test_ptq_calibration_metadata_buffering.py | 22 +++++++++++++++++++
 transformer_engine/pytorch/module/_common.py  |  3 ++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index da9b2c11a9..ef702f064b 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -96,3 +96,25 @@ def test_activation_scale_buffer_uses_decaying_maximum(observed_scale, expected_
         activation_scale_decay=0.5,
     )
     torch.testing.assert_close(scale_buffers[name], torch.tensor([expected_scale]))
+
+
+@pytest.mark.parametrize("activation_scale_decay", (0.0, 0.5))
+@pytest.mark.parametrize("initial_scale", (None, 4.0))
+def test_nan_activation_scale_does_not_update_buffer(
+    activation_scale_decay, initial_scale
+):
+    name = "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+    scale_buffers = {}
+    if initial_scale is not None:
+        scale_buffers[name] = torch.tensor([initial_scale])
+
+    _common._update_scale_buffers(
+        scale_buffers,
+        {name: torch.tensor([float("nan")])},
+        activation_scale_decay=activation_scale_decay,
+    )
+
+    if initial_scale is None:
+        assert name not in scale_buffers
+    else:
+        torch.testing.assert_close(scale_buffers[name], torch.tensor([initial_scale]))
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index f7abb57246..ede7e7a608 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -62,7 +62,8 @@ def _update_scale_buffers(
 ) -> None:
     """Merge observed scaling factors into checkpoint buffers."""
     for buffer_name, scale in scale_updates.items():
-        if scale is None:
+        if scale is None or torch.isnan(scale).any():
+            # Un-initialized scale. Ignore it.
             continue
         if activation_scale_decay > 0.0:
             observed_scale = scale.detach().float()

From c87e8db22c43249beacc7c5719ad8fde2fc08bf1 Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Fri, 11 Sep 2026 08:58:05 -0700
Subject: [PATCH 07/16] Rebase fixes.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 ...test_ptq_calibration_metadata_buffering.py | 31 +++++++++++++++++--
 tests/pytorch/test_torch_compile.py           | 11 +++++++
 .../pytorch/module/grouped_linear.py          | 28 +++++++++++------
 transformer_engine/pytorch/module/linear.py   |  2 ++
 4 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index ef702f064b..0b572ee411 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -60,8 +60,8 @@ def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
         scale_buffers,
         inputs,
         weights,
-        object(),
-        object(),
+        [object(), object()],
+        [object(), object()],
         activation_scale_decay=0.0,
     )
 
@@ -77,6 +77,33 @@ def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
     )
 
 
+def test_grouped_scale_buffers_use_per_gemm_delayed_scaling_amax(monkeypatch):
+    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: "fp8_delayed_scaling")
+    quantizers = [
+        SimpleNamespace(amax=torch.tensor([1.0])),
+        SimpleNamespace(amax=torch.tensor([2.0])),
+    ]
+    scale_buffers = {}
+
+    grouped_linear._update_grouped_scale_buffers(
+        scale_buffers,
+        [object(), object()],
+        [object(), object()],
+        quantizers,
+        quantizers,
+        activation_scale_decay=0.0,
+    )
+
+    torch.testing.assert_close(
+        scale_buffers["input_gemm0_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated"],
+        torch.tensor([1.0]),
+    )
+    torch.testing.assert_close(
+        scale_buffers["input_gemm1_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated"],
+        torch.tensor([2.0]),
+    )
+
+
 @pytest.mark.parametrize(
     ("observed_scale", "expected_scale"),
     (
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index e5d7169da1..91a31d14eb 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2022,6 +2022,7 @@ def fn(inp):
     "fuse_wgrad_accumulation",
     "delayed_wgrad",
     "quantized_input",
+    "scale_buffering",
 ]
 
 
@@ -2032,6 +2033,8 @@ def _fallback_case(case, dtype, device):
         model_kwargs["fuse_wgrad_accumulation"] = True
     elif case == "delayed_wgrad":
         model_kwargs["delay_wgrad_compute"] = True
+    elif case == "scale_buffering":
+        model_kwargs["buffer_quantized_scaling_factors"] = True
     model = te.Linear(64, 32, params_dtype=dtype, device=device, **model_kwargs)
 
     if case == "fp8_output_differentiable":
@@ -2055,6 +2058,14 @@ def fn(inp):
                 return model(inp)
 
         return model, fn, "no_grad", None, "a quantized input tensor"
+    if case == "scale_buffering":
+        fp8_recipe = recipe.Float8CurrentScaling()
+
+        def fn(inp):
+            with te.autocast(recipe=fp8_recipe):
+                return model(inp)
+
+        return model, fn, "bwd", None, "quantized scaling-factor buffering"
     raise ValueError(case)
 
 
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 0974fdaca8..2c178e965c 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -161,19 +161,23 @@ def _update_grouped_scale_buffers(
     scale_buffers: Dict[str, Optional[torch.Tensor]],
     input_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
     weight_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
-    input_quantizer: Optional[Quantizer],
-    weight_quantizer: Optional[Quantizer],
+    input_quantizers: List[Optional[Quantizer]],
+    weight_quantizers: List[Optional[Quantizer]],
     activation_scale_decay: float,
 ) -> None:
     """Update GroupedLinear PTQ calibration buffers with per-GEMM metadata."""
     activation_scale_updates = {}
     for index, tensor in enumerate(input_tensors):
-        scale_buffer = _get_scale_buffer_info(f"input_gemm{index}", tensor, input_quantizer)
+        scale_buffer = _get_scale_buffer_info(
+            f"input_gemm{index}", tensor, input_quantizers[index]
+        )
         if scale_buffer is not None:
             activation_scale_updates[scale_buffer[0]] = scale_buffer[1]
     weight_scale_updates = {}
     for index, tensor in enumerate(weight_tensors):
-        scale_buffer = _get_scale_buffer_info(f"weight_gemm{index}", tensor, weight_quantizer)
+        scale_buffer = _get_scale_buffer_info(
+            f"weight_gemm{index}", tensor, weight_quantizers[index]
+        )
         if scale_buffer is not None:
             weight_scale_updates[scale_buffer[0]] = scale_buffer[1]
     _update_scale_buffers(
@@ -578,12 +582,18 @@ def _forward_grouped_tensor(
             grouped_inputs = grouped_x.quantized_tensors
             if grouped_inputs is None:
                 grouped_inputs = grouped_x.split_into_quantized_tensors()
+            if isinstance(weights_for_gemm, GroupedTensorStorage):
+                grouped_weights = weights_for_gemm.quantized_tensors
+                if grouped_weights is None:
+                    grouped_weights = weights_for_gemm.split_into_quantized_tensors()
+            else:
+                grouped_weights = weights_for_gemm
             _update_grouped_scale_buffers(
                 scale_buffers,
                 grouped_inputs,
-                weights_for_gemm,
-                input_quantizers[0],
-                weight_quantizers[0],
+                grouped_weights,
+                input_quantizers,
+                weight_quantizers,
                 quantized_scaling_factor_buffering_decay,
             )
 
@@ -915,8 +925,8 @@ def forward(
                 scale_buffers,
                 inputmats,
                 weights_fp8,
-                input_quantizers[0],
-                weight_quantizers[0],
+                input_quantizers,
+                weight_quantizers,
                 quantized_scaling_factor_buffering_decay,
             )
 
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index f56be159e5..fa9d159111 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -2708,6 +2708,8 @@ def _compile_eager_fallback_reason(
         prepare_forward. Quantizer checks stay in compile_unsupported_reason."""
         if debug:
             return "debug instrumentation (nvidia-dlfw-inspect)"
+        if self.buffer_quantized_scaling_factors:
+            return "quantized scaling-factor buffering"
         weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
         if is_distributed_weight(weight_tensor):
             return "a DistributedWeight (custom weight parallelism, e.g. GTP)"

From a4915efbd3daef6d68e0b5d4968bebcbdef0f7b1 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 11 Sep 2026 16:00:10 +0000
Subject: [PATCH 08/16] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/pytorch/test_ptq_calibration_metadata_buffering.py | 4 +---
 transformer_engine/pytorch/module/grouped_linear.py      | 4 +---
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index 0b572ee411..008e84108d 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -127,9 +127,7 @@ def test_activation_scale_buffer_uses_decaying_maximum(observed_scale, expected_
 
 @pytest.mark.parametrize("activation_scale_decay", (0.0, 0.5))
 @pytest.mark.parametrize("initial_scale", (None, 4.0))
-def test_nan_activation_scale_does_not_update_buffer(
-    activation_scale_decay, initial_scale
-):
+def test_nan_activation_scale_does_not_update_buffer(activation_scale_decay, initial_scale):
     name = "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
     scale_buffers = {}
     if initial_scale is not None:
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 2c178e965c..a6867b2102 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -168,9 +168,7 @@ def _update_grouped_scale_buffers(
     """Update GroupedLinear PTQ calibration buffers with per-GEMM metadata."""
     activation_scale_updates = {}
     for index, tensor in enumerate(input_tensors):
-        scale_buffer = _get_scale_buffer_info(
-            f"input_gemm{index}", tensor, input_quantizers[index]
-        )
+        scale_buffer = _get_scale_buffer_info(f"input_gemm{index}", tensor, input_quantizers[index])
         if scale_buffer is not None:
             activation_scale_updates[scale_buffer[0]] = scale_buffer[1]
     weight_scale_updates = {}

From f2943b841e4e7ebb70b5526147c40b9bf46cecf3 Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Fri, 11 Sep 2026 12:14:21 -0700
Subject: [PATCH 09/16] Merge calibration tooling into the calibrate() API.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 examples/pytorch/mnist/main.py                |   5 +-
 qa/L0_pytorch_unittest/test.sh                |   1 +
 ...test_ptq_calibration_metadata_buffering.py | 574 ++++++++++++++++--
 tests/pytorch/test_torch_compile.py           |   7 +-
 .../debug/pytorch/debug_quantization.py       |   3 +-
 transformer_engine/pytorch/__init__.py        |   2 +
 .../dot_product_attention.py                  |   5 +-
 transformer_engine/pytorch/graph.py           |  11 +-
 transformer_engine/pytorch/module/_common.py  |  90 +--
 .../pytorch/module/grouped_linear.py          |  88 +--
 .../pytorch/module/layernorm_linear.py        |  66 +-
 .../pytorch/module/layernorm_mlp.py           | 104 ++--
 transformer_engine/pytorch/module/linear.py   |  68 +--
 transformer_engine/pytorch/quantization.py    | 112 +++-
 .../pytorch/quantized_tensor.py               |  58 +-
 .../pytorch/tensor/float8_blockwise_tensor.py |  10 +-
 .../pytorch/tensor/float8_tensor.py           |  58 +-
 .../pytorch/tensor/identity_tensor.py         |   4 +-
 .../pytorch/tensor/mxfp8_tensor.py            |  12 +-
 .../pytorch/tensor/nvfp4_tensor.py            |  32 +-
 transformer_engine/pytorch/tensor/utils.py    |  16 +-
 21 files changed, 976 insertions(+), 350 deletions(-)

diff --git a/examples/pytorch/mnist/main.py b/examples/pytorch/mnist/main.py
index 3754e3643b..24081fb259 100644
--- a/examples/pytorch/mnist/main.py
+++ b/examples/pytorch/mnist/main.py
@@ -76,7 +76,10 @@ def calibrate(model, device, test_loader, fp8):
     with torch.no_grad():
         for data, target in test_loader:
             data, target = data.to(device), target.to(device)
-            with te.autocast(enabled=fp8, calibrating=True):
+            with te.autocast(
+                enabled=fp8,
+                calibration_config=te.QuantizationCalibrationConfig(),
+            ):
                 output = model(data)
 
 
diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh
index a78a99d7f9..a7162cd13d 100644
--- a/qa/L0_pytorch_unittest/test.sh
+++ b/qa/L0_pytorch_unittest/test.sh
@@ -41,6 +41,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_nvfp4.xml $TE_PA
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8.xml $TE_PATH/tests/pytorch/mxfp8 || test_fail "test_mxfp8"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_weight_swizzle_in_layers.xml $TE_PATH/tests/pytorch/test_weight_swizzle_in_layers.py || test_fail "test_weight_swizzle_in_layers.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_quantized_tensor.xml $TE_PATH/tests/pytorch/test_quantized_tensor.py || test_fail "test_quantized_tensor.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ptq_calibration_metadata_buffering.xml $TE_PATH/tests/pytorch/test_ptq_calibration_metadata_buffering.py || test_fail "test_ptq_calibration_metadata_buffering.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_torch_compile.xml $TE_PATH/tests/pytorch/test_torch_compile.py || test_fail "test_torch_compile.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py"
diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index 008e84108d..a83090ab77 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -2,13 +2,222 @@
 #
 # See LICENSE for license information.
 
+import dataclasses
+import inspect
 from types import SimpleNamespace
 
 import pytest
 import torch
 
+from transformer_engine.common.recipe import Float8CurrentScaling
+from transformer_engine.pytorch import is_fp8_available, is_nvfp4_available
+from transformer_engine.pytorch.constants import DType
+from transformer_engine.pytorch.graph import make_graphed_callables
+from transformer_engine.pytorch.module import GroupedLinear, LayerNormLinear, LayerNormMLP, Linear
 from transformer_engine.pytorch.module import _common
 from transformer_engine.pytorch.module import grouped_linear
+from transformer_engine.pytorch.quantization import (
+    FP8GlobalStateManager,
+    FP8GlobalState,
+    TEAutocastState,
+    QuantizationCalibrationConfig,
+    autocast,
+    fp8_autocast,
+)
+from transformer_engine.pytorch.quantized_tensor import Quantizer
+from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer
+from transformer_engine.pytorch.tensor.float8_tensor import (
+    Float8CurrentScalingQuantizer,
+    Float8Quantizer,
+)
+from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
+from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
+from transformer_engine.pytorch.tensor.utils import get_quantization_recipe_name
+
+nvfp4_available, reason_for_no_nvfp4 = is_nvfp4_available(return_reason=True)
+fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True)
+
+
+@pytest.fixture(autouse=True)
+def reset_quantization_state():
+    yield
+    FP8GlobalStateManager.reset()
+
+
+def test_calibration_api_additions_preserve_existing_parameter_order():
+    state_fields = [field.name for field in dataclasses.fields(FP8GlobalState)]
+    assert state_fields[-1] == "calibration_config"
+    assert state_fields[:2] == ["fp8_enabled", "fp8_calibration"]
+
+    assert list(inspect.signature(FP8GlobalStateManager.autocast_enter).parameters) == [
+        "enabled",
+        "calibrating",
+        "fp8_recipe",
+        "fp8_group",
+        "_graph",
+        "calibration_config",
+    ]
+    assert list(inspect.signature(autocast).parameters) == [
+        "enabled",
+        "calibrating",
+        "recipe",
+        "amax_reduction_group",
+        "_graph",
+        "calibration_config",
+    ]
+    assert list(inspect.signature(fp8_autocast).parameters) == [
+        "enabled",
+        "calibrating",
+        "fp8_recipe",
+        "fp8_group",
+        "_graph",
+        "calibration_config",
+    ]
+    assert list(inspect.signature(make_graphed_callables).parameters)[-2:] == [
+        "capture_time_hooks",
+        "calibration_config",
+    ]
+    assert not inspect.signature(FP8GlobalStateManager.get_autocast_state).parameters
+    assert list(inspect.signature(FP8GlobalStateManager.set_autocast_state).parameters) == ["state"]
+    autocast_state = FP8GlobalStateManager.get_autocast_state()
+    assert isinstance(autocast_state, TEAutocastState)
+    assert [field.name for field in dataclasses.fields(autocast_state)] == [
+        "fp8_enabled",
+        "fp8_calibration",
+        "calibration_config",
+        "fp8_recipe",
+        "fp8_distributed_group",
+        "is_first_fp8_module",
+        "fp8_graph_capturing",
+    ]
+
+
+def test_calibrating_argument_enables_default_calibration_config():
+    assert FP8GlobalStateManager.get_calibration_config() is None
+    with autocast(enabled=False, calibrating=True):
+        assert FP8GlobalStateManager.quantization_state.fp8_calibration
+        assert FP8GlobalStateManager.get_calibration_config() == QuantizationCalibrationConfig()
+    assert FP8GlobalStateManager.get_calibration_config() is None
+
+
+def test_explicit_calibration_config_is_active_in_autocast():
+    config = QuantizationCalibrationConfig(activation_scale_decay=0.5)
+    with autocast(enabled=False, calibration_config=config):
+        assert FP8GlobalStateManager.quantization_state.fp8_calibration
+        assert FP8GlobalStateManager.get_calibration_config() is config
+        assert FP8GlobalStateManager.get_autocast_state().calibration_config is config
+
+
+def test_nested_autocast_restores_custom_calibration_config():
+    config = QuantizationCalibrationConfig(activation_scale_decay=0.5)
+    with autocast(enabled=False, calibration_config=config):
+        with autocast(enabled=False):
+            assert FP8GlobalStateManager.get_calibration_config() is None
+        assert FP8GlobalStateManager.get_calibration_config() is config
+
+
+def test_global_calibration_boolean_enables_default_config():
+    qstate = FP8GlobalStateManager.quantization_state
+    qstate.fp8_calibration = True
+    try:
+        assert FP8GlobalStateManager.get_calibration_config() == QuantizationCalibrationConfig()
+    finally:
+        qstate.fp8_calibration = False
+
+
+def test_autocast_enter_preserves_calibrating_boolean_api():
+    saved_state = FP8GlobalStateManager.get_autocast_state()
+    try:
+        FP8GlobalStateManager.autocast_enter(False, True)
+        assert FP8GlobalStateManager.is_fp8_calibration()
+        assert FP8GlobalStateManager.get_calibration_config() == QuantizationCalibrationConfig()
+    finally:
+        FP8GlobalStateManager.autocast_exit(False, False)
+        FP8GlobalStateManager.set_autocast_state(saved_state)
+
+
+def test_calibrating_argument_accepts_explicit_calibration_config():
+    config = QuantizationCalibrationConfig(activation_scale_decay=0.5)
+    with autocast(
+        enabled=False,
+        calibrating=True,
+        calibration_config=config,
+    ):
+        assert FP8GlobalStateManager.get_calibration_config() is config
+
+
+def test_calibration_config_rejects_negative_decay():
+    with pytest.raises(ValueError, match="activation_scale_decay must be non-negative"):
+        QuantizationCalibrationConfig(activation_scale_decay=-0.1)
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize(
+    ("module_name", "expected_buffer_count"),
+    (
+        ("linear", 2),
+        ("layernorm_linear", 2),
+        ("layernorm_mlp", 4),
+        ("grouped_linear", 4),
+    ),
+)
+@pytest.mark.parametrize("enabled", (False, True))
+def test_calibration_config_registers_module_scaling_factor_buffers(
+    module_name, expected_buffer_count, enabled
+):
+    module_kwargs = {
+        "params_dtype": torch.bfloat16,
+        "device": "cuda",
+        "bias": False,
+    }
+    if module_name == "linear":
+        module = Linear(32, 32, **module_kwargs)
+    elif module_name == "layernorm_linear":
+        module = LayerNormLinear(32, 32, **module_kwargs)
+    elif module_name == "layernorm_mlp":
+        module = LayerNormMLP(32, 32, **module_kwargs)
+    else:
+        module = GroupedLinear(2, 32, 32, use_grouped_tensor=False, **module_kwargs)
+
+    inp = torch.randn((16, 32), dtype=torch.bfloat16, device="cuda")
+    with autocast(
+        enabled=enabled,
+        recipe=Float8CurrentScaling(),
+        calibration_config=QuantizationCalibrationConfig(),
+    ):
+        if module_name == "grouped_linear":
+            module(inp, [8, 8])
+        else:
+            module(inp)
+
+    buffers = {
+        name: value
+        for name, value in module.named_buffers()
+        if name.endswith("_te_ptq_calibrated")
+    }
+    assert len(buffers) == expected_buffer_count
+    assert all(torch.isfinite(value).all() for value in buffers.values())
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_calibrating_boolean_registers_scaling_factor_buffers():
+    module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
+    inp = torch.randn((16, 32), dtype=torch.bfloat16, device="cuda")
+
+    with autocast(
+        enabled=False,
+        calibrating=True,
+        recipe=Float8CurrentScaling(),
+    ):
+        module(inp)
+
+    assert len(
+        [
+            name
+            for name, _ in module.named_buffers()
+            if name.endswith("_te_ptq_calibrated")
+        ]
+    ) == 2
 
 
 @pytest.mark.parametrize(
@@ -21,31 +230,73 @@
     ),
 )
 def test_scale_buffer_info_selects_recipe_metadata(
-    monkeypatch, recipe, metadata_name, expected_value
+    recipe, metadata_name, expected_value
 ):
-    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: recipe)
     tensor = SimpleNamespace(
         _scale_inv=torch.tensor([0.25], dtype=torch.float32),
         _amax_rowwise=torch.tensor([2688.0 if recipe == "nvfp4" else 1344.0], dtype=torch.float32),
     )
-    quantizer = SimpleNamespace(amax=torch.tensor([448.0], dtype=torch.float32))
+    if recipe == "fp8_delayed_scaling":
+        quantizer = object.__new__(Float8Quantizer)
+        quantizer.amax = torch.tensor([0.0], dtype=torch.float32)
+        tensor = torch.tensor([expected_value])
+    elif recipe == "fp8_current_scaling":
+        quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    else:
+        quantizer = object.__new__(NVFP4Quantizer)
+        quantizer.row_scaled_nvfp4 = recipe == "nvfp4_rowwise"
 
-    buffer_name, value = _common._get_scale_buffer_info("input", tensor, quantizer)
+    assert get_quantization_recipe_name(quantizer) == recipe
+    quantizer.calibrate(tensor)
+    buffers = _common._get_scale_buffer_info("input", quantizer)
+    buffer_name = f"input_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
+    value = buffers[buffer_name]
 
-    assert buffer_name == f"input_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
     torch.testing.assert_close(value, torch.tensor([expected_value]))
+    assert value is quantizer._calibration_state[metadata_name]
 
 
 @pytest.mark.parametrize("recipe", ("mxfp8", "fp8_block_scaling"))
-def test_scale_buffer_info_skips_non_global_scaling_recipes(monkeypatch, recipe):
-    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: recipe)
+def test_scale_buffer_info_skips_non_global_scaling_recipes(recipe):
     tensor = SimpleNamespace(_rowwise_scale_inv=torch.ones(2, 2))
+    quantizer_cls = MXFP8Quantizer if recipe == "mxfp8" else Float8BlockQuantizer
+    quantizer = object.__new__(quantizer_cls)
+
+    assert get_quantization_recipe_name(quantizer) == recipe
+    quantizer.calibrate(tensor)
+    assert not _common._get_scale_buffer_info("input", quantizer)
+
+
+def test_custom_quantizer_defaults_to_no_calibration_metadata():
+    quantizer = Quantizer(rowwise=True, columnwise=False)
+
+    assert get_quantization_recipe_name(quantizer) == ""
+    assert not _common._get_scale_buffer_info("input", quantizer)
+
+
+def test_resolve_calibration_quantizer_prefers_tensor_owner_and_unwraps_parent():
+    parent_quantizer = object()
+    tensor_quantizer = SimpleNamespace(parent_quantizer=parent_quantizer)
+    tensor = SimpleNamespace(_quantizer=tensor_quantizer)
+
+    assert (
+        _common._resolve_calibration_quantizer(tensor, object())
+        is parent_quantizer
+    )
+
+
+def test_quantizer_calibration_state_is_keyed_by_quantized_metadata():
+    quantizer = Quantizer(rowwise=True, columnwise=False)
+
+    quantizer._update_calibration_value("amax", torch.tensor([2.0]), decay=0.0)
+    quantizer._update_calibration_value("scale_inv", torch.tensor([0.5]), decay=0.0)
 
-    assert _common._get_scale_buffer_info("input", tensor, object()) is None
+    assert set(quantizer._calibration_state) == {"amax", "scale_inv"}
+    torch.testing.assert_close(quantizer._calibration_state["amax"], torch.tensor([2.0]))
+    torch.testing.assert_close(quantizer._calibration_state["scale_inv"], torch.tensor([0.5]))
 
 
-def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
-    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: "fp8_current_scaling")
+def test_grouped_scale_buffers_are_per_gemm():
     inputs = [
         SimpleNamespace(_scale_inv=torch.tensor([0.25])),
         SimpleNamespace(_scale_inv=torch.tensor([0.5])),
@@ -54,15 +305,29 @@ def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
         SimpleNamespace(_scale_inv=torch.tensor([0.75])),
         SimpleNamespace(_scale_inv=torch.tensor([1.0])),
     ]
+    input_quantizers = [
+        object.__new__(Float8CurrentScalingQuantizer),
+        object.__new__(Float8CurrentScalingQuantizer),
+    ]
+    weight_quantizers = [
+        object.__new__(Float8CurrentScalingQuantizer),
+        object.__new__(Float8CurrentScalingQuantizer),
+    ]
     scale_buffers = {}
 
+    grouped_linear._calibrate_grouped_tensors(
+        inputs,
+        weights,
+        input_quantizers,
+        weight_quantizers,
+        activation_scale_decay=0.0,
+    )
     grouped_linear._update_grouped_scale_buffers(
         scale_buffers,
         inputs,
         weights,
-        [object(), object()],
-        [object(), object()],
-        activation_scale_decay=0.0,
+        input_quantizers,
+        weight_quantizers,
     )
 
     assert set(scale_buffers) == {
@@ -75,23 +340,57 @@ def test_grouped_scale_buffers_are_per_gemm(monkeypatch):
         scale_buffers["input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"],
         torch.tensor([0.5]),
     )
+    assert (
+        scale_buffers[
+            "input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+        ]
+        is input_quantizers[1]._calibration_state["scale_inv"]
+    )
 
 
-def test_grouped_scale_buffers_use_per_gemm_delayed_scaling_amax(monkeypatch):
-    monkeypatch.setattr(_common, "get_quantization_recipe_name", lambda _: "fp8_delayed_scaling")
-    quantizers = [
-        SimpleNamespace(amax=torch.tensor([1.0])),
-        SimpleNamespace(amax=torch.tensor([2.0])),
-    ]
+def test_grouped_calibration_applies_decay_only_to_activations():
+    input_quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    input_quantizer._calibration_state = {"scale_inv": torch.tensor([4.0])}
+    weight_quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    weight_quantizer._calibration_state = {"scale_inv": torch.tensor([4.0])}
+
+    grouped_linear._calibrate_grouped_tensors(
+        [SimpleNamespace(_scale_inv=torch.tensor([1.0]))],
+        [SimpleNamespace(_scale_inv=torch.tensor([1.0]))],
+        [input_quantizer],
+        [weight_quantizer],
+        activation_scale_decay=0.5,
+    )
+
+    torch.testing.assert_close(
+        input_quantizer._calibration_state["scale_inv"], torch.tensor([2.0])
+    )
+    torch.testing.assert_close(
+        weight_quantizer._calibration_state["scale_inv"], torch.tensor([1.0])
+    )
+
+
+def test_grouped_scale_buffers_use_per_gemm_delayed_scaling_amax():
+    quantizers = []
+    for amax in (1.0, 2.0):
+        quantizer = object.__new__(Float8Quantizer)
+        quantizer.amax = torch.tensor([amax])
+        quantizers.append(quantizer)
     scale_buffers = {}
 
+    grouped_linear._calibrate_grouped_tensors(
+        [torch.tensor([1.0]), torch.tensor([2.0])],
+        [torch.tensor([1.0]), torch.tensor([2.0])],
+        quantizers,
+        quantizers,
+        activation_scale_decay=0.0,
+    )
     grouped_linear._update_grouped_scale_buffers(
         scale_buffers,
-        [object(), object()],
-        [object(), object()],
+        [torch.tensor([1.0]), torch.tensor([2.0])],
+        [torch.tensor([1.0]), torch.tensor([2.0])],
         quantizers,
         quantizers,
-        activation_scale_decay=0.0,
     )
 
     torch.testing.assert_close(
@@ -115,31 +414,230 @@ def test_grouped_scale_buffers_use_per_gemm_delayed_scaling_amax(monkeypatch):
 )
 def test_activation_scale_buffer_uses_decaying_maximum(observed_scale, expected_scale):
     name = "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
-    scale_buffers = {name: torch.tensor([4.0])}
-
-    _common._update_scale_buffers(
-        scale_buffers,
-        {name: torch.tensor([observed_scale])},
-        activation_scale_decay=0.5,
+    quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    initial_buffer = torch.tensor([4.0])
+    quantizer._calibration_state = {"scale_inv": initial_buffer}
+    quantizer.calibrate(
+        SimpleNamespace(_scale_inv=torch.tensor([observed_scale])),
+        decay=0.5,
     )
-    torch.testing.assert_close(scale_buffers[name], torch.tensor([expected_scale]))
+    buffers = _common._get_scale_buffer_info("fc1_input", quantizer)
+    value = buffers[name]
+
+    torch.testing.assert_close(value, torch.tensor([expected_scale]))
+    assert value is initial_buffer
+    assert value is quantizer._calibration_state["scale_inv"]
+
+
+def test_zero_decay_keeps_observed_metadata_reference():
+    observed_scale = torch.tensor([2.0])
+    quantizer = object.__new__(Float8CurrentScalingQuantizer)
+
+    quantizer.calibrate(SimpleNamespace(_scale_inv=observed_scale), decay=0.0)
+
+    assert quantizer._calibration_state["scale_inv"] is observed_scale
+
+
+def test_decaying_calibration_rejects_metadata_shape_change():
+    quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    quantizer._calibration_state = {"scale_inv": torch.ones(1)}
+
+    with pytest.raises(RuntimeError, match="calibration value shape changed"):
+        quantizer.calibrate(SimpleNamespace(_scale_inv=torch.ones(2)), decay=0.5)
 
 
 @pytest.mark.parametrize("activation_scale_decay", (0.0, 0.5))
 @pytest.mark.parametrize("initial_scale", (None, 4.0))
 def test_nan_activation_scale_does_not_update_buffer(activation_scale_decay, initial_scale):
-    name = "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
-    scale_buffers = {}
+    quantizer = object.__new__(Float8CurrentScalingQuantizer)
     if initial_scale is not None:
-        scale_buffers[name] = torch.tensor([initial_scale])
+        quantizer._calibration_state = {"scale_inv": torch.tensor([initial_scale])}
 
-    _common._update_scale_buffers(
-        scale_buffers,
-        {name: torch.tensor([float("nan")])},
-        activation_scale_decay=activation_scale_decay,
+    quantizer.calibrate(
+        SimpleNamespace(_scale_inv=torch.tensor([float("nan")])),
+        decay=activation_scale_decay,
     )
+    result = _common._get_scale_buffer_info("fc1_input", quantizer)
 
     if initial_scale is None:
-        assert name not in scale_buffers
+        assert not result
+        assert not quantizer._calibration_state
     else:
-        torch.testing.assert_close(scale_buffers[name], torch.tensor([initial_scale]))
+        value = result[
+            "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+        ]
+        torch.testing.assert_close(value, torch.tensor([initial_scale]))
+
+
+def test_current_scaling_calibrates_from_high_precision_tensor():
+    quantizer = Float8CurrentScalingQuantizer(
+        fp8_dtype=DType.kFloat8E4M3,
+        device=torch.device("cpu"),
+    )
+
+    quantizer.calibrate(torch.tensor([-112.0, 224.0]))
+    value = quantizer._calibration_state["scale_inv"]
+
+    torch.testing.assert_close(value, torch.tensor([0.5]))
+
+
+@pytest.mark.parametrize("input_value", (0.0, float("inf")))
+@pytest.mark.parametrize("force_pow_2_scales", (False, True))
+def test_current_scaling_calibration_handles_non_finite_scale(
+    input_value, force_pow_2_scales
+):
+    quantizer = Float8CurrentScalingQuantizer(
+        fp8_dtype=DType.kFloat8E4M3,
+        device=torch.device("cpu"),
+        force_pow_2_scales=force_pow_2_scales,
+    )
+
+    quantizer.calibrate(torch.tensor([input_value]))
+
+    torch.testing.assert_close(quantizer._calibration_state["scale_inv"], torch.ones(1))
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("fp8_dtype", (DType.kFloat8E4M3, DType.kFloat8E5M2))
+def test_delayed_scaling_high_precision_calibration_matches_quantization(fp8_dtype):
+    torch.manual_seed(123)
+    tensor = torch.randn((32, 32), dtype=torch.bfloat16, device="cuda")
+
+    active_quantizer = Float8Quantizer(
+        scale=torch.ones(1, dtype=torch.float32, device="cuda"),
+        amax=torch.zeros(1, dtype=torch.float32, device="cuda"),
+        fp8_dtype=fp8_dtype,
+        rowwise=True,
+        columnwise=False,
+    )
+    quantized_tensor = active_quantizer(tensor)
+    active_quantizer.calibrate(quantized_tensor)
+
+    calibration_quantizer = Float8Quantizer(
+        scale=torch.ones(1, dtype=torch.float32, device="cuda"),
+        amax=torch.zeros(1, dtype=torch.float32, device="cuda"),
+        fp8_dtype=fp8_dtype,
+        rowwise=True,
+        columnwise=False,
+    )
+    calibration_quantizer.calibrate(tensor)
+    assert (
+        calibration_quantizer.copy()._calibration_state
+        is calibration_quantizer._calibration_state
+    )
+
+    torch.testing.assert_close(
+        active_quantizer._calibration_state["amax"],
+        active_quantizer.amax,
+        atol=0.0,
+        rtol=0.0,
+    )
+    torch.testing.assert_close(
+        calibration_quantizer._calibration_state["amax"],
+        active_quantizer.amax,
+        atol=0.0,
+        rtol=0.0,
+    )
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("fp8_dtype", (DType.kFloat8E4M3, DType.kFloat8E5M2))
+@pytest.mark.parametrize("force_pow_2_scales", (False, True))
+@pytest.mark.parametrize("amax_epsilon", (0.0, 4.0))
+def test_current_scaling_high_precision_calibration_matches_quantization(
+    fp8_dtype, force_pow_2_scales, amax_epsilon
+):
+    torch.manual_seed(123)
+    tensor = torch.randn((32, 32), dtype=torch.bfloat16, device="cuda")
+    quantizer_kwargs = {
+        "fp8_dtype": fp8_dtype,
+        "device": torch.device("cuda"),
+        "rowwise": True,
+        "columnwise": False,
+        "force_pow_2_scales": force_pow_2_scales,
+        "amax_epsilon": amax_epsilon,
+    }
+
+    active_quantizer = Float8CurrentScalingQuantizer(**quantizer_kwargs)
+    quantized_tensor = active_quantizer(tensor)
+    active_quantizer.calibrate(quantized_tensor)
+
+    calibration_quantizer = Float8CurrentScalingQuantizer(**quantizer_kwargs)
+    calibration_quantizer.calibrate(tensor)
+
+    torch.testing.assert_close(
+        active_quantizer._calibration_state["scale_inv"],
+        quantized_tensor._scale_inv,
+        atol=0.0,
+        rtol=0.0,
+    )
+    torch.testing.assert_close(
+        calibration_quantizer._calibration_state["scale_inv"],
+        quantized_tensor._scale_inv,
+        atol=0.0,
+        rtol=0.0,
+    )
+
+
+@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4)
+@pytest.mark.parametrize(
+    ("row_scaled_nvfp4", "with_rht", "with_post_rht_amax", "with_random_sign_mask"),
+    (
+        (False, False, False, False),
+        (False, True, False, False),
+        (False, True, True, False),
+        (False, True, True, True),
+        (True, False, False, False),
+    ),
+)
+def test_nvfp4_high_precision_calibration_matches_quantization(
+    row_scaled_nvfp4, with_rht, with_post_rht_amax, with_random_sign_mask
+):
+    torch.manual_seed(123)
+    tensor = torch.randn((32, 32), dtype=torch.bfloat16, device="cuda")
+    quantizer_kwargs = {
+        "rowwise": True,
+        "columnwise": not row_scaled_nvfp4,
+        "with_rht": with_rht,
+        "with_post_rht_amax": with_post_rht_amax,
+        "row_scaled_nvfp4": row_scaled_nvfp4,
+        "with_random_sign_mask": with_random_sign_mask,
+    }
+
+    active_quantizer = NVFP4Quantizer(**quantizer_kwargs)
+    quantized_tensor = active_quantizer(tensor)
+    active_quantizer.calibrate(quantized_tensor)
+    calibration_quantizer = NVFP4Quantizer(**quantizer_kwargs)
+    calibration_quantizer.calibrate(tensor)
+    assert (
+        calibration_quantizer.copy()._calibration_state
+        is calibration_quantizer._calibration_state
+    )
+
+    expected_metadata_name = "amax_rowwise" if row_scaled_nvfp4 else "amax"
+    active_amax = active_quantizer._calibration_state[expected_metadata_name]
+    calibrated_amax = calibration_quantizer._calibration_state[expected_metadata_name]
+    torch.testing.assert_close(
+        active_amax,
+        quantized_tensor._amax_rowwise,
+        atol=0.0,
+        rtol=0.0,
+    )
+    torch.testing.assert_close(
+        calibrated_amax,
+        quantized_tensor._amax_rowwise,
+        atol=0.0,
+        rtol=0.0,
+    )
+
+
+def test_shallow_quantizer_copy_shares_calibration_state():
+    quantizer = Float8CurrentScalingQuantizer(
+        fp8_dtype=DType.kFloat8E4M3,
+        device=torch.device("cpu"),
+    )
+    copied_quantizer = quantizer.copy()
+    copied_quantizer.calibrate(torch.tensor([-112.0, 224.0]))
+    value = copied_quantizer._calibration_state["scale_inv"]
+
+    assert quantizer._calibration_state["scale_inv"] is value
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 91a31d14eb..5ace9253f4 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2033,8 +2033,6 @@ def _fallback_case(case, dtype, device):
         model_kwargs["fuse_wgrad_accumulation"] = True
     elif case == "delayed_wgrad":
         model_kwargs["delay_wgrad_compute"] = True
-    elif case == "scale_buffering":
-        model_kwargs["buffer_quantized_scaling_factors"] = True
     model = te.Linear(64, 32, params_dtype=dtype, device=device, **model_kwargs)
 
     if case == "fp8_output_differentiable":
@@ -2062,7 +2060,10 @@ def fn(inp):
         fp8_recipe = recipe.Float8CurrentScaling()
 
         def fn(inp):
-            with te.autocast(recipe=fp8_recipe):
+            with te.autocast(
+                recipe=fp8_recipe,
+                calibration_config=te.QuantizationCalibrationConfig(),
+            ):
                 return model(inp)
 
         return model, fn, "bwd", None, "quantized scaling-factor buffering"
diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py
index 107dd7a373..873da92d7a 100644
--- a/transformer_engine/debug/pytorch/debug_quantization.py
+++ b/transformer_engine/debug/pytorch/debug_quantization.py
@@ -429,8 +429,9 @@ def any_feature_enabled(self) -> bool:
                 return True
         return False
 
-    def calibrate(self, tensor: torch.Tensor):
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0):
         """Calibration override, should not be invoked."""
+        del tensor, decay
         raise RuntimeError("[NVTORCH-INSPECT ERROR] Calibration with debug is not supported")
 
     def update_quantized(
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 576fb57c5e..11ef1bf5a9 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -55,6 +55,8 @@
 from transformer_engine.pytorch.quantization import QuantizerRole
 from transformer_engine.pytorch.quantization import QuantizerRequest
 from transformer_engine.pytorch.quantization import DelayedScalingRequest
+from transformer_engine.pytorch.quantization import TEAutocastState
+from transformer_engine.pytorch.quantization import QuantizationCalibrationConfig
 from transformer_engine.pytorch.utils import get_cudnn_version
 from transformer_engine.pytorch.utils import get_device_compute_capability
 from transformer_engine.pytorch.utils import is_bf16_available
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 658dab5d88..ba865543a7 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -1121,7 +1121,10 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         # With quantization off the base class does all that is needed.
         qstate = FP8GlobalStateManager.quantization_state
         if torch.compiler.is_compiling() and not (
-            qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters
+            qstate.fp8_enabled
+            or qstate.fp8_calibration
+            or qstate.calibration_config is not None
+            or qstate.fp8_parameters
         ):
             super().init_fp8_metadata(num_gemms=num_gemms)
             return
diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py
index 04fa56721d..bad73a02e6 100644
--- a/transformer_engine/pytorch/graph.py
+++ b/transformer_engine/pytorch/graph.py
@@ -20,6 +20,7 @@
 from .quantization import (
     autocast,
     FP8GlobalStateManager,
+    QuantizationCalibrationConfig,
     get_default_fp8_recipe,
 )
 from .distributed import get_all_rng_states, graph_safe_rng_available
@@ -1445,6 +1446,7 @@ def make_graphed_callables(
     pre_warmup_hook: Optional[Callable] = None,
     post_warmup_hook: Optional[Callable] = None,
     capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]] = None,
+    calibration_config: Optional[QuantizationCalibrationConfig] = None,
 ) -> Union[Callable, Tuple[Callable, ...]]:
     """
     Make CUDA graph version of Transformer Engine modules
@@ -1518,10 +1520,10 @@ def make_graphed_callables(
              whether or not to enable low precision quantization (FP8/FP4).
              If tuple, the length must match the number of modules.
     calibrating: bool, default = False
-                 calibration mode allows collecting statistics such as amax and scale
-                 data of quantized tensors even when executing without quantization enabled.
-                 This is useful for saving an inference ready checkpoint while training
-                 using a higher precision.
+                 Enables calibration with the default configuration.
+    calibration_config: QuantizationCalibrationConfig, default = None
+                 Custom configuration for collecting checkpointable quantization scaling
+                 factors. Providing a config also enables calibration.
     recipe: recipe.Recipe, default = None
             recipe used for low precision quantization.
     amax_reduction_group: torch._C._distributed_c10d.ProcessGroup, default = None
@@ -1668,6 +1670,7 @@ def call_func(self, *args, **kwargs):
                 recipe=recipe,
                 amax_reduction_group=amax_reduction_group,
                 _graph=True,
+                calibration_config=calibration_config,
             ):
                 outputs = old_call_funcs[block_cls](self, *args, **kwargs)
             return outputs
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index ede7e7a608..b4ef1619fe 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -14,83 +14,27 @@
 from ..constants import TE_DType
 from ..export import is_in_onnx_export_mode
 from ..tensor.hybrid_tensor import HybridQuantizer
-from ..tensor.utils import get_quantization_recipe_name
 from ..utils import get_default_init_method
 
 
-def _get_scale_buffer_info(
-    tensor_name: str,
-    tensor: Any,
-    quantizer: Any,
-) -> Optional[Tuple[str, Optional[torch.Tensor]]]:
-    """Get the calibration buffer name and value for a quantized tensor."""
-    recipe = get_quantization_recipe_name(quantizer)
+def _resolve_calibration_quantizer(tensor: Any, quantizer: Any) -> Any:
+    """Get the quantizer that owns calibration state for a tensor."""
+    quantizer = getattr(tensor, "_quantizer", None) or quantizer
+    return getattr(quantizer, "parent_quantizer", quantizer)
+
+
+def _get_scale_buffer_info(tensor_name: str, quantizer: Any) -> Dict[str, torch.Tensor]:
+    """Get checkpoint-buffer aliases from quantizer calibration state."""
+    if quantizer is None:
+        return {}
+    recipe = quantizer.get_quantization_recipe_name()
     if not recipe:
-        return None
-
-    if recipe == "fp8_delayed_scaling":
-        metadata_name = "amax"
-        metadata = getattr(quantizer, "amax", None)
-    elif recipe == "fp8_current_scaling":
-        metadata_name = "scale_inv"
-        metadata = getattr(tensor, "_scale_inv", None)
-    elif recipe == "nvfp4":
-        metadata_name = "amax"
-        metadata = getattr(tensor, "_amax_rowwise", None)
-    elif recipe == "nvfp4_rowwise":
-        metadata_name = "amax_rowwise"
-        metadata = getattr(tensor, "_amax_rowwise", None)
-    elif recipe == "mxfp8":
-        # MXFP8 only exposes blockwise E8M0-encoded inverse scales, not a
-        # global FP32 scaling factor suitable for PTQ checkpoint export.
-        return None
-    elif recipe == "fp8_block_scaling":
-        # FP8 block scaling only exposes blockwise inverse scales, not a
-        # global FP32 scaling factor suitable for PTQ checkpoint export.
-        return None
-    else:
-        raise ValueError(f"Unsupported quantization recipe {recipe!r}")
-
-    buffer_name = f"{tensor_name}_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
-    return buffer_name, metadata
-
-
-def _update_scale_buffers(
-    scale_buffers: Dict[str, Optional[torch.Tensor]],
-    scale_updates: Dict[str, Optional[torch.Tensor]],
-    activation_scale_decay: float = 0.0,
-) -> None:
-    """Merge observed scaling factors into checkpoint buffers."""
-    for buffer_name, scale in scale_updates.items():
-        if scale is None or torch.isnan(scale).any():
-            # Un-initialized scale. Ignore it.
-            continue
-        if activation_scale_decay > 0.0:
-            observed_scale = scale.detach().float()
-            scale_buffer = scale_buffers.get(buffer_name)
-            if scale_buffer is not None and scale_buffer.shape != observed_scale.shape:
-                raise RuntimeError(
-                    "Quantized scaling-factor buffer shape changed from "
-                    f"{tuple(scale_buffer.shape)} to {tuple(observed_scale.shape)}"
-                )
-            if scale_buffer is None:
-                # Initialize the rolling activation scaling factor.
-                # Requires CUDA graph warmup step.
-                scale_buffer = torch.zeros_like(observed_scale)
-                scale_buffers[buffer_name] = scale_buffer
-            # Track a decaying maximum so early-training activation
-            # outliers do not permanently determine the inference scale.
-            scale_buffer.mul_(activation_scale_decay)
-            torch.maximum(
-                scale_buffer,
-                observed_scale,
-                out=scale_buffer,
-            )
-        else:
-            # Without scale history, keep a reference to the current metadata
-            # without allocating or copying a separate buffer.
-            # Requires CUDA graph warmup step.
-            scale_buffers[buffer_name] = scale.detach()
+        return {}
+    return {
+        # Standard naming for PTQ calibration data.
+        f"{tensor_name}_tensor_{metadata_name}_{recipe}_te_ptq_calibrated": value
+        for metadata_name, value in getattr(quantizer, "_calibration_state", {}).items()
+    }
 
 
 def set_quantizer_amax_reduction_group(quantizer, amax_reduction_group) -> None:
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index a6867b2102..e01580ce19 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -33,7 +33,7 @@
 )
 from ._common import (
     _get_scale_buffer_info,
-    _update_scale_buffers,
+    _resolve_calibration_quantizer,
     can_reconstruct_wgrad_input_from_original,
     WeightGradStore,
 )
@@ -163,31 +163,38 @@ def _update_grouped_scale_buffers(
     weight_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
     input_quantizers: List[Optional[Quantizer]],
     weight_quantizers: List[Optional[Quantizer]],
-    activation_scale_decay: float,
 ) -> None:
-    """Update GroupedLinear PTQ calibration buffers with per-GEMM metadata."""
+    """Update GroupedLinear PTQ buffers from per-GEMM calibration state."""
     activation_scale_updates = {}
     for index, tensor in enumerate(input_tensors):
-        scale_buffer = _get_scale_buffer_info(f"input_gemm{index}", tensor, input_quantizers[index])
-        if scale_buffer is not None:
-            activation_scale_updates[scale_buffer[0]] = scale_buffer[1]
+        quantizer = _resolve_calibration_quantizer(tensor, input_quantizers[index])
+        activation_scale_updates.update(_get_scale_buffer_info(f"input_gemm{index}", quantizer))
     weight_scale_updates = {}
     for index, tensor in enumerate(weight_tensors):
-        scale_buffer = _get_scale_buffer_info(
-            f"weight_gemm{index}", tensor, weight_quantizers[index]
+        quantizer = _resolve_calibration_quantizer(tensor, weight_quantizers[index])
+        weight_scale_updates.update(
+            _get_scale_buffer_info(f"weight_gemm{index}", quantizer)
         )
-        if scale_buffer is not None:
-            weight_scale_updates[scale_buffer[0]] = scale_buffer[1]
-    _update_scale_buffers(
-        scale_buffers,
-        activation_scale_updates,
-        activation_scale_decay,
-    )
-    _update_scale_buffers(
-        scale_buffers,
-        weight_scale_updates,
-        activation_scale_decay=0.0,
-    )
+    scale_buffers.update(activation_scale_updates)
+    scale_buffers.update(weight_scale_updates)
+
+
+def _calibrate_grouped_tensors(
+    input_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
+    weight_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
+    input_quantizers: List[Optional[Quantizer]],
+    weight_quantizers: List[Optional[Quantizer]],
+    activation_scale_decay: float,
+) -> None:
+    """Calibrate GroupedLinear input and weight quantizers once per GEMM."""
+    for tensor, quantizer in zip(input_tensors, input_quantizers):
+        quantizer = _resolve_calibration_quantizer(tensor, quantizer)
+        if quantizer is not None:
+            quantizer.calibrate(tensor, decay=activation_scale_decay)
+    for tensor, quantizer in zip(weight_tensors, weight_quantizers):
+        quantizer = _resolve_calibration_quantizer(tensor, quantizer)
+        if quantizer is not None:
+            quantizer.calibrate(tensor)
 
 
 class _GroupedLinear(torch.autograd.Function):
@@ -586,13 +593,19 @@ def _forward_grouped_tensor(
                     grouped_weights = weights_for_gemm.split_into_quantized_tensors()
             else:
                 grouped_weights = weights_for_gemm
+            _calibrate_grouped_tensors(
+                grouped_inputs,
+                grouped_weights,
+                input_quantizers,
+                weight_quantizers,
+                quantized_scaling_factor_buffering_decay,
+            )
             _update_grouped_scale_buffers(
                 scale_buffers,
                 grouped_inputs,
                 grouped_weights,
                 input_quantizers,
                 weight_quantizers,
-                quantized_scaling_factor_buffering_decay,
             )
 
         if is_grad_enabled:
@@ -918,6 +931,16 @@ def forward(
         else:
             weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]
 
+        if scale_buffers is not None:
+            _calibrate_grouped_tensors(
+                inputmats,
+                weights_fp8,
+                input_quantizers,
+                weight_quantizers,
+                quantized_scaling_factor_buffering_decay,
+            )
+
+        # Buffer scaling metadata only when requested.
         if scale_buffers is not None:
             _update_grouped_scale_buffers(
                 scale_buffers,
@@ -925,7 +948,6 @@ def forward(
                 weights_fp8,
                 input_quantizers,
                 weight_quantizers,
-                quantized_scaling_factor_buffering_decay,
             )
 
         # Initialize biases
@@ -963,11 +985,6 @@ def forward(
             use_split_accumulator=use_split_accumulator,
         )
 
-        if fp8_calibration:
-            for i in range(num_gemms):
-                input_quantizers[i].calibrate(inputmats[i])
-                weight_quantizers[i].calibrate(weights[i])
-
         if cpu_offloading:
             mark_not_offload(*weights_fp8, *weights)
 
@@ -1663,10 +1680,6 @@ class GroupedLinear(TransformerEngineBaseModule):
                        and saving the original input tensor may reduce the memory usage.
                        Requires input quantizers that can safely reproduce their results from the
                        original input. Cannot work with FP8 DelayedScaling recipe.
-    buffer_quantized_scaling_factors : bool, default = False
-                       If set to ``True``, maintain nonpersistent input and weight quantization
-                       metadata buffers for inference checkpoint export. Buffers store metadata
-                       per grouped GEMM, using inverse scales directly for FP8 current scaling.
     single_grouped_weight : bool, default = False
                        If set to ``True``, grouped weights are stored as a single grouped parameter
                        instead of one parameter per GEMM.
@@ -1720,8 +1733,6 @@ def __init__(
         single_grouped_bias: bool = False,
         name: Optional[str] = None,
         use_grouped_tensor: Optional[bool] = None,
-        buffer_quantized_scaling_factors: bool = False,
-        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -1755,8 +1766,6 @@ def __init__(
                 f"use_grouped_tensor must be a bool or None, got {type(use_grouped_tensor)}."
             )
         self.use_grouped_tensor = use_grouped_tensor
-        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
-        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
         single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags(
             single_grouped_weight, single_grouped_bias
         )
@@ -2353,8 +2362,9 @@ def forward(
                     if cache_weight
                     else [None] * num_gemms
                 )
+            calibration_config = FP8GlobalStateManager.get_calibration_config()
             scale_buffers = None
-            if self.buffer_quantized_scaling_factors:
+            if calibration_config is not None:
                 scale_buffers = {
                     name: value
                     for name, value in self._buffers.items()
@@ -2387,7 +2397,11 @@ def forward(
                 use_grouped_bias,
                 self.use_grouped_tensor,
                 scale_buffers,
-                self.quantized_scaling_factor_buffering_decay,
+                (
+                    calibration_config.activation_scale_decay
+                    if calibration_config is not None
+                    else 0.0
+                ),
             )
             out, new_workspaces = linear_fn(
                 *autograd_ctx,
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 2f3df6fe87..9b72c5056d 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -70,7 +70,7 @@
 from ..graph import is_graph_capturing
 from ._common import (
     _get_scale_buffer_info,
-    _update_scale_buffers,
+    _resolve_calibration_quantizer,
     apply_normalization,
     noop_cat,
     set_quantizer_amax_reduction_group,
@@ -383,28 +383,31 @@ def forward(
             bias_dtype = torch.bfloat16
         bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias
 
-        # Calibrate quantizers if needed
-        if not fp8 and fp8_calibration:
-            if input_quantizer is not None:
-                input_quantizer.calibrate(ln_out_total)
-            if weight_quantizer is not None:
-                weight_quantizer.calibrate(weight)
+        input_calibration_quantizer = _resolve_calibration_quantizer(
+            ln_out_total, input_quantizer
+        )
+        weight_calibration_quantizer = _resolve_calibration_quantizer(
+            weightmat, weight_quantizer
+        )
 
+        # Calibrate quantizers if needed.
         if scale_buffers is not None:
-            input_scale_buffer = _get_scale_buffer_info("input", ln_out_total, input_quantizer)
-            if input_scale_buffer is not None:
-                _update_scale_buffers(
-                    scale_buffers,
-                    {input_scale_buffer[0]: input_scale_buffer[1]},
-                    quantized_scaling_factor_buffering_decay,
-                )
-            weight_scale_buffer = _get_scale_buffer_info("weight", weightmat, weight_quantizer)
-            if weight_scale_buffer is not None:
-                _update_scale_buffers(
-                    scale_buffers,
-                    {weight_scale_buffer[0]: weight_scale_buffer[1]},
-                    activation_scale_decay=0.0,
+            if input_calibration_quantizer is not None:
+                input_calibration_quantizer.calibrate(
+                    ln_out_total,
+                    decay=quantized_scaling_factor_buffering_decay,
                 )
+            if weight_calibration_quantizer is not None:
+                weight_calibration_quantizer.calibrate(weightmat)
+
+        # Buffer scaling metadata only when requested.
+        if scale_buffers is not None:
+            scale_buffers.update(
+                _get_scale_buffer_info("input", input_calibration_quantizer)
+            )
+            scale_buffers.update(
+                _get_scale_buffer_info("weight", weight_calibration_quantizer)
+            )
 
         # Choose whether to use GEMM kernel with split accumulator
         use_split_accumulator = _2X_ACC_FPROP
@@ -1324,15 +1327,6 @@ class LayerNormLinear(TransformerEngineBaseModule):
                    This can help in latency bound communication situations.
                    Requires PyTorch version 2.7.0 or higher. When set to ``None``, standard all-reduce
                    is used.
-    buffer_quantized_scaling_factors : bool, default = False
-                       If set to ``True``, maintain nonpersistent input and weight quantization
-                       metadata buffers for inference checkpoint export. Per-tensor buffers
-                       store raw global amaxes, except FP8 current scaling buffers, which store
-                       inverse scales directly.
-    quantized_scaling_factor_buffering_decay : float, default = 0.0
-                       Decay applied to buffered activation scaling factors before incorporating
-                       each new observation. Defaults to 0.0, in which case only the most recent
-                       scaling factor is buffered.
     """
 
     def __init__(
@@ -1365,8 +1359,6 @@ def __init__(
         delay_wgrad_compute: bool = False,
         symmetric_ar_type: Optional[str] = None,
         name: Optional[str] = None,
-        buffer_quantized_scaling_factors: bool = False,
-        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -1629,9 +1621,6 @@ def __init__(
                 if name in self.weight_names or name in self.bias_names:
                     param.skip_backward_post_hook = True
 
-        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
-        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
-
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""
         super().set_meta_tensor(fwd, recipe)
@@ -1802,8 +1791,9 @@ def forward(
                 self._fp8_workspaces.get(cache_name) if cache_name is not None else None
             )
 
+            calibration_config = FP8GlobalStateManager.get_calibration_config()
             scale_buffers = None
-            if self.buffer_quantized_scaling_factors:
+            if calibration_config is not None:
                 scale_buffers = {
                     name: value
                     for name, value in self._buffers.items()
@@ -1850,7 +1840,11 @@ def forward(
                 self.symmetric_ar_type,
                 debug,
                 self.is_fsdp2,
-                self.quantized_scaling_factor_buffering_decay,
+                (
+                    calibration_config.activation_scale_decay
+                    if calibration_config is not None
+                    else 0.0
+                ),
                 scale_buffers,
             )
             out, ln_out, new_weight_workspace = fwd_fn(
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 0ca4e353dc..2f5438293c 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -77,7 +77,7 @@
 from ..tensor.identity_tensor import IdentityQuantizer
 from ._common import (
     _get_scale_buffer_info,
-    _update_scale_buffers,
+    _resolve_calibration_quantizer,
     apply_normalization,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
@@ -568,21 +568,32 @@ def _forward(
         if fc2_bias is not None:
             fc2_bias = cast_if_needed(fc2_bias, bias_dtype)
 
-        # Calibrate quantizers if needed
-        if not fp8 and fp8_calibration:
-            if fc1_input_quantizer is not None:
-                fc1_input_quantizer.calibrate(ln_out_total)
-            if fc1_weight_quantizer is not None:
-                fc1_weight_quantizer.calibrate(fc1_weight)
+        fc1_input_calibration_quantizer = _resolve_calibration_quantizer(
+            ln_out_total, fc1_input_quantizer
+        )
+        fc1_weight_calibration_quantizer = _resolve_calibration_quantizer(
+            fc1_weight_final, fc1_weight_quantizer
+        )
+
+        # Calibrate FC1 quantizers if needed.
+        should_calibrate = scale_buffers is not None
+        if should_calibrate:
+            if fc1_input_calibration_quantizer is not None:
+                fc1_input_calibration_quantizer.calibrate(
+                    ln_out_total,
+                    decay=quantized_scaling_factor_buffering_decay,
+                )
+            if fc1_weight_calibration_quantizer is not None:
+                fc1_weight_calibration_quantizer.calibrate(fc1_weight_final)
 
-        fc1_input_scale_buffer = None
-        fc1_weight_scale_buffer = None
+        fc1_input_scale_buffer = {}
+        fc1_weight_scale_buffer = {}
         if scale_buffers is not None:
             fc1_input_scale_buffer = _get_scale_buffer_info(
-                "fc1_input", ln_out_total, fc1_input_quantizer
+                "fc1_input", fc1_input_calibration_quantizer
             )
             fc1_weight_scale_buffer = _get_scale_buffer_info(
-                "fc1_weight", fc1_weight_final, fc1_weight_quantizer
+                "fc1_weight", fc1_weight_calibration_quantizer
             )
 
         # ------------------------------------------------------
@@ -675,9 +686,14 @@ def _forward(
                 else:
                     act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params)
 
-        if not fp8 and fp8_calibration:
-            if fc2_input_quantizer is not None:
-                fc2_input_quantizer.calibrate(act_out)
+        fc2_input_calibration_quantizer = _resolve_calibration_quantizer(
+            act_out, fc2_input_quantizer
+        )
+        if should_calibrate and fc2_input_calibration_quantizer is not None:
+            fc2_input_calibration_quantizer.calibrate(
+                act_out,
+                decay=quantized_scaling_factor_buffering_decay,
+            )
 
         # we want to skip fc2 computation if we are checkpointing and recomputing,
         # otherwise we compute fc2
@@ -694,38 +710,27 @@ def _forward(
             ):  # we can safely get rid of these if this is the case
                 clear_tensor_data(fc1_out)
 
-            if not fp8 and fp8_calibration:
-
-                if fc2_weight_quantizer is not None:
-                    fc2_weight_quantizer.calibrate(fc2_weight)
+            fc2_weight_calibration_quantizer = _resolve_calibration_quantizer(
+                fc2_weight_final, fc2_weight_quantizer
+            )
+            if should_calibrate and fc2_weight_calibration_quantizer is not None:
+                fc2_weight_calibration_quantizer.calibrate(fc2_weight_final)
 
             if scale_buffers is not None:
                 activation_scale_updates = {}
                 weight_scale_updates = {}
-                if fc1_input_scale_buffer is not None:
-                    activation_scale_updates[fc1_input_scale_buffer[0]] = fc1_input_scale_buffer[1]
-                if fc1_weight_scale_buffer is not None:
-                    weight_scale_updates[fc1_weight_scale_buffer[0]] = fc1_weight_scale_buffer[1]
+                activation_scale_updates.update(fc1_input_scale_buffer)
+                weight_scale_updates.update(fc1_weight_scale_buffer)
                 fc2_input_scale_buffer = _get_scale_buffer_info(
-                    "fc2_input", act_out, fc2_input_quantizer
+                    "fc2_input", fc2_input_calibration_quantizer
                 )
-                if fc2_input_scale_buffer is not None:
-                    activation_scale_updates[fc2_input_scale_buffer[0]] = fc2_input_scale_buffer[1]
+                activation_scale_updates.update(fc2_input_scale_buffer)
                 fc2_weight_scale_buffer = _get_scale_buffer_info(
-                    "fc2_weight", fc2_weight_final, fc2_weight_quantizer
-                )
-                if fc2_weight_scale_buffer is not None:
-                    weight_scale_updates[fc2_weight_scale_buffer[0]] = fc2_weight_scale_buffer[1]
-                _update_scale_buffers(
-                    scale_buffers,
-                    activation_scale_updates,
-                    quantized_scaling_factor_buffering_decay,
-                )
-                _update_scale_buffers(
-                    scale_buffers,
-                    weight_scale_updates,
-                    activation_scale_decay=0.0,
+                    "fc2_weight", fc2_weight_calibration_quantizer
                 )
+                weight_scale_updates.update(fc2_weight_scale_buffer)
+                scale_buffers.update(activation_scale_updates)
+                scale_buffers.update(weight_scale_updates)
 
             # Configure Userbuffers reduce-scatter if needed
             ub_obj_fc2out = None
@@ -1993,15 +1998,6 @@ class LayerNormMLP(TransformerEngineBaseModule):
                 whether to use selective activation checkpointing, where activations are not saved for bwd,
                 and instead are recomputed (skipping fc2, as it is not needed for backward). Trades compute
                 for memory. default is false, in which activations are saved in fwd. not supported for onnx forward
-    buffer_quantized_scaling_factors : bool, default = False
-                       If set to ``True``, maintain nonpersistent activation and weight quantization
-                       metadata buffers for both internal linear layers for inference checkpoint
-                       export. Per-tensor buffers store raw global amaxes, except FP8 current
-                       scaling buffers, which store inverse scales directly.
-    quantized_scaling_factor_buffering_decay : float, default = 0.0
-                       Decay applied to buffered activation scaling factors before incorporating
-                       each new observation. Defaults to 0.0, in which case only the most recent
-                       scaling factor is buffered.
     """
 
     def __init__(
@@ -2038,8 +2034,6 @@ def __init__(
         delay_wgrad_compute: bool = False,
         symmetric_ar_type: Optional[str] = None,
         checkpoint: bool = False,
-        buffer_quantized_scaling_factors: bool = False,
-        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -2061,9 +2055,6 @@ def __init__(
         self.zero_centered_gamma = zero_centered_gamma
         self.symmetric_ar_type = symmetric_ar_type
         self.checkpoint = checkpoint
-        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
-        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
-
         # GEMM-GELU fusion is currently only supported with split GEMM-AG overlap
         self.gemm_gelu_fusion = (
             bool(int(os.getenv("NVTE_GEMM_GELU_FUSION", "0")))
@@ -2443,8 +2434,9 @@ def forward(
                 self._fp8_workspaces.get(cache_name_fc2) if cache_name_fc2 is not None else None
             )
 
+            calibration_config = FP8GlobalStateManager.get_calibration_config()
             scale_buffers = None
-            if self.buffer_quantized_scaling_factors:
+            if calibration_config is not None:
                 scale_buffers = {
                     name: value
                     for name, value in self._buffers.items()
@@ -2501,7 +2493,11 @@ def forward(
                 self.checkpoint,
                 debug,
                 self.is_fsdp2,
-                self.quantized_scaling_factor_buffering_decay,
+                (
+                    calibration_config.activation_scale_decay
+                    if calibration_config is not None
+                    else 0.0
+                ),
                 scale_buffers,
             )
             out, ln_out, new_fc1_ws, new_fc2_ws = fwd_fn(
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index fa9d159111..7a08b2ffc3 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -33,7 +33,7 @@
 )
 from ._common import (
     _get_scale_buffer_info,
-    _update_scale_buffers,
+    _resolve_calibration_quantizer,
     can_reconstruct_wgrad_input_from_original,
     noop_cat,
     set_quantizer_amax_reduction_group,
@@ -608,29 +608,29 @@ def _linear_forward_impl(
         bias_dtype = torch.bfloat16
     bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias
 
-    # Calibrate quantizers if needed
-    if not fp8 and args.fp8_calibration:
-        if input_quantizer is not None:
-            input_quantizer.calibrate(inputmat_total)
-        if weight_quantizer is not None:
-            weight_quantizer.calibrate(weight)
+    input_calibration_quantizer = _resolve_calibration_quantizer(
+        inputmat_total, input_quantizer
+    )
+    weight_calibration_quantizer = _resolve_calibration_quantizer(weightmat, weight_quantizer)
 
-    # Capture scaling metadata while it is still available.
+    # Calibrate quantizers if needed.
     if args.scale_buffers is not None:
-        input_scale_buffer = _get_scale_buffer_info("input", inputmat_total, input_quantizer)
-        if input_scale_buffer is not None:
-            _update_scale_buffers(
-                args.scale_buffers,
-                {input_scale_buffer[0]: input_scale_buffer[1]},
-                args.quantized_scaling_factor_buffering_decay,
-            )
-        weight_scale_buffer = _get_scale_buffer_info("weight", weightmat, weight_quantizer)
-        if weight_scale_buffer is not None:
-            _update_scale_buffers(
-                args.scale_buffers,
-                {weight_scale_buffer[0]: weight_scale_buffer[1]},
-                activation_scale_decay=0.0,
+        if input_calibration_quantizer is not None:
+            input_calibration_quantizer.calibrate(
+                inputmat_total,
+                decay=args.quantized_scaling_factor_buffering_decay,
             )
+        if weight_calibration_quantizer is not None:
+            weight_calibration_quantizer.calibrate(weightmat)
+
+    # Buffer scaling metadata only when requested.
+    if args.scale_buffers is not None:
+        args.scale_buffers.update(
+            _get_scale_buffer_info("input", input_calibration_quantizer)
+        )
+        args.scale_buffers.update(
+            _get_scale_buffer_info("weight", weight_calibration_quantizer)
+        )
 
     # Choose whether to use GEMM kernel with split accumulator
     use_split_accumulator = _2X_ACC_FPROP
@@ -2017,18 +2017,6 @@ class Linear(TransformerEngineBaseModule):
                        and saving the original input tensor may reduce the memory usage.
                        Requires an input quantizer that can safely reproduce its result from the
                        original input. Cannot work with FP8 DelayedScaling recipe.
-    buffer_quantized_scaling_factors : bool, default = False
-                       If set to ``True``, maintain nonpersistent input and weight quantization
-                       metadata buffers for inference checkpoint export. Per-tensor buffers
-                       store raw global amaxes, except FP8 current scaling buffers, which store
-                       inverse scales directly.
-                       Each buffer is materialized only when its tensor uses a quantizer with a
-                       per-tensor scaling factor. Used to propagate scaling factors from training
-                       into inference.
-    quantized_scaling_factor_buffering_decay : float, default = 0.0
-                       Decay applied to buffered activation scaling factors before incorporating
-                       each new observation. Defaults to 0.0, in which case only the most recent
-                       scaling factor is buffered.
     """
 
     def __init__(
@@ -2058,8 +2046,6 @@ def __init__(
         symmetric_ar_type: Optional[str] = None,
         save_original_input: bool = False,
         name: Optional[str] = None,
-        buffer_quantized_scaling_factors: bool = False,
-        quantized_scaling_factor_buffering_decay: float = 0.0,
     ) -> None:
         super().__init__(name)
 
@@ -2287,9 +2273,6 @@ def __init__(
                 if name in self.weight_names or name in self.bias_names:
                     param.skip_backward_post_hook = True
 
-        self.buffer_quantized_scaling_factors = buffer_quantized_scaling_factors
-        self.quantized_scaling_factor_buffering_decay = quantized_scaling_factor_buffering_decay
-
     def get_quantizer_roles(
         self,
         *,
@@ -2519,8 +2502,9 @@ def forward(
                 bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None
             )
             wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None
+            calibration_config = FP8GlobalStateManager.get_calibration_config()
             scale_buffers = None
-            if self.buffer_quantized_scaling_factors:
+            if calibration_config is not None:
                 scale_buffers = {
                     name: value
                     for name, value in self._buffers.items()
@@ -2585,7 +2569,9 @@ def forward(
                 # Inference Scaling Factor Calibration Buffering
                 scale_buffers=scale_buffers,
                 quantized_scaling_factor_buffering_decay=(
-                    self.quantized_scaling_factor_buffering_decay
+                    calibration_config.activation_scale_decay
+                    if calibration_config is not None
+                    else 0.0
                 ),
                 # misc
                 cpu_offloading=is_cpu_offload_enabled(),
@@ -2708,7 +2694,7 @@ def _compile_eager_fallback_reason(
         prepare_forward. Quantizer checks stay in compile_unsupported_reason."""
         if debug:
             return "debug instrumentation (nvidia-dlfw-inspect)"
-        if self.buffer_quantized_scaling_factors:
+        if FP8GlobalStateManager.get_calibration_config() is not None:
             return "quantized scaling-factor buffering"
         weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
         if is_distributed_weight(weight_tensor):
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 34ae20b498..be553b0d87 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -385,6 +385,37 @@ def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, s
     return check_nvfp4_support()[0]
 
 
+@dataclass(frozen=True, slots=True)
+class QuantizationCalibrationConfig:
+    """Configuration for collecting checkpointable quantization scaling factors.
+
+    Parameters
+    ----------
+    activation_scale_decay : float, default = 0.0
+        Decay applied to buffered activation scaling factors before incorporating
+        each new observation. With zero decay, only the latest value is retained.
+    """
+
+    activation_scale_decay: float = 0.0
+
+    def __post_init__(self) -> None:
+        if self.activation_scale_decay < 0.0:
+            raise ValueError("activation_scale_decay must be non-negative")
+
+
+@dataclass(frozen=True, slots=True)
+class TEAutocastState:
+    """Snapshot of process-global quantization autocast state."""
+
+    fp8_enabled: bool
+    fp8_calibration: bool
+    calibration_config: Optional[QuantizationCalibrationConfig]
+    fp8_recipe: Optional[Recipe]
+    fp8_distributed_group: Optional[dist_group_type]
+    is_first_fp8_module: bool
+    fp8_graph_capturing: bool
+
+
 @dataclass(slots=True)
 class FP8GlobalState:
     """Mutable process-global FP8 state stored on an instance.
@@ -410,6 +441,7 @@ class FP8GlobalState:
         default_factory=dict
     )
     skip_fp8_weight_update_tensor: Optional[torch.Tensor] = None
+    calibration_config: Optional[QuantizationCalibrationConfig] = None
 
 
 class FP8GlobalStateManager:
@@ -571,9 +603,22 @@ def is_fp8_enabled(cls) -> bool:
 
     @classmethod
     def is_fp8_calibration(cls) -> bool:
-        """Is FP8 calibration"""
+        """Whether quantization calibration is enabled."""
         return cls.quantization_state.fp8_calibration
 
+    @classmethod
+    def get_calibration_config(cls) -> Optional[QuantizationCalibrationConfig]:
+        """Get the active quantization calibration configuration."""
+        qstate = cls.quantization_state
+        if not qstate.fp8_calibration:
+            # User-declined calibration using the legacy fp8_calibration config.
+            return None
+        if qstate.calibration_config is not None:
+            # User-provided calibration config.
+            return qstate.calibration_config
+        # Default (fp8_calibration=True) config.
+        return QuantizationCalibrationConfig()
+
     @classmethod
     def with_fp8_parameters(cls) -> bool:
         """Should the parameters be stored as FP8"""
@@ -616,30 +661,30 @@ def get_fp8_group(cls) -> Union[dist_group_type, None]:
         return cls.quantization_state.fp8_distributed_group
 
     @classmethod
-    def get_autocast_state(cls) -> tuple:
+    def get_autocast_state(cls) -> TEAutocastState:
         """Snapshot the autocast-related fields of the quantization state."""
         qstate = cls.quantization_state
-        return (
-            qstate.fp8_enabled,
-            qstate.fp8_calibration,
-            qstate.fp8_recipe,
-            qstate.fp8_distributed_group,
-            qstate.is_first_fp8_module,
-            qstate.fp8_graph_capturing,
+        return TEAutocastState(
+            fp8_enabled=qstate.fp8_enabled,
+            fp8_calibration=qstate.fp8_calibration,
+            calibration_config=qstate.calibration_config,
+            fp8_recipe=qstate.fp8_recipe,
+            fp8_distributed_group=qstate.fp8_distributed_group,
+            is_first_fp8_module=qstate.is_first_fp8_module,
+            fp8_graph_capturing=qstate.fp8_graph_capturing,
         )
 
     @classmethod
-    def set_autocast_state(cls, state: tuple) -> None:
+    def set_autocast_state(cls, state: TEAutocastState) -> None:
         """Restore a previously saved autocast state snapshot."""
         qstate = cls.quantization_state
-        (
-            qstate.fp8_enabled,
-            qstate.fp8_calibration,
-            qstate.fp8_recipe,
-            qstate.fp8_distributed_group,
-            qstate.is_first_fp8_module,
-            qstate.fp8_graph_capturing,
-        ) = state
+        qstate.fp8_enabled = state.fp8_enabled
+        qstate.fp8_calibration = state.fp8_calibration
+        qstate.calibration_config = state.calibration_config
+        qstate.fp8_recipe = state.fp8_recipe
+        qstate.fp8_distributed_group = state.fp8_distributed_group
+        qstate.is_first_fp8_module = state.is_first_fp8_module
+        qstate.fp8_graph_capturing = state.fp8_graph_capturing
 
     @staticmethod
     def reduce_tensor_across_group_op_max(tensor: torch.Tensor, group: dist_group_type) -> None:
@@ -734,9 +779,13 @@ def autocast_enter(
         fp8_recipe: Optional[Recipe] = None,
         fp8_group: Optional[dist_group_type] = None,
         _graph: bool = False,
+        calibration_config: Optional[QuantizationCalibrationConfig] = None,
     ) -> None:
         """Set state and tracking variables for entry into FP8 region."""
 
+        if calibrating and calibration_config is None:
+            calibration_config = QuantizationCalibrationConfig()
+
         fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe
         autocast_key = cls.get_unique_autocast_key(fp8_recipe, fp8_group)
         qstate = cls.quantization_state
@@ -746,7 +795,8 @@ def autocast_enter(
         )
 
         qstate.fp8_enabled = enabled
-        qstate.fp8_calibration = calibrating
+        qstate.fp8_calibration = calibration_config is not None
+        qstate.calibration_config = calibration_config
         qstate.fp8_recipe = fp8_recipe
         qstate.fp8_distributed_group = fp8_group
         qstate.fp8_graph_capturing = _graph
@@ -944,6 +994,7 @@ def fp8_autocast(
     fp8_recipe: Optional[Recipe] = None,
     fp8_group: Optional[dist_group_type] = None,
     _graph: bool = False,
+    calibration_config: Optional[QuantizationCalibrationConfig] = None,
 ) -> "autocast":
     """
     .. warning::
@@ -966,6 +1017,7 @@ def fp8_autocast(
         recipe=fp8_recipe,
         amax_reduction_group=fp8_group,
         _graph=_graph,
+        calibration_config=calibration_config,
     )
 
 
@@ -997,10 +1049,12 @@ class autocast:
     enabled : bool, default = True
              whether or not to enable low precision quantization (FP8/FP4).
     calibrating : bool, default = False
-                 calibration mode allows collecting statistics such as amax and scale
-                 data of quantized tensors even when executing without quantization enabled.
-                 This is useful for saving an inference ready checkpoint while training
-                 using a higher precision.
+                 Enables calibration with the default configuration. Calibration
+                 collects and buffers quantized scaling factors even when executing
+                 without quantization enabled.
+    calibration_config : QuantizationCalibrationConfig, default = None
+                  Custom configuration for collecting checkpointable quantization scaling
+                  factors. Providing a config also enables calibration.
     recipe : recipe.Recipe, default = None
             recipe used for low precision quantization.
     amax_reduction_group : torch._C._distributed_c10d.ProcessGroup, default = None
@@ -1012,7 +1066,7 @@ class autocast:
     # to avoid overheads.
     __slots__ = (
         "_enabled",
-        "_calibrating",
+        "_calibration_config",
         "_recipe",
         "_amax_reduction_group",
         "_graph",
@@ -1026,9 +1080,14 @@ def __init__(
         recipe: Optional["Recipe"] = None,
         amax_reduction_group: Optional["dist_group_type"] = None,
         _graph: bool = False,
+        calibration_config: Optional[QuantizationCalibrationConfig] = None,
     ) -> None:
         self._enabled = enabled
-        self._calibrating = calibrating
+        self._calibration_config = (
+            QuantizationCalibrationConfig()
+            if calibrating and calibration_config is None
+            else calibration_config
+        )
         self._recipe = recipe
         self._amax_reduction_group = amax_reduction_group
         self._graph = _graph
@@ -1046,7 +1105,8 @@ def __enter__(self) -> "autocast":
         self._fp8_state = FP8GlobalStateManager.get_autocast_state()
         FP8GlobalStateManager.autocast_enter(
             enabled=self._enabled,
-            calibrating=self._calibrating,
+            calibrating=False,
+            calibration_config=self._calibration_config,
             fp8_recipe=self._recipe,
             fp8_group=self._amax_reduction_group,
             _graph=self._graph,
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 7149a5a163..284f17ed2c 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -404,6 +404,7 @@ def __init__(self, *, rowwise: bool, columnwise: bool) -> None:
         self.columnwise_usage = columnwise
         self.internal = False
         self.optimize_for_gemm = False
+        self._calibration_state: Dict[str, torch.Tensor] = {}
 
     def __repr__(self):
         return (
@@ -564,13 +565,60 @@ def create_metadata(
             "nontensor_kwargs": meta["nontensor_kwargs"],
         }
 
-    def calibrate(self, tensor: torch.Tensor) -> None:
-        """Calibrate quantizer state
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+        """Observe a tensor and update persistent calibration state."""
+        pass
 
-        Updates quantization state as if quantizing a tensor, but
-        without actually performing the quantization.
+    def get_quantization_recipe_name(self) -> str:
+        """Get the stable name of the quantization recipe."""
+        return ""
 
-        """
+    def _update_calibration_value(
+        self,
+        metadata_name: str,
+        observed_value: Optional[torch.Tensor],
+        *,
+        decay: float,
+    ) -> None:
+        """Merge an observation into quantizer-owned calibration state."""
+        if observed_value is None or torch.isnan(observed_value).any():
+            # Un-initialized scale. Ignore it.
+            return
+        observed_value = observed_value.detach()
+        calibration_state = getattr(self, "_calibration_state", None)
+        if calibration_state is None:
+            calibration_state = {}
+            self._calibration_state = calibration_state
+        calibration_value = calibration_state.get(metadata_name)
+        if decay > 0.0:
+            if calibration_value is not None and calibration_value.shape != observed_value.shape:
+                raise RuntimeError(
+                    "Quantizer calibration value shape changed from "
+                    f"{tuple(calibration_value.shape)} to {tuple(observed_value.shape)}"
+                )
+            if calibration_value is None:
+                # Initialize the rolling activation scaling factor.
+                # Requires CUDA graph warmup step.
+                calibration_value = torch.zeros_like(observed_value)
+                calibration_state[metadata_name] = calibration_value
+            # Track a decaying maximum so early-training activation
+            # outliers do not permanently determine the inference scale.
+            calibration_value.mul_(decay)
+            torch.maximum(calibration_value, observed_value, out=calibration_value)
+        else:
+            # Without scale history, keep a reference to the current metadata
+            # without allocating or copying a separate buffer.
+            # Requires CUDA graph warmup step, and only access this value
+            # at an appropriate time (e.g. checkpointing) if captured by CG.
+            calibration_state[metadata_name] = observed_value
+
+    def _share_calibration_state_with(self, quantizer: "Quantizer") -> None:
+        """Make a shallow quantizer copy share persistent calibration state."""
+        calibration_state = getattr(self, "_calibration_state", None)
+        if calibration_state is None:
+            calibration_state = {}
+            self._calibration_state = calibration_state
+        quantizer._calibration_state = calibration_state
 
     def set_usage(
         self, *, rowwise: Optional[bool] = None, columnwise: Optional[bool] = None
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index 6105b18a73..ac3dde7fbf 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -243,11 +243,19 @@ def is_quantizable(self, inp: torch.Tensor) -> bool:
             return False
         return True
 
-    def calibrate(self, tensor: torch.Tensor) -> None:
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
         # NOTE: This interface is specific to requirements like delayed scaling
         # where state from an estimator influences distribution parameters.
+        # NOTE(@cspades): Currently, PTQ calibration requirements don't need
+        # non-global / blockwise scaling factors, which are usually computed
+        # on-the-fly during inference. Implement this interface for future
+        # applications of blockwise scaling factor calibration.
         pass
 
+    def get_quantization_recipe_name(self) -> str:
+        """Get the stable name of the quantization recipe."""
+        return "fp8_block_scaling"
+
     def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
         return Float8BlockScaling
 
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index cf37c36c59..509a88b613 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -17,14 +17,14 @@
 )
 from ..utils import canonicalize_process_group, devices_match, is_non_tn_fp8_gemm_supported
 from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func
-from ..quantized_tensor import QuantizedTensor, Quantizer
+from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 from ..dynamo import register_value_opaque_quantizer
 from ._quantization_helpers import (
     _IdentityFunc,
     _resolve_view_shape,
     safe_quantized_repr,
 )
-from ..constants import dist_group_type, DType
+from ..constants import dist_group_type, DType, TE_DType_To_Torch
 
 aten = torch.ops.aten
 
@@ -93,6 +93,7 @@ def copy(self) -> Float8Quantizer:
             columnwise=self.columnwise_usage,
         )
         quantizer.internal = self.internal
+        self._share_calibration_state_with(quantizer)
 
         return quantizer
 
@@ -124,9 +125,21 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor:
         """Quantize tensor implementation"""
         return tex.quantize(tensor, self)
 
-    def calibrate(self, tensor: torch.Tensor) -> None:
-        amin, amax = tensor.aminmax()
-        self.amax.copy_(torch.max(-amin, amax))
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+        if isinstance(tensor, (QuantizedTensor, QuantizedTensorStorage)):
+            observed_amax = self.amax
+        else:
+            # If quantized amax metadata does not yet exist or calibrate() is called directly,
+            # then recompute the absmax without quantization. This path is
+            # not performant and SHOULD NOT be called within training or inference.
+            amin, amax = tensor.aminmax()
+            observed_amax = torch.max(-amin, amax).reshape(1)
+            self.amax.copy_(observed_amax)
+        self._update_calibration_value("amax", observed_amax, decay=decay)
+
+    def get_quantization_recipe_name(self) -> str:
+        """Get the stable name of the quantization recipe."""
+        return "fp8_delayed_scaling"
 
     def get_columnwise_shape(self, rowwise_data_shape: Iterable[int]) -> Tuple[int, ...]:
         """Calculate the shape of the columnwise data for Float8 1D blockwise quantization."""
@@ -276,6 +289,7 @@ def copy(self) -> Float8CurrentScalingQuantizer:
         )
         quantizer.internal = self.internal
         quantizer.optimize_for_gemm = self.optimize_for_gemm
+        self._share_calibration_state_with(quantizer)
 
         return quantizer
 
@@ -315,9 +329,37 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor:
         """Quantize tensor implementation"""
         return tex.quantize(tensor, self)
 
-    def calibrate(self, tensor: torch.Tensor) -> None:
-        # current scaling don't need to calibrate
-        return
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+        """Compute and calibrate (decaying amax) quantization metadata."""
+        scale_inv = getattr(tensor, "_scale_inv", None)
+        if scale_inv is None:
+            # If the scale_inv does not yet exist or calibrate() is called directly,
+            # then recompute the absmax / scale without quantization. This path is
+            # not performant and SHOULD NOT be called within training or inference.
+            amin, amax = tensor.aminmax()
+            amax = torch.maximum(-amin, amax).float().reshape(1)
+            if self.with_amax_reduction and torch.distributed.is_initialized():
+                torch.distributed.all_reduce(
+                    amax,
+                    op=torch.distributed.ReduceOp.MAX,
+                    group=self._canonicalized_amax_reduction_group(),
+                )
+            if self.amax_epsilon > 0.0:
+                amax.clamp_min_(self.amax_epsilon)
+            fp8_max = torch.finfo(TE_DType_To_Torch[self.dtype]).max
+            scale = fp8_max / amax
+            finite_scale = torch.finfo(tensor.dtype).max
+            scale.nan_to_num_(nan=float("nan"), posinf=finite_scale, neginf=finite_scale)
+            if self.force_pow_2_scales:
+                _, exponent = torch.frexp(scale)
+                scale = torch.ldexp(torch.ones_like(scale), exponent - 1)
+            scale.masked_fill_(torch.isinf(amax) | (amax == 0), 1.0)
+            scale_inv = torch.reciprocal(scale)
+        self._update_calibration_value("scale_inv", scale_inv, decay=decay)
+
+    def get_quantization_recipe_name(self) -> str:
+        """Get the stable name of the quantization recipe."""
+        return "fp8_current_scaling"
 
     def create_tensor_from_data(
         self,
diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py
index ec171564fe..0ff8c01550 100644
--- a/transformer_engine/pytorch/tensor/identity_tensor.py
+++ b/transformer_engine/pytorch/tensor/identity_tensor.py
@@ -162,9 +162,9 @@ def update_quantized(
         dst._dtype = data.dtype
         return dst
 
-    def calibrate(self, tensor: torch.Tensor) -> None:
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
         # No state to calibrate.
-        return
+        pass
 
     def _get_compatible_recipe(self):
         # Only reachable via CustomRecipe (qfactory returns IdentityQuantizer).
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index 267806e43e..e0c409f11d 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -132,10 +132,18 @@ def is_quantizable(self, inp: torch.Tensor) -> bool:
             return False
         return True
 
-    def calibrate(self, tensor: torch.Tensor) -> None:
-        # TODO(ksivamani): No calibration needed for mxfp8?
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+        """Calibrate an MXFP8 tensor."""
+        # NOTE(@cspades): Currently, PTQ calibration requirements don't need
+        # non-global / blockwise scaling factors, which are usually computed
+        # on-the-fly during inference. Implement this interface for future
+        # applications of MXFP8 calibration.
         pass
 
+    def get_quantization_recipe_name(self) -> str:
+        """Get the stable name of the quantization recipe."""
+        return "mxfp8"
+
     def get_scale_shape(
         self,
         shape: Iterable[int],
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 5e537c3ff0..fef3e10293 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -247,6 +247,7 @@ def copy(self) -> NVFP4Quantizer:
         )
         quantizer.internal = self.internal
         quantizer.optimize_for_gemm = self.optimize_for_gemm
+        self._share_calibration_state_with(quantizer)
 
         return quantizer
 
@@ -339,8 +340,35 @@ def convert_shape_for_fp4(shape: Iterable[int]) -> Tuple[int, ...]:
         shape[-1] = shape[-1] // 2
         return tuple(shape)
 
-    def calibrate(self, tensor: torch.Tensor) -> None:
-        pass  # Calibration is no-op
+    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+        metadata_name = "amax_rowwise" if self.row_scaled_nvfp4 else "amax"
+        observed_amax = getattr(tensor, "_amax_rowwise", None)
+        if observed_amax is None:
+            # If quantized amax metadata does not yet exist or calibrate() is called directly,
+            # then recompute the absmax without quantization. This path is
+            # not performant and SHOULD NOT be called within training or inference.
+            calibration_input = tensor
+            if self.with_rht and self.with_post_rht_amax:
+                original_shape = calibration_input.shape
+                calibration_input = (
+                    calibration_input.reshape(-1, 16).to(torch.bfloat16) @ self.rht_matrix
+                ).reshape(original_shape)
+            if self.row_scaled_nvfp4:
+                amin, amax = calibration_input.aminmax(dim=-1)
+            else:
+                amin, amax = calibration_input.aminmax()
+            observed_amax = torch.maximum(-amin, amax).reshape(-1).float()
+            if self.with_amax_reduction and torch.distributed.is_initialized():
+                torch.distributed.all_reduce(
+                    observed_amax,
+                    op=torch.distributed.ReduceOp.MAX,
+                    group=self._canonicalized_amax_reduction_group(),
+                )
+        self._update_calibration_value(metadata_name, observed_amax, decay=decay)
+
+    def get_quantization_recipe_name(self) -> str:
+        """Get the stable name of the quantization recipe."""
+        return "nvfp4_rowwise" if self.row_scaled_nvfp4 else "nvfp4"
 
     def _canonicalized_amax_reduction_group(self) -> dist_group_type:
         """Get process group for amax reduction"""
diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py
index f0b9898612..fef222dcff 100644
--- a/transformer_engine/pytorch/tensor/utils.py
+++ b/transformer_engine/pytorch/tensor/utils.py
@@ -32,21 +32,7 @@ def get_quantization_recipe_name(quantizer: Optional[Quantizer]) -> str:
     quantizer = getattr(quantizer, "parent_quantizer", quantizer)
     if quantizer is None:
         return ""
-    if isinstance(quantizer, Float8Quantizer):
-        return "fp8_delayed_scaling"
-    if isinstance(quantizer, Float8CurrentScalingQuantizer):
-        return "fp8_current_scaling"
-    if isinstance(quantizer, MXFP8Quantizer):
-        return "mxfp8"
-    if isinstance(quantizer, Float8BlockQuantizer):
-        return "fp8_block_scaling"
-    if isinstance(quantizer, NVFP4Quantizer):
-        if quantizer.row_scaled_nvfp4:
-            return "nvfp4_rowwise"
-        return "nvfp4"
-    # Custom recipes may provide arbitrary Quantizer implementations without a
-    # stable recipe name or globally checkpointable scaling metadata.
-    return ""
+    return quantizer.get_quantization_recipe_name()
 
 
 def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor):

From 39da2806a3f433a0b878e79b101dfa671d16b79a Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 11 Sep 2026 19:18:32 +0000
Subject: [PATCH 10/16] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ...test_ptq_calibration_metadata_buffering.py | 46 ++++++-------------
 .../pytorch/module/grouped_linear.py          |  4 +-
 .../pytorch/module/layernorm_linear.py        | 16 ++-----
 transformer_engine/pytorch/module/linear.py   | 12 ++---
 4 files changed, 21 insertions(+), 57 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index a83090ab77..7ad71ec659 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -191,9 +191,7 @@ def test_calibration_config_registers_module_scaling_factor_buffers(
             module(inp)
 
     buffers = {
-        name: value
-        for name, value in module.named_buffers()
-        if name.endswith("_te_ptq_calibrated")
+        name: value for name, value in module.named_buffers() if name.endswith("_te_ptq_calibrated")
     }
     assert len(buffers) == expected_buffer_count
     assert all(torch.isfinite(value).all() for value in buffers.values())
@@ -211,13 +209,10 @@ def test_calibrating_boolean_registers_scaling_factor_buffers():
     ):
         module(inp)
 
-    assert len(
-        [
-            name
-            for name, _ in module.named_buffers()
-            if name.endswith("_te_ptq_calibrated")
-        ]
-    ) == 2
+    assert (
+        len([name for name, _ in module.named_buffers() if name.endswith("_te_ptq_calibrated")])
+        == 2
+    )
 
 
 @pytest.mark.parametrize(
@@ -229,9 +224,7 @@ def test_calibrating_boolean_registers_scaling_factor_buffers():
         ("nvfp4_rowwise", "amax_rowwise", 1344.0),
     ),
 )
-def test_scale_buffer_info_selects_recipe_metadata(
-    recipe, metadata_name, expected_value
-):
+def test_scale_buffer_info_selects_recipe_metadata(recipe, metadata_name, expected_value):
     tensor = SimpleNamespace(
         _scale_inv=torch.tensor([0.25], dtype=torch.float32),
         _amax_rowwise=torch.tensor([2688.0 if recipe == "nvfp4" else 1344.0], dtype=torch.float32),
@@ -279,10 +272,7 @@ def test_resolve_calibration_quantizer_prefers_tensor_owner_and_unwraps_parent()
     tensor_quantizer = SimpleNamespace(parent_quantizer=parent_quantizer)
     tensor = SimpleNamespace(_quantizer=tensor_quantizer)
 
-    assert (
-        _common._resolve_calibration_quantizer(tensor, object())
-        is parent_quantizer
-    )
+    assert _common._resolve_calibration_quantizer(tensor, object()) is parent_quantizer
 
 
 def test_quantizer_calibration_state_is_keyed_by_quantized_metadata():
@@ -341,9 +331,7 @@ def test_grouped_scale_buffers_are_per_gemm():
         torch.tensor([0.5]),
     )
     assert (
-        scale_buffers[
-            "input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
-        ]
+        scale_buffers["input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"]
         is input_quantizers[1]._calibration_state["scale_inv"]
     )
 
@@ -362,9 +350,7 @@ def test_grouped_calibration_applies_decay_only_to_activations():
         activation_scale_decay=0.5,
     )
 
-    torch.testing.assert_close(
-        input_quantizer._calibration_state["scale_inv"], torch.tensor([2.0])
-    )
+    torch.testing.assert_close(input_quantizer._calibration_state["scale_inv"], torch.tensor([2.0]))
     torch.testing.assert_close(
         weight_quantizer._calibration_state["scale_inv"], torch.tensor([1.0])
     )
@@ -463,9 +449,7 @@ def test_nan_activation_scale_does_not_update_buffer(activation_scale_decay, ini
         assert not result
         assert not quantizer._calibration_state
     else:
-        value = result[
-            "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
-        ]
+        value = result["fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"]
         torch.testing.assert_close(value, torch.tensor([initial_scale]))
 
 
@@ -483,9 +467,7 @@ def test_current_scaling_calibrates_from_high_precision_tensor():
 
 @pytest.mark.parametrize("input_value", (0.0, float("inf")))
 @pytest.mark.parametrize("force_pow_2_scales", (False, True))
-def test_current_scaling_calibration_handles_non_finite_scale(
-    input_value, force_pow_2_scales
-):
+def test_current_scaling_calibration_handles_non_finite_scale(input_value, force_pow_2_scales):
     quantizer = Float8CurrentScalingQuantizer(
         fp8_dtype=DType.kFloat8E4M3,
         device=torch.device("cpu"),
@@ -522,8 +504,7 @@ def test_delayed_scaling_high_precision_calibration_matches_quantization(fp8_dty
     )
     calibration_quantizer.calibrate(tensor)
     assert (
-        calibration_quantizer.copy()._calibration_state
-        is calibration_quantizer._calibration_state
+        calibration_quantizer.copy()._calibration_state is calibration_quantizer._calibration_state
     )
 
     torch.testing.assert_close(
@@ -610,8 +591,7 @@ def test_nvfp4_high_precision_calibration_matches_quantization(
     calibration_quantizer = NVFP4Quantizer(**quantizer_kwargs)
     calibration_quantizer.calibrate(tensor)
     assert (
-        calibration_quantizer.copy()._calibration_state
-        is calibration_quantizer._calibration_state
+        calibration_quantizer.copy()._calibration_state is calibration_quantizer._calibration_state
     )
 
     expected_metadata_name = "amax_rowwise" if row_scaled_nvfp4 else "amax"
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index e01580ce19..42f9bf88f9 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -172,9 +172,7 @@ def _update_grouped_scale_buffers(
     weight_scale_updates = {}
     for index, tensor in enumerate(weight_tensors):
         quantizer = _resolve_calibration_quantizer(tensor, weight_quantizers[index])
-        weight_scale_updates.update(
-            _get_scale_buffer_info(f"weight_gemm{index}", quantizer)
-        )
+        weight_scale_updates.update(_get_scale_buffer_info(f"weight_gemm{index}", quantizer))
     scale_buffers.update(activation_scale_updates)
     scale_buffers.update(weight_scale_updates)
 
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 9b72c5056d..39758531a5 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -383,12 +383,8 @@ def forward(
             bias_dtype = torch.bfloat16
         bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias
 
-        input_calibration_quantizer = _resolve_calibration_quantizer(
-            ln_out_total, input_quantizer
-        )
-        weight_calibration_quantizer = _resolve_calibration_quantizer(
-            weightmat, weight_quantizer
-        )
+        input_calibration_quantizer = _resolve_calibration_quantizer(ln_out_total, input_quantizer)
+        weight_calibration_quantizer = _resolve_calibration_quantizer(weightmat, weight_quantizer)
 
         # Calibrate quantizers if needed.
         if scale_buffers is not None:
@@ -402,12 +398,8 @@ def forward(
 
         # Buffer scaling metadata only when requested.
         if scale_buffers is not None:
-            scale_buffers.update(
-                _get_scale_buffer_info("input", input_calibration_quantizer)
-            )
-            scale_buffers.update(
-                _get_scale_buffer_info("weight", weight_calibration_quantizer)
-            )
+            scale_buffers.update(_get_scale_buffer_info("input", input_calibration_quantizer))
+            scale_buffers.update(_get_scale_buffer_info("weight", weight_calibration_quantizer))
 
         # Choose whether to use GEMM kernel with split accumulator
         use_split_accumulator = _2X_ACC_FPROP
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 7a08b2ffc3..303df94ed3 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -608,9 +608,7 @@ def _linear_forward_impl(
         bias_dtype = torch.bfloat16
     bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias
 
-    input_calibration_quantizer = _resolve_calibration_quantizer(
-        inputmat_total, input_quantizer
-    )
+    input_calibration_quantizer = _resolve_calibration_quantizer(inputmat_total, input_quantizer)
     weight_calibration_quantizer = _resolve_calibration_quantizer(weightmat, weight_quantizer)
 
     # Calibrate quantizers if needed.
@@ -625,12 +623,8 @@ def _linear_forward_impl(
 
     # Buffer scaling metadata only when requested.
     if args.scale_buffers is not None:
-        args.scale_buffers.update(
-            _get_scale_buffer_info("input", input_calibration_quantizer)
-        )
-        args.scale_buffers.update(
-            _get_scale_buffer_info("weight", weight_calibration_quantizer)
-        )
+        args.scale_buffers.update(_get_scale_buffer_info("input", input_calibration_quantizer))
+        args.scale_buffers.update(_get_scale_buffer_info("weight", weight_calibration_quantizer))
 
     # Choose whether to use GEMM kernel with split accumulator
     use_split_accumulator = _2X_ACC_FPROP

From b00f1f88d1e5ab816fa4331def7ed49e4baa68a2 Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Fri, 11 Sep 2026 12:58:49 -0700
Subject: [PATCH 11/16] Tweaks.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 ...test_ptq_calibration_metadata_buffering.py | 239 +++++++++++++-----
 tests/pytorch/test_torch_compile.py           |  17 +-
 .../debug/pytorch/debug_quantization.py       |   3 +-
 transformer_engine/pytorch/module/_common.py  |  20 +-
 .../pytorch/module/grouped_linear.py          |  73 +++---
 .../pytorch/module/layernorm_linear.py        |  45 ++--
 .../pytorch/module/layernorm_mlp.py           |  91 +++----
 transformer_engine/pytorch/module/linear.py   |  57 +++--
 transformer_engine/pytorch/quantization.py    |  10 +-
 .../pytorch/quantized_tensor.py               |  25 +-
 .../pytorch/tensor/float8_blockwise_tensor.py |   2 +-
 .../pytorch/tensor/float8_tensor.py           |  26 +-
 .../pytorch/tensor/hybrid_tensor.py           |   4 +
 .../pytorch/tensor/identity_tensor.py         |   2 +-
 .../pytorch/tensor/mxfp8_tensor.py            |   2 +-
 .../pytorch/tensor/nvfp4_tensor.py            |  21 +-
 16 files changed, 407 insertions(+), 230 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index 7ad71ec659..b1f9b28cbf 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -9,7 +9,7 @@
 import pytest
 import torch
 
-from transformer_engine.common.recipe import Float8CurrentScaling
+from transformer_engine.common.recipe import DelayedScaling, Float8CurrentScaling
 from transformer_engine.pytorch import is_fp8_available, is_nvfp4_available
 from transformer_engine.pytorch.constants import DType
 from transformer_engine.pytorch.graph import make_graphed_callables
@@ -24,12 +24,14 @@
     autocast,
     fp8_autocast,
 )
-from transformer_engine.pytorch.quantized_tensor import Quantizer
+from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer
 from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer
 from transformer_engine.pytorch.tensor.float8_tensor import (
     Float8CurrentScalingQuantizer,
     Float8Quantizer,
 )
+from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer
+from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer
 from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
 from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
 from transformer_engine.pytorch.tensor.utils import get_quantization_recipe_name
@@ -44,6 +46,21 @@ def reset_quantization_state():
     FP8GlobalStateManager.reset()
 
 
+def _make_test_quantizer(quantizer_cls):
+    """Construct a quantizer without allocating recipe-specific CUDA state."""
+    quantizer = object.__new__(quantizer_cls)
+    Quantizer.__init__(quantizer, rowwise=True, columnwise=False)
+    return quantizer
+
+
+def _make_test_quantized_storage(**metadata):
+    """Construct bare quantized storage carrying only the requested metadata."""
+    storage = QuantizedTensorStorage()
+    for name, value in metadata.items():
+        setattr(storage, name, value)
+    return storage
+
+
 def test_calibration_api_additions_preserve_existing_parameter_order():
     state_fields = [field.name for field in dataclasses.fields(FP8GlobalState)]
     assert state_fields[-1] == "calibration_config"
@@ -101,7 +118,7 @@ def test_calibrating_argument_enables_default_calibration_config():
 
 
 def test_explicit_calibration_config_is_active_in_autocast():
-    config = QuantizationCalibrationConfig(activation_scale_decay=0.5)
+    config = QuantizationCalibrationConfig(transformer_engine_calibration_decay=0.5)
     with autocast(enabled=False, calibration_config=config):
         assert FP8GlobalStateManager.quantization_state.fp8_calibration
         assert FP8GlobalStateManager.get_calibration_config() is config
@@ -109,7 +126,7 @@ def test_explicit_calibration_config_is_active_in_autocast():
 
 
 def test_nested_autocast_restores_custom_calibration_config():
-    config = QuantizationCalibrationConfig(activation_scale_decay=0.5)
+    config = QuantizationCalibrationConfig(transformer_engine_calibration_decay=0.5)
     with autocast(enabled=False, calibration_config=config):
         with autocast(enabled=False):
             assert FP8GlobalStateManager.get_calibration_config() is None
@@ -137,7 +154,7 @@ def test_autocast_enter_preserves_calibrating_boolean_api():
 
 
 def test_calibrating_argument_accepts_explicit_calibration_config():
-    config = QuantizationCalibrationConfig(activation_scale_decay=0.5)
+    config = QuantizationCalibrationConfig(transformer_engine_calibration_decay=0.5)
     with autocast(
         enabled=False,
         calibrating=True,
@@ -147,8 +164,11 @@ def test_calibrating_argument_accepts_explicit_calibration_config():
 
 
 def test_calibration_config_rejects_negative_decay():
-    with pytest.raises(ValueError, match="activation_scale_decay must be non-negative"):
-        QuantizationCalibrationConfig(activation_scale_decay=-0.1)
+    with pytest.raises(
+        ValueError,
+        match="transformer_engine_calibration_decay must be non-negative",
+    ):
+        QuantizationCalibrationConfig(transformer_engine_calibration_decay=-0.1)
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
@@ -197,6 +217,65 @@ def test_calibration_config_registers_module_scaling_factor_buffers(
     assert all(torch.isfinite(value).all() for value in buffers.values())
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_linear_calibration_config_applies_activation_decay_only():
+    module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
+    calibration_config = QuantizationCalibrationConfig(
+        transformer_engine_calibration_decay=0.5
+    )
+    buffer_suffix = "_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+
+    with autocast(
+        enabled=False,
+        recipe=Float8CurrentScaling(),
+        calibration_config=calibration_config,
+    ):
+        module(torch.full((16, 32), 448.0, dtype=torch.bfloat16, device="cuda"))
+
+    torch.testing.assert_close(
+        module.get_buffer(f"input{buffer_suffix}"), torch.ones(1, device="cuda")
+    )
+    weight_scale = module.get_buffer(f"weight{buffer_suffix}").clone()
+
+    with autocast(
+        enabled=False,
+        recipe=Float8CurrentScaling(),
+        calibration_config=calibration_config,
+    ):
+        module(torch.full((16, 32), 112.0, dtype=torch.bfloat16, device="cuda"))
+
+    torch.testing.assert_close(
+        module.get_buffer(f"input{buffer_suffix}"), torch.full((1,), 0.5, device="cuda")
+    )
+    torch.testing.assert_close(module.get_buffer(f"weight{buffer_suffix}"), weight_scale)
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_linear_calibration_config_buffers_delayed_scaling_amax():
+    module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
+
+    with autocast(
+        enabled=False,
+        recipe=DelayedScaling(),
+        calibration_config=QuantizationCalibrationConfig(),
+    ):
+        module(torch.full((16, 32), 2.0, dtype=torch.bfloat16, device="cuda"))
+
+    calibration_buffers = {
+        name: value
+        for name, value in module.named_buffers()
+        if name.endswith("_te_ptq_calibrated")
+    }
+    assert set(calibration_buffers) == {
+        "input_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated",
+        "weight_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated",
+    }
+    torch.testing.assert_close(
+        calibration_buffers["input_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated"],
+        torch.full((1,), 2.0, device="cuda"),
+    )
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_calibrating_boolean_registers_scaling_factor_buffers():
     module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
@@ -225,23 +304,23 @@ def test_calibrating_boolean_registers_scaling_factor_buffers():
     ),
 )
 def test_scale_buffer_info_selects_recipe_metadata(recipe, metadata_name, expected_value):
-    tensor = SimpleNamespace(
+    tensor = _make_test_quantized_storage(
         _scale_inv=torch.tensor([0.25], dtype=torch.float32),
         _amax_rowwise=torch.tensor([2688.0 if recipe == "nvfp4" else 1344.0], dtype=torch.float32),
     )
     if recipe == "fp8_delayed_scaling":
-        quantizer = object.__new__(Float8Quantizer)
+        quantizer = _make_test_quantizer(Float8Quantizer)
         quantizer.amax = torch.tensor([0.0], dtype=torch.float32)
         tensor = torch.tensor([expected_value])
     elif recipe == "fp8_current_scaling":
-        quantizer = object.__new__(Float8CurrentScalingQuantizer)
+        quantizer = _make_test_quantizer(Float8CurrentScalingQuantizer)
     else:
-        quantizer = object.__new__(NVFP4Quantizer)
+        quantizer = _make_test_quantizer(NVFP4Quantizer)
         quantizer.row_scaled_nvfp4 = recipe == "nvfp4_rowwise"
 
     assert get_quantization_recipe_name(quantizer) == recipe
     quantizer.calibrate(tensor)
-    buffers = _common._get_scale_buffer_info("input", quantizer)
+    buffers = _common._get_calibration_metadata_buffers("input", quantizer)
     buffer_name = f"input_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
     value = buffers[buffer_name]
 
@@ -253,18 +332,28 @@ def test_scale_buffer_info_selects_recipe_metadata(recipe, metadata_name, expect
 def test_scale_buffer_info_skips_non_global_scaling_recipes(recipe):
     tensor = SimpleNamespace(_rowwise_scale_inv=torch.ones(2, 2))
     quantizer_cls = MXFP8Quantizer if recipe == "mxfp8" else Float8BlockQuantizer
-    quantizer = object.__new__(quantizer_cls)
+    quantizer = _make_test_quantizer(quantizer_cls)
 
     assert get_quantization_recipe_name(quantizer) == recipe
     quantizer.calibrate(tensor)
-    assert not _common._get_scale_buffer_info("input", quantizer)
+    assert not _common._get_calibration_metadata_buffers("input", quantizer)
 
 
 def test_custom_quantizer_defaults_to_no_calibration_metadata():
     quantizer = Quantizer(rowwise=True, columnwise=False)
 
     assert get_quantization_recipe_name(quantizer) == ""
-    assert not _common._get_scale_buffer_info("input", quantizer)
+    assert not _common._get_calibration_metadata_buffers("input", quantizer)
+
+
+def test_hybrid_quantizer_rejects_calibration():
+    quantizer = HybridQuantizer(
+        rowwise_quantizer=IdentityQuantizer(),
+        columnwise_quantizer=IdentityQuantizer(),
+    )
+
+    with pytest.raises(NotImplementedError, match="not yet supported for HybridQuantizer"):
+        quantizer.calibrate(torch.ones(1))
 
 
 def test_resolve_calibration_quantizer_prefers_tensor_owner_and_unwraps_parent():
@@ -278,76 +367,80 @@ def test_resolve_calibration_quantizer_prefers_tensor_owner_and_unwraps_parent()
 def test_quantizer_calibration_state_is_keyed_by_quantized_metadata():
     quantizer = Quantizer(rowwise=True, columnwise=False)
 
-    quantizer._update_calibration_value("amax", torch.tensor([2.0]), decay=0.0)
-    quantizer._update_calibration_value("scale_inv", torch.tensor([0.5]), decay=0.0)
+    quantizer._update_calibration_value(
+        "amax", torch.tensor([2.0]), calibration_decay=0.0
+    )
+    quantizer._update_calibration_value(
+        "scale_inv", torch.tensor([0.5]), calibration_decay=0.0
+    )
 
     assert set(quantizer._calibration_state) == {"amax", "scale_inv"}
     torch.testing.assert_close(quantizer._calibration_state["amax"], torch.tensor([2.0]))
     torch.testing.assert_close(quantizer._calibration_state["scale_inv"], torch.tensor([0.5]))
 
 
-def test_grouped_scale_buffers_are_per_gemm():
+def test_grouped_calibration_metadata_buffers_are_per_gemm():
     inputs = [
-        SimpleNamespace(_scale_inv=torch.tensor([0.25])),
-        SimpleNamespace(_scale_inv=torch.tensor([0.5])),
+        _make_test_quantized_storage(_scale_inv=torch.tensor([0.25])),
+        _make_test_quantized_storage(_scale_inv=torch.tensor([0.5])),
     ]
     weights = [
-        SimpleNamespace(_scale_inv=torch.tensor([0.75])),
-        SimpleNamespace(_scale_inv=torch.tensor([1.0])),
+        _make_test_quantized_storage(_scale_inv=torch.tensor([0.75])),
+        _make_test_quantized_storage(_scale_inv=torch.tensor([1.0])),
     ]
     input_quantizers = [
-        object.__new__(Float8CurrentScalingQuantizer),
-        object.__new__(Float8CurrentScalingQuantizer),
+        _make_test_quantizer(Float8CurrentScalingQuantizer),
+        _make_test_quantizer(Float8CurrentScalingQuantizer),
     ]
     weight_quantizers = [
-        object.__new__(Float8CurrentScalingQuantizer),
-        object.__new__(Float8CurrentScalingQuantizer),
+        _make_test_quantizer(Float8CurrentScalingQuantizer),
+        _make_test_quantizer(Float8CurrentScalingQuantizer),
     ]
-    scale_buffers = {}
+    calibration_buffers = {}
 
     grouped_linear._calibrate_grouped_tensors(
         inputs,
         weights,
         input_quantizers,
         weight_quantizers,
-        activation_scale_decay=0.0,
+        transformer_engine_calibration_decay=0.0,
     )
-    grouped_linear._update_grouped_scale_buffers(
-        scale_buffers,
+    grouped_linear._update_grouped_calibration_metadata_buffers(
+        calibration_buffers,
         inputs,
         weights,
         input_quantizers,
         weight_quantizers,
     )
 
-    assert set(scale_buffers) == {
+    assert set(calibration_buffers) == {
         "input_gemm0_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
         "input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
         "weight_gemm0_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
         "weight_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated",
     }
     torch.testing.assert_close(
-        scale_buffers["input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"],
+        calibration_buffers["input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"],
         torch.tensor([0.5]),
     )
     assert (
-        scale_buffers["input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"]
+        calibration_buffers["input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"]
         is input_quantizers[1]._calibration_state["scale_inv"]
     )
 
 
 def test_grouped_calibration_applies_decay_only_to_activations():
-    input_quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    input_quantizer = _make_test_quantizer(Float8CurrentScalingQuantizer)
     input_quantizer._calibration_state = {"scale_inv": torch.tensor([4.0])}
-    weight_quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    weight_quantizer = _make_test_quantizer(Float8CurrentScalingQuantizer)
     weight_quantizer._calibration_state = {"scale_inv": torch.tensor([4.0])}
 
     grouped_linear._calibrate_grouped_tensors(
-        [SimpleNamespace(_scale_inv=torch.tensor([1.0]))],
-        [SimpleNamespace(_scale_inv=torch.tensor([1.0]))],
+        [_make_test_quantized_storage(_scale_inv=torch.tensor([1.0]))],
+        [_make_test_quantized_storage(_scale_inv=torch.tensor([1.0]))],
         [input_quantizer],
         [weight_quantizer],
-        activation_scale_decay=0.5,
+        transformer_engine_calibration_decay=0.5,
     )
 
     torch.testing.assert_close(input_quantizer._calibration_state["scale_inv"], torch.tensor([2.0]))
@@ -356,23 +449,41 @@ def test_grouped_calibration_applies_decay_only_to_activations():
     )
 
 
-def test_grouped_scale_buffers_use_per_gemm_delayed_scaling_amax():
+def test_calibration_supports_legacy_custom_quantizer_signature():
+    class LegacyCustomQuantizer:
+        def calibrate(self, tensor):
+            self.observed_tensor = tensor
+
+    tensor = torch.tensor([1.0])
+    quantizer = LegacyCustomQuantizer()
+    grouped_linear._calibrate_grouped_tensors(
+        [tensor],
+        [],
+        [quantizer],
+        [],
+        transformer_engine_calibration_decay=0.5,
+    )
+
+    assert quantizer.observed_tensor is tensor
+
+
+def test_grouped_calibration_metadata_uses_per_gemm_delayed_scaling_amax():
     quantizers = []
     for amax in (1.0, 2.0):
-        quantizer = object.__new__(Float8Quantizer)
+        quantizer = _make_test_quantizer(Float8Quantizer)
         quantizer.amax = torch.tensor([amax])
         quantizers.append(quantizer)
-    scale_buffers = {}
+    calibration_buffers = {}
 
     grouped_linear._calibrate_grouped_tensors(
         [torch.tensor([1.0]), torch.tensor([2.0])],
         [torch.tensor([1.0]), torch.tensor([2.0])],
         quantizers,
         quantizers,
-        activation_scale_decay=0.0,
+        transformer_engine_calibration_decay=0.0,
     )
-    grouped_linear._update_grouped_scale_buffers(
-        scale_buffers,
+    grouped_linear._update_grouped_calibration_metadata_buffers(
+        calibration_buffers,
         [torch.tensor([1.0]), torch.tensor([2.0])],
         [torch.tensor([1.0]), torch.tensor([2.0])],
         quantizers,
@@ -380,11 +491,11 @@ def test_grouped_scale_buffers_use_per_gemm_delayed_scaling_amax():
     )
 
     torch.testing.assert_close(
-        scale_buffers["input_gemm0_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated"],
+        calibration_buffers["input_gemm0_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated"],
         torch.tensor([1.0]),
     )
     torch.testing.assert_close(
-        scale_buffers["input_gemm1_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated"],
+        calibration_buffers["input_gemm1_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated"],
         torch.tensor([2.0]),
     )
 
@@ -400,14 +511,14 @@ def test_grouped_scale_buffers_use_per_gemm_delayed_scaling_amax():
 )
 def test_activation_scale_buffer_uses_decaying_maximum(observed_scale, expected_scale):
     name = "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
-    quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    quantizer = _make_test_quantizer(Float8CurrentScalingQuantizer)
     initial_buffer = torch.tensor([4.0])
     quantizer._calibration_state = {"scale_inv": initial_buffer}
     quantizer.calibrate(
-        SimpleNamespace(_scale_inv=torch.tensor([observed_scale])),
-        decay=0.5,
+        _make_test_quantized_storage(_scale_inv=torch.tensor([observed_scale])),
+        calibration_decay=0.5,
     )
-    buffers = _common._get_scale_buffer_info("fc1_input", quantizer)
+    buffers = _common._get_calibration_metadata_buffers("fc1_input", quantizer)
     value = buffers[name]
 
     torch.testing.assert_close(value, torch.tensor([expected_scale]))
@@ -417,33 +528,41 @@ def test_activation_scale_buffer_uses_decaying_maximum(observed_scale, expected_
 
 def test_zero_decay_keeps_observed_metadata_reference():
     observed_scale = torch.tensor([2.0])
-    quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    quantizer = _make_test_quantizer(Float8CurrentScalingQuantizer)
 
-    quantizer.calibrate(SimpleNamespace(_scale_inv=observed_scale), decay=0.0)
+    quantizer.calibrate(
+        _make_test_quantized_storage(_scale_inv=observed_scale),
+        calibration_decay=0.0,
+    )
 
     assert quantizer._calibration_state["scale_inv"] is observed_scale
 
 
 def test_decaying_calibration_rejects_metadata_shape_change():
-    quantizer = object.__new__(Float8CurrentScalingQuantizer)
+    quantizer = _make_test_quantizer(Float8CurrentScalingQuantizer)
     quantizer._calibration_state = {"scale_inv": torch.ones(1)}
 
     with pytest.raises(RuntimeError, match="calibration value shape changed"):
-        quantizer.calibrate(SimpleNamespace(_scale_inv=torch.ones(2)), decay=0.5)
+        quantizer.calibrate(
+            _make_test_quantized_storage(_scale_inv=torch.ones(2)),
+            calibration_decay=0.5,
+        )
 
 
-@pytest.mark.parametrize("activation_scale_decay", (0.0, 0.5))
+@pytest.mark.parametrize("transformer_engine_calibration_decay", (0.0, 0.5))
 @pytest.mark.parametrize("initial_scale", (None, 4.0))
-def test_nan_activation_scale_does_not_update_buffer(activation_scale_decay, initial_scale):
-    quantizer = object.__new__(Float8CurrentScalingQuantizer)
+def test_nan_activation_scale_does_not_update_buffer(
+    transformer_engine_calibration_decay, initial_scale
+):
+    quantizer = _make_test_quantizer(Float8CurrentScalingQuantizer)
     if initial_scale is not None:
         quantizer._calibration_state = {"scale_inv": torch.tensor([initial_scale])}
 
     quantizer.calibrate(
-        SimpleNamespace(_scale_inv=torch.tensor([float("nan")])),
-        decay=activation_scale_decay,
+        _make_test_quantized_storage(_scale_inv=torch.tensor([float("nan")])),
+        calibration_decay=transformer_engine_calibration_decay,
     )
-    result = _common._get_scale_buffer_info("fc1_input", quantizer)
+    result = _common._get_calibration_metadata_buffers("fc1_input", quantizer)
 
     if initial_scale is None:
         assert not result
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 5ace9253f4..456c83b186 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2066,7 +2066,7 @@ def fn(inp):
             ):
                 return model(inp)
 
-        return model, fn, "bwd", None, "quantized scaling-factor buffering"
+        return model, fn, "bwd", None, "Transformer Engine calibration metadata buffering"
     raise ValueError(case)
 
 
@@ -2125,6 +2125,21 @@ def make_inp():
                 model.weight.grad, model_ref.weight.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL
             )
     torch.testing.assert_close(out.detach(), out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
+    if case == "scale_buffering":
+        calibration_buffers_ref = {
+            name: value
+            for name, value in model_ref.named_buffers()
+            if name.endswith("_te_ptq_calibrated")
+        }
+        calibration_buffers = {
+            name: value
+            for name, value in model.named_buffers()
+            if name.endswith("_te_ptq_calibrated")
+        }
+        assert calibration_buffers
+        assert calibration_buffers.keys() == calibration_buffers_ref.keys()
+        for name, value in calibration_buffers.items():
+            torch.testing.assert_close(value, calibration_buffers_ref[name])
 
     torch._dynamo.reset()
     compiled_fg = torch.compile(fn, fullgraph=True)
diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py
index 873da92d7a..552209ad9a 100644
--- a/transformer_engine/debug/pytorch/debug_quantization.py
+++ b/transformer_engine/debug/pytorch/debug_quantization.py
@@ -429,9 +429,8 @@ def any_feature_enabled(self) -> bool:
                 return True
         return False
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0):
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0):
         """Calibration override, should not be invoked."""
-        del tensor, decay
         raise RuntimeError("[NVTORCH-INSPECT ERROR] Calibration with debug is not supported")
 
     def update_quantized(
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index b4ef1619fe..a5718f5c05 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -5,6 +5,8 @@
 """Internal function used by multiple modules."""
 
 import dataclasses
+import functools
+import inspect
 import queue
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -17,13 +19,27 @@
 from ..utils import get_default_init_method
 
 
+@functools.lru_cache(maxsize=None)
+def _supports_calibration_decay(quantizer_type: type) -> bool:
+    """Whether a quantizer's calibrate override accepts calibration_decay."""
+    try:
+        parameters = inspect.signature(quantizer_type.calibrate).parameters
+    except (TypeError, ValueError):
+        return False
+    return "calibration_decay" in parameters or any(
+        parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
+    )
+
+
 def _resolve_calibration_quantizer(tensor: Any, quantizer: Any) -> Any:
     """Get the quantizer that owns calibration state for a tensor."""
     quantizer = getattr(tensor, "_quantizer", None) or quantizer
     return getattr(quantizer, "parent_quantizer", quantizer)
 
 
-def _get_scale_buffer_info(tensor_name: str, quantizer: Any) -> Dict[str, torch.Tensor]:
+def _get_calibration_metadata_buffers(
+    tensor_name: str, quantizer: Any
+) -> Dict[str, torch.Tensor]:
     """Get checkpoint-buffer aliases from quantizer calibration state."""
     if quantizer is None:
         return {}
@@ -33,7 +49,7 @@ def _get_scale_buffer_info(tensor_name: str, quantizer: Any) -> Dict[str, torch.
     return {
         # Standard naming for PTQ calibration data.
         f"{tensor_name}_tensor_{metadata_name}_{recipe}_te_ptq_calibrated": value
-        for metadata_name, value in getattr(quantizer, "_calibration_state", {}).items()
+        for metadata_name, value in quantizer._calibration_state.items()
     }
 
 
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 42f9bf88f9..4a5daecfc7 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -32,8 +32,9 @@
     _get_high_precision_init_val,
 )
 from ._common import (
-    _get_scale_buffer_info,
+    _get_calibration_metadata_buffers,
     _resolve_calibration_quantizer,
+    _supports_calibration_decay,
     can_reconstruct_wgrad_input_from_original,
     WeightGradStore,
 )
@@ -157,24 +158,24 @@ def is_module_grouped_tensor_path_supported(
     return False
 
 
-def _update_grouped_scale_buffers(
-    scale_buffers: Dict[str, Optional[torch.Tensor]],
+def _update_grouped_calibration_metadata_buffers(
+    calibration_buffers: Dict[str, Optional[torch.Tensor]],
     input_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
     weight_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
     input_quantizers: List[Optional[Quantizer]],
     weight_quantizers: List[Optional[Quantizer]],
 ) -> None:
     """Update GroupedLinear PTQ buffers from per-GEMM calibration state."""
-    activation_scale_updates = {}
     for index, tensor in enumerate(input_tensors):
         quantizer = _resolve_calibration_quantizer(tensor, input_quantizers[index])
-        activation_scale_updates.update(_get_scale_buffer_info(f"input_gemm{index}", quantizer))
-    weight_scale_updates = {}
+        calibration_buffers.update(
+            _get_calibration_metadata_buffers(f"input_gemm{index}", quantizer)
+        )
     for index, tensor in enumerate(weight_tensors):
         quantizer = _resolve_calibration_quantizer(tensor, weight_quantizers[index])
-        weight_scale_updates.update(_get_scale_buffer_info(f"weight_gemm{index}", quantizer))
-    scale_buffers.update(activation_scale_updates)
-    scale_buffers.update(weight_scale_updates)
+        calibration_buffers.update(
+            _get_calibration_metadata_buffers(f"weight_gemm{index}", quantizer)
+        )
 
 
 def _calibrate_grouped_tensors(
@@ -182,13 +183,19 @@ def _calibrate_grouped_tensors(
     weight_tensors: List[Union[torch.Tensor, QuantizedTensorStorage]],
     input_quantizers: List[Optional[Quantizer]],
     weight_quantizers: List[Optional[Quantizer]],
-    activation_scale_decay: float,
+    transformer_engine_calibration_decay: float,
 ) -> None:
     """Calibrate GroupedLinear input and weight quantizers once per GEMM."""
     for tensor, quantizer in zip(input_tensors, input_quantizers):
         quantizer = _resolve_calibration_quantizer(tensor, quantizer)
         if quantizer is not None:
-            quantizer.calibrate(tensor, decay=activation_scale_decay)
+            if _supports_calibration_decay(type(quantizer)):
+                quantizer.calibrate(
+                    tensor,
+                    calibration_decay=transformer_engine_calibration_decay,
+                )
+            else:
+                quantizer.calibrate(tensor)
     for tensor, quantizer in zip(weight_tensors, weight_quantizers):
         quantizer = _resolve_calibration_quantizer(tensor, quantizer)
         if quantizer is not None:
@@ -471,8 +478,8 @@ def _forward_grouped_tensor(
         save_original_input: bool,
         single_grouped_weight: bool,
         single_grouped_bias: bool,
-        scale_buffers: Optional[Dict[str, Optional[torch.Tensor]]],
-        quantized_scaling_factor_buffering_decay: float,
+        calibration_buffers: Optional[Dict[str, Optional[torch.Tensor]]],
+        transformer_engine_calibration_decay: float,
         weights: Tuple[torch.Tensor, ...],
         biases: Tuple[torch.Tensor, ...],
         out: Optional[torch.Tensor] = None,
@@ -581,7 +588,7 @@ def _forward_grouped_tensor(
             use_split_accumulator=use_split_accumulator,
         )
 
-        if scale_buffers is not None:
+        if calibration_buffers is not None:
             grouped_inputs = grouped_x.quantized_tensors
             if grouped_inputs is None:
                 grouped_inputs = grouped_x.split_into_quantized_tensors()
@@ -596,10 +603,10 @@ def _forward_grouped_tensor(
                 grouped_weights,
                 input_quantizers,
                 weight_quantizers,
-                quantized_scaling_factor_buffering_decay,
+                transformer_engine_calibration_decay,
             )
-            _update_grouped_scale_buffers(
-                scale_buffers,
+            _update_grouped_calibration_metadata_buffers(
+                calibration_buffers,
                 grouped_inputs,
                 grouped_weights,
                 input_quantizers,
@@ -727,8 +734,8 @@ def forward(
             single_grouped_weight,
             single_grouped_bias,
             use_grouped_tensor,
-            scale_buffers,
-            quantized_scaling_factor_buffering_decay,
+            calibration_buffers,
+            transformer_engine_calibration_decay,
         ) = non_tensor_args
         recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
         backward_override = recipe.backward_override if recipe is not None else None
@@ -881,8 +888,8 @@ def forward(
                 save_original_input=save_original_input,
                 single_grouped_weight=single_grouped_weight,
                 single_grouped_bias=single_grouped_bias,
-                scale_buffers=scale_buffers,
-                quantized_scaling_factor_buffering_decay=(quantized_scaling_factor_buffering_decay),
+                calibration_buffers=calibration_buffers,
+                transformer_engine_calibration_decay=transformer_engine_calibration_decay,
                 weights=weights,
                 biases=biases,
                 out=out,
@@ -929,19 +936,17 @@ def forward(
         else:
             weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]
 
-        if scale_buffers is not None:
+        # Calibrate quantizers and buffer their metadata when requested.
+        if calibration_buffers is not None:
             _calibrate_grouped_tensors(
                 inputmats,
                 weights_fp8,
                 input_quantizers,
                 weight_quantizers,
-                quantized_scaling_factor_buffering_decay,
+                transformer_engine_calibration_decay,
             )
-
-        # Buffer scaling metadata only when requested.
-        if scale_buffers is not None:
-            _update_grouped_scale_buffers(
-                scale_buffers,
+            _update_grouped_calibration_metadata_buffers(
+                calibration_buffers,
                 inputmats,
                 weights_fp8,
                 input_quantizers,
@@ -2361,9 +2366,9 @@ def forward(
                     else [None] * num_gemms
                 )
             calibration_config = FP8GlobalStateManager.get_calibration_config()
-            scale_buffers = None
+            calibration_buffers = None
             if calibration_config is not None:
-                scale_buffers = {
+                calibration_buffers = {
                     name: value
                     for name, value in self._buffers.items()
                     if name.endswith("_te_ptq_calibrated")
@@ -2394,9 +2399,9 @@ def forward(
                 self.single_grouped_weight,
                 use_grouped_bias,
                 self.use_grouped_tensor,
-                scale_buffers,
+                calibration_buffers,
                 (
-                    calibration_config.activation_scale_decay
+                    calibration_config.transformer_engine_calibration_decay
                     if calibration_config is not None
                     else 0.0
                 ),
@@ -2412,10 +2417,10 @@ def forward(
                 *bias_tensors,
             )
 
-            if scale_buffers is not None:
+            if calibration_buffers is not None:
                 # Assign scaling-factor calibration buffers to the model.
                 # Materializing a new buffer requires a CUDA graph warmup step.
-                for name, value in scale_buffers.items():
+                for name, value in calibration_buffers.items():
                     if value is not None:
                         if name in self._buffers:
                             setattr(self, name, value)
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 39758531a5..49cb6839f1 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -69,8 +69,9 @@
 from ..jit import no_torch_dynamo
 from ..graph import is_graph_capturing
 from ._common import (
-    _get_scale_buffer_info,
+    _get_calibration_metadata_buffers,
     _resolve_calibration_quantizer,
+    _supports_calibration_decay,
     apply_normalization,
     noop_cat,
     set_quantizer_amax_reduction_group,
@@ -165,8 +166,8 @@ def forward(
             symmetric_ar_type,
             debug,
             is_fsdp2,
-            quantized_scaling_factor_buffering_decay,
-            scale_buffers,
+            transformer_engine_calibration_decay,
+            calibration_buffers,
         ) = non_tensor_args
         if fp8:
             backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override
@@ -386,20 +387,24 @@ def forward(
         input_calibration_quantizer = _resolve_calibration_quantizer(ln_out_total, input_quantizer)
         weight_calibration_quantizer = _resolve_calibration_quantizer(weightmat, weight_quantizer)
 
-        # Calibrate quantizers if needed.
-        if scale_buffers is not None:
+        # Calibrate quantizers and buffer their metadata when requested.
+        if calibration_buffers is not None:
             if input_calibration_quantizer is not None:
-                input_calibration_quantizer.calibrate(
-                    ln_out_total,
-                    decay=quantized_scaling_factor_buffering_decay,
-                )
+                if _supports_calibration_decay(type(input_calibration_quantizer)):
+                    input_calibration_quantizer.calibrate(
+                        ln_out_total,
+                        calibration_decay=transformer_engine_calibration_decay,
+                    )
+                else:
+                    input_calibration_quantizer.calibrate(ln_out_total)
             if weight_calibration_quantizer is not None:
                 weight_calibration_quantizer.calibrate(weightmat)
-
-        # Buffer scaling metadata only when requested.
-        if scale_buffers is not None:
-            scale_buffers.update(_get_scale_buffer_info("input", input_calibration_quantizer))
-            scale_buffers.update(_get_scale_buffer_info("weight", weight_calibration_quantizer))
+            calibration_buffers.update(
+                _get_calibration_metadata_buffers("input", input_calibration_quantizer)
+            )
+            calibration_buffers.update(
+                _get_calibration_metadata_buffers("weight", weight_calibration_quantizer)
+            )
 
         # Choose whether to use GEMM kernel with split accumulator
         use_split_accumulator = _2X_ACC_FPROP
@@ -1784,9 +1789,9 @@ def forward(
             )
 
             calibration_config = FP8GlobalStateManager.get_calibration_config()
-            scale_buffers = None
+            calibration_buffers = None
             if calibration_config is not None:
-                scale_buffers = {
+                calibration_buffers = {
                     name: value
                     for name, value in self._buffers.items()
                     if name.endswith("_te_ptq_calibrated")
@@ -1833,11 +1838,11 @@ def forward(
                 debug,
                 self.is_fsdp2,
                 (
-                    calibration_config.activation_scale_decay
+                    calibration_config.transformer_engine_calibration_decay
                     if calibration_config is not None
                     else 0.0
                 ),
-                scale_buffers,
+                calibration_buffers,
             )
             out, ln_out, new_weight_workspace = fwd_fn(
                 *autograd_ctx,
@@ -1850,8 +1855,8 @@ def forward(
                 non_tensor_args,
             )
 
-            if scale_buffers is not None:
-                for name, value in scale_buffers.items():
+            if calibration_buffers is not None:
+                for name, value in calibration_buffers.items():
                     if value is not None:
                         if name in self._buffers:
                             setattr(self, name, value)
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 2f5438293c..1077a06ad2 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -76,8 +76,9 @@
 from ..tensor.hybrid_tensor import HybridQuantizer
 from ..tensor.identity_tensor import IdentityQuantizer
 from ._common import (
-    _get_scale_buffer_info,
+    _get_calibration_metadata_buffers,
     _resolve_calibration_quantizer,
+    _supports_calibration_decay,
     apply_normalization,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
@@ -249,8 +250,8 @@ def _forward(
             checkpoint,
             debug,
             is_fsdp2,
-            quantized_scaling_factor_buffering_decay,
-            scale_buffers,
+            transformer_engine_calibration_decay,
+            calibration_buffers,
             recompute_for_bwd,
         ) = non_tensor_args
         if fp8:
@@ -353,10 +354,8 @@ def _forward(
                 "checkpoint": checkpoint,
                 "debug": debug,
                 "is_fsdp2": is_fsdp2,
-                "quantized_scaling_factor_buffering_decay": (
-                    quantized_scaling_factor_buffering_decay
-                ),
-                "scale_buffers": scale_buffers,
+                "transformer_engine_calibration_decay": transformer_engine_calibration_decay,
+                "calibration_buffers": calibration_buffers,
                 "recompute_for_bwd": True,  # set this to true for recomputation phase
             }
         # Make sure input dimensions are compatible
@@ -575,24 +574,26 @@ def _forward(
             fc1_weight_final, fc1_weight_quantizer
         )
 
-        # Calibrate FC1 quantizers if needed.
-        should_calibrate = scale_buffers is not None
+        fc1_input_calibration_buffers = {}
+        fc1_weight_calibration_buffers = {}
+
+        # Calibrate FC1 quantizers and collect their metadata when requested.
+        should_calibrate = calibration_buffers is not None
         if should_calibrate:
             if fc1_input_calibration_quantizer is not None:
-                fc1_input_calibration_quantizer.calibrate(
-                    ln_out_total,
-                    decay=quantized_scaling_factor_buffering_decay,
-                )
+                if _supports_calibration_decay(type(fc1_input_calibration_quantizer)):
+                    fc1_input_calibration_quantizer.calibrate(
+                        ln_out_total,
+                        calibration_decay=transformer_engine_calibration_decay,
+                    )
+                else:
+                    fc1_input_calibration_quantizer.calibrate(ln_out_total)
             if fc1_weight_calibration_quantizer is not None:
                 fc1_weight_calibration_quantizer.calibrate(fc1_weight_final)
-
-        fc1_input_scale_buffer = {}
-        fc1_weight_scale_buffer = {}
-        if scale_buffers is not None:
-            fc1_input_scale_buffer = _get_scale_buffer_info(
+            fc1_input_calibration_buffers = _get_calibration_metadata_buffers(
                 "fc1_input", fc1_input_calibration_quantizer
             )
-            fc1_weight_scale_buffer = _get_scale_buffer_info(
+            fc1_weight_calibration_buffers = _get_calibration_metadata_buffers(
                 "fc1_weight", fc1_weight_calibration_quantizer
             )
 
@@ -690,10 +691,13 @@ def _forward(
             act_out, fc2_input_quantizer
         )
         if should_calibrate and fc2_input_calibration_quantizer is not None:
-            fc2_input_calibration_quantizer.calibrate(
-                act_out,
-                decay=quantized_scaling_factor_buffering_decay,
-            )
+            if _supports_calibration_decay(type(fc2_input_calibration_quantizer)):
+                fc2_input_calibration_quantizer.calibrate(
+                    act_out,
+                    calibration_decay=transformer_engine_calibration_decay,
+                )
+            else:
+                fc2_input_calibration_quantizer.calibrate(act_out)
 
         # we want to skip fc2 computation if we are checkpointing and recomputing,
         # otherwise we compute fc2
@@ -713,24 +717,21 @@ def _forward(
             fc2_weight_calibration_quantizer = _resolve_calibration_quantizer(
                 fc2_weight_final, fc2_weight_quantizer
             )
-            if should_calibrate and fc2_weight_calibration_quantizer is not None:
-                fc2_weight_calibration_quantizer.calibrate(fc2_weight_final)
-
-            if scale_buffers is not None:
-                activation_scale_updates = {}
-                weight_scale_updates = {}
-                activation_scale_updates.update(fc1_input_scale_buffer)
-                weight_scale_updates.update(fc1_weight_scale_buffer)
-                fc2_input_scale_buffer = _get_scale_buffer_info(
-                    "fc2_input", fc2_input_calibration_quantizer
+            if should_calibrate:
+                if fc2_weight_calibration_quantizer is not None:
+                    fc2_weight_calibration_quantizer.calibrate(fc2_weight_final)
+                calibration_buffers.update(fc1_input_calibration_buffers)
+                calibration_buffers.update(fc1_weight_calibration_buffers)
+                calibration_buffers.update(
+                    _get_calibration_metadata_buffers(
+                        "fc2_input", fc2_input_calibration_quantizer
+                    )
                 )
-                activation_scale_updates.update(fc2_input_scale_buffer)
-                fc2_weight_scale_buffer = _get_scale_buffer_info(
-                    "fc2_weight", fc2_weight_calibration_quantizer
+                calibration_buffers.update(
+                    _get_calibration_metadata_buffers(
+                        "fc2_weight", fc2_weight_calibration_quantizer
+                    )
                 )
-                weight_scale_updates.update(fc2_weight_scale_buffer)
-                scale_buffers.update(activation_scale_updates)
-                scale_buffers.update(weight_scale_updates)
 
             # Configure Userbuffers reduce-scatter if needed
             ub_obj_fc2out = None
@@ -2435,9 +2436,9 @@ def forward(
             )
 
             calibration_config = FP8GlobalStateManager.get_calibration_config()
-            scale_buffers = None
+            calibration_buffers = None
             if calibration_config is not None:
-                scale_buffers = {
+                calibration_buffers = {
                     name: value
                     for name, value in self._buffers.items()
                     if name.endswith("_te_ptq_calibrated")
@@ -2494,11 +2495,11 @@ def forward(
                 debug,
                 self.is_fsdp2,
                 (
-                    calibration_config.activation_scale_decay
+                    calibration_config.transformer_engine_calibration_decay
                     if calibration_config is not None
                     else 0.0
                 ),
-                scale_buffers,
+                calibration_buffers,
             )
             out, ln_out, new_fc1_ws, new_fc2_ws = fwd_fn(
                 *autograd_ctx,
@@ -2514,8 +2515,8 @@ def forward(
                 non_tensor_args,
             )
 
-            if scale_buffers is not None:
-                for name, value in scale_buffers.items():
+            if calibration_buffers is not None:
+                for name, value in calibration_buffers.items():
                     if value is not None:
                         if name in self._buffers:
                             setattr(self, name, value)
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 303df94ed3..9ce1541341 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -32,8 +32,9 @@
     _2X_ACC_WGRAD,
 )
 from ._common import (
-    _get_scale_buffer_info,
+    _get_calibration_metadata_buffers,
     _resolve_calibration_quantizer,
+    _supports_calibration_decay,
     can_reconstruct_wgrad_input_from_original,
     noop_cat,
     set_quantizer_amax_reduction_group,
@@ -181,9 +182,9 @@ class LinearFwdArgs:
     fuse_wgrad_accumulation: bool
     wgrad_store: Optional[Any]
 
-    # Inference Scaling Factor Calibration Buffering
-    scale_buffers: Optional[Dict[str, Optional[torch.Tensor]]]
-    quantized_scaling_factor_buffering_decay: float
+    # Transformer Engine calibration metadata buffering
+    calibration_buffers: Optional[Dict[str, Optional[torch.Tensor]]]
+    transformer_engine_calibration_decay: float
 
     # --- Misc ---
     cpu_offloading: bool
@@ -379,7 +380,7 @@ def _linear_forward_impl(
     ctx_attrs)``. ``new_weight_workspace`` is the freshly produced FP8 weight
     workspace (returned alongside ``out`` so the caller can refresh its
     cache). The last two are ``None`` when gradients are disabled.
-    Scaling-factor checkpoint buffers are updated through ``args.scale_buffers``.
+    Calibration metadata buffers are updated through ``args.calibration_buffers``.
     """
 
     weight = args.weight
@@ -611,20 +612,24 @@ def _linear_forward_impl(
     input_calibration_quantizer = _resolve_calibration_quantizer(inputmat_total, input_quantizer)
     weight_calibration_quantizer = _resolve_calibration_quantizer(weightmat, weight_quantizer)
 
-    # Calibrate quantizers if needed.
-    if args.scale_buffers is not None:
+    # Calibrate quantizers and buffer their metadata when requested.
+    if args.calibration_buffers is not None:
         if input_calibration_quantizer is not None:
-            input_calibration_quantizer.calibrate(
-                inputmat_total,
-                decay=args.quantized_scaling_factor_buffering_decay,
-            )
+            if _supports_calibration_decay(type(input_calibration_quantizer)):
+                input_calibration_quantizer.calibrate(
+                    inputmat_total,
+                    calibration_decay=args.transformer_engine_calibration_decay,
+                )
+            else:
+                input_calibration_quantizer.calibrate(inputmat_total)
         if weight_calibration_quantizer is not None:
             weight_calibration_quantizer.calibrate(weightmat)
-
-    # Buffer scaling metadata only when requested.
-    if args.scale_buffers is not None:
-        args.scale_buffers.update(_get_scale_buffer_info("input", input_calibration_quantizer))
-        args.scale_buffers.update(_get_scale_buffer_info("weight", weight_calibration_quantizer))
+        args.calibration_buffers.update(
+            _get_calibration_metadata_buffers("input", input_calibration_quantizer)
+        )
+        args.calibration_buffers.update(
+            _get_calibration_metadata_buffers("weight", weight_calibration_quantizer)
+        )
 
     # Choose whether to use GEMM kernel with split accumulator
     use_split_accumulator = _2X_ACC_FPROP
@@ -2497,9 +2502,9 @@ def forward(
             )
             wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None
             calibration_config = FP8GlobalStateManager.get_calibration_config()
-            scale_buffers = None
+            calibration_buffers = None
             if calibration_config is not None:
-                scale_buffers = {
+                calibration_buffers = {
                     name: value
                     for name, value in self._buffers.items()
                     if name.endswith("_te_ptq_calibrated")
@@ -2560,10 +2565,10 @@ def forward(
                 # weight-grad scheduling
                 fuse_wgrad_accumulation=self.fuse_wgrad_accumulation,
                 wgrad_store=wgrad_store,
-                # Inference Scaling Factor Calibration Buffering
-                scale_buffers=scale_buffers,
-                quantized_scaling_factor_buffering_decay=(
-                    calibration_config.activation_scale_decay
+                # Buffering TE calibration metadata, e.g. scaling factors.
+                calibration_buffers=calibration_buffers,
+                transformer_engine_calibration_decay=(
+                    calibration_config.transformer_engine_calibration_decay
                     if calibration_config is not None
                     else 0.0
                 ),
@@ -2590,10 +2595,10 @@ def forward(
                     weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled
                 )
 
-            if scale_buffers is not None:
-                # Assign the scaling factor calibration buffers to model.
+            if calibration_buffers is not None:
+                # Assign Transformer Engine calibration metadata buffers to the model.
                 # Requires CUDA graph warmup step.
-                for name, value in scale_buffers.items():
+                for name, value in calibration_buffers.items():
                     if value is not None:
                         if name in self._buffers:
                             setattr(self, name, value)
@@ -2689,7 +2694,7 @@ def _compile_eager_fallback_reason(
         if debug:
             return "debug instrumentation (nvidia-dlfw-inspect)"
         if FP8GlobalStateManager.get_calibration_config() is not None:
-            return "quantized scaling-factor buffering"
+            return "Transformer Engine calibration metadata buffering"
         weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
         if is_distributed_weight(weight_tensor):
             return "a DistributedWeight (custom weight parallelism, e.g. GTP)"
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index be553b0d87..580b6ada84 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -391,16 +391,16 @@ class QuantizationCalibrationConfig:
 
     Parameters
     ----------
-    activation_scale_decay : float, default = 0.0
+    transformer_engine_calibration_decay : float, default = 0.0
         Decay applied to buffered activation scaling factors before incorporating
         each new observation. With zero decay, only the latest value is retained.
     """
 
-    activation_scale_decay: float = 0.0
+    transformer_engine_calibration_decay: float = 0.0
 
     def __post_init__(self) -> None:
-        if self.activation_scale_decay < 0.0:
-            raise ValueError("activation_scale_decay must be non-negative")
+        if self.transformer_engine_calibration_decay < 0.0:
+            raise ValueError("transformer_engine_calibration_decay must be non-negative")
 
 
 @dataclass(frozen=True, slots=True)
@@ -1050,7 +1050,7 @@ class autocast:
              whether or not to enable low precision quantization (FP8/FP4).
     calibrating : bool, default = False
                  Enables calibration with the default configuration. Calibration
-                 collects and buffers quantized scaling factors even when executing
+                 collects and buffers Transformer Engine calibration metadata even when executing
                  without quantization enabled.
     calibration_config : QuantizationCalibrationConfig, default = None
                   Custom configuration for collecting checkpointable quantization scaling
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 284f17ed2c..7e12651d33 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -565,8 +565,12 @@ def create_metadata(
             "nontensor_kwargs": meta["nontensor_kwargs"],
         }
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
-        """Observe a tensor and update persistent calibration state."""
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
+        """Observe a tensor and update persistent calibration state.
+
+        ``calibration_decay`` decays the historical maximum before incorporating
+        the current observation. A value of zero retains only the current metadata.
+        """
         pass
 
     def get_quantization_recipe_name(self) -> str:
@@ -578,19 +582,16 @@ def _update_calibration_value(
         metadata_name: str,
         observed_value: Optional[torch.Tensor],
         *,
-        decay: float,
+        calibration_decay: float,
     ) -> None:
         """Merge an observation into quantizer-owned calibration state."""
         if observed_value is None or torch.isnan(observed_value).any():
             # Un-initialized scale. Ignore it.
             return
         observed_value = observed_value.detach()
-        calibration_state = getattr(self, "_calibration_state", None)
-        if calibration_state is None:
-            calibration_state = {}
-            self._calibration_state = calibration_state
+        calibration_state = self._calibration_state
         calibration_value = calibration_state.get(metadata_name)
-        if decay > 0.0:
+        if calibration_decay > 0.0:
             if calibration_value is not None and calibration_value.shape != observed_value.shape:
                 raise RuntimeError(
                     "Quantizer calibration value shape changed from "
@@ -603,7 +604,7 @@ def _update_calibration_value(
                 calibration_state[metadata_name] = calibration_value
             # Track a decaying maximum so early-training activation
             # outliers do not permanently determine the inference scale.
-            calibration_value.mul_(decay)
+            calibration_value.mul_(calibration_decay)
             torch.maximum(calibration_value, observed_value, out=calibration_value)
         else:
             # Without scale history, keep a reference to the current metadata
@@ -614,11 +615,7 @@ def _update_calibration_value(
 
     def _share_calibration_state_with(self, quantizer: "Quantizer") -> None:
         """Make a shallow quantizer copy share persistent calibration state."""
-        calibration_state = getattr(self, "_calibration_state", None)
-        if calibration_state is None:
-            calibration_state = {}
-            self._calibration_state = calibration_state
-        quantizer._calibration_state = calibration_state
+        quantizer._calibration_state = self._calibration_state
 
     def set_usage(
         self, *, rowwise: Optional[bool] = None, columnwise: Optional[bool] = None
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index ac3dde7fbf..3d338ecbd5 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -243,7 +243,7 @@ def is_quantizable(self, inp: torch.Tensor) -> bool:
             return False
         return True
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
         # NOTE: This interface is specific to requirements like delayed scaling
         # where state from an estimator influences distribution parameters.
         # NOTE(@cspades): Currently, PTQ calibration requirements don't need
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index 509a88b613..d9d45396bc 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -125,8 +125,9 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor:
         """Quantize tensor implementation"""
         return tex.quantize(tensor, self)
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
         if isinstance(tensor, (QuantizedTensor, QuantizedTensorStorage)):
+            # Retrieve the quantization metadata from quantized storage.
             observed_amax = self.amax
         else:
             # If quantized amax metadata does not yet exist or calibrate() is called directly,
@@ -135,7 +136,9 @@ def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
             amin, amax = tensor.aminmax()
             observed_amax = torch.max(-amin, amax).reshape(1)
             self.amax.copy_(observed_amax)
-        self._update_calibration_value("amax", observed_amax, decay=decay)
+        self._update_calibration_value(
+            "amax", observed_amax, calibration_decay=calibration_decay
+        )
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""
@@ -329,13 +332,14 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor:
         """Quantize tensor implementation"""
         return tex.quantize(tensor, self)
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
-        """Compute and calibrate (decaying amax) quantization metadata."""
-        scale_inv = getattr(tensor, "_scale_inv", None)
-        if scale_inv is None:
-            # If the scale_inv does not yet exist or calibrate() is called directly,
-            # then recompute the absmax / scale without quantization. This path is
-            # not performant and SHOULD NOT be called within training or inference.
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
+        """Compute and calibrate quantization metadata."""
+        if isinstance(tensor, (QuantizedTensor, QuantizedTensorStorage)):
+            # Retrieve the quantization metadata from quantized storage.
+            scale_inv = tensor._scale_inv
+        else:
+            # Direct calibration of a non-quantized tensor must reconstruct the metadata.
+            # This path is not performant and SHOULD NOT be called within training or inference.
             amin, amax = tensor.aminmax()
             amax = torch.maximum(-amin, amax).float().reshape(1)
             if self.with_amax_reduction and torch.distributed.is_initialized():
@@ -355,7 +359,9 @@ def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
                 scale = torch.ldexp(torch.ones_like(scale), exponent - 1)
             scale.masked_fill_(torch.isinf(amax) | (amax == 0), 1.0)
             scale_inv = torch.reciprocal(scale)
-        self._update_calibration_value("scale_inv", scale_inv, decay=decay)
+        self._update_calibration_value(
+            "scale_inv", scale_inv, calibration_decay=calibration_decay
+        )
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""
diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py
index 26d0798b92..bff83089b8 100644
--- a/transformer_engine/pytorch/tensor/hybrid_tensor.py
+++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py
@@ -157,6 +157,10 @@ def copy(self) -> "HybridQuantizer":
         quantizer.optimize_for_gemm = self.optimize_for_gemm
         return quantizer
 
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
+        """Reject calibration until child metadata ownership is supported."""
+        raise NotImplementedError("Calibration is not yet supported for HybridQuantizer")
+
     @property
     def with_amax_reduction(self) -> bool:
         """Whether either sub-quantizer has cross-rank amax reduction enabled."""
diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py
index 0ff8c01550..02d5ea7bd1 100644
--- a/transformer_engine/pytorch/tensor/identity_tensor.py
+++ b/transformer_engine/pytorch/tensor/identity_tensor.py
@@ -162,7 +162,7 @@ def update_quantized(
         dst._dtype = data.dtype
         return dst
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
         # No state to calibrate.
         pass
 
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index e0c409f11d..dde7b02207 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -132,7 +132,7 @@ def is_quantizable(self, inp: torch.Tensor) -> bool:
             return False
         return True
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
         """Calibrate an MXFP8 tensor."""
         # NOTE(@cspades): Currently, PTQ calibration requirements don't need
         # non-global / blockwise scaling factors, which are usually computed
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index fef3e10293..3a695ae82b 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -22,7 +22,7 @@
 )
 
 from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func
-from ..quantized_tensor import QuantizedTensor, Quantizer
+from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 from ..dynamo import register_value_opaque_quantizer
 from ._quantization_helpers import _IdentityFunc, safe_quantized_repr
 
@@ -340,13 +340,14 @@ def convert_shape_for_fp4(shape: Iterable[int]) -> Tuple[int, ...]:
         shape[-1] = shape[-1] // 2
         return tuple(shape)
 
-    def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
+    def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
         metadata_name = "amax_rowwise" if self.row_scaled_nvfp4 else "amax"
-        observed_amax = getattr(tensor, "_amax_rowwise", None)
-        if observed_amax is None:
-            # If quantized amax metadata does not yet exist or calibrate() is called directly,
-            # then recompute the absmax without quantization. This path is
-            # not performant and SHOULD NOT be called within training or inference.
+        if isinstance(tensor, (QuantizedTensor, QuantizedTensorStorage)):
+            # Retrieve the quantization metadata from quantized storage.
+            observed_amax = tensor._amax_rowwise
+        else:
+            # Direct calibration of a non-quantized tensor must reconstruct the metadata.
+            # This path is not performant and SHOULD NOT be called within training or inference.
             calibration_input = tensor
             if self.with_rht and self.with_post_rht_amax:
                 original_shape = calibration_input.shape
@@ -364,7 +365,11 @@ def calibrate(self, tensor: torch.Tensor, *, decay: float = 0.0) -> None:
                     op=torch.distributed.ReduceOp.MAX,
                     group=self._canonicalized_amax_reduction_group(),
                 )
-        self._update_calibration_value(metadata_name, observed_amax, decay=decay)
+        self._update_calibration_value(
+            metadata_name,
+            observed_amax,
+            calibration_decay=calibration_decay,
+        )
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""

From 0a64a735c2d22ad516dece87ff20b2e414027516 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 11 Sep 2026 22:44:03 +0000
Subject: [PATCH 12/16] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ...test_ptq_calibration_metadata_buffering.py | 23 +++++++------------
 transformer_engine/pytorch/module/_common.py  |  4 +---
 .../pytorch/module/layernorm_mlp.py           |  4 +---
 .../pytorch/tensor/float8_blockwise_tensor.py |  4 +---
 .../pytorch/tensor/float8_tensor.py           |  8 ++-----
 .../pytorch/tensor/hybrid_tensor.py           |  4 ++--
 6 files changed, 15 insertions(+), 32 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index b1f9b28cbf..81e9dbe6e3 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -220,9 +220,7 @@ def test_calibration_config_registers_module_scaling_factor_buffers(
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_linear_calibration_config_applies_activation_decay_only():
     module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
-    calibration_config = QuantizationCalibrationConfig(
-        transformer_engine_calibration_decay=0.5
-    )
+    calibration_config = QuantizationCalibrationConfig(transformer_engine_calibration_decay=0.5)
     buffer_suffix = "_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
 
     with autocast(
@@ -262,9 +260,7 @@ def test_linear_calibration_config_buffers_delayed_scaling_amax():
         module(torch.full((16, 32), 2.0, dtype=torch.bfloat16, device="cuda"))
 
     calibration_buffers = {
-        name: value
-        for name, value in module.named_buffers()
-        if name.endswith("_te_ptq_calibrated")
+        name: value for name, value in module.named_buffers() if name.endswith("_te_ptq_calibrated")
     }
     assert set(calibration_buffers) == {
         "input_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated",
@@ -346,14 +342,15 @@ def test_custom_quantizer_defaults_to_no_calibration_metadata():
     assert not _common._get_calibration_metadata_buffers("input", quantizer)
 
 
-def test_hybrid_quantizer_rejects_calibration():
+def test_hybrid_quantizer_calibration_is_noop():
     quantizer = HybridQuantizer(
         rowwise_quantizer=IdentityQuantizer(),
         columnwise_quantizer=IdentityQuantizer(),
     )
 
-    with pytest.raises(NotImplementedError, match="not yet supported for HybridQuantizer"):
-        quantizer.calibrate(torch.ones(1))
+    quantizer.calibrate(torch.ones(1))
+
+    assert not quantizer._calibration_state
 
 
 def test_resolve_calibration_quantizer_prefers_tensor_owner_and_unwraps_parent():
@@ -367,12 +364,8 @@ def test_resolve_calibration_quantizer_prefers_tensor_owner_and_unwraps_parent()
 def test_quantizer_calibration_state_is_keyed_by_quantized_metadata():
     quantizer = Quantizer(rowwise=True, columnwise=False)
 
-    quantizer._update_calibration_value(
-        "amax", torch.tensor([2.0]), calibration_decay=0.0
-    )
-    quantizer._update_calibration_value(
-        "scale_inv", torch.tensor([0.5]), calibration_decay=0.0
-    )
+    quantizer._update_calibration_value("amax", torch.tensor([2.0]), calibration_decay=0.0)
+    quantizer._update_calibration_value("scale_inv", torch.tensor([0.5]), calibration_decay=0.0)
 
     assert set(quantizer._calibration_state) == {"amax", "scale_inv"}
     torch.testing.assert_close(quantizer._calibration_state["amax"], torch.tensor([2.0]))
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index a5718f5c05..4dd76b3a9e 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -37,9 +37,7 @@ def _resolve_calibration_quantizer(tensor: Any, quantizer: Any) -> Any:
     return getattr(quantizer, "parent_quantizer", quantizer)
 
 
-def _get_calibration_metadata_buffers(
-    tensor_name: str, quantizer: Any
-) -> Dict[str, torch.Tensor]:
+def _get_calibration_metadata_buffers(tensor_name: str, quantizer: Any) -> Dict[str, torch.Tensor]:
     """Get checkpoint-buffer aliases from quantizer calibration state."""
     if quantizer is None:
         return {}
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 1077a06ad2..6cd8398254 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -723,9 +723,7 @@ def _forward(
                 calibration_buffers.update(fc1_input_calibration_buffers)
                 calibration_buffers.update(fc1_weight_calibration_buffers)
                 calibration_buffers.update(
-                    _get_calibration_metadata_buffers(
-                        "fc2_input", fc2_input_calibration_quantizer
-                    )
+                    _get_calibration_metadata_buffers("fc2_input", fc2_input_calibration_quantizer)
                 )
                 calibration_buffers.update(
                     _get_calibration_metadata_buffers(
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index 3d338ecbd5..adee186886 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -244,12 +244,10 @@ def is_quantizable(self, inp: torch.Tensor) -> bool:
         return True
 
     def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
-        # NOTE: This interface is specific to requirements like delayed scaling
-        # where state from an estimator influences distribution parameters.
         # NOTE(@cspades): Currently, PTQ calibration requirements don't need
         # non-global / blockwise scaling factors, which are usually computed
         # on-the-fly during inference. Implement this interface for future
-        # applications of blockwise scaling factor calibration.
+        # applications of blockwise FP8 calibration.
         pass
 
     def get_quantization_recipe_name(self) -> str:
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index d9d45396bc..d90cc424f9 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -136,9 +136,7 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
             amin, amax = tensor.aminmax()
             observed_amax = torch.max(-amin, amax).reshape(1)
             self.amax.copy_(observed_amax)
-        self._update_calibration_value(
-            "amax", observed_amax, calibration_decay=calibration_decay
-        )
+        self._update_calibration_value("amax", observed_amax, calibration_decay=calibration_decay)
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""
@@ -359,9 +357,7 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
                 scale = torch.ldexp(torch.ones_like(scale), exponent - 1)
             scale.masked_fill_(torch.isinf(amax) | (amax == 0), 1.0)
             scale_inv = torch.reciprocal(scale)
-        self._update_calibration_value(
-            "scale_inv", scale_inv, calibration_decay=calibration_decay
-        )
+        self._update_calibration_value("scale_inv", scale_inv, calibration_decay=calibration_decay)
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""
diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py
index bff83089b8..13cee56592 100644
--- a/transformer_engine/pytorch/tensor/hybrid_tensor.py
+++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py
@@ -158,8 +158,8 @@ def copy(self) -> "HybridQuantizer":
         return quantizer
 
     def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
-        """Reject calibration until child metadata ownership is supported."""
-        raise NotImplementedError("Calibration is not yet supported for HybridQuantizer")
+        """HybridQuantizer calibrate() has not yet been implemented."""
+        pass
 
     @property
     def with_amax_reduction(self) -> bool:

From 19cf8bbc03baf46db29a282322da2bfaaab51de2 Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Fri, 11 Sep 2026 17:11:24 -0700
Subject: [PATCH 13/16] Fix activation recompute.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 ...test_ptq_calibration_metadata_buffering.py | 120 ++++++++++++++++++
 .../pytorch/attention/fused_mla_q_uproj.py    |   1 +
 transformer_engine/pytorch/module/_common.py  |  11 ++
 .../pytorch/module/grouped_linear.py          |   5 +-
 .../pytorch/module/layernorm_linear.py        |   5 +-
 .../pytorch/module/layernorm_mlp.py           |   7 +-
 transformer_engine/pytorch/module/linear.py   |   3 +-
 .../pytorch/quantized_tensor.py               |   1 -
 .../pytorch/tensor/float8_blockwise_tensor.py |   2 +-
 .../pytorch/tensor/hybrid_tensor.py           |   1 -
 .../pytorch/tensor/identity_tensor.py         |   3 +-
 .../pytorch/tensor/mxfp8_tensor.py            |   1 -
 12 files changed, 148 insertions(+), 12 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index 81e9dbe6e3..af938cf747 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -8,14 +8,17 @@
 
 import pytest
 import torch
+from torch.utils.checkpoint import checkpoint
 
 from transformer_engine.common.recipe import DelayedScaling, Float8CurrentScaling
 from transformer_engine.pytorch import is_fp8_available, is_nvfp4_available
+from transformer_engine.pytorch import distributed as te_distributed
 from transformer_engine.pytorch.constants import DType
 from transformer_engine.pytorch.graph import make_graphed_callables
 from transformer_engine.pytorch.module import GroupedLinear, LayerNormLinear, LayerNormMLP, Linear
 from transformer_engine.pytorch.module import _common
 from transformer_engine.pytorch.module import grouped_linear
+from transformer_engine.pytorch.module import layernorm_mlp
 from transformer_engine.pytorch.quantization import (
     FP8GlobalStateManager,
     FP8GlobalState,
@@ -109,6 +112,16 @@ def test_calibration_api_additions_preserve_existing_parameter_order():
     ]
 
 
+def test_activation_recompute_detection_uses_te_marker(monkeypatch):
+    monkeypatch.setattr(
+        te_distributed,
+        "in_fp8_activation_recompute_phase",
+        lambda: True,
+    )
+
+    assert _common._is_in_activation_recompute_phase()
+
+
 def test_calibrating_argument_enables_default_calibration_config():
     assert FP8GlobalStateManager.get_calibration_config() is None
     with autocast(enabled=False, calibrating=True):
@@ -248,6 +261,113 @@ def test_linear_calibration_config_applies_activation_decay_only():
     torch.testing.assert_close(module.get_buffer(f"weight{buffer_suffix}"), weight_scale)
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("use_reentrant", (False, True))
+def test_external_activation_recomputation_does_not_update_calibration(use_reentrant):
+    module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
+    calibration_config = QuantizationCalibrationConfig(
+        transformer_engine_calibration_decay=0.5
+    )
+    recipe = Float8CurrentScaling()
+    buffer_name = "input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+
+    with (
+        torch.no_grad(),
+        autocast(enabled=False, recipe=recipe, calibration_config=calibration_config),
+    ):
+        module(torch.full((16, 32), 448.0, dtype=torch.bfloat16, device="cuda"))
+
+    def checkpointed_forward(inp):
+        with autocast(enabled=False, recipe=recipe, calibration_config=calibration_config):
+            return module(inp)
+
+    inp = torch.full(
+        (16, 32),
+        56.0,
+        dtype=torch.bfloat16,
+        device="cuda",
+        requires_grad=True,
+    )
+    out = checkpoint(checkpointed_forward, inp, use_reentrant=use_reentrant)
+    scale_after_forward = module.get_buffer(buffer_name).clone()
+    torch.testing.assert_close(scale_after_forward, torch.full((1,), 0.5, device="cuda"))
+
+    out.sum().backward()
+
+    torch.testing.assert_close(module.get_buffer(buffer_name), scale_after_forward)
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_layernorm_mlp_recomputation_does_not_update_calibration(monkeypatch):
+    recomputation_states = []
+    is_in_recompute = _common._is_in_activation_recompute_phase
+
+    def record_recomputation_state():
+        state = is_in_recompute()
+        recomputation_states.append(state)
+        return state
+
+    monkeypatch.setattr(
+        layernorm_mlp,
+        "_is_in_activation_recompute_phase",
+        record_recomputation_state,
+    )
+    module = LayerNormMLP(
+        32,
+        32,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        bias=False,
+        checkpoint=True,
+    )
+    calibration_config = QuantizationCalibrationConfig(
+        transformer_engine_calibration_decay=0.5
+    )
+    torch.manual_seed(123)
+    calibration_input = torch.randn((16, 32), dtype=torch.bfloat16, device="cuda")
+    with torch.no_grad():
+        module.layer_norm_weight.fill_(448.0)
+    with (
+        torch.no_grad(),
+        autocast(
+            enabled=False,
+            recipe=Float8CurrentScaling(),
+            calibration_config=calibration_config,
+        ),
+    ):
+        module(calibration_input)
+
+    buffer_name = "fc1_input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
+    previous_scale = module.get_buffer(buffer_name).clone()
+    with torch.no_grad():
+        module.layer_norm_weight.fill_(56.0)
+    inp = calibration_input.detach().clone().requires_grad_()
+
+    with autocast(
+        enabled=False,
+        recipe=Float8CurrentScaling(),
+        calibration_config=calibration_config,
+    ):
+        out = module(inp)
+
+    calibration_buffers = {
+        name: value.clone()
+        for name, value in module.named_buffers()
+        if name.endswith("_te_ptq_calibrated")
+    }
+    assert calibration_buffers
+    torch.testing.assert_close(
+        calibration_buffers[buffer_name],
+        previous_scale * 0.5,
+    )
+
+    out.sum().backward()
+
+    assert any(recomputation_states)
+    for name, value in calibration_buffers.items():
+        torch.testing.assert_close(module.get_buffer(name), value)
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_linear_calibration_config_buffers_delayed_scaling_amax():
     module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
index b15e53e6ee..f45a40ecb9 100644
--- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
+++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
@@ -203,6 +203,7 @@ def run(
 
         from cuda.bindings import driver as cuda
 
+        # pylint: disable-next=c-extension-no-member
         stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream)
         wrapper = cls._kernel()
 
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 4dd76b3a9e..c870546941 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -31,6 +31,17 @@ def _supports_calibration_decay(quantizer_type: type) -> bool:
     )
 
 
+def _is_in_activation_recompute_phase() -> bool:
+    """Whether a forward is running during activation recomputation."""
+    from ..distributed import in_fp8_activation_recompute_phase
+
+    if in_fp8_activation_recompute_phase():
+        return True
+    # Special hidden PyTorch AutoGrad identifier for activation recompute.
+    current_graph_task_id = getattr(torch._C, "_current_graph_task_id", None)
+    return current_graph_task_id is not None and current_graph_task_id() != -1
+
+
 def _resolve_calibration_quantizer(tensor: Any, quantizer: Any) -> Any:
     """Get the quantizer that owns calibration state for a tensor."""
     quantizer = getattr(tensor, "_quantizer", None) or quantizer
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 4a5daecfc7..dcf58ef8e2 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -33,6 +33,7 @@
 )
 from ._common import (
     _get_calibration_metadata_buffers,
+    _is_in_activation_recompute_phase,
     _resolve_calibration_quantizer,
     _supports_calibration_decay,
     can_reconstruct_wgrad_input_from_original,
@@ -588,7 +589,7 @@ def _forward_grouped_tensor(
             use_split_accumulator=use_split_accumulator,
         )
 
-        if calibration_buffers is not None:
+        if calibration_buffers is not None and not _is_in_activation_recompute_phase():
             grouped_inputs = grouped_x.quantized_tensors
             if grouped_inputs is None:
                 grouped_inputs = grouped_x.split_into_quantized_tensors()
@@ -937,7 +938,7 @@ def forward(
             weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]
 
         # Calibrate quantizers and buffer their metadata when requested.
-        if calibration_buffers is not None:
+        if calibration_buffers is not None and not _is_in_activation_recompute_phase():
             _calibrate_grouped_tensors(
                 inputmats,
                 weights_fp8,
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 49cb6839f1..37beb4c647 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -70,6 +70,7 @@
 from ..graph import is_graph_capturing
 from ._common import (
     _get_calibration_metadata_buffers,
+    _is_in_activation_recompute_phase,
     _resolve_calibration_quantizer,
     _supports_calibration_decay,
     apply_normalization,
@@ -130,7 +131,7 @@ def forward(
             eps,
             is_first_microbatch,
             fp8,
-            fp8_calibration,
+            _fp8_calibration,
             wgrad_store,
             fuse_wgrad_accumulation,
             input_quantizer,
@@ -388,7 +389,7 @@ def forward(
         weight_calibration_quantizer = _resolve_calibration_quantizer(weightmat, weight_quantizer)
 
         # Calibrate quantizers and buffer their metadata when requested.
-        if calibration_buffers is not None:
+        if calibration_buffers is not None and not _is_in_activation_recompute_phase():
             if input_calibration_quantizer is not None:
                 if _supports_calibration_decay(type(input_calibration_quantizer)):
                     input_calibration_quantizer.calibrate(
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 6cd8398254..1e9de95c13 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -77,6 +77,7 @@
 from ..tensor.identity_tensor import IdentityQuantizer
 from ._common import (
     _get_calibration_metadata_buffers,
+    _is_in_activation_recompute_phase,
     _resolve_calibration_quantizer,
     _supports_calibration_decay,
     apply_normalization,
@@ -578,7 +579,11 @@ def _forward(
         fc1_weight_calibration_buffers = {}
 
         # Calibrate FC1 quantizers and collect their metadata when requested.
-        should_calibrate = calibration_buffers is not None
+        should_calibrate = (
+            calibration_buffers is not None
+            and not _is_in_activation_recompute_phase()
+            and not is_recomputation
+        )
         if should_calibrate:
             if fc1_input_calibration_quantizer is not None:
                 if _supports_calibration_decay(type(fc1_input_calibration_quantizer)):
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 9ce1541341..098d75b998 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -33,6 +33,7 @@
 )
 from ._common import (
     _get_calibration_metadata_buffers,
+    _is_in_activation_recompute_phase,
     _resolve_calibration_quantizer,
     _supports_calibration_decay,
     can_reconstruct_wgrad_input_from_original,
@@ -613,7 +614,7 @@ def _linear_forward_impl(
     weight_calibration_quantizer = _resolve_calibration_quantizer(weightmat, weight_quantizer)
 
     # Calibrate quantizers and buffer their metadata when requested.
-    if args.calibration_buffers is not None:
+    if args.calibration_buffers is not None and not _is_in_activation_recompute_phase():
         if input_calibration_quantizer is not None:
             if _supports_calibration_decay(type(input_calibration_quantizer)):
                 input_calibration_quantizer.calibrate(
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 7e12651d33..038f0b53ee 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -571,7 +571,6 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
         ``calibration_decay`` decays the historical maximum before incorporating
         the current observation. A value of zero retains only the current metadata.
         """
-        pass
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index adee186886..2aac4ce6c0 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -244,11 +244,11 @@ def is_quantizable(self, inp: torch.Tensor) -> bool:
         return True
 
     def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
+        """Float8BlockQuantizer does not yet support calibration."""
         # NOTE(@cspades): Currently, PTQ calibration requirements don't need
         # non-global / blockwise scaling factors, which are usually computed
         # on-the-fly during inference. Implement this interface for future
         # applications of blockwise FP8 calibration.
-        pass
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""
diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py
index 13cee56592..874d101c96 100644
--- a/transformer_engine/pytorch/tensor/hybrid_tensor.py
+++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py
@@ -159,7 +159,6 @@ def copy(self) -> "HybridQuantizer":
 
     def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
         """HybridQuantizer calibrate() has not yet been implemented."""
-        pass
 
     @property
     def with_amax_reduction(self) -> bool:
diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py
index 02d5ea7bd1..cb6838d116 100644
--- a/transformer_engine/pytorch/tensor/identity_tensor.py
+++ b/transformer_engine/pytorch/tensor/identity_tensor.py
@@ -163,8 +163,7 @@ def update_quantized(
         return dst
 
     def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) -> None:
-        # No state to calibrate.
-        pass
+        """No-op since identity quantization has no calibration state."""
 
     def _get_compatible_recipe(self):
         # Only reachable via CustomRecipe (qfactory returns IdentityQuantizer).
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index dde7b02207..d6920f3aea 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -138,7 +138,6 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
         # non-global / blockwise scaling factors, which are usually computed
         # on-the-fly during inference. Implement this interface for future
         # applications of MXFP8 calibration.
-        pass
 
     def get_quantization_recipe_name(self) -> str:
         """Get the stable name of the quantization recipe."""

From 1842f94630dd85e7c25079a80cc4587a2437b2d5 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 12 Sep 2026 00:13:01 +0000
Subject: [PATCH 14/16] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/pytorch/test_ptq_calibration_metadata_buffering.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index af938cf747..4f5cdda2f2 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -265,9 +265,7 @@ def test_linear_calibration_config_applies_activation_decay_only():
 @pytest.mark.parametrize("use_reentrant", (False, True))
 def test_external_activation_recomputation_does_not_update_calibration(use_reentrant):
     module = Linear(32, 32, params_dtype=torch.bfloat16, device="cuda", bias=False)
-    calibration_config = QuantizationCalibrationConfig(
-        transformer_engine_calibration_decay=0.5
-    )
+    calibration_config = QuantizationCalibrationConfig(transformer_engine_calibration_decay=0.5)
     recipe = Float8CurrentScaling()
     buffer_name = "input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated"
 
@@ -320,9 +318,7 @@ def record_recomputation_state():
         bias=False,
         checkpoint=True,
     )
-    calibration_config = QuantizationCalibrationConfig(
-        transformer_engine_calibration_decay=0.5
-    )
+    calibration_config = QuantizationCalibrationConfig(transformer_engine_calibration_decay=0.5)
     torch.manual_seed(123)
     calibration_input = torch.randn((16, 32), dtype=torch.bfloat16, device="cuda")
     with torch.no_grad():

From 41923e4aa6db70331ee2145e5c76bda02bdf65ed Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Sat, 12 Sep 2026 18:50:58 -0700
Subject: [PATCH 15/16] Minor fixes, still build quantizers when calibrating
 BF16/FP32, some calibration parity issues with RHT, and torch compile
 support.

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 .../pytorch/test_ptq_calibration_metadata_buffering.py |  5 +++--
 tests/pytorch/test_torch_compile.py                    |  1 +
 transformer_engine/pytorch/dynamo/quantizer_opaque.py  |  3 +++
 transformer_engine/pytorch/module/grouped_linear.py    |  3 ++-
 transformer_engine/pytorch/module/layernorm_linear.py  |  7 +++++--
 transformer_engine/pytorch/module/layernorm_mlp.py     |  1 +
 transformer_engine/pytorch/module/linear.py            |  7 +++++--
 transformer_engine/pytorch/tensor/float8_tensor.py     |  2 +-
 transformer_engine/pytorch/tensor/nvfp4_tensor.py      | 10 ++--------
 9 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/tests/pytorch/test_ptq_calibration_metadata_buffering.py b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
index 4f5cdda2f2..8c17fb02aa 100644
--- a/tests/pytorch/test_ptq_calibration_metadata_buffering.py
+++ b/tests/pytorch/test_ptq_calibration_metadata_buffering.py
@@ -316,6 +316,7 @@ def record_recomputation_state():
         params_dtype=torch.bfloat16,
         device="cuda",
         bias=False,
+        activation="relu",
         checkpoint=True,
     )
     calibration_config = QuantizationCalibrationConfig(transformer_engine_calibration_decay=0.5)
@@ -644,7 +645,8 @@ def test_zero_decay_keeps_observed_metadata_reference():
         calibration_decay=0.0,
     )
 
-    assert quantizer._calibration_state["scale_inv"] is observed_scale
+    calibrated_scale = quantizer._calibration_state["scale_inv"]
+    assert calibrated_scale.data_ptr() == observed_scale.data_ptr()
 
 
 def test_decaying_calibration_rejects_metadata_shape_change():
@@ -793,7 +795,6 @@ def test_current_scaling_high_precision_calibration_matches_quantization(
     ("row_scaled_nvfp4", "with_rht", "with_post_rht_amax", "with_random_sign_mask"),
     (
         (False, False, False, False),
-        (False, True, False, False),
         (False, True, True, False),
         (False, True, True, True),
         (True, False, False, False),
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 456c83b186..071b47b399 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -1442,6 +1442,7 @@ def test_quantizer_value_object(factory):
     rebuilt = eval(repr_str, dict(globals_))  # pylint: disable=eval-used
     assert rebuilt == a and rebuilt is not a
     assert hash(rebuilt) == hash(a)
+    assert rebuilt._calibration_state == {}
     # The deprecated amax-reduction group is never part of the value.
     assert getattr(rebuilt, "amax_reduction_group", None) is None
 
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index da595b1070..10a1cc0cc6 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -45,6 +45,9 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
         if name == "dtype":
             value = DType.cast(value)
         object.__setattr__(obj, name, value)
+    # Runtime calibration metadata is intentionally excluded from value semantics,
+    # but rebuilt quantizers still need an initialized state for copy()/quantize().
+    object.__setattr__(obj, "_calibration_state", {})
     # Restore non-value derived state that ``__init__`` would normally build but
     # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor).
     finalize = getattr(obj, "_rebuild_derived_state", None)
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index dcf58ef8e2..16760af961 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -2592,7 +2592,8 @@ def _get_quantizers(self):
             [None] * self.num_gemms,
             [None] * self.num_gemms,
         )
-        if self.fp8:
+        # Calibration-only mode needs quantizers to observe non-quantized tensors.
+        if self.fp8 or self.fp8_calibration:
             input_quantizers = [
                 self.quantizers["scaling_fwd"][
                     self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"]
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 37beb4c647..f0d85f7d32 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -1884,10 +1884,13 @@ def forward(
         return out
 
     def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
-        if not self.fp8:
+        if not self.fp8 and not self.fp8_calibration:
+            # Calibration-only mode needs quantizers to observe non-quantized tensors,
+            # so return None only when neither FP8 nor calibration is active.
             return [None] * 6
 
-        self._warn_missing_output_quantizer_role(fp8_output, fp8_grad)
+        if self.fp8:
+            self._warn_missing_output_quantizer_role(fp8_output, fp8_grad)
 
         grad_input_quantizer = None
         grad_weight_quantizer = None
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 1e9de95c13..a57398930d 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -2566,6 +2566,7 @@ def _get_quantizers(self, fp8_output, is_grad_enabled):
             fc2_grad_output_quantizer,
         ) = [None] * 10
         fc1_weight_quantizer, fc2_weight_quantizer = self._get_weight_quantizers()
+        # Calibration-only mode needs quantizers to observe non-quantized tensors.
         if self.fp8 or self.fp8_calibration:
             fc1_input_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT]
             fc1_input_quantizer.internal = True
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 098d75b998..86d990a7ea 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -2621,10 +2621,13 @@ def forward(
         return out
 
     def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
-        if not self.fp8:
+        if not self.fp8 and not self.fp8_calibration:
+            # Calibration-only mode needs quantizers to observe non-quantized tensors,
+            # so return None only when neither FP8 nor calibration is active.
             return [None] * 6
 
-        self._warn_missing_output_quantizer_role(fp8_output, fp8_grad)
+        if self.fp8:
+            self._warn_missing_output_quantizer_role(fp8_output, fp8_grad)
 
         grad_input_quantizer = None
         grad_weight_quantizer = None
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index d90cc424f9..6297544026 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -134,7 +134,7 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
             # then recompute the absmax without quantization. This path is
             # not performant and SHOULD NOT be called within training or inference.
             amin, amax = tensor.aminmax()
-            observed_amax = torch.max(-amin, amax).reshape(1)
+            observed_amax = torch.max(-amin, amax).float().reshape(1)
             self.amax.copy_(observed_amax)
         self._update_calibration_value("amax", observed_amax, calibration_decay=calibration_decay)
 
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 3a695ae82b..c3307a517a 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -348,16 +348,10 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
         else:
             # Direct calibration of a non-quantized tensor must reconstruct the metadata.
             # This path is not performant and SHOULD NOT be called within training or inference.
-            calibration_input = tensor
-            if self.with_rht and self.with_post_rht_amax:
-                original_shape = calibration_input.shape
-                calibration_input = (
-                    calibration_input.reshape(-1, 16).to(torch.bfloat16) @ self.rht_matrix
-                ).reshape(original_shape)
             if self.row_scaled_nvfp4:
-                amin, amax = calibration_input.aminmax(dim=-1)
+                amin, amax = tensor.aminmax(dim=-1)
             else:
-                amin, amax = calibration_input.aminmax()
+                amin, amax = tensor.aminmax()
             observed_amax = torch.maximum(-amin, amax).reshape(-1).float()
             if self.with_amax_reduction and torch.distributed.is_initialized():
                 torch.distributed.all_reduce(

From dea6c30e18e62c5b58bf21b823d88397ec0a561f Mon Sep 17 00:00:00 2001
From: Cory Ye <cye@nvidia.com>
Date: Sat, 12 Sep 2026 19:21:50 -0700
Subject: [PATCH 16/16] nit: fix comment for bf16/fp32 calibration

Signed-off-by: Cory Ye <cye@nvidia.com>
---
 transformer_engine/pytorch/tensor/float8_tensor.py | 6 +++---
 transformer_engine/pytorch/tensor/nvfp4_tensor.py  | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index 6297544026..5ac2846e09 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -131,8 +131,8 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
             observed_amax = self.amax
         else:
             # If quantized amax metadata does not yet exist or calibrate() is called directly,
-            # then recompute the absmax without quantization. This path is
-            # not performant and SHOULD NOT be called within training or inference.
+            # then recompute the absmax without quantization.
+            # This path is NOT performant and should only be used for non-quantized Tensor calibration.
             amin, amax = tensor.aminmax()
             observed_amax = torch.max(-amin, amax).float().reshape(1)
             self.amax.copy_(observed_amax)
@@ -337,7 +337,7 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
             scale_inv = tensor._scale_inv
         else:
             # Direct calibration of a non-quantized tensor must reconstruct the metadata.
-            # This path is not performant and SHOULD NOT be called within training or inference.
+            # This path is NOT performant and should only be used for non-quantized Tensor calibration.
             amin, amax = tensor.aminmax()
             amax = torch.maximum(-amin, amax).float().reshape(1)
             if self.with_amax_reduction and torch.distributed.is_initialized():
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index c3307a517a..08b5be2bfb 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -347,7 +347,7 @@ def calibrate(self, tensor: torch.Tensor, *, calibration_decay: float = 0.0) ->
             observed_amax = tensor._amax_rowwise
         else:
             # Direct calibration of a non-quantized tensor must reconstruct the metadata.
-            # This path is not performant and SHOULD NOT be called within training or inference.
+            # This path is NOT performant and should only be used for non-quantized Tensor calibration.
             if self.row_scaled_nvfp4:
                 amin, amax = tensor.aminmax(dim=-1)
             else: