Зачем нужен очередной inference-движок?
vLLM хорош, но он не понимает Liquid Foundational Models. Hugging Face Transformers тормозит на батчах разной длины. А стандартный PyTorch с model.generate() просто сжирает всю память, если попробовать обрабатывать несколько запросов одновременно. Знакомо?
Я взял LFM2-350M - свежую модель от Liquid AI - и понял, что существующие инструменты с ней не дружат. Пришлось писать свой движок. Результат: 50-кратное ускорение на батчах из 8 запросов на одной RTX 3090.
Не путайте LFM с обычными трансформерами. У них другой механизм внимания, и стандартные оптимизации тут не работают.
Архитектура LFM: что ломает привычные оптимизации
Liquid Foundational Models используют не стандартный multi-head attention, а что-то среднее между линейными вниманиями и state space моделями. Ключевая фишка - динамическое перераспределение вычислительных ресурсов между токенами.
В обычном трансформере каждый токен получает одинаковое «внимание». В LFM - нет. Модель сама решает, какие части контекста важнее, и тратит на них больше вычислений. Звучит круто, пока не пытаешься сделать батчинг.
Гибридный KV-кэш: память против скорости
Стандартный подход: хранить ключи и значения для всех слоев и всех токенов. На LFM2-350M с контекстом 8192 токенов это ~3.5 ГБ только на кэш. Неприемлемо.
Мой вариант: гибридный кэш. Первые N слоев (где внимание наиболее «жидкое») кэшируем полностью. Остальные слои - вычисляем на лету, но с оптимизацией.
| Стратегия | Память | Скорость | Применимость |
|---|---|---|---|
| Полный кэш (vLLM-style) | 3.5 ГБ | Быстро | Не подходит для LFM |
| Без кэша (наивный) | 0 ГБ | Очень медленно | Только для тестов |
| Гибридный (моя реализация) | 1.2 ГБ | Быстро | Идеально для LFM |
Магия в том, чтобы найти точку перелома. Для LFM2-350M это 12 слоев из 24. Кэшируем первые 12, остальные вычисляем с переиспользованием промежуточных результатов.
Ragged prefill: когда запросы разной длины
Представьте: один пользователь спрашивает «Привет», другой отправляет техническую документацию на 2000 токенов. В стандартном батчинге придется дополнять короткие запросы до длины самого длинного. Тратить вычисления на padding-токены - преступление.
Решение: ragged prefill. Обрабатываем каждый запрос до его реальной длины, но делаем это параллельно на GPU. Секрет в правильной работе с CUDA streams и группировке операций.
Если не знакомы с низкоуровневой оптимизацией на CUDA, почитайте статью про кастомные CUDA ядра. Без этого дальше будет сложно.
1Группируем запросы по длинам
Вместо одного большого батча создаем микробатчи: запросы длиной 1-10 токенов, 11-50, 51-200, 201+. Для каждой группы - свой граф вычислений.
2Используем асинхронные копирования
Пока GPU обрабатывает одну группу, CPU готовит следующую. Звучит просто, но в PyTorch по умолчанию все синхронно. Придется лезть в torch.cuda.stream.
3Динамическое планирование
Новые запросы приходят во время обработки старых. Движок должен решать: добавить к текущему батчу или подождать следующего цикла. Я использую простую эвристику: если новый запрос короче среднего по батчу в 3 раза - ждем.
Оптимизация декодирования: где теряется 90% времени
Prefill - это разово. Декодирование - токен за токеном, для каждого запроса в батче. Наивная реализация делает отдельный вызов модели для каждого токена каждого запроса. Безумие.
Мой подход: групповое декодирование. Все запросы, которые готовы генерировать следующий токен (не достигли max_length, не сгенерировали stop_token), объединяются в один вызов.
Но с LFM есть нюанс: из-за «жидкого» внимания разные запросы требуют разного количества вычислений даже для одного токена. Приходится балансировать между группировкой и индивидуальной обработкой.
Цифры, а не слова
Тестировал на RTX 3090 (24 ГБ), LFM2-350M, контекст 8192:
- Один запрос: 45 токенов/сек (базовый PyTorch)
- Батч из 8 одинаковых запросов: 6 токенов/сек (наивный батчинг)
- Батч из 8 разных запросов (ragged): 2200+ токенов/сек (мой движок)
50-кратное ускорение - не маркетинг. Это следствие правильной работы с аппаратурой.
Сравнение с альтернативами
vLLM: Не поддерживает LFM. Можно попробовать добавить, но придется переписывать ядро внимания. Если готовы к этому - вариант.
Text Generation Inference (Hugging Face): Лучше, чем чистый Transformers, но все еще отстает на ragged батчах. Плюс - проще развернуть.
TensorRT-LLM: Максимальная производительность, но подготовка модели занимает часы. Для экспериментов - слишком долго.
Мой движок: 2000 строк на PyTorch, работает из коробки с LFM, но только на CUDA. Для AMD или CPU придется дописывать.
Если работаете на AMD, посмотрите материал про HyperNova на AMD GPU. Там свои оптимизации.
Кому это нужно?
Разработчикам, которые:
- Работают с нестандартными архитектурами (LFM, State Space, MLP-Mixers)
- Нужен максимальный throughput на одной карте
- Готовы потратить неделю на настройку, чтобы потом экономить часы вычислений
- Хотят понять, как inference-движки работают изнутри
Если вы просто хотите запустить Llama или Mistral - берите vLLM. Не усложняйте.
Но если экспериментируете с архитектурами вроде Genesis-152M-Instruct или той же LFM - свой движок становится необходимостью.
Что в итоге?
Писать inference-движок с нуля - не магия. Это инженерная работа: понять архитектуру модели, найти узкие места, применить правильные оптимизации.
Гибридный кэш и ragged prefill - не уникальные техники. Их используют в vLLM, TGI, TensorRT-LLM. Просто для LFM они реализованы иначе.
Самый важный урок: не пытайтесь скопировать оптимизации из одного движка в другой. Поймите, почему они работают в исходном контексте, и адаптируйте под свою задачу.
И да, 50-кратное ускорение - это не предел. С кастомными CUDA ядрами для attention-слоев LFM можно выжать еще 2-3x. Но это уже для фанатов.
P.S. Если собираетесь разворачивать что-то подобное в продакшене, сначала почитайте про стратегии масштабирования. Одна карта - это только начало.
