Add standalone layer-wise KV-cache AutoQuant with forward KL

[meenchen]

What does this PR do?

Type of change: new feature.

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

• introduces mtq.auto_quantize_kv_cache with one supported K/V format selected per eligible attention layer;
• solves a width-weighted additive recipe under constraints.kv_effective_bits;
• preserves existing non-KV quantizer execution while freezing unrelated calibration state;
• supports persistent/exportable FP8 K/V, NVFP4 K/V, and FP8-K/NVFP4-V candidates;
• saves resumable search state and a JSON-safe sensitivity report;
• exports the exact selected per-layer K/V mapping through unified Hugging Face export; 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, hybrid full-attention mixers are discovered through their K/V quantizer boundary, and nonattention/Mamba modules remain outside the search. Ambiguous aliases, unsupported distributed execution, structural algorithms, invalid storage declarations, nonpersistent scales, and unsupported K/V pairs fail closed.

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

Usage

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

• 364 focused tests passed across KV AutoQuant, recipe loading, unified export, and the actual hf_ptq.py public-API path.
• The shipped standalone recipe runs end to end on a tiny offline Qwen fixture and preserves exportable scale state.
• Changed-file pre-commit hooks passed, including recipe validation, Ruff, mypy, Bandit, Markdown, and YAML checks.

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)

Additional Information

• This is split from the combined ground-truth implementation in draft PR Add layer-wise KV-cache AutoQuant with forward KL #2211 to reduce review scope; composition is isolated in Compose GEMM and KV-cache AutoQuant workflows #2273.
• The standalone core tree contains no composed GEMM→KV recipe schema or orchestration.
• Full FP8 K/V and full NVFP4 K/V require a downstream consumer that reads kv_cache_quantized_layers; FP8-K/NVFP4-V additionally requires a compatible asymmetric runtime kernel.

[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

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{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-2272/
Built to branch gh-pages at 2026-08-27 19:57 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 91.27660% with 41 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.71%. Comparing base (fbcdc16) to head (efc25b3).
⚠️ Report is 16 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/kv_cache_auto_quant.py	94.88%	16 Missing ⚠️
modelopt/torch/export/quant_utils.py	77.61%	15 Missing ⚠️
modelopt/torch/export/unified_export_hf.py	64.70%	6 Missing ⚠️
modelopt/torch/export/model_utils.py	71.42%	2 Missing ⚠️
...delopt/torch/export/unified_export_hf_streaming.py	50.00%	1 Missing ⚠️
modelopt/torch/quantization/model_quant.py	97.05%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2272      +/-   ##
==========================================
- Coverage   78.94%   76.71%   -2.24%     
==========================================
  Files         522      524       +2     
  Lines       60550    63901    +3351     
==========================================
+ Hits        47803    49019    +1216     
- Misses      12747    14882    +2135     

Flag	Coverage Δ
unit	56.03% <91.27%> (+0.47%)	⬆️

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.