fix(quantization): warn once when calibration runs with a KV cache

CHANGELOG.rst

@@ -43,6 +43,7 @@ Changelog
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
 - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs.
+- Fix layerwise calibration miscalibrating every module downstream of attention within a decoder layer; models using an explicit sliding-window mask raised a shape mismatch instead. Re-run calibration for any recipe with ``layerwise.enable: true``, including the shipped ``nvfp4_experts_only-kv_fp8_layerwise*`` and ``nvfp4_default-kv_none-gptq`` recipes — under ``gptq`` and ``awq_lite`` the exported weights change too, not just activation scales.
 
 0.46 (2026-08-17)
 ^^^^^^^^^^^^^^^^^

modelopt/torch/quantization/model_calib.py

@@ -34,6 +34,7 @@
 from modelopt.torch.quantization.utils.layerwise_calib import (
     LayerActivationCollector,
     _CheckpointState,
+    _with_empty_kv_cache,
 )
 from modelopt.torch.utils import print_rank_0, warn_rank_0
 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master
@@ -2125,22 +2126,7 @@ def _set_layer_status(status: str):
 
             def _layer_forward_loop(m, _inputs=layer_inputs):
                 for args, kwargs_input in _inputs:
-                    # Reset past_key_values to prevent the KV cache from
-                    # accumulating across multiple forward replays (e.g.
-                    # max_calibrate then Hessian collection in GPTQ).
-                    # The layer doesn't need stale KV data — each replay
-                    # should start with a fresh cache.
-                    if (
-                        "past_key_values" in kwargs_input
-                        and kwargs_input["past_key_values"] is not None
-                    ):
-                        kwargs_input = dict(kwargs_input)
-                        cache = kwargs_input["past_key_values"]
-                        if hasattr(cache, "reset"):
-                            cache.reset()
-                        else:
-                            kwargs_input["past_key_values"] = None
-                    m(*args, **kwargs_input)
+                    m(*args, **_with_empty_kv_cache(kwargs_input))
 
             is_last = layer_idx + 1 >= num_layers
 

modelopt/torch/quantization/utils/layerwise_calib.py

@@ -51,6 +51,21 @@ class _EarlyStopForwardError(Exception):
     """Raised to halt the forward pass after capturing layer inputs."""
 
 
+def _with_empty_kv_cache(kwargs_input: dict) -> dict:
+    """Drop any attention cache, so a replay does not attend over its own earlier writes.
+
+    Both replay paths -- the patched ``run`` forward below and ``calib_func``'s replay in
+    ``layerwise_calibrate`` -- consume the same captured kwargs, so each clears the cache
+    itself rather than relying on the other not having written to it.
+
+    Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full
+    length, leaving the replay attending over an all-zero cache instead of none.
+    """
+    if kwargs_input.get("past_key_values") is None:
+        return kwargs_input
+    return {**kwargs_input, "past_key_values": None}
+
+
 @dataclass
 class _LayerCalibState:
     """Mutable per-layer state used during layerwise calibration.
@@ -236,7 +251,7 @@ def _patched_forward(self, *args, **kwargs):
                     f"Layer {info.name} is in 'run' mode but has no cached inputs to replay."
                 )
                 real_args, real_kwargs = info.cached_inputs.popleft()
-                output = self._original_forward(*real_args, **real_kwargs)
+                output = self._original_forward(*real_args, **_with_empty_kv_cache(real_kwargs))
                 info.output_meta = LayerActivationCollector._extract_output_meta(output)
                 return output
 

tests/unit/torch/quantization/test_layerwise_calibrate.py

@@ -22,6 +22,7 @@
 import pytest
 import torch
 import torch.nn as nn
+from _test_utils.torch.transformers_models import get_tiny_llama
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.model_calib import layerwise_calibrate
@@ -830,6 +831,49 @@ def fwd(m):
     _assert_amax_close(_collect_amax(model_lw), seq_amax, "layerwise vs sequential")
 
 
+def test_layerwise_replay_does_not_attend_over_its_own_kv_cache():
+    """A layer replayed during calibration must not see the keys and values its
+    own earlier run wrote.
+
+    The equivalence test above uses a model with no attention cache, so it cannot
+    reach this: only the captured ``past_key_values`` makes a replay stateful.
+    With the cache left in place, ``o_proj`` saw an all-zero attention output on
+    every layer but the last (the one with no preceding capture pass) and its
+    input amax collapsed to exactly 0.0, while ``down_proj`` picked up a
+    plausible but wrong value from the residual alone.
+    """
+    calib_data = [torch.randint(0, 32, (2, 8)) for _ in range(2)]
+
+    def fwd(m):
+        for batch in calib_data:
+            m(batch)
+
+    def calibrate(algorithm):
+        torch.manual_seed(0)
+        model = get_tiny_llama(num_hidden_layers=4).eval()
+        cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+        cfg["algorithm"] = algorithm
+        mtq.quantize(model, cfg, forward_loop=fwd)
+        return model
+
+    sequential = calibrate({"method": "max"})
+    layerwise = calibrate({"method": "max", "layerwise": {"enable": True}})
+
+    expected = _collect_amax(sequential)
+    assert expected, "sequential calibration populated no amax values"
+    layerwise_amax = _collect_amax(layerwise)
+    _assert_amax_close(layerwise_amax, expected, "layerwise vs sequential (KV cache)")
+
+    # Pinned separately from the comparison: a future regression that made both
+    # paths collapse to zero would still satisfy the equality above.
+    collapsed = [
+        name
+        for name, amax in layerwise_amax.items()
+        if name.endswith("input_quantizer") and not torch.count_nonzero(amax)
+    ]
+    assert not collapsed, f"activation amax collapsed to zero: {collapsed}"
+
+
 def test_layerwise_no_qdq_captures_inputs_before_calib_func_mutates_weights(monkeypatch):
     """A destructive ``calib_func`` (zeros weights) must not affect what is
     captured for downstream layers under ``qdq_from_prev=False`` — otherwise