ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить chunked prefill для long context (32k токенов, TTFT -60%)

1. Цель задачи

Освоить технику chunked prefill (разбиение этапа prefill на фрагменты) для снижения времени до первого токена (TTFT) при работе с длинными промптами (до 32k токенов) в LLM инференсе. Научиться выбирать параметры чанкования, измерять влияние на TTFT и память, а также убедиться, что качество генерации не деградирует.

Ключевой результат Настройка inference-сервера (vLLM) с включённым chunked prefill, при которой TTFT для промпта длиной 32k токенов снижается не менее чем на 60% по сравнению с baseline (без chunked prefill) без потери качества ответов.

2. Исходные данные

Перед началом необходимо иметь:

Если нет реального GPU с 24 ГБ — симулируем:

1. Использовать модель меньшего размера (например, Qwen2.5-1.5B-Instruct) и пропорционально уменьшить длину контекста (8k токенов).
2. Целевой TTFT -60% остаётся, но абсолютные значения будут другими.
3. Измерять TTFT через time python-скрипта или vllm.entrypoints.openai.api_server с клиентом.

3. Технологический стек

4. Этапы выполнения

Этап 1: Подготовка окружения и baseline (1 час)

Действия

1. Установить vLLM и зависимости

   pip install vllm transformers pandas matplotlib aiohttp
   

   Проверить версию: python -c "import vllm; print(vllm.version)" (должна быть ≥0.6.0).

2. Скачать модель (если HuggingFace — авторизуйтесь с токеном).

   huggingface-cli login
   huggingface-cli download meta-llama/Llama-3.1-8B-Instruct --local-dir ./model
   

3. Запустить vLLM server без chunked prefill (baseline).

   python -m vllm.entrypoints.openai.api_server \
       --model ./model \
       --max-model-len 32768 \
       --max-num-seqs 2 \
       --gpu-memory-utilization 0.9 \
       --disable-chunked-prefill
   

   Ключевой параметр: --disable-chunked-prefill (в vLLM по умолчанию chunked prefill включён для контекстов >32768, но мы его принудительно отключаем).

4. Подготовить тестовый промпт длиной 32k токенов.
   Создать Python-скрипт generate_prompt.py:

   from transformers import AutoTokenizer
   tokenizer = AutoTokenizer.from_pretrained("./model")
   base_text = "Напиши подробное эссе о пользе чтения. " * 2000
   tokens = tokenizer.encode(base_text)
   while len(tokens) < 32000:
       tokens += [tokenizer.pad_token_id] * 100
   prompt = tokenizer.decode(tokens[:32000])
   with open("prompt_32k.txt", "w") as f:
       f.write(prompt)
   

   Убедиться, что длина ровно 32k (+- несколько токенов).

5. Измерить baseline TTFT для 32k.
   Запустить клиент:

   import asyncio, time, aiohttp
   async def measure_ttft(prompt):
       async with aiohttp.ClientSession() as session:
           payload = {"model": "./model", "messages": [{"role": "user", "content": prompt}], "max_tokens": 1}
           t0 = time.time()
           async with session.post("http://localhost:8000/v1/chat/completions", json=payload) as resp:
               data = await resp.json()
               t1 = time.time()
       return t1 - t0, data
   ttft, _ = asyncio.run(measure_ttft(open("prompt_32k.txt").read()))
   print(f"Baseline TTFT: {ttft:.2f} сек")
   

   Записать значение. Ожидаем порядка 2-5 секунд на A100 (зависит от модели и GPU).

Ожидаемый результат этапа Работающий vLLM сервер (без chunked prefill), измеренный baseline TTFT для 32k.

Этап 2: Включение chunked prefill с разными параметрами (1,5 часа)

Действия

1. Остановить сервер (Ctrl+C).

2. Запустить vLLM с chunked prefill, используя минимальный размер чанка по умолчанию

   # chunked prefill включён по умолчанию для max_model_len > 32768, 
   # но для явного контроля используем --enable-chunked-prefill и --max-num-batched-tokens
   python -m vllm.entrypoints.openai.api_server \
       --model ./model \
       --max-model-len 32768 \
       --max-num-seqs 2 \
       --gpu-memory-utilization 0.9 \
       --enable-chunked-prefill \
       --max-num-batched-tokens 4096
   

   Параметр --max-num-batched-tokens определяет размер чанка (количество токенов prefill на один batch). Значение 4096 — разумный старт.

3. Повторить замер TTFT (тот же скрипт).
   Записать значение. Ожидаем снижения TTFT в 2-3 раза (50-70%).

4. Оптимизировать размер чанка — провести эксперименты с --max-num-batched-tokens:

   Значение	Ожидаемый эффект на TTFT	Потребление памяти
   2048	Ещё меньше TTFT, но больше накладных расходов	Низкая
   4096	Хороший баланс	Средняя
   8192	TTFT снижается меньше, но overhead меньше	Высокая
   16384	Почти как baseline	Очень высокая

   Для каждого значения:

   • Перезапустить сервер (можно скриптом).
   • Замерить TTFT (3 повторения, усреднить).
   • Записать nvidia-smi в момент выполнения (память).

5. Проверить влияние на качество генерации.
   Взять 5 вопросов, требующих длинного ответа (например, "Объясни квантовую запутанность"), отправить их модели с chunking (чанк 4096) и без chunking, сравнить ответы (субъективно на логику и факты). Отсутствие деградации — обязательное требование.

Ожидаемый результат этапа Таблица зависимости TTFT от размера чанка, выбранное оптимальное значение (максимальное снижение TTFT при сохранении качества и без OOM).

Этап 3: Глубокий анализ и профилирование (1 час)

Действия

1. Включить логирование vLLM
   Добавить --log-requests и --log-stats при запуске. Собирать время prefill и decode отдельно.

2. Измерить детально время prefill vs decode
   Использовать встроенные метрики vLLM (через /metrics). Написать скрипт, который:

   • Отправляет запрос с stream=True.
   • Замеряет время получения первого токена (TTFT).
   • Анализирует time_to_first_token из логов.

3. Сравнить теоретическую модель
   TTFT_baseline ≈ prefill_time (весь промпт за один проход)
   TTFT_chunked ≈ (ceil(seq_len / chunk_size)) * (chunk_time) + (first_decode_time)
   Построить график зависимости TTFT от chunk_size. Вывести расчётные кривые.

4. Проверить влияние на throughput
   Сделать 10 запросов параллельно (batch=10) с длиной 8k токенов (чтобы не OOM), замерить TTFT и tokens per second. Сравнить baseline и chunked prefill.

Ожидаемый результат этапа График TTFT(chunk_size) + выводы о trade-off между TTFT и throughput.

Этап 4: Тонкая настройка и итоговая конфигурация (30 минут)

Действия

1. На основе этапов 2 и 3 выбрать итоговую конфигурацию
   Рекомендуемый старт: --max-num-batched-tokens 4096 (или 2048, если память позволяет). Дополнительно можно настроить --gpu-memory-utilization (0.85–0.95).

2. Зафиксировать config в файле vllm_chunked_config.yaml:

   model: ./model
   max-model-len: 32768
   max-num-seqs: 2
   gpu-memory-utilization: 0.9
   enable-chunked-prefill: true
   max-num-batched-tokens: 4096
   

3. Повторить финальные замеры TTFT (3 прогона с промптом 32k) и проверить что целевые -60% достигнуты.

4. Проверить корректность ответов на том же наборе из 5 вопросов — ответы должны быть адекватными.

Ожидаемый результат этапа Итоговый конфигурационный файл и финальное измерение TTFT с улучшением ≥60%.

Этап 5: Документирование и отчёт (30 минут)

Действия

1. Подготовить отчёт (Markdown или Colab Notebook). Включить:

   • Baseline TTFT → значение.
   • Таблица TTFT для разных chunk_size.
   • График зависимости TTFT от chunk_size.
   • Выбранный chunk_size и обоснование.
   • Скриншоты nvidia-smi при разных настройках.
   • Выводы о применимости для long context (потенциальный компромисс по throughput).

2. Написать короткую памятку "Как включить chunked prefill в production" (4-5 шагов).

Ожидаемый результат этапа Файл report.md и vllm_chunked_config.yaml.

5. Критерии приемки (Definition of Done)

• Baseline TTFT для промпта 32k измерен и записан.
• Найдена конфигурация chunked prefill, при которой TTFT снижается ≥60% относительно baseline.
• Все измерения проведены как минимум для 3 различных размеров чанка (2048, 4096, 8192).
• Качество ответов не ухудшилось (субъективная оценка по 5 вопросам из тестового набора).
• При оптимальной конфигурации не возникает Out-Of-Memory (память GPU <90% utilisation).
• Создан отчёт с таблицами, графиками и итоговой конфигурацией YAML.
• Код замеров (Python-скрипты) приложен к репозиторию.

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

Основные артефакты

1. vllm_chunked_config.yaml — конфигурация для запуска vLLM с оптимальным chunked prefill.
2. report.md — подробный отчёт с:
   • Значением baseline TTFT.
   • Таблицей с TTFT для разных chunk_size.
   • Графиком (PNG или встроенная mermaid).
   • Выбранным chunk_size и обоснованием.
   • Описанием влияния на память и throughput.
3. benchmark.py — скрипт для воспроизводимых замеров TTFT.
4. test_quality.py — скрипт с 5 тестовыми вопросами и проверкой логики ответов (опционально, если ответы сравниваются с эталоном).

Дополнительно (опционально): Настроить простой дашборд (например, Streamlit) для визуализации метрик в реальном времени.

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание: Если выполняется впервые, заложите дополнительно 1 час на отладку (установка, проблемы с памятью).

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я запустил vLLM сервер без chunked prefill и измерил baseline TTFT для 32k.
• Я протестировал хотя бы три разных значения --max-num-batched-tokens (2048, 4096, 8192) и записал результаты.
• Я убедился, что при оптимальном чанке TTFT снизился ≥60% относительно baseline.
• Я проверил, что ответы модели не стали бессмысленными или неполными (провёл 5 тестовых вопросов).
• Я задокументировал все шаги в отчёте и приложил конфигурационные файлы.