-
Notifications
You must be signed in to change notification settings - Fork 0
Add RotorQuant KV cache backend with deferred prefill on Metal #103
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 1 commit
4557598
53a0907
9c93b6e
ff21598
f40e9d4
aba98c2
8545421
032593b
f587d2b
a43ed82
7ba35fc
f41b97d
3d526b7
3100207
bae6455
c5ec472
13a5f47
04893db
6c4ba6f
d612d14
a50a537
856b30f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -449,16 +449,20 @@ export function SettingsPanel({ open, onClose }: SettingsPanelProps) { | |
| filled | ||
| content={ | ||
| `• Default — No cache quantization. Best baseline quality, highest memory use.\n` + | ||
| `• OptiQ — Rotation-based quantization via mlx-optiq. Best long-context quality.\n` + | ||
| `• RotorQuant Adaptive — IsoQuant 3-bit with deferred prefill, FP16 edge layers. Recommended.\n` + | ||
| `• RotorQuant — IsoQuant 3-bit on all KV layers. Most aggressive compression with deferred prefill.\n` + | ||
| `• OptiQ — Rotation-based quantization via mlx-optiq. Good long-context quality, no GQA support.\n` + | ||
| `• TurboQuant Adaptive — Quantizes middle KV layers, keeps edge layers in FP16. Proven stable.\n` + | ||
| `• TurboQuant — Quantizes all KV layers. Most aggressive compression, higher quality risk.\n` + | ||
| `• TurboQuant — Quantizes all KV layers. Most aggressive non-rotorquant compression.\n` + | ||
| `• MLX Quantized — MLX's built-in cache quantization.\n\n` + | ||
| `Takes effect on next model launch. Incompatible models fall back to Default automatically.` | ||
|
ttupper92618 marked this conversation as resolved.
Outdated
|
||
| } | ||
|
Comment on lines
451
to
457
|
||
| /> | ||
| </FieldLabel> | ||
| <Select value={kvBackend} onChange={(e) => setKvBackend(e.target.value)} disabled={!!envOverride}> | ||
| <option value="default">Default (no quantization)</option> | ||
| <option value="rotorquant_adaptive">RotorQuant Adaptive (recommended)</option> | ||
| <option value="rotorquant">RotorQuant</option> | ||
| <option value="optiq">OptiQ (rotation-based)</option> | ||
| <option value="turboquant_adaptive">TurboQuant Adaptive</option> | ||
| <option value="turboquant">TurboQuant</option> | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| """RotorQuant KV cache backend (IsoQuant variant). | ||
|
|
||
| Pure-MLX port of the IsoQuant 3-bit KV cache compression from | ||
| johndpope/llama-cpp-turboquant (MIT) and scrya-com/rotorquant (MIT). | ||
|
|
||
| Key properties vs the older TurboQuant native backend: | ||
| - Block-diagonal 4D quaternion rotations instead of randomized Hadamard | ||
| - Norm-correction trick for unbiased magnitudes after centroid quantization | ||
| - Optional deferred prefill: K/V stay in fp16 during prompt processing, | ||
| flushed to compressed storage on the first decode-shaped call | ||
| - GQA-native (compression is per-(kv_head, token), Q heads fan out at SDPA) | ||
|
|
||
| The backend stores indices and norms; ``update_and_fetch`` returns fully | ||
| dequantized fp16 K/V to standard ``mx.fast.scaled_dot_product_attention``. | ||
| The centroid-space attention optimization from OptiQ is intentionally | ||
| deferred to a follow-up; v1 prioritizes correctness and the deferred-prefill | ||
| quality win over the rotated-space SDPA perf win. | ||
| """ | ||
|
|
||
| from exo.worker.engines.mlx.rotorquant.cache import ( | ||
| RotorQuantKVCache, | ||
| make_rotorquant_adaptive_cache, | ||
| make_rotorquant_cache_from_template, | ||
| ) | ||
|
|
||
| __all__ = [ | ||
| "RotorQuantKVCache", | ||
| "make_rotorquant_adaptive_cache", | ||
| "make_rotorquant_cache_from_template", | ||
| ] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This bullet links to
docs/kv-cache-backends.md, but that page currently doesn’t mention RotorQuant (it still lists only default/mlx_quantized/turboquant/optiq). Either updatedocs/kv-cache-backends.mdin this PR to include RotorQuant, or change the link to the up-to-date page underwebsite/docs/kv-cache-backends.md/ the published docs URL so readers don’t land on stale information.