Lightning Attention 2（ライトニングアテンションツー）

Lightning Attention 2 とは

Lightning Attention 2はQin et al.（OpenNLPLab、2024）が発表した線形アテンションの高効率実装。訓練フェーズと推論フェーズで異なる計算戦略を自動切替することで実用的なスループットを達成した。

I/O-Aware チャンク並列訓練

訓練時の核心技術：

系列を M個のチャンクに分割:
  チャンク内 (intra-chunk): 並列計算
    → FlashAttention的なHBM-SRAM転送最適化
  チャンク間 (inter-chunk): 累積和
    → KᵀV状態を左から累積（線形RNN的動作）

このI/O-Aware分割により計算量O(n)・HBM転送効率はFlashAttentionに近いレベルを達成。

推論時のRNNモード自動切替

if training:
    output = chunk_parallel_forward(Q, K, V, chunk_size=256)
else:
    # RNN再帰モード: KVキャッシュをO(1)状態行列に圧縮
    state = prev_state  # 固定サイズ (d×d)
    kv = k.outer(v)
    state = state + kv
    output = q @ state
    return output, state

推論時はKVキャッシュが固定サイズ(d×d)の状態行列1つに圧縮されるため、長文脈でもメモリが増加しない。

TransNormerLLM-15B での採用

2024年に公開されたTransNormerLLM-15Bはフル線形アテンションを採用した最大規模のモデルの一つ：

• Llama 2-13Bと同等のPerplexity
• 1Mトークン以上のコンテキストで計算量がほぼ線形
• A100 80GB×8でのスループットがFlashAttention-2採用LLaMAに迫る

Triton カーネル実装

pip install lightning-attn2

from lightning_attn.ops import lightning_attn_func
output = lightning_attn_func(q, k, v, slope=slopes)

H100/A100/RTX 4090での最適化済みカーネルが提供されている。