[cuda] Nestest Superscale Quantization for CUDA Weights (int4 / int6)#20571
Conversation
…ow super-scale)
Decode for gemma4-31B on CUDA is weight-bandwidth-bound: the int4/int6
weight-only matvecs are ~72% of per-token decode time and already run at
~89% of the RTX 5090 HBM roofline using the same dp4a algorithm as
llama.cpp. The only lever left is fewer bytes/token.
ET previously stored per-group (group_size=32) scale AND zero as bf16:
int4 = 5.0 bits/weight (20% metadata overhead), int6 = 7.0 bits/weight.
llama.cpp's Q4_K/Q6_K store far less metadata (4.5 / 6.5625 bpw).
This change re-encodes the quant metadata to match llama.cpp's byte
density WITHOUT touching the dp4a inner loop or the 4/6-bit weight codes:
- per-group scale/zero: bf16 -> uint8 codes
- per-row bf16 "super-scale" (step = row_max/255) restores dynamic range
(plain fp8/int8 alone fails: Q4_K scales span ~1e-4..8e-2 -> subnormal
blowup, 15-18 dB SNR; the two-level encoding keeps 46.7-48.1 dB == bf16)
- int6 is symmetric (no zero), int8 scale codes + per-row step (mirrors Q6_K)
Result -> 4.77 bpw (llama.cpp parity). Touches the packer
(coalesced_int4_tensor / dp4a_planar_int6_tensor), the decode shims
(int4/int6_plain_mm .cu/.cuh/.h, +steps arg), the dispatch ops, the AOTI
ABI sig in cuda_backend, and the gemma4_31b concat source transform.
e2e (gemma4-31B, 128k+TurboQuant export, CUDA-graph decode, 3-rep median):
ctx baseline -> this
512 57.16 -> 60.78 (+6.3%)
2048 56.39 -> 59.79 (+6.0%)
8192 55.60 -> 58.88 (+5.9%)
32768 55.29 -> 58.60 (+6.0%)
VRAM peak 25.06 -> 23.29 GiB; .ptd 26.18 -> 24.28 GB. Accuracy: Paris
coherent, dequant SNR ~baseline. Win holds at long context under CUDA
graph (verified, not a microbenchmark estimate).
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
The compact-metadata change (bf16 -> uint8/int8 codes + per-row bf16 super-scale) migrated the int4/int6 plain_mm kernels and Python dispatch but left the tests on the old signatures, breaking the CUDA shim CMake build and the dispatch unit tests. - test_aoti_torch_cuda_int4_plain_mm.cpp: regenerate vectors as uint8 scale/zero codes + [N,2] bf16 steps; update to the 7-arg signature. - test_aoti_torch_cuda_int6_plain_mm.cpp: regenerate vectors as int8 scale codes + [N,1] bf16 steps; update to the 6-arg signature. - test_int4_dispatch.py / test_int6_dispatch.py: add the steps arg to the custom-op recorders, build tensors via the new constructors, and compare against the decoded (code * step) scale instead of raw codes. Vectors are generated from the production pack path (CudaCoalescedInt4Tensor / pack_int6 + _encode_int8_per_row) with expected[] from the export-path _dequant_matmul references. Test Plan: - Built the CUDA shim gtests and ran on GPU: int4 7/7, int6 3/3 pass. - python -m pytest backends/cuda/tests/test_int4_dispatch.py backends/cuda/tests/test_int6_dispatch.py -> 33 passed. - lintrunner init && lintrunner -a: no lint issues.
Migrate the CUDA int4 (Q4_K) path from a per-row bf16 scale-step to a
per-256-super-block fp16 scale-step + uint8 scale code, mirroring GGUF Q4_K's
per-super-block scale granularity. Improves whole-weight dequant SNR to
45.15 dB (vs 44.37 dB), with NEW >= OLD generation quality.
The naive per-group load of the separate scale_step tensor cost ~8% decode
(latency-bound W4A8 matvec). Fixed with a llama.cpp-style super-block-cooperative
kernel: 8-lane warp subgroups each cover one super-block per iteration; only the
subgroup leader loads+converts the fp16 scale_step and __shfl-broadcasts it to
its 7 followers (8x fewer loads, fp16 convert off the dp4a path), register-only
(no smem, no occupancy cliff). dp4a inner product + zero decode unchanged.
Same-box decode vs the prior encoding: +1.7%..+2.9% (512/8K, bf16+tq4), VRAM
neutral. Op-level: packer SNR 45.15 dB, int4 gtest 5/5, dispatch 20/20, int6 3/3.
Also adds aoti_torch_dtype_float16 (ScalarType::Half) to the AOTI shims, required
by the fp16 scale_step tensor.
Files: kernel (int4_plain_mm.{cu,cuh,h}), packer (coalesced_int4_tensor.py),
dispatch (int4_dispatch.py), ABI (cuda_backend.py), fp16 shim
(common_shims{,_slim}.{cpp,h}), tests (int4 gtest + test_int4_dispatch.py),
gate/up fusion consumer (cuda_source_transformations.py).
…code
Migrate the CUDA int6 (Q6_K) path from a per-row bf16 scale-step to a
per-256-super-block fp16 scale-step + int8 scale code, mirroring GGUF Q6_K's
per-super-block scale granularity. Improves whole-weight dequant SNR to
~49.2 dB (vs ~45.7-46.6 dB for the per-row step), measured against the exact
GGUF Q6_K decode reference (extension/llm/export/gguf.py) across 3 shapes.
NEW >= OLD everywhere (+3.3 dB). Weight layout (planar ql/qh), int8 code width,
and Q6_K symmetry (no zero) are UNCHANGED.
Kernel: keep decode perf-neutral by amortizing metadata per uint4 rather than
the T3 warp-shuffle used for int4. The int4 shuffle only pays off because
int4's OLD baseline had a distant per-group step load to amortize; int6's OLD
step was already a per-row loop-invariant, so a shuffle recovers nothing (an
int6 T3 variant measured ~2-7% SLOWER, tied by a plain per-super-block __ldg).
The actual lever is the SAME code-load hoist int4's commit landed: load the
per-group int8 codes once per group per uint4 (gs=16 => 2 codes/uint4, held in
registers across the 4 dp4a words) instead of 4x/uint4 per word, plus a single
per-uint4 __ldg of the fp16 step (8 lanes share each tiny L1-resident address).
SASS confirms LDG.E.S8 drops 4 -> 2 per uint4. dp4a inner product + planar
weight decode unchanged; register-only (no smem, no occupancy cliff).
Op-level decode vs the prior per-row encoding (A100, M=1, same-box, 3 reps):
avg -3.0% (faster in 4/5 of {4096,8192,11008,14336}x{4096,8192,11008}, worst
+1.0%) -- neutral-to-better. The hoist's extra ILP raises the matvec from
REG:40 to REG:48 (theoretical occupancy 75% -> 62% on sm_80) with STACK:0 (zero
spills); measured wall-clock is faster regardless (ILP > occupancy on this
latency-bound matvec), and __launch_bounds__(256,6) to reclaim REG:40 caused
spills, so REG:48 is kept intentionally.
Gates: packer round-trip SNR +3.3 dB vs exact GGUF Q6_K; int6 gtest 4/4
(single/multi super-block step indexing + wide-N warp reduction + null args);
python int6 dispatch 20/20-style green (13 tests). ABI: steps arg is now the
per-256 fp16 tensor [N, K/256] (was per-row bf16 [N, 1]); reuses the float16
AOTI shim added for int4.
Files: kernel (int6_plain_mm.{cuh,h}), packer (dp4a_planar_int6_tensor.py),
dispatch (int6_dispatch.py), tests (int6 gtest + test_int6_dispatch.py).
…+0.74dB SNR (45.89 vs 45.15), op-level decode neutral (-0.44%), +0.06 bpw (4.625) Promote the int4 zero step from per-row bf16 to per-256-super-block fp16 (the SNR-max zero256 encoding, B2), decoded via z_pack: the 8-lane subgroup leader packs BOTH fp16 steps (scale low16, zero high16) into the ONE existing 32-bit warp-shuffle word and broadcasts it once — no extra shuffle vs the committed scale-only baseline (the naive second __shfl / z_2shuf regressed +7.4%). - packer (coalesced_int4_tensor.py): zero_step now [N,K/256] fp16 via _encode_uint8_per_super (same per-256 grid as scale); dequantize + from_int4 updated. bpw 4.5625 -> 4.625. - decode (int4_plain_mm.cuh): per-256 fp16 zero_step packed with scale_step into one shuffle word (z_pack); dp4a inner loop unchanged. cuobjdump: REG:32 STACK:0 SHARED:0 == committed (no occupancy change). - ABI (int4_plain_mm.cu/.h): zero_step dtype Half, shape [N,K/256]. C-shim arg count unchanged (AtenTensorHandle), so cuda_backend.py needs no change. - dispatch (int4_dispatch.py): prefill _dequant_matmul zero uses per-256 step. - tests: int4 gtest regenerated for fp16 [N,K/256] zero_step + new PackedShuffleMultiSuper case (K=1024/n_super=4, packed-shuffle-word correctness); test_int4_dispatch.py per-256 zero decode. Gates (same A100 box): whole-weight SNR 45.89 dB (calibration: committed 45.15); int4 gtest 6/6 + python dispatch 20/20 + int6 gtest 4/4 (untouched) green; op-level decode z_pack -0.44% time-weighted (3 solo reps, within ~1-2% bar). int6 path untouched. e2e export intentionally deferred (running separately).
Gasoonjia
changed the title
|
@claude review this PR |
|
Claude finished @digantdesai's task in 7m 15s —— View job PR Review: Superscale Quantization for CUDA Weights (int4 / int6)Reviewed the full diff against Correctness
Things done well
Test coverageSolid — CIThe listed failures ( None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge. |
|
|
||
| _lib.define( | ||
| "int4_plain_mm(Tensor self, Tensor qdata, Tensor scale, Tensor zero, int group_size) -> Tensor" | ||
| "int4_plain_mm(Tensor self, Tensor qdata, Tensor scale, Tensor scale_step, Tensor zero, Tensor zero_step, int group_size) -> Tensor" |
There was a problem hiding this comment.
what about non k-quant users?
| TEST_F(AOTITorchInt4PlainMMTest, SingleSuperBlock) { | ||
| int64_t M = 1, K = 256, N = 8, gs = 32; | ||
| int64_t ng = K / gs; // 8 | ||
| int64_t n_super = K / 256; // 1 | ||
| // clang-format off | ||
| uint8_t qdata_host[] = { |
There was a problem hiding this comment.
check in the script to recreate these
| "scale", | ||
| "scale_step", | ||
| "zero_point", | ||
| "zero_step", |
There was a problem hiding this comment.
s/zero_step/zero_point_step - Nit
| scale_step: torch.Tensor, | ||
| zero_point: torch.Tensor, | ||
| zero_step: torch.Tensor, |
There was a problem hiding this comment.
Is this Tensor only for k-quants?
digantdesai
left a comment
digantdesai
left a comment
There was a problem hiding this comment.
Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.
|
Thanks @digantdesai for review. Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion. |
What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no? |
digantdesai
left a comment
digantdesai
left a comment
There was a problem hiding this comment.
stamping it to unblock you
…oncat form for pre-PEP701 black)
for kernel-base test perf is even better cuz now bpw is lower. |
2 participants
Quant schema update
Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine
scale/zero) instead of storing it in bf16. Dequant stays affine:w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout:group_size = 32(int4) /16(int6), and a 256-weight super-block that shares one fp16 metadata step.group_size=32, super-block=256 (8 groups).scale= per-group uint8 code x per-256 fp16 step;zero= per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.group_size=16, super-block=256 (16 groups).scale= per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.Why this design
(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.
(b) int6 keeps its int8 scale code. int6 is symmetric, so
scaleis its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.Ablation study
int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).
*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).
Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-
dgranularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.e2e results
gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.
A100, turboquant OFF / bf16 KV
Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.
Prefill (tok/s, higher better)
Decode (tok/s, higher better)
Peak VRAM (GB, lower better)
A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.
RTX 5090, turboquant ON / tq4 KV cache
Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.
Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)
5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)