[cuda backend] int5 quantization support#20757
Draft
Gasoonjia wants to merge 4 commits into
Draft
Conversation
Add a real 5-bit weight path to ExecuTorch, mirroring the existing INT4
(asymmetric zero) and INT6 (bit-plane split) CUDA paths and reusing the
branch's compact int8-encoded scale/zero metadata scheme (uint8 codes +
per-row bf16 super-scale). A GGUF Q5_K weight now loads as a genuine
CudaDp4aPlanarInt5Tensor via the shared CUDA packer, with no lossy
Q5_K->bf16->INT4 round-trip.
- extension/llm/export/gguf.py: Q5_K 176-byte block decode (_q5_k_fields)
+ to_intx_unpacked_to_int8_tensor Q5_K branch (center [0,31]->[-16,15],
fold affine min into zero-point, target_dtype=torch.int5).
- backends/cuda/dp4a_planar_int5_tensor.py: CudaDp4aPlanarInt5Tensor
(ql K/2 low-nibble + qh K/8 high-1-bit planes, asymmetric zero; uint8
scale/zero codes + steps[N,2] bf16 super-scale) + pack/unpack.
- backends/cuda/runtime/shims/int5_plain_mm.{cu,cuh,h}: W5A8 dp4a matvec
kernel + AOTI C shim (aoti_torch_cuda_int5_plain_mm).
- backends/cuda/quantize_op_dispatch/int5_dispatch.py + __init__.py:
int5_plain_mm custom op + F.linear dispatch (decode M<=4 -> custom op,
prefill -> inline dequant).
- backends/cuda/cuda_backend.py: int5_plain_mm AOTI ABI registration.
- backends/cuda/CMakeLists.txt: build int5_plain_mm.cu.
- examples/models/gemma4_31b/quant/pack_cuda.py: route Q5_K
ExportableGGUFTensor -> CudaDp4aPlanarInt5Tensor.
- INT5_NOTES.md: design + validation summary.
Validation (et1, GPU 0): Q5_K decode BIT-EXACT vs gguf.dequantize
(max_abs=0 over 1.34B elems); packed int5 dequant SNR 46.92 dB @ 5.505
bpw; pack/unpack round-trip bit-exact; eager F.linear dispatch maxrel=0
(M=1/4/16); standalone on-GPU kernel numeric test PASS (max_abs=0.25,
tol 0.5); aoti_cuda_shims cmake build RC=0.
…int4): +1.52dB SNR, op-level tests Migrate the INT5 (GGUF Q5_K) metadata encoding from the old compact per-row bf16 super-scale (uint8 code + steps[N,2] bf16) to the per-256-super-block fp16 z_pack scheme already used by INT4 (D ef56bd9). Q5_K is asymmetric gs=32 (super-block=256=8 groups), structurally identical to INT4, so this mirrors the INT4 z_pack packer + kernel exactly: - dp4a_planar_int5_tensor.py: scale/zero each -> per-group uint8 code + per-256 fp16 step (scale_step / zero_step, [N, K/256]); add _encode_uint8_per_super; tensor_data_names now ql/qh/scale/scale_step/zero_point/zero_step; dequantize reconstructs code*step[g//8]. 5-bit ql/qh WEIGHT layout unchanged. - int5_plain_mm.cuh: kernel decodes both steps via a single 8-lane subgroup leader that PACKS both fp16 steps into ONE 32-bit warp-shuffle word (z_pack) and broadcasts with ONE __shfl_sync -- no extra shuffle vs the per-row baseline; warp-aligned trip count; register-only (40 regs, 0 spills). Adds SUPER_BLOCK_I5 + K%256 check + n_super arg. - int5_plain_mm.{cu,h}: ABI self,ql,qh,scale,scale_step,zero,zero_step,gs,ret (fp16 steps). cuda_backend.py: int5 C-shim now 7 AtenTensorHandle. - int5_dispatch.py: op schema 8-arg; per-256 broadcast in _dequant_matmul_int5. Op-level validation (GPU, et1): - Packer SNR (NEW per-256 vs OLD per-row, ref = exact Q5_K affine dequant): +1.52 dB mean (48.65 -> 50.17 dB), new >= old on every shape (like int4 +0.74, int6 +3.3 dB). - C++ gtest test_aoti_torch_cuda_int5_plain_mm.cpp (new; covers z_pack packed shuffle word, single/multi super-block, wide N, 4-super-block warp broadcast): 5/5 PASS. - test_int5_dispatch.py (new): 16/16 PASS. - cuobjdump: int5_w5a8_matvec_kernel 40 reg / 0 spill / 0 smem, full occupancy.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20757
Note: Links to docs will display an error until the docs builds have been completed. ❌ 22 New Failures, 1 Cancelled Job, 1 Unrelated FailureAs of commit b02be86 with merge base d54a0c0 ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
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
1 participant
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
Adds a genuine INT5 (GGUF Q5_K) weight path for the CUDA backend, and brings its metadata quantization in line with the int4/int6 nested-quant scheme (per-256 fp16 "z_pack" superscale).
What INT5 is
ET INT5 is an asymmetric 5-bit affine format (has a zero-point, like Q4_K),
group_size = 32, super-block = 256 weights = 8 groups. Dequant:w = (u - zero) * scale, withuthe 5-bit weight code. The 5-bit weight is stored planar: a 4-bit low planeql[N, K/2]+ a 1-bit high planeqh[N, K/8]. This weight layout is unchanged by the z_pack commit.Metadata scheme
We followed gguf Q5_K super-scale quantization Instead of storing per-group scale/zero in bf16, the metadata is further-quantized to mirror GGUF's Q5_K per-super-block granularity, identical in structure to the int4path (int5 is asymmetric gs=32, structurally the same as int4):
scale= per-group uint8 code x per-256 fp16 stepzero= per-group uint8 code x per-256 fp16 stepop-level experiment
Whole-weight SNR vs the exact GGUF Q5_K affine dequant (synthetic bf16 weights quantized to Q5_K, gs=32 ):
Mean +1.48 dB (per 256 ~50.2 dB vs per row ~48.7 dB); the gain grows with K, consistent with per-256 tracking Q5_K's per-super-block granularity. Same lesson as int4/int6: the accuracy lever is step granularity (per-row -> per-256), and z_pack makes the finer step perf-neutral by amortizing it through the existing warp shuffle.
Testing
test_aoti_torch_cuda_int5_plain_mm: 5/5 passed on GPU (SingleSuperBlock, MultiSuperBlock, WideN, PackedShuffleMultiSuper z_pack stress at K=1024/n_super=4, NullInputHandling). Recompiled from scratch; executed on GPU (non-zero kernel times).test_int5_dispatch.py: 16/16 passed (incl. real CUDA F.linear on GPU).