Skip to content

Compose GEMM and KV-cache AutoQuant workflows - #2273

Draft
meenchen wants to merge 1 commit into
agent/kv-cache-autoquant-corefrom
agent/kv-autoquant-composition
Draft

Compose GEMM and KV-cache AutoQuant workflows#2273
meenchen wants to merge 1 commit into
agent/kv-cache-autoquant-corefrom
agent/kv-autoquant-composition

Conversation

@meenchen

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: new feature.

Stacked on #2272. Adds composition of existing GEMM quantization with the standalone KV-cache AutoQuant API:

  • fixed FP8 GEMM PTQ followed by mixed-KV AutoQuantize;
  • gradient-based NVFP4/FP8 GEMM AutoQuantize followed by independent mixed-KV AutoQuantize;
  • an optional kv_auto_quantize recipe stage with independent method, constraints, candidates, and checkpoint path;
  • ordered hf_ptq.py orchestration that keeps the selected weight/activation QDQ behavior active while freezing its calibration state during KV scoring;
  • fail-closed validation of the actual converted model when a preceding stage leaves K/V quantizers enabled; and
  • unified export of a uniform-weight or mixed-weight checkpoint together with the selected per-layer KV map.

The two stages remain independent searches and reuse the public APIs from #2272. This PR does not change either solver, scoring protocol, or checkpoint schema.

Usage

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

Weight AutoQuantize followed by 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

Testing

  • 378 focused tests passed across both composed recipes, KV AutoQuant, recipe loading, and unified export.
  • All three shipped standalone/composed recipe paths invoke the actual public KV AutoQuant API on tiny offline Qwen fixtures.
  • Changed-file pre-commit hooks passed. The isolated recipe validator was rerun successfully after its initial network-only dependency bootstrap failure.
  • The complete tracked tree is byte-identical to the preserved combined ground-truth head f94d74a98 (tree 2564939297095db02c0ffea5848a0c1f99b18d6d).

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?: N/A (covered by Add standalone layer-wise KV-cache AutoQuant with forward KL #2272)
  • Did you get Claude approval on this PR?: ❌ (draft)

Additional Information

  • This PR intentionally targets agent/kv-cache-autoquant-core so reviewers see only the nine composition-specific files.
  • --auto_quantize_checkpoint and --kv_auto_quantize_checkpoint are intentionally separate. Changing the preceding GEMM configuration requires a fresh KV checkpoint path.
  • Downstream deployment of uniform-weight plus mixed-KV metadata remains gated until the runtime's uniform-weight ModelOpt configuration consumes kv_cache_quantized_layers.

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 27, 2026

Copy link
Copy Markdown

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

coderabbitai Bot commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant