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-ჯერ კითხულობს და 16-ჯერ წერს. ReplaySSM-ი საკონტროლო წერტილს ფანჯარაზე მხოლოდ ერთხელ წერს, რაც დომინანტურ მდგომარეობის ტრაფიკს დაახლოებით ნახევრად ამცირებს. NVIDIA Blackwell Ultra B300 GPU-ებზე batch ზომით 256 დეკოდირების ბირთვი 1,46-ჯერიდან 1,80-ჯერ უფრო სწრაფად მუშაობს, ხოლო დეკოდირების ნაბიჯი ოთხ შეფასებულ მოდელზე 1,20-ჯერიდან 1,47-ჯერ უფრო სწრაფად. batch ზომით 1-ზე სხვა შეყოვნების წყაროები დომინირებს და ReplaySSM-ი ამ ტესტებში სიჩქარის მატებას არ იძლევა.

სპეკულაციური დეკოდირებისთვის ReplaySSM-ი დრაფტ ტოკენების შეყვანებს გადამოწმებამდე ბუფერს უმატებს და მიღებულ პრეფიქსს ბუფერის მაჩვენებლის წინ წაწევით ამტკიცებს, მდგომარეობის snapshot-ების შენახვისა და აღდგენის გარეშე. ეს შესაძლებელს ხდის მხოლოდ გამოსავალზე დაფუძნებულ შეფასებას, რომელიც ორ პატარა მატრიცულ ნამრავლს ქმნის შუალედური მდგომარეობის აგების გარეშე. GSM8K-ის პრომპტებზე 120B კლასის მოდელებით ReplaySSM-ი batch ზომით 512-ზე vLLM-ის სპეკულაციური გამტარუნარიანობის 1,85-ჯერ და 2,10-ჯერ მეტს იძლევა, მიღების თანაბარი სიგრძეებით.

კვანტიზაციის შეცდომის შეზღუდვა

კვანტიზებული SSM-ის მდგომარეობა რეკურენტულობაში ბრუნდება, ამიტომ დამრგვალების შეცდომა სიგრძესთან ერთად გროვდება, განსხვავებით KV cache-ის ჩანაწერისგან. ReplaySSM-ი რეკურენტულობისთვის მაღალი სიზუსტის საკონტროლო წერტილს და წაკითხვებისთვის დაბალი სიზუსტის ასლს ინახავს, ამიტომ წაკითხვის შეცდომა რეკურენტულ ციკლში არასდროს ვრცელდება. Qwen3.5-4B-ზე 128K ტოკენზე NVFP4 წაკითხვები FP32 რეკურენტულობით გამოსავლის KL დივერგენციას 0,136-დან 0,0038-მდე ამცირებს. GPQA-Diamond-ზე ოთხი მოდელიდან სამი FP32-ის სიზუსტის დაახლოებით 2 პროცენტული პუნქტის ფარგლებში რჩება 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