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 HF recipe path now also composes GEMM and KV optimization as ordered stages on the same in-memory model:

• fixed FP8 GEMM PTQ followed by KL-divergence mixed-KV AutoQuantize;
• gradient-based NVFP4/FP8 GEMM AutoQuantize followed by an independent KL-divergence mixed-KV AutoQuantize stage.

The selected GEMM weight/activation QDQ remains enabled during KV scoring while its calibration state is frozen, so KV sensitivity is measured against the model that will be exported. Weight and KV searches retain independent constraints, methods, candidates, reports, and resumable checkpoint paths.

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 internally while reusing the common LPS solver, calibration/quantizer utilities, and safe_load/safe_save primitives. YAML composition sequences the independent stages without merging their solver or checkpoint contracts.

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 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 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 use disabled_layers, which removes them from scoring and budget accounting;
• top-level dynamic FP8 candidates are rejected, and calibrated/restored candidates must have persistent export scales;
• projection output quantizers remain separate from the per-layer K/V map.

Composition status and open review gates

The producer can export a uniform FP8-weight or mixed-weight checkpoint together with schema-v1 kv_cache_quantized_layers metadata. Two composed Qwen3-8B checkpoints were generated from the published branch and their producer-side checkpoint inventories and mappings validated.

The following review gates remain open at head 452328d:

1. The linked vLLM consumer currently reads the layer-wise KV map only through ModelOptMixedPrecisionConfig; uniform FP8/NVFP4 weight dispatch must also consume the map before composed uniform-weight checkpoints are deployment-ready.
2. A resumed KV search must bind its compatibility signature to the preceding GEMM quantization/selection state, not only to KV candidates and layer geometry.
3. Composed recipes need model-side validation against the converted model's fully qualified K/V quantizers so a GEMM-stage rule cannot enable K/V through a path that evades schema-only probes.

Until these gates are addressed, the composed recipes are producer-validation functionality and should not be treated as end-to-end deployment-ready with the linked consumer.

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.

For full FP8 K/V and full NVFP4 K/V, the consumer can dispatch each selected layer through existing uniform kernels once every ModelOpt weight-config path resolves the layer map. The current linked PR still needs the uniform-weight configuration fix described above.

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

KV-only search:

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

Fixed FP8 GEMM PTQ followed by KV AutoQuantize:

python examples/hf_ptq/hf_ptq.py \
  --pyt_ckpt_path Qwen/Qwen3-8B \
  --recipe general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
  --kv_auto_quantize_checkpoint /path/to/kv_autoquant.pth \
  --export_path /path/to/qwen3-8b-fp8-and-mixed-kv

Gradient GEMM AutoQuantize followed by independent KV AutoQuantize:

python examples/hf_ptq/hf_ptq.py \
  --pyt_ckpt_path Qwen/Qwen3-8B \
  --recipe general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
  --auto_quantize_checkpoint /path/to/weight_autoquant.pth \
  --kv_auto_quantize_checkpoint /path/to/kv_autoquant.pth \
  --export_path /path/to/qwen3-8b-autoquant-and-mixed-kv

Direct public API 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,
)

Testing

• Published composition amendment:
  • recipe and hf_ptq.py tests: 307 passed;
  • export tests: 37 passed;
  • AutoQuantize/KV tests: 125 passed;
  • recipe documentation/preset tests: 10 passed;
  • tiny offline end-to-end composed search/export tests: 3 passed.
• Architecture-focused offline coverage includes:
  • 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.
• Review regressions cover unequal-width asymmetric cost validation, KV recipe exclusion wiring, projection-output classification, dynamic-FP8 scale rejection, and minimum-Transformers Qwen3-VL compatibility.
• Current GitHub CI is green.
• Pre-commit passed on all changed files, including Ruff, mypy, recipe validation, Bandit, and 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.
• Producer-side unified export supports uniform FP8 weights or mixed quantized weights together with a mixed per-layer KV map; linked runtime-consumer support remains gated as documented above.
• This PR implements the forward-KL KV search method. Gradient-based KV sensitivity is not included.
• Checkpoint generation and runtime smoke testing are downstream validation activities; their custom launch/test harnesses are not included in this PR.

[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.