Скрытый узкий горлышко в инференсе LLM

Инференс больших языковых моделей делится на быстрый prefill для входа и постепенную autoregressive генерацию токенов. Длина входа известна заранее, но длина вывода — нет: это может быть одно слово или длинное сочинение. Эта неопределенность тормозит работу на GPU с ограниченной KV-кеш-памятью. Чтобы избежать переполнений, планировщики часто резервируют память под наихудший сценарий, что приводит к низкой загрузке и высоким задержкам.

Почему пессимизм вреден

Консервативные алгоритмы вроде Amax предполагают, что каждое задание развернется до предсказанного верхнего предела. Это предотвращает сбои, но ухудшает эффективность: в кеше хранится меньше запросов, параллелизм падает, задержка растет. На реальных наборах (например, LMSYS-Chat-1M) Amax показывает резкое ухудшение при росте неопределенности, иногда давая задержки в 5 раз выше оптимума. При миллиардах запросов это означает заметные лишние затраты на вычисления и энергию.

Amin: оптимизм с адаптацией

Группа исследователей из Пекинского университета, Стэнфорда и HKUST предложила Amin. Вместо учета верхних оценок Amin стартует оптимистично, полагая для каждого запроса минимально предсказанную длину. Это увеличивает первоначальные батчи и лучше заполняет KV-кеш.

Амин добавляет два ключевых механизма, чтобы не рисковать работой:

Динамическое уточнение: по мере генерации токенов Amin обновляет псевдонижнюю границу для каждого задания. Если уже сгенерировано 100 токенов, нижняя граница становится как минимум 100, что улучшает последующую упаковку.
Упорядоченная эвикция: при нехватке памяти Amin сортирует активные задания по текущим псевдонижним границам и выпихивает в первую очередь те, что меньше продвинулись, случайно решая равенства. Это защищает задания, близкие к завершению, и снижает потерю работы из-за рестартов.

Amin специально игнорирует верхние оценки, так как их трудно предсказывать надежно, в то время как нижние границы проще получить. Алгоритм работает за O(M log M) на шаг, где M — размер KV-кеша.

Теоретические гарантии и эмпирика

Amin подкреплен математическими доказательствами и экспериментами. Авторы анализируют конкурентное отношение Amin по сравнению с планировщиком с ретроспективным знанием длин вывода. Amin достигает O(log(alpha^-1)) конкурентного отношения, где alpha — отношение нижней к верхней границе, характеризующее неопределенность. В то время как консервативные алгоритмы могут иметь бесконтрольно растущие потери при увеличении неопределенности, Amin остается логарифмическим в худшем случае.

Для практичных распределений получаются небольшие коэффициенты:

Двухточечные распределения: отношение не более 1.5.
Геометрические распределения: не более 1.7.
Линейно-взвешенные геометрические: около 1.56.

На численных тестах с выборками из LMSYS-Chat-1M Amin часто близок к теоретическому оптимуму. При грубых предсказаниях Amin сравнялся с ретроспективным оптимумом, тогда как Amax отстал примерно в 2 раза. При шумных предсказаниях Amin давал до 5x меньшую задержку по сравнению с Amax.

Что это даёт на практике

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

Практические замечания

Amin требует нижней оценки длины вывода для каждого запроса, что получить проще, чем точный верхний предел.
Политика эвикции сохраняет прогресс, минимизируя потери сгенерированных токенов.
Сложность O(M log M) приемлема для типичных размеров KV-кеша.

Если вы разворачиваете LLM в продакшене, стоит прочитать статью и посмотреть псевдокод — это может дать быстрый путь к заметному ускорению ваших inference-пайплайнов.