flashinfer
v0.2.0.post1
| Blog | Documentation | Slack | Forum de discussion |
FlashInfer est une bibliothèque et un générateur de noyau pour les modèles de grands langues qui fournit une implémentation haute performance des noyaux GPU LLM tels que FlashAtttention, SparsEttension, PageAttension, échantillonnage, etc. FlashInfer se concentre sur le service et l'inférence LLM, et offre des performances de pointe dans divers scénarios.
Consultez notre blog V0.2 pour de nouvelles fonctionnalités!
Les caractéristiques principales de FlashInfer comprennent:
plan / run de l'étape du calcul de l'attention où nous planifions le calcul des entrées de longueur variable à l'étape plan pour atténuer le problème de la charge de charge.FlashInfer prend en charge les API Pytorch, TVM et C ++ (en tête uniquement) et peut être facilement intégrée dans les projets existants.
L'utilisation de notre API Pytorch est le moyen le plus simple de commencer:
Nous fournissons des roues prédéfinies pour Linux. Vous pouvez installer FlashInfer avec la commande suivante:
# 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.htmlNous proposons également des roues nocturnes pour essayer les dernières fonctionnalités de la branche principale:
pip install flashinfer -i https://flashinfer.ai/whl/nightly/cu124/torch2.4Alternativement, vous pouvez construire FlashInfer à partir de la source:
git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install -e . -v Par défaut, FlashInfer utilise une compilation JIT-Time (JIT) pour ses noyaux. Pour pré-compiler les noyaux essentiels, définissez la variable d'environnement FLASHINFER_ENABLE_AOT=1 avant d'exécuter la commande d'installation:
FLASHINFER_ENABLE_AOT=1 pip install -e . -vPour plus de détails, reportez-vous à l'installation à partir de la documentation source.
Vous trouverez ci-dessous un exemple minimal d'utilisation des noyaux d'attention de décodage / annexe / préfabrication de FlashInfer:
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 maskConsultez la documentation pour l'utilisation des noyaux de décodage / annexe / préfilleurs et des noyaux en cascade de préfixe partagé.
Nous profilons les performances du noyau FlashInfer avec NVBench et vous pouvez compiler et exécuter les repères avec les commandes suivantes:
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 Vous pouvez exécuter ./bench_{single/batch}_{prefill/decode} ./bench_single_prefill ./bench_{single/batch}_{prefill/decode} --help
FlashInfer fournit également des liaisons API et TVM C ++, veuillez vous référer à la documentation pour plus de détails.
Nous sommes ravis de partager que FlashInfer est adopté par de nombreux projets de pointe, y compris, mais sans s'y limiter:
FlashInfer est inspiré par les projets FlashAtttention 1 et 2, Vllm, Stream-K, Cutlass et AITEmplate.
