NVIDIA və Dao AI Lab qibrid modellərin dekodlaşdırmağını sürətləndirmək üçün ReplaySSM yaratdı

ReplaySSM necə işləyir

Qibrid modellər vəziyyət fazası modelinin (SSM) təbəqələrini bir neçə attention təbəqəsi ilə birləşdirir və NVIDIA Nemotron 3, Qwen3.5, Kimi-K3 və GLM-5.3-Flash tamamilə bu dizaynı izləyir. KV cache əvəzinə SSM təbəqələri ölçüsü sabit olan vəziyyət istifadə edir, lakin xidmət göstərmə partiya ölçülərində dərhal SSM təbəqələri və attention deyil, 100K token kontekstə qədər əsas gecikmə barriyeri olmaq olar. Adi dekodlaşdırma hər addımda tam vəziyyəti oxur və yazır.

ReplaySSM vəziyyəti nəzarət nöqtəsi və son giriş qeydlərinin buferi kimi təmsil edir. Bir çox addımda, cəmi non-flush addımda, o, cari girişi buferə əlavə edir və nəticəni nöqtədən və buferdən vəziyyəti saxlamadan hesablayır. 16 qeyd dolduqda, flush addımı buferdəki girişləri nöqtədə birləşdirir, yeni nöqtəni bir dəfə yazır və buferi təmizləyir. Çıxışlar adi dekodlaşdırmaya uyğundur, yalnız üzən nöqtə fərqləri istisna olmaqla.

Vəziyyət trafikinin azaldılması və rollback-in sadələşdirilməsi

16 token pəncərəsi üçün adi dekodlaşdırma vəziyyəti 16 dəfə oxur və 16 dəfə yazır. ReplaySSM nöqtəni pəncərə başına yalnız bir dəfə yazır, bu da dominat vəziyyət trafikini təxminən yarıya qədər azaldır. NVIDIA Blackwell Ultra B300 GPU-larında 256 ölçülü partiya ilə dekodlaşdırma nüvələri 1,46 dəfədən 1,80 dəfə qədər daha sürətli işləyir, dekodlaşdırma addımları isə dörd qiymətləndirilmiş model üzrə 1,20 dəfədən 1,47 dəfə qədər daha sürətli. 1 ölçülü partiyada digər gecikmə mənbələri dominantdır və ReplaySSM bu testlərdə sürət artırımı təqdim etmir.

Spekulativ dekodlaşdırma üçün ReplaySSM qaralam token girişlərini yoxlamadan əvvəl buferə əlavə edir və alınan prefiksi bufer göstəricisinə qədər təsdiq edir, vəziyyət snapshot-larını saxlamağı və bərpasını istisna etməklə. Bu, yalnız çıxışa əsaslanan qiymətləşdirməni mümkün edir, ki, bu da ara vəziyyət qurulmadan iki kiçik matris çarpımı yaradır. GSM8K promptlarında 120B sinif modelləri ilə ReplaySSM 512 ölçülü partiyada vLLM spekulativ ötürücünü 1,85 dəfə və 2,10 dəfə artırır, qəbul uzunluqları müqaisə olunur.

Kvantlaşdırma xətasının məhdudlaşdırılması

Kvantlaşdırılmış SSM vəziyyəti rekurrent olur, buna görə yuvarıqlaşdırma xətası uzunluğa görə toplanır, KV cache qeydlərindən fərqli olaraq. Rekursiya üçün ReplaySSM yüksək dəqiqlikli nöqtə və oxuma üçün aşağı dəqiqlikli nüsxəsini saxlayır, buna görə oxuma xətası heç vaxt rekurrent dövrdə yayılmir. Qwen3.5-4B-də 128K token üzrə NVFP4 oxumaları FP32 rekursiyası ilə müqaisədə çıxış KL divergensiyasını 0,136-dan 0,0038-ə qədər azaldır. GPQA-Diamond-da dörd modeldən üçü NVFP4 oxumaları və FP16 və ya FP8 rekursiyası ilə FP32 dəqiqliyindən təxminən 2 faiz nöqtəsi daxilinde qalır, tam NVFP4 nöqtəsi isə 3-dən 31 nöqtəyə qədər itirir.

Mövcudluq

ReplaySSM NVIDIA və Princeton Universitetinin Dao AI Lab-inin birgə işidir, ki, "ReplaySSM: Cache SSM Inputs, Not State" əsərində ətraflı təsvir edilib və kod GitHub-da mövcuddur. O, TensorRT-LLM, vLLM, SGLang və FlashInfer-ə inteqrasiya edilib və istehsalda GLM-5.3-Flash və Kimi-K3-də işləyir, Kimi-K3 isə qaralam token SSM girişlərini yoxlama üçün KDA təbəqələrində keşləyir və ReplaySSM-ni rəqabətçi yanaşma kimi qeyd edir.

Mənbələr: Nvidia Dev