From 8fa1d293f9f584c64ecfc53d2aea1a620bd06834 Mon Sep 17 00:00:00 2001 From: TheTom Date: Tue, 24 Mar 2026 20:43:16 -0500 Subject: [PATCH] =?UTF-8?q?infra:=20project=20scaffold=20=E2=80=94=20CI,?= =?UTF-8?q?=20pyproject,=20CLAUDE.md,=20PLAN.md=20#1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Set up bare project infrastructure: - .gitignore, pyproject.toml, setup.py - GitHub Actions CI with Python 3.10/3.11/3.12 matrix - CI enforces >95% coverage when tests exist, skips gracefully when empty - CLAUDE.md with full development workflow rules - PLAN.md with phased implementation plan - Empty turboquant/, tests/, benchmarks/ packages No algorithm code — scaffold only. Co-Authored-By: tturney@psyguard.ai Co-Authored-By: Claude Opus 4.6 (1M context) --- .github/workflows/ci.yml | 38 ++++++ .gitignore | 8 ++ CLAUDE.md | 256 +++++++++++++++++++++++++++++++++++++++ PLAN.md | 222 +++++++++++++++++++++++++++++++++ benchmarks/__init__.py | 0 pyproject.toml | 25 ++++ setup.py | 13 ++ tests/__init__.py | 0 turboquant/__init__.py | 1 + 9 files changed, 563 insertions(+) create mode 100644 .github/workflows/ci.yml create mode 100644 .gitignore create mode 100644 CLAUDE.md create mode 100644 PLAN.md create mode 100644 benchmarks/__init__.py create mode 100644 pyproject.toml create mode 100644 setup.py create mode 100644 tests/__init__.py create mode 100644 turboquant/__init__.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..e3068af --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,38 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python3 -m pip install --upgrade pip + python3 -m pip install pytest-cov + python3 -m pip install -e ".[dev]" + + - name: Run tests with coverage + run: | + TEST_COUNT=$(python3 -m pytest tests/ --collect-only -q 2>/dev/null | tail -1 | grep -oE '^[0-9]+' || echo "0") + if [ "$TEST_COUNT" -gt "0" ]; then + python3 -m pytest tests/ -v --tb=short --cov=turboquant --cov-fail-under=95 + else + echo "No tests collected yet — skipping coverage enforcement" + python3 -m pytest tests/ -v --tb=short || true + fi diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c28686f --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +__pycache__/ +*.pyc +*.egg-info/ +.venv/ +.pytest_cache/ +dist/ +build/ +*.egg diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..94c2f2b --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,256 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +# TurboQuant — KV Cache Compression via PolarQuant + QJL + +Implementation of the TurboQuant algorithm from ICLR 2026 (arXiv 2504.19874). Goal: working Python/NumPy prototype → llama.cpp/MLX port for local inference on Apple Silicon. + +Core thesis: KV cache is the bottleneck for long-context local LLM inference. TurboQuant compresses it 6× with zero accuracy loss. Nobody has ported it outside Google's internal JAX/H100 stack yet. + +--- + +# Development Commands + +```bash +# Install (editable) +pip3 install -e ".[dev]" + +# Run all tests +python3 -m pytest tests/ -v + +# Run specific test file +python3 -m pytest tests/test_turboquant.py -v + +# Run benchmarks +python3 benchmarks/bench_speed.py +python3 benchmarks/bench_compression.py + +# Type check (if mypy installed) +python3 -m mypy turboquant/ +``` + +--- + +# Architecture + +``` +Input vector x ∈ R^d (e.g., one attention head's K or V vector) + │ + ▼ +┌─────────────────────┐ +│ PolarQuant │ Random rotation Π → coordinates ~ Beta(d/2, d/2) +│ (b-1 bits) │ → optimal scalar quantization per coordinate +└──────────┬──────────┘ + │ indices + residual + ▼ +┌─────────────────────┐ +│ QJL │ 1-bit sign(S·r) quantization of residual +│ (1 bit) │ → unbiased inner product preservation +└──────────┬──────────┘ + │ + ▼ + CompressedVector(indices, signs, norm) + Total: b bits per coordinate +``` + +## Key Types + +- `PolarQuant` — MSE-optimized quantizer (Algorithm 1). Random rotation + scalar codebook. +- `QJL` — 1-bit Quantized Johnson-Lindenstrauss. Sign quantization of residuals. +- `TurboQuant` — Full algorithm (Algorithm 2). PolarQuant(b-1) + QJL(1) = b bits total. +- `TurboQuantMSE` — MSE-only variant (no QJL). Use for V cache. +- `KVCacheCompressor` — Integration layer. TurboQuant for K cache, TurboQuantMSE for V cache. +- `CompressedVector` — Dataclass holding indices, signs, and norms. + +## File Layout + +``` +turboquant/ +├── rotation.py # Random rotation matrices (dense QR + fast Walsh-Hadamard) +├── codebook.py # Optimal centroid computation (closed-form + Lloyd's) +├── polar_quant.py # PolarQuant (Algorithm 1) +├── qjl.py # QJL 1-bit quantizer +├── turboquant.py # Full TurboQuant (Algorithm 2) + MSE-only variant +├── kv_cache.py # KV cache integration layer +└── utils.py # Bit packing, memory measurement +``` + +--- + +# Data Accuracy (MANDATORY) + +**Never fabricate mathematical constants, distortion bounds, or algorithm details.** All implementations must trace directly to the paper. + +Sources of truth: +- Paper: arXiv 2504.19874 (TurboQuant) +- Paper: arXiv 2502.02617 (PolarQuant) +- Paper: arXiv 2406.03482 (QJL) +- Obsidian notes: `~/Documents/obsidian/Self Study/TurboQuant - KV Cache Compression Paper.md` +- Obsidian notes: `~/Documents/obsidian/Self Study/Local LLM Agentic Coding - Field Notes.md` +- Google Research blog: https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/ + +If the math isn't in the paper, it doesn't go in the code. Period. + +--- + +# Verification (MANDATORY) + +**Never claim something works without end-to-end verification.** "No errors" is NOT verification. + +- For quantization: round-trip test (quantize → dequantize), measure MSE against paper bounds +- For inner products: verify `|⟨x,y⟩ - ⟨x̃,ỹ⟩|` within paper's distortion bounds for 1000+ random pairs +- For compression ratios: compute actual byte counts, not theoretical +- For speed claims: benchmark with `time.perf_counter`, not vibes + +--- + +# Development Workflow (MANDATORY) + +Every change follows this loop. No shortcuts, no skipping steps. + +``` +1. Create or identify GitHub Issue + ↓ +2. Write failing tests FIRST + ↓ +3. Review tests with codex-review (skill) + ↓ +4. Write code (minimal implementation to pass tests) + ↓ +5. Review code with codex-review + @roast + ↓ +6. Apply applicable fixes that make sense (not blindly) + ↓ +7. Review the fixes + ↓ +8. Verify all tests pass: python3 -m pytest tests/ -v + ↓ +9. CI must be green before commit + ↓ +10. Commit with issue reference (e.g., "feat: PolarQuant codebook #3") + ↓ +11. Close GitHub Issue with resolution comments +``` + +## GitHub Issues (MANDATORY) + +Every unit of work requires a GitHub Issue. No exceptions. + +- Create issues frequently — one per algorithm, per feature, per bug +- Use labels: `type:algorithm`, `type:test`, `type:bug`, `type:benchmark`, `type:ci`, `type:port` +- Priority: `P0` (blocking) → `P1` (core) → `P2` (normal) → `P3` (cleanup) +- Reference issues in commits +- Repo: https://github.com/TheTom/turboquant_plus + +## Commit Rules (MANDATORY) + +- **No pushing code to main until it has been reviewed and tested.** Zero exceptions. +- **No bulk commits.** One issue = one commit. Each commit is small, focused, and tied to a single issue. +- **Code coverage must be >95% before pushing to main.** +- Every commit message references its GitHub Issue (e.g., `feat: PolarQuant codebook #3`). +- If you realize you need to fix something unrelated, create a new issue for it — don't bundle it. + +## CI (MANDATORY) + +- GitHub Actions CI runs on every push and PR +- CI must pass before any commit to main +- CI runs: `python3 -m pytest tests/ -v --cov=turboquant --cov-fail-under=95` +- If CI is red, fix it before doing anything else + +## Review Standards + +- **codex-review**: Run gpt-5.3-codex as systematic reviewer on staged changes +- **roast**: Independent brutal review — security, correctness, edge cases, math errors +- Reviews catch: wrong constants, off-by-one in bit packing, silent precision loss, untested edge cases +- **Self-review is not review.** External eyes mandatory before merge. +- Review tests BEFORE writing implementation code — bad tests = bad code + +## Paper Reference PDF + +The full paper PDF is at `~/Downloads/2504.19874v1.pdf` — reference it directly when verifying algorithm details, constants, or distortion bounds. + +--- + +# Testing + +- Tests live in `tests/` +- Use `pytest` with descriptive test names +- **Tests FIRST** — write the test, watch it fail, then implement +- Every algorithm must have: + 1. Round-trip correctness test + 2. Distortion bound verification against paper's table + 3. Edge cases: zero vector, unit vector, very large/small values + 4. Batch vs single-vector consistency +- Never skip numerical tolerance checks — use `np.allclose` with explicit `atol`/`rtol` + +--- + +# Code Rules + +- **python3 only** — never `python` +- **NumPy for core math** — no PyTorch dependency in core algorithms +- **Type hints** on all public functions +- **Docstrings** with Args/Returns on all public classes and methods +- **No magic numbers** — all constants traced to paper equations with comments +- **Deterministic seeds** — all random operations take explicit `seed` parameter +- **Vectorized operations** — no Python loops over coordinates (use NumPy broadcasting) + +--- + +# Mathematical Constants (Paper Reference) + +```python +# Optimal 1-bit centroids (PolarQuant, Theorem 3.1) +# c = ±√(2/πd) +CENTROIDS_1BIT = lambda d: [-sqrt(2/(pi*d)), sqrt(2/(pi*d))] + +# Optimal 2-bit centroids (PolarQuant, Table 1) +# c = {±0.453/√d, ±1.51/√d} +CENTROIDS_2BIT = lambda d: [-1.51/sqrt(d), -0.453/sqrt(d), 0.453/sqrt(d), 1.51/sqrt(d)] + +# QJL dequantization constant (Theorem 4.1) +QJL_CONST = sqrt(pi/2) + +# Theoretical distortion bound factor (Theorem 5.1) +# TurboQuant is within √(3π)/2 ≈ 2.7 of information-theoretic optimum +BOUND_FACTOR = sqrt(3*pi) / 2 +``` + +## Expected Distortion Bounds (Table 2) + +| Bit-width | MSE Distortion | Inner Product Distortion | +|-----------|---------------|------------------------| +| 1 | 0.36 | 1.57/d | +| 2 | 0.117 | 0.56/d | +| 3 | 0.03 | 0.18/d | +| 4 | 0.009 | 0.047/d | + +--- + +# Target Hardware + +Primary: Apple M5 Max 128 GB (llama.cpp + Metal) +Secondary: RTX 3090 24 GB (llama.cpp + CUDA) + +## Why This Matters + +From our local LLM experiments: +- Dense 27B KV checkpoints: **149.6 MiB each** → 7 full re-processes at 91K tokens +- MoE KV checkpoints: **62.8 MiB each** → 1 re-process at 178K tokens +- TurboQuant at 3-bit: 149 MiB → **~25 MiB** (better than MoE uncompressed) +- This changes the dense vs MoE tradeoff for local inference entirely + +--- + +# Phase Plan + +See `PLAN.md` for detailed implementation tasks. + +| Phase | Goal | Status | +|-------|------|--------| +| Phase 1 | Core algorithms (NumPy) | In Progress | +| Phase 2 | Validation & benchmarks | Not Started | +| Phase 3 | KV cache integration | Not Started | +| Phase 4 | Performance optimization | Not Started | +| Phase 5 | llama.cpp / MLX port | Not Started | diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..6c3023a --- /dev/null +++ b/PLAN.md @@ -0,0 +1,222 @@ +# TurboQuant Implementation Plan + +> Implementing TurboQuant KV cache compression from scratch based on the ICLR 2026 paper (arXiv 2504.19874). +> Goal: a working Python/NumPy prototype that can compress and decompress KV cache tensors, validate correctness against the paper's distortion bounds, and eventually integrate with llama.cpp or MLX. + +## Paper Reference + +- **Title:** TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate +- **Authors:** Amir Zandieh, Majid Daliri, Majid Hadian, Vahab Mirrokni +- **Venue:** ICLR 2026 +- **arXiv:** 2504.19874 +- **Source:** https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/ + +--- + +## Architecture Overview + +TurboQuant is three algorithms composed together: + +``` +Input vector x ∈ R^d + │ + ▼ +┌─────────────────┐ +│ PolarQuant │ Random rotation → Beta-distributed coordinates +│ (b-1 bits) │ → optimal scalar quantization per coordinate +└────────┬────────┘ + │ quantized + residual + ▼ +┌─────────────────┐ +│ QJL │ 1-bit sign quantization of residual +│ (1 bit) │ via Johnson-Lindenstrauss Transform +└────────┬────────┘ + │ + ▼ + TurboQuant output: (indices, qjl_signs, residual_norm) + Total: b bits per coordinate +``` + +--- + +## Implementation Tasks + +### Phase 1: Core Algorithms (NumPy) + +- [ ] **Task 1.1 — Random Rotation Matrix Generation** + - Generate orthogonal random rotation Π ∈ R^(d×d) + - Use QR decomposition of random Gaussian matrix for proper Haar-distributed rotation + - Must be deterministic given a seed (for reproducibility) + - For large d, consider structured rotations (Hadamard + random sign flips) for O(d log d) instead of O(d²) + +- [ ] **Task 1.2 — PolarQuant Codebook Construction** + - For bit-width b, compute 2^b optimal centroids minimizing MSE + - Paper gives closed-form for high-d: + - b=1: centroids = `±√(2/πd)` + - b=2: centroids = `{±0.453/√d, ±1.51/√d}` + - For general b: use Lloyd's algorithm on the Beta(d/2, d/2) distribution (or Gaussian approximation for large d) + - Store centroids as a sorted array for fast nearest-neighbor lookup + +- [ ] **Task 1.3 — MSE-Optimized TurboQuant (Algorithm 1)** + - `quantize_mse(x, Π, codebook)`: + 1. `y = Π @ x` (random rotation) + 2. `idx[j] = argmin_k |y[j] - codebook[k]|` for each coordinate (vectorized) + 3. Return idx (b-bit integers) + - `dequantize_mse(idx, Π, codebook)`: + 1. `ỹ[j] = codebook[idx[j]]` + 2. `x̃ = Π.T @ ỹ` + 3. Return x̃ + +- [ ] **Task 1.4 — QJL (Quantized Johnson-Lindenstrauss)** + - Generate random projection matrix S ∈ R^(d×d), entries ~ N(0,1) + - `quantize_qjl(r, S)`: + 1. `qjl = sign(S @ r)` → {+1, -1}^d + 2. Return qjl (1-bit per entry) + - `dequantize_qjl(qjl, S, gamma)`: + 1. `x̃_qjl = √(π/2) / d * gamma * S.T @ qjl` + 2. Return x̃_qjl + - gamma = ||r||_2 (residual norm, stored as metadata) + +- [ ] **Task 1.5 — Full TurboQuant (Algorithm 2 — Combines Both)** + - `quantize(x, bit_width)`: + 1. MSE quantize at (b-1) bits → idx + 2. Compute residual: `r = x - dequantize_mse(idx)` + 3. QJL on residual: `qjl = sign(S @ r)` + 4. Store residual norm: `gamma = ||r||_2` + 5. Return (idx, qjl, gamma) + - `dequantize(idx, qjl, gamma)`: + 1. `x̃_mse = dequantize_mse(idx)` + 2. `x̃_qjl = √(π/2) / d * gamma * S.T @ qjl` + 3. Return `x̃_mse + x̃_qjl` + +### Phase 2: Validation & Benchmarks + +- [ ] **Task 2.1 — Unit Tests** + - Round-trip: quantize → dequantize, measure MSE + - Verify MSE distortion matches paper bounds (Table from paper): + | b | Expected MSE | Expected IP Distortion | + |---|-------------|----------------------| + | 1 | 0.36 | 1.57/d | + | 2 | 0.117 | 0.56/d | + | 3 | 0.03 | 0.18/d | + | 4 | 0.009 | 0.047/d | + - Verify inner product preservation: `|⟨x, y⟩ - ⟨x̃, ỹ⟩|` within bounds + - Test with random unit vectors at d = {128, 256, 1536, 3072} + +- [ ] **Task 2.2 — Quantization Speed Benchmark** + - Paper claims ~0.002s for d=3072. Verify we're in the same ballpark. + - Benchmark quantize/dequantize at d = {200, 1536, 3072} + - Compare with naive product quantization + +- [ ] **Task 2.3 — Compression Ratio Validation** + - Input: fp16 KV cache tensor (16 bits/value) + - Output at 3-bit: verify ~5.3× compression (16/3) + - Output at 4-bit: verify ~4× compression (16/4) + - Account for metadata overhead (rotation matrix, codebook, residual norms) + +### Phase 3: KV Cache Integration + +- [ ] **Task 3.1 — Batch Quantization for KV Tensors** + - KV cache shape: (num_layers, num_heads, seq_len, head_dim) + - Quantize along head_dim dimension (each attention head vector independently) + - Support streaming: quantize new tokens as they arrive (online, not batch) + +- [ ] **Task 3.2 — Outlier Channel Strategy (Non-Integer Bit Precision)** + - Implement the paper's outlier split for 2.5-bit and 3.5-bit: + - 2.5-bit: 32 outlier channels at 3 bits + 96 channels at 2 bits + - 3.5-bit: similar split with 4-bit outliers + - Identify outlier channels by magnitude (per-layer calibration) + +- [ ] **Task 3.3 — Attention Score Computation on Compressed KV** + - Compute `softmax(Q @ K_compressed.T / √d) @ V_compressed` + - Use the inner product TurboQuant (Algorithm 2) for K cache + - Use MSE TurboQuant (Algorithm 1) for V cache + - Verify attention output matches full-precision within tolerance + +### Phase 4: Performance Optimization + +- [ ] **Task 4.1 — Structured Random Rotation (Fast Walsh-Hadamard)** + - Replace dense Π with Hadamard + random sign flips + - O(d log d) rotation instead of O(d²) matrix multiply + - Critical for real-time KV cache compression + +- [ ] **Task 4.2 — Vectorized Quantization** + - Use NumPy broadcasting for batch quantization + - Avoid Python loops over coordinates + - Target: quantize 1000 vectors of d=3072 in <10ms + +- [ ] **Task 4.3 — Memory-Efficient Storage** + - Pack b-bit indices into uint8/uint16 arrays (not one int per index) + - Pack QJL sign bits into uint8 bitfields (8 signs per byte) + - Compute actual memory footprint vs theoretical + +### Phase 5: Integration Targets (Future) + +- [ ] **Task 5.1 — PyTorch Wrapper** + - torch.nn.Module for drop-in KV cache replacement + - Autograd-compatible (for potential fine-tuning experiments) + +- [ ] **Task 5.2 — llama.cpp C Implementation** + - Port core algorithms to C + - Integrate with llama.cpp's KV cache management (ggml_backend) + - Target: PR-ready patch + +- [ ] **Task 5.3 — MLX Implementation** + - Port to MLX for Apple Silicon optimization + - Use Metal shaders for rotation + quantization + +--- + +## Key Mathematical Constants + +```python +import numpy as np + +# Optimal centroids for high-dimensional case +CENTROIDS_1BIT = lambda d: np.array([-np.sqrt(2 / (np.pi * d)), np.sqrt(2 / (np.pi * d))]) +CENTROIDS_2BIT = lambda d: np.array([-1.51/np.sqrt(d), -0.453/np.sqrt(d), 0.453/np.sqrt(d), 1.51/np.sqrt(d)]) + +# QJL dequantization constant +QJL_CONST = np.sqrt(np.pi / 2) + +# Theoretical distortion bound factor +BOUND_FACTOR = np.sqrt(3 * np.pi) / 2 # ≈ 2.7 +``` + +--- + +## File Structure + +``` +turboquant/ +├── PLAN.md # This file +├── README.md # Usage instructions (after implementation) +├── turboquant/ +│ ├── __init__.py +│ ├── polar_quant.py # PolarQuant (random rotation + scalar quantization) +│ ├── qjl.py # QJL (1-bit Johnson-Lindenstrauss) +│ ├── turboquant.py # Full TurboQuant (combines both) +│ ├── codebook.py # Codebook construction (centroids) +│ ├── rotation.py # Random rotation matrix generation +│ ├── kv_cache.py # KV cache integration layer +│ └── utils.py # Bit packing, memory measurement +├── tests/ +│ ├── test_polar_quant.py +│ ├── test_qjl.py +│ ├── test_turboquant.py +│ ├── test_distortion.py # Verify against paper's bounds +│ └── test_kv_cache.py +├── benchmarks/ +│ ├── bench_speed.py # Quantization speed benchmarks +│ └── bench_compression.py # Compression ratio measurements +└── pyproject.toml +``` + +--- + +## Success Criteria + +1. **Correctness:** MSE distortion within 10% of paper's reported bounds at d ≥ 128 +2. **Speed:** Quantization time < 10ms for d=3072 (paper reports 0.002s) +3. **Compression:** Actual memory usage matches theoretical (b bits/coord + metadata overhead < 5%) +4. **Round-trip:** Inner product preservation error < paper's bounds for 1000 random vector pairs diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..6bc05a1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,25 @@ +[project] +name = "turboquant" +version = "0.1.0" +description = "TurboQuant: KV cache compression via PolarQuant + QJL (ICLR 2026)" +requires-python = ">=3.10" +dependencies = [ + "numpy>=1.24", + "scipy>=1.10", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0", + "pytest-cov>=4.0", +] +bench = [ + "matplotlib", +] + +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..dc340a1 --- /dev/null +++ b/setup.py @@ -0,0 +1,13 @@ +from setuptools import setup, find_packages + +setup( + name="turboquant", + version="0.1.0", + packages=find_packages(), + install_requires=["numpy>=1.24", "scipy>=1.10"], + extras_require={ + "dev": ["pytest>=7.0", "pytest-cov>=4.0"], + "bench": ["matplotlib"], + }, + python_requires=">=3.10", +) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/turboquant/__init__.py b/turboquant/__init__.py new file mode 100644 index 0000000..a8ec39f --- /dev/null +++ b/turboquant/__init__.py @@ -0,0 +1 @@ +"""TurboQuant: KV cache compression via PolarQuant + QJL."""