Add Aumann-Shapley sensitivity scoring method to auto_quantize

[joshua-hill]

What does this PR do?

Type of change: new feature

Adds method="aumann_shapley" to mtq.auto_quantize: label-free sensitivity scoring via
Aumann-Shapley path-integral damage attributions, so every allocation carries a
predicted_damage quote in calibration-KL units (nats) rather than only an unitless score.

Damage is measured as KL divergence against the model's own reference outputs — the reference
keeps any fixed or forced-single-format groups quantized, so scores are incremental KL relative
to the resolved baseline ({"type": "unquantized"} when nothing is pinned). At each midpoint
node t = (k + 1/2) / num_path_nodes of the joint quantization path, every scored module
propagates y + t * (Q(y) - y) using detached local-replay differences, and one backward pass
accumulates <dL/dy, Q(y) - y> per (group, format). The path integral is required because a KL
loss has dKL = 0 exactly at the unquantized point — that is precisely why the existing
gradient method must square its Taylor term into a Fisher proxy and therefore needs labels,
while this method does not.

A measured aggressive corner anchors a coverage link damage = c * (1 - exp(-sum(b))); a
fixed-point inversion turns attributions into per-group log-headroom written into
candidate_stats["scores"], so the standard solve is the coverage-optimal allocation. Solver
scores are projected onto the monotone compression ladder so quotes stay conservative. The quote
is an internal-model estimate, not a bound on realized deployment KL.

Cost per batch is one reference forward, one corner forward, and one forward+backward per
(candidate format, path node) — independent of how many configurations the solver later
considers.

Supporting changes:

• New auto_quantize(method_options=...) argument (optional), validated against each
  searcher's declared method_options_keys so core inputs cannot be overridden. Supported keys:
  num_path_nodes, damage_link, solver (lp exact / dp deterministic grid-approximate),
  and max_predicted_damage (minimize weight cost subject to a damage bound; mutually exclusive
  with an effective_bits constraint).
• Scoring methods are now a registry (AUTO_QUANTIZE_SEARCHERS in
  modelopt.torch.quantization.algorithms) instead of hardcoded dispatch; methods register
  themselves. gradient / kl_div behavior is unchanged.
• Internal format tables are keyed by QuantRecipe.checkpoint_signature, so identical custom
  formats under different auto-generated names resolve to one format. A scoring signature in the
  search state rejects checkpoint resumes that would change what stored scores mean, while still
  allowing solver-only re-solves.
• Separate commit: group MLA q_a_proj / kv_a_proj_with_mqa projections in auto_quantize.
  TRT-LLM fuses these low-rank input projections (DeepSeek / GLM lineage) into a single
  fused_qkv_a_proj_with_mqa GEMM, so both shards must be assigned one shared format; without
  this rule any scoring method can emit checkpoints the runtime cannot fuse.

An efficient implementation of the estimator in https://arxiv.org/abs/2607.12266
(overview,
thread), validated empirically
against it; implementation details are documented in the module docstring.
Vocab-sharded (Megatron-TP) losses raise NotImplementedError pending an autograd-correct
vocab-parallel log-softmax.

Usage

import modelopt.torch.quantization as mtq

model, search_state = mtq.auto_quantize(
    model,
    constraints={"effective_bits": 4.8},
    quantization_formats=["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG"],
    data_loader=calib_loader,
    forward_step=lambda model, batch: model(**batch),
    method="aumann_shapley",  # label-free; no `labels` needed in the batch
    method_options={"num_path_nodes": 1, "damage_link": "coverage", "solver": "lp"},
)

# Every allocation carries a damage quote in calibration-KL units (nats).
print(search_state["predicted_damage"])

Bound the predicted damage instead of the bit budget (mutually exclusive with effective_bits):

model, search_state = mtq.auto_quantize(
    model,
    constraints={},
    quantization_formats=["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG"],
    data_loader=calib_loader,
    forward_step=forward_step,
    method="aumann_shapley",
    method_options={"max_predicted_damage": 0.05},
)

Also selectable from examples/hf_ptq/hf_ptq.py via --auto_quantize_method aumann_shapley, and
from an AutoQuantize recipe via auto_quantize_method: aumann_shapley.

Testing

• tests/unit/torch/quantization/test_autoquant_shapley.py (new): config parity with the standard
  builder (dict-for-dict), the path-integral completeness diagnostic, corner anchoring under
  incomplete attributions, damage-bound certification, solver optimality contracts against brute
  force, custom-format identity, heterogeneous-ladder flagging, exact-zero and tiny-attribution
  inversion behavior, scoring-signature resume guards, and method-option validation.
• tests/unit/torch/quantization/test_autoquant.py: method parametrizations extended (21 new
  cases); the existing gradient / kl_div cases pass unmodified, confirming no behavior change
  for the current methods.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• 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: ✅
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌

Additional Information

Draft — opened for early feedback on the method-registry and method_options plumbing before
final review. Rebased onto main.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: edd424de-3df3-4664-8da3-1e4bb0c11ed3

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}