Thinformer Express: Streaming Causal Attention

• Thinformer Express is a causal attention approximation algorithm that converts non-causal models into streaming ones through a proven sub-Gaussian thinning and meta-procedure.
• It achieves an approximation error of O(log^(3/2)n/s) with O(s) memory and compressed computational overhead, supporting efficient long-context processing.
• Its Triton-optimized implementation delivers significant speedups over methods like FlashAttention 2 and HyperAttention for prefill, KV-cache compression, and compute-constrained decoding.

Thinformer Express is a causal attention approximation algorithm that combines the state-of-the-art Thinformer thinning procedure with the Express meta-procedure for converting non-causal attention approximations into streaming, causal ones, while maintaining strong theoretical error guarantees. The method achieves uniform $\varepsilon = O(\log^{3/2} n / s)$ approximation error, $O(s)$ memory, and $O(s^2 \log^2 n)$ compression overhead for a sequence of length $n$. Its highly optimized implementation in Triton delivers significant speedups compared to prior state-of-the-art methods such as FlashAttention 2 and HyperAttention, supporting efficient deployment for long-context prefill, KV-cache compression, and memory-/compute-constrained decoding workloads (Gong et al., 9 Jun 2026).

1. Problem Statement and Theoretical Foundations

Causal attention, foundational to sequence modeling, requires restricting each position $j$ to attend only to prior or current positions ($i\leq j$). For query–key–value triplets $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$, the exact masked (causal) attention output for $q_j$ is

$$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$

Exact computation requires $\Theta(d n^2)$ time and $O(s)$0 memory, making scaling to long sequences intractable for standard hardware constraints.

Coreset-based attention approximation addresses this by maintaining a compressed, weighted subset of key–value pairs, enabling an approximate attention output

$O(s)$1

with $O(s)$2 provably small for all $O(s)$3.

The key theoretical device is sub-Gaussian thinning: a randomized algorithm produces a weighted coreset $O(s)$4 with guarantees

$O(s)$5

for all test functions $O(s)$6 in a suitable RKHS, with high probability. Thinformer provides non-causal thinning with $O(s)$7. The Express meta-procedure converts such thinning to the causal (streaming) domain with only an $O(s)$8 inflation in $O(s)$9, yielding Thinformer Express.

2. Algorithmic Structure and Execution Phases

Thinformer Express maintains a weighted cache $O(s^2 \log^2 n)$0 of target size $O(s^2 \log^2 n)$1, supporting the following workflow for each token:

• Coreset retrieval: Obtain $O(s^2 \log^2 n)$2.
• Approximate attention formation: Compute

$O(s^2 \log^2 n)$3

• Update: Call $O(s^2 \log^2 n)$4.

Express operates internally in three phases:

1. Exact: Accumulates the first $O(s^2 \log^2 n)$5 pairs exactly.
2. Thin: Processes subsequent input in blocks ($O(s^2 \log^2 n)$6 in size), performing stratified subsampling followed by ThinformerHalve to reduce block size to $O(s^2 \log^2 n)$7; summaries are appended to cache.
3. Halve: When cache exceeds $O(s^2 \log^2 n)$8, recursively apply halving (down to $O(s^2 \log^2 n)$9 points), incrementing $n$0.

This scheme guarantees at most $n$1 points in cache and streaming update complexity $n$2 halving calls per token. Thinformer’s kernel, quadratic in block size $n$3, leads to total compression cost $n$4.

3. Approximation Guarantees

The approximation error for Thinformer Express is rigorously characterized. Fix $n$5 and run Thinformer Express with cache size $n$6. With probability at least $n$7, for every $n$8,

$n$9

where $j$0, $j$1, $j$2, and $j$3 is a small absolute constant. In big-O notation,

$j$4

uniformly over all tokens. This guarantee arises by chaining sub-Gaussian error bounds across $j$5 thinning stages and subsequently applying an attention-output stability lemma (Gong et al., 9 Jun 2026).

4. Resource Complexity and Asymptotics

Thinformer Express achieves asymptotically optimal scaling for memory and runtime overheads in attention computation.

Resource	Scaling	Notes
Memory (KV storage)	$j$6	At most $j$7 weighted points
Per-token query	$j$8	Full attention over cache
Compression overhead	$j$9	Dominated by $i\leq j$0 small halving calls per token

Total streaming compression time over $i\leq j$1 tokens satisfies $i\leq j$2. Memory is independent of sequence length $i\leq j$3; only computational overhead grows slowly as $i\leq j$4.

5. Implementation Optimizations

The combination of Triton and algorithmic design underpins Thinformer Express’s efficiency. Major optimizations include:

• Tiling: Key–query exponentiation and halving inner loops are partitioned into $i\leq j$5 tiles to maximize on-chip cache utilization.
• Fused operations: $i\leq j$6 multiplications are fused with summation, preventing formation of $i\leq j$7 or $i\leq j$8 matrices in high-bandwidth memory (HBM).
• Parallelism: Same-size halving tasks are executed in parallel during offline prefill phases.
• Double indirection: Keys and values remain contiguous in HBM, while coresets store references in shared memory, eliminating scatter/gather overhead.

Performance gains include:

• Unmasked prefill ($i\leq j$9K, $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$0): Torch-compiled Thinformer $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$1, Triton Thinformer $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$2 faster than FlashAttention 2.
• Masked prefill (ChatGLM2-6B-32K, $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$3K, $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$4 or 1024): Thinformer Express up to $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$5 faster than FlashAttention 2. HyperAttention is out-of-memory above $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$6K.

6. Empirical Evaluation and Benchmarks

Rigorous empirical evaluations span diverse large-language-model workloads:

(a) Long-context prefill, masked (ChatGLM2-6B-32K, $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$7):

Length $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$8	FlashAttn2	HyperAttn	Thinf-Expr
64K	1×	5×	12×
128K	1×	10×	25×
256K	1×	--	45×
384K	1×	--	60×
512K	1×	--	82×

(b) KV-cache compression (Llama 3.1 8B, LongBench-E): Wrapping compressors (SnapKV, StreamingLLM, PyramidKV) with Express yields $$(q_i, k_i, v_i)\in \mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}^d$$9–$q_j$0 reduction in attention time at preserved end-task accuracy.

(c) Memory-constrained decoding (MATH-500, DeepSeek-R1-Distill-LLama-8B):

Cache size	Exact Acc.	Thinf-Expr Acc.
1000	28.4%	28.3%
2000	31.2%	31.1%
4000	34.7%	34.6%

Cache size	Exact Mem.	Thinf-Expr Mem.
4000	100%	61%

(d) Compute-constrained decoding (same settings):

Time	Exact Acc.	Thinf-Expr Acc.
1.0×	34.7%	34.7%
0.75×	28.4%	28.4%

Time	Exact Cost	Thinf-Expr Cost
Decoding	100%	56%

Thinf-Expr matches exact accuracy with $q_j$1 of KV memory and $q_j$2 of computational cost.

7. Deployment Considerations and Parameter Selection

• Cache size $q_j$3: Set $q_j$4 to target desired maximum approximation error $q_j$5. In practice, $q_j$6–2048 suffices for $q_j$7 up to several hundred thousand.
• Sequence length $q_j$8: Memory usage is $q_j$9, independent of $$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$0. The only $$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$1-dependence is in the compression overhead ($$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$2 factor), negligible when $$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$3.
• Hardware: The Triton reference implementation utilizes on-chip shared memory and contiguous HBM allocation. The algorithm is AISA-friendly. Warp specialization and FP8 are not used but can be integrated.
• Pipeline integration:
  • Prefill: Thinformer Express is applied in all layers to compress long input contexts prior to decoding.
  • KV-cache: Generation-time storage only requires the $$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$4-sized Express cache, not the full $$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$5.
  • Decoding: Express.update is called only after computing $$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$6; halving costs per-token ($$o_j = \frac{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})\,v_i}{\frac{1}{j}\sum_{i=1}^j \exp(q_j^\top k_i/\sqrt{d})}.$$7) are negligible compared to attention.

Thinformer Express thus provides an end-to-end solution for causal attention with proven error bounds, optimal memory scaling, and practical speedups for long-context and resource-constrained neural language modeling (Gong et al., 9 Jun 2026).