feat(kvcache): add reliable quantized KV cache - #309
Open
plsgivemeachane wants to merge 5 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).
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.
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.