Add layer-wise KV-cache AutoQuant with forward KL

CHANGELOG.rst

@@ -9,6 +9,7 @@ Changelog
 *Quantization*
 
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint.
+- Add layer-wise KV-cache AutoQuantize through ``mtq.auto_quantize_kv_cache`` and ``constraints.kv_effective_bits``. It measures isolated full-vocabulary forward KL for caller-supplied K/V formats, solves a width-weighted additive storage-constrained recipe across eligible layers, preserves search-disabled layers in their existing format, exports the selected per-attention mapping in unified HF checkpoints, and writes a JSON sensitivity report alongside the checkpoint. A cast-mode FP8/NVFP4 recipe at 5.4 bits/scalar is included.
 
 *Megatron Framework (M-LM / M-Bridge)*
 

examples/hf_ptq/README.md

@@ -422,8 +422,37 @@ leaving the original recipe unchanged.
 For models without backprop support (e.g. Llama-4), use the `kl_div` scoring method — see the shipped
 `general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits` recipe.
 
-KV cache is applied as a uniform post-step, not part of the per-layer search. An AutoQuantize recipe
-falls back to `--kv_cache_qformat` (default `fp8_cast`) unless it sets an explicit `kv_cache` field.
+Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall back to
+`--kv_cache_qformat` (default `fp8_cast`) unless they set an explicit `kv_cache` field.
+
+KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their
+`candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes
+packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are
+preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice.
+The shipped canary recipe searches calibrated FP8 K/V (8.0 bits/scalar) and packed NVFP4 K/V
+(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-1.7B \
+  --recipe general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
+  --auto_quantize_checkpoint /path/to/kv_autoquant.pth \
+  --export_path /path/to/qwen3-1.7b-mixed-kv
+```
+
+Each candidate uses max calibration so its persistent K/V scales are present in the unified HF
+checkpoint. Unified export records the selected formats in `kv_cache_quantized_layers` and writes
+the JSON-safe sensitivity report to `kv_cache_auto_quantize_report.json`;
+`--auto_quantize_checkpoint` stores the resumable raw search state.
+
+> [!NOTE]
+> Layer-wise KV checkpoints require the companion
+> [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):
@@ -460,6 +489,8 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop)
 
 ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point.
 
+> *AutoQuantize recipes are not supported with `--use_fsdp2` and are rejected before model loading. Distributed sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Use a PTQ recipe with FSDP2.*
+
 ### Usage
 
 #### Slurm (recommended)

examples/hf_ptq/hf_ptq.py

@@ -96,6 +96,23 @@
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
 RAND_SEED = 1234
+_FSDP2_AUTOQUANT_ERROR = (
+    "AutoQuantize does not support --use_fsdp2 until distributed sensitivity scoring, "
+    "selection, and checkpoint writes are synchronized across ranks."
+)
+
+
+def _select_unpadded_logits(logits: torch.Tensor, batch: dict[str, Any]) -> torch.Tensor:
+    """Return logits only for token positions selected by ``attention_mask``."""
+    attention_mask = batch.get("attention_mask")
+    if attention_mask is None:
+        return logits
+    if logits.shape[:-1] != attention_mask.shape:
+        raise ValueError(
+            "AutoQuant KL logits and attention_mask must have matching token dimensions; "
+            f"got {tuple(logits.shape[:-1])} and {tuple(attention_mask.shape)}."
+        )
+    return logits[attention_mask.bool()]
 
 
 def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool:
@@ -338,6 +355,24 @@ def _mtq_candidate_formats(formats) -> list[dict]:
     return quantization_formats
 
 
+def _mtq_kv_candidate_formats(formats) -> list[tuple[dict, str]]:
+    """Translate format-agnostic KV candidates while preserving useful preset names."""
+    candidates = []
+    for idx, fmt in enumerate(formats):
+        quant_cfg = fmt.model_dump(exclude_none=True)
+        candidate_quantizers = quant_cfg.get("quant_cfg", [])
+        name = None
+        for preset_name, preset in KV_QUANT_CFG_CHOICES.items():
+            normalized_preset_quantizers = (
+                type(fmt)(**preset).model_dump(exclude_none=True).get("quant_cfg", [])
+            )
+            if normalized_preset_quantizers == candidate_quantizers:
+                name = preset_name
+                break
+        candidates.append((quant_cfg, name or f"KV_CACHE_FORMAT_{idx}"))
+    return candidates
+
+
 def _mtq_inputs_from_auto_quantize_config(
     aq_config, args: argparse.Namespace, fixed_quantize_config=None
 ) -> dict:
@@ -349,6 +384,16 @@ def _mtq_inputs_from_auto_quantize_config(
     to ``--kv_cache_qformat`` when the recipe omits it.
     """
     constraints = aq_config.constraints.model_dump(exclude_none=True)
+    is_kv_search = aq_config.constraints.kv_effective_bits is not None
+    if is_kv_search:
+        return {
+            "search_domain": "kv_cache",
+            "constraints": {"kv_effective_bits": constraints["kv_effective_bits"]},
+            "quantization_formats": _mtq_kv_candidate_formats(aq_config.candidate_formats),
+            "disabled_layers": aq_config.disabled_layers,
+            "method": aq_config.auto_quantize_method,
+            "score_size": aq_config.score_size,
+        }
     # cost_excluded_layers (sibling of disabled_layers) maps to the mtq cost key: these layers are
     # kept out of the bit-budget denominator (cost_weight 0) — e.g. VL vision towers — distinct from
     # disabled_layers, which removes them from the search.
@@ -380,6 +425,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,
@@ -414,11 +460,7 @@ def auto_quantize(
     )
 
     if args.use_fsdp2:
-        warnings.warn(
-            "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet "
-            "(distributed calibration, sensitivity scoring, and recipe/checkpoint "
-            "synchronization across ranks); use at your own risk."
-        )
+        raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR)
 
     inputs = _mtq_inputs_from_auto_quantize_config(
         aq_config, args, fixed_quantize_config=fixed_quantize_config
@@ -461,14 +503,33 @@ def forward_step(model, batch):
             output = model(**inputs_)
             if is_base_model:
                 assert full_model is not None
-                return full_model.lm_head(output.last_hidden_state)
-            return output.logits
+                logits = full_model.lm_head(output.last_hidden_state)
+            else:
+                logits = output.logits
+            return _select_unpadded_logits(logits, batch)
 
     else:
         raise ValueError(
             f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'"
         )
 
+    if inputs["search_domain"] == "kv_cache":
+        language_model, _ = mtq.auto_quantize_kv_cache(
+            language_model,
+            constraints=inputs["constraints"],
+            data_loader=calib_dataloader,
+            forward_step=forward_step,
+            quantization_formats=inputs["quantization_formats"],
+            num_calib_steps=len(calib_dataloader),
+            num_score_steps=min(
+                len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)
+            ),
+            verbose=True,
+            disabled_layers=inputs["disabled_layers"],
+            checkpoint=args.auto_quantize_checkpoint,
+        )
+        return language_model
+
     language_model, _ = mtq.auto_quantize(
         language_model,
         constraints=inputs["constraints"],
@@ -512,6 +573,8 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool:
 def load_model(args: argparse.Namespace):
     # If low memory mode is enabled, we compress the model while loading the HF checkpoint.
     calibration_only = False
+    if args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
+        raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR)
     if args.use_fsdp2:
         hf_config = AutoConfig.from_pretrained(
             args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code

modelopt/recipe/config.py

@@ -149,10 +149,22 @@ def _validate_active_moe_expert_ratio(cls, v: float | None) -> float | None:
 class AutoQuantizeConstraints(ModeloptBaseConfig):
     """LP search constraints + cost model; matches the ``mtq.auto_quantize`` constraints dict."""
 
-    effective_bits: float = ModeloptField(
-        default=4.8,
+    effective_bits: float | None = ModeloptField(
+        default=None,
         title="Effective bits per weight",
-        description="Average weight-storage bits target for the LP, in (0, 16].",
+        description=(
+            "Average weight-storage bits target for the LP, in (0, 16]. Defaults to 4.8 "
+            "when neither bit constraint is specified."
+        ),
+    )
+    kv_effective_bits: float | None = ModeloptField(
+        default=None,
+        title="Effective bits per KV-cache scalar",
+        description=(
+            "Average KV-cache storage bits target across eligible layers for layer-wise KV "
+            "AutoQuant, in (0, 16]. Exactly one of effective_bits and kv_effective_bits may "
+            "be set."
+        ),
     )
     cost_model: Literal["weight", "active_moe"] = ModeloptField(
         default="weight",
@@ -165,13 +177,36 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
         description="Extra cost-model parameters; omit for the 'weight' cost model.",
     )
 
-    @field_validator("effective_bits")
+    @model_validator(mode="before")
     @classmethod
-    def _validate_effective_bits(cls, v: float) -> float:
-        if not (0 < v <= 16):
+    def _default_weight_constraint(cls, data):
+        if isinstance(data, dict):
+            data = dict(data)
+            if "effective_bits" not in data and "kv_effective_bits" not in data:
+                data["effective_bits"] = 4.8
+        return data
+
+    @field_validator("effective_bits", "kv_effective_bits")
+    @classmethod
+    def _validate_effective_bits(cls, v: float | None) -> float | None:
+        if v is not None and not (0 < v <= 16):
             raise ValueError(f"effective_bits must be in (0, 16], got {v}")
         return v
 
+    @model_validator(mode="after")
+    def _exactly_one_bit_constraint(self):
+        if (self.effective_bits is None) == (self.kv_effective_bits is None):
+            raise ValueError(
+                "Exactly one of effective_bits and kv_effective_bits must be specified."
+            )
+        if self.kv_effective_bits is not None and (
+            self.cost_model != "weight" or self.cost is not None
+        ):
+            raise ValueError(
+                "KV-cache AutoQuant does not support weight or active-MoE cost settings."
+            )
+        return self
+
 
 class AutoQuantizeModuleSearchSpace(ModeloptBaseConfig):
     """Candidate formats selectable for modules matching one or more name patterns."""
@@ -272,6 +307,20 @@ def _has_search_space(self):
                 "auto_quantize requires candidate_formats or at least one module_search_spaces "
                 "entry. For uniform quantization, use a PTQ recipe instead."
             )
+        if self.constraints.kv_effective_bits is not None:
+            if self.auto_quantize_method != "kl_div":
+                raise ValueError(
+                    "KV-cache AutoQuant currently requires auto_quantize_method=kl_div."
+                )
+            if self.module_search_spaces:
+                raise ValueError(
+                    "KV-cache AutoQuant uses one candidate space for all eligible attention "
+                    "layers; module_search_spaces is not supported."
+                )
+            if self.kv_cache is not None:
+                raise ValueError(
+                    "KV-cache AutoQuant candidate_formats replace the uniform kv_cache post-step."
+                )
         return self
 
 

modelopt/torch/export/convert_hf_config.py

@@ -269,6 +269,14 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An
     if kv_cache_quant_algo:
         if kv_cache_quant_algo == "FP8":
             new_config["kv_cache_scheme"] = {"dynamic": False, "num_bits": 8, "type": "float"}
+        elif kv_cache_quant_algo == "MIXED_PRECISION":
+            new_config["kv_cache_quant_algo"] = kv_cache_quant_algo
+            new_config["kv_cache_quantized_layers"] = original_quantization_details.get(
+                "kv_cache_quantized_layers", {}
+            )
+            new_config["kv_cache_schema_version"] = original_quantization_details.get(
+                "kv_cache_schema_version", 1
+            )
         else:
             # TODO: Handle other kv cache quantization algorithms
             new_config["kv_cache_scheme"] = kv_cache_quant_algo

modelopt/torch/export/model_config.py

@@ -45,6 +45,7 @@
 QUANTIZATION_FP8_PC_PT = "fp8_pc_pt"
 
 KV_CACHE_FP8 = "FP8"
+KV_CACHE_FP8_K_NVFP4_V = "FP8_K_NVFP4_V"
 KV_CACHE_INT8 = "INT8"
 KV_CACHE_NVFP4 = "NVFP4"
 KV_CACHE_NVFP4_AFFINE = "NVFP4_AFFINE"

modelopt/torch/export/quant_aware_conversion.py

@@ -278,7 +278,7 @@ def _map(name: str) -> str:
 
 
 def revert_quant_config_names(quantization: dict, mapper) -> None:
-    """Revert ``exclude_modules`` / ``quantized_layers`` keys to hub names, in place.
+    """Revert layer-reference keys to hub names, in place.
 
     ``mapper`` is the callable from :func:`build_reverse_name_mapper` (a no-op when
     ``None``). Applies to the ModelOpt ``{"quantization": {...}}`` sub-dict before it is
@@ -293,6 +293,11 @@ def revert_quant_config_names(quantization: dict, mapper) -> None:
     quantized_layers = quantization.get("quantized_layers")
     if isinstance(quantized_layers, dict) and quantized_layers:
         quantization["quantized_layers"] = {mapper(k): v for k, v in quantized_layers.items()}
+    kv_cache_quantized_layers = quantization.get("kv_cache_quantized_layers")
+    if isinstance(kv_cache_quantized_layers, dict) and kv_cache_quantized_layers:
+        quantization["kv_cache_quantized_layers"] = {
+            mapper(k): v for k, v in kv_cache_quantized_layers.items()
+        }
 
 
 def _assert_experts_pre_expanded(