starred/spectralquant
1
0 Fork 0
mirror of https://github.com/Dynamis-Labs/spectralquant.git synced 2026-10-21 17:52:16 -05:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
spectralquant/experiments/shaped_cache_sweep.py
T
Ashwin Gopinath 1254311a9c Add shaped cache sweep experiment, results, and figures 
- experiments/shaped_cache_sweep.py
- results/shaped_cache/ (Qwen-1.5B, Qwen-7B, Llama-8B, combined)
- figures/ (pareto, kv_sensitivity, heatmap)
- gen_shaped_cache_figures.py
2026-04-07 01:30:31 +05:30

808 lines
31 KiB
Python
Raw Permalink Blame History

"""
Asymmetric Shaped KV Cache: Exploiting Key/Value Spectral Asymmetry
====================================================================
Keys have d_eff ~ 4, values have d_eff ~ 40-55. Instead of storing all 128
dimensions for both, store only m dimensions for keys and p dimensions for
values after spectral rotation. This is pure truncation in the eigenbasis --
the spectral rotation ensures we keep the highest-variance directions.
Sweep:
m (key dims): {2, 4, 8, 16, 32}
p (value dims): {8, 16, 32, 48, 64, 96, 128}
quant: {FP16 (baseline), 4-bit, 3-bit, 2-bit}
models: Qwen 2.5-1.5B, Qwen 2.5-7B, Llama 3.1-8B
Metrics at each (m, p, quant) point:
- Attention output cosine similarity vs FP16 (no compression)
- Compression ratio (bits per token vs FP16 baseline)
- Key reconstruction cosine similarity
- Value reconstruction cosine similarity
Theory: at m=4, p=32, 2-bit quant, we get:
keys: 4 dims * 2 bits = 8 bits + 16 bits norm = 24 bits
values: 32 dims * 3 bits = 96 bits + 16 bits norm = 112 bits
total = 136 bits/token vs 4096 bits FP16 = 30.1x compression
Usage:
python shaped_cache_sweep.py [--quick] [--device cuda]
Estimated runtime: ~60 min full, ~10 min quick (on B200)
"""
import sys, os, math, time, json, logging, argparse, gc
from pathlib import Path
from itertools import product
import torch
import torch.nn.functional as F
import numpy as np
# -- project paths --
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT / "src"))
sys.path.insert(0, str(PROJECT_ROOT / "baseline" / "turboquant_cutile"))
from turboquant_cutile import TurboQuantEngine, LloydMaxCodebook
log = logging.getLogger("shaped_cache")
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
)
RESULTS_DIR = PROJECT_ROOT / "results" / "shaped_cache"
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
HF_TOKEN = os.environ.get("HF_TOKEN")
# ============================================================================
# CONFIGURATION
# ============================================================================
MODELS = [
("Qwen/Qwen2.5-1.5B-Instruct", "Qwen-1.5B"),
("Qwen/Qwen2.5-7B-Instruct", "Qwen-7B"),
]
if HF_TOKEN:
MODELS.append(("meta-llama/Llama-3.1-8B-Instruct", "Llama-8B"))
M_VALUES = [2, 4, 8, 16, 32]
P_VALUES = [8, 16, 32, 48, 64, 96, 128]
QUANT_BITS = [0, 4, 3, 2]
N_CALIB = 64
N_EVAL = 32
SEQ_LEN = 512
# Quick overrides
N_CALIB_QUICK = 16
N_EVAL_QUICK = 8
M_VALUES_QUICK = [4, 8, 16]
P_VALUES_QUICK = [16, 32, 64, 128]
QUANT_BITS_QUICK = [0, 3, 2]
# ============================================================================
# UTILITIES
# ============================================================================
def save_result(filename, data):
path = RESULTS_DIR / filename
with open(path, "w") as f:
json.dump(data, f, indent=2, default=str)
log.info("Saved: %s", path)
def quantize_nearest(x: torch.Tensor, centroids: torch.Tensor) -> torch.Tensor:
"""Scalar nearest-centroid quantization."""
c = centroids.to(x.device)
diffs = x.unsqueeze(-1) - c
return c[diffs.abs().argmin(dim=-1).long()]
def solve_lloyd_max_for_sigma(sigma: float, bits: int, max_iter=200, tol=1e-10):
"""Solve Lloyd-Max for N(0, sigma^2)."""
from scipy import integrate
n_levels = 1 << bits
pdf = lambda x: (
(1.0 / (math.sqrt(2 * math.pi) * sigma))
* math.exp(-x * x / (2 * sigma * sigma))
)
lo, hi = -3.5 * sigma, 3.5 * sigma
centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
for _ in range(max_iter):
boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)]
edges = [lo * 3] + boundaries + [hi * 3]
new_centroids = []
for i in range(n_levels):
a, b = edges[i], edges[i + 1]
num, _ = integrate.quad(lambda x: x * pdf(x), a, b)
den, _ = integrate.quad(pdf, a, b)
new_centroids.append(num / den if den > 1e-15 else centroids[i])
if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol:
break
centroids = new_centroids
return torch.tensor(centroids, dtype=torch.float32)
def load_model_tokenizer(model_name: str, device: str):
"""Load a HuggingFace model + tokenizer."""
from transformers import AutoModelForCausalLM, AutoTokenizer
needs_token = any(x in model_name for x in ["llama", "Llama", "gemma", "Gemma"])
token = HF_TOKEN if needs_token else None
log.info("Loading %s ...", model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map=device,
token=token,
)
model.eval()
cfg = model.config
n_layers = cfg.num_hidden_layers
n_kv = getattr(cfg, "num_key_value_heads", cfg.num_attention_heads)
hd = getattr(cfg, 'head_dim', cfg.hidden_size // cfg.num_attention_heads)
# Detect actual KV head_dim from a forward pass
try:
test_ids = tokenizer("test", return_tensors="pt").to(device)
with torch.no_grad():
test_out = model(**test_ids, use_cache=True)
kv = test_out.past_key_values
try:
actual_kv_hd = kv.key_cache[0].shape[-1]
except:
try:
actual_kv_hd = kv[0][0].shape[-1]
except:
actual_kv_hd = hd
if actual_kv_hd != hd:
log.info("KV head_dim=%d differs from Q head_dim=%d, using KV head_dim", actual_kv_hd, hd)
hd = actual_kv_hd
except Exception as e:
log.warning("Could not detect KV head_dim: %s", e)
n_params = sum(p.numel() for p in model.parameters()) / 1e9
log.info("%s: %d layers, %d kv_heads, head_dim=%d, %.1fB params",
model_name, n_layers, n_kv, hd, n_params)
return model, tokenizer, n_layers, n_kv, hd
def extract_kv_layer(kv, l):
"""Extract (key, value) tensors for layer l from past_key_values."""
try:
return kv.key_cache[l], kv.value_cache[l]
except Exception:
pass
try:
return kv[l][0], kv[l][1]
except Exception:
pass
entry = list(kv)[l]
return entry[0], entry[1]
# ============================================================================
# CALIBRATION: per-(layer, head) eigenvectors for KEYS and VALUES separately
# Uses the same method as neurips_models_asymmetry.py (accumulate X^T X)
# ============================================================================
def calibrate_keys_and_values(model, tokenizer, n_calib, device, n_layers, n_kv, hd):
"""
Calibrate SEPARATELY for keys and values.
Accumulates X^T X (second moment, not centered covariance) which is the
correct formulation for the participation ratio in this context.
Returns:
eigen_keys: {(l,h): {"evec": Tensor, "ev": Tensor, "d_eff": int}}
eigen_vals: {(l,h): {"evec": Tensor, "ev": Tensor, "d_eff": int}}
"""
from datasets import load_dataset
ds = load_dataset("wikitext", "wikitext-103-raw-v1",
split=f"train[:{n_calib * 5}]")
cov_keys = {
(l, h): {"xtx": torch.zeros(hd, hd, dtype=torch.float64), "n": 0}
for l in range(n_layers) for h in range(n_kv)
}
cov_vals = {
(l, h): {"xtx": torch.zeros(hd, hd, dtype=torch.float64), "n": 0}
for l in range(n_layers) for h in range(n_kv)
}
nd = 0
t0 = time.time()
for item in ds:
text = item.get("text", "")
if len(text.strip()) < 100:
continue
enc = tokenizer(text, return_tensors="pt",
max_length=512, truncation=True).to(device)
if enc["input_ids"].shape[1] < 16:
continue
with torch.no_grad():
out = model(**enc, use_cache=True)
kv = out.past_key_values
for l in range(n_layers):
try:
k_l, v_l = extract_kv_layer(kv, l)
k_l = k_l.float().cpu()
v_l = v_l.float().cpu()
except Exception:
continue
for h in range(min(n_kv, k_l.shape[1])):
X_key = k_l[0, h, :, :].double() # (seq, hd)
cov_keys[(l, h)]["xtx"] += X_key.T @ X_key
cov_keys[(l, h)]["n"] += X_key.shape[0]
X_val = v_l[0, h, :, :].double()
cov_vals[(l, h)]["xtx"] += X_val.T @ X_val
cov_vals[(l, h)]["n"] += X_val.shape[0]
nd += 1
if nd >= n_calib:
break
if nd % 20 == 0:
log.info(" Calibration: %d/%d (%.0fs)", nd, n_calib, time.time() - t0)
def _eigendecompose(cov_dict):
eigen = {}
for l in range(n_layers):
for h in range(n_kv):
n = cov_dict[(l, h)]["n"]
if n == 0:
eigen[(l, h)] = {
"evec": torch.eye(hd),
"ev": torch.ones(hd),
"d_eff": hd,
}
continue
C = (cov_dict[(l, h)]["xtx"] / n).float()
ev, evec = torch.linalg.eigh(C)
# Sort descending
ev = ev.flip(0).clamp(min=0)
evec = evec.flip(1)
d_eff = max(2, min(
int(round((ev.sum() ** 2 / (ev ** 2).sum()).item())),
hd - 2
))
eigen[(l, h)] = {"evec": evec, "ev": ev, "d_eff": d_eff}
return eigen
eigen_keys = _eigendecompose(cov_keys)
eigen_vals = _eigendecompose(cov_vals)
mean_deff_k = float(np.mean([eigen_keys[(l, h)]["d_eff"]
for l in range(n_layers) for h in range(n_kv)]))
mean_deff_v = float(np.mean([eigen_vals[(l, h)]["d_eff"]
for l in range(n_layers) for h in range(n_kv)]))
log.info("Calibration done in %.1fs. mean d_eff: keys=%.1f, values=%.1f",
time.time() - t0, mean_deff_k, mean_deff_v)
return eigen_keys, eigen_vals
# ============================================================================
# SHAPED CACHE ENGINE
# ============================================================================
class ShapedCacheEngine:
"""Asymmetric shaped KV cache.
After spectral rotation (V^T), truncate to keep only the top-m dimensions
for keys and top-p dimensions for values. Optionally quantize the retained
dimensions.
Args:
key_eigvecs: (head_dim, head_dim) key eigenvectors sorted by eigenvalue desc
key_eigvals: (head_dim,) key eigenvalues sorted desc
val_eigvecs: (head_dim, head_dim) value eigenvectors sorted by eigenvalue desc
val_eigvals: (head_dim,) value eigenvalues sorted desc
m: number of key dimensions to retain
p: number of value dimensions to retain
quant_bits: 0 for FP16 (no quantization), 2/3/4 for scalar quantization
head_dim: original head dimension
"""
def __init__(self, key_eigvecs, key_eigvals, val_eigvecs, val_eigvals,
m, p, quant_bits=0, head_dim=128):
self.head_dim = head_dim
self.m = min(m, head_dim)
self.p = min(p, head_dim)
self.quant_bits = quant_bits
Vk = key_eigvecs.float()
Vv = val_eigvecs.float()
assert Vk.shape == (head_dim, head_dim)
assert Vv.shape == (head_dim, head_dim)
# Truncated rotation matrices
self.Vk_trunc = Vk[:, :self.m].contiguous() # (hd, m)
self.Vv_trunc = Vv[:, :self.p].contiguous() # (hd, p)
self.Vk_trunc_T = self.Vk_trunc.T.contiguous() # (m, hd)
self.Vv_trunc_T = self.Vv_trunc.T.contiguous() # (p, hd)
# Build per-regime quantization codebooks
self.key_centroids = None
self.val_centroids = None
if quant_bits > 0:
ev_k = key_eigvals.float().clamp(min=0)
sigma_k = float(ev_k[:self.m].mean().clamp(min=1e-8).sqrt().item())
self.key_centroids = solve_lloyd_max_for_sigma(sigma_k, quant_bits)
ev_v = val_eigvals.float().clamp(min=0)
sigma_v = float(ev_v[:self.p].mean().clamp(min=1e-8).sqrt().item())
self.val_centroids = solve_lloyd_max_for_sigma(sigma_v, quant_bits)
def compress_keys(self, K):
"""Compress keys: rotate + truncate to m dims + optional quantize."""
K_f = K.float()
vec_norms = torch.norm(K_f, dim=-1, keepdim=True)
K_normed = K_f / (vec_norms + 1e-8)
rotated = K_normed @ self.Vk_trunc.to(K.device) # (seq, m)
if self.quant_bits > 0 and self.key_centroids is not None:
rotated = quantize_nearest(rotated, self.key_centroids)
return {"rotated": rotated, "vec_norms": vec_norms.squeeze(-1)}
def compress_values(self, V):
"""Compress values: rotate + truncate to p dims + optional quantize."""
V_f = V.float()
vec_norms = torch.norm(V_f, dim=-1, keepdim=True)
V_normed = V_f / (vec_norms + 1e-8)
rotated = V_normed @ self.Vv_trunc.to(V.device) # (seq, p)
if self.quant_bits > 0 and self.val_centroids is not None:
rotated = quantize_nearest(rotated, self.val_centroids)
return {"rotated": rotated, "vec_norms": vec_norms.squeeze(-1)}
def decompress_keys(self, compressed_k):
"""Reconstruct keys from compressed representation."""
rotated = compressed_k["rotated"].float()
norms = compressed_k["vec_norms"].float()
device = rotated.device
return rotated @ self.Vk_trunc_T.to(device) * norms.unsqueeze(-1)
def decompress_values(self, compressed_v):
"""Reconstruct values from compressed representation."""
rotated = compressed_v["rotated"].float()
norms = compressed_v["vec_norms"].float()
device = rotated.device
return rotated @ self.Vv_trunc_T.to(device) * norms.unsqueeze(-1)
def attention_output(self, Q, compressed_k, compressed_v):
"""Compute full attention: softmax(Q @ K_hat^T / sqrt(d)) @ V_hat."""
K_hat = self.decompress_keys(compressed_k)
V_hat = self.decompress_values(compressed_v)
scale = 1.0 / math.sqrt(self.head_dim)
scores = Q.float() @ K_hat.T * scale
weights = F.softmax(scores, dim=-1)
return weights @ V_hat
def bits_per_token(self):
"""Compressed bits per token for this (m, p, quant) config."""
if self.quant_bits > 0:
key_bits = self.m * self.quant_bits + 16 # m quantized dims + FP16 norm
val_bits = self.p * self.quant_bits + 16 # p quantized dims + FP16 norm
else:
key_bits = self.m * 16 + 16 # m FP16 dims + FP16 norm
val_bits = self.p * 16 + 16 # p FP16 dims + FP16 norm
return key_bits + val_bits
def compression_ratio(self):
"""Compression ratio vs FP16 KV cache."""
fp16_bits = 2 * self.head_dim * 16
return fp16_bits / self.bits_per_token()
# ============================================================================
# EVALUATION
# ============================================================================
def evaluate_config(engine, Q_all, K_all, V_all, fp16_output):
"""Evaluate a single (m, p, quant) configuration."""
compressed_k = engine.compress_keys(K_all)
compressed_v = engine.compress_values(V_all)
K_hat = engine.decompress_keys(compressed_k)
V_hat = engine.decompress_values(compressed_v)
key_cos = F.cosine_similarity(K_all.float(), K_hat.float(), dim=-1).mean().item()
val_cos = F.cosine_similarity(V_all.float(), V_hat.float(), dim=-1).mean().item()
attn_out = engine.attention_output(Q_all, compressed_k, compressed_v)
attn_cos = F.cosine_similarity(fp16_output.float(), attn_out.float(), dim=-1).mean().item()
return {
"key_cos_sim": key_cos,
"val_cos_sim": val_cos,
"attn_cos_sim": attn_cos,
"compression_ratio": engine.compression_ratio(),
"bits_per_token": engine.bits_per_token(),
"m": engine.m,
"p": engine.p,
"quant_bits": engine.quant_bits,
}
def run_sweep_for_model(model_name, short_name, m_values, p_values, quant_bits_list,
n_calib, n_eval, device):
"""Run the full (m, p, quant) sweep for one model."""
log.info("=" * 70)
log.info("MODEL: %s", model_name)
log.info("=" * 70)
model, tokenizer, n_layers, n_kv, hd = load_model_tokenizer(model_name, device)
# Calibrate eigenbasis for keys AND values separately
log.info("Calibrating eigenbasis from %d sequences...", n_calib)
eigen_keys, eigen_vals = calibrate_keys_and_values(
model, tokenizer, n_calib, device, n_layers, n_kv, hd
)
# Get d_eff stats
deff_keys = [eigen_keys[(l, h)]["d_eff"] for l in range(n_layers) for h in range(n_kv)]
deff_vals = [eigen_vals[(l, h)]["d_eff"] for l in range(n_layers) for h in range(n_kv)]
mean_deff_keys = float(np.mean(deff_keys))
mean_deff_vals = float(np.mean(deff_vals))
log.info("d_eff: keys=%.2f, values=%.2f, asymmetry=%.1fx",
mean_deff_keys, mean_deff_vals, mean_deff_vals / max(mean_deff_keys, 0.01))
# Representative layers + heads for evaluation
layer_indices = sorted(set([
n_layers // 4,
n_layers // 2,
3 * n_layers // 4,
]))
head_indices = list(range(min(n_kv, 4)))
log.info("Evaluating on layers=%s, heads=%s", layer_indices, head_indices)
# Collect evaluation KV data (separate forward passes)
log.info("Collecting evaluation data (%d sequences)...", n_eval)
from datasets import load_dataset
ds = load_dataset("wikitext", "wikitext-103-raw-v1",
split=f"train[{n_calib * 5}:{(n_calib + n_eval) * 5}]")
# Store eval KV per (layer, head)
eval_keys = {(l, h): [] for l in layer_indices for h in head_indices}
eval_vals = {(l, h): [] for l in layer_indices for h in head_indices}
nd = 0
for item in ds:
text = item.get("text", "")
if len(text.strip()) < 100:
continue
enc = tokenizer(text, return_tensors="pt",
max_length=512, truncation=True).to(device)
if enc["input_ids"].shape[1] < 32:
continue
with torch.no_grad():
out = model(**enc, use_cache=True)
kv = out.past_key_values
for l in layer_indices:
try:
k_l, v_l = extract_kv_layer(kv, l)
k_l = k_l.float().cpu()
v_l = v_l.float().cpu()
except Exception:
continue
for h in head_indices:
if h < k_l.shape[1]:
eval_keys[(l, h)].append(k_l[0, h])
eval_vals[(l, h)].append(v_l[0, h])
nd += 1
if nd >= n_eval:
break
log.info("Collected eval data from %d sequences", nd)
# Run the sweep
all_results = []
total_configs = len(m_values) * len(p_values) * len(quant_bits_list)
config_idx = 0
for m, p, qb in product(m_values, p_values, quant_bits_list):
config_idx += 1
if m > p:
continue # keys should use fewer dims than values
if m > hd or p > hd:
continue
if config_idx % 20 == 1:
log.info("Config %d/%d: m=%d, p=%d, quant=%s",
config_idx, total_configs, m, p, f"{qb}-bit" if qb > 0 else "FP16")
config_metrics = []
for l in layer_indices:
for h in head_indices:
kv_k = eval_keys.get((l, h), [])
kv_v = eval_vals.get((l, h), [])
if not kv_k or not kv_v:
continue
K_all = torch.cat(kv_k, dim=0)[:512].to(device).float()
V_all = torch.cat(kv_v, dim=0)[:512].to(device).float()
Q_all = K_all.clone() # Use keys as queries for self-attention test
if K_all.shape[0] < 32:
continue
# FP16 ground truth
scale = 1.0 / math.sqrt(hd)
fp16_scores = Q_all @ K_all.T * scale
fp16_weights = F.softmax(fp16_scores, dim=-1)
fp16_output = fp16_weights @ V_all
engine = ShapedCacheEngine(
key_eigvecs=eigen_keys[(l, h)]["evec"].to(device),
key_eigvals=eigen_keys[(l, h)]["ev"].to(device),
val_eigvecs=eigen_vals[(l, h)]["evec"].to(device),
val_eigvals=eigen_vals[(l, h)]["ev"].to(device),
m=m, p=p, quant_bits=qb, head_dim=hd,
)
metrics = evaluate_config(engine, Q_all, K_all, V_all, fp16_output)
metrics["layer"] = l
metrics["head"] = h
config_metrics.append(metrics)
if config_metrics:
avg = {
"m": m, "p": p,
"quant_bits": qb,
"quant_label": f"{qb}-bit" if qb > 0 else "FP16",
"key_cos_sim": float(np.mean([r["key_cos_sim"] for r in config_metrics])),
"val_cos_sim": float(np.mean([r["val_cos_sim"] for r in config_metrics])),
"attn_cos_sim": float(np.mean([r["attn_cos_sim"] for r in config_metrics])),
"attn_cos_sim_std": float(np.std([r["attn_cos_sim"] for r in config_metrics])),
"key_cos_sim_std": float(np.std([r["key_cos_sim"] for r in config_metrics])),
"val_cos_sim_std": float(np.std([r["val_cos_sim"] for r in config_metrics])),
"compression_ratio": config_metrics[0]["compression_ratio"],
"bits_per_token": config_metrics[0]["bits_per_token"],
"n_evals": len(config_metrics),
}
all_results.append(avg)
# Compute baselines: TQ 3-bit, SQ 3-bit (full-dim)
log.info("Computing baselines...")
baselines = compute_baselines(
eigen_keys, eigen_vals, eval_keys, eval_vals,
layer_indices, head_indices, n_kv, hd, device
)
# Per-layer d_eff breakdown
deff_by_layer_k = {}
deff_by_layer_v = {}
for l in range(n_layers):
dk = [eigen_keys[(l, h)]["d_eff"] for h in range(n_kv)]
dv = [eigen_vals[(l, h)]["d_eff"] for h in range(n_kv)]
deff_by_layer_k[l] = float(np.mean(dk))
deff_by_layer_v[l] = float(np.mean(dv))
# Eigenvalue spectra for representative layers (for plotting)
spectra = {}
for l in layer_indices:
for h in head_indices[:1]: # just head 0
spectra[f"layer_{l}_head_{h}_keys"] = eigen_keys[(l, h)]["ev"][:32].tolist()
spectra[f"layer_{l}_head_{h}_values"] = eigen_vals[(l, h)]["ev"][:32].tolist()
result = {
"model": model_name,
"short_name": short_name,
"head_dim": hd,
"n_layers": n_layers,
"n_kv_heads": n_kv,
"mean_d_eff_keys": mean_deff_keys,
"mean_d_eff_values": mean_deff_vals,
"asymmetry_ratio": mean_deff_vals / max(mean_deff_keys, 0.01),
"d_eff_by_layer_keys": deff_by_layer_k,
"d_eff_by_layer_values": deff_by_layer_v,
"eigenvalue_spectra": spectra,
"sweep_results": all_results,
"baselines": baselines,
"config": {
"m_values": m_values,
"p_values": p_values,
"quant_bits": quant_bits_list,
"n_calib": n_calib,
"n_eval": n_eval,
"layers_evaluated": layer_indices,
"heads_evaluated": head_indices,
},
}
save_result(f"shaped_cache_{short_name}.json", result)
del model
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return result
def compute_baselines(eigen_keys, eigen_vals, eval_keys, eval_vals,
layer_indices, head_indices, n_kv, hd, device):
"""Compute TQ and SQ baselines for comparison."""
from spectralquant.engine import SpectralQuantEngine
baselines = {}
for method_name, use_spectral in [("TQ_3bit", False), ("SQ_noQJL_3bit", True)]:
cos_sims = []
for l in layer_indices:
for h in head_indices:
kv_k = eval_keys.get((l, h), [])
kv_v = eval_vals.get((l, h), [])
if not kv_k or not kv_v:
continue
K_all = torch.cat(kv_k, dim=0)[:512].to(device).float()
V_all = torch.cat(kv_v, dim=0)[:512].to(device).float()
Q_all = K_all.clone()
if K_all.shape[0] < 32:
continue
scale = 1.0 / math.sqrt(hd)
fp16_scores = Q_all @ K_all.T * scale
fp16_weights = F.softmax(fp16_scores, dim=-1)
fp16_output = fp16_weights @ V_all
try:
if use_spectral:
evec = eigen_keys[(l, h)]["evec"]
ev = eigen_keys[(l, h)]["ev"]
d_eff = eigen_keys[(l, h)]["d_eff"]
VT = evec.T.contiguous().to(device).float()
Vm = evec.to(device).float()
k_n = torch.norm(K_all, dim=-1, keepdim=True)
K_rot = (K_all / (k_n + 1e-8)) @ VT
v_n = torch.norm(V_all, dim=-1, keepdim=True)
V_rot = (V_all / (v_n + 1e-8)) @ VT
cb_k = LloydMaxCodebook(hd, 2)
cb_v = LloydMaxCodebook(hd, 3)
K_hat = quantize_nearest(K_rot, cb_k.centroids.to(device))
V_hat = quantize_nearest(V_rot, cb_v.centroids.to(device))
K_mse = (K_hat @ Vm) * k_n
V_rec = (V_hat @ Vm) * v_n
scores = Q_all @ K_mse.T * scale
output = F.softmax(scores, dim=-1) @ V_rec
else:
tq = TurboQuantEngine(d=hd, mse_bits=2, seed=42, device=device)
ck = tq.compress_keys_pytorch(K_all.half())
cv = tq.compress_values_pytorch(V_all.half())
scores = tq.attention_scores_pytorch(Q_all.half(), ck)
weights = F.softmax(scores.float(), dim=-1)
V_hat = tq.decompress_values_pytorch(cv)
output = weights @ V_hat.float()
cos = F.cosine_similarity(
fp16_output.float(), output.float(), dim=-1
).mean().item()
cos_sims.append(cos)
except Exception as e:
log.warning("%s baseline failed l=%d h=%d: %s", method_name, l, h, e)
if cos_sims:
baselines[method_name] = {
"attn_cos_sim": float(np.mean(cos_sims)),
"attn_cos_sim_std": float(np.std(cos_sims)),
"compression_ratio": 5.95 if use_spectral else 5.02,
"n_evals": len(cos_sims),
}
log.info(" Baseline %s: cos=%.4f +/- %.4f (n=%d)",
method_name, baselines[method_name]["attn_cos_sim"],
baselines[method_name]["attn_cos_sim_std"],
len(cos_sims))
return baselines
# ============================================================================
# MAIN
# ============================================================================
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--quick", action="store_true")
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
quick = args.quick
device = args.device
m_vals = M_VALUES_QUICK if quick else M_VALUES
p_vals = P_VALUES_QUICK if quick else P_VALUES
qb_vals = QUANT_BITS_QUICK if quick else QUANT_BITS
n_calib = N_CALIB_QUICK if quick else N_CALIB
n_eval = N_EVAL_QUICK if quick else N_EVAL
log.info("=" * 70)
log.info("SHAPED CACHE SWEEP: %s mode", "QUICK" if quick else "FULL")
log.info("m_values=%s, p_values=%s, quant=%s", m_vals, p_vals, qb_vals)
log.info("n_calib=%d, n_eval=%d", n_calib, n_eval)
log.info("=" * 70)
all_model_results = {}
t0 = time.time()
for model_name, short_name in MODELS:
try:
result = run_sweep_for_model(
model_name, short_name, m_vals, p_vals, qb_vals,
n_calib, n_eval, device,
)
all_model_results[short_name] = result
except Exception as e:
log.error("FAILED on %s: %s", model_name, e)
import traceback
traceback.print_exc()
elapsed = time.time() - t0
combined = {
"experiment": "shaped_cache_sweep",
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
"elapsed_seconds": elapsed,
"mode": "quick" if quick else "full",
"models": all_model_results,
}
save_result("shaped_cache_combined.json", combined)
# Print summary
log.info("\n" + "=" * 70)
log.info("SHAPED CACHE SWEEP COMPLETE (%.1f min)", elapsed / 60)
log.info("=" * 70)
for short_name, result in all_model_results.items():
log.info("\n--- %s (d_eff: keys=%.1f, values=%.1f, asymmetry=%.1fx) ---",
short_name, result["mean_d_eff_keys"], result["mean_d_eff_values"],
result["asymmetry_ratio"])
sweep = result["sweep_results"]
if sweep:
# Pareto analysis
high_quality = [r for r in sweep if r["attn_cos_sim"] > 0.95]
if high_quality:
best = max(high_quality, key=lambda r: r["compression_ratio"])
log.info(" Best >0.95 quality: m=%d, p=%d, %s, cos=%.4f, %.1fx",
best["m"], best["p"], best["quant_label"],
best["attn_cos_sim"], best["compression_ratio"])
high_compress = [r for r in sweep if r["compression_ratio"] > 10]
if high_compress:
best_hc = max(high_compress, key=lambda r: r["attn_cos_sim"])
log.info(" Best >10x compress: m=%d, p=%d, %s, cos=%.4f, %.1fx",
best_hc["m"], best_hc["p"], best_hc["quant_label"],
best_hc["attn_cos_sim"], best_hc["compression_ratio"])
best_overall = max(sweep, key=lambda r: r["attn_cos_sim"])
log.info(" Best quality: m=%d, p=%d, %s, cos=%.4f, %.1fx",
best_overall["m"], best_overall["p"], best_overall["quant_label"],
best_overall["attn_cos_sim"], best_overall["compression_ratio"])
baselines = result.get("baselines", {})
for bname, bdata in baselines.items():
log.info(" Baseline %s: cos=%.4f, %.2fx",
bname, bdata["attn_cos_sim"], bdata["compression_ratio"])
if __name__ == "__main__":
main()
Reference in new issue View Git Blame Copy Permalink