### Install TurboQuant dependencies

Source: https://github.com/tonbistudio/turboquant-pytorch/blob/master/README.md

Commands to install the required Python packages and the specific CUDA-enabled version of PyTorch for GPU acceleration.

```bash
pip install -r requirements.txt
pip install torch --index-url https://download.pytorch.org/whl/cu128
```

--------------------------------

### Load Model with 4-bit Quantization and KV Cache Compression Example

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

This example demonstrates loading a model with 4-bit quantization and then using TurboQuant to compress its KV cache for inference.

```APIDOC
## Load Model with 4-bit Quantization and KV Cache Compression

### Description
This example shows how to load a pre-trained causal language model with 4-bit quantization using `BitsAndBytesConfig` and then prepare it for inference with TurboQuant's KV cache compression.

### Method
```python
# Import necessary libraries
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from turboquant import TurboQuantCompressorV2, TurboQuantCompressorMSE, TurboQuantKVCache
import torch.nn.functional as F

# Define model name
model_name = "Qwen/Qwen2.5-3B-Instruct"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Configure 4-bit quantization
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4"
)

# Load model with quantization and automatic device mapping
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.float16,
)
model.eval() # Set model to evaluation mode

# Build prompt and prepare inputs
prompt = "Your long context here..."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Perform forward pass to generate KV cache
with torch.no_grad():
    outputs = model(**inputs, use_cache=True)

# Extract KV cache
cache = outputs.past_key_values
n_layers = len(cache.layers)

# --- TurboQuant Integration Example ---

# Example: Validate each layer with different bit settings
for bits in [2, 3, 4]:
    cosine_sims = []

    for layer_idx in range(n_layers):
        # Extract keys for the current layer
        keys = cache.layers[layer_idx].keys      # Shape: (1, heads, seq, head_dim)
        B, H, S, D = keys.shape

        # Initialize TurboQuant compressor for keys
        # Using TurboQuantCompressorV2 for keys as it's designed for attention scores
        key_comp = TurboQuantCompressorV2(D, bits, seed=layer_idx * 1000, device="cuda")
        compressed_k = key_comp.compress(keys)

        # Prepare query for attention score calculation (using the last token)
        query = keys[:, :, -1:, :]  # Shape: (1, heads, 1, head_dim)

        # Calculate real attention scores
        real_scores = torch.matmul(query.float(), keys.float().transpose(-2, -1))
        
        # Calculate attention scores using compressed keys
        # This bypasses full decompression for efficiency
        tq_scores = key_comp.asymmetric_attention_scores(query, compressed_k)

        # Compute cosine similarity per head to evaluate compression fidelity
        for h in range(H):
            cos = F.cosine_similarity(
                real_scores[0, h].flatten().unsqueeze(0),
                tq_scores[0, h].flatten().unsqueeze(0)
            ).item()
            cosine_sims.append(cos)

    # Calculate and print average cosine similarity for the current bit setting
    avg_cos = sum(cosine_sims) / len(cosine_sims)
    print(f"TQ-{bits}bit: Avg attention cosine similarity = {avg_cos:.6f}")
    # Expected results: ~0.985 (2-bit), ~0.995 (3-bit), ~0.998 (4-bit)
```

### Parameters
None for this example script.

### Request Example
N/A (This is a Python script, not an API endpoint).

### Response
```
TQ-2bit: Avg attention cosine similarity = 0.985xxx
TQ-3bit: Avg attention cosine similarity = 0.995xxx
TQ-4bit: Avg attention cosine similarity = 0.998xxx
```
```

--------------------------------

### Complete Real Model Validation Example

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Demonstrates validating TurboQuant against a real language model's KV cache, comparing attention scores between original and compressed representations.

```APIDOC
## Complete Real Model Validation Example

This example demonstrates validating TurboQuant against a real language model's KV cache, comparing attention scores between original and compressed representations.

```python
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from turboquant.compressors import TurboQuantCompressorV2, TurboQuantCompressorMSE

# ... (rest of the example code)
```
```

--------------------------------

### Initialize and Use TurboQuantKVCache

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Demonstrates initializing a KV cache wrapper for transformer layers, appending key-value pairs, and computing attention scores. It also shows how to retrieve compression statistics and decompressed values.

```python
import torch
from turboquant import TurboQuantKVCache

d_key = 128
d_value = 128
cache = TurboQuantKVCache(d_key, d_value, bits=3, seed=42, device="cuda")

seq_len = 2048
keys = torch.randn(seq_len, d_key, device="cuda")
values = torch.randn(seq_len, d_value, device="cuda")
cache.append(keys, values)

query = torch.randn(1, d_key, device="cuda")
scores = cache.attention_scores(query)
print(f"Attention scores shape: {scores.shape}")

cached_values = cache.get_values()
usage = cache.memory_usage_bits()
print(f"Compression ratio: {usage['compression_ratio']:.2f}x")
print(f"Cached tokens: {len(cache)}")
```

--------------------------------

### Load Model with 4-bit Quantization and KV Cache Compression Validation

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

This snippet demonstrates loading a causal language model with 4-bit quantization using BitsAndBytesConfig. It then builds a prompt, runs a forward pass to generate KV cache, and iterates through different bit configurations (2, 3, 4) to compress the keys. For each bit configuration and each layer, it calculates the cosine similarity between real and TurboQuant-compressed attention scores to validate the compression quality. The average cosine similarity is printed for each bit configuration.

```python
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
import torch.nn.functional as F
from turbomath.turboquant import TurboQuantCompressorV2

# Load model with 4-bit quantization
model_name = "Qwen/Qwen2.5-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_quant_type="nf4"
    ),
    device_map="auto",
    torch_dtype=torch.float16,
)
model.eval()

# Build prompt and run forward pass
prompt = "Your long context here..."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model(**inputs, use_cache=True)

cache = outputs.past_key_values
n_layers = len(cache.layers)

# Validate each layer
for bits in [2, 3, 4]:
    cosine_sims = []

    for layer_idx in range(n_layers):
        keys = cache.layers[layer_idx].keys      # (1, heads, seq, head_dim)
        B, H, S, D = keys.shape

        # Compress
        key_comp = TurboQuantCompressorV2(D, bits, seed=layer_idx * 1000, device="cuda")
        compressed_k = key_comp.compress(keys)

        # Compare attention scores
        query = keys[:, :, -1:, :]  # Last token as query

        real_scores = torch.matmul(query.float(), keys.float().transpose(-2, -1))
        tq_scores = key_comp.asymmetric_attention_scores(query, compressed_k)

        # Compute cosine similarity per head
        for h in range(H):
            cos = F.cosine_similarity(
                real_scores[0, h].flatten().unsqueeze(0),
                tq_scores[0, h].flatten().unsqueeze(0)
            ).item()
            cosine_sims.append(cos)

    avg_cos = sum(cosine_sims) / len(cosine_sims)
    print(f"TQ-{bits}bit: Avg attention cosine similarity = {avg_cos:.6f}")
    # Expected: ~0.985 (2-bit), ~0.995 (3-bit), ~0.998 (4-bit)

```

--------------------------------

### TurboQuantKVCache - KV Cache Wrapper

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Provides a high-level interface for compressing LLM key-value caches using TurboQuantProd for keys and TurboQuantMSE for values.

```APIDOC
## TurboQuantKVCache - Drop-in KV Cache Wrapper

`TurboQuantKVCache` provides a high-level interface for compressing LLM key-value caches. Keys use `TurboQuantProd` (inner products needed for attention), while values use `TurboQuantMSE` (MSE reconstruction sufficient for weighted averaging).

### Example Usage

```python
import torch
from turboquant import TurboQuantKVCache

# Initialize cache for typical transformer dimensions
d_key = 128
d_value = 128
cache = TurboQuantKVCache(d_key, d_value, bits=3, seed=42, device="cuda")

# Simulate appending KV pairs from transformer layers
seq_len = 2048
keys = torch.randn(seq_len, d_key, device="cuda")
values = torch.randn(seq_len, d_value, device="cuda")
cache.append(keys, values)

# Compute attention scores for a query
query = torch.randn(1, d_key, device="cuda")
scores = cache.attention_scores(query)  # (seq_len,) attention logits
print(f"Attention scores shape: {scores.shape}")

# Retrieve decompressed values for weighted sum
cached_values = cache.get_values()  # (seq_len, d_value)

# Check compression statistics
usage = cache.memory_usage_bits()
print(f"Compression ratio: {usage['compression_ratio']:.2f}x")
print(f"Compressed: {usage['total_bits'] / 8 / 1024:.1f} KB")
print(f"FP16 equivalent: {usage['fp16_bits'] / 8 / 1024:.1f} KB")

# Cache length
print(f"Cached tokens: {len(cache)}")
```
```

--------------------------------

### TurboQuant Compressor Usage

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Details on using TurboQuant compressors for keys and values, and the `TurboQuantKVCache` wrapper.

```APIDOC
## TurboQuant Compressor Usage

### Description
This section details how to use TurboQuant's compression algorithms for KV caches. It covers using `TurboQuantCompressorV2` for keys and `TurboQuantCompressorMSE` for values, as well as the convenience wrapper `TurboQuantKVCache`.

### Key Compression (`TurboQuantCompressorV2`)
- **Purpose**: Compresses keys for attention score calculation.
- **Usage**: Requires unbiased inner products for accurate attention.
- **Method**: `key_comp.compress(keys)` to compress, `key_comp.asymmetric_attention_scores(query, compressed_keys)` to compute attention scores directly from compressed keys.

### Value Compression (`TurboQuantCompressorMSE`)
- **Purpose**: Compresses values, where MSE reconstruction is sufficient.
- **Usage**: Suitable when direct attention score calculation isn't the primary concern for values.

### Combined KV Cache (`TurboQuantKVCache`)
- **Purpose**: A drop-in wrapper that automatically handles compression and decompression for both keys and values.
- **Usage**: Simplifies integration by managing the compression logic internally.

### Example Snippet (Conceptual)
```python
# Assuming 'keys' and 'values' tensors are extracted from the model's KV cache

# --- Using individual compressors ---
# For keys (e.g., 3-bit)
key_compressor = TurboQuantCompressorV2(dimension=D, bits=3, seed=42, device='cuda')
compressed_keys = key_compressor.compress(keys)
# Use key_compressor.asymmetric_attention_scores(...) for inference

# For values (e.g., 4-bit, assuming MSE is sufficient)
value_compressor = TurboQuantCompressorMSE(dimension=D, bits=4, seed=42, device='cuda')
compressed_values = value_compressor.compress(values)
# Decompress values if needed: decompressed_values = value_compressor.decompress(compressed_values)

# --- Using the convenience wrapper ---
# kv_cache_wrapper = TurboQuantKVCache(model, bits=3) # Example initialization
# compressed_kv = kv_cache_wrapper.compress(original_kv_cache)
# model_output = model(..., past_key_values=compressed_kv)
```

### Parameters
- **`dimension` (int)**: The dimension of the keys/values to be compressed.
- **`bits` (int)**: The target number of bits per coordinate (e.g., 2, 3, 4).
- **`seed` (int, optional)**: Seed for reproducibility.
- **`device` (str, optional)**: The device to perform computations on ('cuda' or 'cpu').

### Request Example
N/A (This is a conceptual code example).

### Response
N/A (This is a conceptual code example).
```

--------------------------------

### Run Real Model KV Cache Validation

Source: https://github.com/tonbistudio/turboquant-pytorch/blob/master/README.md

Executes the validation script against a real language model (Qwen2.5-3B-Instruct). It captures the KV cache, applies TurboQuant compression at various bit-widths, and calculates attention score fidelity metrics.

```bash
python -m turboquant.validate
```

--------------------------------

### Compress Keys with TurboQuantCompressorV2

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Shows how to compress 4D transformer key tensors and perform asymmetric attention calculations directly on the compressed representation to save memory while maintaining accuracy.

```python
import torch
from turboquant.compressors import TurboQuantCompressorV2

head_dim = 128
bits = 3
compressor = TurboQuantCompressorV2(head_dim, bits, seed=42, device="cuda")

batch_size = 1
num_heads = 8
seq_len = 4096
keys = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda")

compressed = compressor.compress(keys)
query_len = 1
queries = torch.randn(batch_size, num_heads, query_len, head_dim, device="cuda")

scores = compressor.asymmetric_attention_scores(queries, compressed)
print(f"Score shape: {scores.shape}")
attention_weights = torch.softmax(scores / (head_dim ** 0.5), dim=-1)
```

--------------------------------

### TurboQuantProd: Stage 1 + Stage 2 Unbiased Inner Products in PyTorch

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Implements the full TurboQuant algorithm by combining MSE quantization (Stage 1) with Quantized Johnson-Lindenstrauss (QJL) residual correction (Stage 2). This approach enables unbiased inner product estimation, crucial for accurate attention computation. It uses a specified bit budget for MSE quantization and QJL signs, along with storage for residual norms.

```python
import torch
from turboquant import TurboQuantProd

# Initialize with 3-bit total budget (2-bit MSE + 1-bit QJL)
quantizer = TurboQuantProd(d=128, bits=3, qjl_dim=128, seed=42, device="cuda")

# Generate query and key vectors
n_keys = 1000
keys = torch.randn(n_keys, 128, device="cuda")
keys = keys / torch.norm(keys, dim=-1, keepdim=True)

query = torch.randn(1, 128, device="cuda")
query = query / torch.norm(query, dim=-1, keepdim=True)

# Compress keys
compressed = quantizer.quantize(keys)
# Returns dict with:
#   'mse_indices': (n_keys, 128) int tensor - MSE codebook indices
#   'qjl_signs': (n_keys, 128) sign tensor - QJL projection signs
#   'residual_norm': (n_keys,) float tensor - L2 norm of residuals

# Compute unbiased inner product estimates
query_expanded = query.expand(n_keys, -1)
estimated_ips = quantizer.inner_product(query_expanded, compressed)

# Compare with true inner products
true_ips = (query_expanded * keys).sum(dim=-1)

# Verify unbiasedness
bias = (estimated_ips - true_ips).mean().item()
correlation = torch.corrcoef(torch.stack([true_ips, estimated_ips]))[0, 1].item()
print(f"Bias: {bias:+.6f} (should be ~0)")
print(f"Correlation: {correlation:.4f}")
```

--------------------------------

### TurboQuantMSE: Stage 1 MSE Quantizer in PyTorch

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Implements Stage 1 of the TurboQuant algorithm, which involves random orthogonal rotation followed by per-coordinate Lloyd-Max quantization. This quantizer is optimized for Mean Squared Error (MSE) vector reconstruction but may introduce bias in inner product estimation. It provides methods for rotation, quantization, and dequantization.

```python
import torch
from turboquant import TurboQuantMSE

# Initialize MSE quantizer for 128-dim vectors with 3-bit precision
quantizer = TurboQuantMSE(d=128, bits=3, seed=42, device="cuda")

# Generate random unit vectors
n_vectors = 1000
x = torch.randn(n_vectors, 128, device="cuda")
x = x / torch.norm(x, dim=-1, keepdim=True)

# Quantize and reconstruct
x_reconstructed, indices = quantizer(x)

# Measure reconstruction quality
mse = ((x - x_reconstructed) ** 2).sum(dim=-1).mean().item()
print(f"Total MSE: {mse:.6f}")

# Individual operations available
rotated = quantizer.rotate(x)           # Apply random rotation: y = Pi @ x
indices = quantizer.quantize(x)         # Get codebook indices (no reconstruction)
x_hat = quantizer.dequantize(indices)   # Reconstruct from indices

# Theoretical bound from paper: D_mse <= sqrt(3)*pi/2 * (1/4^bits)
import math
theoretical_bound = math.sqrt(3) * math.pi / 2 * (1 / (4 ** 3))
print(f"Theoretical upper bound: {theoretical_bound:.6f}")
print(f"Ratio (measured/bound): {mse / theoretical_bound:.3f}")  # Should be < 1.0
```

--------------------------------

### Run Synthetic Algorithm Validation

Source: https://github.com/tonbistudio/turboquant-pytorch/blob/master/README.md

Executes the synthetic test suite to validate the core TurboQuant algorithm, including Lloyd-Max codebook generation, MSE measurement, and QJL correction. This script does not require a GPU, though one can be used for optional speed benchmarking.

```bash
python -m turboquant.test_turboquant
```

--------------------------------

### Lloyd-Max Codebook Solver

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Computes optimal scalar quantization centroids and boundaries for coordinate distributions using iterative Lloyd-Max optimization.

```APIDOC
## solve_lloyd_max

### Description
Computes optimal scalar quantization centroids for the coordinate distribution arising after random rotation of unit vectors.

### Method
Python Function

### Parameters
#### Arguments
- **d** (int) - Required - Dimensionality of the vectors.
- **bits** (int) - Required - Number of bits for quantization.
- **use_exact** (bool) - Optional - Whether to use exact computation.
- **max_iter** (int) - Optional - Maximum iterations for optimization.
- **tol** (float) - Optional - Tolerance for convergence.

### Response
- **centroids** (Tensor) - The computed quantization centroids.
- **boundaries** (Tensor) - The computed quantization boundaries.
```

--------------------------------

### Lloyd-Max Codebook Solver in PyTorch

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Computes optimal scalar quantization centroids using iterative Lloyd-Max optimization for a given bit budget. This function is crucial for Stage 1 of the TurboQuant algorithm. It can be used directly or via the `LloydMaxCodebook` class for quantizing and reconstructing scalar values.

```python
import torch
from turboquant.lloyd_max import solve_lloyd_max, LloydMaxCodebook

# Solve Lloyd-Max quantizer for 128-dimensional vectors with 3-bit precision
d = 128
bits = 3
centroids, boundaries = solve_lloyd_max(d, bits, use_exact=False, max_iter=200, tol=1e-10)

print(f"Centroids ({2**bits} levels): {centroids.tolist()}")
print(f"Boundaries: {boundaries.tolist()}")

# Or use the convenience class that precomputes everything
codebook = LloydMaxCodebook(d=128, bits=3)
print(f"Codebook: {codebook}")
# Output: LloydMaxCodebook(d=128, bits=3, levels=8, distortion_per_coord=0.000264)

# Quantize scalar values
x = torch.randn(1000, d) / torch.sqrt(torch.tensor(d, dtype=torch.float32))
indices = codebook.quantize(x)
x_reconstructed = codebook.dequantize(indices)

mse_per_coord = ((x - x_reconstructed) ** 2).mean().item()
print(f"MSE per coordinate: {mse_per_coord:.6f}")
```

--------------------------------

### Compress Values with TurboQuantCompressorMSE

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Provides an implementation for compressing value vectors using MSE. This approach is suitable for weighted averaging in attention mechanisms where per-vector error is minimized.

```python
import torch
from turboquant.compressors import TurboQuantCompressorMSE

head_dim = 128
bits = 3
value_compressor = TurboQuantCompressorMSE(head_dim, bits, seed=42, device="cuda")

batch_size = 1
num_heads = 8
seq_len = 4096
values = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda")

compressed_v = value_compressor.compress(values)
values_reconstructed = value_compressor.decompress(compressed_v)

attention_weights = torch.softmax(torch.randn(1, 8, 1, seq_len, device="cuda"), dim=-1)
attention_output = torch.matmul(attention_weights, values_reconstructed)
print(f"Attention output shape: {attention_output.shape}")
```

--------------------------------

### TurboQuantMSE - Stage 1 Quantizer

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Implements MSE-optimal vector reconstruction using random orthogonal rotation followed by per-coordinate Lloyd-Max quantization.

```APIDOC
## TurboQuantMSE

### Description
Provides MSE-optimal vector reconstruction. Note that this method introduces bias in inner product estimation.

### Parameters
#### Constructor Arguments
- **d** (int) - Required - Vector dimensionality.
- **bits** (int) - Required - Quantization bit budget.
- **seed** (int) - Optional - Random seed for rotation.
- **device** (str) - Optional - Computation device (e.g., 'cuda').

### Methods
- **rotate(x)**: Applies random rotation.
- **quantize(x)**: Returns codebook indices.
- **dequantize(indices)**: Reconstructs vectors from indices.
```

--------------------------------

### TurboQuantProd - Unbiased Inner Product Estimation

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Combines MSE quantization with QJL residual correction to provide unbiased inner product estimation for compressed representations.

```APIDOC
## TurboQuantProd

### Description
Combines MSE quantization (bits-1) with QJL residual correction (1 bit) to enable unbiased inner product estimation.

### Parameters
#### Constructor Arguments
- **d** (int) - Required - Vector dimensionality.
- **bits** (int) - Required - Total bit budget.
- **qjl_dim** (int) - Required - Dimension for QJL projection.

### Response
- **compressed** (dict) - Contains 'mse_indices', 'qjl_signs', and 'residual_norm'.
```

--------------------------------

### TurboQuantCompressorV2 - Asymmetric Attention

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Designed for production use with real transformer models, handling 4D tensors and computing attention scores directly from compressed representations using the asymmetric estimator.

```APIDOC
## TurboQuantCompressorV2 - Asymmetric Attention for Real Models

`TurboQuantCompressorV2` is designed for production use with real transformer models. It handles 4D tensors (batch, heads, seq, head_dim), preserves vector norms, and computes attention scores directly from compressed representations using the asymmetric estimator.

### Example Usage

```python
import torch
from turboquant.compressors import TurboQuantCompressorV2

# Initialize for standard transformer head dimension
head_dim = 128
bits = 3
compressor = TurboQuantCompressorV2(head_dim, bits, seed=42, device="cuda")

# Compress key states from a transformer layer
# Shape: (batch, num_kv_heads, seq_len, head_dim)
batch_size = 1
num_heads = 8
seq_len = 4096
keys = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda")

compressed = compressor.compress(keys)
# Returns dict with:
#   'k_mse': (B, H, S, D) fp16 - MSE reconstruction in original space
#   'qjl_signs': (B, H, S, D) int8 - QJL sign bits {-1, +1}
#   'residual_norm': (B, H, S) fp16 - residual norms ||r||
#   'shape': tuple - original shape

# Compute attention scores directly from compressed keys
# Queries: (batch, num_heads, query_len, head_dim)
query_len = 1  # Single query token (next-token generation)
queries = torch.randn(batch_size, num_heads, query_len, head_dim, device="cuda")

# Asymmetric attention: <q, k> = <q, k_mse> + ||r|| * sqrt(pi/2)/m * <S@q, signs>
scores = compressor.asymmetric_attention_scores(queries, compressed)
# Output shape: (batch, num_heads, query_len, seq_len) = (1, 8, 1, 4096)

print(f"Score shape: {scores.shape}")
print(f"Score range: [{scores.min():.3f}, {scores.max():.3f}]")

# Apply softmax and compute attention output
attention_weights = torch.softmax(scores / (head_dim ** 0.5), dim=-1)
```
```

--------------------------------

### TurboQuantCompressorMSE - Value Compression

Source: https://context7.com/tonbistudio/turboquant-pytorch/llms.txt

Provides MSE-only compression for value vectors, suitable for weighted averaging in attention mechanisms.

```APIDOC
## TurboQuantCompressorMSE - Value Compression

`TurboQuantCompressorMSE` provides MSE-only compression for value vectors. QJL correction is unnecessary for values since the weighted averaging in attention (softmax @ V) averages out per-vector reconstruction errors.

### Example Usage

```python
import torch
from turboquant.compressors import TurboQuantCompressorMSE

# Initialize value compressor
head_dim = 128
bits = 3
value_compressor = TurboQuantCompressorMSE(head_dim, bits, seed=42, device="cuda")

# Compress value states
batch_size = 1
num_heads = 8
seq_len = 4096
values = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda")

compressed_v = value_compressor.compress(values)
# Returns dict with:
#   'indices': (N, D) uint8 - codebook indices
#   'vec_norms': (N,) fp16 - original vector norms
#   'shape': tuple - original shape

# Decompress for attention output computation
values_reconstructed = value_compressor.decompress(compressed_v)
# Shape: (batch, num_heads, seq_len, head_dim)

# Compute attention output: attention_weights @ values
attention_weights = torch.softmax(torch.randn(1, 8, 1, seq_len, device="cuda"), dim=-1)
attention_output = torch.matmul(attention_weights, values_reconstructed)
print(f"Attention output shape: {attention_output.shape}")  # (1, 8, 1, 128)
```
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.