sam-gate 0.1.2

SAM-Gate: Semantic-Aware Memory Gate for heterogeneous KV-cache compression in transformer models

SAM-Gate

Semantic-Aware Memory Gate — adaptive KV-cache compression for transformer models, guided by the intrinsic geometry of the attention flow.

SAM-Gate estimates the harmonic curvature f(t) per layer at inference time and assigns each layer to a compression regime:

Regime	Condition	KV bits	Heads	Window
Flat	f(t) < f_flat_max	int4	40%	short
Transition	f_flat_max ≤ f(t) < f_obs_min	int8	70%	medium
Obstructed	f(t) ≥ f_obs_min	fp16	100%	long

The KV window is fixed regardless of context length — memory is O(1), not O(T).

Benchmark results (Qwen2.5-3B-Instruct, RTX 4070 8GB)

Engine	ctx=512 KV	ctx=1024 KV	TPS (ctx=512)
Baseline HF	18828 KB	37260 KB	30.4
SAM-Gate	4608 KB	4608 KB	8.5
RCI (int8)	9702 KB	19206 KB	9.2

SAM-Gate KV footprint does not grow with context. At ctx=1024, baseline uses 8x more memory than SAM-Gate.

Install

# Base install (Windows / Mac / Linux — uses PyTorch SDPA)
pip install sam-gate

# Linux / WSL with NVIDIA GPU — enables Flash Attention (2-3x TPS gain)
pip install sam-gate[flash]

# With SpectralRCI support
pip install sam-gate[spectral]

Note: flash-attn requires Linux or WSL2 with CUDA 11.7+. On Windows, SAM-Gate automatically falls back to PyTorch SDPA — fully functional, ~50% lower TPS at long contexts.

Quick start

from transformers import AutoModelForCausalLM, AutoTokenizer
from sam_gate import attach_semantic_hooks, SAMConfig

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B-Instruct")

cfg = SAMConfig(
    max_ctx_flat  = 64,
    max_ctx_trans = 128,
    max_ctx_obs   = 128,
)

kv_caches = {}
handles = attach_semantic_hooks(model, kv_caches, cfg=cfg)

inputs = tokenizer("Explain how attention works.", return_tensors="pt").to("cuda")
out = model.generate(**inputs, max_new_tokens=200, use_cache=False)

for h in handles:
    h.remove()

print(tokenizer.decode(out[0], skip_special_tokens=True))
print(f"KV used: {sum(c.ram_bytes() for c in kv_caches.values()) / 1024:.1f} KB")

Calibration

Thresholds f_flat_max and f_obs_min are model-specific. Before using SAM-Gate on a new model, run calibration to observe the real f(t) distribution:

python -m sam_gate.sam --model Qwen/Qwen2.5-3B-Instruct --device cuda --calibrate --verbose

Then set f_flat_max and f_obs_min in SAMConfig based on the observed values per layer.

Configuration reference

from sam_gate import SAMConfig

cfg = SAMConfig(
    tau           = 0.05,    # harmonic kernel threshold
    flat_bits     = 4,       # int4 in flat regime
    trans_bits    = 8,       # int8 in transition
    obs_bits      = 16,      # fp16 in obstructed
    flat_heads    = 0.5,     # fraction of query heads in flat regime
    trans_heads   = 0.75,    # fraction in transition
    f_flat_max    = 1e-2,    # calibrated for Qwen2.5-3B-Instruct
    f_obs_min     = 50.0,    # calibrated for Qwen2.5-3B-Instruct
    max_ctx_flat  = 64,      # KV window in flat regime (tokens)
    max_ctx_trans = 128,     # KV window in transition (tokens)
    max_ctx_obs   = 128,     # KV window in obstructed (tokens)
)

How it works

SAM-Gate estimates the Morse functional f(t) via Hutchinson probing — O(K·H·d) instead of O(d³) eigendecomposition:

Δ_t  = Σ_i (T_i - I)*(T_i - I)     ← semantic Laplacian
f(t) = Σ_{i<j} ||[T_i, T_j]||²_F   ← harmonic curvature

At decode time, the policy is cached per layer — zero overhead from the prober after prefill.

The KV cache uses a dense ring buffer (fp16, fixed capacity) initialized once from the quantized prefill chunks. At each decode step, reconstruct_for_attn is O(1) — a single slice, no dequantization loop.

License

MIT