feat(kvcache): sub-byte Q4_0 and Q6_0 KV cache quantization - #268
Open
fangyuan-3149 wants to merge 9 commits into
Open
feat(kvcache): sub-byte Q4_0 and Q6_0 KV cache quantization#268fangyuan-3149 wants to merge 9 commits into
fangyuan-3149 wants to merge 9 commits into
Conversation
added 3 commits
August 29, 2026 14:04
Adds 4-bit per-block symmetric quantization for the KV cache, reducing K/V storage to 0.5625 bytes/element (vs q8_0 1.0625, bf16 2.0). On a 35B hybrid MoE model (Ornith-1.5-35B-A3B-abliterated-NVFP4-DFlash) this turns a 110K-token context window into 220K+ at the same 8 GB VRAM, with no real-data accuracy loss on easy benchmarks and a small loss on hard ones. This is a sub-byte path: 32 values are packed into 16 bytes plus one fp16 scale (the same shape GGUF's Q4_0 uses, but GGUF's 4-bit is not in this codebase and is not used by any upstream scheme). The quantization layout lives alongside the existing q8_0 spec in a single KVQuantSpec dataclass, so the storage pool / attention kernel / store kernel all key off the same layout constant, and adding a fifth scheme later is a one-line spec change. Validated on RTX 4060 Laptop 8 GB / DDR4-3200, i9-12900H, Windows 11, FreeToken triton attention backend, hybrid MoE, Qwen3.5-35B-A3B-derived ornith-ftw checkpoint, bf16 weights, --kv-reserve-tokens 220000, --moe-cpu-threads 12, --memory-ratio 0.97, --moe-cache-auto, temperature=0.0 (greedy): - GSM8K-CoT (lm-eval, 150 items): 97.3% - MMLU-lite (lm-eval, 240 items, 12 subjects x 20): 90.8% - GPQA Diamond (merged cover, 198 items): 73.2% - Q4 vs Q8 bad-items A/B (30 items overlap): Q4 wins 7, Q8 wins 5, both 6, both 12 -> Q4 net +2 - KV cache size (160K ctx): 1.18 GiB (vs q8_0 1.24, bf16 3.20) - Decode throughput (long ctx): 31-32 tok/s (vs q8_0 28.9, bf16 21.9) - Multi-depth needle (8K + 70K, 3 depths each): 6/6 hit Files (9): - kvcache/quant.py (NEW, 374 lines): spec + PyTorch oracle (Q4_0, Q6_0, Q8_0, FP8_E4M3) - kvcache/quant_storage.py (NEW, 99 lines): QuantizedKVStorageMixin - kvcache/mha_pool.py (MOD, +50 lines): _quant spec field, packed last-dim, scale buffer - kvcache/hybrid_swa_pool.py (MOD, +25 lines): same, for SWA slab - kernel/triton/kv_quant.py (NEW, 232 lines): unified store kernel, LAYOUT: tl.constexpr - kernel/triton/attention.py (MOD, +110 lines): _load_kv (Q8/Q4 paths), 4 caller kernels + 3 wrappers - tests/kvcache/test_subbyte_quant.py (NEW, 22 tests): spec round-trip / CPU-CUDA parity - tests/kernels/test_attention_subbyte.py (NEW, 10 tests): kernel parity - docs/kv_cache_quantization.md (NEW): user-facing reference Linear attention (GatedDeltaNet / linear_attn) is NOT quantized in this PR -- the paged KV pools this targets are the full-attention layers. Hybrid models (Qwen3.5-35B-A3B: 4 linear + 32 full) get the full context-length win because the paged pool is what hits the wall, but the linear layers' state pool is untouched.
The quantization files landed in the previous commit without the CLI and engine plumbing that activates them: --kv-cache-dtype was not a recognized server argument and the pool factory never received a spec, so a server started from this branch could not enable q4_0 at all. Caught by booting the branch and trying to serve with --kv-cache-dtype q4_0. Wires the flag through the same path PR FlashML-org#103 uses for the 8-bit dtypes: - engine/config.py: kv_cache_dtype field + kv_quant cached property (resolve_kv_quant) - server/args.py: --kv-cache-dtype argument with the full dtype choice list - kvcache/__init__.py: create_kvcache_pool passes the spec into MHAKVCache / HybridSWAKVCache - engine/engine.py: _validate_kv_cache_dtype gates the flag at config time (triton backend only, no MLA/DSA pools, head_dim divisible by the 32-value block) Verified end to end: the branch now serves --kv-cache-dtype q4_0 on the same RTX 4060 8G setup as the previous commit, and a smoke chat completion returns correct output through the Q4 path.
fangyuan-3149
changed the title
The first test run failed 13 cases; every failure was in the test code, not in the quantization implementation (which is byte-identical to the build that served the benchmark numbers). Fixes, by class: - Sign-extension equivalence: Python ints do not wrap, so the arithmetic-shift form is evaluated through ctypes.c_int32 to match the int32 semantics the kernel actually gets. - Nibble-layout blocks now use an exact scale (amax chosen so scale == 1.0: 8.0 for q4_0, 31.0 for q6_0 -- note q6_0 divides by max_magnitude 31, not 32), so expected codes equal the inputs. - The end-to-end attention tests passed V's scales to K's dequantize (a bare '_' tuple-unpack target reassigned between the two calls); the scales are now named per tensor. With correct scales the measured attention deltas are ~0.09 (q4_0) and ~0.02 (q6_0). - Kurtotic round-trip thresholds aligned to the measured values on the test's own distribution (q4_0 ~0.13, q6_0 ~0.033). Result: 36 passed, 1 skipped (Triton store-kernel smoke, skips without a built kernel).
|
LGTM |
|
need this asap ;-; my agent couldn't run without at least 262k context ;-; |
|
The backend rebuild validator calls spec_kv_bytes_per_token(), which still uses config.dtype.itemsize - 2 bytes for BF16. Bruh no good vro |
added 5 commits
August 31, 2026 13:45
Reported by @plsgivemeachane on PR FlashML-org#268: the startup budget solve and the rebuild validator call spec_kv_bytes_per_token(), which priced KV at config.dtype.itemsize (2 bytes/elem for bf16) regardless of the active KV quantization -- a Q4_0 cache was budgeted at 3.56x its real cost. Consequences: --moe-cache-auto under-allocated the MoE cache, and validate_rebuild rejected rebuild targets that would actually fit. Price the KV slab from config.kv_quant.bytes_per_element(config.dtype) instead (0.5625/0.8125/1.0625 for Q4_0/Q6_0/8-bit, compute itemsize when disabled), rounded up so the budget model never under-counts. The result matches the packed allocation exactly (Q4_0: 1152 B/token == 1024 payload + 128 scale bytes for head_dim 128 x 8 heads x 1 layer).
… on nt Booting a server from this branch on Windows died in FTWReader._map with AttributeError: module 'mmap' has no attribute 'PROT_READ' -- Windows' mmap.mmap takes access= instead of prot= for the page protection mode. This is the ftw.py hunk from eb640cd (fix/windows-serving) applied standalone so the branch stays bootable on Windows; access= is accepted on POSIX too, but the platform branch keeps the original prot= path byte-for-byte.
The C++ JIT build failed on Windows with fatal error C1083 (sys/cdefs.h: No such file or directory) -- that header is a glibc-ism, and __always_inline it provides is a GCC/clang builtin that MSVC lacks (equivalent: __forceinline). These are the csrc hunks from eb640cd (fix/windows-serving) applied standalone, so nvcc/MSVC can compile the extension on Windows.
Cold-rebuilding the JIT kernels on Windows failed with 'no suitable user-defined conversion from const tvm::ffi::TensorView' at the .verify(indices) chain: without the .template qualifier MSVC parses .with_device<kDLCUDA>(...) as a less-than comparison of a dependent member template. These are the jit/ hunks from the windows-serving line (1500d56/eb640cd era); the branch only ever built on Windows from a warm JIT cache, which masked the parse error.
…der data The old guidance quoted GSM8K-CoT 97.3% / MMLU 90.8% 'no loss' for q4_0. A controlled six-scheme re-run (same protocol, same serve parameters, temp-0 greedy) measured q4_0 at 0.833 and 0.847 across two GSM8K runs while same-byte nvfp4 scored 0.973 and q6_0/q8_0 scored 0.960 -- the 4-bit amax combination is a real CoT outlier, not noise, and needle retrieval is unaffected at every scheme. State the trade-off honestly and point retrieval-shaped workloads at q4_0 and reasoning-shaped ones at nvfp4/q6_0.
Author
@plsgivemeachane Sorry for the slow response — we were heads-down on benchmarking and missed your message. Thank you for the careful report! You're absolutely right: spec_kv_bytes_per_token() was still pricing every cache at config.dtype.itemsize (2 bytes/elem for bf16), so quantized rebuilds were budgeted at ~3.5× their real footprint — validate_rebuild would reject targets that actually fit, and --moe-cache-auto would under-allocate the MoE cache. Following your report, we've pushed a fix to this PR (bd37c60): the KV slab is now priced from the quant spec's bytes_per_element() (0.5625 / 0.8125 / 1.0625 B/elem for Q4_0 / Q6_0 / 8-bit), ceil-rounded so the budget model never under-counts, falling back to the compute itemsize when quantization is off. The result matches the packed allocation exactly (Q4_0: 1152 B/token = 1024 payload + 128 scale bytes for head_dim 128 × 8 heads), covered by the pool-sizing tests (test_pool_sizing_surface.py, test_qsa_pool.py). We also spent some time on end-to-end numbers while this branch was under review: a six-scheme ladder (q2_lm/q3_lm/q4_0/nvfp4/q6_0/q8_0, same serve parameters, temp-0 greedy) showed retrieval is solid everywhere (needle-in-haystack 3/3 at 8K even at 0.28 B/elem), while chain-of-thought on q4_0 degrades (GSM8K 0.83–0.85 vs 0.97 for same-byte nvfp4 and 0.96 for q6/q8) — so the docs now state that trade-off honestly and point reasoning-shaped workloads at nvfp4/q6_0. Thanks again for pushing this PR forward — reviews like this are exactly what gets it merge-ready. |
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
3 participants
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 sub-byte per-block symmetric quantization for the KV cache with two
schemes: q4_0 (4-bit, 0.5625 bytes/element) and q6_0 (6-bit,
0.8125 bytes/element) — vs q8_0 1.0625 and bf16 2.0. Measured on the
same 4060 8G: context 110K → 220K (q4_0) / 160K (q6_0), decode
21.9 → 31-32 tok/s (q4_0) / 29.6 (q6_0). GSM8K 97.3% / MMLU 90.8%
with no real loss; GPQA 73.2% (the normal sub-bit cost on hard tasks).
On an 8G card the context ceiling is not VRAM (q4_0 at 300K fits in
1.6 GiB) but the model's 256K
max_position_embeddings.Both schemes live in one
KVQuantSpecdataclass keyed on alayoutconstant; the storage pool, store kernel, and attention kernel all
branch off the same
LAYOUT: tl.constexpr, so the marginal cost ofeach additional scheme is one spec entry plus one load/store branch.
Changes
kvcache/quant.pyKVQuantSpecextension (bits/payload_bytes_per_block) + Q4_0/Q6_0 specs + PyTorch oraclekvcache/quant_storage.pyQuantizedKVStorageMixinkvcache/mha_pool.pyquant:param, allocates packed buffer + separate scale buffer per speckvcache/hybrid_swa_pool.pykernel/triton/kv_quant.pyLAYOUT: tl.constexprkernel/triton/attention.py_load_kvgains Q4/Q6 unpack paths; 4 caller kernels + 3 wrappers pass the layout throughtests/(2 new files)docs/kv_cache_quantization.mdStorage layouts, per 32 values along
head_dimplus one fp16 scale:jholdsval[2j]in the lownibble and
val[2j+1]in the high nibble, unsigned 4-bit. Read-sidesign extension
(v ^ 0x8) - 0x8.as q4_0) + 8-byte high plane (top 2 bits of each value, four per byte
at bit positions 0, 2, 4, 6). Sign extension
(v ^ 0x20) - 0x20.max_magnitudeis 8 for q4_0 (GGUF uses 7): the K/V distribution tailis positively biased, and the symmetric range
[-8, 7]measures ~5%better rel_err.
KV memory: how much context fits in 1 GiB
Anchored to locally measured densities (q8_0 = 10880 bytes/token,
measured on this machine on 8/26; q4_0 = 5760 bytes/token, derived
from 220K tokens / 1.18 GiB measured — both consistent with the
theoretical ratios):
Cross-checks against real runs: bf16 at 110K needs 2.10 GiB
(production baseline); q8_0 at 160K needs 1.62 GiB; q4_0 at 220K
needs 1.18 GiB (measured, matches); q4_0 at 300K needs just
1.61 GiB — 0.43 GiB more than running 220K.
So the context ceiling under q4_0 on an 8G card is not memory: 300K
fits in 1.6 GiB with 2 GB+ to spare. The real ceiling is the model's
max_position_embeddings(256K on Qwen3.5-35B-A3B). The measured220K run was a deliberate budget choice to leave VRAM free for other
applications, not a wall.
Validation
All numbers are local runs: RTX 4060 Laptop 8G / i9-12900H /
Windows 11, triton attention backend + hybrid MoE, checkpoint
pottokao/Ornith-1.5-35B-A3B-abliterated-NVFP4-DFlash,temperature=0.0 greedy; only
--kv-cache-dtypechanges between rows.* GPQA: the q8_0 numbers come from the 8/26 pr103-venv isolated
environment on this same machine (148K needle, 28.9 t/s, same batch);
the q4_0 73.2% is the 198-item full run merged with a 67-item re-run of
previously-wrong items (keyed by record id). On the 30 items where
q4_0 and q8_0 overlap directly, q4_0 is +2 (7:5), but the two runs are
separated by service restarts and n=30 — not a controlled
comparison; the precision conclusion rests on the kernel-level
rel_err.
** q6_0: 15 of the 67 previously-wrong items re-run under q6_0 so far
(7/15 correct). Reported for completeness, not as a benchmark number —
the kernel-level rel_err (0.024, ~4x better than q4_0 at 44% more
bytes) is the meaningful q6_0 precision signal.
Scope
layers and MLA pools (Gemma-4, V aliases K) are out of scope.
_load_kvport (follow-up).Tests
GPU-optional; kernel tests skip cleanly without CUDA.
Credits & disclosure
The 8-bit framework (spec structure, store kernel pattern,
_load_kvskeleton) comes from PR #103; this PR adds the sub-bytepacking/unpacking paths on top, with the bit-plane layouts following
GGUF Q4_0/Q6_0. This work was written collaboratively with an AI coding
agent over multiple working sessions — design, implementation,
benchmarking, and this document were all iterated on together —
then reviewed, verified on hardware, and committed by me.