9.4 פרויקט reduction, scan ו softmax פרויקט
פרויקט - reduction, scan ו-softmax¶
זהו הפרויקט המסכם השני של הקורס, והוא בן-הזוג המשלים של פרויקט ה-GEMM (9.3). שם עסקנו בבעיה compute-bound (compute-bound) שכל עניינה arithmetic intensity ו-Tensor Cores; כאן נעסוק בשלוש אבני-הבניין המקבילות (parallel primitives) של כמעט כל kernel אמיתי - reduction (צמצום), scan (סכום-רץ / prefix sum) ו-softmax - וכולן memory-bound (memory-bound) או שגר קרוב לכך. הבעיה המרכזית כאן איננה לדחוס עוד FLOPs, אלא לתאם 32, 256 או אלפי threads כך שיצרכו את רוחב הפס של ה-HBM (זיכרון ה-GPU) עד תום, בלי warp divergence (התפצלות ה-warp, פרק 8.3), בלי bank conflicts (התנגשות בנקים, פרק 8.4) ותוך coalescing (איחוד גישות) מלא. נבנה כל primitive בשלבים, נמדוד בכל שלב את רוחב הפס המושג כאחוז מה-peak של ה-HBM, ונסביר כל קפיצה במושג שלמדנו: SIMT (פרק 2), shared memory (פרק 3), warp primitives כמו __shfl_down_sync, וה-SFU (יחידת הפונקציות המיוחדות) עבור exp. בסוף softmax נגיע לניסוח ה-online/streaming שהוא לב-ליבו של FlashAttention, ונבין למה בלעדיו לא הייתה קיימת attention יעילה בזיכרון.
מה נבנה - the three primitives¶
הפרויקט מחולק לשלושה חלקים עצמאיים, כל אחד עם סולם (ladder) משלו וקריטריוני קבלה משלו. את חלק A (reduction) בונים במלואו כי הוא מלמד את כל טכניקות התיאום; חלקים B ו-C נשענים עליו.
part primitive the ladder the key metric
────────────────────────────────────────────────────────────────────────────
A reduction 0 naive divergent -> 1 non-divergent -> % of HBM peak
(sum) 2 sequential addressing -> 3 first-add- (memory-bound)
on-load -> 4 warp-shuffle -> 5 multi-block
B scan Hillis-Steele within block -> composition correctness + BW
(prefix sum) across blocks (block sums + scan + addition)
C softmax stable (max, exp, sum, normalize) -> correctness (rtol)
online/streaming (FlashAttention) + single pass
────────────────────────────────────────────────────────────────────────────
ה-GPU הרץ לדוגמה הוא H100 (SXM), כרגיל בקורס. עבור reduction נעבוד עם וקטור של N = 2^28 = 268,435,456 איברי float (בדיוק 1 GiB, 1,073,741,824 בתים) - גדול מספיק כדי שרוחב הפס, ולא הoverhead של הlaunch, ישלוט בזמן.
הרקע המספרי - why bandwidth, not FLOPs¶
ב-GEMM מדדנו GFLOP/s כי הבעיה compute-bound. reduction הפוכה לחלוטין: לסכום N איברים דורש N-1 חיבורים בלבד - פעולת חישוב אחת לכל איבר - אבל קוראים N איברים מ-global memory. הarithmetic intensity היא בערך 1/4 FLOP לבית, נמוכה בסדרי גודל מנקודת ה-ridge של ה-H100. לכן התקרה איננה חישובית אלא רוחב הפס של ה-HBM: כמה מהר אפשר לשפוך את הווקטור דרך ה-SMs. ב-H100 SXM עם HBM3 רוחב הפס הוא כ-3.35 TB/s = 3350 GB/s, וזו התקרה שכל מדידה תימדד מולה.
הגבול התחתון לזמן הוא זמן הקריאה של הווקטור פעם אחת:
bytes = N * sizeof(float) = 2^28 * 4 = 1,073,741,824 B ≈ 1.0737 GB
t_min (peak HBM) = bytes / 3.35e12 ≈ 0.320 ms
BW_achieved = bytes / (time_in_seconds) [GB/s]
% HBM peak = 100 * BW_achieved / 3350
כל reduction נכון קורא כל איבר בדיוק פעם אחת, ולכן bytes קבוע בכל השלבים - בדיוק כמו ש-2N^3 היה קבוע ב-GEMM. מה שמשתנה הוא כמה קרוב לתקרת ה-3350 GB/s אנחנו מגיעים. reduction מכוונן היטב מגיע ל-85%-95% מה-peak; זו המטרה של החלק.
תשתית ההתחלה - starter scaffolding¶
שמרו את השלד כ-primitives.cu. הוא כולל את מקרו ה-CUDA_CHECK, מד-זמן מבוסס cudaEvent (זהה לזה מ-9.3), ומחשבון רוחב-פס ואחוז-peak. השלימו את ה-kernels שלב אחר שלב.
#include <cstdio>
#include <cstdlib>
#include <cmath>
#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)
// theoretical bandwidth of HBM3 on H100 SXM, in GB/s
static const double H100_HBM_PEAK_GBPS = 3350.0;
// ------- cudaEvent-based timer: warmup + reps, average time per kernel -------
template <typename LaunchFn>
float timeKernel(LaunchFn launch, int reps = 50) {
cudaEvent_t start, stop;
CUDA_CHECK(cudaEventCreate(&start));
CUDA_CHECK(cudaEventCreate(&stop));
launch(); // warmup: takes JIT/caches out of the picture
CUDA_CHECK(cudaGetLastError());
CUDA_CHECK(cudaDeviceSynchronize());
CUDA_CHECK(cudaEventRecord(start));
for (int r = 0; r < reps; r++) launch();
CUDA_CHECK(cudaEventRecord(stop));
CUDA_CHECK(cudaEventSynchronize(stop));
float ms = 0.0f;
CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
CUDA_CHECK(cudaEventDestroy(start));
CUDA_CHECK(cudaEventDestroy(stop));
return ms / reps;
}
// ------- bandwidth report for reduction (each element read exactly once) -------
void reportBW(const char* name, size_t bytes, float ms) {
double gbps = (double)bytes / (ms / 1000.0) / 1e9;
double pct = 100.0 * gbps / H100_HBM_PEAK_GBPS;
printf(" %-22s %8.4f ms %8.1f GB/s %5.1f%% HBM\n", name, ms, gbps, pct);
}
שתי הערות שחשוב להפנים לפני שממשיכים:
- reduction היא memory-bound, ולכן המדד הוא GB/s ולא GFLOP/s. אם תדווחו GFLOP/s תקבלו מספר קטן ומטעה שאינו אומר דבר על איכות ה-kernel. המדד הנכון הוא: כמה מרוחב הפס של ה-HBM ניצלתם.
- צמצום דורש ערך התחלה נייטרלי (identity). עבור סכום זהו
0.0f, עבור max זהו-INFINITY. הthreads שנופלים מחוץ לתחום (i >= n) חייבים לתרום את ה-identity, אחרת יזהמו את התוצאה.
חלק A, שלב 0 - naive interleaved עם divergence¶
הגישה הקלאסית: כל block טוען את מנתו ל-shared memory, ואז מצמצם בעץ בינארי עם "כתובות משתלבות" (interleaved addressing). הגרסה הראשונה, התמימה, מבצעת את החיבור בכל צעד רק בthreads ש-tid שלהם מתחלק ב-2s:
__global__ void reduce0(const float* in, float* out, int n) {
extern __shared__ float sdata[];
unsigned tid = threadIdx.x;
unsigned i = blockIdx.x * blockDim.x + tid;
sdata[tid] = (i < n) ? in[i] : 0.0f; // identity for threads out of range
__syncthreads();
for (unsigned s = 1; s < blockDim.x; s *= 2) {
if (tid % (2 * s) == 0) // <-- source of the divergence
sdata[tid] += sdata[tid + s];
__syncthreads();
}
if (tid == 0) out[blockIdx.x] = sdata[0]; // one partial sum per block
}
הבעיה היא בדיוק ב-tid % (2*s) == 0. בתוך warp (קבוצת 32 threads) התנאי הזה מפצל את הthreads: בצעד הראשון רק הthreads הזוגיים עובדים, בשני רק כל רביעי, וכו'. זו warp divergence קלאסית (פרק 8.3) - החומרה מריצה את ה-warp פעמיים, פעם למסלול העובד ופעם לריק, ומחצית מ-issue slots מתבזבזת. בנוסף, האופרטור % על ה-CUDA Cores יקר יחסית.
קריטריון קבלה: הסכום שווה לייחוס על ה-CPU (בסובלנות rtol=1e-4, כי FP32 צובר שגיאת עיגול), ורשום GB/s ו-%HBM. אין דרישת ביצועים; זהו הבסיס.
בערך 18% מרוחב הפס - הרבה מתחת לתקרה. הבא נעיף את ה-divergence.
חלק A, שלב 1 - non-divergent interleaved¶
התיקון: במקום לבדוק tid % (2s), נחשב אינדקס רציף index = 2*s*tid ונבדוק אם הוא בתחום. עכשיו הthreads הפעילים בכל צעד הם ה-blockDim.x/(2s) הראשונים - כלומר warps שלמים עובדים או warps שלמים לא עובדים, בלי פיצול בתוך ה-warp:
for (unsigned s = 1; s < blockDim.x; s *= 2) {
unsigned index = 2 * s * tid;
if (index < blockDim.x) // whole warps, no divergence
sdata[index] += sdata[index + s];
__syncthreads();
}
הסרנו את ה-divergence, אבל יצרנו בעיה חדשה: הגישה sdata[index] עם index = 2*s*tid היא מדורגת (strided) ב-shared memory. כבר בצעד השני, threads מרובים נופלים על אותו bank מבין 32 הבנקים - bank conflict (פרק 8.4) - והגישות מסתדרות טורית. כך שהרווחנו על החישוב אבל נחסמנו על ה-shared memory.
פי 1.6 מהר יותר. עדיין רחוק; מקור החסימה עכשיו הוא bank conflicts.
חלק A, שלב 2 - sequential addressing (bank-conflict-free)¶
הרעיון המרכזי של Mark Harris: להפוך את כיוון הלולאה. במקום stride שגדל, מתחילים ב-stride של חצי-block ומקטינים אותו. עכשיו הthreads הפעילים בכל צעד הם ה-s הראשונים, כל אחד ניגש ל-sdata[tid] וְל-sdata[tid+s] - כתובות רצופות, כל thread ל-bank שונה, אפס bank conflicts:
for (unsigned s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) // threads 0..s-1: contiguous addresses
sdata[tid] += sdata[tid + s];
__syncthreads();
}
ההפרש דק בקוד אך עמוק ברעיון: sdata[tid] += sdata[tid+s] מבטיח שכל 32 threads ה-warp הפעיל ניגשים ל-32 בנקים שונים (word w ממופה ל-bank w mod 32), ולכן כל שלב הוא transaction בודד של shared memory.
פי 1.5 נוסף. עברנו 40% מה-peak, אבל עדיין מחצית מהthreads בטלים כבר בצעד הראשון של הלולאה. את זה נתקן בשלב הבא, שכבר בטעינה מנצל את הthreads הבטלים.
חלק A, שלב 3 - first add during global load¶
הבזבוז הגדול: בצעד הראשון של הלולאה חצי מהthreads כבר לא עושים כלום. במקום שכל block יטען blockDim.x איברים, נשק חצי מהבלוקים, וכל thread יטען שני איברים מ-global memory ויחבר אותם כבר בטעינה. כך אף thread לא נולד בטל, ומספר ה-blocks (וה-overhead שלהם) נחתך בחצי:
__global__ void reduce3(const float* in, float* out, int n) {
extern __shared__ float sdata[];
unsigned tid = threadIdx.x;
unsigned i = blockIdx.x * (blockDim.x * 2) + tid; // each block: 2*blockDim
float a = (i < n) ? in[i] : 0.0f;
float b = (i + blockDim.x < n) ? in[i + blockDim.x] : 0.0f;
sdata[tid] = a + b; // first add on load
__syncthreads();
for (unsigned s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) sdata[tid] += sdata[tid + s];
__syncthreads();
}
if (tid == 0) out[blockIdx.x] = sdata[0];
}
שתי הקריאות in[i] ו-in[i + blockDim.x] נשארות מאוחדות (coalesced): בתוך warp, tid רץ, ולכן שתי הקבוצות הן שני מקטעים רצופים של 128 בתים כל אחד. שימו לב שכל איבר עדיין נקרא בדיוק פעם אחת, כך ש-bytes לא השתנה - רק חתכנו את הoverhead והעלנו את ה-issue efficiency.
קפיצה גדולה - פי 1.7, ל-72% מה-peak. כעת רוחב הפס כבר קרוב לתקרה, ומה שנשאר הוא הoverhead של הצעדים האחרונים בעץ, שבהם רק warp בודד עובד אך עדיין משלם __syncthreads.
חלק A, שלב 4 - warp-shuffle unrolling של ה-warp האחרון¶
כש-s <= 32, רק warp אחד נשאר פעיל. בתוך warp בודד אין צורך ב-__syncthreads (הthreads כבר בלוק-סטפ תחת SIMT), ואין צורך אפילו ב-shared memory: נשתמש ב-warp primitive __shfl_down_sync, שמעביר ערך מ-register של thread אחד ל-register של thread אחר באותו warp בלי לגעת ב-shared memory כלל. זה מוחק גם את ה-barriers וגם את גישות ה-shared memory של חמשת הצעדים האחרונים:
__global__ void reduce4(const float* in, float* out, int n) {
extern __shared__ float sdata[];
unsigned tid = threadIdx.x;
unsigned i = blockIdx.x * (blockDim.x * 2) + tid;
float a = (i < n) ? in[i] : 0.0f;
float b = (i + blockDim.x < n) ? in[i + blockDim.x] : 0.0f;
sdata[tid] = a + b;
__syncthreads();
// reduce in shared until 64 elements remain (two warps)
for (unsigned s = blockDim.x / 2; s > 32; s >>= 1) {
if (tid < s) sdata[tid] += sdata[tid + s];
__syncthreads();
}
// the last warp: 64 -> 1 with shuffle, no shared and no __syncthreads
if (tid < 32) {
float v = sdata[tid] + sdata[tid + 32];
for (int offset = 16; offset > 0; offset >>= 1)
v += __shfl_down_sync(0xffffffff, v, offset); // mask: all 32 threads
if (tid == 0) out[blockIdx.x] = v;
}
}
הפרמטר הראשון של __shfl_down_sync הוא מסכת ה-warp הפעיל, 0xffffffff = כל 32 הthreads משתתפים. הקריאה מחזירה לthread tid את הערך של thread tid + offset; threads בקצה, שעבורם tid + offset חורג מהטווח, מקבלים בחזרה את הערך של עצמם (כל 32 הthreads משתתפים במסכה 0xffffffff, ולכן ההתנהגות דטרמיניסטית ולא לא-מוגדרת). כתוצאה מכך התוצאות החלקיות שלהם חסרות משמעות, ורק tid == 0 מחזיק את הסכום הכולל הנכון - ולכן חשוב שרק tid == 0 יכתוב את התוצאה. דורש blockDim.x >= 64.
87% מרוחב הפס של ה-HBM. זו התקרה המעשית עבור reduction ב-FP32 - קרוב מאוד ל-0.320 ms התיאורטיים. הקפיצות מ-18% ל-87% כולן נבעו מתיאום נכון של הthreads, לא משינוי כמות העבודה.
חלק A, שלב 5 - multi-block: pass שני מול atomics¶
כל ה-kernels לעיל מייצרים gridDim.x סכומים חלקיים, אחד ל-block, כי אין סנכרון בין blocks (פרק 2). צריך לצמצם את המערך החלקי הזה לסקלר בודד. שתי דרכים סטנדרטיות:
דרך א - pass שני. קוראים ל-reduce4 שוב על מערך הסכומים החלקיים. אם gridDim.x נשאר גדול, חוזרים על כך רקורסיבית עד שנשאר איבר אחד. פשוט, דטרמיניסטי, וללא תלות ב-atomics:
// pass 1: N -> gridDim.x partials, pass 2: gridDim.x -> 1
reduce4<<<grid, block, block.x * sizeof(float)>>>(d_in, d_partial, N);
reduce4<<<1, block, block.x * sizeof(float)>>>(d_partial, d_out, grid.x);
דרך ב - atomicAdd. במקום לכתוב סכום חלקי למערך, כל block מוסיף אטומית לסקלר יחיד:
// at the end of reduce4, instead of out[blockIdx.x]:
if (tid == 0) atomicAdd(out, v); // out must be initialized to 0
הפתרון ה-atomic חוסך kernel שני והוא מהיר כשמספר ה-blocks סביר, אבל שימו לב לשתי מלכודות: (1) חובה לאתחל את out ל-0 לפני הlaunch (cudaMemset); (2) atomicAdd על float איננו אסוציאטיבי בסדר קבוע, ולכן התוצאה עשויה להשתנות מעט בין הרצות בביט האחרון - עובדה שחשוב לזכור כשבודקים נכונות בסובלנות ולא בשוויון מדויק.
טבלת התוצאות של חלק A (N = 2^28, דוגמה מ-H100; ערכיכם ישתנו):
| שלב | kernel | ms | GB/s | % HBM | speedup |
|---|---|---|---|---|---|
| 0 | naive divergent | 1.7890 | 600 | 17.9% | 1.0x |
| 1 | non-divergent | 1.1300 | 950 | 28.4% | 1.6x |
| 2 | sequential addressing | 0.7670 | 1400 | 41.8% | 2.3x |
| 3 | first-add-on-load | 0.4470 | 2402 | 71.7% | 4.0x |
| 4 | warp-shuffle | 0.3700 | 2902 | 86.6% | 4.8x |
חלק B - prefix sum / scan¶
scan (סריקה / prefix sum) מחשב לכל איבר את סכום כל קודמיו. ב-inclusive scan האיבר ה-i מקבל את x[0]+...+x[i]. נממש את אלגוריתם Hillis-Steele בתוך block: ב-log2(blockDim.x) צעדים, בכל צעד כל thread מוסיף לעצמו את הערך במרחק offset אחורה, וה-offset מוכפל. שימו לב לתבנית double-buffering הזעירה: קוראים ל-temp[tid-offset] למשתנה, __syncthreads, ורק אז כותבים - כדי למנוע read-after-write hazard בין הthreads על אותו מערך shared.
// inclusive scan of a single block; blockSums (optional) accumulates the sum of the whole block
__global__ void scanBlock(const float* in, float* out, float* blockSums, int n) {
extern __shared__ float temp[];
int tid = threadIdx.x;
int gid = blockIdx.x * blockDim.x + tid;
temp[tid] = (gid < n) ? in[gid] : 0.0f;
__syncthreads();
for (int offset = 1; offset < blockDim.x; offset <<= 1) {
float add = (tid >= offset) ? temp[tid - offset] : 0.0f;
__syncthreads(); // everyone read before anyone writes
temp[tid] += add;
__syncthreads();
}
if (gid < n) out[gid] = temp[tid];
if (blockSums && tid == blockDim.x - 1) // the full sum of the block
blockSums[blockIdx.x] = temp[tid];
}
Hillis-Steele איננו work-efficient: הוא מבצע O(n log n) פעולות לעומת O(n) באלגוריתם Blelloch (סריקת up-sweep/down-sweep). עבור block בגודל 256-1024 הפשטות מנצחת, ולכן נשתמש בו; מי שרוצה יממש את Blelloch ה-work-efficient כבונוס וישווה.
הרכבה על-פני blocks. scan הוא גלובלי מטבעו - האיבר האחרון תלוי בכולם - ולכן block בודד לא מספיק. התבנית בת-שלושת-השלבים:
1. scanBlock on each block separately -> partial out + blockSums[b] = sum of block b
2. exclusive scan on blockSums -> offsets[b] = sum of all blocks before b
3. addition: every element in block b gets += offsets[b]
// step 3: add to every element the offset of its block (block 0 gets nothing)
__global__ void addOffsets(float* out, const float* offsets, int n) {
int gid = blockIdx.x * blockDim.x + threadIdx.x;
if (blockIdx.x > 0 && gid < n)
out[gid] += offsets[blockIdx.x];
}
קריטריון קבלה לחלק B: (1) scanBlock בודד נכון מול ייחוס CPU (std::partial_sum) עבור n <= blockDim.x, בסובלנות rtol=1e-4; (2) ה-scan המורכב נכון עבור n גדול משתרע על-פני מספר blocks; (3) בדקו את מקרה הקצה n שאינו כפולה של blockDim.x. דווחו גם את רוחב הפס של הצעד הראשון (scan קורא וכותב, ולכן bytes = 2 * N * 4).
חלק C, שלב 1 - stable softmax¶
softmax הופך וקטור לוגיטים (logits) להתפלגות הסתברות: softmax(x)_i = exp(x_i) / sum_j exp(x_j). המימוש הנאיבי exp(x_i) מתפוצץ: כבר ב-x_i = 89 הערך exp(x) גולש מעבר לטווח של float ל-inf, וה-inf/inf נותן NaN. הפתרון הסטנדרטי הוא stable softmax - מחסירים את המקסימום של השורה לפני ה-exp. מתמטית זה זהה (המונה והמכנה מוכפלים ב-exp(-m)), אך מספרית כל הארגומנטים ל-exp הופכים ל-<= 0, כך ש-exp בטווח (0, 1] ואף פעם לא גולש:
נעבד שורה אחת לכל block, בשלושה מעברים: max, סכום-exp, נרמול. את ה-exp נחשב עם __expf - הגרסה ה-intrinsic הרצה על ה-SFU (Special Function Unit, יחידת הפונקציות המיוחדות של ה-SM), מהירה בהרבה מ-expf המדויקת, במחיר ~2 ULP דיוק - מקובל לחלוטין ל-softmax:
__global__ void softmaxRows(const float* in, float* out, int rows, int cols) {
extern __shared__ float sh[];
int row = blockIdx.x;
int tid = threadIdx.x;
const float* x = in + (size_t)row * cols;
float* y = out + (size_t)row * cols;
// pass 1: max of the row (grid-stride over the columns, then reduce in shared)
float localMax = -INFINITY;
for (int j = tid; j < cols; j += blockDim.x)
localMax = fmaxf(localMax, x[j]);
sh[tid] = localMax;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) sh[tid] = fmaxf(sh[tid], sh[tid + s]);
__syncthreads();
}
float m = sh[0];
__syncthreads();
// pass 2: sum of exp(x-m), the same reduction pattern
float localSum = 0.0f;
for (int j = tid; j < cols; j += blockDim.x)
localSum += __expf(x[j] - m); // SFU: fast exp
sh[tid] = localSum;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) sh[tid] += sh[tid + s];
__syncthreads();
}
float invDenom = 1.0f / sh[0];
// pass 3: normalize and write (coalesced)
for (int j = tid; j < cols; j += blockDim.x)
y[j] = __expf(x[j] - m) * invDenom;
}
הגישות x[j] עם j = tid, tid+blockDim, ... מאוחדות (coalesced) לכל warp. שני הצמצומים בתוך השורה משתמשים בדיוק בתבנית sequential-addressing מחלק A - שם לב איך primitive אחד (reduction) הופך לאבן-בניין בתוך אחר.
קריטריון קבלה: הפלט נכון מול ייחוס CPU (stable) בסובלנות rtol=1e-4; כל שורה מסתכמת ל-1.0 (בדקו |sum - 1| < 1e-4); הריצו גם קלט עם ערכים גדולים (למשל x בטווח [80, 100]) ווודאו שאין NaN - מבחן שהגרסה הנאיבית בלי חיסור ה-max נכשלת בו.
חלק C, שלב 2 - online / streaming softmax (FlashAttention)¶
ה-stable softmax דורש שלושה מעברים על השורה: אחד למצוא את המקסימום, שני לסכום ה-exp, שלישי לנרמל. ב-attention השורה היא ציון-הדמיון של token אחד מול כל השאר, והיא עלולה להיות ארוכה מאוד (הקשר של עשרות אלפי tokens). שלושה מעברים פירושם להחזיק את כל השורה בזיכרון, או לקרוא אותה שלוש פעמים מ-HBM - בדיוק מה ש-FlashAttention בא לחסוך. הפתרון הוא online softmax: לחשב max וסכום במעבר יחיד, תוך שמירת "מקסימום רץ" (running max) ו-"סכום רץ" (running sum), ותיקון הסכום בכל פעם שמתגלה מקסימום חדש.
התובנה: אם ראינו עד כה מקסימום m וסכום l = sum exp(x - m), ומגיע איבר חדש x_new, המקסימום מתעדכן ל-m' = max(m, x_new). אבל l חושב ביחס ל-m הישן, ולכן צריך לכייל אותו בגורם exp(m - m') לפני שמוסיפים את התרומה החדשה:
הביטוי הזה זהה מתמטית לחלוטין לחישוב הרגיל, אבל דורש רק מעבר אחד. הנה המעבר היחיד לכל thread, ואחריו מיזוג (merge) של החלקיים בין הthreads - כי פעולת ה-merge הזו היא אסוציאטיבית וקומוטטיבית, בדיוק כמו חיבור, ולכן משתלבת בעץ הצמצום המוכר:
// (m, l) partial for each thread over its subset of columns, single pass
struct ml { float m, l; };
__device__ ml combine(ml a, ml b) { // associative merge of two partials
float m = fmaxf(a.m, b.m);
float l = a.l * __expf(a.m - m) + b.l * __expf(b.m - m);
return {m, l};
}
__global__ void softmaxOnline(const float* in, float* out, int rows, int cols) {
extern __shared__ ml shml[];
int row = blockIdx.x, tid = threadIdx.x;
const float* x = in + (size_t)row * cols;
float* y = out + (size_t)row * cols;
// single pass: running max + running sum over the thread's columns
ml acc = { -INFINITY, 0.0f };
for (int j = tid; j < cols; j += blockDim.x) {
float xj = x[j];
float mNew = fmaxf(acc.m, xj);
acc.l = acc.l * __expf(acc.m - mNew) + __expf(xj - mNew);
acc.m = mNew;
}
shml[tid] = acc;
__syncthreads();
// reduction with the merge operator (the same sequential-addressing tree from part A)
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) shml[tid] = combine(shml[tid], shml[tid + s]);
__syncthreads();
}
float m = shml[0].m, invDenom = 1.0f / shml[0].l;
for (int j = tid; j < cols; j += blockDim.x) // final write pass
y[j] = __expf(x[j] - m) * invDenom;
}
מדוע זה משנה ל-attention. ב-FlashAttention לא מחזיקים את מטריצת ה-attention המלאה (בגודל seq x seq) בזיכרון כלל; מעבדים אותה בtiles (tiles), וכל tile מעדכן את ה-(m, l) הרצים ואת מצבור הפלט בגורם התיקון exp(m - m'). כך הזיכרון יורד מ-O(seq^2) ל-O(seq), וזו הסיבה שאפשר להריץ attention על הקשרים ארוכים. ה-primitive של online softmax שבניתם כאן הוא בדיוק הליבה המתמטית של אותו kernel; מה שמוסיף FlashAttention הוא שילובו עם ה-GEMM של Q*K^T ושל *V באותו מעבר, ועל Hopper גם wgmma ו-warpgroups (פרק 2.4). המעבר השלישי (הכתיבה) עדיין קורא את x פעם נוספת; ב-FlashAttention אמיתי גם הוא נמזג פנימה, אך לצורך הפרויקט שני-מעברים-לוגית (חישוב + כתיבה) מספיק כדי להדגים את הרעיון.
קריטריון קבלה לחלק C: (1) softmaxOnline נותן בדיוק אותו פלט כמו softmaxRows (rtol=1e-4); (2) מדדו את שני ה-kernels והשוו - ה-online מבצע מעבר קריאה אחד פחות על השורה; (3) הסבירו בפסקה מדוע ה-merge אסוציאטיבי (ולכן חוקי בעץ צמצום ובכל סדר של tiles).
קריטריוני קבלה - acceptance criteria¶
הפרויקט מושלם כאשר כל השורות מתקיימות:
[ ] A: reduce0..reduce4 return the correct sum against the CPU (rtol=1e-4)
[ ] A: ms / GB/s / %HBM table complete for the five stages, N=2^28
[ ] A: reduce4 reaches >= 80% of the HBM bandwidth
[ ] A: multi-block implemented both ways (second pass + atomicAdd), out initialized to 0
[ ] B: a single scanBlock correct against partial_sum; the composed scan correct for N spanning many blocks
[ ] B: the edge case N not a multiple of blockDim is handled
[ ] C: softmaxRows correct (rtol=1e-4), every row sums to 1, no NaN on large input
[ ] C: softmaxOnline identical to softmaxRows, in a single compute pass
[ ] all CUDA calls wrapped in CUDA_CHECK; an unverified kernel does not report
הערה שחוזרת מ-9.3: אף פעם אל תדווחו GB/s של kernel שלא עבר את בודק הנכונות. ב-reduction קל במיוחד "להאיץ" תוך שבירה שקטה - __syncthreads חסר בין הצעדים יאיץ מעט ויחזיר סכום שגוי ולא-דטרמיניסטי. הריצו את הבודק אחרי כל שינוי.
סיכום ולאן ממשיכים - what comes next¶
הפרויקט הזה סוגר את המעגל של הקורס יחד עם GEMM (9.3): שם הבעיה הייתה compute-bound וטיפסנו לכיוון ה-Tensor Cores; כאן הבעיה memory-bound וטיפסנו לכיוון רוחב הפס של ה-HBM. את מה שלמדנו בפרקים 2, 3 ו-8 ראינו כאן משלם ביצועים מדיד:
- reduction: חמישה שלבים לקחו אותנו מ-18% ל-87% מרוחב הפס, וכל קפיצה הייתה תיאום, לא חישוב - הסרת divergence (8.3), הסרת bank conflicts (8.4), ניצול threads בטלים כבר בטעינה, ו-warp primitives (
__shfl_down_sync) שמוחקים barriers ו-shared memory ב-warp האחרון. - scan: primitive גלובלי שבנוי משלוש שכבות - Hillis-Steele בתוך block, סריקה על סכומי ה-blocks, והוספה - שממחיש למה חוסר-סנכרון בין blocks (פרק 2) מכתיב ארכיטקטורה רב-שלבית.
- softmax: מ-stable (חיסור max נגד גלישה,
expעל ה-SFU) אל online/streaming - הליבה המתמטית של FlashAttention, שבזכות ה-running max וה-running sum האסוציאטיביים מוריד את זיכרון ה-attention מ-O(seq^2)ל-O(seq).
הצעד הבא הטבעי הוא לחבר את שני הפרויקטים: GEMM + online softmax = attention. עם ה-primitives של 9.3 ו-9.4 ביד, ועם warpgroups ו-wgmma מפרק 2.4 ו-CUTLASS/CuTe מפרק 6, יש לכם את כל אבני-הבניין להבין - ואף לכתוב גרסה מצומצמת של - kernel FlashAttention אמיתי.