NVIDIA ve Dao AI Lab, hibrit model çözümlemesini hızlandırmak için ReplaySSM'i sundu

ReplaySSM nasıl çalışır?

Hibrit modeller, durum uzayı modeli (SSM) katmanlarını birkaç attention katmanıyla değiştirerek birleştirir ve NVIDIA Nemotron 3, Qwen3.5, Kimi-K3 ile GLM-5.3-Flash tam da bu tasarımı izler. SSM katmanları KV cache yerine sabit boyutlu durum kullanır; ancak servis batch boyutlarında, 100K token bağlamına kadar gecikmenin ön engeli attention değil SSM katmanları olabilir. Normal çözümleme her adımda tam durumu okur ve yazar.

ReplaySSM durumunu bir denetim noktası ve son girdilerin kayıtları arabelleği olarak temsil eder. Çoğu adımda, yani sözde non-flush adımda, arabelleğe yeni girdileri ekler ve çıktıyı denetim noktasından ve arabellekten durum depolama olmadan hesaplar. 16 kayıtlık arabellek dolduğunda, flush adımı arabellekteki girdileri denetim noktasında birleştirir, yeni denetim noktasını bir kez yazar arabelleği temizler. Çıktılar, kayan nokta farkları dışında normal çözümlemede eşleşir.

Durum trafiğini azaltma ve rollback'ın basitleştirilmesi

16 token penceresinde normal çözümleme durumu 16 kez okur ve 16 kez yazar. ReplaySSM ise denetim noktasını pencerede yalnızca bir kez yazar; bu da baskın durum trafiğini yaklaşık yarıya indirir. NVIDIA Blackwell Ultra B300 GPU'larında batch boyutu 256'da çözümleme çekirdeği 1,46 kata göre 1,80 kata daha hızlı çalışır, çözümleme adımı ise dört değerlendirilmiş modelde 1,20 kata göre 1,47 kata daha hızlıdır. Batch boyutu 1'de farklı gecikme kaynakları baskındır ve ReplaySSM bu testlerde hız artışı sağlamaz.

Speculative çözümleme için ReplaySSM, draft token girdilerini doğrulamadan arabelleğe ekler ve elde edilen öneki arabellek göstergesinin önüne ekler; durum snapshot'larını depolamadan ve geri yüklemek zorunda kalmadan. Bu, yalnızca çıktıya dayalı değerlendirmeyi mümkün kılar; iki küçük matris çarpımıyla ara durum oluşturulmasını sağlar. GSM8K istemlerinde, 120B sınıf modellerle ReplaySSM, batch boyutu 512'de vLLM'in speculative üretkenliğini aynı ortalama uzunlukta 1,85 ve 2,10 kat aşar.

Kuantizasyon hatasının sınırlandırılması

Kuantize edilmiş SSM durumu tekrarlanır; bu yüzden yuvarlama hatası uzunlukla birikir ve KV cache kaydından farklıdır. ReplaySSM, tekrarlama için yüksek doğruluklu bir denetim noktası ve okumalar için düşük doğruluklu bir kopya saklar; bu nedenle okuma hatası tekrarlanan döngüde asla yayılmaz. Qwen3.5-4B'de 128K token'de NVFP4 okumaları, FP32 tekrarlı çıktı KL diverjansını 0,136'dan 0,0038'e düşürür. GPQA-Diamond'da dört modelden üçü, NVFP4 okumaları ve FP16 veya FP8 tekrarlı ile FP32 doğruluğunun yaklaşık 2 puanında kalır; oysa tam NVFP4 denetim noktası 3'ten 31 puana kadar kaybeder.

Erişilebilirlik

ReplaySSM, NVIDIA ve Princeton Üniversitesi Dao AI Lab'ın ortak ürünüdür; ayrıntıları «ReplaySSM: Cache SSM Inputs, Not State» başlıklı çalışmada açıklanmıştır, kod ise GitHub'dadır. TensorRT-LLM, vLLM, SGLang ve FlashInfer ile entegre edilmiştir ve üretimde GLM-5.3-Flash ile Kimi-K3'te çalışır; Kimi-K3 ise draft token girdilerini doğrulamak için KDA katmanlarında önbelleğe alır ve ReplaySSM'i rekabetçi bir yaklaşım olarak nitelendirir.

Kaynaklar: Nvidia Dev