NVIDIA объяснила, почему длинный контекст превращает attention в главное узкое место LLM

При росте контекста с 4K до 128K токенов доля attention во времени prefill у DeepSeek-R1 увеличивается с 18% до 85%. Поэтому ускорять нужно уже не один CUDA-кернел — архитектуру модели стоит проектировать с учётом того, как её выполняет.

NVIDIA разделяет инференс на две разные задачи:

- prefill обрабатывает весь промпт параллельно, зависит от вычислительной мощности и растёт примерно как O(n²);

- decode генерирует по одному токену, постоянно читает KV-кэш и чаще упирается в пропускную способность памяти.

Из анализа получились четыре практических правила:

1. Увеличивать число query-heads на один KV-head.

Высокий group size почти не влияет на prefill, но заметно ускоряет decode за счёт сокращения чтения KV-кэша.

2. Использовать head dimension 128 или 256.

Такие размеры лучше совпадают с тайлами GPU и 128-байтовыми операциями памяти. Размер 64 может фактически платить цену за тайл шириной 128.

3. Сокращать эффективный KV-кэш.

Подойдут компрессия, sliding-window и sparse attention, а также гибридные архитектуры, где глобальный KV-кэш хранится только в части слоёв.

4. Не задавать tensor parallelism выше числа KV-heads.

Иначе KV-кэш начинает дублироваться между GPU, увеличивая расход памяти и трафик без полезного ускорения.

Смысл co-design простой: модель с более удачной структурой attention может отвечать быстрее и обслуживать больше запросов на том же железе.

Для long-context и агентных систем архитектура attention становится частью инфраструктурной оптимизации - наравне с TensorRT-LLM, кэшированием и распределением нагрузки.

Источник:

https://developer.nvidia.com/blog/co-designing-ai-model-attention-for-fast-interactive-long-context-inference/