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 AutoQuantizeKLDivSearcher and QuantRecipeHparam operate on linear-module weight/input/output quantizer recipes. Their grouping follows runtime-fusion rules and their cost is derived from parameter numel. KV search instead owns one joint K/V choice at an attention boundary, costs resident K/V scalar widths, and checkpoints calibrated K/V scale tensors together with per-layer search progress. Generalizing the existing hparam would therefore change its grouping, cost, and checkpoint contracts and risk existing weight AutoQuant behavior; subclassing BaseSearcher would not remove those KV-specific responsibilities. The KV path remains separate while reusing the common LPS solver, calibration/quantizer utilities, and safe_load/safe_save primitives, so the weight search API and checkpoints stay unchanged.

Model-family support status

Support remains architecture-driven: there are no model-name checks, checkpoint paths, fixed layer lists, or campaign flags.

Model family	Current PR support
Qwen3-8B and Qwen3.8-27B	Plain causal-decoder K/V attention is discovered through the existing quantizer boundary and exercised through the public API on an offline Qwen fixture.
Qwen3.6-35B-A3B	Causal language-model extraction is supported without traversing unrelated model roots. Unified export now handles config.architectures = None, as observed with this family.
NVIDIA Nemotron 3 Nano and NVIDIA Nemotron 3 Lightning	Hybrid decoder support is structure-based: full-attention mixers exposing supported K/V quantizers are eligible, while Mamba and other nonattention mixers remain outside the search.

Additional fail-closed behavior:

• competing or aliased language-model roots are rejected instead of being selected by traversal order;
• attention boundaries registered through aliases are rejected instead of being searched twice;
• disabled_layers continue to preserve fixed or unsupported layers;
• the supported asymmetric candidate is reported and exported consistently as fp8_k_nvfp4_v; FP8 and NVFP4 preset names are unchanged.
• asymmetric FP8-K/NVFP4-V candidates fail closed when an eligible layer has unequal K/V widths, because the declared average bits-per-scalar would not be an exact storage cost;
• KV recipes reject cost_excluded_layers; non-KV-cache modules such as MTP must use disabled_layers, which is passed to the public KV AutoQuant API and removes them from both scoring and budget accounting.

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
• The latest architecture amendment adds focused offline regressions for:
  • an end-to-end examples/hf_ptq invocation of the public KV AutoQuant API;
  • plain and conditional-generation Qwen causal-attention discovery;
  • hybrid full-attention versus nonattention mixer discovery;
  • null architecture metadata and ambiguous/aliased model boundaries;
  • semantic FP8-K/NVFP4-V candidate naming and JSON-safe export state.
• 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.
• Review-amendment regressions: 38 passed
  • asymmetric FP8-K/NVFP4-V cost validation on unequal K/V widths
  • KV recipe exclusion schema and recipe-to-public-API wiring
• Pre-commit on all changed files: passed
  • Ruff check/format
  • mypy
  • recipe validation
  • Bandit
  • Markdown/RST/YAML checks
• Credential history audit passed.
• All 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.
• Checkpoint generation, runtime smoke testing, and accuracy evaluation are downstream validation activities; this PR contains the ModelOpt checkpoint-producer implementation only.

[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-2211/
Built to branch gh-pages at 2026-08-25 21:56 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 91.48073% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.70%. Comparing base (fbcdc16) to head (452328d).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/export/quant_utils.py	76.81%	16 Missing ⚠️
modelopt/torch/quantization/kv_cache_auto_quant.py	94.88%	16 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    #2211      +/-   ##
==========================================
- Coverage   78.94%   78.70%   -0.25%     
==========================================
  Files         522      524       +2     
  Lines       60550    62255    +1705     
==========================================
+ Hits        47803    48996    +1193     
- Misses      12747    13259     +512     

Flag	Coverage Δ
unit	56.05% <91.48%> (+0.49%)	⬆️

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 review (gpt-5.6-sol) — DM the bot to share feedback.

Re-review: the feature addresses a real gap—selecting one calibrated K/V format per attention layer under a KV-storage budget and exporting the heterogeneous mapping—but all three prior blocking concerns remain.

• Design (approval-blocking): the updated PR body explains why the existing weight cost model cannot be used unchanged, but it still does not explain why the existing AutoQuantizeKLDivSearcher isolated forward-KL flow, QuantRecipeHparam choice swapping/grouping, and BaseSearcher checkpoint lifecycle cannot be generalized for K/V quantizer attributes and a pluggable width-based cost model. The already-imported PuLP alternative is exposed through LPS and is appropriately reused. Before approval, please document why extending/subclassing those in-repo abstractions is not viable instead of adding a parallel 684-line search/checkpoint implementation.
• Critical correctness: asymmetric FP8-K/NVFP4-V accounting still averages K/V bits and then multiplies by total width, so unequal K/V widths can violate the advertised kv_effective_bits constraint. Removing this candidate from the shipped recipe limits exposure but does not fix the public API, which still accepts and advertises it; the new test still uses equal K/V widths.
• Critical recipe/API mismatch: the shipped recipe still imports base_cost_excluded_layers, but the KV recipe conversion returns before forwarding that field. *mtp* is not in base_disabled_layers, so the recipe can search and budget MTP attention despite appearing to exclude it.

The new source/test headers match LICENSE_HEADER, and test coverage is otherwise substantial. At 22 files and +2,357/-67, the PR remains difficult to validate as one unit; splitting search/API from export metadata would reduce risk.

[meenchen]

Addressed the latest review in commit 34a249597:

• asymmetric FP8-K/NVFP4-V candidates now fail closed on eligible layers with unequal K/V widths, before model mutation; the focused regression was red at the prior PR head and is green with the fix;
• the shipped KV recipe moves *mtp* into disabled_layers, removes the unsupported cost_excluded_layers field, and KV recipe validation now rejects that field instead of silently dropping it;
• recipe-to-hf_ptq.py coverage verifies that *mtp* reaches the public KV AutoQuant API through disabled_layers;
• the PR description now explains the boundary with AutoQuantizeKLDivSearcher, QuantRecipeHparam, and BaseSearcher: their linear-module grouping, parameter-numel cost, and checkpoint contracts do not represent a joint attention-layer K/V decision without changing existing weight AutoQuant behavior. The KV path reuses LPS, calibration/quantizer utilities, and safe checkpoint I/O while keeping those existing contracts unchanged.

Validation: 38 focused changed-area tests passed; all changed-file pre-commit hooks passed, including recipe validation.

[cjluo-nv]

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

Re-review/design review: this solves a real gap—choosing one calibrated K/V format per attention boundary under a width-weighted cache budget and exporting the heterogeneous mapping. I checked the existing alternatives: AutoQuantizeKLDivSearcher/_AutoQuantizeBaseSearcher already provide isolated forward-KL search flow, QuantRecipeHparam provides candidate quantizer swapping/grouping, and BaseSearcher provides checkpoint lifecycle; the already-imported PuLP path is exposed through LPS and is reused here. Pydantic/OmegaConf cover schema/config composition but not the model-side search itself.

💬 Author replied in commit 34a249597 that the existing abstractions are tied to linear-module grouping, parameter-numel cost, and weight-search checkpoint contracts, and that generalizing them risks changing existing AutoQuant behavior. The PR body now records that tradeoff, so the deterministic design gate is substantially addressed; however, a parallel 714-line search/checkpoint implementation beside those existing abstractions still warrants owner/architect sign-off rather than automated approval.

💬 Addressed in 34a249597: asymmetric FP8-K/NVFP4-V now fails before mutation when any eligible layer has unequal K/V widths, with a focused unequal-width regression. This resolves the prior storage-accounting correctness issue.

💬 Addressed in 34a249597: the KV recipe removes unsupported cost_excluded_layers, moves *mtp* into disabled_layers, validation rejects the unsupported field, and recipe-to-public-API coverage verifies forwarding. This resolves the prior recipe/API mismatch.

Test coverage is substantial and the new-file headers match LICENSE_HEADER. The remaining concern is reviewability and architectural ownership: at 22 files and +2,426/-71, splitting the search/API work from export metadata would materially reduce risk if practical.

[meenchen]

Addressed the latest findings in signed+DCO commit c06fb0ee5:

• HF export now builds kv_cache_quantized_layers only from attention modules with both enabled K/V quantizers. Projection output_quantizer remains available through the existing Megatron fallback but can no longer misclassify uniform FP8 projection-output recipes as mixed KV.
• KV AutoQuant now rejects top-level dynamic FP8 before model mutation and verifies that every calibrated or restored candidate has persistent _amax state before saving/using search state.
• The tiny Qwen3-VL fixture now supplies a valid mRoPE configuration ([1, 1, 2] for head-dim 8), fixing the Transformers 4.57.6 minimum-version failure.
• AutoQuantizeConstraints.effective_bits again has the schema-visible default 4.8; selecting kv_effective_bits explicitly clears that weight default, preserving the exactly-one constraint.

Evidence:

• All four focused regressions failed at prior head 34a249597 and pass with this commit.
• 9/9 export tests passed.
• 30/30 KV AutoQuant tests passed.
• 298/298 recipe + hf_ptq.py tests passed.
• The exact Qwen3-VL KV AutoQuant test passed under Transformers 4.57.6.
• All changed-file pre-commit hooks passed, including ruff, mypy, bandit, and recipe validation.