NVIDIA та Dao AI Lab створили ReplaySSM для прискорення декодування гібридних моделей

Як працює ReplaySSM

Гібридні моделі поєднують шари моделі стану (SSM) з кількома шарами attention у чергування, і саме цьому дизайну слідують NVIDIA Nemotron 3, Qwen3.5, Kimi-K3 та GLM-5.3-Flash. Шари SSM замість KV cache використовують стан фіксованого розміру; однак при розмірах batch для обслуговування саме шари SSM, а не attention, можуть стати головним бар’єром затримки аж до 100K токенів контексту. Звичайне декодування на кожному кроці читає і зберігає весь стан.

ReplaySSM представляє стан як контрольну точку та буфер останніх записів вхідних даних. На більшості кроків, тобто non-flush кроках, він додає поточні входи в буфер і обчислює вихід із контрольної точки та буфера без збереження стану. Коли буфер на 16 записів заповнюється, flush крок об’єднує входи з буфера в контрольній точці, один раз записує нову контрольну точку і очищує буфер. Результати збігаються зі звичайним декодуванням, крім дрібних розбіжностей контрольних точок.

Скорочення трафіку стану та спрощення rollback

На вікні 16 токенів звичайне декодування читає і пише стан по 16 разів. ReplaySSM записує контрольну точку на вікні лише раз, що приблизно вдвічі зменшує домінуючий трафік стану. На NVIDIA Blackwell Ultra B300 GPU із batch 256 ядро декодування працює від 1,46 до 1,80 раза швидше, а крок декодування на чотирьох оцінених моделях — від 1,20 до 1,47 раза швидше. Для batch 1 домінують інші джерела затримки, і ReplaySSM не дає прискорення в цих тестах.

Для спекулятивного декодування ReplaySSM додає чернеткові токени вхідних даних у буфер до перевірки та швидко фіксує отриманий префікс перед індикатором буфера, без збереження чи відновлення знімків стану. Це дозволяє покладатися лише на оцінку вихідних даних, яка створює два малих матричні добутки без побудови проміжного стану. На промптах GSM8K із моделями класу 120B ReplaySSM дає 1,85 і 2,10 рази більше спекулятивної пропускної здатності vLLM при batch 512 і однаковій довжині прийняття.

Обмеження помилки квантування

Стан квантованого SSM повертається в рекурентність, тому помилка округлення накопичується зі зростанням довжини, на відміну від запису KV cache. ReplaySSM зберігає контрольну точку високої точності для рекурентності та копію низької точності для читання, тому помилка читання ніколи не поширюється в рекурентному циклі. На Qwen3.5-4B за 128K токенів читання NVFP4 із рекурентністю FP32 знижує KL дивергенцію виходу з 0,136 до 0,0038. На GPQA-Diamond три з чотирьох моделей залишаються приблизно в межах 2 процентних пунктів точності FP32 із читанням NVFP4 та рекурентністю FP16 або FP8, тоді як контрольна точка NVFP4 загалом втрачає від 3 до 31 пунктів.

Доступність

ReplaySSM — спільна розробка NVIDIA та Dao AI Lab Принстонського університету, детально описана в роботі ReplaySSM: Cache SSM Inputs, Not State; код є на GitHub. Вона інтегрована в TensorRT-LLM, vLLM, SGLang і FlashInfer, а в продуктових рішеннях працює в GLM-5.3-Flash та Kimi-K3; Kimi-K3 кешує входи токенів-чернеток SSM у шарах KDA для перевірки та називає ReplaySSM конкурентним підходом.

Джерела: Nvidia Dev