Проблема узкого места в гибридных моделях
Гибридные архитектуры, сочетающие слои внимания (Attention) и пространственные модели состояний (SSM, например, Mamba-2, Gated DeltaNet), становятся стандартом для современных LLM, таких как NVIDIA Nemotron 3, Qwen3.5, Kimi-K3 и GLM-5.3-Flash. Однако при длине контекста до 100K токенов именно SSM-слои становятся узким местом при декодировании (генерации токенов по одному).
Причина — в высокой пропускной способности памяти (HBM). Каждый шаг декодирования требует загрузки и записи полного состояния модели (до 4 MiB на последовательность в FP32), что ограничивает скорость вычислений. В отличие от KV-кэша внимания, где данные читаются многократно, состояние SSM обновляется на каждом шаге, создавая «бутылочное горлышко» пропускной способности.
Решение ReplaySSM: кэш входов, а не состояния
Технология ReplaySSM (совместная работа NVIDIA и Dao AI Lab при Принстоне) меняет парадигму хранения состояния. Вместо того чтобы обновлять и записывать полное состояние $S_t$ на каждом шаге, метод разделяет процесс:
- Чекпоинт (Checkpoint): Сохраняется в HBM редко (раз в окно).
- Буфер входов: Недавние входные данные ($v_t, k_t$) накапливаются в быстром буфере.
- Вычисление: Выход рассчитывается из чекпоинта и буфера без записи нового состояния в память.
Это позволяет снизить объем трафика памяти с состоянием примерно на 50%.
Ключевые метрики производительности
Тестирование проводилось на GPU NVIDIA Blackwell Ultra B300 (batch size 256) и H100. Результаты демонстрируют значительный прирост скорости по сравнению с базовыми реализациями в vLLM:
| Метрика / Сценарий | Ускорение (Speedup) | Детали |
|---|---|---|
| Декодирование SSM-слоев | 1.46x – 1.80x | На GPU Blackwell Ultra B300, batch size 256. Устранение узкого места HBM. |
| Спекулятивное декодирование | 1.85x – 2.10x | На моделях класса 120B. Упрощение отката (rollback) за счет буферизации черновиков. |
| Точность (Quantization) | ~2 п.п. отклонения | Использование low-precision для чтения и high-precision для чекпоинта сохраняет точность на уровне FP32 (GPQA-Diamond). |
Почему это важно для индустрии
ReplaySSM решает три критические проблемы SSM-архитектур:
- Скорость: Делает SSM-слои быстрее слоев внимания, что критично для длинных контекстов.
- Спекулятивное декодирование: Ранее откат при ошибке в спекулятивном декодировании требовал создания снимков состояния (snapshots), что было дорого. Теперь достаточно добавить черновики в буфер, что ускоряет процесс почти в 2 раза.
- Квантование: Метод позволяет использовать низкую точность (low-precision) для чтения состояния, минимизируя накопление ошибок округления, так как ошибка не возвращается в рекуррентное обновление.
Реализация ReplaySSM уже интегрирована в TensorRT-LLM, vLLM и FlashInfer и используется в продакшене в моделях GLM-5.3-Flash и Kimi-K3. Исходный код доступен на GitHub.
Источник: NVIDIA dev blog ↗