From 196bdf5a0e7f444b9bd9d3c004ea56e6b36e6421 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Thu, 10 Sep 2026 23:12:16 -0700
Subject: [PATCH 1/3] Compose GEMM and KV-cache AutoQuantize workflows

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md                     |  29 ++-
 examples/hf_ptq/hf_ptq.py                     | 219 +++++++++++++-----
 modelopt/recipe/config.py                     |  32 ++-
 modelopt/torch/export/quant_utils.py          |   9 +-
 modelopt/torch/quantization/model_quant.py    |  25 +-
 ...n_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml |  64 +++++
 ...n_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml |  74 ++++++
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 191 ++++++++++++++-
 tests/unit/recipe/test_loader.py              |  88 +++++++
 .../torch/export/test_get_quantization.py     |  39 ++++
 .../quantization/test_kv_cache_auto_quant.py  |  34 ++-
 11 files changed, 712 insertions(+), 92 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 6ed60f9c23b..b29f799e4fa 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -203,7 +203,7 @@ python hf_ptq.py \
 
 Built-in recipes are located in `modelopt_recipes/general/ptq/` for model-agnostic recipes and in `modelopt_recipes/huggingface/<model_type>/ptq/` for recipes tuned to a specific Hugging Face `model_type` (see [`modelopt_recipes/huggingface/README.md`](../../modelopt_recipes/huggingface/README.md)). You can also provide a path to your own custom YAML recipe file or directory. See the [recipe documentation](https://nvidia.github.io/Model-Optimizer) for details on the YAML schema and available recipes.
 
-> *When `--recipe` is specified, `--qformat` is ignored. KV cache handling depends on the recipe type: a **PTQ** recipe bakes KV cache into its config and ignores `--kv_cache_qformat`; an **AutoQuantize** recipe falls back to `--kv_cache_qformat` unless it sets an explicit `kv_cache` field.*
+> *When `--recipe` is specified, `--qformat` is ignored. KV cache handling depends on the recipe type: a **PTQ** recipe bakes KV cache into its config and ignores `--kv_cache_qformat`; an **AutoQuantize** recipe falls back to `--kv_cache_qformat` unless it sets an explicit `kv_cache` field or `kv_auto_quantize` follow-up.*
 
 #### KV Cache Quantization
 
@@ -480,6 +480,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 use the same `mtq.auto_quantize` API and set
 `constraints.cost_model: kv_cache` with an `effective_bits` target. Their
 `candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
@@ -493,14 +506,14 @@ companion vLLM implementation does not support that asymmetric per-layer format:
 python hf_ptq.py \
   --pyt_ckpt_path Qwen/Qwen3.8-27B \
   --recipe general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
-  --auto_quantize_checkpoint /path/to/kv_autoquant.pth \
+  --kv_auto_quantize_checkpoint /path/to/kv_autoquant.pth \
   --export_path /path/to/qwen3.8-27b-mixed-kv
 ```
 
 Each candidate uses an explicit constant scale, avoiding an additional calibration pass while
 keeping persistent K/V scales in the unified HF checkpoint. Unified export records the selected
 formats in `kv_cache_quantized_layers`. `mtq.auto_quantize` returns the sensitivity scores and
-selected recipe in its search state; `--auto_quantize_checkpoint` stores that resumable state,
+selected recipe in its search state; `--kv_auto_quantize_checkpoint` stores that resumable state,
 including the candidate quantizer tensors needed for replay.
 
 KV sensitivity scoring runs one reference forward plus one forward per eligible-layer candidate
@@ -513,11 +526,13 @@ times the model vocabulary because reference and candidate log probabilities are
 > [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.
-> 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 companion consumer also does not yet apply the layer map when uniform FP8/NVFP4 weights are
+> present; export warns for that composed combination. Do not deploy either unsupported case. Full
+> FP8 K/V and full NVFP4 K/V use existing vLLM kernels once the relevant metadata path is available.
 
-The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an
-interrupted search (skips re-scoring):
+`--auto_quantize_checkpoint` saves/restores weight-search state. Every KV-domain search uses
+`--kv_auto_quantize_checkpoint`; a KV-primary recipe temporarily accepts the former flag as a
+deprecated fallback. Composed recipes therefore keep the weight and KV search states separate:
 
 ```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 94ca3b39b39..5e62cfcd698 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -19,6 +19,7 @@
 import random
 import time
 import warnings
+from fnmatch import fnmatch
 from pathlib import Path
 from typing import Any
 
@@ -368,7 +369,10 @@ def _mtq_candidate_formats(formats) -> list[dict]:
 
 
 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.
 
@@ -397,7 +401,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
@@ -433,6 +439,34 @@ def _mtq_inputs_from_auto_quantize_config(
     }
 
 
+def _quantize_config_explicitly_enables_kv(quant_cfg: dict[str, Any]) -> bool:
+    """Detect explicit K/V rules while preserving their ordered override semantics."""
+    enabled = dict.fromkeys(("k_bmm_quantizer", "v_bmm_quantizer"), False)
+    for entry in quant_cfg["quant_cfg"]:
+        pattern = entry["quantizer_name"]
+        if pattern != "*" and "bmm_quantizer" not in pattern:
+            continue
+        suffix = pattern.rsplit(".", 1)[-1]
+        for name in enabled:
+            if fnmatch(name, suffix):
+                enabled[name] = entry["enable"]
+    return any(enabled.values())
+
+
+def _resolve_kv_auto_quantize_checkpoint(args: argparse.Namespace) -> str | None:
+    """Resolve the dedicated KV checkpoint flag with a one-release legacy fallback."""
+    if args.kv_auto_quantize_checkpoint is not None:
+        return args.kv_auto_quantize_checkpoint
+    if args.auto_quantize_checkpoint is not None:
+        warnings.warn(
+            "Using --auto_quantize_checkpoint for a KV-cache search is deprecated; use "
+            "--kv_auto_quantize_checkpoint instead.",
+            FutureWarning,
+        )
+        return args.auto_quantize_checkpoint
+    return None
+
+
 def auto_quantize(
     args: argparse.Namespace,
     language_model: torch.nn.Module,
@@ -440,6 +474,8 @@ def auto_quantize(
     aq_config,
     full_model: torch.nn.Module | None = None,
     fixed_quantize_config=None,
+    allow_uniform_kv: bool = True,
+    checkpoint: str | None = None,
 ):
     """Recipe-driven auto_quantize, organized around an AutoQuantizeConfig.
 
@@ -456,13 +492,15 @@ def auto_quantize(
     )
 
     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,
     )
     if args.use_fsdp2:
         if inputs["search_domain"] == "kv_cache":
             raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR)
         warnings.warn(_FSDP2_AUTOQUANT_WARNING)
-
     # base-model lm_head handling (mirrors the CLI helper)
     is_base_model = (
         full_model is not None
@@ -524,7 +562,7 @@ def forward_step(model, batch):
         "verbose": True,
         "disabled_layers": inputs["disabled_layers"],
         "method": inputs["method"],
-        "checkpoint": args.auto_quantize_checkpoint,
+        "checkpoint": checkpoint,
     }
     if inputs["search_domain"] == "weight":
         auto_quantize_kwargs.update(
@@ -570,9 +608,9 @@ def _recipe_is_kv_auto_quantize(recipe: str | None) -> bool:
     if recipe is None:
         return False
     loaded_recipe = load_recipe(recipe)
-    return (
-        isinstance(loaded_recipe, ModelOptAutoQuantizeRecipe)
-        and loaded_recipe.auto_quantize.constraints.cost_model == "kv_cache"
+    return isinstance(loaded_recipe, ModelOptAutoQuantizeRecipe) and any(
+        stage is not None and stage.constraints.cost_model == "kv_cache"
+        for stage in (loaded_recipe.auto_quantize, loaded_recipe.kv_auto_quantize)
     )
 
 
@@ -865,6 +903,103 @@ 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 args.layerwise_export and not mtp_layer_prefixes:
+        mtp_layer_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path)
+    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 args.layerwise_export:
+        assert_layerwise_export_compatible(args, full_model, quant_cfg.get("algorithm"))
+        quant_cfg = set_layerwise_export_dir(quant_cfg, args.export_path)
+        print(f"Layerwise export enabled: writing quantized shards to {args.export_path}")
+        quant_cfg, moved = default_layerwise_resume_dir(quant_cfg, args.export_path)
+        if moved:
+            print(
+                "Layerwise checkpoint_dir co-located with the export path so a resumed run "
+                "finds its manifest next to the shards it must not overwrite."
+            )
+
+    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.cost_model == "kv_cache"
+    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)
+        if _quantize_config_explicitly_enables_kv(quant_cfg):
+            raise ValueError(
+                "The fixed quantize stage explicitly enables K/V quantizers before KV-cache "
+                "AutoQuantize. Disable them in the fixed stage."
+            )
+        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=(
+            _resolve_kv_auto_quantize_checkpoint(args)
+            if primary_is_kv
+            else args.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=args.kv_auto_quantize_checkpoint,
+        )
+
+
 def assert_layerwise_export_compatible(args, full_model, algorithm) -> None:
     """Refuse layerwise export before calibration starts, not after the run is paid for.
 
@@ -1261,10 +1396,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
 
     layerwise_cfgs = recipe_layerwise_blocks(recipe)
     is_layerwise = any(cfg.get("enable", False) for cfg in layerwise_cfgs)
@@ -1367,16 +1500,16 @@ def quantize_main(
     )
 
     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:
@@ -1411,44 +1544,7 @@ def quantize_main(
                     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 args.layerwise_export and not mtp_layer_prefixes:
-            # Only the FSDP2 loader flags these before quantization, and the exclusions must
-            # be in quant_cfg before mtq.quantize converts the first layer.
-            mtp_layer_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path)
-        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}")
-
-        # Before resolve_checkpoint_dir, which hashes the config: with the placeholder
-        # still in it, two --export_path values would share one checkpoint dir.
-        if args.layerwise_export:
-            assert_layerwise_export_compatible(args, full_model, quant_cfg.get("algorithm"))
-            quant_cfg = set_layerwise_export_dir(quant_cfg, args.export_path)
-            print(f"Layerwise export enabled: writing quantized shards to {args.export_path}")
-            # The shards are only a resume artifact if the manifest that names the resume
-            # point survives alongside them; see default_layerwise_resume_dir.
-            quant_cfg, moved = default_layerwise_resume_dir(quant_cfg, args.export_path)
-            if moved:
-                print(
-                    "Layerwise checkpoint_dir co-located with the export path so a resumed "
-                    "run finds its manifest next to the shards it must not overwrite."
-                )
-
-        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(
@@ -1685,10 +1781,21 @@ def parse_args() -> argparse.Namespace:
         type=str,
         default=None,
         help=(
-            "Path to checkpoint file for saving/restoring auto_quantize search state "
+            "Path to checkpoint file for saving/restoring weight AutoQuantize search state "
             "(sensitivity scores, costs, etc.). Used with an AutoQuantize --recipe."
         ),
     )
+    parser.add_argument(
+        "--kv_auto_quantize_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Path for saving/restoring any KV-cache AutoQuantize search checkpoint. Use a new "
+            "path whenever the preceding weight/activation quantization stage changes. "
+            "KV-primary recipes temporarily accept --auto_quantize_checkpoint as a deprecated "
+            "fallback."
+        ),
+    )
     parser.add_argument(
         "--moe_calib_experts_ratio",
         type=float,
diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index 28de372adad..1621f425e8f 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -313,9 +313,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(
@@ -323,22 +323,42 @@ 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.cost_model == "kv_cache"
+        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.cost_model != "kv_cache":
+                raise ValueError("kv_auto_quantize must use cost_model=kv_cache.")
+            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."
+                )
         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."
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index accedfb3a90..674d33a54c7 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -1731,9 +1731,12 @@ def get_quant_config(
     )
     if needs_layerwise_kv_metadata:
         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 AutoQuantize recipe."
+            warn(
+                "The exported checkpoint combines uniform quantized weights with a mixed-precision "
+                "KV-cache layer map. Released runtimes do not yet consume "
+                "kv_cache_quantized_layers for uniform-weight ModelOpt checkpoints; deployment "
+                "remains unsupported until the runtime adds that metadata path.",
+                stacklevel=2,
             )
         # KV metadata is orthogonal to weight metadata. In particular, a KV-only search
         # must preserve BF16 weights instead of synthesizing a weight quantization algorithm.
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 5206b691b23..3a0f9a86a14 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -325,11 +325,6 @@ def _auto_quantize_kv_cache(
             "KV-cache AutoQuantize is single-process only; distributed scoring, selection, "
             "and checkpoint writes are not synchronized."
         )
-    if is_quantized(model):
-        raise NotImplementedError(
-            "KV-cache AutoQuantize requires an unquantized model; composing it after GEMM "
-            "PTQ or AutoQuantize is not supported yet."
-        )
     if method not in (None, "kl_div"):
         raise ValueError("cost_model='kv_cache' requires method='kl_div'.")
     if fixed_quantization_config is not None or module_search_spaces:
@@ -347,6 +342,21 @@ def _auto_quantize_kv_cache(
     if data_loader is None or forward_step is None:
         raise ValueError("data_loader and forward_step must be provided for KV-cache AutoQuantize.")
 
+    converted_for_search = not is_quantized(model)
+    if not converted_for_search:
+        enabled_kv_quantizers = [
+            name
+            for name, module in model.named_modules(remove_duplicate=False)
+            if name.endswith(("k_bmm_quantizer", "v_bmm_quantizer"))
+            and getattr(module, "is_enabled", False)
+        ]
+        if enabled_kv_quantizers:
+            raise ValueError(
+                "The preceding quantization stage left K/V quantizers enabled: "
+                f"{enabled_kv_quantizers}. Disable them before running KV-cache AutoQuantize; "
+                "clearing them now would not undo prior calibration or sensitivity measurements."
+            )
+
     processed_kv_formats: list[tuple[dict[str, Any], str | None]] = []
     for candidate in quantization_formats:
         if isinstance(candidate, tuple):
@@ -375,8 +385,9 @@ def _auto_quantize_kv_cache(
         num_calib_steps,
         num_score_steps,
     )
-    model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
-    set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
+    if converted_for_search:
+        model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
+        set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
     searcher = AutoQuantizeKVSearcher()
     searcher.search(
         model,
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..98fc35184ed
--- /dev/null
+++ b/modelopt_recipes/general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -0,0 +1,64 @@
+# 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.
+
+# 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
+  fp8: configs/numerics/fp8
+  nvfp4: configs/numerics/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:
+    effective_bits: 5.4
+    cost_model: kv_cache
+
+  candidate_formats:
+    - quant_cfg:
+        - quantizer_name: "*[kv]_bmm_quantizer"
+          cfg:
+            $import: fp8
+            constant_amax: 448.0
+      algorithm:
+      effective_bits: 8.0
+    - quant_cfg:
+        - quantizer_name: "*[kv]_bmm_quantizer"
+          cfg:
+            $import: nvfp4
+            constant_amax: 448.0
+      algorithm:
+      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..328c3d429a7
--- /dev/null
+++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml
@@ -0,0 +1,74 @@
+# 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.
+
+# 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/numerics/fp8
+  kv_nvfp4: configs/numerics/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:
+    effective_bits: 5.4
+    cost_model: kv_cache
+
+  candidate_formats:
+    - quant_cfg:
+        - quantizer_name: "*[kv]_bmm_quantizer"
+          cfg:
+            $import: kv_fp8
+            constant_amax: 448.0
+      algorithm:
+      effective_bits: 8.0
+    - quant_cfg:
+        - quantizer_name: "*[kv]_bmm_quantizer"
+          cfg:
+            $import: kv_nvfp4
+            constant_amax: 448.0
+      algorithm:
+      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 c3e04f253a4..163809d9b61 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 import tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
@@ -116,6 +120,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)
@@ -127,12 +142,14 @@ def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch):
     model = get_tiny_qwen3(num_hidden_layers=1)
     aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize
     args = SimpleNamespace(
+        qformat="fp8",
         calib_with_images=False,
         inference_pipeline_parallel=1,
         use_fsdp2=False,
         kv_cache_qformat="none",
         batch_size=1,
         auto_quantize_checkpoint=None,
+        kv_auto_quantize_checkpoint=None,
     )
     data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
 
@@ -145,6 +162,175 @@ def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch):
     assert attention.v_bmm_quantizer.amax == 448.0
 
 
+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"])],
+    )
+    kv_aq = AutoQuantizeConfig(
+        constraints=AutoQuantizeConstraints(effective_bits=8.0, cost_model="kv_cache"),
+        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(
+            auto_quantize_checkpoint="weight-search.pth",
+            kv_auto_quantize_checkpoint="kv-search.pth",
+        ),
+        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"] == "weight-search.pth"
+    assert calls[1]["checkpoint"] == "kv-search.pth"
+
+
+def test_hf_ptq_runs_fixed_ptq_before_kv_autoquantize(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(
+        tensor_quant,
+        "dynamic_block_quantize_op",
+        lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs),
+    )
+    recipe = load_recipe("general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits")
+    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,
+        layerwise_export=False,
+        specdec_offline_dataset=None,
+    )
+
+    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_autoquantize_checkpoint_uses_dedicated_flag_with_legacy_fallback(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    args = SimpleNamespace(
+        auto_quantize_checkpoint="legacy.pth",
+        kv_auto_quantize_checkpoint="kv.pth",
+    )
+
+    assert hf_ptq._resolve_kv_auto_quantize_checkpoint(args) == "kv.pth"
+
+    args.kv_auto_quantize_checkpoint = None
+    with pytest.warns(FutureWarning, match="deprecated"):
+        assert hf_ptq._resolve_kv_auto_quantize_checkpoint(args) == "legacy.pth"
+
+
+def test_fixed_ptq_kv_precheck_does_not_widen_scoped_gemm_rule(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    fixed = QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "model.layers.*.mlp.*",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+            }
+        ],
+        algorithm="max",
+    )
+
+    assert not hf_ptq._quantize_config_explicitly_enables_kv(fixed.model_dump())
+
+
+def test_fixed_ptq_then_kv_rejects_explicit_kv_before_calibration(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    fixed = QuantizeConfig(
+        quant_cfg=[
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "model.layers.*.self_attn.*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+            },
+        ],
+        algorithm="max",
+    )
+    kv_aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize
+    recipe = ModelOptAutoQuantizeRecipe(quantize=fixed, auto_quantize=kv_aq)
+    args = SimpleNamespace(
+        auto_quantize_checkpoint=None,
+        kv_auto_quantize_checkpoint=None,
+        pyt_ckpt_path="dummy",
+        cast_mxfp4_to_nvfp4=False,
+        layerwise_export=False,
+    )
+    monkeypatch.setattr(
+        hf_ptq,
+        "mono_quantize",
+        lambda *_args, **_kwargs: pytest.fail("fixed PTQ must not start"),
+    )
+
+    with pytest.raises(ValueError, match="fixed quantize stage explicitly enables K/V"):
+        hf_ptq._run_auto_quantize_recipe(
+            args, recipe, torch.nn.Module(), torch.nn.Module(), None, False, [], False
+        )
+
+
+def test_composed_kv_autoquant_rejects_enabled_actual_kv_quantizers(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    model = get_tiny_qwen3(num_hidden_layers=1)
+    hf_ptq.mtq.quantize(
+        model,
+        {
+            "quant_cfg": [
+                {
+                    "quantizer_name": "*[kv]_bmm_quantizer",
+                    "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                }
+            ],
+            "algorithm": None,
+        },
+    )
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=False,
+        kv_cache_qformat="none",
+        batch_size=1,
+    )
+    aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize
+
+    with pytest.raises(ValueError, match="preceding quantization stage left K/V"):
+        hf_ptq.auto_quantize(args, model, [], aq, full_model=model)
+
+
 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)
@@ -281,6 +467,9 @@ def test_fsdp2_preload_guard_distinguishes_weight_and_kv_autoquant(monkeypatch):
     assert hf_ptq._recipe_is_kv_auto_quantize(
         "general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits"
     )
+    assert hf_ptq._recipe_is_kv_auto_quantize(
+        "general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits"
+    )
     assert not hf_ptq._recipe_is_kv_auto_quantize("general/auto_quantize/nvfp4_fp8_at_5p4bits")
 
 
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index 7f9766323dc..13168e0f694 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -25,6 +25,7 @@
 from pathlib import Path
 
 import pytest
+from pydantic import ValidationError
 
 import modelopt.torch.quantization.config as qcfg
 from modelopt.recipe.config import (
@@ -1955,8 +1956,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",
@@ -2000,6 +2003,91 @@ def test_load_recipe_kv_autoquantize_contract():
         assert fmt.algorithm is None
 
 
+@pytest.mark.parametrize(
+    ("recipe_path", "kv_stage"),
+    [
+        (
+            "general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits",
+            "auto_quantize",
+        ),
+        (
+            "general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits",
+            "kv_auto_quantize",
+        ),
+    ],
+)
+def test_builtin_composed_kv_recipes_use_calibration_free_cast_candidates(recipe_path, kv_stage):
+    aq = getattr(load_recipe(recipe_path), kv_stage)
+
+    assert aq is not None
+    assert aq.constraints.cost_model == "kv_cache"
+    assert all(candidate.algorithm is None for candidate in aq.candidate_formats)
+    assert all(
+        candidate.quant_cfg[0].cfg.constant_amax == 448.0 for candidate in aq.candidate_formats
+    )
+
+
+def _weight_autoquantize_test_config(**updates):
+    config = AutoQuantizeConfig(
+        constraints=AutoQuantizeConstraints(effective_bits=8.0),
+        candidate_formats=[qcfg.QuantizeConfig(quant_cfg=[], algorithm="max")],
+    )
+    return config.model_copy(update=updates)
+
+
+def _kv_autoquantize_test_config(**updates):
+    config = AutoQuantizeConfig(
+        constraints=AutoQuantizeConstraints(effective_bits=8.0, cost_model="kv_cache"),
+        candidate_formats=[
+            qcfg.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",
+    )
+    return config.model_copy(update=updates)
+
+
+def test_autoquantize_recipe_rejects_second_kv_search():
+    with pytest.raises(ValidationError, match=r"cannot follow.*already searches the KV cache"):
+        ModelOptAutoQuantizeRecipe(
+            auto_quantize=_kv_autoquantize_test_config(),
+            kv_auto_quantize=_kv_autoquantize_test_config(),
+        )
+
+
+def test_autoquantize_recipe_rejects_non_kv_followup():
+    with pytest.raises(ValidationError, match="must use cost_model=kv_cache"):
+        ModelOptAutoQuantizeRecipe(
+            auto_quantize=_weight_autoquantize_test_config(),
+            kv_auto_quantize=_weight_autoquantize_test_config(),
+        )
+
+
+def test_autoquantize_recipe_rejects_uniform_and_searched_kv():
+    uniform_kv = qcfg.QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+            }
+        ],
+        algorithm=None,
+    )
+    with pytest.raises(ValidationError, match=r"must omit.*uniform auto_quantize.kv_cache"):
+        ModelOptAutoQuantizeRecipe(
+            auto_quantize=_weight_autoquantize_test_config(kv_cache=uniform_kv),
+            kv_auto_quantize=_kv_autoquantize_test_config(),
+        )
+
+
 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 8b9670e5576..5f196ac64b7 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -220,6 +220,45 @@ def test_uniform_vlm_export_ignores_disabled_vision_attention():
     assert "kv_cache_quantized_layers" not in quantization
 
 
+def test_uniform_weight_quantization_exports_mixed_kv_cache_map():
+    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), "constant_amax": 1.0},
+            }
+        ],
+    )
+    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)},
+                    "constant_amax": 1.0,
+                },
+            }
+        ],
+    )
+
+    with pytest.warns(UserWarning, match="uniform quantized weights.*mixed-precision KV-cache"):
+        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_quant_config_tolerates_ambiguous_language_model_roots():
     model = torch.nn.Module()
     model.model = torch.nn.Module()
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 ff87d2c8f7c..8c2ace6b76c 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -761,7 +761,7 @@ def test_public_kv_autoquant_rejects_distributed_execution_before_mutation(monke
     assert {name: type(module) for name, module in model.named_modules()} == original_types
 
 
-def test_public_kv_autoquant_rejects_preceding_quantization_before_search():
+def test_public_kv_autoquant_preserves_preceding_weight_quantization():
     model = get_tiny_llama(num_hidden_layers=2)
     model = mtq.quantize(
         model,
@@ -770,7 +770,7 @@ def test_public_kv_autoquant_rejects_preceding_quantization_before_search():
                 {"quantizer_name": "*", "enable": False},
                 {
                     "quantizer_name": "*.weight_quantizer",
-                    "cfg": {"num_bits": (4, 3), "axis": None},
+                    "cfg": {"num_bits": (4, 3), "axis": None, "constant_amax": 1.0},
                     "enable": True,
                 },
             ],
@@ -788,13 +788,23 @@ def test_public_kv_autoquant_rejects_preceding_quantization_before_search():
         "effective_bits": 8.0,
     }
 
-    with pytest.raises(NotImplementedError, match="requires an unquantized model"):
-        mtq.auto_quantize(
-            model,
-            {"effective_bits": 8.0, "cost_model": "kv_cache"},
-            [candidate],
-            [],
-            lambda *_: pytest.fail("Validation must fail before search."),
-            num_calib_steps=1,
-            num_score_steps=1,
-        )
+    data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}]
+    weight_quantizer = model.model.layers[0].self_attn.q_proj.weight_quantizer
+
+    model, _ = mtq.auto_quantize(
+        model,
+        {"effective_bits": 8.0, "cost_model": "kv_cache"},
+        [candidate],
+        data,
+        lambda search_model, batch: search_model(**batch).logits,
+        num_calib_steps=1,
+        num_score_steps=1,
+    )
+
+    assert weight_quantizer.is_enabled
+    assert weight_quantizer.num_bits == (4, 3)
+    assert weight_quantizer.amax == 1.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
+    )

From fef442ac57ec3c11d4e8a879ec6da3d242b1e830 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Fri, 11 Sep 2026 16:18:09 -0700
Subject: [PATCH 2/3] Fix composed KV AutoQuantize validation

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/hf_ptq.py                     | 10 ++--
 .../torch/quantization/kv_cache_auto_quant.py | 38 ++++++++++++++-
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 28 +++++++++++
 .../quantization/test_kv_cache_auto_quant.py  | 46 +++++++++++++++++++
 4 files changed, 117 insertions(+), 5 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 5e62cfcd698..f998bd9039e 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -443,18 +443,20 @@ def _quantize_config_explicitly_enables_kv(quant_cfg: dict[str, Any]) -> bool:
     """Detect explicit K/V rules while preserving their ordered override semantics."""
     enabled = dict.fromkeys(("k_bmm_quantizer", "v_bmm_quantizer"), False)
     for entry in quant_cfg["quant_cfg"]:
+        if entry.get("parent_class") is not None:
+            continue
         pattern = entry["quantizer_name"]
         if pattern != "*" and "bmm_quantizer" not in pattern:
             continue
-        suffix = pattern.rsplit(".", 1)[-1]
         for name in enabled:
-            if fnmatch(name, suffix):
+            qualified_name = f"model.layers.0.self_attn.{name}"
+            if fnmatch(name, pattern) or fnmatch(qualified_name, pattern) or pattern.endswith(name):
                 enabled[name] = entry["enable"]
     return any(enabled.values())
 
 
 def _resolve_kv_auto_quantize_checkpoint(args: argparse.Namespace) -> str | None:
-    """Resolve the dedicated KV checkpoint flag with a one-release legacy fallback."""
+    """Resolve a KV-primary checkpoint with a one-release legacy fallback."""
     if args.kv_auto_quantize_checkpoint is not None:
         return args.kv_auto_quantize_checkpoint
     if args.auto_quantize_checkpoint is not None:
@@ -996,6 +998,8 @@ def _run_auto_quantize_recipe(
             aq_config=followup_kv,
             full_model=full_model,
             allow_uniform_kv=False,
+            # The weight search owns --auto_quantize_checkpoint, so a follow-up KV search must
+            # never use the KV-primary legacy fallback and collide with the weight state.
             checkpoint=args.kv_auto_quantize_checkpoint,
         )
 
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index 760a7683c51..ef758f7ebcd 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import copy
 import fnmatch
 import math
 from contextlib import contextmanager
@@ -435,6 +436,7 @@ def _search_signature(
     layers: list[tuple[str, nn.Module, int]],
     num_calib_steps: int,
     num_score_steps: int,
+    preceding_quantizers: list[dict[str, Any]],
 ) -> dict[str, Any]:
     return {
         "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
@@ -455,11 +457,41 @@ def _search_signature(
             }
             for name, module, _ in layers
         ],
+        "preceding_quantizers": preceding_quantizers,
     }
 
 
 def _checkpoint_state_is_compatible(state: dict[str, Any], signature: dict[str, Any]) -> bool:
-    return state.get("search_signature") == signature
+    checkpoint_signature = state.get("search_signature")
+    if checkpoint_signature == signature:
+        return True
+    if not isinstance(checkpoint_signature, dict) or signature["preceding_quantizers"]:
+        return False
+
+    # Checkpoints written before composed GEMM -> KV searches had no preceding quantizers.
+    # Preserve their compatibility with an unquantized model while rejecting them for a
+    # quantized baseline, whose sensitivity scores depend on that baseline.
+    legacy_signature = signature.copy()
+    legacy_signature.pop("preceding_quantizers")
+    return checkpoint_signature == legacy_signature
+
+
+def _preceding_quantizer_signature(model: nn.Module) -> list[dict[str, Any]]:
+    """Describe enabled non-K/V formats that affect KV sensitivity scores."""
+    return sorted(
+        (
+            {
+                "name": name,
+                "num_bits": copy.deepcopy(module.num_bits),
+                "block_sizes": copy.deepcopy(module.block_sizes),
+            }
+            for name, module in model.named_modules(remove_duplicate=False)
+            if isinstance(module, TensorQuantizer)
+            and module.is_enabled
+            and not name.endswith(_KV_QUANTIZER_ATTRS)
+        ),
+        key=lambda entry: entry["name"],
+    )
 
 
 def _quantizer_state_dict(
@@ -691,13 +723,15 @@ def before_search(self) -> None:
             layers,
             self.config["num_calib_steps"],
             self.config["num_score_steps"],
+            _preceding_quantizer_signature(self.model),
         )
         if self.search_signature is not None and not _checkpoint_state_is_compatible(
             self.state_dict(), signature
         ):
             raise ValueError(
                 "KV-cache AutoQuantize checkpoint does not match the current candidates, scoring "
-                "setup, or eligible layers. Use a different checkpoint path."
+                "setup, eligible layers, or preceding non-K/V quantizer configuration. Use a "
+                "different checkpoint path."
             )
         self.search_signature = signature
         self._hparams = [
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 163809d9b61..13c29bb1f62 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -270,6 +270,34 @@ def test_fixed_ptq_kv_precheck_does_not_widen_scoped_gemm_rule(monkeypatch):
     assert not hf_ptq._quantize_config_explicitly_enables_kv(fixed.model_dump())
 
 
+@pytest.mark.parametrize(
+    "kv_pattern",
+    [
+        "model.layers.*.self_attn.*[kv]_bmm_quantizer",
+        "*self_attn*k_bmm_quantizer",
+    ],
+)
+def test_fixed_ptq_kv_precheck_ignores_unrelated_parent_scoped_rules(monkeypatch, kv_pattern):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    fixed = QuantizeConfig(
+        quant_cfg=[
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": kv_pattern,
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+            },
+            {
+                "parent_class": "nn.Embedding",
+                "quantizer_name": "*",
+                "enable": False,
+            },
+        ],
+        algorithm="max",
+    )
+
+    assert hf_ptq._quantize_config_explicitly_enables_kv(fixed.model_dump())
+
+
 def test_fixed_ptq_then_kv_rejects_explicit_kv_before_calibration(monkeypatch):
     hf_ptq = _import_hf_ptq(monkeypatch)
     fixed = QuantizeConfig(
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 8c2ace6b76c..0cc27f1b969 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -808,3 +808,49 @@ def test_public_kv_autoquant_preserves_preceding_weight_quantization():
         layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled
         for layer in model.model.layers
     )
+
+
+def test_kv_autoquant_checkpoint_rejects_changed_preceding_quantization(
+    tmp_path, nvfp4_fake_quant_stub
+):
+    def quantized_model(bits):
+        return mtq.quantize(
+            get_tiny_llama(num_hidden_layers=1),
+            {
+                "quant_cfg": [
+                    {"quantizer_name": "*", "enable": False},
+                    {
+                        "quantizer_name": "*.weight_quantizer",
+                        "cfg": _quantizer_cfg(bits, constant_amax=1.0),
+                    },
+                ],
+                "algorithm": None,
+            },
+        )
+
+    candidate = _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump()
+    data = [{"input_ids": torch.randint(0, 16, (1, 8))}]
+    checkpoint = str(tmp_path / "kv_search.pth")
+
+    mtq.auto_quantize(
+        quantized_model((4, 3)),
+        {"effective_bits": 8.0, "cost_model": "kv_cache"},
+        [candidate],
+        data,
+        lambda model, batch: model(**batch).logits,
+        num_calib_steps=1,
+        num_score_steps=1,
+        checkpoint=checkpoint,
+    )
+
+    with pytest.raises(ValueError, match="preceding non-K/V quantizer"):
+        mtq.auto_quantize(
+            quantized_model((2, 1)),
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
+            [candidate],
+            data,
+            lambda *_: pytest.fail("A stale checkpoint must be rejected before scoring."),
+            num_calib_steps=1,
+            num_score_steps=1,
+            checkpoint=checkpoint,
+        )

From 970ce89fafbafbaf5dd1172728825df322e56429 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Fri, 11 Sep 2026 17:19:46 -0700
Subject: [PATCH 3/3] Strengthen composed KV checkpoint validation

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/hf_ptq.py                     | 24 +++--
 .../torch/quantization/kv_cache_auto_quant.py | 50 +++++++++-
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 59 ++++++++++++
 .../quantization/test_kv_cache_auto_quant.py  | 94 ++++++++++++++++---
 4 files changed, 201 insertions(+), 26 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index f998bd9039e..04725da5dd7 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -441,18 +441,28 @@ def _mtq_inputs_from_auto_quantize_config(
 
 def _quantize_config_explicitly_enables_kv(quant_cfg: dict[str, Any]) -> bool:
     """Detect explicit K/V rules while preserving their ordered override semantics."""
-    enabled = dict.fromkeys(("k_bmm_quantizer", "v_bmm_quantizer"), False)
+    names = ("k_bmm_quantizer", "v_bmm_quantizer")
+    enabled_by_parent = {None: dict.fromkeys(names, False)}
     for entry in quant_cfg["quant_cfg"]:
-        if entry.get("parent_class") is not None:
-            continue
         pattern = entry["quantizer_name"]
         if pattern != "*" and "bmm_quantizer" not in pattern:
             continue
-        for name in enabled:
-            qualified_name = f"model.layers.0.self_attn.{name}"
-            if fnmatch(name, pattern) or fnmatch(qualified_name, pattern) or pattern.endswith(name):
+        basename_pattern = pattern.rsplit(".", 1)[-1]
+        matched_names = [
+            name for name in names if fnmatch(name, basename_pattern) or pattern.endswith(name)
+        ]
+        if not matched_names:
+            continue
+
+        parent_class = entry.get("parent_class")
+        if parent_class is None:
+            scopes = enabled_by_parent.values()
+        else:
+            scopes = [enabled_by_parent.setdefault(parent_class, enabled_by_parent[None].copy())]
+        for enabled in scopes:
+            for name in matched_names:
                 enabled[name] = entry["enable"]
-    return any(enabled.values())
+    return any(any(enabled.values()) for enabled in enabled_by_parent.values())
 
 
 def _resolve_kv_auto_quantize_checkpoint(args: argparse.Namespace) -> str | None:
diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py
index ef758f7ebcd..82846cb1bc1 100644
--- a/modelopt/torch/quantization/kv_cache_auto_quant.py
+++ b/modelopt/torch/quantization/kv_cache_auto_quant.py
@@ -17,8 +17,9 @@
 
 from __future__ import annotations
 
-import copy
 import fnmatch
+import hashlib
+import json
 import math
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, cast
@@ -476,14 +477,55 @@ def _checkpoint_state_is_compatible(state: dict[str, Any], signature: dict[str,
     return checkpoint_signature == legacy_signature
 
 
+def _fingerprint_value(value: Any) -> Any:
+    """Convert quantizer configuration and tensor state into a stable JSON value."""
+    if isinstance(value, torch.Tensor):
+        if value.device.type == "meta":
+            raise ValueError("Cannot fingerprint a meta-device preceding quantizer state.")
+        tensor = value.detach().contiguous().cpu()
+        raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes()
+        return {
+            "dtype": str(tensor.dtype),
+            "shape": list(tensor.shape),
+            "sha256": hashlib.sha256(raw).hexdigest(),
+        }
+    if hasattr(value, "model_dump"):
+        return _fingerprint_value(value.model_dump(mode="json"))
+    if isinstance(value, dict):
+        return [
+            [_fingerprint_value(key), _fingerprint_value(item)]
+            for key, item in sorted(
+                value.items(), key=lambda entry: (type(entry[0]).__qualname__, repr(entry[0]))
+            )
+        ]
+    if isinstance(value, (list, tuple)):
+        return [_fingerprint_value(item) for item in value]
+    if isinstance(value, (torch.dtype, torch.device)):
+        return str(value)
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    raise TypeError(f"Unsupported preceding quantizer state value: {type(value).__qualname__}.")
+
+
+def _quantizer_fingerprint(module: TensorQuantizer) -> str:
+    payload = {
+        "type": f"{type(module).__module__}.{type(module).__qualname__}",
+        "properties": module.get_modelopt_state(properties_only=True),
+        "state_dict": module.state_dict(),
+    }
+    serialized = json.dumps(
+        _fingerprint_value(payload), sort_keys=True, separators=(",", ":"), allow_nan=False
+    )
+    return hashlib.sha256(serialized.encode()).hexdigest()
+
+
 def _preceding_quantizer_signature(model: nn.Module) -> list[dict[str, Any]]:
-    """Describe enabled non-K/V formats that affect KV sensitivity scores."""
+    """Fingerprint enabled non-K/V configuration and state that affect KV scores."""
     return sorted(
         (
             {
                 "name": name,
-                "num_bits": copy.deepcopy(module.num_bits),
-                "block_sizes": copy.deepcopy(module.block_sizes),
+                "fingerprint": _quantizer_fingerprint(module),
             }
             for name, module in model.named_modules(remove_duplicate=False)
             if isinstance(module, TensorQuantizer)
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 13c29bb1f62..5453d14c94e 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -208,6 +208,47 @@ def test_hf_ptq_runs_weight_then_kv_autoquantize_stages(monkeypatch):
     assert calls[1]["checkpoint"] == "kv-search.pth"
 
 
+def test_hf_ptq_runs_real_weight_then_kv_autoquantize_stages(monkeypatch):
+    """Exercise the shipped gradient-weight -> KL-div KV composition without mocked stages."""
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(
+        tensor_quant,
+        "dynamic_block_quantize_op",
+        lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs),
+    )
+    recipe = load_recipe(
+        "general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits"
+    )
+    model = get_tiny_qwen3(num_hidden_layers=1)
+    input_ids = torch.arange(8).unsqueeze(0) % model.config.vocab_size
+    data = [{"input_ids": input_ids, "labels": input_ids.clone()}]
+    args = SimpleNamespace(
+        qformat="fp8",
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=False,
+        kv_cache_qformat="none",
+        batch_size=1,
+        auto_quantize_checkpoint=None,
+        kv_auto_quantize_checkpoint=None,
+    )
+
+    hf_ptq._run_auto_quantize_recipe(args, recipe, model, model, None, False, data, False)
+
+    enabled_weight_quantizers = [
+        module
+        for name, module in model.named_modules()
+        if name.endswith("weight_quantizer") and getattr(module, "is_enabled", False)
+    ]
+    assert enabled_weight_quantizers
+    assert all(module.num_bits in ((2, 1), (4, 3)) for module in enabled_weight_quantizers)
+    attention = model.model.layers[0].self_attn
+    assert attention.k_bmm_quantizer.is_enabled
+    assert attention.v_bmm_quantizer.is_enabled
+    assert attention.k_bmm_quantizer.num_bits in ((2, 1), (4, 3))
+    assert attention.v_bmm_quantizer.num_bits in ((2, 1), (4, 3))
+
+
 def test_hf_ptq_runs_fixed_ptq_before_kv_autoquantize(monkeypatch):
     hf_ptq = _import_hf_ptq(monkeypatch)
     monkeypatch.setattr(
@@ -275,6 +316,7 @@ def test_fixed_ptq_kv_precheck_does_not_widen_scoped_gemm_rule(monkeypatch):
     [
         "model.layers.*.self_attn.*[kv]_bmm_quantizer",
         "*self_attn*k_bmm_quantizer",
+        "*.language_model.*.attention.*_bmm_quantizer",
     ],
 )
 def test_fixed_ptq_kv_precheck_ignores_unrelated_parent_scoped_rules(monkeypatch, kv_pattern):
@@ -298,6 +340,23 @@ def test_fixed_ptq_kv_precheck_ignores_unrelated_parent_scoped_rules(monkeypatch
     assert hf_ptq._quantize_config_explicitly_enables_kv(fixed.model_dump())
 
 
+def test_fixed_ptq_kv_precheck_detects_parent_scoped_kv_rule(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    fixed = QuantizeConfig(
+        quant_cfg=[
+            {"quantizer_name": "*", "enable": False},
+            {
+                "parent_class": "LlamaAttention",
+                "quantizer_name": "*_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+            },
+        ],
+        algorithm="max",
+    )
+
+    assert hf_ptq._quantize_config_explicitly_enables_kv(fixed.model_dump())
+
+
 def test_fixed_ptq_then_kv_rejects_explicit_kv_before_calibration(monkeypatch):
     hf_ptq = _import_hf_ptq(monkeypatch)
     fixed = QuantizeConfig(
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 0cc27f1b969..f4bb658706b 100644
--- a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
+++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py
@@ -810,30 +810,88 @@ def test_public_kv_autoquant_preserves_preceding_weight_quantization():
     )
 
 
+def _quantized_weight_baseline(bits, *, constant_amax=1.0, axis=None):
+    quantizer_cfg = _quantizer_cfg(bits, constant_amax=constant_amax)
+    quantizer_cfg["axis"] = axis
+    return mtq.quantize(
+        get_tiny_llama(num_hidden_layers=1),
+        {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {
+                    "quantizer_name": "*.weight_quantizer",
+                    "cfg": quantizer_cfg,
+                },
+            ],
+            "algorithm": None,
+        },
+    )
+
+
 def test_kv_autoquant_checkpoint_rejects_changed_preceding_quantization(
     tmp_path, nvfp4_fake_quant_stub
 ):
-    def quantized_model(bits):
-        return mtq.quantize(
-            get_tiny_llama(num_hidden_layers=1),
-            {
-                "quant_cfg": [
-                    {"quantizer_name": "*", "enable": False},
-                    {
-                        "quantizer_name": "*.weight_quantizer",
-                        "cfg": _quantizer_cfg(bits, constant_amax=1.0),
-                    },
-                ],
-                "algorithm": None,
-            },
+    candidate = _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump()
+    data = [{"input_ids": torch.randint(0, 16, (1, 8))}]
+    checkpoint = str(tmp_path / "kv_search.pth")
+
+    mtq.auto_quantize(
+        _quantized_weight_baseline((4, 3)),
+        {"effective_bits": 8.0, "cost_model": "kv_cache"},
+        [candidate],
+        data,
+        lambda model, batch: model(**batch).logits,
+        num_calib_steps=1,
+        num_score_steps=1,
+        checkpoint=checkpoint,
+    )
+
+    mtq.auto_quantize(
+        _quantized_weight_baseline((4, 3)),
+        {"effective_bits": 8.0, "cost_model": "kv_cache"},
+        [candidate],
+        data,
+        lambda *_: pytest.fail("An identical preceding state must restore without rescoring."),
+        num_calib_steps=1,
+        num_score_steps=1,
+        checkpoint=checkpoint,
+    )
+
+    with pytest.raises(ValueError, match="preceding non-K/V quantizer"):
+        mtq.auto_quantize(
+            _quantized_weight_baseline((2, 1)),
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
+            [candidate],
+            data,
+            lambda *_: pytest.fail("A stale checkpoint must be rejected before scoring."),
+            num_calib_steps=1,
+            num_score_steps=1,
+            checkpoint=checkpoint,
         )
 
+
+@pytest.mark.parametrize(
+    ("first_kwargs", "second_kwargs", "mutate_second_amax"),
+    [
+        ({"constant_amax": 1.0}, {"constant_amax": 2.0}, False),
+        (
+            {"constant_amax": 1.0, "axis": None},
+            {"constant_amax": 1.0, "axis": 0},
+            False,
+        ),
+        ({"constant_amax": 1.0}, {"constant_amax": 1.0}, True),
+    ],
+    ids=("constant-amax", "axis", "calibrated-amax"),
+)
+def test_kv_autoquant_checkpoint_rejects_changed_preceding_state(
+    tmp_path, first_kwargs, second_kwargs, mutate_second_amax
+):
     candidate = _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump()
     data = [{"input_ids": torch.randint(0, 16, (1, 8))}]
     checkpoint = str(tmp_path / "kv_search.pth")
 
     mtq.auto_quantize(
-        quantized_model((4, 3)),
+        _quantized_weight_baseline((4, 3), **first_kwargs),
         {"effective_bits": 8.0, "cost_model": "kv_cache"},
         [candidate],
         data,
@@ -843,9 +901,15 @@ def quantized_model(bits):
         checkpoint=checkpoint,
     )
 
+    second_model = _quantized_weight_baseline((4, 3), **second_kwargs)
+    if mutate_second_amax:
+        for name, module in second_model.named_modules():
+            if name.endswith("weight_quantizer") and module.is_enabled:
+                module.amax = module.amax * 2
+
     with pytest.raises(ValueError, match="preceding non-K/V quantizer"):
         mtq.auto_quantize(
-            quantized_model((2, 1)),
+            second_model,
             {"effective_bits": 8.0, "cost_model": "kv_cache"},
             [candidate],
             data,