Add layer-wise KV-cache AutoQuant with forward KL

[meenchen]

What does this PR do?

Type of change: new feature.

Adds layer-wise KV-cache AutoQuantize with isolated forward-KL sensitivity:

• introduces mtq.auto_quantize_kv_cache, using BF16/no-KV-quant logits as the reference and quantizing one eligible attention layer at a time;
• solves a width-weighted additive recipe under constraints.kv_effective_bits, with caller-declared packed cost per K/V scalar;
• calibrates each candidate without recalibrating unrelated weight or fixed-layer quantizers, supports resumable checkpoints, and restores model state on failure;
• adds a 5.4-bit FP8 K/V, FP8-K/NVFP4-V, and NVFP4 K/V recipe for examples/hf_ptq;
• exports heterogeneous per-layer KV formats and a JSON-safe sensitivity report in unified Hugging Face checkpoints.

The existing weight-only AutoQuant search groups decisions and derives cost from weights, so it cannot safely represent one joint K/V decision per attention layer or heterogeneous KV export. This adds a separate KV-specific search path and versioned export metadata instead of overloading the weight cost model.

Relationship to vLLM runtime support

This PR is the checkpoint producer: it searches the recipe and writes schema-v1 kv_cache_quantized_layers metadata. The companion vllm-project/vllm#52813 is the checkpoint consumer: it reads that mapping and dispatches each attention layer.

Together, the two PRs support layer-wise mixtures of full FP8 K/V and full NVFP4 K/V without any new kernel code; vLLM #52813 uses the existing uniform FP8 and NVFP4 implementations for each selected layer.

FP8-K/NVFP4-V within one layer is a separate capability. This PR can search and export that format, but vLLM #52813 deliberately rejects it because it requires the independent mixed-K/V attention-kernel implementation. Neither PR bundles that kernel work.

Usage

formats = [
    ({**mtq.FP8_KV_CFG, "effective_bits": 8.0}, "fp8"),
    ({**mtq.NVFP4_KV_CFG, "effective_bits": 4.5}, "nvfp4"),
]
model, report = mtq.auto_quantize_kv_cache(
    model,
    constraints={"kv_effective_bits": 5.4},
    quantization_formats=formats,
    data_loader=calibration_loader,
    forward_step=lambda model, batch: model(**batch).logits[:, -128:],
    num_calib_steps=64,
    num_score_steps=64,
)

The shipped three-format recipe can also be run with:

python examples/hf_ptq/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

Testing

• Focused changed-area suite: 309 passed
  • KV search and solver behavior
  • recipe loading and HF PTQ argument construction
  • mixed-KV quantization metadata and AutoQuant report export
  • quant-aware config-name mapping integration
• Broader changed-area run on current main: 319 passed, 4 environment-gated failures
  • the four failures are unchanged upstream scoped-prefix tests requiring Transformers APIs (PrefixChange / scope_prefix) newer than the locally available transformers==5.4.0;
  • the PR-specific quant-aware mapping test passes.
• Pre-commit on all changed files: passed
  • Ruff check/format
  • mypy
  • recipe validation
  • Bandit
  • Markdown/RST/YAML checks
• Credential history audit passed.
• All four commits are signed and include DCO sign-off.

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?: ✅
• Did you get Claude approval on this PR?: ❌ (draft; review not requested yet)

Additional Information

• Search scoring requires full-vocabulary logits, though callers may return only the token positions they want scored.
• disabled_layers are preserved in their existing KV format and excluded from the searched-layer bit budget.
• Mixed-KV export with a uniform quantized-weight format currently fails explicitly; BF16 weights and existing mixed-weight export are supported.
• This PR implements the forward-KL search method. Gradient-based KV sensitivity is not included.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b1a21fff-d8f9-44df-a509-ecf54c73c87d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2211/
Built to branch gh-pages at 2026-08-20 05:46 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 90.69767% with 40 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.02%. Comparing base (fbcdc16) to head (faa1f95).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/kv_cache_auto_quant.py	94.34%	16 Missing ⚠️
modelopt/torch/export/quant_utils.py	76.92%	15 Missing ⚠️
modelopt/torch/export/unified_export_hf.py	64.70%	6 Missing ⚠️
modelopt/recipe/config.py	95.45%	1 Missing ⚠️
...delopt/torch/export/unified_export_hf_streaming.py	50.00%	1 Missing ⚠️
modelopt/torch/quantization/model_quant.py	96.96%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2211      +/-   ##
==========================================
+ Coverage   78.94%   79.02%   +0.07%     
==========================================
  Files         522      523       +1     
  Lines       60550    60953     +403     
==========================================
+ Hits        47803    48168     +365     
- Misses      12747    12785      +38     

Flag	Coverage Δ
unit	55.90% <90.69%> (+0.34%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

Bot review (gpt-5.6-sol) — DM the bot to share feedback.

The feature solves a real gap: choosing one calibrated K/V format per attention layer under a KV-storage budget and exporting that heterogeneous mapping. However, the architectural choice is not fully justified. The repo already has AutoQuantizeKLDivSearcher for isolated forward-KL scoring, QuantRecipeHparam for swapping per-choice quantizers/grouped decisions, and BaseSearcher for checkpoint/signature lifecycle, with the existing PuLP-backed LPS already used here. The PR body explains why the current weight cost model cannot be used unchanged, but not why those abstractions cannot be generalized or subclassed with K/V quantizer attributes and a pluggable width-based cost model instead of adding a parallel 673-line search/checkpoint engine. Please document that tradeoff before approval.

I also found a storage-accounting bug for the supported asymmetric FP8-K/NVFP4-V candidate: its effective bits are averaged equally between K and V, while layers are weighted by the sum of potentially unequal K/V widths. This can violate kv_effective_bits; the existing width test even demonstrates unequal K/V projections, but no asymmetric test covers the case. The shipped recipe also imports cost_excluded_layers, while the KV recipe-to-API path silently drops that field, including *mtp*.

Test coverage is otherwise substantial, and the new source headers match LICENSE_HEADER. At +2,179/-61 across 20 files, this is also a high-risk review size; splitting the search/API work from export metadata where practical would make validation easier.

[cjluo-nv]

Bot comment.

This cost is exact only when K and V have equal widths. _kv_scalar_weight() explicitly supports unequal k_proj.out_features/v_proj.out_features (and the unit test uses 24 vs 16), but an FP8-K/NVFP4-V candidate is charged (k_width + v_width) * (8 + 4.5) / 2 instead of k_width * 8 + v_width * 4.5. That can select a recipe whose actual packed storage violates kv_effective_bits. Please either reject asymmetric candidates on unequal-width layers or build per-layer/per-candidate costs from the separate K and V widths, and add an asymmetric unequal-width solver test.

[cjluo-nv]

Bot comment.

This field is silently ignored for KV searches: _mtq_inputs_from_auto_quantize_config() returns early in the KV branch and never forwards cost_excluded_layers. In particular, *mtp* is present in this imported unit but not in base_disabled_layers, so an MTP attention layer can remain eligible and count toward the KV budget despite this recipe appearing to exclude it. Please remove the dead field and explicitly disable the intended layers, reject cost_excluded_layers in KV recipe validation, or define and forward its KV-search semantics.