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 כותב את נקודת הבקרה לחלון רק פעם אחת, מה שמפחית את תעבורת המצב הדומיננטית בכחצי בערך. על GPU-ים של NVIDIA Blackwell Ultra B300 עם גודל batch של 256 גרעין פענוח רץ מהיר יותר מפי 1,46 לפי 1,80, וצעד הפענוח על ארבעה מודלים מוערכים מהיר יותר מפי 1,20 לפי 1,47. בגודל batch של 1 מקורות עיכוב אחרים דומיננטיים ו-ReplaySSM אינו מספק האצת מהירות בבדיקות אלו.

לפענוח ספקולטיבי, ReplaySSM מוסיף קלט של טוקני דרפט למאגר עד לאימות ומעגן את הקידומת שהתקבלה במקום המאגר באמצעות דחיפה, ללא שמירה ושחזור של תמונות מצב. זה מאפשר הערכה המבוססת רק על פלט, אשר יוצרת מכפלה מטריצלית קטנה בלי לבנות מצב ביניים. על מבוני GSM8K עם מודלים בקטגוריית 120B, ReplaySSM עם גודל batch של 512 מספק 1,85 ו-2,10 פעמים יותר תפוקה ספקולטיבית מ-vLLM, באורכי קבלה שווים.

הגבלת טעות הכמותיות

מצב 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, אשר מנתב קלט SSM של טוקני דרפט לאימות בשכבות KDA ומציין את ReplaySSM כגישה מתחרה.

מקורות: Nvidia Dev