[OMNIML-5570, OMNIML-5569] 1/2 Add layer-wise KV-cache AutoQuant with forward KL (#2272)

### What does this PR do?

Type of change: new feature.

Adds standalone layer-wise KV-cache AutoQuantize through the existing
public
`mtq.auto_quantize` API:

- dispatches KV search with
`constraints={"effective_bits": ..., "cost_model": "kv_cache"}` and
forward-KL
  sensitivity;
- selects one supported K/V format for every eligible causal-attention
layer;
- supports persistent/exportable FP8 K/V, NVFP4 K/V, and FP8-K/NVFP4-V
candidates;
- solves a K/V-width- and scale-storage-aware additive recipe with the
existing
  PuLP-backed constrained solver;
- uses `BaseSearcher` lifecycle and safe checkpoint restore/save
machinery;
- preserves existing non-KV execution while isolating K/V candidate
calibration;
- returns standard AutoQuantize state that can be re-solved at another
KV budget;
- produces a complete KV-only replay config that disables every non-KV
quantizer;
- saves JSON-safe sensitivity metadata and the exact selected layer
mapping; and
- invokes the public API from `examples/hf_ptq/hf_ptq.py` through a
standalone
  calibration-free recipe.

The implementation is architecture-driven. Plain and
conditional-generation Qwen
causal attention is supported, VLM vision attention is excluded through
the existing
language-model extraction boundary, hybrid full-attention mixers are
discovered through
their paired K/V quantizers, and nonattention/Mamba modules remain
outside the search.
Ambiguous language-model roots, unsupported distributed execution,
structural
algorithms, invalid storage declarations, nonpersistent scales, and
unsupported K/V
pairs fail closed.

KV-only unified HF exports leave weight-quantization fields unset.
Uniform all-FP8 or
all-NVFP4 selections retain their legacy KV scheme while also carrying
the complete
`kv_cache_quantized_layers` map and schema version; genuinely
layer-mixed selections use
the KV-side `MIXED_PRECISION` marker plus the same map. This keeps
weight-loader metadata
accurate and prevents disabled vision attention from making uniform
language-model KV
quantization appear partially quantized.

GEMM PTQ/AutoQuantize followed by KV AutoQuantize is intentionally
excluded and proposed
separately in stacked PR #2273.

### Why KV search has a dedicated backend

The user-facing entry point remains `mtq.auto_quantize`; no separate
public KV search API
is introduced. `AutoQuantizeKVSearcher` extends `BaseSearcher` and
reuses its reset,
checkpoint load/save, and search lifecycle, along with existing Pydantic
configuration,
calibration, safe checkpoint I/O, and PuLP-backed selection utilities.

The backend remains KV-specific because a decision owns paired K/V
quantizers on one
attention layer, its cost depends on separate K/V widths and data/scale
storage, BF16 is a
scoring reference but not a deployable solver choice, and the
optimization objective is
additive isolated forward KL under a KV-storage constraint. These
contracts do not match
the weight-domain hparam grouping, parameter-count cost, or
threshold-selection behavior
of the existing weight AutoQuant searchers. Keeping the specialization
behind the shared
API avoids changing established weight-search solver and scoring
behavior.

### Usage

```bash
python examples/hf_ptq/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 \
  --export_path /path/to/qwen3.8-27b-mixed-kv
```

The search checkpoint is compatible only with the same model,
eligible-layer geometry,
candidate configurations, and scoring setup. Use a distinct checkpoint
path after any of
those inputs change.

KV-cache AutoQuantize rejects `--use_fsdp2` before model loading because
its sensitivity
scoring, selection, and checkpoint writes are single-process. Existing
weight
AutoQuantize retains its previous experimental FSDP2 warning and
behavior.

### Testing

- Focused coverage exercises candidate validation/calibration, paired
K/V scoring and
storage accounting, solving, checkpoint resume, failure atomicity,
disabled layers,
fresh-model replay, Qwen/VLM/hybrid boundaries, JSON-safe reports, and
unified export.
- Uniform FP8/NVFP4 KV-only exports retain the legacy KV scheme and
complete layer map
without claiming a weight algorithm; disabled VLM vision attention is
excluded from
  causal-KV eligibility.
- The shipped recipe runs end to end on a tiny offline Qwen fixture and
preserves
  exportable scale state.
- After merging current `main`: 432 focused recipe/KV/export/hf_ptq
tests passed, with one
unrelated optional-dependency skip; changed-file pre-commit hooks
passed.

### Deployment gate

The producer schema is covered here. Runtime consumption of
`kv_cache_quantized_layers` is tracked in vLLM PR
https://github.com/vllm-project/vllm/pull/52813. Do not treat a produced
checkpoint as
runtime-supported until that consumer lands and the target K/V kernels
are available.

### Before your PR is "Ready for review"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow
  guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update Changelog?: ✅

### Additional information

- This is split from the combined ground-truth implementation in draft
PR #2211 to reduce
  review scope; composition is isolated in #2273.
- The standalone core tree contains no composed GEMM→KV recipe schema or
orchestration.
- No model-name checks, checkpoint-specific layer lists, campaign data
contracts, cluster
  launch logic, or runtime-kernel implementations are included.

---------

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

CHANGELOG.rst

@@ -44,6 +44,7 @@ Changelog
 - Add opt-in FP8 Vision Encoder recipes under the ``qwen3_vl`` and ``qwen3_5`` model types. The vision-only recipe keeps the language model and KV cache in high precision; the joint recipe quantizes Vision Encoder and language-model Linears and uses FP8 KV-cache cast. Both quantize primary and deepstack merger Linears where present, while leaving patch embedding and vision-attention BMMs in high precision. Exported checkpoints require an inference runtime that supports quantized Vision Encoder Linears.
 - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported.
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint.
+- Add layer-wise KV-cache AutoQuantize through ``mtq.auto_quantize`` with ``cost_model="kv_cache"``. It measures isolated full-vocabulary forward KL for caller-supplied K/V formats, solves a width-weighted additive storage-constrained recipe across eligible layers, preserves search-disabled layers in their existing format, exports the selected per-attention mapping in unified HF checkpoints, returns sensitivity metadata in the search state, and supports resumable search checkpoints. A cast-mode FP8/NVFP4 recipe at 5.4 bits/scalar is included.
 
 *Speculative Decoding*
 

examples/hf_ptq/README.md

@@ -477,8 +477,44 @@ 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 use the same `mtq.auto_quantize` API and set
+`constraints.cost_model: kv_cache` with an `effective_bits` target. Their
+`candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
+packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are
+preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice.
+The shipped recipe searches FP8-cast K/V (8.0 bits/scalar) and NVFP4-cast K/V
+(4.5 bits/scalar) at 5.4 bits/scalar. It intentionally excludes FP8-K/NVFP4-V because the
+companion vLLM implementation does not support that asymmetric per-layer format:
+
+```bash
+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 \
+  --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,
+including the candidate quantizer tensors needed for replay.
+
+KV sensitivity scoring runs one reference forward plus one forward per eligible-layer candidate
+for every scoring step. Choose the recipe's `auto_quantize.score_size` with the number of eligible
+layers and candidates in mind. Peak scoring memory also grows with the number of selected tokens
+times the model vocabulary because reference and candidate log probabilities are computed in FP32.
+
+> [!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.
+> 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):
@@ -515,6 +551,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.
 
+> *KV-cache AutoQuantize recipes are not supported with `--use_fsdp2` and are rejected before model loading. Distributed KV sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Existing weight AutoQuantize recipes retain their previous experimental warning with FSDP2.*
+
 ### Usage
 
 #### Slurm (recommended)

examples/hf_ptq/hf_ptq.py

@@ -103,6 +103,28 @@ from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloa
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
 RAND_SEED = 1234
+_FSDP2_KV_AUTOQUANT_ERROR = (
+    "KV-cache AutoQuantize does not support --use_fsdp2 until distributed sensitivity scoring, "
+    "selection, and checkpoint writes are synchronized across ranks."
+)
+_FSDP2_AUTOQUANT_WARNING = (
+    "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet "
+    "(distributed calibration, sensitivity scoring, and recipe/checkpoint "
+    "synchronization across ranks); use at your own risk."
+)
+
+
+def _select_unpadded_logits(logits: torch.Tensor, batch: dict[str, Any]) -> torch.Tensor:
+    """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(
+            "AutoQuantize 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:
@@ -135,7 +157,7 @@ def extract_and_prepare_language_model_from_vl(full_model):
     Returns:
         tuple: (language_model, model_type) or (None, None) if not a VLM
     """
-    language_model_lineage = get_language_model_from_vl(full_model)
+    language_model_lineage = get_language_model_from_vl(full_model, strict=True)
     if language_model_lineage is not None:
         language_model = language_model_lineage.pop(-1)
         ancestors = language_model_lineage
@@ -255,7 +277,7 @@ def make_calib_dataloader(
         )
     else:
         assert tokenizer is not None and isinstance(
-            tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)
+            tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast
         ), "The PreTrainedTokenizer must be set"
         # Labels are only needed for gradient-based auto_quantize
         include_labels = autoquant_gradient_recipe
@@ -356,6 +378,18 @@ 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.cost_model == "kv_cache"
+    if is_kv_search:
+        return {
+            "search_domain": "kv_cache",
+            "constraints": constraints,
+            "quantization_formats": [
+                fmt.model_dump(exclude_none=True) for fmt in 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.
@@ -387,6 +421,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,
@@ -420,16 +455,13 @@ def auto_quantize(
         "Auto Quantization is not supported for pipeline parallel size > 1"
     )
 
-    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."
-        )
-
     inputs = _mtq_inputs_from_auto_quantize_config(
         aq_config, args, fixed_quantize_config=fixed_quantize_config
     )
+    if args.use_fsdp2:
+        if inputs["search_domain"] == "kv_cache":
+            raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR)
+        warnings.warn(_FSDP2_AUTOQUANT_WARNING)
 
     # base-model lm_head handling (mirrors the CLI helper)
     is_base_model = (
@@ -468,30 +500,47 @@ def auto_quantize(
             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
+            if inputs["search_domain"] == "kv_cache":
+                return _select_unpadded_logits(logits, batch)
+            return logits
 
     else:
         raise ValueError(
             f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'"
         )
 
+    auto_quantize_kwargs: dict[str, Any] = {
+        "constraints": inputs["constraints"],
+        "data_loader": calib_dataloader,
+        "forward_step": forward_step,
+        "quantization_formats": inputs["quantization_formats"],
+        "num_calib_steps": len(calib_dataloader),
+        "num_score_steps": min(
+            len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)
+        ),
+        "verbose": True,
+        "disabled_layers": inputs["disabled_layers"],
+        "method": inputs["method"],
+        "checkpoint": args.auto_quantize_checkpoint,
+    }
+    if inputs["search_domain"] == "weight":
+        auto_quantize_kwargs.update(
+            {
+                "loss_func": loss_func,
+                "fixed_quantization_config": inputs["fixed_quantization_config"],
+                "module_search_spaces": inputs["module_search_spaces"],
+            }
+        )
+
     language_model, _ = mtq.auto_quantize(
         language_model,
-        constraints=inputs["constraints"],
-        data_loader=calib_dataloader,
-        forward_step=forward_step,
-        loss_func=loss_func,
-        quantization_formats=inputs["quantization_formats"],
-        fixed_quantization_config=inputs["fixed_quantization_config"],
-        module_search_spaces=inputs["module_search_spaces"],
-        num_calib_steps=len(calib_dataloader),
-        num_score_steps=min(len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)),
-        verbose=True,
-        disabled_layers=inputs["disabled_layers"],
-        method=inputs["method"],
-        checkpoint=args.auto_quantize_checkpoint,
+        **auto_quantize_kwargs,
     )
+    if inputs["search_domain"] == "kv_cache":
+        return language_model
 
     # KV cache quantization is uniform; applied after the LP search.
     kv_cache_quant_cfg = inputs["kv_cache_quant_cfg"]
@@ -516,6 +565,17 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool:
     return recipe is not None and isinstance(load_recipe(recipe), ModelOptAutoQuantizeRecipe)
 
 
+def _recipe_is_kv_auto_quantize(recipe: str | None) -> bool:
+    """True if ``recipe`` resolves to a KV AutoQuantize recipe (peeked before model load)."""
+    if recipe is None:
+        return False
+    loaded_recipe = load_recipe(recipe)
+    return (
+        isinstance(loaded_recipe, ModelOptAutoQuantizeRecipe)
+        and loaded_recipe.auto_quantize.constraints.cost_model == "kv_cache"
+    )
+
+
 def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None:
     """Require image-text calibration when a PTQ recipe enables visual input quantizers."""
     if not isinstance(recipe, ModelOptPTQRecipe) or args.calib_with_images:
@@ -539,6 +599,8 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None:
 def load_model(args: argparse.Namespace):
     # If low memory mode is enabled, we compress the model while loading the HF checkpoint.
     calibration_only = False
+    if args.use_fsdp2 and _recipe_is_kv_auto_quantize(args.recipe):
+        raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR)
     if args.use_fsdp2:
         hf_config = AutoConfig.from_pretrained(
             args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code
@@ -722,7 +784,7 @@ def sparsity_main(
     # Different calibration datasets are also available, e.g., "pile" and "wikipedia"
     # Please also check the docstring for the datasets available
     assert tokenizer is not None and isinstance(
-        tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)
+        tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast
     ), "The PreTrainedTokenizer must be set"
     calib_dataloader = get_dataset_dataloader(
         dataset_name=args.dataset,
@@ -1189,7 +1251,7 @@ def quantize_main(
     if args.recipe is not None:
         print(f"Use recipe {args.recipe} for quantization")
         recipe = load_recipe(args.recipe)
-        if not isinstance(recipe, (ModelOptPTQRecipe, ModelOptAutoQuantizeRecipe)):
+        if not isinstance(recipe, ModelOptPTQRecipe | ModelOptAutoQuantizeRecipe):
             raise TypeError(
                 f"Expected PTQ or AutoQuantize recipe, but got {type(recipe).__name__} "
                 f"from {args.recipe}"
@@ -1450,9 +1512,10 @@ def parse_args() -> argparse.Namespace:
         help=(
             "PTQ or AutoQuantize recipe YAML file or name without suffix (e.g. "
             "general/ptq/nvfp4_default-kv_fp8_cast, general/auto_quantize/nvfp4_fp8_at_4p8bits). "
-            "KV cache source depends on the recipe type: PTQ recipes bake KV cache into quant_cfg "
-            "and --kv_cache_qformat is ignored; AutoQuantize recipes fall back to --kv_cache_qformat "
-            "unless the recipe sets an explicit kv_cache field."
+            "KV cache behavior depends on the recipe type: PTQ recipes configure it in quant_cfg "
+            "and ignore --kv_cache_qformat; weight AutoQuantize recipes use their kv_cache setting "
+            "or fall back to --kv_cache_qformat; KV-cache AutoQuantize recipes select per-layer K/V "
+            "formats from candidate_formats and ignore --kv_cache_qformat."
         ),
         default=None,
     )

modelopt/recipe/config.py

@@ -151,13 +151,18 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
 
     effective_bits: float = ModeloptField(
         default=4.8,
-        title="Effective bits per weight",
-        description="Average weight-storage bits target for the LP, in (0, 16].",
+        title="Effective bits",
+        description=(
+            "Average storage-bits target for the selected cost model, in (0, 16]. Defaults to 4.8."
+        ),
     )
-    cost_model: Literal["weight", "active_moe"] = ModeloptField(
+    cost_model: Literal["weight", "active_moe", "kv_cache"] = ModeloptField(
         default="weight",
         title="Cost model",
-        description="'weight' counts all weights equally; 'active_moe' scales routed-expert weights.",
+        description=(
+            "'weight' counts all weights equally; 'active_moe' scales routed-expert weights; "
+            "'kv_cache' accounts for paired K/V-cache storage."
+        ),
     )
     cost: AutoQuantizeCost | None = ModeloptField(
         default=None,
@@ -172,6 +177,12 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
             raise ValueError(f"effective_bits must be in (0, 16], got {v}")
         return v
 
+    @model_validator(mode="after")
+    def _validate_cost_settings(self):
+        if self.cost_model == "kv_cache" and self.cost is not None:
+            raise ValueError("KV-cache AutoQuant does not accept weight cost settings.")
+        return self
+
 
 class AutoQuantizeModuleSearchSpace(ModeloptBaseConfig):
     """Candidate formats selectable for modules matching one or more name patterns."""
@@ -272,6 +283,25 @@ class AutoQuantizeConfig(ModeloptBaseConfig):
                 "auto_quantize requires candidate_formats or at least one module_search_spaces "
                 "entry. For uniform quantization, use a PTQ recipe instead."
             )
+        if self.constraints.cost_model == "kv_cache":
+            if self.auto_quantize_method != "kl_div":
+                raise ValueError(
+                    "KV-cache AutoQuant currently requires auto_quantize_method=kl_div."
+                )
+            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."
+                )
+            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
 
 

modelopt/torch/export/convert_hf_config.py

@@ -269,9 +269,18 @@ 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 in ("MIXED_PRECISION", "FP8_K_NVFP4_V"):
+            new_config["kv_cache_quant_algo"] = kv_cache_quant_algo
         else:
             # TODO: Handle other kv cache quantization algorithms
             new_config["kv_cache_scheme"] = kv_cache_quant_algo
+        if "kv_cache_quantized_layers" in original_quantization_details:
+            new_config["kv_cache_quantized_layers"] = original_quantization_details[
+                "kv_cache_quantized_layers"
+            ]
+            new_config["kv_cache_schema_version"] = original_quantization_details.get(
+                "kv_cache_schema_version", 1
+            )
 
     producer_info = input_config.get("producer")
     if producer_info:

modelopt/torch/export/layerwise_export.py

@@ -38,7 +38,12 @@ from .layer_utils import is_moe, sync_moe_gate_up_amax
 from .model_utils import TiedWeightMap, get_language_model_from_vl
 from .quant_aware_conversion import build_reverse_name_mapper, revert_quant_config_names
 from .quant_format import FUSION_FREE_FORMATS, QUANTIZATION_NVFP4
-from .quant_utils import _postprocess_single_tensor, get_quant_config, get_quantization_format
+from .quant_utils import (
+    _get_kv_cache_postprocess_config,
+    _postprocess_single_tensor,
+    get_quant_config,
+    get_quantization_format,
+)
 from .registry import ExportContext, PrepareMoEInputsRegistry
 from .unified_export_hf import (
     _add_mtp_exclusions,
@@ -247,7 +252,7 @@ class LayerwiseExporter:
         # it goes, so by finalize() the model would look unquantized.
         self._quant_config = get_quant_config(model, is_modelopt_qlora=self._ctx.is_modelopt_qlora)
         # Not get_kv_cache_dtype: it does not recurse, so on the root it answers None.
-        self._kv_cache_format = self._quant_config["quantization"]["kv_cache_quant_algo"]
+        self._kv_cache_format = _get_kv_cache_postprocess_config(self._quant_config["quantization"])
 
         self._name_mapper = None
         try:

modelopt/torch/export/model_utils.py

@@ -121,7 +121,7 @@ def is_multimodal_model(model):
     )
 
 
-def get_language_model_from_vl(model) -> list[nn.Module] | None:
+def get_language_model_from_vl(model, *, strict: bool = False) -> list[nn.Module] | None:
     """Extract the language model lineage from a Vision-Language Model (VLM).
 
     This function handles the common patterns for accessing the language model component
@@ -130,6 +130,12 @@ def get_language_model_from_vl(model) -> list[nn.Module] | None:
 
     Args:
         model: The VLM model instance to extract the language model from
+        strict: Raise if distinct direct and nested language-model roots are present. Generic
+            export callers retain the historical nested-root preference by default; search
+            boundaries can opt in to fail-closed ambiguity handling.
+
+    Raises:
+        ValueError: If ``strict`` is True and the model exposes competing language-model roots.
 
     Returns:
         list: the lineage path towards the language model
@@ -140,12 +146,22 @@ 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
-    if hasattr(model, "model") and hasattr(model.model, "language_model"):
-        return [model, model.model, model.model.language_model]
-
-    if hasattr(model, "language_model"):
-        return [model, model.language_model]
+    nested_parent = getattr(model, "model", None)
+    nested_language_model = getattr(nested_parent, "language_model", None)
+    direct_language_model = getattr(model, "language_model", None)
+    if (
+        strict
+        and nested_language_model is not None
+        and direct_language_model is not None
+        and nested_language_model is not direct_language_model
+    ):
+        raise ValueError(
+            "Found multiple language-model roots; refusing to select one by traversal order."
+        )
+    if nested_language_model is not None:
+        return [model, nested_parent, nested_language_model]
+    if direct_language_model is not None:
+        return [model, direct_language_model]
 
     # 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

modelopt/torch/export/quant_aware_conversion.py

@@ -278,7 +278,7 @@ def build_reverse_name_mapper(model):
 
 
 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(

modelopt/torch/export/quant_format.py

@@ -43,6 +43,7 @@ QUANTIZATION_FP8_PC_PT = "fp8_pc_pt"
 FUSION_FREE_FORMATS = frozenset({QUANTIZATION_FP8, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL})
 
 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"

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
 
@@ -48,9 +47,10 @@ from modelopt.torch.quantization.utils import (
 from modelopt.torch.utils import clear_cuda_cache
 
 from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer
-from .model_utils import TiedWeightMap
+from .model_utils import TiedWeightMap, get_language_model_from_vl
 from .quant_format import (
     KV_CACHE_FP8,
+    KV_CACHE_FP8_K_NVFP4_V,
     KV_CACHE_INT8,
     KV_CACHE_NVFP4,
     KV_CACHE_NVFP4_AFFINE,
@@ -76,6 +76,11 @@ from .quant_format import (
 logger = logging.getLogger(__name__)
 
 
+def _has_large_fp8_scale(value: torch.Tensor) -> bool:
+    """Return whether an FP8 scale exceeds the recommended threshold."""
+    return bool(torch.any(value > 0.5))
+
+
 def maybe_transpose_expert_weight_dimensions(
     weight: torch.Tensor,
     weight_scale: torch.Tensor | None = None,
@@ -307,27 +312,38 @@ def get_kv_cache_scaling_factor(
         for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer")
     ]
 
-    # For FP8, we recommend default kv cache scaling factor to be 1.
-    if clamp_fp8_scales and 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 not clamp_fp8_scales:
+        fp8_indices = ()
+    elif 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 _has_large_fp8_scale(factor):
+            warn(
+                "Warning: Large KV activation detected. "
+                "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.
@@ -342,6 +358,29 @@ 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
+            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
         for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer", "output_quantizer"):
@@ -931,7 +970,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.
@@ -965,12 +1004,15 @@ 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_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:
+                if layer_quantization == KV_CACHE_FP8 and _has_large_fp8_scale(value):
                     logger.warning(
                         "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                     )
@@ -981,10 +1023,40 @@ 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 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 quantization
+
+
+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,
-    quantization: str | None,
+    quantization: str | dict[str, dict[str, str]] | None,
     is_modelopt_qlora: bool = False,
     tied_map: "TiedWeightMap | None" = None,
 ) -> dict:
@@ -993,7 +1065,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
@@ -1031,15 +1104,18 @@ def postprocess_state_dict(
                 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 = _resolve_kv_cache_format_for_key(key, quantization)
+                    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 _has_large_fp8_scale(value):
                         logger.warning(
                             "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                         )
@@ -1531,7 +1607,7 @@ def get_quant_config(
     block_size = None
 
     # Create base config
-    quant_config = {
+    quant_config: dict[str, Any] = {
         "producer": {
             "name": "modelopt",
             "version": __version__,
@@ -1549,7 +1625,14 @@ 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]] = {}
+    language_model_lineage = get_language_model_from_vl(model)
+    language_model_modules = (
+        None
+        if language_model_lineage is None
+        else {id(module) for module in language_model_lineage[-1].modules()}
+    )
     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))
@@ -1602,21 +1685,34 @@ 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
-        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
-        ):
+        is_language_model_module = (
+            language_model_modules is None or id(module) in language_model_modules
+        )
+        k_quantizer = getattr(module, "k_bmm_quantizer", None)
+        v_quantizer = getattr(module, "v_bmm_quantizer", None)
+        output_quantizer = getattr(module, "output_quantizer", None)
+        # Projection modules also expose output_quantizer; it is a KV fallback only on an
+        # attention boundary that has at least one K/V quantizer attribute.
+        is_kv_boundary = k_quantizer is not None or v_quantizer is not None
+        enabled_kv_quantizers = [
+            quantizer
+            for quantizer in (k_quantizer, v_quantizer, output_quantizer)
+            if is_kv_boundary and quantizer is not None and quantizer.is_enabled
+        ]
+        if enabled_kv_quantizers and is_language_model_module:
             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)
+                # Per-layer mixed-KV metadata is defined only for an actual enabled K/V pair.
+                # Keep the output-quantizer/single-sided Megatron fallback top-level only.
+                if (
+                    k_quantizer is not None
+                    and v_quantizer is not None
+                    and k_quantizer.is_enabled
+                    and v_quantizer.is_enabled
+                ):
+                    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
@@ -1629,8 +1725,26 @@ 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
+    weight_quant_algo = quant_config["quantization"].get("quant_algo")
+    needs_layerwise_kv_metadata = bool(kv_cache_quantized_layers) and (
+        weight_quant_algo is None or len(kv_cache_formats) > 1
+    )
+    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."
+            )
+        # KV metadata is orthogonal to weight metadata. In particular, a KV-only search
+        # must preserve BF16 weights instead of synthesizing a weight quantization algorithm.
+        quant_config["quantization"]["kv_cache_quant_algo"] = (
+            next(iter(kv_cache_formats)) if len(kv_cache_formats) == 1 else "MIXED_PRECISION"
+        )
+        quant_config["quantization"]["kv_cache_quantized_layers"] = kv_cache_quantized_layers
+        quant_config["quantization"]["kv_cache_schema_version"] = 1
+    elif len(kv_cache_formats) == 1:
+        # Preserve the pre-AutoQuantize uniform KV schema, including partial coverage.
+        quant_config["quantization"]["kv_cache_quant_algo"] = next(iter(kv_cache_formats))
 
     return quant_config
 

modelopt/torch/export/unified_export_hf.py

@@ -16,6 +16,7 @@
 """Code that export quantized Hugging Face models for deployment."""
 
 import contextlib
+import copy
 import importlib
 import json
 import re
@@ -110,6 +111,7 @@ from .quant_format import (
     QUANTIZATION_W4A16_NVFP4,
 )
 from .quant_utils import (
+    _get_kv_cache_postprocess_config,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -1060,12 +1062,12 @@ 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"]
+    kv_cache_postprocess_config = _get_kv_cache_postprocess_config(quant_config["quantization"])
 
     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,
     )
@@ -1564,7 +1566,16 @@ def _write_hf_export_config(
         json.dump(config_data, file, indent=4)
 
 
-def _revert_quant_config_names_best_effort(model: nn.Module, hf_quant_config: dict | None) -> None:
+def _revert_hf_quant_config_names(hf_quant_config: dict, name_mapper: Callable[[str], str]) -> dict:
+    """Return a name-reverted copy, leaving the input untouched if mapping fails."""
+    mapped_quant_config = copy.deepcopy(hf_quant_config)
+    revert_quant_config_names(mapped_quant_config.get("quantization", {}), name_mapper)
+    return mapped_quant_config
+
+
+def _revert_quant_config_names_best_effort(
+    model: nn.Module, hf_quant_config: dict | None
+) -> dict | None:
     """Rename the quant config's modules back to their original checkpoint names.
 
     On failure it warns and keeps the current names, so the config still matches the weights.
@@ -1572,12 +1583,13 @@ def _revert_quant_config_names_best_effort(model: nn.Module, hf_quant_config: di
     try:
         name_mapper = build_reverse_name_mapper(model)
         if name_mapper is not None and hf_quant_config:
-            revert_quant_config_names(hf_quant_config.get("quantization", {}), name_mapper)
+            return _revert_hf_quant_config_names(hf_quant_config, name_mapper)
     except Exception as exc:
         warnings.warn(
             f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor "
             "names may not match the original HF hub checkpoint."
         )
+    return hf_quant_config
 
 
 def export_hf_checkpoint(
@@ -1673,7 +1685,7 @@ def export_hf_checkpoint(
             )
             if getattr(model, "hf_quantizer", None) is not None:
                 model.hf_quantizer = None
-            _revert_quant_config_names_best_effort(model, hf_quant_config)
+            hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config)
         elif is_fsdp2_sharded:
             # FSDP2 multi-rank: stream each rank's owned units straight to its own shard files, so
             # a rank buffers its own share of the model rather than the whole checkpoint, and the
@@ -1693,7 +1705,7 @@ def export_hf_checkpoint(
             if rank == 0:
                 if save_modelopt_state and ModeloptStateManager.is_converted(model):
                     torch.save(modelopt_state(model), export_dir / _MODELOPT_STATE_SAVE_NAME)
-                _revert_quant_config_names_best_effort(model, hf_quant_config)
+                hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config)
         else:
             post_state_dict, hf_quant_config = _export_transformers_checkpoint(
                 model, dtype, **kwargs
@@ -1716,13 +1728,19 @@ def export_hf_checkpoint(
             # do it here. The same rename is applied to the quant-config module references
             # (exclude_modules / quantized_layers keys) so a deployment loader matches them against
             # the reverted hub-named modules (otherwise an excluded BF16 layer is loaded as quantized
-            # and fails). Best-effort: any failure (an op we cannot reverse yet, transformers API
-            # drift, unexpected shapes) falls back to the in-memory names.
+            # 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.
             try:
                 name_mapper = build_reverse_name_mapper(model)
-                export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict)
+                mapped_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict)
+                mapped_quant_config = hf_quant_config
                 if name_mapper is not None and hf_quant_config:
-                    revert_quant_config_names(hf_quant_config.get("quantization", {}), name_mapper)
+                    mapped_quant_config = _revert_hf_quant_config_names(
+                        hf_quant_config, name_mapper
+                    )
+                export_state_dict = mapped_state_dict
+                hf_quant_config = mapped_quant_config
             except Exception as exc:
                 warnings.warn(
                     f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor "

modelopt/torch/export/unified_export_hf_streaming.py

@@ -44,7 +44,11 @@ from modelopt.torch.utils import distributed as _dist
 
 from .model_utils import get_export_units
 from .quant_aware_conversion import _build_reverse_rules, 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,
@@ -282,7 +286,7 @@ def _make_tensor_sink(
     name_mapper,
     tied_alias_keys: set[str],
     kv_cache_max_bound: float,
-    kv_cache_format: str | None,
+    kv_cache_format: str | dict[str, dict[str, str]] | None,
     is_modelopt_qlora: bool,
 ):
     """Build the per-tensor step both streaming exporters use.
@@ -373,7 +377,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.
@@ -639,7 +643,7 @@ def _export_fsdp2_checkpoint_streaming(
     dtype, tied_map, quant_config = _prepare_model_for_export(model, dtype, is_modelopt_qlora)
 
     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 weights are dropped by name: with one unit in hand at a time there is no whole-dict
     # view to compare storage against. tied_map covers dict-style and MoE ties.

modelopt/torch/quantization/_auto_quantize_cost.py

@@ -31,6 +31,7 @@ ACTIVE_MOE_EXPERT_RATIO_KEY: Final = "active_moe_expert_ratio"
 EXCLUDED_MODULE_NAME_PATTERNS_KEY: Final = "excluded_module_name_patterns"
 COST_MODEL_WEIGHT: Final = "weight"
 COST_MODEL_ACTIVE_MOE: Final = "active_moe"
+COST_MODEL_KV_CACHE: Final = "kv_cache"
 
 _ROUTED_MOE_EXPERT_NAME_RE = re.compile(r"(^|\.)experts(\.|$)")
 _ACTIVE_MOE_TOP_K_ATTRS = (
@@ -68,7 +69,7 @@ def _iter_model_configs(model: nn.Module):
 def _get_first_numeric_config_attr(config: Any, attr_names: tuple[str, ...]) -> float | None:
     for attr_name in attr_names:
         value = getattr(config, attr_name, None)
-        if isinstance(value, (int, float)) and not isinstance(value, bool):
+        if isinstance(value, int | float) and not isinstance(value, bool):
             return float(value)
     return None
 
@@ -185,7 +186,7 @@ class ActiveMoECostModel(AutoQuantizeCostModel):
                 )
 
         if not (
-            isinstance(active_moe_expert_ratio, (int, float))
+            isinstance(active_moe_expert_ratio, int | float)
             and not isinstance(active_moe_expert_ratio, bool)
             and 0.0 < active_moe_expert_ratio <= 1.0
         ):
@@ -206,9 +207,22 @@ class ActiveMoECostModel(AutoQuantizeCostModel):
         return base_weight
 
 
+class KVCacheCostModel(AutoQuantizeCostModel):
+    """Account for the exact storage of paired K/V-cache candidates."""
+
+    name = COST_MODEL_KV_CACHE
+    supported_cost_keys = frozenset()
+
+    @staticmethod
+    def candidate_cost(k_width: int, v_width: int, k_bits: float, v_bits: float) -> float:
+        """Return candidate storage in bits for one attention layer and token."""
+        return k_width * k_bits + v_width * v_bits
+
+
 _COST_MODELS: Final = {
     COST_MODEL_WEIGHT: WeightCostModel(),
     COST_MODEL_ACTIVE_MOE: ActiveMoECostModel(),
+    COST_MODEL_KV_CACHE: KVCacheCostModel(),
 }
 
 

modelopt/torch/quantization/kv_cache_auto_quant.py

@@ -0,0 +1,938 @@
+# 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 AutoQuantize using isolated forward KL sensitivity."""
+
+from __future__ import annotations
+
+import fnmatch
+import math
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from modelopt.torch.opt.hparam import Hparam
+from modelopt.torch.opt.searcher import LPS, BaseSearcher, SearchConfig, SearchStateDict
+from modelopt.torch.utils import print_rank_0
+
+from ._auto_quantize_cost import (
+    COST_MODEL_KV_CACHE,
+    KVCacheCostModel,
+    get_auto_quantize_cost_model,
+    normalize_auto_quantize_constraints,
+)
+from .config import QuantizeConfig, QuantizerAttributeConfig
+from .conversion import set_quantizer_by_cfg
+from .nn import TensorQuantizer
+
+__all__ = ["AutoQuantizeKVSearcher"]
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+_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:
+    quantizer = TensorQuantizer()
+    quantizer.disable()
+    return quantizer
+
+
+def _candidate_quantizers(config: QuantizeConfig) -> dict[str, TensorQuantizer]:
+    """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(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:
+            raise ValueError(
+                f"KV-cache candidate must enable {attr}; got {type(quantizer).__name__}."
+            )
+    return quantizers
+
+
+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 = _KV_CANDIDATE_NAMES | _NON_KV_PROBE_NAMES
+    for entry in config.quant_cfg:
+        if entry.parent_class is not None:
+            raise ValueError(
+                "KV-cache AutoQuantize candidates do not support parent_class filters."
+            )
+        quantizer_configs = entry.cfg if isinstance(entry.cfg, list) else [entry.cfg]
+        if any(
+            isinstance(quantizer_config, QuantizerAttributeConfig)
+            and not isinstance(quantizer_config.calibrator, str)
+            for quantizer_config in quantizer_configs
+        ):
+            raise ValueError(
+                "KV-cache AutoQuantize candidates require string calibrators so search state "
+                "remains JSON-safe and replayable."
+            )
+        matches = {name for name in probe_names if fnmatch.fnmatch(name, entry.quantizer_name)}
+        if not matches:
+            raise ValueError(
+                "KV-cache AutoQuantize 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 AutoQuantize candidates may configure only k_bmm_quantizer and "
+                f"v_bmm_quantizer; pattern {entry.quantizer_name!r} also matches "
+                f"{sorted(non_kv_matches)}."
+            )
+        matched_names.update(matches)
+    if matched_names != _KV_CANDIDATE_NAMES:
+        raise ValueError(
+            "KV-cache AutoQuantize 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 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 AutoQuantize 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 AutoQuantize candidates must use unified-export-compatible per-tensor FP8 "
+        "or block-16 dynamic NVFP4 quantizers."
+    )
+
+
+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"]
+    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 AutoQuantize candidates."
+        )
+
+    algorithm_method = _algorithm_method(config)
+    for attr, quantizer in quantizers.items():
+        if quantizer._dynamic:
+            raise ValueError(
+                f"KV-cache AutoQuantize 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(
+                f"KV-cache AutoQuantize 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 AutoQuantize 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 AutoQuantize 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 AutoQuantize 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 | None]],
+    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) - {"effective_bits", "cost_model", "cost"}
+        or constraints.get("cost_model") != COST_MODEL_KV_CACHE
+        or "effective_bits" not in constraints
+        or constraints.get("cost") not in (None, {})
+    ):
+        raise ValueError(
+            "KV-cache AutoQuantize requires an effective_bits target with "
+            f"cost_model='kv_cache'; got {constraints}."
+        )
+    target_bits = float(constraints["effective_bits"])
+    if not (0 < target_bits <= 16):
+        raise ValueError(f"effective_bits must be in (0, 16], got {target_bits}.")
+    if num_calib_steps <= 0:
+        raise ValueError("num_calib_steps must be positive.")
+    if num_score_steps <= 0:
+        raise ValueError("num_score_steps must be positive.")
+
+    candidates = []
+    seen_names = set()
+    for index, (raw_config, name) in enumerate(quantization_formats):
+        config = QuantizeConfig(**raw_config)
+        _validate_kv_only_config(config)
+        if name is None:
+            k_bits, v_bits = _candidate_kv_bits(config)
+            name = "fp8_k_nvfp4_v" if (k_bits, v_bits) == (8.0, 4.5) else f"KV_CACHE_FORMAT_{index}"
+        if name in seen_names:
+            raise ValueError(f"Duplicate KV-cache AutoQuantize candidate name: {name!r}.")
+        candidates.append((name, config))
+        seen_names.add(name)
+    if not candidates:
+        raise ValueError("KV-cache AutoQuantize 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)
+    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 _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 AutoQuantize 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]]:
+    patterns = [disabled_layers] if isinstance(disabled_layers, str) else disabled_layers or []
+    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)))
+    if not layers:
+        raise ValueError("KV-cache AutoQuantize 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)
+
+
+@contextmanager
+def _freeze_existing_quantizers(model: nn.Module, candidate_quantizers: list[TensorQuantizer]):
+    """Freeze calibration without changing existing quantizers' execution mode."""
+    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_calib))
+        module.disable_calib()
+    try:
+        yield
+    finally:
+        for quantizer, if_calib in states:
+            quantizer._if_calib = if_calib
+
+
+def _get_logits(
+    forward_step: Callable[[nn.Module, Any], torch.Tensor],
+    model: nn.Module,
+    data: Any,
+    *,
+    validate_finite: bool = True,
+) -> torch.Tensor:
+    logits = forward_step(model, data)
+    if not isinstance(logits, torch.Tensor):
+        raise TypeError("KV-cache AutoQuantize forward_step must return a logits tensor.")
+    if logits.ndim < 2 or logits.shape[-1] == 0:
+        raise ValueError(
+            "KV-cache AutoQuantize forward_step must return logits with a non-empty vocabulary "
+            "dimension."
+        )
+    if validate_finite and not torch.isfinite(logits).all():
+        raise ValueError("KV-cache AutoQuantize encountered NaN or Inf logits.")
+    return logits
+
+
+def _forward_kl(logits_quant: torch.Tensor, log_prob_ref: torch.Tensor) -> torch.Tensor:
+    """Compute the summed forward KL used by KV sensitivity scoring."""
+    return F.kl_div(
+        torch.log_softmax(logits_quant.float(), dim=-1),
+        log_prob_ref,
+        reduction="sum",
+        log_target=True,
+    )
+
+
+def _solve_additive_recipe(
+    layer_names: list[str],
+    layer_widths: list[tuple[int, int]],
+    candidate_names: list[str],
+    candidate_kv_bits: list[tuple[float, float]],
+    scores: list[list[float]],
+    target_bits: float,
+    verbose: bool,
+) -> tuple[list[int], str]:
+    cost_model = get_auto_quantize_cost_model(COST_MODEL_KV_CACHE)
+    assert isinstance(cost_model, KVCacheCostModel)
+    denominator = float(sum(k_width + v_width for k_width, v_width in layer_widths))
+    candidate_costs = [
+        [
+            cost_model.candidate_cost(k_width, v_width, k_bits, v_bits) / 16.0
+            for k_bits, v_bits in candidate_kv_bits
+        ]
+        for k_width, v_width in layer_widths
+    ]
+    max_cost = denominator * target_bits / 16.0
+    lps = LPS(
+        name="KVCacheAutoQuantize",
+        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(min(costs) for costs in candidate_costs) * 16.0 / denominator
+        raise ValueError(
+            f"KV-cache AutoQuantize could not satisfy 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 AutoQuantize 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]],
+    num_calib_steps: int,
+    num_score_steps: int,
+) -> dict[str, Any]:
+    return {
+        "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
+        "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,
+                "k_width": _projection_width(module, "k"),
+                "v_width": _projection_width(module, "v"),
+            }
+            for name, module, _ 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 _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 AutoQuantize candidate {candidate_name!r} for "
+                        f"{layer_name!r}/{attr} has no persistent export scale after calibration."
+                    )
+
+
+class QuantKVRecipeHparam(Hparam):
+    """One paired K/V format decision for an attention layer."""
+
+    def __init__(
+        self,
+        name: str,
+        module: nn.Module,
+        candidates: list[tuple[str, QuantizeConfig]],
+    ) -> None:
+        super().__init__(range(len(candidates)), original=0)
+        self.name = name
+        self.module = module
+        self.candidates = candidates
+        self.reference_quantizers = {attr: _disabled_quantizer() for attr in _KV_QUANTIZER_ATTRS}
+        self.candidate_quantizers = {
+            index: _candidate_quantizers(config) for index, (_, config) in enumerate(candidates)
+        }
+        k_width = _projection_width(module, "k")
+        v_width = _projection_width(module, "v")
+        assert k_width is not None and v_width is not None
+        self.k_width = k_width
+        self.v_width = v_width
+        self.use_reference()
+
+    @property
+    def active(self) -> int:
+        """Return the selected candidate index."""
+        assert isinstance(self._active, int)
+        return self._active
+
+    @active.setter
+    def active(self, value: int | None) -> None:
+        if value is None:
+            assert isinstance(self.original, int)
+            value = self.original
+        value = int(value)
+        assert value in self.choices
+        self._active = value
+        _apply_layer_quantizers(self.module, self.candidate_quantizers[value])
+
+    def use_reference(self) -> None:
+        """Use BF16/no-quant K/V as the scoring reference, never as a solver choice."""
+        _apply_layer_quantizers(self.module, self.reference_quantizers)
+
+    def candidate_name(self, index: int) -> str:
+        return self.candidates[index][0]
+
+    def candidate_bits(self, index: int) -> tuple[float, float]:
+        return _candidate_kv_bits(self.candidates[index][1])
+
+    def candidate_cost(self, index: int, cost_model: KVCacheCostModel) -> float:
+        k_bits, v_bits = self.candidate_bits(index)
+        return cost_model.candidate_cost(self.k_width, self.v_width, k_bits, v_bits)
+
+
+class AutoQuantizeKVSearcher(BaseSearcher):
+    """KV-cache AutoQuantize backend using the shared search/checkpoint lifecycle."""
+
+    method_name = "kl_div"
+
+    @property
+    def default_search_config(self) -> SearchConfig:
+        """Return KV-specific fields layered on the shared search configuration."""
+        config = super().default_search_config
+        config.update(
+            {
+                "quantization_formats": [],
+                "forward_step": None,
+                "num_calib_steps": 512,
+                "num_score_steps": 128,
+                "disabled_layers": None,
+            }
+        )
+        return config
+
+    @property
+    def default_state_dict(self) -> SearchStateDict:
+        """Return the checkpointed KV search state."""
+        return {
+            "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION,
+            "method": self.method_name,
+            "cost_model": COST_MODEL_KV_CACHE,
+            "search_signature": None,
+            "calibration_complete": False,
+            "num_calib_steps": 0,
+            "score_reduction": "mean_per_scored_token",
+            "num_score_steps": 0,
+            "num_scored_tokens": 0,
+            "candidates": [],
+            "layers": {},
+            "quantizer_state": {},
+            "requested_constraints": {},
+            "best": {
+                "recipe": {},
+                "constraints": {},
+                "score": float("inf"),
+                "is_satisfied": False,
+                "solver_status": None,
+            },
+        }
+
+    def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig:
+        """Validate the data inputs required by isolated forward-KL scoring."""
+        config = super().sanitize_search_config(config)
+        if config["data_loader"] is None:
+            raise ValueError("data_loader must be provided for KV-cache AutoQuantize.")
+        if config["forward_step"] is None:
+            raise ValueError("forward_step must be provided for KV-cache AutoQuantize.")
+        return config
+
+    def load_search_checkpoint(self) -> bool:
+        """Load compatible fields before validating the resolved KV search signature."""
+        return super().load_search_checkpoint(strict=False)
+
+    @property
+    def _candidate_quantizer_map(
+        self,
+    ) -> dict[str, dict[str, dict[str, TensorQuantizer]]]:
+        return {
+            hparam.name: {
+                hparam.candidate_name(index): quantizers
+                for index, quantizers in hparam.candidate_quantizers.items()
+            }
+            for hparam in self._hparams
+        }
+
+    def _calibrate_candidates(self) -> None:
+        from .model_quant import calibrate
+
+        data_loader = self.config["data_loader"]
+        forward_step = self.config["forward_step"]
+        num_calib_steps = self.config["num_calib_steps"]
+        for candidate_index, (_, config) in enumerate(self._candidates):
+            for hparam in self._hparams:
+                hparam.active = candidate_index
+
+            if config.algorithm is not None:
+
+                def calibration_loop(calibration_model):
+                    for step, data in enumerate(data_loader):
+                        if step >= num_calib_steps:
+                            break
+                        _get_logits(forward_step, calibration_model, data)
+
+                active_quantizers = [
+                    quantizer
+                    for hparam in self._hparams
+                    for quantizer in hparam.candidate_quantizers[candidate_index].values()
+                ]
+                calibration_proxy = nn.Module()
+                calibration_proxy.quantizers = nn.ModuleList(active_quantizers)
+                with _freeze_existing_quantizers(self.model, active_quantizers):
+                    calibrate(
+                        calibration_proxy,
+                        algorithm=config.algorithm,
+                        forward_loop=lambda _: calibration_loop(self.model),
+                    )
+
+            for hparam in self._hparams:
+                hparam.use_reference()
+
+        candidate_quantizers = self._candidate_quantizer_map
+        _validate_persistent_candidate_scales(candidate_quantizers)
+        self.quantizer_state = _quantizer_state_dict(candidate_quantizers)
+        self.calibration_complete = True
+        self.num_calib_steps = num_calib_steps
+        self.save_search_checkpoint(verbose=self.config["verbose"])
+
+    def before_search(self) -> None:
+        """Resolve attention decisions and calibrate or restore candidate scales."""
+        super().before_search()
+        self.constraints = normalize_auto_quantize_constraints(self.model, self.constraints)
+        target_bits, self._candidates = _validate_search_inputs(
+            self.constraints,
+            self.config["quantization_formats"],
+            self.config["num_calib_steps"],
+            self.config["num_score_steps"],
+        )
+        self._target_bits = target_bits
+        layers = _eligible_layers(self.model, self.config["disabled_layers"])
+        _validate_candidate_cost_geometry(self._candidates, layers)
+        signature = _search_signature(
+            self._candidates,
+            layers,
+            self.config["num_calib_steps"],
+            self.config["num_score_steps"],
+        )
+        if self.search_signature is not None and not _checkpoint_state_is_compatible(
+            self.state_dict(), signature
+        ):
+            raise ValueError(
+                "KV-cache AutoQuantize checkpoint does not match the current candidates, scoring "
+                "setup, or eligible layers. Use a different checkpoint path."
+            )
+        self.search_signature = signature
+        self._hparams = [
+            QuantKVRecipeHparam(name, module, self._candidates) for name, module, _ in layers
+        ]
+        self._cost_model = get_auto_quantize_cost_model(COST_MODEL_KV_CACHE)
+        assert isinstance(self._cost_model, KVCacheCostModel)
+        self.candidates = [
+            {
+                "name": name,
+                "effective_bits": config.effective_bits,
+                "k_bits": _candidate_kv_bits(config)[0],
+                "v_bits": _candidate_kv_bits(config)[1],
+                # Preserve integer block-size keys for replay. Callable calibrators are rejected,
+                # so this Python-mode representation is still safe for the exported JSON report.
+                "config": config.model_dump(mode="python", exclude_none=True),
+            }
+            for name, config in self._candidates
+        ]
+        self.model.eval()
+
+        if self.calibration_complete:
+            if not self.quantizer_state:
+                raise ValueError(
+                    "KV-cache AutoQuantize checkpoint is missing calibrated quantizer state. "
+                    "Use a different checkpoint path."
+                )
+            candidate_quantizers = self._candidate_quantizer_map
+            _restore_quantizer_state_dict(candidate_quantizers, self.quantizer_state)
+            _validate_persistent_candidate_scales(candidate_quantizers)
+            if self.config["verbose"]:
+                print_rank_0("KV-cache AutoQuantize restored calibration from checkpoint.")
+        else:
+            self._calibrate_candidates()
+
+    def _estimate_sensitivity_scores(self) -> None:
+        candidate_names = [name for name, _ in self._candidates]
+        score_sums: dict[str, dict[str, torch.Tensor | None]] = {
+            hparam.name: dict.fromkeys(candidate_names) for hparam in self._hparams
+        }
+        scored_tokens = 0
+        scored_steps = 0
+        all_logits_finite: torch.Tensor | None = None
+        iterator = tqdm(
+            self.config["data_loader"],
+            total=self.config["num_score_steps"],
+            desc="Estimating KV-cache KL sensitivity",
+            disable=not self.config["verbose"],
+        )
+        for data in iterator:
+            if scored_steps >= self.config["num_score_steps"]:
+                break
+            logits_ref = _get_logits(
+                self.config["forward_step"], self.model, data, validate_finite=False
+            )
+            batch_logits_finite = torch.isfinite(logits_ref).all()
+            log_prob_ref = torch.log_softmax(logits_ref.float(), dim=-1)
+            scored_tokens += logits_ref.numel() // logits_ref.shape[-1]
+
+            for hparam in self._hparams:
+                for candidate_index, candidate_name in enumerate(candidate_names):
+                    hparam.active = candidate_index
+                    logits_quant = _get_logits(
+                        self.config["forward_step"], self.model, data, validate_finite=False
+                    )
+                    batch_logits_finite.logical_and_(torch.isfinite(logits_quant).all())
+                    if logits_quant.shape != logits_ref.shape:
+                        raise ValueError(
+                            "KV-cache AutoQuantize forward_step returned different reference and "
+                            f"candidate logits shapes: {tuple(logits_ref.shape)} and "
+                            f"{tuple(logits_quant.shape)}."
+                        )
+                    score = _forward_kl(logits_quant, log_prob_ref)
+                    previous_score = score_sums[hparam.name][candidate_name]
+                    score_sums[hparam.name][candidate_name] = (
+                        score if previous_score is None else previous_score + score
+                    )
+                    hparam.use_reference()
+            if all_logits_finite is None:
+                all_logits_finite = batch_logits_finite
+            else:
+                all_logits_finite.logical_and_(batch_logits_finite)
+            scored_steps += 1
+
+        if scored_steps == 0 or scored_tokens == 0:
+            raise ValueError("KV-cache AutoQuantize data_loader produced no scoring batches.")
+        assert all_logits_finite is not None
+        if not all_logits_finite:
+            raise ValueError("KV-cache AutoQuantize encountered NaN or Inf logits.")
+
+        self.layers = {}
+        for hparam in self._hparams:
+            layer_score_sums = []
+            for candidate_name in candidate_names:
+                score_sum = score_sums[hparam.name][candidate_name]
+                if score_sum is None:
+                    raise RuntimeError(
+                        "KV-cache AutoQuantize did not collect a score for "
+                        f"{hparam.name!r}/{candidate_name!r}."
+                    )
+                layer_score_sums.append(score_sum)
+            layer_scores = (torch.stack(layer_score_sums) / scored_tokens).tolist()
+            scores = {}
+            for candidate_name, score in zip(candidate_names, layer_scores):
+                if not math.isfinite(score):
+                    raise ValueError(
+                        "KV-cache AutoQuantize produced a non-finite KL score for "
+                        f"{hparam.name!r}/{candidate_name!r}."
+                    )
+                scores[candidate_name] = score
+            self.layers[hparam.name] = {
+                "k_width": hparam.k_width,
+                "v_width": hparam.v_width,
+                "scores": scores,
+            }
+        self.num_score_steps = scored_steps
+        self.num_scored_tokens = scored_tokens
+        self.save_search_checkpoint(verbose=self.config["verbose"])
+
+    def _solve(self) -> None:
+        candidate_names = [name for name, _ in self._candidates]
+        candidate_kv_bits = [_candidate_kv_bits(config) for _, config in self._candidates]
+        layers = cast("dict[str, dict[str, Any]]", self.layers)
+        scores = [
+            [layers[hparam.name]["scores"][name] for name in candidate_names]
+            for hparam in self._hparams
+        ]
+        selections, status = _solve_additive_recipe(
+            [hparam.name for hparam in self._hparams],
+            [(hparam.k_width, hparam.v_width) for hparam in self._hparams],
+            candidate_names,
+            candidate_kv_bits,
+            scores,
+            self._target_bits,
+            self.config["verbose"],
+        )
+        denominator = float(sum(hparam.k_width + hparam.v_width for hparam in self._hparams))
+        cost_model = self._cost_model
+        assert isinstance(cost_model, KVCacheCostModel)
+        total_cost = sum(
+            (
+                hparam.candidate_cost(selected, cost_model)
+                for hparam, selected in zip(self._hparams, selections)
+            ),
+            start=0.0,
+        )
+        achieved_bits = total_cost / denominator
+        selected_score = sum(
+            layer_scores[selected] for selected, layer_scores in zip(selections, scores)
+        )
+        recipe = {}
+        for hparam, selected in zip(self._hparams, selections):
+            hparam.active = selected
+            selected_name = hparam.candidate_name(selected)
+            self.layers[hparam.name]["selected"] = selected_name
+            recipe[hparam.name] = selected_name
+            if self.config["verbose"]:
+                print_rank_0(f"KV-cache AutoQuantize selected {selected_name} for {hparam.name}.")
+        self.requested_constraints = {
+            "effective_bits": self._target_bits,
+            "cost_model": COST_MODEL_KV_CACHE,
+        }
+        self.best = {
+            "recipe": recipe,
+            "constraints": {
+                "effective_bits": achieved_bits,
+                "cost_model": COST_MODEL_KV_CACHE,
+            },
+            "score": selected_score,
+            "is_satisfied": achieved_bits <= self._target_bits + 1e-12,
+            "solver_status": status,
+        }
+
+    @torch.inference_mode()
+    def run_search(self) -> None:
+        """Score candidates when needed, solve the budget, and apply the selection."""
+        if not self.layers:
+            self._estimate_sensitivity_scores()
+        self._solve()
+        self.save_search_checkpoint(verbose=self.config["verbose"])
+
+
+def _config_entry_dict(entry: Any) -> dict[str, Any]:
+    if hasattr(entry, "model_dump"):
+        return entry.model_dump(mode="python", exclude_none=True)
+    return dict(entry)
+
+
+def get_kv_cache_auto_quantize_config(
+    search_state: dict[str, Any],
+    constraints: dict[str, Any] | None = None,
+    verbose: bool = False,
+) -> dict[str, Any]:
+    """Build a flat K/V quantization config, optionally re-solving at a new target."""
+    requested = constraints or search_state.get("requested_constraints")
+    if not isinstance(requested, dict):
+        raise ValueError("KV-cache AutoQuantize search state has no requested constraints.")
+    target_bits = float(requested["effective_bits"])
+    if requested.get("cost_model", COST_MODEL_KV_CACHE) != COST_MODEL_KV_CACHE:
+        raise ValueError("KV-cache search state can only be re-solved with cost_model='kv_cache'.")
+
+    candidates = search_state["candidates"]
+    candidate_names = [candidate["name"] for candidate in candidates]
+    candidate_kv_bits = [
+        (float(candidate["k_bits"]), float(candidate["v_bits"])) for candidate in candidates
+    ]
+    layers = search_state["layers"]
+    layer_names = list(layers)
+    selections, _ = _solve_additive_recipe(
+        layer_names,
+        [(layers[name]["k_width"], layers[name]["v_width"]) for name in layer_names],
+        candidate_names,
+        candidate_kv_bits,
+        [
+            [layers[name]["scores"][candidate] for candidate in candidate_names]
+            for name in layer_names
+        ],
+        target_bits,
+        verbose,
+    )
+
+    quant_cfg: list[dict[str, Any]] = [{"quantizer_name": "*", "enable": False}]
+    needs_max_calibration = False
+    for layer_name, selected in zip(layer_names, selections):
+        selected_config = candidates[selected]["config"]
+        # ``algorithm=None`` is omitted from the JSON-safe search state, while calibrated
+        # candidates retain their validated ``max`` algorithm.
+        needs_max_calibration |= selected_config.get("algorithm") is not None
+        config = QuantizeConfig(**selected_config)
+        for entry in config.quant_cfg:
+            entry_dict = _config_entry_dict(entry)
+            pattern = entry_dict["quantizer_name"]
+            for attr in _KV_QUANTIZER_ATTRS:
+                if fnmatch.fnmatch(f"{_KV_CANDIDATE_HOLDER_NAME}.{attr}", pattern):
+                    resolved = dict(entry_dict)
+                    resolved["quantizer_name"] = f"{layer_name}.{attr}"
+                    quant_cfg.append(resolved)
+    return {"quant_cfg": quant_cfg, "algorithm": "max" if needs_max_calibration else None}

modelopt/torch/quantization/model_quant.py

@@ -19,7 +19,7 @@ import fnmatch
 import inspect
 import os
 import warnings
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Sequence
 from contextlib import contextmanager
 from typing import Any, cast
 
@@ -38,9 +38,12 @@ from modelopt.torch.quantization.conversion import (
 )
 from modelopt.torch.utils import atomic_print
 
+from ._auto_quantize_cost import COST_MODEL_KV_CACHE
 from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe
 from .algorithms import get_auto_quantize_config as _get_auto_quantize_config
 from .config import QuantizeAlgoCfgType
+from .kv_cache_auto_quant import AutoQuantizeKVSearcher, get_kv_cache_auto_quantize_config
+from .kv_cache_auto_quant import _validate_search_inputs as _validate_kv_cache_search_inputs
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, SequentialQuantizer, TensorQuantizer
 from .utils import is_quantized
@@ -269,10 +272,133 @@ _AUTO_QUANTIZE_SUPPORTED_ALGORITHMS = {
 }
 
 
+def _process_quantization_formats(formats, custom_name_prefix):
+    """Resolve search formats and preserve explicitly supplied display names."""
+    processed = []
+    for index, item in enumerate(formats):
+        if item is None:
+            continue
+        if isinstance(item, tuple):
+            if len(item) != 2:
+                raise ValueError("Named quantization formats must be (config, name) pairs.")
+            quant_cfg, name = item
+            if not isinstance(name, str) or not name:
+                raise ValueError("Quantization format names must be non-empty strings.")
+        else:
+            quant_cfg = item
+            name = QuantRecipe.get_auto_name_for_config(quant_cfg)
+            if name is None:
+                name = f"{custom_name_prefix}_{index}"
+                warnings.warn(
+                    "Received custom quantization formats for search, auto_quantize results may "
+                    f"not be optimal. This config will be displayed as {name}"
+                )
+        processed.append((quant_cfg, name))
+    return processed
+
+
+def _auto_quantize_kv_cache(
+    model: nn.Module,
+    constraints: dict[str, Any],
+    quantization_formats: Sequence[dict[str, Any] | str | tuple[dict[str, Any], str]],
+    *,
+    data_loader: Iterable | None,
+    forward_step: Callable[[nn.Module, Any], Any | torch.Tensor] | None,
+    loss_func: Callable[[Any, Any], torch.Tensor] | None,
+    forward_backward_step: Callable[[nn.Module, Any], Any] | None,
+    disabled_layers: list[str] | str | None,
+    num_calib_steps: int,
+    num_score_steps: int,
+    verbose: bool,
+    method: str | None,
+    checkpoint: str | None,
+    module_search_spaces: list[dict[str, Any]] | None,
+    fixed_quantization_config: dict[str, Any] | str | None,
+):
+    """Run the KV-cache-specific AutoQuantize validation and search lifecycle."""
+    if (
+        torch.distributed.is_available()
+        and torch.distributed.is_initialized()
+        and torch.distributed.get_world_size() > 1
+    ):
+        raise RuntimeError(
+            "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:
+        raise ValueError(
+            "KV-cache AutoQuantize does not support fixed_quantization_config or "
+            "module_search_spaces."
+        )
+    if loss_func is not None or forward_backward_step is not None:
+        raise ValueError(
+            "KV-cache AutoQuantize uses forward KL and does not accept loss_func or "
+            "forward_backward_step."
+        )
+    if not quantization_formats:
+        raise ValueError("cost_model='kv_cache' requires a non-empty quantization_formats list.")
+    if data_loader is None or forward_step is None:
+        raise ValueError("data_loader and forward_step must be provided for KV-cache AutoQuantize.")
+
+    processed_kv_formats: list[tuple[dict[str, Any], str | None]] = []
+    for candidate in quantization_formats:
+        if isinstance(candidate, tuple):
+            raw_config, name = candidate
+            if not isinstance(name, str) or not name:
+                raise ValueError("KV-cache AutoQuantize candidate names must be non-empty strings.")
+        elif isinstance(candidate, str):
+            if not hasattr(mtq, candidate):
+                raise ValueError(f"Unknown KV-cache quantization format: {candidate!r}.")
+            raw_config, name = getattr(mtq, candidate), candidate
+        elif isinstance(candidate, dict):
+            raw_config = candidate
+            name = QuantRecipe.get_auto_name_for_config(candidate)
+        else:
+            raise TypeError(
+                "KV-cache quantization formats must be config dictionaries, preset names, "
+                "or (config, name) tuples."
+            )
+        if not isinstance(raw_config, dict):
+            raise TypeError("KV-cache AutoQuantize formats must resolve to config dictionaries.")
+        processed_kv_formats.append((raw_config, name))
+
+    _validate_kv_cache_search_inputs(
+        constraints,
+        processed_kv_formats,
+        num_calib_steps,
+        num_score_steps,
+    )
+    model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry)
+    set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
+    searcher = AutoQuantizeKVSearcher()
+    searcher.search(
+        model,
+        cast("ConstraintsDict", constraints),
+        config={
+            "quantization_formats": processed_kv_formats,
+            "data_loader": data_loader,
+            "forward_step": forward_step,
+            "num_calib_steps": num_calib_steps,
+            "num_score_steps": num_score_steps,
+            "disabled_layers": disabled_layers,
+            "verbose": verbose,
+            "checkpoint": checkpoint,
+        },
+    )
+    return model, searcher.state_dict()
+
+
 def auto_quantize(
     model: nn.Module,
     constraints: dict[str, Any] | None = None,
-    quantization_formats: list[dict[str, Any] | str] | None = None,
+    quantization_formats: Sequence[dict[str, Any] | str | tuple[dict[str, Any], str]] | None = None,
     data_loader: Iterable | None = None,
     forward_step: Callable[[nn.Module, Any], Any | torch.Tensor] | None = None,
     loss_func: Callable[[Any, Any], torch.Tensor] | None = None,
@@ -281,7 +407,7 @@ def auto_quantize(
     num_calib_steps: int = 512,
     num_score_steps: int = 128,
     verbose: bool = False,
-    method: str = "gradient",
+    method: str | None = None,
     checkpoint: str | None = None,
     module_search_spaces: list[dict[str, Any]] | None = None,
     fixed_quantization_config: dict[str, Any] | str | None = None,
@@ -305,9 +431,11 @@ def auto_quantize(
         model: A pytorch model with quantizer modules.
         constraints: Constraints for the search. ``effective_bits`` specifies the effective number
             of bits for the quantized model and defaults to 4.8. ``cost_model`` selects the metric
-            used for the effective-bits constraint and currently supports ``"weight"`` (default)
-            and ``"active_moe"``. Additional cost-model parameters are provided through the nested
-            ``cost`` dict.
+            used for the effective-bits constraint and supports ``"weight"`` (default),
+            ``"active_moe"``, and ``"kv_cache"``. The KV-cache cost model dispatches to isolated
+            forward-KL scoring over paired K/V formats; BF16/no-quant is its scoring reference but
+            is never solver-selectable. Additional cost-model parameters are provided through the
+            nested ``cost`` dict.
 
             Here is an example for valid ``effective_bits`` argument:
 
@@ -326,7 +454,10 @@ def auto_quantize(
                     },
                 }
 
-        quantization_formats: A list of quantization format config dictionaries or string names to search for.
+                # For paired K/V-cache formats with exact K/V storage accounting
+                constraints = {"effective_bits": 5.4, "cost_model": "kv_cache"}
+
+        quantization_formats: A sequence of quantization format config dictionaries or string names to search for.
             Each config dictionary should be valid as a ``config`` argument in
             :meth:`quantize <modelopt.torch.quantization.model_quant.quantize>`.
             The supported quantization format names are as listed by :attr:`modelopt.torch.quantization.config.choices`.
@@ -509,30 +640,40 @@ def auto_quantize(
         might not be readily deployable to TensorRT-LLM yet.
 
     """
+    if quantization_formats is not None:
+        if isinstance(quantization_formats, str) or not isinstance(quantization_formats, Sequence):
+            raise TypeError("`quantization_formats` must be a sequence of formats.")
+        quantization_formats = list(quantization_formats)
 
-    def _process_quantization_formats(formats, custom_name_prefix):
-        processed = []
-        for i, quant_cfg in enumerate(formats):
-            if quant_cfg is None:
-                continue
+    is_kv_search = constraints is not None and constraints.get("cost_model") == COST_MODEL_KV_CACHE
+    if is_kv_search:
+        assert constraints is not None
+        assert quantization_formats is not None
+        return _auto_quantize_kv_cache(
+            model,
+            constraints,
+            quantization_formats,
+            data_loader=data_loader,
+            forward_step=forward_step,
+            loss_func=loss_func,
+            forward_backward_step=forward_backward_step,
+            disabled_layers=disabled_layers,
+            num_calib_steps=num_calib_steps,
+            num_score_steps=num_score_steps,
+            verbose=verbose,
+            method=method,
+            checkpoint=checkpoint,
+            module_search_spaces=module_search_spaces,
+            fixed_quantization_config=fixed_quantization_config,
+        )
 
-            name = QuantRecipe.get_auto_name_for_config(quant_cfg)
-            if name is None:
-                name = f"{custom_name_prefix}_{i}"
-                warnings.warn(
-                    "Received custom quantization formats for search, auto_quantize results may "
-                    f"not be optimal. This config will be displayed as {name}"
-                )
-            processed.append((quant_cfg, name))
-        return processed
+    method = method or "gradient"
 
     if fixed_quantization_config is None and quantization_formats is None:
         quantization_formats = [mtq.NVFP4_AWQ_LITE_CFG, mtq.FP8_DEFAULT_CFG]
     elif fixed_quantization_config is not None and quantization_formats is None:
         quantization_formats = []
 
-    if not isinstance(quantization_formats, list):
-        raise TypeError("`quantization_formats` must be a list.")
     if fixed_quantization_config is not None and quantization_formats:
         raise ValueError(
             "`fixed_quantization_config` cannot be combined with global "
@@ -692,6 +833,8 @@ def get_auto_quantize_config(search_state, constraints=None, verbose=False):
             # fresh_model = load_model(...)
             # fresh_model = mtq.quantize(fresh_model, config, forward_loop=calibrate_loop)
     """
+    if search_state.get("cost_model") == COST_MODEL_KV_CACHE:
+        return get_kv_cache_auto_quantize_config(search_state, constraints, verbose=verbose)
     return _get_auto_quantize_config(search_state, constraints, verbose=verbose)
 
 

modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml

@@ -0,0 +1,56 @@
+# 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 search over FP8-cast and NVFP4-cast at 5.4 bits/scalar.
+# Isolated full-vocabulary forward KL is measured with every other eligible
+# attention layer kept in BF16.
+
+# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
+imports:
+  base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
+  fp8: configs/numerics/fp8
+  nvfp4: configs/numerics/nvfp4
+
+metadata:
+  recipe_type: auto_quantize
+  description: Layer-wise FP8-cast/NVFP4-cast KV-cache AutoQuantize at 5.4 bits using forward KL.
+
+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*"

tests/_test_utils/torch/transformers_models.py

@@ -224,6 +224,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]},
     }
     # Transformers 4.x names this field rope_scaling; 5.x renamed it to rope_parameters.
     # Supplying it avoids the 4.57 Qwen3-VL constructor dereferencing a None rope config.

tests/examples/hf_ptq/test_hf_ptq_args.py

@@ -21,11 +21,14 @@ from pathlib import Path
 from types import SimpleNamespace
 
 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
 from modelopt.recipe.presets import QUANT_CFG_CHOICES
+from modelopt.torch.quantization import tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
@@ -56,6 +59,19 @@ def _parse_hf_ptq_args(monkeypatch, *args):
     return hf_ptq, parsed_args
 
 
+def test_recipe_help_distinguishes_weight_and_kv_autoquant(monkeypatch, capsys):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--help"])
+
+    with pytest.raises(SystemExit) as exc_info:
+        hf_ptq.parse_args()
+
+    assert exc_info.value.code == 0
+    help_text = " ".join(capsys.readouterr().out.split())
+    assert "weight AutoQuantize recipes use their kv_cache setting" in help_text
+    assert "KV-cache AutoQuantize recipes select per-layer K/V formats" in help_text
+
+
 def test_autoquant_recipe_builds_mtq_inputs(monkeypatch):
     """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs."""
     hf_ptq, args = _parse_hf_ptq_args(
@@ -84,6 +100,190 @@ 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"] == {"effective_bits": 5.4, "cost_model": "kv_cache"}
+    assert inputs["method"] == "kl_div"
+    assert [config["effective_bits"] for config in inputs["quantization_formats"]] == [8.0, 4.5]
+    assert aq.cost_excluded_layers == []
+    assert "*mtp*" in inputs["disabled_layers"]
+    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)
+    monkeypatch.setattr(
+        tensor_quant,
+        "dynamic_block_quantize_op",
+        lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs),
+    )
+    model = get_tiny_qwen3(num_hidden_layers=1)
+    aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        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 == (2, 1)
+    assert attention.v_bmm_quantizer.num_bits == (2, 1)
+    assert attention.k_bmm_quantizer.amax == 448.0
+    assert attention.v_bmm_quantizer.amax == 448.0
+
+
+def test_kv_autoquant_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_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)})
+
+
+@pytest.mark.parametrize(
+    ("search_domain", "expected_shape"),
+    [("weight", (2, 4, 3)), ("kv_cache", (4, 3))],
+)
+def test_kl_padding_exclusion_is_scoped_to_kv_autoquant(monkeypatch, search_domain, expected_shape):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    inputs = {
+        "search_domain": search_domain,
+        "constraints": {"effective_bits": 8.0},
+        "quantization_formats": [],
+        "fixed_quantization_config": None,
+        "module_search_spaces": [],
+        "disabled_layers": [],
+        "kv_cache_quant_cfg": None,
+        "method": "kl_div",
+        "score_size": 1,
+    }
+    monkeypatch.setattr(
+        hf_ptq, "_mtq_inputs_from_auto_quantize_config", lambda *_args, **_kwargs: inputs
+    )
+    logits = torch.arange(2 * 4 * 3).reshape(2, 4, 3).float()
+    batch = {
+        "input_ids": torch.ones(2, 4, dtype=torch.long),
+        "attention_mask": torch.tensor([[1, 1, 0, 0], [0, 1, 1, 0]]),
+    }
+
+    class Model(torch.nn.Module):
+        def forward(self, **_kwargs):
+            return SimpleNamespace(logits=logits, loss=torch.tensor(0.0))
+
+    observed = {}
+
+    def fake_auto_quantize(search_model, **kwargs):
+        observed["shape"] = tuple(kwargs["forward_step"](search_model, batch).shape)
+        return search_model, {}
+
+    monkeypatch.setattr(hf_ptq.mtq, "auto_quantize", fake_auto_quantize)
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=False,
+        batch_size=1,
+        auto_quantize_checkpoint=None,
+    )
+    model = Model()
+
+    hf_ptq.auto_quantize(args, model, [batch], SimpleNamespace(), full_model=model)
+
+    assert observed["shape"] == expected_shape
+
+
+def test_kv_autoquant_rejects_fsdp2(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(
+        hf_ptq,
+        "_mtq_inputs_from_auto_quantize_config",
+        lambda *_args, **_kwargs: {"search_domain": "kv_cache"},
+    )
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=True,
+    )
+
+    with pytest.raises(NotImplementedError, match="KV-cache AutoQuantize does not support"):
+        hf_ptq.auto_quantize(args, torch.nn.Module(), [], SimpleNamespace())
+
+
+def test_weight_autoquant_retains_fsdp2_warning(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    model = torch.nn.Module()
+    inputs = {
+        "search_domain": "weight",
+        "constraints": {"effective_bits": 8.0},
+        "quantization_formats": [],
+        "fixed_quantization_config": None,
+        "module_search_spaces": [],
+        "disabled_layers": [],
+        "kv_cache_quant_cfg": None,
+        "method": "gradient",
+        "score_size": 1,
+    }
+    monkeypatch.setattr(
+        hf_ptq, "_mtq_inputs_from_auto_quantize_config", lambda *_args, **_kwargs: inputs
+    )
+    monkeypatch.setattr(
+        hf_ptq.mtq, "auto_quantize", lambda search_model, **_kwargs: (search_model, {})
+    )
+    args = SimpleNamespace(
+        calib_with_images=False,
+        inference_pipeline_parallel=1,
+        use_fsdp2=True,
+        batch_size=1,
+        auto_quantize_checkpoint=None,
+    )
+
+    with pytest.warns(UserWarning, match="use at your own risk"):
+        assert hf_ptq.auto_quantize(args, model, [], SimpleNamespace()) is model
+
+
+def test_fsdp2_kv_autoquant_rejected_before_model_load(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+    monkeypatch.setattr(hf_ptq, "_recipe_is_kv_auto_quantize", lambda _: True)
+    monkeypatch.setattr(
+        hf_ptq.AutoConfig,
+        "from_pretrained",
+        lambda *_args, **_kwargs: pytest.fail("The model config must not be loaded."),
+    )
+
+    with pytest.raises(NotImplementedError, match="KV-cache AutoQuantize does not support"):
+        hf_ptq.load_model(SimpleNamespace(use_fsdp2=True, recipe="autoquant"))
+
+
+def test_fsdp2_preload_guard_distinguishes_weight_and_kv_autoquant(monkeypatch):
+    hf_ptq = _import_hf_ptq(monkeypatch)
+
+    assert hf_ptq._recipe_is_kv_auto_quantize(
+        "general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits"
+    )
+    assert not hf_ptq._recipe_is_kv_auto_quantize("general/auto_quantize/nvfp4_fp8_at_5p4bits")
+
+
 def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch):
     """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns
     key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict."""

tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py

@@ -88,6 +88,16 @@ def test_image_calibration_model_target_follows_recipe(
     )
 
 
+def test_language_model_extraction_fails_closed_on_competing_roots(hf_ptq):
+    full_model = torch.nn.Module()
+    full_model.model = torch.nn.Module()
+    full_model.model.language_model = torch.nn.Module()
+    full_model.language_model = torch.nn.Module()
+
+    with pytest.raises(ValueError, match="multiple language-model roots"):
+        hf_ptq.extract_and_prepare_language_model_from_vl(full_model)
+
+
 def test_image_calibration_uses_full_vlm_forward(hf_ptq, monkeypatch):
     args = SimpleNamespace(
         qformat="fp8",

tests/gpu/torch/export/test_export.py

@@ -489,7 +489,7 @@ def test_get_scaling_factor(
             {
                 "exclude_modules": ["linears.0", "linears.2"],
                 "quant_algo": "FP8",
-                "kv_cache_quant_algo": "FP8",
+                "kv_cache_quant_algo": None,
             },
         ),
         (
@@ -497,7 +497,7 @@ def test_get_scaling_factor(
             {
                 "exclude_modules": ["linears.0", "linears.2"],
                 "quant_algo": "FP8",
-                "kv_cache_quant_algo": "INT8",
+                "kv_cache_quant_algo": None,
             },
         ),
     ],

tests/gpu/torch/export/test_layerwise_export.py

@@ -145,6 +145,41 @@ def _kv_cache_cfg():
     )
 
 
+def test_layerwise_export_resolves_mixed_kv_cache_formats(tmp_path):
+    model = get_tiny_llama(num_hidden_layers=2).eval()
+    model.config.architectures = ["LlamaForCausalLM"]
+    mtq.quantize(
+        model,
+        {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {
+                    "quantizer_name": "model.layers.0.self_attn.*[kv]_bmm_quantizer",
+                    "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                    "enable": True,
+                },
+                {
+                    "quantizer_name": "model.layers.1.self_attn.*[kv]_bmm_quantizer",
+                    "cfg": {
+                        "num_bits": (2, 1),
+                        "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                        "constant_amax": 1.0,
+                    },
+                    "enable": True,
+                },
+            ],
+            "algorithm": None,
+        },
+    )
+    exporter = LayerwiseExporter(model, tmp_path)
+    exporter.bind(list(model.model.layers))
+
+    assert exporter._kv_cache_format == {
+        "model.layers.0.self_attn": {"quant_algo": "FP8"},
+        "model.layers.1.self_attn": {"quant_algo": "NVFP4"},
+    }
+
+
 def _nvfp4_cfg():
     """NVFP4 with o_proj left unquantized.
 

tests/unit/recipe/test_loader.py

@@ -28,6 +28,9 @@ import pytest
 
 import modelopt.torch.quantization.config as qcfg
 from modelopt.recipe.config import (
+    AutoQuantizeConfig,
+    AutoQuantizeConstraints,
+    AutoQuantizeCost,
     ModelOptAutoQuantizeRecipe,
     ModelOptDFlashRecipe,
     ModelOptEagleRecipe,
@@ -1799,6 +1802,22 @@ def test_load_recipe_autoquantize_minimal(tmp_path):
     assert aq.module_search_spaces == []
 
 
+def test_autoquantize_constraints_use_effective_bits_for_kv_cost_model():
+    assert AutoQuantizeConstraints.model_fields["effective_bits"].default == 4.8
+    assert AutoQuantizeConstraints().effective_bits == 4.8
+
+    constraints = AutoQuantizeConstraints(effective_bits=5.4, cost_model="kv_cache")
+    assert constraints.effective_bits == 5.4
+    assert constraints.cost_model == "kv_cache"
+
+    with pytest.raises(ValueError, match="does not accept weight cost settings"):
+        AutoQuantizeConstraints(
+            effective_bits=5.4,
+            cost_model="kv_cache",
+            cost=AutoQuantizeCost(active_moe_expert_ratio=0.5),
+        )
+
+
 def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path):
     """cost_model + cost.active_moe_expert_ratio parse and dump to the mtq constraints dict shape."""
     recipe_file = tmp_path / "aq.yml"
@@ -1938,6 +1957,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",
@@ -1949,11 +1969,45 @@ 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.cost_model == "kv_cache":
+        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():
+    recipe = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits")
+    aq = recipe.auto_quantize
+
+    assert aq.constraints.effective_bits == 5.4
+    assert aq.constraints.cost_model == "kv_cache"
+    assert aq.auto_quantize_method == "kl_div"
+    assert "*mtp*" in aq.disabled_layers
+    assert aq.cost_excluded_layers == []
+    assert all(candidate.algorithm is None for candidate in aq.candidate_formats)
+    assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5]
+    for fmt in aq.candidate_formats:
+        for entry in fmt.quant_cfg:
+            assert entry.quantizer_name == "*[kv]_bmm_quantizer"
+            assert not entry.cfg.use_constant_amax
+            assert entry.cfg.constant_amax == 448.0
+        assert fmt.algorithm is None
+
+
+def test_kv_autoquantize_rejects_cost_excluded_layers():
+    with pytest.raises(ValueError, match=r"cost_excluded_layers.*disabled_layers"):
+        AutoQuantizeConfig(
+            constraints=AutoQuantizeConstraints(effective_bits=8.0, cost_model="kv_cache"),
+            candidate_formats=[qcfg.QuantizeConfig(quant_cfg=[], effective_bits=8.0)],
+            auto_quantize_method="kl_div",
+            cost_excluded_layers=["*mtp*"],
+        )
 
 
 @pytest.mark.parametrize("recipe_path", _BUILTIN_PTQ_RECIPES)

tests/unit/torch/export/test_convert_hf_config.py

@@ -0,0 +1,112 @@
+# 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
+
+from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
+from modelopt.torch.export.unified_export_hf import _revert_hf_quant_config_names
+
+
+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": {
+                "kv_cache_quant_algo": "MIXED_PRECISION",
+                "kv_cache_quantized_layers": layer_map,
+                "kv_cache_schema_version": 1,
+            },
+        }
+    )
+
+    assert converted["quant_method"] == "modelopt"
+    assert "quant_algo" not in converted
+    assert "config_groups" not in converted
+    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_convert_uniform_kv_cache_config_preserves_layer_map():
+    layer_map = {"model.layers.0.self_attn": {"quant_algo": "FP8"}}
+    converted = convert_hf_quant_config_format(
+        {
+            "producer": {"name": "modelopt", "version": "test"},
+            "quantization": {
+                "kv_cache_quant_algo": "FP8",
+                "kv_cache_quantized_layers": layer_map,
+                "kv_cache_schema_version": 1,
+            },
+        }
+    )
+
+    assert "quant_algo" not in converted
+    assert "config_groups" not in converted
+    assert converted["kv_cache_scheme"] == {
+        "dynamic": False,
+        "num_bits": 8,
+        "type": "float",
+    }
+    assert converted["kv_cache_quantized_layers"] == layer_map
+    assert converted["kv_cache_schema_version"] == 1
+
+
+def test_convert_uniform_asymmetric_kv_cache_preserves_semantic_algo():
+    layer_map = {"model.layers.0.self_attn": {"quant_algo": "FP8_K_NVFP4_V"}}
+    converted = convert_hf_quant_config_format(
+        {
+            "producer": {"name": "modelopt", "version": "test"},
+            "quantization": {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {},
+                "kv_cache_quant_algo": "FP8_K_NVFP4_V",
+                "kv_cache_quantized_layers": layer_map,
+                "kv_cache_schema_version": 1,
+            },
+        }
+    )
+
+    assert converted["quant_algo"] == "MIXED_PRECISION"
+    assert converted["kv_cache_quant_algo"] == "FP8_K_NVFP4_V"
+    assert "kv_cache_scheme" not in converted
+    assert converted["kv_cache_quantized_layers"] == layer_map
+
+
+def test_reverse_quant_config_name_mapping_is_atomic():
+    quant_config = {
+        "quantization": {
+            "exclude_modules": ["model.good"],
+            "kv_cache_quantized_layers": {"model.bad": {"quant_algo": "FP8"}},
+        }
+    }
+
+    def failing_mapper(name):
+        if name == "model.bad":
+            raise RuntimeError("unsupported mapping")
+        return f"hub.{name}"
+
+    with pytest.raises(RuntimeError, match="unsupported mapping"):
+        _revert_hf_quant_config_names(quant_config, failing_mapper)
+
+    assert quant_config == {
+        "quantization": {
+            "exclude_modules": ["model.good"],
+            "kv_cache_quantized_layers": {"model.bad": {"quant_algo": "FP8"}},
+        }
+    }

tests/unit/torch/export/test_get_quantization.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import fnmatch
+from typing import cast
 
 import pytest
 import torch
@@ -26,17 +27,30 @@ from _test_utils.torch.export.utils import (
 
 import modelopt.torch.export.unified_export_megatron as unified_export_megatron
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
 from modelopt.torch.export.quant_format import (
+    KV_CACHE_FP8,
+    KV_CACHE_FP8_K_NVFP4_V,
+    KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A8_AWQ,
 )
 from modelopt.torch.export.quant_utils import (
+    _has_large_fp8_scale,
     get_kv_cache_scaling_factor,
     get_quant_config,
     get_quantization_format,
+    postprocess_state_dict,
 )
-from modelopt.torch.quantization.nn import NVFP4StaticQuantizer
+from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer
+
+
+class _FakeAttention(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.k_bmm_quantizer = TensorQuantizer()
+        self.v_bmm_quantizer = TensorQuantizer()
 
 
 class _FakeKVCacheQuantizer(torch.nn.Module):
@@ -63,6 +77,16 @@ def test_get_kv_cache_scaling_factor_can_disable_fp8_clamping():
     assert all(torch.equal(scale, torch.tensor([0.5])) for scale in unclamped_scales)
 
 
+def test_large_fp8_scale_check_is_device_agnostic():
+    class CudaLikeScale:
+        device = torch.device("cuda")
+
+        def __gt__(self, _threshold):
+            return torch.tensor([True])
+
+    assert _has_large_fp8_scale(cast("torch.Tensor", CudaLikeScale()))
+
+
 @pytest.mark.parametrize("clamp_kv_cache_scales", [True, False])
 def test_export_mcore_gpt_to_hf_passes_kv_cache_clamping_option(
     monkeypatch, tmp_path, clamp_kv_cache_scales
@@ -119,6 +143,366 @@ 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_uniform_vlm_export_ignores_disabled_vision_attention():
+    model = ToyModel()
+    weight_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,
+            },
+        ],
+        "algorithm": "max",
+    }
+    mtq.quantize(
+        model,
+        weight_config,
+        lambda quantized_model: quantized_model(torch.randn(1, 4, 10)),
+    )
+    model.language_model = torch.nn.Module()
+    model.language_model.attention = _FakeAttention()
+    model.vision_attention = _FakeAttention()
+    mtq.set_quantizer_by_cfg(
+        model,
+        [
+            {"quantizer_name": "*[kv]_bmm_quantizer", "enable": False},
+            {
+                "quantizer_name": "language_model.*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                "enable": True,
+            },
+        ],
+    )
+
+    quantization = get_quant_config(model)["quantization"]
+
+    assert quantization["quant_algo"] == "FP8"
+    assert quantization["kv_cache_quant_algo"] == "FP8"
+    assert "kv_cache_quantized_layers" not in quantization
+
+
+def test_quant_config_tolerates_ambiguous_language_model_roots():
+    model = torch.nn.Module()
+    model.model = torch.nn.Module()
+    model.model.language_model = torch.nn.Module()
+    model.language_model = torch.nn.Module()
+    model.model.language_model.attention = _FakeAttention()
+    mtq.set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
+
+    quantization = get_quant_config(model)["quantization"]
+
+    assert quantization["kv_cache_quant_algo"] is None
+
+
+@pytest.mark.parametrize("enabled_quantizer", ["k_bmm_quantizer", "output_quantizer"])
+def test_legacy_single_sided_kv_detection_is_preserved(enabled_quantizer):
+    model = torch.nn.Module()
+    model.attention = _FakeAttention()
+    model.attention.output_quantizer = TensorQuantizer()
+    mtq.set_quantizer_by_cfg(
+        model,
+        [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": f"*.{enabled_quantizer}",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                "enable": True,
+            },
+        ],
+    )
+
+    quantization = get_quant_config(model)["quantization"]
+
+    assert quantization["kv_cache_quant_algo"] == "FP8"
+    assert "kv_cache_quantized_layers" not in quantization
+
+
+@pytest.mark.parametrize(
+    ("quantizer_cfg", "expected_format"),
+    [
+        ({"num_bits": (4, 3), "constant_amax": 1.0}, "FP8"),
+        (
+            {
+                "num_bits": (2, 1),
+                "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                "constant_amax": 1.0,
+            },
+            "NVFP4",
+        ),
+    ],
+)
+def test_uniform_kv_only_export_preserves_bf16_weight_metadata(quantizer_cfg, expected_format):
+    model = torch.nn.Module()
+    model.attention = _FakeAttention()
+    mtq.set_quantizer_by_cfg(
+        model,
+        [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": quantizer_cfg,
+                "enable": True,
+            },
+        ],
+    )
+
+    hf_quant_config = get_quant_config(model)
+    quantization = hf_quant_config["quantization"]
+    assert quantization["quant_algo"] is None
+    assert quantization["kv_cache_quant_algo"] == expected_format
+    assert quantization["quantized_layers"] == {}
+    assert quantization["kv_cache_quantized_layers"] == {
+        "attention": {"quant_algo": expected_format}
+    }
+
+    converted = convert_hf_quant_config_format(hf_quant_config)
+    assert "quant_algo" not in converted
+    assert "config_groups" not in converted
+    assert converted["kv_cache_scheme"] == (
+        {"dynamic": False, "num_bits": 8, "type": "float"}
+        if expected_format == "FP8"
+        else expected_format
+    )
+    assert converted["kv_cache_quantized_layers"] == {"attention": {"quant_algo": expected_format}}
+
+
+def test_uniform_kv_export_retains_scheme_and_layer_map_without_search_marker():
+    model = torch.nn.Module()
+    model.attention = _FakeAttention()
+    mtq.set_quantizer_by_cfg(
+        model,
+        [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                "enable": True,
+            },
+        ],
+    )
+    quantization = get_quant_config(model)["quantization"]
+
+    assert quantization["quant_algo"] is None
+    assert quantization["quantized_layers"] == {}
+    assert quantization["kv_cache_quant_algo"] == "FP8"
+    assert quantization["kv_cache_quantized_layers"] == {"attention": {"quant_algo": "FP8"}}
+
+    converted = convert_hf_quant_config_format({"quantization": quantization})
+    assert "quant_algo" not in converted
+    assert "config_groups" not in converted
+    assert converted["kv_cache_scheme"] == {
+        "dynamic": False,
+        "num_bits": 8,
+        "type": "float",
+    }
+    assert converted["kv_cache_quantized_layers"] == {"attention": {"quant_algo": "FP8"}}
+
+
+def test_partial_uniform_kv_with_uniform_weights_preserves_legacy_schema():
+    model = ToyModel()
+    mtq.quantize(
+        model,
+        {
+            "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,
+                },
+            ],
+            "algorithm": "max",
+        },
+        lambda quantized_model: quantized_model(torch.randn(1, 4, 10)),
+    )
+    model.attn0 = _FakeAttention()
+    model.attn1 = _FakeAttention()
+    mtq.set_quantizer_by_cfg(
+        model,
+        [
+            {"quantizer_name": "*[kv]_bmm_quantizer", "enable": False},
+            {
+                "quantizer_name": "attn0.*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                "enable": True,
+            },
+        ],
+    )
+
+    quantization = get_quant_config(model)["quantization"]
+
+    assert quantization["quant_algo"] == "FP8"
+    assert quantization["kv_cache_quant_algo"] == "FP8"
+    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):
+            super().__init__()
+            self.k_bmm_quantizer = TensorQuantizer()
+            self.v_bmm_quantizer = TensorQuantizer()
+
+    model = torch.nn.Module()
+    model.attn0 = FakeAttention()
+    model.attn1 = FakeAttention()
+    model.attn2 = 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,
+                },
+            }
+        ],
+    )
+    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"] is None
+    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"},
+        "attn2": {"quant_algo": "FP8_K_NVFP4_V"},
+    }
+
+
+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]),
+        "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},
+        "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)
+
+    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]),
+        "attn2.k_proj.k_scale": torch.tensor([1.0]),
+        "attn2.v_proj.v_scale": torch.tensor([0.25]),
+    }
+
+
 class _FakeTopKRouter(torch.nn.Module):
     """Mimics a transformers>=5.0 MoE router: owns a ``weight`` but is NOT an ``nn.Linear``.
 

tests/unit/torch/export/test_offload_export.py

@@ -37,8 +37,12 @@ from _test_utils.torch.quantization.tied_modules import (
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.model_utils import TiedWeightMap
-from modelopt.torch.export.quant_format import KV_CACHE_FP8
-from modelopt.torch.export.quant_utils import _postprocess_single_tensor
+from modelopt.torch.export.quant_format import KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_NVFP4
+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 (
     _parse_shard_size,
@@ -376,6 +380,53 @@ def test_postprocess_kv_scale_renamed_and_divided():
     assert abs(val.item() - 0.5) < 1e-5
 
 
+@pytest.mark.parametrize(
+    ("layer_name", "quant_algo", "side", "resolved_format"),
+    [
+        ("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_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(
+        original_key,
+        torch.tensor(224.0),
+        448.0,
+        postprocess_config,
+    )
+
+    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)
+
+
 def test_postprocess_scale_squeezed():
     """3D scale tensors with shape[0]==1 are squeezed."""
     t = torch.ones(1, 4, 4)

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)

tests/unit/torch/export/test_unified_export_hf.py

@@ -25,7 +25,11 @@ from _test_utils.torch.quantization.tied_modules import (
 )
 
 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,
@@ -35,6 +39,52 @@ from modelopt.torch.export.unified_export_hf import _resolve_export_dtype
 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)
+
+
+def test_language_model_extraction_accepts_aliased_compatibility_property():
+    """A top-level compatibility property may alias the standardized nested LM root."""
+    model = torch.nn.Module()
+    model.model = torch.nn.Module()
+    model.model.language_model = torch.nn.Module()
+    model.language_model = model.model.language_model
+
+    assert get_language_model_from_vl(model) == [
+        model,
+        model.model,
+        model.model.language_model,
+    ]
+
+
+def test_language_model_extraction_preserves_nested_preference_by_default():
+    """Generic callers retain the historical nested-root preference."""
+    model = torch.nn.Module()
+    model.model = torch.nn.Module()
+    model.model.language_model = torch.nn.Module()
+    model.language_model = torch.nn.Module()
+
+    assert get_language_model_from_vl(model) == [
+        model,
+        model.model,
+        model.model.language_model,
+    ]
+
+
+def test_language_model_extraction_rejects_competing_roots_when_strict():
+    """Search callers can fail closed instead of selecting a competing root."""
+    model = torch.nn.Module()
+    model.model = torch.nn.Module()
+    model.model.language_model = torch.nn.Module()
+    model.language_model = torch.nn.Module()
+
+    with pytest.raises(ValueError, match="multiple language-model roots"):
+        get_language_model_from_vl(model, strict=True)
+
+
 @pytest.mark.parametrize(
     ("configured_dtype", "dtype", "expected_dtype", "warning_count"),
     [

tests/unit/torch/quantization/test_autoquant.py

@@ -533,18 +533,24 @@ def test_auto_quantize_fixed_ptq_baseline_requires_explicit_search_spaces():
         )
 
 
-@pytest.mark.parametrize("formats", ["FP8_DEFAULT_CFG", mtq.FP8_DEFAULT_CFG, ()])
-def test_auto_quantize_rejects_non_list_global_formats(formats):
-    with pytest.raises(TypeError, match="`quantization_formats` must be a list"):
+@pytest.mark.parametrize("formats", ["FP8_DEFAULT_CFG", mtq.FP8_DEFAULT_CFG])
+def test_auto_quantize_rejects_non_sequence_global_formats(formats):
+    with pytest.raises(TypeError, match="`quantization_formats` must be a sequence"):
         mtq.auto_quantize(TransformerBlock(), quantization_formats=formats)
 
 
-@pytest.mark.parametrize("formats", [[], [None]])
+@pytest.mark.parametrize("formats", [[], [None], ()])
 def test_auto_quantize_rejects_empty_global_formats(formats):
     with pytest.raises(ValueError, match="`quantization_formats` must"):
         mtq.auto_quantize(TransformerBlock(), quantization_formats=formats)
 
 
+def test_process_quantization_formats_preserves_explicit_tuple_name():
+    assert model_quant._process_quantization_formats(
+        [(mtq.FP8_DEFAULT_CFG, "named_fp8")], "CUSTOM"
+    ) == [(mtq.FP8_DEFAULT_CFG, "named_fp8")]
+
+
 @pytest.mark.parametrize("formats", ["FP8_DEFAULT_CFG", mtq.FP8_DEFAULT_CFG, ()])
 def test_auto_quantize_rejects_non_list_module_formats(formats):
     with pytest.raises(

tests/unit/torch/quantization/test_kv_cache_auto_quant.py

@@ -0,0 +1,800 @@
+# 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 json
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+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
+from modelopt.torch.quantization import model_quant, tensor_quant
+from modelopt.torch.quantization.calib import MaxCalibrator
+from modelopt.torch.quantization.config import QuantizeConfig
+from modelopt.torch.quantization.kv_cache_auto_quant import (
+    QuantKVRecipeHparam,
+    _candidate_quantizers,
+    _eligible_layers,
+    _kv_scalar_weight,
+    _solve_additive_recipe,
+    _validate_kv_only_config,
+)
+from modelopt.torch.quantization.nn import TensorQuantizer
+
+
+@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": _quantizer_cfg(bits, constant_amax=constant_amax),
+            }
+        ],
+        algorithm=algorithm,
+        effective_bits=effective_bits,
+    )
+
+
+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))
+
+    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,
+            )
+        )
+
+
+@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})
+
+    with pytest.raises(ValueError, match="only non-structural calibration algorithms"):
+        _validate_kv_only_config(config)
+
+
+def test_kv_candidate_rejects_non_json_calibrator():
+    config = QuantizeConfig(
+        quant_cfg=[
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "calibrator": MaxCalibrator},
+            }
+        ],
+        algorithm="max",
+        effective_bits=8.0,
+    )
+
+    with pytest.raises(ValueError, match="string calibrators"):
+        _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"],
+        layer_widths=[(128, 128), (128, 128)],
+        candidate_names=["fp8", "nvfp4"],
+        candidate_kv_bits=[(8.0, 8.0), (4.5, 4.5)],
+        scores=[[0.0, 10.0], [0.0, 1.0]],
+        target_bits=6.25,
+        verbose=False,
+    )
+
+    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
+
+
+@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 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)
+
+
+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__()
+        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_rejects_missing_scale_after_calibration(monkeypatch):
+    model = _ToyKVModel()
+    monkeypatch.setattr(model_quant, "calibrate", lambda *_args, **_kwargs: None)
+
+    with pytest.raises(ValueError, match="no persistent export scale after calibration"):
+        mtq.auto_quantize(
+            model,
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
+            [(_kv_config((4, 3), 8.0).model_dump(), "fp8")],
+            [torch.randn(1, 2, 8)],
+            lambda search_model, batch: search_model(batch),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+    assert hasattr(model, "_modelopt_state")
+
+
+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()
+    data = [torch.randn(2, 3, 8), torch.randn(2, 3, 8)]
+    candidates = [
+        (
+            {
+                "quant_cfg": [
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 8.0,
+            },
+            "fp8",
+        ),
+        (
+            {
+                "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,
+                        },
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 4.5,
+            },
+            "nvfp4",
+        ),
+    ]
+
+    model, state = mtq.auto_quantize(
+        model,
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
+        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"]["constraints"]["effective_bits"] == pytest.approx(6.25)
+    assert state["best"]["is_satisfied"]
+    json.dumps({key: value for key, value in state.items() if key != "quantizer_state"})
+    assert not hasattr(model, "_modelopt_kv_cache_auto_quantize_state")
+    assert model.training
+    assert {layer["selected"] for layer in state["layers"].values()} == {
+        "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 = (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 = mtq.auto_quantize(
+        restored_model,
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
+        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["best"] == state["best"]
+    assert restored_state["layers"] == state["layers"]
+    assert not restored_model.training
+    for layer_name, layer_state in restored_state["layers"].items():
+        layer = restored_model.get_submodule(layer_name)
+        expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1)
+        assert layer.k_bmm_quantizer.num_bits == expected_bits
+
+    resolved_config = mtq.get_auto_quantize_config(state, {"effective_bits": 4.5})
+    assert resolved_config["algorithm"] is None
+    assert resolved_config["quant_cfg"][0] == {"quantizer_name": "*", "enable": False}
+    assert {entry["quantizer_name"] for entry in resolved_config["quant_cfg"]} == {
+        "*",
+        "attn0.k_bmm_quantizer",
+        "attn0.v_bmm_quantizer",
+        "attn1.k_bmm_quantizer",
+        "attn1.v_bmm_quantizer",
+    }
+
+    re_solved_model = _ToyKVModel()
+    _, re_solved_state = mtq.auto_quantize(
+        re_solved_model,
+        {"effective_bits": 4.5, "cost_model": "kv_cache"},
+        candidates,
+        data,
+        lambda *_: pytest.fail("Changing only the budget must re-solve without re-scoring."),
+        num_calib_steps=2,
+        num_score_steps=2,
+        checkpoint=str(tmp_path / "kv_search.pth"),
+    )
+    assert re_solved_state["best"]["constraints"]["effective_bits"] == pytest.approx(4.5)
+    assert {layer["selected"] for layer in re_solved_state["layers"].values()} == {"nvfp4"}
+
+
+def test_kv_autoquant_replay_preserves_selected_calibration_algorithm():
+    candidate = _kv_config((4, 3), 8.0, algorithm="max")
+    state = {
+        "cost_model": "kv_cache",
+        "requested_constraints": {"effective_bits": 8.0, "cost_model": "kv_cache"},
+        "candidates": [
+            {
+                "name": "fp8",
+                "k_bits": 8.0,
+                "v_bits": 8.0,
+                "config": candidate.model_dump(mode="python", exclude_none=True),
+            }
+        ],
+        "layers": {
+            "attention": {
+                "k_width": 8,
+                "v_width": 8,
+                "scores": {"fp8": 0.0},
+            }
+        },
+        "best": {"recipe": {"attention": "fp8"}},
+    }
+
+    resolved_config = mtq.get_auto_quantize_config(state)
+
+    assert resolved_config["algorithm"] == "max"
+
+
+def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quant_stub):
+    model = _ToyKVModel()
+    candidate = _asymmetric_kv_config().model_dump(exclude_none=True)
+
+    model, state = mtq.auto_quantize(
+        model,
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
+        [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"]["constraints"]["effective_bits"] == pytest.approx(6.25)
+    for layer_state in state["layers"].values():
+        assert layer_state["selected"] == "fp8_k_nvfp4_v"
+    for layer in (model.attn0, model.attn1):
+        assert layer.k_bmm_quantizer.num_bits == (4, 3)
+        assert layer.v_bmm_quantizer.num_bits == (2, 1)
+    exported = get_quant_config(model)["quantization"]
+    assert exported["quant_algo"] is None
+    assert exported["kv_cache_quant_algo"] == "FP8_K_NVFP4_V"
+    assert {layer["quant_algo"] for layer in exported["kv_cache_quantized_layers"].values()} == {
+        "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"):
+        mtq.auto_quantize(
+            model,
+            {"effective_bits": 6.25, "cost_model": "kv_cache"},
+            [(_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v")],
+            [torch.randn(1, 2, 8)],
+            lambda search_model, batch: search_model(batch),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+
+def test_kv_autoquant_runtime_failure_leaves_converted_model():
+    model = _ToyKVModel()
+    candidates = [
+        (
+            {
+                "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,
+                        },
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 4.5,
+            },
+            "nvfp4",
+        )
+    ]
+
+    with pytest.raises(ValueError, match="non-empty vocabulary dimension"):
+        mtq.auto_quantize(
+            model,
+            {"effective_bits": 4.5, "cost_model": "kv_cache"},
+            candidates,
+            [torch.randn(2, 3, 8)],
+            lambda *_: torch.ones(8),
+            num_calib_steps=1,
+            num_score_steps=1,
+        )
+
+    assert hasattr(model, "_modelopt_state")
+
+
+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)]
+    candidates = [
+        (
+            {
+                "quant_cfg": [
+                    {
+                        "quantizer_name": "*[kv]_bmm_quantizer",
+                        "cfg": {"num_bits": (4, 3)},
+                    },
+                ],
+                "algorithm": "max",
+                "effective_bits": 8.0,
+            },
+            "fp8",
+        ),
+        (
+            {
+                "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,
+                        },
+                    }
+                ],
+                "algorithm": None,
+                "effective_bits": 4.5,
+            },
+            "nvfp4",
+        ),
+    ]
+
+    model, state = mtq.auto_quantize(
+        model,
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
+        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"]["constraints"]["effective_bits"] == pytest.approx(6.25)
+    assert all(
+        layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled
+        for layer in model.model.layers
+    )
+    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)
+    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["quant_algo"] is None
+    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(
+        restored_model,
+        {"effective_bits": 6.25, "cost_model": "kv_cache"},
+        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["best"] == state["best"]
+    assert restored_state["layers"] == state["layers"]
+    assert any(
+        hasattr(layer.self_attn.k_bmm_quantizer, "_amax")
+        for layer in restored_model.model.layers
+        if layer.self_attn.k_bmm_quantizer.num_bits == (4, 3)
+    )
+
+    replay_config = mtq.get_auto_quantize_config(state)
+    replay_model = get_tiny_llama(num_hidden_layers=2)
+    replay_model = mtq.quantize(
+        replay_model,
+        replay_config,
+        lambda quantized_model: quantized_model(**data[0]),
+    )
+    assert all(
+        layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled
+        for layer in replay_model.model.layers
+    )
+    replay_non_kv_quantizers = [
+        quantizer
+        for name, quantizer in replay_model.named_modules()
+        if isinstance(quantizer, TensorQuantizer)
+        and not name.endswith(("k_bmm_quantizer", "v_bmm_quantizer"))
+    ]
+    assert replay_non_kv_quantizers
+    assert all(not quantizer.is_enabled for quantizer in replay_non_kv_quantizers)
+
+
+def test_kv_recipe_hparam_accepts_numpy_solver_index():
+    module = torch.nn.Module()
+    module.k_bmm_quantizer = TensorQuantizer()
+    module.v_bmm_quantizer = TensorQuantizer()
+    module.k_proj = torch.nn.Linear(4, 4)
+    module.v_proj = torch.nn.Linear(4, 4)
+    candidates = [("fp8", _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0))]
+
+    hparam = QuantKVRecipeHparam("attention", module, candidates)
+    hparam.active = np.int64(0)
+
+    assert hparam.active == 0
+    assert isinstance(hparam.active, int)
+
+
+@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(
+        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,
+        disabled_layers=disabled_layers,
+    )
+
+    assert set(state["layers"]) == {expected_layer}
+    report = {key: value for key, value in state.items() if key != "quantizer_state"}
+    assert json.loads(json.dumps(report))["layers"][expected_layer]["selected"] == "fp8"
+    exported = get_quant_config(model)["quantization"]
+    assert exported["quant_algo"] is None
+    assert exported["quantized_layers"] == {}
+    assert exported["kv_cache_quant_algo"] == "FP8"
+    assert set(exported["kv_cache_quantized_layers"]) == {expected_layer}
+
+
+def test_public_kv_autoquant_validation_fails_before_conversion():
+    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), 8.0).model_dump()
+    invalid_candidate["algorithm"] = "svdquant"
+
+    with pytest.raises(ValueError, match="only non-structural calibration algorithms"):
+        mtq.auto_quantize(
+            model,
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
+            [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
+
+
+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(
+                model,
+                {"effective_bits": 8.0, "cost_model": "kv_cache"},
+                [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(
+            model,
+            {"effective_bits": 8.0, "cost_model": "kv_cache"},
+            [],
+            [],
+            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_rejects_preceding_quantization_before_search():
+    model = get_tiny_llama(num_hidden_layers=2)
+    model = mtq.quantize(
+        model,
+        {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {
+                    "quantizer_name": "*.weight_quantizer",
+                    "cfg": {"num_bits": (4, 3), "axis": None},
+                    "enable": True,
+                },
+            ],
+            "algorithm": None,
+        },
+    )
+    candidate = {
+        "quant_cfg": [
+            {
+                "quantizer_name": "*[kv]_bmm_quantizer",
+                "cfg": {"num_bits": (4, 3), "constant_amax": 1.0},
+            }
+        ],
+        "algorithm": None,
+        "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,
+        )