לדלג לתוכן

6.4 CuTe ו CuTe DSL layouts ו tiles הרצאה

לאורך פרק 6 עלינו במעלה מחסנית הספריות: בשיעורים 6.1 ו-6.2 ראינו את cuBLAS ואת cuDNN - ספריות "תקרא לי kernel מהיר, אני רק מפעיל אותו", סגורות ומונחות-היוריסטיקה; ובשיעור 6.3 עברנו ל-CUTLASS, שהיא לא ספריית שגרות מוכנות אלא ערכת כלים לבנייה של kernels משלכם, ספריית תבניות C++ header-only שבה אתם מרכיבים את ה-GEMM בעצמכם מתוך mainloop ו-epilogue. אבל השארנו שאלה פתוחה: מה בדיוק מאפשר ל-CUTLASS 3.x לתאר tiling מורכב, מיפוי threads לנתונים ותכנות של Tensor Cores בלי לשלם על ההפשטה בזמן ריצה, ובלי לכתוב מחדש את ה-kernel לכל דור חומרה? התשובה היא שכבה אחת עמוקה יותר, השכבה שעליה CUTLASS יושב: CuTe (מלשון CUDA Templates), אלגברת פריסות (layout algebra) שמתארת גם נתונים וגם threads באותה מכונה בדיוק. בשיעור הזה, שסוגר את פרק הספריות, נבין את הרעיון המרכזי - Layout כ-Shape ועוד Stride - נראה איך פריסות היררכיות מקודדות tiling, איך הכול נפתר בזמן קומפילציה באפס runtime overhead, נכיר את CuTe DSL שמביא את אותו מודל ל-Python, ולבסוף נעלה מדרגה למודל התכנות מבוסס-tiles (cuTile) - חלופה ל-SIMT שנולדה מתוך מגמת חומרה, ונסתכל על FlashAttention-4 כחוד החנית.

מ-CUTLASS ל-CuTe - הרעיון של אלגברת פריסות

כשכתבנו kernel של כפל מטריצות tiled בפרק 3 (שיעור 3.3 והפרויקט ב-3.5), עשינו ביד את החשבון המייגע: איזה thread קורא איזה איבר, מה ה-offset שלו ב-shared memory, איך פורסים בלוק גדול לtiles קטנים. כל שורה כזו של אריתמטיקת אינדקסים היא מקום שבו קל לטעות, וקשה מאוד לבדוק נכונות. CuTe מציעה טענה שאפתנית: כל הדפוסים האלה - פריסת נתונים בזיכרון, פריסת threads על פני הנתונים, ה-tiling עצמו - הם אובייקט מתמטי אחד ויחיד ששמו Layout, ולפריסות האלה יש אלגברה שלמה שאפשר להרכיב ולתמרן בה בדיוק כמו שמרכיבים פונקציות.

ה-glossary מנסח זאת חד: CuTe היא ספריית C++ header-only בתוך CUTLASS "לתיאור ולתמרון של טנזורים של נתונים ושל threads". שימו לב למילים "נתונים ושל threads" - זו הנקודה כולה. אותה מכונת Layout שמתארת איפה יושב איבר A[i][j] בזיכרון, מתארת גם איזה thread אחראי לו. שני העולמות, שבפרק 3 טיפלנו בהם בנפרד וביד, הופכים כאן לאובייקט אחד. וכיוון שהכול מקודד במערכת הטיפוסים של C++, הcompiler יכול לאמת חלק גדול מהנכונות עוד לפני שהתוכנית רצה.

   two ways to describe a "4x8 matrix in memory":

   chapter 3 (by hand):         CuTe (Layout):
   idx = row * 8 + col          Layout = Shape (4,8)
   for (row...) for (col...)              Stride (8,1)
     load A[idx] ...            layout(row,col) = row*8 + col*1
                                (checked by the compiler, zero runtime overhead)

מבנה הפריסה - Layout, Shape ו-Stride

ה-Layout הוא לב מערכת הטיפוסים של CuTe, והוא בנוי משני רכיבים, כל אחד הוא tuple (רשומה סדורה):

  • Shape - האורכים (extents) של הטנזור לאורך כל ממד. Shape = (4, 8) אומר מטריצה בת 4 שורות ו-8 עמודות.
  • Stride - דפוס המעבר (traversal): כמה איברים בזיכרון צריך לדלג כדי להתקדם באחד לאורך אותו ממד. Stride = (8, 1) אומר: צעד בשורה = 8 איברים, צעד בעמודה = 1 איבר.

הפונקציה שממפה קואורדינטה לוגית ל-offset בזיכרון היא פשוט המכפלה הפנימית של הקואורדינטה ב-Stride:

   offset(i, j) = i * stride_0 + j * stride_1

   Shape (4,8), Stride (8,1)  ->  row-major (contiguous rows):
     offset(2,3) = 2*8 + 3*1 = 19
     offset(0,0) = 0 ,  offset(3,7) = 3*8+7 = 31

   Shape (4,8), Stride (1,4)  ->  column-major (contiguous columns):
     offset(2,3) = 2*1 + 3*4 = 14

זהו הרעיון כולו: אותו Shape עם Stride שונה נותן פריסת זיכרון שונה. החלפת (8,1) ב-(1,4) היא ההבדל בין row-major (הקונבנציה של C/C++/NumPy) ל-column-major (הקונבנציה של Fortran/BLAS, ולכן של cuBLAS - ראו את מלכודת ה-leading dimension משיעור 6.1). ב-CuTe אף איבר לא זז בזיכרון; רק ה-Stride משתנה, וכל אריתמטיקת האינדקסים נגררת אחריו אוטומטית.

בקוד CuTe C++ נבנה את זה כך:

#include <cute/tensor.hpp>
using namespace cute;

// a fully static Layout: the sizes are known at compile time (Int<>)
auto rowmajor = make_layout(make_shape (Int<4>{}, Int<8>{}),
                            make_stride(Int<8>{}, Int<1>{}));

// equivalent: make_layout(Shape<_4,_8>{}, Stride<_8,_1>{})
print(rowmajor);            // (_4,_8):(_8,_1)
print(size(rowmajor));      // 32   = 4*8, the number of elements
print(rowmajor(2, 3));      // 19   = the offset we computed by hand

הטיפוסים Int<4> ו-_8 הם קבועים בזמן קומפילציה. כשכל הגדלים סטטיים כך, כל קריאה כמו rowmajor(2,3) מתקפלת בזמן compilation לקבוע 19 - אין שום חישוב בזמן ריצה. אפשר גם לערבב: ממדים דינמיים (למשל int M שנקבע בזמן ריצה) לצד ממדים סטטיים, וה-CuTe תדע לאילו לייצר חישוב ולאילו קבוע.

הטנזור - Tensor - פריסה ועוד מצביע

ה-Layout מתאר איך לגשת, אבל לא לאיפה. כשמצרפים ל-Layout מצביע (pointer) לזיכרון, מקבלים את ההפשטה המלאה - Tensor:

   Tensor  =  Layout  +  pointer
             (how to access)   (where the data is)

float* ptr = /* global / shared / register memory */;
Tensor A = make_tensor(ptr, rowmajor);   // Layout (_4,_8):(_8,_1)

A(2, 3) = 1.5f;   // writes to ptr[19], exactly per the offset from the Layout
float x = A(0, 0);// reads from ptr[0]

ה-Tensor של CuTe מפריד לגמרי בין "מה מבנה הגישה" (ה-Layout, שהוא ידע בזמן קומפילציה) לבין "איפה הבתים" (ה-pointer, שיכול להצביע ל-global memory, ל-shared memory או אפילו לאוגרים). זה מה שמאפשר לכתוב פונקציה גנרית אחת שעובדת על טנזור בזיכרון גלובלי ועל טנזור בזיכרון משותף באותו קוד בדיוק - ההבדל היחיד הוא מאיפה בא ה-pointer.

וכאן חוזרת התובנה הגדולה של CuTe: threads הם Tensor בדיוק כמו נתונים. נניח שבלוק של 128 threads צריך לעבד מטריצה 4x8. אפשר לתאר את פריסת ה-threads על פני הנתונים כ-Layout בפני עצמו, ואז "לחתוך" מהטנזור של הנתונים את החלק ששייך ל-thread מסוים - הכול באלגברת ה-Layout, בלי אריתמטיקת אינדקסים ידנית. במקום לכתוב int idx = threadIdx.x; if (idx < N) ... ולחשב offsets בראש, מרכיבים Layout של threads עם Layout של נתונים ומקבלים את המיפוי מוכן ומאומת.

פריסות היררכיות ו-tiling

עד כה ה-Shape שלנו היה שטוח: (4, 8). הכוח האמיתי של CuTe מגיע כשה-Shape עצמו הופך להיררכי - tuple של tuples. פריסה מקוננת היא בדיוק הדרך שבה CuTe מקודדת tiling: פירוק מטריצה גדולה לtiles קטנים.

נניח שאנחנו רוצים לתאר מטריצה 8x8, אבל מחולקת לtiles של 4x4 (ארבעה tiles בסך הכול). במקום Shape שטוח (8,8), נכתוב Shape מקונן שמפריד את הציר ל"בתוך הtile" ו"בין הtiles":

   8x8 matrix, 4x4 tiles:

   Shape  ((4,2),(4,2))   meaning:
     row axis = (4, 2)    -> 4 rows within a tile, 2 tiles tall
     column axis = (4, 2) -> 4 columns within a tile, 2 tiles wide

   every coordinate becomes hierarchical:  A((r_in, r_tile), (c_in, c_tile))
   and the nested Stride determines how a whole tile jumps in memory vs. a step within it.

הרעיון שהקואורדינטה עצמה נעשית מקוננת - (r_in, r_tile) במקום סתם r - הוא מה שהופך את ה-tiling מחישוב ידני לתכונה של הטיפוס. וכאן נכנסת האלגברה: ה-glossary אומר שפריסות "מרכיבות קטגוריה עם אלגברה עשירה". בפועל, CuTe נותנת פעולות שמייצרות מפריסה אחת פריסה אחרת:

פעולה מה היא עושה לשם מה
composition הרכבת שתי פריסות (כמו הרכבת פונקציות) מיפוי דרך שכבת אינדוקס נוספת
logical_divide / tiling חלוקת ציר ל"בתוך tile" ו"בין tiles" לקודד tiling
logical_product שכפול פריסה על פני רשת לפרוש tile על כל המטריצה
complement מילוי ה"חורים" לכיסוי מלא של הזיכרון בנייה נכונה של פריסות מקוננות

הודות לאלגברה הזו, פעולות שבפרק 3 כתבנו ביד - "קח את הtile של הבלוק הזה", "חלק אותו בין ה-threads" - הופכות לקריאה אחת כמו local_tile(A, tile_shape, coord) או local_partition(tile, thread_layout, tid), שמחזירה Tensor חדש שכבר מצביע למקום הנכון. אין for ידני, אין חשבון offsets, ובזכות הטיפוסים הסטטיים - אין runtime overhead.

הcompilation בזמן קומפילציה - אפס runtime overhead

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

  1. הכול נפתר בזמן קומפילציה. כשה-Shape וה-Stride הם קבועי compile-time, כל אריתמטיקת ה-offset מתקפלת לקבועים בזמן הcompilation. הגישה לזיכרון "נושאת אפס runtime overhead נוסף" - הקוד שנוצר זהה לזה שהייתם כותבים ביד עם אינדקסים קבועים, רק שאתם לא כתבתם אותו ולא טעיתם בו.
  2. תבנית אחת, ריבוי ארכיטקטורות. אותה תבנית CuTe מתקמפלת ל-kernels ממוטבים על פני דורות SM שונים (Ampere, Ada, Hopper, Blackwell). מכיוון שהפריסות הן פרמטרים של הטיפוס, שינוי גודל tile או מיפוי threads לדור חומרה חדש הוא החלפת פרמטר תבנית - לא כתיבה מחדש של ה-kernel.
  3. הcompiler בודק נכונות. ה-glossary אומר ש-CuTe "מקודדת מטא-נתונים כמו ארגון הזיכרון, גישות strided ו-tiling כך שהcompiler יכול לבדוק היבטים רבים של נכונות". Stride לא-תואם או tiling שאינו מכסה נכון את הזיכרון יכולים להיתפס כשגיאת קומפילציה, לא כבאג שקט בזמן ריצה.

הדוגמה הקנונית שה-glossary מציין היא kernel לשחלוף מטריצה (matrix transpose) שכתוב ב-CuTe ורץ על H100 (Hopper). השחלוף הוא בדיוק המקום שבו אלגברת הפריסות זוהרת: שחלוף לוגי הוא לא הזזת נתונים אלא החלפת ה-Stride. מטריצה row-major (_M,_N):(_N,_1) הופכת ל-column-major (_N,_M):(_1,_N) פשוט על ידי היפוך זוגות Shape/Stride - אותם בתים בזיכרון, פריסה אחרת. (זו בדיוק הזהות שראינו ב-6.1 כדי להימנע משחלוף פיזי מול cuBLAS: מאגר row-major זהה ביט-בביט לשחלוף שלו כשמפרשים אותו column-major.)

// transpose as a Layout: zero memory copies, just swapping Shape/Stride
auto A = make_layout(make_shape (M, N), make_stride(N, Int<1>{})); // row-major
auto A_T = make_layout(make_shape (N, M), make_stride(Int<1>{}, N)); // = transpose(A)
// A_T(j, i) points to the same byte as A(i, j)

התובנה החשובה: אין כאן קסם ריצה. יש כאן מטא-תכנות של ה-kernel ברמה גבוהה מאוד, שמתקפל בזמן compilation לקוד ה-SASS (ראו 4.2) שהייתם כותבים ביד אילו הייתם וירטואוזים אינסופיים ולא עייפים לעולם. תפיסה מוטעית נפוצה היא ש"הפשטת טנזור ברמה כזו חייבת לעלות בביצועים"; ב-CuTe היא לא עולה דבר, כי כל ההפשטה חיה בטיפוסים ומתה בזמן הקומפילציה.

שפה ייעודית ב-Python - CuTe DSL

לכתוב תבניות C++ של CuTe זה חזק אבל תלול. CuTe DSL היא שפה ייעודית (Domain-Specific Language) ב-Python שחושפת בדיוק את אותו מודל - אותן פריסות, אותם טנזורים - אבל בתחביר Python, עם JIT compilation (Just-In-Time) לקוד מקומי מהיר. חשוב להסיר מיד תפיסה מוטעית: CuTe DSL אינו Python מפורש בזמן ריצה. הקוד נלכד, עובר trace ומקומפל למטה עד SASS; ה-Python הוא רק החזית התחבירית.

המודל של CuTe DSL בנוי מארבע הפשטות ליבה:

  1. פריסות - Layouts - מתארות ארגון נתונים בזיכרון וגם על פני threads (בדיוק כמו CuTe C++).
  2. טנזורים - Tensors - מצרפים מצביע/איטרטור לנתונים עם מטא-נתוני Layout.
  3. אטומים - Atoms - ייצוג של פעולת חומרה יסודית בודדת: כפל-וצבירה של מטריצות (MMA - matrix multiply-accumulate) על Tensor Core, או פעולת העתקה (copy). ה-atom הוא אבן הבניין הקטנה ביותר שממופה ישירות להוראת חומרה.
  4. פעולות מרוצפות - Tiled operations - מתארות איך פורשים atom בודד על פני blocks ו-warps: איך לוקחים MMA קטן של Tensor Core אחד ומרחיבים אותו לפעולה קולקטיבית שלמה על פני הבלוק.

ההיררכיה הזו - atom בודד -> tiled operation על פני warps -> קולקטיב על פני הבלוק - היא בדיוק הדרך שבה מתכנתים Tensor Cores בפועל, ותואמת את שלוש הרמות של CUTLASS שראינו ב-6.3 (device / kernel / collective).

שני הדקורטורים (decorators) שמגדירים את מודל התכנות:

import cutlass.cute as cute

@cute.jit                      # entry point on the host, callable from Python
def gemm_launch(mA, mB, mC):
    # run Layout computations on the host, then launch the device's kernel
    gemm_kernel(mA, mB, mC).launch(
        grid=[grid_m, grid_n, 1],
        block=[128, 1, 1],       # 128 threads = 4 warps = one warpgroup
    )

@cute.kernel                   # the device's kernel, launched only from @cute.jit
def gemm_kernel(mA, mB, mC):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    # ... Layout algebra: local_tile, local_partition, then an MMA atom ...

ההבחנה בין השניים היא לב מודל התכנות, ומקבילה ל-@cute.jit כ"host" ו-@cute.kernel כ"device":

  • @cute.jit מכריז על פונקציה שעוברת JIT compilation וניתנת לקריאה מ-Python (או מפונקציות CuTe DSL אחרות). זו נקודת הכניסה מצד הhost.
  • @cute.kernel מגדיר kernel של הdevice שאינו ניתן לקריאה ישירה מ-Python - חייבים להשיק אותו מתוך פונקציית @cute.jit. זה בדיוק היחס שראינו לאורך הקורס בין קוד host ל-kernel: את ה-kernel לא "קוראים", אותו "משיקים".

הpipeline הcompilation (JIT lowering) - שרשרת רב-שלבית:

   Python code
      │  (converted to an abstract syntax tree)
      ▼
   AST  ─── trace with proxy arguments ───►  MLIR
                                            │  (Multi-Level Intermediate Representation)
                                            ▼
                                           PTX   (see 4.1)
                                            │
                                            ▼
                                           SASS  (device-specific, see 4.2)
                                            │
                                            ▼
                                           runs on the GPU

זה בדיוק אותו "מותן צר" (narrow waist) שראינו בפרק 4: CuTe DSL אינו ממציא נתיב חדש לחומרה, הוא רק מזין את אותו pipeline PTX -> SASS מלמעלה, דרך MLIR. לכן פלט הביצועים שווה-ערך למה שהייתם מקבלים מ-CUTLASS C++.

המגבלה החדה שחובה לזכור: רק תת-קבוצה של סמנטיקת Python נתמכת בקוד המקומפל ב-JIT. אי אפשר להשתמש ב-Python שרירותי בתוך @cute.kernel - לא כל בנייה, לא כל ספרייה, לא כל דפוס בקרת-זרימה. חושבים על זה כעל "Python בטעם C": התחביר מוכר, אבל מה שמותר בפנים מוגבל למה שאפשר לתרגם ל-MLIR ולמטה. CuTe DSL קרוב ברוחו ל-DSLs אחרים לכתיבת kernels ב-Python (למשל Triton, שנפגוש שוב מיד), אבל הוא ספציפית מודל ה-CuTe/CUTLASS - פריסות, טנזורים, atoms ו-tiled ops - שנותן שליטה עדינה יותר בפריסות וב-atoms של החומרה ממה ש-DSL ברמה גבוהה יותר בדרך כלל חושף.

השוואה מסכמת של שלושת המסלולים לאותו מודל:

היבט CuTe (C++) CuTe DSL (Python) cuTile
שפה תבניות CUDA C++ Python עם JIT Python (וניסיוני BASIC/Rust)
מתי נפתרות הפריסות קומפילציה (nvcc) JIT (Python->MLIR->PTX->SASS) JIT (Tile IR)
הפשטה עיקרית Layout + Tensor Layout, Tensor, Atom, Tiled op structured pointer + tile
יחידת התכנות thread/warp (SIMT) thread/warp (SIMT) tile block (thread-of-execution יחיד)
בשל / ניסיוני בשל, בסיס CUTLASS חדש, בפיתוח פעיל ניסיוני, תפקידו עוד פתוח

מודל התכנות מבוסס-tiles - cuTile

ספריות CuTe ו-CuTe DSL נותנות כלים חזקים, אבל הן עדיין חיות בתוך מודל ה-SIMT (Single Instruction, Multiple Threads) שראינו בשיעור 2.1: אתם עדיין חושבים במונחי thread, warp ו-block, גם אם אלגברת ה-Layout מקלה על אריתמטיקת האינדקסים. מודל התכנות מבוסס-tiles של CUDA (cuTile, מלשון tile = אריח) הוא צעד רדיקלי יותר - חלופה ל-SIMT עצמו.

ההבדל היסודי: ב-SIMT אתם כותבים תוכנית מנקודת מבט של thread בודד ("מה ה-thread הזה עושה עם האיבר שלו"), ואותה הוראה מונפקת במקביל ל-32 threads ב-warp. ב-cuTile אתם כותבים תוכנית מנקודת מבט של tile שלם: מופע של התוכנית - "tile-kernel" - רץ על פני רשת של tile blocks, וכל tile block הוא thread ביצוע יחיד שפועל על בלוק שלם של נתונים בבת אחת. יחידת התכנות עברה מ"thread אחד לאיבר אחד" ל"מופע תוכנית אחד לtile".

   SIMT (the classic model):     cuTile (tile-based):
   grid of blocks of threads      grid of tile blocks
   each thread -> one element     each tile block -> one whole tile (matrix)
   thinking at thread level       thinking at tile level
   the warp is a hidden impl. detail   hardware is mapped below the tile abstraction

המניע הוא מגמת חומרה, לא נוחות. ה-glossary חד בעניין: מודל ה-SIMT הקלאסי הוא "התאמה גרועה ל-GPUs של ארכיטקטורות ה-SM האחרונות, שבהן הרוב המכריע של רוחב-הפס האריתמטי נמצא ב-Tensor Cores". וה-Tensor Cores הם חיה אחרת: הם יודעים לבצע רק כפל מטריצות, וצריך לתכנת אותם עם הוראות ברמת ה-thread ואסינכרוניות, ולא באסינכרוניות ברמת ה-warp שבה מתכנתים את שאר החומרה (זכרו את ה-wgmma ברמת ה-warpgroup - 128 threads - מ-2.4, שדורש כבר היום in-line PTX). כשעיקר כוח החישוב יושב ביחידה שבטבעה עובדת על מטריצות/tiles, הגיוני להרים את ההפשטה מ"thread לאיבר" ל"תוכנית לtile", כך שהמודל התכנותי יתאים לאופן שבו החומרה באמת רוצה שיאכילו אותה.

structured pointers - מצביעים מובְנים. אבן הבניין של הנתונים ב-cuTile היא "structured pointer": מצביע שמשלב בתוכו גם את המידע על המערך - ה-extent הכולל שלו (shape) ודפוסי הגישה (stride). כלומר מצביע + shape + stride. אם זה מצלצל מוכר - זה בכוונה: ה-glossary מציין במפורש את הדמיון למערכת הטיפוסים של CuTe ל-Layouts ו-Tensors. structured pointer הוא בעצם אותו רעיון של Tensor = pointer + Layout, מוגש כאבן בניין מרכזית של מודל התכנות ולא כספריית תבניות.

מצב ומעמד:

  • מודל cuTile זמין כיום דרך cuTile Python, ובצורה ניסיונית גם דרך cuTile BASIC ו-cuTile Rust.
  • ייצוג הביניים שלו הוא Tile IR (במקביל ל-PTX/SASS של המסלול הקלאסי, ול-MLIR של CuTe DSL).
  • קרבה רעיונית: cuTile קרוב ברוחו ל-DSLs מבוססי-tiles כמו Triton (וכמובן לפריסות של CuTe/CUTLASS).
  • הסתייגות כנה שה-glossary עצמו רושם: "עד כמה הוא יחליף את מודל ה-CUDA SIMT הקיים עדיין לא ברור". זה לא תחליף מוגמר ל-SIMT אלא כיוון פתוח, תגובה למגמת חומרה שעוד לא הכריעה את דרכה.

חזית הביצועים - FlashAttention-4

כדי לראות עד לאן המודל הזה מגיע, ה-glossary מפנה לכתיבת kernels של FlashAttention-4 ב-CuTe DSL - נכון להיום, חוד החנית של kernels ל-attention. FlashAttention היא משפחת ה-kernels שמחשבת attention בלי לממש (materialize) את מטריצת ה-scores המלאה בזיכרון גלובלי, אלא בזרימה (streaming) על פני tiles תוך שמירת ביניים ב-shared memory - בדיוק סוג ה-tiling שאלגברת ה-Layout של CuTe נועדה לתאר. הגרסה הרביעית ממנפת שלוש טכניקות שכולן מבטאות את מה שראינו בשיעור:

  • הִתמחות warps בpipeline - pipelined warp specialization. במקום שכל ה-warps יעשו את אותו דבר, warps שונים מקבלים תפקידים שונים בpipeline: חלקם "מפיקים" (producers) שמזרימים tiles מ-global memory ל-shared memory, וחלקם "צרכנים" (consumers) שמזינים את ה-Tensor Cores. החפיפה בין הזזת הנתונים לחישוב היא מה שמחזיק את ה-Tensor Cores עסוקים - הסתרת latency ברמת הpipeline.
  • פעולות Tensor Core. ה-atoms של MMA שדיברנו עליהם, ממופים להוראות כפל-המטריצות של החומרה (זכרו את 1.5) - שם יושב רוב רוחב-הפס האריתמטי, וזה בדיוק המניע ל-cuTile.
  • פעולות Tensor Memory. דור Blackwell (B200) הוסיף Tensor Memory (TMEM) - זיכרון על-שבבי ייעודי שמשמש כמצבור (accumulator) עבור ה-Tensor Cores של הדור החדש. שימו לב למלכודת השמות: ה-Tensor Memory (TMEM) הזה שונה מה-Tensor Memory Accelerator (TMA) - ה-TMA הוא מנוע העתקה אסינכרוני בין global ל-shared memory, ולמרות שמו הוא אינו מאיץ את ה-Tensor Memory. שני דברים שונים לגמרי עם שם דומה במכוון להשלמת הבלבול.

הנקודה החינוכית: FlashAttention-4 היא ההוכחה ש-CuTe DSL אינו צעצוע לימודי אלא כלי מספיק חזק לבטא את ה-kernel המהיר ביותר בעולם ל-attention, כולל ניצול תכונות החומרה החדשות ביותר. וזה סוגר את הקשת של פרק 6: התחלנו מ-cuBLAS/cuDNN ("תן לי kernel מוכן"), עברנו ל-CUTLASS ("תן לי הפשטות להרכיב kernel"), ירדנו ל-CuTe שמתחת ("אלגברת הפריסות שמאפשרת את ההרכבה"), עלינו ל-CuTe DSL (אותו מודל ב-Python), והצצנו ל-cuTile (מעבר ל-SIMT) - וכולם, בסופו של דבר, נשפכים לאותו PTX -> SASS שפרק 4 לימד אותנו לקרוא.

סיכום

  • CuTe היא ספריית C++ header-only בתוך CUTLASS לתיאור טנזורים של נתונים ושל threads באותה מכונה בדיוק; היא הבסיס ש-CUTLASS 3.x יושב עליו, ופותרת את אריתמטיקת האינדקסים והבטיחות שכתבנו ביד בפרק 3.
  • ה-Layout הוא לב מערכת הטיפוסים: Shape (האורכים לאורך כל ממד) ועוד Stride (כמה איברים לדלג לצעד בכל ממד), ושניהם tuples; ה-offset הוא המכפלה הפנימית של הקואורדינטה ב-Stride, ולכן החלפת Stride לבדה הופכת row-major ל-column-major בלי להזיז ביט.
  • Tensor הוא Layout ועוד מצביע לזיכרון - הפרדה מלאה בין "איך ניגשים" (ידע קומפילציה) ל"איפה הבתים" (pointer), מה שמאפשר קוד גנרי אחד על global, shared ואוגרים.
  • פריסות היררכיות (tuple של tuples) מקודדות tiling, ולפריסות יש אלגברה עשירה (composition, logical_divide/product, complement) שהופכת "קח את הtile של הבלוק וחלק בין ה-threads" לקריאה אחת מאומתת.
  • הכול נפתר בזמן קומפילציה: אפס runtime overhead, תבנית אחת שמתקמפלת ל-kernels ממוטבים על פני דורות SM (Ampere עד Blackwell), וה-compiler בודק נכונות של strides ו-tiling; הדוגמה הקנונית היא kernel שחלוף מטריצה על H100, שהוא רק היפוך Shape/Stride.
  • CuTe DSL מביא את אותו מודל ל-Python עם ארבע הפשטות: Layouts, Tensors, Atoms (פעולת חומרה יסודית - MMA או copy) ו-Tiled operations (פרישת atom על blocks ו-warps).
  • מודל התכנות של CuTe DSL מבוסס שני דקורטורים: @cute.jit היא נקודת כניסה בhost הניתנת לקריאה מ-Python, ו-@cute.kernel היא kernel של הdevice שמושק רק מתוך @cute.jit; הcompilation הוא JIT דרך Python -> AST -> (trace) -> MLIR -> PTX -> SASS, ורק תת-קבוצה של Python נתמכת בקוד ה-JIT.
  • מודל התכנות מבוסס-tiles (cuTile) הוא חלופה ל-SIMT שבה tile-kernels רצים על רשת של tile blocks (כל אחד thread ביצוע יחיד) מעל structured pointers (מצביע + shape + stride); הוא נולד ממגמת החומרה של ריכוז רוחב-הפס האריתמטי ב-Tensor Cores, ייצוגו הוא Tile IR, הוא קרוב ל-Triton, ותפקידו מול ה-SIMT הקלאסי עדיין פתוח.
  • FlashAttention-4, שנכתב ב-CuTe DSL, הוא חוד החנית: הִתמחות warps בpipeline (pipelined warp specialization), פעולות Tensor Core, ופעולות Tensor Memory (TMEM של Blackwell, שאין לבלבל עם ה-TMA שאינו מאיץ אותה).