Descargar flashinfer - Descargar código fuente de flashinfer

flashinfer

Pitón

v0.2.0.post1

Descargar

Flashinfer

Biblioteca de kernel para servir LLM

FlashInfer es un generador de biblioteca y núcleo para modelos de idiomas grandes que proporciona una implementación de alto rendimiento de los núcleos de GPU de LLM, como flashatent, spareattention, pageattention, muestreo y más. FlashInfer se centra en la servicio e inferencia de LLM, y ofrece un rendimiento de última generación en diversos escenarios.

¡Consulte nuestro blog de lanzamiento V0.2 para ver nuevas funciones!

Las características principales de FlashInfer incluyen:

Kernels de atención escasos/densos eficientes : atención eficiente individual/por lotes para el almacenamiento de KV escaso (paginado)/denso en núcleos CUDA y núcleos de tensor (tanto FA2 y FA3). La atención de vector-spario puede lograr el 90% del ancho de banda de los núcleos densos con el mismo tamaño de problema.
Programación equilibrada de carga : Flashinfer Decoulples plan / Etapa run del cálculo de atención donde programamos el cálculo de las entradas de longitud variable en la etapa plan para aliviar el problema de la carga de carga.
Eficiencia de la memoria : FlashInfer ofrece atención en cascada para el cache de KV jerical, e implementa la fusión de la sede para acelerar la atención grupal y los núcleos eficientes para la atención de baja precisión y la atención de la cuerda fusionada para el cache KV comprimido.
Atención personalizable : traiga sus propias variantes de atención a través de la compilación JIT.
CUDAGRAPH Y TORCH.COMPILE COMPATIBLION : Los núcleos Flashinfer pueden ser capturados por CudaGraphs y Torch.com Pile para inferencia de baja latencia.
Operadores específicos de LLM eficientes : kernel fusionado de alto rendimiento para el muestreo Top-P, Top-K/Min-P sin la necesidad de clasificar.

API de Pytorch, TVM y C ++ (solo de encabezado) de soporte de FlashInfer, y se pueden integrar fácilmente en proyectos existentes.

Noticias

[16 de diciembre de 2024] Publicación de blog Flashinfer 0.2 - Kernels eficientes y personalizables para la inferencia de LLM para servir
[Septiembre 2024] Hemos lanzado un espacio de trabajo Slack para usuarios y desarrolladores de Flashinfer. ¡Únase a nosotros para obtener apoyo oportuno, discusiones, actualizaciones y intercambio de conocimientos!
[31 de enero de 2024] Inferencia en cascada de publicación de blog: decodificación de lotes de prefijo compartido de memoria eficiente
[31 de enero de 2024] Publicación de blog que acelera las autoatenciones para LLM sirviendo con flashinfer

Empezando

Usar nuestra API de Pytorch es la forma más fácil de comenzar:

Instalación

Proporcionamos ruedas preconstruidas para Linux. Puede instalar FlashInfer con el siguiente 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.html

También ofrecemos ruedas nocturnas para probar las últimas características de la rama principal:

pip install flashinfer -i https://flashinfer.ai/whl/nightly/cu124/torch2.4

Alternativamente, puede construir flashinfer desde la fuente:

git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install -e . -v

Por defecto, FlashInfer utiliza la compilación de justo en el tiempo (JIT) para sus núcleos. Para compilar previamente los núcleos esenciales, establezca la variable de entorno FLASHINFER_ENABLE_AOT=1 antes de ejecutar el comando de instalación:

FLASHINFER_ENABLE_AOT=1 pip install -e . -v

Para obtener más detalles, consulte la instalación de la documentación de origen.

Probándolo

A continuación se muestra un ejemplo mínimo de usar los núcleos de atención de decodificación/append/de preflexión 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 mask

Consulte la documentación para obtener el uso de núcleos de decodificación/appendición/preflamiento de lotes y núcleos en cascada de prefijo compartido.

Ejecutar puntos de referencia

Perfil FlashInfer Kernel Performance con NVBench y puede compilar y ejecutar los puntos de referencia con los siguientes 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

Puede ejecutar ./bench_{single/batch}_{prefill/decode} para comparar el rendimiento (por ejemplo ./bench_single_prefill para una atención previa a la requisito único). ./bench_{single/batch}_{prefill/decode} --help le mostrará las opciones disponibles.