В чем разница между prefill и decode stage в LLM инференсе?

Краткий тезис

Inference LLM состоит из двух принципиально разных этапов: prefill (обработка входного промпта) и decode (генерация ответных токенов). Prefill работает параллельно над всеми токенами промпта, является compute-bound (ограничен вычислениями) и определяет TTFT (Time To First Token). Decode работает последовательно, генерирует по одному токену, является memory-bound (ограничен пропускной способностью памяти) и определяет TPOT (Time Per Output Token). Оптимизации для этих этапов радикально различаются: для prefill — FlashAttention и параллельные вычисления, для decode — speculative decoding, KV cache и эффективное использование памяти.

1. Термин: LLM Inference (инференс большой языковой модели)

LLM Inference — процесс получения ответа модели на заданный входной текст. Большинство современных LLM (например, GPT, LLaMA, Mistral) являются авторегрессионными: они генерируют последовательность токенов, предсказывая следующий токен на основе предыдущих. Этот процесс делится на две стадии: prefill (первичная обработка промпта) и decode (генерация выходных токенов).

Почему это важно:

• Разные стадии имеют разные профили нагрузок на вычислительное оборудование (GPU/TPU).
• Оптимизация каждой стадии критически влияет на latency (задержку) и throughput (пропускную способность) системы.
• Понимание разницы позволяет выбирать правильные техники ускорения и эффективно распределять ресурсы.

2. Prefill stage — параллельная обработка промпта

Prefill (также называемый prefill phase или initialization) — этап, на котором модель получает весь входной промпт (последовательность токенов) и вычисляет первый выходной токен. На этом этапе:

• Все токены промпта обрабатываются параллельно (матрица внимания вычисляется для всех пар токентов).
• Вычислительная сложность внимания: O(n²) для полного внимания (n — длина промпта).
• Compute-bound: основным узким местом являются вычисления (умножение матриц), а не доступ к памяти.
• Время выполнения называется TTFT (Time To First Token) — задержка до генерации первого токена ответа.

Детали вычислений в prefill:

• Вычисляются Key (K) и Value (V) для всех токенов промпта.
• Вычисляется Query (Q) для последнего токена (или для всех, но используется только последняя позиция для генерации первого выходного токена).
• Матрица внимания softmax(Q K^T / sqrt(d)) требует O(n²) операций.
• Полученные KV-пары кэшируются в KV cache для использования на decode stage.

Пример (промпт из 1024 токенов):

• GPU выполняет одно большое матричное умножение для всех токенов.
• TTFT может составлять десятки миллисекунд (зависит от размера модели и длины промпта).

3. Decode stage — последовательная генерация ответа

Decode (также decoding phase или generation phase) — этап, на котором модель генерирует каждый последующий токен ответа по одному за раз. После prefill модель уже имеет KV cache для токенов промпта. При decode:

• На каждом шаге обрабатывается только последний сгенерированный токен.
• Внимание вычисляется между новым токеном (как Query) и всеми предыдущими токенами (Keys и Values из KV cache).
• Вычислительная сложность на шаг: O(n) (линейно зависит от длины уже сгенерированной последовательности).
• Memory-bound: основным узким местом является чтение KV cache из памяти (memory bandwidth), а не арифметические операции.
• Время на один токен называется TPOT (Time Per Output Token). Общее время дешифровки = TPOT × (количество выходных токенов).

Детали decode:

• На каждом шаге вычисляется только один новый набор Key и Value для нового токена.
• KV cache растёт на один элемент на каждом шаге.
• Чтение KV cache (размер которого O(n) токенов × d_model × layers) доминирует над вычислениями.
• Для ускорения используют speculative decoding, quantization, paged attention.

Пример (генерация 100 токенов):

• TPOT может составлять 5–15 мс на токен (зависит от размера модели и аппаратного обеспечения).
• Общая задержка decode = 100 × TPOT.

4. Сравнение prefill и decode

5. Prefill: compute-bound и FlashAttention

Compute-bound означает, что скорость выполнения ограничена вычислительной мощностью GPU (FLOP/s), а не пропускной способностью памяти (GB/s). Prefill выполняет много операций на каждый байт данных (операции/байт высокое), поэтому утилизация вычислительных блоков GPU высока.

FlashAttention — ключевая оптимизация для prefill:

• Разбивает внимание на блоки, которые помещаются в быструю память (SRAM) GPU.
• Уменьшает чтение/запись из HBM (High Bandwidth Memory) без потери точности.
• Ускоряет prefill в 2–5 раз.

Другие оптимизации prefill:

• Tensor parallelism (распределение слоёв между несколькими GPU).
• Fused kernels (объединение операций layer-norm, softmax, matmul в один kernel).
• Prefix caching (сохранение KV cache для повторяющихся частей промпта).

6. Decode: memory-bound и специфичные оптимизации

Memory-bound означает, что узким местом является пропускная способность памяти (memory bandwidth). На decode каждому шагу требуется прочитать весь текущий KV cache из HBM в GPU. Вычислений на байт данных мало (операции/байт низкое), поэтому GPU недогружена арифметикой.

Основной показатель — memory bandwidth utilization. Чем выше процент от пиковой memory bandwidth, тем лучше.

Оптимизации decode:

7. Влияние prefill и decode на пользовательский опыт

• TTFT определяет, как быстро пользователь видит первый токен ответа. Высокий TTFT делает систему медлительной.
• TPOT определяет скорость генерации текста (токенов в секунду). Низкий TPOT важен для потоковых чатов.
• Trade-off: Prefill можно ускорить, уменьшая качество (например, обрезая промпт), но это снижает качество ответа. Decode можно ускорить, используя speculative decoding, но это требует дополнительной модели.

Пример метрик для production:

• Target TTFT < 200 мс для интерактивных приложений.
• Target TPOT < 10 мс для плавного потока (10 токенов/сек — комфортное чтение).

8. Сравнение с non-autoregressive inference

В отличие от авторегрессионных LLM, существуют non-autoregressive модели (например, Masked Language Models, BERT, или T5 с параллельным декодированием). Они генерируют все выходные токены за один проход (нет decode stage). Однако для генеративных задач авторегрессия остаётся стандартом из-за качества, поэтому разделение на prefill и decode актуально.

9. Итог: ключевые различия

• Prefill — параллельный, compute-bound, TTFT, использует FlashAttention.
• Decode — последовательный, memory-bound, TPOT, использует KV cache и speculative decoding.
• Оптимизации prefill нацелены на увеличение FLOP utilisation; decode — на снижение memory footprint и увеличение memory bandwidth utilisation.
• Для практических систем баланс между prefill и decode важен: можно разделять вычислительные ресурсы (например, выделенные GPU для prefill и decode) или использовать гибридные подходы (например, chunked prefill — частичное предварительное заполнение во время decode).

Пет-проект для закрепления

Задача: Сравнить TTFT и TPOT для LLM на разных длинах промпта и продемонстрировать влияние prefill и decode.

Инструменты: Python, Hugging Face Transformers, PyTorch, библиотека datasets, можно дополнительно nvtop для мониторинга GPU.

Шаги:

1. Загрузите небольшую LLM (например, TinyLlama-1.1B или phi-2).
2. Напишите скрипт, который генерирует промпты разной длины (от 128 до 4096 токенов).
3. Используя model.generate(...) с output_attentions=False, замерьте:
   • Время до первого токена (засеките момент до вызова generate и момент получения первого токена).
   • Время каждого последующего токена (можно извлечь через token_by_token или по разнице во времени).
4. Постройте графики зависимости TTFT от длины промпта и TPOT от номера шага (или общей длины генерации).
5. Проведите эксперимент с включённым use_cache=True и use_cache=False, чтобы увидеть влияние KV cache на decode.
6. Добавьте FlashAttention (через attn_implementation="flash_attention_2") для моделей, которые его поддерживают, и сравните prefill-время.

Ожидаемый результат:

• Вы увидите, что TTFT растёт квадратично с длиной промпта (без FlashAttention) или более линейно (с FlashAttention).
• TPOT будет почти постоянным (независимо от длины генерации), но существенно ускорится при use_cache=True.
• Графики и выводы сформируют интуитивное понимание разницы между compute-bound и memory-bound.

Связь с другими вопросами

Навигация

• Предыдущий: 435
• Следующий: 437
• Индекс: 00. Индекс разборов