Перейти к содержимому
НейроКуб
Открытый материал · 3 мин
Даниил Баланс

Prefill быстрее почти в 20 раз. Почему это не обещает такой же скорости ответа

На примере K-Search разбираем, как оценить перенос GPU-кода на MLX: отдельно проверить корректность, prefill, decode и условия сравнения.

Ускорение ИИ-кода в двадцать раз может почти не изменить скорость появления следующих слов. Чтобы оценить перенос вычислений с CUDA на MLX, сначала разделите обработку входного текста — prefill — и последовательную генерацию — decode. Иначе один удачный участок скроет ограничение всего сценария.

Что стоит за большим множителем

В июльском разборе K-Search исследователи показали перенос знаний об оптимизации CUDA-ядер на Apple Silicon. Языковая модель предлагает изменения, а кандидаты компилируются и измеряются на оборудовании. Исходная работа о K-Search описывает поиск с отдельным планированием и реализацией; это препринт, а не независимая проверка всех результатов переноса.

В таблице для mamba-370m, f16, на M1 Max с 64 ГБ памяти prefill при длине 4096 токенов даёт 6743 токена/с у mlx-mamba против 339 у mlx-lm: примерно 19,9×. Для decode приведены 152 против 116 токенов/с — около 1,31×. Это данные авторов за июль, а не рейтинг актуальных версий библиотек.

График NeuroKub построен по таблице Shiyi Cao и соавторов в BAIR Blog. Он показывает отдельный режим decode; обложка — prefill при L=4096. Измерения редакция не повторяла.

Авторы объясняют большой выигрыш prefill параллельным вычислением рекуррентного состояния. Во время decode доступен один новый токен, поэтому та же возможность распараллеливания исчезает. Перенос оптимизации меняет конкретную операцию, а не автоматически ускоряет каждый этап.

Как проверить свой перенос

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

Затем отдельно измерьте prefill на нескольких длинах и decode на одинаковой длине продолжения. Отделите первый запуск от повторных: пользователь может платить временем за подготовку, которую удобный средний показатель скрывает. Запишите не только лучший прогон, но и разброс, а также пиковую память.

Условный критерий выбора: если приложение читает большой документ и выдаёт короткую метку, выигрыш prefill может быть важен. Если оно долго пишет после короткого запроса, ключевым станет decode. Это пример выбора метрики, а не прогноз конкретного ускорения.

Принимайте перенос по времени до нужного результата при сохранённой корректности. Большой множитель полезен только рядом с названием операции, конфигурацией и подходящей базой сравнения. Если эти условия не совпадают с вашей задачей, число остаётся поводом для пилота, а не обещанием.

Комментарии

Пока нет комментариев. Начните обсуждение.