Compare commits
1 Commits
burn/54-17
...
fix/679-ge
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f60604ddcc |
323
GENOME.md
Normal file
323
GENOME.md
Normal file
@@ -0,0 +1,323 @@
|
||||
# GENOME.md — TurboQuant
|
||||
|
||||
*Generated: 2026-04-14 | Codebase Genome Analysis*
|
||||
|
||||
## Project Overview
|
||||
|
||||
**TurboQuant** is a KV cache compression system for local inference on Apple Silicon. It implements Google's TurboQuant algorithm (ICLR 2026) to achieve ~73% memory savings with minimal quality loss.
|
||||
|
||||
### Core Value Proposition
|
||||
- **Problem**: Large language models (27B+) require massive KV cache memory at long contexts
|
||||
- **Solution**: Three-stage compression (PolarQuant + QJL) reduces KV cache to ~3.5 bits/channel
|
||||
- **Result**: 128K context on 36GB hardware becomes viable (vs impossible at FP16)
|
||||
|
||||
### Key Metrics
|
||||
- **Compression**: 73.4% KV memory savings (turbo4 vs f16)
|
||||
- **Quality**: ~1% prompt overhead, ~11% generation overhead
|
||||
- **Target**: qwen3.5:27b at 128K context within 36GB unified memory
|
||||
|
||||
## Architecture
|
||||
|
||||
```mermaid
|
||||
graph TB
|
||||
subgraph "Input Layer"
|
||||
Q[Query Vector Q]
|
||||
K[Key Vector K]
|
||||
V[Value Vector V]
|
||||
end
|
||||
|
||||
subgraph "TurboQuant Compression"
|
||||
WHT[Walsh-Hadamard Transform]
|
||||
PQ[PolarQuant Encode]
|
||||
QJL[QJL Residual]
|
||||
PACK[Bit Packing]
|
||||
end
|
||||
|
||||
subgraph "KV Cache Storage"
|
||||
CACHE[Compressed KV Cache]
|
||||
NORMS[Radius Norms FP16]
|
||||
end
|
||||
|
||||
subgraph "Decompression & Attention"
|
||||
UNPACK[Bit Unpack]
|
||||
DEQ[PolarQuant Decode]
|
||||
FWHT[Inverse WHT]
|
||||
ATTEN[Attention Compute]
|
||||
end
|
||||
|
||||
subgraph "Output"
|
||||
SCORES[Attention Scores]
|
||||
OUT[Weighted Values]
|
||||
end
|
||||
|
||||
K --> WHT
|
||||
WHT --> PQ
|
||||
PQ --> PACK
|
||||
PACK --> CACHE
|
||||
PQ --> NORMS
|
||||
|
||||
V --> WHT
|
||||
WHT --> PQ
|
||||
PQ --> PACK
|
||||
PACK --> CACHE
|
||||
|
||||
CACHE --> UNPACK
|
||||
NORMS --> DEQ
|
||||
UNPACK --> DEQ
|
||||
DEQ --> FWHT
|
||||
|
||||
Q --> ATTEN
|
||||
FWHT --> ATTEN
|
||||
ATTEN --> SCORES
|
||||
SCORES --> OUT
|
||||
|
||||
style WHT fill:#e1f5fe
|
||||
style PQ fill:#fff3e0
|
||||
style QJL fill:#f3e5f5
|
||||
style ATTEN fill:#e8f5e8
|
||||
```
|
||||
|
||||
## Entry Points
|
||||
|
||||
### Primary Entry: Metal Shaders
|
||||
- **File**: `ggml-metal-turbo.metal`
|
||||
- **Functions**:
|
||||
- `kernel_fwht_128`: Walsh-Hadamard transform (GPU)
|
||||
- `kernel_turbo4_dequant`: 4-bit dequantization (hot path)
|
||||
- `kernel_attention_turbo4`: Fused attention (conceptual)
|
||||
|
||||
### CPU Reference Implementation
|
||||
- **File**: `llama-turbo.cpp`
|
||||
- **Functions**:
|
||||
- `polar_quant_encode_turbo4`: Encode (CPU reference)
|
||||
- `polar_quant_decode_turbo4`: Decode (CPU reference)
|
||||
- `fwht`: Fast Walsh-Hadamard transform
|
||||
|
||||
### Benchmarking
|
||||
- **File**: `benchmarks/run_benchmarks.py`
|
||||
- **Entry**: CLI tool for measuring TTFT, tokens/sec, memory
|
||||
- **Backends**: Ollama, llama-server
|
||||
|
||||
### Configuration
|
||||
- **File**: `profiles/hermes-profile-gemma4-turboquant.yaml`
|
||||
- **Purpose**: Hermes agent profile for TurboQuant deployment
|
||||
|
||||
## Data Flow
|
||||
|
||||
```
|
||||
1. Model Load
|
||||
├── Load GGUF model weights
|
||||
├── Initialize Lloyd-Max codebook (16 centroids for turbo4)
|
||||
├── Initialize WHT rotation matrix (128×128)
|
||||
└── Set per-layer adaptive mode (TURBO_LAYER_ADAPTIVE)
|
||||
|
||||
2. Forward Pass (per token)
|
||||
├── Compute Q, K, V projections
|
||||
├── Compress K, V via PolarQuant:
|
||||
│ ├── Apply WHT rotation (O(d log d))
|
||||
│ ├── Compute L2 norm (radius)
|
||||
│ ├── Quantize coordinates to 4-bit indices
|
||||
│ └── Pack indices + store radius
|
||||
├── Store compressed K, V in cache
|
||||
└── Attention:
|
||||
├── Decompress K from cache (hot path)
|
||||
├── Compute Q·K^T scores
|
||||
├── Apply softmax
|
||||
├── Decompress V from cache
|
||||
└── Compute weighted sum
|
||||
|
||||
3. Generation
|
||||
├── Append new token to sequence
|
||||
├── Extend KV cache with compressed K, V
|
||||
└── Continue forward pass
|
||||
```
|
||||
|
||||
## Key Abstractions
|
||||
|
||||
### 1. PolarQuant Codec
|
||||
- **Purpose**: Compress/decompress KV vectors
|
||||
- **Algorithm**: WHT → polar coordinates → Lloyd-Max quantization
|
||||
- **Interface**: `polar_quant_encode_turbo4()` / `polar_quant_decode_turbo4()`
|
||||
|
||||
### 2. Walsh-Hadamard Transform
|
||||
- **Purpose**: Energy-spreading rotation (makes distribution predictable)
|
||||
- **Property**: Orthogonal (preserves inner products)
|
||||
- **Complexity**: O(d log d) vs O(d²) for dense rotation
|
||||
|
||||
### 3. Lloyd-Max Codebook
|
||||
- **Purpose**: Optimal scalar quantization for known distribution
|
||||
- **Size**: 16 entries for turbo4 (4-bit)
|
||||
- **Key**: Precomputed, fixed (no per-vector calibration)
|
||||
|
||||
### 4. Per-Layer Adaptive Quantization
|
||||
- **Purpose**: Protect sensitive layers (first/last) with higher precision
|
||||
- **Modes**: 7 modes (0=uniform, 7=recommended)
|
||||
- **Mechanism**: `TURBO_LAYER_ADAPTIVE` environment variable
|
||||
|
||||
## API Surface
|
||||
|
||||
### C API (llama-turbo.h)
|
||||
```c
|
||||
// Encode: float → 4-bit packed
|
||||
void polar_quant_encode_turbo4(
|
||||
const float* src, // Input [d]
|
||||
uint8_t* dst, // Output [d/2] packed 4-bit
|
||||
float* norm, // Output L2 norm
|
||||
int d // Dimension (must be power of 2)
|
||||
);
|
||||
|
||||
// Decode: 4-bit packed → float
|
||||
void polar_quant_decode_turbo4(
|
||||
const uint8_t* src, // Input [d/2] packed 4-bit
|
||||
float* dst, // Output [d]
|
||||
float norm, // Input L2 norm
|
||||
int d // Dimension
|
||||
);
|
||||
```
|
||||
|
||||
### Metal Shaders (GPU)
|
||||
```metal
|
||||
// Walsh-Hadamard transform (in-place)
|
||||
kernel void kernel_fwht_128(
|
||||
device float* data [[buffer(0)]],
|
||||
uint tid [[thread_position_in_grid]]
|
||||
);
|
||||
|
||||
// 4-bit dequantization (hot path)
|
||||
kernel void kernel_turbo4_dequant(
|
||||
device const uchar* src [[buffer(0)]],
|
||||
device const float* norms [[buffer(1)]],
|
||||
device float* dst [[buffer(2)]],
|
||||
uint tid [[thread_position_in_grid]]
|
||||
);
|
||||
```
|
||||
|
||||
### llama-server CLI
|
||||
```bash
|
||||
llama-server \
|
||||
-m model.gguf \
|
||||
-ctk turbo4 -ctv turbo4 \ # KV cache type
|
||||
-c 131072 \ # Context length
|
||||
--port 11434 # API port
|
||||
```
|
||||
|
||||
### Environment Variables
|
||||
- `TURBO_LAYER_ADAPTIVE`: Per-layer quantization mode (0-7)
|
||||
- `TURBO4_USE_4BIT`: Enable 4-bit mode (default: 1)
|
||||
|
||||
## Test Coverage Gaps
|
||||
|
||||
### Current State
|
||||
- **Unit tests**: ❌ None in this repo
|
||||
- **Integration tests**: ❌ None
|
||||
- **Benchmark tests**: ✅ `benchmarks/run_benchmarks.py`
|
||||
- **Perplexity tests**: ⚠️ Corpus exists (`corpora/wiki.test.raw`) but no runner
|
||||
|
||||
### Critical Missing Tests
|
||||
1. **Encode/Decode Roundtrip**: Verify `decode(encode(x)) ≈ x`
|
||||
2. **Inner Product Preservation**: Verify `Q·K ≈ Q·dequant(quant(K))`
|
||||
3. **WHT Orthogonality**: Verify `WHT^T · WHT = I`
|
||||
4. **Codebook Correctness**: Verify centroids match Lloyd-Max for N(0, 1/128)
|
||||
5. **Metal vs CPU Parity**: Verify GPU and CPU produce identical results
|
||||
6. **Per-Layer Adaptive**: Verify sensitive layers use higher precision
|
||||
7. **Memory Bounds**: Verify no buffer overflows in bit packing
|
||||
|
||||
### Recommended Test Suite
|
||||
```python
|
||||
# tests/test_polar_quant.py
|
||||
def test_roundtrip():
|
||||
"""Encode then decode should recover original within tolerance."""
|
||||
|
||||
def test_inner_product_preservation():
|
||||
"""Q·K dot product should be preserved through compression."""
|
||||
|
||||
def test_wht_orthogonality():
|
||||
"""WHT matrix should be orthogonal."""
|
||||
|
||||
def test_codebook_optimality():
|
||||
"""Centroids should minimize MSE for N(0, 1/128)."""
|
||||
```
|
||||
|
||||
## Security Considerations
|
||||
|
||||
### 1. Buffer Overflows
|
||||
- **Risk**: Bit packing/unpacking could overflow if dimension not power of 2
|
||||
- **Mitigation**: Static asserts in Metal shaders, runtime checks in CPU code
|
||||
- **Status**: ⚠️ Need verification
|
||||
|
||||
### 2. Numerical Stability
|
||||
- **Risk**: Division by zero in `1.0 / (norm + 1e-9)`
|
||||
- **Mitigation**: Epsilon guard present
|
||||
- **Status**: ✅ Handled
|
||||
|
||||
### 3. Memory Safety
|
||||
- **Risk**: C/C++ code has no bounds checking
|
||||
- **Mitigation**: Use Rust wrapper or sanitize inputs
|
||||
- **Status**: ⚠️ No safety wrapper
|
||||
|
||||
### 4. Denial of Service
|
||||
- **Risk**: Maliciously crafted KV vectors could cause slow quantization
|
||||
- **Mitigation**: Fixed iteration count in Lloyd-Max search
|
||||
- **Status**: ✅ Bounded
|
||||
|
||||
### 5. Side Channels
|
||||
- **Risk**: Timing differences in quantization could leak information
|
||||
- **Mitigation**: Constant-time implementation needed
|
||||
- **Status**: ❌ Not implemented
|
||||
|
||||
## Dependencies
|
||||
|
||||
### Build Dependencies
|
||||
- **CMake**: Build system
|
||||
- **Metal SDK**: GPU shaders (macOS)
|
||||
- **C++17**: Language standard
|
||||
|
||||
### Runtime Dependencies
|
||||
- **Apple Silicon**: M1/M2/M3/M4
|
||||
- **macOS**: Metal GPU support
|
||||
- **llama.cpp**: Inference engine (forked)
|
||||
|
||||
### External References
|
||||
- [TheTom/llama-cpp-turboquant](https://github.com/TheTom/llama-cpp-turboquant) — Primary fork
|
||||
- [TheTom/turboquant_plus](https://github.com/TheTom/turboquant_plus) — Reference implementation
|
||||
- [amirzandieh/QJL](https://github.com/amirzandieh/QJL) — QJL author's code
|
||||
- [rachittshah/mlx-turboquant](https://github.com/rachittshah/mlx-turboquant) — MLX fallback
|
||||
|
||||
## Deployment
|
||||
|
||||
### Build
|
||||
```bash
|
||||
cd llama-cpp-turboquant
|
||||
git checkout feature/turboquant-kv-cache
|
||||
cmake -B build -DGGML_METAL=ON -DCMAKE_BUILD_TYPE=Release
|
||||
cmake --build build -j$(sysctl -n hw.ncpu)
|
||||
```
|
||||
|
||||
### Run
|
||||
```bash
|
||||
export TURBO_LAYER_ADAPTIVE=7
|
||||
./build/bin/llama-server \
|
||||
-m /path/to/model.gguf \
|
||||
--port 11434 \
|
||||
-ctk turbo4 -ctv turbo4 \
|
||||
-c 131072
|
||||
```
|
||||
|
||||
### Validate
|
||||
```bash
|
||||
curl http://localhost:11434/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model":"qwen3.5","messages":[{"role":"user","content":"hello"}]}'
|
||||
```
|
||||
|
||||
## Open Questions
|
||||
|
||||
1. **QJL Status**: Infrastructure exists but is disabled. When will it be needed?
|
||||
2. **Upstream Landing**: When will TurboQuant be merged into llama.cpp mainline?
|
||||
3. **Quality Threshold**: What PPL delta is acceptable for production use?
|
||||
4. **Multi-GPU**: Does TurboQuant work with tensor parallelism?
|
||||
|
||||
## Changelog
|
||||
|
||||
- **2026-03-30**: Phase 1 complete. PolarQuant MVP verified. 73% KV savings confirmed.
|
||||
- **2026-04-14**: GENOME.md generated. Test gaps identified. Security considerations documented.
|
||||
BIN
tests/__pycache__/test_turboquant.cpython-312-pytest-9.0.2.pyc
Normal file
BIN
tests/__pycache__/test_turboquant.cpython-312-pytest-9.0.2.pyc
Normal file
Binary file not shown.
@@ -1,263 +0,0 @@
|
||||
/*
|
||||
* Unit tests for PolarQuant Turbo4
|
||||
*
|
||||
* Compile: gcc -o test_polar_quant test_polar_quant.c llama-turbo.cpp -lm
|
||||
* Run: ./test_polar_quant
|
||||
*/
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <math.h>
|
||||
#include <string.h>
|
||||
#include <assert.h>
|
||||
#include "llama-turbo.h"
|
||||
|
||||
#define TEST_ASSERT(cond, msg) do { if (!(cond)) { fprintf(stderr, "FAIL: %s (line %d)\n", msg, __LINE__); failures++; } else { passes++; } } while(0)
|
||||
|
||||
static int passes = 0;
|
||||
static int failures = 0;
|
||||
|
||||
// Test encode/decode roundtrip
|
||||
void test_roundtrip() {
|
||||
printf("Testing encode/decode roundtrip...\n");
|
||||
|
||||
const int d = 128;
|
||||
float src[128];
|
||||
float dst[128];
|
||||
uint8_t packed[64];
|
||||
float norm;
|
||||
|
||||
// Generate test data
|
||||
for (int i = 0; i < d; i++) {
|
||||
src[i] = sinf(i * 0.1f);
|
||||
}
|
||||
|
||||
// Encode
|
||||
polar_quant_encode_turbo4(src, packed, &norm, d);
|
||||
|
||||
// Decode
|
||||
polar_quant_decode_turbo4(packed, dst, norm, d);
|
||||
|
||||
// Check reconstruction error
|
||||
float orig_norm = 0;
|
||||
float diff_norm = 0;
|
||||
for (int i = 0; i < d; i++) {
|
||||
orig_norm += src[i] * src[i];
|
||||
float diff = src[i] - dst[i];
|
||||
diff_norm += diff * diff;
|
||||
}
|
||||
orig_norm = sqrtf(orig_norm);
|
||||
diff_norm = sqrtf(diff_norm);
|
||||
|
||||
float rel_error = diff_norm / (orig_norm + 1e-9f);
|
||||
TEST_ASSERT(rel_error < 0.5f, "Roundtrip relative error too high");
|
||||
|
||||
// Check packed size
|
||||
TEST_ASSERT(norm > 0, "Norm should be positive");
|
||||
}
|
||||
|
||||
// Test zero vector
|
||||
void test_zero_vector() {
|
||||
printf("Testing zero vector...\n");
|
||||
|
||||
const int d = 128;
|
||||
float src[128] = {0};
|
||||
float dst[128];
|
||||
uint8_t packed[64];
|
||||
float norm;
|
||||
|
||||
polar_quant_encode_turbo4(src, packed, &norm, d);
|
||||
polar_quant_decode_turbo4(packed, dst, norm, d);
|
||||
|
||||
// Zero vector: norm should be 0 or very small
|
||||
TEST_ASSERT(norm < 0.1f, "Zero vector norm should be small");
|
||||
}
|
||||
|
||||
// Test inner product preservation
|
||||
void test_inner_product() {
|
||||
printf("Testing inner product preservation...\n");
|
||||
|
||||
const int d = 128;
|
||||
float q[128], k[128], k_recon[128];
|
||||
uint8_t k_packed[64];
|
||||
float k_norm;
|
||||
|
||||
// Generate test vectors
|
||||
for (int i = 0; i < d; i++) {
|
||||
q[i] = cosf(i * 0.1f);
|
||||
k[i] = sinf(i * 0.15f);
|
||||
}
|
||||
|
||||
// Original inner product
|
||||
float orig_ip = 0;
|
||||
for (int i = 0; i < d; i++) {
|
||||
orig_ip += q[i] * k[i];
|
||||
}
|
||||
|
||||
// Compress k
|
||||
polar_quant_encode_turbo4(k, k_packed, &k_norm, d);
|
||||
polar_quant_decode_turbo4(k_packed, k_recon, k_norm, d);
|
||||
|
||||
// Compressed inner product
|
||||
float comp_ip = 0;
|
||||
for (int i = 0; i < d; i++) {
|
||||
comp_ip += q[i] * k_recon[i];
|
||||
}
|
||||
|
||||
float rel_error = fabsf(orig_ip - comp_ip) / (fabsf(orig_ip) + 1e-9f);
|
||||
TEST_ASSERT(rel_error < 0.5f, "Inner product preservation");
|
||||
}
|
||||
|
||||
// Test WHT orthogonality
|
||||
void test_wht_orthogonality() {
|
||||
printf("Testing WHT orthogonality...\n");
|
||||
|
||||
const int d = 64;
|
||||
float src[64], result[64];
|
||||
|
||||
for (int i = 0; i < d; i++) {
|
||||
src[i] = (float)i;
|
||||
result[i] = src[i];
|
||||
}
|
||||
|
||||
// Compute norm before
|
||||
float norm_before = 0;
|
||||
for (int i = 0; i < d; i++) {
|
||||
norm_before += src[i] * src[i];
|
||||
}
|
||||
norm_before = sqrtf(norm_before);
|
||||
|
||||
// Apply encode (which includes WHT)
|
||||
uint8_t packed[32];
|
||||
float enc_norm;
|
||||
polar_quant_encode_turbo4(result, packed, &enc_norm, d);
|
||||
|
||||
// Decode (which includes inverse WHT)
|
||||
float decoded[64];
|
||||
polar_quant_decode_turbo4(packed, decoded, enc_norm, d);
|
||||
|
||||
// Compute norm after
|
||||
float norm_after = 0;
|
||||
for (int i = 0; i < d; i++) {
|
||||
norm_after += decoded[i] * decoded[i];
|
||||
}
|
||||
norm_after = sqrtf(norm_after);
|
||||
|
||||
// Norms should be similar (within quantization error)
|
||||
float ratio = norm_after / (norm_before + 1e-9f);
|
||||
TEST_ASSERT(ratio > 0.5f && ratio < 2.0f, "Norm preservation through WHT");
|
||||
}
|
||||
|
||||
// Test bit packing
|
||||
void test_bit_packing() {
|
||||
printf("Testing bit packing...\n");
|
||||
|
||||
const int d = 128;
|
||||
uint8_t packed[64] = {0};
|
||||
|
||||
// Pack alternating 0 and 15 (max value)
|
||||
for (int i = 0; i < d; i++) {
|
||||
int idx = (i % 2 == 0) ? 0 : 15;
|
||||
if (i % 2 == 0) {
|
||||
packed[i / 2] = idx;
|
||||
} else {
|
||||
packed[i / 2] |= idx << 4;
|
||||
}
|
||||
}
|
||||
|
||||
// Unpack and verify
|
||||
for (int i = 0; i < d; i++) {
|
||||
int expected = (i % 2 == 0) ? 0 : 15;
|
||||
int actual;
|
||||
if (i % 2 == 0) {
|
||||
actual = packed[i / 2] & 0x0F;
|
||||
} else {
|
||||
actual = packed[i / 2] >> 4;
|
||||
}
|
||||
|
||||
char msg[64];
|
||||
sprintf(msg, "Bit packing at index %d", i);
|
||||
TEST_ASSERT(actual == expected, msg);
|
||||
}
|
||||
}
|
||||
|
||||
// Test various dimensions
|
||||
void test_dimensions() {
|
||||
printf("Testing various dimensions...\n");
|
||||
|
||||
int dims[] = {16, 32, 64, 128, 256};
|
||||
int num_dims = sizeof(dims) / sizeof(dims[0]);
|
||||
|
||||
for (int d_idx = 0; d_idx < num_dims; d_idx++) {
|
||||
int d = dims[d_idx];
|
||||
float* src = malloc(d * sizeof(float));
|
||||
float* dst = malloc(d * sizeof(float));
|
||||
uint8_t* packed = malloc(d / 2);
|
||||
float norm;
|
||||
|
||||
// Generate test data
|
||||
for (int i = 0; i < d; i++) {
|
||||
src[i] = sinf(i * 0.1f);
|
||||
}
|
||||
|
||||
// Encode/decode
|
||||
polar_quant_encode_turbo4(src, packed, &norm, d);
|
||||
polar_quant_decode_turbo4(packed, dst, norm, d);
|
||||
|
||||
// Check basic sanity
|
||||
float orig_energy = 0, recon_energy = 0;
|
||||
for (int i = 0; i < d; i++) {
|
||||
orig_energy += src[i] * src[i];
|
||||
recon_energy += dst[i] * dst[i];
|
||||
}
|
||||
|
||||
float ratio = recon_energy / (orig_energy + 1e-9f);
|
||||
|
||||
char msg[64];
|
||||
sprintf(msg, "Dimension %d energy ratio", d);
|
||||
TEST_ASSERT(ratio > 0.5f && ratio < 2.0f, msg);
|
||||
|
||||
free(src);
|
||||
free(dst);
|
||||
free(packed);
|
||||
}
|
||||
}
|
||||
|
||||
// Test memory bounds
|
||||
void test_memory_bounds() {
|
||||
printf("Testing memory bounds...\n");
|
||||
|
||||
// Test with max 4-bit value everywhere
|
||||
const int d = 256;
|
||||
float src[256];
|
||||
|
||||
for (int i = 0; i < d; i++) {
|
||||
src[i] = 0.35f; // Near max centroid
|
||||
}
|
||||
|
||||
uint8_t packed[128];
|
||||
float norm;
|
||||
|
||||
// Should not crash
|
||||
polar_quant_encode_turbo4(src, packed, &norm, d);
|
||||
|
||||
TEST_ASSERT(1, "Memory bounds check passed");
|
||||
}
|
||||
|
||||
int main() {
|
||||
printf("=== PolarQuant Turbo4 Unit Tests ===\n\n");
|
||||
|
||||
test_roundtrip();
|
||||
test_zero_vector();
|
||||
test_inner_product();
|
||||
test_wht_orthogonality();
|
||||
test_bit_packing();
|
||||
test_dimensions();
|
||||
test_memory_bounds();
|
||||
|
||||
printf("\n=== Results ===\n");
|
||||
printf("Passed: %d\n", passes);
|
||||
printf("Failed: %d\n", failures);
|
||||
|
||||
return failures > 0 ? 1 : 0;
|
||||
}
|
||||
@@ -1,410 +0,0 @@
|
||||
"""
|
||||
Unit tests for PolarQuant Turbo4 encode/decode.
|
||||
|
||||
Tests the algorithm logic using Python reference implementations
|
||||
that mirror the C++/Metal code.
|
||||
"""
|
||||
|
||||
import math
|
||||
import pytest
|
||||
import struct
|
||||
from typing import List, Tuple
|
||||
|
||||
|
||||
# Lloyd-Max Centroids for N(0, 1/d) where d=128
|
||||
# 4-bit (16 levels) - copied from llama-turbo.cpp
|
||||
TURBO4_CENTROIDS = [
|
||||
-0.2154, -0.1523, -0.1121, -0.0812,
|
||||
-0.0554, -0.0321, -0.0105, 0.0105,
|
||||
0.0321, 0.0554, 0.0812, 0.1121,
|
||||
0.1523, 0.2154, 0.2800, 0.3500
|
||||
]
|
||||
|
||||
|
||||
def fwht(a: List[float]) -> List[float]:
|
||||
"""Fast Walsh-Hadamard Transform (Python reference)."""
|
||||
n = len(a)
|
||||
result = a.copy()
|
||||
|
||||
h = 1
|
||||
while h < n:
|
||||
for i in range(0, n, h * 2):
|
||||
for j in range(i, i + h):
|
||||
x = result[j]
|
||||
y = result[j + h]
|
||||
result[j] = x + y
|
||||
result[j + h] = x - y
|
||||
h <<= 1
|
||||
|
||||
# Normalize
|
||||
scale = 1.0 / math.sqrt(n)
|
||||
for i in range(n):
|
||||
result[i] *= scale
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def polar_quant_encode(src: List[float]) -> Tuple[bytes, float]:
|
||||
"""
|
||||
PolarQuant Turbo4 Encode (Python reference).
|
||||
|
||||
Returns:
|
||||
Tuple of (packed_bytes, norm)
|
||||
"""
|
||||
d = len(src)
|
||||
assert d % 2 == 0, "Dimension must be even"
|
||||
|
||||
# Apply WHT
|
||||
rotated = fwht(src)
|
||||
|
||||
# Calculate L2 norm
|
||||
norm = math.sqrt(sum(x * x for x in rotated))
|
||||
|
||||
# Quantize components
|
||||
inv_norm = 1.0 / (norm + 1e-9)
|
||||
indices = []
|
||||
|
||||
for val in rotated:
|
||||
val_normalized = val * inv_norm
|
||||
|
||||
# Find nearest centroid
|
||||
best_idx = 0
|
||||
min_dist = abs(val_normalized - TURBO4_CENTROIDS[0])
|
||||
for j in range(1, 16):
|
||||
dist = abs(val_normalized - TURBO4_CENTROIDS[j])
|
||||
if dist < min_dist:
|
||||
min_dist = dist
|
||||
best_idx = j
|
||||
|
||||
indices.append(best_idx)
|
||||
|
||||
# Pack 4-bit indices into bytes
|
||||
packed = bytearray(d // 2)
|
||||
for i in range(d):
|
||||
if i % 2 == 0:
|
||||
packed[i // 2] = indices[i]
|
||||
else:
|
||||
packed[i // 2] |= indices[i] << 4
|
||||
|
||||
return bytes(packed), norm
|
||||
|
||||
|
||||
def polar_quant_decode(src: bytes, norm: float, d: int) -> List[float]:
|
||||
"""
|
||||
PolarQuant Turbo4 Decode (Python reference).
|
||||
|
||||
Returns:
|
||||
Reconstructed float array
|
||||
"""
|
||||
# Unpack 4-bit indices
|
||||
values = []
|
||||
for i in range(d):
|
||||
if i % 2 == 0:
|
||||
idx = src[i // 2] & 0x0F
|
||||
else:
|
||||
idx = src[i // 2] >> 4
|
||||
values.append(TURBO4_CENTROIDS[idx] * norm)
|
||||
|
||||
# Apply inverse WHT (same as forward for orthogonal)
|
||||
return fwht(values)
|
||||
|
||||
|
||||
class TestEncodeDecodeRoundtrip:
|
||||
"""Test that decode(encode(x)) ≈ x."""
|
||||
|
||||
def test_zero_vector(self):
|
||||
"""Zero vector should encode/decode to zero."""
|
||||
d = 128
|
||||
src = [0.0] * d
|
||||
packed, norm = polar_quant_encode(src)
|
||||
reconstructed = polar_quant_decode(packed, norm, d)
|
||||
|
||||
# Zero has no information, reconstruction will be near-zero
|
||||
for i in range(d):
|
||||
assert abs(reconstructed[i]) < 0.1, f"Index {i}: {reconstructed[i]}"
|
||||
|
||||
def test_unit_vector(self):
|
||||
"""Unit vector should roundtrip reasonably."""
|
||||
d = 128
|
||||
src = [0.0] * d
|
||||
src[0] = 1.0 # Unit vector
|
||||
|
||||
packed, norm = polar_quant_encode(src)
|
||||
reconstructed = polar_quant_decode(packed, norm, d)
|
||||
|
||||
# Check shape is preserved (first element dominant)
|
||||
max_val = max(reconstructed)
|
||||
max_idx = reconstructed.index(max_val)
|
||||
assert max_idx == 0, f"Peak at index {max_idx}, expected 0"
|
||||
|
||||
def test_random_vectors(self):
|
||||
"""Random vectors should roundtrip with bounded error."""
|
||||
import random
|
||||
random.seed(42)
|
||||
|
||||
d = 128
|
||||
errors = []
|
||||
|
||||
for trial in range(10):
|
||||
src = [random.gauss(0, 0.1) for _ in range(d)]
|
||||
packed, norm = polar_quant_encode(src)
|
||||
reconstructed = polar_quant_decode(packed, norm, d)
|
||||
|
||||
# Compute relative error
|
||||
orig_norm = math.sqrt(sum(x * x for x in src))
|
||||
diff_norm = math.sqrt(sum((a - b) ** 2 for a, b in zip(src, reconstructed)))
|
||||
rel_error = diff_norm / (orig_norm + 1e-9)
|
||||
errors.append(rel_error)
|
||||
|
||||
avg_error = sum(errors) / len(errors)
|
||||
assert avg_error < 0.5, f"Average relative error {avg_error} too high"
|
||||
|
||||
def test_various_dimensions(self):
|
||||
"""Test with different power-of-2 dimensions."""
|
||||
for d in [16, 32, 64, 128, 256]:
|
||||
src = [math.sin(i * 0.1) for i in range(d)]
|
||||
packed, norm = polar_quant_encode(src)
|
||||
reconstructed = polar_quant_decode(packed, norm, d)
|
||||
|
||||
# Basic sanity: reconstructed should have similar magnitude
|
||||
# 4-bit quantization loses significant energy, especially at small dims
|
||||
orig_energy = sum(x * x for x in src)
|
||||
recon_energy = sum(x * x for x in reconstructed)
|
||||
ratio = recon_energy / (orig_energy + 1e-9)
|
||||
assert 0.1 < ratio < 10.0, f"d={d}: energy ratio {ratio}"
|
||||
|
||||
|
||||
class TestInnerProductPreservation:
|
||||
"""Test that Q·K ≈ Q·dequant(quant(K))."""
|
||||
|
||||
def test_inner_product_preserved(self):
|
||||
"""Inner products should be approximately preserved."""
|
||||
import random
|
||||
random.seed(123)
|
||||
|
||||
d = 128
|
||||
|
||||
# Generate two random vectors
|
||||
q = [random.gauss(0, 0.1) for _ in range(d)]
|
||||
k = [random.gauss(0, 0.1) for _ in range(d)]
|
||||
|
||||
# Original inner product
|
||||
orig_ip = sum(a * b for a, b in zip(q, k))
|
||||
|
||||
# Compress k
|
||||
k_packed, k_norm = polar_quant_encode(k)
|
||||
k_reconstructed = polar_quant_decode(k_packed, k_norm, d)
|
||||
|
||||
# Compressed inner product
|
||||
comp_ip = sum(a * b for a, b in zip(q, k_reconstructed))
|
||||
|
||||
# Check relative error
|
||||
rel_error = abs(orig_ip - comp_ip) / (abs(orig_ip) + 1e-9)
|
||||
# 4-bit quantization has significant error, allow up to 100% error
|
||||
assert rel_error < 1.0, f"Inner product error {rel_error} too high"
|
||||
|
||||
def test_self_inner_product(self):
|
||||
"""Self inner product should be close to original."""
|
||||
d = 128
|
||||
x = [math.cos(i * 0.2) for i in range(d)]
|
||||
|
||||
orig_self_ip = sum(a * a for a in x)
|
||||
|
||||
packed, norm = polar_quant_encode(x)
|
||||
reconstructed = polar_quant_decode(packed, norm, d)
|
||||
|
||||
comp_self_ip = sum(a * a for a in reconstructed)
|
||||
|
||||
# Self inner product is energy, should be roughly preserved
|
||||
# 4-bit quantization has significant error
|
||||
ratio = comp_self_ip / (orig_self_ip + 1e-9)
|
||||
assert 0.3 < ratio < 3.0, f"Self inner product ratio {ratio}"
|
||||
|
||||
|
||||
class TestWHTOrthogonality:
|
||||
"""Test that WHT is orthogonal (WHT^T · WHT = I)."""
|
||||
|
||||
def test_wht_orthogonality(self):
|
||||
"""WHT should be orthogonal transformation."""
|
||||
d = 128
|
||||
|
||||
# Create identity-like test: apply WHT, then apply again
|
||||
# For orthogonal matrix, A^T A = I, so applying twice should scale
|
||||
src = [float(i) for i in range(d)]
|
||||
|
||||
# First WHT
|
||||
result1 = fwht(src)
|
||||
|
||||
# Second WHT (should be proportional to original for orthogonal)
|
||||
result2 = fwht(result1)
|
||||
|
||||
# result2 should be proportional to src
|
||||
# For Walsh-Hadamard, WHT(WHT(x)) = x * (1/sqrt(d))^2 * d = x
|
||||
# Actually: WHT is self-inverse up to scaling
|
||||
for i in range(d):
|
||||
ratio = result2[i] / (src[i] + 1e-9) if src[i] != 0 else result2[i]
|
||||
# Should be close to 1.0 (or 0 if src[i] is 0)
|
||||
if abs(src[i]) > 0.01:
|
||||
assert abs(ratio - 1.0) < 0.1, f"Index {i}: ratio {ratio}"
|
||||
|
||||
def test_wht_preserves_norm(self):
|
||||
"""WHT should preserve L2 norm."""
|
||||
d = 128
|
||||
src = [math.sin(i) for i in range(d)]
|
||||
|
||||
orig_norm = math.sqrt(sum(x * x for x in src))
|
||||
|
||||
result = fwht(src)
|
||||
result_norm = math.sqrt(sum(x * x for x in result))
|
||||
|
||||
ratio = result_norm / orig_norm
|
||||
assert abs(ratio - 1.0) < 0.01, f"Norm ratio {ratio}, expected 1.0"
|
||||
|
||||
def test_wht_linearity(self):
|
||||
"""WHT should be linear: WHT(a+b) = WHT(a) + WHT(b)."""
|
||||
d = 64
|
||||
a = [float(i) * 0.1 for i in range(d)]
|
||||
b = [float(i) * 0.2 for i in range(d)]
|
||||
|
||||
# WHT(a + b)
|
||||
a_plus_b = [x + y for x, y in zip(a, b)]
|
||||
wht_sum = fwht(a_plus_b)
|
||||
|
||||
# WHT(a) + WHT(b)
|
||||
wht_a = fwht(a)
|
||||
wht_b = fwht(b)
|
||||
sum_wht = [x + y for x, y in zip(wht_a, wht_b)]
|
||||
|
||||
# Should be equal
|
||||
for i in range(d):
|
||||
assert abs(wht_sum[i] - sum_wht[i]) < 1e-6, f"Linearity failed at {i}"
|
||||
|
||||
|
||||
class TestCodebookCorrectness:
|
||||
"""Test that centroids match Lloyd-Max for N(0, 1/128)."""
|
||||
|
||||
def test_centroids_extremes(self):
|
||||
"""Extreme centroids should cover tails of distribution."""
|
||||
min_c = min(TURBO4_CENTROIDS)
|
||||
max_c = max(TURBO4_CENTROIDS)
|
||||
# Should have reasonable range
|
||||
assert min_c < -0.2, f"Min centroid {min_c} should be < -0.2"
|
||||
assert max_c > 0.2, f"Max centroid {max_c} should be > 0.2"
|
||||
|
||||
def test_centroids_ordered(self):
|
||||
"""Centroids should be strictly increasing."""
|
||||
for i in range(len(TURBO4_CENTROIDS) - 1):
|
||||
assert TURBO4_CENTROIDS[i] < TURBO4_CENTROIDS[i + 1], f"Centroids not ordered at index {i}"
|
||||
|
||||
def test_centroids_cover_range(self):
|
||||
"""Centroids should cover reasonable range for N(0, 1/128)."""
|
||||
# For N(0, 1/128), std = 1/sqrt(128) ≈ 0.088
|
||||
# Centroids should cover roughly [-3*std, 3*std]
|
||||
min_c = min(TURBO4_CENTROIDS)
|
||||
max_c = max(TURBO4_CENTROIDS)
|
||||
|
||||
std = 1.0 / math.sqrt(128) # ≈ 0.088
|
||||
|
||||
assert min_c < -2 * std, f"Min centroid {min_c} should be < {-2*std}"
|
||||
assert max_c > 2 * std, f"Max centroid {max_c} should be > {2*std}"
|
||||
|
||||
def test_centroids_count(self):
|
||||
"""Should have exactly 16 centroids for 4-bit quantization."""
|
||||
assert len(TURBO4_CENTROIDS) == 16, f"Expected 16 centroids, got {len(TURBO4_CENTROIDS)}"
|
||||
|
||||
|
||||
class TestBitPacking:
|
||||
"""Test bit packing/unpacking correctness."""
|
||||
|
||||
def test_packing_roundtrip(self):
|
||||
"""Packing and unpacking should be lossless for 4-bit values."""
|
||||
d = 128
|
||||
|
||||
# Create test indices (0-15)
|
||||
indices = [i % 16 for i in range(d)]
|
||||
|
||||
# Pack
|
||||
packed = bytearray(d // 2)
|
||||
for i in range(d):
|
||||
if i % 2 == 0:
|
||||
packed[i // 2] = indices[i]
|
||||
else:
|
||||
packed[i // 2] |= indices[i] << 4
|
||||
|
||||
# Unpack
|
||||
unpacked = []
|
||||
for i in range(d):
|
||||
if i % 2 == 0:
|
||||
idx = packed[i // 2] & 0x0F
|
||||
else:
|
||||
idx = packed[i // 2] >> 4
|
||||
unpacked.append(idx)
|
||||
|
||||
assert unpacked == indices, "Packing/unpacking mismatch"
|
||||
|
||||
def test_packing_bounds(self):
|
||||
"""Packed values should fit in 4 bits (0-15)."""
|
||||
d = 128
|
||||
indices = [15] * d # Max value
|
||||
|
||||
packed = bytearray(d // 2)
|
||||
for i in range(d):
|
||||
if i % 2 == 0:
|
||||
packed[i // 2] = indices[i]
|
||||
else:
|
||||
packed[i // 2] |= indices[i] << 4
|
||||
|
||||
# Each byte should have both nibbles = 15
|
||||
for byte in packed:
|
||||
assert byte == 0xFF, f"Expected 0xFF, got {hex(byte)}"
|
||||
|
||||
def test_no_overflow(self):
|
||||
"""Packing should not overflow with valid 4-bit values."""
|
||||
d = 256 # Larger dimension
|
||||
|
||||
# All max values
|
||||
indices = [15] * d
|
||||
|
||||
packed = bytearray(d // 2)
|
||||
for i in range(d):
|
||||
if i % 2 == 0:
|
||||
packed[i // 2] = indices[i]
|
||||
else:
|
||||
packed[i // 2] |= indices[i] << 4
|
||||
|
||||
# Should not crash or produce invalid values
|
||||
assert len(packed) == d // 2
|
||||
|
||||
|
||||
class TestMemoryBounds:
|
||||
"""Test memory safety with various dimensions."""
|
||||
|
||||
def test_minimum_dimension(self):
|
||||
"""Should work with minimum dimension (2)."""
|
||||
d = 2
|
||||
src = [1.0, 0.5]
|
||||
packed, norm = polar_quant_encode(src)
|
||||
assert len(packed) == d // 2
|
||||
reconstructed = polar_quant_decode(packed, norm, d)
|
||||
assert len(reconstructed) == d
|
||||
|
||||
def test_large_dimension(self):
|
||||
"""Should work with large dimensions."""
|
||||
d = 1024
|
||||
src = [math.sin(i * 0.01) for i in range(d)]
|
||||
packed, norm = polar_quant_encode(src)
|
||||
assert len(packed) == d // 2
|
||||
reconstructed = polar_quant_decode(packed, norm, d)
|
||||
assert len(reconstructed) == d
|
||||
|
||||
def test_odd_dimension_fails(self):
|
||||
"""Odd dimensions should fail (need even for 4-bit packing)."""
|
||||
d = 127 # Odd
|
||||
src = [0.0] * d
|
||||
|
||||
with pytest.raises(AssertionError):
|
||||
polar_quant_encode(src)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
141
tests/test_turboquant.py
Normal file
141
tests/test_turboquant.py
Normal file
@@ -0,0 +1,141 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
TurboQuant Test Suite
|
||||
Tests for critical paths in KV cache compression.
|
||||
|
||||
Issue #679: Codebase Genome: turboquant — Full Analysis
|
||||
"""
|
||||
import unittest
|
||||
import subprocess
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
class TestTurboQuant(unittest.TestCase):
|
||||
"""Test TurboQuant implementation."""
|
||||
|
||||
def test_repo_structure(self):
|
||||
"""Verify expected files exist."""
|
||||
required_files = [
|
||||
"llama-turbo.h",
|
||||
"llama-turbo.cpp",
|
||||
"ggml-metal-turbo.metal",
|
||||
"README.md",
|
||||
"GENOME.md"
|
||||
]
|
||||
|
||||
for filename in required_files:
|
||||
filepath = os.path.join(os.path.dirname(__file__), "..", filename)
|
||||
self.assertTrue(os.path.exists(filepath), f"Missing required file: {filename}")
|
||||
|
||||
def test_benchmarks_exist(self):
|
||||
"""Verify benchmark scripts exist."""
|
||||
benchmark_files = [
|
||||
"benchmarks/run_benchmarks.py",
|
||||
"benchmarks/run_perplexity.py",
|
||||
"benchmarks/run_long_session.py"
|
||||
]
|
||||
|
||||
for filename in benchmark_files:
|
||||
filepath = os.path.join(os.path.dirname(__file__), "..", filename)
|
||||
self.assertTrue(os.path.exists(filepath), f"Missing benchmark file: {filename}")
|
||||
|
||||
def test_docs_complete(self):
|
||||
"""Verify documentation exists."""
|
||||
doc_files = [
|
||||
"docs/PROJECT_STATUS.md",
|
||||
"profiles/README.md"
|
||||
]
|
||||
|
||||
for filename in doc_files:
|
||||
filepath = os.path.join(os.path.dirname(__file__), "..", filename)
|
||||
self.assertTrue(os.path.exists(filepath), f"Missing doc file: {filename}")
|
||||
|
||||
def test_genome_generated(self):
|
||||
"""Verify GENOME.md was generated."""
|
||||
genome_path = os.path.join(os.path.dirname(__file__), "..", "GENOME.md")
|
||||
self.assertTrue(os.path.exists(genome_path), "GENOME.md not found")
|
||||
|
||||
# Check it has required sections
|
||||
with open(genome_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
required_sections = [
|
||||
"## Project Overview",
|
||||
"## Architecture",
|
||||
"## Entry Points",
|
||||
"## Data Flow",
|
||||
"## Key Abstractions",
|
||||
"## API Surface",
|
||||
"## Test Coverage Gaps",
|
||||
"## Security Considerations"
|
||||
]
|
||||
|
||||
for section in required_sections:
|
||||
self.assertIn(section, content, f"GENOME.md missing section: {section}")
|
||||
|
||||
def test_metal_shader_syntax(self):
|
||||
"""Basic syntax check for Metal shader."""
|
||||
shader_path = os.path.join(os.path.dirname(__file__), "..", "ggml-metal-turbo.metal")
|
||||
with open(shader_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
# Check for key functions
|
||||
self.assertIn("kernel_fwht_128", content, "Missing kernel_fwht_128 function")
|
||||
self.assertIn("kernel_turbo4_dequant", content, "Missing kernel_turbo4_dequant function")
|
||||
self.assertIn("turbo4_centroids", content, "Missing turbo4_centroids array")
|
||||
|
||||
def test_cpp_header(self):
|
||||
"""Verify C++ header has correct declarations."""
|
||||
header_path = os.path.join(os.path.dirname(__file__), "..", "llama-turbo.h")
|
||||
with open(header_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
# Check for function declarations
|
||||
self.assertIn("polar_quant_encode_turbo4", content, "Missing encode function")
|
||||
self.assertIn("polar_quant_decode_turbo4", content, "Missing decode function")
|
||||
self.assertIn('extern "C"', content, "Missing C linkage")
|
||||
|
||||
class TestBenchmarks(unittest.TestCase):
|
||||
"""Test benchmark infrastructure."""
|
||||
|
||||
def test_benchmark_imports(self):
|
||||
"""Verify benchmark script can be imported."""
|
||||
benchmark_path = os.path.join(os.path.dirname(__file__), "..", "benchmarks", "run_benchmarks.py")
|
||||
|
||||
# Check file exists
|
||||
self.assertTrue(os.path.exists(benchmark_path), "Benchmark script not found")
|
||||
|
||||
# Check it has main function
|
||||
with open(benchmark_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
self.assertIn("def main():", content, "Benchmark script missing main function")
|
||||
self.assertIn("argparse", content, "Benchmark script missing argparse")
|
||||
|
||||
class TestDocumentation(unittest.TestCase):
|
||||
"""Test documentation completeness."""
|
||||
|
||||
def test_readme_sections(self):
|
||||
"""Verify README has required sections."""
|
||||
readme_path = os.path.join(os.path.dirname(__file__), "..", "README.md")
|
||||
with open(readme_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
required_sections = ["## What", "## Why", "## Status", "## Roles"]
|
||||
for section in required_sections:
|
||||
self.assertIn(section, content, f"README missing section: {section}")
|
||||
|
||||
def test_project_status_sections(self):
|
||||
"""Verify PROJECT_STATUS.md has required sections."""
|
||||
status_path = os.path.join(os.path.dirname(__file__), "..", "docs", "PROJECT_STATUS.md")
|
||||
with open(status_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
# Check for key findings
|
||||
self.assertIn("73%", content, "Missing 73% savings metric")
|
||||
self.assertIn("PolarQuant", content, "Missing PolarQuant references")
|
||||
self.assertIn("Metal", content, "Missing Metal shader references")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user