לדלג לתוכן

6.3 CUTLASS בניית GEMM בתבניות תרגול

תרגול - CUTLASS - בניית GEMM בתבניות

בתרגול הזה תבנו GEMM אמיתי מתבניות CUTLASS, תריצו אותו על Tensor Cores, ותרגישו על בשרכם את ההבדל בין "לקרוא ל-kernel" (cuBLAS) ל"לבנות kernel" (CUTLASS). תתחילו מהורדת ה-repo ובניית דוגמה, תאמתו נכונות מול cuBLAS, תכווננו את צורת ה-tile ותמדדו את ההשפעה, תמזגו epilogue של scale ו-ReLU ותאמתו שהמיזוג נכון, ולבסוף תשוו את שלושת העולמות: CUTLASS, cuBLAS, וה-kernel ה-tiled שכתבתם בעצמכם בפרויקט 3.5. עבדו לפי הסדר - כל תרגיל בונה על קודמו. הריצו הכל על GPU אמיתי; ה-H100 הוא כרטיס הייחוס, אבל כל תרגיל כאן ירוץ על כל GPU מ-Ampere ומעלה (Turing עם התאמות קטנות) - רק המספרים המוחלטים ישתנו. אם אין לכם כרטיס, השתמשו במכונת ענן עם GPU, כפי שראינו בשיעור 0.2 (שימו לב ש-T4 החינמי של Colab הוא Turing/sm_75, ותצטרכו להתאים את ה-InstructionShape).

הכנה

הורידו את CUTLASS ובנו את הפרופיילר ואת הדוגמאות. זו הכנה חד-פעמית שתשמש את כל התרגילים:

git clone --depth 1 https://github.com/NVIDIA/cutlass ~/cutlass
export CUTLASS=~/cutlass
cd $CUTLASS && mkdir -p build && cd build
cmake .. -DCUTLASS_NVCC_ARCHS=90a          # 90a = H100; use 80 for A100, 75 for T4
make cutlass_profiler -j$(nproc)           # takes a few minutes

לאורך התרגול נקמפל קבצי .cu משלנו כך (CUTLASS היא header-only, אין לינקוג' אליה; רק ל-cuBLAS צריך ‎-lcublas):

nvcc -O3 -std=c++17 -arch=sm_90a \
     -I$CUTLASS/include -I$CUTLASS/tools/util/include \
     prog.cu -o prog -lcublas

מקרו הבדיקה שנשתמש בו לאורך התרגול (זהה לזה שראינו בכל הקורס), וכן עוטף סטטוס ל-CUTLASS:

#include <cstdio>
#include <cstdlib>
#include <cuda_runtime.h>

#define CUDA_CHECK(call)                                                      \
    do {                                                                      \
        cudaError_t err_ = (call);                                            \
        if (err_ != cudaSuccess) {                                            \
            fprintf(stderr, "CUDA error at %s:%d\n  '%s'\n  -> %s\n",         \
                    __FILE__, __LINE__, #call, cudaGetErrorString(err_));     \
            exit(EXIT_FAILURE);                                               \
        }                                                                     \
    } while (0)

#define CUTLASS_CHECK(status)                                                 \
    do {                                                                      \
        cutlass::Status s_ = (status);                                        \
        if (s_ != cutlass::Status::kSuccess) {                                \
            fprintf(stderr, "CUTLASS error at %s:%d -> %s\n",                 \
                    __FILE__, __LINE__, cutlassGetStatusString(s_));          \
            exit(EXIT_FAILURE);                                               \
        }                                                                     \
    } while (0)

תרגיל 1 - בניית GEMM ברמת ה-device והרצה

  1. צרו קובץ cutlass_gemm.cu. הגדירו טיפוס Gemm עם cutlass::gemm::device::Gemm, בטיפוסי FP16 לקלט (A שורות, B עמודות), פלט FP16 שורות, צבירת FP32, OpClassTensorOp, arch::Sm80, ThreadblockShape <128,128,32>, WarpShape <64,64,32>, InstructionShape <16,8,16>, epilogue LinearCombination, ו-3 stages (בדיוק כמו בהרצאה).
  2. בhost, הקצו cutlass::HostTensor ל-A, B, C, D בגודל 4096x4096. מלאו את A ו-B בערכים אקראיים קטנים (למשל בעזרת cutlass::reference::host::TensorFillRandomUniform או ידנית), וקראו ל-sync_device().
  3. בנו Gemm::Arguments עם גודל הבעיה {M,N,K}, ה-device_ref() של כל טנזור, ו-{alpha=1, beta=0}. הריצו gemm_op(args) ועטפו ב-CUTLASS_CHECK.
  4. סנכרנו, קראו D.sync_host(), והדפיסו כמה יסודות מ-D כדי לוודא שהם אינם אפס או NaN.
  5. הריצו והדפיסו הודעת הצלחה.

רמז: אל תשכחו את -I$CUTLASS/tools/util/include - שם יושבים ה-host_tensor.h וכלי המילוי. אם ה-cutlass::Status חוזר kErrorNotSupported, כנראה שצירוף הצורות/הארכיטקטורה לא נתמך; ודאו ש-InstructionShape תואם לטיפוס (‎16x8x16 ל-FP16 על Ampere). ה-beta=0 אומר ש-C לא נקרא, אז אין צורך לאתחל אותו לערך משמעותי.

תרגיל 2 - אימות נכונות מול cuBLAS

  1. הוסיפו לתוכנית קריאה ל-cuBLAS שמחשבת את אותו GEMM. זכרו מ-6.1: cuBLAS היא column-major. כדי לחשב D = A*B כש-A,B,D שמורים row-major, השתמשו בזהות D^T = B^T * A^T והזינו את המטריצות בסדר מוחלף (B ואז A) עם CUBLAS_OP_N, כך שהתוצאה ב-column-major זהה ל-D ה-row-major שלכם.
  2. הריצו את cuBLAS ל-buffer נפרד D_ref. השתמשו ב-cublasGemmEx עם CUDA_R_16F לקלט ו-CUDA_R_32F ל-compute type (צבירה ב-FP32, בדיוק כמו CUTLASS).
  3. השוו יסוד-יסוד את D של CUTLASS מול D_ref של cuBLAS. כי מדובר ב-FP16, אל תדרשו שוויון מדויק - חשבו את השגיאה היחסית המקסימלית ודרשו שהיא קטנה מסף (למשל 1e-2).
  4. הדפיסו את השגיאה המקסימלית ואת המסקנה (עבר / נכשל).

רמז: הזהות C^T = B^T A^T עובדת כי buffer שמור row-major זהה ביט-אחר-ביט ל-transpose שלו הנקרא column-major (בדיוק התעלול מ-6.1). בפועל: cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_16F, N, A, CUDA_R_16F, K, &beta, Dref, CUDA_R_16F, N, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT). שימו לב לסדר: N ואז M, ו-B לפני A. אל תשכחו -lcublas.

תרגיל 3 - כיוונון צורת ה-tile ומדידת ההשפעה

  1. הפכו את הצורות לפרמטר: הגדירו כמה טיפוסי Gemm שנבדלים רק ב-ThreadblockShape וב-WarpShape. נסו לפחות שלושה צירופים: <64,64,32>/<32,32,32>, <128,128,32>/<64,64,32>, ו-<256,128,32>/<64,64,32>.
  2. עטפו כל הרצה במדידת cudaEvent (זמן GPU, אחרי חימום, מינימום על 20 חזרות), כפי שעשינו ב-5.5.
  3. חשבו לכל וריאנט את ה-TFLOPS: 2*M*N*K / (זמן_בשניות) / 1e12. דווחו טבלה: צורה, זמן (ms), TFLOPS, אחוז מהשיא (שיא FP16 של H100 ≈ 990 TFLOPS).
  4. חשבו ידנית לכל וריאנט את מספר ה-warps ל-block ואת ה-shared memory ל-block (3 stages), והסבירו את הקשר בין הצורה לביצועים שמדדתם.
  5. בונוס: הריצו cutlass_profiler --operation=Gemm --m=4096 --n=4096 --k=4096 --A=f16:row --B=f16:col --C=f16 --accum=f32 והשוו את הווריאנט המנצח שלו לצורות שבחרתם ידנית.

רמז: מספר ה-warps ל-block = (TB_M/Warp_M) * (TB_N/Warp_N). shared memory ל-stage ב-FP16 = (TB_M*TB_K + TB_K*TB_N) * 2 בתים. tile גדול יותר נותן reuse טוב יותר אבל פחות blocks ל-SM (occupancy יורדת); לצורה ריבועית גדולה כמו 4096^3 ה-reuse בדרך כלל שולט, אז tile גדול ינצח. אם וריאנט מחזיר kErrorNotSupported, דלגו עליו - לא כל צירוף חוקי.

תרגיל 4 - מיזוג epilogue: scale ו-ReLU

  1. שנו את ה-Gemm המנצח מתרגיל 3 כך שה-epilogue יהיה cutlass::epilogue::thread::LinearCombinationRelu במקום LinearCombination (הכניסו #include <cutlass/epilogue/thread/linear_combination_relu.h>).
  2. הריצו עם alpha=0.5, beta=0 על אותם A,B. עכשיו התוצאה הצפויה היא D = ReLU(0.5 * A*B).
  3. אמתו את המיזוג: חשבו רפרנס בשני שלבים נפרדים - GEMM רגיל (CUTLASS או cuBLAS) ל-D_tmp, ואז kernel נפרד או קוד host שמחשב ReLU(0.5*D_tmp). השוו יסוד-יסוד ל-D הממוזג.
  4. מדדו את הזמן של הגרסה הממוזגת מול הזמן של שני השלבים הנפרדים (GEMM + kernel ReLU נפרד). דווחו את החיסכון.
  5. חשבו כמה תעבורת global memory נחסכה בזכות המיזוג (רמז: מטריצה שלמה שלא נכתבה ולא נקראה שוב).

רמז: LinearCombinationRelu מקבל את אותם {alpha, beta} ומחיל ReLU עם threshold=0 כברירת מחדל אחרי alpha*acc + beta*C. החיסכון בזיכרון: מטריצת 4096x4096 ב-FP16 = 32 MiB; המיזוג חוסך כתיבה של D_tmp וקריאה שלו = כ-64 MiB פחות תעבורת HBM. שימו לב שהחיסכון בזמן יהיה קטן יחסית לחלק ה-GEMM (שהוא compute-bound), אבל ה-ReLU הנפרד היה memory-bound לחלוטין - וכולו נעלם.

תרגיל 5 - שלוש דרכים לאותו GEMM: CUTLASS מול cuBLAS מול tiled שלכם

  1. קחו את ה-kernel ה-tiled עם shared memory שכתבתם בפרויקט 3.5 (הגרסה ב-FP32, על CUDA Cores). הריצו אותו על אותה בעיה (התאימו את הגודל אם צריך, למשל 2048^3 כדי שירוץ בזמן סביר).
  2. הריצו על אותה בעיה בדיוק את שלושת המימושים: (א) ה-kernel ה-tiled שלכם, (ב) cuBLAS (cublasSgemm ל-FP32), (ג) CUTLASS (FP32 על CUDA Cores עם OpClassSimt, או FP16 על Tensor Cores - ציינו מה מדדתם).
  3. מדדו את שלושתם עם cudaEvent, חשבו TFLOPS ואחוז מהשיא לכל אחד, ודווחו טבלה.
  4. נתחו: מדוע ה-kernel שלכם רחוק מ-cuBLAS/CUTLASS? מה חסר בו (double buffering, pipeline רב-שלבי, Tensor Cores)? מדוע CUTLASS ב-FP16/Tensor Core מנצח את שניהם ב-FP32?
  5. כתבו פסקה: באיזה מצב הייתם בוחרים בכל אחד מהשלושה בעבודה אמיתית?

רמז: השוואה הוגנת דורשת אותו טיפוס נתונים ואותה צורה. ה-kernel שלכם מ-3.5 הוא ב-FP32 על CUDA Cores, ולכן השוו אותו ל-cuBLAS/CUTLASS ב-FP32 (שיא FP32 של H100 ≈ 67 TFLOPS, ראו 1.5), ובנפרד הראו כמה Tensor Cores ב-FP16 קופצים (שיא ≈ 990 TFLOPS, פי ~15). הפער בין ה-kernel שלכם ל-cuBLAS ב-FP32 מגלה כמה שווה הpipeline הרב-שלבי וה-double buffering שאין לכם.

תרגיל 6 (בונוס) - סקירת CUTLASS 3.x עם CollectiveBuilder

  1. פתחו את הדוגמה examples/48_hopper_warp_specialized_gemm (או דומה) ב-repo של CUTLASS, וקמפלו אותה מתוך תיקיית ה-build (make 48_hopper_warp_specialized_gemm).
  2. אתרו בקוד את שלושת המרכיבים מההרצאה: את ה-TileShape וה-ClusterShape (מטיפוסי CuTe Shape<_128,_256,_64>), את ה-CollectiveBuilder שבונה את ה-mainloop, ואת בחירת ה-KernelScheduleAuto.
  3. הריצו את הדוגמה על ה-H100 שלכם ורשמו את ה-TFLOPS שהיא מדווחת. השוו לתרגיל 3 (ה-API הקלאסי 2.x): כמה יותר קרוב לשיא מגיע נתיב ה-WGMMA/TMA של Hopper?
  4. הסבירו במילים: מהו warp specialization, ומדוע הוא מסתיר latency טוב יותר על Hopper מאשר הpipeline של Ampere שראינו ב-2.x?
  5. קשרו להרצאה: מדוע NVIDIA ממליצה דווקא על CUTLASS (ולא cuDNN) ל-Transformers על Hopper/Blackwell?

רמז: אם הדוגמה לא מתקמפלת, ודאו ש--DCUTLASS_NVCC_ARCHS=90a (עם ה-a) הועבר ל-cmake - נתיב ה-WGMMA/TMA דורש את הפיצ'רים של sm_90a ולא יתקמפל ל-sm_90 רגיל. warp specialization מחלק את ה-warps ב-block ל-producer (טעינה עם TMA) ו-consumer (חישוב עם WGMMA), כך שהטעינה והחישוב חופפים בpipeline עמוק. את הפירוט המלא של CuTe ו-layouts נלמד בשיעור 6.4.