Transformersで連続バッチ処理の性能問題が報告、未使用のKVキャッシュ容量も生成を遅らせる可能性

Hugging Faceプロジェクトのメンバーであるalbertvillanovaは10月20日、再現ケースを報告した。Transformersのgenerate_batchを連続バッチモードで実行し、Flash Attentionを使うと、実際のリクエストがキャッシュを十分に使い切らなくても、KVキャッシュ容量の増加に伴って実行時間が延びる可能性がある。報告は安定版5.20.2と開発版5.21.0.dev0を対象としており、現時点でissueは未解決のままだ。問題報告

テストではH100 80GBとQwen2.5-0.5B-Instructを使い、最大128トークンを生成するプロンプトを64件入力した。利用可能メモリの50%をキャッシュに割り当てた場合、ウォームアップ後の実行時間は1回あたり約11秒だった。割り当てを90%に増やすと約18.7秒となり、num_blocks=64を指定した場合は約1.9秒だった。分析トレースでは、GPU時間の大部分がprefillのattention呼び出しに費やされ、そのK/V入力形状がキャッシュ容量全体をカバーしていた一方、decode呼び出しのコストは低かった。これは報告者による特定の環境での分析結果であり、すべてのモデルやattentionバックエンドに当てはまるとは限らない。テストとトレースのデータ

この性能差は強化学習の学習にも及ぶ。別のTRL issueでは、Transformersの連続バッチ処理を有効にすると、GRPOとRLOOが生成のたびに、デフォルトで利用可能なGPUメモリの半分をキャッシュに割り当てると報告された。小規模なGRPOテストを5ステップ実行したところ、デフォルト設定では59秒かかったが、64ブロックを使うと15秒に短縮した。これは特定の生成設定での比較であり、学習フレームワーク全体に一般化できる高速化率ではない。TRLの報告

公式ドキュメントでは、キャッシュ容量、バッチあたりのトークン予算、スケジューリングを別々の制御項目として説明している。連続バッチ処理は各ステップでリクエストを調整し、GPU使用率を高める。今回のケースは、容量設定も性能テストに含める必要があることを示している。現時点の結果からは、短いプロンプト、小規模モデル、低並列度のワークロードでは、過剰な事前割り当ての影響を特に確認する価値がある。連続バッチ処理のドキュメント

今後は、prefill経路が有効なキャッシュ領域だけを処理するよう変更されるか、またTRLの割り当て方針が見直されるかを注視したい。ブロック数を減らす方法は、手元で比較を始める際の選択肢になる。ただし、長いシーケンスや並列処理時の容量は改めて検証する必要がある。報告ではSDPAの処理時間はキャッシュサイズによって大きく変わらなかったものの、小さなキャッシュを使ったFlash Attentionより遅かったため、最適な代替手段とみなすことはできない。