{"id":49028922,"url":"https://github.com/arozanov/turboquant-mlx","last_synced_at":"2026-05-05T13:01:00.927Z","repository":{"id":347568083,"uuid":"1194350559","full_name":"arozanov/turboquant-mlx","owner":"arozanov","description":"TurboQuant KV cache compression for MLX with fused Metal kernels. 4.6x compression at 98% FP16 speed.","archived":false,"fork":false,"pushed_at":"2026-04-17T09:26:05.000Z","size":98,"stargazers_count":88,"open_issues_count":1,"forks_count":16,"subscribers_count":1,"default_branch":"main","last_synced_at":"2026-04-17T10:35:00.335Z","etag":null,"topics":["apple-silicon","kv-cache","llm","metal","mlx","quantization","turboquant"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/arozanov.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-03-28T08:27:42.000Z","updated_at":"2026-04-17T09:26:08.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/arozanov/turboquant-mlx","commit_stats":null,"previous_names":["arozanov/turboquant-mlx"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/arozanov/turboquant-mlx","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arozanov%2Fturboquant-mlx","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arozanov%2Fturboquant-mlx/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arozanov%2Fturboquant-mlx/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arozanov%2Fturboquant-mlx/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/arozanov","download_url":"https://codeload.github.com/arozanov/turboquant-mlx/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/arozanov%2Fturboquant-mlx/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":32650449,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-05T11:29:49.557Z","status":"ssl_error","status_checked_at":"2026-05-05T11:29:48.587Z","response_time":54,"last_error":"SSL_connect returned=1 errno=0 peeraddr=140.82.121.5:443 state=error: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["apple-silicon","kv-cache","llm","metal","mlx","quantization","turboquant"],"created_at":"2026-04-19T09:00:36.932Z","updated_at":"2026-05-05T13:01:00.920Z","avatar_url":"https://github.com/arozanov.png","language":"Python","funding_links":[],"categories":["Rising projects"],"sub_categories":[],"readme":"# turboquant-mlx\n\n[TurboQuant](https://arxiv.org/abs/2504.19874) KV cache compression for [MLX](https://github.com/ml-explore/mlx) on Apple Silicon.\n\nPolarQuant (randomized Hadamard rotation + Lloyd-Max quantization) compresses KV cache values to 3-bit with fused Metal kernels. Drop-in replacement for mlx-lm's KVCache.\n\n## Key Finding\n\nK and V quantization behave very differently:\n\n- **K quantization destroys greedy decode** at 4-bit and below (even MLX's native `kv_bits=4`). Softmax is sensitive to small score perturbations.\n- **V quantization is safe** at 3-bit. Weighted interpolation tolerates noise.\n\nThis means mixed-precision is the right approach: K at 8-bit (preserves attention) + V at 4-bit or lower (saves memory).\n\n## Results (Qwen 2.5 7B, 32K context)\n\n| Config | Active Memory | Savings | Decode | Quality |\n|--------|--------------|---------|--------|---------|\n| Baseline fp16 | 6.21 GB | -- | 35.75 tok/s | correct |\n| **K8 + V4 mixed-quant** | **5.08 GB** | **-1.13 GB (-18%)** | 25.84 tok/s | **identical** |\n| K8 + V2 mixed-quant | 4.97 GB | -1.24 GB (-20%) | 25.52 tok/s | identical |\n\nQuality is verified identical: greedy decode produces the same text as baseline.\n\n## Quick Start\n\n### Mixed-precision quantized cache (recommended)\n\nUses Apple's native `mx.quantized_matmul` for both K and V. Requires the [mlx-lm fork](https://github.com/arozanov/mlx-lm/tree/feature/turboquant-kv-cache) with `mixed_quantized_scaled_dot_product_attention`.\n\n```python\nfrom mlx_lm import load, stream_generate\nfrom mlx_lm.models.cache import make_prompt_cache\nfrom mlx_lm.models.mixed_quant_cache import MixedQuantKVCache\n\nmodel, tokenizer = load(\"mlx-community/Qwen2.5-7B-Instruct-4bit\")\nn_layers = len(model.model.layers)\ncache = make_prompt_cache(model)\n\n# Generate with fp16 cache for prefill, then convert\nfor i, response in enumerate(stream_generate(model, tokenizer, prompt=prompt, max_tokens=256, prompt_cache=cache)):\n    if i == 0:  # after prefill, convert to mixed-quant\n        for j in range(n_layers):\n            cache[j] = MixedQuantKVCache.from_kvcache(cache[j], k_bits=8, v_bits=4)\n    print(response.text, end=\"\", flush=True)\n```\n\n### V-only TurboQuant cache\n\nWorks with stock mlx-lm (no fork needed). Keeps K in fp16, compresses V with PolarQuant 3-bit.\n\n```python\nfrom turboquant_mlx.v_only_cache import VOnlyTurboQuantCache\n\ncache = [VOnlyTurboQuantCache(bits=3) for _ in range(n_layers)]\n# Use as normal mlx-lm prompt_cache\n```\n\n## Features\n\n- **Mixed-precision KV cache**: K at 8-bit + V at 4-bit via Apple's `mx.quantized_matmul`\n- **V-only TurboQuant**: PolarQuant 3-bit V compression, quality-preserving\n- **Fused Metal kernels**: pre-rotated Q scoring (`prerot_fused_qk_scores`), sparse V attention (`sparse_v_matvec`)\n- **Butterfly-pulled-out optimization**: WHT linearity lets us accumulate weighted centroids first, butterfly once at end (4.5x speedup on V-attention kernel)\n- **SIMD-group reductions**: `simd_sum` replaces tree reduction in QK scoring (1.85x kernel speedup)\n- **Flash-attention scaffold**: single-kernel fused SDPA over packed K/V (correct, scaffold for future optimization)\n- **GQA-aware kernels**: `n_rep` parameter avoids `mx.repeat` allocation on GQA models\n\n## How It Works\n\n```\nQuantize (fused Metal kernel):\n  Input KV vector x (head_dim=128)\n  -\u003e norm = ||x||, x_unit = x / norm\n  -\u003e rotate: y = WHT(signs * x_unit)  (O(d log d), Gaussianizes coordinates)\n  -\u003e quantize: idx = nearest_centroid(y)  (Lloyd-Max codebook, 8 levels for 3-bit)\n  -\u003e pack: 10 x 3-bit indices per uint32\n\nDequant (parallel Metal kernel, d threads cooperating):\n  centroids[indices] -\u003e parallel WHT butterfly -\u003e * signs -\u003e * norm -\u003e output\n\nButterfly-pulled-out (sparse_v_matvec):\n  sum_pos w[pos] * butterfly(c[idx_pos])\n    = butterfly(sum_pos w[pos] * c[idx_pos])   # WHT is linear!\n  -\u003e accumulate per-thread (no barriers), one butterfly at end\n```\n\n## Project Structure\n\n```\nturboquant_mlx/\n  cache.py               TurboQuantKVCache (packed K/V with fused Metal encode/decode)\n  v_only_cache.py        VOnlyTurboQuantCache (fp16 K + TQ 3-bit V)\n  metal_kernels_v4.py    Pre-rotated Q kernels (prerotate_query, prerot_fused_qk_scores)\n  sparse_v.py            Sparse V attention with butterfly-pulled-out trick\n  flash_attention.py     Single-kernel fused SDPA scaffold\n  fused_attention.py     Composed fused attention (prerot Q + sparse V)\n  patch.py               Monkey-patch mlx-lm SDPA for fused/hybrid paths\n  hybrid_cache.py        Experimental: Apple K8 + TQ V3 (scaffold)\n  hybrid_attention.py    Experimental: mixed Apple + TQ SDPA\n  rotation.py            Walsh-Hadamard Transform (pure MLX)\n  quantizer.py           PolarQuant: rotation + Lloyd-Max codebook\n  kernels.py             Packed dequant + fused QK Metal kernels\n  metal.py               Fused quantize + dequant Metal kernels\n  packing.py             Bit-packing utilities\n  adaptive.py            Layer-adaptive cache factory\n\nscripts/\n  bench_sparse_v.py      Sparse V kernel microbenchmark\n  bench_real_model.py    End-to-end model benchmark (4 paths)\n  bench_long_context.py  Long-context memory comparison\n\ntests/\n  test_core.py           Core algorithm (10 tests)\n  test_prerot.py         Pre-rotated Q kernel correctness (9 tests)\n  test_sparse_v.py       Sparse V correctness + GQA (8 tests)\n  test_fused_attn.py     End-to-end fused attention (6 tests)\n  test_flash_attention.py Flash-attention correctness (7 tests)\n  test_v_only_cache.py   V-only cache, adaptive cache, serialization (11 tests)\n```\n\n## Install\n\n```bash\ngit clone https://github.com/arozanov/turboquant-mlx.git\ncd turboquant-mlx\npip install -e .\n```\n\nFor mixed-quant cache (K8+V4), also install the mlx-lm fork:\n```bash\npip install -e ../mlx-lm  # or wherever the fork lives\n```\n\n## Run Tests\n\n```bash\npytest tests/ -v\n# 51 tests, all passing\n```\n\n## Server Integration\n\nThe [mlx-lm fork](https://github.com/arozanov/mlx-lm/tree/feature/turboquant-kv-cache) adds KV cache quantization, disk persistence, and MoE support to `mlx_lm.server`.\n\n```bash\npip install --force-reinstall --no-cache-dir git+https://github.com/arozanov/mlx-lm.git@feature/turboquant-kv-cache\n```\n\n### Server flags\n\n| Flag | Description |\n|------|-------------|\n| `--kv-cache-quantization K,V` | Quantize KV cache: K at K-bit, V at V-bit (e.g. `8,4`) |\n| `--quantized-kv-start N` | Only quantize caches with at least N tokens (skip short prefills) |\n| `--prompt-cache-dir PATH` | Persist prompt caches to disk, survives server restarts |\n| `--no-batch` | Disable batch processing, use single-serve mode |\n\n### Example\n\n```bash\nmlx_lm.server \\\n  --model mlx-community/Qwen2.5-7B-Instruct-4bit \\\n  --kv-cache-quantization 8,4 \\\n  --quantized-kv-start 1024 \\\n  --prompt-cache-dir ~/.cache/mlx_kv_cache \\\n  --no-batch\n```\n\nDisk cache saves KV caches to disk on every insert. On server restart, caches are loaded from disk on cache miss (lazy loading). Works with MoE models (GLM-5.1, Kimi-K2.6, DeepSeek V3) that use CacheList.\n\n## References\n\n- **TurboQuant**: [arXiv 2504.19874](https://arxiv.org/abs/2504.19874)\n- **PolarQuant**: [arXiv 2502.02617](https://arxiv.org/abs/2502.02617)\n- **MLX**: [github.com/ml-explore/mlx](https://github.com/ml-explore/mlx)\n\n## License\n\nApache License 2.0\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farozanov%2Fturboquant-mlx","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Farozanov%2Fturboquant-mlx","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farozanov%2Fturboquant-mlx/lists"}