CuTeDSL ב-Perplexity

ב-Perplexity אנו משתמשים במנוע ההגשה הפנימי המותאם לזמן ריצ' (ROSE) כדי להריץ מודלים החל מטמיעה (embeddings) ועד למודלי שפה גדולים (LLMs) בעלי טריליוני פרמטרים. מנוע הריצה שלנו, היושב מאחורי כל ממשקי ה-API כולל Sonar, Search ו-Embeddings, נבנה כדי לאפשר לחוקרים לחקור את המודלים והארכיטקטורות המתאימים ביותר, תוך פריסה מהירה של המועמדים הטובים ביותר לייצור. כדי להשיג זאת, עלינו להיות מסוגלים ליישם, פרוס ולייעל הסקת מסקנות עבור ארכיטקטורות מתקדמות בקצב מהיר. בפוסט זה בבלוג, אנו בוחנים כיצד ליבת CuTeDSL המשתלבת במנוע ההסקות שלנו מאפשרת לנו לבנות את ליבות ה-GPU הנדרשות כדי להביא את המודלים לביצועי שיא על גבי GPUs של NVIDIA.

מודלים מותאמים אישית, הבנויים על ידי צוותים שונים ב-Perplexity, מתארחים בתוך החברה על גבי מעבדי GPU מסוג NVIDIA Hopper ו-Blackwell תוך שימוש במנוע ההסקות ROSE שנבנה על ידי צוות הסקת המסקנות של AI. מנוע ROSE פותח במקור כדי לשרת טוב יותר מודלי Llama מותאמים אישית עבור פענוח שפה וסיווג כאחד, עם ממשק תואם לשרת ההסקות של NVIDIA Triton. עם הזמן, הוא צמח למנוע הסקות מלא המשרת מודלי LLM משמעותיים, וכן מודלים מבוססי טרנספורמר עבור דירוג, סיווג, מתן ציונים ועוד עבור מגוון רחב של צוותים כדי להפעיל את המוצרים של Perplexity. מטרת העיצוב של ROSE היא להיות פשוטה: לאורך זמן, משקיעים מאמצים כדי להגיע לביצועי שיא, אך העדיפות העליונה היא להסתגל לארכיטקטורות ולמקרי השימוש המתפתחים כל הזמן עבור המודלים שלנו.

מרחב התצורה של הקרנל across data types, head dimensions, compile-time constants, and launch settings.

ברמה גבוהה, ROSE מורכב ממנוע המתאים מודל לממשק שמשתמשים יכולים לאינטראקציה עמו. המנועים ניתנים לשימוש חוזר במידה רבה, ומטפלים בתזמון בקשות, יצירת אבְצָות (batching), ושליפת טוקנים מפוענחים, הטמעות או ציונים מהמודלים הבסיסיים. אתחול ההתקנים, ערוצי התקשורת בין תהליכים ובין צמתים, וטעינת המשקולות מטופלים כולם על ידם. עבור LLMs, המנוע אחראי על יצירת אבְצָות, חלוקה למקטעים (chunking), ודגימה. הוא אחראי גם על הקצאת אחסון KV, הן עבור תשומת לב מלאה והן עבור לינארית, וכן התאמת קידומות (prefix matching). מנועי הטמעה פשוטים יותר, שכן הם מבצעים רק אבְצָות מקוונת. הביצוע ממוטב כך שהמנוע מכין את סבב העבודה הבא בזמן שהמודל רץ על המאיץ במקביל.

ROSE חושף אוסף עשיר של שכבות מותאמות אישית למודלים, העוטפות ליבות GPU שנבנו באמצעות CuTeDSL, Triton, CUDA, CUTLASS ו-cuBLAS. עבור מימוש כפל מטריצות ותשומת לב (attention), אנו נסמכים בעיקר על ליבות NVIDIA מבית CUTLASS ו-cuBLAS. שאר הצינור (pipeline), כולל ליבות הטמעה, נורמליזציה, MoE והפעלה (activation), מהווים חלק מספרית הליבות של ROSE. כדי לתמוך בצורה הטובה ביותר בחומרה שעליה אנו מריצים הסקות, ROSE מצויד בפרימיטיבי תקשורת בין-התקנים משלו עבור ניתוב MoE, שליחה (dispatch) ושילוב (combine). בעוד שליבות אחרות פשוטות יחסית, עליהן להיות מותאמות במיוחד כדי להשיג ביצועי שיא בכל ממדי הנסתר וההתקנים השונים שעליהם הן רצות. בנוסף, אנו נסמכים על התאמה אישית בזמן הידור כדי להתאים לכל הכפתורים והכיוונונים של מודלים שונים (הטיית משקולות לעומת היעדר הטיית משקולות עבור נורמליזציה מסוג RMS, פונקציות הפעלה וכו'). כתוצאה מכך, החלטנו להשתמש ב-DSLs כדי לממש אותם: תחילה בנינו את ROSE באמצעות Triton. עם זאת, מאז אימצנו את CuTeDSL כסביבה העיקרית לתכנות GPU.

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

CuTeDSL הוא DSL מבוסס פייתון עבור מעבדי GPU של NVIDIA שנבנו על גבי אלגברה פריסתית של CuTe ו-MLIR. הוא מתהדר בהידור בזמן אמת (just-in-time) ל-PTX ממוטב ביותר, תוך שמירה על שליטה בפרימיטיבי חומרה ברמה נמוכה. מטרתו להציע את אותן תכונות וביצועים של CUTLASS C++, ללא זמני ההידור היקרים המעכבים את מהירות הפיתוח. מבין השפות שסקרנו או עבדנו איתן (כולל CUDA, cuTile, TVM IR ו-Triton), אנו מחשיבים אותו כמתאים ביותר לפיתוח הליבות שלנו.

בסך הכל, אמנם לא ברמה גבוהה ותמציתית כמו Triton או cuTile, מצאנו שהאלגברה של פריסת CuTe וההרכבים שלה (הנחשפים דרך cute.Tensor) היו תחליפים סביר. בנוסף, CuTeDSL מציע גישה בלתי מוגבלת לפרימיטיבי חומרה, ומספק לנו את רמת השליטה הנדרשת לנו כדי להשיג ביצועי שיא בכל מקרי הקצה. עם CUDA, הזדמנויות להתאמה אישית בזמן הידור היו מוגבלות, שכן כל פרמטר נוסף הוסיף עלות בלתי סבירה לזמני ההידור או מורכבות למערכת הבנייה. בניגוד לכך, בעזרת DSL המהודר בזמן אמת, אלא נוכל להתאים אישית בצורה אגרסיבית יותר מבלי להקריב את קצב האיטרציות. בניגוד ל-TVM IR, אשר השתמשה ב-CUDA כשלב ביניים, מצאנו שצינור הנתונים מבוסס ה-MLIR שהוריד ל-PTX היה חזק יותר וקל יותר לניפוי שגיאות מאשר כזה שפלט CUDA כשלב ביניים. בנוסף, באמצעות שימוש בפעולות מותאמות אישית לעטיפת פקודות PTX, גילינו שאנו יכולים למנף במלואן את כל היכולות של חומרת הבסיס.

עיקר עלות ההידור של ליבות CUDA מושקעת בדרך כלל בשלב פריסת התבנית של הליבות, שכן אופטימיזציה והידור משלב ה-PTX ואילך יכולים להסתיים תוך שניות ספורות. עם פרמטרים רבים בזמן הידור, שלב הידור זה יכול להתפוצץ למספר דקות עבור ליבה שאינה טריוויאלית. צוואר בקבוק זה מודגש בשל העובדה שהזדמנויות להידור מקבילים צצות רק מאוחר יותר, ברמות ה-IR וה-PTX, כאשר הטיפול בתבניות C++ מתרחש באופן סדרתי. אמנם קיימים טריקים לטיפול בכך, על ידי העברת מופעי תבניות ליחידות הידור נפרדות, אך הם דורשים התאמות משמעותיות במערכות הבנייה ושימוש במקרואים מגושמים. בניגוד לכך, ליבות CuTeDSL המהודרות בזמן אמת (JIT) אינן דורשות כללי בנייה מורכבים ואינן פוגעות בזמני ההידור, שכן מחזור הפיתוח יכול להתמקד בנקודה יחידה במרחב התצורה, בעוד ש בדיקות יחידה קפדניות בהמשך יכולות לכסות את כל מרחב התכנון של הליבה.

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

עלות ההידור הכבדה של CUDA מודגשת בשל הקושי לעקוב אחר מקור השגיאות כאשר נעשה שימוש בניתוב מבוסס-תבניות והתאמה אישית, כפי שקורה ב-CuTe וב-CUTLASS. באמצעות DSL, ניתן להצביע על שגיאות בשורת הקוד שיצרה אותן, מה שמקל משמעותית על ניפוי שגיאות. בנוסף, פשוט הרבה יותר למפות את הקוד שנוצר להצהרת ה-DSL שיצרה אותו באמצעות שלבי MLIR ביניים או הסתמכות על מידע ניפוי שגיאות המצורף ל-PTX שנוצר. הודות להיבטים אלו של ה-DSL, מצאנו את עצמנו פרודוקטיביים הרבה יותר בעת פיתוח ליבות.

בעוד שליבות חולקות דמיונות רבים בין מודלים, הן עדיין חייבות להיות מצוידות במגוון כפתורי תצורה שעשויים להשפיע על הביצועים. פרמטרים רלוונטיים כוללים:

ממד נסתר / ממד ראש / מספר ראשי תשומת לב (Hidden Dimension/Head Dimension/Number of Attention Heads)

סוגי נתונים של קלט/פלט

סכימת קוונטיזציה (Quantization)

קבועים בזמן הידור (הטיות, פונקציות הפעלה וכו')

ארכיטקטורת ROSE, כולל המנוע שלה, המודל, ספריית הקרנל והשירותים.

התאמה אישית בזמן הידור שימושית הן לביצועים והן לניפוי שגיאות. במיוחד במעבד צנרת (pipelined), כמו במקרה של מעבדי GPU, ענף שאינו קיים עדיף על ענף שתמיד מבוצע או לעולם אינו מבוצע. בנוסף, הידור ליבה עבור קבוצה מסוימת של כפתורי תצורה מביא לפחות פקודות PTX/SASS, שקל יותר לנתח ולנפות שגיאות בעת שימוש ב-cuda-gdb במהלך חקירה.

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

בדרך כלל, יש לבצע אופטימיזציה לביצועי ליבות הסקה עבור שני מקרי שימוש נפרדים:

Prefill: גודל אבְצָה קטן (מעט מקטעים, רצפים), טוקנים רבים לכל רצף. יש למקסם את התפוקה עבור הטמעות וגם עבור טוקני הקלט של LLMs. השהיה (latency) של עשרות עד מאות אלפיות השנייה.

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

עבור GEMM ותשומת לב (attention), שלב המילוי המוקדם (prefill) פועל על מספר גדול של טוקנים והביצועים מוגבלים על ידי תפוקת החישוב של ההתקן. פענוח כולל את אותה פעולה על גדלי אבְצָה קטנים, כאשר זמני הִרְצוּת נשלטים על ידי עלות טעינת ההפעלה והמשקולות מ-DRAM לתוך מעבדי הזרם המרובים (streaming multiprocessors) וליבות הטווח, המוגבלים במקום זאת על ידי רוחב פס הזיכרון של ההתקן. אף על פי שגודל האבְצָה או המספר הכולל של הטוקנים באבְצָה הוא מדד ברור יחסית, סף המעבר בין מילוי מוקדם לפענוח הוא עדין יותר מפני שרוחב פס הזיכרון נקבע גם על ידי הממד הנסתר של המודל וסוג הנתונים של ההפעלה. כתוצאה מכך, ליבות מפוצלות למימושי מילוי מוקדם ופענוח אם הבדלים אלגוריתמיים או הקצאות רשת/בלוק יתנו יתרון לאחד מהם.

הפרמטר שונה אוניברסלית בין ליבות מילוי מוקדם ופענוח הוא מספר ה-warps בבלוק (מספר השרשורים, מחולק בגודל ה-warp של 32 שרשורים לכל warp). רצוי להפוך פרמטר זה לקבוע בזמן הידור, מכיוון שהוא יכול לסייע בהקצאת אגירת נתונים (registers). ב-CUDA, פעולה זו מתבצעת בדרך כלל על ידי ציון מספר השרשורים באמצעות __launch_bounds__ ואספקת מספר ה-warps כפרמטר תבנית:

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

מספר האגירים (registers) הזמינים לכל השרשורים בבלוק הוא מגבלת חומרה, 64K ב-Hopper וב-Blackwell. אם מספר השרשורים ידוע מראש, המהדר יודע בדיוק כמה אגירים הוא יכול להקצות לכל שרשור ויכול להצליח להדר ליבות ממוטבות באגרסיביות רבה יותר. עם זאת, הדבר כרוך בעלות בזמן הידור, שכן כעת יש להציג מחדש את הליבה עבור גדלי warp שונים. עם CuTeDSL, זו אינה בעיה עוד, מכיוון שאיננו צריכים לשלם את עלות ההידור עבור כל מרחב התצורה מראש. ה-DSL מצפה גם שגודל הבלוק יהיה ביטוי קבוע, תוך ביצוע משתמע של האופטימיזציה המרומזת על ידי __latch_bounds__ ב-CUDA.

מלבד הפיכת מספר ה-warps לקבוע, רצוי לשנות אותו בין prefill ל-decode. כדי לנצל במלואו את רוחב הפס של הזיכרון ב-GPU, אין די תמיד להקצות מספר גבוה של warps לכל block, שכן ניתן להשיג רוחב פס מלא רק באמצעות שימוש בכמה שיותר מן ה-Streaming Multiprocessors‏ (SMs) הזמינים. לכן, עבור מספר נתון של tokens וסוגי נתונים וממדים חבויים ידועים, אנו בוחרים תחילה ממד וקטורי אידיאלי לכל thread ‏(8 עבור bfloat16, 4 עבור float32) כדי להבטיח שהקוד שנוצר משתמש בהוראות load/store ברוחב 128-bit ‏(256-bit ב-Blackwell). לאחר שאנו יודעים את גודל הווקטור, אנו יכולים לחשב את מספר ה-warps הנדרש על ידי הליבה. תחילה אנו מנסים למקסם את מספר ה-SMs: אם יש פחות warps מ-SMs, נבחר warp יחיד לכל block. זוהי בדרך כלל גרסת ה-decode של הליבה. אם יש מספיק warps, אנו מגדילים בהדרגה את מספר ה-warps ל-8, 16 או 32, ומבצעים benchmark של גבול עליון שמעבר לו בדרך כלל כבר לא רואים שיפור בביצועים.

אף על פי ש-Triton ביצע חלק מאופטימיזציות אלו מתחת למכסה המנוע, מצאנו שאנו יכולים להשיג תוצאות דומות באמצעות שימוש זהיר באלגברה של פריסת CuTe, מבלי להיות מוגבלים על ידי החסרונות של מהדר ברמה גבוהה. במיוחד כאשר מעורבים צמצומים אופקיים תת-warp, מצאנו ש-CuTeDSL הציע לנו את השליטה העדינה הנדרשת כדי ליצור קוד מיטבי.

במצבים מסוימים, כדי לחלק עבודה על פני כמה שיותר בלוקים, ניתן להשתמש ב-cooperative_groups::this_grid().sync() כדי לסנכרן את כל הבלוקים המשתתפים. חלות מגבלות: מכיוון שלא ניתן לבטל את תזמון הבלוק מ-SM מרגע שהוא מתחיל להתבצע, כדי להבטיח התקדמות קדימה, מספר הבלוקים אינו יכול לעלות על מספר ה-SMs הזמינים בהתקן, כדי להבטיח שבכל נקודה כל הבלוקים יוכלו לעדכן ולדגום את מיקומי הזיכרון הגלובלי המעורבים בסנכרון. מגבלה זו נאכפת על ידי cudaLaunchCooperativeKernel. עלות מחסום רשת (grid barrier) היא כ-3us ב-Hopper, בתוספת עלות פינוי כל פעולות הזיכרון הקודמות באמצעות מחסום זיכרון משתמע במצבים מסוימים. בסך הכל, הדבר בר השוואה לעלות פיצול ליבה לשתי הפעלות, מה שמוביל להחלטת עיצוב רגישה לביצועים. למרות העובדה שהפעלת תוכנה תלויה (PDL) יכולה להפחית את עלויות ההפעלה מ-Hopper ואילך, עבור ליבות פענוח, הימנעות מהעלות כדאית.

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

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

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

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

הליבה הראשונה שבנינו ב-CuTeDSL הייתה ליבת נורמליזציה מסוג RMS מותאמת אישית, הממוטבת לממדים קטנים. נורמליזציה של QK, המשמשת בשכבת תשומת הלב העצמית של מודלי Qwen האחרונים, מיושמת על בסיס כל ראש, על 64, 128 או 256 אלמנטים. בניגוד לנורמליזציה רגילה של RMS, הפועלת על ממדים נסתרים העולים על אלפים, הדורשת צמצום אופקי דרך זיכרון משותף ואגירים כאחד, נורמליזציה של QK דורשת צמצום רק בין warps, עם האפשרות לדחוס מספר ראשים לתוך warp בודד. זה הפך למקרה שימוש שבו יכולנו לבצע אופטימיזציה אגרסיבית בזמן הידור כדי למצוא את הקצאות הטוקן ל-warp הטובות ביותר עבור כל ממדי הראש וסוגי הנתונים.

נורמליזציה של RMS היא פעולה ישרה למדי, אם כי היא דורשת סכום אופקי כדי לחשב את נורמת הריבוע של וקטור:

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

מכיוון שראשי attention הם בדרך כלל קטנים ופעולות לכל thread ממוקדות מאוד בוקטוריזציה, warp יחיד יכול לחשב את הנורמה הריבועית. כדי למקסם את רוחב הפס של הזיכרון, כל thread צריך לעבד 128 ביט של קלט, שהם 8 ערכי bfloat16. כתוצאה מכך, half-warp מספיק לעיבוד ראש, וניתן לקבוע את הנורמה הריבועית באמצעות צמצום אופקי שמבוצע על ידי shuffle_sync_bfly. החל מממד הראש וסוג הנתונים, אנו יכולים לקבוע את מספר הראשים לכל warp ולציין את הארגומנטים ואת מספר סבבי הצמצום בזמן ההידור.

לאחר מכן נוכל לבנות ליבה מותאמת אישית עם מעט יחסית שורות קוד, תוך הסתמכות על טעינות/שמירות שעוברות וקטוריזציה אוטומטית ושימוש בפעולות ספציפיות לארכיטקטורה ברמה נמוכה יותר עבור צמצום ה-warp. הליבה קובעת באיזה ראש (Q או K) לעבוד ומזהה את המשקולות התואמות ליישום. היא טעינה וקטור שלם מהראש ומסכמת תחילה באופן מקומי. לאחר מכן, בהתבסס על מספר השרשורים לכל ראש, נורמת הריבוע מחושבת. לבסוף, נורמת ה-RMS, שעוברת וקטוריזציה אוטומטית, נקבעת ונכתבת חזרה לזיכרון. פרמטרים ספציפיים למודל, כגון WEIGHT_BIAS ו-EPS, עוברים קיפול קבועים (constant-folded).

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

בהשוואה לליבה רגילה של RMS norm, וריאנט מיוחד זה מספק תפוקה והשהיה טובות פי 2-3 הן ב-prefill והן ב-decode.

שליחה ושילוב של MoE משנים את סדר טוקני הקלט למומחים שאליהם הם מנותבים, ומבצעים צבירה משוקללת לאחר שליבות GEMM של MoE מבצעות את כפל המטריצות הרלוונטי. הביצועים מעניינים במיוחד, שכן מוצרים שונים נסמכים על מודלי MoE בגדלים שונים. בהתבסס על סוגי הפריסה, אנו הבחנים בין ליבות שונות למיטוב:

TP=1: ליבה מקומית עבור GPU יחיד. השליחה מערבבת טוקנים למומחים הרלוונטיים והשילוב צובר אותם. משמש בעיקר לפריסות Qwen3 30B על H100.

DP=1, TP ≠ 1: ליבת טנזור-מקבילית, צומת-יחיד. הקלט משוכפל, השליחה בוחרת את הטוקנים עבור הדרגה הנוכחית, והשילוב צובר אותם בכל הדרגות. צבירה דרך NVLink או NVLS. משמש לממלאים מוקדמים גדולים.

DP ≠ 1, NVLink: ליבה מקבילית נתונים עם מקבילות טנזור אופציונלית. כל דרגה שולחת טוקנים לכל דרגה אחרת המארחת מומחים ספציפיים. מידע הניתוב מצטבר, טוקנים מועברים דרך NVIDIA NVLink. משמש לפריסות H200 או GB200 מקבילי נתונים.

DP ≠ 1, InfiniBand: ליבה מקבילית נתונים ומקבילית טנזור עמית-לעמית הפועלת הן דרך NVLink והן דרך NVIDIA Quantum InfiniBand. משמש לפריסות Hopper רב-צמתיות.

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

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

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

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

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

עבור all-reduce, אנו משתמשים ב-NVLink רגיל אם רק שתי דרגות משתתפות, ועוברים ל-NVLS כאשר מעורבות 4 דרגות או יותר עבור עלייה של 10-20% בתפוקה. הודות ל-DSL, אנו יכולים לשתף את רוב לוגיקת הצמצום והחלפות מאגרי העמיתים בין ליבות ה-MoE וליבות ה-all-reduce הרגילות. זהו מצב שבו שליטה הדוקה בחומרה היא מועילה, שכן נדרשות פעולות DSL על multimem.ld_reduce ו-multimem.st עבור NVLS. מימוש הצמצום עצמו עוקב במידה רבה אחר דוגמת ה-all_reduce מתוך בחינות NCCL.

ליבות מקביליות נתונים נסמכות על NVLink כדי לצבור טוקנים עבור התקנים מרובים. בשלב השליחה, לפני החלפת טוקנים, עמיתים מתנתבים תחילה מקומית כדי לזהות את ההסט היעד בכל דרגה עבור כל טוקן שהם שולחים. לאחר מכן, הם מתחילים להעתיק את הטוקנים לתוך מאגרי קבלה שהוקצו לכל שולח כדי להימנע מקונפליקטים. כתיבות שחרור-אחסון (release-store) לכל טוקן משמשות לאותת על השלמה. בצד המקבל, מידע הניתוב מצטבר כדי להיות מסוגל לדחוס את כל הטוקנים למאגר יחיד עבור MoE GEMM. מקבלים דוגמים את הדגל שנכתב על ידי עמיתים עם טעינת-רכישה (acquire-load) כדי להמתין להשלמת הכתיבות.

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

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

ליבות העמית-לעמית שלנו השיגו בעבר ביצועים מובילים בתחום על מעבדי GPU מסוג Hopper המזווגים עם מתאמי ConnectX-7 InfiniBand. מאז העברנו את ליבות השילוב ל-CuTeDSL לרווחי ביצועים נוספים. מכיוון שגרסאות הCUDA כבר הסתמכו רבות על התאמה אישית בזמן הידור, הן התגלו כצווארי הבקבוק בזמן הידור במאגר שלנו, מה שחייב אותנו להשמיט פרמטרים מסוימים. על ידי מעבר ל-DSL המהודר ב-JIT, ביצענו התאמה אישית לפרמטרים נוספים וכווננו הקצאות רשת, מה שהציג הבדלים נוספים בין מילוי מוקדם לפענוח עבור תוספת ביצועים של כ-3% בהשהיה.

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

ל-GPU מסוג Blackwell יש עד 160 SMs, בניגוד ל-132 של Hopper. עבור רוב הליבות, זהו קבוע בזמן הידור ולוגיקה בקוד המפעיל בצד המארח לוקחת זאת בחשבון כדי לקבוע את הליבה המיטבית (מילוי מוקדם לעומת פענוח) להרצה. בנוסף, Blackwell מציג גם פקודות טעינה של 256 סיביות, אותן ניתן לפלוט בקלות על ידי הגדלת רוחב הווקטור (הממד האחרון של טנזורי CuTe) כדי לשפר את ניצול רוחב פס הזיכרון.

Blackwell מציגה תמיכה חומרתית במספר משמעותי של סכמות כימות (MXFP8, MXFP4, NVFP4). התמיכה בהן יושמה בעיקר על ידי התאמה של ליבות קיימות עבור קנה מידה לפי בלוקים ב-FP8. תמיכה ב-JIT היא קריטית כאן, שכן ניתן לתמוך במצבים חדשים מבלי לנפח את מרחב הקונפיגורציה בזמן ההידור. בעבר מצאנו ש-Triton מסורבלת בביטוי צמצומים אופקיים על פני מעט אלמנטים, מעל מחצית או רבע warp. שליטה עדינה יותר בצמצומים הייתה קריטית להשגת ביצועים מיטביים בעת חישוב ערך מקסימלי מוחלט על פני 16 או 32 אלמנטים עבור NVFP4 ו-MXFP8/MXFP4, בהתאמה.

כשפה חדשה יחסית, מודלי שפה גדולים (LLMs) נחשפו מעט מאוד לקוד CuTeDSL במהלך האימון המוקדם שלהם. בנוסף, הסגנון של הקוד הנוכחי שלנו, עם ליבות הממוטבות להפעלות ונורמליזציות, היה שונה מעט מדוגמאות CUTLASS או Quack הקיימות, שהיו ממוטבות בכבדות עבור פקודות TMA וכפל מטריצות ספציפיות לארכיטקטורה. הגעה לנקודה שבה LLMs יכלו לסייע בכתיבת ליבות הייתה תהליך הדרגתי.

בנינו באופן ידני מאגר ראשוני של ליבות ופונקציות עזר, תוך מעקב אחר דפוסים עקביים, טרם יכולנו להסתמך על סוכנים (agents) להמרת ליבות. ככל שהמאגר התקדם, יכולנו לבצע תרגומים גסים מ-CUDA ל-CuTeDSL, עם עבודה ידנית משמעותית שבאה בעקבותיה. בפרט, מצאנו שמודלים התקשו בהמרות מצביע לטנזור (pointer-to-tensor) ווקטוריזציה, ונתקעו בניסיון לתקן שגיאות סוגים כאשר כיוונו לגרום לבדיקה לעבור הידור. הם הציגו הבנה מועטה של האופי מבוסס-הרשת של הליבות, ובלבלו בין סנכרון שרשורים לרשת. באופן המאכזב ביותר, הייתה פיתוי להזות תרגומים שאינם קיימים של ספריות ספציפיות ל-CUDA במקום לבנות כנגד החלפות הCuTeDSL שלנו עצמנו. לאורך זמן, ככל שמאגר ליבות הCuTeDSL גדל, מצאנו ש-LLMs יכולים לבנות את רוב קוד השלד ולמלא את החלקיקים החסרים, מה שמאפשר לנו להתמקד בעיצוב ברמה גבוהה.

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

בזמן כתיבת שורות אלו, פער הביצועים בין ליבות CUTLASS C++ ו-CuTeDSL הוא מינימלי, מה שפותח הזדמנויות מרגשות להרכיב ולמזג ליבות GEMM ביצועים גבוהים עם תקשורת בין-התקנים, נורמליזציה או הפעלות. על ידי מחויבות ל-DSL זה, נוכל לחלץ ביצועים נוספים מההתקנים שלנו כדי לשפר את התפוקה בכל עומסי העבודה שלנו, החל מטמיעות ועד ל-LLMs.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

הסקת מסקנות ב-Perplexity AI

למה CuTeDSL?

חוויית מפתח

מרחב תצורה רחב

מילוי מראש ופענוח התמחות

סנכרון רשת (Grid synchronization)

ליבות CuTeDSL

נורמליזציה של QK

שליחה/שילוב של תערובת מומחים (Mixture-of-Experts Dispatch/Combine)

מקומי (Local)

צומת יחיד (Singlenode)

מקבילי נתונים (Data Parallel)

עמית-לעמית (Peer-to-Peer)

תמיכה ב-Blackwell

יצירת ליבות בסיוע AI

סיכום ועבודה עתידית

הפניות

ציטוט