Гибридные модели чередуют слои пространства состояний (SSM) с небольшим числом слоёв внимания. Так устроены NVIDIA Nemotron 3, Qwen3.5–Qwen3.8, Kimi-K3 и GLM-5.3-Flash, где в роли SSM-слоя выступают Mamba-2, Gated DeltaNet (GDN) или Kimi Delta Attention (KDA). Слои SSM задуманы как быстрые: вместо растущего с контекстом KV-кэша они держат состояние фиксированного размера. Но при сервисных размерах батча именно SSM-слои, а не внимание, становятся узким местом по задержке — вплоть до 100K токенов контекста. Кроме того, они мешают двум надёжным ускорителям декодирования: спекулятивному декодированию и квантизации.
ReplaySSM возвращает SSM-слоям скорость. Метод представляет состояние SSM как контрольную точку (checkpoint) плюс буфер недавних входов. Это разделяет вычисление выхода и обновление состояния и решает три задачи: трафик памяти состояний, откат при спекулятивном декодировании и накопление ошибки в квантованном состоянии. Работа выполнена совместно NVIDIA и Dao AI Lab Принстонского университета; статья называется «ReplaySSM: Cache SSM Inputs, Not State», код доступен на GitHub. Метод интегрирован в NVIDIA TensorRT-LLM, vLLM и FlashInfer, а также используется в продакшене GLM-5.3-Flash и Kimi-K3.
Почему SSM-слои становятся узким местом декодирования
Возьмём слой Mamba-2 из NVIDIA Nemotron 3 Ultra. Он декодирует один токен по рекуррентной формуле:
S_t = α_t S_{t-1} + v_t k_t^T, o_t = S_t q_t
Здесь S_t — состояние, α_t — зависящий от данных коэффициент затухания, а v_t, k_t и q_t — входы токена. На каждом шаге состояние S_{t-1} загружается из памяти с высокой пропускной способностью (HBM), к нему применяется затухание и обновление ранга 1, по q_t вычисляется выход, а S_t записывается обратно. GDN и KDA добавляют дельта-коррекцию, но их обновления записываются в той же форме внешнего произведения, так что всё сказанное ниже применимо и к ним.
При декодировании состояние выполняет две работы: читается для получения выхода текущего токена и обновляется, становясь состоянием следующего. Обычное декодирование делает и то и другое сразу: на каждом шаге обновляет состояние, читает из него выход и сохраняет его для следующего шага. Все три проблемы проистекают из этого совмещения.
Трафик памяти состояний. Шаг декодирования SSM загружает и сохраняет полное состояние, выполняя над ним мало арифметики, поэтому ядро ограничено пропускной способностью HBM, а не вычислениями. При сервисных размерах батча это делает SSM-слои крупнейшей частью шага декодирования — вплоть до 100K токенов на Nemotron-3-Super-NVFP4 при размере батча 256 на одном GPU NVIDIA Blackwell Ultra B300. Причин две. Во-первых, слоёв SSM в пять раз больше, чем слоёв внимания. Во-вторых, состояние одного слоя в FP32 занимает 4 MiB на последовательность, тогда как одна запись KV в FP8 — 0.5 KiB на токен, то есть состояние SSM-слоя равно KV-кэшу слоя внимания на 8K токенов.
Откат при спекулятивном декодировании. Драфтер предлагает T токенов, они проверяются за один проход по целевой модели, и происходит откат к последнему принятому токену. Внимание поддерживает откат естественно — сдвигом указателя KV-кэша. Но когда SSM уже свернул черновики в своё состояние, у обновления нет простого обратного преобразования, и откат становится дорогим. Движки вроде vLLM обходят это, сохраняя снимок состояния после каждого чернового токена и выбирая тот, что соответствует принятому префиксу. Это одно чтение и T записей на шаг по памяти, а снимки ограничивают число запросов, помещающихся в память GPU. В итоге спекулятивное декодирование, надёжно ускоряющее трансформеры, при сервисных батчах почти не помогает SSM.
Квантизация состояния. Квантованное состояние даёт выход и служит входом для следующего обновления, поэтому его ошибка округления возвращается в рекуррентность на каждом шаге. У KV-кэша такого цикла нет: квантованная запись читается много раз, но никогда не записывается обратно в себя, поэтому KV-кэши терпят меньшую точность, чем состояния SSM. Например, Qwen3.5-122B-A10B-NVFP4 использует 4-битный формат NVIDIA NVFP4 для весов, KV-кэш в FP8 и состояние SSM в FP32.
Как ReplaySSM представляет состояние
Рекуррентность SSM можно развернуть на окно из h токенов:
S_t = α_t S_{t-1} + v_t k_t^T = γ_h S_ckpt + Σ_{j=1..h} γ_{j+1:h} v_j k_j^T
Здесь S_ckpt — состояние до окна, а γ — накопленное затухание. Обычное декодирование идёт по левой части шаг за шагом. ReplaySSM использует правую: представляет состояние как контрольную точку плюс небольшой буфер записей недавних входов (α_j, v_j, k_j) для Mamba-2.
На большинстве шагов, которые авторы называют не-сбросовыми (non-flush), метод добавляет текущие входы в буфер и вычисляет выход из контрольной точки и буфера, не сохраняя состояние. Когда буфер из L записей заполняется, шаг сброса (flush) сворачивает буферизованные входы в контрольную точку, один раз записывает новую контрольную точку в HBM и очищает буфер. Выходы совпадают с обычным декодированием с точностью до различий в плавающей точке.
Суть подхода укладывается в три вопроса. Что делает ReplaySSM? Кэширует недавние входы SSM и представляет состояние как контрольную точку плюс эти входы. Что это меняет? Разделяет вычисление выхода и обновление состояния: выходы вычисляются каждый шаг, а контрольная точка обновляется раз за окно. Что это даёт? Снижает трафик памяти состояний, упрощает спекулятивный откат и поддерживает чтение состояния в низкой точности без распространения ошибки округления через рекуррентные обновления.
Как оценивали ReplaySSM
Авторы взяли четыре гибридные модели из двух семейств — Nemotron 3 со слоями Mamba-2 и Qwen3.5 со слоями GDN — от плотных моделей на 4B на одном GPU NVIDIA H100 Tensor Core до смешанных экспертных (MoE) моделей класса 120B на одном GPU NVIDIA Blackwell Ultra B300.
| Модель | Слой SSM | Веса | GPU | Слоёв SSM / внимания |
|---|---|---|---|---|
| Nemotron-3-Nano-4B | Mamba-2 | BF16 | 1×H100 | 21 / 4 |
| Nemotron-3-Super-120B-A12B | Mamba-2 | NVFP4 | 1×B300 | 40 / 8 |
| Qwen3.5-4B | GDN | BF16 | 1×H100 | 24 / 8 |
| Qwen3.5-122B-A10B | GDN | NVFP4 | 1×B300 | 36 / 12 |
ReplaySSM реализовали в vLLM и сравнили с собственными ядрами vLLM в их лучшей конфигурации запуска. В замерах производительности состояние SSM держали в FP32 — точности по умолчанию для этих моделей — и использовали окно L = 16, лучшее по результатам для всех четырёх моделей. В тестах спекулятивного декодирования черновые токены генерировались головой многотокенного предсказания (MTP) каждой модели.
Выгода 1: меньше трафика памяти состояний
На окне из L токенов обычное декодирование читает состояние L раз и L раз записывает. ReplaySSM по-прежнему читает контрольную точку на каждом шаге, потому что она нужна для выхода, но записывает её только один раз — при сбросе. У не-сбросового шага одна задача — вычисление выхода, и он записывает лишь новую запись входа: несколько векторов вместо полной матрицы. Это примерно вдвое сокращает доминирующий трафик состояний.
При сервисных размерах батча пропускная способность памяти ограничивает производительность обычного ядра, поэтому сокращение трафика ускоряет декодирование. При размере батча 1 доминируют другие источники задержки, и в этих тестах ReplaySSM ускорения не даёт. При размере батча 256 ядро работает в 1.46–1.80 раза быстрее, а шаг декодирования — в 1.20–1.47 раза быстрее на всех четырёх моделях. Сквозной выигрыш растёт с размером батча, потому что растёт доля SSM в шаге.
Выгода 2: проще откат и быстрее проверка
При спекулятивном декодировании ReplaySSM добавляет входы всех T черновых токенов в буфер входов до проверки, а после неё фиксирует принятый префикс сдвигом указателя буфера. Как и в KV-кэше, записи за указателем игнорируются, а следующее добавление их перезаписывает. Никаких снимков состояния на каждый черновой токен и никакого состояния для восстановления. Сброс сворачивает в контрольную точку только зафиксированные записи, поэтому никогда не ждёт принятой длины.
У этого разделения есть и второй эффект. Обычное декодирование обязано выдавать на токен две вещи — выход и обновлённое состояние, потому что следующий токен потребляет состояние. Не-сбросовый шаг ReplaySSM выдаёт только выход, а вычисление только выхода открывает другой способ расчёта. Для одного обновления из нулевого состояния:
o = (v k^T) q = v (k^T q)
Выражение слева строит состояние размера d_v × d_k и читает его. Выражение справа формирует скаляр и масштабирует вектор, не строя состояние. Для T черновиков против h зафиксированных записей та же идея даёт:
O = S_ckpt Q D + V (W ⊙ (K^T Q))
где D содержит накопленные затухания позиций черновиков, а W — причинно замаскированная матрица затуханий. Это называется вычислением только выхода (output-only evaluation). Контрольная точка используется один раз для всех T запросов, а всё остальное — два небольших матричных произведения по кэшированным ключам, значениям и запросам черновиков. Промежуточное состояние не формируется, что экономит память на кристалле и FLOPs, а эти вычисления могут ускоряться тензорными ядрами. Кроме того, K^T Q разделяется всеми головами значений в группе. Это то же блочное вычисление, которое уже используют ядра префилла, где в роли блока выступают буфер и черновики. GDN тоже может использовать эту форму. Вычисление только выхода применимо и к обычному декодированию, но важнее всего оно при проверке, где T умножает арифметику на шаг.
На промптах GSM8K с T = 4 на двух моделях класса 120B ReplaySSM при размере батча 512 даёт 1.85× и 2.10× пропускной способности спекулятивного декодирования vLLM при совпадающих длинах принятия. Выигрыш складывается из двух источников. Быстрее проверка: ядро vLLM хранит один снимок на токен, поэтому его задержка растёт с T и при T = 8 почти в пять раз превышает задержку обычного шага декодирования vLLM; ReplaySSM на большинстве шагов состояние не пишет, так что проверяет восемь токенов примерно за время одного стандартного шага, включая сброс. Выше параллелизм: снимки состояний черновиков занимают память GPU и ограничивают число одновременных запросов. Обычное спекулятивное декодирование перестаёт давать выигрыш после размера батча 128, а на Qwen3.5-122B-A10B-NVFP4 при размере батча 512 падает ниже стандартного декодирования, тогда как ReplaySSM продолжает масштабироваться с размером батча.
Выгода 3: меньше накопление ошибки за счёт раздельной точности
Состояние в более низкой точности требует меньше байт, что снижает задержку, когда производительность ядра ограничена пропускной способностью памяти. На Mamba-2 не-сбросовое ядро работает в 3.8 раза быстрее при чтении состояния в NVFP4, чем при чтении в FP32.
ReplaySSM меняет квантизацию двумя способами. Во-первых, контрольная точка записывается раз за окно, а не раз за шаг, поэтому одна низкоточная контрольная точка округляется один раз в L шагов. Такой подход с периодическими контрольными точками разбирается в отчёте Nemotron 3 Ultra: он снижает частоту округления, но ошибка округления всё равно попадает в последующие обновления. Во-вторых, не-сбросовые шаги читают состояние только для вычисления выходов, а сбросовые обновляют контрольную точку, используемую в последующих обновлениях состояния. Эти две задачи могут использовать разную точность. Авторы держат высокоточную контрольную точку для рекуррентности и низкоточную копию для не-сбросовых чтений. Сброс начинается с высокоточной контрольной точки, включает буферизованные входы и независимо приводит результат к обоим форматам. Низкоточная копия никогда не читается обратно в рекуррентность, поэтому её ошибка округления влияет только на выходы текущего окна и не распространяется по рекуррентности. Запись дополнительной низкоточной копии при сбросе сокращает объём данных, читаемых на каждом не-сбросовом шаге. Принятие, откат и управление буфером не меняются.
Ошибку выхода измеряли как расхождение Кульбака — Лейблера (KL) от декодирования с состоянием в FP32 при teacher forcing на тексте PG-19 вплоть до 128K сгенерированных токенов. Общая контрольная точка в FP8 или NVFP4 накапливает ошибку с длиной. Сохранение рекуррентности в FP32 убирает большую часть роста: при 128K токенов на Qwen3.5-4B чтение в NVFP4 с рекуррентностью в FP32 снижает KL с 0.136 до 0.0038. Остаётся ошибка чтения, которая с длиной не растёт. Рекуррентность в FP16 со стохастическим округлением (SR) близко следует за FP32, так что полуточной контрольной точки достаточно. Лучшая квантизация копии для чтения — например, отдельные масштабы по измерениям состояния и канала — могла бы ещё уменьшить остаточную ошибку чтения.
Авторы также проверили, проявляется ли разрыв в KL в точности на прикладных задачах. На GPQA-Diamond — бенчмарке с самыми длинными генерациями из трёх запущенных — общая контрольная точка в NVFP4 теряет от 3 до 31 процентного пункта точности на четырёх моделях. Чтение в NVFP4 с рекуррентностью в FP16 или FP8 возвращает большую часть потерь и удерживает три из четырёх моделей в пределах примерно 2 процентных пунктов от FP32.
| GPQA-Diamond, изменение от FP32 (пункты) | Точность FP32 | Общая контрольная точка NVFP4 | Чтение NVFP4, рекуррентность FP16+SR | Чтение NVFP4, рекуррентность FP8+SR |
|---|---|---|---|---|
| Nemotron-3-Nano-4B | 56.7 | −31.3 | −6.7 | −5.6 |
| Nemotron-3-Super-NVFP4 | 79.3 | −12.5 | −1.4 | −1.4 |
| Qwen3.5-4B | 76.3 | −3.2 | +0.7 | +1.5 |
| Qwen3.5-122B-A10B-NVFP4 | 82.7 | −5.2 | −2.0 | −1.7 |
На GPQA-Diamond общая контрольная точка на Nemotron-3-Super-NVFP4 повторяет точность FP32 на коротких генерациях, но расходится с ней по мере роста длины; конфигурация с раздельной точностью остаётся близкой к FP32 на всём протяжении. На MMLU-Pro и Math-L5 генерации обычно короче, чем на GPQA-Diamond: три из четырёх моделей теряют не более 3 процентных пунктов точности с общей контрольной точкой NVFP4, а конфигурации с раздельной точкой в среднем остаются в пределах половины процентного пункта от FP32. В статье есть результаты по всем трём бенчмаркам и тест RULER на длинном входе.
Как начать работать с ReplaySSM
ReplaySSM уже в продакшене. GLM-5.3-Flash обслуживается с ReplaySSM, а инфраструктурный отчёт Z.ai описывает его добавление в ядро декодирования KDA. Kimi-K3 кэширует входы SSM черновых токенов для проверки в своих слоях KDA и называет ReplaySSM параллельным подходом. Метод также интегрирован в открытые фреймворки инференса: TensorRT-LLM, vLLM, SGLang и FlashInfer.
Материалы для старта: статья «ReplaySSM: Cache SSM Inputs, Not State», код на github.com/Johnny-Liou/ReplaySSM, PR 14203 и PR 16464 для TensorRT-LLM, PR 48018, PR 52506 и PR 52928 для vLLM, PR 3975 и PR 4081 для FlashInfer. Более подробный разбор обычного и спекулятивного декодирования есть в предыдущем посте Dao AI Lab.
Если вы разворачиваете инференс локально, приёмы экономии памяти состояний стоит сопоставить с тем, как устроен запуск моделей в Hermes Desktop, а перед загрузкой чужих весов — прогнать четыре проверки из практики Unsloth Studio. Для сервинга на нескольких GPU пригодится разбор TensorRT multi-device в Dynamo-Triton.
Частые вопросы
Что такое гибридная модель и зачем в ней SSM-слои?
Это модель, где слои пространства состояний (SSM) чередуются с небольшим числом слоёв внимания. SSM-слой использует состояние фиксированного размера вместо KV-кэша, растущего с длиной контекста, что и задумано для быстрого декодирования.
Почему SSM-слои вдруг становятся медленными?
При сервисных размерах батча шаг декодирования упирается в пропускную способность памяти, а не в вычисления. SSM-слой на каждом шаге загружает и сохраняет полное состояние, делая над ним мало арифметики, поэтому вплоть до 100K токенов именно эти слои, а не внимание, определяют задержку.
В чём главная идея ReplaySSM?
Состояние представляется как контрольная точка плюс буфер недавних входов. На большинстве шагов выход считается из контрольной точки и буфера без записи состояния, а раз за окно буфер сворачивается в контрольную точку. Это примерно вдвое сокращает трафик памяти состояний.
Какое ускорение даёт метод?
При размере батча 256 ядро работает в 1.46–1.80 раза быстрее, а шаг декодирования — в 1.20–1.47 раза быстрее на четырёх протестированных моделях. При размере батча 1 ускорения нет: там доминируют другие источники задержки.
Как ReplaySSM помогает спекулятивному декодированию?
Вместо снимка состояния на каждый черновой токен он добавляет входы черновиков в буфер и фиксирует принятый префикс сдвигом указателя. На моделях класса 120B при размере батча 512 это даёт 1.85× и 2.10× пропускной способности спекулятивного декодирования vLLM при совпадающих длинах принятия.
Можно ли держать состояние SSM в низкой точности?
Да, если разделить точности: высокоточная контрольная точка для рекуррентности и низкоточная копия для чтения. Низкоточная копия не читается обратно в рекуррентность, поэтому её ошибка округления не накапливается с длиной. На Mamba-2 чтение состояния в NVFP4 ускоряет не-сбросовое ядро в 3.8 раза по сравнению с FP32.
Какой размер окна L использовать?
В замерах авторов L = 16 оказался лучшим по результатам для всех четырёх моделей — и Nemotron 3, и Qwen3.5.
Что теряется в точности на прикладных задачах?
На GPQA-Diamond общая контрольная точка NVFP4 теряет от 3 до 31 процентного пункта. Раздельная точность возвращает большую часть потерь: три из четырёх моделей остаются в пределах примерно 2 процентных пунктов от FP32.
Где уже используется ReplaySSM?
В продакшене — в GLM-5.3-Flash и Kimi-K3 (в слоях KDA). Метод интегрирован в TensorRT-LLM, vLLM, SGLang и FlashInfer, код выложен на GitHub.
Подходит ли метод для моделей на GDN, а не только на Mamba-2?
Да. GDN и KDA добавляют дельта-коррекцию, но их обновления записываются в той же форме внешнего произведения, что и у Mamba-2, поэтому подход применим и к ним; вычисление только выхода GDN тоже поддерживает.
Источник: developer.nvidia.com