nexusquant

NexusQuant: E8 Lattice KV Cache Compression

Training-free KV cache compression for LLM inference. Uses E8 lattice vector quantization + Hadamard rotation. Calibration-free.

Headline

+0.276% wikitext PPL on Mistral-7B-v0.1 (K3V2 pb=0, n=60, KIVI-iso). 5.83x compression measured on Mistral-Inst-v0.3 (same K3V2 config). NIAH retrieval preserved through 32K context on A100. Validated on 10 architecture families.

Head-to-head (Llama-3.1-8B-Instruct, 4K, n=30, iso-bpe)

NQ-vs-TQ McNemar b=30, c=0 (30 discordant pairs, all favor NexusQuant); FP16-vs-TQ floor b=29, c=0.

Install

pip install "nexusquant-kv[hf]"

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
from nexusquant import compress_kv_cache

model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
inputs = tokenizer("The KV cache stores key and value tensors for each attention layer.", return_tensors="pt").to(model.device)

cache = DynamicCache()
with torch.no_grad():
    out = model(inputs["input_ids"][:, :-1], use_cache=True, past_key_values=cache)

compressed = compress_kv_cache(out.past_key_values, mode="quant_only", key_bits=3, value_bits=2)
with torch.no_grad():
    gen = model.generate(inputs["input_ids"], past_key_values=compressed, max_new_tokens=200, do_sample=False)

print(tokenizer.decode(gen[0], skip_special_tokens=True))

Links

• Live demo
• GitHub
• Papers: Method, Architectures, Retrieval, Downstream, Competitors, System
• llama.cpp PR
• vLLM PR