flashinfer
v0.2.0.post1
|ブログ|ドキュメント|スラック|ディスカッションフォーラム|
FlashInferは、フラッシュアット、スパルシート、パジーテンション、サンプリングなど、LLM GPUカーネルの高性能実装を提供する大規模な言語モデル用のライブラリおよびカーネルジェネレーターです。 FlashInferはLLMのサービングと推論に焦点を当てており、多様なシナリオ全体で最先端のパフォーマンスを提供します。
新しい機能については、V0.2リリースブログを確認してください！
FlashInferのコア機能は次のとおりです。
plan runで計画段階の計算をスケジュールして、 plan段階で変数長入力の計算をスケジュールして、負荷の均衡の問題を軽減するためにスケジュールします。FlashInferは、Pytorch、TVM、C ++（ヘッダーのみ）APIをサポートし、既存のプロジェクトに簡単に統合できます。
Pytorch APIを使用することは、開始する最も簡単な方法です。
Linux用の事前に構築されたホイールを提供します。次のコマンドでFlashInferをインストールできます。
# For CUDA 12.4 & torch 2.4
pip install flashinfer -i https://flashinfer.ai/whl/cu124/torch2.4
# For other CUDA & torch versions, please check https://docs.flashinfer.ai/installation.htmlまた、メインブランチから最新の機能を試すために、夜間に構築されたホイールも提供しています。
pip install flashinfer -i https://flashinfer.ai/whl/nightly/cu124/torch2.4または、SourceからFlashInferを作成できます。
git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install -e . -vデフォルトでは、FlashInferはカーネルにJust-in-Time（JIT）コンピレーションを使用します。必須カーネルを事前にコンパイルするには、インストールコマンドを実行する前に、環境変数FLASHINFER_ENABLE_AOT=1を設定します。
FLASHINFER_ENABLE_AOT=1 pip install -e . -v詳細については、ソースドキュメントのインストールを参照してください。
以下は、Flashinferのシングルリクエストデコード/追加/Prefill注意カーネルを使用する最小限の例です。
import torch
import flashinfer
kv_len = 2048
num_kv_heads = 32
head_dim = 128
k = torch . randn ( kv_len , num_kv_heads , head_dim ). half (). to ( 0 )
v = torch . randn ( kv_len , num_kv_heads , head_dim ). half (). to ( 0 )
# decode attention
num_qo_heads = 32
q = torch . randn ( num_qo_heads , head_dim ). half (). to ( 0 )
o = flashinfer . single_decode_with_kv_cache ( q , k , v ) # decode attention without RoPE on-the-fly
o_rope_on_the_fly = flashinfer . single_decode_with_kv_cache ( q , k , v , pos_encoding_mode = "ROPE_LLAMA" ) # decode with LLaMA style RoPE on-the-fly
# append attention
append_qo_len = 128
q = torch . randn ( append_qo_len , num_qo_heads , head_dim ). half (). to ( 0 ) # append attention, the last 128 tokens in the KV-Cache are the new tokens
o = flashinfer . single_prefill_with_kv_cache ( q , k , v , causal = True ) # append attention without RoPE on-the-fly, apply causal mask
o_rope_on_the_fly = flashinfer . single_prefill_with_kv_cache ( q , k , v , causal = True , pos_encoding_mode = "ROPE_LLAMA" ) # append attention with LLaMA style RoPE on-the-fly, apply causal mask
# prefill attention
qo_len = 2048
q = torch . randn ( qo_len , num_qo_heads , head_dim ). half (). to ( 0 ) # prefill attention
o = flashinfer . single_prefill_with_kv_cache ( q , k , v , causal = False ) # prefill attention without RoPE on-the-fly, do not apply causal maskバッチデコード/付録/Prefillカーネルと共有Prefixカスケードカーネルの使用法のドキュメントをご覧ください。
NVBenchを使用してFlashinferカーネルのパフォーマンスをプロファイルし、次のコマンドでベンチマークをコンパイルして実行できます。
mkdir build
cp cmake/config.cmake build # you can modify the config.cmake to enable/disable benchmarks and change CUDA architectures
cd build
cmake ..
make -j12 ./bench_{single/batch}_{prefill/decode}を実行して、パフォーマンスをベンチマークします（シングルリケストのPrefillの注意のための./bench_single_prefillなど）。 ./bench_{single/batch}_{prefill/decode} --helpは、利用可能なオプションが表示されます。
FlashInferはC ++ APIおよびTVMバインディングも提供しています。詳細については、ドキュメントを参照してください。
Flashinferが以下を含むがこれらに限定されない多くの最先端のプロジェクトで採用されていることを共有できることに興奮しています。
FlashInferは、Flashattention 1＆2、Vllm、Stream-K、Cutlass、AITemplateプロジェクトに触発されています。
