KV-cache и speculative decoding на собеседовании Data Scientist

Зачем разбирать на собесе

Production LLM — performance критичен. На собесе DS / MLE по LLM: «как работает KV-cache», «зачем speculative decoding», «FlashAttention».

Зачем нужен KV-cache

В Transformer decoder для каждого нового token нужны Q, K, V для attention.

Без cache. На каждом шаге пересчитываем K, V для всех previous токенов.

Token N: compute K, V for tokens [1..N]  (O(N) work)

На каждый next token — повторный work. O(N²) total.

С KV-cache. K, V для предыдущих токенов сохранены. На next token — только K, V для current.

Token N: cached K, V for [1..N-1] + new K, V for token N

O(N) total, на порядок быстрее на длинных sequences.

Размер KV-cache

Size = 2 (K + V) × num_layers × num_heads × head_dim × seq_length × precision

Для Llama-2 70B (80 layers, 8 heads, 128 head_dim, FP16):

• Per token: 2 × 80 × 8 × 128 × 2 = 327 KB.
• 4096 tokens: ~1.3 GB.
• 32k tokens: ~10 GB.

При длинном context — KV-cache становится bottleneck (больше, чем weights в memory).

Optimizations:

• GQA (Grouped Query Attention). Меньше K/V heads чем Q heads. Используется в Llama-2 70B.
• MQA (Multi-Query Attention). Один K/V для всех Q heads.
• Sliding window attention. Сохраняем только последние N токенов.
• PagedAttention (vLLM). Страничная организация KV cache, позволяет sharing.

Speculative decoding

Идея. Маленькая «draft model» предсказывает несколько next tokens forward. Большая модель проверяет batch и принимает / отклоняет.

Draft (small): generate "the cat sits on the mat"
Big (large): verify all 7 tokens parallel
- Accept: "the cat sits on" (4 tokens)
- Reject: " the mat" (need different)
- Big generates correct: " a"
- Continue with draft from "the cat sits on a"

Скорость. Большая модель делает forward на batch, что параллелится. Effective speedup 2-3×.

Tools: vLLM, TensorRT-LLM, SGLang.

FlashAttention

FlashAttention (Dao 2022). Re-implementation attention с tiling и без storing intermediate big matrix.

Стандарт. Compute attention(Q, K, V) =softmax(QK^T / sqrt(d)) V.

QK^T — N×N matrix. Огромная память для длинного контекста.

FlashAttention.

• Compute по тilesам, не all at once.
• Recompute backward (вместо storing).
• IO-aware (учитывает GPU SRAM vs HBM).

Результат. 2-4× быстрее, поддерживает длинный контекст без OOM.

FlashAttention v2/v3 — улучшения для специфичных GPU.

Batching и continuous batching

Static batching. Все requests batched вместе, ждут пока самый медленный закончится → bad utilization.

Continuous batching (vLLM). Batches dynamically. Когда request finished, его слот занят новым. Continuous flow.

Throughput. В 5-10× выше static batching на mixed workloads.

Частые ошибки

Не использовать KV-cache. Slow по умолчанию некоторые implementations.

KV-cache в FP32. Wastes memory. FP16 / BF16 / INT8 нормально для cache.

Speculative без verification. Получаешь quality drop. Big model должна проверять.

FlashAttention v1 на новых GPU. v2/v3 быстрее на H100. Update.

Single-request workload. Без batching — теряешь throughput. vLLM / TGI / TRT-LLM серверы — стандарт.

Связанные темы

• Transformer на собесе DS
• BERT vs GPT на собесе DS
• Attention механизм для DS
• Квантование LLM на собесе DS
• Подготовка к собесу Data Scientist

FAQ

KV-cache работает в encoder?

Encoder обрабатывает все токены параллельно — нет concept «previous tokens». KV-cache — только для autoregressive decoding.

Это официальная информация?

Нет. Статья основана на работах Dao 2022 (FlashAttention), Leviathan 2023 (speculative decoding), документации vLLM.

Тренируйте Data Science — откройте тренажёр с 1500+ вопросами для собесов.