Repository navigation
Enable register-codebook W4A4 decode on Ada GPUs and add INT8 output heads - #13
Merged
Merged
Conversation
…ows_int8/matmul_int8_rows) Opt-in per-row absmax INT8 for output heads kept outside the W4A4 packed format: Int8RowLinear / quantize_output_heads in Python, plus two CUDA operators in native 1.0.5 (quantize_rows_int8, matmul_int8_rows: one warp per row, DP4A, 1..8 activation rows). The native kernel and the torch._int_mm fallback share the epilogue float(sum) * (x_scale * w_scale) and derive scales with a true IEEE division, so both paths return identical values (verified on RTX 4090: 81 native + 10 Python tests, bit-exact against _int_mm).
The CUDA source now carries 14 kernel launches with launch checks (the two INT8 row operators added two); the guard also names both new kernels.
iamwavecut
changed the title
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.
Ada GPUs currently use the shared-codebook W4A4 decode kernel even though the register-codebook path reduces decode time on RTX 4000 Ada. Enable the existing register implementation for SM89 with the same shape guards as SM120, leaving quantization, checkpoint layout, and other architecture dispatch unchanged.
Prepare OrbitQuant 0.9.8 and native 1.0.5, including managed-cache version isolation. Add a CUDA profiler regression that verifies the register kernel actually launches, and cover rejection/preservation of native 1.0.4 artifacts.
Validation: the integrated native build passed 260 GPU tests (30 platform skips) on RTX 4000 Ada / Torch 2.10.0+cu128. Ten complete Russian songs (warmup/control/candidate/candidate/control at CFG 1 and 1.5) preserved all PCM outputs. Mean semantic generation improved 1.33% and 6.39%; whole-song means were 39.882 → 39.244 s and 44.634 → 42.700 s, including unrelated VAE variation. Ruff and all 46 targeted package/provision tests passed after generated kernel-builder files were moved outside the source tree. The CI CPU suite also passed. GPU source/build and full metrics are durably archived; the disposable GPU has been deleted.
Add opt-in INT8 row-quantized output heads:
orbitquant.Int8RowLinear/quantize_output_heads/absmax_int8_scalesin Python and two CUDA operators in native 1.0.5 (quantize_rows_int8,matmul_int8_rows: one warp per output row, DP4A accumulation, 1..8 activation rows, BF16/FP16 output, optional bias). The native kernel and thetorch._int_mmfallback share the epiloguefloat(sum) * (x_scale * w_scale)and derive scales with a true IEEE division (PyTorch's Python-scalar divisor path multiplies by a reciprocal and is 1 ulp off in some rows), so both paths return identical values. Heads stay in source precision unless the caller opts in.INT8-head validation on RTX 4090 / Torch 2.10.0+cu128: the full native suite passed (333 passed, 30 platform skips),
tests/test_int8_rows.py81/81 bit-exact againsttorch._int_mmandtests/test_int8_head.py10/10. On YuE2-3B the semantic head (32,769 classes) measured a mean forward KL of 0.0015 (top-1 agreement 0.914) and the ABC head 0.0002 over 128 forced positions; the DP4A head GEMV reached 937 GB/s. The PR Build kernels matrix (19 variants + manifest) and the CPU CI suite pass.