feat(kv-cache): add q6 paged KV quantization - #267
Open
TESTYEE-09 wants to merge 1 commit into
Open
Conversation
Adds an intermediate 6-bit paged KV storage option. Large-context Qwen 3.8
27B on a 32 GB Mac has no comfortable setting near 262k: q8 is still tight
at the top of the context and q4 quantizes on a much coarser grid. q6 is a
point in between. It is additive — no default changes, and q8/q4 numerics
are untouched.
Same symmetric per-token-per-head quantizer as the existing modes, only the
clip bound moves (qmax 31, one fp16 scale per row). Codes are biased by +32
into a 6-bit field and packed four per three bytes, little-endian:
byte0 = u0 | (u1 & 0x03) << 6
byte1 = (u1 >> 2) | (u2 & 0x0F) << 4
byte2 = (u2 >> 4) | (u3 & 0x3F) << 2
So packed_dim(D, 6) = D * 3/4, and q6 requires D % 4 == 0 — satisfied
naturally by Qwen 3.8's 256-wide K/V heads (256 -> 192 bytes). Pack and
unpack are vectorized MLX ops over groups of four, no Python element loop.
Two existing expressions would have mis-read a 6-bit mode, and both are now
explicit tables that raise on unknown widths: PagedKVQuantConfig.bits read
`4 if mode == "q4" else 8`, and quantize_symmetric's qmax read
`127 if bits == 8 else 7`. Either would have put q6 on the wrong grid while
allocating q6-sized pages. The cache allocation was already width-generic
through packed_dim; q6 shares q4's uint8 packed-storage branch, now
commented as covering every sub-byte width. No native paged-attention
kernel consumes the packed bytes on this path — plain KV quant dequantizes
and runs MLX SDPA — so no kernel needed a 6-bit case.
Calculated K+V storage per token per KV head at head_dim 256: 388 B for q6
against 516 (q8), 260 (q4) and 1024 (fp16). For a filled 262,144-token
Qwen 3.8 full-attention cache (16 layers, 4 KV heads) that is ~6.06 GiB
against ~8.06 / ~4.06 / ~16.00 GiB. Storage arithmetic, not measured
process RAM.
Aliases stay narrow (6, 6bit, int6, uint6); q6_k and q6_0 are rejected. The
packing is MTPLX's own and is not llama.cpp Q6_K or any other external
format. q6 is lossy — no quality claim is made relative to q8.
NOT RUN: no test in this commit has been executed, and nothing was
benchmarked, on any machine. Quant/dequant timing is unmeasured, and q6's
byte-crossing packing is the most likely place for it to cost decode
throughput relative to q8. Design notes and the review checklist are in
docs/specs/2026-08-16-q6-paged-kv-design.md.
Owner
|
The two latent bit-width fixes in kv_quant.py (the bits property and the clip-bound map) are correct defensive hardening and I will land them with credit in the next minor. The q6 surface itself stays out for now: it collides with the 2.8.0 kv-quant rework and the PR is self-declared unexecuted. If you rebase the q6 lane onto 2.8.0 and run the checks in CONTRIBUTING.md, I will look again. |
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
2 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
q6as an intermediate paged KV-cache quantization mode between the existingq8andq4modes.The implementation uses the existing symmetric KV quantization architecture with
qmax = 31and packs four 6-bit values into three bytes.For a 256-dimensional KV head:
This is an internal MTPLX q6 representation and is not llama.cpp Q6_K or another external quantization format.
Motivation
For large-context models on memory-constrained Apple Silicon, there is a fairly large memory gap between q8 and q4 KV storage. q6 provides an intermediate memory/precision option.
For Qwen3.8-27B with 256-dimensional KV heads, calculated K+V storage per token/KV-head is:
At a completely filled 262,144-token context with 16 full-attention layers and 4 KV heads, calculated KV storage is approximately:
These are calculated cache-storage figures, not measured process RAM.
Implementation
q6mode normalization/config supportThe audit also found two existing assumptions that effectively treated any non-q4 width as q8, or any non-q8 width as q4. Those are now explicit bit-width mappings so q6 cannot silently use the wrong quantization grid.
Verification
None of the checks in
CONTRIBUTING.mdwere run against this commit, and neither was the new q6 test suite. This branch is submitted for maintainer review as an unexecuted artifact, by deliberate choice of the submitter — not because the checklist passed or was overlooked. Please treat the validation status below as "written, not executed".The commit adds q6-specific tests covering:
To exercise them:
Two tolerance bounds in the new tests (the paged-cache round-trip tolerance and the float reconstruction-error bound) were chosen analytically rather than from observed output, so they are the most likely place for a first run to need adjustment.
Benchmark Evidence
None. q6 quantization/dequantization has not been benchmarked, and no runtime performance claim is made in this PR.
The implementation detail most worth review is the cross-byte packing cost: unlike q8 (one code per byte) and q4 (two per byte), q6 codes straddle byte boundaries, so pack/unpack each cost several mask/shift operations plus a stack and reshape per group of four values. Dequantization runs on every attention call over the active window, so that is the most likely place for q6 to cost decode throughput relative to q8.
Notes
Based on
963b923(2.7.1); the branch is a single commit on top of it.