Современные крупные языковые модели (LLM) стремительно меняют ландшафт искусственного интеллекта, предоставляя передовые возможности для обработки и генерации текста. Однако высокая производительность таких моделей неразрывно связана с эффективным использованием ресурсов вычислительной техники, особенно графических процессоров (GPU). Компании Nvidia удалось сделать значительный шаг вперед в данном направлении, выпустив открытую и мощную библиотеку FlashInfer — инструмент, обеспечивающий высокоскоростное и гибкое выполнение инференса на базе LLM. FlashInfer представляет собой кастомизируемую и оптимизированную библиотеку вычислительных ядер, специально разработанную для построения эффективных движков обслуживания запросов больших языковых моделей. Одна из ключевых особенностей FlashInfer заключается в инновационном подходе к оптимизации хранения кеша ключ-значение (KV-cache), которая достигается с помощью блочно-разреженных и компонуемых форматов.
Это позволяет существенно улучшить память, снизить избыточные обращения и, как следствие, повысить общую скорость обработки запросов. Архитектура FlashInfer строится вокруг динамической адаптивности — библиотека учитывает разнообразие размеров последовательностей, правил маскирования и схем позиционного кодирования, характерных для современных LLM-запросов. Вместо фиксированных реализаций FlashInfer использует шаблонный подход к вниманию (attention) с поддержкой компиляции при выполнении (JIT), что позволяет подстраивать вычислительные ядра под конкретные сценарии и модели, без необходимости пересборки всей системы или реализации специализированного кода на CUDA C++. Еще одним важным аспектом FlashInfer является его интеграция с ведущими фреймворками для обслуживания LLM, такими как MLC Engine, SGLang и vLLM, а также возможностью использоваться с кастомными движками. Благодаря этому обеспечивается высокая скорость внедрения новых алгоритмов и аппаратных решений, а также удобство разработчиков при интеграции новейших вычислительных ядер прямо в существующую инфраструктуру без значительных временных затрат.
Разработка FlashInfer началась как совместный исследовательский проект между Полом Алленом Школой компьютерных наук и инженерии Университета Карнеги-Меллона и компанией OctoAI (сейчас принадлежащей Nvidia). Проект был нацелен на создание универсальной, высокопроизводительной и расширяемой библиотеки ядров инференса для LLM, способной быстро адаптироваться под разнообразные техники, в том числе методы повторного использования кеша ключ-значение, что значительно повышает эффективность работы при генерации последовательностей. Система делит рабочие нагрузки на четыре основных класса операторов: Attention, GEMM (общие матричные умножения), коммуникация и сэмплирование токенов. Это разделение позволяет не только добиться высокой производительности каждого этапа, но и поддерживать гибкость при работе с различными моделями и их конфигурациями. Каждый оператор представлен в виде легковесного, высокопроизводительного коллектива, обеспечивающего простую интеграцию с любыми движками с минимальными изменениями кода.
В частности, для операторов внимания предусмотрено единое хранение кеша в формате блочно-разреженных матриц, что позволяет оптимально работать с динамическими вариантами длины последовательностей и различиями в маскировании. Для повышения производительности FlashInfer предлагает шаблоны и JIT-ядра с множеством переключателей, которые позволяют настроить работу в зависимости от типа внимания, группировки ключей и запросов, мультиголовной архитектуры и других параметров. Разработчики также внедрили инспектор–исполнительный интерфейс, совместимый с PyTorch, который перед запуском анализирует форму входных данных и шаблоны совместного префикса, тем самым передавая наиболее оптимизированные вычислительные ядра. Что касается матричных операций и коммуникаций, FlashInfer поддерживает как традиционные GEMM и GEMV вычисления, так и более сложные сценарии, включая grouped GEMM для множества мелких матричных умножений в одном вызове, а также all-to-all коммуникации, требуемые сложными архитектурами моделей с экспертами (mixture-of-experts) и методами LoRA. FlashInfer отбирает самые быстрые из доступных open-source или фирменных Nvidia ядер, включая использование вычислений с низкой точностью (fp4, fp8) на тензорных ядрах, и предоставляет их через унифицированный API, что исключает необходимость изменения логики приложений при обновлении аппаратуры или библиотек.
Для этапа генерации следующего токена FlashInfer предлагает инновационный метод выборки без глобальной сортировки лексикона. Традиционные методы с сортировкой значений логитов оказываются громоздкими и замедляют обработку, особенно при больших размерах словарей. В отличие от них, FlashInfer использует алгоритм отбора с отклонениями (rejection-based sampling), который динамически отсеивает маловероятные токены, обеспечивая существенное сокращение задержек при выборе токенов без потери точности и корректности результата. Такая архитектура дает универсальность и устойчивость будущих разработок: изменение стиля хранения кеша, внедрение новых механизмов внимания, обработка батчей произвольной длины и достижение новых целевых показателей задержек возможно без необходимости переписывать ядра или переключаться на центральные процессоры. Благодаря этому все критические участки пути инференса остаются на GPU, что обеспечивает максимальную производительность и минимальные задержки.
Для разработчиков FlashInfer предоставляет готовый пакет, доступный на PyPI, с простым для освоения API, который разделяет этапы планирования и выполнения. Это позволяет сначала произвести компиляцию и настройку вычислительных ядер под конкретные задачи и модели, а затем многократно использовать подготовленные конфигурации при генерации последовательностей. Такой подход повышает скорость и гибкость разработки, а также снижает нагрузку на команду инженеров. В части логит-процессинга FlashInfer предлагает модульный интерфейс для построения цепочек операций, включая масштабирование температуры, вычисление softmax, фильтрацию по топ-p и собственно выборку токенов. Этот подход облегчает создание пользовательских конфигураций с минимальными затратами времени и усилий.
FlashInfer получила признание сообщества и технических экспертов, подтверждением чему служит получение награды лучшей статьи на конференции MLSys 2025 за работу «FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving». Nvidia активно внедряет разработанные ядра в свои продукты, включая TensorRT-LLM, и поддерживает их интеграцию с популярными фреймворками и пользовательскими решениями. Подводя итог, FlashInfer от Nvidia — это значимый шаг в развитии технологий высокопроизводительного выполнения больших языковых моделей. Ее архитектура, ориентированная на скорость, масштабируемость и гибкость, позволяет максимально эффективно использовать возможности современных графических процессоров, предоставляя разработчикам удобный инструмент для быстрой интеграции и оптимизации новых методов инференса. Такой инструмент особенно актуален для компаний и исследовательских групп, работающих с генеративным ИИ и стремящихся снизить затраты на вычисления, повысить пропускную способность систем и сократить задержки при обслуживании пользователей.
Для тех, кто заинтересован в использовании FlashInfer, доступна подробная документация и репозиторий на GitHub, что позволяет быстро начать работу с библиотекой и глубже изучить ее возможности. В совокупности с мощным аппаратным обеспечением Nvidia FlashInfer становится незаменимым инструментом для достижения новых высот в области искусственного интеллекта и больших языковых моделей.
