Lightning Attention: GPU-Optimized Linear Attention

Updated 27 January 2026

Lightning Attention is a GPU-optimized, tile-based causal linear attention mechanism that partitions input sequences into blocks to ensure constant memory usage regardless of sequence length.
It eliminates the need for serial cumulative summation by updating a d×d accumulator during inter-block processing, thereby achieving high and consistent GPU throughput.
Integrated within architectures like TransNormerLLM and MiniMax, it delivers competitive language modeling and multimodal performance while reducing FLOPs and memory costs.

Lightning Attention is a GPU-optimized, tile-based, causal linear attention mechanism that realizes the theoretical promise of $O(n d^2)$ time and constant memory with respect to sequence length in LLMs. By partitioning input sequences into blocks and separating intra-block masked attention from inter-block linear accumulation, Lightning Attention eliminates the need for serial cumulative summation ("cumsum"), thus ensuring constant tokens-per-GPU-second (TGS) throughput as context lengths scale from thousands to millions of tokens. It has been instantiated in highly efficient architectures such as TransNormerLLM, MiniMax-01, and MiniMax-M1, demonstrating state-of-the-art long-context scaling and performance competitive with full softmax attention on both language modeling and multimodal tasks.

1. Formal Definition and Core Mechanism

Let %%%%1%%%% denote the usual query, key, and value matrices for a sequence of $n$ tokens and hidden dimension $d$ . In the causal setting, attention output $O \in \mathbb{R}^{n \times d}$ is generally given by

$O = \left[(Q K^\top) \odot M \right] V,$

where $M_{ts} = 1$ if $t \geq s$ , and $0$ otherwise encodes the lower-triangular causal mask.

Lightning Attention introduces a blockwise partitioning:

Choose block size $B$ (in practice $B \approx d$ ).
Divide $Q, K, V$ into $T = n/B$ nonoverlapping blocks: $Q = [Q_1; \ldots; Q_T]$ , $K = [K_1; \ldots; K_T]$ , $V = [V_1; \ldots; V_T]$ with each $B \times d$ .

For each block $t$ ( $1 \leq t \leq T$ ), the output is split into

Intra-block term (local masked attention):

$O^{\text{intra}}_t = \left[( Q_t K_t^\top ) \odot M \right] V_t$

Inter-block term (linear kernel accumulation):

$O^{\text{inter}}_t = Q_t \; KV_{t-1}, \quad \text{where } KV_{t-1} = \sum_{i=1}^{t-1} K_i^\top V_i$

Block output:

$O_t = O^{\text{intra}}_t + O^{\text{inter}}_t$

No per-token prefix-sum is needed: intra-block computation is standard masked matmul for local context, and inter-block uses a $d \times d$ accumulator for efficient global context (Qin et al., 2024, Qin et al., 2024, MiniMax et al., 14 Jan 2025).

2. Tiling, Algorithmic Eliminations of Cumsum, and Implementation

In standard linear attention, causal computation requires, for each token $t$ , a sequential update:

$kv_t = kv_{t-1} + k_t v_t^\top, \qquad o_t = q_t kv_t$

This demands a full-sequence prefix-sum (i.e., cumsum), which has $O(n)$ serialized steps and inhibits GPU parallelization.

Lightning Attention instead:

Tiles the sequence into blocks of size $B$ , computes full masked attention within each block, and represents all historical contributions from earlier blocks by sequentially updating a $d \times d$ accumulator $KV$ .
For each block, $KV$ is updated with $K_t^\top V_t$ and the new block is processed independently on-chip (SRAM).
All blockwise computations are overlapped with memory transfers, saturating compute bandwidth and achieving full hardware efficiency (Qin et al., 2024, Qin et al., 2024).

Forward pass pseudocode:

Input: Q, K, V ∈ ℝ^{n×d}, block size B
Divide into T = n // B blocks
KV = zeros(d, d)
for t in 1…T:
    Load Q_t, K_t, V_t to SRAM
    O_intra = (Q_t @ K_t.T) * M @ V_t
    O_inter = Q_t @ KV
    KV += K_t.T @ V_t
    O_t = O_intra + O_inter
    Write O_t back to DRAM
return concatenated O

Backward pass is analogous: gradients accumulate over blocks and the full-sequence cumsum is never required (Qin et al., 2024, Qin et al., 2024).

3. Mathematical Properties and Theoretical Analysis

Lightning Attention and its linear forms admit a precise algebraic geometry characterization. In the fully algebraic setting (without normalization), the single-layer attention map is

$\phi_{Q,K,V}(X)_i = \sum_{j=1}^t \langle x_j, A x_i \rangle \, V x_j$

with $A = K^\top Q$ . The neuromanifold $\mathcal{M}$ of all such maps is a determinantal variety whose dimension, identifiability, and singular loci have been explicitly described (Henry et al., 2024):

Dimension: For $d, d', a \in \mathbb{N}$ , if $a \leq d$ ,

$\dim \mathcal{M}_{d,d',a} = 2ad - a^2 + d'd - 1.$

Generic identifiability: In the unnormalized case, fibers are generically one-dimensional up to overall $(A, V) \mapsto (\lambda A, \lambda^{-1} V)$ scaling; for softmax-normalized attention, parameterization is generically injective.
Singular/boundary loci: Points where $A$ and $V$ both have rank $1$ lie on the algebraic boundary or are singular. These loci inform about function space complexity and where training may stall.

For deep (multi-layer) architectures, additional gauge symmetries arise, but essentially the structure remains highly constrained and well-understood (Henry et al., 2024).

4. Systems-Level Scalability, Kernel Fusion, and Memory Efficiency

Lightning Attention is distinguished by strict $O(n d^2)$ time and $O(n d)$ end-to-end memory, independent of sequence length $n$ . FlashAttention-2 and naive linear attention are $O(n^2 d)$ and cannot maintain throughput as $n \to 10^5$ – $10^6$ .

System-level optimizations include:

Tile-based kernel launches: Each $B \times d$ block is processed in SRAM with maximally fused kernels.
Overlap of computation and IO: Double-buffering and block pipelining hide global memory latency.
LASP+: Parallel prefix-sum within Context-Parallel GPU groups using AllGather for inter-node scaling (MiniMax et al., 14 Jan 2025).
VarLen ring: Efficient packing of sequences in multimodal or varied-length contexts (MiniMax et al., 14 Jan 2025).

This enables consistent throughput up to $n = 4 \times 10^6$ tokens, using 8 H800/H20 GPUs per 1M-token training batch, and inferred performance matches claimed theoretical scaling (MiniMax et al., 14 Jan 2025, MiniMax et al., 16 Jun 2025).

5. Architectural Integration: Hybrid Patterns, Gated Modules, and MoE

In current state-of-the-art large language and vision-LLMs, Lightning Attention is deployed with:

Hybrid stacking: Sequences of 7 Lightning Attention blocks are followed by 1 softmax block for global context anchoring (MiniMax-01/M1: $[LA \rightarrow \mathrm{MoE}]^7 \rightarrow [\mathrm{Softmax} \rightarrow \mathrm{FFN}]$ ) (MiniMax et al., 14 Jan 2025, MiniMax et al., 16 Jun 2025).
Gated mixing: GLA (Gated Linear Attention) and SGLU (Simple Gated Linear Unit) are used for token and channel mixing, with $\phi$ typically set as Swish or ELU-based mapping.
Normalization: SRMSNorm is used for stability and speed, with negligible perplexity difference to LayerNorm or RMSNorm.
MoE integration: Each Lightning Attention block feeds into a Mixture-of-Experts FFN, with router-based sparse activation and token-expert sharding across GPUs (MiniMax et al., 14 Jan 2025).
Relative positional encoding: Exponential-decay LRPE-d encoding, $a_{ts} = q_t^T k_s \lambda^{t-s} e^{i\theta(t-s)}$ , is fully compatible with Lightning-style tile-based block updates (Qin et al., 2024).

This pattern delivers stable long-context behavior and allows per-token compute in MoE layers to remain sublinear in model size.

6. Empirical Performance, Benchmarks, and Limitations

Empirical evidence across several models demonstrates that Lightning Attention:

Maintains constant training/inference throughput (TGS) as context increases; e.g., $33$k tokens/GPU/sec on a $3$B model for any context length (MiniMax et al., 14 Jan 2025).
Enables training and inference on up to $4$ million tokens at batch and production scale (MiniMax et al., 14 Jan 2025, MiniMax et al., 16 Jun 2025).
Reduces FLOPs and memory cost by $3$– $4\times$ compared to dense softmax models on $64$k–$100$k generations (MiniMax et al., 16 Jun 2025).
Preserves accuracy: On WikiText-103 (44M), TNL achieves test PPL $24.03$ vs.\ Transformer $24.78$, exceeding prior efficient models (Qin et al., 2024). In large-scale LLMs and vision-LLMs, performance on OpenAI MRCR, LongBench v2, MMLU, and C-Eval matches or outperforms LLaMA and DeepSeek baselines at 1M context windows (MiniMax et al., 14 Jan 2025, MiniMax et al., 16 Jun 2025).
At the architectural level, inference throughput for $7$B models with $512+1024$ sequences is up to $11\times$ that of Transformer+Flash2 (Qin et al., 2024).
For pure long-context retrieval, hybrid LA+Softmax blocks outperform pure linear blocks, though pure Lightning shows some retrieval trade-off (MiniMax et al., 14 Jan 2025).
Empirical performance is robust across activation, gating, block size, and positional encoding ablations; SRMSNorm consistently yields fastest implementation (Qin et al., 2024).

7. Limitations, Open Challenges, and Future Directions

Some limitations and outstanding issues include:

Block size $B$ : A trade-off exists between local detail (larger $B$ for fidelity) and global throughput; tunable per deployment (Qin et al., 2024, Qin et al., 2024).
Retrieval accuracy: Pure LA is weaker than softmax for cross-attention, necessitating periodic softmax anchoring in deep models (MiniMax et al., 14 Jan 2025, MiniMax et al., 16 Jun 2025).
Hardware limitations: Maximum $n$ is bounded by total device memory for $Q,K,V,O$ (global $O(nd)$ storage); Lightning Attention does not eliminate this constraint.
Kernel-level optimization: Further kernel fusion, sequence-parallelism, and dynamic block sizing are open paths to even better HW utilization; adaptation to new architectures (e.g. Hopper) is ongoing (MiniMax et al., 14 Jan 2025).
Theory: Geometric/identifiability theory for normalized multi-layer attention is conjectural beyond $L=1$ (Henry et al., 2024).

Prospective research directions include direct elimination of softmax blocks via improved global aggregation, adaptive or content-based block sizing, and further hybridization with structured sparsity or token-expert routing (MiniMax et al., 14 Jan 2025).

References: