לדלג לתוכן

1.5 Tensor Cores והאצת מטריצות תרגול

תרגול - Tensor Cores והאצת מטריצות

בתרגול הזה תפגשו את ליבות הטנזור מכמה כיוונים: קודם תפעילו כפל מטריצות מואץ דרך cuBLAS ותוודאו שהוא רץ, אחר כך תקמפלו kernel של WMMA ותצודו במו עיניכם את הוראות ה-HMMA בתוך ה-SASS, אז תפרקו הוראת ליבת טנזור בחישוב ידני של MAC-ים ותחשבו את יחס התפוקה מול CUDA Cores, ולבסוף תסבירו במילים שלכם את המלכוד המושגי בין ה-TMA לזיכרון הטנזור. עבדו לפי הסדר - כל תרגיל בונה על הקודם. הריצו על GPU אמיתי; ה-H100 הוא כרטיס הייחוס שלנו, אבל הכל ירוץ על כל כרטיס עם ליבות טנזור (מ-Volta והלאה, כולל T4 חינמי ב-Google Colab). כרטיסים ישנים ללא ליבות טנזור לא יריצו את תרגילים 1 ו-2.

הכנה

ודאו שיש לכם CUDA Toolkit מותקן (nvcc --version) ו-GPU עם ליבות טנזור. בררו את ה-compute capability:

nvidia-smi --query-gpu=name,compute_cap --format=csv

לאורך התרגול נקמפל ל-sm_90a (H100). החליפו בארכיטקטורה שלכם - למשל sm_75 ל-T4, sm_80 ל-A100, sm_89 ל-RTX 4090. ה-WMMA API מוגדר בכותרת <mma.h> ונמצא במרחב השמות nvcuda::wmma.

תרגיל 1 - GEMM מואץ דרך cuBLAS

צרו קובץ gemm_tc.cu המבצע C = A * B על מטריצות מרובעות בגודל M = N = K = 1024, עם קלט FP16 וצבירה FP32, דרך cublasGemmEx.

  1. הקצו על ה-device שלוש מטריצות: A ו-B מסוג half (FP16), ו-C מסוג float (FP32). מלאו את A ו-B בערכים פשוטים (למשל כולם 1.0).
  2. קראו ל-cublasGemmEx עם computeType = CUBLAS_COMPUTE_32F ו-algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP.
  3. העתיקו את C חזרה לhost ובדקו איבר אחד. אם A ו-B מלאים ב-1.0, כל איבר של C צריך לצאת K = 1024 (סכום של 1024 מכפלות 1*1).
  4. קמפלו עם קישור ל-cuBLAS והריצו:
nvcc -arch=sm_90a -o gemm_tc gemm_tc.cu -lcublas
./gemm_tc

רמז: cuBLAS היא column-major, אבל כשכל האיברים שווים ל-1.0 והמטריצות מרובעות, סדר השורות/עמודות לא משפיע על התוצאה - כל איבר יוצא K. עטפו כל קריאת cuBLAS ב-בדיקת שגיאה (הסטטוס cublasStatus_t צריך להיות CUBLAS_STATUS_SUCCESS). אל תשכחו cudaMalloc נפרד לכל מטריצה ו-cudaMemcpy להזנת הקלט.

תרגיל 2 - ציד הוראות HMMA ב-SASS

כתבו kernel מינימלי שמשתמש ב-WMMA API ישירות, קמפלו אותו, ומצאו את הוראות ליבת הטנזור בקוד המכונה.

  1. צרו wmma_min.cu עם kernel שמבצע wmma::mma_sync בודד על tile 16x16x16, קלט half וצבירה float. השיקו block אחד עם 32 threads (warp אחד):
#include <mma.h>
using namespace nvcuda;

__global__ void wmmaKernel(const half* A, const half* B, float* C) {
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;
    wmma::fill_fragment(c_frag, 0.0f);
    wmma::load_matrix_sync(a_frag, A, 16);
    wmma::load_matrix_sync(b_frag, B, 16);
    wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
    wmma::store_matrix_sync(C, c_frag, 16, wmma::mem_row_major);
}
  1. קמפלו: nvcc -arch=sm_90a -c wmma_min.cu -o wmma_min.o.
  2. שלפו את ה-SASS וסננו את הוראות ליבת הטנזור:
cuobjdump -sass wmma_min.o | grep -i hmma
  1. ספרו: כמה הוראות HMMA יש? מה השם המדויק שלהן (איזה קידוד tile)?

רמז: השתמשו ב-grep -i כי הרישיות משתנה. הוראת wmma::mma_sync אחת ב-m16n16k16 מתפרקת לכמה הוראות HMMA קטנות - זה בדיוק מה שראינו בהרצאה. אם grep hmma לא מחזיר כלום, ודאו שקימפלתם לארכיטקטורה עם ליבות טנזור (לא, למשל, sm_52) ושהשארתם את התוצאה של ה-fragment נכתבת לזיכרון (אחרת הcompiler עלול לחסל את החישוב כ-dead code).

תרגיל 3 - חשבונאות MAC ויחס התפוקה

כאן אין קוד - רק עיפרון, נייר ומחשבון. עבדו על ההוראה HMMA16.16816.F32 מההרצאה (tile m=16, n=8, k=16, קלט FP16, צבירה FP32).

  1. חשבו כמה פעולות כפל-וצבירה (MAC) מבצעת ההוראה הזו: m * n * k.
  2. חשבו כמה MAC-ים נופלים על כל thread ב-warp (חלקו ב-32).
  3. תרגמו את מספר ה-MAC-ים למספר FLOP-ים (רמז: MAC אחד = כפל + חיבור).
  4. כתבו את יחס התפוקה בין ליבת טנזור ל-CUDA Core לפי כלל האצבע של ה-glossary, והשוו: הוראת FFMA בודדת של CUDA Core מבצעת כמה MAC לכל thread? כמה פעמים יותר אריתמטיקה-לכל-הוראה יש בליבת הטנזור?

רמז: 16 * 8 * 16 - חשבו את זה בשלבים (16 * 8 = 128, ואז 128 * 16). לכל thread: חלקו את התוצאה ב-32. הוראת FFMA של CUDA Core היא MAC אחד לכל thread, אז היחס לכל-הוראה-לכל-thread הוא בדיוק מספר ה-MAC-לכל-thread שקיבלתם. כלל האצבע הכולל בתפוקה (FLOPS) הוא פי ~100.

תרגיל 4 - הוכחה שליבות הטנזור באמת עבדו

הריצו את gemm_tc מתרגיל 1 תחת הפרופיילר ואַמתו שהמנוע שרץ בפועל הוא ליבת הטנזור ולא ה-CUDA Cores.

  1. הריצו את ה-GEMM תחת ncu (Nsight Compute) ובקשו את מטריקת מונה הוראות ליבת הטנזור:
ncu --metrics sm__inst_executed_pipe_tensor.sum ./gemm_tc
  1. קִראו את הערך: אם הוא גדול מאפס, הוראות עברו דרך pipe ליבת הטנזור.
  2. לשם השוואה, הריצו kernel של כפל מטריצות רגיל שכתבתם (למשל מפרק קודם, על CUDA Cores בלבד) תחת אותה מטריקה, וראו שהערך שם הוא אפס.

רמז: אם ncu דורש הרשאות, הריצו עם sudo או הגדירו NVIDIA_DEV_ALLOW_PROFILING. אם המטריקה לא קיימת בגרסת ה-ncu שלכם, נסו ncu --set full ./gemm_tc וחפשו בפלט את שורת ה-"Tensor" תחת ניצול הpipes (Pipe Utilization). ערך חיובי ב-sm__inst_executed_pipe_tensor.sum הוא ההוכחה החד-משמעית שליבות הטנזור הופעלו.

תרגיל 5 - TMA מול זיכרון טנזור: הסבר מושגי

בלי קוד. כתבו תשובה קצרה וברורה (כמה משפטים לכל סעיף) שמראה שהבנתם את ההבחנה שההרצאה הזהירה מפניה.

  1. מהו ה-Tensor Memory Accelerator (TMA)? לאן הוא מעביר נתונים, ומאיפה?
  2. מהו זיכרון הטנזור (Tensor Memory / TMEM)? באיזה דור הוא קיים, ומה הוא מחזיק בעיקר?
  3. הסבירו את המלכוד: מדוע המשפט "ה-TMA מאיץ את זיכרון הטנזור" שגוי? איך הנתונים באמת מגיעים לזיכרון הטנזור?
  4. שרטטו (בטקסט) את מסלול הנתונים מ-HBM ועד זיכרון הטנזור, וסמנו איזה חץ ה-TMA "קופץ".

רמז: ה-TMA מעביר מ-global memory / זיכרון ה-GPU אל ה-shared memory / cache L1, תוך עקיפת האוגרים. זיכרון הטנזור הוא מאגר על-SM נפרד (Blackwell / B200) שמחזיק בעיקר מצברים. השם משותף אבל הדברים שונים: ה-TMA לעולם אינו כותב ישירות ל-TMEM - הנתונים מגיעים לשם רק דרך ליבות הטנזור / הוראות MMA.

תרגיל 6 (בונוס) - מ-wmma ל-wgmma

חקרו את ההבדל בין הוראת ליבת הטנזור ברמת ה-warp לבין זו שברמת ה-warpgroup.

  1. הסבירו במילים שלכם: מדוע Hopper הוסיפה את ה-wgmma.mma_async על גבי ה-wmma הקיים? מהם שני היתרונות (גודל tile, אסינכרוניות)?
  2. warp אחד הוא 32 threads. כמה threads יש ב-warpgroup, וכמה warps הוא מכיל?
  3. חשבו: אם הוראת wgmma מבצעת tile בגודל m64 n256 k16 (דוגמה טיפוסית ל-Hopper), כמה MAC-ים זה בהוראה אחת? השוו למספר ה-MAC-ים של ה-HMMA16.16816 מתרגיל 3.

רמז: warpgroup = ארבעה warps = 4 * 32 = 128 threads. עבור m64 n256 k16 חשבו 64 * 256 * 16. תגלו שהtile של ה-warpgroup גדול בסדרי גודל, וזה בדיוק העניין: tiles גדולים יותר רווים את ליבות הטנזור. את השימוש המעשי בכל זה (Flash Attention, CUTLASS) נראה בפרקים 6 עד 8.