flashinfer
v0.2.0.post1
| Blog | Dokumentasi | Slack | Forum Diskusi |
FlashInfer adalah generator perpustakaan dan kernel untuk model bahasa besar yang menyediakan implementasi berkinerja tinggi dari kernel GPU LLM seperti flashattention, sparseattention, pageattention, sampling, dan banyak lagi. FlashInfer berfokus pada penyajian dan inferensi LLM, dan memberikan kinerja canggih di berbagai skenario.
Periksa blog rilis v0.2 kami untuk fitur baru!
Fitur inti dari FlashInfer meliputi:
plan / run Stage of Attention Computation Di mana kami menjadwalkan perhitungan input panjang variabel dalam tahap plan untuk mengurangi masalah penindasan beban.FlashInfer mendukung API Pytorch, TVM dan C ++ (hanya header), dan dapat dengan mudah diintegrasikan ke dalam proyek yang ada.
Menggunakan API Pytorch kami adalah cara termudah untuk memulai:
Kami menyediakan roda prebuilt untuk Linux. Anda dapat menginstal FlashInfer dengan perintah berikut:
# 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.htmlKami juga menawarkan roda malam yang dibangun untuk mencoba fitur terbaru dari cabang utama:
pip install flashinfer -i https://flashinfer.ai/whl/nightly/cu124/torch2.4Atau, Anda dapat membangun flashInfer dari sumber:
git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install -e . -v Secara default, FlashInfer menggunakan kompilasi just-in-time (JIT) untuk kernelnya. Untuk pra-kompilasi kernel penting, atur variabel lingkungan FLASHINFER_ENABLE_AOT=1 sebelum menjalankan perintah instalasi:
FLASHINFER_ENABLE_AOT=1 pip install -e . -vUntuk detail lebih lanjut, lihat instalasi dari dokumentasi sumber.
Di bawah ini adalah contoh minimal dari menggunakan flashInfer's single-request decode/append/prefill perhatian kernel:
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 maskLihat dokumentasi untuk penggunaan decode batch/append/prefill kernel dan kernel cascading prefix bersama.
Kami Profil Kinerja Kernel FlashInfer dengan NVBench dan Anda dapat mengkompilasi dan menjalankan tolok ukur dengan perintah berikut:
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 Anda dapat menjalankan ./bench_{single/batch}_{prefill/decode} untuk membandingkan kinerja (mis ./bench_single_prefill untuk perhatian prefill single-request). ./bench_{single/batch}_{prefill/decode} --help akan menunjukkan kepada Anda opsi yang tersedia.
FlashInfer juga menyediakan C ++ API dan TVM Bindings, silakan merujuk ke dokumentasi untuk detail lebih lanjut.
Kami sangat senang berbagi bahwa FlashInfer sedang diadopsi oleh banyak proyek mutakhir, termasuk tetapi tidak terbatas pada:
FlashInfer terinspirasi oleh proyek Flashattention 1 & 2, VLLM, Stream-K, Cutlass dan AITemplate.
