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 from a fresh ``layerwise.checkpoint_dir`` for any recipe with ``layerwise.enable: true``, including the shipped ``nvfp4_experts_only-kv_fp8_layerwise*`` and ``nvfp4_default-kv_none-gptq`` recipes; algorithms that derive weight scales or weight updates from activations (``gptq``, ``awq_lite``, ``awq_clip``, ``local_hessian``, ``smoothquant``) change their exported weights too, not just activation scales.
 
 0.46 (2026-08-17)
 ^^^^^^^^^^^^^^^^^

modelopt/torch/quantization/model_calib.py

@@ -2125,21 +2125,6 @@ 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)
 
             is_last = layer_idx + 1 >= num_layers

modelopt/torch/quantization/utils/layerwise_calib.py

@@ -51,6 +51,27 @@ class _EarlyStopForwardError(Exception):
     """Raised to halt the forward pass after capturing layer inputs."""
 
 
+def _is_kv_cache(obj: Any) -> bool:
+    """Duck-typed ``transformers.Cache``, avoiding a transformers import here.
+
+    Matched by shape, not name: the keyword is ``past_key_values`` on HF-native layers
+    but ``past_key_value`` on older remote-code ones (Kimi-K2), and may be positional.
+    """
+    return hasattr(obj, "update") and hasattr(obj, "get_seq_length")
+
+
+def _with_empty_kv_cache(args: tuple, kwargs_input: dict) -> tuple[tuple, dict]:
+    """Strip the attention cache, so a replay cannot attend over its own earlier writes.
+
+    Not ``Cache.reset()``: it zeroes the keys and values but keeps them at full length,
+    so the replay attends over an all-zero cache instead of none.
+    """
+    return (
+        tuple(None if _is_kv_cache(a) else a for a in args),
+        {k: (None if _is_kv_cache(v) else v) for k, v in kwargs_input.items()},
+    )
+
+
 @dataclass
 class _LayerCalibState:
     """Mutable per-layer state used during layerwise calibration.
@@ -241,7 +262,7 @@ def _patched_forward(self, *args, **kwargs):
                 return output
 
             if info.mode == "capture":
-                info.collected_inputs.append((args, kwargs))
+                info.collected_inputs.append(_with_empty_kv_cache(args, kwargs))
                 raise _EarlyStopForwardError()
 
             return self._original_forward(*args, **kwargs)
@@ -435,6 +456,8 @@ def get_first_layer_inputs(
             for i in range(start_layer):
                 self._swap_to_dummy(i)
             layer = self._decoder_layers[start_layer]
+            # Not fed by capture: an older checkpoint can still carry a live cache.
+            resumed_inputs = [_with_empty_kv_cache(a, kw) for a, kw in resumed_inputs]
             layer._layerwise_calib.collected_inputs = resumed_inputs
             layer._layerwise_calib.mode = "original"
             return resumed_inputs

tests/unit/torch/quantization/test_layerwise_calibrate.py

@@ -22,11 +22,18 @@
 import pytest
 import torch
 import torch.nn as nn
+from _test_utils.torch.transformers_models import get_tiny_llama
+from transformers.cache_utils import DynamicCache
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.model_calib import layerwise_calibrate
 from modelopt.torch.quantization.nn import TensorQuantizer
-from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector, _SkipLayer
+from modelopt.torch.quantization.utils.layerwise_calib import (
+    LayerActivationCollector,
+    _is_kv_cache,
+    _SkipLayer,
+    _with_empty_kv_cache,
+)
 
 
 class _DecoderBlock(nn.Module):
@@ -830,6 +837,45 @@ 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():
+    """Same equivalence as above, on a model that actually has a KV cache.
+
+    The toy model used above has none, so it cannot go stale -- which is why a
+    replay attending over its own writes shipped, collapsing ``o_proj``'s input
+    amax to 0.0 on every layer but the last.
+    """
+    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
@@ -1036,6 +1082,58 @@ def crashing_torch_save(obj, path, *args, **kwargs):
     assert manifest["last_completed_layer"] == 1, f"manifest leaked mid-window state: {manifest}"
 
 
+def test_capture_stores_no_kv_cache():
+    """Captured inputs must carry no cache.
+
+    Asserted structurally because the amax comparison cannot see it: a retained cache
+    only accumulates across replays, and max calibration's max absorbs that on a small
+    model, so amaxes match while the invariant is broken.
+    """
+    model = get_tiny_llama(num_hidden_layers=3).eval()
+    collector = LayerActivationCollector(model)
+    collector._patch_all_layers(decoder_layers=model.model.layers)
+    try:
+        captured = collector.get_input_activations(
+            model.model.layers[0], lambda m: m(torch.randint(0, 32, (2, 8)))
+        )
+    finally:
+        collector._unpatch_all_layers()
+
+    assert captured, "nothing was captured"
+    live = [
+        k
+        for args, kwargs in captured
+        for k, v in [*enumerate(args), *kwargs.items()]
+        if _is_kv_cache(v)
+    ]
+    assert not live, f"captured inputs still hold a KV cache at {live}"
+
+
+def test_with_empty_kv_cache_matches_by_shape_not_by_name():
+    """The cache must be cleared however it reaches the layer.
+
+    Kimi-K2's remote code passes it as ``past_key_value`` (see
+    ``modelopt/torch/speculative/utils.py``), so matching only the plural keyword
+    would silently no-op on a model layerwise calibration exists for.
+    """
+    cache = DynamicCache()
+    hidden = torch.randn(1, 4)
+
+    for args, kwargs in (
+        ((), {"past_key_values": cache}),
+        ((), {"past_key_value": cache}),
+        ((hidden, cache), {}),
+    ):
+        out_args, out_kwargs = _with_empty_kv_cache(args, kwargs)
+        assert not any(_is_kv_cache(a) for a in out_args)
+        assert not any(_is_kv_cache(v) for v in out_kwargs.values())
+
+    # Non-cache values pass through untouched.
+    out_args, out_kwargs = _with_empty_kv_cache((hidden,), {"attention_mask": None})
+    assert out_args[0] is hidden
+    assert out_kwargs == {"attention_mask": None}
+
+
 def test_layerwise_checkpoint_mismatch_save_every_raises(monkeypatch, tmp_path):
     """Resuming with a different ``save_every`` than the checkpoint was produced
     with must raise — the on-disk window layout assumes a fixed value.