mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Wei-Ming Chen
59d93af064
### What does this PR do?
Type of change: new feature.
Adds standalone layer-wise KV-cache AutoQuantize through the existing
public
`mtq.auto_quantize` API:
- dispatches KV search with
`constraints={"effective_bits": ..., "cost_model": "kv_cache"}` and
forward-KL
sensitivity;
- selects one supported K/V format for every eligible causal-attention
layer;
- supports persistent/exportable FP8 K/V, NVFP4 K/V, and FP8-K/NVFP4-V
candidates;
- solves a K/V-width- and scale-storage-aware additive recipe with the
existing
PuLP-backed constrained solver;
- uses `BaseSearcher` lifecycle and safe checkpoint restore/save
machinery;
- preserves existing non-KV execution while isolating K/V candidate
calibration;
- returns standard AutoQuantize state that can be re-solved at another
KV budget;
- produces a complete KV-only replay config that disables every non-KV
quantizer;
- saves JSON-safe sensitivity metadata and the exact selected layer
mapping; 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, 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 leave weight-quantization fields unset.
Uniform all-FP8 or
all-NVFP4 selections retain their legacy KV scheme while also carrying
the complete
`kv_cache_quantized_layers` map and schema version; genuinely
layer-mixed selections use
the KV-side `MIXED_PRECISION` marker plus the same map. This keeps
weight-loader metadata
accurate 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 API
is introduced. `AutoQuantizeKVSearcher` extends `BaseSearcher` and
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
```bash
python examples/hf_ptq/hf_ptq.py \
--pyt_ckpt_path Qwen/Qwen3.8-27B \
--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.8-27b-mixed-kv
```
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_fsdp2` before model loading because
its sensitivity
scoring, selection, and checkpoint writes are single-process. Existing
weight
AutoQuantize retains its previous experimental FSDP2 warning and
behavior.
### Testing
- Focused coverage exercises candidate validation/calibration, paired
K/V scoring and
storage accounting, solving, checkpoint resume, failure atomicity,
disabled layers,
fresh-model replay, Qwen/VLM/hybrid boundaries, JSON-safe reports, and
unified export.
- Uniform FP8/NVFP4 KV-only exports retain the legacy KV scheme and
complete layer map
without claiming a weight algorithm; disabled VLM vision attention is
excluded from
causal-KV eligibility.
- The shipped recipe runs end to end on a tiny offline Qwen fixture and
preserves
exportable scale state.
- After merging current `main`: 432 focused recipe/KV/export/hf_ptq
tests passed, with one
unrelated optional-dependency skip; changed-file pre-commit hooks
passed.
### Deployment gate
The producer schema is covered here. Runtime consumption of
`kv_cache_quantized_layers` is tracked in vLLM PR
https://github.com/vllm-project/vllm/pull/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"
- 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?: ✅
### Additional information
- This is split from the combined ground-truth implementation in draft
PR #2211 to reduce
review scope; composition is isolated in #2273.
- The standalone core tree contains no composed GEMM→KV recipe schema or
orchestration.
- No model-name checks, checkpoint-specific layer lists, campaign data
contracts, cluster
launch logic, or runtime-kernel implementations are included.
---------
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>