feat(kvcache): add reliable quantized KV cache - #309
Open
plsgivemeachane wants to merge 7 commits into
Open
Conversation
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.
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).
fangyuan-3149
pushed a commit
to fangyuan-3149/FreeToken
that referenced
this pull request
Sep 2, 2026
… wrapper The attention wrapper never forwarded k_scale/v_scale/layout to the paged kernels, so quantized pools served attention against the raw byte payload: sub-byte caches trip the physical/logical head_dim assert immediately, and Q8 -- whose physical shape equals the logical one -- silently feeds raw int8 codes into attention without their dequantization scales. Wrapper-local wiring ported from plsgivemeachane's aad9a16 on PR FlashML-org#309, with the author's report as the original bug source; the serving stack used here had been masked from the same defect by a stale pre-PR wrapper copy in the local venv. Also asserts the pool's physical head_dim against the active quant spec so a wiring regression cannot go silent again.
fangyuan-3149
pushed a commit
to fangyuan-3149/FreeToken
that referenced
this pull request
Sep 2, 2026
l_i = 1 is int32 in Triton; when a later block redefines it as fp32 the merged-type check rejects the kernel outright (CompilationError: initial value for l_i is of type int32[], but the then block redefines it as fp32[]) -- caught by test_paged_triton_attention_with_sinks_matches_reference, which failed to compile before this change. One-line fix ported from plsgivemeachane's 716c1a2 on PR FlashML-org#309.
|
Thanks for building on this work with the attribution intact — a 262K-context user becoming a contributor is exactly the outcome we hoped for when opening #268. Status sync for maintainers: the core of both PRs is now the same lineage (this PR's first three commits carry fangyuan-3149's authorship), and the two things this PR added on top were genuinely good catches — the missing scale/layout forwarding in the attention wrapper (the silent Q8 raw-int8 path) and the fp32 sink accumulator. Both are now in #268 with credit (fc0771a, 3096737), alongside bd37c60 for the sizing issue you reported there. What #268 additionally carries at this point: a six-scheme accuracy/speed ladder (q2_lm → q8_0, same protocol, greedy), needle-in-haystack from 8K through 220K, 220K–300K context measurements on an 8 GB consumer GPU, Windows cold-start fixes, and the mag=8 rounding tweak (~5% rel_err over the GGUF clamp). Given the overlap, running two PRs in parallel probably serves nobody — we're happy either way: stack your remaining deltas (the test names are nicer than ours in places) on #268 as follow-up commits, or let the maintainers pick a direction. Thanks again for the energy here. |
|
Second this |
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.
Closes #262. Related to #280. Supersedes #268 while preserving and crediting @fangyuan-3149 for the original work.
Summary
--kv-cache-dtypethrough serving and pool allocation.The Q8 corruption occurred because attention read raw int8 values without their dequantization scales.
Validation
92 passed, 1 skipped.git diff --checkpasses.The exact launch command and checkpoint used for the RTX 5060 smoke test were not recorded.