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

CHANGELOG.rst

@@ -40,6 +40,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.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 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
+            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,
@@ -1187,7 +1249,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}"
@@ -1448,9 +1510,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 @@ def _validate_effective_bits(cls, v: float) -> float:
             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 @@ 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.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