לדלג לתוכן

9.3 פרויקט אופטימיזציית GEMM צעד אחר צעד פרויקט

פרויקט - אופטימיזציית GEMM צעד אחר צעד

זהו פרויקט הגמר (capstone) של הקורס - הוא מגלם בקוד אחד את כל מה שלמדנו. בפרויקט 3.5 טיפסנו את שלושת השלבים הראשונים של סולם ה-GEMM: naive, ואז coalesced (איחוד גישות), ואז tiled עם shared memory (זיכרון משותף), והגענו לכ-26% מה-peak של ה-H100 ל-FP32, מול cuBLAS שיושבת על 86%. מאז למדנו את התיאוריה שמסבירה למה כל קפיצה קרתה: מודל ה-Roofline והנקודה שבה עוברים מ-memory-bound ל-compute-bound (פרק 7), את ה-occupancy, ה-pipe utilization וה-register pressure (פרק 8), ואת ה-Tensor Cores (ליבות טנזור) שבהן מרוכז רוב רוחב הפס האריתמטי של החומרה המודרנית (פרק 6). ובשיעור 9.1 קיבלנו את הכלי שסוגר את הלולאה: NSight Compute (ncu), פרופיילר ברמת ה-kernel שאומר לנו, אחרי כל שלב, מהו צוואר הבקבוק הבא. הפעם נבנה את הסולם המלא - נמשיך מ-tiled אל register tiling (חד-ממדי ואז דו-ממדי), ולבסוף אל Tensor Cores - וכל צעד יונע לא מ"אינטואיציה" אלא ממדידה קונקרטית ב-ncu שמצביעה על צוואר הבקבוק. המטרה: להגיע ל-70% ומעלה מה-peak של CUDA Cores עם kernel שכתבנו בעצמנו, ואז לשבור את התקרה הזו לגמרי במעבר ל-Tensor Cores.

מפת הסולם המלא - the full GEMM ladder

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

   stage kernel               what it fixes              justification tool (ncu)  target %peak (FP32)
   ────────────────────────────────────────────────────────────────────────────────────────────
   1     naive                baseline                    -                          ~3%
   2     coalesced            thread->data mapping         sectors/request drops 32->4  ~10%
   3     tiled (shared)       reuse through shared          DRAM throughput drops        ~26%
   4a    1D blocktiling       each thread computes an output column   shared throughput / AI      ~48%
   4b    2D blocktiling       each thread computes an 8x8 tile      FMA pipe util, registers    ~70%
   ────────────────────────────────────────────────────────────────────────────────────────────
   ceiling  cuBLAS SGEMM (FP32)  reference ceiling          -                           ~86%
   ────────────────────────────────────────────────────────────────────────────────────────────
   5     wmma (FP16/TC)       move to Tensor Cores         tensor pipe util            (peak different!)
   ceiling  cublasGemmEx (FP16)  tensor reference ceiling   -                           (peak different!)

מטריצות ההרצה ריבועיות N x N מסוג float (FP32), עם N = 4096 כברירת מחדל למדידה. שימו לב ש-4096 מתחלק ב-64, ב-128 וב-8 - כל גדלי הtiles בפרויקט - כך שנוכל להשמיט בדיקות גבול בתוך ה-kernels המתקדמים ולהתמקד באלגוריתם. ה-GPU הרץ לדוגמה הוא H100 (SXM), כרגיל.

התשתית והמדדים - scaffolding and metrics

נמשיך מאותה תשתית של 3.5: מקרו ה-CUDA_CHECK, מימוש ייחוס על ה-CPU שצובר ב-double, בודק נכונות עם סובלנות יחסית, ומד-זמן מבוסס cudaEvent עם warmup. לא נחזור על כל הקוד כאן (הוא ב-3.5); נזכיר רק את מחשבון המדדים, ונרחיב אותו כדי שיתמוך גם ב-peak של ה-Tensor Cores שנצטרך בשלב 5.

// the two theoretical peaks of H100 (SXM), in GFLOP/s:
static const double H100_FP32_PEAK_GFLOPS   =  66900.0;   // CUDA Cores, FP32 (FFMA)
static const double H100_FP16_TC_PEAK_GFLOPS = 989000.0;  // Tensor Cores, FP16 dense

// reports N, time, GFLOP/s and percent of a given peak. peak differs for CUDA Cores vs Tensor Cores.
void report(const char* name, int N, float ms, double peak_gflops) {
    double flops  = 2.0 * (double)N * (double)N * (double)N;   // 2*N^3
    double gflops = flops / (ms / 1000.0) / 1e9;
    double pct    = 100.0 * gflops / peak_gflops;
    printf("  %-18s N=%d  %8.3f ms  %10.1f GFLOP/s  %6.2f%% peak\n",
           name, N, ms, gflops, pct);
}

שני ה-peaks הם עוגני ההשוואה. את שלבים 1 עד 4 נמדוד מול H100_FP32_PEAK_GFLOPS (הגזירה: 132 SM x 128 CUDA Core x 2 FLOP/FFMA x 1.98 GHz ≈ 66.9 TFLOPS, ראו 3.5). את שלב 5 נמדוד מול H100_FP16_TC_PEAK_GFLOPS (רוחב הפס האריתמטי הצפוף של ה-Tensor Cores ל-FP16, כ-989 TFLOPS), כי אלו יחידות חישוב אחרות עם roofline אחר לגמרי - כפי שהדגיש פרק 7, קו-הגג נגזר פעם אחת לכל תת-מערכת. השוואת kernel של Tensor Core מול ה-peak של CUDA Cores היא טעות מושגית.

כמות העבודה קבועה בכל השלבים: 2 * 4096^3 ≈ 137.44 GFLOP. מה שמשתנה הוא רק המהירות. וכלל הברזל של הפרויקט חוזר: אף פעם אל תדווחו GFLOP/s של kernel שלא עבר את בודק הנכונות.

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

שלבים 1-3 - סיכום מהיר עם ncu

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

שלב 1 - naive. כל thread מחשב איבר פלט אחד; המיפוי השגוי threadIdx.x -> row יוצר גישות לא-מאוחדות. בסיס ההשוואה, ~3% מה-peak.

שלב 2 - coalesced. מחליפים כך ש-threadIdx.x -> col; threads סמוכים ב-warp קוראים כתובות רצופות. ההצדקה ב-ncu אינה "הרגשה" אלא מטריקה יחידה וחד-משמעית - מספר ה-sectors לכל בקשת load מ-global memory:

ncu -k mmNaive --launch-skip 1 --launch-count 1 \
  --metrics l1tex__average_t_sectors_per_request_pipe_lsu_mem_global_op_ld.ratio \
  ./gemm
ncu -k mmCoalesced --launch-skip 1 --launch-count 1 \
  --metrics l1tex__average_t_sectors_per_request_pipe_lsu_mem_global_op_ld.ratio \
  ./gemm
  mmNaive:      l1tex__average_t_sectors_per_request_..._ld.ratio  ≈ 32.0
  mmCoalesced:  l1tex__average_t_sectors_per_request_..._ld.ratio  ≈  4.0

זו הראיה: warp של 32 threads שקורא 32 floats רצופים נוגע בדיוק ב-128 בתים = 4 sectors של 32 בתים - היחס האידיאלי 4.0. ב-naive, כל אחד מ-32 הthreads נוגע ב-sector אחר, ומקבלים 32.0 - פי 8 תעבורה מבוזבזת. ה-ncu אף מדפיס בקטע "Memory Workload Analysis" אזהרה מפורשת בשם "Uncoalesced Global Accesses".

שלב 3 - tiled (shared memory). מבצעים staging של tiles A ו-B ל-shared memory (ראו קוד מלא ב-3.5, TILE=32), וכל איבר משמש ב-TILE פעולות FMA לפני שנזרק. ההצדקה ב-ncu היא נפילת התעבורה ל-DRAM, שנמדדת ב"GPU Speed Of Light" כאחוז מרוחב הפס:

ncu -k "mmCoalesced|mmTiled" --launch-skip 1 --launch-count 1 \
  --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\
dram__throughput.avg.pct_of_peak_sustained_elapsed ./gemm
                sm__throughput(%)   dram__throughput(%)
  mmCoalesced        ~18%                 ~85%       <- memory-bound: DRAM saturated, SM starved
  mmTiled            ~40%                 ~30%       <- we moved the bottleneck from DRAM to the SM

התמונה ברורה: coalesced הוא memory-bound (ה-DRAM ב-85%, ה-SM מורעב), ואילו tiled מזיז את הצוואר - עכשיו ה-DRAM נושם וה-SM עסוק יותר. אבל 40% ניצול SM עדיין רחוק מרווי. מדוע? ה-kernel מבזבז זמן על __syncthreads ועל תעבורת shared memory יחסית לחישוב, כי כל thread מחשב עדיין רק איבר פלט אחד. זה בדיוק מה ששלב 4 מתקן.

התוצאות עד כאן (H100, N=4096, לדוגמה):

  naive              N=4096    68.000 ms      2021.2 GFLOP/s    3.02% peak
  coalesced          N=4096    20.000 ms      6872.0 GFLOP/s   10.27% peak
  tiled              N=4096     8.000 ms     17180.0 GFLOP/s   25.68% peak

שלב 4א - register tiling חד-ממדי - 1D blocktiling

הרעיון המרכזי: עד עכשיו כל thread חישב איבר פלט אחד, ולכן קרא מ-shared memory שני ערכים (As ו-Bs) לכל FMA. אם thread יחשב עמודה שלמה של TM איברי פלט, הוא יכול לטעון ערך אחד של B מ-shared memory ל-register, ולהשתמש בו מול TM ערכים שונים של A - כלומר TM פעולות FMA לכל load מ-shared memory. חתכנו את תעבורת ה-shared memory פי TM, בדיוק כמו ש-shared memory חתכה את תעבורת ה-global memory פי TILE. זו אותה מנטרה של פרק 7: להעלות arithmetic intensity (arithmetic intensity) על ידי ניצול-חוזר, שלב אחד עמוק יותר בהיררכיית הזיכרון.

הפרמטרים: הtile block בגודל BM x BN = 64 x 64, מימד פנימי BK = 8, וכל thread מחשב TM = 8 איברים. מספר הthreads ל-block הוא (BM * BN) / TM = 64 * 64 / 8 = 512.

#define BM 64
#define BN 64
#define BK 8
#define TM 8
// assumes N divides evenly by BM, BN, BK (true for N=4096). no boundary checks.

__global__ void mm1DBlocktiling(const float* A, const float* B, float* C, int N) {
    __shared__ float As[BK][BM];    // tile A stored transposed (transposed): As[k][m]
    __shared__ float Bs[BK][BN];    // tile B: Bs[k][n]

    const int blockRow = blockIdx.y, blockCol = blockIdx.x;

    // thread position within the 64x64 output tile:
    const int threadCol = threadIdx.x % BN;   // 0..63  -> output column
    const int threadRow = threadIdx.x / BN;   // 0..7   -> group of TM rows

    // indices for cooperative loading (all 512 threads load a tile together):
    const int innerRowA = threadIdx.x / BK, innerColA = threadIdx.x % BK;  // 64x8
    const int innerRowB = threadIdx.x / BN, innerColB = threadIdx.x % BN;  //  8x64

    const float* Aptr = A + blockRow * BM * N;              // head of A's row panel
    const float* Bptr = B + blockCol * BN;                  // head of B's column panel
    float* Cptr = C + blockRow * BM * N + blockCol * BN;

    float threadResults[TM] = {0.0f};                       // TM accumulators in registers

    for (int bk = 0; bk < N; bk += BK) {
        // loading the tiles into shared memory (A is stored transposed so later access is efficient):
        As[innerColA][innerRowA] = Aptr[innerRowA * N + innerColA];
        Bs[innerRowB][innerColB] = Bptr[innerRowB * N + innerColB];
        __syncthreads();

        Aptr += BK;                                         // advance to the next K tile
        Bptr += BK * N;

        for (int k = 0; k < BK; k++) {
            float Btmp = Bs[k][threadCol];                  // one load from shared...
            for (int m = 0; m < TM; m++)                    // ...for TM FMA operations
                threadResults[m] += As[k][threadRow * TM + m] * Btmp;
        }
        __syncthreads();
    }

    for (int m = 0; m < TM; m++)
        Cptr[(threadRow * TM + m) * N + threadCol] = threadResults[m];
}

הlaunch:

dim3 block(512);                                            // (BM*BN)/TM = 512 threads
dim3 grid(N / BN, N / BM);
auto launch = [&]{ mm1DBlocktiling<<<grid, block>>>(dA, dB, dC, N); };
float ms = timeKernel(launch);
report("1D blocktiling", N, ms, H100_FP32_PEAK_GFLOPS);

החשבון של הarithmetic intensity. בגרסת ה-tiled, הלולאה הפנימית הריצה FMA אחד לכל שני loads מ-shared memory. עכשיו, בלולאה הפנימית, טוענים Btmp פעם אחת ומריצים איתו TM = 8 פעולות FMA (מול TM ערכי As שכבר ב-shared). היחס FMA-ל-load מ-shared זינק פי 8. פחות לחץ על ה-shared memory משמעו שה-warp scheduler מוצא יותר עבודה אריתמטית מוכנה בין הגישות, וה-FMA pipe מתמלא.

קריטריון קבלה לשלב 4א: checkResult (עם rtol=1e-3) עובר, וה-GFLOP/s לפחות מכפיל את שלב ה-tiled. מה ncu אומר עכשיו:

  mmTiled           sm__throughput ~40%    shared throughput ~65%   FMA pipe ~35%
  mm1DBlocktiling   sm__throughput ~62%    shared throughput ~45%   FMA pipe ~55%
  1D blocktiling     N=4096     4.300 ms     31963.0 GFLOP/s   47.78% peak

הגענו ל-48% - כמעט פי 2 מ-tiled. אבל ה-FMA pipe עדיין ב-55% בלבד, וה-shared memory throughput עדיין גבוה יחסית. הצוואר הבא: אפשר להעלות את הarithmetic intensity שלב נוסף, ולנצל-חוזר ערכים בשני הצירים בבת אחת. זה שלב 4ב.

שלב 4ב - register tiling דו-ממדי - 2D blocktiling

בשלב 4א thread ניצל-חוזר ערך של B מול TM ערכי A. עכשיו נעשה זאת סימטרית: כל thread יחשב tile פלט מלא של TM x TN = 8 x 8 = 64 איברים. לכל צעד k, ה-thread טוען TM ערכי A ו-TN ערכי B ל-registers, ומ-TM + TN = 16 loads מ-shared memory הוא מפיק TM * TN = 64 פעולות FMA. יחס של 4 FMA לכל load מ-shared - שוב קפיצה בarithmetic intensity. זו הצורה שבה kernels אמיתיים (וגם cuBLAS) בנויים ביסודם.

הפרמטרים: BM x BN = 128 x 128, BK = 8, TM = TN = 8. מספר הthreads ל-block הוא (BM * BN) / (TM * TN) = 128 * 128 / 64 = 256.

#define BM 128
#define BN 128
#define BK 8
#define TM 8
#define TN 8

__global__ void mm2DBlocktiling(const float* A, const float* B, float* C, int N) {
    __shared__ float As[BK][BM];    // A transposed: As[k][m]
    __shared__ float Bs[BK][BN];    // Bs[k][n]

    const int blockRow = blockIdx.y, blockCol = blockIdx.x;
    const int numThreads = (BM * BN) / (TM * TN);          // 256

    // thread position in the 16x16 grid of thread-tiles:
    const int threadCol = threadIdx.x % (BN / TN);         // 0..15
    const int threadRow = threadIdx.x / (BN / TN);         // 0..15

    // loading indices + strides (256 threads load 128x8 and 8x128 in several passes):
    const int innerRowA = threadIdx.x / BK, innerColA = threadIdx.x % BK;
    const int strideA   = numThreads / BK;                 // 32 rows per pass
    const int innerRowB = threadIdx.x / BN, innerColB = threadIdx.x % BN;
    const int strideB   = numThreads / BN;                 // 2 rows per pass

    const float* Aptr = A + blockRow * BM * N;
    const float* Bptr = B + blockCol * BN;
    float* Cptr = C + blockRow * BM * N + blockCol * BN;

    float threadResults[TM * TN] = {0.0f};                 // 64 accumulators in registers
    float regM[TM], regN[TN];                              // register cache for step k

    for (int bk = 0; bk < N; bk += BK) {
        for (int off = 0; off < BM; off += strideA)        // loading tile A (transposed)
            As[innerColA][innerRowA + off] = Aptr[(innerRowA + off) * N + innerColA];
        for (int off = 0; off < BK; off += strideB)        // loading tile B
            Bs[innerRowB + off][innerColB] = Bptr[(innerRowB + off) * N + innerColB];
        __syncthreads();

        Aptr += BK;
        Bptr += BK * N;

        for (int k = 0; k < BK; k++) {
            for (int i = 0; i < TM; i++) regM[i] = As[k][threadRow * TM + i];
            for (int i = 0; i < TN; i++) regN[i] = Bs[k][threadCol * TN + i];
            for (int m = 0; m < TM; m++)                   // 8x8 outer product
                for (int n = 0; n < TN; n++)
                    threadResults[m * TN + n] += regM[m] * regN[n];
        }
        __syncthreads();
    }

    for (int m = 0; m < TM; m++)
        for (int n = 0; n < TN; n++)
            Cptr[(threadRow * TM + m) * N + threadCol * TN + n] = threadResults[m * TN + n];
}

הlaunch:

dim3 block(256);
dim3 grid(N / BN, N / BM);
auto launch = [&]{ mm2DBlocktiling<<<grid, block>>>(dA, dB, dC, N); };
float ms = timeKernel(launch);
report("2D blocktiling", N, ms, H100_FP32_PEAK_GFLOPS);

כאן register pressure הופך למרכזי - קישור ישיר ל-8.5. ספרו את ה-registers שהצהרנו במפורש: threadResults[64] + regM[8] + regN[8] = 80 floats, ומעליהם אינדקסים ומצביעים. נאמת את המספר האמיתי מהcompiler:

nvcc -O3 -arch=sm_90a -Xptxas -v -lcublas -o gemm gemm.cu
  ptxas info : Used 88 registers, used 0 stack, ... for mm2DBlocktiling

זהו התרחיש הקלאסי של פרק 8.5. עם 88 registers ל-thread ו-256 threads ל-block, ה-block דורש 88 * 256 = 22528 registers, וקובץ האוגרים של ה-SM ב-H100 מכיל 65536 registers - כלומר לכל היותר 65536 / 22528 = 2 blocks בו-זמנית, 512 threads = 16 warps, שהם רק 512 / 2048 = 25% occupancy. ובכל זאת זה מהיר יותר. זה בדיוק הלקח של 8.5: occupancy נמוך אינו רע כשה-kernel compute-bound עם ILP (instruction-level parallelism) גבוה - 64 פעולות FMA בלתי-תלויות בלולאה הפנימית ממלאות את ה-FMA pipe גם עם מעט warps. מה שכן חייבים לוודא הוא שאין register spilling: אם ptxas היה מדפיס spill stores, היינו מגבילים עם __launch_bounds__ או -maxrregcount, כפי שראינו ב-8.5, ובוחנים אם ה-tradeoff משתלם.

קריטריון קבלה לשלב 4ב: checkResult עובר, וה-kernel מגיע ל-לפחות 80% מהביצועים של cuBLAS (זמן_cuBLAS / זמן_2D >= 0.8). מה ncu אומר:

  mm2DBlocktiling   sm__throughput ~85%   FMA pipe ~80%   achieved occupancy ~25%   spills 0
  2D blocktiling     N=4096     2.900 ms     47393.0 GFLOP/s   70.83% peak

71% מה-peak, ו-47393 / 57267 ≈ 83% מ-cuBLAS - עברנו את סף הקבלה. ה-FMA pipe ב-80%, ה-occupancy נמוך ולא אכפת לנו. הגענו קרוב לתקרה של CUDA Cores. הצוואר החדש שה-ncu מראה הוא ה-FMA pipe עצמו - יחידות ה-FP32 רוויות. אין דרך להאיץ הלאה בלי לשנות יחידת חישוב. וזו בדיוק הנקודה שבה נכנסים ה-Tensor Cores.

שלב 5 - שבירת התקרה עם Tensor Cores - the wmma API

הגענו לגבול הפיזי של ה-CUDA Cores. כפי שהדגיש הגלוסארי לאורך כל הקורס: ה-matmul המהירים ביותר אינם רצים על CUDA Cores אלא על Tensor Cores, שרוחב הפס האריתמטי שלהן גבוה בסדר גודל. ל-H100 יש 4 Tensor Cores לכל SM (אחד לכל warp scheduler), וה-peak הצפוף שלהן ל-FP16 הוא כ-989 TFLOPS - פי כ-15 מ-66.9 TFLOPS של ה-FP32. המחיר: ה-Tensor Cores מבצעות רק כפל-מטריצות על tiles קטנים בדיוק כפול קבוע, והקלט הוא בדרך כלל בדיוק נמוך יותר (FP16 / BF16 / TF32) עם צבירה ב-FP32.

הדרך הנקייה לתכנת אותן ישירות ב-CUDA C++ היא ה-API בשם wmma (Warp Matrix Multiply-Accumulate), מתוך <mma.h>. היחידה כאן היא ה-warp, לא ה-thread הבודד: warp שלם משתף פעולה על tile פלט של 16 x 16, והנתונים חיים ב-fragment שמפוזר על פני 32 הthreads. זו בדיוק החריגה שהגלוסארי ציין בהיררכיית הזיכרון - עבור כפל-מטריצות ב-Tensor Core, ה-registers משותפים בין threads ה-warp ואינם עוד private לחלוטין ל-thread.

#include <mma.h>
using namespace nvcuda;

// valid fragment sizes for FP16 on H100: 16 x 16 x 16
static const int WM = 16, WN = 16, WK = 16;

// A, B in FP16 input; accumulation and output in FP32. row-major.
__global__ void mmWMMA(const half* A, const half* B, float* C, int N) {
    // each warp computes one 16x16 output tile:
    int warpM = (blockIdx.y * blockDim.y + threadIdx.y);
    int warpN = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;

    wmma::fragment<wmma::matrix_a, WM, WN, WK, half, wmma::row_major> aFrag;
    wmma::fragment<wmma::matrix_b, WM, WN, WK, half, wmma::row_major> bFrag;
    wmma::fragment<wmma::accumulator, WM, WN, WK, float>             cFrag;
    wmma::fill_fragment(cFrag, 0.0f);

    for (int k = 0; k < N; k += WK) {                   // accumulating along the K axis in tiles of 16
        int aRow = warpM * WM, aCol = k;
        int bRow = k,          bCol = warpN * WN;
        wmma::load_matrix_sync(aFrag, A + aRow * N + aCol, N);   // stride = N
        wmma::load_matrix_sync(bFrag, B + bRow * N + bCol, N);
        wmma::mma_sync(cFrag, aFrag, bFrag, cFrag);             // D = A*B + C
    }

    int cRow = warpM * WM, cCol = warpN * WN;
    wmma::store_matrix_sync(C + cRow * N + cCol, cFrag, N, wmma::mem_row_major);
}

הlaunch - חשוב לזכור שהיחידה היא warp, ולכן ציר ה-x מחולק ב-warpSize:

dim3 block(128, 4);                                    // 128*4 = 512 threads = 16 warps
int warpsX = block.x / 32;                             // 4 warps along x
dim3 grid((N + WM * warpsX  - 1) / (WM * warpsX),
          (N + WN * block.y - 1) / (WN * block.y));
auto launch = [&]{ mmWMMA<<<grid, block>>>(dA_h, dB_h, dC, N); };
float ms = timeKernel(launch);
report("wmma (FP16/TC)", N, ms, H100_FP16_TC_PEAK_GFLOPS);   // Tensor Cores peak!

סובלנות הנכונות משתנה - וזו נקודה חינוכית. הקלט עובר עיגול ל-FP16 (כ-3 ספרות עשרוניות מובהקות), ולכן התוצאה תסטה מהייחוס ב-FP32 הרבה יותר מ-1e-3. זה אינו באג - זו הגדרת ה"נכונות עד מספר ביטים" שפרק 7 פתח בה: ב-GPU נכונות מוגדרת עד דיוק מסוים. השוו מול הייחוס עם rtol של כ-2e-2, והבינו שאתם באמת מחשבים בדיוק נמוך יותר תמורת מהירות. אם צריך דיוק FP32 בלי לוותר על Tensor Cores, המסלול הוא TF32 (19 ביט: 10 ביטי מנטיסה כמו FP16, אבל עם טווח האקספוננט המלא של FP32, ולכן עמיד בהרבה מ-FP16 לגלישה/תת-גלישה לאורך ציר K).

קריטריון קבלה לשלב 5: checkResult עובר עם הסובלנות המורחבת, וה-GFLOP/s עובר בבירור את התקרה של CUDA Cores (66.9 TFLOPS). מה ncu אומר - עכשיו מסתכלים על ה-tensor pipe, לא ה-FMA pipe:

ncu -k mmWMMA --launch-skip 1 --launch-count 1 \
  --metrics sm__pipe_tensor_cycles_active.avg.pct_of_peak_sustained_active ./gemm
  wmma (FP16/TC)     N=4096     0.720 ms    190889.0 GFLOP/s   19.30% (tensor peak)

190.9 TFLOPS - פי 2.85 מהתקרה המלאה של כל ה-CUDA Cores, למרות שזה רק 19% מה-peak של ה-Tensor Cores. הפער נובע מכך ש-wmma נאיבי הוא memory-bound מחדש (קורא FP16 ישירות מ-global memory בלי staging של shared memory ובלי double-buffering). זה הצעד הבא ב"סולם של הסולם", ובדיוק מה שספריות בשלות עושות. נשווה לתקרה המקצועית.

התקרה של ה-Tensor Cores - cublasGemmEx

כמו שב-FP32 השווינו ל-cublasSgemm, כאן משווים ל-cublasGemmEx, שבוחרת kernel של Tensor Core מכוונן ידנית. זו ה-API הכללית של cuBLAS שבה מציינים במפורש את טיפוסי הקלט/פלט ואת מנוע החישוב. מלכוד ה-column-major זהה ל-3.5 - מחליפים את סדר A ו-B:

#include <cublas_v2.h>
const float alpha = 1.0f, beta = 0.0f;                 // alpha/beta in FP32 (compute type)

auto launch = [&]{
    cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                 N, N, N, &alpha,
                 dB_h, CUDA_R_16F, N,                  // cuBLAS's "A" = our B (FP16)
                 dA_h, CUDA_R_16F, N,                  // cuBLAS's "B" = our A (FP16)
                 &beta,
                 dC,   CUDA_R_32F, N,                  // FP32 output
                 CUBLAS_COMPUTE_32F,                   // accumulation in FP32
                 CUBLAS_GEMM_DEFAULT_TENSOR_OP);       // requires Tensor Cores
};
float ms = timeKernel(launch);
report("cublasGemmEx (TC)", N, ms, H100_FP16_TC_PEAK_GFLOPS);
  cublasGemmEx (TC)  N=4096     0.300 ms    458133.0 GFLOP/s   46.33% (tensor peak)

458 TFLOPS, 46% מה-peak של ה-Tensor Cores, ופי כ-6.85 מהתקרה המלאה של CUDA Cores. הפער בין ה-wmma הנאיבי שלנו (19%) לבין cuBLAS (46%) הוא בדיוק אותו סוג פער שראינו ב-FP32: staging חכם דרך shared memory, double-buffering, ו-swizzling של הtiles כדי למנוע bank conflicts. מי שרוצה לסגור אותו לא כותב wmma ידני אלא עובר ל-CUTLASS או ל-CuTe DSL (פרק 6) - ספריות התבניות שנבנו בדיוק כדי לתכנת Tensor Cores קרוב ל-peak, וש-cuBLAS עצמה בנויה על רעיונותיהן.

טבלת התוצאות המלאה - the full results table

הדליוורבל המרכזי. שתי טבלאות, כי יש שתי תקרות (שני roofline-ים). מלאו מהמדידות שלכם; הערכים כאן הם דוגמה מ-H100 עם N = 4096.

טבלה 1 - סולם ה-CUDA Cores (מול 66,900 GFLOP/s):

N kernel ms GFLOP/s % FP32 peak speedup מול naive
4096 naive 68.000 2021 3.0% 1.0x
4096 coalesced 20.000 6872 10.3% 3.4x
4096 tiled (shared) 8.000 17180 25.7% 8.5x
4096 1D blocktiling 4.300 31963 47.8% 15.8x
4096 2D blocktiling 2.900 47393 70.8% 23.4x
4096 cuBLAS SGEMM 2.400 57267 85.6% 28.3x

טבלה 2 - סולם ה-Tensor Cores (מול 989,000 GFLOP/s, קלט FP16):

N kernel ms GFLOP/s % TC peak פי-כמה מתקרת FP32
4096 wmma (FP16) 0.720 190889 19.3% 2.85x
4096 cublasGemmEx (FP16) 0.300 458133 46.3% 6.85x

לצד הטבלאות, כתבו פסקת הסבר לכל מעבר, כשהמשפט המחייב הוא המטריקה מ-ncu שהצדיקה את הצעד:

  • coalesced -> tiled: ה-DRAM היה ב-85% (memory-bound); shared memory חתך את התעבורה ל-global memory, וה-DRAM ירד ל-30%.
  • tiled -> 1D: ה-shared throughput היה הצוואר; ניצול-חוזר של ערך B מול TM ערכי A חתך את תעבורת ה-shared פי TM, וה-FMA pipe עלה מ-35% ל-55%.
  • 1D -> 2D: ניצול-חוזר דו-צירי הביא 64 FMA ל-16 loads מ-shared; ה-FMA pipe הגיע ל-80% - יחידות ה-FP32 רוויות, למרות occupancy של 25% (ILP מפצה, כפי ש-8.5 חוזה).
  • 2D -> Tensor Cores: ה-FMA pipe היה במקסימום; אין דרך למעלה בלי יחידת חישוב אחרת. המעבר ל-Tensor Cores שבר את תקרת ה-CUDA Cores כולה, במחיר ירידה בדיוק ל-FP16.

קריטריוני קבלה - acceptance criteria

הפרויקט מושלם כשכל השורות מתקיימות:

   [ ] all stages 1-4 pass checkResult (rtol=1e-3) against the reference; stage 5 passes with rtol=2e-2
   [ ] table 1 (CUDA Cores) complete for all 5 kernels + cuBLAS: N, ms, GFLOP/s, %peak, speedup
   [ ] table 2 (Tensor Cores) complete for wmma and cublasGemmEx, against the correct peak (989 TFLOPS)
   [ ] stage 2: ncu shows sectors/request dropping from ~32 (naive) to ~4 (coalesced)
   [ ] stage 3: ncu shows dram__throughput dropping from ~85% to ~30%
   [ ] stage 4b: 2D blocktiling reaches >=80% of cuBLAS (cuBLAS_time/2D_time >= 0.8)
   [ ] stage 4b: registers/thread reported from -Xptxas -v and spills=0 confirmed
   [ ] stage 5: wmma clearly beats 66.9 TFLOPS (the CUDA Cores ceiling)
   [ ] explanation paragraph for each transition, where the justification is a concrete ncu metric (not "a feeling")
   [ ] all CUDA/cuBLAS calls wrapped in CUDA_CHECK

פקודות הפרופיילינג המרכזיות שתריצו (מ-9.1), ריכוז:

# build with a registers report:
nvcc -O3 -arch=sm_90a -Xptxas -v -lcublas -o gemm gemm.cu

# automatic roofline for all kernels (ncu identifies where each one sits under the roof):
ncu --set roofline -o gemm_roofline ./gemm

# full report for manual bottleneck analysis (skip the warmup with launch-skip):
ncu --set full --launch-skip 1 --launch-count 1 -o gemm_full ./gemm

# focused metrics for the transition between two specific stages:
ncu -k "mm1DBlocktiling|mm2DBlocktiling" --launch-skip 1 --launch-count 1 \
  --metrics sm__pipe_fma_cycles_active.avg.pct_of_peak_sustained_active,\
sm__warps_active.avg.pct_of_peak_sustained_active,\
launch__registers_per_thread ./gemm

הערה מסכמת על נכונות מול מהירות: זכרו ש-__syncthreads חסר, בדיקת גבול שגויה, או stride לא נכון ל-load_matrix_sync ייצרו kernel מהיר ושגוי. הריצו checkResult אחרי כל שינוי, ואל תרשמו לעולם מספר ביצועים של kernel שלא אימת נכונות - זה הכלל הראשון של הפרויקט, והוא נשאר הכלל הראשון גם עכשיו כשאנחנו נוגעים ב-Tensor Cores.

לאן ממשיכים - what comes next

בפרויקט אחד עברנו מ-3% מה-peak ועד לשבירת תקרת ה-CUDA Cores כולה: פי 23 מ-naive עם kernel שכתבנו בעצמנו (2D blocktiling, 71% מה-peak), ואז פי 2.85 נוספים מעבר לתקרה הזו במעבר ל-Tensor Cores. וכל צעד לא היה ניחוש - הוא הונע ממטריקה ב-ncu שהצביעה על צוואר הבקבוק הבא, בדיוק כמו שתורת האילוצים מלמדת. זה הלב של הקורס כולו: ה-kernel, ה-warp, ה-coalescing, ה-shared memory, ה-registers, ה-occupancy, ה-arithmetic intensity, וה-Tensor Cores אינם מושגים נפרדים - הם מדרגות בסולם אחד, וכל אחת שווה קפיצת ביצועים מדידה.

מכאן שני כיוונים להעמקה. הראשון, לסגור את הפער ב-Tensor Cores: לעטוף את ה-wmma ב-staging של shared memory ו-double-buffering, ואז לעבור ל-CUTLASS או ל-CuTe DSL (פרק 6), שם ה-layouts וה-atoms מאפשרים להתקרב ל-peak האמיתי בלי לכתוב SASS ביד. השני, להכליל את הסולם למקרים אמיתיים: מטריצות לא-ריבועיות, N שאינו מתחלק בtiles (החזרת בדיקות הגבול), fusion של epilogue (הוספת bias והפעלת ReLU בתוך אותו kernel כדי לחסוך מעבר ל-global memory), ובחינה כיצד המסקנות משתנות על A100 (Ampere) מול B200 (Blackwell), שבהם ה-ridge point וה-peaks שונים. בכל אחד מהכיוונים, המתודולוגיה נשארת זהה: למדוד, למצוא את הצוואר, להרים אותו, ולחזור.