🌟 FlashInfer: библиотека ускорения LLM-инференса на GPU

🌟 FlashInfer: библиотека ускорения LLM-инференса на GPU.

FlashInfer - это библиотека для ускорения работы с LLM, созданная NVIDIA, чтобы объединить скорость обработки на GPU и гибкость для разработчиков. Еt главная цель — сократить время вывода текста, одновременно позволяя инженерам быстро внедрять новые алгоритмы и адаптировать решения под разные задачи.

Ее архитектура спроектирована так, чтобы оставаться актуальной при появлении новых алгоритмов: будь то методы повторного использования кэша или эксперименты с форматами внимания. Плюс к этому, библиотека легковесна, она не требует установки лишних зависимостей, а ее API напоминает стандартные инструменты PyTorch.

FlashInfer базируется на 2 принципах : эффективное управление памятью и динамическое планирование вычислений. Библиотека оптимизирует хранение KV-cache через блочно-разреженные структуры, уменьшая объем лишних обращений к памяти.

Это особенно важно при обработке запросов с разной длиной текста. Также используется технология JIT-компиляции, которая на лету генерирует оптимизированные CUDA-ядра под конкретную задачу.

Архитектура FlashInfer разбита на 4 модуля: Attention, GEMM, Communication и Token sampling.

🟢«Attention» работает с любыми схемами маскирования и позиционного кодирования, используя унифицированное представление кэша как разреженной матрицы.

🟢GEMM и Communication отвечают за матричные операции, включая сложные сценарии вроде grouped-GEMM (множество мелких умножений за один вызов). Для распределенных систем реализованы алгоритмы all-reduce и all-to-all, что критично для MoE-моделей.

🟢"Token sampling" ускоряет генерацию текста, заменяя традиционные сортировки вероятностей на rejection-based алгоритмы, отсекающие маловероятные варианты на лету.

FlashInfer поддерживает PyTorch через собственные операторы и DLPack API, тем самым упрощает внедрение в фреймворки vLLM и SGLang. Благодаря разделению процесса на этапы «планирования» и «запуска» библиотека минимизирует задержки: на первом шаге выбирается оптимальное ядро под параметры запроса, а затем оно переиспользуется для последующих аналогичных задач.

📌 Лицензирование: Apache 2.0 License.

🟡Статья

🟡Документация

🟡Arxiv

🖥GitHub

@machinelearning

#AI #ML #LLM #FlashInfer #NVIDIA