flashinfer
v0.2.0.post1
| Blog | Documentação | Slack | Fórum de Discussão |
O FlashInfer é um gerador de biblioteca e kernel para grandes modelos de linguagem que fornece implementação de alto desempenho dos kernels de GPU LLM, como Flashattion, SparsaSeation, Pageation, amostragem e muito mais. O Flashinfer se concentra na porção e inferência do LLM e oferece desempenho de ponta em diversos cenários.
Verifique nosso blog de lançamento em V0.2 para novos recursos!
As principais características do flashInfer incluem:
plan / run Fase de Atenção Computação, onde agendamos o cálculo de entradas de comprimento de variável no estágio plan para aliviar a questão da carga-equilíbrio.O FlashInfer suporta APIs Pytorch, TVM e C ++ (somente cabeçalho) e pode ser facilmente integrado aos projetos existentes.
Usar nossa API Pytorch é a maneira mais fácil de começar:
Fornecemos rodas pré -construídas para Linux. Você pode instalar o flashInfer com o seguinte comando:
# 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.htmlTambém oferecemos rodas noturnas para experimentar os últimos recursos da filial principal:
pip install flashinfer -i https://flashinfer.ai/whl/nightly/cu124/torch2.4Como alternativa, você pode construir flashinfer a partir da fonte:
git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install -e . -v Por padrão, o FlashInfer usa compilação just-in-time (JIT) para seus kernels. Para pré-compilar kernels essenciais, defina a variável de ambiente FLASHINFER_ENABLE_AOT=1 antes de executar o comando de instalação:
FLASHINFER_ENABLE_AOT=1 pip install -e . -vPara mais detalhes, consulte a instalação da documentação de origem.
Abaixo está um exemplo mínimo de usar os kernels de atenção/preenchimento de Atenção de Atenção 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 maskConfira a documentação para o uso dos kernels de decodificação/apêndice/preenchimento em lote e kernels em cascata de prefixo compartilhado.
Performinamos o desempenho do kernel Flashinfer com o NVBEnch e você pode compilar e executar os benchmarks com os seguintes comandos:
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 Você pode executar ./bench_{single/batch}_{prefill/decode} para comparar o desempenho (por ./bench_single_prefill ./bench_{single/batch}_{prefill/decode} --help mostrará as opções disponíveis.
O FlashInfer também fornece ligações de API e TVM C ++, consulte a documentação para obter mais detalhes.
Estamos entusiasmados ao compartilhar que o Flashinfer está sendo adotado por muitos projetos de ponta, incluindo, entre outros,:
O FlashInfer é inspirado nos projetos Flashattion 1 e 2, VLLM, Stream-K, Cutlass e Aitemplate.
