feat(kvcache): add ISO3/ISO4 KV-cache quantization (--kv-cache-iso)

[AsmanovLev]

What

Optional IsoQuant quantization of the paged KV cache for plain
full-attention models, ported bit-exact from the llama.cpp fork
llama-cpp-turbo-planar-iso (ggml ISO3_0/ISO4_0):

• --kv-cache-iso iso3: 3.125 bits/value (50 B per 128-value block) - 5.1x smaller KV than bf16
• --kv-cache-iso iso4: 4.25 bits/value (68 B per block) - 3.8x smaller

Per 128-value block of a head vector: L2-normalize, rotate each 4D group by a
fixed unit quaternion, quantize to the nearest Lloyd-Max centroid, store the
corrected norm ||x||/||centroids|| (norm-preserving dequant).

The flag forces --attention-backend iso: a new ISOKVCache pool (packed
uint8 slabs, drop-in MHAKVCache replacement incl. hybrid GDN layer remap) plus
custom CUDA kernels. Decode is quantize-on-write; extend attends the packed
prefix + bf16 new tokens and packs them after attention (deferred: prefill
never consumes its own quantized KV). Extend runs the packed prefix through a
one-time dequant into a bf16 scratch served by the existing tiled triton
kernel (custom CUDA extend kernel remains as fallback for empty prefix /
oversized scratch, FREETOKEN_ISO_SCRATCH_MB, default 128 MiB).

Why

On small-VRAM GPUs the KV cache - not the weights - caps the usable context.

Hardware:

• RTX 3060 Laptop 6 GB (sm_86), driver 610.57.04, CUDA 13.3
• I5-10500H, 32 GB DDR4

Checkpoint:

• aabbas/Huihui-Ornith-1.5-35B-A3B-abliterated-NVFP4-freetoken
  (35B-A3B Qwen3.5-MoE hybrid; GPU-resident weights 4.03 GiB, experts offloaded)
  Command:

  ft serve --model-path <ckpt> --kv-cache-iso iso3 --memory-ratio 1.0 \
    --max-running-req 1 --num-tokens 81920 --moe-cache-size 256 \
    --disable-moe-prefill-overlap --cuda-graph-max-bs 0 --cache-type naive

	main (bf16 KV)	branch (iso3)
KV bytes/token (10 full-attn layers)	20,480	4,000
This model on 6 GB	does not start (negative KV budget)	65k-80k token context
Decode	-	~7-9 tok/s
Prefill	-	~213 tok/s, flat with prefix length

Prefill A/B within the branch (custom CUDA extend kernel vs dequant+triton,
same 30k-token prompt): 3m42s -> 2m21s; throughput no longer degrades with
prefix length (~146 -> ~213 tok/s at 30k).

Correctness

• quantize/dequantize bit-exact vs the fork's unmodified C reference
  (golden vectors: tests/kernels/iso_golden.npz); CUDA pack/unpack kernels
  byte-identical to the torch reference
• attention kernels vs dequantized-oracle attention: cosine > 0.9999
• 31 new tests (tests/kernels/test_iso_{quant,store_kernel,attention}.py,
  tests/kvcache/test_iso_pool.py); full tests/kvcache + tests/engine +
  tests/kernels suite passes (597 passed, 21 skipped)

Caveats

• Quality is model-dependent: models with strong outlier K channels (dense
  Qwen3) degrade noticeably at iso3 (the 4D-group rotation cannot gaussianize
  channel spikes). Documented in docs/cli.md. Validated well-behaved on the
  Qwen3.5-MoE hybrid above (coherent 30k-context answers).
• Plain full-attention pools only (no SWA/MLA/DSA/BSA/QSA/DSV4), bf16 KV
  input, head_dim % 128 == 0 (validated at config time with clear errors).

[Copilot]

🔵 Needs a closer look

It introduces new CUDA kernels and a new attention backend in a core execution path, which warrants careful human validation despite strong test coverage.

Pull request overview

This PR adds optional ISO3/ISO4 quantization for the paged KV cache, including a new ISO attention backend and CUDA pack/dequant/attention kernels, exposed via a new --kv-cache-iso server flag and documented in the CLI docs.

Changes:

• Introduces ISOKVCache (packed uint8 KV pool) and ISO3/ISO4 reference quantization utilities.
• Adds ISO decode/extend attention CUDA kernels and wires them into a new iso attention backend with config-time validation/forcing.
• Adds targeted test coverage for quantization bit-exactness, store kernels, attention correctness, and pool sizing/accessors; updates CLI docs and server args.

File summaries

File	Description
tests/kvcache/test_iso_pool.py	New unit tests for ISOKVCache sizing, allocation, remapping, and store packing.
tests/kernels/test_iso_store_kernel.py	New CUDA tests validating store/dequant kernels vs reference.
tests/kernels/test_iso_quant.py	New tests for ISO reference quantize/dequant and golden-vector bit-exactness.
tests/kernels/test_iso_attention.py	New CUDA tests validating ISO attention decode/extend vs dequantized oracle.
python/freetoken/server/args.py	Adds --kv-cache-iso CLI flag with choices and help text.
python/freetoken/kvcache/iso_pool.py	Adds ISOKVCache packed KV pool implementation.
python/freetoken/kvcache/base.py	Improves KV-cache “not enough memory” assertion with detailed sizing context.
python/freetoken/kvcache/init.py	Threads kv_cache_iso through KV pool creation and selects ISOKVCache when enabled.
python/freetoken/kernel/iso.py	Adds ISO3/ISO4 reference quantization + JIT bindings for store/dequant and attention kernels.
python/freetoken/kernel/csrc/jit/iso_store.cu	Adds CUDA pack/unpack (quantize-on-write) kernels for ISO KV.
python/freetoken/kernel/csrc/jit/iso_common.cuh	Adds shared constants and device helpers for ISO quant/dequant.
python/freetoken/kernel/csrc/jit/iso_attention.cu	Adds CUDA ISO attention decode/extend kernels (packed-prefix + bf16-extend).
python/freetoken/engine/engine.py	Adds config-time validation for kv_cache_iso and forces/guards attention backend selection.
python/freetoken/engine/config.py	Adds kv_cache_iso field to EngineConfig.
python/freetoken/attention/iso.py	Adds IsoAttentionBackend integrating packed KV + CUDA kernels + dequant-to-triton fast path.
python/freetoken/attention/init.py	Registers the new iso attention backend (FULL attention only).
docs/cli.md	Documents --kv-cache-iso and adds iso to attention-backend list.

Review details

• Files reviewed: 17/18 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.