Context Parallelismとは？長いpromptを複数GPUへ分ける仕組み

長いpromptをLLMへ入れると、model weightがGPUへ収まっていてもPrefillの計算とmemoryが限界へ達する。Tensor Parallelismはhidden dimensionやheadを分けるが、各GPUが長いsequence全体を扱う構成では、context長そのものによる負荷は残る。

Context Parallelismは、sequenceをtoken方向に分け、複数GPUが別々のchunkを担当する方法である。ただしSelf-Attentionでは各tokenが過去のtokenを参照するため、文章を単純に分割して独立計算することはできない。GPU間でKeyとValueを交換しながら、元と同じAttentionを計算する必要がある。

分けるのはmodelではなくsequenceである

128K tokenのpromptをContext Parallel size 4で処理する単純な例では、各GPUが32K token分の入力とactivationを担当する。weightはContext Parallel group内で複製される場合がある。

ほかのparallelismとの違いは分割軸にある。

方法	主な分割軸	解決したい制約
Tensor Parallelism	hidden dimension、head、matrix	大きいlayerとweight
Pipeline Parallelism	layerの深さ	深いmodelのweight
Data Parallelism	request、batch	request throughput
Context Parallelism	sequence length	長いcontextのactivationとAttention

実際にはTPとCPを組み合わせる。TP groupが一つのtoken位置のlayer計算を分担し、CP groupが異なるtoken範囲を担当する。

local Queryは全contextのKeyとValueを必要とする

Attentionでは、あるtokenのQueryが参照可能なKeyとValueへscoreを計算する。GPU 0が前半token、GPU 1が後半tokenを持つだけでは、後半のQueryが前半のKVを参照できない。

そこで各GPUはlocal Queryを保持し、ほかのGPUが持つKV blockを順番に受け取ってAttentionを更新する。Ring方式では、KV blockを隣のrankへ送り、受け取ったblockを計算しながら次へ回す。全blockを処理すれば、各GPUのlocal Queryについて全contextを反映した出力が得られる。

Context Parallelism for Scalable Million-Token Inferenceは、KVをringで回すpass-KVとQueryを回すpass-Qを使い分けるexactな方式を示した。近似的に一部contextを捨てるのではなく、通常のAttentionと同じ結果を目指す。

causal maskで仕事量が偏る

autoregressive LLMのPrefillでは、tokenは未来を参照できない。sequenceを連続chunkへ分けると、先頭側のQueryが参照する過去は少なく、末尾側は多い。

GPUごとに同じtoken数を割り当てても、Attentionの有効計算量は均等にならない。先頭chunkを担当するGPUが早く終わり、末尾chunkを担当するGPUが長く計算すれば、同期点でidleが生じる。

実装には、sequenceの前後を同じrankへ組み合わせる、causal maskで不要なblock計算を避ける、ring順序を調整するなどのload balanceが必要になる。単純なsequence length ÷ GPU数だけでは均等にならない。

communicationとAttention計算を重ねる

各GPUがKV blockの到着を待ってから計算し、終わってから次へ送るとnetwork待ちが直列に加わる。実装では、現在のblockでAttentionを計算しながら、次のblockを非同期に受信する。

効果は次の条件に左右される。

KV blockのbyte数
GPU間linkのbandwidthとlatency
Attention kernelの計算時間
CP groupのrank数
GQAやMQAによるKV head数
token chunkの大きさ
node内とnode間のtopology

計算時間の中へ通信を隠せる範囲ではscaleしやすい。chunkが小さい、networkが遅い、rankが多すぎる場合は通信が表へ出る。

長いPrefillで効果が出やすい

長いpromptのPrefillでは、多数tokenのAttentionとFFNをまとめて計算する。sequence方向に十分な仕事があるため、複数GPUへ分ける余地が大きい。

短いpromptでは、分割後のlocal token数が小さくなり、kernel効率が落ちる。communicationの固定costも相対的に大きい。すべてのrequestへCPを適用せず、context長がthresholdを超えた場合だけ別poolへroutingする設計が考えられる。

Prefill・Decode分離を使う場合、長文Prefill poolだけCPを大きくし、Decode poolは別のparallel構成にできる。ただしPrefill後のKV CacheをDecode側のlayoutへ変換・転送するcostを測る。

Decodeでは並列化の性質が変わる

Decodeは通常、1 requestにつき新しいQueryが1 tokenずつ増える。一方、過去のKV Cacheはcontext全体に広がる。

KV CacheをCP rankへ分割したままなら、各rankが担当KVとの部分Attentionを計算し、最終結果を集約する。Prefillのように多くのQuery tokenを分ける余地は小さく、collectiveやring通信の比率が上がりやすい。

長いcontextでは各Decode stepが読むKV量も大きいため、KVを複数GPUへ分散する価値はある。しかし短いcontextや高concurrencyでは、requestをData Parallel replicaへ分ける方が効率的な場合もある。Prefill speedupをそのままDecodeへ当てはめない。

Sequence Parallelismと同じではない

名称が似ているが、Megatron-LMのSequence ParallelismとContext Parallelismは範囲が違う。

Sequence ParallelismはTensor Parallelと組み合わせ、LayerNormやDropoutなど一部activationをsequence方向へshardする。Attentionに必要なfull contextの扱いはTP側へ戻す。

Context Parallelismはnetwork inputとすべてのactivationをsequence方向へ分け、AttentionでもKVをrank間で交換する。長いsequenceそのものを分散して扱う点が中心である。

frameworkごとに用語が異なる場合があるため、機能名だけでなく、AttentionのQ、K、Vのどれがどのrankへあり、いつ通信するかを確認する。

KV Cache layoutとPrefix Cachingをそろえる

CPでは、requestのKV Cacheがsequence chunkごとに複数rankへ分散する。CP size、chunk規則、TP size、KV dtypeが違うdeploymentへ、そのままblockを移せるとは限らない。

Prefix Cacheを使う場合、同じprefixのblock hashだけでなく、各blockのowner rankとlayoutを管理する。Cache-aware routingで同じCP groupへrequestを送っても、一部rankだけblockを失っていればcomplete hitではない。

Offloadingでは各rankがlocal KVをCPUやSSDへ移す方法と、共有storageへ一つのlogical sequenceとして保存する方法がある。復元時にrank mappingが変わる場合はreshardが必要になる。

context length上限と有効な長さは別である

Context Parallelismは計算とmemoryを分散するsystem技術であり、modelが学習した長さやposition encodingの限界を変えるものではない。

1M tokenを処理できるhardware構成でも、modelが遠い位置の情報を正確に使えるとは限らない。Needle retrieval、複数根拠の統合、時系列、長文要約などで品質を別に測る。

「OOMせず入力できた」と「長文を理解できた」を同じ成功指標にしない。

failure domainがCP groupへ広がる

一つのrequestが複数rankのcontext chunkへ依存するため、rankが一つ停止すると、そのrequest全体を継続できない場合がある。ringの一箇所が切れれば通信も止まる。

collective timeout、rank heartbeat、GPU error、network errorを監視する。途中rankの失敗後に別rankへchunkを再配置できない場合は、健全なgroupでPrefillを最初からやり直す。

長文Prefillの再実行は高価である。retry回数を制限し、deadline、queue、admission controlと組み合わせる。障害時に巨大requestの再試行が一斉集中しないようにする。

監視するmetric

requestごとのcontext長とlocal token数
CP size、TP size、rank placement
rankごとのAttention、FFN、通信時間
KV send・receive byte数と待機時間
communicationとcomputeのoverlap率
causal mask後のrank別有効block数
Prefill latencyとTTFT
DecodeのITLとKV Cache read量
rank別KV Cache使用率、free block、eviction
Prefix Cacheのcomplete hitとpartial hit
collective timeout、rank failure、再実行時間
context長別の長文task品質

平均Prefill latencyだけでなく、最も遅いrankと最も速いrankの差を見る。imbalanceが大きければGPUを増やしても待ち時間が残る。

負荷試験

8K、32K、128Kなどcontext長を段階的に増やす
CP size 1、2、4、8を同じmodelで比較する
node内とnode間のringを比較する
TPとCPの組み合わせを変える
causal promptとbidirectional workloadを分ける
長文requestと短文requestを同じqueueへ混ぜる
Prefill・Decode分離時のKV転送を含める
Prefix Cacheをcold、complete hit、partial hitに分ける
ringのrankとnetwork linkを一つずつ停止する
長文retrievalと要約品質をcontext長別に測る

速度だけでなく、GPU当たりtoken throughput、最大安定context、TTFT、ITL、再試行cost、品質を比較する。

Context Parallelismは長文を独立chunkへ切る技術ではない

Context ParallelismはsequenceをGPUへ分けるが、Attentionの依存関係は維持する。local Queryが全contextのKVを参照できるよう、ringなどでblockを交換し、部分結果を正しく集約する。

長いPrefillのmemoryとlatencyを分散できる一方、network、causal load imbalance、KV layout、failure domainが新しい制約になる。短文を含む平均値ではなく、対象context長と本番topologyでCP sizeを決める必要がある。