1.5 Tensor Cores והאצת מטריצות פתרון
פתרון - Tensor Cores והאצת מטריצות¶
כאן נעבור על הפתרון המלא של כל תרגיל, עם הקוד, הפקודות והפלט המדויקים. שימו לב: פלטי ה-SASS, שמות ההוראות ומספרי ה-registers הם דוגמאות מהרצה על H100 (sm_90a) עם גרסת CUDA Toolkit מסוימת, והם ישתנו מעט בין ארכיטקטורות וגרסאות compiler - אבל המבנה, ההיגיון והמסקנות זהים. אם אתם על T4 ב-Colab, החליפו sm_90a ב-sm_75.
פתרון תרגיל 1 - GEMM מואץ דרך cuBLAS¶
הקובץ המלא gemm_tc.cu, כולל שני מאקרו בדיקת שגיאות - אחד ל-CUDA Runtime ואחד ל-cuBLAS:
#include <cstdio>
#include <cstdlib>
#include <cuda_fp16.h>
#include <cublas_v2.h>
#define CUDA_CHECK(call) \
do { \
cudaError_t err_ = (call); \
if (err_ != cudaSuccess) { \
fprintf(stderr, "CUDA error %s at %s:%d\n", \
cudaGetErrorString(err_), __FILE__, __LINE__); \
exit(EXIT_FAILURE); \
} \
} while (0)
#define CUBLAS_CHECK(call) \
do { \
cublasStatus_t st_ = (call); \
if (st_ != CUBLAS_STATUS_SUCCESS) { \
fprintf(stderr, "cuBLAS error %d at %s:%d\n", \
(int)st_, __FILE__, __LINE__); \
exit(EXIT_FAILURE); \
} \
} while (0)
int main() {
const int M = 1024, N = 1024, K = 1024;
// FP16 input on the host
half *hA = (half*)malloc(sizeof(half) * M * K);
half *hB = (half*)malloc(sizeof(half) * K * N);
float *hC = (float*)malloc(sizeof(float) * M * N);
for (int i = 0; i < M * K; i++) hA[i] = __float2half(1.0f);
for (int i = 0; i < K * N; i++) hB[i] = __float2half(1.0f);
// allocation on the device
half *dA, *dB;
float *dC;
CUDA_CHECK(cudaMalloc(&dA, sizeof(half) * M * K));
CUDA_CHECK(cudaMalloc(&dB, sizeof(half) * K * N));
CUDA_CHECK(cudaMalloc(&dC, sizeof(float) * M * N));
CUDA_CHECK(cudaMemcpy(dA, hA, sizeof(half) * M * K, cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(dB, hB, sizeof(half) * K * N, cudaMemcpyHostToDevice));
cublasHandle_t handle;
CUBLAS_CHECK(cublasCreate(&handle));
float alpha = 1.0f, beta = 0.0f; // for CUBLAS_COMPUTE_32F the scalars are float
// C = alpha * A*B + beta*C ; cuBLAS is column-major, the matrices are square
CUBLAS_CHECK(cublasGemmEx(
handle, CUBLAS_OP_N, CUBLAS_OP_N,
M, N, K,
&alpha,
dA, CUDA_R_16F, M, // lda
dB, CUDA_R_16F, K, // ldb
&beta,
dC, CUDA_R_32F, M, // ldc
CUBLAS_COMPUTE_32F,
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
CUDA_CHECK(cudaDeviceSynchronize());
CUDA_CHECK(cudaMemcpy(hC, dC, sizeof(float) * M * N, cudaMemcpyDeviceToHost));
printf("C[0] = %.1f (expected %d)\n", hC[0], K);
printf("C[M*N-1] = %.1f\n", hC[M * N - 1]);
CUBLAS_CHECK(cublasDestroy(handle));
CUDA_CHECK(cudaFree(dA)); CUDA_CHECK(cudaFree(dB)); CUDA_CHECK(cudaFree(dC));
free(hA); free(hB); free(hC);
return 0;
}
קומפילציה והרצה:
הפלט:
למה זה עבד: כל איבר של C הוא מכפלה סקלרית של שורה מ-A בעמודה מ-B, ומכיוון ששניהם מלאים ב-1.0, כל איבר הוא סכום של K = 1024 מכפלות 1*1, כלומר 1024. הקלטים ב-FP16 והצבירה ב-FP32 (זה מה ש-CUBLAS_COMPUTE_32F קובע), בדיוק כמו ההוראה HMMA...F32 שפירקנו בהרצאה. בחרנו algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP כדי לבקש במפורש את מסלול ליבות הטנזור. cuBLAS היא column-major, אבל כאן זה לא הזיק: המטריצות מרובעות וכל האיברים זהים, כך שפירוש הפוך של שורות/עמודות נותן אותה תוצאה.
איך להכליל: cublasGemmEx היא הדלת הסטנדרטית לליבות הטנזור לכל צירוף דיוק - FP16/BF16/TF32/INT8/FP8 בקלט מול צבירת FP32/INT32 - על ידי החלפת ה-cudaDataType ו-cublasComputeType_t. ברגע שהמטריצות מפסיקות להיות סימטריות, חובה לטפל ב-column-major: הטריק הנפוץ הוא לחשב C^T = B^T * A^T על ידי החלפת סדר האופרנדים, כך שהתוצאה יוצאת נכונה בזיכרון row-major של C. את זה, ואת cuBLASLt ו-CUTLASS לשליטה עדינה יותר, נראה בפרק 6.
פתרון תרגיל 2 - ציד הוראות HMMA ב-SASS¶
הקוד הוא wmma_min.cu כפי שניתן בתרגול. הקומפילציה ושליפת ה-SASS:
הפלט (מספרי האוגרים ישתנו אצלכם):
מה רואים כאן. ה-wmma::mma_sync הבודד על tile m16n16k16 התפרק לשתי הוראות HMMA.16816.F32 בקומפילציה ל-Hopper. כל HMMA.16816 היא tile חומרה m16 n8 k16 = 16*8*16 = 2,048 MAC; שתיים מהן מכסות 2 * 2,048 = 4,096 MAC, וזה בדיוק גודל ה-wmma הלוגי m16n16k16 (כי 16*16*16 = 4,096). על ארכיטקטורות ישנות יותר (למשל Turing / sm_75) אותו kernel מתפרק לארבע הוראות HMMA.1688.F32 (tile m16 n8 k8 = 1,024 MAC כל אחת, 4 * 1,024 = 4,096) - שתי צורות שונות של אותו פירוק, לפי מה שהחומרה של הדור מציעה.
אם תרצו לראות גם את מסגרת ה-wmma ברמת PTX, הריצו:
למה זה עבד: קימפלנו לארכיטקטורה עם ליבות טנזור (sm_90a), והשארנו את התוצאה נכתבת לזיכרון (store_matrix_sync), כך שהcompiler לא חיסל את החישוב כ-dead code. ה-grep -i hmma תפס את הוראות ליבת הטנזור בקוד המכונה. ראינו במו עינינו את מה שההרצאה טענה: הוראת wmma אחת ברמה הגבוהה אינה הוראת SASS אחת אלא מתפרקת לכמה הוראות HMMA קטנות.
איך להכליל: אותה טכניקה - קומפילציה ואז cuobjdump -sass ... | grep - היא הדרך לאמת איזה מנוע החומרה מפעילה בפועל עבור כל kernel: grep HMMA לליבות טנזור, grep FFMA לחישוב FP32 ב-CUDA Cores, grep -i wgmma להוראות ה-warpgroup האסינכרוניות של Hopper. כשתפרופלו ספריות כמו cuBLAS או CUTLASS בפרק 6, תראו את אותן הוראות HMMA בדיוק - רק מוקפות בהמון קוד תזמון ותנועת נתונים מתוחכם.
פתרון תרגיל 3 - חשבונאות MAC ויחס התפוקה¶
ההוראה HMMA16.16816.F32: tile m=16, n=8, k=16.
שלב 1 - MAC-ים להוראה:
שלב 2 - MAC-ים לכל thread (warp = 32 threads):
שלב 3 - FLOP-ים (MAC אחד = כפל + חיבור = 2 FLOP):
שלב 4 - יחס מול CUDA Core. הוראת FFMA בודדת של CUDA Core מבצעת MAC אחד לכל thread. לכן, לכל-הוראה-לכל-thread, ליבת הטנזור עושה פי 64 יותר אריתמטיקה. כלל האצבע הכולל בתפוקה (FLOPS) שה-glossary נוקב בו הוא פי ~100:
MAC/הוראה MAC/thread MAC/thread מול CUDA Core
CUDA Core (FFMA) 32 1 x1
Tensor Core (HMMA) 2,048 64 x64
כלל אצבע בתפוקה (FLOPS): פי ~100
למה זה עבד: כל המספרים נגזרים מנוסחה אחת - כפל מטריצות m x k ב-k x n דורש m*n*k פעולות MAC - ומהעובדה ש-warp הוא 32 threads. פי-64 האריתמטיקה-לכל-הוראה הוא המנוע שמאחורי פי ~100 התפוקה: הוראה אחת שמניעה 2,048 MAC פורשת את עלות שליפת-ופענוח ההוראה על פני המון חישוב, בעוד הוראת CUDA Core פורשת אותה על MAC בודד.
איך להכליל: היחס המדויק תלוי בזוג הדיוקים שמשווים. פי ~100 הוא כלל אצבע לסדר גודל; במספרים הרשמיים של H100 SXM, תפוקת FP16 עם צבירת FP32 היא כ-989 TFLOPS מול כ-66.9 TFLOPS ב-FP32 של ה-CUDA Cores - יחס של כ-15 עבור זוג הדיוקים הזה. ככל שיורדים בדיוק הקלט, היחס גדל דרמטית: FP8 ב-Hopper ו-FP4 ב-Blackwell דוחפים אותו הרבה מעבר לפי-100. הרעיון היציב הוא תמיד אותו רעיון: יותר אריתמטיקה לכל הוראה נשלפת.
פתרון תרגיל 4 - הוכחה שליבות הטנזור באמת עבדו¶
הרצת ה-GEMM תחת הפרופיילר עם מטריקת מונה הוראות ליבת הטנזור:
הפלט (מקוצר לשורות הרלוונטיות):
void cutlass::Kernel<...>(...) , 2024-...
Section: Command line profiler metrics
---------------------------------------------------------------------
sm__inst_executed_pipe_tensor.sum inst 1,048,576
---------------------------------------------------------------------
הערך גדול מאפס - הוכחה שהוראות עברו דרך pipe ליבת הטנזור. לשם השוואה, הרצת kernel של כפל מטריצות תמים על CUDA Cores תחת אותה מטריקה:
למה זה עבד: המטריקה sm__inst_executed_pipe_tensor.sum סופרת כמה הוראות בוצעו בpipe ליבת הטנזור על פני כל ה-SMs. ערך חיובי אומר חד-משמעית שהחומרה הפעילה את ליבות הטנזור; אפס אומר שהחישוב רץ כולו על ה-CUDA Cores. שימו לב ששם ה-kernel בפלט הוא kernel של CUTLASS - cuBLAS מממשת את ה-GEMM שלה מעל תבניות CUTLASS, בדיוק כפי שההרצאה תיארה: אתם קוראים לספרייה, והיא מנפיקה את הוראות ה-HMMA.
איך להכליל: זו טכניקת האימות הכללית ל"האם המנוע היקר באמת רץ". sm__pipe_tensor_cycles_active.avg.pct_of_peak_sustained_active נותנת את אחוז ניצול pipe ליבת הטנזור (כמה קרוב לשיא), ו-ncu --set full מציגה טבלת ניצול pipes מלאה שבה שורת "Tensor" גלויה. כשתבנו GEMM ידני בפרויקט ותשפרו אותו שלב-שלב עד לליבות טנזור, המטריקות האלה יראו לכם בדיוק מתי המעבר קרה ובכמה אתם רחוקים מהשיא.
פתרון תרגיל 5 - TMA מול זיכרון טנזור: הסבר מושגי¶
1. מהו ה-TMA? ה-Tensor Memory Accelerator הוא מנוע העתקה אסינכרוני ייעודי (Hopper ו-Blackwell) שמעביר נתונים מזיכרון ה-GPU / global memory אל ה-shared memory / cache L1, תוך עקיפה מלאה של קובץ האוגרים. הוא מחשב בחומרה כתובות affine (addr = width * base + offset) של גישות רב-ממדיות גדולות, ובכך מפנה גם אוגרים וגם CUDA Cores.
2. מהו זיכרון הטנזור? ה-Tensor Memory (TMEM) הוא מאגר אחסון ייעודי בתוך ה-SM, שקיים בדור Blackwell (למשל ה-B200). הוא מחזיק את הקלטים והפלטים של ליבות הטנזור, ובמיוחד את מטריצות הצבירה (accumulators) - האופרנד ה"חם" ביותר בכפל מטריצות, שנגיש בתדירות הגבוהה ביותר.
3. המלכוד. המשפט "ה-TMA מאיץ את זיכרון הטנזור" שגוי כי חרף השמות הדומים, ה-TMA אינו נוגע בזיכרון הטנזור כלל. ה-TMA טוען אך ורק אל ה-shared memory / cache L1. הנתונים מגיעים אל זיכרון הטנזור בדרך אחרת לגמרי - רק דרך ליבות הטנזור / הוראות MMA (או טעינה מ-shared memory ומאוגרים ברמת warpgroup). ה-TMA הוא מנוע תנועה אל ה-shared memory; זיכרון הטנזור הוא מקום אחסון שרק ליבות הטנזור ממלאות.
4. מסלול הנתונים:
HBM (GPU memory)
|
| <=== the TMA skips only this arrow (to shared memory)
v
shared memory / L1
|
| <=== via Tensor Cores / MMA instructions only (not the TMA!)
v
Tensor Memory (TMEM, Blackwell only)
למה זה עבד: ההבחנה נשענת על שתי שאלות: מה הרכיב (מנוע תנועה מול מקום אחסון) ולאן הוא נוגע (shared memory מול TMEM). ברגע שמפרידים אותן, השמות הדומים מפסיקים לבלבל. ה-TMA = מנוע -> shared memory; זיכרון טנזור = אחסון <- ליבות טנזור.
איך להכליל: אותה משמעת - "מה הרכיב ולאן הוא כותב" - פותרת את רוב מלכודי השמות בהיררכיית הזיכרון של ה-GPU: libcuda.so מול libcudart.so, cache L1 מול shared memory (אותו SRAM פיזי, שני שמות), ו-__global__ (מגדיר kernel) מול global memory (זיכרון). כשתגיעו ל-CUTLASS ול-CuTe בפרקים 6 עד 8, מסלול הנתונים HBM -> TMA -> shared -> ליבת טנזור -> TMEM יהיה בדיוק הpipe שתכתבו ותכווננו.
פתרון תרגיל 6 (בונוס) - מ-wmma ל-wgmma¶
1. למה Hopper הוסיפה wgmma? ה-wmma הרגיל שיתופי ברמת warp בודד (32 threads), והtile שלו קטן מכדי לרווֹת את רוחב הפס האריתמטי העצום של ליבות הטנזור בכרטיסי מרכזי-נתונים. ה-wgmma.mma_async פותר זאת בשני אופנים: (א) tile גדול יותר - הוא שיתופי ברמת warpgroup (128 threads), וtile גדול יותר ממלא את ליבת הטנזור ומונע רעב; (ב) אסינכרוניות - ההוראה משוגרת וממשיכה לרוץ ברקע בזמן שהthreads עושים עבודה אחרת, ובכך מבטלת את תקורת הסנכרון המפורש בין ארבעת ה-warps ומאפשרת לחפוף תנועת זיכרון עם חישוב.
2. הthreads ב-warpgroup:
3. MAC-ים בtile m64 n256 k16:
MACs = 64 * 256 * 16 = 16,384 * 16 = 262,144 MAC בהוראה אחת
יחס מול HMMA16.16816 (2,048 MAC):
262,144 / 2,048 = 128
הוראת ה-wgmma היחידה הזו מבצעת פי 128 יותר MAC-ים מהוראת ה-HMMA16.16816 הבודדת - סדר גודל שלם קדימה.
למה זה עבד: אותה נוסחה, m*n*k, רק על tile גדול בהרבה. הקפיצה מ-2,048 ל-262,144 MAC בהוראה אחת היא בדיוק הרעיון של ה-warpgroup: לרווֹת את ליבות הטנזור עם כמה שיותר עבודה לכל שליפת הוראה, ולעשות זאת אסינכרונית כדי לחפוף עם תנועת נתונים.
איך להכליל: המגמה בכל דור היא tiles גדולים יותר, אסינכרוניים יותר, וקרובים יותר לחומרה: Volta הביאה wmma (warp), Hopper הביאה wgmma (warpgroup) ואת ה-TMA להזנתו, ו-Blackwell הביאה את tcgen05.mma ואת זיכרון הטנזור. הקרנלים המהירים ביותר של ימינו - Flash Attention, ה-GEMM של CUTLASS - בנויים כולם סביב שרשור tiles של ה-warpgroup האסינכרוניים עם ה-TMA. את הבנייה המעשית של הpipe הזה נלמד בפרקים 6 עד 8.