ELSA：不靠 Tensor Core 的 exact attention，讓 FP32 與邊緣 GPU 重新有優化空間

這篇 Threads 分享 ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers。官方 arXiv 是 2604.23798，CVPR Findings 2026，作者來自陽明交通大學 Advanced Computer Vision Laboratory。

ELSA 的定位很清楚：

不需要 Tensor Core、不需要重新訓練 Foundation Model，作為 drop-in replacement 的 exact attention kernel，在 FP32 與受限硬體上補 FlashAttention 沒照顧好的缺口。

為什麼這重要？

FlashAttention-2 / 3 很強，但它的高效路徑依賴 Ampere / Hopper 的 Tensor Core 指令，例如 HMMA / GMMA。對某些硬體或場景，這會形成限制：

• Jetson TX2 這類 edge device 沒有對應 Tensor Core 能力。
• AMD GPU / 舊伺服器 GPU 不能直接走 NVIDIA Tensor Core 路徑。
• 高解析影像、醫療影像、高光譜遙測、科學計算等場景常需要 FP32 precision。
• FlashAttention 的 FP32 fallback 可能回到較未最佳化的 SIMD / math path，speedup 消失。

ELSA 的核心想法：把 online softmax attention 改寫成 prefix scan。

傳統 online softmax 需要依序維護 running maximum，下一步依賴上一步，parallel depth 是 O(n)。ELSA 把 softmax 狀態寫成 monoid triple：

(m, S, W)

並定義 merge operator，讓任意相鄰 block 的狀態可以精確合併。這樣每個 block 可以先平行計算，再用 reduction tree / prefix scan 合併，parallel depth 從 O(n) 降到 O(log n)。

官方 arXiv abstract 的三個要點：

1. 保留 exact softmax semantics

在 real arithmetic 下保留 exact softmax semantics，FP32 relative error bound 是 O(u log n)。這比 sequential depth 造成的 O(nu) 誤差更有利。

2. O(n) extra memory + O(log n) parallel depth

把 online softmax update cast 成 associative monoid 上的 prefix scan，不需要 O(n²) score matrix，額外記憶體 O(n)。

3. Tensor-Core independent

實作包含 Triton 與 CUDA C++，可部署為 drop-in replacement，不需重新訓練或修改權重。

Threads 裡補充了實作方式：

• 同時載入多個 tile 到 SRAM block。
• block 內用 Hillis–Steele scan。
• block 間用 Blelloch sweep / reduction 整合跨 block 狀態。
• 不需要 Tensor Core。

官方 performance highlights：

• A100 FP32 benchmarks 1K–16K tokens：比 memory-efficient SDPA 快 1.3–3.5×。
• BERT FP32：1.97–2.27×。
• Jetson TX2：比 Math kernel 快約 1.5–1.6×。
• LLaMA-13B offloading：32K+ tokens 下 throughput 提升 17.8–20.2%。
• Project page 提到 CLIP ViT-L/14 strict FP32 full image encoder latency 降 3.7%，attention module level 1.46–2.15×；full pipeline gain 被非 attention compute 攤薄。

GitHub repo：ming053l/ELSA

查詢時資訊：

• 描述：[CVPR 2026 Findings] ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers
• code / README 提供 Triton / CUDA kernels、PyTorch module、timm ViT / Swin patch、ElsaViT、ElsaSwinTransformerV2、HuggingFace LLaMA patch examples。