אז הצ׳אט כן יודע קודה??

ביום חמישי בלילה מצאתי את עצמי קוראת מחקר שיצא באותו שבוע:

CUDA-LLM: LLMs Can Write Efficient CUDA Kernels

רציתי להבין איפה בדיוק עובר הקו בין קוד שעובד לבין קוד שמבין את הGPU.

העבודה הזו מציעה framework בשם FSR - ראשי תיבת של Feature Search and Reinforcement

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

ככה עד שמגיעים לקרנל שגם נכון וגם מהיר.

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

הם בנו את FSR סביב השאלה: איך גורמים לLLM לא רק לכתוב קוד לGPU אלא לכתוב קוד שעובד והוא גם ביצועיסטי.

לי זה די מזכיר את האלגוריתם הגנטי שמיממשתי בהאקתון עם תכנות גנטי (האקתון) לאלגו טריידינג.

המערכת מקבלת שלושה קלטים שנכנסים יחד לprompt:

1. תיאור מילולי של המשימה באנגלית. לדוגמה: Write a CUDA kernel for matrix transpose using shared memory

2. הקוד שמריץ את הkernel מהCPU (קוד הhost). כולל את ההגדרות של הקלטים, סוגי הנתונים, והlaunch עצמו (קריטי כי קובע את ממשק הפונקציה שהמודל צריך לכתוב)

3. פרטים על הGPU: כמות SMs, גודל compute capability ו-shared memory ועוד.

   
   

   

התהליך עצמו מחולק לשלושה שלבים קבועים:

אחרי שמאתחל N קרנלים..

1. Compilation Verifier

בודק אם הkernel עובר קומפילציה (פשוט nvcc או משהו, ומחזיר success/fail עם הודעת שגיאה)

2. Function Validator

מריץ את הkernel עם inputs קבועים, משווה לתוצאה ידועה מראש. אם יוצא משהו שונה, הkernel נפסל.

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

3. Performance Profiler

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

אם הkernel נפסל

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

ואז פשוט מריצים סיבוב נוסף. שוב N קרנלים, ואז תיקוף וחיפוש וכו. ממשיכים עד שלפחות אחד עובר. ורק אז (!!) עוברים לשלב ביצועים.

הביצועים נבדקים רק אחרי שהקרנל נכון

לא מתחילים למדוד זמנים על קוד שלא מחזיר את מה שצריך.

כשהkernelים שעברו מתמודדים אחד מול השני, נבחר הקרנל הכי מהיר - והוא זה שישמש כבסיס לסיבוב הבא. אבל הפעם, הprompt לא יכיל הודעת שגיאה - אלא תובנות ביצועיות: hints, latency, profiles

ככה נבנה loop כפול:

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

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

הPROMPTS

מוגדרים מראש עם תבניות מדויקות.

אם הקוד לא מתקמפל אז שולחים הודעת שגיאה

אם הקוד מתקמפל אבל לא נכון, שולחים mismatch בתוצאה

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

אז הprompt כולל קוד מלא, device type, שם הkernel והגבלות על מה לא (!!) לשנות.

אתחול

לפחות kernel אחד עבר את שני השלבים הראשונים, הפרופיילר בוחר את המהיר ביותר ויוצר refined prompt

אף קרנל מתוך הN הקיימים לא עבר, הFSR יוצר prompt חדש לכל מועמד. אם מועמד נכשל ספציפית בשלב הקומפילציה, זה מבנה הprompt

אם אף אחד מהקרנלים המועמדים לא עובר את שני השלבים ומועמד נכשל ספציפית בפונקציונליות בגלל execution שגוי, זה מבנה הprompt

יש mismatch עם הreference output + אף קרנל לא עבר את 2 השלבים הראשונים

פרומפט של משימה

משימות או לא להיות

סוגי המשימות שקיבלו די מגוונת אבל לא גדולה מדי. בדקו אותה על שני כרטיסים שונים.

יש משימות כמו dotproduct, transpose, sigmoid וגם prefixsum, histogramming, reduction. שילבו גם משימות של attention, convolution, matrix multiplication

גם משימות כמו Monte Carlo integration ורגרסיית OLS בוצעו

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

כן ראו שיפורים על פני קוד שנכתב ידנית

בtranspose ראו פי 104, בdotproduct פי 102, ובmonte carlo פי 179 (!) ובמקרים אחרים היה שיפור מתון יותר אבל עקבי. הם השוו מול baseline של CUDA ידני (לא מול קוד מג׳ונרט אלא מימוש כתוב על ידי אדם). הם השוו מול מימושים ידניים ב־CUDA, אבל לא מול ספריות קיימות כמו TVM, Triton או Cutlass. היה מעניין לראות איך הם היו עומדים מול קוד production ברמה גבוהה.

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

האלגוריתם

האלגוריתם