[論文レビュー] FlashDecoding++: Faster Large Language Model Inference on GPUs

Ke Hong, Guohao Dai+7 more|arXiv (Cornell University)|Nov 2, 2023

Topic Modeling被引用数 14

ひとこと要約

FlashDecoding++ は、注意の softmax 同期を最適化し、フラット GEMM の活用を最大化し、データフロー適応性を高めることで、NVIDIA および AMD GPU 上で Hugging Face や他のベースラインと比較して大幅なスピードアップを実現する高速な LLM 推論エンジンです。

ABSTRACT

As the Large Language Model (LLM) becomes increasingly important in various domains. However, the following challenges still remain unsolved in accelerating LLM inference: (1) Synchronized partial softmax update. The softmax operation requires a synchronized update operation among each partial softmax result, leading to ~20% overheads for the attention computation in LLMs. (2) Under-utilized computation of flat GEMM. The shape of matrices performing GEMM in LLM inference is flat, leading to under-utilized computation and >50% performance loss after padding zeros in previous designs. (3) Performance loss due to static dataflow. Kernel performance in LLM depends on varied input data features, hardware configurations, etc. A single and static dataflow may lead to a 50.25% performance loss for GEMMs of different shapes in LLM inference. We present FlashDecoding++, a fast LLM inference engine supporting mainstream LLMs and hardware back-ends. To tackle the above challenges, FlashDecoding++ creatively proposes: (1) Asynchronized softmax with unified max value. FlashDecoding++ introduces a unified max value technique for different partial softmax computations to avoid synchronization. (2) Flat GEMM optimization with double buffering. FlashDecoding++ points out that flat GEMMs with different shapes face varied bottlenecks. Then, techniques like double buffering are introduced. (3) Heuristic dataflow with hardware resource adaptation. FlashDecoding++ heuristically optimizes dataflow using different hardware resource considering input dynamics. Due to the versatility of optimizations in FlashDecoding++, FlashDecoding++ can achieve up to 4.86x and 2.18x speedup on both NVIDIA and AMD GPUs compared to Hugging Face implementations. FlashDecoding++ also achieves an average speedup of 1.37x compared to state-of-the-art LLM inference engines on mainstream LLMs.

研究の動機と目的

• 実運用での高い計算コストにより、より速い LLM 推論を促進する。
• デコードフェーズの GEMM/アテンションのボトルネックを特定し、的確な最適化を提案する。
• 異なるハードウェアバックエンドや入力ダイナミクスに適応する統一フレームワークを開発する。
• 複数のモデルと GPU にわたり、最先端の推論エンジンに対する性能向上を示す。

提案手法

• アテンションの部分的な同期を排除するため、統一された最大値を持つ非同期 softmax を導入する。
• フラット GEMM を 8 でパディングしてダブルバッファを適用し、メモリ待機を隠すよう最適化する。
• ワークロードの転換点に基づいて CUDA Core と Tensor Core の実装を選択するヒューリスティックなデータフローを開発する。
• GEMV/GEMM の形状を最適なカーネルへマッピングするオフラインのルックアップベースの決定フローをプロファイルして構築する。
• NVIDIA および AMD GPU 上で Llama2, OPT, ChatGLM2 モデルを用い、複数のベースラインと比較して評価する。

実験結果

リサーチクエスチョン

• RQ1 統一された最大値を持つ非同期 softmax は、部分的 softmax 計算における同期オーバヘッドを排除できるか？
• RQ2 パディングによる無駄を生むことなく、デコードフェーズの LLM ワークロードに対してフラット GEMM をどのように最適化して利用率を向上させるか？
• RQ3 ヒューリスティックなデータフローは、変化する入力ダイナミクスとハードウェア構成に応じて GEMV/GEMM ワークロードを最も適切なハードウェアカーネルへ適応的にルーティングし、スループットを最大化できるか？
• RQ4 主流の GPU 上で Hugging Face および他の最先端 LLM 推論エンジンと比較した FlashDecoding++ の性能向上はどの程度か？

主な発見

• デコードフェーズで NVIDIA および AMD GPU 上で Hugging Face に対して最大 4.86× の加速を達成。
• AMD GPU 上で Hugging Face に対して最大 3.93× の加速を達成。
• さまざまなモデルに対して、最先端の LLM 推論エンジンである FlashDecoding を平均 1.37× 上回る加速を提供。
• FlashDecoding と比較して Tesla A100 で平均 1.37× のスピードアップを達成。
• Llama2、OPT、ChatGLM2 を含む複数のモデルで競争力のある性能を提供。
• 非対称 softmax、ダブルバッファ付きフラット GEMM、ハードウェア適応性のあるデータフローがバックエンドを跨いで効果を発揮することを証明。