האצת Sonar באמצעות ספקולציה

פענוח ספקולטיבי מאיץ את מהירות היצירה של מודלי שפה גדולים (LLMs) על ידי שימוש במודל טיוטה מהיר וקטן המייצר מועמדי השלמה שעוברים אימות על ידי מודל היעד הגדול יותר. במסגרת זו, במקום שריצה של מודל היעד היקר תייצר טוקן יחיד, מונפקים טוקנים מרובים בצעד יחיד. כאן אנו מציגים את פרטי היישום של סוגים שונים של פענוח ספקולטיבי, המיושמים ב-Perplexity כדי לצמצם את השהיית בין-הטוקנים (inter-token latency) במודלי Sonar.

פענוח ספקולטיבי

פענוח ספקולטיבי מנצל את המבנה של שפות טבעיות ואת האופי האוטו-רגרסיבי של טרנספורמרים כדי להאיץ יצירת טוקנים. אף על פי שמודלים גדולים יותר, כגון Llama-70B, נושאים ידע רב יותר מאשר קטנים מהם, כגון Llama-1B, במטלות פשוטות מסוימות הביצועים שלהם דומים. חפיפה זו אכן מעידה כי רצפים מסוימים מיוצרים טוב יותר על ידי המודלים הזולים יותר, תוך השארת בעיות מורכבות למודלים הגדולים יותר. האתגר טמון בקביעת אילו השלמות טובות יותר והאם היצירה של המודל הקטן היא באותה איכות כמו זו של המודל הגדול.

למרבה המזל, LLMs הם טרנספורמרים אוטו-רגרסיביים: כאשר ניתנת להם רצף של טוקנים, הם מפיקים את התפלגות ההסתברות של הטוקן הבא. בנוסף, הלוגיטים שנגזרים מהתכונות הביניים המשויכות לטוקנים ברצף הקלט גם מצביעים על הסבירות לכך שהמודל יפיק בדיוק את הטוקנים האלה. תכונה זו מאפשרת ספקולציה: אם רצף של טוקנים נוצר על ידי מודל קטן יותר, החל מקידומת קלט, אפשר להעביר אותו דרך המודל הגדול יותר כדי לקבוע עד כמה הוא תואם למודל היעד. כל קידומת של המועמדים מקבלת ציון הסתברות, והארוכה ביותר שמעל לסף קבלה נבחרת. כבונוס, מודל היעד גם מספק טוקן עוקב בחינם: אם מודל טיוטה מייצר n טוקנים, ניתן לפלוט עד n + 1 בצעד אחד.

הגדלה

בזמן הסקה (inference), תהליך הדגימה הספקולטיבית מתחלק בקירוב ל-4 שלבים:

הכנה מראש (Prefill): יש להריץ הן את מודל היעד והן את מודל הטיוטה על רצף הקלט כדי לאכלס את ערכי מטמון ה-KV (KV cache). בעוד שחלק מהשיטות, כגון Medusa, משתמשות בשכבות צפופות פשוטות יותר לניבוי, בפוסט זה אנו מתמקדים בטיוטות מבוססות-טרנספורמר הזקוקות למטמוני ה-KV שלהן.

יצירת טיוטה (Draft generation): מודל הטיוטה מבצע איטרציות כדי לייצר מספר טוקנים קבועים. רצף הטיוטה יכול להיות ליניארי או שהמודל יכול לחקור מבנה דמוי עץ עד לעומק נתון (EAGLE, Medusa). כאן, אנו מתמקדים ברצפים ליניאריים.

קבלה (Acceptance): מודל היעד רץ על רצף הטיוטה ובונה לוגיטים התואמים לכל טוקן טיוטה. אורכו של הרצף הקביל הארוך ביותר מנוסח ונקבע.

יצירת יעד (Target generation): מאחר שמודל היעד יצר לוגיטים, במיקום שאינו תואם או בסוף הרצף, הלוגיטים תואמים לטוקן עוקב. ניתן לדגום לוגיטים אלו כדי לספק טוקן איתן ממודל היעד, ולחתום על הרצף.

קיימות שיטות שונות ליישום פענוח ספקולטיבי. בפוסט זה, נתרכז בשיטות שבהן השתמשנו כדי להאיץ את מודלי Sonar באמצעות מודל 1B פנימי, וכן במנגנוני הניבוי שאנו בונים כדי להאיץ מודלים בקנה מידה של DeepSeek.

יעד-טיוטה (Target-Draft)

ניתן להשיג פענוח ספקולטיבי על ידי צימוד LLM קטן קיים כמודל טיוטה למודל יעד כדי ליצור רצפים מועמדים. בייצור, האצנו את Sonar באמצעות מודל Llama-1B שעבר כוונון עדין על אותו מערך נתונים כמו מודל היעד. אף על פי שהגישה הזו לא דרשה אימון של טיוטה מאפס, המודל הקטן עדיין משתמש בקיבולת משמעותית של מטמון KV ומוסיף תקורה קלה של prefill, מה שמגדיל את TTFT.

במסגרת זו, המפענח מבצע ספקולציה רק על אבות-טיפוס (batches) המיועדים לפענוח בלבד, ומייצר טוקנים באמצעות דגימה סטנדרטית במהלך הכנה מראש או באבות-טיפוס מעורבים של הכנה מראש-פענוח. בשלב ההכנה מראש, לוגיטים היעד נדגמים מיד כדי להכין מראש גם את הטוקן שנוצר לאחרונה במטמון ה-KV של הטיוטה. הטיוטה טרם נדגמת, אך הלוגיטים שהיא מייצרת מועברים לשלב הפענוח.

הגדלה

בפענוח, מודל הטיוטה מקודם ודוגם את הטוקן העליון בכל שלב. לאחר הגעה לאורך הטיוטה הרצוי, הטוקנים מורצים דרך מודל היעד כדי לייצר את הלוגיטים שעליהם מתבסס הדוגם בזיהוי אורך הרצף המתקבל. הקבלה נקבעת על ידי השוואת התפלגויות ההסתברות המלאות מהטיוטה ומהיעד. מכיוון שמודל היעד תמיד פולט קבוצת לוגיטים אחת העוקבת אחר רצף הטיוטה שהתקבל, הוא נדגם כדי לייצר פלט נוסף. מכיוון שמודל הטיוטה טרם ראה את הטוקן שהתקבל, הוא מורץ מחדש כדי לאכלס את ערכי מטמון ה-KV התואמים לו כהכנה לצעד הפענוח הבא, תוך העברת הלוגיטים שוב.

EAGLE

EAGLE הוא מתווה פענוח ספקולטיבי החוקר מספר רצפי טיוטה שנוצרים באמצעות מעבר דמוי עץ על טוקני טיוטה סבירים. עץ קבוע (EAGLE) או בעל מבנה דינמי (EAGLE-2) נחקר באמצעות הפעלות עוקבות של טוקני הטיוטה, תוך התחשבות במועמדי Top-K בכל צומת במקום מעקב אחר הטוקן בעל הציון הגבוה ביותר ברצף ליניארי. הרצפים מקבלים לאחר מכן ציונים והרצף המתאים הארוך ביותר נבחר להמשך, תוך הוספת טוקן נוסף ממודל היעד.

הגדלה

כדי להשיג ניבוי מדויק יותר, מודל טיוטה מסוג EAGLE אינו מנבא על בסיס טוקנים בלבד, אלא משתמש גם במאפייני היעד (מצבים נסתרים של השכבה האחרונה - last layer hidden states) של מודל היעד. החיסרון של EAGLE הוא הצורך לאמן מודלי טיוטה קטנים ומותאמים אישית המדויקים די הצורך כדי לייצר מועמדים מתאימים במסגרת תקציב השהייה (latency budget) נמוך. בדרך כלל, מודל טיוטה הוא שכבת טרנספורמר יחידה זהה לשכבת מפענח של המודל המקורי, המצומדת הידוק למודל היעד באמצעות קישור להטבעות (embeddings) ולתחזיות ה-lm_head שלו. מכיוון שהדרישה לקיבולת מטמון KV פחותה, ל-EAGLE יש טביעת רגל זיכרון קטנה יותר.

כדי לאמת רצפים דמויי עץ במודל היעד, יש להשתמש במסכות קשב מותאמות אישית. למרבה הצער, שימוש במסכת קשב מותאמת אישית לרצף שלם מאט באופן משמעותי את הקשב עבור אורכי קלט ריאליסטיים (עד 50%), ומבטל חלק מההאצה הניתנת להשגה באמצעות ספקולציה. טרם פרסנו חקירת עץ מלאה לייצור מסיבה זו, ובמקום זאת אנו מתרכזים במקרה המיוחד של ניבוי טוקן יחיד באמצעות מתווה דמוי MTP המוצג בדוח הטכני של DeepSeek-V3.

MTP

מתווה זה דומה לפענוח טיוטה-יעד, למעט השימוש במצבים נסתרים (hidden states) לצד טוקנים לצורך ניבוי. יש לבצע מעט יותר עבודה הן בשלב ההכנה מראש והן בשלב הפענוח בהשוואה לספקולציית טיוטה-יעד רגילה. מודל הטיוטה משתמש גם בטוקנים וגם במצבים נסתרים: טוקן t_{i+1} נדגם מהלוגיטים L_i התואמים לטוקן t_i, אשר בתורם מופקים מהמצבים הנסתרים H_i. בהתאם לכך, יש להזיז את מאגרי טוקני הקלט צעד אחד שמאלה ביחס לווקטורי המצבים הנסתרים שנוצרו כפלט על ידי מודל היעד. האיור להלן מציין את ההתאמות המשמשות לאימון, וכן את ההזזה במהלך ההסקה.

הגדלה

זרימת הפענוח דומה מאוד לפענוח טיוטה-יעד, למעט העובדה שגם המצבים החבויים וגם הלוגיטים מועברים הלאה. המימוש שלנו חולק את כל ה-sampling ועיבוד הלוגיטים המשויכים, ומייחד רק את קריאות ה-forward של המודל. כאשר מנבאים כמה טוקנים, מודל הטיוטה משתמש במצבים החבויים של הטיוטה לצורך הניבוי, ובמקביל מאכלס גם רשומות מטמון KV על סמך התכונות שלו עצמו. בטווח הארוך, הדבר עלול לפגוע בדיוק. לאחר מכן, כאשר מריצים את מודל הטיוטה כדי לאכלס את רשומת מטמון ה-KV לצורך חיזוי היעד, אנו מריצים אותו על כל הרצף תוך שימוש במצבים החבויים המדויקים יותר של היעד כקלט. מאחר שמודלי הטיוטה האלה קטנים, העלות הנוספת של עיבוד הטוקנים הנוספים זניחה.

אימון ראשי MTP

כדי להפיק תועלת מ-MTP, בנינו את התשתיות הנדרשות לאימון ראשי MTP המצורפים למודלים המכווננים שלנו על גבי מערכי הנתונים של Perplexity, הפועלים על צומת יחיד עם 8 התקני H100. בתוך כשיום אחד, ביכולתנו לבנות ראשים למודלים הנעים מ-Llama-1B ועד Llama-70B ו-DeepSeek V2-Lite. עבור מודלים גדולים יותר, אנו מסתמכים על ראשי MTP שנבנו במהלך תהליך הכוונון העדין.

המטרה של אימון MTP היא להתאים את המצבים החבויים של הטיוטה ואת הלוגיטים שנאמדו מהמצבים החבויים של היעד ללוגיטים והמצבים החבויים של הטוקן הבא של היעד. מאחר שהסקה עבור מצבים חבויים יקרה, אנו מחשבים אותם מראש באמצעות המימוש שלנו של מודל היעד, המותאם להסקה, כדי להשתמש בהם במהלך האימון. עם זאת, כדי לאמת את מימוש ה-MTP בהסקה ולהבטיח שהבדלים מספריים עקב קוונטיזציה או אופטימיזציות אינם פוגעים בתוצאות, לצורך הערכת אובדן האימות והדיוק אנו עושים שימוש מלא מחדש במימוש ההסקה של מודל היעד ושל מודל הטיוטה כאחד.

בעת מעבר ממערך הנתונים ShareGPT ששימש במאמר המקורי אל דגימות גדולות יותר, הבחנו כי ארכיטקטורת ראש ה-MTP שתווארה ויישמה במאמר EAGLE כשלה באימון עבור מודלים בגודל 70B. בניגוד ל-ShareGPT שהכיל מספר רב של רצפים קצרים יותר, אנו מאמנים על מספר קטן במתן משמעותית של הנחיות ארוכות יותר. מכיוון שראשי ה-EAGLE המקוריים סטטו מעט במבנה שלהם מטרנספורמר טיפוסי, החזרנו שכבות נורמליזציית RMS (RMS Normalization) שהוסרו. גילינו שהדבר לא רק אפשר לאימון להתכנס, אלא גם שיפר את דיוק הראשים בכמה נקודות אחוז.

הגדלה

לא זו בלבד שנורמות השכבות מקילות על האימון, אלא שהחזרת הנורמות היא גם אינטואיטיבית מבחינה מתמטית. ראשי MTP עושים שימוש חוזר בהטבעות ובחזיות הלוגיטים של מודל היעד, שכן גודלם עשוי להיות משמעותי (כ-2 GB עבור Llama 70B). במהלך האימון אלה קפואים, והציפייה היא ששכבת ה-MTP תלמד להטביע ניבויים באותו מרחב וקטורי שלמדה שכבת התחזית של המודל המקורי במהלך האימון. על ידי הסרת הנורמות, מצופה מ-MLP יחיד ללמוד את אותה פונקציה כמו MLP שאחריו נורמה, דבר המקשה על ההתאמה בין המצבים הנסתרים של מודל הטיוטה למודל היעד.

הסקה עם פענוח ספקולטיבי

במנוע ההסקה, על מנת לייצר טוקנים עבור רצפי קלט, יש לקבצם תחילה באבות-טיפוס בגודל סביר, ולאחר מכן יש להקצות עמודים במטמון ה-KV עבור הטוקנים הבאים. טוקני הקלט ומידע עמודי ה-KV נארזים לאחר מכן לתוך מאגר המשודר לכל הדרגות המקבילות המריצות את לב המודל. לבסוף, מטא-הנתונים מועתקים לזכרון ה-GPU והמודל מופעל כדי לייצר את הלוגיטים שמהם נדגם הטוקן הבא.

בניגוד ליישומים מסוימים המצמדים באופן רופף שרת הסקה של טיוטה ויעד באמצעות מעטפת (wrapper) המתזמרת בקשות ביניהם, זוגות הטיוטה-יעד שלנו מצומדים בהדוק ועוברים יחד את שלבי היצירה בתיאום מושלם. תזמון אבות-טיפוס והקצאת עמודי KV משותפים בין המודלים עבור כל סוגי הפענוח הספקולטיבי: הדבר מאחד את הלוגיקה המגשרת בין מודל לבין שרת ההסקה הכולל, שכן כולם חושפים את אותו ממשק.

זמן ריצת ההסקה ב-Perplexity מבוסס סביב FlashInfer, הקובע את מטא-הנתונים שיש לבנות כדי להגדיר ולתזמן את גרעין הקשב (attention kernel). בהינתן רצפי קלט המהווים אב-טיפוס (batch), עבור הכנה מראש, פענוח או אימות, יש לבצע עבודה בצד ה-CPU כדי להקצות מאגרים ביניים ולאכלס מאגרי קבועים מסוימים המשמשים בקשב. עבודה זו מתווספת לעלות תזמון האבות-טיפוס והקצאת עמודי ה-KV, הגוררים גם הם השהיות שיש להסוות כדי למקסם את ניצול ה-GPU.

בעוד שמקבילים במלואם את העבודה בצד ה-CPU ובצד ה-GPU עבור הסקה ללא ספקולציה, מצאנו כי איזון ה-CPU-GPU בפענוח ספקולטיבי מורכב יותר. האתגר העיקרי נובע מכך שמספר הטוקנים שהתקבלו קובע את אורך הרצף עבור ריצה עוקבת, מה שיוצר נקודת סנכרון קשה-להימנעות בין ה-GPU ל-CPU. התנסנו במתווי תזמון שונים כדי להסוות בצורה הטובה ביותר את ההשהיה של עבודת ה-CPU.

תזמון יעד-טיוטה

למרות שהוא קטן ממודל היעד, כאשר מודל LLM שלם משמש כטיוטה, הוא עדיין מייצר השהיה ניכרת ב-GPU, ומספק מרחב תמרון להסתרת פעולות CPU יקרות. מכיוון שמודלים קטנים אינם מרוויחים מקבילות טנזור (tensor parallelism), יש אי-התאמה בין מספר הדרגות שאליהן מפוצל מודל יעד לעומת טיוטה. ביישום שלנו, מודל הטיוטה רץ אך ורק בדרגת המוביל של קבוצת TP.

הגדלה

כפי שצוין קודם, צעד פענוח מעביר לוגיטים לריצה הבאה. הדבר מאפשר לנו לחפוף בין הפעלה אחת של מודל הטיוטה לבין עבודת תזמון אבות-הטיפוס בצד ה-CPU. לאחר הרכבת אב-הטיפוס, קריאות חוזרות ונשנות לדוגם ולטיוטה מפיקות את טוקני הטיוטה. במקביל, אב-הטיפוס לאימות מורכב עבור מודל היעד ומסונכרן עם העובדים המקבילים. לוגיטי היעד עוברים אימות ודגימה כדי לקבוע את אורכי הרצף שהתקבלו. בשלב זה, נדרש סנכרון GPU-ל-CPU על מנת לקבוע את אורכי הרצף הבאים. מכיוון שמודל הטיוטה פועל רק בצומת המוביל, אב-הטיפוס שלו מוגדר באופן סדרתי והרצתו מופעלת כדי לאכלס את ערכי מטמון ה-KV שלו עם הטוקן הנוסף שמודל היעד ייצר. הלוגיטים שיופקו על ידי ריצת טיוטה זו בריצה הנוכחית ישמשו לדגימת טוקן הטיוטה הראשון בריצה העוקבת. והחשוב מכל, בזמן שהטיוטה רצה, ניתן לתזמן את אב-הטיפוס הבא.

תזמון MTP עבור טוקן יחיד

אף על פי שזמן הריצה עדיין אינו מספק חקירת עץ טיוטה בסגנון Eagle, יישמנו מקרה מיוחד של מתווה זה, תוך התחשבות ברצף ליניארי של טוקני טיוטה המופקים על ידי מודל בגודל של שכבת מפענח טרנספורמר יחידה. מתווה זה יכול לשמש לניבוי טיוטה באמצעות משקלי הקוד הפתוח של DeepSeek R1. מקרה המשנה של ניבוי טוקן יחיד מעניין, שכן שכבות MTP גדולות משיגות שיעורי קבלה גבוהים די הצורך כדי להצדיק את התקורה שלהן.

תזמון MTP מורכב מעט יותר, שכן מודל הטיוטה מהיר הרבה יותר, ומסתיר פחות השהיה בצד ה-CPU. בנוסף, הטיוטה מפוצלת לצד מודל היעד, מה שמחייב העברות זיכרון משותפות עבור מידע אבות-טיפוס. ריצה מתחילה בהעברת מידע אבות-טיפוס ודגימת הטוקן הראשון מלוגיטים מועברים (carry-over logits), בדומה למתווה הקודם. לאחר מכן, היעד רץ כדי לאמת טוקנים, ומעבד 2 * D טוקנים, כאשר D הוא גודל אב-טיפוס הפענוח. הדבר אידיאלי עבור יצירת אב-טיפוס מיקרו (micro-batching) במודלי מומחים משולבים (Mixture-of-Experts - MoE) על פני חיבורי רשת איטיים יותר כגון InfiniBand, שכן אב-הטיפוס מתחלק באופן שווה לשני חצאים. המצבים הנסתרים של היעד עוברים לריצת הטיוטה הבאה, בעוד שהלוגיטים מועברים לדוגם לצורך אימות.

הגדלה

על ידי ביצוע כמות מוגבלת של עבודה נוספת ב-GPU, אנו נמנעים מסנכרון GPU-ל-CPU לאחר קבלת רצף הטיוטה. לאחר שטוקני הקלט של היעד מוזזים, גרעין (kernel) משבץ את טוקני היעד הבאים במיקומים התואמים להם. הטיוטה מורצת מחדש לאחר מכן עם אותם נתוני אב-טיפוס כמו היעד, מאכלסת את ערכי מטמון ה-KV ובונה את הלוגיטים והמצבים הנסתרים לריצה הבאה, תוך ביצוע עבודה עודפת מסוימת על טוקנים שלא התקבלו. במצבים אלו, ההשהיה של העבודה הלא-מנוצלת ניתנת למדידה בקושי עקב גודלו הקטן של מודל הטיוטה. במקביל לריצת הטיוטה, אורכי הרצפים נקבעים ב-CPU ותזמון אב-הטיפוס הבא מופעל, מבלי להצטרך להמתין לסיום עבודת ה-GPU.

התקורה של עבודה נוספת בשכבת הטיוטה אינה ניכרת בקשב, אך שכבות MLP בעייתיות יותר. מכיוון שהוראות הכפלת מטריצות (matrix multiplication) מרופדות לגבול של 64 לאורך ממד מספר הטוקנים, אם הכפלה אינה דורשת משמעותית יותר בלוקים, התקורה מוסתרת. עבור רצפי טיוטה ארוכים יותר התקורה יקרה יותר והמתווה המשמש למודלי טיוטה-יעד רגילים עובד טוב יותר.

מקורות

הסקה מהירה מטרנספורמרים באמצעות פענוח ספקולטיבי

EAGLE: דגימה ספקולטיבית דורשת חשיבה מחדש על אי-ודאות תכונות

EAGLE-2: פענוח מהיר יותר של מודלי שפה עם עצי טיוטה דינמיים

EAGLE-3: הגדלת היקף האצת ההסקה של מודלי שפה גדולים באמצעות מבחן זמן-אימון

Medusa: מסגרת האצת הסקה פשוטה ל-LLM עם ראשי פענוח מרובים

נורמליזציית שכבות

FlashInfer: מנוע קשב יעיל וניתן להתאמה אישית להגשת הסקת LLM

הדוח הטכני של DeepSeek-V3

ציטוט

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}