NVIDIA и Dao AI Lab at Princeton представили ReplaySSM, метод для более быстрого декодирования в гибридных моделях, сообщил NVIDIA Technical Blog. Декодирование здесь означает выдачу следующего токена, то есть фрагмента текста.

Главное

  • ReplaySSM хранит состояние модели как контрольную точку и буфер недавних входов.
  • Метод уже интегрирован в NVIDIA TensorRT-LLM, vLLM и FlashInfer.
  • При размере пакета 256 шаг декодирования ускорился в 1.20 до 1.47 раза на четырех моделях.

Что изменили в работе моделей

Гибридные модели чередуют слои SSM и несколько слоев внимания. Слой внимания использует KV-кэш, который растет вместе с длиной контекста.

Так устроены NVIDIA Nemotron 3, Qwen3.5 до Qwen3.8, Kimi-K3 и GLM-5.3-Flash. В них используются Mamba-2, Gated DeltaNet или Kimi Delta Attention как SSM-слои.

По данным компании, SSM-слои должны ускорять декодирование. Но при обслуживании больших пакетов запросов они могут стать узким местом до контекста в 100 тыс. токенов.

ReplaySSM меняет способ хранения состояния. Вместо постоянной перезаписи полного состояния он использует контрольную точку и небольшой буфер последних входов.

В обычном подходе каждый шаг читает состояние, считает выход и записывает новое состояние. ReplaySSM считает выход на каждом шаге, но обновляет контрольную точку один раз за окно.

Что показали тесты

Метод оценивали на четырех гибридных моделях из двух семейств. Это Nemotron 3 со слоями Mamba-2 и Qwen3.5 со слоями GDN.

В тест вошли Nemotron-3-Nano-4B, Nemotron-3-Super-120B-A12B, Qwen3.5-4B и Qwen3.5-122B-A10B. Использовались один NVIDIA H100 Tensor Core GPU и один NVIDIA Blackwell Ultra B300 GPU.

Для всех четырех моделей применялось окно L = 16. Это размер буфера входных записей, который оказался лучшим в этих тестах.

При размере пакета 1 ускорения в этих тестах не было. При размере пакета 256 ядро работало быстрее в 1.46 до 1.80 раза. Весь шаг декодирования ускорился в 1.20 до 1.47 раза.

Для Nemotron-3-Super-NVFP4 при размере пакета 256 на одном NVIDIA Blackwell Ultra B300 GPU SSM-слои были крупнейшей частью шага декодирования до 100 тыс. токенов контекста.

Где уже используется ReplaySSM

ReplaySSM интегрирован в NVIDIA TensorRT-LLM, vLLM и FlashInfer. Он также работает в production в GLM-5.3-Flash и Kimi-K3.

Работа доступна на arXiv, а код размещен на GitHub. В статье метод описан под названием «ReplaySSM: Cache SSM Inputs, Not State».

Отдельно авторы указали три задачи, которые решает подход. Это снижение обмена состоянием с памятью, более простой откат при проверке черновых токенов и работа с низкой точностью без накопления ошибки в состоянии.

Контекстом для работы стали ограничения SSM-слоев. В обычном декодировании состояние читается и записывается на каждом шаге. При проверке черновых токенов нужно сохранять снимки состояния, а округление в сжатом состоянии может повторно попадать в следующие шаги.