לדלג לתוכן

6.3 CUTLASS בניית GEMM בתבניות הרצאה

בשיעור 6.1 ראינו את cuBLAS: ספרייה סגורה שנותנת לנו kernel מוכן ל-GEMM (כפל מטריצות), בוחרת בעצמה, על סמך היוריסטיקות פנימיות שאיננו רואים, איזה מבין עשרות ה-kernel-ים שלה להריץ, ומחזירה ביצועים מצוינים - כל עוד הצורה, טיפוסי הנתונים ודור החומרה נופלים בדיוק לתוך מה שהיא כוונה עבורו. אבל מה קורה כשאנחנו רוצים משהו שהספרייה הסגורה לא נותנת: GEMM עם epilogue מיוחד שממזג פעולה נוספת לתוך אותו kernel, צורת מטריצה חריגה שה-heuristic של cuBLAS לא מכוונת אליה, או טיפוס נתונים חדש שרק החומרה האחרונה תומכת בו. כאן נכנסת CUTLASS - לא ספרייה שקוראים לה, אלא ערכת כלים שבונים איתה. בשיעור הזה נבין את ההבחנה העמוקה הזו בין "לקרוא ל-kernel" ל"לבנות kernel", נכיר את שלוש הרמות של CUTLASS (device, kernel, collective), נפרק את ה-GEMM ל-mainloop ול-epilogue, ונכתוב GEMM אמיתי ברמת ה-device שרץ על Tensor Cores (ליבות טנזור, ראו 1.5), נכוון את צורת ה-tile שלו ונמזג לתוכו scale ו-ReLU. נסגור בנתיב המודרני של Hopper - CUTLASS 3.x מעל CuTe - שאותו נפרוט לפרטים בשיעור 6.4, ובכלי ה-profiler שמאפשר לסרוק אלפי וריאנטים בלי לקמפל ידנית.

מ-cuBLAS ל-CUTLASS - לקרוא קרנל מול לבנות קרנל

ההבחנה המרכזית שכל השיעור נשען עליה היא זו: cuBLAS ו-cuDNN אומרים לכם "תן לי בעיה, אני אחזיר kernel מהיר"; CUTLASS אומרת "תן לי אבני בניין, אני ארכיב kernel מהיר בעצמי". ה-glossary מגדיר את CUTLASS כ"ספרייה של הפשטות למימוש אלגברה לינארית בעלת ביצועים גבוהים ב-kernel-ים של CUDA" - שימו לב לניסוח: הפשטות למימוש, לא שגרות מוכנות לקריאה.

מדוע בכלל צריך את זה? כי להגיע לביצועי-שיא על Tensor Cores הוא קשה להחריד. צריך לתזמן ידנית את טעינת ה-tiles מ-global memory (זיכרון גלובלי) ל-shared memory (זיכרון משותף) בpipeline רב-שלבי (multi-stage) שמסתיר latency, לנהל את ה-double buffering, למפות בדיוק את ה-threads ל-fragments של רגיסטרים, ולהזין את יחידות ה-MMA (Matrix Multiply-Accumulate) בקצב שלא מרעיב אותן. CUTLASS אורזת את כל המומחיות הזו כתבניות C++ שאפשר להרכיב מחדש. זו הסיבה שהיא header-only (ספרייה בכותרות בלבד): כל הקוד הוא תבניות שמתמחות (specialize) בזמן קומפילציה לטיפוסים, לצורות ולדור ה-SM (מעבד רב-זרמי) הקונקרטיים שלכם.

הטכניקה שמאפשרת את זה נקראת פולימורפיזם פרמטרי - parametric polymorphism: מימוש מתבנת יחיד שמתמחה בזמן קומפילציה למאות טיפוסים, צורות וארכיטקטורות. אתם כותבים Gemm<half_t, RowMajor, ...> והcompiler מייצר מזה kernel קונקרטי; מחליפים טיפוס אחד, ומקבלים kernel אחר לגמרי - אבל מאותו קוד מקור. זה ההפך הגמור מ-cuBLAS, שבה הבחירה נעשית בזמן ריצה ומוסתרת מכם.

היבט cuBLAS / cuDNN CUTLASS
מודל שימוש קוראים לפונקציה מוכנה מרכיבים kernel מתבניות
בחירת kernel היוריסטיקה פנימית, בזמן ריצה, סגורה אתם בוחרים, בזמן קומפילציה, גלוי
קוד פתוח לא (backend סגור) כן, מלא, על GitHub
מתי מנצחת הצורה סטנדרטית, רוצים "פשוט שיעבוד" epilogue מותאם, צורה חריגה, טיפוס חדש
שכבה application layer toolkit / abstractions

הערה על השם: ראשי התיבות הם "CUDA Templates for Linear Algebra Subroutines and Solvers", אבל ה-glossary מדגיש שבפועל CUTLASS עוסקת בעיקר ב-GEMM, כלומר ברמה השלישית של תקן ה-BLAS (BLAS Level 3, פעולות מטריצה-מטריצה). ה-"Solvers" הוא שאיפה, לא המוקד.

שלוש הרמות - device, kernel ו-collective

CUTLASS בנויה בשלוש רמות היררכיות, שממפות ישירות להיררכיית החומרה שלמדנו בפרק 2:

   ┌──────────────────────────────────────────────────────────┐
   │  device        the whole-GPU level                       │
   │  ─ cutlass::gemm::device::Gemm<...>                        │
   │  ─ the API the host calls: allocates, launches the grid, │
   │    syncs                                                  │
   │  ┌────────────────────────────────────────────────────┐  │
   │  │  kernel      what runs as a grid on the SMs         │  │
   │  │  ─ the __global__ itself: swizzle, loop over tiles  │  │
   │  │  ┌──────────────────────────────────────────────┐  │  │
   │  │  │  collective   usually the thread block          │  │  │
   │  │  │  ─ mainloop + epilogue                         │  │  │
   │  │  │  ─ this is where all the multiply-accumulate    │  │  │
   │  │  │    happens on the Tensor Cores                  │  │  │
   │  │  └──────────────────────────────────────────────┘  │  │
   │  └────────────────────────────────────────────────────┘  │
   └──────────────────────────────────────────────────────────┘
  • רמת ה-device היא ה-API שהhost (ה-CPU) קורא לו. אובייקט כמו cutlass::gemm::device::Gemm<...> הוא כמו קריאת cuBLAS מבחוץ: מעבירים לו מצביעים, גדלים וסקלרים, והוא מקצה זיכרון-עזר, מחשב את תצורת ה-grid ומשיק את ה-kernel. זו הרמה שנעבוד בה ברוב השיעור, כי היא הכי קרובה למה שכבר מוכר לנו.
  • רמת ה-kernel היא פונקציית ה-__global__ עצמה - מה שרץ כ-grid (רשת) על ה-SM-ים. היא אחראית על ה-swizzle (מיפוי בלוקים ל-tiles של הפלט) ועל הלולאה החיצונית על ה-tiles.
  • רמת ה-collective היא, בלשון ה-glossary, "בדרך כלל ה-thread block". זה הלב: כאן מתבצע הכפל-צבירה בפועל על ה-Tensor Cores, מול ה-tiles שנטענו ל-shared memory. הרמה הזו מתחלקת לשני חלקים - mainloop ו-epilogue - שהם הסעיף הבא.

חשוב להבין למה ההיררכיה הזו מקבילה לחומרה: ה-device הוא ה-GPU כולו, ה-kernel הוא ה-grid, וה-collective הוא ה-block שרץ על SM יחיד ומשתף ביניהם shared memory. ההרכבה הזו היא בדיוק מה שמאפשר להחליף אסטרטגיית tiling ברמת ה-collective בלי לגעת ב-API ברמת ה-device.

הלב של ה-collective - mainloop ו-epilogue

ה-glossary מפרק את ה-GEMM ברמת ה-collective לשני חלקים, וההבחנה הזו היא אולי הרעיון החשוב ביותר בשיעור:

  • ה-mainloop הוא האלגוריתם המרכזי ואסטרטגיות ה-tiling: הלולאה שרצה על ממד ה-K, בכל צעד טוענת tile של A ו-tile של B ל-shared memory, ומזינה אותם ל-Tensor Cores שמבצעים כפל-צבירה (multiply-accumulate) לתוך רגיסטרים. זה ה-tiling עם shared memory שראינו ב-3.3 ובפרויקט 3.5, אבל בגרסה שמנוהלת בקפדנות מרבית ומשתמשת ב-Tensor Cores במקום ב-CUDA Cores.
  • ה-epilogue הוא העיבוד שאחרי הכפל: הפעלת סקלרים (למשל alpha ו-beta ב-D = alpha*A*B + beta*C), פונקציות אקטיבציה לא-לינאריות (ReLU, GELU), המרות טיפוס, וכתיבה חזרה ל-global memory. ה-epilogue ממוזג (fused) לתוך אותו kernel של ה-GEMM.

מדוע המיזוג הזה כה חשוב? בדיוק מהסיבה שראינו ב-cuDNN בשיעור 6.2: אם תריצו GEMM ואז kernel נפרד ל-ReLU, התוצר של ה-GEMM ייכתב ל-global memory, ואז ייקרא ממנה שוב - תעבורת זיכרון כפולה על מטריצה שלמה. עם epilogue ממוזג, תוצאת הצבירה יושבת ברגיסטרים, ה-ReLU חל עליה שם, ורק התוצאה הסופית נכתבת פעם אחת. חסכנו קריאה-כתיבה שלמה של המטריצה. עבור GEMM שהוא memory-bound בזנב שלו, זה שיפור ישיר.

   without fusion (two kernels):
     GEMM ─► D_tmp to global ─► [read] ─► ReLU kernel ─► D to global
             (write N*M)        (read N*M)              (write N*M)

   with a fused epilogue (one kernel):
     GEMM ─► accumulator in registers ─► ReLU in place ─► D to global
                                                        (write N*M only)

ה-mainloop קובע כמה מהר תגיעו לשיא ה-Tensor Cores; ה-epilogue קובע כמה עבודה נוספת תוכלו למזג בחינם. CUTLASS מפרידה ביניהם כך שתוכלו להחליף epilogue (מ-LinearCombination פשוט ל-LinearCombinationRelu) בלי לגעת ב-mainloop.

דוגמה קונקרטית - GEMM ברמת ה-device

נכתוב עכשיו GEMM אמיתי ומלא ב-CUTLASS, בטיפוסי FP16 עם צבירה ב-FP32 (הדפוס הנפוץ ב-deep learning), שרץ על ה-Tensor Cores. נשתמש ב-API הקלאסי של device::Gemm, כי הוא חושף במפורש את צורות ה-tile שאותן נכוון בסעיף הבא. שימו לב שהטיפוס Gemm כולו מוגדר בזמן קומפילציה מפרמטרי תבנית - זה בדיוק הפולימורפיזם הפרמטרי בפעולה:

#include <cutlass/cutlass.h>
#include <cutlass/gemm/device/gemm.h>
#include <cutlass/util/host_tensor.h>
#include <cutlass/util/reference/device/gemm.h>

// D = alpha * A * B + beta * C
using Gemm = cutlass::gemm::device::Gemm<
    cutlass::half_t, cutlass::layout::RowMajor,      // A: FP16, rows
    cutlass::half_t, cutlass::layout::ColumnMajor,   // B: FP16, columns
    cutlass::half_t, cutlass::layout::RowMajor,      // C/D: FP16, rows
    float,                                           // accumulator: FP32
    cutlass::arch::OpClassTensorOp,                  // use Tensor Cores
    cutlass::arch::Sm80,                             // hardware generation (Ampere; also runs on H100)
    cutlass::gemm::GemmShape<128, 128, 32>,          // threadblock tile  (M,N,K)
    cutlass::gemm::GemmShape<64, 64, 32>,            // warp tile         (M,N,K)
    cutlass::gemm::GemmShape<16, 8, 16>,             // MMA instruction   (M,N,K)
    cutlass::epilogue::thread::LinearCombination<
        cutlass::half_t,                             // output type
        128 / cutlass::sizeof_bits<cutlass::half_t>::value,  // vector width = 8
        float, float>,                               // accumulator, compute
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3>;                                              // number of pipeline stages

עכשיו הצד הhost שמריץ אותו. שימו לב שאין <<<...>>> - אובייקט ה-Gemm הוא שמשיק את ה-kernel פנימית:

int main() {
    int M = 4096, N = 4096, K = 4096;
    float alpha = 1.0f, beta = 0.0f;

    // HostTensor allocates on both the host and the device and syncs between them
    cutlass::HostTensor<cutlass::half_t, cutlass::layout::RowMajor>    A({M, K});
    cutlass::HostTensor<cutlass::half_t, cutlass::layout::ColumnMajor> B({K, N});
    cutlass::HostTensor<cutlass::half_t, cutlass::layout::RowMajor>    C({M, N});
    cutlass::HostTensor<cutlass::half_t, cutlass::layout::RowMajor>    D({M, N});

    // ... fill A,B with values, then A.sync_device(), B.sync_device() ...

    Gemm gemm_op;
    Gemm::Arguments args(
        {M, N, K},              // problem size
        A.device_ref(),         // TensorRef including the leading dimension
        B.device_ref(),
        C.device_ref(),
        D.device_ref(),         // output
        {alpha, beta});         // epilogue parameters

    cutlass::Status status = gemm_op(args);   // this is where the kernel is launched
    if (status != cutlass::Status::kSuccess) {
        std::cerr << "CUTLASS GEMM failed: "
                  << cutlassGetStatusString(status) << "\n";
        return 1;
    }
    cudaDeviceSynchronize();
    D.sync_host();              // copies the result back to the host
    return 0;
}

הקומפילציה (CUTLASS היא header-only, לכן רק כוללים את הנתיבים, בלי לינקוג' לספרייה):

export CUTLASS=$HOME/cutlass          # root of the repo you cloned
nvcc -O3 -std=c++17 -arch=sm_90a \
     -I$CUTLASS/include \
     -I$CUTLASS/tools/util/include \
     gemm.cu -o gemm

שימו לב לשתי דקויות שנוגעות ישירות לשיעור 6.1: ראשית, ה-TensorRef שמחזיר device_ref() נושא בתוכו את ה-leading dimension - הפרמטר שראינו ש-cuBLAS דורש במפורש, וכאן הוא נגזר אוטומטית מה-Layout. שנית, שמנו את A כ-RowMajor אבל את B כ-ColumnMajor. זו אינה שרירותיות: הבחירה של layout לכל אופרנד היא חלק מהתבנית, וה-mainloop מתמחה בהתאם. בניגוד ל-cuBLAS שכופה column-major על הכל ומאלץ אתכם בתעלול ה-C^T = B^T A^T, כאן ה-layout הוא פרמטר גלוי שאתם שולטים בו לכל מטריצה בנפרד.

הערה על דור החומרה: כתבנו arch::Sm80 (Ampere). ה-kernel הזה, על נתיב ה-Tensor Core של Ampere (הוראת mma.sync בצורת 16x8x16 עם cp.async לטעינה א-סינכרונית), רץ מצוין גם על ה-H100 שלנו (compute capability 9.0), כי H100 תואם לאחור. את הנתיב הילידי של Hopper - WGMMA ו-TMA - נראה בסעיף על CUTLASS 3.x.

כיוונון צורת ה-tile - threadblock, warp ו-instruction shape

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

  • ThreadblockShape - ה-tile של הפלט שכל thread block מחשב. במקטע שלנו <128, 128, 32>: כל block מחשב tile 128x128 של D, ומעבד את ממד ה-K בפרוסות של 32.
  • WarpShape - ה-tile שכל warp (קבוצת 32 threads) מחשב בתוך ה-block. <64, 64, 32> אומר שכל warp אחראי על 64x64 מתוך ה-128x128.
  • InstructionShape - צורת הוראת ה-MMA היחידה של החומרה. <16, 8, 16> היא הוראת ה-Tensor Core של Ampere ל-FP16 (16x8x16 - ראו 1.5).

מכאן נגזר מספר ה-warps ל-block, וזה קריטי ל-occupancy (תפוסה):

   warps per block = (TB_M / Warp_M) x (TB_N / Warp_N)
                    = (128 / 64)     x (128 / 64)
                    = 2 x 2 = 4 warps = 128 threads

וכמה shared memory כל stage צורך (tile של A בגודל 128x32 ו-tile של B בגודל 32x128, ב-FP16):

   per stage:  (128*32 + 32*128) elements * 2 bytes
            =  (4096 + 4096) * 2  =  16384 bytes = 16 KiB
   with 3 stages:  3 * 16 KiB = 48 KiB shared memory per block

ה-H100 מאפשר עד כ-228 KiB shared memory ל-SM (ראו 1.6), כך שיש מרווח גדול; אפשר להעלות ל-4 או 5 stages כדי להעמיק את הpipeline ולהסתיר טוב יותר את ה-latency של טעינת ה-tiles הבאים. הנה טבלת האיזון:

שינוי השפעה חיובית סיכון
threadblock גדול יותר (256x128) פחות tiles, יותר reuse של נתונים פחות blocks ל-SM, occupancy יורדת
threadblock קטן יותר (64x64) יותר blocks במקביל, occupancy עולה פחות reuse, יותר תעבורת זיכרון
warp tile גדול יותר פחות סנכרון בין warps register pressure עולה, spilling
יותר stages (5 במקום 3) הסתרת latency טובה יותר shared memory גדל, occupancy יורדת

התובנה המרכזית: אין צורת tile אחת אופטימלית. הבחירה תלויה בגודל הבעיה, בטיפוס הנתונים ובדור ה-SM - וזה בדיוק מה ש-cuBLAS מחביאה מאחורי ה-heuristic הסגור שלה. ב-CUTLASS אתם רואים את הידיות וסורקים אותן בעצמכם. עבור GEMM גדול וריבועי (למשל 4096^3) tile גדול כמו 128x256 מנצח כי ה-reuse שולט; עבור GEMM "רזה" (M קטן, batched decode של LLM) tile קטן וצר מנצח כי הוא מפזר עבודה על יותר SM-ים.

חשבון ביצועים לדוגמה: GEMM של 4096^3 הוא 2*M*N*K = 2*4096^3 ≈ 1.374e11 פעולות = 137.4 GFLOP. שיא ה-Tensor Core של H100 SXM ב-FP16 עם צבירת FP32 הוא כ-990 TFLOPS. הגבול התיאורטי התחתון לזמן:

   t_min = 137.4e9 FLOP / 990e12 FLOP/s = 0.139 ms

kernel של CUTLASS מכוונן היטב יגיע ל-70%-85% מהשיא (כ-700-840 TFLOPS) - קרוב מאוד ל-cuBLAS על אותה צורה, ולעתים מנצח אותו כשהצורה חריגה.

הוספת epilogue - scale ו-ReLU

הכוח האמיתי של CUTLASS מתגלה כשממזגים עבודה ל-epilogue. ב-mainloop שלנו כבר יש את הצבירה; עכשיו נחליף את ה-LinearCombination (שמחשב רק alpha*acc + beta*C) ב-LinearCombinationRelu, שמוסיף אקטיבציית ReLU על התוצאה - הכל בתוך אותו kernel, בלי מעבר נוסף בזיכרון:

#include <cutlass/epilogue/thread/linear_combination_relu.h>

using EpilogueReLU = cutlass::epilogue::thread::LinearCombinationRelu<
    cutlass::half_t,                                          // output type
    128 / cutlass::sizeof_bits<cutlass::half_t>::value,       // vector width = 8
    float,                                                    // accumulator
    float>;                                                   // compute

using GemmReLU = cutlass::gemm::device::Gemm<
    cutlass::half_t, cutlass::layout::RowMajor,
    cutlass::half_t, cutlass::layout::ColumnMajor,
    cutlass::half_t, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<128, 128, 32>,
    cutlass::gemm::GemmShape<64, 64, 32>,
    cutlass::gemm::GemmShape<16, 8, 16>,
    EpilogueReLU,                                             // <<< only this changed
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3>;

ההרצה כמעט זהה; רק פרמטרי ה-epilogue משתנים. ה-LinearCombinationRelu מקבל alpha, beta ו-threshold (סף ה-ReLU, ברירת המחדל 0):

GemmReLU gemm_op;
GemmReLU::Arguments args(
    {M, N, K},
    A.device_ref(), B.device_ref(), C.device_ref(), D.device_ref(),
    {alpha, beta});          // ReLU is applied automatically after alpha*acc + beta*C
cutlass::Status status = gemm_op(args);

מה שקורה בפועל בתוך ה-epilogue, ליסוד יחיד של הפלט:

   acc  = the accumulator from the mainloop (in FP32, in registers)
   x    = alpha * acc + beta * C[i][j]        (LinearCombination)
   D[i][j] = max(x, threshold)                (the added ReLU)   -> written once

הרווח: המטריצה השלמה של התוצאה לא נכתבת ל-global memory ואז נקראת שוב ל-ReLU נפרד. היא נשארת ברגיסטרים, ה-ReLU חל שם, ורק ה-D הסופי נכתב. עבור מטריצת 4096x4096 ב-FP16 חסכנו קריאה וכתיבה של 32 MiB - מעבר שלם על ה-HBM (זיכרון ה-GPU). זה בדיוק העיקרון של operation fusion שראינו ב-cuDNN (6.2), אבל כאן אתם מרכיבים אותו, ולא היוריסטיקה סגורה.

CUTLASS מספקת ספריית epilogues מוכנה - LinearCombinationGELU, LinearCombinationSigmoid, epilogues עם bias, עם clamp להמרת טיפוסים, ועוד - וכן מאפשרת לכתוב epilogue משלכם. זה בדיוק המקום שבו CUTLASS מנצחת את cuBLAS: כשצריך למזג פעולה שהספרייה הסגורה לא מציעה.

הנתיב המודרני - CUTLASS 3.x מעל CuTe

כל מה שראינו עד כה הוא ה-API הקלאסי (2.x). מ-CUTLASS 3.x ואילך הליבה בנויה מעל CuTe (ספריית התבניות לתיאור layouts ו-tensors, נושא שיעור 6.4), וזה הנתיב שבו מגיעים לביצועי-שיא על Hopper ו-Blackwell. במקום לפרט ידנית ThreadblockShape ו-WarpShape, משתמשים ב-CollectiveBuilder שבוחר את פרטי ה-collective האופטימליים לצורת ה-tile ולדור החומרה:

#include <cutlass/gemm/collective/collective_builder.hpp>
#include <cutlass/epilogue/collective/collective_builder.hpp>
#include <cutlass/gemm/device/gemm_universal_adapter.h>

using namespace cute;

using TileShape    = Shape<_128, _256, _64>;   // CTA tile: a CuTe layout
using ClusterShape = Shape<_2, _1, _1>;        // cluster of thread blocks (new in Hopper)

// the CollectiveBuilder builds the mainloop (WGMMA + TMA) automatically
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
    cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
    cutlass::half_t, cutlass::layout::RowMajor,    8,
    cutlass::half_t, cutlass::layout::ColumnMajor, 8,
    float,
    TileShape, ClusterShape,
    cutlass::gemm::collective::StageCountAuto,      // number of stages chosen automatically
    cutlass::gemm::collective::KernelScheduleAuto   // automatic warp-specialization
>::CollectiveOp;

שני הרעיונות החדשים כאן שייכים ל-Hopper וקושרים לשיעורים קודמים:

  • WGMMA (הוראת ה-warpgroup MMA) פועלת על warpgroup שלם - ארבעה warps, 128 threads (ראו 2.4) - במקום על warp בודד כמו ב-Ampere. היא גדולה יותר, ולכן יעילה יותר ל-tiles גדולים.
  • TMA (Tensor Memory Accelerator, מאיץ זיכרון הטנזור) מבצע העתקות tile שלמות בין global ל-shared memory בהוראה אחת, ומשחרר את ה-threads מחישוב כתובות. הערה חשובה מ-1.5: על אף שמו, ה-TMA אינו מאיץ את Tensor Memory - הוא מאיץ את תנועת ה-tiles.

הנתיב הזה משתמש גם ב-warp specialization: חלק מה-warps ב-block מוקדשים ל-producer (טעינה עם TMA) וחלק ל-consumer (חישוב עם WGMMA), בpipeline עמוק שמסתיר latency טוב יותר מכל מה שאפשר על Ampere. זו הסיבה ש-NVIDIA ממליצה על CUTLASS דווקא ל-Transformers על Hopper ו-Blackwell (כפי שציינו ב-6.2 על cuDNN): את ה-kernels הכי מהירים לחומרה הכי חדשה בונים ב-CUTLASS 3.x, לא קוראים מספרייה סגורה.

הרעיון החשוב לזכור: 3.x לא מבטלת את 2.x - היא מוסיפה שכבת הפשטה (CuTe) שבה ה-layout של הנתונים וה-layout של ה-threads מתוארים באותה מכונה מתמטית, מה שהופך את ה-tiling לניתן-הרכבה. את זה נלמד לעומק בשיעור 6.4.

בנייה, פרופיילר, ומתי להשתמש

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

git clone --depth 1 https://github.com/NVIDIA/cutlass
cd cutlass && mkdir build && cd build
cmake .. -DCUTLASS_NVCC_ARCHS=90a       # 90a = H100 architecture
make cutlass_profiler -j$(nproc)
# measures all GEMM variants in FP16 for a given size, reports GFLOP/s for each
./tools/profiler/cutlass_profiler \
     --operation=Gemm --m=4096 --n=4096 --k=4096 \
     --A=f16:row --B=f16:col --C=f16 --accum=f32

הפרופיילר מריץ את כל צורות ה-tile ואת כל ה-schedule-ים הרלוונטיים ומדווח מי המנצח - זו הדרך המהירה למצוא את הווריאנט הטוב ביותר לצורה שלכם, ואז לצרוב אותו בקוד. זהו גם מקור לימוד: קוד ה-kernel-ים הפתוח של CUTLASS נחשב לרפרנס, וה-glossary מציין שהוא "מוזכר בהרחבה בפיתוח kernels" - כלומר אנשים מעתיקים וממנו לומדים כשהם כותבים kernel-ים מהירים.

הנכס השני הוא הידע מתי בכלל להשתמש. ההחלטה:

מצב הכלי הנכון
צורה סטנדרטית, רוצים "פשוט מהיר" cuBLAS (6.1) - קריאה אחת, בלי לבנות
primitive של רשת (conv, attention) cuDNN (6.2)
GEMM עם epilogue מותאם (bias+act ממוזג) CUTLASS
צורה חריגה שה-heuristic של cuBLAS מפספסת CUTLASS, סרקו עם ה-profiler
טיפוס נתונים בלעדי לחומרה חדשה (FP8) CUTLASS 3.x
לומדים איך kernel מהיר בנוי קראו את מקורות CUTLASS

הכלל: אל תתחילו מ-CUTLASS אם cuBLAS פותרת את הבעיה. CUTLASS היא לא תחליף "drop-in" לספריות הסגורות; היא ערכת בנייה, ומשתמשים בה כשצריכים משהו מותאם או חדיש שהספרייה הסגורה לא נותנת. פתחו איתה כשה-profiler מראה ש-cuBLAS משאירה ביצועים על השולחן בצורה שלכם, או כשמיזוג epilogue יחסוך לכם מעבר זיכרון שלם.

סיכום

  • CUTLASS היא ערכת כלים לבניית kernel-ים, בניגוד ל-cuBLAS/cuDNN שהם ספריות שמקוראים להן; היא header-only, קוד פתוח, ומשתמשת בפולימורפיזם פרמטרי כדי שמימוש מתבנת יחיד יתמחה בזמן קומפילציה לטיפוסים, צורות ודורות SM רבים.
  • הספרייה בנויה בשלוש רמות - device (ה-API שהhost קורא לו), kernel (ה-grid שרץ על ה-SM-ים), ו-collective (בדרך כלל ה-thread block) - המקבילות להיררכיית החומרה.
  • ברמת ה-collective, ה-GEMM מתפרק ל-mainloop (הלולאה על K עם tiling וכפל-צבירה על Tensor Cores) ול-epilogue (scale, אקטיבציות, המרות), כשה-epilogue ממוזג לתוך אותו kernel וחוסך מעבר שלם בזיכרון.
  • מוקד הספרייה הוא GEMM (BLAS Level 3), ומגיעים לביצועי-שיא רק דרך תכנות זהיר של ה-Tensor Cores - בדיוק המומחיות ש-CUTLASS אורזת לתבניות ניתנות-להרכבה.
  • ב-API הקלאסי (2.x), device::Gemm חושף את ThreadblockShape, WarpShape ו-InstructionShape כפרמטרי תבנית, ואת ה-layout של כל אופרנד בנפרד; מספר ה-warps וה-shared memory ל-block נגזרים ישירות מהצורות האלה ומכתיבים את ה-occupancy.
  • כיוונון צורת ה-tile הוא איזון: tile גדול נותן reuse אבל מוריד occupancy, יותר stages מסתירים latency אבל צורכים shared memory; אין צורה אחת אופטימלית, וזו הידית ש-cuBLAS מחביאה מאחורי heuristic סגור.
  • החלפת ה-epilogue מ-LinearCombination ל-LinearCombinationRelu מוסיפה ReLU ממוזג בשינוי פרמטר תבנית יחיד, בלי לגעת ב-mainloop, וחוסכת קריאה-כתיבה של המטריצה השלמה - זהו operation fusion שאתם מרכיבים בעצמכם.
  • CUTLASS 3.x בנויה מעל CuTe (נושא 6.4) ומשתמשת ב-CollectiveBuilder, WGMMA (על warpgroup שלם), TMA (העתקת tiles, שלמרות שמו אינו מאיץ Tensor Memory) ו-warp specialization - זהו הנתיב שבו NVIDIA ממליצה לבנות Transformers על Hopper ו-Blackwell.
  • ה-cutlass_profiler (נבנה עם CMake, ‎-DCUTLASS_NVCC_ARCHS=90a ל-H100) סורק אלפי וריאנטים ומדווח GFLOP/s לכל אחד, וקוד המקור הפתוח של הספרייה הוא רפרנס לימודי לכתיבת kernel-ים מהירים.
  • מתי CUTLASS: כשצריך epilogue מותאם, צורה חריגה, או טיפוס חדש שהספרייה הסגורה לא נותנת; כשלא, cuBLAS פשוט יותר ומספיק.