今回の論文
今回取り上げるのは、Ke Hong、Guohao Dai、Jiaming Xu らによる論文「FlashDecoding++: Faster Large Language Model Inference on GPUs」です。2023年11月2日に arXiv へ投稿され、その後 MLSys 2024 でも公開されました。研究分野は LLM サービング、GPU カーネル最適化、推論高速化です。URL は https://arxiv.org/abs/2311.01282 です。
この論文を選んだ理由は、モデル自体を変えずに、推論ランタイムの詰まり方を3つの低レイヤーボトルネックに分解して直している からです。量子化や speculative decoding のような大きな設計変更ではなく、既存の LLM 推論基盤をどう速くするかに直結しており、開発現場での再利用価値が高いです。
研究分野としては、GPU 上の LLM 推論システム、カーネル最適化、推論エンジン実装です。
どんな技術か
FlashDecoding++ は、LLM 推論の prefill と decode の両方で起きる GPU 非効率を減らす推論エンジンです。特に decode フェーズでは、1 トークンずつ生成する都合で行列演算の形が細長くなり、GPU が得意な並列計算を活かしにくくなります。
この論文のポイントは、遅さをひとまとめにせず、次の3つに分けて対処していることです。
- partial softmax を更新するたびに同期が入る
- decode 時の Flat GEMM が細すぎて計算器を遊ばせる
- 入力長やバッチサイズが変わっても同じデータフローを使ってしまう
FlashDecoding++ は、この3点に対してそれぞれ別の最適化を入れています。要するに、「attention を速くする」ではなく、softmax、GEMM、実行パス選択を個別に詰める推論ランタイム です。
課題
この技術が解決しようとしている課題は、LLM 推論が GPU 上で動いていても、実際にはかなり無駄が残っていることです。
何が難しいのかというと、LLM 推論は prefill と decode で計算の性質が大きく違うからです。prefill は長い入力をまとめて処理するので大きな GEMM が多く、GPU の並列性を使いやすいです。一方の decode は 1 トークンずつ進むため、行列の片側が極端に小さい Flat GEMM や GEMV が増え、GPU の計算ユニットが埋まりにくくなります。
既存手法ではどこに限界があるのかというと、たとえば attention の部分では partial softmax をタイル単位で計算した後に前回結果との同期更新が必要でした。論文では、これが Llama2-7B の attention 計算で約 20% のオーバーヘッドになると報告しています。また decode フェーズでは、cuBLAS や CUTLASS のような汎用ライブラリが Flat GEMM に強く最適化されているわけではなく、ゼロ埋めで大きなサイズへ合わせる実装では 50% 超の計算ロスが起きうると指摘しています。
なぜこの課題を解く必要があるのかというと、実際の AI システムではこの手の無駄がそのまま原価と待ち時間になるからです。社内チャット、コード補完、顧客向け生成 API、RAG の最終応答生成などでは、モデル精度を変えずに 1.2 倍から 1.5 倍速くなるだけでも UX と GPU コストの両方に効きます。
特に問題になるのは、バッチサイズや入力長が一定ではないサービスです。短い問い合わせと長文要約が同じ基盤に流れると、固定最適化では勝ちにくくなります。FlashDecoding++ はそこを「実行条件に応じてデータフローを切り替えるべき問題」として扱っています。
用語解説
- Prefill
- 入力プロンプト全体をまとめて読み、最初の出力トークンを出すまでの段階です。FlashDecoding++ では decode だけでなく prefill 側の softmax 処理も最適化対象になっています。
- Decode フェーズ
- 1トークンずつ自己回帰で続きを生成する段階です。行列の形が細くなりやすく、GPU 利用率が落ちやすいため、この論文の中心的な改善対象です。
- Flat GEMM
- 片側の次元が小さく、細長い形になった行列積です。decode 時の投影計算でよく現れ、通常の大きな GEMM 向け実装では効率が悪くなりやすいため、FlashDecoding++ は専用最適化を入れています。
- Partial Softmax
- attention 行列をタイルごとに分割して softmax を計算するやり方です。長い系列を扱うために必要ですが、従来はタイル間の同期更新がボトルネックになっていました。
- データフロー
- どの演算をどの順で、どの GPU リソースで処理するかという実行パスです。FlashDecoding++ では入力長やバッチサイズ、GPU 特性に応じて最適なデータフローが変わる前提で設計されています。
技術の仕組み
FlashDecoding++ は 1 つの大きなアルゴリズムというより、LLM 推論カーネルを 3 方向から改善する設計です。
基本アイデア
著者らは、LLM 推論の遅さを「attention」「線形層」「実行パス」の混ざった問題としてではなく、別々のボトルネックとして切り分けています。
1つ目は partial softmax の同期です。2つ目は decode 時に発生する Flat GEMM の低利用率です。3つ目は静的データフローです。FlashDecoding++ は、それぞれに独立した改善を入れたうえで、全体として推論エンジンに統合しています。
非同期 softmax と unified max value
attention では、QK^T で得たスコアに softmax をかけてから V を掛けます。系列長が長いと行列全体を一度に扱えないため、タイルごとに partial softmax を計算します。
従来は、各タイルの計算結果をまとめるために「前までの最大値・総和」と同期しながら更新する必要がありました。この同期が入ると、GPU 上でタイル並列に計算していても待ち時間が発生します。
FlashDecoding++ は、各 partial softmax が共通で参照できる unified max value を使うことで、partial 結果を個別に処理しやすくしています。これにより、前のタイル結果を待ちながら更新する同期操作を減らし、さらに fine-grained pipelining を組み合わせています。
論文では、この工夫により prefill で 1.18 倍、decode で 1.14 倍の改善が出たと報告しています。つまり softmax だけでも小さくない改善余地があるということです。
Flat GEMM 最適化
decode フェーズでは、バッチサイズが小さいほど投影計算が GEMV や Flat GEMM に近づきます。ここで汎用 GEMM 実装をそのまま使うと、小さすぎる行列を大きいブロックへ合わせるためにゼロ埋めが多くなり、演算器が空回りします。
FlashDecoding++ はこの問題に対して、従来のように 64 へ大きくそろえるのではなく、小さい粒度へ寄せたうえで、Flat GEMM の形状に応じて処理を切り替えます。
ゼロ埋めを減らす
論文では、従来実装が 64 単位で合わせていたところを、より小さい 8 単位程度へ抑える方針を取っています。これにより、特に小バッチ時の無駄な計算を減らせます。
double buffering を使う
Flat GEMM は単純に軽いわけではなく、メモリロード待ちで止まりやすいケースもあります。そこで FlashDecoding++ は double buffering を入れ、次に使うデータのロードと現在の計算を重ねます。論文では、この最適化で Flat GEMM 単体に対して最大 52% の高速化を報告しています。
hardware-aware な heuristic dataflow
3つ目の工夫は、どの実行パスを選ぶかを固定しないことです。バッチサイズが小さいとメモリ帯域が支配的になり、大きいと計算資源側が支配的になります。また NVIDIA と AMD でも Tensor Core 相当の使い方やキャッシュ特性が異なります。
FlashDecoding++ は、入力長、バッチサイズ、GPU のリソース特性を見て、GEMM、Flat GEMM、GEMV のどの流れが有利かを heuristic に選びます。論文では、静的データフローより最大 29% 速くなったとしています。
この発想の重要な点は、単一の最速カーネルは存在しない と認めていることです。実運用でも、短文チャットと長文生成が同じ設定で最適化できるとは限りません。FlashDecoding++ はそこを実行時選択で吸収しています。
処理の流れ
- prefill では partial softmax を unified max value ベースで処理し、同期コストを減らします。
- decode では行列形状を見て Flat GEMM 向け最適化を適用します。
- GPU と入力条件に応じて、どのデータフローを使うかを選びます。
- それらを統合した推論エンジンとして、prefill と decode の両方で throughput と latency を改善します。
つまり FlashDecoding++ は、単一の数学的トリックよりも、GPU 上の LLM 推論をシステムとして最適化する設計 と見ると理解しやすいです。
実験と結果
論文では、FlashDecoding++ が本当に速いのか、どの GPU・モデルで効くのか、各最適化に個別の効果があるのかを検証しています。
何を検証したのか
主な検証ポイントは次の3つです。
- 既存の LLM 推論エンジンよりどれだけ高速か
- softmax、Flat GEMM、heuristic dataflow それぞれに改善効果があるか
- NVIDIA と AMD の両方で効果が出るか
単に end-to-end で速いかだけでなく、個々のボトルネック改善も見ているのが特徴です。
どんなデータセットや評価指標を使ったのか
この論文はベンチマークデータセットの精度評価ではなく、推論性能評価が中心です。対象モデルは Llama2-7B、Llama2-13B、OPT-6.7B、ChatGLM2-6B で、Tesla A100、RTX 3090、AMD RX7900XTX、AMD MI210 など複数 GPU で比較しています。
評価指標は主に、decode フェーズの throughput、prefill の speedup、first-token latency、each-token latency です。入力長は 128、1k、8k、32k など、バッチサイズは 1、2、4、8 が中心です。つまり、実サービスでよくある「短文少量」と「長文多バッチ」の両方をかなり意識した評価です。
既存エンジンに対する速度改善
論文では、decode フェーズで Hugging Face 実装に対して最大 4.86 倍、prefill でも最大 1.40 倍の高速化を報告しています。さらに A100 上では FlashDecoding に対して平均 1.37 倍速いとしています。
ここで重要なのは、比較対象が単純な naive 実装だけではない点です。vLLM、DeepSpeed、TensorRT-LLM、OpenPPL、FlashDecoding など、すでに最適化された推論エンジンと比べても改善が出ています。つまり「基礎最適化を全部入れた後の最後の 20〜40%」を狙うタイプの研究です。
個別最適化の結果
非同期 softmax
論文では unified max value による非同期 softmax と fine-grained pipelining によって、prefill で 1.18 倍、decode で 1.14 倍の改善が報告されています。softmax の同期だけでも、意外と大きなボトルネックだったことがわかります。
Flat GEMM
Flat GEMM 最適化では、特に小さな M のケース、つまり小バッチ・逐次生成寄りの条件で効いています。論文では最大 52% の改善に加えて、小さい M=1,2 のケースでも 9% と 23% の改善があったとしています。インタラクティブなチャットや補完に効きやすい結果です。
heuristic dataflow
heuristic dataflow は、Tesla A100 で平均 10%、RTX 3090 で平均 20%、最大 29% の speedup を示しています。入力長 1024 の decode 条件で見ても、固定データフローより実行条件に応じて切り替えたほうが有利だと読めます。
結果から何が言えるのか
この結果から言えるのは、LLM 推論高速化では「新しい生成アルゴリズムを入れる」だけが手ではないということです。FlashDecoding++ は、モデルや出力分布を変えずに、カーネルの同期、行列形状、実行パス選択を最適化するだけで大きな改善が出る ことを示しています。
また、効果が NVIDIA だけでなく AMD 側にも広がっているのは重要です。GPU の選択肢が増える今、特定ベンダー専用の最適化だけでは足りないことがよくわかります。
何に使える？
FlashDecoding++ が向いているのは、同じモデル品質のまま、推論基盤の速度と原価を改善したいケースです。
チャット API や社内アシスタント
1トークンごとの待ち時間が効くチャット系では、decode 最適化の価値が大きいです。特にバッチサイズ 1 近辺のワークロードでは Flat GEMM 改善が効きやすく、体感応答速度を上げやすいです。
コード補完や開発支援
IDE 補完やターミナル支援では、数十ミリ秒単位の差が使い勝手に直結します。FlashDecoding++ のような低レイヤー最適化は、モデルを小さくするのではなく、同じモデルのまま補完レスポンスを速くする 方向で使えます。
自前 GPU サービングの原価圧縮
法人向け SaaS や社内基盤で自前 GPU を使っているなら、1.2 倍から 1.4 倍の throughput 改善でもコスト影響は大きいです。RAG やエージェント基盤の最終生成部分は呼び出し回数が多いため、こうした最適化の積み重ねが効きます。
AMD GPU を含むマルチハードウェア運用
NVIDIA だけでなく AMD でも改善を出しているので、将来的に GPU 調達の柔軟性を持たせたい基盤とも相性があります。ハードウェア差分を吸収する heuristic dataflow の発想は、ベンダー混在環境で特に役立ちます。
開発や事業へのヒント
この論文から得られるヒントは、推論基盤の差分はまだ十分に事業価値になる、ということです。
モデル選定だけでなくカーネル設計も競争力になる
多くのプロダクトは「どのモデルを使うか」で議論が止まりがちですが、実際の UX と原価は推論ランタイムでも大きく変わります。FlashDecoding++ は、同じ Llama2 を使っていてもサービング層で差が出ることを示しています。
小規模チームでも計測文化を持つべき
論文の本質は、理屈で最適化を語るだけでなく、入力長、バッチサイズ、GPU ごとに実測している点です。自分で AI アプリを作るなら、モデル精度ベンチマークだけでなく、first-token latency と token/s を継続測定する仕組みを持つ価値があります。
heuristic の価値を過小評価しない
研究では end-to-end の最適解を数式で出したくなりますが、実際の GPU サービングでは heuristic による切り替えが強いことがあります。FlashDecoding++ は、完全な理論最適より、実行条件ごとの現実的な分岐戦略が勝つ領域 があると教えてくれます。
既存プロダクト改善にも入れやすい
この論文の良いところは、モデル品質や API 仕様を大きく変えないことです。量子化やモデル入れ替えより導入リスクを抑えつつ、速度改善を狙える可能性があります。既存の RAG、エージェント、社内生成ツールのバックエンド改善案として考えやすいです。
限界
まず、FlashDecoding++ はかなり実装が重いです。softmax カーネル、Flat GEMM、データフロー分岐まで含めて最適化するため、単純なアプリケーションコードだけでは再現しにくく、ランタイムや CUDA/HIP の理解が必要です。
次に、効果はハードウェア依存です。論文でも AMD では double buffering がそのまま効きにくく、L1 キャッシュの小ささゆえに heuristic 依存が強いと述べています。つまり、ある GPU で効いた最適化が別 GPU で同じように効くとは限りません。
また、この研究は推論速度の論文であり、モデル品質自体を上げるものではありません。事実性、推論能力、ドメイン適合性に課題があるなら、RAG や fine-tuning と別に考える必要があります。
さらに、論文の主戦場は比較的低レイヤーなカーネル最適化です。もしマネージド API を使っていて内部ランタイムを触れないなら、そのまま導入することは難しいです。自前サービングか、推論エンジンを差し替えられる環境で価値が出やすいです。
最後に、評価は主に throughput と latency であり、運用上重要な連続バッチング、長時間稼働、マルチテナント競合、メモリ断片化などまでは深く検証していません。実サービス導入ではその点を別途確認する必要があります。
よくある質問
Q. FlashDecoding++ は speculative decoding のような生成アルゴリズムですか？
A. いいえ、主眼は生成アルゴリズムではなく推論エンジン最適化です。モデルの出力分布を変えるより、GPU 上で softmax、GEMM、実行パスをどう回すかを改善しています。
Q. どんなワークロードで特に効きますか？
A. 小バッチの逐次生成、つまりチャットやコード補完のような decode 寄りワークロードで特に効きやすいです。論文でも Flat GEMM 最適化は小さい M で効果が大きいとされています。
Q. vLLM や TensorRT-LLM を使っていても学べることはありますか？
A. あります。FlashDecoding++ の価値は、既存エンジンを置き換えることだけではなく、どこがボトルネックになりやすいかを明示している点です。softmax 同期、行列形状、動的データフローという視点は他の推論基盤にも応用できます。
Q. 自前で全部実装しないと意味がありませんか？
A. 完全再実装は重いですが、考え方だけでも有益です。たとえば小バッチ向け専用パスを分ける、GPU ごとにカーネルを切り替える、prefill と decode を別ベンチマークで測る、といった改善は小規模チームでも取り入れやすいです。
Q. RAG やエージェント開発にも関係ありますか？
A. はい。RAG やエージェントの品質改善そのものではありませんが、最終的な生成レイヤーの速度改善には効きます。検索やツール呼び出しを速くしても、最後の生成が遅ければ体感は改善しにくいため、基盤最適化として重要です。
今日の学び
この論文は、LLM 推論が遅い理由を softmax 同期、Flat GEMM の低利用率、静的データフローという3つの具体的な課題に分けて扱いました。そして FlashDecoding++ は、それぞれに対して非同期 softmax、double buffering、hardware-aware な heuristic を入れることで、推論エンジン全体を速くしています。
ここから得られるヒントは、AI プロダクトの差はモデルの賢さだけでなく、同じモデルをどれだけ効率よく動かせるか でも生まれるということです。特に自前 GPU サービングを行うなら、低レイヤー最適化は十分にプロダクト価値になります。