APB - 清華大学、テンセント、およびその他の機関が共同開発した、分散型長コンテキスト推論フレームワーク。 - AiBoss

APBとは何ですか？

APB（Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs）は、清華大学などが共同で提案した分散型長文コンテキスト推論フレームワークです。疎なアテンションメカニズムとシーケンシャル並列推論により、長文テキストを処理する大規模モデルの効率性のボトルネックを効果的に解消します。APBは、より小さなアンカーブロックとパッシングブロックを採用し、クエリ認識型コンテキスト圧縮技術と組み合わせることで、重要な情報を正確に伝達しながら計算オーバーヘッドを削減し、長距離の意味的依存関係の効率的な処理を実現します。128Kのテキストにおいて、APBの推論速度はFlash Attentionの約10倍、NVIDIAのStar Attentionの約1.6倍と、優れた性能を発揮します。また、さまざまな分散環境やモデルサイズに対応できる優れた互換性も備えています。

APBの主な機能

長文コンテキスト推論を加速するAPBは、マルチホスト近似アテンション機構により推論速度を大幅に向上させ、Flash Attention、Ring Attention、Star Attentionと比較して、それぞれ最大9.2倍、4.2倍、1.6倍の速度向上を実現しています。シーケンス並列化と近似アテンションを組み合わせることで、APBはタスクパフォーマンスを維持しながら、計算および通信のオーバーヘッドを大幅に削減します。
高効率分散コンピューティング：
- コンテキストセグメンテーション入力シーケンスは複数のホストに均等に分散され、各ホストのローカルコンテキストブロックの前にアンカーブロックが追加されることで、入力シーケンスの最初の部分の可視性が維持されます。
- ブロック圧縮各ホストでは、Locretの保持ヘッダーを使用してKVキャッシュを圧縮し、通信および計算のオーバーヘッドを削減します。
- 通信メカニズムAllGather通信メカニズムは、圧縮されたコンテキストブロックをすべてのホストに送信し、先行するホストから重要なキーバリューキャッシュユニットを渡すためのパッシングブロックを構築します。
- 計算するアテンションは、アンカーブロック、トランジットブロック、およびローカルコンテキストブロックを組み合わせることによって、各ホスト上で計算されます。トランジットブロックはアテンション計算後に破棄され、以降の計算には関与しません。
非常に適応力が高いAPBは複数のモデルと並列構成をサポートし、さまざまな分散環境やモデルサイズに対応でき、優れた拡張性を備えています。アンカーブロックとトランジットブロックのサイズを調整することで、APBは長さの異なる入力シーケンスに対して最適なパフォーマンスを実現できます。
タスクのパフォーマンスを維持する長文コンテキスト推論タスクにおいて、APBはFull Attentionと同等の性能を発揮し、一部のタスクではFull Attentionよりも高速です。クエリ認識型コンテキスト圧縮により、APBはクエリ関連のコンテキスト情報をより正確に識別して伝達することができ、タスクのパフォーマンスを維持または向上させます。

APBの技術原則

スパースアテンションメカニズムAPBフレームワークは、計算コストを削減することで推論速度を向上させるために、スパースアテンションメカニズムを統合しています。スパースアテンションは、以下の方法で実装されています。
- 小型アンカーブロックStar Attentionと比較して、APBはアンカーブロックのサイズをコンテキストブロックの1/4または1/8に縮小するため、追加の計算オーバーヘッドを削減できます。
- Passing block長距離の意味的依存関係の問題に対処するため、APBは重要な情報を伝送するためのパッシングブロックを構築します。パッシングブロックは、先行デバイスからの重要なキーと値のペアで構成され、各コンテキストブロックは圧縮されて後続のGPUに送信され、パッシングブロックを構築します。
- クエリ認識型コンテキスト圧縮APBはアンカーブロックの先頭にクエリを埋め込むことで、コンテキストコンプレッサーがクエリの内容を確認し、クエリに関連するキーと値のペアをより正確に識別し、通信メカニズムを介して後続のデバイスに送信できるようにします。
シーケンス並列推論APBフレームワークは、シーケンス並列処理方式を採用し、長文テキストを複数のGPUに均等に分散させて並列処理を行います。また、ローカルなキーバリューキャッシュ圧縮と簡素化されたGPU間通信メカニズムにより、長文テキストにおける長距離の意味的依存関係の問題を解決します。

APSプロジェクトの住所

GitHubリポジトリ：https://github.com/thunlp/APB
arXiv技術論文：https://arxiv.org/pdf/2502.12085

APBの適用シナリオ

長文推論非常に長い入力シーケンスを処理する必要のあるアプリケーション。例えば、長文テキストの生成や長文テキストによる質問応答など。
マルチエージェントコラボレーション複数のエージェントが協力して長いコンテキスト情報を処理する必要があるシナリオ。
大規模モデルサービス分散環境において、長いコンテキストを効率的に処理する必要のあるサービスをモデル化する。
知識グラフの構築知識グラフの構築には、大量のテキストデータを処理して知識を抽出し統合する必要があります。APBフレームワークは、効率的なコンテキスト圧縮および伝送メカニズムにより、知識グラフ構築の効率を大幅に向上させることができます。
リアルタイム対話型システムリアルタイム対話型システムは、ユーザー入力を迅速に処理し、正確な応答を生成する必要があります。APBフレームワークは、効率的なコンテキスト圧縮および伝送メカニズムを通じて、リアルタイム対話型システムの効率を大幅に向上させることができます。