Как построить экономные по памяти трансформеры с xFormers

Почему обычные батчи «переполняются» нулями?

Большинство инференс‑сервисов работают со «пакетами» запросов фиксированной длины. Чтобы поместить в один батч запросы разной длины, добавляют паддинг‑токены. Выглядит удобно, а на деле — каждый лишний ноль съедает драгоценную видеопамять и замедляет вычисления.

• Память: каждый токен в FP16 занимает 2 Б, а при длине 2048 токенов и 8‑м головах это уже сотни мегабайт.
• Скорость: операции над паддингом всё равно вычисляются, хотя их «значимость» нулевая.
• Стоимость: из‑за лишних вычислений растёт время ответа и потребление GPU‑часов.

Решение — избавиться от паддинга полностью, используя variable‑length packed batch. В этом подходе данные «упаковываются» в один тензор без пропусков, а маска указывает, где заканчивается каждый запрос.

Как упаковать батч без потерь памяти

Рассмотрим минимальный, но полностью рабочий пример на PyTorch + XFormers (или flash‑attention, если ваш бекенд поддерживает).

print("\n" + "="*70 + "\n4. Variable-length packed batch — no padding waste\n" + "="*70)
seqlens = [37, 120, 8, 200]
total = sum(seqlens)
H, K = 8, 64
q = torch.randn(1, total, H, K, device=device, dtype=torch.float16)
k = torch.randn(1, total, H, K, device=device, dtype=torch.float16)
v = torch.randn(1, total, H, K, device=device, dtype=torch.float16)
bias = ab.BlockDiagonalMask.from_seqlens(seqlens)
out_packed = xops.memory_efficient_attention(q, k, v, attn_bias=bias)
print("packed shape :", tuple(out_packed.shape), "(all", total, "tokens, no pad)")
splits = bias.split(out_packed)
print("recovered segments :", [tuple(t.shape) for t in splits])

Ключевые шаги:

• seqlens – список реальных длин запросов.
• BlockDiagonalMask генерирует маску без паддинга, которую понимает memory_efficient_attention.
• После вычисления результат можно «разрезать» обратно по оригинальным сегментам.

Эти строки работают «из коробки» в последних версиях XFormers и flash-attention. При отсутствии поддержки система просто бросит исключение, которое легко отловить и переключиться на обычный паддинг‑батч.

Группированный запрос (Grouped‑Query Attention, GQA)

Модели Llama, Mistral и их клоны используют GQA, чтобы сократить размер KV‑кэша. Идея проста: несколько «query‑голов» делятся на группы, а каждая группа делится одним набором KV‑голов.

B, M, K = 2, 256, 64
n_q_heads, n_kv_heads = 8, 2
G, Hq = n_kv_heads, n_q_heads // n_kv_heads
qg = torch.randn(B, M, G, Hq, K, device=device, dtype=torch.float16)
kg = torch.randn(B, M, G, 1,  K, device=device, dtype=torch.float16)
vg = torch.randn(B, M, G, 1,  K, device=device, dtype=torch.float16)
out_gqa = xops.memory_efficient_attention(qg, kg, vg)
print("GQA output shape :", tuple(out_gqa.shape), "= [B, M, G, Hq, K]")

Плюсы GQA:

• Меньший KV‑кэш: вместо 8 KV‑голов у нас только 2, экономим до 75 % памяти.
• Быстрее при длинных контекстах: кэш‑операции линейны от количества KV‑голов.
• Совместимость: большинство современных LLM‑инференс‑движков (vLLM, Text Generation Inference) уже поддерживают эту схему.

Практические рекомендации

• Обновляйте библиотеки. Поддержка packed‑batch и GQA появилась только в версиях XFormers 0.0.20+ и flash‑attention 2.0+. Проверьте pip list и при необходимости сделайте pip install -U xformers flash-attn.
• Тестируйте на небольших примерах. Перед переходом на продакшн запустите сравнение: паддинг‑батч vs. packed‑batch vs. GQA. Измерьте torch.cuda.max_memory_allocated() и latency.
• Обрабатывайте исключения. Не все GPU‑драйверы и версии CUDA поддерживают нужные ядра. Окружите вызовы try/except, как в примерах выше, и падайте обратно к «традиционному» режиму.
• Следите за типами данных. Для максимальной пропускной способности используйте torch.float16 или torch.bfloat16 (если ваш бекенд это позволяет).
• Документируйте seqlens. Храните массив реальных длин в метаданных запроса, чтобы в дальнейшем можно быстро собрать/разобрать packed‑батч.

Куда смотреть дальше?

• Официальная реализация memory_efficient_attention в XFormers.
• Репозиторий flash-attention — быстрые ядра для GQA и packed‑batch.
• Статья от разработчиков vLLM о zero‑padding inference: github.com/vllm-project/vllm.