From e9c1f4af2199d55eec9e4f0e239a6e4a62508e27 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Mon, 26 Oct 2026 06:30:53 +0000
Subject: [PATCH 1/4] Accumulate GPTQ Hessians from BF16/FP16 inputs with a
 float32-output GEMM

update_hessian cast every input batch to float32 and ran a float32 GEMM,
which runs on CUDA cores. BF16 and FP16 products are exact in float32, so
a tensor-core GEMM that takes the inputs directly and writes float32
differs only in how it rounds the sums: about 1e-5 of the Hessian's
diagonal, against 1e-6 for the float32 GEMM and 1e-2 for a BF16 output.
On GB300 it is 9-16x faster for 2560- and 4096-wide expert Hessians and
needs no float32 copy of the inputs. A torch without float32-output
GEMMs, CPU inputs and float32 inputs keep the float32 path.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 .../torch/quantization/utils/calib_utils.py   | 46 +++++++++++++++++--
 tests/gpu/torch/quantization/test_gptq.py     | 22 +++++++++
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index 4895a31103b..f5a1fe65940 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -38,6 +38,7 @@
 
 """GPTQ helper and Hessian utilities for calibration."""
 
+import functools
 from contextlib import contextmanager
 
 import torch
@@ -120,7 +121,7 @@ def update_hessian(input, hessian, n_samples):
     Note: input must be non-empty (batch_size > 0); a zero-sized input causes division by zero.
     """
     # Flatten to 2D (total_tokens, features) first, so batch_size counts tokens
-    input_flat = input.reshape(-1, input.shape[-1]).float()
+    input_flat = input.reshape(-1, input.shape[-1])
     batch_size = input_flat.shape[0]
 
     if batch_size == 0:  # in MOEs some experts receive no tokens
@@ -130,16 +131,51 @@ def update_hessian(input, hessian, n_samples):
     # once per expert per forward, and separate scale and add passes each rewrite all of H.
     scale = n_samples / (n_samples + batch_size)
     n_samples += batch_size
+    alpha = 2 / n_samples
+    if (
+        input_flat.is_cuda
+        and input_flat.dtype in (torch.bfloat16, torch.float16)
+        and _has_float32_out_gemm()
+    ):
+        # BF16 and FP16 products are exact in float32, so a tensor-core GEMM writing float32 differs
+        # from a float32 GEMM only in how it rounds the sums, by ~1e-5 of the Hessian's diagonal.
+        # It is 9-16x faster on GB300 and needs no float32 copy of the inputs.
+        if hessian.device == input_flat.device:
+            torch.addmm(
+                hessian,
+                input_flat.t(),
+                input_flat,
+                beta=scale,
+                alpha=alpha,
+                out_dtype=torch.float32,
+                out=hessian,
+            )
+        else:
+            gram = torch.mm(input_flat.t(), input_flat, out_dtype=torch.float32)
+            hessian.mul_(scale).add_(gram.to(hessian.device), alpha=alpha)
+        return hessian, n_samples
+
+    input_flat = input_flat.float()
     if hessian.device == input_flat.device:
-        hessian.addmm_(input_flat.t(), input_flat, beta=scale, alpha=2 / n_samples)
+        hessian.addmm_(input_flat.t(), input_flat, beta=scale, alpha=alpha)
     else:
-        hessian.mul_(scale).add_(
-            (input_flat.t() @ input_flat).to(hessian.device), alpha=2 / n_samples
-        )
+        hessian.mul_(scale).add_((input_flat.t() @ input_flat).to(hessian.device), alpha=alpha)
 
     return hessian, n_samples
 
 
+@functools.cache
+def _has_float32_out_gemm() -> bool:
+    """Whether this torch's CUDA GEMMs take BF16/FP16 inputs and write float32 (``out_dtype``)."""
+    try:
+        x = torch.zeros(1, 1, device="cuda", dtype=torch.bfloat16)
+        h = torch.mm(x, x, out_dtype=torch.float32)
+        torch.addmm(h, x, x, out_dtype=torch.float32, out=h)
+    except (TypeError, RuntimeError):
+        return False
+    return True
+
+
 def compute_hessian_inverse(hessian, weight, perc_damp):
     """Compute damped upper-Cholesky inverse Hessian.
 
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index 286883969f9..ebb3e3a7568 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -81,6 +81,28 @@ def test_hessian_of_a_full_gpu_falls_back_to_the_cpu(need_2_gpus, monkeypatch, b
     assert helper.hessian.device == torch.device("cpu")
 
 
+def test_bf16_inputs_accumulate_the_float32_hessian():
+    """BF16 inputs give the float32 Hessian, on the inputs' GPU or the CPU, without a float32
+    copy of the inputs."""
+    torch.manual_seed(0)
+    batches = [torch.randn(4096, 1024, device="cuda", dtype=torch.bfloat16) for _ in range(3)]
+    inputs = torch.cat(batches).double()
+    expected = 2 * inputs.t() @ inputs / inputs.shape[0]
+
+    for device in ("cuda", "cpu"):
+        hessian, n_samples = torch.zeros(1024, 1024, device=device), 0
+        for batch in batches:
+            torch.cuda.synchronize()
+            torch.cuda.reset_peak_memory_stats()
+            base = torch.cuda.memory_allocated()
+            hessian, n_samples = update_hessian(batch, hessian, n_samples)
+            growth = torch.cuda.max_memory_allocated() - base
+        # Tensor cores round float32 sums their own way: ~1e-5 of the diagonal (2), not BF16's 1e-2.
+        torch.testing.assert_close(hessian.cuda().double(), expected, rtol=0, atol=1e-4)
+        if device == "cuda":
+            assert growth < batch.numel() * 4
+
+
 @pytest.mark.parametrize(
     ("measured_fraction", "stays"),
     [(None, False), (0.05, True), (0.25, False)],

From 7a35f3f0919db2ef43960cc1f54ee88fc85d024e Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Tue, 27 Oct 2026 04:29:14 +0000
Subject: [PATCH 2/4] Make the bfloat16 Hessian GEMM opt-in with
 hessian_input_dtype

GPTQ's algorithm config gains hessian_input_dtype ("float32" or
"bfloat16", default "float32"). The default keeps the float32 copy and
float32 GEMM; "bfloat16" feeds the activations to the GEMM in bfloat16
and accumulates in float32, so recipes choose speed over the exact
float32 sums. Where CUDA GEMMs do not take out_dtype, "bfloat16" falls
back to float32 with a warning instead of being ignored. The out_dtype
check is now cached per device rather than for the current GPU only.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 CHANGELOG.rst                                 |  1 +
 modelopt/torch/quantization/config.py         | 10 ++++
 modelopt/torch/quantization/model_calib.py    |  4 ++
 .../torch/quantization/utils/calib_utils.py   | 48 ++++++++++++++-----
 tests/gpu/torch/quantization/test_gptq.py     | 34 +++++++++++--
 5 files changed, 80 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index a073a683af9..20f3a743c44 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -15,6 +15,7 @@ Changelog
 - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time.
 - GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch to what one decoder layer's forward fits in GPU memory, up to 32,768 tokens (64 samples at ``--calib_seq 512``), instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching.
 - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``.
+- Add ``hessian_input_dtype`` to the GPTQ algorithm config. ``bfloat16`` accumulates the float32 Hessian with a tensor-core GEMM on bfloat16 activations, up to 16x faster on GB300 and within about 1e-5 of the Hessian's diagonal; the default ``float32`` keeps the existing float32 GEMM.
 - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype.
 
 *Speculative Decoding*
diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py
index a43e031f135..c12de90e0be 100644
--- a/modelopt/torch/quantization/config.py
+++ b/modelopt/torch/quantization/config.py
@@ -1250,6 +1250,16 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig):
         (the GGML IQ1 and IQ2 formats) weight it by the square root of the Hessian diagonal,
         llama.cpp's imatrix statistic. Other formats ignore it.""",
     )
+    hessian_input_dtype: Literal["float32", "bfloat16"] = ModeloptField(
+        default="float32",
+        title="Dtype the Hessian GEMM reads the activations in.",
+        description="""The Hessian is accumulated and stored in float32 either way. ``float32``
+        copies each batch of activations to float32 for the GEMM. ``bfloat16`` feeds them to a
+        tensor-core GEMM in bfloat16 that accumulates in float32: for bfloat16 activations it
+        forms the same products and only rounds the sums differently, by about 1e-5 of the
+        Hessian's diagonal, and it is 9-16x faster on GB300. It needs CUDA GEMMs that take
+        ``out_dtype``; without them the Hessian falls back to float32 with a warning.""",
+    )
 
 
 class ImatrixCalibConfig(QuantizeAlgorithmConfig):
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index ac20729da19..7cef364e8f7 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -2425,6 +2425,7 @@ def gptq(
     block_size: int = 128,
     fused: bool = False,
     importance_weighted: bool = False,
+    hessian_input_dtype: str = "float32",
 ):
     """GPTQ quantization.
 
@@ -2463,6 +2464,8 @@ def gptq(
         fused: If True, use fused Triton kernel for NVFP4 static quantization.
         importance_weighted: If True, formats whose codebook search takes a per-input-column
             importance weight it by the square root of the Hessian diagonal.
+        hessian_input_dtype: ``"float32"`` or ``"bfloat16"``, the dtype the Hessian GEMM reads
+            the activations in; see :class:`GPTQCalibConfig <modelopt.torch.quantization.config.GPTQCalibConfig>`.
     """
     total_start = time.time()
 
@@ -2513,6 +2516,7 @@ def _make_gptq_handle(name, m, weight_quantizer=None, weight=None, log_mse=True)
             offload_to_cpu=True,
             fused=fused,
             importance_weighted=importance_weighted,
+            hessian_input_dtype=hessian_input_dtype,
             weight_quantizer=weight_quantizer,
             weight=weight,
             log_mse=log_mse,
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index f5a1fe65940..63e617e1b17 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -44,7 +44,12 @@
 import torch
 import torch.distributed as dist
 
-from modelopt.torch.utils import get_available_gpu_memory, get_process_gpus, print_rank_0
+from modelopt.torch.utils import (
+    get_available_gpu_memory,
+    get_process_gpus,
+    print_rank_0,
+    warn_rank_0,
+)
 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method
 
 # A GPU takes a new Hessian only if enough memory stays available after it for the calibration
@@ -108,13 +113,16 @@ def _hessian_offload_device(nbytes: int) -> torch.device:
     return torch.device("cpu")
 
 
-def update_hessian(input, hessian, n_samples):
+def update_hessian(input, hessian, n_samples, input_dtype=torch.float32):
     """Update hessian matrix with new input samples using incremental formula.
 
     Args:
         input: Input tensor (batch_size, ..., features)
         hessian: Current Hessian matrix to update in-place
         n_samples: Number of samples already processed
+        input_dtype: The dtype the GEMM reads ``input`` in: ``torch.float32``, or
+            ``torch.bfloat16`` for a tensor-core GEMM that accumulates in float32.
+
     Returns:
         Tuple of (updated_hessian, new_sample_count)
 
@@ -133,13 +141,14 @@ def update_hessian(input, hessian, n_samples):
     n_samples += batch_size
     alpha = 2 / n_samples
     if (
-        input_flat.is_cuda
-        and input_flat.dtype in (torch.bfloat16, torch.float16)
-        and _has_float32_out_gemm()
+        input_dtype == torch.bfloat16
+        and input_flat.is_cuda
+        and _has_float32_out_gemm(input_flat.device)
     ):
-        # BF16 and FP16 products are exact in float32, so a tensor-core GEMM writing float32 differs
-        # from a float32 GEMM only in how it rounds the sums, by ~1e-5 of the Hessian's diagonal.
-        # It is 9-16x faster on GB300 and needs no float32 copy of the inputs.
+        # BF16 products are exact in float32, so a tensor-core GEMM writing float32 differs from a
+        # float32 GEMM only in how it rounds the sums, by ~1e-5 of the Hessian's diagonal. It is
+        # 9-16x faster on GB300 and needs no float32 copy of the inputs.
+        input_flat = input_flat.to(torch.bfloat16)
         if hessian.device == input_flat.device:
             torch.addmm(
                 hessian,
@@ -154,6 +163,8 @@ def update_hessian(input, hessian, n_samples):
             gram = torch.mm(input_flat.t(), input_flat, out_dtype=torch.float32)
             hessian.mul_(scale).add_(gram.to(hessian.device), alpha=alpha)
         return hessian, n_samples
+    if input_dtype == torch.bfloat16:
+        _warn_bfloat16_hessian_unavailable(input_flat.device)
 
     input_flat = input_flat.float()
     if hessian.device == input_flat.device:
@@ -165,10 +176,10 @@ def update_hessian(input, hessian, n_samples):
 
 
 @functools.cache
-def _has_float32_out_gemm() -> bool:
-    """Whether this torch's CUDA GEMMs take BF16/FP16 inputs and write float32 (``out_dtype``)."""
+def _has_float32_out_gemm(device: torch.device) -> bool:
+    """Whether GEMMs on CUDA ``device`` take bfloat16 inputs and write float32 (``out_dtype``)."""
     try:
-        x = torch.zeros(1, 1, device="cuda", dtype=torch.bfloat16)
+        x = torch.zeros(1, 1, device=device, dtype=torch.bfloat16)
         h = torch.mm(x, x, out_dtype=torch.float32)
         torch.addmm(h, x, x, out_dtype=torch.float32, out=h)
     except (TypeError, RuntimeError):
@@ -176,6 +187,14 @@ def _has_float32_out_gemm() -> bool:
     return True
 
 
+@functools.cache
+def _warn_bfloat16_hessian_unavailable(device: torch.device) -> None:
+    warn_rank_0(
+        f"GPTQ hessian_input_dtype='bfloat16' needs CUDA GEMMs that take out_dtype, which {device} "
+        "lacks with this torch; computing its Hessians in float32."
+    )
+
+
 def compute_hessian_inverse(hessian, weight, perc_damp):
     """Compute damped upper-Cholesky inverse Hessian.
 
@@ -238,6 +257,7 @@ def __init__(
         weight_quantizer=None,
         weight=None,
         log_mse=True,
+        hessian_input_dtype="float32",
     ):
         """Initialize GPTQHelper with module state and Hessian storage.
 
@@ -245,11 +265,13 @@ def __init__(
         one expert's slice of an HF fused-MoE parameter; by default the helper owns
         ``module.weight`` and ``module.weight_quantizer``. ``log_mse=False`` skips the logged
         Hessian-weighted error, which costs two weight-by-Hessian matmuls.
+        ``hessian_input_dtype`` is the dtype the Hessian GEMM reads the inputs in, by name.
         """
         self.module = module
         self.name = name
         self.fused = fused
         self.importance_weighted = importance_weighted
+        self.hessian_input_dtype = getattr(torch, hessian_input_dtype)
         self._weight_quantizer = weight_quantizer
         self.log_mse = log_mse
         weight = module.weight if weight is None else weight
@@ -295,7 +317,9 @@ def hessian_forward(self, input, *args, **kwargs):
     def accumulate(self, inputs):
         """Add a batch of this weight's inputs, after input quantization, to the Hessian."""
         inputs = inputs.to_local() if hasattr(inputs, "to_local") else inputs
-        self.hessian, self.n_samples = update_hessian(inputs, self.hessian, self.n_samples)
+        self.hessian, self.n_samples = update_hessian(
+            inputs, self.hessian, self.n_samples, self.hessian_input_dtype
+        )
 
     def cleanup(self):
         """Unpatch the module's forward method."""
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index ebb3e3a7568..d2733293234 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -81,13 +81,14 @@ def test_hessian_of_a_full_gpu_falls_back_to_the_cpu(need_2_gpus, monkeypatch, b
     assert helper.hessian.device == torch.device("cpu")
 
 
-def test_bf16_inputs_accumulate_the_float32_hessian():
-    """BF16 inputs give the float32 Hessian, on the inputs' GPU or the CPU, without a float32
-    copy of the inputs."""
+def test_bf16_hessian_input_gives_the_float32_hessian():
+    """``input_dtype=torch.bfloat16`` gives the float32 Hessian, on the inputs' GPU or the CPU,
+    and with float32-output GEMMs makes no float32 copy of the inputs."""
     torch.manual_seed(0)
     batches = [torch.randn(4096, 1024, device="cuda", dtype=torch.bfloat16) for _ in range(3)]
     inputs = torch.cat(batches).double()
     expected = 2 * inputs.t() @ inputs / inputs.shape[0]
+    fast = calib_utils._has_float32_out_gemm(batches[0].device)
 
     for device in ("cuda", "cpu"):
         hessian, n_samples = torch.zeros(1024, 1024, device=device), 0
@@ -95,14 +96,37 @@ def test_bf16_inputs_accumulate_the_float32_hessian():
             torch.cuda.synchronize()
             torch.cuda.reset_peak_memory_stats()
             base = torch.cuda.memory_allocated()
-            hessian, n_samples = update_hessian(batch, hessian, n_samples)
+            hessian, n_samples = update_hessian(batch, hessian, n_samples, torch.bfloat16)
             growth = torch.cuda.max_memory_allocated() - base
         # Tensor cores round float32 sums their own way: ~1e-5 of the diagonal (2), not BF16's 1e-2.
         torch.testing.assert_close(hessian.cuda().double(), expected, rtol=0, atol=1e-4)
-        if device == "cuda":
+        if device == "cuda" and fast:
             assert growth < batch.numel() * 4
 
 
+def test_gptq_hessian_input_dtype_reaches_the_hessian(monkeypatch):
+    """A recipe's ``hessian_input_dtype: bfloat16`` drops the float32 copy of each batch that the
+    default float32 Hessian GEMM makes."""
+    if not calib_utils._has_float32_out_gemm(torch.device("cuda", torch.cuda.current_device())):
+        pytest.skip("needs CUDA GEMMs that take out_dtype")
+    inputs = torch.randn(4, 1024, 2048, device="cuda", dtype=torch.bfloat16)
+    peaks = {}
+    for hessian_input_dtype in ("float32", "bfloat16"):
+        monkeypatch.setattr(calib_utils, "_calib_forward_peak", {})
+        model = torch.nn.Linear(2048, 16, bias=False, device="cuda", dtype=torch.bfloat16)
+        quant_cfg = {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}},
+            ],
+            "algorithm": {"method": "gptq", "hessian_input_dtype": hessian_input_dtype},
+        }
+        mtq.quantize(model, quant_cfg, forward_loop=lambda m: m(inputs))
+        peaks[hessian_input_dtype] = calib_utils._calib_forward_peak[torch.cuda.current_device()]
+    # The float32 GEMM reads a float32 copy of the batch: 4 bytes per element.
+    assert peaks["float32"] - peaks["bfloat16"] >= inputs.numel() * 2
+
+
 @pytest.mark.parametrize(
     ("measured_fraction", "stays"),
     [(None, False), (0.05, True), (0.25, False)],

From bf0cec2d6c69449fbac9713403afd8f9fb0a9af8 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Tue, 27 Oct 2026 04:52:17 +0000
Subject: [PATCH 3/4] Offer hessian_input_dtype for local-Hessian calibration
 too

local_hessian's per-block Hessians (16x16 blocks of X^T X) took the same
float32 copy and float32 GEMM as GPTQ's. With hessian_input_dtype:
bfloat16 they use a batched tensor-core GEMM that reads bfloat16 and
accumulates in float32: 6-8x faster at 32K tokens x 2-4K features on an
RTX PRO 6000, with no 256-512 MiB float32 copy, and up to about 1.5e-4
of the block diagonal from the float64 sums (float32: 2e-5). The field
moves to a mixin both configs share, and one helper decides the GEMM and
warns when it is unavailable. The GPTQ error figure is corrected to up to
about 3e-5 at 32K-token batches.

imatrix keeps its float32 statistic: it is a per-column sum of squares,
not a GEMM, so tensor cores do not apply.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 CHANGELOG.rst                                 |  2 +-
 modelopt/torch/quantization/config.py         | 33 +++++---
 modelopt/torch/quantization/model_calib.py    | 24 ++++--
 .../torch/quantization/utils/calib_utils.py   | 31 +++++---
 .../quantization/test_local_hessian_gpu.py    | 77 +++++++++++++++++++
 5 files changed, 138 insertions(+), 29 deletions(-)
 create mode 100644 tests/gpu/torch/quantization/test_local_hessian_gpu.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 20f3a743c44..08b1293ea12 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -15,7 +15,7 @@ Changelog
 - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time.
 - GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch to what one decoder layer's forward fits in GPU memory, up to 32,768 tokens (64 samples at ``--calib_seq 512``), instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching.
 - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``.
-- Add ``hessian_input_dtype`` to the GPTQ algorithm config. ``bfloat16`` accumulates the float32 Hessian with a tensor-core GEMM on bfloat16 activations, up to 16x faster on GB300 and within about 1e-5 of the Hessian's diagonal; the default ``float32`` keeps the existing float32 GEMM.
+- Add ``hessian_input_dtype`` to the GPTQ and ``local_hessian`` algorithm configs. ``bfloat16`` accumulates the float32 Hessian with a tensor-core GEMM on bfloat16 activations, several times faster and without a float32 copy of each batch; the default ``float32`` keeps the existing float32 GEMM.
 - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype.
 
 *Speculative Decoding*
diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py
index c12de90e0be..88cab998218 100644
--- a/modelopt/torch/quantization/config.py
+++ b/modelopt/torch/quantization/config.py
@@ -999,7 +999,26 @@ class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
     )
 
 
-class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
+class _HessianInputDtypeConfig(ModeloptBaseConfig):
+    """The ``hessian_input_dtype`` knob, shared by GPTQ and local-Hessian calibration."""
+
+    hessian_input_dtype: Literal["float32", "bfloat16"] = ModeloptField(
+        default="float32",
+        title="Dtype the Hessian GEMM reads the activations in.",
+        description="""The Hessian is accumulated and stored in float32 either way. ``float32``
+        copies each batch of activations to float32 for the GEMM. ``bfloat16`` feeds them to a
+        tensor-core GEMM in bfloat16 that accumulates in float32: for bfloat16 activations it
+        forms the same products and only rounds the sums differently. At 32K tokens per batch
+        that is up to about 3e-5 of the Hessian's diagonal for GPTQ's full Hessian and 1.5e-4
+        for local Hessian's 16-column blocks, and the GEMM is 6-16x faster. It needs CUDA GEMMs
+        that take ``out_dtype``; without them the Hessian falls back to float32 with a
+        warning.""",
+    )
+
+
+class LocalHessianCalibConfig(
+    _SharedStatesConfig, _HessianInputDtypeConfig, QuantizeAlgorithmConfig
+):
     """Configuration for local Hessian-weighted MSE calibration.
 
     This algorithm uses activation information to optimize per-block scales for weight
@@ -1210,7 +1229,7 @@ class SVDQuantConfig(QuantizeAlgorithmConfig):
     )
 
 
-class GPTQCalibConfig(QuantizeAlgorithmConfig):
+class GPTQCalibConfig(_HessianInputDtypeConfig, QuantizeAlgorithmConfig):
     """The config for GPTQ quantization.
 
     GPTQ minimizes the layer-wise quantization error by using second-order (Hessian) information
@@ -1250,16 +1269,6 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig):
         (the GGML IQ1 and IQ2 formats) weight it by the square root of the Hessian diagonal,
         llama.cpp's imatrix statistic. Other formats ignore it.""",
     )
-    hessian_input_dtype: Literal["float32", "bfloat16"] = ModeloptField(
-        default="float32",
-        title="Dtype the Hessian GEMM reads the activations in.",
-        description="""The Hessian is accumulated and stored in float32 either way. ``float32``
-        copies each batch of activations to float32 for the GEMM. ``bfloat16`` feeds them to a
-        tensor-core GEMM in bfloat16 that accumulates in float32: for bfloat16 activations it
-        forms the same products and only rounds the sums differently, by about 1e-5 of the
-        Hessian's diagonal, and it is 9-16x faster on GB300. It needs CUDA GEMMs that take
-        ``out_dtype``; without them the Hessian falls back to float32 with a warning.""",
-    )
 
 
 class ImatrixCalibConfig(QuantizeAlgorithmConfig):
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 7cef364e8f7..047e4443189 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -74,6 +74,7 @@
     _GPTQ_HELPER_REGISTRY,
     _IMATRIX_ENCODER_REGISTRY,
     GPTQHelper,
+    bfloat16_hessian_gemm,
     track_calib_forward_memory,
 )
 
@@ -957,10 +958,13 @@ class _LocalHessianAccumulator:
     scale; the buffer is allocated lazily so never-routed experts cost nothing.
     """
 
-    def __init__(self, cout: int, cin: int, block_size: int):
+    def __init__(
+        self, cout: int, cin: int, block_size: int, input_dtype: torch.dtype = torch.float32
+    ):
         self.cout = cout
         self.cin = cin
         self.block_size = block_size
+        self.input_dtype = input_dtype
         self.num_blocks_per_cin = cin // block_size
         # Not block-divisible -> no Hessian (falls back to plain MSE).
         self.is_enabled = cin % block_size == 0
@@ -973,10 +977,15 @@ def accumulate(self, input_tensor: torch.Tensor) -> None:
         """Accumulate ``XᵀX`` per block from an activation of shape ``(..., cin)``."""
         if not self.is_enabled:
             return
-        # fp32 GEMM avoids bf16/fp16 precision loss; (cin, tokens) -> (n_blocks, bs, tokens).
-        x = input_tensor.reshape(-1, self.cin).to(torch.float32).T
+        # (cin, tokens) -> (n_blocks, bs, tokens). A float32 GEMM avoids bf16/fp16 precision loss;
+        # a bfloat16 one, accumulating in float32, is faster and needs no float32 copy.
+        bf16 = bfloat16_hessian_gemm(self.input_dtype, input_tensor)
+        x = input_tensor.reshape(-1, self.cin).to(torch.bfloat16 if bf16 else torch.float32).T
         x = x.reshape(self.num_blocks_per_cin, self.block_size, -1)
-        hessian_batch = x @ x.transpose(-1, -2)
+        if bf16:
+            hessian_batch = torch.bmm(x, x.transpose(-1, -2), out_dtype=torch.float32)
+        else:
+            hessian_batch = x @ x.transpose(-1, -2)
         if self.hessian_per_block is None:
             self.hessian_per_block = hessian_batch
         else:
@@ -1178,6 +1187,7 @@ def local_hessian_calibrate(
     block_size: int = 16,
     debug: bool = False,
     shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None,
+    hessian_input_dtype: str = "float32",
 ):
     """Calibrate weight quantizers by minimizing the Hessian-weighted error.
 
@@ -1206,6 +1216,8 @@ def local_hessian_calibrate(
         block_size: Block size for local Hessian computation (default: 16).
         debug: If True, retain the per-quantizer Hessian accumulators on the model
             (``model._local_hessian_accumulators``) for inspection.
+        hessian_input_dtype: ``"float32"`` or ``"bfloat16"``, the dtype the Hessian GEMM reads
+            the activations in.
 
     See :class:`LocalHessianCalibConfig <modelopt.torch.quantization.config.LocalHessianCalibConfig>`
     for details on the configuration options.
@@ -1227,7 +1239,9 @@ def capture(weight_quantizer, weight, input_tensor):
         input_local = input_tensor.to_local() if hasattr(input_tensor, "to_local") else input_tensor
         acc = accumulators.get(id(weight_quantizer))
         if acc is None:
-            acc = _LocalHessianAccumulator(weight.shape[0], weight.shape[1], block_size)
+            acc = _LocalHessianAccumulator(
+                weight.shape[0], weight.shape[1], block_size, getattr(torch, hessian_input_dtype)
+            )
             accumulators[id(weight_quantizer)] = acc
         acc.accumulate(input_local)
 
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index 63e617e1b17..8b045de3f15 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -140,14 +140,10 @@ def update_hessian(input, hessian, n_samples, input_dtype=torch.float32):
     scale = n_samples / (n_samples + batch_size)
     n_samples += batch_size
     alpha = 2 / n_samples
-    if (
-        input_dtype == torch.bfloat16
-        and input_flat.is_cuda
-        and _has_float32_out_gemm(input_flat.device)
-    ):
+    if bfloat16_hessian_gemm(input_dtype, input_flat):
         # BF16 products are exact in float32, so a tensor-core GEMM writing float32 differs from a
-        # float32 GEMM only in how it rounds the sums, by ~1e-5 of the Hessian's diagonal. It is
-        # 9-16x faster on GB300 and needs no float32 copy of the inputs.
+        # float32 GEMM only in how it rounds the sums: up to ~3e-5 of the Hessian's diagonal at 32K
+        # tokens per batch. It is 9-16x faster on GB300 and needs no float32 copy of the inputs.
         input_flat = input_flat.to(torch.bfloat16)
         if hessian.device == input_flat.device:
             torch.addmm(
@@ -163,8 +159,6 @@ def update_hessian(input, hessian, n_samples, input_dtype=torch.float32):
             gram = torch.mm(input_flat.t(), input_flat, out_dtype=torch.float32)
             hessian.mul_(scale).add_(gram.to(hessian.device), alpha=alpha)
         return hessian, n_samples
-    if input_dtype == torch.bfloat16:
-        _warn_bfloat16_hessian_unavailable(input_flat.device)
 
     input_flat = input_flat.float()
     if hessian.device == input_flat.device:
@@ -175,6 +169,20 @@ def update_hessian(input, hessian, n_samples, input_dtype=torch.float32):
     return hessian, n_samples
 
 
+def bfloat16_hessian_gemm(input_dtype: torch.dtype, inputs: torch.Tensor) -> bool:
+    """Whether a Hessian GEMM over ``inputs`` runs on bfloat16 with float32 output.
+
+    True when ``input_dtype`` asks for bfloat16 and the inputs' GPU has such GEMMs; when they are
+    missing, warns once per device and returns False, so the caller computes in float32.
+    """
+    if input_dtype != torch.bfloat16:
+        return False
+    if inputs.is_cuda and _has_float32_out_gemm(inputs.device):
+        return True
+    _warn_bfloat16_hessian_unavailable(inputs.device)
+    return False
+
+
 @functools.cache
 def _has_float32_out_gemm(device: torch.device) -> bool:
     """Whether GEMMs on CUDA ``device`` take bfloat16 inputs and write float32 (``out_dtype``)."""
@@ -182,6 +190,7 @@ def _has_float32_out_gemm(device: torch.device) -> bool:
         x = torch.zeros(1, 1, device=device, dtype=torch.bfloat16)
         h = torch.mm(x, x, out_dtype=torch.float32)
         torch.addmm(h, x, x, out_dtype=torch.float32, out=h)
+        torch.bmm(x[None], x[None], out_dtype=torch.float32)
     except (TypeError, RuntimeError):
         return False
     return True
@@ -190,8 +199,8 @@ def _has_float32_out_gemm(device: torch.device) -> bool:
 @functools.cache
 def _warn_bfloat16_hessian_unavailable(device: torch.device) -> None:
     warn_rank_0(
-        f"GPTQ hessian_input_dtype='bfloat16' needs CUDA GEMMs that take out_dtype, which {device} "
-        "lacks with this torch; computing its Hessians in float32."
+        f"hessian_input_dtype='bfloat16' needs CUDA GEMMs that take out_dtype, which {device} lacks "
+        "with this torch; computing its Hessians in float32."
     )
 
 
diff --git a/tests/gpu/torch/quantization/test_local_hessian_gpu.py b/tests/gpu/torch/quantization/test_local_hessian_gpu.py
new file mode 100644
index 00000000000..25318cff2cb
--- /dev/null
+++ b/tests/gpu/torch/quantization/test_local_hessian_gpu.py
@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""GPU tests for local-Hessian calibration's bfloat16 Hessian GEMM."""
+
+import pytest
+import torch
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.model_calib import _LocalHessianAccumulator
+from modelopt.torch.quantization.utils import calib_utils
+
+
+def _fast():
+    return calib_utils._has_float32_out_gemm(torch.device("cuda", torch.cuda.current_device()))
+
+
+def test_bf16_local_hessian_input_gives_the_float32_hessian():
+    """``input_dtype=torch.bfloat16`` gives the float32 per-block Hessian, and with float32-output
+    GEMMs makes no float32 copy of the inputs."""
+    torch.manual_seed(0)
+    batches = [torch.randn(4096, 1024, device="cuda", dtype=torch.bfloat16) for _ in range(3)]
+    inputs = torch.cat(batches).double().T.reshape(64, 16, -1)
+    expected = inputs @ inputs.transpose(-1, -2)
+    accumulator = _LocalHessianAccumulator(8, 1024, 16, torch.bfloat16)
+    for batch in batches:
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        base = torch.cuda.memory_allocated()
+        accumulator.accumulate(batch)
+        growth = torch.cuda.max_memory_allocated() - base
+    error = (accumulator.hessian_per_block.double() - expected).abs().max()
+    # Tensor cores round float32 sums their own way: ~1e-4 of the diagonal, not BF16's 1e-2.
+    assert error < 1e-3 * expected.diagonal(dim1=-2, dim2=-1).mean()
+    if _fast():
+        assert growth < batch.numel() * 4
+
+
+def test_local_hessian_input_dtype_reaches_the_hessian():
+    """A recipe's ``hessian_input_dtype: bfloat16`` reaches local Hessian's accumulator: its Hessian
+    matches the float32 one to bfloat16-GEMM rounding, without being the same sums."""
+    if not _fast():
+        pytest.skip("needs CUDA GEMMs that take out_dtype")
+    torch.manual_seed(0)
+    inputs = torch.randn(4, 2048, 1024, device="cuda", dtype=torch.bfloat16)
+    hessians = {}
+    for hessian_input_dtype in ("float32", "bfloat16"):
+        model = torch.nn.Linear(1024, 32, bias=False, device="cuda", dtype=torch.bfloat16)
+        quant_cfg = {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}},
+            ],
+            "algorithm": {
+                "method": "local_hessian",
+                "hessian_input_dtype": hessian_input_dtype,
+                "debug": True,
+            },
+        }
+        mtq.quantize(model, quant_cfg, forward_loop=lambda m: m(inputs))
+        (accumulator,) = model._local_hessian_accumulators.values()
+        hessians[hessian_input_dtype] = accumulator.hessian_per_block
+    diagonal = hessians["float32"].diagonal(dim1=-2, dim2=-1).mean()
+    assert not torch.equal(hessians["float32"], hessians["bfloat16"])
+    assert (hessians["float32"] - hessians["bfloat16"]).abs().max() < 1e-3 * diagonal

From eabd003ba82082c1bcb5907b3a87c5c5b22b4a18 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Tue, 27 Oct 2026 06:42:12 +0000
Subject: [PATCH 4/4] Test the bfloat16 Hessian fallback on CPU

On CPU, hessian_input_dtype: bfloat16 falls back to the float32 Hessian.
New CPU tests for GPTQ and local Hessian check that it then gives the
same result as float32 and warns exactly once over several batches.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 tests/unit/torch/quantization/test_gptq.py    | 32 +++++++++++++++++++
 .../torch/quantization/test_local_hessian.py  | 24 ++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tests/unit/torch/quantization/test_gptq.py b/tests/unit/torch/quantization/test_gptq.py
index 5744348d67e..6e5b01190ec 100644
--- a/tests/unit/torch/quantization/test_gptq.py
+++ b/tests/unit/torch/quantization/test_gptq.py
@@ -15,6 +15,8 @@
 
 """CPU unit tests for GPTQ utilities."""
 
+import copy
+import warnings
 from functools import partial
 
 import pytest
@@ -153,6 +155,36 @@ def forward_loop(model):
     assert len(calls) == forwards
 
 
+def test_bf16_hessian_input_on_cpu_gives_the_float32_result_and_warns_once():
+    """CPU GEMMs take no ``out_dtype``, so ``hessian_input_dtype: bfloat16`` falls back to the
+    float32 Hessian, warning once however many batches it reads."""
+    calib_utils._warn_bfloat16_hessian_unavailable.cache_clear()
+    torch.manual_seed(0)
+    model = torch.nn.Linear(32, 8, bias=False)
+    batches = [torch.randn(4, 32) for _ in range(3)]
+    weights = {}
+    for hessian_input_dtype in ("float32", "bfloat16"):
+        quantized = copy.deepcopy(model)
+        config = {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}},
+            ],
+            "algorithm": {
+                "method": "gptq",
+                "block_size": 16,
+                "hessian_input_dtype": hessian_input_dtype,
+            },
+        }
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            mtq.quantize(quantized, config, lambda m: [m(batch) for batch in batches])
+        fallbacks = [w for w in caught if "hessian_input_dtype" in str(w.message)]
+        assert len(fallbacks) == (hessian_input_dtype == "bfloat16")
+        weights[hessian_input_dtype] = quantized.weight
+    assert torch.equal(weights["float32"], weights["bfloat16"])
+
+
 def _gptq_hessian_dp_sync_helper(rank, size, token_counts):
     torch.manual_seed(0)
     inputs = [torch.randn(n, 8) for n in token_counts]
diff --git a/tests/unit/torch/quantization/test_local_hessian.py b/tests/unit/torch/quantization/test_local_hessian.py
index d8b4c4eecec..8cf89707126 100644
--- a/tests/unit/torch/quantization/test_local_hessian.py
+++ b/tests/unit/torch/quantization/test_local_hessian.py
@@ -39,6 +39,7 @@
     _QUANT_FUNCTIONAL_BACKENDS,
     register_quant_backend,
 )
+from modelopt.torch.quantization.utils import calib_utils
 
 # Weight-only INT8 per-channel; calibration is re-run explicitly per test.
 INT8_WEIGHT_CFG = {
@@ -132,6 +133,29 @@ def test_refines_amax_beyond_max_and_plain_mse(self):
         assert any(not torch.allclose(lh[n], max_amax[n]) for n in lh)  # refined past max-cal
         assert any(not torch.allclose(lh[n], mse[n]) for n in lh)  # Hessian changed the choice
 
+    def test_bf16_hessian_input_on_cpu_gives_the_float32_result_and_warns_once(self):
+        """CPU GEMMs take no ``out_dtype``, so ``hessian_input_dtype="bfloat16"`` falls back to
+        the float32 Hessians, warning once however many batches they read."""
+        calib_utils._warn_bfloat16_hessian_unavailable.cache_clear()
+        amaxes = {}
+        for hessian_input_dtype in ("float32", "bfloat16"):
+            torch.manual_seed(0)
+            model = SimpleLinear()
+            mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=_make_forward_loop())
+            with warnings.catch_warnings(record=True) as caught:
+                warnings.simplefilter("always")
+                local_hessian_calibrate(
+                    model,
+                    _make_forward_loop(),
+                    fp8_scale_sweep=False,
+                    hessian_input_dtype=hessian_input_dtype,
+                )
+            fallbacks = [w for w in caught if "hessian_input_dtype" in str(w.message)]
+            assert len(fallbacks) == (hessian_input_dtype == "bfloat16")
+            amaxes[hessian_input_dtype] = _weight_amaxes(model)
+        for name, amax in amaxes["float32"].items():
+            assert torch.equal(amax, amaxes["bfloat16"][name])
+
     def test_warns_with_module_name_when_cin_not_divisible(self):
         class _OddModel(nn.Module):
             def __init__(self):