Add WMSE weight scale search calibration

[realAsma]

What

Adds wmse, an NVFP4 weight-scale-search calibration algorithm that weights the
block quantization error by a per-input-channel importance
Imp_i = ||X[:, i]||^2, instead of the full per-cin-block Hessian X^T X used
by local_hessian. This is the ScaleSweep objective (arXiv 2606.07618,
eqs. 12-13) — the same quadratic form with H replaced by diag(Imp).

Stacked on #1976. Base branch is asma/activation-quant-aware-scale-setting,
so the diff here is WMSE-only. Review #1976 first.

How

Rather than duplicating the calibration pipeline, the shared
accumulate → metric → MSE-search body is extracted into
_hessian_weighted_calibrate, parameterized by the weighting term.
local_hessian_calibrate and the new wmse_calibrate are thin wrappers that
differ only in a diagonal flag:

• The accumulator's diagonal path accumulates x.square().sum(-1) into a
  [n_blocks, block_size] buffer instead of [n_blocks, block_size, block_size].
• The block loss dispatches on tensor rank: einsum("cnb,nb->cn") for the
  diagonal case, einsum("cnb,nbd,cnd->cn") for the full Hessian.
• The Triton FP8 scale-sweep kernel _fp8_scale_sweep_hessian_kernel gains a
  matching DIAGONAL: tl.constexpr specialization that replaces tl.dot with a
  weighted sum, so the diagonal case does not materialize diag_embed(Imp) on
  the hot path across the 126 candidate scales.

Also adds the ABLATE_W4A4_WMSE / ABLATE_W4A16_WMSE eval configs and the
--calib_dataset / --calib_seqlen plumbing in examples/llm_eval that they
need, so the ablation below is reproducible.

Refactor safety — local_hessian is bit-identical

The shared-path refactor touches local_hessian, so it was verified rather than
assumed. Qwen3-1.7B was calibrated in two separate processes — pre-refactor and
post-refactor — from identical seeded input (matching input_ids SHA-256),
through the real 28-layer layerwise NVFP4 path with the FP8 scale sweep on.

Fingerprint: 784 tensors per run covering quantizer amax / _amax / _scale /
_double_scale / _pre_quant_scale, all non-recursive buffers, and the final
fake-quantized weights, compared by value and by int32 bit pattern.

• local_hessian: 784/784 bit-identical.
• mse control: 784/784 bit-identical.
• Non-vacuity check: July-local_hessian vs July-mse differs in 588/784
  tensors, so the fingerprint does detect real objective changes.

Proven for the dense single-GPU path; fused-MoE and multi-rank are covered by
unit tests only.

Ablation

10 cells (5 models x NVFP4 W4A4 / W4A16), calibration matched exactly to the
existing local-Hessian campaign: nemotron-post-training-v3, --calib_size 512,
--calib_seqlen 2048, --calib_batch_size 1, layerwise calibration on, FP8
scale sweep on, WORLD_SIZE=1, one GPU per job, no chat template. Baseline rows
are the stored local-Hessian campaign numbers, not re-runs.

avg acc % is the unweighted mean of mmlu.acc, arc_challenge.acc_norm,
arc_easy.acc_norm, boolq.acc, hellaswag.acc_norm, piqa.acc_norm, winogrande.acc.
Deltas are versus local_hessian.

cell	method	avg acc %	vs LH	wikitext ppl	vs LH
qwen3-1.7B W4A4	max	57.87	-2.13	24.389	+1.597
qwen3-1.7B W4A4	mse	59.68	-0.32	23.635	+0.843
qwen3-1.7B W4A4	local_hessian	60.00	—	22.792	—
qwen3-1.7B W4A4	wmse	60.16	+0.17	22.609	-0.184
qwen3-1.7B W4A16	max	58.76	-3.20	23.235	+1.409
qwen3-1.7B W4A16	mse	62.22	+0.26	22.457	+0.632
qwen3-1.7B W4A16	local_hessian	61.96	—	21.826	—
qwen3-1.7B W4A16	wmse	62.33	+0.37	21.630	-0.196
qwen3-8B W4A4	max	72.39	-0.12	13.314	+0.110
qwen3-8B W4A4	mse	72.17	-0.34	13.120	-0.084
qwen3-8B W4A4	local_hessian	72.51	—	13.204	—
qwen3-8B W4A4	wmse	72.71	+0.20	13.236	+0.032
qwen3-8B W4A16	max	73.52	+0.01	12.936	+0.093
qwen3-8B W4A16	mse	73.41	-0.11	12.716	-0.127
qwen3-8B W4A16	local_hessian	73.51	—	12.843	—
qwen3-8B W4A16	wmse	73.67	+0.16	12.909	+0.066
qwen3.5-2B W4A4	max	58.54	-0.15	17.162	+0.543
qwen3.5-2B W4A4	mse	58.05	-0.63	16.765	+0.146
qwen3.5-2B W4A4	local_hessian	58.68	—	16.620	—
qwen3.5-2B W4A4	wmse	58.83	+0.15	16.569	-0.051
qwen3.5-2B W4A16	max	61.23	-0.21	15.849	+0.403
qwen3.5-2B W4A16	mse	61.03	-0.41	15.530	+0.085
qwen3.5-2B W4A16	local_hessian	61.44	—	15.445	—
qwen3.5-2B W4A16	wmse	61.22	-0.22	15.383	-0.062
qwen3.5-4B W4A4	max	69.14	-2.41	12.483	+0.313
qwen3.5-4B W4A4	mse	71.25	-0.29	12.237	+0.067
qwen3.5-4B W4A4	local_hessian	71.54	—	12.169	—
qwen3.5-4B W4A4	wmse	71.03	-0.52	12.138	-0.031
qwen3.5-4B W4A16	max	70.84	-1.92	11.792	+0.238
qwen3.5-4B W4A16	mse	72.76	+0.00	11.609	+0.055
qwen3.5-4B W4A16	local_hessian	72.76	—	11.554	—
qwen3.5-4B W4A16	wmse	72.71	-0.05	11.509	-0.045
qwen3.5-9B W4A4	max	73.41	-0.28	10.082	+0.185
qwen3.5-9B W4A4	mse	73.76	+0.07	9.975	+0.078
qwen3.5-9B W4A4	local_hessian	73.69	—	9.897	—
qwen3.5-9B W4A4	wmse	73.47	-0.21	9.890	-0.007
qwen3.5-9B W4A16	max	74.29	-0.55	9.697	+0.165
qwen3.5-9B W4A16	mse	75.03	+0.18	9.612	+0.080
qwen3.5-9B W4A16	local_hessian	74.85	—	9.532	—
qwen3.5-9B W4A16	wmse	74.63	-0.22	9.514	-0.018

Reading:

• WMSE is statistically indistinguishable from local_hessian on accuracy —
  every cell is within +-0.52 points, 6 of 10 up, 4 of 10 down, mean delta
  -0.02 points.
• WMSE has the lowest wikitext perplexity of the four methods in 8 of 10 cells,
  and beats local_hessian in 8 of 10 (mean -0.05).
• Both beat max on perplexity in every cell, and on accuracy in 9 of 10.
  Against mse, WMSE has lower perplexity in 8 of 10; the two exceptions are
  both Qwen3-8B cells, where mse edges it out.
• So the diagonal approximation loses essentially nothing versus the full block
  Hessian, at 1/16 the accumulator memory and no matmul in the sweep.

Tests

• 16 new unit tests in tests/unit/torch/quantization/test_wmse.py.
• test_local_hessian.py: no regression (22 tests).
• Full unit suite: 862 passed, 7 skipped.
• GPU kernel suite test_nvfp4_fp8_sweep_kernel.py: 54 passed, including new
  DIAGONAL specialization cases.
• pre-commit clean.

🤖 Generated with Claude Code

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