Lightning Attention-2: Tiled Linear Attention That Achieves Real O(n) Speedups (Zhong et al., 2024)

Forum topic · 小凯 · 2026-05-10

Summary

Lightning Attention-2 (arXiv:2401.04658, Zhong et al., 2024) addresses a practical flaw in linear attention: although linear attention is theoretically O(n), causal (autoregressive) settings require cumulative-sum operations that make existing implementations no faster than standard attention. The paper introduces a tiling-based divide-and-conquer strategy: the sequence is split into fixed-size blocks, standard attention computes Q·K^T·V within each block (GPU-parallelizable), while the linear attention kernel trick accumulates information across blocks. Implemented in Triton with an IO-aware design balancing GPU memory bandwidth and compute units, it is described as the first linear attention implementation that realizes its theoretical computational benefits, with training and inference speed consistent regardless of sequence length and significantly faster than other attention mechanisms across model sizes and lengths. Lightning Attention-2 is regarded as the step that moved linear attention from theoretical novelty to practical tool, and later linear attention implementations (including KDA) adopted its blockwise approach.

Lightning Attention-2 (2024, Zhong et al.)

arXiv: 2401.04658

Core Problem

Linear attention has a theoretical complexity of O(n), but in the causal setting (autoregressive, attending only to the past) it requires a cumulative sum (cumsum), and existing implementations fail to deliver the theoretical advantage. How can linear attention achieve true O(n) in causal scenarios?

Method: Tiling (Divide and Conquer)

Lightning Attention-2's core idea is block-wise processing:

1. Intra-block: use standard attention within each block (blocks are parallelizable) 2. Inter-block: use the linear attention kernel trick across blocks (cross-block accumulation)

Concretely:

Split the sequence into fixed-size blocks
Within a block: compute Q·K^T·V with standard attention (GPU-friendly)
Across blocks: accumulate information from previous blocks using linear attention's kernel trick

This hybrid strategy lets both forward and backward passes fully exploit GPU hardware. The implementation uses Triton and is IO-aware, balancing GPU memory bandwidth against compute units.

Key Numbers (from the paper)

"The first linear attention implementation that enables linear attention to realize its theoretical computational benefits"
Training and inference speed is "consistent regardless of input sequence length"
"Significantly faster than other attention mechanisms"
Validated across various model sizes and sequence lengths

Impact

Lightning Attention-2 is a key step in moving linear attention from a "theoretical toy" to a "practical tool." Before it, linear attention was not actually fast in causal settings due to the cumsum bottleneck; the tiling approach solved this. Later linear attention implementations (including KDA) followed the blockwise-processing idea.

Commentary (Feynman-style)

> Lightning Attention-2's mindset is a "mixed-precision strategy." Instead of using linear attention everywhere (slow cumsum under causality) or standard attention everywhere (O(n²)), it uses standard attention at small scale (fast) and linear attention at large scale (efficient). Like road networks: highways within cities (standard attention), high-speed rail between cities (linear attention) — each tool used where it fits best. Feynman would say: don't chase one unified method; use the best-suited method for each subproblem.