6.3 CUTLASS בניית GEMM בתבניות פתרון
פתרון - CUTLASS - בניית GEMM בתבניות¶
הערה: הזמנים, ה-TFLOPS והאחוזים מהשיא בפלטים למטה הם דוגמאות ממכונת H100 SXM עם CUDA 12.4 ו-CUTLASS 3.5, וישתנו אצלכם לפי הכרטיס, גרסת ה-toolkit, גרסת ה-CUTLASS ומצב המערכת. מה שלא ישתנה הוא המבנה והמגמות: ש-CUTLASS מגיע קרוב מאוד ל-cuBLAS על צורה סטנדרטית, ש-tile גדול מנצח על GEMM ריבועי גדול, שמיזוג epilogue חוסך מעבר זיכרון שלם, ושה-Tensor Cores ב-FP16 קופצים בסדר גודל מעל FP32 על CUDA Cores. כל הקטעים מתקמפלים עם הפקודה מה-### הכנה (-std=c++17 -arch=sm_90a והנתיבים ל-include). החליפו את הארכיטקטורה לפי הכרטיס שלכם.
פתרון תרגיל 1 - בניית GEMM ברמת ה-device והרצה¶
#include <iostream>
#include <cutlass/cutlass.h>
#include <cutlass/gemm/device/gemm.h>
#include <cutlass/util/host_tensor.h>
#include <cutlass/util/reference/host/tensor_fill.h>
#define CUTLASS_CHECK(status) \
do { \
cutlass::Status s_ = (status); \
if (s_ != cutlass::Status::kSuccess) { \
std::cerr << "CUTLASS error " << cutlassGetStatusString(s_) \
<< " at " << __FILE__ << ":" << __LINE__ << "\n"; \
std::exit(1); \
} \
} while (0)
using Gemm = 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>,
cutlass::epilogue::thread::LinearCombination<
cutlass::half_t, 128 / cutlass::sizeof_bits<cutlass::half_t>::value,
float, float>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
3>;
int main() {
int M = 4096, N = 4096, K = 4096;
float alpha = 1.0f, beta = 0.0f;
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});
cutlass::reference::host::TensorFillRandomUniform(
A.host_view(), 1, cutlass::half_t(1), cutlass::half_t(-1), 0);
cutlass::reference::host::TensorFillRandomUniform(
B.host_view(), 2, cutlass::half_t(1), cutlass::half_t(-1), 0);
A.sync_device();
B.sync_device();
Gemm gemm_op;
Gemm::Arguments args(
{M, N, K},
A.device_ref(), B.device_ref(), C.device_ref(), D.device_ref(),
{alpha, beta});
CUTLASS_CHECK(gemm_op(args));
cudaDeviceSynchronize();
D.sync_host();
std::cout << "OK. D[0][0]=" << float(D.host_view().at({0, 0}))
<< " D[100][100]=" << float(D.host_view().at({100, 100})) << "\n";
return 0;
}
הפלט הצפוי:
למה זה עבד: הטיפוס Gemm כולו נבנה בזמן קומפילציה מפרמטרי התבנית - הcompiler התמחה מהם kernel קונקרטי שמשתמש ב-Tensor Cores (מ-OpClassTensorOp) עם הוראת MMA בצורת 16x8x16. אובייקט ה-gemm_op הוא רק עטיפה ברמת ה-device שמקצה זיכרון-עזר ומשיקה את ה-grid; אין <<<...>>> כי CUTLASS מחביאה אותו. ה-HostTensor דואג לשני העותקים (הhost והdevice) ולסנכרון ביניהם, וה-device_ref() נושא את ה-leading dimension שנגזר מה-Layout - בדיוק הפרמטר שב-cuBLAS היינו מעבירים ביד. הערכים אינם אפס או NaN, מה שמאשר שה-mainloop רץ וצבר נכון.
איך להכליל: זהו השלד לכל GEMM ב-CUTLASS 2.x - מגדירים טיפוס Gemm מהפרמטרים, בונים Arguments, קוראים ל-gemm_op(args). כדי לשנות טיפוס נתונים (BF16, TF32, FP8), משנים את טיפוסי הקלט ואת ה-InstructionShape המתאים; כדי לשנות layout, משנים את פרמטר ה-Layout לכל אופרנד. אותו קוד מקור מייצר kernel אחר לגמרי לכל צירוף - זה הפולימורפיזם הפרמטרי בפעולה.
פתרון תרגיל 2 - אימות נכונות מול cuBLAS¶
#include <cublas_v2.h>
#include <cmath>
// ... define Gemm and run it as in exercise 1, the result is in D on the device ...
int main() {
int M = 4096, N = 4096, K = 4096;
float alpha = 1.0f, beta = 0.0f;
// ... allocation, filling, and running CUTLASS into D (device) as in exercise 1 ...
// cuBLAS into a separate reference buffer. row-major -> we use D^T = B^T * A^T
cutlass::HostTensor<cutlass::half_t, cutlass::layout::RowMajor> Dref({M, N});
cublasHandle_t h; cublasCreate(&h);
cublasGemmEx(h,
CUBLAS_OP_N, CUBLAS_OP_N,
N, M, K, // note: N then M (the result is swapped)
&alpha,
B.device_data(), CUDA_R_16F, N,
A.device_data(), CUDA_R_16F, K,
&beta,
Dref.device_data(), CUDA_R_16F, N,
CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
cudaDeviceSynchronize();
D.sync_host();
Dref.sync_host();
float max_rel = 0.0f;
for (int i = 0; i < M; i++)
for (int j = 0; j < N; j++) {
float a = float(D.host_view().at({i, j}));
float b = float(Dref.host_view().at({i, j}));
float rel = std::fabs(a - b) / (std::fabs(b) + 1e-3f);
if (rel > max_rel) max_rel = rel;
}
std::cout << "max relative error = " << max_rel
<< (max_rel < 1e-2f ? " -> PASS\n" : " -> FAIL\n");
cublasDestroy(h);
return 0;
}
הפלט הצפוי:
למה זה עבד: cuBLAS היא column-major (6.1), אבל ה-buffers שלנו row-major. הזהות D^T = B^T * A^T פותרת זאת בלי להעביר נתונים: buffer row-major זהה ביט-אחר-ביט ל-transpose שלו הנקרא column-major. לכן העברנו את B לפני A ואת הממדים בסדר N, M, K עם OP_N - וקיבלנו את אותו D. השגיאה אינה אפס מדויק כי שני המימושים צוברים ב-FP32 אבל בסדר פעולות שונה ומעגלים ל-FP16; שגיאה יחסית של ~0.3% היא בדיוק מה שמצפים מ-FP16, והיא הרבה מתחת לסף. זה מאשר שה-CUTLASS שבנינו מחשב את אותו דבר כמו הספרייה הסגורה.
איך להכליל: תמיד אמתו kernel חדש מול מקור אמין (cuBLAS או רפרנס על CPU) לפני שסומכים עליו. ב-FP16/BF16 השוו בשגיאה יחסית עם סף, לא בשוויון מדויק; ב-FP32 הסף יכול להיות הדוק בהרבה. תעלול ה-C^T=B^T A^T הוא הדרך הכללית לגשר בין קוד row-major לספרייה column-major בלי transpose פיזי, וכדאי לשמור אותו בארגז הכלים.
פתרון תרגיל 3 - כיוונון צורת ה-tile ומדידת ההשפעה¶
// define a few Gemm types that differ only in shapes, and benchmark each
template <typename GemmT>
float bench(int M, int N, int K,
cutlass::half_t* dA, cutlass::half_t* dB,
cutlass::half_t* dC, cutlass::half_t* dD) {
GemmT op;
typename GemmT::Arguments args({M, N, K},
{dA, K}, {dB, K}, {dC, N}, {dD, N}, {1.0f, 0.0f});
if (op.can_implement(args) != cutlass::Status::kSuccess) return -1.0f;
op.initialize(args);
cudaEvent_t s, e; cudaEventCreate(&s); cudaEventCreate(&e);
op(); // warmup
cudaDeviceSynchronize();
float best = 1e30f;
for (int r = 0; r < 20; r++) {
cudaEventRecord(s);
op();
cudaEventRecord(e); cudaEventSynchronize(e);
float ms = 0; cudaEventElapsedTime(&ms, s, e);
if (ms < best) best = ms;
}
cudaEventDestroy(s); cudaEventDestroy(e);
return best;
}
// three types; only the two GemmShape lines differ between them
template <typename TB, typename WP>
using G = 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,
TB, WP, cutlass::gemm::GemmShape<16, 8, 16>,
cutlass::epilogue::thread::LinearCombination<
cutlass::half_t, 8, float, float>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>;
using S = cutlass::gemm::GemmShape<0,0,0>; // just for readability
// ... in main, define:
// using G1 = G<GemmShape<64,64,32>, GemmShape<32,32,32>>;
// using G2 = G<GemmShape<128,128,32>,GemmShape<64,64,32>>;
// using G3 = G<GemmShape<256,128,32>,GemmShape<64,64,32>>;
// and print time, TFLOPS, and percent of peak for each.
הפלט הצפוי (H100 SXM, שיא FP16 ≈ 990 TFLOPS):
TB shape Warp shape time(ms) TFLOPS %peak warps smem(KiB)
64x64x32 32x32x32 0.412 333.6 33.7% 4 24
128x128x32 64x64x32 0.198 694.1 70.1% 4 48
256x128x32 64x64x32 0.171 803.8 81.2% 8 96
חשבון ה-warps וה-smem שהודפס, ליד כל שורה:
64x64: warps = (64/32)*(64/32) = 4 smem = 3*(64*32+32*64)*2 = 24 KiB
128x128: warps = (128/64)*(128/64) = 4 smem = 3*(128*32+32*128)*2 = 48 KiB
256x128: warps = (256/64)*(128/64) = 8 smem = 3*(256*32+32*128)*2 = 72 KiB
(שימו לב: 256x128 מדווח 96 KiB בפלט כי הוא כולל גם את שטח ה-epilogue וה-alignment; החשבון הגולמי של ה-mainloop לבדו נותן 72 KiB.)
למה זה עבד: ה-tile של 64x64 קטן מדי - כל block מחשב מעט פלט, ה-reuse של הנתונים ב-shared memory נמוך, וה-kernel מבזבז bandwidth על טעינות חוזרות; הגענו רק ל-34% מהשיא. הגדלה ל-128x128 מכפילה את ה-reuse (כל יסוד שנטען משמש יותר פעמים) וקופצת ל-70%. 256x128 ממשיך לשפר כי על צורה ריבועית גדולה כמו 4096^3 ה-reuse שולט, וה-8 warps ל-block מספיקים כדי להסתיר latency למרות שיש פחות blocks ל-SM. זה בדיוק האיזון מההרצאה: tile גדול = reuse טוב אבל occupancy נמוכה, ולצורה הזו ה-reuse מנצח.
איך להכליל: אין צורת tile אחת אופטימלית - היא תלויה בצורה, בטיפוס ובדור. סרקו כמה צורות ומדדו, בדיוק כמו כאן, או תנו ל-cutlass_profiler לסרוק אלפי וריאנטים ולדווח את המנצח. עבור GEMM "רזה" (M קטן, כמו decode של LLM) הכלל מתהפך: tile קטן וצר מפזר עבודה על יותר SM-ים ומנצח. השתמשו ב-can_implement כדי לדלג בבטחה על צירופים לא-חוקיים בזמן סריקה.
פתרון תרגיל 4 - מיזוג epilogue: scale ו-ReLU¶
#include <cutlass/epilogue/thread/linear_combination_relu.h>
using EpilogueReLU = cutlass::epilogue::thread::LinearCombinationRelu<
cutlass::half_t, 8, float, float>;
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<256, 128, 32>,
cutlass::gemm::GemmShape<64, 64, 32>,
cutlass::gemm::GemmShape<16, 8, 16>,
EpilogueReLU, // <<< only this changed
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>;
// run: alpha=0.5, beta=0 -> D = ReLU(0.5 * A*B)
GemmReLU op;
GemmReLU::Arguments args({M, N, K},
A.device_ref(), B.device_ref(), C.device_ref(), D.device_ref(),
{0.5f, 0.0f});
CUTLASS_CHECK(op(args));
cudaDeviceSynchronize();
D.sync_host();
// verification: two-step reference - regular GEMM into Dtmp, then ReLU(0.5*Dtmp) on the host
// (Dtmp comes from Gemm without an epilogue, from exercise 1)
float max_err = 0.0f;
for (int i = 0; i < M; i++)
for (int j = 0; j < N; j++) {
float raw = float(Dtmp.host_view().at({i, j})); // A*B
float expect = std::max(0.5f * raw, 0.0f); // ReLU(0.5*A*B)
float got = float(D.host_view().at({i, j}));
max_err = std::max(max_err, std::fabs(got - expect)
/ (std::fabs(expect) + 1e-3f));
}
std::cout << "fused vs two-step max rel err = " << max_err << "\n";
הפלט הצפוי:
fused vs two-step max rel err = 0.00291 -> PASS
fused (GEMM+ReLU one kernel): 0.174 ms
two-step (GEMM then ReLU): 0.198 + 0.061 = 0.259 ms
saved: 0.085 ms (~33%), and ~64 MiB of HBM traffic
למה זה עבד: החלפנו פרמטר תבנית יחיד - את ה-epilogue מ-LinearCombination ל-LinearCombinationRelu - וקיבלנו GEMM שממזג scale ו-ReLU לתוך אותו kernel. ה-mainloop לא השתנה כלל. תוצאת הצבירה נשארה ברגיסטרים ב-FP32, ה-0.5*acc וה-max(.,0) חלו שם, ורק ה-D הסופי נכתב פעם אחת. ה-two-step היה כותב מטריצת ביניים שלמה (32 MiB) ל-HBM וקורא אותה שוב ל-kernel ה-ReLU - שני מעברים שנעלמו. השגיאה מול הרפרנס דו-השלבי זניחה (רעש FP16), מה שמאשר שהמיזוג נכון. שימו לב שהחיסכון בזמן (~33%) מגיע כמעט כולו מהעלמת ה-ReLU הנפרד, שהיה memory-bound לחלוטין - ה-GEMM עצמו לא הואץ.
איך להכליל: זהו לב הכוח של CUTLASS - operation fusion שאתם מרכיבים בעצמכם. אותה טכניקה חלה על bias (LinearCombinationBias), GELU, sigmoid, clamp להמרת טיפוסים, ואפילו epilogue שכתבתם ידנית. הכלל: כל פעולה אלמנטית שרצה מיד אחרי GEMM היא מועמדת למיזוג - היא כמעט חינמית כשהתוצאה כבר ברגיסטרים, ויקרה כ-kernel נפרד כי היא memory-bound. זה בדיוק מה ש-cuDNN עושה פנימית (6.2), אבל כאן השליטה בידיכם.
פתרון תרגיל 5 - שלוש דרכים לאותו GEMM¶
// (a) the tiled kernel from 3.5 (FP32, CUDA Cores) - called as-is
// (b) cuBLAS FP32:
cublasSgemm(h, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K,
&alpha, dB, N, dA, K, &beta, dC, N);
// (c) CUTLASS FP32 on CUDA Cores (OpClassSimt), and separately FP16 on Tensor Cores
using GemmF32 = cutlass::gemm::device::Gemm<
float, cutlass::layout::RowMajor,
float, cutlass::layout::ColumnMajor,
float, cutlass::layout::RowMajor,
float, cutlass::arch::OpClassSimt, cutlass::arch::Sm80,
cutlass::gemm::GemmShape<128, 128, 8>,
cutlass::gemm::GemmShape<64, 64, 8>,
cutlass::gemm::GemmShape<1, 1, 1>,
cutlass::epilogue::thread::LinearCombination<float, 1, float, float>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 2>;
// each is measured with cudaEvent (minimum over 20 runs, after warmup), at M=N=K=2048
הפלט הצפוי (H100 SXM, M=N=K=2048; שיא FP32 ≈ 67 TFLOPS, שיא FP16-TC ≈ 990):
implementation time(ms) TFLOPS %peak
(a) your tiled (3.5, FP32) 3.91 4.39 6.6% (vs FP32 peak)
(b) cuBLAS SGEMM (FP32) 0.281 61.1 91.2% (vs FP32 peak)
(c) CUTLASS SIMT (FP32) 0.334 51.4 76.7% (vs FP32 peak)
(c') CUTLASS TensorOp (FP16) 0.026 661.0 66.8% (vs FP16 peak)
למה זה עבד: ה-kernel ה-tiled מ-3.5 מגיע רק ל-~6.6% מהשיא כי חסרים בו כל הטריקים שהופכים GEMM למהיר: אין double buffering (הוא נעצר בכל צעד לחכות לטעינה), אין pipeline רב-שלבי שמסתיר latency, ואין שימוש ב-Tensor Cores כלל. cuBLAS ו-CUTLASS ב-FP32 מגיעים ל-77%-91% כי יש בהם pipeline עמוק ו-double buffering. אבל הקפיצה האמיתית היא ב-(ג'): אותו GEMM ב-FP16 על Tensor Cores רץ פי ~11 מהר יותר מ-cuBLAS ב-FP32, כי החומרה עצמה מספקת פי ~15 throughput ל-FP16-TC (990 מול 67 TFLOPS, ראו 1.5). הצורה 2048^3 קטנה יחסית, ולכן ה-TensorOp מגיע "רק" ל-67% מהשיא שלו - על 4096^3 הוא היה מתקרב יותר.
איך להכליל: ההשוואה מלמדת שלושה לקחים. ראשית, kernel נאיבי שכתבתם ביד רחוק מאוד מספרייה מקצועית - הפער הוא ערך הpipeline וה-double buffering. שנית, על צורה סטנדרטית cuBLAS ו-CUTLASS קרובים מאוד, אז השתמשו ב-cuBLAS אם אין לכם צורך מיוחד. שלישית, בחירת הטיפוס (FP16-TC מול FP32) חשובה יותר מכל אופטימיזציה אחרת - הזיזה פי 11. בעבודה אמיתית: cuBLAS לצורה סטנדרטית, CUTLASS ל-epilogue מותאם או צורה חריגה, וה-kernel שלכם רק כתרגיל לימודי - אף פעם לא בפרודקשן.
פתרון תרגיל 6 (בונוס) - סקירת CUTLASS 3.x עם CollectiveBuilder¶
cd $CUTLASS/build
make 48_hopper_warp_specialized_gemm -j$(nproc)
./examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm \
--m=4096 --n=4096 --k=4096
הפלט הצפוי:
Disposition: Passed
Problem Size: 4096x4096x4096
Avg runtime: 0.153 ms
GFLOPS: 898216 (≈ 898 TFLOPS, ~90.7% of peak)
המרכיבים בקוד (מ-48_hopper_warp_specialized_gemm.cu):
TileShape = Shape<_128, _256, _64> // CTA tile in CuTe types
ClusterShape = Shape<_2, _1, _1> // cluster of thread blocks (Hopper)
CollectiveMainloop = CollectiveBuilder<
arch::Sm90, OpClassTensorOp, ...,
TileShape, ClusterShape,
StageCountAuto, KernelScheduleAuto> // builds WGMMA+TMA automatically
למה זה עבד: נתיב ה-3.x מגיע ל-~91% מהשיא, לעומת ~81% שקיבלנו מ-2.x בתרגיל 3, כי הוא משתמש בפיצ'רים הילידיים של Hopper: WGMMA שפועלת על warpgroup שלם (128 threads, 2.4) במקום על warp בודד, ו-TMA שמעתיק tiles שלמים בין global ל-shared בהוראה אחת בלי שה-threads יחשבו כתובות. ה-CollectiveBuilder בחר את פרטי ה-collective ואת ה-schedule אוטומטית. warp specialization מחלק את ה-warps ב-block לשני תפקידים - producer שטוען עם TMA ו-consumer שמחשב עם WGMMA - כך שהטעינה והחישוב חופפים בpipeline עמוק יותר מכל מה שאפשר על Ampere, ולכן ה-latency מוסתר טוב יותר וה-Tensor Cores כמעט לא מורעבים.
איך להכליל: זו הסיבה ש-NVIDIA ממליצה על CUTLASS דווקא ל-Transformers על Hopper ו-Blackwell (6.2): את ה-kernels הכי מהירים לחומרה הכי חדשה בונים ב-CUTLASS 3.x מעל CuTe, כי רק שם יש גישה מסודרת ל-WGMMA, TMA ו-warp specialization. ה-CollectiveBuilder נותן את רוב הרווח בלי לכוונן ידנית, וכשצריך שליטה מלאה יורדים לרמת ה-CuTe עצמה - את זה נלמד בשיעור 6.4. הכלל: על חומרה חדשה, התחילו מ-3.x; ה-2.x נשאר שימושי ל-Ampere ולהבנת המושגים.