Az NVIDIA ReplaySSM módszere gyorsíthatja a hibrid modellek dekódolását

A hibrid modellekben az SSM is szűk keresztmetszet lehet

A hibrid modellek állapotterű modellrétegeket, vagyis SSM-eket, néhány figyelmi réteggel felváltva használnak. A forrás példaként említi az NVIDIA Nemotron 3, a Qwen3.5 és Qwen3.8, a Kimi-K3, valamint a GLM-5.3-Flash modelleket. Ezekben többek között Mamba-2, Gated DeltaNet vagy Kimi Delta Attention működik SSM-rétegként.

Az SSM-rétegek fix méretű állapotot használnak, szemben a kontextussal együtt növekvő kulcsérték-gyorsítótárral, vagyis KV cache-sel. Ez elvileg gyors dekódolást tesz lehetővé. Az NVIDIA szerint kiszolgálási kötegméretek mellett azonban éppen az SSM-rétegek válhatnak a dekódolás késleltetésének fő forrásává, akár 100 ezer tokenes kontextusig.

A hagyományos megoldás minden lépésben beolvassa a teljes állapotot, frissíti, majd visszaírja a nagy sávszélességű memóriába. A forrás szerint ez főként memória-sávszélességhez kötött művelet, kevés számítással. A Nemotron-3-Super-NVFP4 esetében, 256-os kötegméretnél, egy NVIDIA Blackwell Ultra B300 GPU-n az SSM-rétegek nagyobb részt adhatnak a dekódolási lépésből, mint a figyelmi rétegek.

A ReplaySSM ellenőrzési pontot és bemeneti puffert használ

A ReplaySSM az SSM állapotát egy ellenőrzési pontként tárolt állapot és a legutóbbi bemenetek kis puffere formájában reprezentálja. A módszer így szétválasztja a kimenet kiszámítását az állapot frissítésétől.

A legtöbb lépésben az aktuális bemenet bekerül a pufferbe, a modell pedig az ellenőrzési pont és a puffer alapján számítja ki a kimenetet. Amikor a puffer megtelik, a rendszer összevonja a benne lévő bemeneteket az ellenőrzési ponttal, egyszer írja vissza az új állapotot a memóriába, majd kiüríti a puffert. A forrás szerint a kimenetek a lebegőpontos eltéréseket leszámítva megegyeznek a hagyományos dekódolás eredményeivel.

Az NVIDIA és a Dao AI Lab korábbi ReplaySSM-munkájára építő megoldás három problémát céloz. Csökkenti az állapot memóriaforgalmát, egyszerűsíti a spekulatív dekódolás visszagörgetését, és mérsékli a kis pontosságú állapot tárolásából származó hibafelhalmozódást.

Gyorsulás a Blackwell Ultra B300 GPU-n

Az NVIDIA által közölt összefoglaló szerint a ReplaySSM nagyjából felére csökkenti az SSM állapotának memóriaforgalmát. A módszerrel az SSM-dekódolási kernelek 1,46-szoros és 1,80-szoros közötti gyorsulást értek el 256-os kötegméretnél, NVIDIA Blackwell Ultra B300 GPU-kon.

A spekulatív dekódolásnál a rendszernek ellenőriznie kell a tervezett tokeneket, majd el kell vetnie azokat, amelyek nem illeszkednek az elfogadott előtaghoz. A figyelmi rétegek ezt a KV cache mutatójának mozgatásával kezelik, az SSM-eknél viszont a már összegzett állapotnak nincs egyszerű inverze. A ReplaySSM a vázlatként generált bemeneteket a pufferhez fűzi, ezért nincs szükség minden token után külön állapotmentésre.

Az NVIDIA szerint ez 120 milliárd paraméterosztályú modelleknél 1,85-szoros és 2,10-szeres közötti spekulatív áteresztőképességet eredményezett. A módszert négy hibrid modellen vizsgálták, köztük Mamba-2- és Gated DeltaNet-rétegeket használó Nemotron 3 és Qwen3.5 változatokon.

Elérhető a kutatás és több fejlesztői eszközben is

Az NVIDIA és a Princeton Dao AI Lab közös tanulmánya, a ReplaySSM: Cache SSM Inputs, Not State, az arXivon érhető el, a megvalósítás kódja pedig GitHubon található. A ReplaySSM beépült a TensorRT-LLM, a vLLM és a FlashInfer rendszerekbe. A forrás szerint a technológia éles környezetben is fut a GLM-5.3-Flash és a Kimi-K3 modellekben.

A kis pontosságú állapot kezelésénél a módszer nagy pontosságú ellenőrzési pontot használ a rekurencia frissítéséhez, miközben a beolvasásokhoz alacsonyabb pontosságú másolat alkalmazható. Az NVIDIA értékelése szerint a négy vizsgált modellből három nagyjából 2 százalékponton belül maradt az FP32 pontosságához képest a GPQA-Diamond teszten. A fejlesztők így olyan dekódolási megoldást kapnak, amely a vállalat szerint a hibrid modellek SSM-rétegeit gyorsabbá, a spekulatív dekódolást pedig kezelhetőbbé teszi.