ReplaySSM: Революция в декодировании гибридных моделей

В мире генеративного искусственного интеллекта скорость и эффективность — это не просто метрики, это фундаментальные ограничения, определяющие возможность практического применения больших языковых моделей (LLM). Особенно остро эта проблема стоит при работе с гибридными архитектурами, которые сочетают в себе классические слои внимания (attention) и новые слои пространственных состояний (State-Space Models, SSM). Несмотря на то, что SSM-слои обещают линейную сложность и высокую скорость декодирования, на практике они часто становятся «узким горлышком» из-за особенностей управления памятью. Именно эту проблему решает новая технология ReplaySSM, разработанная совместно NVIDIA и лабораторией Dao AI при Принстонском университете. В этой статье мы подробно разберем, как работает эта архитектура, почему она превосходит традиционные подходы и что это значит для разработчиков AI-решений в России и мире.

Гибридные модели, такие как NVIDIA Nemotron 3, Qwen3.5, Kimi-K3 и GLM-5.3-Flash, используют SSM-слои (например, Mamba-2, Gated DeltaNet или Kimi Delta Attention) для обработки длинных контекстов. Однако стандартные методы декодирования заставляют эти слои тратить огромное количество времени на чтение и запись состояния в высокоскоростную память (HBM). ReplaySSM предлагает радикально новый взгляд на хранение состояния: вместо того чтобы обновлять и сохранять полное состояние на каждом шаге, система кэширует недавние входные данные. Это простое, но гениальное изменение позволяет ускорить декодирование до 1.8 раза, упростить спекулятивное декодирование и снизить накопление ошибок при квантовании.

01Почему SSM-слои становятся узким местом декодирования?

Чтобы понять инновационность ReplaySSM, нужно сначала разобраться в механике работы стандартных SSM-слоев. Возьмем для примера слой Mamba-2, используемый в модели NVIDIA Nemotron 3. Процесс декодирования одного токена описывается рекуррентным уравнением:

S_t = \alpha_t S_{t-1} + v_t k_t^\top, \qquad o_t = S_t q_t

Здесь S_t — это состояние модели на шаге t, \alpha_t — зависящая от данных величина затухания, а v_t, k_t и q_t — входные векторы токена. На каждом шаге декодирования система должна загрузить предыдущее состояние S_{t-1} из высокоскоростной памяти (HBM), применить к нему операцию затухания и обновить его, вычислить выходной токен o_t и записать новое состояние S_t обратно в память.

Проблема заключается в том, что состояние выполняет две функции одновременно: оно используется для вычисления текущего выхода и обновляется для подготовки к следующему шагу. В традиционном подходе эти операции слиты. На каждом шаге происходит чтение полного состояния из памяти, выполнение относительно простых арифметических операций и запись результата обратно. Это создает ситуацию, когда производительность ограничена не вычислительной мощностью GPU, а пропускной способностью памяти (memory bandwidth bound).

При размерах батча, типичных для обслуживания (serving batch sizes), SSM-слои, а не слои внимания, становятся основным источником задержек при работе с контекстом до 100 тысяч токенов. Два фактора усугубляют эту проблему:

1. Количество слоев: В гибридных моделях количество SSM-слоев может в пять раз превышать количество слоев внимания.
2. Размер состояния: Состояние одного слоя в формате FP32 занимает 4 МБ на последовательность. Для сравнения, одна запись KV-кэша в формате FP8 занимает всего 0.5 КБ. Таким образом, состояние одного SSM-слоя по объему сопоставимо с KV-кэшем слоя внимания для 8 тысяч токенов.

02Три главные проблемы традиционного декодирования SSM

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

1. Избыточный трафик памяти

Поскольку на каждом шаге декодирования состояние загружается и сохраняется полностью, а арифметических операций над ним выполняется мало, ядро (kernel) ограничено пропускной способностью HBM. При размере батча 256 на GPU NVIDIA Blackwell Ultra B300 трафик памяти для SSM-слоев становится доминирующим компонентом времени декодирования.

2. Сложность отката при спекулятивном декодировании

Спекулятивное декодирование — это техника ускорения, при которой «черновик» (drafter) предлагает несколько токенов, а целевая модель проверяет их за один проход. Если предсказание неверно, система должна откатиться к последнему принятому токену. В трансформерах это легко: достаточно просто сдвинуть указатель в KV-кэше. Однако в SSM, если черновые токены уже суммированы в состоянии, обратная операция не имеет простого математического инверса. Системы вроде vLLM обходят это, сохраняя снимок (snapshot) состояния после каждого чернового токена. Это требует одного чтения и T записей на шаг по пути, ограниченному памятью, и сильно ограничивает количество запросов, которые могут поместиться в памяти GPU. В результате спекулятивное декодирование, которое дает огромный прирост скорости для трансформеров, практически не помогает для SSM при больших батчах.

3. Накопление ошибок при квантовании

В квантованных моделях состояние используется как для вычисления выхода, так и как вход для следующего обновления. Ошибка округления, возникающая при квантовании, возвращается в рекуррентное обновление на каждом шаге. В отличие от этого, KV-кэш не имеет такой обратной связи: квантованная запись читается много раз, но никогда не записывается обратно в себя. Поэтому KV-кэши могут использовать более низкую точность (например, FP8), тогда как SSM-состояние часто приходится держать в FP32, чтобы избежать накопления ошибок. Например, модель Qwen3.5-122B использует NVFP4 для весов, FP8 для KV-кэша, но FP32 для SSM-состояния.

03Как работает ReplaySSM: архитектура кэширования входов

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

S_t = \gamma_h S_{ckpt} + \sum_{j=1}^{h} \gamma_{j+1:h} v_j k_j^\top

Здесь S_{ckpt} — это состояние (чекпоинт) до начала окна, а \gamma — кумулятивное затухание. ReplaySSM представляет состояние как комбинацию этого чекпоинта и небольшого буфера недавних записей входов (\alpha_j, v_j, k_j).

Процесс работы делится на два типа шагов:

1. Не-флеш-шаги (Non-flush steps): На большинстве шагов система просто добавляет текущие входы в буфер и вычисляет выход, используя чекпоинт и буфер. Состояние в память не записывается. Это экономит трафик.
2. Флеш-шаги (Flush steps): Когда буфер заполняется, система «схлопывает» (folds) накопленные входы в чекпоинт, записывает новый чекпоинт в HBM один раз и очищает буфер.

Результаты вычислений совпадают с традиционным декодированием с точностью до плавающей запятой. Но преимущества этого подхода колоссальны. Давайте разберем их подробнее.

04Преимущество 1: Снижение трафика памяти в два раза

За окно из L токенов традиционное декодирование читает состояние L раз и записывает его L раз. ReplaySSM читает чекпоинт на каждом шаге (так как он нужен для выхода), но записывает его только один раз — во время флеш-шага. На не-флеш-шаге система записывает только новые записи входов (несколько векторов), а не полную матрицу состояния. Это примерно вдвое снижает доминирующий трафик памяти.

Поскольку при больших батчах производительность ограничена пропускной способностью памяти, снижение трафика напрямую ускоряет декодирование. При размере батча 256 на GPU NVIDIA Blackwell Ultra B300 ядро ReplaySSM работает в 1.46x – 1.80x быстрее, а шаг декодирования — в 1.20x – 1.47x быстрее для четырех протестированных моделей. При размере батча 1 прироста нет, так как другие источники задержек доминируют, но именно при обслуживании (serving) разница критична.

05Преимущество 2: Упрощение отката и ускорение спекулятивного декодирования

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

Это разделение открывает возможность для «вычислений только выхода» (output-only evaluation). В традиционном подходе нужно вычислить и выход, и новое состояние. В ReplaySSM на не-флеш-шаге вычисляется только выход. Для обновления из нулевого состояния это позволяет использовать оптимизацию: o = (v k^\top)q = v(k^\top q). Вместо построения большой матрицы состояния и ее чтения, система формирует скаляр и масштабирует вектор. Это экономит память и вычисления, позволяя Tensor Cores работать эффективнее.

Результаты тестов на моделях класса 120B (Nemotron-3-Super и Qwen3.5-122B) показывают, что ReplaySSM обеспечивает пропускную способность спекулятивного декодирования в 1.85x – 2.10x выше, чем стандартное vLLM, при длине черновика T=4. Этот прирост обусловлен двумя факторами:

1. Быстрее проверка: Ядро vLLM хранит снимок для каждого токена, поэтому его задержка растет с T. При T=8 задержка vLLM почти в пять раз превышает задержку стандартного шага. ReplaySSM проверяет восемь токенов примерно за время одного стандартного шага.
2. Выше конкурентность: Снимки состояния черновиков потребляют память GPU, ограничивая конкурентность. Традиционное спекулятивное декодирование перестает давать прирост уже при батче 128, а на Qwen3.5-122B падает ниже стандартного декодирования при батче 512. ReplaySSM продолжает масштабироваться с ростом батча.

06Преимущество 3: Ограничение накопления ошибок через раздельные точности

Использование состояния более низкой точности (например, NVFP4) снижает задержку, так как требуется меньше байт для чтения. На Mamba-2 не-флеш-ядро работает в 3.8 раза быстрее при чтении NVFP4 состояния по сравнению с FP32. Однако, как мы выяснили, квантованное состояние накапливает ошибки.

ReplaySSM решает эту проблему, используя разные точности для разных задач. Система хранит высокоточный чекпоинт (например, FP32) для рекуррентного обновления и низко-точную копию (например, NVFP4) для чтения на не-флеш-шагах. Во время флеш-шага система берет высокоточный чекпоинт, добавляет буферизованные входы и независимо преобразует результат в оба формата. Низкоточная копия никогда не читается обратно в рекуррентное обновление, поэтому ее ошибка округления влияет только на выходы текущего окна и не распространяется дальше.

Эксперименты на тексте PG-19 показали, что раздельные точности устраняют большинство накопленных ошибок. Для модели Qwen3.5-4B при длине 128K токенов чтение NVFP4 с рекуррентным обновлением FP32 сократило расхождение KL (Kullback-Leibler divergence) с 0.136 до 0.0038 по сравнению с традиционным FP32. Оставшаяся ошибка — это ошибка чтения, которая не растет с длиной контекста. Использование FP16 для рекуррентного чекпоинта с стохастическим округлением также показывает результаты, близкие к FP32.

07Что это значит на практике

Технология ReplaySSM — это не просто академическое исследование, а готовое решение, интегрированное в ключевые инструменты индустрии. Код ReplaySSM доступен на GitHub, а сама технология уже интегрирована в NVIDIA TensorRT-LLM, vLLM и FlashInfer. Более того, гибридные модели, использующие SSM-слои, такие как GLM-5.3-Flash и Kimi-K3, уже работают в продакшене с применением этих принципов.

Для разработчиков и компаний, работающих с AI в России, это открывает новые возможности. Учитывая ограничения на доступ к новейшим GPU и необходимость оптимизации ресурсов, использование гибридных моделей с ReplaySSM позволяет:

• Снижать затраты на inference: Ускорение декодирования в 1.5-2 раза означает, что один и тот же GPU может обслуживать в 1.5-2 раза больше запросов.
• Работать с длинными контекстами: SSM-слои эффективнее обрабатывают длинные последовательности, а ReplaySSM устраняет их главные недостатки.
• Использовать квантование без потери качества: Возможность использовать низко-точное состояние для чтения и высокоточное для обновления позволяет экономить память и пропускную способность, не жертвуя точностью модели.

ReplaySSM демонстрирует, что переосмысление фундаментальных алгоритмов хранения данных может привести к прорывам в производительности, сопоставимым с улучшением аппаратного обеспечения. Для тех, кто хочет углубиться в тему, рекомендуется изучить полный текст статьи на arXiv и попробовать код на GitHub. Это шаг вперед в эволюции эффективного и быстрого искусственного интеллекта.

Источник: NVIDIA Developer ↗