עשרות שנות אופטימיזציה של CUDA עוברות אוטומטית לשבבי אפל

מעבירים ידע של עשרות שנים ב-CUDA לשבבי אפל בלי לכתוב שורת קוד ידנית. צוות מחקר הרחיב את K-Search, מסגרת אבולוציונית לאופטימיזציה של קרנלים שפותחה במקור ב-UC Berkeley Sky Lab, כדי שתיקח קרנלים קיימים מ-CUDA ותתרגם אותם ל-MLX, מסגרת למידת המכונה של אפל לשבבי ה-Silicon שלה. במקום לגלות מחדש אופטימיזציות מאפס לכל פלטפורמת חומרה, השיטה משתמשת בקוד ה-CUDA הקיים כבסיס ידע ומתאימה אותו לארכיטקטורה החדשה.

איך עובד K-Search

K-Search מקבל קרנל נאיבי ומפרט חומרה, ומריץ לולאת אופטימיזציה איטרטיבית: מודל שפה מנתח אילו אופטימיזציות כדאי לנסות, מודל כתיבת קוד מייצר קרנלים מועמדים, והם מהודרים ונמדדים על חומרה אמיתית. המדידות נזרעות חזרה לתהליך החיפוש, שממשיך ללטש כיוונים מבטיחים ולזרוק סמטרים עד שהביצועים מתכנסים. החיפוש מתבסס על מסמך Spec שמקודד חוקי חומרה, דפוסי אופטימיזציה ומגבלות מתמטיות, כדי למנוע מהקוד המיוצר להזות פרימיטיבים לא חוקיים ולוודא שהמועמדים יעברו הידור וירוץ ביעילות. בהרצות שלהם מילא מודל אחד (Gemini 3.5 Pro Preview) את שני התפקידים: הוא החזיק את מצב הניתוח וגם כתב את הקרנלים.

מה מרוויחים בפועל

התוצאות מראות שהשיטה מגיעה לביצועים כמעט ברמת מומחה על שבבי אפל, עם האצה של 0.97x בהשוואה לקרנל ה-Attention המובנה של MLX, ועד פי 20 על קרנל Mamba SSM בהשוואה למימוש mlx-lm של הקהילה. המחקר מפרט כמה מהרווח מגיע משכבת התרגום עצמה וכמה מלולאת החיפוש. למרות שהעבודה מתמקדת ב-MLX, השיטה אינה ספציפית ל-MLX וחלה על כל סביבה שבה אפשר להעביר ידע מ-CUDA.

פער העומק של MLX

MLX זוכה לאימוץ נרחב מאז סוף 2023. עם שבבי אפל במאות מיליוני מכשירי MacBook ו-Mac Studio, MLX מאפשר הרצת מודלים מקומית בלי עלויות ענן, וארכיטקטורת הזיכרון המאוחד הופכת אותו לאטרקטיבי למודלים בינוניים (7B עד 70B פרמטרים על שבבי סדרת M). מתחת למומנטום הזה יש פער משמעותי: קרנלים קריטיים לביצועים שמערכת ה-NVIDIA מתייחסת אליהם כמובנים מאליהם, כמו paged attention, קרנלי סריקה מאופטמיזציה של SSM, או fused MoE routing, או שלא קיימים ב-MLX או שהם נאיבים בלי כיוונון ספציפי לחומרה. MLX מריץ מודלים נכון אבל משאיר הרבה ביצועים על הרצפה. זה הפער שמניע את העבודה הזו.