Скачать flashinfer - flashinfer исходный код скачать

flashinfer

Питон

v0.2.0.post1

Скачать

Flashinfer

Библиотека ядра для LLM обслуживающей

FlashInfer-это библиотечный и ядро генератор для моделей крупных языков, который обеспечивает высокопроизводительную реализацию ядра графических процессоров LLM, таких как вспышка, разрезание, Pageattention, выборка и многое другое. FlashInfer фокусируется на LLM обслуживании и выводе, а также обеспечивает современные результаты в разных сценариях.

Проверьте наш блог V0.2 для новых функций!

Основные особенности FlashInfer включают:

Эффективные ядра разреженного/плотного внимания : эффективное одно/партитное внимание для разреженных (страховых)/плотных шаблонов KV на ядрах CUDA и тензора (как FA2, так и FA3). Внимание вектора-пассажиров может достичь 90% полосы пропускания плотных ядров с тем же размером проблемы.
Сбалансированная нагрузка. Планирование : FlashInfer DeCous plan / run Stage Вычисление внимания, где мы планируем вычисление входов с переменной длиной на этапе plan , чтобы облегчить проблему с имбалсом нагрузки.
Эффективность памяти : FlashInfer предлагает каскадное внимание для Hierical KV-Cache и реализует слияние директоров для ускорения сгруппированного внимания и эффективных ядер для низкого внимания и внимания с сжатым KV-Cache.
Настраиваемое внимание : приведите свои собственные варианты внимания через JIT-компиляцию.
Cudagraph и Torch.compile Совместимость : ядра Flashinfer могут быть захвачены Cudagraphs и Torch.compile для вывода с низкой задержкой.
Эффективные операторы, специфичные для LLM : высокопроизводительное плавленое ядро для отбора проб Top-K/Min-P без необходимости сортировки.

API API-файлы FlashInfer поддерживают Pytorch, TVM и C ++ (только для заголовка) и могут быть легко интегрированы в существующие проекты.

Новости

[16 дек.
[Сентябрь 2024] Мы запустили Slack Workspace для пользователей и разработчиков Flashinfer. Присоединяйтесь к нам для своевременной поддержки, обсуждений, обновлений и обмена знаниями!
[31 января 2024 г.] Каскадный вывод о поведении в блоге.
[31 января 2024 г.] Пост в блоге ускоряет самообслуживания для LLM, обслуживающего FlashInfer

Начиная

Использование нашего Pytorch API - самый простой способ начать:

Установка

Мы предоставляем предварительно построенные колеса для Linux. Вы можете установить FlashInfer со следующей командой:

 # 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

Мы также предлагаем ночные колеса, чтобы попробовать последние функции из главной ветви:

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

В качестве альтернативы, вы можете построить FlashInfer из источника:

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

По умолчанию FlashInfer использует Just-Time (JIT) компиляцию для своих ядров. Чтобы предварительно выполнять необходимые ядра, установите переменную среды FLASHINFER_ENABLE_AOT=1 перед запуском команды установки:

FLASHINFER_ENABLE_AOT=1 pip install -e . -v

Для получения более подробной информации обратитесь к установке из исходной документации.

Пробовать это

Ниже приведен минимальный пример использования одноразового декодирования/приложения/приложения/преподавателя ядра.

 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

Ознакомьтесь с документацией для использования партийных декодирования/приложения/префиллевых ядров и каскадных ядер общего предварительного профиля.

Запустите тесты

Мы профиляем производительность ядра FlashInfer с помощью NVBench, и вы можете скомпилировать и запускать тесты со следующими командами:

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

Вы можете запустить ./bench_{single/batch}_{prefill/decode} , чтобы сравнить производительность (например, ./bench_single_prefill для предварительного внимания с одним запросом). ./bench_{single/batch}_{prefill/decode} --help