Long-context reasoning models like DeepSeek-R1 and Qwen3 generate tens of thousands of tokens per response, and the KV cache grows linearly with sequence length. At 32K tokens, a single request on a 32B-parameter model can consume the entire VRAM of a consumer GPU (24 GB), making long-context inference impractical without compression. Existing KV-cache compression methods (SnapKV, H2O, R-KV) score tokens using “post-RoPE” attention patterns computed over a small local window. This introduces two problems: the scores depend on the current query (requiring recomputation at every step) and the local window misses tokens whose importance only becomes apparent later, causing permanent eviction of useful context.
TriAttention, introduced by Mao et al. (arXiv:2604.04921, 2026), takes a fundamentally different approach. It exploits an empirical observation: in the pre-RoPE (before Rotary Position Embedding is applied) space, Query and Key vectors cluster tightly around stable, fixed centers—a property the authors call Q/K concentration. When concentration is high, attention logits between any query and any key reduce to a trigonometric function of their positional distance, making future attention patterns predictable from position alone without observing the actual query. This allows TriAttention to score and evict KV-cache entries offline, achieving 2.5x throughput improvement and 10.7x KV memory reduction on reasoning benchmarks while matching full-attention accuracy.
The practical impact is significant: TriAttention enables a 32B reasoning model to run on a single 24 GB RTX 4090 that would otherwise run out of memory under full attention, and it integrates directly with vLLM for production serving.
In standard transformer attention, Query and Key vectors are first computed from the hidden state, then Rotary Position Embedding (RoPE) is applied to inject positional information. TriAttention examines these vectors before RoPE is applied (the “pre-RoPE” space) and observes that across positions, Q and K vectors in each attention head cluster around a fixed, non-zero center. The degree of clustering is measured by the Mean Resultant Length (R), a circular statistics metric ranging from 0 (uniform spread) to 1 (perfect concentration). Empirically, approximately 90% of attention heads in Qwen3-8B exhibit R > 0.95, meaning near-perfect clustering.
When Q and K vectors are concentrated around stable centers, the dot product (attention logit) between a query at position i and a key at position j depends primarily on the positional distance Δ = i − j rather than on the specific token content. The logit can be approximated as a trigonometric series: logit(Δ) ≈ Σ[axcos(ωxΔ) + bxsin(ωxΔ)], where the coefficients are derived from the Q/K centers and the RoPE frequencies. This means attention patterns are predictable from positional geometry alone—no live query needed.
Not all attention heads exhibit high Q/K concentration. For heads with lower R values, the trigonometric approximation is inaccurate. TriAttention falls back to a Norm-Based Score that uses the magnitude (L2 norm) of Key vectors as an importance signal. Keys with larger norms tend to attract more attention regardless of position. The two scores are combined using a weighted average, with the weight determined by each head’s R value: high-R heads rely on the trigonometric score; low-R heads rely on the norm score.
The critical advantage of TriAttention over prior methods is that Q/K centers and their derived trigonometric coefficients are computed offline (once, from a calibration set) and remain fixed during inference. This eliminates the per-step scoring overhead that methods like SnapKV and H2O incur, making compression essentially free at inference time.
Before deployment, TriAttention runs a calibration pass over a small dataset to compute, for each attention head: (1) the mean Q center and K center in pre-RoPE space, (2) the Mean Resultant Length R, and (3) the trigonometric coefficients derived from the centers and RoPE frequencies. These statistics are saved as a “sparse stats” file that the runtime loads at startup.
# Generate calibration statistics (run once per model)
python scripts/cli.py calibrate \
--model Qwen3-8B \
--dataset calibration_set \
--output stats/qwen3_8b_sparse_stats.ptDuring inference, the KV cache grows until it reaches a configured budget (e.g., 2048 tokens). At that point, TriAttention scores every cached token using the precomputed trigonometric and norm-based scores, evicts the lowest-scoring entries, and continues generation. This compression step triggers every DIVIDE_LENGTH tokens (default: 128). A recent-token window (default: 128 tokens) is always protected from eviction to preserve local context.
# Configure TriAttention runtime via environment variables
export TRIATTN_RUNTIME_KV_BUDGET=2048 # Max cached tokens
export TRIATTN_RUNTIME_DIVIDE_LENGTH=128 # Compression interval
export TRIATTN_RUNTIME_WINDOW_SIZE=128 # Protected recent window
export TRIATTN_RUNTIME_SPARSE_STATS_PATH=stats/qwen3_8b_sparse_stats.pt
export TRIATTN_RUNTIME_PRUNING_MODE=per_head # Per-head selection
# Launch vLLM with TriAttention monkeypatch
python -c "
from triattention.vllm.runtime.integration_monkeypatch import (
install_vllm_integration_monkeypatches
)
install_vllm_integration_monkeypatches(
patch_scheduler=True, patch_worker=True
)
from vllm import LLM, SamplingParams
llm = LLM(
model='Qwen/Qwen3-8B',
dtype='bfloat16',
max_model_len=32768,
enforce_eager=True,
)
output = llm.generate(['Solve: ...'], SamplingParams(max_tokens=16384))
"For each attention head h and each cached position j, the combined importance score is: S(j) = Rh · Strig(j) + (1 - Rh) · Snorm(j), where Rh is the head’s Mean Resultant Length. High-concentration heads (R ≈ 1) rely almost entirely on the trigonometric prediction; low-concentration heads rely on the key-norm signal. Tokens with the lowest combined scores across all heads are evicted from the cache.
For production deployment, TriAttention integrates with vLLM’s server mode to expose a standard chat/completions endpoint:
# Serve with TriAttention compression enabled
export TRIATTN_RUNTIME_KV_BUDGET=12000 # Higher budget for chat
vllm serve Qwen/Qwen3-8B \
--dtype bfloat16 \
--max-model-len 32768 \
--enforce-eager \
--enable-prefix-caching false \
--max-num-batched-tokens 1024Mathematical reasoning models (e.g., DeepSeek-R1-32B) produce chain-of-thought sequences exceeding 32K tokens. Without compression, this requires multiple high-end GPUs. TriAttention’s 10.7x KV memory reduction enables these models to run on a single RTX 4090 (24 GB), democratizing access to long-context reasoning capabilities.
In production serving environments, KV-cache memory is the primary constraint on batch size. By compressing each request’s cache to a fixed budget, TriAttention allows more concurrent requests per GPU. The 2.5x throughput improvement on reasoning workloads translates directly to lower per-token serving costs.
RAG pipelines that inject long retrieved documents into the context window benefit from TriAttention’s ability to retain retrieval-relevant tokens while evicting padding and boilerplate. On the RULER retrieval benchmark, TriAttention scores 66.1 versus SnapKV’s 55.6, indicating better preservation of retrieval-relevant context.
Chat applications accumulate context across turns. TriAttention’s trigonometric scoring naturally prioritizes structurally important tokens (turn boundaries, key facts) over filler, making it well-suited for compressing long conversation histories without losing critical context.
TriAttention modifies vLLM’s scheduler and worker modules at runtime via monkeypatching rather than forking the framework. This allows it to track upstream vLLM releases without maintaining a separate fork. The install_vllm_integration_monkeypatches function must be called before any vLLM class is instantiated, as it replaces internal scheduling and attention computation functions.
The PRUNING_MODE setting controls whether token eviction decisions are made independently per attention head (per_head) or globally across all heads (global). Per-head mode allows different heads to retain different tokens, preserving head-specific attention patterns. Global mode evicts the same tokens across all heads, simplifying memory management at the cost of potentially losing tokens important to specific heads.
For interactive chat workloads (as opposed to single-shot reasoning), the recommended configuration uses a higher KV budget (12,000 tokens), disables prefix caching (incompatible with dynamic eviction), and limits prefill chunk size to 1,024 tokens to reduce latency spikes during compression steps.
The scoring algorithm is simple enough to port to other inference runtimes. Community implementations exist for llama.cpp (C/ggml, targeting AMD GPUs) and MLX (Apple Silicon). The core requirement is access to pre-RoPE Q/K vectors and the ability to intercept the KV-cache eviction policy.
--enable-prefix-caching false when using TriAttention.