Что такое chunked prefill и зачем он нужен?

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

Chunked prefill — это техника оптимизации инференса LLM, при которой длинный входной промпт разбивается на несколько частей (чанков), и этап prefill (одновременная обработка всех токенов промпта) выполняется не целиком, а порциями, чередуясь с этапами decode (генерация одного нового токена). Основная цель — снизить latency первого токена (TTFT) для очень длинных промптов (например, >8K токенов) ценой небольшого снижения общего throughput, что критично для интерактивных приложений вроде чат-ботов.

1. Термины: Prefill, Decode и TTFT

Чтобы понять prefill|chunked prefill, нужно чётко разделить два этапа генерации в LLM.

Prefill (или encoding): Этап, на котором модель обрабатывает все токены входного промпта за один проход. Вычисления на этом этапе — это матричное умножение (matmul) с участием всех токенов сразу, что позволяет эффективно использовать GPU (высокий compute utilization). Результат — KV cache (ключи и значения для каждого слоя и каждого токена), который сохраняется для последующего этапа decode.

Decode (или generation): Этап, на котором модель генерирует по одному новому токену за раз. На каждом шаге decode вычисляется только один новый токен, а KV cache из prefill используется для attention без пересчёта старых токенов. Этот этап memory-bound (ограничен пропускной способностью памяти), так как доминирует операция чтения KV cache из памяти.

TTFT (Time To First Token): Время от момента получения запроса до генерации первого нового токена. На TTFT напрямую влияет длительность этапа prefill: чем длиннее промпт, тем больше времени занимает prefill, и тем выше TTFT.

2. Проблема длинных промптов

Когда промпт становится очень длинным (например, 16K, 32K или 128K токенов), этап prefill сталкивается с двумя проблемами:

• Квадратичная сложность attention Вычислительная сложность self-attention растёт как O(n²), где n — длина промпта. Для n=32K это ~1 млрд операций attention за один проход.
• Ограничение памяти GPU KV cache для длинного промпта занимает огромный объём памяти (например, для модели 7B с 32 слоями и n=32K — это ~16 ГБ в FP16). Это может превысить объём HBM (High Bandwidth Memory) одного GPU.
• Блокировка decode Пока выполняется prefill, модель не может начать генерацию (decode). Пользователь ждёт первый токен всё дольше.

В результате TTFT для длинных промптов может достигать десятков секунд, что неприемлемо для интерактивных сценариев.

3. Идея chunked prefill

Chunked prefill предлагает разбить длинный промпт на несколько последовательных частей (чанков) и обрабатывать их по одному, чередуя prefill каждого чанка с decode предыдущих токенов.

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

Пример работы

1. Промпт длиной 12K токенов разбивается на 3 чанка по 4K.
2. Выполняется prefill первого чанка (4K токенов) → получаем KV cache для первых 4K.
3. Начинается decode: генерируется первый токен ответа.
4. Между decode-шагами (или в фоне) выполняется prefill второго чанка (4K) → KV cache расширяется до 8K.
5. Продолжается decode, затем prefill третьего чанка.
6. После обработки всех чанков генерация идёт в обычном decode-режиме.

4. Как это работает: пошаговый алгоритм

Рассмотрим реализацию chunked prefill на примере псевдокода:

def chunked_prefill(model, prompt_tokens, chunk_size=4096):
    # Разбиваем промпт на чанки
    chunks = [prompt_tokens[i:i+chunk_size] 
              for i in range(0, len(prompt_tokens), chunk_size)]
    
    kv_cache = None
    generated_tokens = []
    
    # Обрабатываем первый чанк (обычный prefill)
    kv_cache = model.prefill(chunks[0], kv_cache=None)
    
    # Начинаем decode
    next_token = model.decode(kv_cache)
    generated_tokens.append(next_token)
    
    # Для оставшихся чанков чередуем prefill и decode
    for chunk in chunks[1:]:
        # Prefill следующего чанка (добавляем к существующему KV cache)
        kv_cache = model.prefill(chunk, kv_cache=kv_cache)
        
        # Продолжаем decode (один или несколько шагов)
        for _ in range(len(chunk) // decode_steps_per_chunk):
            next_token = model.decode(kv_cache)
            generated_tokens.append(next_token)
    
    # После всех чанков — обычный decode до конца
    while not eos:
        next_token = model.decode(kv_cache)
        generated_tokens.append(next_token)
    
    return generated_tokens

Важный нюанс Prefill каждого следующего чанка выполняется не «в лоб», а с учётом уже существующего KV cache. Это означает, что при обработке второго чанка модель видит все предыдущие токены (из первого чанка) через KV cache, а новые токены второго чанка обрабатываются как prefill.

5. Математическая модель: влияние на TTFT и throughput

Обозначим:

• n — длина промпта (токенов)
• m — количество генерируемых токенов
• c — размер чанка (токенов)
• t_prefill(n) — время prefill для n токенов (растёт ~O(n²) из-за attention)
• t_decode — время одного decode-шага (константа для данной модели)

Без chunked prefill

• TTFT = t_prefill(n)
• Total latency = t_prefill(n) + m * t_decode

С chunked prefill (k чанков, k = n/c):

• TTFT = t_prefill(c) (только первый чанк!)
• Total latency ≈ t_prefill(c) + (m + n - c) * t_decode + (k-1) * t_prefill(c) / concurrency_factor

Ключевое TTFT снижается с t_prefill(n) до t_prefill(c). Для n=16K и c=4K это может быть снижение в 4-16 раз (из-за квадратичной сложности).

Плата Общее время генерации может немного вырасти, так как prefill чанков выполняется не за один проход, а с переключениями контекста. Throughput (токенов/сек) снижается на 5-15% в зависимости от реализации.

6. Компромиссы и границы применимости

Когда использовать

• Длинные промпты (>8K токенов) в интерактивных сценариях (чат-боты, ассистенты)
• Системы реального времени, где важен низкий TTFT
• RAG-системы с длинными контекстами (много документов в промпте)

Когда не использовать

• Короткие промпты (<2K) — выигрыш в TTFT минимален, а накладные расходы ощутимы
• Batch-обработка (offline) — там важнее throughput, а не latency
• Модели с поддержкой continuous batching (например, vLLM) — они уже решают эту проблему иначе

7. Сравнение с альтернативами

Chunked prefill часто комбинируют с continuous batching: в паузах между decode одного запроса можно выполнять prefill чанков другого запроса.

8. Реализация в популярных фреймворках

vLLM Использует continuous batching и prefix caching, что частично решает ту же проблему. Chunked prefill в чистом виде не реализован, но есть механизм preemption (вытеснение), который может прервать decode для обработки prefill другого запроса.

TensorRT-LLM Поддерживает in-flight batching и chunked prefill через опцию --chunked_prefill. Позволяет указать максимальный размер чанка.

Hugging Face TGI Реализует continuous batching и prefix caching, но chunked prefill как отдельная опция отсутствует.

Пример конфигурации в TensorRT-LLM

{
  "builder_config": {
    "max_batch_size": 64,
    "max_input_len": 32768,
    "max_output_len": 2048,
    "chunked_prefill": true,
    "chunk_size": 4096
  }
}

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

Задача Реализовать симуляцию chunked prefill на Python и сравнить TTFT для разных размеров чанков.

Инструменты Python, NumPy, Matplotlib (для визуализации), простая имплементация attention (без реальной модели).

Шаги:

1. Создайте класс SimulatedLLM, который эмулирует время выполнения prefill и decode:
   • prefill(n) — время пропорционально n² (имитация квадратичной сложности attention)
   • decode() — константное время
2. Реализуйте функцию run_inference(prompt_len, chunk_size, num_gen_tokens), которая симулирует:
   • Обычный инференс (один prefill + decode)
   • Chunked prefill (разбивка на чанки, чередование)
3. Замерьте TTFT и общее время для разных prompt_len (4K, 8K, 16K, 32K) и chunk_size (1K, 2K, 4K).
4. Постройте графики зависимости TTFT от длины промпта для разных стратегий.

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

• График, показывающий, что TTFT при chunked prefill остаётся низким даже для длинных промптов.
• Понимание, что общее время генерации незначительно растёт (на 5-10%).
• Вывод: chunked prefill — это trade-off между latency и throughput.

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

11. Навигация

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

Навигация

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