リング注意と文脈並列

なぜ長文脈は1枚のGPUに載らないのか

Self-Attention の計算 で見たとおり、アテンションは各クエリが全キーを参照します。系列長を n、ヘッド次元を d とすると、素朴な実装はスコア行列を n × n で作るため、メモリが系列長に対し二乗で膨らみます。FlashAttention はこの n × n の中間行列を実体化せず、追加メモリを O(n) まで落としました。しかし O(n) でも、n が数十万〜数百万トークンに達すると、入力・出力のアクティベーションや各層の中間結果が積み上がり、単一デバイスのメモリを超えます。

問題の本質は、分散学習 で扱ったデータ並列・テンソル並列・パイプライン並列の3軸が、いずれも 系列（トークン）軸を割らない 点にあります。バッチを割っても層内の行列を割っても、1本のシーケンス全体は依然として各デバイスに丸ごと乗ります。そこで登場するのが、系列そのものを分割する第4の軸です。

系列を割ったときの障壁：アテンションだけが「全体」を要る

系列を P 台のデバイスへ割り、各デバイス i がトークン区間 i 分の Q_i・K_i・V_i を持つとします。ここで層ごとの処理を見分けると、大半の演算はローカルで完結 します。埋め込み・全結合・LayerNorm・活性化 はトークンごとに独立なので、各デバイスは自分の区間だけで計算できます。

厄介なのはアテンションだけです。クエリ Q_i は自分の区間のキーだけでなく 全区間のキー K（= K_0 から K_{P-1} まで） を参照しなければなりません。つまり各デバイスは、自分が持っていない他デバイスのKVを何らかの形で手に入れる必要があります。

素朴な解は「全デバイスのKVを一度に集める」All-Gather です。しかしこれは全KVを各デバイスに実体化するため、メモリが再び全系列長 n に比例 してしまい、系列を割った意味が消えます。長文脈では K・V そのものが巨大なので、これは致命的です。求められるのは、全KVを同時に持たずに、全KVを順に見る 仕組みです。

リング注意：KVを隣へ回しながら少しずつ計算する

リング注意（Ring Attention）の着想は、分散学習 の Ring-AllReduce と同じ「リング」です。P 台のデバイスを論理的な輪に並べ、各デバイスは 自分の Q_i を固定したまま、KVブロックだけをリング上で順送り します。

各ステップで、デバイス i は「いま手元にあるKVブロック」と自分の Q_i でアテンションを計算し、結果を積算します。同時に、そのKVブロックを次の隣デバイスへ送り、前の隣デバイスから新しいKVブロックを受け取ります。これを P ステップ繰り返すと、各デバイスの Q_i は 全デバイス由来のKVを一巡 し、フルアテンションが完成します。

デバイス i（担当クエリ Q_i、初期KVは自分の K_i, V_i）:
  o, l, m を初期化                       ← オンライン softmax の統計量
  現在KV ← (K_i, V_i)
  for step in 0 .. P-1:
      # 計算と通信を同時に走らせる
      非同期送信: 現在KV を 次の隣 (i+1) へ
      非同期受信: 次に使うKV を 前の隣 (i-1) から
      # 受信を待つ間にローカル計算
      ブロック単位で S = Q_i · 現在KV.Kᵀ / √d
      o, l, m を オンライン softmax で更新（下記）
      通信の完了を待つ
      現在KV ← 受信したKV
  出力 O_i = o / l                        ← P 巡した時点で全KVを反映済み

重要なのは、各デバイスが 同時に保持するKVは1ブロック分だけ という点です。全KVを一度に集める All-Gather と違い、リングを回すことで メモリはブロック長に依存し、全系列長 n には依存しません。

オンライン softmax：ブロックを分割しても厳密に一致する

「一度に1ブロックずつ」処理して、なぜフルアテンションと同じ結果になるのか。ここを支えるのが、FlashAttention と同じ オンライン softmax です。softmax の分母は全キーにわたる指数和が要るため、本来は全KVを見終えるまで正規化を確定できません。しかし行ごとの最大値 m と指数和 l、部分出力 o を持ち回り、ブロックごとに補正すれば逐次集約できます。

新しいKVブロックのスコア s_j に対して:
  m_new = max(m, ブロック内スコアの最大値)
  α     = exp(m - m_new)                    ← 旧統計量の再スケール係数
  l     = l * α + Σ exp(s_j - m_new)
  o     = o * α + Σ exp(s_j - m_new) * v_j
  m     = m_new
最後に: O_i = o / l                          ← 全ブロック走査後に一度だけ正規化

補正係数 α = exp(m - m_new) が肝で、後から大きいスコアが出て最大値が更新されたとき、過去に積んだ l と o をスケールし直します。この再スケールにより、KVを何ブロックに割って、どんな順に見ても、全キーを一度に見た softmax と厳密に一致 します。リング注意は近似アテンションではなく、線形アテンション のような精度トレードオフを持ちません。出力が数学的に等価なので、既存モデルの学習・推論にそのまま差し込めます。

通信と計算の重ね合わせ：リングが速い理由

リング注意の実用性は 通信と計算のオーバーラップ に懸かっています。各ステップでデバイスがやることは2つ――(1) 手元のKVで Q_i のアテンションを計算する、(2) KVブロックを隣へ送り次を受け取る。この2つを 同時に走らせる のが要点です。

1ブロックのアテンション計算に要する時間を T_compute、1ブロックのKV転送に要する時間を T_comm とすると、両者を重ねたステップ時間は max(T_compute, T_comm) に近づきます。もし T_comm が T_compute 以下 なら、通信は計算の裏に完全に隠れ、リングを回す追加コストは実質ゼロになります。ブロック長を大きく取るほど T_compute が伸びて通信を吸収しやすくなる一方、瞬間メモリは増えるため、ここが設計上のつまみです。

シーケンス並列との違いと、実務での位置づけ

用語が紛らわしいのが「シーケンス並列」です。文脈により2つの意味があります。狭義のシーケンス並列は、テンソル並列と組で使い、LayerNorm や Dropout などテンソル並列が割り残す部分をトークン方向に割ってアクティベーションメモリを節約する技法を指します。一方、リング注意が実現するのは アテンションまで含めてシーケンス全体を割る 文脈並列で、こちらの方が守備範囲が広い概念です。

実際の超長文脈システムでは、これらを掛け合わせます。ノード内はテンソル並列、その上に文脈並列でシーケンスを割り、さらにデータ並列で束ねる、という多軸構成です。文脈並列は KV キャッシュ と組めば、推論時の超長プロンプト処理（prefill）でも効きます。KVキャッシュが1本のシーケンスの過去KVを保持するのに対し、文脈並列はその1本自体を複数デバイスへ割るので、両者はレイヤーが違い併用できます。

まとめ：系列を割るという発想

リング注意と文脈並列が示すのは、「どの軸で割り、その結果どんな通信が要るか」 という分散学習の一般問題に、系列軸という新しい答えを加えたことです。

この見立てを持つと、「文脈長を伸ばすとOOMする」「デバイスを足しても長文が扱えない」といった壁が、分割軸と通信パターンの言葉で 説明できます。アテンション本体の効率化は FlashAttention と KV キャッシュ で、並列の全体像は 分散学習 で押さえると、超長文脈を学習・推論可能にする工夫が一本の線でつながります。