לדלג לתוכן

5.5 CUDA Graphs לצמצום overhead הרצאה

לאורך פרק 5 מיפינו את מחסנית התוכנה בצד הhost (host, ה-CPU): את ה-CUDA Runtime API ואת הספרייה שמממשת אותו libcudart.so, את היחס בינו לבין ה-Driver API, ובשיעור 5.3 את ה-streams (זרמים) - התורים המסודרים שמאפשרים לחפוף עבודה על ה-GPU. בכל השיעורים האלה הנחנו הנחה שקטה ולא מדויקת: שlaunch kernel היא זולה. היא איננה. כל launch בודדת עולה לhost בסדר גודל של מיקרו-שניות של עבודת CPU - זיהוי הפונקציה, אריזת הארגומנטים, הגדרת תצורת הlaunch והגשתה לדרייבר - ובעומסים שמשגרים מאות kernels בתוך מילישניות בודדות, כמו inference של מודל שפה בזמן-אמת (low-latency LLM inference), הoverhead הזו הופכת לנתח משמעותי מזמן הריצה הכולל. בשיעור הזה, שסוגר את פרק הhost, נכיר את הכלי שנועד בדיוק לבעיה הזו: גרף CUDA (CUDA Graph), מבנה שלוכד DAG שלם של launches ועבודה נלווית ומגיש אותו לdevice בפעולת host אחת, שאפשר להריץ מחדש (replay) שוב ושוב בזול. נראה את דפוס לכידת ה-stream (stream capture), נבנה גרף גם ידנית, נבין למה הוא מוזיל overhead, ונעמוד על המגבלה החדה שלו - המצביעים הצרובים בו. את התיקון הקונקרטי הזה נזכיר שוב כשנדבר על profiling בשיעור 7.4 ועל צנרת ה-inference בשיעור 9.2.

הבעיה - overhead הlaunch של הhost - launch overhead

נתחיל מהמספר שמניע את כל השיעור. כשאתם כותבים myKernel<<<grid, block>>>(args), ה-<<< >>> הוא סוכר תחבירי שה-nvcc מתרגם לרצף קריאות אל ה-Runtime API (בעיקר cudaLaunchKernel). הרצף הזה עושה עבודה אמיתית על ה-CPU: הוא אורז את הארגומנטים למערך, מעתיק את תצורת הlaunch (grid, block, shared memory, stream), פונה לדרייבר, והדרייבר מצדו כותב פקודה לתור החומרה. כל זה עולה בסדר גודל של 3 עד 10 מיקרו-שניות של זמן CPU לlaunch בודדת, גם אם ה-kernel עצמו רץ על הdevice במיקרו-שנייה אחת.

כל עוד ה-kernels גדולים - רצים מילישניות שלמות - הoverhead הזו זניחה. הבעיה מתלקחת כשה-kernels קטנים ורבים. נעשה חשבון על צעד decode יחיד של מודל שפה בסגנון Transformer:

   Transformer decode step (one token), a 32-layer model:

   per layer, roughly: QKV projection, RoPE, attention,
   output projection, two MLP GEMMs, two normalizations,
   and additional elements  ->  roughly 12-16 kernels per layer

   32 layers x ~15 kernels  ≈  480 launches per token
   480 launches x 5 microseconds  ≈  2.4 milliseconds of CPU overhead alone

אם יעד ה-latency שלכם הוא, נניח, 10 מילישניות לטוקן, אז 2.4 מילישניות של launch overhead הן כ-24% מהתקציב - נשרפים עוד לפני שה-GPU חישב ביט אחד. גרוע מכך: אם ה-kernels קצרים מהזמן שלוקח לhost לשגר את הבא בתור, ה-GPU מורעב (starved) - הוא מסיים את ה-kernel הנוכחי ויושב בטל בזמן שה-CPU עדיין טורח על הlaunch הבאה. במצב כזה התוכנית שלכם overhead-bound: לא compute-bound ולא memory-bound, אלא חסומה ביכולת של הhost להזין את הdevice בעבודה. הגדלת ה-batch או שיפור ה-kernel לא יעזרו; הצוואר הוא ה-CPU.

מדוע streams לבדם לא פותרים את זה. בשיעור 5.3 ראינו ש-streams מאפשרים חפיפה, ואפשר לחשוב שהם הפתרון. הם לא. streams מסירים את overhead הסנכרון (הhost לא צריך לחכות לסיום), אבל הם לא מסירים את overhead הlaunch - הhost עדיין חייב להנפיק כל אחת מ-480 הlaunches בנפרד, אחת-אחת, ולשלם את המיקרו-שניות שלה. ה-stream קובע את הסדר; הוא לא מקבץ את ההגשה. בדיוק את הקיבוץ הזה גרף CUDA מספק.

מהו גרף CUDA - CUDA Graph

גרף CUDA הוא, בלשון ה-glossary, "גרף של kernel launches ועבודה אחרת שהhost יכול להגיש לdevice בבת אחת". במקום שהhost ינפיק מאות launches בזו אחר זו, אנחנו מתארים את כל גרף העבודה פעם אחת, ואז מגישים ומריצים אותו מחדש כיחידה אחת. המבנה הוא DAG (Directed Acyclic Graph, גרף מכוון חסר-מעגלים): הצמתים (nodes) הם פעולות - kernel launch, העתקת זיכרון, memset, ואפילו גרף-בן - והקשתות (edges) הן תלויות: קשת מ-A ל-B אומרת ש-B לא יתחיל לפני ש-A הסתיים.

   The conceptual difference:

   Regular (per-launch):         Graph (submitted at once):
   host: launch K0  ─┐            host: launch(graphExec) ─┐
   host: launch K1  ─┤ each one            (one call)       │
   host: launch K2  ─┤ costs                                │
   host: launch K3  ─┤ ~5us                 K0 ─► K1 ─► K2 ─► K3
   host: launch K4  ─┘ to host             (the driver runs the whole DAG)

חשוב להבחין בין שלושה שלבים בחיי גרף, כי כל אחד משלם overhead שונה:

  • תיאור (definition) - בונים את הטופולוגיה: אילו צמתים, ואילו תלויות ביניהם. זה נעשה פעם אחת. אפשר לבנות אותה בלכידת stream (הדרך הנפוצה) או ידנית (צומת-צומת).
  • הגשם למכונה (instantiation) - cudaGraphInstantiate הופך את הגרף הטופולוגי לגרף בר-הרצה (cudaGraphExec_t). זה השלב היקר - הדרייבר מאמת את ה-DAG, ממפה משאבים ומכין את כל מבני הפקודות מראש. גם הוא נעשה, בדרך כלל, פעם אחת.
  • הרצה (launch / replay) - cudaGraphLaunch מגיש את הגרף בר-ההרצה לdevice. זו הפעולה שחוזרת שוב ושוב, וזו שהיא זולה: קריאת host אחת שמשגרת את כל ה-DAG.

הרווח כולו נובע מכך שאת העבודה היקרה - זיהוי, תצורה, אימות, מיפוי - עושים פעם אחת ב-instantiation, ומשלמים אותה מופחתת (amortized) על פני אלפי הרצות. ההרצה עצמה כמעט לא עולה לhost.

לכידת stream - stream capture

הדרך הנפוצה והקלה ביותר לבנות גרף היא לכידת stream: אנחנו לא בונים צמתים ביד, אלא מנפיקים את העבודה כרגיל על stream, ומבקשים מ-CUDA להקליט את מה שהנפקנו לתוך גרף במקום להריץ אותו. ה-glossary מציג את השלד המדויק:

cudaStreamBeginCapture(stream);
kernelGemm<<<{32,20},64,19200,stream>>>(a, b, c);
kernelEpilogue<<<{256,2},{8,32},0,stream>>>(c, c);
cudaStreamEndCapture(stream, &graph);

cudaGraphInstantiate(&graphExec, graph, flags);
cudaGraphLaunch(graphExec, stream);

שימו לב לתצורות הlaunch הקונקרטיות שנלכדות: kernelGemm<<<{32,20},64,19200,stream>>> הוא grid דו-ממדי של {32,20} בלוקים, block של 64 threads, ו-19200 בתים של shared memory דינמי; kernelEpilogue<<<{256,2},{8,32},0,stream>>> הוא grid {256,2}, block דו-ממדי {8,32}, ואפס shared memory דינמי. הפרמטר הרביעי בכל launch הוא ה-stream - וזה קריטי: הלכידה עובדת דרך ה-stream. בין cudaStreamBeginCapture ל-cudaStreamEndCapture, הlaunches לא רצות - הן מוקלטות כצמתים בגרף, והתלויות ביניהן נגזרות אוטומטית מהסדר וממבנה ה-stream.

בואו נכתוב תוכנית מלאה ובת-הרצה שמדגימה את המחזור המלא, כולל מקרו הבדיקה CUDA_CHECK שראינו לאורך הקורס והרצה חוזרת בלולאה:

#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)

__global__ void tiny(float* d) { d[0] += 1.0f; }   // deliberately tiny kernel

int main(void) {
    float* d;
    CUDA_CHECK(cudaMalloc(&d, sizeof(float)));
    CUDA_CHECK(cudaMemset(d, 0, sizeof(float)));

    cudaStream_t stream;
    CUDA_CHECK(cudaStreamCreate(&stream));

    const int N = 512;                 // 512 small launches in a dependent chain
    cudaGraph_t     graph;
    cudaGraphExec_t graphExec;

    // 1. capture: issue on a non-default stream, doesn't run - it's recorded
    CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
    for (int i = 0; i < N; i++)
        tiny<<<1, 1, 0, stream>>>(d);
    CUDA_CHECK(cudaStreamEndCapture(stream, &graph));

    // 2. instantiation: the expensive step, once
    CUDA_CHECK(cudaGraphInstantiate(&graphExec, graph, 0));

    // 3. repeated launch: cheap, one host call for all 512 launches
    for (int iter = 0; iter < 1000; iter++) {
        CUDA_CHECK(cudaGraphLaunch(graphExec, stream));
    }
    CUDA_CHECK(cudaStreamSynchronize(stream));

    float host = 0.0f;
    CUDA_CHECK(cudaMemcpy(&host, d, sizeof(float), cudaMemcpyDeviceToHost));
    printf("d[0] = %.1f (expected %d)\n", host, N * 1000);   // 512000

    CUDA_CHECK(cudaGraphExecDestroy(graphExec));
    CUDA_CHECK(cudaGraphDestroy(graph));
    CUDA_CHECK(cudaStreamDestroy(stream));
    CUDA_CHECK(cudaFree(d));
    return 0;
}

כמה דקויות ומלכודות שחייבים להכיר בלכידה:

  • חובה stream לא-default. אי אפשר ללכוד על ה-legacy default stream (stream 0). לכידה דורשת stream מפורש (שנוצר ב-cudaStreamCreate) או את ה-per-thread default stream. זו אותה הסיבה מ-5.3 שבה נמנענו מה-default stream לחפיפה.
  • מצב הלכידה (capture mode). ל-cudaStreamBeginCapture יש פרמטר מצב. cudaStreamCaptureModeGlobal הוא הברירה הבטוחה; ThreadLocal ו-Relaxed משנים כמה בקפדנות CUDA אוסר פעולות "מסוכנות" מ-threads אחרים בזמן הלכידה.
  • אסור לסנכרן בתוך לכידה. כל פעולה שחוסמת את הhost או שואלת סטטוס - cudaMemcpy חוסמת, cudaDeviceSynchronize, cudaStreamSynchronize, cudaMalloc רגיל - פוסלת את הלכידה ומחזירה שגיאה. בתוך הלכידה משתמשים רק בגרסאות ה-async (למשל cudaMemcpyAsync), וכל ההקצאות צריכות להיעשות לפני הלכידה (או עם cudaMallocAsync, שנלכד כצומת זיכרון).
  • תמיד להריץ warm-up. ה-instantiation הראשון וה-JIT של PTX ל-SASS (ראו שיעור 4.3) עולים overhead חד-פעמית. מדדו את ה-replay אחרי הרצת חימום, בדיוק כמו שעשינו עם kernels ב-5.3.

בנייה ידנית של גרף - manual graph construction

לכידת stream נוחה, אבל לפעמים רוצים שליטה מפורשת בטופולוגיה - למשל לבנות תלויות שלא נובעות מסדר ה-stream, או לייצר את הגרף פרוגרמטית. אז בונים את הגרף צומת-צומת. הרעיון: יוצרים גרף ריק ב-cudaGraphCreate, מוסיפים צמתים כשלכל צומת מעבירים במפורש את מערך התלויות שלו, ואז מגישים למכונה כרגיל.

cudaGraph_t graph;
CUDA_CHECK(cudaGraphCreate(&graph, 0));

// kernel node parameters
cudaKernelNodeParams p = {};
p.func           = (void*)tiny;         // pointer to the kernel function
p.gridDim        = dim3(1, 1, 1);
p.blockDim       = dim3(1, 1, 1);
p.sharedMemBytes = 0;
void* args[]     = { &d };               // array of pointers to the arguments
p.kernelParams   = args;
p.extra          = nullptr;

// node A - no dependencies (source)
cudaGraphNode_t a;
CUDA_CHECK(cudaGraphAddKernelNode(&a, graph, nullptr, 0, &p));

// node B - depends on A (pass A as a dependency)
cudaGraphNode_t b;
cudaGraphNode_t depsB[] = { a };
CUDA_CHECK(cudaGraphAddKernelNode(&b, graph, depsB, 1, &p));

cudaGraphExec_t graphExec;
CUDA_CHECK(cudaGraphInstantiate(&graphExec, graph, 0));
CUDA_CHECK(cudaGraphLaunch(graphExec, stream));

לצד cudaGraphAddKernelNode קיימות cudaGraphAddMemcpyNode, cudaGraphAddMemsetNode, cudaGraphAddHostNode (להרצת callback על ה-CPU כחלק מהגרף) ו-cudaGraphAddChildGraphNode (הטמעת תת-גרף). את התלויות אפשר גם להוסיף בדיעבד עם cudaGraphAddDependencies. בפועל, רוב הקוד משתמש בלכידת stream כי היא קצרה בהרבה ופחות שברירית; הבנייה הידנית שמורה למקרים שדורשים שליטה מדויקת או ייצור דינמי של הטופולוגיה. חשוב להכיר את שתי הדרכים, אבל להעדיף לכידה.

מדוע גרף מוזיל overhead - anatomy of the saving

נחזור למספרים ונבין מה בדיוק נחסך. בגרסה הרגילה, הhost משלם עבור כל אחת מ-N הlaunches את מלוא המחיר: אריזת ארגומנטים, תצורה, פנייה לדרייבר. בגרסת הגרף, את כל העבודה הזו עשינו פעם אחת ב-instantiation; ה-cudaGraphLaunch רק אומר לדרייבר "הרץ את ה-DAG שכבר הכנת", וזו קריאה בעלות של launch בודדת בערך, ללא תלות ב-N. נשווה:

מדד לולאת launches (per-launch) גרף (replay)
עלות host לכל הרצה N x ~5us ~2us (קבוע, לא תלוי ב-N)
מתי משלמים על תצורה/אימות בכל launch, כל פעם פעם אחת ב-instantiate
רגישות למספר ה-kernels לינארית ב-N כמעט אפס
מה מגביל בסוף overhead הhost (אם kernels קטנים) זמן ההרצה על הdevice

חשבון ממודל ל-N=512 kernels זעירים, ~5 מיקרו-שניות overhead לlaunch:

   per-launch:  512 x 5us  =  2560 us  =  2.56 ms  (host overhead per run)
   graph replay:               ~2us    +  device run time

   If the 512 kernels run on the device in ~0.6 ms total (in a dependent chain),
   then: per-launch ≈ 2.56 ms   (host-bound, the GPU is starved)
         graph      ≈ 0.60 ms   (device-bound, the CPU is almost uninvolved)
         speedup    ≈ 4.3x

שימו לב לתובנה החדה: הגרף לא מאיץ את ה-kernels. ה-512 launches עדיין רצות על הdevice באותו זמן בדיוק (בשרשרת תלויה, בזו אחר זו). מה שהגרף עושה הוא להסיר את overhead הhost, ובכך להוריד את הרצפה מ-max(host_overhead, device_time) ל-device_time בלבד. לכן הרווח הגדול ביותר מגיע כשהיינו overhead-bound מלכתחילה: הרבה kernels קטנים. אם ה-kernels גדולים וזמן הdevice שולט, הגרף כמעט לא יעזור - הoverhead ממילא הייתה זניחה. זו בדיוק הסיבה שה-glossary מצביע על inference של LLM בזמן-אמת כמקרה הקנוני: שם, במיוחד ב-batch קטן, ה-kernels קצרים והoverhead שולטת.

המגבלה - מצביעים צרובים - baked-in pointers

לגרף CUDA יש מגבלה אחת חדה שחובה להבין, וה-glossary מדגיש אותה: ה-kernels וה-I/O שלהם מזוהים לפי מצביעי זיכרון גולמיים. כשלכדנו את tiny<<<...>>>(d), הכתובת הקונקרטית של d - נניח 0x564603AFCC00 - נצרבה לתוך הגרף בר-ההרצה כארגומנט. הגרף לא זוכר "המשתנה d"; הוא זוכר את המספר 0x564603AFCC00. מכאן נובעות שתי מסקנות מעשיות:

  • הגרף אינו ניתן לשמירה או להעברה (not serializable / not portable). אי אפשר לכתוב graphExec לדיסק ולטעון אותו בתהליך אחר, כי הכתובות הצרובות לא יהיו תקפות במרחב הכתובות החדש. הדרך היחידה "להעביר" גרף היא checkpoint ו-restore מלאים של כל זיכרון הhost והdevice, כך שהכתובות הצרובות יישארו תקפות - פתרון כבד שכמעט לא משמש בפועל.
  • חובה יציבות ההקצאות (stable allocations). ה-buffers שהגרף נוגע בהם חייבים להישאר מוקצים באותן כתובות לאורך כל ההרצות החוזרות. אם תשחררו buffer ותקצו אותו מחדש, ה-allocator עלול להחזיר כתובת אחרת - והגרף ימשיך לכתוב לכתובת הישנה, שכעת חופשית או שייכת למשהו אחר. זהו use-after-free קלאסי, שקט וקטלני.
   The scenario that breaks a graph:

   cudaMalloc(&d, ...)  ->  d = 0x...CC00   (captured into the graph)
   ... instantiate, a few valid replays ...
   cudaFree(d)                               (the address was freed)
   cudaMalloc(&d2, ...) ->  d2 = 0x...E100   (a new address!)
   cudaGraphLaunch(...)                       (still writes to 0x...CC00 - bug)

לכן הכלל: הקצו את ה-buffers פעם אחת, לכדו את הגרף, ואל תשחררו/תקצו מחדש כל עוד הגרף בשימוש. אם הטופולוגיה נשארת אבל רק המצביעים או הפרמטרים משתנים, אין צורך ללכוד מחדש מאפס - cudaGraphExecUpdate מעדכן את הפרמטרים של גרף בר-הרצה קיים בזול, כל עוד המבנה זהה. זה מאפשר, למשל, להחליף את ה-buffer של הקלט בין הרצה להרצה בלי לשלם שוב על instantiation. זו הדרך הנכונה לטפל בשינוי מצביעים, במקום לשחרר ולהקצות.

גרפים ב-PyTorch - torch.cuda.graph

כמעט אף אחד לא כותב את קריאות ה-cudaStreamBeginCapture ידנית בעבודת deep learning; משתמשים בעטיפה שה-glossary מציין - מנהל ההקשר (context manager) torch.cuda.graph. הוא לוכד את פעולות ה-CUDA ש-PyTorch מנפיק בתוך הבלוק, ובונה מהן גרף בר-הרצה. הדפוס הקנוני ל-inference, שמכבד את כל המגבלות שראינו (יציבות מצביעים, warm-up, stream נפרד):

import torch

model = build_model().cuda().eval()

# 1. static buffers - fixed addresses the graph will bake in. Do not replace them!
static_input  = torch.zeros(1, seq_len, hidden, device="cuda")
static_output = None

# 2. warm-up on a side stream (required: run a few iterations before capture)
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        static_output = model(static_input)
torch.cuda.current_stream().wait_stream(s)

# 3. capture: the whole block becomes a runnable graph
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = model(static_input)

# 4. repeated run per token/request
def infer(new_input):
    static_input.copy_(new_input)   # write into the same buffer (same address)
    g.replay()                       # one launch for the whole forward pass
    return static_output.clone()     # copy out before the next step overwrites it

הנקודה המרכזית שמחברת הכל: static_input ו-static_output הם buffers קבועים. אנחנו לא יוצרים tensor חדש בכל צעד (זה יקצה כתובת חדשה ויפר את הגרף); במקום זאת אנחנו כותבים לתוך ה-buffer הקבוע עם copy_, מריצים g.replay(), ומעתיקים את התוצאה החוצה. זו בדיוק אותה מגבלת "יציבות ההקצאות" מהסעיף הקודם, מתורגמת לעולם ה-tensors. עבור מקרים פשוטים, torch.cuda.make_graphed_callables עוטף פונקציה או מודול ומחזיר גרסה מגורפת אוטומטית. בצעד decode של LLM, שבו אותו forward pass של השכבות רץ שוב ושוב לכל טוקן עם אותם משקלים ואותם buffers, גרף CUDA הוא כמעט מובן מאליו - וזה מה שנפרוט לפרטים בשיעור 9.2.

סיכום

  • כל kernel launch בודדת עולה לhost בסדר גודל של 3-10 מיקרו-שניות של עבודת CPU (אריזת ארגומנטים, תצורה, פנייה לדרייבר), גם כשה-kernel עצמו רץ במיקרו-שנייה על הdevice.
  • כשמשגרים מאות kernels קטנים בתוך מילישניות - כמו בצעד decode של LLM, שיכול להגיע לכ-480 launches לטוקן - הoverhead המצטברת הופכת את התוכנית לoverhead-bound וה-GPU מורעב.
  • streams לבדם לא פותרים את זה: הם מסירים overhead סנכרון אבל לא launch overhead, כי הhost עדיין מנפיק כל launch בנפרד.
  • גרף CUDA לוכד DAG של launches ועבודה נלווית ומגיש אותו בפעולת host אחת שאפשר להריץ מחדש בזול; חיי הגרף הם תיאור, instantiation (השלב היקר, פעם אחת), ו-replay (זול, חוזר).
  • לכידת stream (cudaStreamBeginCapture -> הנפקת ה-kernels על stream -> cudaStreamEndCapture -> cudaGraphInstantiate -> cudaGraphLaunch בלולאה) היא הדרך הנפוצה; היא דורשת stream לא-default, אוסרת פעולות סנכרון בתוך הלכידה, והlaunches מוקלטות ולא רצות.
  • אפשר גם לבנות גרף ידנית צומת-צומת עם cudaGraphCreate ו-cudaGraphAddKernelNode תוך העברת מערכי תלויות מפורשים, אבל בפועל מעדיפים לכידה.
  • הגרף לא מאיץ את ה-kernels - הוא מסיר את overhead הhost בלבד, ולכן הרווח הגדול ביותר מגיע דווקא כשהיינו overhead-bound מלכתחילה (הרבה kernels קטנים).
  • מגבלה חדה: ה-kernels וה-I/O שלהם צרובים בגרף לפי מצביעי זיכרון גולמיים, ולכן גרף אינו ניתן לשמירה או להעברה ומחייב יציבות מוחלטת של ההקצאות; שחרור והקצאה מחדש של buffer הופכים את הגרף ל-use-after-free.
  • כשרק הפרמטרים משתנים והטופולוגיה נשארת, cudaGraphExecUpdate מעדכן גרף בר-הרצה קיים בזול במקום ללכוד מחדש.
  • ב-PyTorch, מנהל ההקשר torch.cuda.graph עוטף את כל המנגנון; משתמשים ב-buffers סטטיים קבועים, כותבים אליהם ב-copy_ ומריצים replay(), וזה הכלי הטבעי לצעד inference חוזר של LLM.