[4/4] Run the local-Hessian choice search as one CUDA kernel by cjluo-nv · Pull Request #3218 · NVIDIA/Model-Optimizer

cjluo-nv · 2026-11-04T19:28:50Z

What does this PR do?

Type of change: new feature (performance)

[4/4] of the GGML local-Hessian series. It depends on [3/4] #3209. Merge order: #3158 → #3159 → #3209 → #3218.

After the one-pass encode, search_choices scores every choice of every group under the block Hessian and refines the groups in two sweeps. In PyTorch that is a few small kernels per group and sweep, for every 256-column block: local_hessian searches every block of every weight, and GPTQ with local_hessian_scales searches one block at a time. On MoE models, with thousands of small expert weights, launch overhead dominated calibration.

choice_search.cu runs the whole search for one block in a single launch:

Layout: one CUDA block per row, one thread per column.
Shared memory: the candidate errors, current errors and gradient.
The Hessian: read from L2 by column. It is symmetric by construction (U_gg⁻¹ U_gg⁻ᵀ).
Reductions: per-group sums are warp-segment shuffles.

Its picks match the PyTorch search, which remains the CPU path.

Per 256-column block on a GB300:

Rows	Format	[2/4]	this PR
512	IQ2_XXS	4.8 ms	0.8 ms
512	IQ2_S	9.1 ms	0.8 ms
512	IQ1_S	4.4 ms	0.7 ms
4096	IQ2_XXS	4.8 ms	3.0 ms
4096	IQ2_S	9.0 ms	3.1 ms
4096	IQ1_S	5.3 ms	4.0 ms

At 4,096 rows the one-pass encode is now the larger part: about 2 ms, against about 1 ms for the search kernel.

End to end, GPTQ with local_hessian_scales relative to plain GPTQ:

Workload	[1/4] search	[2/4]	this PR
16 expert-sized 1536×4096 weights	7.0×	2.9×	1.6×

On the real models (4× GB300, 1,024 samples), calibration with local_hessian_scales:

Model	plain GPTQ	[2/4]	this PR
GLM-5.3-Flash	25.4 min	108 min	28.1 min
Qwen3.8-Flash-Next	9.8 min	50 min	15.5 min

KL against BF16 is unchanged within noise: GLM 0.625 vs 0.624, Flash-Next 0.249 vs 0.247.

local_hessian on GLM-5.3-Flash, whose earlier PyTorch search already batched every block of a weight together, also gets faster: encode 9.1 min vs 11.1 min, calibration 18.6 min vs 20.8 min. KL is unchanged at 0.720.

Memory. The kernel scores the candidate errors in place, without [3/4]'s einsum temporaries. So search_scratch, which _update_room adds for an update GPTQ may move to another GPU, now counts only the candidate errors and payloads. Peak memory of one IQ2_XXS update moved to a second GB200, with the option on:

Weight `[out, in]`	[3/4]	this PR (room)
4096×2048	329 MiB	199 MiB (360)
16384×2048	1,123 MiB	604 MiB (865)
17408×5120	1,675 MiB	1,125 MiB (2,089)
65536×256	3,431 MiB	1,355 MiB (1,636)

Testing

tests/gpu/torch/quantization/test_ggml_gptq_cuda.py::test_fused_choice_search_matches_the_pytorch_search: on a compensated GPTQ block Hessian, the fused picks equal the PyTorch search's on all five formats. It fails when the kernel skips its refinement sweeps.
Fused and PyTorch picks are also identical on 8,192 rows per format over four random 256-column blocks.
test_gptq.py::test_an_update_that_may_move_fits_its_room, parametrized by local_hessian_scales in [3/4], passes with the lowered estimate.
I ran the GGML unit and GPU suites, test_autoquant.py and the GPU test_gptq.py on this branch on GB200 (oci-jhb) in the calibration container: 46 + 314 + 96 passed.

Before your PR is "Ready for review"

Is this change backward compatible?: ✅ (internal speedup; picks unchanged)
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 (speedup of 1/4 and 3/4, same release)
Did you get Claude approval on this PR?: ❌ (pending)

Additional Information

Series (#3157, fixed per-group choices, is merged): [1/4] #3158 local_hessian for GGML IQ · [2/4] #3159 one-pass every-choice encode · [3/4] #3209 GPTQ local-Hessian scale picks · [4/4] this PR.

🤖 Generated with Claude Code

coderabbitai · 2026-11-04T19:28:58Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review: 3

main
release/.*
feature/.*

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

Source mode: Auto
Repository: NVIDIA/Model-Optimizer
Branch: chenjiel/ggml-choice-search-kernel
File: .coderabbit.yaml
Commit: e7c6633
Layers: Repository YAML
Review profile: CHILL
Plan: Enterprise
Run ID: 779b62b6-4061-4aa1-a5f5-77f116e92bba

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

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

### What does this PR do? Type of change: new feature **[1/5]** of the GGML local-Hessian series. Merge order: #3157 → #3158 → #3159 → #3209 → #3218. Each GGML IQ format (IQ2_XXS, IQ2_XS, IQ2_S, IQ1_S, IQ1_M) picks one choice per group of every 256-value block: a local scale, plus the delta sign for IQ1_S. It picks the one with the least (importance-weighted) squared error. `quantize_<format>(..., choices=...)` now takes those picks as input, in both the CUDA kernels and the PyTorch encoders. The codebook entries and signs are still searched under the given choice. This lets a caller pick the scales by another metric. [2/5] uses it to pick them by the local Hessian. When `choices` is given, the kernels skip the scoring sweep and keep the entry search, so the added CUDA code is small (about 13 lines per encoder family). Each `GGMLFormat` now records `choice_groups` and `num_choices`. ### Usage ```python from modelopt.torch.quantization.ggml.registry import GGML_FORMAT_REGISTRY fmt = GGML_FORMAT_REGISTRY["iq2_xs"] choices = torch.full((*weight.shape[:-1], weight.shape[-1] // 256, fmt.choice_groups), 7) packed, shape = fmt.quantize(weight, choices=choices) # every group takes local scale 7 ``` ### Testing - `tests/unit/torch/quantization/test_iq_formats.py`: a fixed choice reproduces, group by group, the payload encoded with that choice everywhere. - `tests/gpu/torch/quantization/test_iq_formats_cuda.py`: the CUDA encoders match the PyTorch encoders byte for byte under random fixed choices. - I ran `test_iq_formats.py` and `test_iq_formats_cuda.py` locally on one GPU (199 passed). I ran the GGML unit and GPU suites (imatrix, GPTQ, backend, IQ formats) on GB200 in the calibration container. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (new keyword argument, default unchanged) - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (enabling API; the user-facing entries are in 2/5 and 4/5) - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information Series: [1/5] this PR · [2/5] #3158 local_hessian for GGML IQ · [3/5] #3159 one-pass every-choice encode · [4/5] #3209 GPTQ local-Hessian scale picks · [5/5] #3218 fused choice-search kernel. 🤖 Generated with [Claude Code](https://claude.com/claude-code)   ## Summary by CodeRabbit * **New Features** * IQ1_S, IQ1_M, IQ2_S, IQ2_XS, and IQ2_XXS quantization now accept optional per-group choices to control scale or encoding selection. * When choices are omitted, quantization continues to select choices automatically as before. * Supplied choices are checked for compatibility with the weights and supported choice range, and are used by both CUDA and PyTorch quantization paths.  Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

…ix (#3158) ### What does this PR do? Type of change: new feature **[1/4]** of the GGML local-Hessian series. It builds on #3157, which is merged. Merge order: #3158 → #3159 → #3209 → #3218. Before this PR, `local_hessian` left GGML weights unchanged, because they have no amax to search. It now encodes each GGML IQ weight once: - **One pass, no layerwise calibration:** like imatrix, so it fits flows that need a single-pass algorithm, such as AutoQuantize. - **Codebook search:** keeps imatrix's weighting, the square root of the Hessian diagonal. - **Choices:** each group's local scale, plus the delta sign for IQ1_S, is picked by the layer-output error `dw H dw^T` over the 256-column diagonal block of the input Hessian. Before, it was picked by the importance-weighted squared error. - **Damping:** the Hessian blocks are damped by the new `perc_damp` (default 0.3), as GPTQ damps. Undamped, the picks overfit the calibration set and did worse end to end. - **Pinning:** the codes are pinned as imatrix pins them. A weight that already holds codes (GPTQ's) is left alone, so `[gptq, local_hessian]` works like `[gptq, imatrix]`. - **Codes kept until QAD changes a block:** under #3197 a pinned IQ weight is re-encoded under its pinned scale and importance, which would re-pick the searched local scales. So the encoder also pins a per-block fingerprint of the weight. Blocks still equal to it keep the searched codes (PTQ, eval, export), and a block QAD changed is re-encoded as after GPTQ, with only the super-block scale frozen. GPTQ and imatrix pins carry no fingerprint and behave as before. The fingerprint shares the payload's buffer prefix, so checkpointing, Megatron sharding and export treat it as part of the pin. The search lives in a new `choice_search` module: 1. Encode each block under every choice. 2. Give each group its best choice under its own `H_gg` block. 3. Revisit the groups twice, each taking the choice that most lowers the block's error with the others held. [3/4] reuses this search inside GPTQ, and [2/4] and [4/4] make it fast. Other changes: - An encoder registry for `local_hessian`, mirroring imatrix's. - The imatrix target selection and per-GPU encode loop are factored out of `imatrix_calibrate` and shared by both algorithms. - Like imatrix, `local_hessian` now skips the max-calibration forward when no activation quantizer needs one. **Results.** Three models with the following weights quantized, same 1,024 calibration samples per arm: - **Qwen3.8-27B:** the language-model MLP. - **GLM-5.3-Flash:** all routed experts. - **Qwen3.8-Flash-Next:** the routed-expert `gate_up`. SciCode subtask accuracy (%), 8 runs per arm, INT8 fake-quant serving (BF16 for Flash-Next), 64K cap; local_hessian minus imatrix: | Model | IQ2_XXS | IQ2_XS | IQ2_S | |---|---|---|---| | Qwen3.8-27B | 9.5 → 16.0 (**+6.5 ± 0.8**) | 20.6 → 26.5 (**+5.9 ± 1.0**) | 32.0 → 34.0 (+2.0 ± 1.1) | | GLM-5.3-Flash | 3.9 → 9.9 (**+6.0 ± 0.5**) | 12.5 → 18.3 (**+5.8 ± 0.5**) | 31.2 → 30.7 (−0.5 ± 0.6) | | Qwen3.8-Flash-Next | 38.4 → 37.2 (−1.2 ± 0.9) | 40.2 → 39.1 (−1.1 ± 1.1) | 41.9 → 42.7 (+0.8 ± 0.8) | - **Never significantly worse than imatrix.** The gains come from fewer runaway (repetition) generations at the two lowest bit widths, which leave fewer subtasks unanswered. Flash-Next never loops. - **KL did not predict the outcome:** wikitext-2 KL against BF16 was lower on both MoE models at IQ2_XS/IQ2_S, but unchanged on 27B, where SciCode gained most. - **Truncation:** 10–51% of generations per arm hit the 64K cap, above the 2% validation gate. Read these as relative comparisons under identical settings. ### Usage ```yaml algorithm: - {method: gptq} - {method: local_hessian} # in place of {method: imatrix}; perc_damp defaults to 0.3 ``` ### Testing - `tests/unit/torch/quantization/test_ggml_imatrix.py`: - The imatrix contract tests are parametrized over both algorithms: pins codes, leaves GPTQ codes alone, refuses a bad width before any state changes, encodes a never-reached weight plainly, runs one forward. - `test_local_hessian_lowers_the_output_error_of_correlated_inputs`: the local-Hessian pick lowers the output error on correlated inputs. - `test_search_under_the_identity_hessian_keeps_the_encoders_own_choices`: under `H = I`, the limit heavy damping tends to, the search reproduces the encoder's own picks. - `test_local_hessian_codes_hold_until_qad_changes_a_block`: the calibrated weight packs to the searched codes, and a row QAD changed is re-encoded exactly as a GPTQ pin is under #3197, its local scales searched again under the pinned super-block scale. It fails without the fingerprint. - I ran the GGML unit and GPU suites on this commit alone on GB300 in the calibration container. I ran the PTQ, KL and SciCode campaign above on oci-jhb. - Not run: multi-rank calibration. `local_hessian` already warns that it is single-rank for now. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (GGML weights were previously untouched by `local_hessian`. NVFP4 and other amax paths are unchanged, except that the max-calibration forward is skipped when no activation quantizer needs it.) - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information Series (#3157, fixed per-group choices, is merged): [1/4] this PR · [2/4] #3159 one-pass every-choice encode · [3/4] #3209 GPTQ local-Hessian scale picks · [4/4] #3218 fused choice-search kernel. 🤖 Generated with [Claude Code](https://claude.com/claude-code)   ## Summary by CodeRabbit * **New Features** * Added local-Hessian calibration for GGML IQ weights, selecting local scales and IQ1_S deltas based on layer-output error. It can run independently or alongside GPTQ. * Added the `perc_damp` setting, defaulting to `0.3`, to control Hessian damping during IQ-weight scale selection. * Calibration preserves existing codes for unchanged weight blocks and re-encodes blocks that have changed. * **Bug Fixes** * Calibration skips an unnecessary forward pass when activation calibration is not needed.  --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

…earch (#3159) ### What does this PR do? Type of change: new feature (performance) **[2/4]** of the GGML local-Hessian series. It depends on [1/4] #3158. Merge order: #3158 → #3159 → #3209 → #3218. The local-Hessian choice search ([1/4]; GPTQ reuses it in [3/4]) needs each 256-column block encoded under every choice. [1/4] runs one forced encode per choice, sweeping the codebook 8–16 times. The IQ encoders now take `every_choice=True` and return all of those payloads at once, shaped `[num_choices, *packed shape]`: - **CUDA kernels:** a new kernel per family keeps each choice's best codebook entry from a single sweep, using a slot-wise block key minimum. - **PyTorch encoders:** these loop over forced choices. The payloads are byte-identical to the forced encodes, so the picks do not change. `search_choices` decodes them in place into its error buffer. **Per 256-column block on a GB300:** | Rows | Format | Search, [1/4] | Search, this PR | |---|---|---|---| | 512 | IQ2_XXS | 13.9 ms | 4.8 ms | | 512 | IQ2_S | 17.7 ms | 9.1 ms | | 512 | IQ1_S | 10.6 ms | 4.4 ms | | 4096 | IQ2_XXS | 17.3 ms | 4.8 ms | | 4096 | IQ2_S | 22.4 ms | 9.0 ms | | 4096 | IQ1_S | 31.9 ms | 5.3 ms | What remains is mostly the scoring and refinement, a few small kernels per group and sweep, which [4/4] fuses into one launch. ### Usage ```python packed, shape = fmt.quantize(weight, importance=imp, every_choice=True) # packed[c] == fmt.quantize(weight, importance=imp, choices=full(c))[0] for every choice c ``` ### Testing - `tests/gpu/torch/quantization/test_iq_formats_cuda.py`: for every format, the every-choice CUDA payload equals each weighted forced encode byte for byte. - Search picks are identical to [1/4] on all five formats on a 1536×1024 weight. The PyTorch every-choice path is exercised by the CPU `local_hessian` unit tests. - I ran the GGML unit and GPU suites on this commit alone on GB300 in the calibration container. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (new keyword argument, default unchanged) - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (speedup of 1/4, same release) - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information Series (#3157, fixed per-group choices, is merged): [1/4] #3158 local_hessian for GGML IQ · [2/4] this PR · [3/4] #3209 GPTQ local-Hessian scale picks · [4/4] #3218 fused choice-search kernel. 🤖 Generated with [Claude Code](https://claude.com/claude-code)   ## Summary by CodeRabbit * **New Features** * IQ1 and IQ2 quantization can now generate a separate packed result for every available choice, applying each choice consistently across all groups. * The all-choices option is available in both CUDA-accelerated and standard quantization paths. Existing fixed-choice behavior remains; the two options cannot be used together. * **Tests** * Added coverage comparing all-choice results with results generated separately for each choice.  --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

After the one-pass encode, search_choices scores every choice of every group under the block Hessian and refines the groups in two sweeps. In PyTorch that is a few small kernels per group and sweep, so on MoE models, where GPTQ searches one 256-column block of thousands of small expert weights at a time, launch overhead dominated calibration. A CUDA kernel now runs the whole search for one block in a single launch, one CUDA block per row with the candidate errors, current errors and gradient in shared memory. Its picks match the PyTorch search, which remains the CPU path. Per 256-column block on a GB300 the search takes 0.8-4 ms instead of 4-9 ms. On 16 expert-sized 1536x4096 weights, GPTQ with local_hessian_scales takes 1.6x plain GPTQ instead of 2.9x. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

cjluo-nv requested review from a team as code owners November 4, 2026 19:28

cjluo-nv requested review from meenchen and removed request for a team November 4, 2026 19:28

cjluo-nv mentioned this pull request Nov 4, 2026

[1/4] Calibrate GGML IQ formats with local_hessian, in place of imatrix #3158

Merged

This was referenced Nov 4, 2026

[2/4] Encode every GGML IQ choice in one pass for the local-Hessian search #3159

Merged

[3/4] Let GPTQ pick GGML IQ local scales by its compensated block Hessian #3209

Open

[1/5] Let the GGML IQ encoders take fixed per-group choices #3157

Merged

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch from a1d9df6 to 586e194 Compare November 4, 2026 21:31

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from 88d8e24 to 044144f Compare November 4, 2026 21:31

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from 044144f to 9f15ebd Compare November 5, 2026 21:10

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch from 586e194 to b301fb5 Compare November 5, 2026 21:10

cjluo-nv changed the title ~~[5/5] Run the local-Hessian choice search as one CUDA kernel~~ [4/4] Run the local-Hessian choice search as one CUDA kernel Nov 5, 2026

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch from b301fb5 to ab58cbd Compare November 5, 2026 22:11

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from 9f15ebd to 8ad11d7 Compare November 5, 2026 22:11

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch from ab58cbd to 9915a47 Compare November 5, 2026 23:52

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from 8ad11d7 to 547bc45 Compare November 5, 2026 23:52

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch from 9915a47 to 51ff6b5 Compare November 6, 2026 07:08

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from 547bc45 to c0a53e4 Compare November 6, 2026 07:08

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from c0a53e4 to 7b422c3 Compare November 6, 2026 17:02

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch 2 times, most recently from 49edc6b to ab9e1ef Compare November 6, 2026 17:29

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from 7b422c3 to 8663072 Compare November 6, 2026 17:29

cjluo-nv force-pushed the chenjiel/ggml-gptq-local-hessian branch from 8663072 to cb8bad9 Compare November 6, 2026 19:09

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch from ab9e1ef to b5334d2 Compare November 6, 2026 19:09

cjluo-nv and others added 2 commits November 6, 2026 20:51

Count only the candidate errors as the fused search's scratch

e7c6633

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

cjluo-nv force-pushed the chenjiel/ggml-choice-search-kernel branch from b5334d2 to e7c6633 Compare November 7, 2026 02:28

Navigation Menu

Repository navigation

[4/4] Run the local-Hessian choice search as one CUDA kernel - #3218

[4/4] Run the local-Hessian choice search as one CUDA kernel#3218
cjluo-nv wants to merge 2 commits into
chenjiel/ggml-gptq-local-hessianfrom
chenjiel/ggml-choice-search-kernel

cjluo-nv commented Nov 4, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Nov 4, 2026 •

edited

Loading

Review skipped

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

cjluo-nv commented Nov 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Testing

Before your PR is "Ready for review"

Additional Information

Uh oh!

coderabbitai Bot commented Nov 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

cjluo-nv commented Nov 4, 2026 •

edited

Loading

coderabbitai Bot commented Nov 4, 2026 •

edited

Loading