From 0cbc2c10a47c74c1f9293a56b7af530d5d32e92e Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Thu, 30 Jul 2026 16:54:55 -0700
Subject: [PATCH 01/15] Add layer-wise KV cache AutoQuant search

Assisted-by: OpenAI Codex
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   1 +
 examples/hf_ptq/README.md                     |  22 +-
 examples/hf_ptq/hf_ptq.py                     |  46 ++
 modelopt/recipe/config.py                     |  60 +-
 modelopt/torch/export/convert_hf_config.py    |   8 +
 .../torch/export/quant_aware_conversion.py    |   7 +-
 modelopt/torch/export/quant_utils.py          |  69 ++-
 modelopt/torch/export/unified_export_hf.py    |  41 +-
 .../torch/quantization/kv_cache_auto_quant.py | 519 ++++++++++++++++++
 modelopt/torch/quantization/model_quant.py    |  69 +++
 .../kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml  |  44 ++
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  17 +
 tests/unit/recipe/test_loader.py              |  16 +
 .../torch/export/test_convert_hf_config.py    |  42 ++
 .../torch/export/test_get_quantization.py     |  71 ++-
 .../export/test_quant_aware_conversion.py     |   2 +
 .../quantization/test_kv_cache_auto_quant.py  | 334 +++++++++++
 17 files changed, 1337 insertions(+), 31 deletions(-)
 create mode 100644 modelopt/torch/quantization/kv_cache_auto_quant.py
 create mode 100644 modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
 create mode 100644 tests/unit/torch/export/test_convert_hf_config.py
 create mode 100644 tests/unit/torch/quantization/test_kv_cache_auto_quant.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 4580a045f39..a1a2c8a4c41 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -9,6 +9,7 @@ Changelog
 *Quantization*
 
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint.
+- Add layer-wise KV-cache AutoQuantize through ``mtq.auto_quantize_kv_cache`` and ``constraints.kv_effective_bits``. It measures isolated full-vocabulary forward KL for caller-supplied K/V formats, solves a width-weighted additive storage-constrained recipe across eligible layers, preserves search-disabled layers in their existing format, exports the selected per-attention mapping in unified HF checkpoints, and writes a JSON sensitivity report alongside the checkpoint. A cast-mode FP8/NVFP4 recipe at 5.4 bits/scalar is included.
 
 *Megatron Framework (M-LM / M-Bridge)*
 
diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index fea69221825..5ce321cd2b9 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -422,8 +422,26 @@ leaving the original recipe unchanged.
 For models without backprop support (e.g. Llama-4), use the `kl_div` scoring method — see the shipped
 `general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits` recipe.
 
-KV cache is applied as a uniform post-step, not part of the per-layer search. An AutoQuantize recipe
-falls back to `--kv_cache_qformat` (default `fp8_cast`) unless it sets an explicit `kv_cache` field.
+Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall back to
+`--kv_cache_qformat` (default `fp8_cast`) unless they set an explicit `kv_cache` field.
+
+KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
+`candidate_formats` are complete K/V cache configs with exact config-level `effective_bits`;
+BF16 is used only as the isolated-KL reference, not as a solver choice. The shipped canary recipe
+searches cast-mode FP8 (8.0 bits/scalar) and packed NVFP4 (4.5 bits/scalar) at 5.4 bits/scalar:
+
+```bash
+python hf_ptq.py \
+  --pyt_ckpt_path Qwen/Qwen3-1.7B \
+  --recipe general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
+  --auto_quantize_checkpoint /path/to/kv_autoquant.pth \
+  --export_path /path/to/qwen3-1.7b-mixed-kv
+```
+
+Cast candidates use constant amax and skip the PTQ calibration forward. Unified HF export records
+the selected formats in `kv_cache_quantized_layers` and writes the JSON-safe sensitivity report to
+`kv_cache_auto_quantize_report.json`; `--auto_quantize_checkpoint` stores the resumable raw search
+state.
 
 The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an
 interrupted search (skips re-scoring):
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index a56a62b54b5..0bfa07cdc32 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -338,6 +338,24 @@ def _mtq_candidate_formats(formats) -> list[dict]:
     return quantization_formats
 
 
+def _mtq_kv_candidate_formats(formats) -> list[tuple[dict, str]]:
+    """Translate format-agnostic KV candidates while preserving useful preset names."""
+    candidates = []
+    for idx, fmt in enumerate(formats):
+        quant_cfg = fmt.model_dump(exclude_none=True)
+        candidate_quantizers = quant_cfg.get("quant_cfg", [])
+        name = None
+        for preset_name, preset in KV_QUANT_CFG_CHOICES.items():
+            normalized_preset_quantizers = (
+                type(fmt)(**preset).model_dump(exclude_none=True).get("quant_cfg", [])
+            )
+            if normalized_preset_quantizers == candidate_quantizers:
+                name = preset_name
+                break
+        candidates.append((quant_cfg, name or f"KV_CACHE_FORMAT_{idx}"))
+    return candidates
+
+
 def _mtq_inputs_from_auto_quantize_config(
     aq_config, args: argparse.Namespace, fixed_quantize_config=None
 ) -> dict:
@@ -349,6 +367,16 @@ def _mtq_inputs_from_auto_quantize_config(
     to ``--kv_cache_qformat`` when the recipe omits it.
     """
     constraints = aq_config.constraints.model_dump(exclude_none=True)
+    is_kv_search = aq_config.constraints.kv_effective_bits is not None
+    if is_kv_search:
+        return {
+            "search_domain": "kv_cache",
+            "constraints": {"kv_effective_bits": constraints["kv_effective_bits"]},
+            "quantization_formats": _mtq_kv_candidate_formats(aq_config.candidate_formats),
+            "disabled_layers": aq_config.disabled_layers,
+            "method": aq_config.auto_quantize_method,
+            "score_size": aq_config.score_size,
+        }
     # cost_excluded_layers (sibling of disabled_layers) maps to the mtq cost key: these layers are
     # kept out of the bit-budget denominator (cost_weight 0) — e.g. VL vision towers — distinct from
     # disabled_layers, which removes them from the search.
@@ -380,6 +408,7 @@ def _mtq_inputs_from_auto_quantize_config(
         for search_space in aq_config.module_search_spaces
     ]
     return {
+        "search_domain": "weight",
         "constraints": constraints,
         "quantization_formats": quantization_formats,
         "fixed_quantization_config": fixed_quantization_config,
@@ -469,6 +498,23 @@ def forward_step(model, batch):
             f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'"
         )
 
+    if inputs["search_domain"] == "kv_cache":
+        language_model, _ = mtq.auto_quantize_kv_cache(
+            language_model,
+            constraints=inputs["constraints"],
+            data_loader=calib_dataloader,
+            forward_step=forward_step,
+            quantization_formats=inputs["quantization_formats"],
+            num_calib_steps=len(calib_dataloader),
+            num_score_steps=min(
+                len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)
+            ),
+            verbose=True,
+            disabled_layers=inputs["disabled_layers"],
+            checkpoint=args.auto_quantize_checkpoint,
+        )
+        return language_model
+
     language_model, _ = mtq.auto_quantize(
         language_model,
         constraints=inputs["constraints"],
diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index 2ca8f4f462b..c4dc89f1841 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -149,10 +149,21 @@ def _validate_active_moe_expert_ratio(cls, v: float | None) -> float | None:
 class AutoQuantizeConstraints(ModeloptBaseConfig):
     """LP search constraints + cost model; matches the ``mtq.auto_quantize`` constraints dict."""
 
-    effective_bits: float = ModeloptField(
-        default=4.8,
+    effective_bits: float | None = ModeloptField(
+        default=None,
         title="Effective bits per weight",
-        description="Average weight-storage bits target for the LP, in (0, 16].",
+        description=(
+            "Average weight-storage bits target for the LP, in (0, 16]. Defaults to 4.8 "
+            "when neither bit constraint is specified."
+        ),
+    )
+    kv_effective_bits: float | None = ModeloptField(
+        default=None,
+        title="Effective bits per KV-cache scalar",
+        description=(
+            "Average KV-cache storage bits target for layer-wise KV AutoQuant, in (0, 16]. "
+            "Exactly one of effective_bits and kv_effective_bits may be set."
+        ),
     )
     cost_model: Literal["weight", "active_moe"] = ModeloptField(
         default="weight",
@@ -165,13 +176,36 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
         description="Extra cost-model parameters; omit for the 'weight' cost model.",
     )
 
-    @field_validator("effective_bits")
+    @model_validator(mode="before")
     @classmethod
-    def _validate_effective_bits(cls, v: float) -> float:
-        if not (0 < v <= 16):
+    def _default_weight_constraint(cls, data):
+        if isinstance(data, dict):
+            data = dict(data)
+            if "effective_bits" not in data and "kv_effective_bits" not in data:
+                data["effective_bits"] = 4.8
+        return data
+
+    @field_validator("effective_bits", "kv_effective_bits")
+    @classmethod
+    def _validate_effective_bits(cls, v: float | None) -> float | None:
+        if v is not None and not (0 < v <= 16):
             raise ValueError(f"effective_bits must be in (0, 16], got {v}")
         return v
 
+    @model_validator(mode="after")
+    def _exactly_one_bit_constraint(self):
+        if (self.effective_bits is None) == (self.kv_effective_bits is None):
+            raise ValueError(
+                "Exactly one of effective_bits and kv_effective_bits must be specified."
+            )
+        if self.kv_effective_bits is not None and (
+            self.cost_model != "weight" or self.cost is not None
+        ):
+            raise ValueError(
+                "KV-cache AutoQuant does not support weight or active-MoE cost settings."
+            )
+        return self
+
 
 class AutoQuantizeModuleSearchSpace(ModeloptBaseConfig):
     """Candidate formats selectable for modules matching one or more name patterns."""
@@ -272,6 +306,20 @@ def _has_search_space(self):
                 "auto_quantize requires candidate_formats or at least one module_search_spaces "
                 "entry. For uniform quantization, use a PTQ recipe instead."
             )
+        if self.constraints.kv_effective_bits is not None:
+            if self.auto_quantize_method != "kl_div":
+                raise ValueError(
+                    "KV-cache AutoQuant currently requires auto_quantize_method=kl_div."
+                )
+            if self.module_search_spaces:
+                raise ValueError(
+                    "KV-cache AutoQuant uses one candidate space for all eligible attention "
+                    "layers; module_search_spaces is not supported."
+                )
+            if self.kv_cache is not None:
+                raise ValueError(
+                    "KV-cache AutoQuant candidate_formats replace the uniform kv_cache post-step."
+                )
         return self
 
 
diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py
index 45fa0c30f3b..534aefe3ba8 100644
--- a/modelopt/torch/export/convert_hf_config.py
+++ b/modelopt/torch/export/convert_hf_config.py
@@ -269,6 +269,14 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An
     if kv_cache_quant_algo:
         if kv_cache_quant_algo == "FP8":
             new_config["kv_cache_scheme"] = {"dynamic": False, "num_bits": 8, "type": "float"}
+        elif kv_cache_quant_algo == "MIXED_PRECISION":
+            new_config["kv_cache_quant_algo"] = kv_cache_quant_algo
+            new_config["kv_cache_quantized_layers"] = original_quantization_details.get(
+                "kv_cache_quantized_layers", {}
+            )
+            new_config["kv_cache_schema_version"] = original_quantization_details.get(
+                "kv_cache_schema_version", 1
+            )
         else:
             # TODO: Handle other kv cache quantization algorithms
             new_config["kv_cache_scheme"] = kv_cache_quant_algo
diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index 2e32f123869..fa594f104c4 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -278,7 +278,7 @@ def _map(name: str) -> str:
 
 
 def revert_quant_config_names(quantization: dict, mapper) -> None:
-    """Revert ``exclude_modules`` / ``quantized_layers`` keys to hub names, in place.
+    """Revert layer-reference keys to hub names, in place.
 
     ``mapper`` is the callable from :func:`build_reverse_name_mapper` (a no-op when
     ``None``). Applies to the ModelOpt ``{"quantization": {...}}`` sub-dict before it is
@@ -293,6 +293,11 @@ def revert_quant_config_names(quantization: dict, mapper) -> None:
     quantized_layers = quantization.get("quantized_layers")
     if isinstance(quantized_layers, dict) and quantized_layers:
         quantization["quantized_layers"] = {mapper(k): v for k, v in quantized_layers.items()}
+    kv_cache_quantized_layers = quantization.get("kv_cache_quantized_layers")
+    if isinstance(kv_cache_quantized_layers, dict) and kv_cache_quantized_layers:
+        quantization["kv_cache_quantized_layers"] = {
+            mapper(k): v for k, v in kv_cache_quantized_layers.items()
+        }
 
 
 def _assert_experts_pre_expanded(
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index c86af3aa9f5..75df85c5cfd 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -1054,7 +1054,7 @@ def _postprocess_single_tensor(
 def postprocess_state_dict(
     state_dict: dict,
     maxbound: float,
-    quantization: str | None,
+    quantization: str | dict[str, dict[str, str]] | None,
     is_modelopt_qlora: bool = False,
     tied_map: "TiedWeightMap | None" = None,
 ) -> dict:
@@ -1063,7 +1063,8 @@ def postprocess_state_dict(
     Args:
         state_dict: The full model state_dict.
         maxbound: The maximum bound value for the output quantizer.
-        quantization: The KV cache quantization format.
+        quantization: The uniform KV cache quantization format, or a per-attention-layer
+            ``{layer_name: {"quant_algo": ...}}`` mapping for mixed precision.
         is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model.
         tied_map: Optional :class:`TiedWeightMap`. When provided, tied-weight
             dedup is authoritative and name-based: a declared alias key whose canonical
@@ -1084,6 +1085,18 @@ def _export_key(key: str) -> str:
 
     post_state_dict = {}
 
+    def _kv_quantization_for_key(key: str) -> str | None:
+        if not isinstance(quantization, dict):
+            return quantization
+        matches = [
+            (layer_name, layer_config.get("quant_algo"))
+            for layer_name, layer_config in quantization.items()
+            if key == layer_name or key.startswith(layer_name + ".")
+        ]
+        if not matches:
+            return None
+        return max(matches, key=lambda item: len(item[0]))[1]
+
     for key, value in state_dict.items():
         # Skip problematic parameters for specific model architectures, e.g., Nemotron Nano VL models
         if key == "vision_model.radio_model.summary_idxs":
@@ -1101,15 +1114,18 @@ def _export_key(key: str) -> str:
                 prefix = key[: -len(old_suffix)]
 
                 if "_amax" in key:
-                    assert quantization in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], (
-                        "Invalid KV cache quantization format."
-                    )
+                    layer_quantization = _kv_quantization_for_key(key)
+                    assert layer_quantization in [
+                        KV_CACHE_FP8,
+                        KV_CACHE_NVFP4,
+                        KV_CACHE_NVFP4_AFFINE,
+                    ], "Invalid KV cache quantization format."
                     assert maxbound > 0, "Maxbound must be greater than zero."
 
                     value = value.float() / maxbound
 
                     # Warn if scale exceeds threshold
-                    if quantization == KV_CACHE_FP8 and value.item() > 0.5:
+                    if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5:
                         logger.warning(
                             "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                         )
@@ -1601,7 +1617,7 @@ def get_quant_config(
     block_size = None
 
     # Create base config
-    quant_config = {
+    quant_config: dict[str, Any] = {
         "producer": {
             "name": "modelopt",
             "version": __version__,
@@ -1619,7 +1635,9 @@ def get_quant_config(
     # It also holds awq_block_size information for applicable layers.
     layer_config_dict = {}
 
-    kv_cache_format = QUANTIZATION_NONE
+    kv_cache_formats: set[str] = set()
+    kv_cache_quantized_layers: dict[str, dict[str, str]] = {}
+    kv_cache_eligible_layers = 0
     for name, module in dict(model.named_modules()).items():
         # Check for standard quantizers or any quantizers from weight attributes
         weight_names = list(weight_attr_names(module))
@@ -1675,18 +1693,22 @@ def get_quant_config(
         not_enabled = SimpleNamespace(is_enabled=False)
 
         # Find kv cache quant format
+        has_kv_quantizers = all(
+            hasattr(module, quantizer_name)
+            for quantizer_name in ("k_bmm_quantizer", "v_bmm_quantizer")
+        )
+        if has_kv_quantizers:
+            kv_cache_eligible_layers += 1
+
         if (
             getattr(module, "k_bmm_quantizer", not_enabled).is_enabled
             or getattr(module, "v_bmm_quantizer", not_enabled).is_enabled
             or getattr(module, "output_quantizer", not_enabled).is_enabled
         ):
             module_kv_quant = get_kv_cache_dtype(module)
-            if kv_cache_format == QUANTIZATION_NONE:
-                kv_cache_format = module_kv_quant
-            else:
-                assert kv_cache_format == module_kv_quant, (
-                    "Do not support mixed precision kv cache quantization"
-                )
+            if module_kv_quant != QUANTIZATION_NONE:
+                kv_cache_formats.add(module_kv_quant)
+                kv_cache_quantized_layers[name] = {"quant_algo": module_kv_quant}
 
     # MoE routers/gates are intentionally kept in original precision. On transformers>=5.0 they
     # are not nn.Linear modules (e.g. TopKRouter), never receive a quantizer, and would otherwise
@@ -1699,8 +1721,23 @@ def get_quant_config(
     # Process per layer quantization config dict
     quant_config["quantization"].update(process_layer_quant_config(layer_config_dict))
 
-    if kv_cache_format is not None:
-        quant_config["quantization"]["kv_cache_quant_algo"] = kv_cache_format
+    all_kv_layers_quantized = (
+        kv_cache_eligible_layers > 0 and len(kv_cache_quantized_layers) == kv_cache_eligible_layers
+    )
+    if len(kv_cache_formats) == 1 and all_kv_layers_quantized:
+        quant_config["quantization"]["kv_cache_quant_algo"] = next(iter(kv_cache_formats))
+    elif kv_cache_quantized_layers:
+        weight_quant_algo = quant_config["quantization"].get("quant_algo")
+        if weight_quant_algo not in (None, "MIXED_PRECISION"):
+            raise NotImplementedError(
+                "Mixed-precision KV-cache export with a uniform quantized-weight format is "
+                "not supported yet. Use BF16 weights or a mixed-weight AutoQuant recipe."
+            )
+        quant_config["quantization"]["quant_algo"] = "MIXED_PRECISION"
+        quant_config["quantization"].setdefault("quantized_layers", {})
+        quant_config["quantization"]["kv_cache_quant_algo"] = "MIXED_PRECISION"
+        quant_config["quantization"]["kv_cache_quantized_layers"] = kv_cache_quantized_layers
+        quant_config["quantization"]["kv_cache_schema_version"] = 1
 
     return quant_config
 
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 77429b1cfaf..814e446ba65 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -15,6 +15,7 @@
 
 """Code that export quantized Hugging Face models for deployment."""
 
+import copy
 import json
 import re
 import tempfile
@@ -1010,12 +1011,18 @@ def _export_transformers_checkpoint(
 
     # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization.
     kv_cache_max_bound = 448
-    kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"]
+    quantization_details = quant_config["quantization"]
+    kv_cache_format = quantization_details["kv_cache_quant_algo"]
+    kv_cache_postprocess_config = (
+        quantization_details["kv_cache_quantized_layers"]
+        if kv_cache_format == "MIXED_PRECISION"
+        else kv_cache_format
+    )
 
     quantized_state_dict = postprocess_state_dict(
         quantized_state_dict,
         kv_cache_max_bound,
-        kv_cache_format,
+        kv_cache_postprocess_config,
         is_modelopt_qlora,
         tied_map=tied_map,
     )
@@ -1461,6 +1468,7 @@ def _write_hf_export_config(
     model: nn.Module,
     hf_quant_config: dict | None,
     export_dir: Path,
+    name_mapper: Callable[[str], str] | None = None,
 ) -> None:
     """Write hf_quant_config.json (if quantized) and embed quantization_config into config.json."""
     quantization_details = (hf_quant_config or {}).get("quantization", {})
@@ -1474,6 +1482,29 @@ def _write_hf_export_config(
             json.dump(hf_quant_config, file, indent=4)
         quantization_config = convert_hf_quant_config_format(hf_quant_config)
 
+        kv_autoquant_report = next(
+            (
+                getattr(module, "_modelopt_kv_cache_auto_quantize_state")
+                for module in model.modules()
+                if hasattr(module, "_modelopt_kv_cache_auto_quantize_state")
+            ),
+            None,
+        )
+        if kv_autoquant_report is not None:
+            kv_autoquant_report = copy.deepcopy(kv_autoquant_report)
+            if name_mapper is not None:
+                kv_autoquant_report["layers"] = {
+                    name_mapper(name): value
+                    for name, value in kv_autoquant_report["layers"].items()
+                }
+                signature_layers = kv_autoquant_report.get("search_signature", {}).get(
+                    "layers", []
+                )
+                for layer in signature_layers:
+                    layer["name"] = name_mapper(layer["name"])
+            with open(f"{export_dir}/kv_cache_auto_quantize_report.json", "w") as file:
+                json.dump(kv_autoquant_report, file, indent=4)
+
     original_config = f"{export_dir}/config.json"
     with open(original_config) as file:
         config_data = json.load(file)
@@ -1571,6 +1602,7 @@ def export_hf_checkpoint(
             )
             if getattr(model, "hf_quantizer", None) is not None:
                 model.hf_quantizer = None
+            name_mapper = None
             try:
                 name_mapper = build_reverse_name_mapper(model)
                 if name_mapper is not None and hf_quant_config:
@@ -1580,7 +1612,7 @@ def export_hf_checkpoint(
                     f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor "
                     "names may not match the original HF hub checkpoint."
                 )
-            _write_hf_export_config(model, hf_quant_config, export_dir)
+            _write_hf_export_config(model, hf_quant_config, export_dir, name_mapper)
             return
 
         post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs)
@@ -1602,6 +1634,7 @@ def export_hf_checkpoint(
         # and fails). Best-effort and atomic: any failure (an op we cannot reverse yet,
         # transformers API drift, unexpected shapes) falls back to the in-memory names for BOTH
         # weights and config so they stay mutually consistent.
+        name_mapper = None
         try:
             name_mapper = build_reverse_name_mapper(model)
             export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict)
@@ -1636,7 +1669,7 @@ def export_hf_checkpoint(
         finally:
             _unpatch_revert_weight_conversion(_patches)
 
-        _write_hf_export_config(model, hf_quant_config, export_dir)
+        _write_hf_export_config(model, hf_quant_config, export_dir, name_mapper)
 
     except Exception as e:
         warnings.warn(
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
new file mode 100644
index 00000000000..14d03713040
--- /dev/null
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -0,0 +1,519 @@
+# 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.
+
+"""Layer-wise KV-cache AutoQuant using isolated forward KL sensitivity."""
+
+from __future__ import annotations
+
+import fnmatch
+import os
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from modelopt.torch.opt.conversion import ModeloptStateManager
+from modelopt.torch.opt.searcher import LPS
+from modelopt.torch.utils import print_rank_0, safe_load, safe_save
+
+from .config import QuantizeConfig
+from .conversion import set_quantizer_by_cfg
+from .nn import TensorQuantizer
+
+__all__ = ["auto_quantize_kv_cache"]
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+_KV_QUANTIZER_ATTRS = ("k_bmm_quantizer", "v_bmm_quantizer")
+_KV_AUTOQUANT_SCHEMA_VERSION = 1
+
+
+def _disabled_quantizer() -> TensorQuantizer:
+    quantizer = TensorQuantizer()
+    quantizer.disable()
+    return quantizer
+
+
+def _candidate_quantizers(module: nn.Module, config: QuantizeConfig) -> dict[str, TensorQuantizer]:
+    original = {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS}
+    try:
+        for attr in _KV_QUANTIZER_ATTRS:
+            setattr(module, attr, _disabled_quantizer())
+        set_quantizer_by_cfg(module, config.quant_cfg)
+        quantizers = {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS}
+        for attr, quantizer in quantizers.items():
+            if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled:
+                raise ValueError(
+                    f"KV-cache candidate must enable {attr}; got {type(quantizer).__name__}."
+                )
+        return quantizers
+    finally:
+        for attr, quantizer in original.items():
+            setattr(module, attr, quantizer)
+
+
+def _validate_kv_only_config(config: QuantizeConfig) -> None:
+    if config.effective_bits is None:
+        raise ValueError(
+            "Each KV-cache AutoQuant candidate must declare config-level effective_bits."
+        )
+    allowed_names = set(_KV_QUANTIZER_ATTRS)
+    matched_names: set[str] = set()
+    probe_names = {
+        *allowed_names,
+        "q_bmm_quantizer",
+        "p_bmm_quantizer",
+        "input_quantizer",
+        "weight_quantizer",
+        "output_quantizer",
+    }
+    for entry in config.quant_cfg:
+        pattern = entry.quantizer_name
+        matches = {name for name in probe_names if fnmatch.fnmatch(name, pattern)}
+        if matches - allowed_names:
+            raise ValueError(
+                "KV-cache AutoQuant candidates may configure only k_bmm_quantizer and "
+                f"v_bmm_quantizer; pattern {pattern!r} also matches {sorted(matches - allowed_names)}."
+            )
+        matched_names.update(matches)
+    if matched_names != allowed_names:
+        raise ValueError(
+            "KV-cache AutoQuant candidates must completely configure both "
+            "k_bmm_quantizer and v_bmm_quantizer."
+        )
+
+
+def _projection_width(module: nn.Module, side: str) -> int | None:
+    projection = getattr(module, f"{side}_proj", None)
+    out_features = getattr(projection, "out_features", None)
+    if isinstance(out_features, int) and out_features > 0:
+        return out_features
+
+    config = getattr(module, "config", None)
+    num_kv_heads = getattr(config, "num_key_value_heads", None)
+    head_dim = getattr(module, "head_dim", None) or getattr(config, "head_dim", None)
+    if not isinstance(head_dim, int):
+        hidden_size = getattr(config, "hidden_size", None)
+        num_heads = getattr(config, "num_attention_heads", None)
+        if isinstance(hidden_size, int) and isinstance(num_heads, int) and num_heads > 0:
+            head_dim = hidden_size // num_heads
+    if isinstance(num_kv_heads, int) and isinstance(head_dim, int):
+        return num_kv_heads * head_dim
+    return None
+
+
+def _kv_scalar_weight(module: nn.Module, name: str) -> int:
+    k_width = _projection_width(module, "k")
+    v_width = _projection_width(module, "v")
+    if k_width is None or v_width is None:
+        raise ValueError(
+            "Cannot determine exact KV width for eligible attention layer "
+            f"{name!r}. Expected k_proj/v_proj.out_features or config "
+            "num_key_value_heads plus head_dim."
+        )
+    return k_width + v_width
+
+
+def _eligible_layers(
+    model: nn.Module, disabled_layers: list[str] | str | None
+) -> list[tuple[str, nn.Module, int]]:
+    patterns = [disabled_layers] if isinstance(disabled_layers, str) else disabled_layers or []
+    layers = []
+    for name, module in model.named_modules():
+        if not all(hasattr(module, attr) for attr in _KV_QUANTIZER_ATTRS):
+            continue
+        if any(fnmatch.fnmatch(name, pattern) for pattern in patterns):
+            continue
+        layers.append((name, module, _kv_scalar_weight(module, name)))
+    if not layers:
+        raise ValueError("KV-cache AutoQuant found no eligible attention layers.")
+    return layers
+
+
+def _apply_layer_quantizers(module: nn.Module, quantizers: dict[str, TensorQuantizer]) -> None:
+    for attr, quantizer in quantizers.items():
+        setattr(module, attr, quantizer)
+
+
+def _get_logits(forward_step: Callable[[nn.Module, Any], torch.Tensor], model, data):
+    logits = forward_step(model, data)
+    if not isinstance(logits, torch.Tensor):
+        raise TypeError("KV-cache AutoQuant forward_step must return a logits tensor.")
+    if not torch.isfinite(logits).all():
+        raise ValueError("KV-cache AutoQuant encountered NaN or Inf logits.")
+    return logits
+
+
+def _solve_additive_recipe(
+    layer_names: list[str],
+    scalar_weights: list[int],
+    candidate_names: list[str],
+    candidate_bits: list[float],
+    scores: list[list[float]],
+    target_bits: float,
+    verbose: bool,
+) -> tuple[list[int], str]:
+    denominator = float(sum(scalar_weights))
+    candidate_costs = [
+        [weight * bits / 16.0 for bits in candidate_bits] for weight in scalar_weights
+    ]
+    max_cost = denominator * target_bits / 16.0
+    lps = LPS(
+        name="KVCacheAutoQuant",
+        constraints={"kv_cache_size_after_compression": max_cost},
+        constraints_to_candidate_costs={"kv_cache_size_after_compression": candidate_costs},
+        candidate_scores=scores,
+        objective_type="minimize",
+        verbose=verbose,
+    )
+    selections, status = lps()
+    if status != "Optimal":
+        minimum_bits = sum(weight * min(candidate_bits) for weight in scalar_weights) / denominator
+        raise ValueError(
+            f"KV-cache AutoQuant could not satisfy kv_effective_bits={target_bits}; "
+            f"minimum achievable value is {minimum_bits:.4f}. Solver status: {status}."
+        )
+    if len(selections) != len(layer_names):
+        raise RuntimeError(
+            "KV-cache AutoQuant solver returned an invalid selection count: "
+            f"{len(selections)} for {len(layer_names)} layers and candidates {candidate_names}."
+        )
+    return selections, status
+
+
+def _search_signature(
+    candidates: list[tuple[str, QuantizeConfig]],
+    layers: list[tuple[str, nn.Module, int]],
+    target_bits: float,
+    num_calib_steps: int,
+    num_score_steps: int,
+) -> dict[str, Any]:
+    return {
+        "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
+        "kv_effective_bits": target_bits,
+        "num_calib_steps": num_calib_steps,
+        "num_score_steps": num_score_steps,
+        "candidates": [
+            {
+                "name": name,
+                "config": config.model_dump(mode="json", exclude_none=True),
+            }
+            for name, config in candidates
+        ],
+        "layers": [{"name": name, "kv_scalar_weight": weight} for name, _, weight in layers],
+    }
+
+
+def _checkpoint_state_is_compatible(state: dict[str, Any], signature: dict[str, Any]) -> bool:
+    return state.get("search_signature") == signature
+
+
+def _quantizer_state_dict(
+    candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]],
+) -> dict[str, dict[str, dict[str, dict[str, torch.Tensor]]]]:
+    return {
+        layer_name: {
+            candidate_name: {
+                attr: quantizer.state_dict() for attr, quantizer in layer_quantizers.items()
+            }
+            for candidate_name, layer_quantizers in layer_candidates.items()
+        }
+        for layer_name, layer_candidates in candidate_quantizers.items()
+    }
+
+
+def _restore_quantizer_state_dict(
+    candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]],
+    state: dict[str, dict[str, dict[str, dict[str, torch.Tensor]]]],
+) -> None:
+    """Restore calibration buffers into config-created candidate quantizers."""
+    for layer_name, layer_candidates in candidate_quantizers.items():
+        for candidate_name, layer_quantizers in layer_candidates.items():
+            for attr, quantizer in layer_quantizers.items():
+                quantizer_state = state[layer_name][candidate_name][attr]
+                for key, value in quantizer_state.items():
+                    if "." not in key and key not in quantizer._buffers:
+                        quantizer.register_buffer(key, torch.empty_like(value))
+                quantizer.load_state_dict(quantizer_state)
+
+
+def _report_state(state: dict[str, Any]) -> dict[str, Any]:
+    """Return the JSON-safe search report, excluding calibration tensors."""
+    return {key: value for key, value in state.items() if key != "quantizer_state"}
+
+
+@torch.inference_mode()
+def auto_quantize_kv_cache(
+    model: nn.Module,
+    constraints: dict[str, Any],
+    quantization_formats: list[tuple[dict[str, Any], str]],
+    data_loader: Iterable,
+    forward_step: Callable[[nn.Module, Any], torch.Tensor],
+    *,
+    num_calib_steps: int,
+    num_score_steps: int,
+    disabled_layers: list[str] | str | None = None,
+    verbose: bool = False,
+    checkpoint: str | None = None,
+) -> tuple[nn.Module, dict[str, Any]]:
+    """Select one supplied K/V format per attention layer using isolated forward KL.
+
+    Candidate formats are format-agnostic ``QuantizeConfig`` dictionaries. Each must
+    configure K and V together and declare ``effective_bits`` matching their runtime
+    packed storage. Candidate-specific calibration runs before scoring; cast-style
+    constant-amax formats skip calibration forwards through their normal algorithm
+    configuration.
+    """
+    if set(constraints) != {"kv_effective_bits"}:
+        raise ValueError(
+            "KV-cache AutoQuant constraints must contain only kv_effective_bits; "
+            f"got {sorted(constraints)}."
+        )
+    target_bits = float(constraints["kv_effective_bits"])
+    if not (0 < target_bits <= 16):
+        raise ValueError(f"kv_effective_bits must be in (0, 16], got {target_bits}.")
+    if num_calib_steps <= 0:
+        raise ValueError("num_calib_steps must be positive.")
+    if num_score_steps <= 0:
+        raise ValueError("num_score_steps must be positive.")
+
+    candidates = []
+    seen_names = set()
+    for raw_config, name in quantization_formats:
+        if name in seen_names:
+            raise ValueError(f"Duplicate KV-cache AutoQuant candidate name: {name!r}.")
+        config = QuantizeConfig(**raw_config)
+        _validate_kv_only_config(config)
+        candidates.append((name, config))
+        seen_names.add(name)
+    if not candidates:
+        raise ValueError("KV-cache AutoQuant requires at least one candidate format.")
+
+    layers = _eligible_layers(model, disabled_layers)
+    signature = _search_signature(
+        candidates,
+        layers,
+        target_bits,
+        num_calib_steps,
+        num_score_steps,
+    )
+    candidate_names = [name for name, _ in candidates]
+    candidate_bits = []
+    for _, config in candidates:
+        assert config.effective_bits is not None
+        candidate_bits.append(config.effective_bits)
+
+    original_quantizers = {
+        name: {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS}
+        for name, module, _ in layers
+    }
+    disabled_quantizers = {
+        name: {attr: _disabled_quantizer() for attr in _KV_QUANTIZER_ATTRS} for name, _, _ in layers
+    }
+    candidate_quantizers = {
+        name: {
+            candidate_name: _candidate_quantizers(module, config)
+            for candidate_name, config in candidates
+        }
+        for name, module, _ in layers
+    }
+
+    try:
+        for name, module, _ in layers:
+            _apply_layer_quantizers(module, disabled_quantizers[name])
+
+        state: dict[str, Any] | None = None
+        if checkpoint is not None and os.path.exists(checkpoint):
+            restored = safe_load(checkpoint)
+            if _checkpoint_state_is_compatible(restored, signature):
+                state = restored
+                if verbose:
+                    print_rank_0(f"KV-cache AutoQuant restored search state from {checkpoint}.")
+            else:
+                raise ValueError(
+                    "KV-cache AutoQuant checkpoint does not match the current candidates "
+                    "or eligible layers. Use a different checkpoint path."
+                )
+
+        if state is not None and state.get("calibration_complete"):
+            quantizer_state = state.get("quantizer_state")
+            if quantizer_state is None:
+                raise ValueError(
+                    "KV-cache AutoQuant checkpoint is missing calibrated quantizer state. "
+                    "Use a different checkpoint path."
+                )
+            _restore_quantizer_state_dict(candidate_quantizers, quantizer_state)
+        else:
+            from .model_quant import calibrate
+
+            for candidate_name, config in candidates:
+                for layer_name, module, _ in layers:
+                    _apply_layer_quantizers(
+                        module, candidate_quantizers[layer_name][candidate_name]
+                    )
+
+                if config.algorithm is not None:
+
+                    def calibration_loop(calibration_model):
+                        for step, data in enumerate(data_loader):
+                            if step >= num_calib_steps:
+                                break
+                            _get_logits(forward_step, calibration_model, data)
+
+                    modelopt_state = ModeloptStateManager(model).state_dict()
+                    original_mode_count = len(modelopt_state)
+                    try:
+                        calibrate(
+                            model,
+                            algorithm=config.algorithm,
+                            forward_loop=calibration_loop,
+                        )
+                    finally:
+                        del modelopt_state[original_mode_count:]
+
+                for layer_name, module, _ in layers:
+                    _apply_layer_quantizers(module, disabled_quantizers[layer_name])
+
+            state = {
+                "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
+                "search_signature": signature,
+                "calibration_complete": True,
+                "num_calib_steps": num_calib_steps,
+                "quantizer_state": _quantizer_state_dict(candidate_quantizers),
+            }
+            if checkpoint is not None:
+                checkpoint_dir = os.path.dirname(checkpoint)
+                if checkpoint_dir:
+                    os.makedirs(checkpoint_dir, exist_ok=True)
+                safe_save(state, checkpoint)
+
+        assert state is not None
+        if not state.get("layers"):
+            score_sums = {
+                layer_name: dict.fromkeys(candidate_names, 0.0) for layer_name, _, _ in layers
+            }
+            scored_tokens = 0
+            scored_steps = 0
+            iterator = tqdm(
+                data_loader,
+                total=num_score_steps,
+                desc="Estimating KV-cache KL sensitivity",
+                disable=not verbose,
+            )
+            for data in iterator:
+                if scored_steps >= num_score_steps:
+                    break
+                logits_ref = _get_logits(forward_step, model, data)
+                log_prob_ref = torch.log_softmax(logits_ref.float(), dim=-1)
+                scored_tokens += logits_ref.numel() // logits_ref.shape[-1]
+
+                for layer_name, module, _ in layers:
+                    for candidate_name, _ in candidates:
+                        _apply_layer_quantizers(
+                            module, candidate_quantizers[layer_name][candidate_name]
+                        )
+                        logits_quant = _get_logits(forward_step, model, data)
+                        score = F.kl_div(
+                            torch.log_softmax(logits_quant.float(), dim=-1),
+                            log_prob_ref,
+                            reduction="sum",
+                            log_target=True,
+                        )
+                        score_sums[layer_name][candidate_name] += float(score.item())
+                        _apply_layer_quantizers(module, disabled_quantizers[layer_name])
+                scored_steps += 1
+
+            if scored_steps == 0 or scored_tokens == 0:
+                raise ValueError("KV-cache AutoQuant data_loader produced no scoring batches.")
+            scores = [
+                [score_sums[layer_name][candidate] / scored_tokens for candidate in candidate_names]
+                for layer_name, _, _ in layers
+            ]
+            selections, status = _solve_additive_recipe(
+                [name for name, _, _ in layers],
+                [weight for _, _, weight in layers],
+                candidate_names,
+                candidate_bits,
+                scores,
+                target_bits,
+                verbose,
+            )
+            denominator = float(sum(weight for _, _, weight in layers))
+            achieved_bits = (
+                sum(
+                    weight * candidate_bits[selected]
+                    for selected, (_, _, weight) in zip(selections, layers)
+                )
+                / denominator
+            )
+            selected_score = sum(
+                layer_scores[selected] for selected, layer_scores in zip(selections, scores)
+            )
+            state.update(
+                {
+                    "method": "kl_div",
+                    "score_reduction": "mean_per_scored_token",
+                    "constraints": {"kv_effective_bits": target_bits},
+                    "num_score_steps": scored_steps,
+                    "num_scored_tokens": scored_tokens,
+                    "candidates": [
+                        {
+                            "name": name,
+                            "effective_bits": effective_bits,
+                            "config": config.model_dump(mode="json", exclude_none=True),
+                        }
+                        for (name, config), effective_bits in zip(candidates, candidate_bits)
+                    ],
+                    "layers": {
+                        layer_name: {
+                            "kv_scalar_weight": weight,
+                            "scores": dict(zip(candidate_names, layer_scores)),
+                            "selected": candidate_names[selected],
+                        }
+                        for selected, layer_scores, (layer_name, _, weight) in zip(
+                            selections, scores, layers
+                        )
+                    },
+                    "best": {
+                        "effective_bits": achieved_bits,
+                        "score": selected_score,
+                        "is_satisfied": achieved_bits <= target_bits + 1e-12,
+                        "solver_status": status,
+                    },
+                }
+            )
+            if checkpoint is not None:
+                checkpoint_dir = os.path.dirname(checkpoint)
+                if checkpoint_dir:
+                    os.makedirs(checkpoint_dir, exist_ok=True)
+                safe_save(state, checkpoint)
+                if verbose:
+                    print_rank_0(f"Saved KV-cache AutoQuant report to {checkpoint}.")
+
+        for layer_name, module, _ in layers:
+            selected_name = state["layers"][layer_name]["selected"]
+            _apply_layer_quantizers(module, candidate_quantizers[layer_name][selected_name])
+            if verbose:
+                print_rank_0(f"KV-cache AutoQuant selected {selected_name} for {layer_name}.")
+        report = _report_state(state)
+        model._modelopt_kv_cache_auto_quantize_state = report
+        return model, report
+    except Exception:
+        for layer_name, module, _ in layers:
+            _apply_layer_quantizers(module, original_quantizers[layer_name])
+        raise
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 966a3643fe3..4ea9f49dead 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -39,12 +39,14 @@
 from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe
 from .algorithms import get_auto_quantize_config as _get_auto_quantize_config
 from .config import QuantizeAlgoCfgType
+from .kv_cache_auto_quant import auto_quantize_kv_cache as _auto_quantize_kv_cache
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, TensorQuantizer
 from .utils import is_quantized
 
 __all__ = [
     "auto_quantize",
+    "auto_quantize_kv_cache",
     "calibrate",
     "compute_quantization_mse",
     "disable_quantizer",
@@ -656,6 +658,73 @@ def _process_quantization_formats(formats, custom_name_prefix):
     return model, searcher.state_dict()
 
 
+def auto_quantize_kv_cache(
+    model: nn.Module,
+    constraints: dict[str, Any],
+    quantization_formats: list[dict[str, Any] | tuple[dict[str, Any], str]],
+    data_loader: Iterable,
+    forward_step: Callable[[nn.Module, Any], torch.Tensor],
+    *,
+    num_calib_steps: int = 512,
+    num_score_steps: int = 128,
+    disabled_layers: list[str] | str | None = None,
+    verbose: bool = False,
+    checkpoint: str | None = None,
+) -> tuple[nn.Module, dict[str, Any]]:
+    """Search layer-wise KV-cache formats using isolated forward KL sensitivity.
+
+    Unlike weight AutoQuant, BF16/no-quant is only the scoring reference and is
+    not an implicit solver choice. Each supplied format must configure K and V
+    together and declare exact config-level ``effective_bits``. Formats use their
+    own calibration algorithm; cast-mode constant-amax candidates skip calibration
+    forwards and provide the fastest turnaround.
+
+    Args:
+        model: Model whose attention K/V quantizers will be searched.
+        constraints: A ``{"kv_effective_bits": target}`` storage constraint.
+        quantization_formats: Candidate ``QuantizeConfig`` dictionaries, optionally
+            paired with display names.
+        data_loader: Re-iterable calibration and scoring batches.
+        forward_step: Callable returning the full logits tensor for one batch.
+        num_calib_steps: Maximum calibration batches per candidate.
+        num_score_steps: Maximum batches used for isolated forward-KL scoring.
+        disabled_layers: Optional layer-name patterns excluded from the search and
+            preserved in their existing KV-cache format.
+        verbose: Whether to print progress and selected formats.
+        checkpoint: Optional path for resumable calibration and sensitivity state.
+
+    Returns:
+        The converted model with the selected per-layer K/V quantizers and a
+        JSON-safe sensitivity report.
+    """
+    processed_formats = []
+    for idx, candidate in enumerate(quantization_formats):
+        if isinstance(candidate, tuple):
+            raw_config, name = candidate
+        else:
+            raw_config, name = candidate, f"KV_CACHE_FORMAT_{idx}"
+        if not isinstance(raw_config, dict):
+            raise TypeError("KV-cache AutoQuant formats must be config dictionaries.")
+        if not isinstance(name, str) or not name:
+            raise ValueError("KV-cache AutoQuant candidate names must be non-empty strings.")
+        processed_formats.append((raw_config, name))
+
+    if not is_quantized(model):
+        model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
+    return _auto_quantize_kv_cache(
+        model,
+        constraints,
+        processed_formats,
+        data_loader,
+        forward_step,
+        num_calib_steps=num_calib_steps,
+        num_score_steps=num_score_steps,
+        disabled_layers=disabled_layers,
+        verbose=verbose,
+        checkpoint=checkpoint,
+    )
+
+
 def get_auto_quantize_config(search_state, constraints=None, verbose=False):
     """Build a flat quant config from auto_quantize search_state.
 
diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
new file mode 100644
index 00000000000..a9518c86fc0
--- /dev/null
+++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Layer-wise KV-cache search over cast-mode FP8 and NVFP4 at 5.4 bits/scalar.
+# Cast mode fixes amax and skips PTQ calibration; isolated full-vocabulary forward
+# KL is measured with every other eligible attention layer kept in BF16.
+
+# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
+imports:
+  base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
+  base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers
+  kv_fp8_cast: configs/ptq/units/kv_fp8_cast
+  kv_nvfp4_cast: configs/ptq/units/kv_nvfp4_cast
+
+metadata:
+  recipe_type: auto_quantize
+  description: Layer-wise FP8/NVFP4 KV-cache cast search at 5.4 bits using forward KL.
+
+auto_quantize:
+  constraints:
+    kv_effective_bits: 5.4
+
+  candidate_formats:
+    - quant_cfg:
+        - $import: kv_fp8_cast
+      algorithm:
+        method: max
+        skip_forward_without_activation_calib: true
+      effective_bits: 8.0
+    - quant_cfg:
+        - $import: kv_nvfp4_cast
+      algorithm:
+        method: max
+        skip_forward_without_activation_calib: true
+      effective_bits: 4.5
+
+  auto_quantize_method: kl_div
+  score_size: 128
+
+  disabled_layers:
+    - $import: base_disabled_layers
+
+  cost_excluded_layers:
+    - $import: base_cost_excluded_layers
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 6a2c36e4a7c..2911f498202 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -83,6 +83,23 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch):
     assert inputs["quantization_formats"][1] == QUANT_CFG_CHOICES["fp8"]
 
 
+def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
+    hf_ptq, args = _parse_hf_ptq_args(
+        monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "fp8_cast"
+    )
+    aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize
+    inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args)
+
+    assert inputs["search_domain"] == "kv_cache"
+    assert inputs["constraints"] == {"kv_effective_bits": 5.4}
+    assert inputs["method"] == "kl_div"
+    assert [config["effective_bits"] for config, _ in inputs["quantization_formats"]] == [
+        8.0,
+        4.5,
+    ]
+    assert "kv_cache_quant_cfg" not in inputs
+
+
 def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch):
     """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns
     key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict."""
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index 3dfb9906a54..e7220954dd1 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1907,6 +1907,7 @@ def test_load_recipe_autoquantize_fixed_baseline_requires_explicit_search(tmp_pa
     [
         "general/auto_quantize/nvfp4_fp8_at_5p4bits",
         "general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits",
+        "general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits",
         "general/auto_quantize/nvfp4_mse_fp8_at_6p0bits",
         "general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits",
         "general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe",
@@ -1925,6 +1926,21 @@ def test_load_recipe_autoquantize_builtin_general(recipe_path):
     assert recipe.auto_quantize.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"]
 
 
+def test_load_recipe_kv_autoquantize_contract():
+    recipe = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits")
+    aq = recipe.auto_quantize
+
+    assert aq.constraints.effective_bits is None
+    assert aq.constraints.kv_effective_bits == 5.4
+    assert aq.auto_quantize_method == "kl_div"
+    assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5]
+    for fmt in aq.candidate_formats:
+        (entry,) = fmt.quant_cfg
+        assert entry.quantizer_name == "*[kv]_bmm_quantizer"
+        assert entry.cfg.use_constant_amax
+        assert fmt.algorithm["skip_forward_without_activation_calib"]
+
+
 def _all_shipped_ptq_recipe_paths():
     """Every shipped PTQ recipe, discovered from disk rather than a hardcoded list."""
     root = files("modelopt_recipes")
diff --git a/tests/unit/torch/export/test_convert_hf_config.py b/tests/unit/torch/export/test_convert_hf_config.py
new file mode 100644
index 00000000000..ab90168b750
--- /dev/null
+++ b/tests/unit/torch/export/test_convert_hf_config.py
@@ -0,0 +1,42 @@
+# 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.
+
+from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
+
+
+def test_convert_mixed_kv_cache_config_preserves_layer_map():
+    layer_map = {
+        "model.layers.0.self_attn": {"quant_algo": "FP8"},
+        "model.layers.1.self_attn": {"quant_algo": "NVFP4"},
+    }
+    converted = convert_hf_quant_config_format(
+        {
+            "producer": {"name": "modelopt", "version": "test"},
+            "quantization": {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {},
+                "kv_cache_quant_algo": "MIXED_PRECISION",
+                "kv_cache_quantized_layers": layer_map,
+                "kv_cache_schema_version": 1,
+            },
+        }
+    )
+
+    assert converted["quant_method"] == "modelopt"
+    assert converted["quant_algo"] == "MIXED_PRECISION"
+    assert converted["config_groups"] == {}
+    assert converted["kv_cache_quant_algo"] == "MIXED_PRECISION"
+    assert converted["kv_cache_quantized_layers"] == layer_map
+    assert converted["kv_cache_schema_version"] == 1
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index e7ca68d0b69..cbcf97a24a6 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -27,12 +27,14 @@
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.layer_utils import get_quantization_format
 from modelopt.torch.export.model_config import (
+    KV_CACHE_FP8,
+    KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A8_AWQ,
 )
-from modelopt.torch.export.quant_utils import get_quant_config
-from modelopt.torch.quantization.nn import NVFP4StaticQuantizer
+from modelopt.torch.export.quant_utils import get_quant_config, postprocess_state_dict
+from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer
 
 
 @pytest.mark.parametrize(
@@ -64,6 +66,71 @@ def test_nvfp4_static_quantizer_export():
     assert quant_config["quantization"]["group_size"] == 16
 
 
+def test_mixed_kv_cache_quantization_exports_per_layer_map():
+    class FakeAttention(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.k_bmm_quantizer = TensorQuantizer()
+            self.v_bmm_quantizer = TensorQuantizer()
+
+    model = torch.nn.Module()
+    model.attn0 = FakeAttention()
+    model.attn1 = FakeAttention()
+    mtq.set_quantizer_by_cfg(
+        model.attn0,
+        [
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "use_constant_amax": True},
+            }
+        ],
+    )
+    mtq.set_quantizer_by_cfg(
+        model.attn1,
+        [
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {
+                    "num_bits": (2, 1),
+                    "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                    "use_constant_amax": True,
+                },
+            }
+        ],
+    )
+
+    quantization = get_quant_config(model)["quantization"]
+    assert quantization["quant_algo"] == "MIXED_PRECISION"
+    assert quantization["kv_cache_quant_algo"] == "MIXED_PRECISION"
+    assert quantization["quantized_layers"] == {}
+    assert quantization["kv_cache_quantized_layers"] == {
+        "attn0": {"quant_algo": "FP8"},
+        "attn1": {"quant_algo": "NVFP4"},
+    }
+
+
+def test_mixed_kv_cache_postprocess_uses_each_layers_format():
+    state_dict = {
+        "attn0.k_bmm_quantizer._amax": torch.tensor([448.0]),
+        "attn0.v_bmm_quantizer._amax": torch.tensor([224.0]),
+        "attn1.k_bmm_quantizer._amax": torch.tensor([112.0]),
+        "attn1.v_bmm_quantizer._amax": torch.tensor([56.0]),
+    }
+    layer_formats = {
+        "attn0": {"quant_algo": KV_CACHE_FP8},
+        "attn1": {"quant_algo": KV_CACHE_NVFP4},
+    }
+
+    processed = postprocess_state_dict(state_dict, 448.0, layer_formats)
+
+    assert processed == {
+        "attn0.k_proj.k_scale": torch.tensor([1.0]),
+        "attn0.v_proj.v_scale": torch.tensor([0.5]),
+        "attn1.k_proj.k_scale": torch.tensor([0.25]),
+        "attn1.v_proj.v_scale": torch.tensor([0.125]),
+    }
+
+
 class _FakeTopKRouter(torch.nn.Module):
     """Mimics a transformers>=5.0 MoE router: owns a ``weight`` but is NOT an ``nn.Linear``.
 
diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index e3aa22bf3d8..06a1b88f222 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -516,6 +516,7 @@ def test_revert_quant_config_names_mapper():
             "lm_head",
         ],
         "quantized_layers": {"model.layers.0.mlp.experts.0.w1": {"quant_algo": "NVFP4"}},
+        "kv_cache_quantized_layers": {"model.layers.0.mlp.experts.0": {"quant_algo": "FP8"}},
     }
     revert_quant_config_names(quant, mapper)
     assert quant["exclude_modules"] == [
@@ -524,6 +525,7 @@ def test_revert_quant_config_names_mapper():
         "lm_head",
     ]
     assert "model.layers.0.block_sparse_moe.experts.0.w1" in quant["quantized_layers"]
+    assert "model.layers.0.block_sparse_moe.experts.0" in quant["kv_cache_quantized_layers"]
     # mapper(None) is a no-op
     q2 = {"exclude_modules": ["x*"]}
     revert_quant_config_names(q2, None)
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
new file mode 100644
index 00000000000..0dbf792b24d
--- /dev/null
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -0,0 +1,334 @@
+# 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.
+
+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.config import QuantizeConfig
+from modelopt.torch.quantization.kv_cache_auto_quant import (
+    _candidate_quantizers,
+    _kv_scalar_weight,
+    _solve_additive_recipe,
+    _validate_kv_only_config,
+    auto_quantize_kv_cache,
+)
+
+
+def _kv_config(bits, effective_bits):
+    return QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": bits, "use_constant_amax": True},
+            }
+        ],
+        effective_bits=effective_bits,
+    )
+
+
+def test_kv_candidate_requires_exact_bits_and_both_sides():
+    _validate_kv_only_config(_kv_config((4, 3), 8.0))
+
+    with pytest.raises(ValueError, match="config-level effective_bits"):
+        _validate_kv_only_config(
+            QuantizeConfig(
+                quant_cfg=[
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": (4, 3), "use_constant_amax": True},
+                    }
+                ]
+            )
+        )
+    with pytest.raises(ValueError, match="completely configure both"):
+        _validate_kv_only_config(
+            QuantizeConfig(
+                quant_cfg=[
+                    {
+                        "quantizer_name": "*k_bmm_quantizer",
+                        "cfg": {"num_bits": (4, 3), "use_constant_amax": True},
+                    }
+                ],
+                effective_bits=8.0,
+            )
+        )
+
+
+def test_kv_additive_solver_spends_fp8_on_more_sensitive_layer():
+    selections, status = _solve_additive_recipe(
+        layer_names=["layer0", "layer1"],
+        scalar_weights=[256, 256],
+        candidate_names=["fp8", "nvfp4"],
+        candidate_bits=[8.0, 4.5],
+        scores=[[0.0, 10.0], [0.0, 1.0]],
+        target_bits=6.25,
+        verbose=False,
+    )
+
+    assert status == "Optimal"
+    assert selections == [0, 1]
+
+
+def test_kv_scalar_weight_counts_k_and_v_widths():
+    module = nn.Module()
+    module.k_proj = nn.Linear(32, 24, bias=False)
+    module.v_proj = nn.Linear(32, 16, bias=False)
+
+    assert _kv_scalar_weight(module, "attention") == 40
+
+
+def test_kv_candidate_format_can_use_dynamic_amax():
+    module = nn.Module()
+    module.k_bmm_quantizer = nn.Identity()
+    module.v_bmm_quantizer = nn.Identity()
+    config = QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3)},
+            }
+        ],
+        algorithm=None,
+        effective_bits=8.0,
+    )
+
+    quantizers = _candidate_quantizers(module, config)
+
+    assert all(quantizer.is_enabled for quantizer in quantizers.values())
+    assert all(not quantizer._use_constant_amax for quantizer in quantizers.values())
+
+
+class _ToyKVAttention(nn.Module):
+    def __init__(self, width, gain):
+        super().__init__()
+        self.k_proj = nn.Linear(width, width, bias=False)
+        self.v_proj = nn.Linear(width, width, bias=False)
+        self.k_bmm_quantizer = nn.Identity()
+        self.v_bmm_quantizer = nn.Identity()
+        self.gain = gain
+
+    def forward(self, x):
+        return x + self.gain * (self.k_bmm_quantizer(x) + self.v_bmm_quantizer(x))
+
+
+class _ToyKVModel(nn.Module):
+    def __init__(self, width=8):
+        super().__init__()
+        self.attn0 = _ToyKVAttention(width, gain=0.25)
+        self.attn1 = _ToyKVAttention(width, gain=2.0)
+        self.lm_head = nn.Linear(width, width, bias=False)
+
+    def forward(self, x):
+        return self.lm_head(self.attn1(self.attn0(x)))
+
+
+def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
+    torch.manual_seed(123)
+    model = _ToyKVModel()
+    data = [torch.randn(2, 3, 8), torch.randn(2, 3, 8)]
+    candidates = [
+        (
+            {
+                "quant_cfg": [
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": 8, "constant_amax": 1.0},
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 8.0,
+            },
+            "int8",
+        ),
+        (
+            {
+                "quant_cfg": [
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 4.0,
+            },
+            "int4",
+        ),
+    ]
+
+    model, state = auto_quantize_kv_cache(
+        model,
+        {"kv_effective_bits": 6.0},
+        candidates,
+        data,
+        lambda model, batch: model(batch),
+        num_calib_steps=2,
+        num_score_steps=2,
+        checkpoint=str(tmp_path / "kv_search.pth"),
+    )
+
+    assert state["best"]["effective_bits"] == pytest.approx(6.0)
+    assert state["best"]["is_satisfied"]
+    assert {layer["selected"] for layer in state["layers"].values()} == {
+        "int8",
+        "int4",
+    }
+    for layer_name, layer_state in state["layers"].items():
+        layer = model.get_submodule(layer_name)
+        assert layer.k_bmm_quantizer.num_bits == layer.v_bmm_quantizer.num_bits
+        expected_bits = 8 if layer_state["selected"] == "int8" else 4
+        assert layer.k_bmm_quantizer.num_bits == expected_bits
+
+    restored_model = _ToyKVModel()
+    restored_model, restored_state = auto_quantize_kv_cache(
+        restored_model,
+        {"kv_effective_bits": 6.0},
+        candidates,
+        data,
+        lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."),
+        num_calib_steps=2,
+        num_score_steps=2,
+        checkpoint=str(tmp_path / "kv_search.pth"),
+    )
+
+    assert restored_state == state
+    for layer_name, layer_state in restored_state["layers"].items():
+        layer = restored_model.get_submodule(layer_name)
+        expected_bits = 8 if layer_state["selected"] == "int8" else 4
+        assert layer.k_bmm_quantizer.num_bits == expected_bits
+
+
+def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
+    torch.manual_seed(123)
+    model = get_tiny_llama(num_hidden_layers=2)
+    data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))} for _ in range(2)]
+    candidates = [
+        (
+            {
+                "quant_cfg": [
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": 8},
+                    }
+                ],
+                "algorithm": "max",
+                "effective_bits": 8.0,
+            },
+            "int8",
+        ),
+        (
+            {
+                "quant_cfg": [
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 4.0,
+            },
+            "int4",
+        ),
+    ]
+
+    model, state = mtq.auto_quantize_kv_cache(
+        model,
+        {"kv_effective_bits": 6.0},
+        candidates,
+        data,
+        lambda search_model, batch: search_model(**batch).logits,
+        num_calib_steps=2,
+        num_score_steps=2,
+        checkpoint=str(tmp_path / "hf_kv_search.pth"),
+    )
+
+    assert len(state["layers"]) == model.config.num_hidden_layers
+    assert state["best"]["effective_bits"] == pytest.approx(6.0)
+    assert all(
+        layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled
+        for layer in model.model.layers
+    )
+
+    restored_model = get_tiny_llama(num_hidden_layers=2)
+    restored_model, restored_state = mtq.auto_quantize_kv_cache(
+        restored_model,
+        {"kv_effective_bits": 6.0},
+        candidates,
+        data,
+        lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."),
+        num_calib_steps=2,
+        num_score_steps=2,
+        checkpoint=str(tmp_path / "hf_kv_search.pth"),
+    )
+
+    assert restored_state == state
+    assert any(
+        hasattr(layer.self_attn.k_bmm_quantizer, "_amax")
+        for layer in restored_model.model.layers
+        if layer.self_attn.k_bmm_quantizer.num_bits == 8
+    )
+
+
+def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
+    torch.manual_seed(123)
+    model = get_tiny_llama(num_hidden_layers=2)
+    data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
+    fixed_kv_config = {
+        "quant_cfg": [
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": 8, "constant_amax": 1.0},
+            }
+        ],
+        "algorithm": None,
+    }
+    model = mtq.quantize(model, fixed_kv_config)
+    model.model.layers[0].self_attn.q_proj.weight_quantizer.enable()
+
+    model, state = mtq.auto_quantize_kv_cache(
+        model,
+        {"kv_effective_bits": 4.0},
+        [
+            (
+                {
+                    "quant_cfg": [
+                        {
+                            "quantizer_name": "*[kv]_bmm_quantizer",
+                            "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                        }
+                    ],
+                    "algorithm": None,
+                    "effective_bits": 4.0,
+                },
+                "int4",
+            )
+        ],
+        data,
+        lambda search_model, batch: search_model(**batch).logits,
+        num_calib_steps=1,
+        num_score_steps=1,
+        disabled_layers="model.layers.1.self_attn",
+    )
+
+    assert set(state["layers"]) == {"model.layers.0.self_attn"}
+    assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == 4
+    assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled
+    fixed_attention = model.model.layers[1].self_attn
+    assert fixed_attention.k_bmm_quantizer.is_enabled
+    assert fixed_attention.v_bmm_quantizer.is_enabled
+    assert fixed_attention.k_bmm_quantizer.num_bits == 8
+    assert fixed_attention.v_bmm_quantizer.num_bits == 8

From abf5959f969af92aa87643578abc32d07e4cf6b6 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Mon, 3 Aug 2026 00:23:47 -0700
Subject: [PATCH 02/15] Add three-format mixed KV export

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |  3 +-
 modelopt/torch/export/model_config.py         |  1 +
 modelopt/torch/export/quant_utils.py          | 64 ++++++++++++++-----
 .../kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml  | 20 +++++-
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  1 +
 tests/unit/recipe/test_loader.py              | 16 +++--
 .../torch/export/test_get_quantization.py     | 29 +++++++++
 7 files changed, 112 insertions(+), 22 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 5ce321cd2b9..be9dfdfa117 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -428,7 +428,8 @@ Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall
 KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
 `candidate_formats` are complete K/V cache configs with exact config-level `effective_bits`;
 BF16 is used only as the isolated-KL reference, not as a solver choice. The shipped canary recipe
-searches cast-mode FP8 (8.0 bits/scalar) and packed NVFP4 (4.5 bits/scalar) at 5.4 bits/scalar:
+searches cast-mode FP8 K/V (8.0 bits/scalar), FP8-K/NVFP4-V (6.25 bits/scalar),
+and packed NVFP4 K/V (4.5 bits/scalar) at 5.4 bits/scalar:
 
 ```bash
 python hf_ptq.py \
diff --git a/modelopt/torch/export/model_config.py b/modelopt/torch/export/model_config.py
index 5f92cc2e5dc..cf778bf937d 100755
--- a/modelopt/torch/export/model_config.py
+++ b/modelopt/torch/export/model_config.py
@@ -45,6 +45,7 @@
 QUANTIZATION_FP8_PC_PT = "fp8_pc_pt"
 
 KV_CACHE_FP8 = "FP8"
+KV_CACHE_FP8_K_NVFP4_V = "FP8_K_NVFP4_V"
 KV_CACHE_INT8 = "INT8"
 KV_CACHE_NVFP4 = "NVFP4"
 KV_CACHE_NVFP4_AFFINE = "NVFP4_AFFINE"
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 75df85c5cfd..25e930b40fc 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -50,6 +50,7 @@
 from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer
 from .model_config import (
     KV_CACHE_FP8,
+    KV_CACHE_FP8_K_NVFP4_V,
     KV_CACHE_INT8,
     KV_CACHE_NVFP4,
     KV_CACHE_NVFP4_AFFINE,
@@ -389,27 +390,36 @@ def get_kv_cache_scaling_factor(self_attention_module: nn.Module) -> list[torch.
         for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer")
     ]
 
-    # For FP8, we recommend default kv cache scaling factor to be 1.
-    if get_kv_cache_dtype(self_attention_module) == KV_CACHE_FP8:
-        for i, factor in enumerate(scaling_factors):
-            if factor is None:
-                continue
-            if factor.item() > 0.5:
-                warn(
-                    f"Warning: Large KV activation detected: {factor.item()}, "
-                    "Quantized KV cache may lead to higher accuracy drop."
-                )
-            scaling_factors[i] = torch.max(
-                factor, torch.tensor([1.0], dtype=torch.float, device=factor.device)
+    # For FP8, we recommend default KV-cache scaling factor to be 1. The
+    # asymmetric format applies this only to K; V remains NVFP4.
+    kv_cache_dtype = get_kv_cache_dtype(self_attention_module)
+    if kv_cache_dtype == KV_CACHE_FP8:
+        fp8_indices = range(len(scaling_factors))
+    elif kv_cache_dtype == KV_CACHE_FP8_K_NVFP4_V:
+        fp8_indices = (0,)
+    else:
+        fp8_indices = ()
+    for i in fp8_indices:
+        factor = scaling_factors[i]
+        if factor is None:
+            continue
+        if factor.item() > 0.5:
+            warn(
+                f"Warning: Large KV activation detected: {factor.item()}, "
+                "Quantized KV cache may lead to higher accuracy drop."
             )
+        scaling_factors[i] = torch.max(
+            factor, torch.tensor([1.0], dtype=torch.float, device=factor.device)
+        )
     return scaling_factors
 
 
 def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None:
     """Returns the kv_cache dtype.
 
-    If num_bits of output_quantizer is (4, 3) then returns FP8; if it is 8, returns int8,
-    otherwise returns None.
+    K/V quantizers are inspected as a pair so FP8 K with NVFP4 V remains
+    distinguishable from uniform FP8 or NVFP4. The output quantizer is retained
+    as a fallback for the unified Megatron export path.
 
     Args:
         modules: The module or list of modules to inspect.
@@ -424,6 +434,25 @@ def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None:
         modules = [modules]
 
     for module in modules:
+        k_quantizer = getattr(module, "k_bmm_quantizer", None)
+        v_quantizer = getattr(module, "v_bmm_quantizer", None)
+        if (
+            k_quantizer is not None
+            and v_quantizer is not None
+            and k_quantizer.is_enabled
+            and v_quantizer.is_enabled
+        ):
+            k_dtype = _compute_kv_cache_dtype(
+                [k_quantizer.num_bits], hasattr(k_quantizer, "_bias_value")
+            )
+            v_dtype = _compute_kv_cache_dtype(
+                [v_quantizer.num_bits], hasattr(v_quantizer, "_bias_value")
+            )
+            if k_dtype == KV_CACHE_FP8 and v_dtype == KV_CACHE_NVFP4:
+                return KV_CACHE_FP8_K_NVFP4_V
+            if k_dtype == v_dtype:
+                return k_dtype
+
         # Case where the module has both k_bmm_quantizer and v_bmm_quantizer
         # Still check for output quantizer for the unified_megatron_export path
         for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer", "output_quantizer"):
@@ -1117,6 +1146,7 @@ def _kv_quantization_for_key(key: str) -> str | None:
                     layer_quantization = _kv_quantization_for_key(key)
                     assert layer_quantization in [
                         KV_CACHE_FP8,
+                        KV_CACHE_FP8_K_NVFP4_V,
                         KV_CACHE_NVFP4,
                         KV_CACHE_NVFP4_AFFINE,
                     ], "Invalid KV cache quantization format."
@@ -1125,7 +1155,11 @@ def _kv_quantization_for_key(key: str) -> str | None:
                     value = value.float() / maxbound
 
                     # Warn if scale exceeds threshold
-                    if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5:
+                    is_fp8_scale = layer_quantization == KV_CACHE_FP8 or (
+                        layer_quantization == KV_CACHE_FP8_K_NVFP4_V
+                        and key.endswith("k_bmm_quantizer._amax")
+                    )
+                    if is_fp8_scale and value.item() > 0.5:
                         logger.warning(
                             "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                         )
diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
index a9518c86fc0..943a8f0e8c0 100644
--- a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
+++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -1,7 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Layer-wise KV-cache search over cast-mode FP8 and NVFP4 at 5.4 bits/scalar.
+# Layer-wise KV-cache search over full FP8, FP8-K/NVFP4-V, and full NVFP4 at
+# 5.4 bits/scalar.
 # Cast mode fixes amax and skips PTQ calibration; isolated full-vocabulary forward
 # KL is measured with every other eligible attention layer kept in BF16.
 
@@ -11,10 +12,12 @@ imports:
   base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers
   kv_fp8_cast: configs/ptq/units/kv_fp8_cast
   kv_nvfp4_cast: configs/ptq/units/kv_nvfp4_cast
+  fp8: configs/numerics/fp8
+  nvfp4: configs/numerics/nvfp4
 
 metadata:
   recipe_type: auto_quantize
-  description: Layer-wise FP8/NVFP4 KV-cache cast search at 5.4 bits using forward KL.
+  description: Layer-wise three-format KV-cache cast search at 5.4 bits using forward KL.
 
 auto_quantize:
   constraints:
@@ -27,6 +30,19 @@ auto_quantize:
         method: max
         skip_forward_without_activation_calib: true
       effective_bits: 8.0
+    - quant_cfg:
+        - quantizer_name: '*k_bmm_quantizer'
+          cfg:
+            $import: fp8
+            use_constant_amax: true
+        - quantizer_name: '*v_bmm_quantizer'
+          cfg:
+            $import: nvfp4
+            use_constant_amax: true
+      algorithm:
+        method: max
+        skip_forward_without_activation_calib: true
+      effective_bits: 6.25
     - quant_cfg:
         - $import: kv_nvfp4_cast
       algorithm:
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 2911f498202..fc07be31205 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -95,6 +95,7 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
     assert inputs["method"] == "kl_div"
     assert [config["effective_bits"] for config, _ in inputs["quantization_formats"]] == [
         8.0,
+        6.25,
         4.5,
     ]
     assert "kv_cache_quant_cfg" not in inputs
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index e7220954dd1..d51222b1f3b 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1933,11 +1933,19 @@ def test_load_recipe_kv_autoquantize_contract():
     assert aq.constraints.effective_bits is None
     assert aq.constraints.kv_effective_bits == 5.4
     assert aq.auto_quantize_method == "kl_div"
-    assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5]
+    assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 6.25, 4.5]
+    assert [entry.quantizer_name for entry in aq.candidate_formats[1].quant_cfg] == [
+        "*k_bmm_quantizer",
+        "*v_bmm_quantizer",
+    ]
     for fmt in aq.candidate_formats:
-        (entry,) = fmt.quant_cfg
-        assert entry.quantizer_name == "*[kv]_bmm_quantizer"
-        assert entry.cfg.use_constant_amax
+        for entry in fmt.quant_cfg:
+            assert entry.quantizer_name in {
+                "*[kv]_bmm_quantizer",
+                "*k_bmm_quantizer",
+                "*v_bmm_quantizer",
+            }
+            assert entry.cfg.use_constant_amax
         assert fmt.algorithm["skip_forward_without_activation_calib"]
 
 
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index cbcf97a24a6..987987b73ac 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -28,6 +28,7 @@
 from modelopt.torch.export.layer_utils import get_quantization_format
 from modelopt.torch.export.model_config import (
     KV_CACHE_FP8,
+    KV_CACHE_FP8_K_NVFP4_V,
     KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
     QUANTIZATION_NVFP4,
@@ -76,6 +77,7 @@ def __init__(self):
     model = torch.nn.Module()
     model.attn0 = FakeAttention()
     model.attn1 = FakeAttention()
+    model.attn2 = FakeAttention()
     mtq.set_quantizer_by_cfg(
         model.attn0,
         [
@@ -98,6 +100,23 @@ def __init__(self):
             }
         ],
     )
+    mtq.set_quantizer_by_cfg(
+        model.attn2,
+        [
+            {
+                "quantizer_name": "*k_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "use_constant_amax": True},
+            },
+            {
+                "quantizer_name": "*v_bmm_quantizer",
+                "cfg": {
+                    "num_bits": (2, 1),
+                    "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                    "use_constant_amax": True,
+                },
+            },
+        ],
+    )
 
     quantization = get_quant_config(model)["quantization"]
     assert quantization["quant_algo"] == "MIXED_PRECISION"
@@ -106,6 +125,7 @@ def __init__(self):
     assert quantization["kv_cache_quantized_layers"] == {
         "attn0": {"quant_algo": "FP8"},
         "attn1": {"quant_algo": "NVFP4"},
+        "attn2": {"quant_algo": "FP8_K_NVFP4_V"},
     }
 
 
@@ -119,7 +139,14 @@ def test_mixed_kv_cache_postprocess_uses_each_layers_format():
     layer_formats = {
         "attn0": {"quant_algo": KV_CACHE_FP8},
         "attn1": {"quant_algo": KV_CACHE_NVFP4},
+        "attn2": {"quant_algo": KV_CACHE_FP8_K_NVFP4_V},
     }
+    state_dict.update(
+        {
+            "attn2.k_bmm_quantizer._amax": torch.tensor([448.0]),
+            "attn2.v_bmm_quantizer._amax": torch.tensor([112.0]),
+        }
+    )
 
     processed = postprocess_state_dict(state_dict, 448.0, layer_formats)
 
@@ -128,6 +155,8 @@ def test_mixed_kv_cache_postprocess_uses_each_layers_format():
         "attn0.v_proj.v_scale": torch.tensor([0.5]),
         "attn1.k_proj.k_scale": torch.tensor([0.25]),
         "attn1.v_proj.v_scale": torch.tensor([0.125]),
+        "attn2.k_proj.k_scale": torch.tensor([1.0]),
+        "attn2.v_proj.v_scale": torch.tensor([0.25]),
     }
 
 

From c6e9ac680df8d8ed1dc253ec22b3ec50d33fe7aa Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 18 Aug 2026 10:54:00 -0700
Subject: [PATCH 03/15] Harden KV cache AutoQuant search

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |   6 +-
 modelopt/recipe/config.py                     |   5 +-
 modelopt/torch/export/quant_utils.py          |   4 +
 .../torch/quantization/kv_cache_auto_quant.py | 101 +++++++++++++++---
 modelopt/torch/quantization/model_quant.py    |  15 +--
 .../torch/export/test_get_quantization.py     |  30 ++++++
 .../quantization/test_kv_cache_auto_quant.py  |  48 ++++++++-
 7 files changed, 181 insertions(+), 28 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index be9dfdfa117..f3d283753f3 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -426,8 +426,10 @@ Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall
 `--kv_cache_qformat` (default `fp8_cast`) unless they set an explicit `kv_cache` field.
 
 KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
-`candidate_formats` are complete K/V cache configs with exact config-level `effective_bits`;
-BF16 is used only as the isolated-KL reference, not as a solver choice. The shipped canary recipe
+`candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
+packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are
+preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice.
+The shipped canary recipe
 searches cast-mode FP8 K/V (8.0 bits/scalar), FP8-K/NVFP4-V (6.25 bits/scalar),
 and packed NVFP4 K/V (4.5 bits/scalar) at 5.4 bits/scalar:
 
diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index c4dc89f1841..664953a137c 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -161,8 +161,9 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
         default=None,
         title="Effective bits per KV-cache scalar",
         description=(
-            "Average KV-cache storage bits target for layer-wise KV AutoQuant, in (0, 16]. "
-            "Exactly one of effective_bits and kv_effective_bits may be set."
+            "Average KV-cache storage bits target across eligible layers for layer-wise KV "
+            "AutoQuant, in (0, 16]. Exactly one of effective_bits and kv_effective_bits may "
+            "be set."
         ),
     )
     cost_model: Literal["weight", "active_moe"] = ModeloptField(
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 25e930b40fc..2d11516c883 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -452,6 +452,10 @@ def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None:
                 return KV_CACHE_FP8_K_NVFP4_V
             if k_dtype == v_dtype:
                 return k_dtype
+            raise NotImplementedError(
+                "Unsupported mixed K/V cache quantization pair: "
+                f"K uses {k_dtype}, while V uses {v_dtype}."
+            )
 
         # Case where the module has both k_bmm_quantizer and v_bmm_quantizer
         # Still check for output quantizer for the unified_megatron_export path
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 14d03713040..847a6d54c35 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -150,10 +150,41 @@ def _apply_layer_quantizers(module: nn.Module, quantizers: dict[str, TensorQuant
         setattr(module, attr, quantizer)
 
 
-def _get_logits(forward_step: Callable[[nn.Module, Any], torch.Tensor], model, data):
+def _set_only_candidate_quantizers_enabled(
+    model: nn.Module, candidate_quantizers: set[TensorQuantizer]
+) -> list[tuple[TensorQuantizer, bool]]:
+    """Enable the active KV candidates while temporarily disabling every other quantizer."""
+    enable_states = []
+    for module in model.modules():
+        if not isinstance(module, TensorQuantizer):
+            continue
+        enable_states.append((module, module.is_enabled))
+        if module in candidate_quantizers:
+            module.enable()
+        else:
+            module.disable()
+    return enable_states
+
+
+def _restore_quantizer_enable_states(enable_states: list[tuple[TensorQuantizer, bool]]) -> None:
+    for quantizer, is_enabled in enable_states:
+        if is_enabled:
+            quantizer.enable()
+        else:
+            quantizer.disable()
+
+
+def _get_logits(
+    forward_step: Callable[[nn.Module, Any], torch.Tensor], model: nn.Module, data: Any
+) -> torch.Tensor:
     logits = forward_step(model, data)
     if not isinstance(logits, torch.Tensor):
         raise TypeError("KV-cache AutoQuant forward_step must return a logits tensor.")
+    if logits.ndim < 2 or logits.shape[-1] == 0:
+        raise ValueError(
+            "KV-cache AutoQuant forward_step must return logits with a non-empty vocabulary "
+            "dimension."
+        )
     if not torch.isfinite(logits).all():
         raise ValueError("KV-cache AutoQuant encountered NaN or Inf logits.")
     return logits
@@ -274,10 +305,10 @@ def auto_quantize_kv_cache(
     """Select one supplied K/V format per attention layer using isolated forward KL.
 
     Candidate formats are format-agnostic ``QuantizeConfig`` dictionaries. Each must
-    configure K and V together and declare ``effective_bits`` matching their runtime
-    packed storage. Candidate-specific calibration runs before scoring; cast-style
-    constant-amax formats skip calibration forwards through their normal algorithm
-    configuration.
+    configure K and V together and declare ``effective_bits`` matching its packed
+    storage per K-or-V scalar, including scale overhead. Candidate-specific calibration
+    runs with non-KV quantizers disabled before scoring; cast-style constant-amax formats
+    skip calibration forwards through their normal algorithm configuration.
     """
     if set(constraints) != {"kv_effective_bits"}:
         raise ValueError(
@@ -333,6 +364,8 @@ def auto_quantize_kv_cache(
         for name, module, _ in layers
     }
 
+    is_training = model.training
+    model.eval()
     try:
         for name, module, _ in layers:
             _apply_layer_quantizers(module, disabled_quantizers[name])
@@ -340,6 +373,10 @@ def auto_quantize_kv_cache(
         state: dict[str, Any] | None = None
         if checkpoint is not None and os.path.exists(checkpoint):
             restored = safe_load(checkpoint)
+            if not isinstance(restored, dict):
+                raise ValueError(
+                    "KV-cache AutoQuant checkpoint must contain a search-state dictionary."
+                )
             if _checkpoint_state_is_compatible(restored, signature):
                 state = restored
                 if verbose:
@@ -375,16 +412,26 @@ def calibration_loop(calibration_model):
                                 break
                             _get_logits(forward_step, calibration_model, data)
 
-                    modelopt_state = ModeloptStateManager(model).state_dict()
-                    original_mode_count = len(modelopt_state)
+                    active_quantizers = {
+                        quantizer
+                        for layer_name, _, _ in layers
+                        for quantizer in candidate_quantizers[layer_name][candidate_name].values()
+                    }
+                    enable_states = _set_only_candidate_quantizers_enabled(model, active_quantizers)
+                    modelopt_state = None
+                    original_mode_count = 0
                     try:
+                        modelopt_state = ModeloptStateManager(model).state_dict()
+                        original_mode_count = len(modelopt_state)
                         calibrate(
                             model,
                             algorithm=config.algorithm,
                             forward_loop=calibration_loop,
                         )
                     finally:
-                        del modelopt_state[original_mode_count:]
+                        if modelopt_state is not None:
+                            del modelopt_state[original_mode_count:]
+                        _restore_quantizer_enable_states(enable_states)
 
                 for layer_name, module, _ in layers:
                     _apply_layer_quantizers(module, disabled_quantizers[layer_name])
@@ -404,8 +451,8 @@ def calibration_loop(calibration_model):
 
         assert state is not None
         if not state.get("layers"):
-            score_sums = {
-                layer_name: dict.fromkeys(candidate_names, 0.0) for layer_name, _, _ in layers
+            score_sums: dict[str, dict[str, torch.Tensor | None]] = {
+                layer_name: dict.fromkeys(candidate_names) for layer_name, _, _ in layers
             }
             scored_tokens = 0
             scored_steps = 0
@@ -428,22 +475,44 @@ def calibration_loop(calibration_model):
                             module, candidate_quantizers[layer_name][candidate_name]
                         )
                         logits_quant = _get_logits(forward_step, model, data)
+                        if logits_quant.shape != logits_ref.shape:
+                            raise ValueError(
+                                "KV-cache AutoQuant forward_step returned different reference and "
+                                f"candidate logits shapes: {tuple(logits_ref.shape)} and "
+                                f"{tuple(logits_quant.shape)}."
+                            )
                         score = F.kl_div(
                             torch.log_softmax(logits_quant.float(), dim=-1),
                             log_prob_ref,
                             reduction="sum",
                             log_target=True,
                         )
-                        score_sums[layer_name][candidate_name] += float(score.item())
+                        previous_score = score_sums[layer_name][candidate_name]
+                        score_sums[layer_name][candidate_name] = (
+                            score if previous_score is None else previous_score + score
+                        )
                         _apply_layer_quantizers(module, disabled_quantizers[layer_name])
                 scored_steps += 1
 
             if scored_steps == 0 or scored_tokens == 0:
                 raise ValueError("KV-cache AutoQuant data_loader produced no scoring batches.")
-            scores = [
-                [score_sums[layer_name][candidate] / scored_tokens for candidate in candidate_names]
-                for layer_name, _, _ in layers
-            ]
+            scores = []
+            for layer_name, _, _ in layers:
+                layer_scores = []
+                for candidate_name in candidate_names:
+                    score_sum = score_sums[layer_name][candidate_name]
+                    if score_sum is None:
+                        raise RuntimeError(
+                            "KV-cache AutoQuant did not collect a score for "
+                            f"{layer_name!r}/{candidate_name!r}."
+                        )
+                    if not torch.isfinite(score_sum):
+                        raise ValueError(
+                            "KV-cache AutoQuant produced a non-finite KL score for "
+                            f"{layer_name!r}/{candidate_name!r}."
+                        )
+                    layer_scores.append(float(score_sum.item()) / scored_tokens)
+                scores.append(layer_scores)
             selections, status = _solve_additive_recipe(
                 [name for name, _, _ in layers],
                 [weight for _, _, weight in layers],
@@ -517,3 +586,5 @@ def calibration_loop(calibration_model):
         for layer_name, module, _ in layers:
             _apply_layer_quantizers(module, original_quantizers[layer_name])
         raise
+    finally:
+        model.train(is_training)
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 4ea9f49dead..a1a5f43bfd4 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -675,21 +675,24 @@ def auto_quantize_kv_cache(
 
     Unlike weight AutoQuant, BF16/no-quant is only the scoring reference and is
     not an implicit solver choice. Each supplied format must configure K and V
-    together and declare exact config-level ``effective_bits``. Formats use their
-    own calibration algorithm; cast-mode constant-amax candidates skip calibration
-    forwards and provide the fastest turnaround.
+    together and declare config-level ``effective_bits`` equal to their packed
+    storage cost per K-or-V scalar, including scale overhead. The budget is weighted
+    by the K/V projection widths of eligible layers. Formats use their own calibration
+    algorithm; cast-mode constant-amax candidates skip calibration forwards.
 
     Args:
         model: Model whose attention K/V quantizers will be searched.
-        constraints: A ``{"kv_effective_bits": target}`` storage constraint.
+        constraints: A ``{"kv_effective_bits": target}`` storage constraint across
+            eligible layers.
         quantization_formats: Candidate ``QuantizeConfig`` dictionaries, optionally
             paired with display names.
         data_loader: Re-iterable calibration and scoring batches.
-        forward_step: Callable returning the full logits tensor for one batch.
+        forward_step: Callable returning full-vocabulary logits for the token positions
+            to score in one batch.
         num_calib_steps: Maximum calibration batches per candidate.
         num_score_steps: Maximum batches used for isolated forward-KL scoring.
         disabled_layers: Optional layer-name patterns excluded from the search and
-            preserved in their existing KV-cache format.
+            bit budget, and preserved in their existing KV-cache format.
         verbose: Whether to print progress and selected formats.
         checkpoint: Optional path for resumable calibration and sensitivity state.
 
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index 987987b73ac..6c778ec098e 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -129,6 +129,36 @@ def __init__(self):
     }
 
 
+def test_unsupported_asymmetric_kv_cache_pair_fails_export():
+    class FakeAttention(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.k_bmm_quantizer = TensorQuantizer()
+            self.v_bmm_quantizer = TensorQuantizer()
+
+    model = FakeAttention()
+    mtq.set_quantizer_by_cfg(
+        model,
+        [
+            {
+                "quantizer_name": "*k_bmm_quantizer",
+                "cfg": {
+                    "num_bits": (2, 1),
+                    "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                    "use_constant_amax": True,
+                },
+            },
+            {
+                "quantizer_name": "*v_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "use_constant_amax": True},
+            },
+        ],
+    )
+
+    with pytest.raises(NotImplementedError, match="Unsupported mixed K/V cache"):
+        get_quant_config(model)
+
+
 def test_mixed_kv_cache_postprocess_uses_each_layers_format():
     state_dict = {
         "attn0.k_bmm_quantizer._amax": torch.tensor([448.0]),
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index 0dbf792b24d..acc132fade4 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -183,6 +183,7 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
 
     assert state["best"]["effective_bits"] == pytest.approx(6.0)
     assert state["best"]["is_satisfied"]
+    assert model.training
     assert {layer["selected"] for layer in state["layers"].values()} == {
         "int8",
         "int4",
@@ -193,7 +194,7 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
         expected_bits = 8 if layer_state["selected"] == "int8" else 4
         assert layer.k_bmm_quantizer.num_bits == expected_bits
 
-    restored_model = _ToyKVModel()
+    restored_model = _ToyKVModel().eval()
     restored_model, restored_state = auto_quantize_kv_cache(
         restored_model,
         {"kv_effective_bits": 6.0},
@@ -206,12 +207,51 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
     )
 
     assert restored_state == state
+    assert not restored_model.training
     for layer_name, layer_state in restored_state["layers"].items():
         layer = restored_model.get_submodule(layer_name)
         expected_bits = 8 if layer_state["selected"] == "int8" else 4
         assert layer.k_bmm_quantizer.num_bits == expected_bits
 
 
+def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state():
+    model = _ToyKVModel()
+    original_quantizers = {
+        name: (module.k_bmm_quantizer, module.v_bmm_quantizer)
+        for name, module in (("attn0", model.attn0), ("attn1", model.attn1))
+    }
+    candidates = [
+        (
+            {
+                "quant_cfg": [
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 4.0,
+            },
+            "int4",
+        )
+    ]
+
+    with pytest.raises(ValueError, match="non-empty vocabulary dimension"):
+        auto_quantize_kv_cache(
+            model,
+            {"kv_effective_bits": 4.0},
+            candidates,
+            [torch.randn(2, 3, 8)],
+            lambda *_: torch.ones(8),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+    assert model.training
+    for name, module in (("attn0", model.attn0), ("attn1", model.attn1)):
+        assert (module.k_bmm_quantizer, module.v_bmm_quantizer) == original_quantizers[name]
+
+
 def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
     torch.manual_seed(123)
     model = get_tiny_llama(num_hidden_layers=2)
@@ -298,6 +338,7 @@ def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
     }
     model = mtq.quantize(model, fixed_kv_config)
     model.model.layers[0].self_attn.q_proj.weight_quantizer.enable()
+    assert not hasattr(model.model.layers[0].self_attn.q_proj.weight_quantizer, "_amax")
 
     model, state = mtq.auto_quantize_kv_cache(
         model,
@@ -308,10 +349,10 @@ def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
                     "quant_cfg": [
                         {
                             "quantizer_name": "*[kv]_bmm_quantizer",
-                            "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                            "cfg": {"num_bits": 4},
                         }
                     ],
-                    "algorithm": None,
+                    "algorithm": "max",
                     "effective_bits": 4.0,
                 },
                 "int4",
@@ -327,6 +368,7 @@ def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
     assert set(state["layers"]) == {"model.layers.0.self_attn"}
     assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == 4
     assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled
+    assert not hasattr(model.model.layers[0].self_attn.q_proj.weight_quantizer, "_amax")
     fixed_attention = model.model.layers[1].self_attn
     assert fixed_attention.k_bmm_quantizer.is_enabled
     assert fixed_attention.v_bmm_quantizer.is_enabled

From d26d9128256a06ffc51fa08ba514e80655e9464f Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 18 Aug 2026 11:08:21 -0700
Subject: [PATCH 04/15] Test KV AutoQuant report export

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/export/unified_export_hf.py    |  4 +-
 .../torch/export/test_convert_hf_config.py    | 41 +++++++++++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 814e446ba65..37e7f36542b 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1497,9 +1497,7 @@ def _write_hf_export_config(
                     name_mapper(name): value
                     for name, value in kv_autoquant_report["layers"].items()
                 }
-                signature_layers = kv_autoquant_report.get("search_signature", {}).get(
-                    "layers", []
-                )
+                signature_layers = kv_autoquant_report.get("search_signature", {}).get("layers", [])
                 for layer in signature_layers:
                     layer["name"] = name_mapper(layer["name"])
             with open(f"{export_dir}/kv_cache_auto_quantize_report.json", "w") as file:
diff --git a/tests/unit/torch/export/test_convert_hf_config.py b/tests/unit/torch/export/test_convert_hf_config.py
index ab90168b750..a0b8295b485 100644
--- a/tests/unit/torch/export/test_convert_hf_config.py
+++ b/tests/unit/torch/export/test_convert_hf_config.py
@@ -13,7 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+
+import torch
+
 from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
+from modelopt.torch.export.unified_export_hf import _write_hf_export_config
 
 
 def test_convert_mixed_kv_cache_config_preserves_layer_map():
@@ -40,3 +45,39 @@ def test_convert_mixed_kv_cache_config_preserves_layer_map():
     assert converted["kv_cache_quant_algo"] == "MIXED_PRECISION"
     assert converted["kv_cache_quantized_layers"] == layer_map
     assert converted["kv_cache_schema_version"] == 1
+
+
+def test_write_hf_export_config_writes_mapped_kv_autoquant_report(tmp_path):
+    layer_name = "model.layers.0.self_attn"
+    model = torch.nn.Module()
+    model._modelopt_kv_cache_auto_quantize_state = {
+        "layers": {layer_name: {"selected": "fp8"}},
+        "search_signature": {"layers": [{"name": layer_name}]},
+    }
+    quant_config = {
+        "producer": {"name": "modelopt", "version": "test"},
+        "quantization": {
+            "quant_algo": None,
+            "kv_cache_quant_algo": "MIXED_PRECISION",
+            "kv_cache_quantized_layers": {layer_name: {"quant_algo": "FP8"}},
+            "kv_cache_schema_version": 1,
+        },
+    }
+    (tmp_path / "config.json").write_text("{}")
+
+    _write_hf_export_config(
+        model,
+        quant_config,
+        tmp_path,
+        name_mapper=lambda name: f"hub.{name}",
+    )
+
+    report = json.loads((tmp_path / "kv_cache_auto_quantize_report.json").read_text())
+    assert report["layers"] == {f"hub.{layer_name}": {"selected": "fp8"}}
+    assert report["search_signature"]["layers"] == [{"name": f"hub.{layer_name}"}]
+    assert model._modelopt_kv_cache_auto_quantize_state["layers"] == {
+        layer_name: {"selected": "fp8"}
+    }
+    assert (tmp_path / "hf_quant_config.json").is_file()
+    exported_config = json.loads((tmp_path / "config.json").read_text())
+    assert exported_config["quantization_config"]["kv_cache_quant_algo"] == "MIXED_PRECISION"

From 300c95f26a5c5362371c55b5473091c24743b9b7 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 18 Aug 2026 13:39:12 -0700
Subject: [PATCH 05/15] Fix KV AutoQuant state and export handling

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |   8 ++
 examples/hf_ptq/hf_ptq.py                     |  19 ++-
 modelopt/torch/export/quant_utils.py          |  58 +++++---
 modelopt/torch/export/unified_export_hf.py    |   8 +-
 .../export/unified_export_hf_streaming.py     |   8 +-
 .../torch/quantization/kv_cache_auto_quant.py | 133 +++++++++---------
 modelopt/torch/quantization/model_quant.py    |  67 +++++++--
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  11 ++
 .../unit/torch/export/test_offload_export.py  |  32 ++++-
 .../quantization/test_kv_cache_auto_quant.py  | 126 +++++++++++++----
 10 files changed, 332 insertions(+), 138 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index f3d283753f3..b451b104efb 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -446,6 +446,14 @@ the selected formats in `kv_cache_quantized_layers` and writes the JSON-safe sen
 `kv_cache_auto_quantize_report.json`; `--auto_quantize_checkpoint` stores the resumable raw search
 state.
 
+> [!NOTE]
+> Layer-wise KV checkpoints require the companion
+> [vLLM mixed-KV metadata consumer](https://github.com/vllm-project/vllm/pull/52813) or a later
+> vLLM release containing it. The repository's currently pinned vLLM 0.26.0 does not consume
+> `kv_cache_quantized_layers`, so these checkpoints are export-only in that stock environment.
+> Full FP8 K/V and full NVFP4 K/V use existing vLLM kernels; FP8-K/NVFP4-V within one layer also
+> requires the separate mixed-K/V kernel implementation.
+
 The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an
 interrupted search (skips re-scoring):
 
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 0bfa07cdc32..0aca1ffde5a 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -98,6 +98,19 @@
 RAND_SEED = 1234
 
 
+def _select_unpadded_logits(logits: torch.Tensor, batch: dict[str, Any]) -> torch.Tensor:
+    """Return logits only for token positions selected by ``attention_mask``."""
+    attention_mask = batch.get("attention_mask")
+    if attention_mask is None:
+        return logits
+    if logits.shape[:-1] != attention_mask.shape:
+        raise ValueError(
+            "AutoQuant KL logits and attention_mask must have matching token dimensions; "
+            f"got {tuple(logits.shape[:-1])} and {tuple(attention_mask.shape)}."
+        )
+    return logits[attention_mask.bool()]
+
+
 def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool:
     """Return True if this KV cfg pins ``use_constant_amax`` on the bmm quantizer.
 
@@ -490,8 +503,10 @@ def forward_step(model, batch):
             output = model(**inputs_)
             if is_base_model:
                 assert full_model is not None
-                return full_model.lm_head(output.last_hidden_state)
-            return output.logits
+                logits = full_model.lm_head(output.last_hidden_state)
+            else:
+                logits = output.logits
+            return _select_unpadded_logits(logits, batch)
 
     else:
         raise ValueError(
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 2d11516c883..ec451e3f2e9 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -1034,7 +1034,7 @@ def _postprocess_single_tensor(
     key: str,
     value: torch.Tensor,
     kv_cache_max_bound: float,
-    kv_cache_format: str | None,
+    kv_cache_format: str | dict[str, dict[str, str]] | None,
     is_modelopt_qlora: bool = False,
 ) -> tuple[str | None, torch.Tensor | None]:
     """Per-tensor subset of :func:`postprocess_state_dict`, for streaming export.
@@ -1068,12 +1068,20 @@ def _postprocess_single_tensor(
         if key.endswith(old_suffix):
             prefix = key[: -len(old_suffix)]
             if "_amax" in key:
-                assert kv_cache_format in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], (
-                    "Invalid KV cache quantization format."
-                )
+                layer_quantization = _resolve_kv_cache_format_for_key(key, kv_cache_format)
+                assert layer_quantization in [
+                    KV_CACHE_FP8,
+                    KV_CACHE_FP8_K_NVFP4_V,
+                    KV_CACHE_NVFP4,
+                    KV_CACHE_NVFP4_AFFINE,
+                ], "Invalid KV cache quantization format."
                 assert kv_cache_max_bound > 0, "Maxbound must be greater than zero."
                 value = value.float() / kv_cache_max_bound
-                if kv_cache_format == KV_CACHE_FP8 and value.item() > 0.5:
+                is_fp8_scale = layer_quantization == KV_CACHE_FP8 or (
+                    layer_quantization == KV_CACHE_FP8_K_NVFP4_V
+                    and key.endswith("k_bmm_quantizer._amax")
+                )
+                if is_fp8_scale and value.item() > 0.5:
                     logger.warning(
                         "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                     )
@@ -1084,6 +1092,32 @@ def _postprocess_single_tensor(
     return None, None
 
 
+def _resolve_kv_cache_format_for_key(
+    key: str, quantization: str | dict[str, dict[str, str]] | None
+) -> str | None:
+    """Resolve uniform or per-layer KV-cache metadata for one state-dict key."""
+    if not isinstance(quantization, dict):
+        return quantization
+    matches = [
+        (layer_name, layer_config.get("quant_algo"))
+        for layer_name, layer_config in quantization.items()
+        if key == layer_name or key.startswith(layer_name + ".")
+    ]
+    if not matches:
+        return None
+    return max(matches, key=lambda item: len(item[0]))[1]
+
+
+def _get_kv_cache_postprocess_config(
+    quantization_details: dict[str, Any],
+) -> str | dict[str, dict[str, str]] | None:
+    """Return the uniform format or layer map consumed by both HF exporters."""
+    kv_cache_format = quantization_details.get("kv_cache_quant_algo")
+    if kv_cache_format == "MIXED_PRECISION":
+        return quantization_details.get("kv_cache_quantized_layers", {})
+    return kv_cache_format
+
+
 def postprocess_state_dict(
     state_dict: dict,
     maxbound: float,
@@ -1118,18 +1152,6 @@ def _export_key(key: str) -> str:
 
     post_state_dict = {}
 
-    def _kv_quantization_for_key(key: str) -> str | None:
-        if not isinstance(quantization, dict):
-            return quantization
-        matches = [
-            (layer_name, layer_config.get("quant_algo"))
-            for layer_name, layer_config in quantization.items()
-            if key == layer_name or key.startswith(layer_name + ".")
-        ]
-        if not matches:
-            return None
-        return max(matches, key=lambda item: len(item[0]))[1]
-
     for key, value in state_dict.items():
         # Skip problematic parameters for specific model architectures, e.g., Nemotron Nano VL models
         if key == "vision_model.radio_model.summary_idxs":
@@ -1147,7 +1169,7 @@ def _kv_quantization_for_key(key: str) -> str | None:
                 prefix = key[: -len(old_suffix)]
 
                 if "_amax" in key:
-                    layer_quantization = _kv_quantization_for_key(key)
+                    layer_quantization = _resolve_kv_cache_format_for_key(key, quantization)
                     assert layer_quantization in [
                         KV_CACHE_FP8,
                         KV_CACHE_FP8_K_NVFP4_V,
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 37e7f36542b..dd964098c6b 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -101,6 +101,7 @@
     revert_weight_conversion_quant_aware,
 )
 from .quant_utils import (
+    _get_kv_cache_postprocess_config,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -1012,12 +1013,7 @@ def _export_transformers_checkpoint(
     # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization.
     kv_cache_max_bound = 448
     quantization_details = quant_config["quantization"]
-    kv_cache_format = quantization_details["kv_cache_quant_algo"]
-    kv_cache_postprocess_config = (
-        quantization_details["kv_cache_quantized_layers"]
-        if kv_cache_format == "MIXED_PRECISION"
-        else kv_cache_format
-    )
+    kv_cache_postprocess_config = _get_kv_cache_postprocess_config(quantization_details)
 
     quantized_state_dict = postprocess_state_dict(
         quantized_state_dict,
diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py
index 50b5d797c6a..73ca33e2a26 100644
--- a/modelopt/torch/export/unified_export_hf_streaming.py
+++ b/modelopt/torch/export/unified_export_hf_streaming.py
@@ -34,7 +34,11 @@
 from safetensors.torch import save_file
 
 from .quant_aware_conversion import build_reverse_name_mapper
-from .quant_utils import _postprocess_single_tensor, get_quant_config
+from .quant_utils import (
+    _get_kv_cache_postprocess_config,
+    _postprocess_single_tensor,
+    get_quant_config,
+)
 from .registry import ExportContext
 from .unified_export_hf import (
     _add_mtp_exclusions,
@@ -253,7 +257,7 @@ def _export_transformers_checkpoint_streaming(
 
     # --- Per-tensor constants ---
     kv_cache_max_bound = 448
-    kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"]
+    kv_cache_format = _get_kv_cache_postprocess_config(quant_config["quantization"])
 
     # --- Tied alias keys to skip ---
     # data_ptr() is unreliable for disk-offloaded weights, so we use _tied_weights_keys.
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 847a6d54c35..cb2dc81a563 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -26,7 +26,6 @@
 import torch.nn.functional as F
 from tqdm import tqdm
 
-from modelopt.torch.opt.conversion import ModeloptStateManager
 from modelopt.torch.opt.searcher import LPS
 from modelopt.torch.utils import print_rank_0, safe_load, safe_save
 
@@ -97,6 +96,55 @@ def _validate_kv_only_config(config: QuantizeConfig) -> None:
             "k_bmm_quantizer and v_bmm_quantizer."
         )
 
+    algorithm = config.algorithm
+    if algorithm is None:
+        return
+    if isinstance(algorithm, str):
+        algorithm_method = algorithm
+    elif isinstance(algorithm, dict):
+        algorithm_method = algorithm.get("method")
+    else:
+        algorithm_method = getattr(algorithm, "method", None)
+    if algorithm_method != "max":
+        raise ValueError(
+            "KV-cache AutoQuant supports only non-structural calibration algorithms "
+            f"None and 'max'; got {algorithm_method!r}."
+        )
+
+
+def _validate_search_inputs(
+    constraints: dict[str, Any],
+    quantization_formats: list[tuple[dict[str, Any], str]],
+    num_calib_steps: int,
+    num_score_steps: int,
+) -> tuple[float, list[tuple[str, QuantizeConfig]]]:
+    """Validate a KV-cache search before the caller converts the model."""
+    if set(constraints) != {"kv_effective_bits"}:
+        raise ValueError(
+            "KV-cache AutoQuant constraints must contain only kv_effective_bits; "
+            f"got {sorted(constraints)}."
+        )
+    target_bits = float(constraints["kv_effective_bits"])
+    if not (0 < target_bits <= 16):
+        raise ValueError(f"kv_effective_bits must be in (0, 16], got {target_bits}.")
+    if num_calib_steps <= 0:
+        raise ValueError("num_calib_steps must be positive.")
+    if num_score_steps <= 0:
+        raise ValueError("num_score_steps must be positive.")
+
+    candidates = []
+    seen_names = set()
+    for raw_config, name in quantization_formats:
+        if name in seen_names:
+            raise ValueError(f"Duplicate KV-cache AutoQuant candidate name: {name!r}.")
+        config = QuantizeConfig(**raw_config)
+        _validate_kv_only_config(config)
+        candidates.append((name, config))
+        seen_names.add(name)
+    if not candidates:
+        raise ValueError("KV-cache AutoQuant requires at least one candidate format.")
+    return target_bits, candidates
+
 
 def _projection_width(module: nn.Module, side: str) -> int | None:
     projection = getattr(module, f"{side}_proj", None)
@@ -150,30 +198,6 @@ def _apply_layer_quantizers(module: nn.Module, quantizers: dict[str, TensorQuant
         setattr(module, attr, quantizer)
 
 
-def _set_only_candidate_quantizers_enabled(
-    model: nn.Module, candidate_quantizers: set[TensorQuantizer]
-) -> list[tuple[TensorQuantizer, bool]]:
-    """Enable the active KV candidates while temporarily disabling every other quantizer."""
-    enable_states = []
-    for module in model.modules():
-        if not isinstance(module, TensorQuantizer):
-            continue
-        enable_states.append((module, module.is_enabled))
-        if module in candidate_quantizers:
-            module.enable()
-        else:
-            module.disable()
-    return enable_states
-
-
-def _restore_quantizer_enable_states(enable_states: list[tuple[TensorQuantizer, bool]]) -> None:
-    for quantizer, is_enabled in enable_states:
-        if is_enabled:
-            quantizer.enable()
-        else:
-            quantizer.disable()
-
-
 def _get_logits(
     forward_step: Callable[[nn.Module, Any], torch.Tensor], model: nn.Module, data: Any
 ) -> torch.Tensor:
@@ -306,34 +330,13 @@ def auto_quantize_kv_cache(
 
     Candidate formats are format-agnostic ``QuantizeConfig`` dictionaries. Each must
     configure K and V together and declare ``effective_bits`` matching its packed
-    storage per K-or-V scalar, including scale overhead. Candidate-specific calibration
-    runs with non-KV quantizers disabled before scoring; cast-style constant-amax formats
-    skip calibration forwards through their normal algorithm configuration.
+    storage per K-or-V scalar, including scale overhead. Candidate calibration is scoped
+    to the candidate K/V quantizers, while pre-existing fixed quantizers keep executing
+    with frozen state. Cast-style constant-amax formats may skip calibration forwards.
     """
-    if set(constraints) != {"kv_effective_bits"}:
-        raise ValueError(
-            "KV-cache AutoQuant constraints must contain only kv_effective_bits; "
-            f"got {sorted(constraints)}."
-        )
-    target_bits = float(constraints["kv_effective_bits"])
-    if not (0 < target_bits <= 16):
-        raise ValueError(f"kv_effective_bits must be in (0, 16], got {target_bits}.")
-    if num_calib_steps <= 0:
-        raise ValueError("num_calib_steps must be positive.")
-    if num_score_steps <= 0:
-        raise ValueError("num_score_steps must be positive.")
-
-    candidates = []
-    seen_names = set()
-    for raw_config, name in quantization_formats:
-        if name in seen_names:
-            raise ValueError(f"Duplicate KV-cache AutoQuant candidate name: {name!r}.")
-        config = QuantizeConfig(**raw_config)
-        _validate_kv_only_config(config)
-        candidates.append((name, config))
-        seen_names.add(name)
-    if not candidates:
-        raise ValueError("KV-cache AutoQuant requires at least one candidate format.")
+    target_bits, candidates = _validate_search_inputs(
+        constraints, quantization_formats, num_calib_steps, num_score_steps
+    )
 
     layers = _eligible_layers(model, disabled_layers)
     signature = _search_signature(
@@ -412,26 +415,18 @@ def calibration_loop(calibration_model):
                                 break
                             _get_logits(forward_step, calibration_model, data)
 
-                    active_quantizers = {
+                    active_quantizers = [
                         quantizer
                         for layer_name, _, _ in layers
                         for quantizer in candidate_quantizers[layer_name][candidate_name].values()
-                    }
-                    enable_states = _set_only_candidate_quantizers_enabled(model, active_quantizers)
-                    modelopt_state = None
-                    original_mode_count = 0
-                    try:
-                        modelopt_state = ModeloptStateManager(model).state_dict()
-                        original_mode_count = len(modelopt_state)
-                        calibrate(
-                            model,
-                            algorithm=config.algorithm,
-                            forward_loop=calibration_loop,
-                        )
-                    finally:
-                        if modelopt_state is not None:
-                            del modelopt_state[original_mode_count:]
-                        _restore_quantizer_enable_states(enable_states)
+                    ]
+                    calibration_proxy = nn.Module()
+                    calibration_proxy.quantizers = nn.ModuleList(active_quantizers)
+                    calibrate(
+                        calibration_proxy,
+                        algorithm=config.algorithm,
+                        forward_loop=lambda _: calibration_loop(model),
+                    )
 
                 for layer_name, module, _ in layers:
                     _apply_layer_quantizers(module, disabled_quantizers[layer_name])
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index a1a5f43bfd4..14cfb990f0e 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -15,6 +15,7 @@
 
 """User-facing quantization API."""
 
+import copy
 import fnmatch
 import inspect
 import os
@@ -39,6 +40,7 @@
 from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe
 from .algorithms import get_auto_quantize_config as _get_auto_quantize_config
 from .config import QuantizeAlgoCfgType
+from .kv_cache_auto_quant import _validate_search_inputs as _validate_kv_cache_search_inputs
 from .kv_cache_auto_quant import auto_quantize_kv_cache as _auto_quantize_kv_cache
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, TensorQuantizer
@@ -712,21 +714,60 @@ def auto_quantize_kv_cache(
             raise ValueError("KV-cache AutoQuant candidate names must be non-empty strings.")
         processed_formats.append((raw_config, name))
 
-    if not is_quantized(model):
-        model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
-    return _auto_quantize_kv_cache(
-        model,
-        constraints,
-        processed_formats,
-        data_loader,
-        forward_step,
-        num_calib_steps=num_calib_steps,
-        num_score_steps=num_score_steps,
-        disabled_layers=disabled_layers,
-        verbose=verbose,
-        checkpoint=checkpoint,
+    _validate_kv_cache_search_inputs(
+        constraints, processed_formats, num_calib_steps, num_score_steps
     )
 
+    converted_for_search = not is_quantized(model)
+    conversion_snapshot = _snapshot_model_structure(model) if converted_for_search else []
+    try:
+        if converted_for_search:
+            model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
+            set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
+        return _auto_quantize_kv_cache(
+            model,
+            constraints,
+            processed_formats,
+            data_loader,
+            forward_step,
+            num_calib_steps=num_calib_steps,
+            num_score_steps=num_score_steps,
+            disabled_layers=disabled_layers,
+            verbose=verbose,
+            checkpoint=checkpoint,
+        )
+    except Exception:
+        if converted_for_search:
+            _restore_model_structure(conversion_snapshot)
+        raise
+
+
+def _snapshot_model_structure(
+    model: nn.Module,
+) -> list[tuple[nn.Module, type[nn.Module], dict[str, Any]]]:
+    """Capture lightweight module metadata for failure-atomic fresh conversion."""
+    return [
+        (
+            module,
+            type(module),
+            {
+                key: copy.copy(value) if isinstance(value, (dict, list, set)) else value
+                for key, value in module.__dict__.items()
+            },
+        )
+        for module in model.modules()
+    ]
+
+
+def _restore_model_structure(
+    snapshot: list[tuple[nn.Module, type[nn.Module], dict[str, Any]]],
+) -> None:
+    """Undo an in-place quantization conversion without copying model tensors."""
+    for module, original_type, original_state in reversed(snapshot):
+        object.__setattr__(module, "__class__", original_type)
+        module.__dict__.clear()
+        module.__dict__.update(original_state)
+
 
 def get_auto_quantize_config(search_state, constraints=None, verbose=False):
     """Build a flat quant config from auto_quantize search_state.
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index fc07be31205..35e32d16945 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -20,6 +20,7 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 import yaml
 
 from modelopt.recipe import load_recipe
@@ -101,6 +102,16 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
     assert "kv_cache_quant_cfg" not in inputs
 
 
+def test_kv_autoquant_kl_excludes_padding_positions(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    logits = torch.arange(2 * 4 * 3).reshape(2, 4, 3)
+    attention_mask = torch.tensor([[1, 1, 0, 0], [0, 1, 1, 0]])
+
+    selected = hf_ptq._select_unpadded_logits(logits, {"attention_mask": attention_mask})
+
+    assert torch.equal(selected, logits[attention_mask.bool()])
+
+
 def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch):
     """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns
     key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict."""
diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py
index 242a64f762b..2b15ad01c07 100644
--- a/tests/unit/torch/export/test_offload_export.py
+++ b/tests/unit/torch/export/test_offload_export.py
@@ -36,9 +36,12 @@
 )
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.model_config import KV_CACHE_FP8
+from modelopt.torch.export.model_config import KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_NVFP4
 from modelopt.torch.export.model_utils import TiedWeightMap
-from modelopt.torch.export.quant_utils import _postprocess_single_tensor
+from modelopt.torch.export.quant_utils import (
+    _get_kv_cache_postprocess_config,
+    _postprocess_single_tensor,
+)
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
 from modelopt.torch.export.unified_export_hf_streaming import (
     _parse_shard_size,
@@ -319,6 +322,31 @@ def test_postprocess_kv_scale_renamed_and_divided():
     assert abs(val.item() - 0.5) < 1e-5
 
 
+@pytest.mark.parametrize(
+    ("layer_name", "quant_algo"),
+    [
+        ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V),
+        ("model.layers.1.self_attn", KV_CACHE_NVFP4),
+    ],
+)
+def test_postprocess_resolves_mixed_kv_format_per_layer(layer_name, quant_algo):
+    quantization = {
+        "kv_cache_quant_algo": "MIXED_PRECISION",
+        "kv_cache_quantized_layers": {layer_name: {"quant_algo": quant_algo}},
+    }
+    postprocess_config = _get_kv_cache_postprocess_config(quantization)
+
+    key, val = _postprocess_single_tensor(
+        f"{layer_name}.v_bmm_quantizer._amax",
+        torch.tensor(224.0),
+        448.0,
+        postprocess_config,
+    )
+
+    assert key == f"{layer_name}.v_proj.v_scale"
+    assert val.item() == pytest.approx(0.5)
+
+
 def test_postprocess_scale_squeezed():
     """3D scale tensors with shape[0]==1 are squeezed."""
     t = torch.ones(1, 4, 4)
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index acc132fade4..c942eb8786c 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -19,6 +19,7 @@
 from _test_utils.torch.transformers_models import get_tiny_llama
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export.quant_utils import get_quant_config
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.quantization.kv_cache_auto_quant import (
     _candidate_quantizers,
@@ -27,6 +28,7 @@
     _validate_kv_only_config,
     auto_quantize_kv_cache,
 )
+from modelopt.torch.quantization.nn import TensorQuantizer
 
 
 def _kv_config(bits, effective_bits):
@@ -69,6 +71,14 @@ def test_kv_candidate_requires_exact_bits_and_both_sides():
         )
 
 
+@pytest.mark.parametrize("algorithm", ["svdquant", {"method": "smoothquant"}])
+def test_kv_candidate_rejects_structural_or_unscoped_algorithms(algorithm):
+    config = _kv_config((4, 3), 8.0).model_copy(update={"algorithm": algorithm})
+
+    with pytest.raises(ValueError, match="only non-structural calibration algorithms"):
+        _validate_kv_only_config(config)
+
+
 def test_kv_additive_solver_spends_fp8_on_more_sensitive_layer():
     selections, status = _solve_additive_recipe(
         layer_names=["layer0", "layer1"],
@@ -302,6 +312,20 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
         layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled
         for layer in model.model.layers
     )
+    non_kv_quantizers = [
+        quantizer
+        for name, quantizer in model.named_modules()
+        if isinstance(quantizer, TensorQuantizer)
+        and not name.endswith(("k_bmm_quantizer", "v_bmm_quantizer"))
+    ]
+    assert non_kv_quantizers
+    assert all(not quantizer.is_enabled for quantizer in non_kv_quantizers)
+    exported_quantization = get_quant_config(model)["quantization"]
+    assert exported_quantization["quantized_layers"] == {}
+    assert exported_quantization["kv_cache_quantized_layers"]
+    assert set(exported_quantization["kv_cache_quantized_layers"]) <= {
+        f"model.layers.{idx}.self_attn" for idx in range(model.config.num_hidden_layers)
+    }
 
     restored_model = get_tiny_llama(num_hidden_layers=2)
     restored_model, restored_state = mtq.auto_quantize_kv_cache(
@@ -323,6 +347,44 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
     )
 
 
+def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
+    model = get_tiny_llama(num_hidden_layers=1)
+    original_types = {name: type(module) for name, module in model.named_modules()}
+    invalid_candidate = _kv_config((4, 3), 4.5).model_dump()
+    invalid_candidate["algorithm"] = "svdquant"
+
+    with pytest.raises(ValueError, match="only non-structural calibration algorithms"):
+        mtq.auto_quantize_kv_cache(
+            model,
+            {"kv_effective_bits": 4.5},
+            [invalid_candidate],
+            [],
+            lambda *_: pytest.fail("Validation must run before model conversion."),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+    assert not hasattr(model, "_modelopt_state")
+    assert {name: type(module) for name, module in model.named_modules()} == original_types
+
+    valid_candidate = _kv_config((4, 3), 4.5).model_dump()
+    valid_candidate["algorithm"] = None
+    data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
+    with pytest.raises(ValueError, match="non-empty vocabulary dimension"):
+        mtq.auto_quantize_kv_cache(
+            model,
+            {"kv_effective_bits": 4.5},
+            [valid_candidate],
+            data,
+            lambda *_: torch.ones(8),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+    assert not hasattr(model, "_modelopt_state")
+    assert {name: type(module) for name, module in model.named_modules()} == original_types
+
+
 def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
     torch.manual_seed(123)
     model = get_tiny_llama(num_hidden_layers=2)
@@ -337,35 +399,47 @@ def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
         "algorithm": None,
     }
     model = mtq.quantize(model, fixed_kv_config)
-    model.model.layers[0].self_attn.q_proj.weight_quantizer.enable()
-    assert not hasattr(model.model.layers[0].self_attn.q_proj.weight_quantizer, "_amax")
-
-    model, state = mtq.auto_quantize_kv_cache(
-        model,
-        {"kv_effective_bits": 4.0},
-        [
-            (
-                {
-                    "quant_cfg": [
-                        {
-                            "quantizer_name": "*[kv]_bmm_quantizer",
-                            "cfg": {"num_bits": 4},
-                        }
-                    ],
-                    "algorithm": "max",
-                    "effective_bits": 4.0,
-                },
-                "int4",
-            )
-        ],
-        data,
-        lambda search_model, batch: search_model(**batch).logits,
-        num_calib_steps=1,
-        num_score_steps=1,
-        disabled_layers="model.layers.1.self_attn",
+    fixed_weight_quantizer = model.model.layers[0].self_attn.q_proj.weight_quantizer
+    fixed_weight_quantizer.enable()
+    assert not hasattr(fixed_weight_quantizer, "_amax")
+    observed_fixed_states = []
+    hook = fixed_weight_quantizer.register_forward_hook(
+        lambda module, _inputs, _output: observed_fixed_states.append(
+            (module.is_enabled, module._if_quant, module._if_calib)
+        )
     )
 
+    try:
+        model, state = mtq.auto_quantize_kv_cache(
+            model,
+            {"kv_effective_bits": 4.0},
+            [
+                (
+                    {
+                        "quant_cfg": [
+                            {
+                                "quantizer_name": "*[kv]_bmm_quantizer",
+                                "cfg": {"num_bits": 4},
+                            }
+                        ],
+                        "algorithm": "max",
+                        "effective_bits": 4.0,
+                    },
+                    "int4",
+                )
+            ],
+            data,
+            lambda search_model, batch: search_model(**batch).logits,
+            num_calib_steps=1,
+            num_score_steps=1,
+            disabled_layers="model.layers.1.self_attn",
+        )
+    finally:
+        hook.remove()
+
     assert set(state["layers"]) == {"model.layers.0.self_attn"}
+    assert observed_fixed_states
+    assert all(state == (True, True, False) for state in observed_fixed_states)
     assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == 4
     assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled
     assert not hasattr(model.model.layers[0].self_attn.q_proj.weight_quantizer, "_amax")

From b57b4835366351bc5d0770f3d4f12d59579d81f4 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Wed, 19 Aug 2026 12:06:25 -0700
Subject: [PATCH 06/15] fix: harden KV-cache AutoQuant workflows

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     | 10 ++--
 modelopt/torch/export/quant_utils.py          | 38 ++++++--------
 .../torch/quantization/kv_cache_auto_quant.py | 35 +++++++++++--
 .../kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml  | 20 +-------
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  8 ++-
 tests/unit/recipe/test_loader.py              | 12 +----
 .../unit/torch/export/test_offload_export.py  | 35 ++++++++++---
 .../quantization/test_kv_cache_auto_quant.py  | 50 ++++++++++++++++---
 8 files changed, 134 insertions(+), 74 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index b451b104efb..db0538efb3e 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -429,9 +429,9 @@ KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
 `candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
 packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are
 preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice.
-The shipped canary recipe
-searches cast-mode FP8 K/V (8.0 bits/scalar), FP8-K/NVFP4-V (6.25 bits/scalar),
-and packed NVFP4 K/V (4.5 bits/scalar) at 5.4 bits/scalar:
+The shipped canary recipe searches cast-mode FP8 K/V (8.0 bits/scalar) and packed NVFP4 K/V
+(4.5 bits/scalar) at 5.4 bits/scalar. It intentionally excludes FP8-K/NVFP4-V because the
+companion vLLM implementation does not support that asymmetric per-layer format:
 
 ```bash
 python hf_ptq.py \
@@ -451,8 +451,8 @@ state.
 > [vLLM mixed-KV metadata consumer](https://github.com/vllm-project/vllm/pull/52813) or a later
 > vLLM release containing it. The repository's currently pinned vLLM 0.26.0 does not consume
 > `kv_cache_quantized_layers`, so these checkpoints are export-only in that stock environment.
-> Full FP8 K/V and full NVFP4 K/V use existing vLLM kernels; FP8-K/NVFP4-V within one layer also
-> requires the separate mixed-K/V kernel implementation.
+> Do not deploy them with the pinned runtime. Full FP8 K/V and full NVFP4 K/V use existing vLLM
+> kernels once the layer-wise metadata consumer is available.
 
 The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an
 interrupted search (skips re-scoring):
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index ec451e3f2e9..b91ffc4079a 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -1071,17 +1071,12 @@ def _postprocess_single_tensor(
                 layer_quantization = _resolve_kv_cache_format_for_key(key, kv_cache_format)
                 assert layer_quantization in [
                     KV_CACHE_FP8,
-                    KV_CACHE_FP8_K_NVFP4_V,
                     KV_CACHE_NVFP4,
                     KV_CACHE_NVFP4_AFFINE,
                 ], "Invalid KV cache quantization format."
                 assert kv_cache_max_bound > 0, "Maxbound must be greater than zero."
                 value = value.float() / kv_cache_max_bound
-                is_fp8_scale = layer_quantization == KV_CACHE_FP8 or (
-                    layer_quantization == KV_CACHE_FP8_K_NVFP4_V
-                    and key.endswith("k_bmm_quantizer._amax")
-                )
-                if is_fp8_scale and value.item() > 0.5:
+                if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5:
                     logger.warning(
                         "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                     )
@@ -1095,17 +1090,21 @@ def _postprocess_single_tensor(
 def _resolve_kv_cache_format_for_key(
     key: str, quantization: str | dict[str, dict[str, str]] | None
 ) -> str | None:
-    """Resolve uniform or per-layer KV-cache metadata for one state-dict key."""
-    if not isinstance(quantization, dict):
-        return quantization
-    matches = [
-        (layer_name, layer_config.get("quant_algo"))
-        for layer_name, layer_config in quantization.items()
-        if key == layer_name or key.startswith(layer_name + ".")
-    ]
-    if not matches:
+    """Resolve uniform or per-layer metadata to the format for this K or V tensor."""
+    if isinstance(quantization, dict):
+        matches = [
+            (layer_name, layer_config.get("quant_algo"))
+            for layer_name, layer_config in quantization.items()
+            if key == layer_name or key.startswith(layer_name + ".")
+        ]
+        quantization = max(matches, key=lambda item: len(item[0]))[1] if matches else None
+    if quantization == KV_CACHE_FP8_K_NVFP4_V:
+        if key.endswith("k_bmm_quantizer._amax"):
+            return KV_CACHE_FP8
+        if key.endswith("v_bmm_quantizer._amax"):
+            return KV_CACHE_NVFP4
         return None
-    return max(matches, key=lambda item: len(item[0]))[1]
+    return quantization
 
 
 def _get_kv_cache_postprocess_config(
@@ -1172,7 +1171,6 @@ def _export_key(key: str) -> str:
                     layer_quantization = _resolve_kv_cache_format_for_key(key, quantization)
                     assert layer_quantization in [
                         KV_CACHE_FP8,
-                        KV_CACHE_FP8_K_NVFP4_V,
                         KV_CACHE_NVFP4,
                         KV_CACHE_NVFP4_AFFINE,
                     ], "Invalid KV cache quantization format."
@@ -1181,11 +1179,7 @@ def _export_key(key: str) -> str:
                     value = value.float() / maxbound
 
                     # Warn if scale exceeds threshold
-                    is_fp8_scale = layer_quantization == KV_CACHE_FP8 or (
-                        layer_quantization == KV_CACHE_FP8_K_NVFP4_V
-                        and key.endswith("k_bmm_quantizer._amax")
-                    )
-                    if is_fp8_scale and value.item() > 0.5:
+                    if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5:
                         logger.warning(
                             "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                         )
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index cb2dc81a563..19bc7f6b165 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -19,6 +19,7 @@
 
 import fnmatch
 import os
+from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -198,6 +199,29 @@ def _apply_layer_quantizers(module: nn.Module, quantizers: dict[str, TensorQuant
         setattr(module, attr, quantizer)
 
 
+@contextmanager
+def _freeze_existing_quantizers(model: nn.Module, candidate_quantizers: list[TensorQuantizer]):
+    """Run enabled non-candidate quantizers as QDQ without updating their calibration state."""
+    candidate_ids = {id(quantizer) for quantizer in candidate_quantizers}
+    states = []
+    for module in model.modules():
+        if (
+            not isinstance(module, TensorQuantizer)
+            or id(module) in candidate_ids
+            or not module.is_enabled
+        ):
+            continue
+        states.append((module, module._if_quant, module._if_calib))
+        module.enable_quant()
+        module.disable_calib()
+    try:
+        yield
+    finally:
+        for quantizer, if_quant, if_calib in states:
+            quantizer._if_quant = if_quant
+            quantizer._if_calib = if_calib
+
+
 def _get_logits(
     forward_step: Callable[[nn.Module, Any], torch.Tensor], model: nn.Module, data: Any
 ) -> torch.Tensor:
@@ -422,11 +446,12 @@ def calibration_loop(calibration_model):
                     ]
                     calibration_proxy = nn.Module()
                     calibration_proxy.quantizers = nn.ModuleList(active_quantizers)
-                    calibrate(
-                        calibration_proxy,
-                        algorithm=config.algorithm,
-                        forward_loop=lambda _: calibration_loop(model),
-                    )
+                    with _freeze_existing_quantizers(model, active_quantizers):
+                        calibrate(
+                            calibration_proxy,
+                            algorithm=config.algorithm,
+                            forward_loop=lambda _: calibration_loop(model),
+                        )
 
                 for layer_name, module, _ in layers:
                     _apply_layer_quantizers(module, disabled_quantizers[layer_name])
diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
index 943a8f0e8c0..44aa9b29f5d 100644
--- a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
+++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -1,8 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Layer-wise KV-cache search over full FP8, FP8-K/NVFP4-V, and full NVFP4 at
-# 5.4 bits/scalar.
+# Layer-wise KV-cache search over full FP8 and full NVFP4 at 5.4 bits/scalar.
 # Cast mode fixes amax and skips PTQ calibration; isolated full-vocabulary forward
 # KL is measured with every other eligible attention layer kept in BF16.
 
@@ -12,12 +11,10 @@ imports:
   base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers
   kv_fp8_cast: configs/ptq/units/kv_fp8_cast
   kv_nvfp4_cast: configs/ptq/units/kv_nvfp4_cast
-  fp8: configs/numerics/fp8
-  nvfp4: configs/numerics/nvfp4
 
 metadata:
   recipe_type: auto_quantize
-  description: Layer-wise three-format KV-cache cast search at 5.4 bits using forward KL.
+  description: Layer-wise FP8/NVFP4 KV-cache cast search at 5.4 bits using forward KL.
 
 auto_quantize:
   constraints:
@@ -30,19 +27,6 @@ auto_quantize:
         method: max
         skip_forward_without_activation_calib: true
       effective_bits: 8.0
-    - quant_cfg:
-        - quantizer_name: '*k_bmm_quantizer'
-          cfg:
-            $import: fp8
-            use_constant_amax: true
-        - quantizer_name: '*v_bmm_quantizer'
-          cfg:
-            $import: nvfp4
-            use_constant_amax: true
-      algorithm:
-        method: max
-        skip_forward_without_activation_calib: true
-      effective_bits: 6.25
     - quant_cfg:
         - $import: kv_nvfp4_cast
       algorithm:
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 35e32d16945..296c6daaa66 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -96,7 +96,6 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
     assert inputs["method"] == "kl_div"
     assert [config["effective_bits"] for config, _ in inputs["quantization_formats"]] == [
         8.0,
-        6.25,
         4.5,
     ]
     assert "kv_cache_quant_cfg" not in inputs
@@ -112,6 +111,13 @@ def test_kv_autoquant_kl_excludes_padding_positions(monkeypatch):
     assert torch.equal(selected, logits[attention_mask.bool()])
 
 
+def test_kv_autoquant_kl_rejects_misaligned_attention_mask(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+
+    with pytest.raises(ValueError, match="matching token dimensions"):
+        hf_ptq._select_unpadded_logits(torch.zeros(2, 4, 3), {"attention_mask": torch.ones(2, 3)})
+
+
 def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch):
     """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns
     key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict."""
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index d51222b1f3b..ec2e143246b 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1933,18 +1933,10 @@ def test_load_recipe_kv_autoquantize_contract():
     assert aq.constraints.effective_bits is None
     assert aq.constraints.kv_effective_bits == 5.4
     assert aq.auto_quantize_method == "kl_div"
-    assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 6.25, 4.5]
-    assert [entry.quantizer_name for entry in aq.candidate_formats[1].quant_cfg] == [
-        "*k_bmm_quantizer",
-        "*v_bmm_quantizer",
-    ]
+    assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5]
     for fmt in aq.candidate_formats:
         for entry in fmt.quant_cfg:
-            assert entry.quantizer_name in {
-                "*[kv]_bmm_quantizer",
-                "*k_bmm_quantizer",
-                "*v_bmm_quantizer",
-            }
+            assert entry.quantizer_name == "*[kv]_bmm_quantizer"
             assert entry.cfg.use_constant_amax
         assert fmt.algorithm["skip_forward_without_activation_calib"]
 
diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py
index 2b15ad01c07..09551c6bf60 100644
--- a/tests/unit/torch/export/test_offload_export.py
+++ b/tests/unit/torch/export/test_offload_export.py
@@ -41,6 +41,7 @@
 from modelopt.torch.export.quant_utils import (
     _get_kv_cache_postprocess_config,
     _postprocess_single_tensor,
+    _resolve_kv_cache_format_for_key,
 )
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
 from modelopt.torch.export.unified_export_hf_streaming import (
@@ -323,27 +324,49 @@ def test_postprocess_kv_scale_renamed_and_divided():
 
 
 @pytest.mark.parametrize(
-    ("layer_name", "quant_algo"),
+    ("layer_name", "quant_algo", "side", "resolved_format"),
     [
-        ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V),
-        ("model.layers.1.self_attn", KV_CACHE_NVFP4),
+        ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V, "k", KV_CACHE_FP8),
+        ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V, "v", KV_CACHE_NVFP4),
+        ("model.layers.1.self_attn", KV_CACHE_NVFP4, "k", KV_CACHE_NVFP4),
     ],
 )
-def test_postprocess_resolves_mixed_kv_format_per_layer(layer_name, quant_algo):
+def test_postprocess_resolves_mixed_kv_format_per_layer_and_side(
+    layer_name, quant_algo, side, resolved_format
+):
     quantization = {
         "kv_cache_quant_algo": "MIXED_PRECISION",
         "kv_cache_quantized_layers": {layer_name: {"quant_algo": quant_algo}},
     }
     postprocess_config = _get_kv_cache_postprocess_config(quantization)
+    original_key = f"{layer_name}.{side}_bmm_quantizer._amax"
+
+    assert _resolve_kv_cache_format_for_key(original_key, postprocess_config) == resolved_format
 
     key, val = _postprocess_single_tensor(
-        f"{layer_name}.v_bmm_quantizer._amax",
+        original_key,
         torch.tensor(224.0),
         448.0,
         postprocess_config,
     )
 
-    assert key == f"{layer_name}.v_proj.v_scale"
+    assert key == f"{layer_name}.{side}_proj.{side}_scale"
+    assert val.item() == pytest.approx(0.5)
+
+
+@pytest.mark.parametrize(("side", "resolved_format"), [("k", KV_CACHE_FP8), ("v", KV_CACHE_NVFP4)])
+def test_postprocess_resolves_uniform_asymmetric_kv_format(side, resolved_format):
+    original_key = f"model.layers.0.self_attn.{side}_bmm_quantizer._amax"
+
+    assert _resolve_kv_cache_format_for_key(original_key, KV_CACHE_FP8_K_NVFP4_V) == resolved_format
+    key, val = _postprocess_single_tensor(
+        original_key,
+        torch.tensor(224.0),
+        448.0,
+        KV_CACHE_FP8_K_NVFP4_V,
+    )
+
+    assert key == f"model.layers.0.self_attn.{side}_proj.{side}_scale"
     assert val.item() == pytest.approx(0.5)
 
 
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index c942eb8786c..bbd418c24e3 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -20,6 +20,7 @@
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.quant_utils import get_quant_config
+from modelopt.torch.quantization import model_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.quantization.kv_cache_auto_quant import (
     _candidate_quantizers,
@@ -71,7 +72,7 @@ def test_kv_candidate_requires_exact_bits_and_both_sides():
         )
 
 
-@pytest.mark.parametrize("algorithm", ["svdquant", {"method": "smoothquant"}])
+@pytest.mark.parametrize("algorithm", ["svdquant", {"method": "smoothquant"}, {"method": "mse"}])
 def test_kv_candidate_rejects_structural_or_unscoped_algorithms(algorithm):
     config = _kv_config((4, 3), 8.0).model_copy(update={"algorithm": algorithm})
 
@@ -385,7 +386,7 @@ def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
     assert {name: type(module) for name, module in model.named_modules()} == original_types
 
 
-def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
+def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers(monkeypatch):
     torch.manual_seed(123)
     model = get_tiny_llama(num_hidden_layers=2)
     data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
@@ -401,13 +402,38 @@ def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
     model = mtq.quantize(model, fixed_kv_config)
     fixed_weight_quantizer = model.model.layers[0].self_attn.q_proj.weight_quantizer
     fixed_weight_quantizer.enable()
-    assert not hasattr(fixed_weight_quantizer, "_amax")
+    fixed_weight_quantizer.amax = torch.tensor(1.0)
+    fixed_weight_quantizer.disable_quant()
+    fixed_weight_quantizer.disable_calib()
     observed_fixed_states = []
-    hook = fixed_weight_quantizer.register_forward_hook(
+    fixed_hook = fixed_weight_quantizer.register_forward_hook(
         lambda module, _inputs, _output: observed_fixed_states.append(
             (module.is_enabled, module._if_quant, module._if_calib)
         )
     )
+    fixed_qdq_quantizer = model.model.layers[1].self_attn.q_proj.weight_quantizer
+    fixed_qdq_quantizer.enable()
+    fixed_qdq_quantizer.amax = torch.tensor(1.0)
+    observed_qdq_states = []
+    qdq_hook = fixed_qdq_quantizer.register_forward_hook(
+        lambda module, _inputs, _output: observed_qdq_states.append(
+            (module.is_enabled, module._if_quant, module._if_calib)
+        )
+    )
+    calibration_states = []
+    real_calibrate = model_quant.calibrate
+
+    def calibrate_with_state_check(*args, **kwargs):
+        calibration_states.append(
+            (fixed_weight_quantizer._if_quant, fixed_weight_quantizer._if_calib)
+        )
+        result = real_calibrate(*args, **kwargs)
+        calibration_states.append(
+            (fixed_weight_quantizer._if_quant, fixed_weight_quantizer._if_calib)
+        )
+        return result
+
+    monkeypatch.setattr(model_quant, "calibrate", calibrate_with_state_check)
 
     try:
         model, state = mtq.auto_quantize_kv_cache(
@@ -435,14 +461,24 @@ def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers():
             disabled_layers="model.layers.1.self_attn",
         )
     finally:
-        hook.remove()
+        fixed_hook.remove()
+        qdq_hook.remove()
 
     assert set(state["layers"]) == {"model.layers.0.self_attn"}
     assert observed_fixed_states
-    assert all(state == (True, True, False) for state in observed_fixed_states)
+    assert observed_fixed_states[0] == (True, True, False)
+    assert all(state == (True, False, False) for state in observed_fixed_states[1:])
+    assert calibration_states == [(True, False), (True, False)]
     assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == 4
     assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled
-    assert not hasattr(model.model.layers[0].self_attn.q_proj.weight_quantizer, "_amax")
+    assert not fixed_weight_quantizer._if_quant
+    assert not fixed_weight_quantizer._if_calib
+    assert fixed_weight_quantizer.amax.item() == pytest.approx(1.0)
+    assert observed_qdq_states
+    assert all(quantizer_state == (True, True, False) for quantizer_state in observed_qdq_states)
+    assert fixed_qdq_quantizer._if_quant
+    assert not fixed_qdq_quantizer._if_calib
+    assert fixed_qdq_quantizer.amax.item() == pytest.approx(1.0)
     fixed_attention = model.model.layers[1].self_attn
     assert fixed_attention.k_bmm_quantizer.is_enabled
     assert fixed_attention.v_bmm_quantizer.is_enabled

From 722ab98e8920038c8aeb83c80196025d139310bd Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Wed, 19 Aug 2026 16:03:18 -0700
Subject: [PATCH 07/15] fix: guard KV-cache AutoQuant execution

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |  2 +
 examples/hf_ptq/hf_ptq.py                     | 12 +++--
 .../torch/quantization/kv_cache_auto_quant.py | 48 +++++++++----------
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 25 ++++++++++
 .../quantization/test_kv_cache_auto_quant.py  | 25 ++++++----
 5 files changed, 74 insertions(+), 38 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index db0538efb3e..b2b3d1ecf7d 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -489,6 +489,8 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop)
 
 ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point.
 
+> *AutoQuantize recipes are not supported with `--use_fsdp2` and are rejected before model loading. Distributed sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Use a PTQ recipe with FSDP2.*
+
 ### Usage
 
 #### Slurm (recommended)
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 0aca1ffde5a..8a52a78040b 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -96,6 +96,10 @@
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
 RAND_SEED = 1234
+_FSDP2_AUTOQUANT_ERROR = (
+    "AutoQuantize does not support --use_fsdp2 until distributed sensitivity scoring, "
+    "selection, and checkpoint writes are synchronized across ranks."
+)
 
 
 def _select_unpadded_logits(logits: torch.Tensor, batch: dict[str, Any]) -> torch.Tensor:
@@ -456,11 +460,7 @@ def auto_quantize(
     )
 
     if args.use_fsdp2:
-        warnings.warn(
-            "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet "
-            "(distributed calibration, sensitivity scoring, and recipe/checkpoint "
-            "synchronization across ranks); use at your own risk."
-        )
+        raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR)
 
     inputs = _mtq_inputs_from_auto_quantize_config(
         aq_config, args, fixed_quantize_config=fixed_quantize_config
@@ -573,6 +573,8 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool:
 def load_model(args: argparse.Namespace):
     # If low memory mode is enabled, we compress the model while loading the HF checkpoint.
     calibration_only = False
+    if args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
+        raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR)
     if args.use_fsdp2:
         hf_config = AutoConfig.from_pretrained(
             args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 19bc7f6b165..700d743c4b0 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -49,22 +49,25 @@ def _disabled_quantizer() -> TensorQuantizer:
     return quantizer
 
 
-def _candidate_quantizers(module: nn.Module, config: QuantizeConfig) -> dict[str, TensorQuantizer]:
-    original = {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS}
-    try:
-        for attr in _KV_QUANTIZER_ATTRS:
-            setattr(module, attr, _disabled_quantizer())
-        set_quantizer_by_cfg(module, config.quant_cfg)
-        quantizers = {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS}
-        for attr, quantizer in quantizers.items():
-            if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled:
-                raise ValueError(
-                    f"KV-cache candidate must enable {attr}; got {type(quantizer).__name__}."
-                )
-        return quantizers
-    finally:
-        for attr, quantizer in original.items():
-            setattr(module, attr, quantizer)
+def _candidate_quantizers(config: QuantizeConfig) -> dict[str, TensorQuantizer]:
+    """Build a candidate on a holder that exposes only the K/V quantizers.
+
+    Applying an untrusted wildcard config to an attention module can match and mutate
+    unrelated nested quantizers (for example ``q_proj.*_quantizer``). Keeping candidate
+    construction isolated makes such entries no-ops and guarantees that the search can
+    only retain K/V quantizer state.
+    """
+    holder = nn.Module()
+    for attr in _KV_QUANTIZER_ATTRS:
+        setattr(holder, attr, _disabled_quantizer())
+    set_quantizer_by_cfg(holder, config.quant_cfg)
+    quantizers = {attr: getattr(holder, attr) for attr in _KV_QUANTIZER_ATTRS}
+    for attr, quantizer in quantizers.items():
+        if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled:
+            raise ValueError(
+                f"KV-cache candidate must enable {attr}; got {type(quantizer).__name__}."
+            )
+    return quantizers
 
 
 def _validate_kv_only_config(config: QuantizeConfig) -> None:
@@ -201,7 +204,7 @@ def _apply_layer_quantizers(module: nn.Module, quantizers: dict[str, TensorQuant
 
 @contextmanager
 def _freeze_existing_quantizers(model: nn.Module, candidate_quantizers: list[TensorQuantizer]):
-    """Run enabled non-candidate quantizers as QDQ without updating their calibration state."""
+    """Freeze calibration without changing existing quantizers' execution mode."""
     candidate_ids = {id(quantizer) for quantizer in candidate_quantizers}
     states = []
     for module in model.modules():
@@ -211,14 +214,12 @@ def _freeze_existing_quantizers(model: nn.Module, candidate_quantizers: list[Ten
             or not module.is_enabled
         ):
             continue
-        states.append((module, module._if_quant, module._if_calib))
-        module.enable_quant()
+        states.append((module, module._if_calib))
         module.disable_calib()
     try:
         yield
     finally:
-        for quantizer, if_quant, if_calib in states:
-            quantizer._if_quant = if_quant
+        for quantizer, if_calib in states:
             quantizer._if_calib = if_calib
 
 
@@ -385,10 +386,9 @@ def auto_quantize_kv_cache(
     }
     candidate_quantizers = {
         name: {
-            candidate_name: _candidate_quantizers(module, config)
-            for candidate_name, config in candidates
+            candidate_name: _candidate_quantizers(config) for candidate_name, config in candidates
         }
-        for name, module, _ in layers
+        for name, _, _ in layers
     }
 
     is_training = model.training
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 296c6daaa66..5abdfbbf77a 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -118,6 +118,31 @@ def test_kv_autoquant_kl_rejects_misaligned_attention_mask(monkeypatch):
         hf_ptq._select_unpadded_logits(torch.zeros(2, 4, 3), {"attention_mask": torch.ones(2, 3)})
 
 
+def test_autoquant_rejects_fsdp2(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=True,
+    )
+
+    with pytest.raises(NotImplementedError, match="does not support --use_fsdp2"):
+        hf_ptq.auto_quantize(args, torch.nn.Module(), [], SimpleNamespace())
+
+
+def test_fsdp2_autoquant_rejected_before_model_load(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(hf_ptq, "_recipe_is_auto_quantize", lambda _: True)
+    monkeypatch.setattr(
+        hf_ptq.AutoConfig,
+        "from_pretrained",
+        lambda *_args, **_kwargs: pytest.fail("The model config must not be loaded."),
+    )
+
+    with pytest.raises(NotImplementedError, match="does not support --use_fsdp2"):
+        hf_ptq.load_model(SimpleNamespace(use_fsdp2=True, recipe="autoquant"))
+
+
 def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch):
     """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns
     key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict."""
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index bbd418c24e3..7a92af73c27 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -104,9 +104,6 @@ def test_kv_scalar_weight_counts_k_and_v_widths():
 
 
 def test_kv_candidate_format_can_use_dynamic_amax():
-    module = nn.Module()
-    module.k_bmm_quantizer = nn.Identity()
-    module.v_bmm_quantizer = nn.Identity()
     config = QuantizeConfig(
         quant_cfg=[
             {
@@ -118,7 +115,7 @@ def test_kv_candidate_format_can_use_dynamic_amax():
         effective_bits=8.0,
     )
 
-    quantizers = _candidate_quantizers(module, config)
+    quantizers = _candidate_quantizers(config)
 
     assert all(quantizer.is_enabled for quantizer in quantizers.values())
     assert all(not quantizer._use_constant_amax for quantizer in quantizers.values())
@@ -274,7 +271,11 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
                     {
                         "quantizer_name": "*[kv]_bmm_quantizer",
                         "cfg": {"num_bits": 8},
-                    }
+                    },
+                    {
+                        "quantizer_name": "q_proj.*_quantizer",
+                        "cfg": {"num_bits": 2},
+                    },
                 ],
                 "algorithm": "max",
                 "effective_bits": 8.0,
@@ -321,6 +322,9 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
     ]
     assert non_kv_quantizers
     assert all(not quantizer.is_enabled for quantizer in non_kv_quantizers)
+    assert all(
+        layer.self_attn.q_proj.weight_quantizer.num_bits == 8 for layer in model.model.layers
+    )
     exported_quantization = get_quant_config(model)["quantization"]
     assert exported_quantization["quantized_layers"] == {}
     assert exported_quantization["kv_cache_quantized_layers"]
@@ -446,7 +450,11 @@ def calibrate_with_state_check(*args, **kwargs):
                             {
                                 "quantizer_name": "*[kv]_bmm_quantizer",
                                 "cfg": {"num_bits": 4},
-                            }
+                            },
+                            {
+                                "quantizer_name": "q_proj.*_quantizer",
+                                "cfg": {"num_bits": 2},
+                            },
                         ],
                         "algorithm": "max",
                         "effective_bits": 4.0,
@@ -466,9 +474,8 @@ def calibrate_with_state_check(*args, **kwargs):
 
     assert set(state["layers"]) == {"model.layers.0.self_attn"}
     assert observed_fixed_states
-    assert observed_fixed_states[0] == (True, True, False)
-    assert all(state == (True, False, False) for state in observed_fixed_states[1:])
-    assert calibration_states == [(True, False), (True, False)]
+    assert all(state == (True, False, False) for state in observed_fixed_states)
+    assert calibration_states == [(False, False), (False, False)]
     assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == 4
     assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled
     assert not fixed_weight_quantizer._if_quant

From faa1f95438b38d7467baa7382595f6deb730e934 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Wed, 19 Aug 2026 22:41:25 -0700
Subject: [PATCH 08/15] fix: validate KV-cache AutoQuant candidates

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |  10 +-
 .../torch/quantization/kv_cache_auto_quant.py | 141 +++++---
 modelopt/torch/quantization/model_quant.py    |  12 +-
 .../kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml  |  24 +-
 tests/unit/recipe/test_loader.py              |   5 +-
 .../quantization/test_kv_cache_auto_quant.py  | 304 ++++++++++++++----
 6 files changed, 368 insertions(+), 128 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index b2b3d1ecf7d..d77dbd520f4 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -429,7 +429,7 @@ KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
 `candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
 packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are
 preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice.
-The shipped canary recipe searches cast-mode FP8 K/V (8.0 bits/scalar) and packed NVFP4 K/V
+The shipped canary recipe searches calibrated FP8 K/V (8.0 bits/scalar) and packed NVFP4 K/V
 (4.5 bits/scalar) at 5.4 bits/scalar. It intentionally excludes FP8-K/NVFP4-V because the
 companion vLLM implementation does not support that asymmetric per-layer format:
 
@@ -441,10 +441,10 @@ python hf_ptq.py \
   --export_path /path/to/qwen3-1.7b-mixed-kv
 ```
 
-Cast candidates use constant amax and skip the PTQ calibration forward. Unified HF export records
-the selected formats in `kv_cache_quantized_layers` and writes the JSON-safe sensitivity report to
-`kv_cache_auto_quantize_report.json`; `--auto_quantize_checkpoint` stores the resumable raw search
-state.
+Each candidate uses max calibration so its persistent K/V scales are present in the unified HF
+checkpoint. Unified export records the selected formats in `kv_cache_quantized_layers` and writes
+the JSON-safe sensitivity report to `kv_cache_auto_quantize_report.json`;
+`--auto_quantize_checkpoint` stores the resumable raw search state.
 
 > [!NOTE]
 > Layer-wise KV checkpoints require the companion
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 700d743c4b0..3e0700abda3 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import fnmatch
+import math
 import os
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any
@@ -41,6 +42,20 @@
 
 _KV_QUANTIZER_ATTRS = ("k_bmm_quantizer", "v_bmm_quantizer")
 _KV_AUTOQUANT_SCHEMA_VERSION = 1
+_KV_CANDIDATE_HOLDER_NAME = "layer"
+_KV_CANDIDATE_NAMES = {f"{_KV_CANDIDATE_HOLDER_NAME}.{attr}" for attr in _KV_QUANTIZER_ATTRS}
+_NON_KV_PROBE_NAMES = {
+    f"{_KV_CANDIDATE_HOLDER_NAME}.{name}"
+    for name in (
+        "q_bmm_quantizer",
+        "p_bmm_quantizer",
+        "input_quantizer",
+        "output_quantizer",
+        "q_proj.input_quantizer",
+        "q_proj.weight_quantizer",
+        "q_proj.output_quantizer",
+    )
+}
 
 
 def _disabled_quantizer() -> TensorQuantizer:
@@ -50,17 +65,14 @@ def _disabled_quantizer() -> TensorQuantizer:
 
 
 def _candidate_quantizers(config: QuantizeConfig) -> dict[str, TensorQuantizer]:
-    """Build a candidate on a holder that exposes only the K/V quantizers.
-
-    Applying an untrusted wildcard config to an attention module can match and mutate
-    unrelated nested quantizers (for example ``q_proj.*_quantizer``). Keeping candidate
-    construction isolated makes such entries no-ops and guarantees that the search can
-    only retain K/V quantizer state.
-    """
+    """Build a candidate with the same qualified K/V names used by a full model."""
+    _validate_candidate_patterns(config)
+    root = nn.Module()
     holder = nn.Module()
+    root.add_module(_KV_CANDIDATE_HOLDER_NAME, holder)
     for attr in _KV_QUANTIZER_ATTRS:
         setattr(holder, attr, _disabled_quantizer())
-    set_quantizer_by_cfg(holder, config.quant_cfg)
+    set_quantizer_by_cfg(root, config.quant_cfg)
     quantizers = {attr: getattr(holder, attr) for attr in _KV_QUANTIZER_ATTRS}
     for attr, quantizer in quantizers.items():
         if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled:
@@ -70,52 +82,103 @@ def _candidate_quantizers(config: QuantizeConfig) -> dict[str, TensorQuantizer]:
     return quantizers
 
 
-def _validate_kv_only_config(config: QuantizeConfig) -> None:
-    if config.effective_bits is None:
-        raise ValueError(
-            "Each KV-cache AutoQuant candidate must declare config-level effective_bits."
-        )
-    allowed_names = set(_KV_QUANTIZER_ATTRS)
+def _validate_candidate_patterns(config: QuantizeConfig) -> None:
+    """Require every ordered config entry to match only a qualified K/V name."""
     matched_names: set[str] = set()
-    probe_names = {
-        *allowed_names,
-        "q_bmm_quantizer",
-        "p_bmm_quantizer",
-        "input_quantizer",
-        "weight_quantizer",
-        "output_quantizer",
-    }
+    probe_names = _KV_CANDIDATE_NAMES | _NON_KV_PROBE_NAMES
     for entry in config.quant_cfg:
-        pattern = entry.quantizer_name
-        matches = {name for name in probe_names if fnmatch.fnmatch(name, pattern)}
-        if matches - allowed_names:
+        if entry.parent_class is not None:
+            raise ValueError("KV-cache AutoQuant candidates do not support parent_class filters.")
+        matches = {name for name in probe_names if fnmatch.fnmatch(name, entry.quantizer_name)}
+        if not matches:
+            raise ValueError(
+                "KV-cache AutoQuant candidate pattern "
+                f"{entry.quantizer_name!r} does not match a supported qualified K/V quantizer."
+            )
+        non_kv_matches = matches - _KV_CANDIDATE_NAMES
+        if non_kv_matches:
             raise ValueError(
                 "KV-cache AutoQuant candidates may configure only k_bmm_quantizer and "
-                f"v_bmm_quantizer; pattern {pattern!r} also matches {sorted(matches - allowed_names)}."
+                f"v_bmm_quantizer; pattern {entry.quantizer_name!r} also matches "
+                f"{sorted(non_kv_matches)}."
             )
         matched_names.update(matches)
-    if matched_names != allowed_names:
+    if matched_names != _KV_CANDIDATE_NAMES:
         raise ValueError(
             "KV-cache AutoQuant candidates must completely configure both "
             "k_bmm_quantizer and v_bmm_quantizer."
         )
 
+
+def _algorithm_method(config: QuantizeConfig) -> str | None:
     algorithm = config.algorithm
-    if algorithm is None:
-        return
-    if isinstance(algorithm, str):
-        algorithm_method = algorithm
-    elif isinstance(algorithm, dict):
-        algorithm_method = algorithm.get("method")
-    else:
-        algorithm_method = getattr(algorithm, "method", None)
-    if algorithm_method != "max":
+    if algorithm is None or isinstance(algorithm, str):
+        return algorithm
+    if isinstance(algorithm, dict):
+        return algorithm.get("method")
+    return getattr(algorithm, "method", None)
+
+
+def _deployable_kv_bits(quantizer: TensorQuantizer) -> float:
+    """Return storage bits for the narrow K/V formats supported by unified export."""
+    if quantizer.bias is not None:
+        raise ValueError("KV-cache AutoQuant does not support affine candidates yet.")
+    if quantizer.is_fp8:
+        return 8.0
+    if quantizer.is_nvfp4_dynamic and quantizer.block_sizes.get(-1) == 16:
+        return 4.5
+    raise ValueError(
+        "KV-cache AutoQuant candidates must use unified-export-compatible per-tensor FP8 "
+        "or block-16 dynamic NVFP4 quantizers."
+    )
+
+
+def _validate_deployable_candidate(config: QuantizeConfig) -> None:
+    quantizers = _candidate_quantizers(config)
+    k_quantizer = quantizers["k_bmm_quantizer"]
+    v_quantizer = quantizers["v_bmm_quantizer"]
+    k_bits = _deployable_kv_bits(k_quantizer)
+    v_bits = _deployable_kv_bits(v_quantizer)
+    if k_bits != v_bits and (k_bits, v_bits) != (8.0, 4.5):
+        raise ValueError(
+            "Unified export supports only uniform FP8, uniform NVFP4, or FP8-K/NVFP4-V "
+            "KV-cache AutoQuant candidates."
+        )
+
+    algorithm_method = _algorithm_method(config)
+    for attr, quantizer in quantizers.items():
+        will_calibrate = algorithm_method == "max" and not quantizer._use_constant_amax
+        if not hasattr(quantizer, "_amax") and not will_calibrate:
+            raise ValueError(
+                f"KV-cache AutoQuant candidate {attr} has no persistent export scale. "
+                "Use max calibration or constant_amax; dynamic and use_constant_amax-only "
+                "candidates cannot be exported."
+            )
+
+    assert config.effective_bits is not None
+    actual_effective_bits = (k_bits + v_bits) / 2.0
+    if not math.isclose(config.effective_bits, actual_effective_bits, rel_tol=0.0, abs_tol=1e-12):
         raise ValueError(
-            "KV-cache AutoQuant supports only non-structural calibration algorithms "
-            f"None and 'max'; got {algorithm_method!r}."
+            "KV-cache AutoQuant candidate effective_bits does not match its configured K/V "
+            f"storage cost: declared {config.effective_bits}, actual {actual_effective_bits}."
         )
 
 
+def _validate_kv_only_config(config: QuantizeConfig) -> None:
+    if config.effective_bits is None:
+        raise ValueError(
+            "Each KV-cache AutoQuant candidate must declare config-level effective_bits."
+        )
+    algorithm_method = _algorithm_method(config)
+    if algorithm_method != "max":
+        if algorithm_method is not None:
+            raise ValueError(
+                "KV-cache AutoQuant supports only non-structural calibration algorithms "
+                f"None and 'max'; got {algorithm_method!r}."
+            )
+    _validate_deployable_candidate(config)
+
+
 def _validate_search_inputs(
     constraints: dict[str, Any],
     quantization_formats: list[tuple[dict[str, Any], str]],
@@ -357,7 +420,7 @@ def auto_quantize_kv_cache(
     configure K and V together and declare ``effective_bits`` matching its packed
     storage per K-or-V scalar, including scale overhead. Candidate calibration is scoped
     to the candidate K/V quantizers, while pre-existing fixed quantizers keep executing
-    with frozen state. Cast-style constant-amax formats may skip calibration forwards.
+    with frozen state. Persistent ``constant_amax`` formats may skip calibration forwards.
     """
     target_bits, candidates = _validate_search_inputs(
         constraints, quantization_formats, num_calib_steps, num_score_steps
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 14cfb990f0e..8fc9aa37800 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -680,7 +680,7 @@ def auto_quantize_kv_cache(
     together and declare config-level ``effective_bits`` equal to their packed
     storage cost per K-or-V scalar, including scale overhead. The budget is weighted
     by the K/V projection widths of eligible layers. Formats use their own calibration
-    algorithm; cast-mode constant-amax candidates skip calibration forwards.
+    algorithm; candidates with persistent ``constant_amax`` skip calibration forwards.
 
     Args:
         model: Model whose attention K/V quantizers will be searched.
@@ -702,6 +702,16 @@ def auto_quantize_kv_cache(
         The converted model with the selected per-layer K/V quantizers and a
         JSON-safe sensitivity report.
     """
+    if (
+        torch.distributed.is_available()
+        and torch.distributed.is_initialized()
+        and torch.distributed.get_world_size() > 1
+    ):
+        raise RuntimeError(
+            "auto_quantize_kv_cache is single-process only; distributed scoring, selection, "
+            "and checkpoint writes are not synchronized."
+        )
+
     processed_formats = []
     for idx, candidate in enumerate(quantization_formats):
         if isinstance(candidate, tuple):
diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
index 44aa9b29f5d..836273eee53 100644
--- a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
+++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -1,20 +1,20 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Layer-wise KV-cache search over full FP8 and full NVFP4 at 5.4 bits/scalar.
-# Cast mode fixes amax and skips PTQ calibration; isolated full-vocabulary forward
-# KL is measured with every other eligible attention layer kept in BF16.
+# Layer-wise KV-cache search over calibrated FP8 and NVFP4 at 5.4 bits/scalar.
+# Isolated full-vocabulary forward KL is measured with every other eligible
+# attention layer kept in BF16.
 
 # modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
 imports:
   base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
   base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers
-  kv_fp8_cast: configs/ptq/units/kv_fp8_cast
-  kv_nvfp4_cast: configs/ptq/units/kv_nvfp4_cast
+  kv_fp8: configs/ptq/units/kv_fp8
+  kv_nvfp4: configs/ptq/units/kv_nvfp4
 
 metadata:
   recipe_type: auto_quantize
-  description: Layer-wise FP8/NVFP4 KV-cache cast search at 5.4 bits using forward KL.
+  description: Layer-wise calibrated FP8/NVFP4 KV-cache search at 5.4 bits using forward KL.
 
 auto_quantize:
   constraints:
@@ -22,16 +22,12 @@ auto_quantize:
 
   candidate_formats:
     - quant_cfg:
-        - $import: kv_fp8_cast
-      algorithm:
-        method: max
-        skip_forward_without_activation_calib: true
+        - $import: kv_fp8
+      algorithm: max
       effective_bits: 8.0
     - quant_cfg:
-        - $import: kv_nvfp4_cast
-      algorithm:
-        method: max
-        skip_forward_without_activation_calib: true
+        - $import: kv_nvfp4
+      algorithm: max
       effective_bits: 4.5
 
   auto_quantize_method: kl_div
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index ec2e143246b..e4dc4b3f422 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1937,8 +1937,9 @@ def test_load_recipe_kv_autoquantize_contract():
     for fmt in aq.candidate_formats:
         for entry in fmt.quant_cfg:
             assert entry.quantizer_name == "*[kv]_bmm_quantizer"
-            assert entry.cfg.use_constant_amax
-        assert fmt.algorithm["skip_forward_without_activation_calib"]
+            assert not entry.cfg.use_constant_amax
+            assert entry.cfg.constant_amax is None
+        assert fmt.algorithm == "max"
 
 
 def _all_shipped_ptq_recipe_paths():
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index 7a92af73c27..e415873bc99 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -19,8 +19,8 @@
 from _test_utils.torch.transformers_models import get_tiny_llama
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.quant_utils import get_quant_config
-from modelopt.torch.quantization import model_quant
+from modelopt.torch.export.quant_utils import get_kv_cache_dtype, get_quant_config
+from modelopt.torch.quantization import model_quant, tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.quantization.kv_cache_auto_quant import (
     _candidate_quantizers,
@@ -32,14 +32,35 @@
 from modelopt.torch.quantization.nn import TensorQuantizer
 
 
-def _kv_config(bits, effective_bits):
+@pytest.fixture
+def nvfp4_fake_quant_stub(monkeypatch):
+    """Keep CPU search tests independent of the CUDA-only NVFP4 fake-quant kernel."""
+
+    monkeypatch.setattr(
+        tensor_quant,
+        "dynamic_block_quantize_op",
+        lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs),
+    )
+
+
+def _quantizer_cfg(bits, *, constant_amax=None):
+    cfg = {"num_bits": bits}
+    if bits == (2, 1):
+        cfg["block_sizes"] = {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}
+    if constant_amax is not None:
+        cfg["constant_amax"] = constant_amax
+    return cfg
+
+
+def _kv_config(bits, effective_bits, *, algorithm="max", constant_amax=None):
     return QuantizeConfig(
         quant_cfg=[
             {
                 "quantizer_name": "*[kv]_bmm_quantizer",
-                "cfg": {"num_bits": bits, "use_constant_amax": True},
+                "cfg": _quantizer_cfg(bits, constant_amax=constant_amax),
             }
         ],
+        algorithm=algorithm,
         effective_bits=effective_bits,
     )
 
@@ -103,22 +124,50 @@ def test_kv_scalar_weight_counts_k_and_v_widths():
     assert _kv_scalar_weight(module, "attention") == 40
 
 
-def test_kv_candidate_format_can_use_dynamic_amax():
-    config = QuantizeConfig(
-        quant_cfg=[
-            {
-                "quantizer_name": "*[kv]_bmm_quantizer",
-                "cfg": {"num_bits": (4, 3)},
-            }
-        ],
-        algorithm=None,
-        effective_bits=8.0,
-    )
-
+@pytest.mark.parametrize(
+    ("config", "expected_format"),
+    [
+        (_kv_config((4, 3), 8.0), "FP8"),
+        (_kv_config((2, 1), 4.5), "NVFP4"),
+        (_kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0), "FP8"),
+        (_kv_config((2, 1), 4.5, algorithm=None, constant_amax=1.0), "NVFP4"),
+    ],
+)
+def test_kv_candidate_accepts_export_supported_persistent_formats(config, expected_format):
+    _validate_kv_only_config(config)
     quantizers = _candidate_quantizers(config)
+    module = nn.Module()
+    for name, quantizer in quantizers.items():
+        setattr(module, name, quantizer)
 
-    assert all(quantizer.is_enabled for quantizer in quantizers.values())
-    assert all(not quantizer._use_constant_amax for quantizer in quantizers.values())
+    assert get_kv_cache_dtype(module) == expected_format
+
+
+@pytest.mark.parametrize(
+    ("config", "match"),
+    [
+        (_kv_config((4, 3), 8.0, algorithm=None), "no persistent export scale"),
+        (
+            QuantizeConfig(
+                quant_cfg=[
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": (4, 3), "use_constant_amax": True},
+                    }
+                ],
+                algorithm="max",
+                effective_bits=8.0,
+            ),
+            "no persistent export scale",
+        ),
+        (_kv_config(8, 8.0, algorithm=None, constant_amax=1.0), "per-tensor FP8"),
+        (_kv_config(4, 4.0, algorithm=None, constant_amax=1.0), "per-tensor FP8"),
+        (_kv_config((4, 3), 6.0), "does not match its configured K/V storage cost"),
+    ],
+)
+def test_kv_candidate_rejects_non_exportable_or_incorrect_cost(config, match):
+    with pytest.raises(ValueError, match=match):
+        _validate_kv_only_config(config)
 
 
 class _ToyKVAttention(nn.Module):
@@ -145,7 +194,7 @@ def forward(self, x):
         return self.lm_head(self.attn1(self.attn0(x)))
 
 
-def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
+def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fake_quant_stub):
     torch.manual_seed(123)
     model = _ToyKVModel()
     data = [torch.randn(2, 3, 8), torch.randn(2, 3, 8)]
@@ -155,32 +204,40 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
                 "quant_cfg": [
                     {
                         "quantizer_name": "*[kv]_bmm_quantizer",
-                        "cfg": {"num_bits": 8, "constant_amax": 1.0},
+                        "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
                     }
                 ],
                 "algorithm": None,
                 "effective_bits": 8.0,
             },
-            "int8",
+            "fp8",
         ),
         (
             {
                 "quant_cfg": [
                     {
                         "quantizer_name": "*[kv]_bmm_quantizer",
-                        "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                        "cfg": {
+                            "num_bits": (2, 1),
+                            "block_sizes": {
+                                -1: 16,
+                                "type": "dynamic",
+                                "scale_bits": (4, 3),
+                            },
+                            "constant_amax": 1.0,
+                        },
                     }
                 ],
                 "algorithm": None,
-                "effective_bits": 4.0,
+                "effective_bits": 4.5,
             },
-            "int4",
+            "nvfp4",
         ),
     ]
 
     model, state = auto_quantize_kv_cache(
         model,
-        {"kv_effective_bits": 6.0},
+        {"kv_effective_bits": 6.25},
         candidates,
         data,
         lambda model, batch: model(batch),
@@ -189,23 +246,23 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
         checkpoint=str(tmp_path / "kv_search.pth"),
     )
 
-    assert state["best"]["effective_bits"] == pytest.approx(6.0)
+    assert state["best"]["effective_bits"] == pytest.approx(6.25)
     assert state["best"]["is_satisfied"]
     assert model.training
     assert {layer["selected"] for layer in state["layers"].values()} == {
-        "int8",
-        "int4",
+        "fp8",
+        "nvfp4",
     }
     for layer_name, layer_state in state["layers"].items():
         layer = model.get_submodule(layer_name)
         assert layer.k_bmm_quantizer.num_bits == layer.v_bmm_quantizer.num_bits
-        expected_bits = 8 if layer_state["selected"] == "int8" else 4
+        expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1)
         assert layer.k_bmm_quantizer.num_bits == expected_bits
 
     restored_model = _ToyKVModel().eval()
     restored_model, restored_state = auto_quantize_kv_cache(
         restored_model,
-        {"kv_effective_bits": 6.0},
+        {"kv_effective_bits": 6.25},
         candidates,
         data,
         lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."),
@@ -218,10 +275,58 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path):
     assert not restored_model.training
     for layer_name, layer_state in restored_state["layers"].items():
         layer = restored_model.get_submodule(layer_name)
-        expected_bits = 8 if layer_state["selected"] == "int8" else 4
+        expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1)
         assert layer.k_bmm_quantizer.num_bits == expected_bits
 
 
+def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quant_stub):
+    model = _ToyKVModel()
+    candidate = (
+        {
+            "quant_cfg": [
+                {
+                    "quantizer_name": "*[kv]_bmm_quantizer",
+                    "cfg": {
+                        "num_bits": (2, 1),
+                        "block_sizes": {
+                            -1: 16,
+                            "type": "dynamic",
+                            "scale_bits": (4, 3),
+                        },
+                        "constant_amax": 1.0,
+                    },
+                },
+                {
+                    "quantizer_name": "*.k_bmm_quantizer",
+                    "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                },
+            ],
+            "algorithm": None,
+            "effective_bits": 6.25,
+        },
+        "fp8_k_nvfp4_v",
+    )
+
+    model, state = auto_quantize_kv_cache(
+        model,
+        {"kv_effective_bits": 6.25},
+        [candidate],
+        [torch.randn(1, 2, 8)],
+        lambda search_model, batch: search_model(batch),
+        num_calib_steps=1,
+        num_score_steps=1,
+    )
+
+    assert state["candidates"][0]["effective_bits"] == pytest.approx(6.25)
+    assert state["best"]["effective_bits"] == pytest.approx(6.25)
+    for layer_state in state["layers"].values():
+        assert layer_state["selected"] == "fp8_k_nvfp4_v"
+    for layer in (model.attn0, model.attn1):
+        assert layer.k_bmm_quantizer.num_bits == (4, 3)
+        assert layer.v_bmm_quantizer.num_bits == (2, 1)
+    assert get_quant_config(model)["quantization"]["kv_cache_quant_algo"] == "FP8_K_NVFP4_V"
+
+
 def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state():
     model = _ToyKVModel()
     original_quantizers = {
@@ -234,20 +339,28 @@ def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state():
                 "quant_cfg": [
                     {
                         "quantizer_name": "*[kv]_bmm_quantizer",
-                        "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                        "cfg": {
+                            "num_bits": (2, 1),
+                            "block_sizes": {
+                                -1: 16,
+                                "type": "dynamic",
+                                "scale_bits": (4, 3),
+                            },
+                            "constant_amax": 1.0,
+                        },
                     }
                 ],
                 "algorithm": None,
-                "effective_bits": 4.0,
+                "effective_bits": 4.5,
             },
-            "int4",
+            "nvfp4",
         )
     ]
 
     with pytest.raises(ValueError, match="non-empty vocabulary dimension"):
         auto_quantize_kv_cache(
             model,
-            {"kv_effective_bits": 4.0},
+            {"kv_effective_bits": 4.5},
             candidates,
             [torch.randn(2, 3, 8)],
             lambda *_: torch.ones(8),
@@ -260,7 +373,7 @@ def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state():
         assert (module.k_bmm_quantizer, module.v_bmm_quantizer) == original_quantizers[name]
 
 
-def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
+def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_fake_quant_stub):
     torch.manual_seed(123)
     model = get_tiny_llama(num_hidden_layers=2)
     data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))} for _ in range(2)]
@@ -270,36 +383,40 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
                 "quant_cfg": [
                     {
                         "quantizer_name": "*[kv]_bmm_quantizer",
-                        "cfg": {"num_bits": 8},
-                    },
-                    {
-                        "quantizer_name": "q_proj.*_quantizer",
-                        "cfg": {"num_bits": 2},
+                        "cfg": {"num_bits": (4, 3)},
                     },
                 ],
                 "algorithm": "max",
                 "effective_bits": 8.0,
             },
-            "int8",
+            "fp8",
         ),
         (
             {
                 "quant_cfg": [
                     {
                         "quantizer_name": "*[kv]_bmm_quantizer",
-                        "cfg": {"num_bits": 4, "constant_amax": 1.0},
+                        "cfg": {
+                            "num_bits": (2, 1),
+                            "block_sizes": {
+                                -1: 16,
+                                "type": "dynamic",
+                                "scale_bits": (4, 3),
+                            },
+                            "constant_amax": 1.0,
+                        },
                     }
                 ],
                 "algorithm": None,
-                "effective_bits": 4.0,
+                "effective_bits": 4.5,
             },
-            "int4",
+            "nvfp4",
         ),
     ]
 
     model, state = mtq.auto_quantize_kv_cache(
         model,
-        {"kv_effective_bits": 6.0},
+        {"kv_effective_bits": 6.25},
         candidates,
         data,
         lambda search_model, batch: search_model(**batch).logits,
@@ -309,7 +426,7 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
     )
 
     assert len(state["layers"]) == model.config.num_hidden_layers
-    assert state["best"]["effective_bits"] == pytest.approx(6.0)
+    assert state["best"]["effective_bits"] == pytest.approx(6.25)
     assert all(
         layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled
         for layer in model.model.layers
@@ -335,7 +452,7 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
     restored_model = get_tiny_llama(num_hidden_layers=2)
     restored_model, restored_state = mtq.auto_quantize_kv_cache(
         restored_model,
-        {"kv_effective_bits": 6.0},
+        {"kv_effective_bits": 6.25},
         candidates,
         data,
         lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."),
@@ -348,20 +465,20 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path):
     assert any(
         hasattr(layer.self_attn.k_bmm_quantizer, "_amax")
         for layer in restored_model.model.layers
-        if layer.self_attn.k_bmm_quantizer.num_bits == 8
+        if layer.self_attn.k_bmm_quantizer.num_bits == (4, 3)
     )
 
 
 def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
     model = get_tiny_llama(num_hidden_layers=1)
     original_types = {name: type(module) for name, module in model.named_modules()}
-    invalid_candidate = _kv_config((4, 3), 4.5).model_dump()
+    invalid_candidate = _kv_config((4, 3), 8.0).model_dump()
     invalid_candidate["algorithm"] = "svdquant"
 
     with pytest.raises(ValueError, match="only non-structural calibration algorithms"):
         mtq.auto_quantize_kv_cache(
             model,
-            {"kv_effective_bits": 4.5},
+            {"kv_effective_bits": 8.0},
             [invalid_candidate],
             [],
             lambda *_: pytest.fail("Validation must run before model conversion."),
@@ -372,13 +489,12 @@ def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
     assert not hasattr(model, "_modelopt_state")
     assert {name: type(module) for name, module in model.named_modules()} == original_types
 
-    valid_candidate = _kv_config((4, 3), 4.5).model_dump()
-    valid_candidate["algorithm"] = None
+    valid_candidate = _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump()
     data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
     with pytest.raises(ValueError, match="non-empty vocabulary dimension"):
         mtq.auto_quantize_kv_cache(
             model,
-            {"kv_effective_bits": 4.5},
+            {"kv_effective_bits": 8.0},
             [valid_candidate],
             data,
             lambda *_: torch.ones(8),
@@ -390,7 +506,58 @@ def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
     assert {name: type(module) for name, module in model.named_modules()} == original_types
 
 
-def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers(monkeypatch):
+def test_public_kv_autoquant_rejects_unmatched_or_unexportable_candidates_before_conversion():
+    model = get_tiny_llama(num_hidden_layers=1)
+    original_types = {name: type(module) for name, module in model.named_modules()}
+    unmatched = _kv_config((4, 3), 8.0).model_dump()
+    unmatched["quant_cfg"].append({"quantizer_name": "q_proj.*_quantizer", "cfg": {"num_bits": 2}})
+    invalid_candidates = [
+        (unmatched, "does not match a supported qualified K/V quantizer"),
+        (_kv_config((4, 3), 8.0, algorithm=None).model_dump(), "no persistent export scale"),
+        (
+            _kv_config(8, 8.0, algorithm=None, constant_amax=1.0).model_dump(),
+            "per-tensor FP8",
+        ),
+    ]
+
+    for candidate, match in invalid_candidates:
+        with pytest.raises(ValueError, match=match):
+            mtq.auto_quantize_kv_cache(
+                model,
+                {"kv_effective_bits": 8.0},
+                [candidate],
+                [],
+                lambda *_: pytest.fail("Validation must run before model conversion."),
+                num_calib_steps=1,
+                num_score_steps=1,
+            )
+        assert not hasattr(model, "_modelopt_state")
+        assert {name: type(module) for name, module in model.named_modules()} == original_types
+
+
+def test_public_kv_autoquant_rejects_distributed_execution_before_mutation(monkeypatch):
+    model = get_tiny_llama(num_hidden_layers=1)
+    original_types = {name: type(module) for name, module in model.named_modules()}
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
+
+    with pytest.raises(RuntimeError, match="single-process only"):
+        mtq.auto_quantize_kv_cache(
+            model,
+            {"kv_effective_bits": 8.0},
+            [],
+            [],
+            lambda *_: pytest.fail("Distributed validation must fail before search."),
+        )
+
+    assert not hasattr(model, "_modelopt_state")
+    assert {name: type(module) for name, module in model.named_modules()} == original_types
+
+
+def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers(
+    monkeypatch, nvfp4_fake_quant_stub
+):
     torch.manual_seed(123)
     model = get_tiny_llama(num_hidden_layers=2)
     data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
@@ -398,7 +565,7 @@ def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers(monkey
         "quant_cfg": [
             {
                 "quantizer_name": "*[kv]_bmm_quantizer",
-                "cfg": {"num_bits": 8, "constant_amax": 1.0},
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
             }
         ],
         "algorithm": None,
@@ -442,24 +609,27 @@ def calibrate_with_state_check(*args, **kwargs):
     try:
         model, state = mtq.auto_quantize_kv_cache(
             model,
-            {"kv_effective_bits": 4.0},
+            {"kv_effective_bits": 4.5},
             [
                 (
                     {
                         "quant_cfg": [
                             {
                                 "quantizer_name": "*[kv]_bmm_quantizer",
-                                "cfg": {"num_bits": 4},
-                            },
-                            {
-                                "quantizer_name": "q_proj.*_quantizer",
-                                "cfg": {"num_bits": 2},
+                                "cfg": {
+                                    "num_bits": (2, 1),
+                                    "block_sizes": {
+                                        -1: 16,
+                                        "type": "dynamic",
+                                        "scale_bits": (4, 3),
+                                    },
+                                },
                             },
                         ],
                         "algorithm": "max",
-                        "effective_bits": 4.0,
+                        "effective_bits": 4.5,
                     },
-                    "int4",
+                    "nvfp4",
                 )
             ],
             data,
@@ -476,7 +646,7 @@ def calibrate_with_state_check(*args, **kwargs):
     assert observed_fixed_states
     assert all(state == (True, False, False) for state in observed_fixed_states)
     assert calibration_states == [(False, False), (False, False)]
-    assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == 4
+    assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == (2, 1)
     assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled
     assert not fixed_weight_quantizer._if_quant
     assert not fixed_weight_quantizer._if_calib
@@ -489,5 +659,5 @@ def calibrate_with_state_check(*args, **kwargs):
     fixed_attention = model.model.layers[1].self_attn
     assert fixed_attention.k_bmm_quantizer.is_enabled
     assert fixed_attention.v_bmm_quantizer.is_enabled
-    assert fixed_attention.k_bmm_quantizer.num_bits == 8
-    assert fixed_attention.v_bmm_quantizer.num_bits == 8
+    assert fixed_attention.k_bmm_quantizer.num_bits == (4, 3)
+    assert fixed_attention.v_bmm_quantizer.num_bits == (4, 3)

From b907f1cc1764def5aceddf03c11c314aa9e39e24 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Sun, 23 Aug 2026 15:46:30 -0700
Subject: [PATCH 09/15] Fix KV AutoQuant architecture edge cases

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/hf_ptq.py                     | 12 ++++
 modelopt/torch/export/model_utils.py          | 15 +++--
 .../torch/quantization/kv_cache_auto_quant.py | 15 ++++-
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 54 +++++++++++++++
 .../torch/export/test_unified_export_hf.py    | 27 +++++++-
 .../quantization/test_kv_cache_auto_quant.py  | 67 ++++++++++++++++++-
 6 files changed, 181 insertions(+), 9 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 8a52a78040b..6daade6ed24 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -362,6 +362,18 @@ def _mtq_kv_candidate_formats(formats) -> list[tuple[dict, str]]:
         quant_cfg = fmt.model_dump(exclude_none=True)
         candidate_quantizers = quant_cfg.get("quant_cfg", [])
         name = None
+        nvfp4_quantizers = type(fmt)(**KV_QUANT_CFG_CHOICES["nvfp4"]).model_dump(exclude_none=True)[
+            "quant_cfg"
+        ]
+        fp8_quantizers = type(fmt)(**KV_QUANT_CFG_CHOICES["fp8"]).model_dump(exclude_none=True)[
+            "quant_cfg"
+        ]
+        if len(fp8_quantizers) != 1:
+            raise RuntimeError("The FP8 KV preset must contain exactly one quantizer entry.")
+        fp8_k_quantizer = copy.deepcopy(fp8_quantizers[0])
+        fp8_k_quantizer["quantizer_name"] = "*.k_bmm_quantizer"
+        if candidate_quantizers == [*nvfp4_quantizers, fp8_k_quantizer]:
+            name = "fp8_k_nvfp4_v"
         for preset_name, preset in KV_QUANT_CFG_CHOICES.items():
             normalized_preset_quantizers = (
                 type(fmt)(**preset).model_dump(exclude_none=True).get("quant_cfg", [])
diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py
index f27405ec83f..cf89ddc2e5b 100755
--- a/modelopt/torch/export/model_utils.py
+++ b/modelopt/torch/export/model_utils.py
@@ -108,7 +108,7 @@ def is_multimodal_model(model):
     config = model.config
 
     # Check for Nemotron-Parse encoder-decoder architecture
-    architectures = getattr(config, "architectures", [])
+    architectures = getattr(config, "architectures", None) or []
     is_nemotron_parse = any("nemotronparse" in arch.lower() for arch in architectures)
 
     return (
@@ -137,12 +137,17 @@ def get_language_model_from_vl(model) -> list[nn.Module] | None:
         >>> # lineage[0] is vlm_model
         >>> # lineage[1] is vlm_model.language_model
     """
-    # always prioritize model.model.langauge_model
+    candidates = []
     if hasattr(model, "model") and hasattr(model.model, "language_model"):
-        return [model, model.model, model.model.language_model]
-
+        candidates.append([model, model.model, model.model.language_model])
     if hasattr(model, "language_model"):
-        return [model, model.language_model]
+        candidates.append([model, model.language_model])
+    if len(candidates) > 1:
+        raise ValueError(
+            "Found multiple language-model roots; refusing to select one by traversal order."
+        )
+    if candidates:
+        return candidates[0]
 
     # Pattern 3: For encoder-decoder VL models (e.g., Nemotron-Parse), the decoder is the language model.
     # Only match if the model is detected as multimodal to avoid matching non-VLM encoder-decoder
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 3e0700abda3..4d7192df1c5 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -248,10 +248,21 @@ def _eligible_layers(
     model: nn.Module, disabled_layers: list[str] | str | None
 ) -> list[tuple[str, nn.Module, int]]:
     patterns = [disabled_layers] if isinstance(disabled_layers, str) else disabled_layers or []
-    layers = []
-    for name, module in model.named_modules():
+    boundaries = []
+    names_by_identity: dict[int, list[str]] = {}
+    for name, module in model.named_modules(remove_duplicate=False):
         if not all(hasattr(module, attr) for attr in _KV_QUANTIZER_ATTRS):
             continue
+        boundaries.append((name, module))
+        names_by_identity.setdefault(id(module), []).append(name)
+    aliases = [names for names in names_by_identity.values() if len(names) > 1]
+    if aliases:
+        raise ValueError(
+            f"KV-cache attention boundaries are registered through aliases: {aliases}."
+        )
+
+    layers = []
+    for name, module in boundaries:
         if any(fnmatch.fnmatch(name, pattern) for pattern in patterns):
             continue
         layers.append((name, module, _kv_scalar_weight(module, name)))
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 5abdfbbf77a..4762931bb81 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import getpass
 import importlib
 import sys
@@ -22,6 +23,7 @@
 import pytest
 import torch
 import yaml
+from _test_utils.torch.transformers_models import get_tiny_qwen3
 
 from modelopt.recipe import load_recipe
 from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints
@@ -101,6 +103,58 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
     assert "kv_cache_quant_cfg" not in inputs
 
 
+def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch):
+    """The HF entry point runs the real public KV AutoQuant path on an offline Qwen fixture."""
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    model = get_tiny_qwen3(num_hidden_layers=1)
+    aq = AutoQuantizeConfig(
+        constraints=AutoQuantizeConstraints(kv_effective_bits=8.0),
+        candidate_formats=[
+            QuantizeConfig(
+                quant_cfg=[
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                    }
+                ],
+                algorithm=None,
+                effective_bits=8.0,
+            )
+        ],
+        auto_quantize_method="kl_div",
+        score_size=1,
+    )
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=False,
+        kv_cache_qformat="none",
+        batch_size=1,
+        auto_quantize_checkpoint=None,
+    )
+    data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
+
+    hf_ptq.auto_quantize(args, model, data, aq, full_model=model)
+
+    attention = model.model.layers[0].self_attn
+    assert attention.k_bmm_quantizer.num_bits == (4, 3)
+    assert attention.v_bmm_quantizer.num_bits == (4, 3)
+
+
+def test_kv_autoquant_names_asymmetric_export_format(monkeypatch):
+    """The supported FP8-K/NVFP4-V candidate has a stable semantic name."""
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    mixed_config = copy.deepcopy(hf_ptq.KV_QUANT_CFG_CHOICES["nvfp4"])
+    fp8_k_quantizer = copy.deepcopy(hf_ptq.KV_QUANT_CFG_CHOICES["fp8"]["quant_cfg"][0])
+    fp8_k_quantizer["quantizer_name"] = "*.k_bmm_quantizer"
+    mixed_config["quant_cfg"].append(fp8_k_quantizer)
+    mixed_config["effective_bits"] = 6.25
+
+    candidates = hf_ptq._mtq_kv_candidate_formats([QuantizeConfig(**mixed_config)])
+
+    assert candidates[0][1] == "fp8_k_nvfp4_v"
+
+
 def test_kv_autoquant_kl_excludes_padding_positions(monkeypatch):
     hf_ptq = _import_hf_ptq(monkeypatch)
     logits = torch.arange(2 * 4 * 3).reshape(2, 4, 3)
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index b9fa29238d7..ca7147d61da 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -15,6 +15,8 @@
 
 """Tests for tied-weight helpers in unified_export_hf."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 from _test_utils.torch.quantization.tied_modules import (
@@ -23,7 +25,11 @@
 )
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.model_utils import TiedWeightMap
+from modelopt.torch.export.model_utils import (
+    TiedWeightMap,
+    get_language_model_from_vl,
+    is_multimodal_model,
+)
 from modelopt.torch.export.quant_utils import (
     fuse_prequant_layernorm,
     postprocess_state_dict,
@@ -32,6 +38,25 @@
 from modelopt.torch.quantization.nn import TensorQuantizer
 
 
+def test_multimodal_detection_accepts_null_architectures():
+    """Unified export treats absent architecture metadata as an empty list."""
+    model = SimpleNamespace(config=SimpleNamespace(architectures=None))
+
+    assert not is_multimodal_model(model)
+
+
+@pytest.mark.parametrize("aliased", [False, True])
+def test_language_model_extraction_rejects_competing_or_aliased_roots(aliased):
+    """Language-model extraction must not select ambiguous roots by traversal order."""
+    model = torch.nn.Module()
+    model.model = torch.nn.Module()
+    model.model.language_model = torch.nn.Module()
+    model.language_model = model.model.language_model if aliased else torch.nn.Module()
+
+    with pytest.raises(ValueError, match="multiple language-model roots"):
+        get_language_model_from_vl(model)
+
+
 def test_hf_all_tied_weights_keys_contract():
     """Pin the transformers API we build tied_map from, so a version bump fails loud here.
 
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index e415873bc99..8fe4852835d 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -13,10 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+
 import pytest
 import torch
 import torch.nn as nn
-from _test_utils.torch.transformers_models import get_tiny_llama
+from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3, get_tiny_qwen3vl
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.quant_utils import get_kv_cache_dtype, get_quant_config
@@ -24,6 +26,7 @@
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.quantization.kv_cache_auto_quant import (
     _candidate_quantizers,
+    _eligible_layers,
     _kv_scalar_weight,
     _solve_additive_recipe,
     _validate_kv_only_config,
@@ -194,6 +197,29 @@ def forward(self, x):
         return self.lm_head(self.attn1(self.attn0(x)))
 
 
+def test_kv_eligible_layers_supports_hybrid_attention_mixers_only():
+    """Hybrid decoders include attention mixers but exclude nonattention mixers."""
+    model = nn.Module()
+    model.layers = nn.ModuleList([nn.Module(), nn.Module()])
+    model.layers[0].mixer = nn.Linear(8, 8, bias=False)
+    model.layers[1].mixer = _ToyKVAttention(8, gain=1.0)
+
+    layers = _eligible_layers(model, disabled_layers=None)
+
+    assert [(name, width) for name, _, width in layers] == [("layers.1.mixer", 16)]
+
+
+def test_kv_eligible_layers_rejects_aliased_attention_boundary():
+    """An attention object registered at multiple paths must not be selected by traversal order."""
+    model = nn.Module()
+    attention = _ToyKVAttention(8, gain=1.0)
+    model.attention = attention
+    model.attention_alias = attention
+
+    with pytest.raises(ValueError, match="registered through aliases"):
+        _eligible_layers(model, disabled_layers=None)
+
+
 def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fake_quant_stub):
     torch.manual_seed(123)
     model = _ToyKVModel()
@@ -469,6 +495,45 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_
     )
 
 
+@pytest.mark.parametrize(
+    ("model_factory", "expected_layer", "disabled_layers"),
+    [
+        (get_tiny_qwen3, "model.layers.0.self_attn", None),
+        (get_tiny_qwen3vl, "model.language_model.layers.0.self_attn", "*visual*"),
+    ],
+)
+def test_public_kv_autoquant_selects_qwen_causal_attention_only(
+    model_factory, expected_layer, disabled_layers
+):
+    """Plain and conditional Qwen models expose only causal attention to the KV search."""
+    model = model_factory(num_hidden_layers=1)
+    text_config = getattr(model.config, "text_config", model.config)
+    data = [{"input_ids": torch.randint(0, text_config.vocab_size, (1, 8))}]
+    candidate = (
+        _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump(),
+        "fp8",
+    )
+
+    model, state = mtq.auto_quantize_kv_cache(
+        model,
+        {"kv_effective_bits": 8.0},
+        [candidate],
+        data,
+        lambda search_model, batch: search_model(**batch).logits,
+        num_calib_steps=1,
+        num_score_steps=1,
+        disabled_layers=disabled_layers,
+    )
+
+    assert set(state["layers"]) == {expected_layer}
+    assert json.loads(json.dumps(state))["layers"][expected_layer]["selected"] == "fp8"
+    exported = get_quant_config(model)["quantization"]
+    if "kv_cache_quantized_layers" in exported:
+        assert set(exported["kv_cache_quantized_layers"]) == {expected_layer}
+    else:
+        assert exported["kv_cache_quant_algo"] == "FP8"
+
+
 def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
     model = get_tiny_llama(num_hidden_layers=1)
     original_types = {name: type(module) for name, module in model.named_modules()}

From 34a249597ab9b886fa0ffc2709e4332ab261455d Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 25 Aug 2026 11:31:58 -0700
Subject: [PATCH 10/15] Fix KV AutoQuant cost validation

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/recipe/config.py                     |  5 ++
 .../torch/quantization/kv_cache_auto_quant.py | 30 +++++++++
 .../kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml  |  5 +-
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  2 +
 tests/unit/recipe/test_loader.py              | 27 ++++++--
 .../quantization/test_kv_cache_auto_quant.py  | 62 +++++++++++--------
 6 files changed, 98 insertions(+), 33 deletions(-)

diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index 664953a137c..b6a1edb5153 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -321,6 +321,11 @@ def _has_search_space(self):
                 raise ValueError(
                     "KV-cache AutoQuant candidate_formats replace the uniform kv_cache post-step."
                 )
+            if self.cost_excluded_layers:
+                raise ValueError(
+                    "KV-cache AutoQuant does not support cost_excluded_layers; use "
+                    "disabled_layers to exclude non-KV-cache modules from the search."
+                )
         return self
 
 
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 4d7192df1c5..1c9ff64dbf7 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -133,6 +133,14 @@ def _deployable_kv_bits(quantizer: TensorQuantizer) -> float:
     )
 
 
+def _candidate_kv_bits(config: QuantizeConfig) -> tuple[float, float]:
+    quantizers = _candidate_quantizers(config)
+    return (
+        _deployable_kv_bits(quantizers["k_bmm_quantizer"]),
+        _deployable_kv_bits(quantizers["v_bmm_quantizer"]),
+    )
+
+
 def _validate_deployable_candidate(config: QuantizeConfig) -> None:
     quantizers = _candidate_quantizers(config)
     k_quantizer = quantizers["k_bmm_quantizer"]
@@ -244,6 +252,27 @@ def _kv_scalar_weight(module: nn.Module, name: str) -> int:
     return k_width + v_width
 
 
+def _validate_candidate_cost_geometry(
+    candidates: list[tuple[str, QuantizeConfig]],
+    layers: list[tuple[str, nn.Module, int]],
+) -> None:
+    candidate_bits = [_candidate_kv_bits(config) for _, config in candidates]
+    if all(k_bits == v_bits for k_bits, v_bits in candidate_bits):
+        return
+
+    unequal_width_layers = []
+    for name, module, _ in layers:
+        k_width = _projection_width(module, "k")
+        v_width = _projection_width(module, "v")
+        if k_width != v_width:
+            unequal_width_layers.append(f"{name} (K={k_width}, V={v_width})")
+    if unequal_width_layers:
+        raise ValueError(
+            "KV-cache AutoQuant cannot cost asymmetric K/V candidates on layers with unequal "
+            "K/V widths: " + ", ".join(unequal_width_layers) + "."
+        )
+
+
 def _eligible_layers(
     model: nn.Module, disabled_layers: list[str] | str | None
 ) -> list[tuple[str, nn.Module, int]]:
@@ -438,6 +467,7 @@ def auto_quantize_kv_cache(
     )
 
     layers = _eligible_layers(model, disabled_layers)
+    _validate_candidate_cost_geometry(candidates, layers)
     signature = _search_signature(
         candidates,
         layers,
diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
index 836273eee53..db05db9cc70 100644
--- a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
+++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -8,7 +8,6 @@
 # modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
 imports:
   base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
-  base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers
   kv_fp8: configs/ptq/units/kv_fp8
   kv_nvfp4: configs/ptq/units/kv_nvfp4
 
@@ -35,6 +34,4 @@ auto_quantize:
 
   disabled_layers:
     - $import: base_disabled_layers
-
-  cost_excluded_layers:
-    - $import: base_cost_excluded_layers
+    - "*mtp*"
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 4762931bb81..2471bab4189 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -100,6 +100,8 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
         8.0,
         4.5,
     ]
+    assert aq.cost_excluded_layers == []
+    assert "*mtp*" in inputs["disabled_layers"]
     assert "kv_cache_quant_cfg" not in inputs
 
 
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index e4dc4b3f422..a6ec40cf656 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -28,6 +28,8 @@
 
 import modelopt.torch.quantization.config as qcfg
 from modelopt.recipe.config import (
+    AutoQuantizeConfig,
+    AutoQuantizeConstraints,
     ModelOptAutoQuantizeRecipe,
     ModelOptDFlashRecipe,
     ModelOptEagleRecipe,
@@ -1919,11 +1921,16 @@ def test_load_recipe_autoquantize_builtin_general(recipe_path):
     assert isinstance(recipe, ModelOptAutoQuantizeRecipe)
     assert len(recipe.auto_quantize.candidate_formats) >= 2
     assert recipe.auto_quantize.auto_quantize_method in ("gradient", "kl_div")
-    # Both shared base units must be spliced in: the removed --auto_quantize_* CLI shim appended
-    # them unconditionally, so a general recipe is the migration target and must match it. Without
-    # cost_excluded_layers a VL/MTP model counts its vision tower in the effective-bits denominator.
     assert "*output_layer*" in recipe.auto_quantize.disabled_layers
-    assert recipe.auto_quantize.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"]
+    if recipe.auto_quantize.constraints.kv_effective_bits is not None:
+        assert "*mtp*" in recipe.auto_quantize.disabled_layers
+        assert recipe.auto_quantize.cost_excluded_layers == []
+    else:
+        assert recipe.auto_quantize.cost_excluded_layers == [
+            "*visual*",
+            "*mtp*",
+            "*vision_tower*",
+        ]
 
 
 def test_load_recipe_kv_autoquantize_contract():
@@ -1933,6 +1940,8 @@ def test_load_recipe_kv_autoquantize_contract():
     assert aq.constraints.effective_bits is None
     assert aq.constraints.kv_effective_bits == 5.4
     assert aq.auto_quantize_method == "kl_div"
+    assert "*mtp*" in aq.disabled_layers
+    assert aq.cost_excluded_layers == []
     assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5]
     for fmt in aq.candidate_formats:
         for entry in fmt.quant_cfg:
@@ -1942,6 +1951,16 @@ def test_load_recipe_kv_autoquantize_contract():
         assert fmt.algorithm == "max"
 
 
+def test_kv_autoquantize_rejects_cost_excluded_layers():
+    with pytest.raises(ValueError, match=r"cost_excluded_layers.*disabled_layers"):
+        AutoQuantizeConfig(
+            constraints=AutoQuantizeConstraints(kv_effective_bits=8.0),
+            candidate_formats=[qcfg.QuantizeConfig(quant_cfg=[], effective_bits=8.0)],
+            auto_quantize_method="kl_div",
+            cost_excluded_layers=["*mtp*"],
+        )
+
+
 def _all_shipped_ptq_recipe_paths():
     """Every shipped PTQ recipe, discovered from disk rather than a hardcoded list."""
     root = files("modelopt_recipes")
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index 8fe4852835d..b0eadefd30a 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -68,6 +68,23 @@ def _kv_config(bits, effective_bits, *, algorithm="max", constant_amax=None):
     )
 
 
+def _asymmetric_kv_config():
+    return QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": _quantizer_cfg((2, 1), constant_amax=1.0),
+            },
+            {
+                "quantizer_name": "*.k_bmm_quantizer",
+                "cfg": _quantizer_cfg((4, 3), constant_amax=1.0),
+            },
+        ],
+        algorithm=None,
+        effective_bits=6.25,
+    )
+
+
 def test_kv_candidate_requires_exact_bits_and_both_sides():
     _validate_kv_only_config(_kv_config((4, 3), 8.0))
 
@@ -307,31 +324,7 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fa
 
 def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quant_stub):
     model = _ToyKVModel()
-    candidate = (
-        {
-            "quant_cfg": [
-                {
-                    "quantizer_name": "*[kv]_bmm_quantizer",
-                    "cfg": {
-                        "num_bits": (2, 1),
-                        "block_sizes": {
-                            -1: 16,
-                            "type": "dynamic",
-                            "scale_bits": (4, 3),
-                        },
-                        "constant_amax": 1.0,
-                    },
-                },
-                {
-                    "quantizer_name": "*.k_bmm_quantizer",
-                    "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
-                },
-            ],
-            "algorithm": None,
-            "effective_bits": 6.25,
-        },
-        "fp8_k_nvfp4_v",
-    )
+    candidate = (_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v")
 
     model, state = auto_quantize_kv_cache(
         model,
@@ -353,6 +346,25 @@ def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quan
     assert get_quant_config(model)["quantization"]["kv_cache_quant_algo"] == "FP8_K_NVFP4_V"
 
 
+def test_kv_autoquant_rejects_asymmetric_candidate_for_unequal_kv_widths(
+    nvfp4_fake_quant_stub,
+):
+    model = _ToyKVModel()
+    model.attn0.k_proj = nn.Linear(8, 12, bias=False)
+    model.attn0.v_proj = nn.Linear(8, 8, bias=False)
+
+    with pytest.raises(ValueError, match=r"asymmetric K/V candidates.*unequal K/V widths"):
+        auto_quantize_kv_cache(
+            model,
+            {"kv_effective_bits": 6.25},
+            [(_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v")],
+            [torch.randn(1, 2, 8)],
+            lambda search_model, batch: search_model(batch),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+
 def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state():
     model = _ToyKVModel()
     original_quantizers = {

From c06fb0ee559a0d1edf0d4883c9806754e7415b94 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 25 Aug 2026 12:56:48 -0700
Subject: [PATCH 11/15] Fix KV AutoQuant export validation

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/recipe/config.py                     | 14 +++----
 modelopt/torch/export/quant_utils.py          | 18 +++------
 .../torch/quantization/kv_cache_auto_quant.py | 21 ++++++++++
 .../_test_utils/torch/transformers_models.py  |  1 +
 tests/unit/recipe/test_loader.py              |  9 +++++
 .../torch/export/test_get_quantization.py     | 32 +++++++++++++++
 .../quantization/test_kv_cache_auto_quant.py  | 40 +++++++++++++++++++
 7 files changed, 113 insertions(+), 22 deletions(-)

diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index b6a1edb5153..19e9ef7278b 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -150,12 +150,9 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
     """LP search constraints + cost model; matches the ``mtq.auto_quantize`` constraints dict."""
 
     effective_bits: float | None = ModeloptField(
-        default=None,
+        default=4.8,
         title="Effective bits per weight",
-        description=(
-            "Average weight-storage bits target for the LP, in (0, 16]. Defaults to 4.8 "
-            "when neither bit constraint is specified."
-        ),
+        description=("Average weight-storage bits target for the LP, in (0, 16]. Defaults to 4.8."),
     )
     kv_effective_bits: float | None = ModeloptField(
         default=None,
@@ -179,11 +176,10 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
 
     @model_validator(mode="before")
     @classmethod
-    def _default_weight_constraint(cls, data):
-        if isinstance(data, dict):
+    def _select_kv_constraint(cls, data):
+        if isinstance(data, dict) and "kv_effective_bits" in data and "effective_bits" not in data:
             data = dict(data)
-            if "effective_bits" not in data and "kv_effective_bits" not in data:
-                data["effective_bits"] = 4.8
+            data["effective_bits"] = None
         return data
 
     @field_validator("effective_bits", "kv_effective_bits")
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index b91ffc4079a..2d8f94366a6 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -18,7 +18,6 @@
 import logging
 from collections import defaultdict
 from collections.abc import Generator
-from types import SimpleNamespace
 from typing import Any
 from warnings import warn
 
@@ -1744,8 +1743,6 @@ def get_quant_config(
             layer_config_dict[name + ".quantization"] = quantization_format
             layer_config_dict[name + ".awq_block_size"] = block_size
 
-        not_enabled = SimpleNamespace(is_enabled=False)
-
         # Find kv cache quant format
         has_kv_quantizers = all(
             hasattr(module, quantizer_name)
@@ -1753,16 +1750,11 @@ def get_quant_config(
         )
         if has_kv_quantizers:
             kv_cache_eligible_layers += 1
-
-        if (
-            getattr(module, "k_bmm_quantizer", not_enabled).is_enabled
-            or getattr(module, "v_bmm_quantizer", not_enabled).is_enabled
-            or getattr(module, "output_quantizer", not_enabled).is_enabled
-        ):
-            module_kv_quant = get_kv_cache_dtype(module)
-            if module_kv_quant != QUANTIZATION_NONE:
-                kv_cache_formats.add(module_kv_quant)
-                kv_cache_quantized_layers[name] = {"quant_algo": module_kv_quant}
+            if module.k_bmm_quantizer.is_enabled and module.v_bmm_quantizer.is_enabled:
+                module_kv_quant = get_kv_cache_dtype(module)
+                if module_kv_quant != QUANTIZATION_NONE:
+                    kv_cache_formats.add(module_kv_quant)
+                    kv_cache_quantized_layers[name] = {"quant_algo": module_kv_quant}
 
     # MoE routers/gates are intentionally kept in original precision. On transformers>=5.0 they
     # are not nn.Linear modules (e.g. TopKRouter), never receive a quantizer, and would otherwise
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 1c9ff64dbf7..ecf6e52ffbe 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -155,6 +155,11 @@ def _validate_deployable_candidate(config: QuantizeConfig) -> None:
 
     algorithm_method = _algorithm_method(config)
     for attr, quantizer in quantizers.items():
+        if quantizer._dynamic:
+            raise ValueError(
+                f"KV-cache AutoQuant candidate {attr} uses top-level dynamic quantization, "
+                "which does not retain a persistent export scale."
+            )
         will_calibrate = algorithm_method == "max" and not quantizer._use_constant_amax
         if not hasattr(quantizer, "_amax") and not will_calibrate:
             raise ValueError(
@@ -435,6 +440,20 @@ def _restore_quantizer_state_dict(
                 quantizer.load_state_dict(quantizer_state)
 
 
+def _validate_persistent_candidate_scales(
+    candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]],
+) -> None:
+    """Require every calibrated candidate scale to be persistent in its state dict."""
+    for layer_name, layer_candidates in candidate_quantizers.items():
+        for candidate_name, layer_quantizers in layer_candidates.items():
+            for attr, quantizer in layer_quantizers.items():
+                if "_amax" not in quantizer.state_dict():
+                    raise ValueError(
+                        f"KV-cache AutoQuant candidate {candidate_name!r} for "
+                        f"{layer_name!r}/{attr} has no persistent export scale after calibration."
+                    )
+
+
 def _report_state(state: dict[str, Any]) -> dict[str, Any]:
     """Return the JSON-safe search report, excluding calibration tensors."""
     return {key: value for key, value in state.items() if key != "quantizer_state"}
@@ -526,6 +545,7 @@ def auto_quantize_kv_cache(
                     "Use a different checkpoint path."
                 )
             _restore_quantizer_state_dict(candidate_quantizers, quantizer_state)
+            _validate_persistent_candidate_scales(candidate_quantizers)
         else:
             from .model_quant import calibrate
 
@@ -560,6 +580,7 @@ def calibration_loop(calibration_model):
                 for layer_name, module, _ in layers:
                     _apply_layer_quantizers(module, disabled_quantizers[layer_name])
 
+            _validate_persistent_candidate_scales(candidate_quantizers)
             state = {
                 "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
                 "search_signature": signature,
diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py
index cf75e50e107..dc70c180057 100644
--- a/tests/_test_utils/torch/transformers_models.py
+++ b/tests/_test_utils/torch/transformers_models.py
@@ -220,6 +220,7 @@ def get_tiny_qwen3vl(**config_kwargs) -> PreTrainedModel:
         "head_dim": 8,
         "max_position_embeddings": 32,
         "vocab_size": 32,
+        "rope_scaling": {"rope_type": "default", "mrope_section": [1, 1, 2]},
     }
     text_kwargs.update(config_kwargs)
     # Pass as dicts — transformers 5.3.0 Qwen3VLConfig.__init__ only handles
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index a6ec40cf656..3c42e0bfafb 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1770,6 +1770,15 @@ def test_load_recipe_autoquantize_minimal(tmp_path):
     assert aq.module_search_spaces == []
 
 
+def test_autoquantize_constraints_preserve_default_with_kv_override():
+    assert AutoQuantizeConstraints.model_fields["effective_bits"].default == 4.8
+    assert AutoQuantizeConstraints().effective_bits == 4.8
+
+    constraints = AutoQuantizeConstraints(kv_effective_bits=5.4)
+    assert constraints.effective_bits is None
+    assert constraints.kv_effective_bits == 5.4
+
+
 def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path):
     """cost_model + cost.active_moe_expert_ratio parse and dump to the mtq constraints dict shape."""
     recipe_file = tmp_path / "aq.yml"
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index 6c778ec098e..1202b1fcae4 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -67,6 +67,38 @@ def test_nvfp4_static_quantizer_export():
     assert quant_config["quantization"]["group_size"] == 16
 
 
+def test_projection_output_quantizers_are_not_exported_as_kv_cache():
+    model = ToyModel()
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*.weight_quantizer",
+                "cfg": {"num_bits": (4, 3), "axis": None},
+                "enable": True,
+            },
+            {
+                "quantizer_name": "*.input_quantizer",
+                "cfg": {"num_bits": (4, 3), "axis": None},
+                "enable": True,
+            },
+            {
+                "quantizer_name": "*.output_quantizer",
+                "cfg": {"num_bits": (4, 3), "axis": None},
+                "enable": True,
+            },
+        ],
+        "algorithm": "max",
+    }
+    mtq.quantize(model, config, lambda x: x(torch.randn(1, 4, 10)))
+
+    quantization = get_quant_config(model)["quantization"]
+
+    assert quantization["quant_algo"] == "FP8"
+    assert quantization["kv_cache_quant_algo"] is None
+    assert "kv_cache_quantized_layers" not in quantization
+
+
 def test_mixed_kv_cache_quantization_exports_per_layer_map():
     class FakeAttention(torch.nn.Module):
         def __init__(self):
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index b0eadefd30a..f93ce8913af 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -190,6 +190,22 @@ def test_kv_candidate_rejects_non_exportable_or_incorrect_cost(config, match):
         _validate_kv_only_config(config)
 
 
+def test_kv_candidate_rejects_top_level_dynamic_fp8():
+    config = QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "type": "dynamic"},
+            }
+        ],
+        algorithm="max",
+        effective_bits=8.0,
+    )
+
+    with pytest.raises(ValueError, match="top-level dynamic"):
+        _validate_kv_only_config(config)
+
+
 class _ToyKVAttention(nn.Module):
     def __init__(self, width, gain):
         super().__init__()
@@ -214,6 +230,30 @@ def forward(self, x):
         return self.lm_head(self.attn1(self.attn0(x)))
 
 
+def test_kv_autoquant_rejects_missing_scale_after_calibration(monkeypatch):
+    model = _ToyKVModel()
+    original_quantizers = {
+        name: (module.k_bmm_quantizer, module.v_bmm_quantizer)
+        for name, module in (("attn0", model.attn0), ("attn1", model.attn1))
+    }
+    monkeypatch.setattr(model_quant, "calibrate", lambda *_args, **_kwargs: None)
+
+    with pytest.raises(ValueError, match="no persistent export scale after calibration"):
+        auto_quantize_kv_cache(
+            model,
+            {"kv_effective_bits": 8.0},
+            [(_kv_config((4, 3), 8.0).model_dump(), "fp8")],
+            [torch.randn(1, 2, 8)],
+            lambda search_model, batch: search_model(batch),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+    assert model.training
+    for name, module in (("attn0", model.attn0), ("attn1", model.attn1)):
+        assert (module.k_bmm_quantizer, module.v_bmm_quantizer) == original_quantizers[name]
+
+
 def test_kv_eligible_layers_supports_hybrid_attention_mixers_only():
     """Hybrid decoders include attention mixers but exclude nonattention mixers."""
     model = nn.Module()

From efc25b383c91a74f68cd3e05d9860fc8e489b515 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Thu, 27 Aug 2026 12:32:33 -0700
Subject: [PATCH 12/15] Use calibration-free KV AutoQuant candidates

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |  7 +++--
 .../kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml  | 22 +++++++++-----
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 30 +++++++------------
 tests/unit/recipe/test_loader.py              |  5 ++--
 4 files changed, 32 insertions(+), 32 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index d77dbd520f4..4063cea26cd 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -429,7 +429,7 @@ KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
 `candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
 packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are
 preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice.
-The shipped canary recipe searches calibrated FP8 K/V (8.0 bits/scalar) and packed NVFP4 K/V
+The shipped recipe searches FP8-cast K/V (8.0 bits/scalar) and NVFP4-cast K/V
 (4.5 bits/scalar) at 5.4 bits/scalar. It intentionally excludes FP8-K/NVFP4-V because the
 companion vLLM implementation does not support that asymmetric per-layer format:
 
@@ -441,8 +441,9 @@ python hf_ptq.py \
   --export_path /path/to/qwen3-1.7b-mixed-kv
 ```
 
-Each candidate uses max calibration so its persistent K/V scales are present in the unified HF
-checkpoint. Unified export records the selected formats in `kv_cache_quantized_layers` and writes
+Each candidate uses an explicit constant scale, avoiding an additional calibration pass while
+keeping persistent K/V scales in the unified HF checkpoint. Unified export records the selected
+formats in `kv_cache_quantized_layers` and writes
 the JSON-safe sensitivity report to `kv_cache_auto_quantize_report.json`;
 `--auto_quantize_checkpoint` stores the resumable raw search state.
 
diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
index db05db9cc70..ba191e87bc8 100644
--- a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
+++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -1,19 +1,19 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Layer-wise KV-cache search over calibrated FP8 and NVFP4 at 5.4 bits/scalar.
+# Layer-wise KV-cache search over FP8-cast and NVFP4-cast at 5.4 bits/scalar.
 # Isolated full-vocabulary forward KL is measured with every other eligible
 # attention layer kept in BF16.
 
 # modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
 imports:
   base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
-  kv_fp8: configs/ptq/units/kv_fp8
-  kv_nvfp4: configs/ptq/units/kv_nvfp4
+  fp8: configs/numerics/fp8
+  nvfp4: configs/numerics/nvfp4
 
 metadata:
   recipe_type: auto_quantize
-  description: Layer-wise calibrated FP8/NVFP4 KV-cache search at 5.4 bits using forward KL.
+  description: Layer-wise FP8-cast/NVFP4-cast KV-cache search at 5.4 bits using forward KL.
 
 auto_quantize:
   constraints:
@@ -21,12 +21,18 @@ auto_quantize:
 
   candidate_formats:
     - quant_cfg:
-        - $import: kv_fp8
-      algorithm: max
+        - quantizer_name: "*[kv]_bmm_quantizer"
+          cfg:
+            $import: fp8
+            constant_amax: 448.0
+      algorithm:
       effective_bits: 8.0
     - quant_cfg:
-        - $import: kv_nvfp4
-      algorithm: max
+        - quantizer_name: "*[kv]_bmm_quantizer"
+          cfg:
+            $import: nvfp4
+            constant_amax: 448.0
+      algorithm:
       effective_bits: 4.5
 
   auto_quantize_method: kl_div
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 2471bab4189..95eaec41ed5 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -28,6 +28,7 @@
 from modelopt.recipe import load_recipe
 from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints
 from modelopt.recipe.presets import QUANT_CFG_CHOICES
+from modelopt.torch.quantization import tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
@@ -108,24 +109,13 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
 def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch):
     """The HF entry point runs the real public KV AutoQuant path on an offline Qwen fixture."""
     hf_ptq = _import_hf_ptq(monkeypatch)
-    model = get_tiny_qwen3(num_hidden_layers=1)
-    aq = AutoQuantizeConfig(
-        constraints=AutoQuantizeConstraints(kv_effective_bits=8.0),
-        candidate_formats=[
-            QuantizeConfig(
-                quant_cfg=[
-                    {
-                        "quantizer_name": "*[kv]_bmm_quantizer",
-                        "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
-                    }
-                ],
-                algorithm=None,
-                effective_bits=8.0,
-            )
-        ],
-        auto_quantize_method="kl_div",
-        score_size=1,
+    monkeypatch.setattr(
+        tensor_quant,
+        "dynamic_block_quantize_op",
+        lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs),
     )
+    model = get_tiny_qwen3(num_hidden_layers=1)
+    aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize
     args = SimpleNamespace(
         calib_with_images=False,
         inference_pipeline_parallel=1,
@@ -139,8 +129,10 @@ def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch):
     hf_ptq.auto_quantize(args, model, data, aq, full_model=model)
 
     attention = model.model.layers[0].self_attn
-    assert attention.k_bmm_quantizer.num_bits == (4, 3)
-    assert attention.v_bmm_quantizer.num_bits == (4, 3)
+    assert attention.k_bmm_quantizer.num_bits == (2, 1)
+    assert attention.v_bmm_quantizer.num_bits == (2, 1)
+    assert attention.k_bmm_quantizer.amax == 448.0
+    assert attention.v_bmm_quantizer.amax == 448.0
 
 
 def test_kv_autoquant_names_asymmetric_export_format(monkeypatch):
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index 3c42e0bfafb..1a4fabb627b 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1951,13 +1951,14 @@ def test_load_recipe_kv_autoquantize_contract():
     assert aq.auto_quantize_method == "kl_div"
     assert "*mtp*" in aq.disabled_layers
     assert aq.cost_excluded_layers == []
+    assert all(candidate.algorithm is None for candidate in aq.candidate_formats)
     assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5]
     for fmt in aq.candidate_formats:
         for entry in fmt.quant_cfg:
             assert entry.quantizer_name == "*[kv]_bmm_quantizer"
             assert not entry.cfg.use_constant_amax
-            assert entry.cfg.constant_amax is None
-        assert fmt.algorithm == "max"
+            assert entry.cfg.constant_amax == 448.0
+        assert fmt.algorithm is None
 
 
 def test_kv_autoquantize_rejects_cost_excluded_layers():

From d6e0a66b15d29bd3425d7ab1ad353c645fcf15d8 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 1 Sep 2026 13:22:09 -0700
Subject: [PATCH 13/15] Limit FSDP2 rejection to KV AutoQuant

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                 |  2 +-
 examples/hf_ptq/hf_ptq.py                 | 31 ++++++++++---
 tests/examples/hf_ptq/test_hf_ptq_args.py | 56 +++++++++++++++++++++--
 3 files changed, 76 insertions(+), 13 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 1e538e5172c..9e6b031aeeb 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -498,7 +498,7 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop)
 
 ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point.
 
-> *AutoQuantize recipes are not supported with `--use_fsdp2` and are rejected before model loading. Distributed sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Use a PTQ recipe with FSDP2.*
+> *KV-cache AutoQuantize recipes are not supported with `--use_fsdp2` and are rejected before model loading. Distributed KV sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Existing weight AutoQuantize recipes retain their previous experimental warning with FSDP2.*
 
 ### Usage
 
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index f54643b6945..785da653187 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -99,10 +99,15 @@
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
 RAND_SEED = 1234
-_FSDP2_AUTOQUANT_ERROR = (
-    "AutoQuantize does not support --use_fsdp2 until distributed sensitivity scoring, "
+_FSDP2_KV_AUTOQUANT_ERROR = (
+    "KV-cache AutoQuantize does not support --use_fsdp2 until distributed sensitivity scoring, "
     "selection, and checkpoint writes are synchronized across ranks."
 )
+_FSDP2_AUTOQUANT_WARNING = (
+    "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet "
+    "(distributed calibration, sensitivity scoring, and recipe/checkpoint "
+    "synchronization across ranks); use at your own risk."
+)
 
 
 def _select_unpadded_logits(logits: torch.Tensor, batch: dict[str, Any]) -> torch.Tensor:
@@ -474,12 +479,13 @@ def auto_quantize(
         "Auto Quantization is not supported for pipeline parallel size > 1"
     )
 
-    if args.use_fsdp2:
-        raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR)
-
     inputs = _mtq_inputs_from_auto_quantize_config(
         aq_config, args, fixed_quantize_config=fixed_quantize_config
     )
+    if args.use_fsdp2:
+        if inputs["search_domain"] == "kv_cache":
+            raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR)
+        warnings.warn(_FSDP2_AUTOQUANT_WARNING)
 
     # base-model lm_head handling (mirrors the CLI helper)
     is_base_model = (
@@ -585,11 +591,22 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool:
     return recipe is not None and isinstance(load_recipe(recipe), ModelOptAutoQuantizeRecipe)
 
 
+def _recipe_is_kv_auto_quantize(recipe: str | None) -> bool:
+    """True if ``recipe`` resolves to a KV AutoQuantize recipe (peeked before model load)."""
+    if recipe is None:
+        return False
+    loaded_recipe = load_recipe(recipe)
+    return (
+        isinstance(loaded_recipe, ModelOptAutoQuantizeRecipe)
+        and loaded_recipe.auto_quantize.constraints.kv_effective_bits is not None
+    )
+
+
 def load_model(args: argparse.Namespace):
     # If low memory mode is enabled, we compress the model while loading the HF checkpoint.
     calibration_only = False
-    if args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
-        raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR)
+    if args.use_fsdp2 and _recipe_is_kv_auto_quantize(args.recipe):
+        raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR)
     if args.use_fsdp2:
         hf_config = AutoConfig.from_pretrained(
             args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 95eaec41ed5..cd46393196e 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -166,31 +166,77 @@ def test_kv_autoquant_kl_rejects_misaligned_attention_mask(monkeypatch):
         hf_ptq._select_unpadded_logits(torch.zeros(2, 4, 3), {"attention_mask": torch.ones(2, 3)})
 
 
-def test_autoquant_rejects_fsdp2(monkeypatch):
+def test_kv_autoquant_rejects_fsdp2(monkeypatch):
     hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(
+        hf_ptq,
+        "_mtq_inputs_from_auto_quantize_config",
+        lambda *_args, **_kwargs: {"search_domain": "kv_cache"},
+    )
     args = SimpleNamespace(
         calib_with_images=False,
         inference_pipeline_parallel=1,
         use_fsdp2=True,
     )
 
-    with pytest.raises(NotImplementedError, match="does not support --use_fsdp2"):
+    with pytest.raises(NotImplementedError, match="KV-cache AutoQuantize does not support"):
         hf_ptq.auto_quantize(args, torch.nn.Module(), [], SimpleNamespace())
 
 
-def test_fsdp2_autoquant_rejected_before_model_load(monkeypatch):
+def test_weight_autoquant_retains_fsdp2_warning(monkeypatch):
     hf_ptq = _import_hf_ptq(monkeypatch)
-    monkeypatch.setattr(hf_ptq, "_recipe_is_auto_quantize", lambda _: True)
+    model = torch.nn.Module()
+    inputs = {
+        "search_domain": "weight",
+        "constraints": {"effective_bits": 8.0},
+        "quantization_formats": [],
+        "fixed_quantization_config": None,
+        "module_search_spaces": [],
+        "disabled_layers": [],
+        "kv_cache_quant_cfg": None,
+        "method": "gradient",
+        "score_size": 1,
+    }
+    monkeypatch.setattr(
+        hf_ptq, "_mtq_inputs_from_auto_quantize_config", lambda *_args, **_kwargs: inputs
+    )
+    monkeypatch.setattr(
+        hf_ptq.mtq, "auto_quantize", lambda search_model, **_kwargs: (search_model, {})
+    )
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=True,
+        batch_size=1,
+        auto_quantize_checkpoint=None,
+    )
+
+    with pytest.warns(UserWarning, match="use at your own risk"):
+        assert hf_ptq.auto_quantize(args, model, [], SimpleNamespace()) is model
+
+
+def test_fsdp2_kv_autoquant_rejected_before_model_load(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(hf_ptq, "_recipe_is_kv_auto_quantize", lambda _: True)
     monkeypatch.setattr(
         hf_ptq.AutoConfig,
         "from_pretrained",
         lambda *_args, **_kwargs: pytest.fail("The model config must not be loaded."),
     )
 
-    with pytest.raises(NotImplementedError, match="does not support --use_fsdp2"):
+    with pytest.raises(NotImplementedError, match="KV-cache AutoQuantize does not support"):
         hf_ptq.load_model(SimpleNamespace(use_fsdp2=True, recipe="autoquant"))
 
 
+def test_fsdp2_preload_guard_distinguishes_weight_and_kv_autoquant(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+
+    assert hf_ptq._recipe_is_kv_auto_quantize(
+        "general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits"
+    )
+    assert not hf_ptq._recipe_is_kv_auto_quantize("general/auto_quantize/nvfp4_fp8_at_5p4bits")
+
+
 def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch):
     """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns
     key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict."""

From e75d06ea69817f05e216357abaa1a6f1ea7e2805 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 1 Sep 2026 13:37:51 -0700
Subject: [PATCH 14/15] Clarify KV AutoQuant recipe CLI help

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/hf_ptq.py                 |  7 ++++---
 tests/examples/hf_ptq/test_hf_ptq_args.py | 13 +++++++++++++
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 785da653187..36d17ccea69 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -1530,9 +1530,10 @@ def parse_args() -> argparse.Namespace:
         help=(
             "PTQ or AutoQuantize recipe YAML file or name without suffix (e.g. "
             "general/ptq/nvfp4_default-kv_fp8_cast, general/auto_quantize/nvfp4_fp8_at_4p8bits). "
-            "KV cache source depends on the recipe type: PTQ recipes bake KV cache into quant_cfg "
-            "and --kv_cache_qformat is ignored; AutoQuantize recipes fall back to --kv_cache_qformat "
-            "unless the recipe sets an explicit kv_cache field."
+            "KV cache behavior depends on the recipe type: PTQ recipes configure it in quant_cfg "
+            "and ignore --kv_cache_qformat; weight AutoQuantize recipes use their kv_cache setting "
+            "or fall back to --kv_cache_qformat; KV-cache AutoQuantize recipes select per-layer K/V "
+            "formats from candidate_formats and ignore --kv_cache_qformat."
         ),
         default=None,
     )
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index cd46393196e..a3db0be8b67 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -59,6 +59,19 @@ def _parse_hf_ptq_args(monkeypatch, *args):
     return hf_ptq, parsed_args
 
 
+def test_recipe_help_distinguishes_weight_and_kv_autoquant(monkeypatch, capsys):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--help"])
+
+    with pytest.raises(SystemExit) as exc_info:
+        hf_ptq.parse_args()
+
+    assert exc_info.value.code == 0
+    help_text = " ".join(capsys.readouterr().out.split())
+    assert "weight AutoQuantize recipes use their kv_cache setting" in help_text
+    assert "KV-cache AutoQuantize recipes select per-layer K/V formats" in help_text
+
+
 def test_autoquant_recipe_builds_mtq_inputs(monkeypatch):
     """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs."""
     hf_ptq, args = _parse_hf_ptq_args(

From 7a1f35ac0c33471dd510c7d6be38cc7b8cfbc8d1 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 1 Sep 2026 21:06:19 -0700
Subject: [PATCH 15/15] Unify KV-cache search with AutoQuantize

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   2 +-
 examples/hf_ptq/README.md                     |   3 +-
 examples/hf_ptq/hf_ptq.py                     |  64 +-
 modelopt/recipe/config.py                     |  50 +-
 .../torch/quantization/_auto_quantize_cost.py |  18 +-
 .../torch/quantization/kv_cache_auto_quant.py | 728 +++++++++++-------
 modelopt/torch/quantization/model_quant.py    | 208 ++---
 .../kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml  |   3 +-
 tests/examples/hf_ptq/test_hf_ptq_args.py     |   2 +-
 tests/unit/recipe/test_loader.py              |  24 +-
 .../quantization/test_kv_cache_auto_quant.py  | 105 ++-
 11 files changed, 709 insertions(+), 498 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 98339ecf74f..01ebb68405c 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -13,7 +13,7 @@ Changelog
 - Add a calibration-free streaming Kimi-K3 converter and checkpoint-mirror recipe for NVFP4 routed experts with ``input_scale=1.0`` and 128x128 block-FP8 KDA/MLA attention weights. The converter operates shard-by-shard on the source checkpoint's packed MXFP4 experts instead of loading the 2.8T model through the in-memory ``hf_ptq.py`` path.
 - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported.
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint.
-- Add layer-wise KV-cache AutoQuantize through ``mtq.auto_quantize_kv_cache`` and ``constraints.kv_effective_bits``. It measures isolated full-vocabulary forward KL for caller-supplied K/V formats, solves a width-weighted additive storage-constrained recipe across eligible layers, preserves search-disabled layers in their existing format, exports the selected per-attention mapping in unified HF checkpoints, and writes a JSON sensitivity report alongside the checkpoint. A cast-mode FP8/NVFP4 recipe at 5.4 bits/scalar is included.
+- Add layer-wise KV-cache AutoQuantize through ``mtq.auto_quantize`` with ``cost_model="kv_cache"``. It measures isolated full-vocabulary forward KL for caller-supplied K/V formats, solves a width-weighted additive storage-constrained recipe across eligible layers, preserves search-disabled layers in their existing format, exports the selected per-attention mapping in unified HF checkpoints, and writes a JSON sensitivity report alongside the checkpoint. A cast-mode FP8/NVFP4 recipe at 5.4 bits/scalar is included.
 - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Supports FP8 and NVFP4 on single-process models, resident or offloaded; other formats and placements raise ``NotImplementedError`` before calibration starts.
 
 *Megatron Framework (M-LM / M-Bridge)*
diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 9e6b031aeeb..230ee474164 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -433,7 +433,8 @@ For models without backprop support (e.g. Llama-4), use the `kl_div` scoring met
 Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall back to
 `--kv_cache_qformat` (default `fp8_cast`) unless they set an explicit `kv_cache` field.
 
-KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
+KV-cache AutoQuantize recipes use the same `mtq.auto_quantize` API and set
+`constraints.cost_model: kv_cache` with an `effective_bits` target. Their
 `candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
 packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are
 preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice.
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 36d17ccea69..4e6e8cc8534 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -273,7 +273,7 @@ def make_calib_dataloader(
         )
     else:
         assert tokenizer is not None and isinstance(
-            tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)
+            tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast
         ), "The PreTrainedTokenizer must be set"
         # Labels are only needed for gradient-based auto_quantize
         include_labels = autoquant_gradient_recipe
@@ -404,11 +404,11 @@ def _mtq_inputs_from_auto_quantize_config(
     to ``--kv_cache_qformat`` when the recipe omits it.
     """
     constraints = aq_config.constraints.model_dump(exclude_none=True)
-    is_kv_search = aq_config.constraints.kv_effective_bits is not None
+    is_kv_search = aq_config.constraints.cost_model == "kv_cache"
     if is_kv_search:
         return {
             "search_domain": "kv_cache",
-            "constraints": {"kv_effective_bits": constraints["kv_effective_bits"]},
+            "constraints": constraints,
             "quantization_formats": _mtq_kv_candidate_formats(aq_config.candidate_formats),
             "disabled_layers": aq_config.disabled_layers,
             "method": aq_config.auto_quantize_method,
@@ -534,39 +534,35 @@ def forward_step(model, batch):
             f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'"
         )
 
-    if inputs["search_domain"] == "kv_cache":
-        language_model, _ = mtq.auto_quantize_kv_cache(
-            language_model,
-            constraints=inputs["constraints"],
-            data_loader=calib_dataloader,
-            forward_step=forward_step,
-            quantization_formats=inputs["quantization_formats"],
-            num_calib_steps=len(calib_dataloader),
-            num_score_steps=min(
-                len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)
-            ),
-            verbose=True,
-            disabled_layers=inputs["disabled_layers"],
-            checkpoint=args.auto_quantize_checkpoint,
+    auto_quantize_kwargs: dict[str, Any] = {
+        "constraints": inputs["constraints"],
+        "data_loader": calib_dataloader,
+        "forward_step": forward_step,
+        "quantization_formats": inputs["quantization_formats"],
+        "num_calib_steps": len(calib_dataloader),
+        "num_score_steps": min(
+            len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)
+        ),
+        "verbose": True,
+        "disabled_layers": inputs["disabled_layers"],
+        "method": inputs["method"],
+        "checkpoint": args.auto_quantize_checkpoint,
+    }
+    if inputs["search_domain"] == "weight":
+        auto_quantize_kwargs.update(
+            {
+                "loss_func": loss_func,
+                "fixed_quantization_config": inputs["fixed_quantization_config"],
+                "module_search_spaces": inputs["module_search_spaces"],
+            }
         )
-        return language_model
 
     language_model, _ = mtq.auto_quantize(
         language_model,
-        constraints=inputs["constraints"],
-        data_loader=calib_dataloader,
-        forward_step=forward_step,
-        loss_func=loss_func,
-        quantization_formats=inputs["quantization_formats"],
-        fixed_quantization_config=inputs["fixed_quantization_config"],
-        module_search_spaces=inputs["module_search_spaces"],
-        num_calib_steps=len(calib_dataloader),
-        num_score_steps=min(len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)),
-        verbose=True,
-        disabled_layers=inputs["disabled_layers"],
-        method=inputs["method"],
-        checkpoint=args.auto_quantize_checkpoint,
+        **auto_quantize_kwargs,
     )
+    if inputs["search_domain"] == "kv_cache":
+        return language_model
 
     # KV cache quantization is uniform; applied after the LP search.
     kv_cache_quant_cfg = inputs["kv_cache_quant_cfg"]
@@ -598,7 +594,7 @@ def _recipe_is_kv_auto_quantize(recipe: str | None) -> bool:
     loaded_recipe = load_recipe(recipe)
     return (
         isinstance(loaded_recipe, ModelOptAutoQuantizeRecipe)
-        and loaded_recipe.auto_quantize.constraints.kv_effective_bits is not None
+        and loaded_recipe.auto_quantize.constraints.cost_model == "kv_cache"
     )
 
 
@@ -786,7 +782,7 @@ def sparsity_main(
     # Different calibration datasets are also available, e.g., "pile" and "wikipedia"
     # Please also check the docstring for the datasets available
     assert tokenizer is not None and isinstance(
-        tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)
+        tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast
     ), "The PreTrainedTokenizer must be set"
     calib_dataloader = get_dataset_dataloader(
         dataset_name=args.dataset,
@@ -1279,7 +1275,7 @@ def quantize_main(
     if args.recipe is not None:
         print(f"Use recipe {args.recipe} for quantization")
         recipe = load_recipe(args.recipe)
-        if not isinstance(recipe, (ModelOptPTQRecipe, ModelOptAutoQuantizeRecipe)):
+        if not isinstance(recipe, ModelOptPTQRecipe | ModelOptAutoQuantizeRecipe):
             raise TypeError(
                 f"Expected PTQ or AutoQuantize recipe, but got {type(recipe).__name__} "
                 f"from {args.recipe}"
diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index 19e9ef7278b..28de372adad 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -149,24 +149,20 @@ def _validate_active_moe_expert_ratio(cls, v: float | None) -> float | None:
 class AutoQuantizeConstraints(ModeloptBaseConfig):
     """LP search constraints + cost model; matches the ``mtq.auto_quantize`` constraints dict."""
 
-    effective_bits: float | None = ModeloptField(
+    effective_bits: float = ModeloptField(
         default=4.8,
-        title="Effective bits per weight",
-        description=("Average weight-storage bits target for the LP, in (0, 16]. Defaults to 4.8."),
-    )
-    kv_effective_bits: float | None = ModeloptField(
-        default=None,
-        title="Effective bits per KV-cache scalar",
+        title="Effective bits",
         description=(
-            "Average KV-cache storage bits target across eligible layers for layer-wise KV "
-            "AutoQuant, in (0, 16]. Exactly one of effective_bits and kv_effective_bits may "
-            "be set."
+            "Average storage-bits target for the selected cost model, in (0, 16]. Defaults to 4.8."
         ),
     )
-    cost_model: Literal["weight", "active_moe"] = ModeloptField(
+    cost_model: Literal["weight", "active_moe", "kv_cache"] = ModeloptField(
         default="weight",
         title="Cost model",
-        description="'weight' counts all weights equally; 'active_moe' scales routed-expert weights.",
+        description=(
+            "'weight' counts all weights equally; 'active_moe' scales routed-expert weights; "
+            "'kv_cache' accounts for paired K/V-cache storage."
+        ),
     )
     cost: AutoQuantizeCost | None = ModeloptField(
         default=None,
@@ -174,33 +170,17 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
         description="Extra cost-model parameters; omit for the 'weight' cost model.",
     )
 
-    @model_validator(mode="before")
-    @classmethod
-    def _select_kv_constraint(cls, data):
-        if isinstance(data, dict) and "kv_effective_bits" in data and "effective_bits" not in data:
-            data = dict(data)
-            data["effective_bits"] = None
-        return data
-
-    @field_validator("effective_bits", "kv_effective_bits")
+    @field_validator("effective_bits")
     @classmethod
-    def _validate_effective_bits(cls, v: float | None) -> float | None:
-        if v is not None and not (0 < v <= 16):
+    def _validate_effective_bits(cls, v: float) -> float:
+        if not (0 < v <= 16):
             raise ValueError(f"effective_bits must be in (0, 16], got {v}")
         return v
 
     @model_validator(mode="after")
-    def _exactly_one_bit_constraint(self):
-        if (self.effective_bits is None) == (self.kv_effective_bits is None):
-            raise ValueError(
-                "Exactly one of effective_bits and kv_effective_bits must be specified."
-            )
-        if self.kv_effective_bits is not None and (
-            self.cost_model != "weight" or self.cost is not None
-        ):
-            raise ValueError(
-                "KV-cache AutoQuant does not support weight or active-MoE cost settings."
-            )
+    def _validate_cost_settings(self):
+        if self.cost_model == "kv_cache" and self.cost is not None:
+            raise ValueError("KV-cache AutoQuant does not accept weight cost settings.")
         return self
 
 
@@ -303,7 +283,7 @@ def _has_search_space(self):
                 "auto_quantize requires candidate_formats or at least one module_search_spaces "
                 "entry. For uniform quantization, use a PTQ recipe instead."
             )
-        if self.constraints.kv_effective_bits is not None:
+        if self.constraints.cost_model == "kv_cache":
             if self.auto_quantize_method != "kl_div":
                 raise ValueError(
                     "KV-cache AutoQuant currently requires auto_quantize_method=kl_div."
diff --git a/modelopt/torch/quantization/_auto_quantize_cost.py b/modelopt/torch/quantization/_auto_quantize_cost.py
index 297e8cd8dc3..3834e264500 100644
--- a/modelopt/torch/quantization/_auto_quantize_cost.py
+++ b/modelopt/torch/quantization/_auto_quantize_cost.py
@@ -31,6 +31,7 @@
 EXCLUDED_MODULE_NAME_PATTERNS_KEY: Final = "excluded_module_name_patterns"
 COST_MODEL_WEIGHT: Final = "weight"
 COST_MODEL_ACTIVE_MOE: Final = "active_moe"
+COST_MODEL_KV_CACHE: Final = "kv_cache"
 
 _ROUTED_MOE_EXPERT_NAME_RE = re.compile(r"(^|\.)experts(\.|$)")
 _ACTIVE_MOE_TOP_K_ATTRS = (
@@ -68,7 +69,7 @@ def _iter_model_configs(model: nn.Module):
 def _get_first_numeric_config_attr(config: Any, attr_names: tuple[str, ...]) -> float | None:
     for attr_name in attr_names:
         value = getattr(config, attr_name, None)
-        if isinstance(value, (int, float)) and not isinstance(value, bool):
+        if isinstance(value, int | float) and not isinstance(value, bool):
             return float(value)
     return None
 
@@ -185,7 +186,7 @@ def normalize_cost_constraints(
                 )
 
         if not (
-            isinstance(active_moe_expert_ratio, (int, float))
+            isinstance(active_moe_expert_ratio, int | float)
             and not isinstance(active_moe_expert_ratio, bool)
             and 0.0 < active_moe_expert_ratio <= 1.0
         ):
@@ -206,9 +207,22 @@ def module_cost_weight(
         return base_weight
 
 
+class KVCacheCostModel(AutoQuantizeCostModel):
+    """Account for the exact storage of paired K/V-cache candidates."""
+
+    name = COST_MODEL_KV_CACHE
+    supported_cost_keys = frozenset()
+
+    @staticmethod
+    def candidate_cost(k_width: int, v_width: int, k_bits: float, v_bits: float) -> float:
+        """Return candidate storage in bits for one attention layer and token."""
+        return k_width * k_bits + v_width * v_bits
+
+
 _COST_MODELS: Final = {
     COST_MODEL_WEIGHT: WeightCostModel(),
     COST_MODEL_ACTIVE_MOE: ActiveMoECostModel(),
+    COST_MODEL_KV_CACHE: KVCacheCostModel(),
 }
 
 
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index ecf6e52ffbe..dc89a0a628c 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -19,26 +19,32 @@
 
 import fnmatch
 import math
-import os
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from tqdm import tqdm
 
-from modelopt.torch.opt.searcher import LPS
-from modelopt.torch.utils import print_rank_0, safe_load, safe_save
+from modelopt.torch.opt.hparam import Hparam
+from modelopt.torch.opt.searcher import LPS, BaseSearcher, SearchConfig, SearchStateDict
+from modelopt.torch.utils import print_rank_0
 
+from ._auto_quantize_cost import (
+    COST_MODEL_KV_CACHE,
+    KVCacheCostModel,
+    get_auto_quantize_cost_model,
+    normalize_auto_quantize_constraints,
+)
 from .config import QuantizeConfig
 from .conversion import set_quantizer_by_cfg
 from .nn import TensorQuantizer
 
-__all__ = ["auto_quantize_kv_cache"]
+__all__ = ["AutoQuantizeKVSearcher"]
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable
+    from collections.abc import Callable
 
 _KV_QUANTIZER_ATTRS = ("k_bmm_quantizer", "v_bmm_quantizer")
 _KV_AUTOQUANT_SCHEMA_VERSION = 1
@@ -199,14 +205,19 @@ def _validate_search_inputs(
     num_score_steps: int,
 ) -> tuple[float, list[tuple[str, QuantizeConfig]]]:
     """Validate a KV-cache search before the caller converts the model."""
-    if set(constraints) != {"kv_effective_bits"}:
+    if (
+        set(constraints) - {"effective_bits", "cost_model", "cost"}
+        or constraints.get("cost_model") != COST_MODEL_KV_CACHE
+        or "effective_bits" not in constraints
+        or constraints.get("cost") not in (None, {})
+    ):
         raise ValueError(
-            "KV-cache AutoQuant constraints must contain only kv_effective_bits; "
-            f"got {sorted(constraints)}."
+            "KV-cache AutoQuant requires an effective_bits target with "
+            f"cost_model='kv_cache'; got {constraints}."
         )
-    target_bits = float(constraints["kv_effective_bits"])
+    target_bits = float(constraints["effective_bits"])
     if not (0 < target_bits <= 16):
-        raise ValueError(f"kv_effective_bits must be in (0, 16], got {target_bits}.")
+        raise ValueError(f"effective_bits must be in (0, 16], got {target_bits}.")
     if num_calib_steps <= 0:
         raise ValueError("num_calib_steps must be positive.")
     if num_score_steps <= 0:
@@ -349,16 +360,22 @@ def _get_logits(
 
 def _solve_additive_recipe(
     layer_names: list[str],
-    scalar_weights: list[int],
+    layer_widths: list[tuple[int, int]],
     candidate_names: list[str],
-    candidate_bits: list[float],
+    candidate_kv_bits: list[tuple[float, float]],
     scores: list[list[float]],
     target_bits: float,
     verbose: bool,
 ) -> tuple[list[int], str]:
-    denominator = float(sum(scalar_weights))
+    cost_model = get_auto_quantize_cost_model(COST_MODEL_KV_CACHE)
+    assert isinstance(cost_model, KVCacheCostModel)
+    denominator = float(sum(k_width + v_width for k_width, v_width in layer_widths))
     candidate_costs = [
-        [weight * bits / 16.0 for bits in candidate_bits] for weight in scalar_weights
+        [
+            cost_model.candidate_cost(k_width, v_width, k_bits, v_bits) / 16.0
+            for k_bits, v_bits in candidate_kv_bits
+        ]
+        for k_width, v_width in layer_widths
     ]
     max_cost = denominator * target_bits / 16.0
     lps = LPS(
@@ -371,9 +388,9 @@ def _solve_additive_recipe(
     )
     selections, status = lps()
     if status != "Optimal":
-        minimum_bits = sum(weight * min(candidate_bits) for weight in scalar_weights) / denominator
+        minimum_bits = sum(min(costs) for costs in candidate_costs) * 16.0 / denominator
         raise ValueError(
-            f"KV-cache AutoQuant could not satisfy kv_effective_bits={target_bits}; "
+            f"KV-cache AutoQuant could not satisfy effective_bits={target_bits}; "
             f"minimum achievable value is {minimum_bits:.4f}. Solver status: {status}."
         )
     if len(selections) != len(layer_names):
@@ -387,13 +404,11 @@ def _solve_additive_recipe(
 def _search_signature(
     candidates: list[tuple[str, QuantizeConfig]],
     layers: list[tuple[str, nn.Module, int]],
-    target_bits: float,
     num_calib_steps: int,
     num_score_steps: int,
 ) -> dict[str, Any]:
     return {
         "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
-        "kv_effective_bits": target_bits,
         "num_calib_steps": num_calib_steps,
         "num_score_steps": num_score_steps,
         "candidates": [
@@ -403,7 +418,14 @@ def _search_signature(
             }
             for name, config in candidates
         ],
-        "layers": [{"name": name, "kv_scalar_weight": weight} for name, _, weight in layers],
+        "layers": [
+            {
+                "name": name,
+                "k_width": _projection_width(module, "k"),
+                "v_width": _projection_width(module, "v"),
+            }
+            for name, module, _ in layers
+        ],
     }
 
 
@@ -459,277 +481,427 @@ def _report_state(state: dict[str, Any]) -> dict[str, Any]:
     return {key: value for key, value in state.items() if key != "quantizer_state"}
 
 
-@torch.inference_mode()
-def auto_quantize_kv_cache(
-    model: nn.Module,
-    constraints: dict[str, Any],
-    quantization_formats: list[tuple[dict[str, Any], str]],
-    data_loader: Iterable,
-    forward_step: Callable[[nn.Module, Any], torch.Tensor],
-    *,
-    num_calib_steps: int,
-    num_score_steps: int,
-    disabled_layers: list[str] | str | None = None,
-    verbose: bool = False,
-    checkpoint: str | None = None,
-) -> tuple[nn.Module, dict[str, Any]]:
-    """Select one supplied K/V format per attention layer using isolated forward KL.
-
-    Candidate formats are format-agnostic ``QuantizeConfig`` dictionaries. Each must
-    configure K and V together and declare ``effective_bits`` matching its packed
-    storage per K-or-V scalar, including scale overhead. Candidate calibration is scoped
-    to the candidate K/V quantizers, while pre-existing fixed quantizers keep executing
-    with frozen state. Persistent ``constant_amax`` formats may skip calibration forwards.
-    """
-    target_bits, candidates = _validate_search_inputs(
-        constraints, quantization_formats, num_calib_steps, num_score_steps
-    )
+class QuantKVRecipeHparam(Hparam):
+    """One paired K/V format decision for an attention layer."""
+
+    def __init__(
+        self,
+        name: str,
+        module: nn.Module,
+        candidates: list[tuple[str, QuantizeConfig]],
+    ) -> None:
+        super().__init__(range(len(candidates)), original=0)
+        self.name = name
+        self.module = module
+        self.candidates = candidates
+        self.original_quantizers = {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS}
+        self.reference_quantizers = {attr: _disabled_quantizer() for attr in _KV_QUANTIZER_ATTRS}
+        self.candidate_quantizers = {
+            index: _candidate_quantizers(config) for index, (_, config) in enumerate(candidates)
+        }
+        k_width = _projection_width(module, "k")
+        v_width = _projection_width(module, "v")
+        assert k_width is not None and v_width is not None
+        self.k_width = k_width
+        self.v_width = v_width
+        self.use_reference()
+
+    @property
+    def active(self) -> int:
+        """Return the selected candidate index."""
+        assert isinstance(self._active, int)
+        return self._active
+
+    @active.setter
+    def active(self, value: int | None) -> None:
+        if value is None:
+            assert isinstance(self.original, int)
+            value = self.original
+        assert value in self.choices
+        self._active = value
+        _apply_layer_quantizers(self.module, self.candidate_quantizers[value])
+
+    def use_reference(self) -> None:
+        """Use BF16/no-quant K/V as the scoring reference, never as a solver choice."""
+        _apply_layer_quantizers(self.module, self.reference_quantizers)
+
+    def restore_original(self) -> None:
+        """Restore the K/V quantizer objects present before search."""
+        _apply_layer_quantizers(self.module, self.original_quantizers)
+
+    def candidate_name(self, index: int) -> str:
+        return self.candidates[index][0]
+
+    def candidate_bits(self, index: int) -> tuple[float, float]:
+        return _candidate_kv_bits(self.candidates[index][1])
+
+    def candidate_cost(self, index: int, cost_model: KVCacheCostModel) -> float:
+        k_bits, v_bits = self.candidate_bits(index)
+        return cost_model.candidate_cost(self.k_width, self.v_width, k_bits, v_bits)
+
+
+class AutoQuantizeKVSearcher(BaseSearcher):
+    """KV-cache AutoQuant backend using the shared search/checkpoint lifecycle."""
+
+    method_name = "kl_div"
+
+    @property
+    def default_search_config(self) -> SearchConfig:
+        """Return KV-specific fields layered on the shared search configuration."""
+        config = super().default_search_config
+        config.update(
+            {
+                "quantization_formats": [],
+                "forward_step": None,
+                "num_calib_steps": 512,
+                "num_score_steps": 128,
+                "disabled_layers": None,
+            }
+        )
+        return config
+
+    @property
+    def default_state_dict(self) -> SearchStateDict:
+        """Return the checkpointed KV search state."""
+        return {
+            "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
+            "method": self.method_name,
+            "cost_model": COST_MODEL_KV_CACHE,
+            "search_signature": None,
+            "calibration_complete": False,
+            "num_calib_steps": 0,
+            "score_reduction": "mean_per_scored_token",
+            "num_score_steps": 0,
+            "num_scored_tokens": 0,
+            "candidates": [],
+            "layers": {},
+            "quantizer_state": {},
+            "requested_constraints": {},
+            "best": {
+                "recipe": {},
+                "constraints": {},
+                "score": float("inf"),
+                "is_satisfied": False,
+                "solver_status": None,
+            },
+        }
 
-    layers = _eligible_layers(model, disabled_layers)
-    _validate_candidate_cost_geometry(candidates, layers)
-    signature = _search_signature(
-        candidates,
-        layers,
-        target_bits,
-        num_calib_steps,
-        num_score_steps,
-    )
-    candidate_names = [name for name, _ in candidates]
-    candidate_bits = []
-    for _, config in candidates:
-        assert config.effective_bits is not None
-        candidate_bits.append(config.effective_bits)
-
-    original_quantizers = {
-        name: {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS}
-        for name, module, _ in layers
-    }
-    disabled_quantizers = {
-        name: {attr: _disabled_quantizer() for attr in _KV_QUANTIZER_ATTRS} for name, _, _ in layers
-    }
-    candidate_quantizers = {
-        name: {
-            candidate_name: _candidate_quantizers(config) for candidate_name, config in candidates
+    def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig:
+        """Validate the data inputs required by isolated forward-KL scoring."""
+        config = super().sanitize_search_config(config)
+        if config["data_loader"] is None:
+            raise ValueError("data_loader must be provided for KV-cache AutoQuant.")
+        if config["forward_step"] is None:
+            raise ValueError("forward_step must be provided for KV-cache AutoQuant.")
+        return config
+
+    def load_search_checkpoint(self) -> bool:
+        """Load compatible fields before validating the resolved KV search signature."""
+        return super().load_search_checkpoint(strict=False)
+
+    @property
+    def _candidate_quantizer_map(
+        self,
+    ) -> dict[str, dict[str, dict[str, TensorQuantizer]]]:
+        return {
+            hparam.name: {
+                hparam.candidate_name(index): quantizers
+                for index, quantizers in hparam.candidate_quantizers.items()
+            }
+            for hparam in self._hparams
         }
-        for name, _, _ in layers
-    }
 
-    is_training = model.training
-    model.eval()
-    try:
-        for name, module, _ in layers:
-            _apply_layer_quantizers(module, disabled_quantizers[name])
+    def restore_original_quantizers(self) -> None:
+        """Restore pre-search K/V objects after a failed search."""
+        for hparam in getattr(self, "_hparams", []):
+            hparam.restore_original()
+
+    def _calibrate_candidates(self) -> None:
+        from .model_quant import calibrate
+
+        data_loader = self.config["data_loader"]
+        forward_step = self.config["forward_step"]
+        num_calib_steps = self.config["num_calib_steps"]
+        for candidate_index, (_, config) in enumerate(self._candidates):
+            for hparam in self._hparams:
+                hparam.active = candidate_index
+
+            if config.algorithm is not None:
+
+                def calibration_loop(calibration_model):
+                    for step, data in enumerate(data_loader):
+                        if step >= num_calib_steps:
+                            break
+                        _get_logits(forward_step, calibration_model, data)
+
+                active_quantizers = [
+                    quantizer
+                    for hparam in self._hparams
+                    for quantizer in hparam.candidate_quantizers[candidate_index].values()
+                ]
+                calibration_proxy = nn.Module()
+                calibration_proxy.quantizers = nn.ModuleList(active_quantizers)
+                with _freeze_existing_quantizers(self.model, active_quantizers):
+                    calibrate(
+                        calibration_proxy,
+                        algorithm=config.algorithm,
+                        forward_loop=lambda _: calibration_loop(self.model),
+                    )
 
-        state: dict[str, Any] | None = None
-        if checkpoint is not None and os.path.exists(checkpoint):
-            restored = safe_load(checkpoint)
-            if not isinstance(restored, dict):
-                raise ValueError(
-                    "KV-cache AutoQuant checkpoint must contain a search-state dictionary."
-                )
-            if _checkpoint_state_is_compatible(restored, signature):
-                state = restored
-                if verbose:
-                    print_rank_0(f"KV-cache AutoQuant restored search state from {checkpoint}.")
-            else:
-                raise ValueError(
-                    "KV-cache AutoQuant checkpoint does not match the current candidates "
-                    "or eligible layers. Use a different checkpoint path."
-                )
+            for hparam in self._hparams:
+                hparam.use_reference()
+
+        candidate_quantizers = self._candidate_quantizer_map
+        _validate_persistent_candidate_scales(candidate_quantizers)
+        self.quantizer_state = _quantizer_state_dict(candidate_quantizers)
+        self.calibration_complete = True
+        self.num_calib_steps = num_calib_steps
+        self.save_search_checkpoint(verbose=self.config["verbose"])
+
+    def before_search(self) -> None:
+        """Resolve attention decisions and calibrate or restore candidate scales."""
+        super().before_search()
+        self.constraints = normalize_auto_quantize_constraints(self.model, self.constraints)
+        target_bits, self._candidates = _validate_search_inputs(
+            self.constraints,
+            self.config["quantization_formats"],
+            self.config["num_calib_steps"],
+            self.config["num_score_steps"],
+        )
+        self._target_bits = target_bits
+        layers = _eligible_layers(self.model, self.config["disabled_layers"])
+        _validate_candidate_cost_geometry(self._candidates, layers)
+        signature = _search_signature(
+            self._candidates,
+            layers,
+            self.config["num_calib_steps"],
+            self.config["num_score_steps"],
+        )
+        if self.search_signature is not None and not _checkpoint_state_is_compatible(
+            self.state_dict(), signature
+        ):
+            raise ValueError(
+                "KV-cache AutoQuant checkpoint does not match the current candidates, scoring "
+                "setup, or eligible layers. Use a different checkpoint path."
+            )
+        self.search_signature = signature
+        self._hparams = [
+            QuantKVRecipeHparam(name, module, self._candidates) for name, module, _ in layers
+        ]
+        self._cost_model = get_auto_quantize_cost_model(COST_MODEL_KV_CACHE)
+        assert isinstance(self._cost_model, KVCacheCostModel)
+        self.candidates = [
+            {
+                "name": name,
+                "effective_bits": config.effective_bits,
+                "k_bits": _candidate_kv_bits(config)[0],
+                "v_bits": _candidate_kv_bits(config)[1],
+                "config": config.model_dump(mode="python", exclude_none=True),
+            }
+            for name, config in self._candidates
+        ]
+        self.model.eval()
 
-        if state is not None and state.get("calibration_complete"):
-            quantizer_state = state.get("quantizer_state")
-            if quantizer_state is None:
+        if self.calibration_complete:
+            if not self.quantizer_state:
                 raise ValueError(
                     "KV-cache AutoQuant checkpoint is missing calibrated quantizer state. "
                     "Use a different checkpoint path."
                 )
-            _restore_quantizer_state_dict(candidate_quantizers, quantizer_state)
+            candidate_quantizers = self._candidate_quantizer_map
+            _restore_quantizer_state_dict(candidate_quantizers, self.quantizer_state)
             _validate_persistent_candidate_scales(candidate_quantizers)
+            if self.config["verbose"]:
+                print_rank_0("KV-cache AutoQuant restored calibration from checkpoint.")
         else:
-            from .model_quant import calibrate
+            self._calibrate_candidates()
 
-            for candidate_name, config in candidates:
-                for layer_name, module, _ in layers:
-                    _apply_layer_quantizers(
-                        module, candidate_quantizers[layer_name][candidate_name]
+    def _estimate_sensitivity_scores(self) -> None:
+        candidate_names = [name for name, _ in self._candidates]
+        score_sums: dict[str, dict[str, torch.Tensor | None]] = {
+            hparam.name: dict.fromkeys(candidate_names) for hparam in self._hparams
+        }
+        scored_tokens = 0
+        scored_steps = 0
+        iterator = tqdm(
+            self.config["data_loader"],
+            total=self.config["num_score_steps"],
+            desc="Estimating KV-cache KL sensitivity",
+            disable=not self.config["verbose"],
+        )
+        for data in iterator:
+            if scored_steps >= self.config["num_score_steps"]:
+                break
+            logits_ref = _get_logits(self.config["forward_step"], self.model, data)
+            log_prob_ref = torch.log_softmax(logits_ref.float(), dim=-1)
+            scored_tokens += logits_ref.numel() // logits_ref.shape[-1]
+
+            for hparam in self._hparams:
+                for candidate_index, candidate_name in enumerate(candidate_names):
+                    hparam.active = candidate_index
+                    logits_quant = _get_logits(self.config["forward_step"], self.model, data)
+                    if logits_quant.shape != logits_ref.shape:
+                        raise ValueError(
+                            "KV-cache AutoQuant forward_step returned different reference and "
+                            f"candidate logits shapes: {tuple(logits_ref.shape)} and "
+                            f"{tuple(logits_quant.shape)}."
+                        )
+                    score = F.kl_div(
+                        torch.log_softmax(logits_quant.float(), dim=-1),
+                        log_prob_ref,
+                        reduction="sum",
+                        log_target=True,
+                    )
+                    previous_score = score_sums[hparam.name][candidate_name]
+                    score_sums[hparam.name][candidate_name] = (
+                        score if previous_score is None else previous_score + score
+                    )
+                    hparam.use_reference()
+            scored_steps += 1
+
+        if scored_steps == 0 or scored_tokens == 0:
+            raise ValueError("KV-cache AutoQuant data_loader produced no scoring batches.")
+
+        self.layers = {}
+        for hparam in self._hparams:
+            scores = {}
+            for candidate_name in candidate_names:
+                score_sum = score_sums[hparam.name][candidate_name]
+                if score_sum is None:
+                    raise RuntimeError(
+                        "KV-cache AutoQuant did not collect a score for "
+                        f"{hparam.name!r}/{candidate_name!r}."
+                    )
+                if not torch.isfinite(score_sum):
+                    raise ValueError(
+                        "KV-cache AutoQuant produced a non-finite KL score for "
+                        f"{hparam.name!r}/{candidate_name!r}."
                     )
+                scores[candidate_name] = float(score_sum.item()) / scored_tokens
+            self.layers[hparam.name] = {
+                "k_width": hparam.k_width,
+                "v_width": hparam.v_width,
+                "scores": scores,
+            }
+        self.num_score_steps = scored_steps
+        self.num_scored_tokens = scored_tokens
+        self.save_search_checkpoint(verbose=self.config["verbose"])
+
+    def _solve(self) -> None:
+        candidate_names = [name for name, _ in self._candidates]
+        candidate_kv_bits = [_candidate_kv_bits(config) for _, config in self._candidates]
+        layers = cast("dict[str, dict[str, Any]]", self.layers)
+        scores = [
+            [layers[hparam.name]["scores"][name] for name in candidate_names]
+            for hparam in self._hparams
+        ]
+        selections, status = _solve_additive_recipe(
+            [hparam.name for hparam in self._hparams],
+            [(hparam.k_width, hparam.v_width) for hparam in self._hparams],
+            candidate_names,
+            candidate_kv_bits,
+            scores,
+            self._target_bits,
+            self.config["verbose"],
+        )
+        denominator = float(sum(hparam.k_width + hparam.v_width for hparam in self._hparams))
+        cost_model = self._cost_model
+        assert isinstance(cost_model, KVCacheCostModel)
+        total_cost = sum(
+            (
+                hparam.candidate_cost(selected, cost_model)
+                for hparam, selected in zip(self._hparams, selections)
+            ),
+            start=0.0,
+        )
+        achieved_bits = total_cost / denominator
+        selected_score = sum(
+            layer_scores[selected] for selected, layer_scores in zip(selections, scores)
+        )
+        recipe = {}
+        for hparam, selected in zip(self._hparams, selections):
+            hparam.active = selected
+            selected_name = hparam.candidate_name(selected)
+            self.layers[hparam.name]["selected"] = selected_name
+            recipe[hparam.name] = selected_name
+            if self.config["verbose"]:
+                print_rank_0(f"KV-cache AutoQuant selected {selected_name} for {hparam.name}.")
+        self.requested_constraints = {
+            "effective_bits": self._target_bits,
+            "cost_model": COST_MODEL_KV_CACHE,
+        }
+        self.best = {
+            "recipe": recipe,
+            "constraints": {
+                "effective_bits": achieved_bits,
+                "cost_model": COST_MODEL_KV_CACHE,
+            },
+            "score": selected_score,
+            "is_satisfied": achieved_bits <= self._target_bits + 1e-12,
+            "solver_status": status,
+        }
 
-                if config.algorithm is not None:
-
-                    def calibration_loop(calibration_model):
-                        for step, data in enumerate(data_loader):
-                            if step >= num_calib_steps:
-                                break
-                            _get_logits(forward_step, calibration_model, data)
-
-                    active_quantizers = [
-                        quantizer
-                        for layer_name, _, _ in layers
-                        for quantizer in candidate_quantizers[layer_name][candidate_name].values()
-                    ]
-                    calibration_proxy = nn.Module()
-                    calibration_proxy.quantizers = nn.ModuleList(active_quantizers)
-                    with _freeze_existing_quantizers(model, active_quantizers):
-                        calibrate(
-                            calibration_proxy,
-                            algorithm=config.algorithm,
-                            forward_loop=lambda _: calibration_loop(model),
-                        )
+    @torch.inference_mode()
+    def run_search(self) -> None:
+        """Score candidates when needed, solve the budget, and apply the selection."""
+        if not self.layers:
+            self._estimate_sensitivity_scores()
+        self._solve()
+        self.save_search_checkpoint(verbose=self.config["verbose"])
 
-                for layer_name, module, _ in layers:
-                    _apply_layer_quantizers(module, disabled_quantizers[layer_name])
+    def after_search(self) -> None:
+        """Attach the JSON-safe sensitivity report used by unified export."""
+        self.model._modelopt_kv_cache_auto_quantize_state = _report_state(self.state_dict())
 
-            _validate_persistent_candidate_scales(candidate_quantizers)
-            state = {
-                "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
-                "search_signature": signature,
-                "calibration_complete": True,
-                "num_calib_steps": num_calib_steps,
-                "quantizer_state": _quantizer_state_dict(candidate_quantizers),
-            }
-            if checkpoint is not None:
-                checkpoint_dir = os.path.dirname(checkpoint)
-                if checkpoint_dir:
-                    os.makedirs(checkpoint_dir, exist_ok=True)
-                safe_save(state, checkpoint)
-
-        assert state is not None
-        if not state.get("layers"):
-            score_sums: dict[str, dict[str, torch.Tensor | None]] = {
-                layer_name: dict.fromkeys(candidate_names) for layer_name, _, _ in layers
-            }
-            scored_tokens = 0
-            scored_steps = 0
-            iterator = tqdm(
-                data_loader,
-                total=num_score_steps,
-                desc="Estimating KV-cache KL sensitivity",
-                disable=not verbose,
-            )
-            for data in iterator:
-                if scored_steps >= num_score_steps:
-                    break
-                logits_ref = _get_logits(forward_step, model, data)
-                log_prob_ref = torch.log_softmax(logits_ref.float(), dim=-1)
-                scored_tokens += logits_ref.numel() // logits_ref.shape[-1]
-
-                for layer_name, module, _ in layers:
-                    for candidate_name, _ in candidates:
-                        _apply_layer_quantizers(
-                            module, candidate_quantizers[layer_name][candidate_name]
-                        )
-                        logits_quant = _get_logits(forward_step, model, data)
-                        if logits_quant.shape != logits_ref.shape:
-                            raise ValueError(
-                                "KV-cache AutoQuant forward_step returned different reference and "
-                                f"candidate logits shapes: {tuple(logits_ref.shape)} and "
-                                f"{tuple(logits_quant.shape)}."
-                            )
-                        score = F.kl_div(
-                            torch.log_softmax(logits_quant.float(), dim=-1),
-                            log_prob_ref,
-                            reduction="sum",
-                            log_target=True,
-                        )
-                        previous_score = score_sums[layer_name][candidate_name]
-                        score_sums[layer_name][candidate_name] = (
-                            score if previous_score is None else previous_score + score
-                        )
-                        _apply_layer_quantizers(module, disabled_quantizers[layer_name])
-                scored_steps += 1
-
-            if scored_steps == 0 or scored_tokens == 0:
-                raise ValueError("KV-cache AutoQuant data_loader produced no scoring batches.")
-            scores = []
-            for layer_name, _, _ in layers:
-                layer_scores = []
-                for candidate_name in candidate_names:
-                    score_sum = score_sums[layer_name][candidate_name]
-                    if score_sum is None:
-                        raise RuntimeError(
-                            "KV-cache AutoQuant did not collect a score for "
-                            f"{layer_name!r}/{candidate_name!r}."
-                        )
-                    if not torch.isfinite(score_sum):
-                        raise ValueError(
-                            "KV-cache AutoQuant produced a non-finite KL score for "
-                            f"{layer_name!r}/{candidate_name!r}."
-                        )
-                    layer_scores.append(float(score_sum.item()) / scored_tokens)
-                scores.append(layer_scores)
-            selections, status = _solve_additive_recipe(
-                [name for name, _, _ in layers],
-                [weight for _, _, weight in layers],
-                candidate_names,
-                candidate_bits,
-                scores,
-                target_bits,
-                verbose,
-            )
-            denominator = float(sum(weight for _, _, weight in layers))
-            achieved_bits = (
-                sum(
-                    weight * candidate_bits[selected]
-                    for selected, (_, _, weight) in zip(selections, layers)
-                )
-                / denominator
-            )
-            selected_score = sum(
-                layer_scores[selected] for selected, layer_scores in zip(selections, scores)
-            )
-            state.update(
-                {
-                    "method": "kl_div",
-                    "score_reduction": "mean_per_scored_token",
-                    "constraints": {"kv_effective_bits": target_bits},
-                    "num_score_steps": scored_steps,
-                    "num_scored_tokens": scored_tokens,
-                    "candidates": [
-                        {
-                            "name": name,
-                            "effective_bits": effective_bits,
-                            "config": config.model_dump(mode="json", exclude_none=True),
-                        }
-                        for (name, config), effective_bits in zip(candidates, candidate_bits)
-                    ],
-                    "layers": {
-                        layer_name: {
-                            "kv_scalar_weight": weight,
-                            "scores": dict(zip(candidate_names, layer_scores)),
-                            "selected": candidate_names[selected],
-                        }
-                        for selected, layer_scores, (layer_name, _, weight) in zip(
-                            selections, scores, layers
-                        )
-                    },
-                    "best": {
-                        "effective_bits": achieved_bits,
-                        "score": selected_score,
-                        "is_satisfied": achieved_bits <= target_bits + 1e-12,
-                        "solver_status": status,
-                    },
-                }
-            )
-            if checkpoint is not None:
-                checkpoint_dir = os.path.dirname(checkpoint)
-                if checkpoint_dir:
-                    os.makedirs(checkpoint_dir, exist_ok=True)
-                safe_save(state, checkpoint)
-                if verbose:
-                    print_rank_0(f"Saved KV-cache AutoQuant report to {checkpoint}.")
-
-        for layer_name, module, _ in layers:
-            selected_name = state["layers"][layer_name]["selected"]
-            _apply_layer_quantizers(module, candidate_quantizers[layer_name][selected_name])
-            if verbose:
-                print_rank_0(f"KV-cache AutoQuant selected {selected_name} for {layer_name}.")
-        report = _report_state(state)
-        model._modelopt_kv_cache_auto_quantize_state = report
-        return model, report
-    except Exception:
-        for layer_name, module, _ in layers:
-            _apply_layer_quantizers(module, original_quantizers[layer_name])
-        raise
-    finally:
-        model.train(is_training)
+
+def _config_entry_dict(entry: Any) -> dict[str, Any]:
+    if hasattr(entry, "model_dump"):
+        return entry.model_dump(mode="python", exclude_none=True)
+    return dict(entry)
+
+
+def get_kv_cache_auto_quantize_config(
+    search_state: dict[str, Any],
+    constraints: dict[str, Any] | None = None,
+    verbose: bool = False,
+) -> dict[str, Any]:
+    """Build a flat K/V quantization config, optionally re-solving at a new target."""
+    requested = constraints or search_state.get("requested_constraints")
+    if not isinstance(requested, dict):
+        raise ValueError("KV-cache AutoQuant search state has no requested constraints.")
+    target_bits = float(requested["effective_bits"])
+    if requested.get("cost_model", COST_MODEL_KV_CACHE) != COST_MODEL_KV_CACHE:
+        raise ValueError("KV-cache search state can only be re-solved with cost_model='kv_cache'.")
+
+    candidates = search_state["candidates"]
+    candidate_names = [candidate["name"] for candidate in candidates]
+    candidate_kv_bits = [
+        (float(candidate["k_bits"]), float(candidate["v_bits"])) for candidate in candidates
+    ]
+    layers = search_state["layers"]
+    layer_names = list(layers)
+    selections, _ = _solve_additive_recipe(
+        layer_names,
+        [(layers[name]["k_width"], layers[name]["v_width"]) for name in layer_names],
+        candidate_names,
+        candidate_kv_bits,
+        [
+            [layers[name]["scores"][candidate] for candidate in candidate_names]
+            for name in layer_names
+        ],
+        target_bits,
+        verbose,
+    )
+
+    quant_cfg: list[dict[str, Any]] = []
+    for layer_name, selected in zip(layer_names, selections):
+        config = QuantizeConfig(**candidates[selected]["config"])
+        for entry in config.quant_cfg:
+            entry_dict = _config_entry_dict(entry)
+            pattern = entry_dict["quantizer_name"]
+            for attr in _KV_QUANTIZER_ATTRS:
+                if fnmatch.fnmatch(f"{_KV_CANDIDATE_HOLDER_NAME}.{attr}", pattern):
+                    resolved = dict(entry_dict)
+                    resolved["quantizer_name"] = f"{layer_name}.{attr}"
+                    quant_cfg.append(resolved)
+    return {"quant_cfg": quant_cfg, "algorithm": "max"}
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index b6a2079b8d9..f14df6fb466 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -39,18 +39,18 @@
 )
 from modelopt.torch.utils import atomic_print
 
+from ._auto_quantize_cost import COST_MODEL_KV_CACHE
 from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe
 from .algorithms import get_auto_quantize_config as _get_auto_quantize_config
 from .config import QuantizeAlgoCfgType
+from .kv_cache_auto_quant import AutoQuantizeKVSearcher, get_kv_cache_auto_quantize_config
 from .kv_cache_auto_quant import _validate_search_inputs as _validate_kv_cache_search_inputs
-from .kv_cache_auto_quant import auto_quantize_kv_cache as _auto_quantize_kv_cache
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, SequentialQuantizer, TensorQuantizer
 from .utils import is_quantized
 
 __all__ = [
     "auto_quantize",
-    "auto_quantize_kv_cache",
     "calibrate",
     "compute_quantization_mse",
     "disable_quantizer",
@@ -276,7 +276,7 @@ def forward_loop(model) -> None:
 def auto_quantize(
     model: nn.Module,
     constraints: dict[str, Any] | None = None,
-    quantization_formats: list[dict[str, Any] | str] | None = None,
+    quantization_formats: list[dict[str, Any] | str | tuple[dict[str, Any], str]] | None = None,
     data_loader: Iterable | None = None,
     forward_step: Callable[[nn.Module, Any], Any | torch.Tensor] | None = None,
     loss_func: Callable[[Any, Any], torch.Tensor] | None = None,
@@ -285,7 +285,7 @@ def auto_quantize(
     num_calib_steps: int = 512,
     num_score_steps: int = 128,
     verbose: bool = False,
-    method: str = "gradient",
+    method: str | None = None,
     checkpoint: str | None = None,
     module_search_spaces: list[dict[str, Any]] | None = None,
     fixed_quantization_config: dict[str, Any] | str | None = None,
@@ -309,9 +309,11 @@ def auto_quantize(
         model: A pytorch model with quantizer modules.
         constraints: Constraints for the search. ``effective_bits`` specifies the effective number
             of bits for the quantized model and defaults to 4.8. ``cost_model`` selects the metric
-            used for the effective-bits constraint and currently supports ``"weight"`` (default)
-            and ``"active_moe"``. Additional cost-model parameters are provided through the nested
-            ``cost`` dict.
+            used for the effective-bits constraint and supports ``"weight"`` (default),
+            ``"active_moe"``, and ``"kv_cache"``. The KV-cache cost model dispatches to isolated
+            forward-KL scoring over paired K/V formats; BF16/no-quant is its scoring reference but
+            is never solver-selectable. Additional cost-model parameters are provided through the
+            nested ``cost`` dict.
 
             Here is an example for valid ``effective_bits`` argument:
 
@@ -330,6 +332,9 @@ def auto_quantize(
                     },
                 }
 
+                # For paired K/V-cache formats with exact K/V storage accounting
+                constraints = {"effective_bits": 5.4, "cost_model": "kv_cache"}
+
         quantization_formats: A list of quantization format config dictionaries or string names to search for.
             Each config dictionary should be valid as a ``config`` argument in
             :meth:`quantize <modelopt.torch.quantization.model_quant.quantize>`.
@@ -530,6 +535,99 @@ def _process_quantization_formats(formats, custom_name_prefix):
             processed.append((quant_cfg, name))
         return processed
 
+    is_kv_search = constraints is not None and constraints.get("cost_model") == COST_MODEL_KV_CACHE
+    if is_kv_search:
+        if (
+            torch.distributed.is_available()
+            and torch.distributed.is_initialized()
+            and torch.distributed.get_world_size() > 1
+        ):
+            raise RuntimeError(
+                "KV-cache AutoQuant is single-process only; distributed scoring, selection, "
+                "and checkpoint writes are not synchronized."
+            )
+        if method not in (None, "kl_div"):
+            raise ValueError("cost_model='kv_cache' requires method='kl_div'.")
+        if fixed_quantization_config is not None or module_search_spaces:
+            raise ValueError(
+                "KV-cache AutoQuant does not support fixed_quantization_config or "
+                "module_search_spaces."
+            )
+        if loss_func is not None or forward_backward_step is not None:
+            raise ValueError(
+                "KV-cache AutoQuant uses forward KL and does not accept loss_func or "
+                "forward_backward_step."
+            )
+        if not isinstance(quantization_formats, list) or not quantization_formats:
+            raise ValueError(
+                "cost_model='kv_cache' requires a non-empty quantization_formats list."
+            )
+        if data_loader is None or forward_step is None:
+            raise ValueError(
+                "data_loader and forward_step must be provided for KV-cache AutoQuant."
+            )
+        processed_kv_formats = []
+        for index, candidate in enumerate(quantization_formats):
+            if isinstance(candidate, tuple):
+                raw_config, name = candidate
+            elif isinstance(candidate, str):
+                if not hasattr(mtq, candidate):
+                    raise ValueError(f"Unknown KV-cache quantization format: {candidate!r}.")
+                raw_config, name = getattr(mtq, candidate), candidate
+            elif isinstance(candidate, dict):
+                raw_config = candidate
+                name = QuantRecipe.get_auto_name_for_config(candidate) or f"KV_CACHE_FORMAT_{index}"
+            else:
+                raise TypeError(
+                    "KV-cache quantization formats must be config dictionaries, preset names, "
+                    "or (config, name) tuples."
+                )
+            if not isinstance(raw_config, dict):
+                raise TypeError("KV-cache AutoQuant formats must resolve to config dictionaries.")
+            if not isinstance(name, str) or not name:
+                raise ValueError("KV-cache AutoQuant candidate names must be non-empty strings.")
+            processed_kv_formats.append((raw_config, name))
+
+        assert constraints is not None
+        _validate_kv_cache_search_inputs(
+            constraints,
+            processed_kv_formats,
+            num_calib_steps,
+            num_score_steps,
+        )
+        converted_for_search = not is_quantized(model)
+        conversion_snapshot = _snapshot_model_structure(model) if converted_for_search else []
+        is_training = model.training
+        searcher = AutoQuantizeKVSearcher()
+        try:
+            if converted_for_search:
+                model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
+                set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
+            search_config = {
+                "quantization_formats": processed_kv_formats,
+                "data_loader": data_loader,
+                "forward_step": forward_step,
+                "num_calib_steps": num_calib_steps,
+                "num_score_steps": num_score_steps,
+                "disabled_layers": disabled_layers,
+                "verbose": verbose,
+                "checkpoint": checkpoint,
+            }
+            searcher.search(
+                model,
+                cast("ConstraintsDict", constraints),
+                config=search_config,
+            )
+            return model, searcher.state_dict()
+        except Exception:
+            searcher.restore_original_quantizers()
+            model.train(is_training)
+            if converted_for_search:
+                _restore_model_structure(conversion_snapshot)
+            raise
+
+    method = method or "gradient"
+
     if fixed_quantization_config is None and quantization_formats is None:
         quantization_formats = [mtq.NVFP4_AWQ_LITE_CFG, mtq.FP8_DEFAULT_CFG]
     elif fixed_quantization_config is not None and quantization_formats is None:
@@ -663,98 +761,6 @@ def _process_quantization_formats(formats, custom_name_prefix):
     return model, searcher.state_dict()
 
 
-def auto_quantize_kv_cache(
-    model: nn.Module,
-    constraints: dict[str, Any],
-    quantization_formats: list[dict[str, Any] | tuple[dict[str, Any], str]],
-    data_loader: Iterable,
-    forward_step: Callable[[nn.Module, Any], torch.Tensor],
-    *,
-    num_calib_steps: int = 512,
-    num_score_steps: int = 128,
-    disabled_layers: list[str] | str | None = None,
-    verbose: bool = False,
-    checkpoint: str | None = None,
-) -> tuple[nn.Module, dict[str, Any]]:
-    """Search layer-wise KV-cache formats using isolated forward KL sensitivity.
-
-    Unlike weight AutoQuant, BF16/no-quant is only the scoring reference and is
-    not an implicit solver choice. Each supplied format must configure K and V
-    together and declare config-level ``effective_bits`` equal to their packed
-    storage cost per K-or-V scalar, including scale overhead. The budget is weighted
-    by the K/V projection widths of eligible layers. Formats use their own calibration
-    algorithm; candidates with persistent ``constant_amax`` skip calibration forwards.
-
-    Args:
-        model: Model whose attention K/V quantizers will be searched.
-        constraints: A ``{"kv_effective_bits": target}`` storage constraint across
-            eligible layers.
-        quantization_formats: Candidate ``QuantizeConfig`` dictionaries, optionally
-            paired with display names.
-        data_loader: Re-iterable calibration and scoring batches.
-        forward_step: Callable returning full-vocabulary logits for the token positions
-            to score in one batch.
-        num_calib_steps: Maximum calibration batches per candidate.
-        num_score_steps: Maximum batches used for isolated forward-KL scoring.
-        disabled_layers: Optional layer-name patterns excluded from the search and
-            bit budget, and preserved in their existing KV-cache format.
-        verbose: Whether to print progress and selected formats.
-        checkpoint: Optional path for resumable calibration and sensitivity state.
-
-    Returns:
-        The converted model with the selected per-layer K/V quantizers and a
-        JSON-safe sensitivity report.
-    """
-    if (
-        torch.distributed.is_available()
-        and torch.distributed.is_initialized()
-        and torch.distributed.get_world_size() > 1
-    ):
-        raise RuntimeError(
-            "auto_quantize_kv_cache is single-process only; distributed scoring, selection, "
-            "and checkpoint writes are not synchronized."
-        )
-
-    processed_formats = []
-    for idx, candidate in enumerate(quantization_formats):
-        if isinstance(candidate, tuple):
-            raw_config, name = candidate
-        else:
-            raw_config, name = candidate, f"KV_CACHE_FORMAT_{idx}"
-        if not isinstance(raw_config, dict):
-            raise TypeError("KV-cache AutoQuant formats must be config dictionaries.")
-        if not isinstance(name, str) or not name:
-            raise ValueError("KV-cache AutoQuant candidate names must be non-empty strings.")
-        processed_formats.append((raw_config, name))
-
-    _validate_kv_cache_search_inputs(
-        constraints, processed_formats, num_calib_steps, num_score_steps
-    )
-
-    converted_for_search = not is_quantized(model)
-    conversion_snapshot = _snapshot_model_structure(model) if converted_for_search else []
-    try:
-        if converted_for_search:
-            model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
-            set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
-        return _auto_quantize_kv_cache(
-            model,
-            constraints,
-            processed_formats,
-            data_loader,
-            forward_step,
-            num_calib_steps=num_calib_steps,
-            num_score_steps=num_score_steps,
-            disabled_layers=disabled_layers,
-            verbose=verbose,
-            checkpoint=checkpoint,
-        )
-    except Exception:
-        if converted_for_search:
-            _restore_model_structure(conversion_snapshot)
-        raise
-
-
 def _snapshot_model_structure(
     model: nn.Module,
 ) -> list[tuple[nn.Module, type[nn.Module], dict[str, Any]]]:
@@ -764,7 +770,7 @@ def _snapshot_model_structure(
             module,
             type(module),
             {
-                key: copy.copy(value) if isinstance(value, (dict, list, set)) else value
+                key: copy.copy(value) if isinstance(value, dict | list | set) else value
                 for key, value in module.__dict__.items()
             },
         )
@@ -815,6 +821,8 @@ def get_auto_quantize_config(search_state, constraints=None, verbose=False):
             # fresh_model = load_model(...)
             # fresh_model = mtq.quantize(fresh_model, config, forward_loop=calibrate_loop)
     """
+    if search_state.get("cost_model") == COST_MODEL_KV_CACHE:
+        return get_kv_cache_auto_quantize_config(search_state, constraints, verbose=verbose)
     return _get_auto_quantize_config(search_state, constraints, verbose=verbose)
 
 
diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
index ba191e87bc8..28fb8262c08 100644
--- a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
+++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -17,7 +17,8 @@ metadata:
 
 auto_quantize:
   constraints:
-    kv_effective_bits: 5.4
+    effective_bits: 5.4
+    cost_model: kv_cache
 
   candidate_formats:
     - quant_cfg:
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index a3db0be8b67..761d55a386f 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -108,7 +108,7 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
     inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args)
 
     assert inputs["search_domain"] == "kv_cache"
-    assert inputs["constraints"] == {"kv_effective_bits": 5.4}
+    assert inputs["constraints"] == {"effective_bits": 5.4, "cost_model": "kv_cache"}
     assert inputs["method"] == "kl_div"
     assert [config["effective_bits"] for config, _ in inputs["quantization_formats"]] == [
         8.0,
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index 1a4fabb627b..cacc6bcd520 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -30,6 +30,7 @@
 from modelopt.recipe.config import (
     AutoQuantizeConfig,
     AutoQuantizeConstraints,
+    AutoQuantizeCost,
     ModelOptAutoQuantizeRecipe,
     ModelOptDFlashRecipe,
     ModelOptEagleRecipe,
@@ -1770,13 +1771,20 @@ def test_load_recipe_autoquantize_minimal(tmp_path):
     assert aq.module_search_spaces == []
 
 
-def test_autoquantize_constraints_preserve_default_with_kv_override():
+def test_autoquantize_constraints_use_effective_bits_for_kv_cost_model():
     assert AutoQuantizeConstraints.model_fields["effective_bits"].default == 4.8
     assert AutoQuantizeConstraints().effective_bits == 4.8
 
-    constraints = AutoQuantizeConstraints(kv_effective_bits=5.4)
-    assert constraints.effective_bits is None
-    assert constraints.kv_effective_bits == 5.4
+    constraints = AutoQuantizeConstraints(effective_bits=5.4, cost_model="kv_cache")
+    assert constraints.effective_bits == 5.4
+    assert constraints.cost_model == "kv_cache"
+
+    with pytest.raises(ValueError, match="does not accept weight cost settings"):
+        AutoQuantizeConstraints(
+            effective_bits=5.4,
+            cost_model="kv_cache",
+            cost=AutoQuantizeCost(active_moe_expert_ratio=0.5),
+        )
 
 
 def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path):
@@ -1931,7 +1939,7 @@ def test_load_recipe_autoquantize_builtin_general(recipe_path):
     assert len(recipe.auto_quantize.candidate_formats) >= 2
     assert recipe.auto_quantize.auto_quantize_method in ("gradient", "kl_div")
     assert "*output_layer*" in recipe.auto_quantize.disabled_layers
-    if recipe.auto_quantize.constraints.kv_effective_bits is not None:
+    if recipe.auto_quantize.constraints.cost_model == "kv_cache":
         assert "*mtp*" in recipe.auto_quantize.disabled_layers
         assert recipe.auto_quantize.cost_excluded_layers == []
     else:
@@ -1946,8 +1954,8 @@ def test_load_recipe_kv_autoquantize_contract():
     recipe = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits")
     aq = recipe.auto_quantize
 
-    assert aq.constraints.effective_bits is None
-    assert aq.constraints.kv_effective_bits == 5.4
+    assert aq.constraints.effective_bits == 5.4
+    assert aq.constraints.cost_model == "kv_cache"
     assert aq.auto_quantize_method == "kl_div"
     assert "*mtp*" in aq.disabled_layers
     assert aq.cost_excluded_layers == []
@@ -1964,7 +1972,7 @@ def test_load_recipe_kv_autoquantize_contract():
 def test_kv_autoquantize_rejects_cost_excluded_layers():
     with pytest.raises(ValueError, match=r"cost_excluded_layers.*disabled_layers"):
         AutoQuantizeConfig(
-            constraints=AutoQuantizeConstraints(kv_effective_bits=8.0),
+            constraints=AutoQuantizeConstraints(effective_bits=8.0, cost_model="kv_cache"),
             candidate_formats=[qcfg.QuantizeConfig(quant_cfg=[], effective_bits=8.0)],
             auto_quantize_method="kl_div",
             cost_excluded_layers=["*mtp*"],
diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
index f93ce8913af..7df6c2764bd 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -22,15 +22,16 @@
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.quant_utils import get_kv_cache_dtype, get_quant_config
+from modelopt.torch.opt.searcher import BaseSearcher
 from modelopt.torch.quantization import model_quant, tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.quantization.kv_cache_auto_quant import (
+    AutoQuantizeKVSearcher,
     _candidate_quantizers,
     _eligible_layers,
     _kv_scalar_weight,
     _solve_additive_recipe,
     _validate_kv_only_config,
-    auto_quantize_kv_cache,
 )
 from modelopt.torch.quantization.nn import TensorQuantizer
 
@@ -113,6 +114,10 @@ def test_kv_candidate_requires_exact_bits_and_both_sides():
         )
 
 
+def test_kv_autoquant_uses_shared_searcher_lifecycle():
+    assert issubclass(AutoQuantizeKVSearcher, BaseSearcher)
+
+
 @pytest.mark.parametrize("algorithm", ["svdquant", {"method": "smoothquant"}, {"method": "mse"}])
 def test_kv_candidate_rejects_structural_or_unscoped_algorithms(algorithm):
     config = _kv_config((4, 3), 8.0).model_copy(update={"algorithm": algorithm})
@@ -124,9 +129,9 @@ def test_kv_candidate_rejects_structural_or_unscoped_algorithms(algorithm):
 def test_kv_additive_solver_spends_fp8_on_more_sensitive_layer():
     selections, status = _solve_additive_recipe(
         layer_names=["layer0", "layer1"],
-        scalar_weights=[256, 256],
+        layer_widths=[(128, 128), (128, 128)],
         candidate_names=["fp8", "nvfp4"],
-        candidate_bits=[8.0, 4.5],
+        candidate_kv_bits=[(8.0, 8.0), (4.5, 4.5)],
         scores=[[0.0, 10.0], [0.0, 1.0]],
         target_bits=6.25,
         verbose=False,
@@ -239,9 +244,9 @@ def test_kv_autoquant_rejects_missing_scale_after_calibration(monkeypatch):
     monkeypatch.setattr(model_quant, "calibrate", lambda *_args, **_kwargs: None)
 
     with pytest.raises(ValueError, match="no persistent export scale after calibration"):
-        auto_quantize_kv_cache(
+        mtq.auto_quantize(
             model,
-            {"kv_effective_bits": 8.0},
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
             [(_kv_config((4, 3), 8.0).model_dump(), "fp8")],
             [torch.randn(1, 2, 8)],
             lambda search_model, batch: search_model(batch),
@@ -318,9 +323,9 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fa
         ),
     ]
 
-    model, state = auto_quantize_kv_cache(
+    model, state = mtq.auto_quantize(
         model,
-        {"kv_effective_bits": 6.25},
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
         candidates,
         data,
         lambda model, batch: model(batch),
@@ -329,7 +334,7 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fa
         checkpoint=str(tmp_path / "kv_search.pth"),
     )
 
-    assert state["best"]["effective_bits"] == pytest.approx(6.25)
+    assert state["best"]["constraints"]["effective_bits"] == pytest.approx(6.25)
     assert state["best"]["is_satisfied"]
     assert model.training
     assert {layer["selected"] for layer in state["layers"].values()} == {
@@ -343,9 +348,9 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fa
         assert layer.k_bmm_quantizer.num_bits == expected_bits
 
     restored_model = _ToyKVModel().eval()
-    restored_model, restored_state = auto_quantize_kv_cache(
+    restored_model, restored_state = mtq.auto_quantize(
         restored_model,
-        {"kv_effective_bits": 6.25},
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
         candidates,
         data,
         lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."),
@@ -354,21 +359,45 @@ def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fa
         checkpoint=str(tmp_path / "kv_search.pth"),
     )
 
-    assert restored_state == state
+    assert restored_state["best"] == state["best"]
+    assert restored_state["layers"] == state["layers"]
     assert not restored_model.training
     for layer_name, layer_state in restored_state["layers"].items():
         layer = restored_model.get_submodule(layer_name)
         expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1)
         assert layer.k_bmm_quantizer.num_bits == expected_bits
 
+    resolved_config = mtq.get_auto_quantize_config(state, {"effective_bits": 4.5})
+    assert resolved_config["algorithm"] == "max"
+    assert {entry["quantizer_name"] for entry in resolved_config["quant_cfg"]} == {
+        "attn0.k_bmm_quantizer",
+        "attn0.v_bmm_quantizer",
+        "attn1.k_bmm_quantizer",
+        "attn1.v_bmm_quantizer",
+    }
+
+    re_solved_model = _ToyKVModel()
+    _, re_solved_state = mtq.auto_quantize(
+        re_solved_model,
+        {"effective_bits": 4.5, "cost_model": "kv_cache"},
+        candidates,
+        data,
+        lambda *_: pytest.fail("Changing only the budget must re-solve without re-scoring."),
+        num_calib_steps=2,
+        num_score_steps=2,
+        checkpoint=str(tmp_path / "kv_search.pth"),
+    )
+    assert re_solved_state["best"]["constraints"]["effective_bits"] == pytest.approx(4.5)
+    assert {layer["selected"] for layer in re_solved_state["layers"].values()} == {"nvfp4"}
+
 
 def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quant_stub):
     model = _ToyKVModel()
     candidate = (_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v")
 
-    model, state = auto_quantize_kv_cache(
+    model, state = mtq.auto_quantize(
         model,
-        {"kv_effective_bits": 6.25},
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
         [candidate],
         [torch.randn(1, 2, 8)],
         lambda search_model, batch: search_model(batch),
@@ -377,7 +406,7 @@ def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quan
     )
 
     assert state["candidates"][0]["effective_bits"] == pytest.approx(6.25)
-    assert state["best"]["effective_bits"] == pytest.approx(6.25)
+    assert state["best"]["constraints"]["effective_bits"] == pytest.approx(6.25)
     for layer_state in state["layers"].values():
         assert layer_state["selected"] == "fp8_k_nvfp4_v"
     for layer in (model.attn0, model.attn1):
@@ -394,9 +423,9 @@ def test_kv_autoquant_rejects_asymmetric_candidate_for_unequal_kv_widths(
     model.attn0.v_proj = nn.Linear(8, 8, bias=False)
 
     with pytest.raises(ValueError, match=r"asymmetric K/V candidates.*unequal K/V widths"):
-        auto_quantize_kv_cache(
+        mtq.auto_quantize(
             model,
-            {"kv_effective_bits": 6.25},
+            {"effective_bits": 6.25, "cost_model": "kv_cache"},
             [(_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v")],
             [torch.randn(1, 2, 8)],
             lambda search_model, batch: search_model(batch),
@@ -436,9 +465,9 @@ def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state():
     ]
 
     with pytest.raises(ValueError, match="non-empty vocabulary dimension"):
-        auto_quantize_kv_cache(
+        mtq.auto_quantize(
             model,
-            {"kv_effective_bits": 4.5},
+            {"effective_bits": 4.5, "cost_model": "kv_cache"},
             candidates,
             [torch.randn(2, 3, 8)],
             lambda *_: torch.ones(8),
@@ -492,9 +521,9 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_
         ),
     ]
 
-    model, state = mtq.auto_quantize_kv_cache(
+    model, state = mtq.auto_quantize(
         model,
-        {"kv_effective_bits": 6.25},
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
         candidates,
         data,
         lambda search_model, batch: search_model(**batch).logits,
@@ -504,7 +533,7 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_
     )
 
     assert len(state["layers"]) == model.config.num_hidden_layers
-    assert state["best"]["effective_bits"] == pytest.approx(6.25)
+    assert state["best"]["constraints"]["effective_bits"] == pytest.approx(6.25)
     assert all(
         layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled
         for layer in model.model.layers
@@ -528,9 +557,9 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_
     }
 
     restored_model = get_tiny_llama(num_hidden_layers=2)
-    restored_model, restored_state = mtq.auto_quantize_kv_cache(
+    restored_model, restored_state = mtq.auto_quantize(
         restored_model,
-        {"kv_effective_bits": 6.25},
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
         candidates,
         data,
         lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."),
@@ -539,7 +568,8 @@ def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_
         checkpoint=str(tmp_path / "hf_kv_search.pth"),
     )
 
-    assert restored_state == state
+    assert restored_state["best"] == state["best"]
+    assert restored_state["layers"] == state["layers"]
     assert any(
         hasattr(layer.self_attn.k_bmm_quantizer, "_amax")
         for layer in restored_model.model.layers
@@ -566,9 +596,9 @@ def test_public_kv_autoquant_selects_qwen_causal_attention_only(
         "fp8",
     )
 
-    model, state = mtq.auto_quantize_kv_cache(
+    model, state = mtq.auto_quantize(
         model,
-        {"kv_effective_bits": 8.0},
+        {"effective_bits": 8.0, "cost_model": "kv_cache"},
         [candidate],
         data,
         lambda search_model, batch: search_model(**batch).logits,
@@ -578,7 +608,8 @@ def test_public_kv_autoquant_selects_qwen_causal_attention_only(
     )
 
     assert set(state["layers"]) == {expected_layer}
-    assert json.loads(json.dumps(state))["layers"][expected_layer]["selected"] == "fp8"
+    report = model._modelopt_kv_cache_auto_quantize_state
+    assert json.loads(json.dumps(report))["layers"][expected_layer]["selected"] == "fp8"
     exported = get_quant_config(model)["quantization"]
     if "kv_cache_quantized_layers" in exported:
         assert set(exported["kv_cache_quantized_layers"]) == {expected_layer}
@@ -593,9 +624,9 @@ def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
     invalid_candidate["algorithm"] = "svdquant"
 
     with pytest.raises(ValueError, match="only non-structural calibration algorithms"):
-        mtq.auto_quantize_kv_cache(
+        mtq.auto_quantize(
             model,
-            {"kv_effective_bits": 8.0},
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
             [invalid_candidate],
             [],
             lambda *_: pytest.fail("Validation must run before model conversion."),
@@ -609,9 +640,9 @@ def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic():
     valid_candidate = _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump()
     data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
     with pytest.raises(ValueError, match="non-empty vocabulary dimension"):
-        mtq.auto_quantize_kv_cache(
+        mtq.auto_quantize(
             model,
-            {"kv_effective_bits": 8.0},
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
             [valid_candidate],
             data,
             lambda *_: torch.ones(8),
@@ -639,9 +670,9 @@ def test_public_kv_autoquant_rejects_unmatched_or_unexportable_candidates_before
 
     for candidate, match in invalid_candidates:
         with pytest.raises(ValueError, match=match):
-            mtq.auto_quantize_kv_cache(
+            mtq.auto_quantize(
                 model,
-                {"kv_effective_bits": 8.0},
+                {"effective_bits": 8.0, "cost_model": "kv_cache"},
                 [candidate],
                 [],
                 lambda *_: pytest.fail("Validation must run before model conversion."),
@@ -660,9 +691,9 @@ def test_public_kv_autoquant_rejects_distributed_execution_before_mutation(monke
     monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
 
     with pytest.raises(RuntimeError, match="single-process only"):
-        mtq.auto_quantize_kv_cache(
+        mtq.auto_quantize(
             model,
-            {"kv_effective_bits": 8.0},
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
             [],
             [],
             lambda *_: pytest.fail("Distributed validation must fail before search."),
@@ -724,9 +755,9 @@ def calibrate_with_state_check(*args, **kwargs):
     monkeypatch.setattr(model_quant, "calibrate", calibrate_with_state_check)
 
     try:
-        model, state = mtq.auto_quantize_kv_cache(
+        model, state = mtq.auto_quantize(
             model,
-            {"kv_effective_bits": 4.5},
+            {"effective_bits": 4.5, "cost_model": "kv_cache"},
             [
                 (
                     {