Summary (Overview)

• OpenAttention is a novel sparse attention method that approximates the tail distribution of attention weights using proxy tokens rather than discarding or sampling the tail, achieving 99% of dense accuracy at just a 1% read rate on RULER-32K HARD with Llama-3.1-8B.
• The key theoretical insight is that a single count-aware mean proxy yields lower attention weight error than random sampling at low read densities—sampling nearly doubles the weight error of discarding the tail entirely.
• Three novel components are introduced: Query-Geometry Clustering (QGC) for grouping keys under query geometry, Query-Geometry Selector (QGS) for efficient top-k selection, and ErrorGate for budget allocation based on estimated weight error.
• OpenAttention achieves state-of-the-art accuracy across multiple models (Llama-3.1-8B, DeepSeek-R1-Distill-Llama-8B, Mistral-7B), context lengths (16K–96K), and outperforms baselines (HashAttention, vAttention, MagicPIG) by large margins.
• A custom Triton-based kernel, OpenAttention.flash, runs the estimator in one online softmax pass, achieving 5–10× faster decode attention than dense attention at 3% and 1% reads.

---

Introduction and Theoretical Foundation

Background and Motivation

Long-context decoding is bottlenecked by the memory traffic of the key-value (KV) cache. Sparse attention reduces this cost by computing attention over only a subset of the cache. However, existing top-k approaches (Singhania et al., 2024; Yang et al., 2024; Desai et al., 2025) focus solely on predicting the top-k scores accurately and discard the remaining tail. This fundamentally prevents these methods from approximating the dense attention distribution, which is especially critical for aggregation tasks (e.g., FWE).

Sampling-based methods like MagicPIG (Chen et al., 2025) and vAttention (Desai et al., 2026) attempt to estimate the tail via importance sampling or random samples. However, estimation error of sampling grows sharply as the density decreases, and at low reads (<5%), the tail likely holds heavy hitters the selector missed.

Key Observation

The paper's central finding (Figure 3): a random sample complements the tail well at high density but yields higher error than top-k alone at low density. However, folding the same tail into a single proxy with a count-aware mean keeps the error low in every setting examined.

Theoretical analysis (Section 2.2): Let $w_i = e^{z_i}$ with $z_i = d^{-1/2}q^\top k_i$. For a tail $T$ of size $N$, the error formulations are:

$$\hat{a}_i^{\mathrm{d}} = \frac{\mathbf{1}_X(i) w_i}{Z_X}, \qquad \hat{a}_i^{\mathrm{s}} = \frac{[\mathbf{1}_X(i) + \eta \mathbf{1}_B(i)] w_i}{\hat{Z}_{\mathrm{s}}}, \qquad \hat{a}_i^{\mathrm{p}} = \frac{\mathbf{1}_X(i) w_i + \mathbf{1}_T(i) g_T}{\hat{Z}_{\mathrm{p}}}, \tag{1}$$

where $g_T = e^{\bar{z}_T}$ is the count-aware mean, $\eta = N/m$ for sample size $m$, and $\hat{Z}_{\mathrm{s}} = Z_X + \eta \sum_{i \in B} w_j$, $\hat{Z}_{\mathrm{p}} = Z_X + N g_T$.

The key weight-discrepancy results (Eq. 2):

$$\mathcal{E}_{\mathrm{d}} = \tau, \qquad \mathbb{E}[\mathcal{E}_{\mathrm{s}}] = 2\tau\left(1 - \frac{m}{N}\right), \qquad \mathcal{E}_{\mathrm{p}}^{\mathrm{G}} = \tau \operatorname{erf}\left(\frac{\sigma_z}{\sqrt{2}}\right). \tag{2}$$

Critical finding: Sampling approaches twice the discard discrepancy when $m \ll N$, while a single proxy stays below it. This motivates approximating the tail with well-constructed proxies rather than samples.

---

Methodology

3.1 Overall Architecture

OpenAttention clusters the cache once per request into $C$ clusters under query geometry (QGC). Each cluster caches:

• Size $n_b$
• Mean token $(\bar{k}_b, \bar{v}_b)$
• Variance token diag $\mathrm{Var}_b(k)$

For each query, two small matrix products give every $\bar{z}_b$ and $\sigma_b^2$, from which ErrorGate estimates the error each proxy would leave. The exact budget $E$ is split in proportion to this estimate, and each cluster reads its highest-scoring tokens under QGS.

3.2 Query-Geometry Clustering (QGC) and Selection (QGS)

QGC: Uses top-$r$ eigenvectors of the query second-moment matrix $M = \frac{1}{2}\bar{M}_{\mathrm{st}} + \frac{1}{2}\bar{M}_q$. Since $\mathbb{E}_q(z_i - z_j)^2 = s^2(k_i - k_j)^\top M(k_i - k_j)$, minimizing within-cluster logit variance reduces to k-means under metric $M$. Keys are projected as:

$$\Phi = K U_r \Lambda_r^{1/2} \in \mathbb{R}^{n \times r}, \qquad a_i = \arg\min_b \|\phi_i - \mu_b\|^2, \qquad \mu_b = \frac{1}{n_b} \sum_{i: a_i = b} \phi_i, \tag{4}$$

QGS: Scores keys using int8-quantized projections:

$$c_i = \mathrm{Q}_8(U_r^\top k_i), \quad \tilde{z}_i = (U_r^\top q)^\top c_i, \tag{5}$$

3.3 ErrorGate

Estimates cluster discrepancy using a Gaussian model:

$$\hat{\varepsilon}_b = n_b e^{\bar{z}_b} e^{\hat{\sigma}_b^2/2}\big(2\Phi(\hat{\sigma}_b) - 1\big) = \hat{M}_b D_{\mathrm{G}}(\hat{\sigma}_b), \quad \hat{\sigma}_b^2 = s^2 \sum_d q_d^2 \mathrm{Var}_b(k_d). \tag{6}$$

Budget allocation via water-filling:

$$q_b = \min\left(n_b, \lambda \frac{\hat{\varepsilon}_b}{\max_{b'}\hat{\varepsilon}_{b'}}\right), \quad \sum_b q_b = E, \tag{7}$$

3.4 OpenAttention.flash Kernel

A Triton-based FlashAttention-2 decode kernel that processes exact tokens and proxies in one online softmax pass, with pooled sums cached per cluster for efficient subtraction of read tokens.

---

Empirical Validation / Results

4.1 Comparison Across Read Densities (RULER-32K HARD, Llama-3.1-8B)

At 1% read: OpenAttention scores 86.30 vs. dense 87.04 (99% retention), while the best baseline (HashAttention) scores only 48.2. OpenAttention maintains accuracy close to dense at every density from 10% to 1%.

4.2 Ablation Study (5% read, Table 4)

4.3 Cross-Model Evaluation (3% read, Table 5)

4.4 Context Length Scaling (Table 7, 3% read)

4.5 Clustering and Selection Validation

QGC vs. alternatives (Table 1, C=64, 1% read): QGC (16 dim) → 86.30, k-means on keys (128 dim) → 85.42, token position → 84.37.

QGS vs. selectors (Table 2): QGS achieves 0.97 top-32 logit recall at 1/16 read with 0.22 ms selection time—best balance of accuracy vs. cost.

ErrorGate ranking ablation (Table 3): Estimated error → 86.30, mean logit (mass) → 85.99, variance → 84.81.

4.6 Speedup

OpenAttention.flash runs the decode attention stage 10× faster than dense attention at 1% read and 5× faster at 3% read (batch 32, 32K context).

---

Theoretical and Practical Implications

Theoretical Contributions

1. Closed-form comparison (Eq. 2) of tail-treatment strategies shows sampling's fundamental disadvantage at low reads: $\mathbb{E}[\mathcal{E}_{\mathrm{s}}] = 2\tau(1-m/N)$ approaches twice the discard error, while proxy error $\mathcal{E}_{\mathrm{p}}^{\mathrm{G}} = \tau\operatorname{erf}(\sigma_z/\sqrt{2})$ remains bounded below.

2. Exact error decomposition (Eq. 3): $\mathcal{E} = \frac{1}{Z}\sum_{b \in \mathcal{P}} \hat{M}_b D_b$, separating proxy mass ($\hat{M}_b$) from weight spread ($D_b$)—providing a principled framework for cluster-based approximation.

3. Error bound: $\|\hat{o} - o\|_2 \leq (Z/\hat{Z}) \mathcal{E} \max_i \|v_i - o\|_2$, establishing the connection between weight discrepancy and output error.

Practical Implications

• 1% read rate with 99% dense accuracy makes long-context sparse attention practical for deployment on memory-constrained hardware
• The approach is orthogonal to KV cache compression (StreamingLLM, SnapKV, PyramidKV), losing only 0.14 points on average when combined
• Custom kernel accelerates decode attention by 10×, addressing the latency bottleneck critical for interactive applications
• The training-free method requires only once-per-request clustering, making it efficient for dynamic workloads

---

Conclusion

Main Takeaways

OpenAttention demonstrates that approximating the tail distribution with count-aware mean proxies is fundamentally superior to discarding or sampling at low read densities. By exploiting query geometry for clustering and selection, and using estimated weight errors for budget allocation, it achieves:

• 99% dense accuracy at 1% read
• State-of-the-art across models and context lengths
• 5–10× decode attention speedup

Limitations and Future Directions

• ErrorGate approximation: The diagonal variance approximation agrees with full-variance estimates at Spearman 0.92 but cannot recover exact error—reducing this proxy error is a promising future direction
• The budget allocation in Eq. (7) is a heuristic rather than an optimal rule
• The squared logit error metric is not the only objective; output error also depends on normalizer uncertainty and value direction perturbations

Reproducibility

Code is available at https://github.com/hyunwoo137/OpenAttention with project page at https://hyunwoo137.github.io/OpenAttention/. All models and benchmarks are publicly available; the 64K and 96K RULER sets are synthetic.