How 0xsero/turboquant works
Purpose and Entry Points
TurboQuant is a Python/PyTorch GPU library that compresses transformer KV caches during LLM inference. Its production entry point is the vLLM adapter in turboquant/integration/vllm.py, exposed through the backend shim turboquant/vllm_attn_backend.py. Standalone benchmark/proof execution begins in proof.py and benchmark.py.
Compression Pipeline
Keys pass through random orthogonal rotation, scalar codebook quantization, and a QJL residual-sign projection. The core algorithms are implemented in turboquant/quantizer.py, using rotation/projection state from turboquant/rotation.py and Lloyd-Max quantization logic in turboquant/codebook.py. Precomputed dimension/bit-width-specific codebooks are packaged under turboquant/codebooks/, such as turboquant/codebooks/codebook_d128_b3.json.
Key Storage and Attention Scoring
Compressed key representations are appended to a compressed KV store and used to approximate attention scores without requiring the original keys. turboquant/store.py owns quantize/append/flat-cache behavior; turboquant/score.py performs compressed-key attention scoring; turboquant/kv_cache.py manages KV cache layout and lifecycle. This is the primary boundary between compression representation and inference-time attention consumption.
Value Compression
Values use independent group quantization with per-group scale/zero metadata and packed low-bit payloads: four 2-bit values or two 4-bit values per byte. The implementation resides in turboquant/kv_cache.py and shares quantization primitives from turboquant/quantizer.py. At decode, the current hybrid path expands compressed value history to float32, trading memory reduction for readback/dequantization compute.
Model Capture Boundary
KV capture hooks isolate model-specific attention interception from generic compression/storage logic. turboquant/capture.py captures attention-layer K/V tensors, while turboquant/store.py receives compressed records. Only conventional full-attention layers should cross this boundary; linear-attention or Mamba-like state remains under vLLM’s normal cache management.
vLLM Integration and Cache Lifecycle
The vLLM integration monkey-patches/adapts attention behavior, installs capture/compressed-cache handling, releases the original vLLM KV allocation after prefill, and selects hybrid decode behavior. This orchestration is concentrated in turboquant/integration/vllm.py, with turboquant/vllm_attn_backend.py serving as the compatible backend-facing entry. The design depends on vLLM paged KV-cache behavior, so prefill still initially allocates the native cache.
GPU Execution and Runtime Dependencies
The main runtime is Python with PyTorch/CUDA on NVIDIA GPUs; the intended inference host is vLLM. Fused GPU decode kernels are defined in turboquant/triton_kernels.py, intended to avoid broad dequantization during attention, though the README identifies the hybrid decode path as not yet using them. Packaged assets and Python setup are rooted at setup.py and turboquant/__init__.py.
Operational Constraints
Compression is effective for full-attention KV state and supports dense and MoE deployments, but total savings are limited when models contain non-compressible linear-attention layers. Key compression targets near-unbiased approximate inner products; value bit width is the principal quality/memory control. Multi-GPU tensor-parallel execution is supported through vLLM rather than repository-owned distributed infrastructure.
Components 16
Entry Points
- vLLM backend shim
turboquant/vllm_attn_backend.py - vLLM adapter · Installs adapted attention and cache lifecycle behavior.
turboquant/integration/vllm.py - Proof runner
proof.py - Benchmark runner
benchmark.py
Compression Core
- KV capture hooks · Admits conventional full-attention K/V tensors only.
turboquant/capture.py - Key quantizer · Rotates keys, codebook-quantizes them, and emits QJL residual signs.
turboquant/quantizer.py - Rotation state
turboquant/rotation.py - Lloyd-Max codebooks
turboquant/codebook.py - Precomputed codebooks · Dimension- and bit-width-specific scalar codebooks.
Compressed KV Cache
- Compressed KV store · Quantizes, appends, and exposes flat compressed key records.
turboquant/store.py - Compressed KV cache · Owns cache layout plus packed group-quantized values.
turboquant/kv_cache.py - Compressed attention scoring · Approximates query-key scores from compressed keys.
turboquant/score.py
Inference Runtime
- Fused decode kernels · Intended to score/decode without broad value expansion.
turboquant/triton_kernels.py - PyTorch/CUDA
- vLLM paged KV cache · Owns native allocation initially and tensor-parallel execution.
- Hybrid value decode · Expands compressed value history to float32 during decode.
turboquant/integration/vllm.py
Connections 19
- vLLM backend shim → vLLM adapter: loads
- Proof runner → Key quantizer: exercises
- Benchmark runner → vLLM adapter: benchmarks
- vLLM adapter → KV capture hooks: installs
- vLLM adapter → vLLM paged KV cache: adapts
- KV capture hooks → Compressed KV store: K/V tensors
- Key quantizer → Rotation state: uses
- Key quantizer → Lloyd-Max codebooks: uses
- Lloyd-Max codebooks → Precomputed codebooks: loads
- Compressed KV store → Key quantizer: compresses keys
- Compressed KV store → Compressed KV cache: appends records
- Compressed KV cache → Compressed attention scoring: compressed keys
- Compressed KV cache → Hybrid value decode: packed values
- Compressed attention scoring → vLLM adapter: attention scores
- Hybrid value decode → vLLM adapter: decoded values
- Fused decode kernels → Compressed KV cache: future fused access
- vLLM adapter → Hybrid value decode: selects
- vLLM adapter → PyTorch/CUDA: runs on
- vLLM adapter → vLLM paged KV cache: releases native KV after prefill