Technical Summary: EntmaxKV – Support-Aware Decoding for Entmax Attention

1. Problem Statement

Long-context decoding in Large Language Models (LLMs) is increasingly bottlenecked by memory traffic rather than compute. During autoregressive generation, each new token must attend to a Key-Value (KV) cache that grows linearly with context length. Reading this entire cache from high-bandwidth memory (HBM) to on-chip SRAM at every step dominates latency, often accounting for over half of the per-token generation time.

Existing sparse decoding methods attempt to mitigate this by selecting subsets of tokens, blocks, or pages. However, these methods are primarily designed for softmax attention. A fundamental limitation of softmax is its "dense tail": it assigns strictly positive probability to every token, regardless of relevance. Consequently, any truncation of the KV cache under softmax inherently discards non-zero probability mass, making sparse decoding an approximation problem where error is unavoidable unless the full cache is read.

2. Methodology: EntmaxKV

The authors propose EntmaxKV, a sparse decoding framework native to $\alpha$-entmax attention. Unlike softmax, $\alpha$-entmax (for $\alpha > 1$) produces probability distributions with exact zeros. This property transforms the sparse decoding objective from "approximating a dense tail" to "support recovery": if the selected subset of tokens contains the true support (the set of tokens with non-zero probability), the sparse decoding output is mathematically identical to the full-cache output.

EntmaxKV implements this via a paged decoding pipeline with three core components:

A. Query-Aware Page Scoring

The KV cache is partitioned into pages. Before loading full keys and values, the system scores pages using lightweight metadata stored per page:

• Bounding Box Bounds: Coordinate-wise minimum and maximum key vectors ($k_{min}, k_{max}$) allow for a deterministic upper bound on the maximum possible attention score for any token in a page given a query $q$.
• Gaussian Statistics: For a more adaptive approach, the system stores first and second moments (mean and variance) of keys within a page. This allows approximating the distribution of token scores in a page as a Gaussian random variable.

B. Candidate Selection Policies

Two policies are proposed to select which pages to load:

1. Top-k Page Selection: A fixed-budget approach that selects the $k$ pages with the highest upper-bound scores derived from the bounding box metadata.
2. Gaussian-Aware Entmax Selection: A dynamic approach that estimates the global $\alpha$-entmax threshold ($\tau$) using the page-level Gaussian statistics. It solves an approximate normalization equation to find $\hat{\tau}$ and then selects pages where the estimated maximum score exceeds this threshold (adjusted by a safety margin $\Delta$). This adapts the token budget to the specific score distribution of the current query.

C. Sparse Entmax Attention

Once candidate pages are selected, the system loads only the corresponding KV entries. It computes the $\alpha$-entmax attention weights over this reduced set. If the selected set contains the full support of the $\alpha$-entmax distribution, the output is exact.

3. Theoretical Analysis

The paper provides a theoretical framework for the error introduced by truncation:

• Error Bound: The output error $\|o - \tilde{o}\|$ is bounded by $2B\delta$, where $B$ is the bound on value vector norms and $\delta$ is the total probability mass of the dropped tokens.
• Support Recovery: For softmax, $\delta > 0$ whenever any token is dropped, guaranteeing non-zero error. For $\alpha$-entmax, if the selected set $I_{keep}$ contains the support $S$ (i.e., $S \subseteq I_{keep}$), then $\delta = 0$ and the error vanishes exactly.
• Comparison: The paper demonstrates that under the same kept set, $\alpha$-entmax drops significantly less probability mass than softmax, particularly when the candidate set successfully captures the entmax support.

4. Key Results

Experiments were conducted using 1B-parameter models (NAPE) on long-context benchmarks (RULER, Passkey Retrieval, PG19) and efficiency benchmarks.

• Approximation Quality: EntmaxKV consistently drops less probability mass and retains more support tokens than softmax-based sparse decoding at matched KV budgets. The relative output error is substantially lower for entmax, approaching zero as the support is recovered.
• Downstream Performance:
  • Passkey Retrieval: EntmaxKV maintains perfect retrieval accuracy across long contexts (up to 100k tokens) using less than 0.1% of the KV cache, whereas softmax-based methods degrade rapidly.
  • RULER Benchmark: EntmaxKV significantly outperforms softmax sparse decoding, particularly at 16k context lengths, maintaining performance close to full-cache baselines.
  • Language Modeling (PG19): Entmax sparse decoding with a 1024-token budget achieves lower perplexity than softmax decoding with much larger budgets (e.g., 8192 tokens).
• Efficiency:
  • EntmaxKV shifts the bottleneck from full-cache memory traffic to lightweight page selection.
  • At a 1M token context length, the Gaussian-aware variant achieves a 3.36× speedup over full-cache softmax (FlashDecoding) and a 5.43× speedup over full-cache entmax decoding.
  • The Gaussian-aware selector scales better with sequence length than the Top-k variant, avoiding the explicit Top-k bottleneck.

5. Significance and Claims

The paper claims that EntmaxKV fundamentally changes the paradigm of sparse decoding by leveraging the exact sparsity of $\alpha$-entmax.

• From Approximation to Recovery: The primary contribution is reframing sparse decoding as a support recovery problem. The authors argue that if the selection mechanism successfully identifies the support, the method is not an approximation but an exact computation, eliminating the "dense-tail" error inherent to softmax truncation.
• Practical Efficiency: The work bridges the gap between the theoretical benefits of sparse attention and the practical constraints of autoregressive decoding. By exploiting sparsity before loading KV pages, EntmaxKV reduces memory traffic without sacrificing the exactness of the attention mechanism when the support is captured.
• Adaptive Selection: The introduction of the Gaussian-aware selector allows the system to dynamically adjust the candidate budget based on the score distribution, offering a more robust alternative to fixed-budget Top-k selection.

The authors position EntmaxKV as a native solution for $\alpha$-entmax models, distinct from previous work that attempts to force sparsity onto dense softmax distributions. They emphasize that while recent kernels have made $\alpha$-entmax trainable at scale, EntmaxKV is the first to address the autoregressive decoding bottleneck by enabling sparse inference that is both efficient and theoretically exact under support recovery.