Repository navigation
perf(quantizer): vectorize MSE shrink grid and AWQ clip grid evaluation - #3435
Merged
Merged
Conversation
Batch the shrink-grid candidate evaluation in Quantizer.find_params (also shared by qqq.Quantizer) and the AWQ _compute_best_clip grid so the per-candidate kernel-launch storm collapses into a few large ops per chunk. Candidate order, objective values, and first-minimum tie break are identical to the sequential loops, so selected scale/zero/max_val are bitwise-equal; verified 1728/1728 bitwise-equal across bit widths, sym/asym, group sizes, trits, adjacent-zero search, QQQ, and clip paths. Workspace per call is capped by GPTQMODEL_MSE_GRID_BUDGET_MB and GPTQMODEL_CLIP_GRID_BUDGET_MB (default 256 MiB).
The 256 MiB default that amortizes kernel launches on GPU only evicts the L3-resident working set on CPU, turning a 1.3-2x win into a ~0.7x regression at mid shapes. CPU defaults to 16 MiB (GPU keeps 256 MiB); both still honor the *_BUDGET_MB env overrides. Selection semantics are chunk-size independent, so results remain bitwise-identical.
Contributor
|
I'll fix CI failures and address comments from users with write access that start with 'Devin'.
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
8 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Vectorize the MSE shrink grid in
Quantizer.find_params(GPTQ + QQQ) and theAWQ clip grid in
AWQProcessor._compute_best_clip. Each grid step used tolaunch its own clamp/quantize/error/where chain (~a dozen small kernels per
candidate, ~100-200 candidates per call); the grids now evaluate in large
chunked batches, so launch overhead collapses into a few ops per chunk.
The implementation preserves candidate order and first-minimum selection,
including chunk boundaries, adjacent-zero search (2-bit), and the trits
(x > s/2)*s + (x < z/2)*z -> subtract -> abs -> powchain.Exact CPU oracle comparisons pass for the tested inputs; GPU numerical
equivalence and performance have not been validated on this host.
Follow-up review found and fixed premature FP16/BF16 AWQ factor rounding,
an unchunked AWQ reference-product allocation, and MSE reduction-order drift
for sliced transpose views. The new regression cases detect all three bugs
in the pre-review implementation.
What Changed
quantizer.py: new module-level_mse_shrink_grid— chunked evaluator overthe full
(candidates, rows, cols)grid.Quantizer.find_paramsbuilds thecandidate
p/xmin/xmax/scale/zerotensors once (host-float1 - i/gridrounded to
xmin's dtype, identical to the loop) and delegates; the 2-bitadjacent-zero search interleaves candidates in the same order.
qqq.py:Quantizer.find_paramsreuses the same_mse_shrink_grid(guard_zero_scale=False, matching QQQ's unguarded
quantize).awq_processor.py:_compute_best_clipevaluates the shrink-factor grid incandidate chunks with preallocated clamp/quant/out/product scratch buffers;
org_max_valstays the baseline candidate, andargminover stacked errorsreplaces the per-step
whereupdates.GPTQMODEL_MSE_GRID_BUDGET_MBandGPTQMODEL_CLIP_GRID_BUDGET_MBdefault to 256 MiB on GPU and 16 MiB on CPU —on CPU a large transient only evicts the L3-resident working set, so a small
chunk wins there. These are not total allocation caps: AWQ sizes scratch
buffers independently, and the minimum one-candidate/full-token-error or
one-token-product buffers can exceed the target. Documented in the README.
the multiplied bounds, just like the scalar loop. The unclipped reference
product now reuses token-tiled scratch rather than allocating all tokens.
xis a transpose view, including strided slices,the sequential code's
quantizeoutput inherited column-major order andtorch.sumaccumulated in that order; the workspace is allocated in the samephysical order so the reduction walks memory identically (a 1-ULP-different
sum could flip strict-
<ties).Tests
Independent sequential Torch oracles cover MSE, QQQ, and AWQ candidate
selection. New cases cover FP16/BF16/FP32 AWQ bounds (symmetric/asymmetric,
grouped/ungrouped, one-candidate and batched execution), mixed dtypes, empty
grids, ties, overflow, allocation shapes, strided MSE objectives, rounding
thresholds, and QQQ low-precision chunk boundaries. Floating-point assertions
use
rtol=atol=0; no relaxed tolerances hide changed selection.The formerly unused
_DEVICESlist now drives actual CPU/CUDA MSE testclasses. CUDA tests explicitly skip when unavailable.
$ OMP_NUM_THREADS=4 python -m pytest tests/test_find_params_mse_grid.py \ tests/test_qqq_quantizer.py tests/test_awq_clip_consistency.py \ -q --tb=short -o log_cli=false -o addopts='' # 134 passed, 72 skipped (CUDA unavailable) $ ruff check --config format/ruff.toml gptqmodel/quantization/quantizer.py \ gptqmodel/quantization/qqq.py gptqmodel/looper/awq_processor.py \ tests/test_find_params_mse_grid.py tests/test_awq_clip_consistency.py \ tests/test_qqq_quantizer.py # 2 pre-existing findings in awq_processor.py; no new findings. # All changed tests and both quantizer modules are clean.Regression sensitivity: running the new targeted tests with the pre-review
implementations produced 19 failures (16 low-precision AWQ cases, the
reference-product allocation test, and two strided MSE objective cases).
Adding
tests/test_qqq.pyproduced 137 passed, 73 skipped, and 3 setup errors:its
TestGroupSizeintegration fixture/monster/data/model/Llama-3.2-1Bis absent. Those integration cases were not validated.
Historical validation (before follow-up review)
The original out-of-tree harness replayed the sequential loops on the same
inputs: 1728/1728 checks bitwise-equal covering bits 2/3/4/8, sym/asym,
groupsizes -1/16/32/128,
requires_groupwise_processing, trits, 2-bitadjacent-zero search, QQQ (sym/asym × groupsize), clip sym/asym × groupsize,
activation 3-D/4-D, pathological inputs (outliers, zeros, constant channels),
and forced single-candidate chunking. Its coverage missed the FP16/BF16 AWQ
bug found during review; this is not evidence of universal equivalence.
End-to-end: quantized
hf-internal-testing/tiny-random-LlamaForCausalLM(4-bit, group_size=-1, mse=2.4, 8 calibration batches) with stock vs
vectorized
find_params— identical state-dict digest (SHA-256 over allnon-meta tensors).
The original end-to-end digest check has not been rerun after the review fixes.
Historical benchmarks (before follow-up review)
CPU (torch 2.14.1, 8 threads, median of 5 runs per call). These measurements
predate the review fixes; no GPU performance measurements are available.
Review Requirements
Kernel Accuracy Requirements
Current independent CPU oracles check exact selected
scale/zero/max_valvalues and strided MSE objectives, including low-precision arithmetic.
No inference-kernel weight layouts changed; quantized persistent storage is
untouched. GPU numerical validation remains outstanding.
Notes
A green static-analysis CI run does not substitute for GPU numerical tests.
Link to Devin session: https://app.devin.ai/sessions/2eb9e33e49b54734a60b2552d71b0aca
Open in Devin Desktop: https://app.devin.ai/desktop/session/2eb9e33e49b54734a60b2552d71b0aca?variant=devin
Requested by: @Qubitium