flashinfer
v0.2.0.post1
| Blog | Dokumentation | Slack | Diskussionsforum |
FlashInfer ist ein Bibliothek und ein Kernelgenerator für Großsprachenmodelle, die eine leistungsstarke Implementierung von LLM-GPU-Kerneln wie Flashattention, Sparseattention, Pagetentention, Probenahme und vieles mehr bieten. FlashInfer konzentriert sich auf LLM-Servier- und Inferenz- und Inferenz und liefert modernste Leistung in verschiedenen Szenarien.
Weitere Informationen finden Sie in unserem V0.2 -Release -Blog!
Die Kernfunktionen von FlashInfer sind:
plan / run Stufe der Aufmerksamkeitsberechnung, bei der wir die Berechnung der Eingänge der variablen Länge in plan planen, um das Problem der Lastabschlusses zu lindern.FlashInfer unterstützt Pytorch-, TVM- und C ++ (Nur-Header-APIs) und kann leicht in vorhandene Projekte integriert werden.
Die Verwendung unserer Pytorch -API ist der einfachste Weg, um loszulegen:
Wir bieten vorgebaute Räder für Linux. Sie können FlashInfer mit dem folgenden Befehl installieren:
# 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.htmlWir bieten auch nächtliche Räder an, um die neuesten Funktionen aus der Hauptzweig zu probieren:
pip install flashinfer -i https://flashinfer.ai/whl/nightly/cu124/torch2.4Alternativ können Sie FlashInfer aus Quelle erstellen:
git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install -e . -v Standardmäßig verwendet FlashInfer Just-in-Time (JIT) -Kompilation für seine Kerne. Setzen Sie die Umgebungsvariablen variable FLASHINFER_ENABLE_AOT=1 , bevor Sie den Installationsbefehl ausführen:
FLASHINFER_ENABLE_AOT=1 pip install -e . -vWeitere Informationen finden Sie in der Installation aus der Quelldokumentation.
Im Folgenden finden Sie ein minimales Beispiel für die Verwendung von FlashInfer-Eingriffen Decodes/Anhängen/Vorab-Aufmerksamkeitskernel:
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 maskIn den Dokumentationen zur Verwendung von Batch Decodes/Anhang/Pre-Fill-Kerneln und Shared-Prefix-Kernkörnern.
Wir profilieren die FlashInfer -Kernel -Leistung mit NVBench und Sie können die Benchmarks mit den folgenden Befehlen kompilieren und ausführen:
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 Sie können ./bench_{single/batch}_{prefill/decode} ausführen, um die Aufführung zu refernen (z. ./bench_single_prefill für die Aufmerksamkeit der Vorbereitung von Einzelgräben). ./bench_{single/batch}_{prefill/decode} --help zeigt Ihnen die verfügbaren Optionen an.
FlashInfer bietet auch C ++ - API- und TVM -Bindungen. Weitere Informationen finden Sie in der Dokumentation.
Wir freuen uns, mitzuteilen, dass FlashInfer von vielen hochmodernen Projekten übernommen wird, einschließlich, aber nicht beschränkt auf:
FlashInfer ist inspiriert von den Projekte von Flashattention 1 & 2, VLLM, Stream-K, Cutlass und Aitemplate.
