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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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 452328dcec1a8054c2e8cc7589e308683d0a3584 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:05:12 -0700
Subject: [PATCH 12/12] Compose GEMM and KV AutoQuantize recipes

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |  18 ++-
 examples/hf_ptq/hf_ptq.py                     | 149 ++++++++++++++----
 modelopt/recipe/config.py                     |  67 +++++++-
 modelopt/torch/export/quant_utils.py          |  12 +-
 ...n_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml |  45 ++++++
 ...n_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml |  55 +++++++
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 111 ++++++++++++-
 tests/unit/recipe/test_loader.py              |  91 +++++++++++
 .../torch/export/test_get_quantization.py     |  44 ++++++
 9 files changed, 541 insertions(+), 51 deletions(-)
 create mode 100644 modelopt_recipes/general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
 create mode 100644 modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index d77dbd520f4..fa501489b96 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -425,6 +425,19 @@ 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.
 
+To optimize GEMM and KV cache in one invocation, compose ordered stages in the same recipe. A fixed
+`quantize` block followed by a KV-domain `auto_quantize` first calibrates the GEMM weight/activation
+configuration, then searches K/V while the existing GEMM QDQ remains enabled with calibration
+frozen. See `general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits`.
+
+A weight-domain `auto_quantize` can instead add a `kv_auto_quantize` follow-up with its own method,
+constraints, candidates, score size, and disabled layers. This supports, for example, a
+gradient-based GEMM search followed by a KL-divergence KV search; see
+`general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits`. When the
+follow-up is present, the recipe owns KV configuration and suppresses the CLI's uniform
+`--kv_cache_qformat` fallback. Use `--auto_quantize_checkpoint` for the weight search and
+`--kv_auto_quantize_checkpoint` for the KV search.
+
 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
@@ -454,8 +467,9 @@ the JSON-safe sensitivity report to `kv_cache_auto_quantize_report.json`;
 > 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):
+For a single-stage search, `--auto_quantize_checkpoint` saves/restores the search state to resume an
+interrupted search (skips re-scoring). Composed weight-plus-KV recipes additionally use
+`--kv_auto_quantize_checkpoint` for the independent KV search state:
 
 ```bash
 scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_5p4bits \
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 6daade6ed24..ef660110b25 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -386,7 +386,10 @@ def _mtq_kv_candidate_formats(formats) -> list[tuple[dict, str]]:
 
 
 def _mtq_inputs_from_auto_quantize_config(
-    aq_config, args: argparse.Namespace, fixed_quantize_config=None
+    aq_config,
+    args: argparse.Namespace,
+    fixed_quantize_config=None,
+    allow_uniform_kv: bool = True,
 ) -> dict:
     """Map a resolved AutoQuantizeConfig to mtq.auto_quantize inputs.
 
@@ -413,7 +416,9 @@ def _mtq_inputs_from_auto_quantize_config(
         constraints.setdefault("cost", {})["excluded_module_name_patterns"] = (
             aq_config.cost_excluded_layers
         )
-    if aq_config.kv_cache is not None:
+    if not allow_uniform_kv:
+        kv_cache_quant_cfg = None
+    elif aq_config.kv_cache is not None:
         kv_cache_quant_cfg = aq_config.kv_cache.model_dump()
     elif args.kv_cache_qformat == KV_CACHE_NONE:
         kv_cache_quant_cfg = None
@@ -456,6 +461,8 @@ def auto_quantize(
     aq_config,
     full_model: torch.nn.Module | None = None,
     fixed_quantize_config=None,
+    allow_uniform_kv: bool = True,
+    checkpoint_attr: str = "auto_quantize_checkpoint",
 ):
     """Recipe-driven auto_quantize, organized around an AutoQuantizeConfig.
 
@@ -475,8 +482,12 @@ def auto_quantize(
         raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR)
 
     inputs = _mtq_inputs_from_auto_quantize_config(
-        aq_config, args, fixed_quantize_config=fixed_quantize_config
+        aq_config,
+        args,
+        fixed_quantize_config=fixed_quantize_config,
+        allow_uniform_kv=allow_uniform_kv,
     )
+    checkpoint = getattr(args, checkpoint_attr, None)
 
     # base-model lm_head handling (mirrors the CLI helper)
     is_base_model = (
@@ -538,7 +549,7 @@ def forward_step(model, batch):
             ),
             verbose=True,
             disabled_layers=inputs["disabled_layers"],
-            checkpoint=args.auto_quantize_checkpoint,
+            checkpoint=checkpoint,
         )
         return language_model
 
@@ -556,7 +567,7 @@ def forward_step(model, batch):
         verbose=True,
         disabled_layers=inputs["disabled_layers"],
         method=inputs["method"],
-        checkpoint=args.auto_quantize_checkpoint,
+        checkpoint=checkpoint,
     )
 
     # KV cache quantization is uniform; applied after the LP search.
@@ -843,6 +854,81 @@ def mono_quantize(
         warnings.warn("Skipping quantization: model is already quantized.")
 
 
+def _prepare_quant_cfg(
+    args: argparse.Namespace, quant_cfg: dict[str, Any], full_model: torch.nn.Module
+) -> dict[str, Any]:
+    """Apply shared checkpoint-local adjustments to a PTQ configuration."""
+    mtp_layer_prefixes = getattr(full_model, "_mtp_layer_prefixes", None)
+    if mtp_layer_prefixes:
+        quant_cfg = copy.deepcopy(quant_cfg)
+        for prefix in mtp_layer_prefixes:
+            pattern = f"*{prefix}*"
+            quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False})
+            print(f"Excluding MTP layer from quantization: {pattern}")
+
+    if needs_checkpoint_path_update(quant_cfg):
+        quant_cfg, resolved_dir = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path)
+        print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}")
+
+    if args.cast_mxfp4_to_nvfp4:
+        quant_cfg = copy.deepcopy(quant_cfg)
+        force_weight_quantizers_static(quant_cfg["quant_cfg"])
+    return quant_cfg
+
+
+def _run_auto_quantize_recipe(
+    args: argparse.Namespace,
+    recipe: ModelOptAutoQuantizeRecipe,
+    full_model: torch.nn.Module,
+    language_model: torch.nn.Module,
+    model_type: str | None,
+    calibration_only: bool,
+    calib_dataloader: DataLoader,
+    is_nemotron_vl_model: bool,
+) -> None:
+    """Run the recipe's fixed PTQ, weight search, and KV search in order."""
+    primary = recipe.auto_quantize
+    followup_kv = recipe.kv_auto_quantize
+    primary_is_kv = primary.constraints.kv_effective_bits is not None
+    fixed_quantize_config = recipe.quantize
+
+    if primary_is_kv and fixed_quantize_config is not None:
+        quant_cfg = _prepare_quant_cfg(args, fixed_quantize_config.model_dump(), full_model)
+        mono_quantize(
+            args,
+            quant_cfg,
+            full_model,
+            language_model,
+            model_type,
+            calibration_only,
+            calib_dataloader,
+            is_nemotron_vl_model,
+        )
+        fixed_quantize_config = None
+
+    auto_quantize(
+        args,
+        full_model,
+        calib_dataloader,
+        aq_config=primary,
+        full_model=full_model,
+        fixed_quantize_config=fixed_quantize_config,
+        allow_uniform_kv=followup_kv is None,
+        checkpoint_attr="auto_quantize_checkpoint",
+    )
+
+    if followup_kv is not None:
+        auto_quantize(
+            args,
+            full_model,
+            calib_dataloader,
+            aq_config=followup_kv,
+            full_model=full_model,
+            allow_uniform_kv=False,
+            checkpoint_attr="kv_auto_quantize_checkpoint",
+        )
+
+
 def export_quantized(
     args: argparse.Namespace,
     full_model: torch.nn.Module,
@@ -1198,10 +1284,8 @@ def quantize_main(
     # AutoQuantize is recipe-driven: everything downstream reads the resolved AutoQuantizeConfig.
     if isinstance(recipe, ModelOptAutoQuantizeRecipe):
         aq_config = recipe.auto_quantize
-        fixed_quantize_config = recipe.quantize
     else:
         aq_config = None
-        fixed_quantize_config = None
 
     def _is_layerwise(obj):
         if isinstance(obj, ModelOptPTQRecipe):
@@ -1277,7 +1361,11 @@ def _is_layerwise(obj):
         device,
         model_type,
         autoquant_gradient_recipe=(
-            aq_config is not None and aq_config.auto_quantize_method == "gradient"
+            isinstance(recipe, ModelOptAutoQuantizeRecipe)
+            and any(
+                config is not None and config.auto_quantize_method == "gradient"
+                for config in (recipe.auto_quantize, recipe.kv_auto_quantize)
+            )
         ),
     )
 
@@ -1289,16 +1377,16 @@ def _is_layerwise(obj):
     )
 
     if aq_config is not None:
-        # AutoQuantize (recipe-driven). For VL models the search walks the OUTER CausalLM (which
-        # carries lm_head and the LM-head forward path); architecture-specific exclusions come
-        # from aq_config.disabled_layers.
-        auto_quantize(
+        assert isinstance(recipe, ModelOptAutoQuantizeRecipe)
+        _run_auto_quantize_recipe(
             args,
+            recipe,
             full_model,
+            language_model,
+            model_type,
+            calibration_only,
             calib_dataloader,
-            aq_config,
-            full_model=full_model,
-            fixed_quantize_config=fixed_quantize_config,
+            is_nemotron_vl_model,
         )
 
     else:
@@ -1333,25 +1421,7 @@ def _is_layerwise(obj):
                     KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"],
                 )
 
-        # Exclude MTP layers from quantization if detected (e.g., GLM-4.7's layer 92).
-        # These layers are typically speculative decoding layers that should be exported as-is.
-        # Complementary to recipe `*mtp*` wildcards (name-match); this catches MTP layers
-        # identified by index.
-        mtp_layer_prefixes = getattr(full_model, "_mtp_layer_prefixes", None)
-        if mtp_layer_prefixes:
-            quant_cfg = copy.deepcopy(quant_cfg)
-            for prefix in mtp_layer_prefixes:
-                pattern = f"*{prefix}*"
-                quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False})
-                print(f"Excluding MTP layer from quantization: {pattern}")
-
-        if needs_checkpoint_path_update(quant_cfg):
-            quant_cfg, resolved_dir = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path)
-            print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}")
-
-        if args.cast_mxfp4_to_nvfp4:
-            quant_cfg = copy.deepcopy(quant_cfg)
-            force_weight_quantizers_static(quant_cfg["quant_cfg"])
+        quant_cfg = _prepare_quant_cfg(args, quant_cfg, full_model)
 
         if quant_cfg:
             mono_quantize(
@@ -1417,7 +1487,7 @@ def parse_args() -> argparse.Namespace:
             "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."
+            "unless the recipe sets an explicit kv_cache or kv_auto_quantize field."
         ),
         default=None,
     )
@@ -1591,6 +1661,15 @@ def parse_args() -> argparse.Namespace:
             "(sensitivity scores, costs, etc.). Used with an AutoQuantize --recipe."
         ),
     )
+    parser.add_argument(
+        "--kv_auto_quantize_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Path to the independent KV-cache search checkpoint when a recipe runs weight "
+            "AutoQuantize followed by kv_auto_quantize."
+        ),
+    )
     parser.add_argument(
         "--moe_calib_experts_ratio",
         type=float,
diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index 19e9ef7278b..6f233ca7583 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -19,6 +19,7 @@
 
 import warnings
 from enum import Enum
+from fnmatch import fnmatch
 from typing import Literal
 
 from pydantic import Field, field_validator, model_validator
@@ -325,6 +326,17 @@ def _has_search_space(self):
         return self
 
 
+def _quantize_config_enables_kv(config: QuantizeConfig) -> bool:
+    """Return whether ordered quantizer rules leave either K/V quantizer enabled."""
+    probe_names = ("layer.k_bmm_quantizer", "layer.v_bmm_quantizer")
+    enabled = dict.fromkeys(probe_names, False)
+    for entry in config.quant_cfg:
+        for name in probe_names:
+            if fnmatch(name, entry.quantizer_name):
+                enabled[name] = entry.enable
+    return any(enabled.values())
+
+
 class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase):
     """Our config class for AutoQuantize recipes."""
 
@@ -333,9 +345,9 @@ class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase):
     quantize: QuantizeConfig | None = ModeloptField(
         default=None,
         title="Fixed PTQ baseline",
-        description="Optional normal PTQ QuantizeConfig for modules outside the explicit "
-        "AutoQuantize module_search_spaces. Fixed and searched modules are calibrated, scored, "
-        "costed, and exported in one integrated AutoQuantize operation.",
+        description="Optional normal PTQ QuantizeConfig. A weight AutoQuantize stage uses it for "
+        "modules outside explicit module_search_spaces; a KV AutoQuantize stage applies it first "
+        "as the fixed GEMM weight/activation configuration.",
     )
 
     auto_quantize: AutoQuantizeConfig = Field(
@@ -343,26 +355,69 @@ class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase):
         description="AutoQuantize search configuration. Required.",
     )
 
+    kv_auto_quantize: AutoQuantizeConfig | None = ModeloptField(
+        default=None,
+        title="Follow-up KV-cache AutoQuantize config",
+        description="Optional KV-cache search run after the primary weight AutoQuantize search.",
+    )
+
     @model_validator(mode="after")
     def _validate_fixed_and_searched_spaces(self):
+        primary_is_kv = self.auto_quantize.constraints.kv_effective_bits is not None
+        if self.kv_auto_quantize is not None:
+            if primary_is_kv:
+                raise ValueError(
+                    "kv_auto_quantize cannot follow an auto_quantize stage that already searches "
+                    "the KV cache."
+                )
+            if self.kv_auto_quantize.constraints.kv_effective_bits is None:
+                raise ValueError("kv_auto_quantize must use a kv_effective_bits constraint.")
+            if self.auto_quantize.kv_cache is not None:
+                raise ValueError(
+                    "A weight AutoQuantize stage followed by kv_auto_quantize must omit the "
+                    "uniform auto_quantize.kv_cache post-step."
+                )
+            first_stage_candidates = [
+                *self.auto_quantize.candidate_formats,
+                *(
+                    candidate
+                    for search_space in self.auto_quantize.module_search_spaces
+                    for candidate in search_space.candidate_formats
+                ),
+            ]
+            if any(_quantize_config_enables_kv(config) for config in first_stage_candidates):
+                raise ValueError(
+                    "The weight AutoQuantize stage must not enable K/V quantizers when a "
+                    "kv_auto_quantize follow-up is configured."
+                )
+
         has_fixed_baseline = self.quantize is not None
         has_global_search = bool(self.auto_quantize.candidate_formats)
-        if has_fixed_baseline and has_global_search:
+        if not primary_is_kv and has_fixed_baseline and has_global_search:
             raise ValueError(
                 "An AutoQuantize recipe with a fixed quantize baseline must omit top-level "
                 "auto_quantize.candidate_formats and explicitly list searched modules under "
                 "auto_quantize.module_search_spaces."
             )
-        if has_fixed_baseline and not self.auto_quantize.module_search_spaces:
+        if not primary_is_kv and has_fixed_baseline and not self.auto_quantize.module_search_spaces:
             raise ValueError(
                 "An AutoQuantize recipe with a fixed quantize baseline requires at least one "
                 "auto_quantize.module_search_spaces entry."
             )
-        if not has_fixed_baseline and not has_global_search:
+        if not primary_is_kv and not has_fixed_baseline and not has_global_search:
             raise ValueError(
                 "An AutoQuantize recipe without a fixed quantize baseline requires top-level "
                 "auto_quantize.candidate_formats for unmatched modules."
             )
+        if (
+            (primary_is_kv or self.kv_auto_quantize is not None)
+            and self.quantize is not None
+            and _quantize_config_enables_kv(self.quantize)
+        ):
+            raise ValueError(
+                "The fixed quantize stage must not enable K/V quantizers when a KV-cache "
+                "AutoQuantize stage is configured."
+            )
         return self
 
 
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 2d8f94366a6..8cee1c9ed3a 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -1774,13 +1774,11 @@ def get_quant_config(
         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", {})
+        if weight_quant_algo is None:
+            quant_config["quantization"]["quant_algo"] = "MIXED_PRECISION"
+            quant_config["quantization"]["quantized_layers"] = {}
+        elif weight_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
diff --git a/modelopt_recipes/general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
new file mode 100644
index 00000000000..228dd0c19d0
--- /dev/null
+++ b/modelopt_recipes/general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Fixed FP8 GEMM PTQ followed by layer-wise KV-cache AutoQuantize.
+
+# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  kv_fp8: configs/ptq/units/kv_fp8
+  kv_nvfp4: configs/ptq/units/kv_nvfp4
+  w8a8_fp8_fp8: configs/ptq/units/w8a8_fp8_fp8
+
+metadata:
+  recipe_type: auto_quantize
+  description: Fixed FP8 GEMM PTQ followed by mixed FP8/NVFP4 KV-cache search.
+
+quantize:
+  algorithm: max
+  quant_cfg:
+    - $import: base_disable_all
+    - $import: w8a8_fp8_fp8
+    - $import: default_disabled_quantizers
+
+auto_quantize:
+  constraints:
+    kv_effective_bits: 5.4
+
+  candidate_formats:
+    - quant_cfg:
+        - $import: kv_fp8
+      algorithm: max
+      effective_bits: 8.0
+    - quant_cfg:
+        - $import: kv_nvfp4
+      algorithm: max
+      effective_bits: 4.5
+
+  auto_quantize_method: kl_div
+  score_size: 128
+
+  disabled_layers:
+    - $import: base_disabled_layers
+    - "*mtp*"
diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
new file mode 100644
index 00000000000..073d4a9e14f
--- /dev/null
+++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Gradient-based GEMM AutoQuantize followed by layer-wise KV-cache AutoQuantize.
+
+# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
+imports:
+  base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers
+  base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
+  fp8: configs/ptq/presets/model/fp8
+  kv_fp8: configs/ptq/units/kv_fp8
+  kv_nvfp4: configs/ptq/units/kv_nvfp4
+  nvfp4: configs/ptq/presets/model/nvfp4
+
+metadata:
+  recipe_type: auto_quantize
+  description: Gradient GEMM search followed by KL-divergence mixed-KV search at 5.4 bits.
+
+auto_quantize:
+  constraints:
+    effective_bits: 5.4
+
+  candidate_formats:
+    - $import: nvfp4
+    - $import: fp8
+
+  auto_quantize_method: gradient
+  score_size: 128
+
+  disabled_layers:
+    - $import: base_disabled_layers
+
+  cost_excluded_layers:
+    - $import: base_cost_excluded_layers
+
+kv_auto_quantize:
+  constraints:
+    kv_effective_bits: 5.4
+
+  candidate_formats:
+    - quant_cfg:
+        - $import: kv_fp8
+      algorithm: max
+      effective_bits: 8.0
+    - quant_cfg:
+        - $import: kv_nvfp4
+      algorithm: max
+      effective_bits: 4.5
+
+  auto_quantize_method: kl_div
+  score_size: 128
+
+  disabled_layers:
+    - $import: base_disabled_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 2471bab4189..2101529d413 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -26,7 +26,11 @@
 from _test_utils.torch.transformers_models import get_tiny_qwen3
 
 from modelopt.recipe import load_recipe
-from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints
+from modelopt.recipe.config import (
+    AutoQuantizeConfig,
+    AutoQuantizeConstraints,
+    ModelOptAutoQuantizeRecipe,
+)
 from modelopt.recipe.presets import QUANT_CFG_CHOICES
 from modelopt.torch.quantization.config import QuantizeConfig
 
@@ -105,6 +109,17 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch):
     assert "kv_cache_quant_cfg" not in inputs
 
 
+def test_followup_kv_autoquant_suppresses_uniform_kv_fallback(monkeypatch):
+    hf_ptq, args = _parse_hf_ptq_args(
+        monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "fp8_cast"
+    )
+    aq = load_recipe("general/auto_quantize/nvfp4_fp8_at_5p4bits").auto_quantize
+
+    inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args, allow_uniform_kv=False)
+
+    assert inputs["kv_cache_quant_cfg"] is None
+
+
 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)
@@ -143,6 +158,100 @@ def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch):
     assert attention.v_bmm_quantizer.num_bits == (4, 3)
 
 
+def test_hf_ptq_runs_weight_then_kv_autoquantize_stages(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    weight_aq = AutoQuantizeConfig(
+        constraints=AutoQuantizeConstraints(effective_bits=8.0),
+        candidate_formats=[QuantizeConfig(**QUANT_CFG_CHOICES["fp8"])],
+        auto_quantize_method="gradient",
+    )
+    kv_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",
+    )
+    recipe = ModelOptAutoQuantizeRecipe(auto_quantize=weight_aq, kv_auto_quantize=kv_aq)
+    calls = []
+    monkeypatch.setattr(
+        hf_ptq,
+        "auto_quantize",
+        lambda *_args, **kwargs: calls.append(kwargs),
+    )
+
+    hf_ptq._run_auto_quantize_recipe(
+        SimpleNamespace(), recipe, torch.nn.Module(), torch.nn.Module(), None, False, [], False
+    )
+
+    assert [call["aq_config"] for call in calls] == [weight_aq, kv_aq]
+    assert calls[0]["allow_uniform_kv"] is False
+    assert calls[0]["checkpoint_attr"] == "auto_quantize_checkpoint"
+    assert calls[1]["checkpoint_attr"] == "kv_auto_quantize_checkpoint"
+
+
+def test_hf_ptq_runs_fixed_ptq_before_kv_autoquantize(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    fixed = QuantizeConfig(
+        quant_cfg=[
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*q_proj.weight_quantizer",
+                "cfg": {"num_bits": (4, 3), "axis": None, "constant_amax": 1.0},
+            },
+        ],
+        algorithm=None,
+    )
+    kv_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,
+    )
+    recipe = ModelOptAutoQuantizeRecipe(quantize=fixed, auto_quantize=kv_aq)
+    model = get_tiny_qwen3(num_hidden_layers=1)
+    data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
+    args = SimpleNamespace(
+        qformat="fp8",
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=False,
+        batch_size=1,
+        auto_quantize_checkpoint=None,
+        kv_auto_quantize_checkpoint=None,
+        pyt_ckpt_path="dummy",
+        cast_mxfp4_to_nvfp4=False,
+    )
+
+    hf_ptq._run_auto_quantize_recipe(args, recipe, model, model, None, False, data, False)
+
+    attention = model.model.layers[0].self_attn
+    assert attention.q_proj.weight_quantizer.is_enabled
+    assert attention.q_proj.weight_quantizer.num_bits == (4, 3)
+    assert attention.k_bmm_quantizer.is_enabled
+    assert attention.v_bmm_quantizer.is_enabled
+
+
 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)
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index 3c42e0bfafb..a084074eec8 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1916,8 +1916,10 @@ def test_load_recipe_autoquantize_fixed_baseline_requires_explicit_search(tmp_pa
 @pytest.mark.parametrize(
     "recipe_path",
     [
+        "general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits",
         "general/auto_quantize/nvfp4_fp8_at_5p4bits",
         "general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits",
+        "general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_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",
@@ -1960,6 +1962,95 @@ def test_load_recipe_kv_autoquantize_contract():
         assert fmt.algorithm == "max"
 
 
+def test_load_recipe_fixed_ptq_then_kv_autoquantize(tmp_path):
+    recipe_file = tmp_path / "ptq-then-kv.yml"
+    recipe_file.write_text(
+        "metadata:\n  recipe_type: auto_quantize\n"
+        "quantize:\n  algorithm: max\n  quant_cfg:\n"
+        "    - quantizer_name: '*'\n      enable: false\n"
+        "    - quantizer_name: '*.weight_quantizer'\n"
+        "      cfg: {num_bits: [4, 3], axis: null}\n"
+        "auto_quantize:\n  constraints:\n    kv_effective_bits: 8.0\n"
+        "  candidate_formats:\n"
+        "    - algorithm: null\n      effective_bits: 8.0\n      quant_cfg:\n"
+        "        - quantizer_name: '*[kv]_bmm_quantizer'\n"
+        "          cfg: {num_bits: [4, 3], constant_amax: 1.0}\n"
+        "  auto_quantize_method: kl_div\n"
+    )
+
+    recipe = load_recipe(recipe_file)
+
+    assert recipe.quantize is not None
+    assert recipe.auto_quantize.constraints.kv_effective_bits == 8.0
+    assert recipe.kv_auto_quantize is None
+
+
+def test_load_recipe_weight_autoquantize_then_kv_autoquantize(tmp_path):
+    recipe_file = tmp_path / "weight-then-kv.yml"
+    recipe_file.write_text(
+        _AQ_MINIMAL_BODY + "kv_auto_quantize:\n  constraints:\n    kv_effective_bits: 8.0\n"
+        "  candidate_formats:\n"
+        "    - algorithm: null\n      effective_bits: 8.0\n      quant_cfg:\n"
+        "        - quantizer_name: '*[kv]_bmm_quantizer'\n"
+        "          cfg: {num_bits: [4, 3], constant_amax: 1.0}\n"
+        "  auto_quantize_method: kl_div\n"
+    )
+
+    recipe = load_recipe(recipe_file)
+
+    assert recipe.auto_quantize.auto_quantize_method == "gradient"
+    assert recipe.auto_quantize.constraints.effective_bits == 4.8
+    assert recipe.kv_auto_quantize is not None
+    assert recipe.kv_auto_quantize.auto_quantize_method == "kl_div"
+    assert recipe.kv_auto_quantize.constraints.kv_effective_bits == 8.0
+
+
+def test_composed_kv_autoquantize_rejects_preconfigured_kv_quantizers(tmp_path):
+    recipe_file = tmp_path / "invalid.yml"
+    recipe_file.write_text(
+        "metadata:\n  recipe_type: auto_quantize\n"
+        "quantize:\n  algorithm: max\n  quant_cfg:\n"
+        "    - quantizer_name: '*[kv]_bmm_quantizer'\n"
+        "      cfg: {num_bits: [4, 3], constant_amax: 1.0}\n"
+        "auto_quantize:\n  constraints:\n    kv_effective_bits: 8.0\n"
+        "  candidate_formats:\n"
+        "    - algorithm: null\n      effective_bits: 8.0\n      quant_cfg:\n"
+        "        - quantizer_name: '*[kv]_bmm_quantizer'\n"
+        "          cfg: {num_bits: [4, 3], constant_amax: 1.0}\n"
+        "  auto_quantize_method: kl_div\n"
+    )
+
+    with pytest.raises(ValueError, match="must not enable K/V quantizers"):
+        load_recipe(recipe_file)
+
+
+def test_followup_kv_autoquantize_rejects_kv_weight_search_candidate():
+    kv_candidate = qcfg.QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+            }
+        ],
+        effective_bits=8.0,
+    )
+    kv_search = AutoQuantizeConfig(
+        constraints=AutoQuantizeConstraints(kv_effective_bits=8.0),
+        candidate_formats=[kv_candidate],
+        auto_quantize_method="kl_div",
+    )
+    weight_search = AutoQuantizeConfig(
+        constraints=AutoQuantizeConstraints(effective_bits=8.0),
+        candidate_formats=[kv_candidate],
+    )
+
+    with pytest.raises(ValueError, match="weight AutoQuantize stage must not enable K/V"):
+        ModelOptAutoQuantizeRecipe(
+            auto_quantize=weight_search,
+            kv_auto_quantize=kv_search,
+        )
+
+
 def test_kv_autoquantize_rejects_cost_excluded_layers():
     with pytest.raises(ValueError, match=r"cost_excluded_layers.*disabled_layers"):
         AutoQuantizeConfig(
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index 1202b1fcae4..0fe5636f1b8 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -161,6 +161,50 @@ def __init__(self):
     }
 
 
+def test_uniform_weight_quantization_exports_mixed_kv_cache_map():
+    class FakeAttention(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.k_bmm_quantizer = TensorQuantizer()
+            self.v_bmm_quantizer = TensorQuantizer()
+
+    model = ToyModel()
+    mtq.quantize(model, partial_fp8_config, lambda x: x(torch.randn(1, 4, 10)))
+    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"] == "FP8"
+    assert quantization["kv_cache_quant_algo"] == "MIXED_PRECISION"
+    assert quantization["kv_cache_quantized_layers"] == {
+        "attn0": {"quant_algo": "FP8"},
+        "attn1": {"quant_algo": "NVFP4"},
+    }
+
+
 def test_unsupported_asymmetric_kv_cache_pair_fails_export():
     class FakeAttention(torch.nn.Module):
         def __init__(self):