ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers

Abstract: Existing attention accelerators often trade exact softmax semantics, depend on fused Tensor Core kernels, or incur sequential depth that limits FP32 throughput on long sequences. We present \textbf{ELSA}, an algorithmic reformulation of online softmax attention that (i)~preserves exact softmax semantics in real arithmetic with a \emph{provable} $\mathcal{O}(u\log n)$ FP32 relative error bound; (ii)~casts the online softmax update as a prefix scan over an associative monoid $(m,S,W)$, yielding $O(n)$ extra memory and $O(\log n)$ parallel depth; and (iii)~is Tensor-Core independent, implemented in Triton and CUDA C++, and deployable as a \emph{drop-in replacement} requiring no retraining or weight modification. Unlike FlashAttention-2/3, which rely on HMMA/GMMA Tensor Core instructions and provide no compatible FP32 path, ELSA operates identically on A100s and resource-constrained edge devices such as Jetson TX2 -- making it the only hardware-agnostic exact-attention kernel that reduces parallel depth to $O(\log n)$ at full precision. On A100 FP32 benchmarks (1K--16K tokens), ELSA delivers $1.3$--$3.5\times$ speedup over memory-efficient SDPA and $1.97$--$2.27\times$ on BERT; on Jetson TX2, ELSA achieves $1.5$--$1.6\times$ over Math (64--900 tokens), with $17.8$--$20.2\%$ throughput gains under LLaMA-13B offloading at $\ge$32K. In FP16, ELSA approaches hardware-fused baselines at long sequences while retaining full FP32 capability, offering a unified kernel for high-precision inference across platforms. Our code and implementation are available at https://github.com/ming053l/ELSA.

• The paper introduces ELSA, which reformulates online softmax as a prefix scan over a monoid, delivering exact and efficient attention.
• It demonstrates reduced latency and lower memory usage across FP32 and FP16, with provable numerical stability on various hardware.
• ELSA serves as a drop-in, hardware-agnostic replacement for vision transformers, matching accuracy without necessitating retraining.

ELSA: Exact Linear-Scan Attention for Efficient and Hardware-Agnostic Vision Transformers

Introduction and Motivation

The quadratic memory cost of the standard multi-head self-attention (MHSA) operation in Vision Transformers (ViT) presents severe scalability challenges for high-resolution vision tasks, particularly in domains where floating-point (FP32) precision is critical, such as medical imaging or hyperspectral analysis. Most existing accelerators introduce tradeoffs between memory, speed, hardware dependency, and exactness, often requiring either approximation (and thereby retraining), reliance on specialized hardware (e.g., Tensor Cores), or falling back to sequential algorithms with limited throughput at long sequence lengths.

The ELSA (Exact Linear-Scan Attention) method presented in "ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers" (2604.23798) addresses this gap by algorithmically reformulating online softmax attention as a prefix scan over an associative monoid in real arithmetic, yielding an exact, memory- and compute-efficient, and hardware-agnostic solution with provable stability and parallelism.

Theoretical Contributions and Methodology

Monoid-Based Reformulation of Attention

ELSA’s core innovation is the formalization of the online softmax update as a prefix scan over the state triple $(m, S, W)$: the running max, normalized sum, and weighted sum of exponentials across the token sequence. The merge operator $\oplus$ forms a mathematically sound monoid on this triple, ensuring closure, associativity, and identity. This enables parallel prefix scan algorithms (Hillis–Steele for intra-block, Blelloch for inter-block aggregation), reducing scan depth from $O(n)$ sequential (as in ME-SDPA) to $O(\log n)$, without loss of exactness in real arithmetic.

Numerical Stability and Error Bounds

A provable floating-point error bound demonstrates that, under IEEE-754 single precision, ELSA's relative error grows only logarithmically with sequence length: $O(u\log n)$, where $u$ is machine epsilon. This contrasts favorably with standard sequential reductions, especially for very long contexts—crucial for foundation models where auditability and numerical stability are non-negotiable.

Hardware-Agnostic and Drop-In Deployment

Unlike hardware-fused kernels (e.g., FlashAttention-2/3), which depend on fused Tensor Core instructions (HMMA/GMMA) and lack efficient FP32 implementations, ELSA requires only scalar max/exp and vector FMA, and imposes no architectural restrictions or need for retraining, making it suitable for both datacenter GPUs and resource-constrained edge devices.

Empirical Results

Latency, Parallelization, and Memory Efficiency

ELSA demonstrates significant latency improvement and low memory usage across a wide range of sequence lengths and hardware platforms:

• In FP32, ELSA yields 1.3–3.5× speedup over ME-SDPA and up to 2.27× over BERT workloads, with comparable or lower memory usage.
• On edge devices (e.g., Jetson TX2), where Tensor Cores are absent, ELSA achieves up to 38% lower latency relative to the Math baseline, for both FP16 and FP32, confirming true hardware-agnosticism.
• In FP16, ELSA is competitive with hardware-fused FlashAttention while offering a consistent, exact FP32 path.

  Figure 2: Per-head latency scaling with sequence length fits closely to $O(\log n)$ parallel merge depth, verifying ELSA's theoretical prediction.

  

  Figure 4: ELSA achieves significant latency reduction and lowest memory among exact-kernel baselines at long sequence lengths; standard softmax attention fails with OOM at 16K tokens, where ELSA remains stable.

  

Figure 1: Memory footprint scaling with sequence length demonstrates ELSA's advantage in both FP16 and FP32, widening with longer contexts.

Training Dynamics and Model Compatibility

ELSA acts as a drop-in exact replacement for the attention operator in ViT, Swin, and large CLIP models, requiring no change to model weights or retraining. Experiments reveal that convergence and final validation accuracy on ImageNet-1K remain within 0.1% of the Math/FlashAttention baseline, demonstrating that ELSA does not perturb training dynamics.

Figure 3: Training loss and top-1 accuracy curves for ViT and Swin architectures with ELSA closely match or surpass those of the Math baseline.

Precision Validation and Correctness

Detailed numerical verification shows that model outputs under ELSA align with FP64 oracle references up to single-precision machine epsilon, preserving the top attention weight and exhibiting only random, bounded floating-point noise—there is no systematic spatial accumulation of error.

Figure 6: The drift heatmap for $|\delta P|$ shows purely random floating-point noise without systematic error, validating the correctness of the prefix-scan reformulation.

Figure 7: When tokens are divided into blocks and merged via $\oplus$, blockwise drift metrics coincide with the reference implementation, directly validating the correctness of the associative merge.

Implications and Future Directions

ELSA constitutes a strict drop-in, hardware-agnostic, and numerically stable solution for both vision and language transformers, making high-resolution FP32 inference and zero-shot transfer practical on a wide array of hardware. This is especially relevant for foundation models where retraining or adjustment for approximate attention is infeasible and for deployment scenarios (e.g., robotics, medical instruments, on-device AI) with strict hardware and accuracy requirements.

The paper highlights certain limitations: although memory usage and parallel depth are substantially improved, total arithmetic complexity remains $O(n^2)$, and at short sequence lengths or small windowed attention, hardware-optimized fused kernels may still hold a performance advantage. The authors point to future work on incorporating sparsity or low-rank methods for reduced FLOP count, extending multi-GPU and distributed parallelism, and further optimizing for even broader hardware support.

Conclusion

The ELSA method robustly addresses the memory and compute bottlenecks of exact attention for vision transformers by reconstituting softmax as a parallelizable, numerically stable, monoid-based prefix scan. Empirical results across a wide range of architectures, sequence lengths, and hardware platforms validate ELSA’s claims of latency reduction (up to 3.5×), strict FP32 exactness, and minimal additional memory cost, all with no requirement for retraining or specialized hardware. As transformer models scale in both domain and resolution, ELSA's approach provides a viable path for high-precision, efficient inference and deployment across heterogeneous platforms.

For details, see "ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers" (2604.23798).