[OMNIML-5570, OMNIML-5569] 1/2 Add layer-wise KV-cache AutoQuant with forward KL - #2272
[OMNIML-5570, OMNIML-5569] 1/2 Add layer-wise KV-cache AutoQuant with forward KL#2272meenchen wants to merge 23 commits into
Conversation
Assisted-by: OpenAI Codex Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
|
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. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Important Approval pendingCodeRabbit has no unresolved comments, but it has not reviewed the latest commit. Use the checkbox below to review the latest commit. CodeRabbit will approve the changes if it finds no blocking issues.
📝 WalkthroughWalkthroughThe PR integrates layer-wise KV-cache AutoQuantize with the shared ChangesKV-cache AutoQuantize
Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: 🟡 Moderate · up to KV-cache AutoQuantize adds layer-wise selection and export support, but current test coverage can fail and scoring may be substantially slower on CUDA workloads; the CLI documentation also describes outdated behavior. These issues should be resolved before merge. Sequence Diagram(s)sequenceDiagram
participant HFPTQ
participant Model
participant API as mtq.auto_quantize
participant Search as AutoQuantizeKVSearcher
participant Export as UnifiedHFExport
HFPTQ->>Model: prepare masked calibration inputs
HFPTQ->>API: submit KV-cache recipe
API->>Search: validate and run search
Search->>Model: calibrate and score K/V candidates
Search-->>API: return selected quantizers and report
API->>Export: provide mixed KV-cache state
Export->>Model: normalize scales and write metadata
Estimated code review effort: 4 (Complex) | ~60 minutes 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
Full details: Security Anti-PatternsExplanation No stated security anti-pattern was introduced. The diff from the PR base adds no ✨ Finishing Touches 💡 2📝 Generate docstrings 💡
⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2272 +/- ##
==========================================
- Coverage 79.31% 78.96% -0.36%
==========================================
Files 527 528 +1
Lines 61482 62087 +605
==========================================
+ Hits 48765 49026 +261
- Misses 12717 13061 +344
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
The implementation is well tested and the standalone KV search appears internally coherent, but this still needs architectural and compatibility sign-off.
- Problem being solved: choose one deployable K/V-cache format per eligible attention layer using isolated forward-KL sensitivity, subject to a width-weighted storage budget, then checkpoint and export the chosen mapping.
- Existing alternatives: the repository already has
AutoQuantizeKLDivSearcherand_AutoQuantizeBaseSearcherinmodelopt/torch/quantization/algorithms.py, with candidate calibration, isolated KL scoring, solver state, checkpoint compatibility, and recipe application;BaseSearchersupplies the standard search/checkpoint lifecycle; andQuantRecipeHparamplus the existingLPSwrapper already model per-group choices and constrained selection. The project dependencies also already include PuLP, Pydantic, and OmegaConf, and this PR appropriately reuses PuLP/Pydantic rather than adding a dependency. - The PR body explains why KV search is standalone from GEMM→KV composition, but it does not explain why the new 735-line parallel search/checkpoint engine should not extend or generalize the existing AutoQuant searchers. Please have an owner decide whether the different K/V boundaries, calibration isolation, and additive objective justify maintaining both implementations, or document that rationale in the PR body.
- This remains a large change (+2525/-79 across 23 files), although the split from #2211 and the extensive focused tests substantially improve reviewability.
- Compatibility owner call: the PR changes all
hf_ptq.pyAutoQuantize+FSDP2 runs from a warning to an earlyNotImplementedError, not only the new KV search. That is a deliberate-looking and documented safety restriction, but it contradicts the PR's “backward compatible” declaration and is broader than the feature. Please confirm that existing weight-AutoQuant users should be hard-blocked and call the behavior change out explicitly in the changelog/PR metadata.
The new source/test headers match the repository's canonical NVIDIA Apache-2.0 header, so no separate licensing concern remains.
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
|
Addressed the |
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
|
The remaining invariant-list P1 is fixed in |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — full pass (no scoping instructions in the trigger comment).
Scope reviewed: 26 changed files. Opened all of modelopt/ (kv_cache_auto_quant.py, model_quant.py, _auto_quantize_cost.py, export/quant_utils.py, export/unified_export_hf.py, export/unified_export_hf_streaming.py, export/model_utils.py, export/convert_hf_config.py, export/quant_aware_conversion.py, export/model_config.py, recipe/config.py), examples/hf_ptq/hf_ptq.py, the new recipe YAML, plus spot-checks of tests/unit/torch/export/test_get_quantization.py and test_unified_export_hf.py. I did not open CHANGELOG.rst, the READMEs, or the bulk of the new test files. (Note: plugins/mcore_qwen.py appears in a two-dot diff against origin/main but is a main-side change, not part of this PR — not reviewed.)
Findings: CRITICAL 1, IMPORTANT 4, SUGGESTION 1.
CRITICAL
get_language_model_from_vl now raises on every standardized VLM (export/model_utils.py:140). Since transformers 4.52 the ...ForConditionalGeneration classes expose language_model as a backward-compat property returning self.model.language_model, so both probed patterns are true and resolve to the same module. pyproject.toml pins transformers>=4.57, so Qwen2.5-VL / Llava / Gemma3-VL all hit ValueError: Found multiple language-model roots. The deleted "always prioritize model.model.language_model" comment existed for exactly this shape.
This is not confined to the pre-existing Nemotron branch: the PR adds an unconditional get_language_model_from_vl(model) call in get_quant_config (quant_utils.py:1704), which runs on every export_hf_checkpoint. VLM HF export breaks outright, KV AutoQuant or not. The new test test_language_model_extraction_rejects_competing_or_aliased_roots[True] locks in the aliased case as an error, so CI stays green while production fails. Fix: compare identity and only refuse genuinely distinct roots.
IMPORTANT
-
KV-only export schema break (
quant_utils.py:1793).weight_quant_algo is None and kv_cache_quantized_layersalso catches the pre-existing "BF16 weights + uniform FP8 KV" flow, turning a runtime-consumablekv_cache_quant_algo: "FP8"/kv_cache_schemeintoMIXED_PRECISION+kv_cache_quantized_layers— which by the PR own Deployment gate has no runtime consumer yet. Gate the envelope on the KV-AutoQuant marker (orlen(kv_cache_formats) > 1) instead of on absent weight quantization. The adjacentelifat line 1803 also converts a previously-succeeding export (uniform weights + partial KV) into aNotImplementedError. -
Replay config hardcodes
algorithm: "max"(kv_cache_auto_quant.py:929) even though_validate_kv_only_configpermitsalgorithm=Noneand the shipped recipe usesalgorithm: nullplusconstant_amax: 448.0.get_auto_quantize_configtherefore returns a config that demands a calibration pass the searched candidates never ran — so the replayed model is not the model that was scored. Derive the algorithm from the selected candidates. -
kv_cache_auto_quantize_report.jsonmixes name spaces (unified_export_hf.py:1520).best["recipe"]is not passed throughname_mapperwhilelayersandsearch_signature["layers"]are, so on the quant-aware conversion path the report selected-format map cannot be joined against the exported checkpoint layer names. -
_select_unpadded_logitssilently changes existing weightkl_divAutoQuantize (hf_ptq.py:530). Thatforward_stepis shared with pre-existingauto_quantize_method: kl_divrecipes; dropping pad positions shifts per-layer sensitivities and can change the selected recipe. The change is correct, but it needs to be either scoped to the KV domain or called out inCHANGELOG.rst— the PR body claim that it "preserves existing non-KV execution" does not cover it.
SUGGESTION
_has_large_cpu_fp8_scale(quant_utils.py:77) makes the "Large KV activations" warning dead on GPU tensors, i.e. in the common export path. Either accumulate the predicate on-device and sync once, or document the CPU-only gating explicitly.
What looked correct
The KL direction (F.kl_div(log_q, log_p, log_target=True) = forward KL from the BF16 reference), the isolated-scoring protocol (target layer on a candidate, all others at the disabled reference), the 4.5-bit NVFP4 storage accounting (4 data + 8/16 scale), the additive LP formulation and its max_cost = sum(k_w+v_w) * target / 16 budget, the Hparam subclass init ordering (the base sets _active directly, so the overridden setter is not invoked before candidate_quantizers exists), the nn.Module calibration-proxy trick plus _freeze_existing_quantizers, and the fail-closed validation surface (_validate_candidate_patterns, _deployable_kv_bits, persistent-scale checks, checkpoint signature matching) all check out.
Risk: high, driven almost entirely by the two export-path changes bundled with the new feature — the VLM regression and the KV-only schema flip both affect users who never touch KV AutoQuant. The new kv_cache_auto_quant.py module itself is self-contained and reads well.
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
|
/claude review |
There was a problem hiding this comment.
Claude review — KV-cache AutoQuant (core)
Scope: full review per the standard procedure (the trigger comment carried no extra scoping). 26 changed files (+3209/−119); reviewed all 15 modelopt/, examples/, and modelopt_recipes/ files in the scoped diff, plus surrounding context in searcher.py, model_calib.py, mode.py, convert_hf_config.py, and layerwise_export.py. Test files were read only where needed to confirm intent. CodeRabbit's lanes (style, typos, security anti-patterns) intentionally skipped.
Findings: CRITICAL: 1 · IMPORTANT: 3 · SUGGESTION: 2
CRITICAL
[CRITICAL Export] layerwise_export.py was not migrated to the new KV format contract — the third HF exporter asserts on any KV-AutoQuant checkpoint.
This PR introduces _get_kv_cache_postprocess_config() precisely because kv_cache_quant_algo can now be the sentinel "MIXED_PRECISION" rather than a real format, and wires it into two of the three exporters that feed _postprocess_single_tensor / postprocess_state_dict:
unified_export_hf.py:1019✅ updatedunified_export_hf_streaming.py:256✅ updatedlayerwise_export.py:204❌ not updated — stillself._quant_config["quantization"]["kv_cache_quant_algo"]
That raw value is handed to _postprocess_single_tensor at layerwise_export.py:404-405, which calls the new _resolve_kv_cache_format_for_key(key, "MIXED_PRECISION"). Since "MIXED_PRECISION" is neither a dict nor KV_CACHE_FP8_K_NVFP4_V, it is returned unchanged and immediately trips:
assert layer_quantization in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], (
"Invalid KV cache quantization format."
)Failure scenario: run the shipped recipe (general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits), then export through the layerwise exporter. get_quant_config sets kv_cache_quant_algo = "MIXED_PRECISION", the first k_bmm_quantizer._amax tensor reaches _collect, and export dies with a bare AssertionError: Invalid KV cache quantization format. — after the expensive search has already completed. The failure is unconditional for mixed-precision KV, so it is not an edge case; it just isn't covered by the tests in this PR.
Fix is one line, using the helper this PR already added (the key passed to _postprocess_single_tensor at line 404 is the in-memory name and self._name_mapper is applied afterwards, so the per-layer-map lookup resolves correctly):
# modelopt/torch/export/layerwise_export.py
from .quant_utils import (
_get_kv_cache_postprocess_config,
_postprocess_single_tensor,
get_quant_config,
get_quantization_format,
)
...
# Not get_kv_cache_dtype: it does not recurse, so on the root it answers None.
self._kv_cache_format = _get_kv_cache_postprocess_config(
self._quant_config["quantization"]
)IMPORTANT (posted inline)
unified_export_hf.py:1667(and:1635) —name_mappernot reset on the reverse-name-mapping fallback. The comment above the block promises the fallback is atomic across weights and config; now thatname_mapperis also passed to_write_hf_export_config, a mid-tryfailure leaves the KV report on hub names while the checkpoint andhf_quant_config.jsonkeep in-memory names.kv_cache_auto_quant.py:876— the KV-AutoQuant marker is a transient instance attribute._modelopt_kv_cache_auto_quantize_stateis not inmodelopt_state/state_dict(), so aftermto.save()→mto.restore()→ export,is_kv_autoquant_resultisFalse. When the solver picked one format for every layer, the export silently degrades to uniformkv_cache_quant_algo: "FP8"with nokv_cache_quantized_layersmap — the exact ambiguity theMIXED_PRECISIONenvelope exists to prevent. Selection-dependent, so easy to miss.model_quant.py:603— the intentionally-excluded GEMM→KV composition isn't rejected up front. Withis_quantized(model)alreadyTrue, the"*" → enable: Falsereset is skipped, the KL reference is measured against a weight-quantized model (breaking the "isolated forward KL vs. BF16" objective the cost model assumes), and the run only fails atget_quant_config'sNotImplementedErrorafter the full search. Araisein theis_kv_searchguard block makes the documented exclusion cheap and gives #2273 one place to relax.
SUGGESTIONS (posted inline)
quant_utils.py:1791-1816— the first and third branches have byte-identical bodies; collapse to oneneeds_mixed_envelopecondition. Plus a question on whether the weight-sidequant_algo = "MIXED_PRECISION"/ emptyconfig_groupsis actually required for a BF16-weight KV-only checkpoint.model_utils.py:143-151—get_language_model_from_vlnow raises on ambiguous roots, and this PR newly calls it fromget_quant_config, i.e. on every HF export rather than only VLM paths. Suggest keeping the raise scoped to the KV eligibility path, or degrading to "no boundary known" inget_quant_config.
What checked out
Verified as correct, for the record:
- Forward-KL direction.
F.kl_div(input=log_softmax(quant), target=log_prob_ref, log_target=True, reduction="sum")computesKL(P_ref ‖ Q_cand)— forward KL with the BF16 reference asP, matching the stated objective. fp32 upcast before bothlog_softmaxcalls; normalization byscored_tokensmatches the recordedscore_reduction: "mean_per_scored_token". - Solver units.
candidate_costs = (k_width·k_bits + v_width·v_bits)/16againstmax_cost = Σ(k_width+v_width)·target/16, andachieved_bits = total_cost/denominatorin raw bits — dimensionally consistent, and_validate_candidate_cost_geometrycorrectly gates asymmetric K/V candidates on unequal-width layers. - Candidate bit accounting. FP8 → 8.0, block-16 dynamic NVFP4 → 4.5 (4 data + 8/16 block scale);
_validate_deployable_candidatecross-checks the declaredeffective_bitsagainst(k_bits+v_bits)/2, which the shipped recipe's8.0/4.5satisfy. - No calibration bleed into scoring.
finish_stats_collectionends withenable_quant()+disable_calib()on the candidate quantizers, and_freeze_existing_quantizersrestores_if_calibfor everything else, so candidate amaxes are frozen when_estimate_sensitivity_scoresruns — the saved checkpoint scales match what_solveapplies. Isolation is real:use_reference()is reset inside the candidate loop, so exactly one layer is quantized per scored forward. - Constraint normalization round-trips.
normalize_auto_quantize_constraintspops the emptycostkey forKVCacheCostModel, so the second_validate_search_inputscall inbefore_searchsees{effective_bits, cost_model}and its strict key-set check still passes. - Ordered-config replay.
get_kv_cache_auto_quantize_configpreserves candidate entry order when rewriting*[kv]_bmm_quantizerpatterns to{layer}.{k,v}_bmm_quantizer, so the FP8-K override correctly wins over the NVFP4 base entry for the asymmetric candidate. - Name-mapping order in the streaming and layerwise exporters — both resolve the per-layer map against the in-memory key and apply
name_mapperafterwards, which is the correct order. kv_cache_quant_algokey presence — the base config atquant_utils.py:1692already seeds it toNone, so dropping the old unconditional assignment does not introduce aKeyErrorfor unquantized-KV models.
Risk assessment
Medium-high. The search core itself is careful and the algorithm-level math holds up. The risk is concentrated in the export surface: the new kv_cache_quant_algo = "MIXED_PRECISION" sentinel changes an interface that three exporters consume, and one of them was left on the old contract (CRITICAL 1). Findings 2 and 3 are both "works in the single-process hf_ptq flow, degrades or fails late outside it," which is the pattern most likely to reach users after merge. The deployment gate on vllm-project/vllm#52813 is clearly stated in the description and is the right call.
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
|
Addressed the latest Claude review in 1cd026c. The critical layerwise-export gap now uses the same mixed-KV postprocess resolver as resident and streaming export, with focused coverage proving that FP8/NVFP4 layer mappings reach the layerwise exporter rather than the MIXED_PRECISION sentinel. The five inline findings were also fixed, replied to, and resolved: atomic reverse-name fallback, ModelOpt save/restore persistence for the KV report, early rejection of the excluded GEMM-to-KV composition, deduplicated mixed-envelope logic with vLLM consumer verification, and tolerant generic export for competing VLM roots. Validation: 132 passed and 1 optional-dependency skip in the focused KV/export/restore suite; the broader recipe, hf_ptq, and AutoQuant run had 533 passes, with the socket-limited distributed test passing separately and four unrelated tests blocked by the workstation Transformers package skew; changed-file pre-commit passed. |
|
/claude review |
| if needs_mixed_envelope: | ||
| if weight_quant_algo not in (None, "MIXED_PRECISION"): | ||
| raise NotImplementedError( | ||
| "Mixed-precision KV-cache export with a uniform quantized-weight format is " | ||
| "not supported yet. Use BF16 weights or a mixed-weight AutoQuant recipe." | ||
| ) | ||
| # Keep the complete AutoQuant layer map even when every layer selected the same format. | ||
| quant_config["quantization"]["quant_algo"] = "MIXED_PRECISION" | ||
| quant_config["quantization"].setdefault("quantized_layers", {}) | ||
| quant_config["quantization"]["kv_cache_quant_algo"] = "MIXED_PRECISION" | ||
| quant_config["quantization"]["kv_cache_quantized_layers"] = kv_cache_quantized_layers | ||
| quant_config["quantization"]["kv_cache_schema_version"] = 1 |
There was a problem hiding this comment.
[IMPORTANT Export] A KV-only AutoQuant result (BF16 weights) overwrites the weight quant_algo with "MIXED_PRECISION" and an empty quantized_layers.
For the shipped recipe path, no weight is quantized, so weight_quant_algo is None here and the branch is taken. The resulting hf_quant_config.json says:
{"quantization": {"quant_algo": "MIXED_PRECISION", "quantized_layers": {},
"kv_cache_quant_algo": "MIXED_PRECISION", "kv_cache_quantized_layers": {...}}}and convert_hf_quant_config_format then turns that into "config_groups": {} plus "quant_algo": "MIXED_PRECISION" in config.json. Two concrete consequences:
quant_algois the field TRT-LLM and vLLM branch on to select the weight loader. DeclaringMIXED_PRECISIONfor a model with zero quantized weights sends those consumers down the per-layer mixed-weight path, where they will find no entries at all. The PR body documents the KV envelope but not this weight-side claim.- Because
is_kv_autoquant_resultalone forces the envelope, a KV AutoQuant run that happens to select FP8 for every eligible layer exports a different, mutually incompatible schema from the equivalent PTQ checkpoint (kv_cache_quant_algo: "FP8"→kv_cache_scheme).convert_hf_quant_config_formatemits nokv_cache_schemeforMIXED_PRECISION, so today's readers see a checkpoint that carriesk_scale/v_scaletensors but advertises no KV scheme they understand.
Suggested fix: leave quant_algo/quantized_layers untouched when no weight is quantized, and when len(kv_cache_formats) == 1 keep the legacy uniform kv_cache_quant_algo value while also emitting kv_cache_quantized_layers + kv_cache_schema_version for the explicit map. That preserves the "complete layer map" goal from the PR description without breaking consumers that key on the pre-existing fields, and reserves MIXED_PRECISION for genuinely mixed selections.
| def _has_large_cpu_fp8_scale(value: torch.Tensor) -> bool: | ||
| """Check CPU scales only; GPU scales skip this optional warning to avoid host sync.""" | ||
| return value.device.type == "cpu" and bool(torch.any(value > 0.5)) |
There was a problem hiding this comment.
[SUGGESTION] This silently removes the large-KV-activation warning for the common case.
All three call sites (get_kv_cache_scaling_factor at L415, _postprocess_single_tensor at L1088, postprocess_state_dict at L1191) run on tensors that are on the accelerator for a normally-loaded model, so the device.type == "cpu" gate means the accuracy warning now only fires for CPU/offloaded exports. Users who previously got it on a GPU export get nothing.
The sync being avoided is ~2 scalar bool() reads per attention layer (≈56 for a 28-layer model) in an export path that is dominated by safetensors I/O — not a hot path. If you want to keep the .item() out, prefer keeping the comparison on-device and only syncing once:
def _has_large_fp8_scale(value: torch.Tensor) -> bool:
"""One scalar sync per KV scale; negligible against export I/O."""
return bool(torch.any(value > 0.5))If the CPU-only restriction is deliberate, the function name and docstring should say the warning is best-effort and skipped on device tensors, so the loss is visible to future readers.
| "effective_bits": config.effective_bits, | ||
| "k_bits": _candidate_kv_bits(config)[0], | ||
| "v_bits": _candidate_kv_bits(config)[1], | ||
| "config": config.model_dump(mode="python", exclude_none=True), |
There was a problem hiding this comment.
[SUGGESTION] mode="python" here breaks the JSON-safety invariant that _report_state documents.
self.candidates is part of state_dict(), _report_state only strips quantizer_state, and _write_hf_export_config runs json.dump on the result into kv_cache_auto_quantize_report.json. QuantizerAttributeConfig.calibrator is typed str | ConstructorLike, so a candidate that supplies a calibrator constructor (or any other non-JSON-native value) survives a mode="python" dump as a live object and makes that json.dump raise TypeError — at the very end of a long search+export run.
_search_signature (L421) already uses mode="json" for the same object. Aligning them is a one-word change and removes the failure mode:
| "config": config.model_dump(mode="python", exclude_none=True), | |
| "config": config.model_dump(mode="json", exclude_none=True), |
Note get_kv_cache_auto_quantize_config re-parses this via QuantizeConfig(**selected_config), which accepts the JSON-mode form, so the replay path is unaffected.
| kv_autoquant_report["best"]["recipe"] = { | ||
| name_mapper(name): value for name, value in best_recipe.items() | ||
| } | ||
| signature_layers = kv_autoquant_report.get("search_signature", {}).get("layers", []) |
There was a problem hiding this comment.
[SUGGESTION] AttributeError when search_signature is present but None.
AutoQuantizeKVSearcher.default_state_dict declares "search_signature": None, so the key always exists in the report. dict.get(key, {}) returns the stored None (the default only applies to a missing key), and None.get("layers", []) raises. Today after_search always populates it, but any report round-tripped through a state dict that never reached before_search (e.g. a stale/partial checkpoint surfaced through mto.restore) turns a metadata-renaming step into a hard export failure.
| signature_layers = kv_autoquant_report.get("search_signature", {}).get("layers", []) | |
| signature_layers = (kv_autoquant_report.get("search_signature") or {}).get( | |
| "layers", [] | |
| ) |
| converted_for_search = not is_quantized(model) | ||
| conversion_snapshot = _snapshot_model_structure(model) if converted_for_search else [] |
There was a problem hiding this comment.
[SUGGESTION] converted_for_search is invariantly True here — L554-558 already raised NotImplementedError when is_quantized(model), so not is_quantized(model) cannot be False at this point. The three if converted_for_search: guards (L613, L635) and the else [] snapshot branch are therefore dead.
Either drop the flag and always snapshot/convert/restore, or add a short comment stating it is a placeholder for the composed GEMM→KV path in #2273 so the next reader doesn't try to reason about the unquantized-input case.
| for data in iterator: | ||
| if scored_steps >= self.config["num_score_steps"]: | ||
| break | ||
| logits_ref = _get_logits( | ||
| self.config["forward_step"], self.model, data, validate_finite=False | ||
| ) | ||
| batch_logits_finite = torch.isfinite(logits_ref).all() | ||
| log_prob_ref = torch.log_softmax(logits_ref.float(), dim=-1) | ||
| scored_tokens += logits_ref.numel() // logits_ref.shape[-1] | ||
|
|
||
| for hparam in self._hparams: | ||
| for candidate_index, candidate_name in enumerate(candidate_names): | ||
| hparam.active = candidate_index | ||
| logits_quant = _get_logits( | ||
| self.config["forward_step"], self.model, data, validate_finite=False | ||
| ) | ||
| batch_logits_finite.logical_and_(torch.isfinite(logits_quant).all()) |
There was a problem hiding this comment.
[SUGGESTION] Worth documenting the scoring cost, which is num_score_steps × num_eligible_layers × num_candidates full model forwards (plus one reference forward per step).
For the shipped recipe on Qwen3-1.7B (28 layers, 2 candidates, score_size: 128, batch_size 1) that is ~7.2k forwards; a 64-layer model with 3 candidates at the same score_size is ~24.6k. That is inherent to strictly isolated per-layer KL, but the recipe/README currently gives no signal that score_size multiplies by layer count and candidate count, so a user scaling to a larger model will be surprised by the wall-clock.
Also note the peak memory here: log_prob_ref is held in fp32 for the whole layer×candidate loop (num_selected_tokens × vocab × 4 bytes — ~1.2 GB for 2048 tokens at a 152k vocab), and each candidate allocates a second fp32 log-softmax of the same size. A sentence in the recipe comment or examples/hf_ptq/README.md about picking score_size against layer count and vocab size would save users a trial run.
There was a problem hiding this comment.
Claude review — 1 IMPORTANT, 5 SUGGESTION
Scope: the trigger comment was a bare /claude review, so this is a full review. 29 files changed (+3290/-123). I read every modelopt/ and examples/ file in the diff — kv_cache_auto_quant.py in full, plus model_quant.py, _auto_quantize_cost.py, conversion.py, modelopt/recipe/config.py, export/{quant_utils,unified_export_hf,unified_export_hf_streaming,convert_hf_config,layerwise_export,model_utils,model_config,quant_aware_conversion}.py, examples/hf_ptq/hf_ptq.py, and the new recipe YAML. I did not open the 12 test files or CHANGELOG.rst/README.md beyond their diffs. (export/plugins/mcore_qwen.py shows up in a two-dot diff against the current main tip but is not in this PR's file list — base drift, ignored.)
Findings
IMPORTANT (1)
quant_utils.py:1802-1813— a KV-only AutoQuant result (BF16 weights) overwrites the weightquant_algowith"MIXED_PRECISION"and an emptyquantized_layers, whichconvert_hf_quant_config_formatturns into"config_groups": {}.quant_algois the field TRT-LLM/vLLM branch on to pick the weight loader, so this advertises mixed-precision weights for a model with none. Related: becauseis_kv_autoquant_resultalone forces the envelope, a search that selects FP8 for every eligible layer exports a schema that is not interchangeable with the equivalent uniform-FP8 PTQ checkpoint, and nokv_cache_schemeis emitted at all. The PR body documents the KV envelope but not the weight-side claim.
SUGGESTION (5)
quant_utils.py:79-81—_has_large_cpu_fp8_scalegates the large-KV-activation warning ondevice.type == "cpu", so all three call sites lose it for normally-loaded (GPU) exports. The sync it avoids is ~2 scalar reads per attention layer in an I/O-bound path.kv_cache_auto_quant.py:703—model_dump(mode="python")forself.candidatesviolates the JSON-safety invariant_report_statedocuments; acalibratorsupplied as a constructor (allowed by the config schema) makes thekv_cache_auto_quantize_report.jsondump raise._search_signaturealready usesmode="json".unified_export_hf.py:1532—kv_autoquant_report.get("search_signature", {}).get("layers", [])raisesAttributeErrorwhen the key is present-but-None, which is exactly the declared default state.model_quant.py:608-609—converted_for_searchis invariantlyTrue(theis_quantized(model)case already raised above), so its three guards are dead.kv_cache_auto_quant.py:737-753— scoring isnum_score_steps x num_layers x num_candidatesfull forwards (~7.2k for the shipped Qwen3-1.7B recipe, ~24.6k for a 64-layer/3-candidate run) and holds two fp32[tokens, vocab]log-softmax tensors; worth a note in the recipe/README soscore_sizeis picked with layer count and vocab size in mind.
What I traced and found correct
- Forward-KL direction.
F.kl_div(log_softmax(quant), log_prob_ref, log_target=True, reduction="sum")computesKL(p_ref || q_quant)— genuinely forward KL with the BF16 reference asp. fp32 upcast is present on both sides, the reference is recomputed per batch, and every other layer is held at the reference during scoring, so isolation holds. - Cost model and solver.
k_width*k_bits + v_width*v_bitswith the/16.0scaling applied consistently to bothcandidate_costsandmax_cost;achieved_bits = total_cost / sum(k_width+v_width)is the correctly width-weighted average. NVFP4 at 4.5 bits (4 + 8/16 block scale) is right, and_validate_candidate_cost_geometrycorrectly refuses asymmetric K/V candidates on layers with unequal K/V widths, which is exactly where the unweighted(k_bits+v_bits)/2effective_bitsdeclaration would stop being meaningful. - No
.item()in the accumulation loop — scores stay as device tensors and only sync once at the end. Good. - Candidate device handling. Candidate quantizers are built on a CPU proxy module, but
TensorQuantizer._get_amaxalready realigns a config-time_amaxbuffer to the input device, so attaching them to a CUDA model is safe. constant_amaxreplay.max_calibrateskips quantizers with_constant_amaxset (model_calib.py:295,1133,1151), soget_kv_cache_auto_quantize_configreturningalgorithm="max"for a mixed candidate set cannot clobber the pinned constant scales.- Searcher lifecycle. The extra
default_search_configkeys satisfyBaseSearcher.sanitize_search_config's exact-key assertion;Hparam.__init__sets_activedirectly rather than through the property, so theQuantKVRecipeHparamsetter's dependency onself.candidate_quantizersis not hit before assignment. modelopt_statecompat.kv_cache_auto_quantize_stateis an additive, optional metadata key — old checkpoints simply lack it andupdate_quantize_metadatapops it when absent, so no migration is needed.- The
get_quant_configKV-detection rewrite is a net fix, not a regression: the oldgetattr(module, "output_quantizer", ...).is_enabledarm fired on any Linear with output quantization enabled and then calledget_kv_cache_dtypeon that Linear. Requiring a k/v BMM pair is stricter and more correct.kv_cache_quant_algoalso remains present-with-Nonevia the base dict at L1692, so no key disappears fromhf_quant_config.json.
Risk
Moderate. The search core is self-consistent, fails closed on unsupported candidates/architectures/distributed setups, and the algorithm math checks out. The residual risk is concentrated in the producer schema — the one IMPORTANT finding is the only place where an existing consumer can be misled by a field it already reads, and it is worth settling before the vLLM consumer in vllm-project/vllm#52813 lands against it.
3 participants
What does this PR do?
Type of change: new feature.
Adds standalone layer-wise KV-cache AutoQuantize through the existing public
mtq.auto_quantizeAPI:constraints={"effective_bits": ..., "cost_model": "kv_cache"}and forward-KLsensitivity;
PuLP-backed constrained solver;
BaseSearcherlifecycle and safe checkpoint restore/save machinery;examples/hf_ptq/hf_ptq.pythrough a standalonecalibration-free recipe.
The implementation is architecture-driven. Plain and conditional-generation Qwen
causal attention is supported, VLM vision attention is excluded through the existing
language-model extraction boundary, hybrid full-attention mixers are discovered through
their paired K/V quantizers, and nonattention/Mamba modules remain outside the search.
Ambiguous language-model roots, unsupported distributed execution, structural
algorithms, invalid storage declarations, nonpersistent scales, and unsupported K/V
pairs fail closed.
KV-only unified HF exports use a
MIXED_PRECISIONenvelope plus a completekv_cache_quantized_layersmap, including uniform all-FP8 and all-NVFP4 selections.This keeps the checkpoint schema explicit and prevents disabled vision attention from
making uniform language-model KV quantization appear partially quantized.
GEMM PTQ/AutoQuantize followed by KV AutoQuantize is intentionally excluded and proposed
separately in stacked PR #2273.
Why KV search has a dedicated backend
The user-facing entry point remains
mtq.auto_quantize; no separate public KV search APIis introduced.
AutoQuantizeKVSearcherextendsBaseSearcherand reuses its reset,checkpoint load/save, and search lifecycle, along with existing Pydantic configuration,
calibration, safe checkpoint I/O, and PuLP-backed selection utilities.
The backend remains KV-specific because a decision owns paired K/V quantizers on one
attention layer, its cost depends on separate K/V widths and data/scale storage, BF16 is a
scoring reference but not a deployable solver choice, and the optimization objective is
additive isolated forward KL under a KV-storage constraint. These contracts do not match
the weight-domain hparam grouping, parameter-count cost, or threshold-selection behavior
of the existing weight AutoQuant searchers. Keeping the specialization behind the shared
API avoids changing established weight-search solver and scoring behavior.
Usage
The search checkpoint is compatible only with the same model, eligible-layer geometry,
candidate configurations, and scoring setup. Use a distinct checkpoint path after any of
those inputs change.
KV-cache AutoQuantize rejects
--use_fsdp2before model loading because its sensitivityscoring, selection, and checkpoint writes are single-process. Existing weight
AutoQuantize retains its previous experimental FSDP2 warning and behavior.
Testing
storage accounting, solving, checkpoint resume, failure atomicity, disabled layers,
fresh-model replay, Qwen/VLM/hybrid boundaries, JSON-safe reports, and unified export.
envelope; disabled VLM vision attention is excluded from causal-KV eligibility.
exportable scale state.
main: 432 focused recipe/KV/export/hf_ptq tests passed, with oneunrelated optional-dependency skip; changed-file pre-commit hooks passed.
Deployment gate
The producer schema is covered here. Runtime consumption of
kv_cache_quantized_layersis tracked in vLLM PRvllm-project/vllm#52813. Do not treat a produced checkpoint as
runtime-supported until that consumer lands and the target K/V kernels are available.
Before your PR is "Ready for review"
guidance in
CONTRIBUTING.md: N/AAdditional information
review scope; composition is isolated in [OMNIML-5570] 2/2 Compose GEMM and KV-cache AutoQuant workflows #2273.
launch logic, or runtime-kernel implementations are included.