NVIDIA описала четыре подхода к ускорению инференса моделей с длинным контекстом

NVIDIA опубликовала рекомендации по совместному проектированию механизма внимания и схемы распараллеливания для ускорения инференса моделей с длинным контекстом. Компания предлагает отдельно учитывать узкие места обработки входного контекста и генерации токенов, а также согласовывать архитектуру модели с особенностями GPU.

NVIDIA представила технический разбор влияния параметров механизма внимания на производительность инференса. По измерениям компании для DeepSeek-R1, при увеличении входного контекста с 4 тыс. до 128 тыс. токенов доля внимания во времени обработки промпта возрастает с 18% до 85%.

Авторы разделяют инференс на обработку входного контекста, или prefill, и последовательную генерацию токенов, или decode. Первый этап преимущественно ограничен вычислительной производительностью и для плотного внимания масштабируется квадратично с длиной контекста. Второй зависит главным образом от пропускной способности памяти при чтении кеша ключей и значений (KV-кеша), а его стоимость растет линейно.

NVIDIA сформулировала четыре рекомендации:

• увеличивать число голов запросов, совместно использующих одну KV-голову, поскольку это сокращает объем данных, читаемых при генерации токенов;
• выбирать размерность головы внимания 128 или 256 для согласования с размерами вычислительных блоков и операциями памяти GPU;
• уменьшать эффективный объем KV-состояния с помощью сжатия кеша, разреженного внимания, скользящего окна или гибридной архитектуры;
• не задавать степень тензорного параллелизма выше числа KV-голов, чтобы избежать дублирования кеша между GPU. Для моделей с малым числом KV-голов предлагается использовать параллелизм по запросам или разделение KV-кеша между ускорителями.

Анализ ориентирован на GPU NVIDIA и сочетает расчеты форм матричных операций с измерениями ядер, использующих формат FP8 для вычисления внимания и хранения KV-кеша.