"use strict"; const {Fragment: _Fragment, jsx: _jsx, jsxs: _jsxs} = arguments[0]; const {useMDXComponents: _provideComponents} = arguments[0]; function _createMdxContent(props) { const _components = { a: "a", annotation: "annotation", code: "code", figure: "figure", h1: "h1", h2: "h2", h3: "h3", li: "li", math: "math", mfrac: "mfrac", mi: "mi", mo: "mo", mrow: "mrow", msqrt: "msqrt", msub: "msub", msup: "msup", ol: "ol", p: "p", path: "path", pre: "pre", semantics: "semantics", span: "span", strong: "strong", svg: "svg", ul: "ul", ..._provideComponents(), ...props.components }, {Cite, CiteGroup, QuizWizard, ThemedImage} = _components; if (!Cite) _missingMdxReference("Cite", true); if (!CiteGroup) _missingMdxReference("CiteGroup", true); if (!QuizWizard) _missingMdxReference("QuizWizard", true); if (!ThemedImage) _missingMdxReference("ThemedImage", true); return _jsxs(_Fragment, { children: [_jsx(_components.p, { children: "GPUs are massively parallel processors with thousands of cores designed to handle many operations at once. To use a software development analogy, consider a CPU as a single-threaded application that excels at complex, sequential tasks but can only handle a few operations at a time. In contrast, a GPU is like a multi-threaded application with thousands of threads, each handling a small part of the task simultaneously. While the CPU is great for intricate, step-by-step processes, the GPU's strength lies in processing vast amounts of data quickly through parallelism." }), "\n", _jsx(_components.p, { children: "In GPU programming, we explicitly harness this parallel power. For example, a modern consumer GPU like the NVIDIA RTX 4090 contains 16,384 CUDA cores (smaller, specialized compute units) compared to a high-end CPU's 16 to 24 general-purpose cores. Each GPU core is slower than a CPU core, but together their sheer number allows a GPU to perform massive numbers of calculations simultaneously, ideal for tasks like matrix math that we need in deep learning." }), "\n", _jsx(_components.p, { children: "NVIDIA's CUDA (Compute Unified Device Architecture) is a platform and C++ extension that lets us write programs to run on the GPU. CUDA provides a programming model and APIs so developers can write code that runs directly on the GPU, unlocking significant performance gains by offloading parallelizable work from the CPU to the GPU. In essence, you break your problem into many small pieces that can be solved at the same time (like giving each GPU core a tiny task)." }), "\n", _jsxs(_components.p, { children: ["If you've worked with machine learning, you've likely used PyTorch, JAX, or TensorFlow. These frameworks abstract away much of the complexity of GPU programming: you allocate tensors, call operations like ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "tensor.cuda()" }) }) }) }), " or ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "tensor.to(device)" }) }) }) }), ", and under the hood, the framework handles CUDA execution. They rely on highly optimized CUDA libraries (like ", _jsx(_components.a, { href: "https://developer.nvidia.com/cublas", children: "cuBLAS" }), " for matrix multiplication and ", _jsx(_components.a, { href: "https://developer.nvidia.com/cudnn", children: "cuDNN" }), " for deep learning primitives) to achieve high performance while maintaining a user-friendly Pythonic interface."] }), "\n", _jsx(_components.p, { children: "For instance, performing vector addition in PyTorch is as simple as:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "python", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "import" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "# Create two large vectors on the GPU" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "a " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.rand(" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "1000000" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "device" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#9ECBFF" }, children: "'cuda'" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ")" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "b " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.rand(" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "1000000" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "device" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#9ECBFF" }, children: "'cuda'" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ")" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "# Add them elementwise" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "c " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " a " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " b" })] })] }) }) }), "\n", _jsxs(_components.p, { children: ["In this example, ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "a " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " b" })] }) }) }), " appears to be a single mathematical operation, but under the hood, PyTorch launches a GPU kernel (a function running on the GPU) that adds corresponding elements of ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "a" }) }) }) }), " and ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "b" }) }) }) }), " in parallel. PyTorch abstracts away explicit GPU memory management and thread launching, allowing developers to focus on high-level operations."] }), "\n", _jsx(_components.p, { children: "This abstraction is powerful, but understanding CUDA at a lower level can help optimize performance for specialized workloads. Let's explore how CUDA kernels work and how they map computations onto the GPU's hardware." }), "\n", _jsx(_components.h2, { children: "Quiz" }), "\n", _jsx(QuizWizard, { questions: [{ question: "Consider comparing CPUs to databases and GPUs to MapReduce workers, which operation would theoretically perform WORSE on a GPU compared to a CPU?", options: ["Computing the sum of a billion numbers", "Performing a join operation with multiple dependencies", "Applying an element-wise transformation to every element in an array", "Computing element-wise matrix multiplication"], answer: 1, explanation: "GPUs are optimized for simple, parallel tasks, while CPUs handle complex operations with multiple dependencies more efficiently. Therefore, a join operation with multiple dependencies would perform worse on a GPU." }, { question: "Based on the RTX 4090 example given, if each CUDA core operates at 1/4 the speed of a CPU core, approximately how many times faster could a perfectly parallelizable computation theoretically run on the GPU versus a 16-core CPU?", options: ["256x faster", "1024x faster", "4096x faster", "16384x faster"], answer: 0, explanation: "If each GPU core is 1/4 speed, the theoretical speedup is (16,384 × 1/4) ÷ 16 = 256x faster, assuming perfect parallelization." }, { question: "In the PyTorch vector addition example, what critical CUDA operations are being abstracted away that would be required in raw CUDA code?", options: ["Only the kernel launch configuration", "Memory allocation and kernel launch configuration", "Memory allocation, data transfer, and kernel launch configuration", "Memory allocation, data transfer, kernel launch configuration, and thread synchronization"], answer: 2, explanation: "The text shows PyTorch hiding three key operations: memory allocation ('device=cuda'), data transfer (implicit in creation), and kernel launch (hidden in '+'), but thread synchronization isn't mentioned as needed for this simple operation." }, { question: "Given the explanation of how PyTorch abstracts GPU operations, what would happen in the following code if 'device=cuda' was omitted from the vector creation?", options: ["The operation would fail because the tensors are on different devices", "PyTorch would automatically move the tensors to GPU for addition", "The addition would happen on CPU, potentially much slower", "PyTorch would raise a warning and move the operation to GPU"], answer: 2, explanation: "The text explains that PyTorch requires explicit device placement with 'device=cuda' or 'tensor.cuda()' to use GPU. Without this, tensors default to CPU, causing the addition to occur on CPU without automatic GPU migration." }, { question: "Based on the description of GPU vs CPU architecture, which code pattern would be LEAST suitable for GPU acceleration?", options: ["A loop containing a complex if-else tree with different operations per branch", "A loop performing the same mathematical operation on each array element", "A loop computing independent dot products between vector pairs", "A loop applying a fixed transformation matrix to many vectors"], answer: 0, explanation: "The text emphasizes that GPUs excel when many cores perform the same task on different data (like MapReduce workers). Complex branching logic forces divergent execution paths, contradicting the 'same task' parallel execution model described." }, { question: "What would be the key difference in how PyTorch and raw CUDA handle the memory for intermediate results in the expression 'd = (a + b) * c'?", options: ["PyTorch uses unified memory while CUDA requires separate allocations", "PyTorch automatically manages temporary storage while CUDA requires explicit allocation", "PyTorch uses stack memory while CUDA uses heap memory", "PyTorch uses pinned memory while CUDA uses pageable memory"], answer: 1, explanation: "The text explains that PyTorch 'abstracts away explicit GPU memory management,' meaning it automatically handles temporary storage for the result of (a + b) before multiplication, while in CUDA you would need to explicitly allocate space for this intermediate result." }] }), "\n", _jsx(_components.h1, { children: "CUDA Kernels and Threading Model" }), "\n", _jsx(_components.p, { children: "In CUDA, a kernel is a function you write that runs on the GPU. When you launch a kernel, you're not just calling a single function — you're spawning hundreds or thousands of parallel threads that all execute that function simultaneously on different data. This is often called the Single-Instruction Multiple-Thread model (SIMT): one program, many data pieces." }), "\n", _jsx(_components.p, { children: "Let's break that down with an example. Say we want to add two arrays of numbers (A and B) to produce an output array (C). With CUDA C++, we could write a kernel function like this:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "cpp", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "cpp", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "__global__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "void" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " vecAddKernel" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "A, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "B, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "C, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " n) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockIdx.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockDim.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadIdx.x;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " n) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " C[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " A[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " B[idx];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "}" }) })] }) }) }), "\n", _jsxs(_components.p, { children: ["Don't worry about ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "blockIdx" }) }) }) }), " and ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "blockDim" }) }) }) }), " for a moment – essentially, this calculation gives each thread a unique index ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "idx" }) }) }) }), " into the arrays. The kernel uses each thread's index to have that thread add one element from A and B. The ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "__global__" }) }) }) }), " qualifier means this function is a CUDA kernel that runs on the device (GPU) and can be launched from host (CPU) code. The ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " n)" })] }) }) }), " guard is there because we might launch slightly more threads than the array length (often we round up to a convenient number of threads), and those extra threads simply do nothing if their index is out of range."] }), "\n", _jsx(_components.p, { children: "When we initiate a kernel's execution from our main CPU program (called the \"host code\"), we specify the number of threads. The kernel itself will run on the GPU, but we need CPU code to configure and start it. For example:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "cpp", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "cpp", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1000000" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 256" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " numberOfBlocks " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (N " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "/" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "vecAddKernel" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<<<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " numberOfBlocks, threadsPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: ">>>" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_A, d_B, d_C, N);" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Launch configuration" })] })] }) }) }), "\n", _jsxs(_components.p, { children: ["This CPU code configures and starts the GPU kernel, specifying that we want ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "numberOfBlocks " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock" })] }) }) }), " total threads, arranged as ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "numberOfBlocks" }) }) }) }), " blocks of ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "threadsPerBlock" }) }) }) }), " threads each."] }), "\n", _jsx(ThemedImage, { light: "/blog/cuda/flowchart_light.png", dark: "/blog/cuda/flowchart_dark.png", alt: "GPU vs CPU Processing Flow" }), "\n", _jsx(_components.p, { children: "CUDA organizes threads into warps (groups of 32 threads that execute together), which are further grouped into blocks. Each block runs on a Streaming Multiprocessor (SM), which has limited resources like registers and shared memory. The block size affects how these resources are allocated and how many warps can run concurrently (a concept known as occupancy)." }), "\n", _jsx(_components.p, { children: "When threads in a warp encounter an if-statement, if some threads take one path and others take another, execution becomes serialized. The hardware uses mask bits to track which threads should execute each path, ensuring correctness but potentially impacting performance. Here's an example showing warp divergence:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "cpp", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "cpp", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "__global__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "void" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " divergentKernel" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "data, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " n) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockIdx.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockDim.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadIdx.x;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " n) {" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // This condition causes warp divergence because threads within" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // the same warp may take different paths" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (data[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: ">" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0.5" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " data[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 2.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Some threads do this" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " } " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "else" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " data[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // While others do this" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "}" }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Launch configuration considering SM resources" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " maxThreadsPerSM " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1024" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Example resource limit" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " registersPerThread " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 32" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " sharedMemoryPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1024" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // bytes" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Choose block size to maximize occupancy while respecting limits" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 256" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Multiple of warp size (32)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " numberOfBlocks " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (N " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "/" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "divergentKernel" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<<<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "numberOfBlocks, threadsPerBlock" }), _jsx(_components.span, { style: { color: "#F97583" }, children: ">>>" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_data, N);" })] })] }) }) }), "\n", _jsx(_components.p, { children: "The GPU's scheduler manages blocks across available SMs. If there are more blocks than SMs, blocks wait in a queue and are scheduled as resources become available. This scheduling depends on factors like available shared memory and register file size per SM, which limit how many blocks can run concurrently." }), "\n", _jsx(ThemedImage, { light: "/blog/cuda/gpu_schedule_manager_light.png", dark: "/blog/cuda/gpu_schedule_manager_dark.png", alt: "GPU Schedule Manager" }), "\n", _jsx(_components.p, { children: "Understanding these schedule details is crucial when optimizing runtime. For example, choosing an appropriate block size can significantly impact performance by balancing parallelism and resource usage." }), "\n", _jsx(_components.h2, { children: "Quiz" }), "\n", _jsx(QuizWizard, { questions: [{ question: "What role does the thread block size play in CUDA kernel performance?", options: ["It affects register allocation and occupancy on each SM", "It determines the granularity of work distribution", "It influences memory coalescing patterns", "It impacts shared memory bank conflicts"], answer: 0, explanation: "Thread block size directly affects how registers are allocated and the resulting occupancy (number of concurrent warps) on each SM, which is crucial for performance." }, { question: "How do CUDA's execution units handle divergent code paths within a warp?", options: ["They execute both paths simultaneously using predication", "They serialize execution of different paths with mask bits", "They dynamically regroup threads with similar paths", "They use branch prediction to optimize common paths"], answer: 1, explanation: "When threads in a warp take different paths, the execution is serialized using mask bits to track which threads should execute each path, impacting performance." }, { question: "What factors influence the optimal number of thread blocks for a CUDA kernel?", options: ["Available shared memory and register file size per SM", "Memory bandwidth and cache hierarchy utilization", "Warp scheduling and latency hiding capabilities", "Hardware occupancy and thread divergence patterns"], answer: 0, explanation: "The optimal number of blocks is primarily determined by hardware resources like shared memory and registers per SM, which limit how many blocks can run concurrently." }, { question: "How does CUDA's scheduler handle blocks when there are more blocks than available SMs?", options: ["It implements fair scheduling with time slicing", "It prioritizes blocks based on resource requirements", "It uses work-stealing between SMs for load balancing", "It queues blocks and schedules them as resources become available"], answer: 3, explanation: "The scheduler maintains a queue of blocks and launches them on SMs as resources become available, ensuring efficient utilization of the GPU." }, { question: "What is the relationship between grid dimensions and kernel launch configuration?", options: ["Grid dimensions determine SM allocation strategy", "Grid dimensions affect memory access patterns", "Grid dimensions influence warp scheduling", "Grid dimensions control thread organization"], answer: 3, explanation: "Grid dimensions define how threads are organized into blocks and the overall structure of parallel execution, directly impacting work distribution." }, { question: "How does warp size affect CUDA kernel design and optimization?", options: ["It influences memory coalescing and control flow efficiency", "It determines shared memory bank configuration", "It affects register allocation strategies", "It impacts block scheduling decisions"], answer: 0, explanation: "Warp size (32 threads) is crucial for memory coalescing and control flow efficiency, as threads within a warp execute instructions in lockstep." }] }), "\n", _jsx(_components.h1, { children: "Thread Blocks and Grids" }), "\n", _jsx(_components.p, { children: "We've seen how threads are organized into warps (eg. 32 threads each) and then grouped into blocks, and how these blocks are scheduled onto Streaming Multiprocessors (SMs). Now let's zoom out to see the complete picture. A thread block is a group of threads that can cooperate via shared memory and synchronization and that execute on the same SM. You can think of a block as a team of threads working together on a portion of the data. The grid is composed of all the blocks launched by a kernel, representing all the teams tackling the entire problem. In the diagram, you can see how the grid is divided into multiple blocks, each block is subdivided into warps, and each warp contains multiple threads." }), "\n", _jsx(ThemedImage, { light: "/blog/cuda/thread_blocks_light.png", dark: "/blog/cuda/thread_block_dark.png", alt: "CUDA Thread Blocks Organization" }), "\n", _jsxs(_components.p, { children: ["One reason for blocks is practical: there's a hardware limit on how many threads you can launch in a single block (typically 1024 threads maximum in current GPUs). If your problem needs more threads than that, you must split them into multiple blocks. For example, suppose you want to add two arrays of 2048 elements. You could use 2 blocks of 1024 threads each – block 0 handles indices 0–1023 and block 1 handles indices 1024–2047. In general, if you have ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "N" }) }) }) }), " elements and your block can have at most ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "B" }) }) }) }), " threads, you'd launch ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "ceil(N" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "/" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "B)" })] }) }) }), " blocks, so that all elements get covered. In our earlier vector addition example, we launched 3907 blocks of 256 threads to handle 1,000,000 elements. The collection of all these blocks (3907 of them) is the grid."] }), "\n", _jsxs(_components.p, { children: ["Another reason for blocks is that they allow ", _jsx(_components.strong, { children: "scaling and scheduling" }), " flexibility. The GPU has a certain number of SMs (say your GPU has 20 SMs). Each SM can run a few blocks at a time (depending on resource availability). If you launch 100 blocks and only 20 can run concurrently (one per SM), the GPU will start 20 blocks in parallel, and as soon as one finishes, it schedules the next block on that free SM. From the programmer's perspective, all 100 blocks collectively compute the result (the end result is as if 100 * blockSize threads ran). But the GPU handles distributing those blocks onto its hardware resources. This means you don't have to worry if your problem launches more threads than the GPU can physically execute at once – the runtime will time-slice blocks as needed. Blocks also provide a natural way to distribute work across GPUs or to limit how much work runs in parallel (which can sometimes help with resource constraints like shared memory or registers)."] }), "\n", _jsx(_components.p, { children: "Threads within a block have special capabilities that make them more powerful when working together. They can share a fast, on-chip memory space and can synchronize their execution using barriers. This enables efficient cooperation for tasks that require communication between threads. However, this cooperation is limited to threads within the same block – threads in different blocks must use slower global memory to communicate and cannot directly synchronize with each other. This design choice enables flexible scheduling of blocks across SMs and allows CUDA programs to scale across different GPU architectures with varying numbers of SMs." }), "\n", _jsx(_components.p, { children: "When choosing the number of threads per block, developers typically use powers of 2 (like 128, 256, or 512) to align with the warp size (32 threads) and hardware characteristics. This choice affects various performance aspects including resource utilization, memory access patterns, and the GPU's ability to hide memory latency through thread scheduling." }), "\n", _jsx(_components.p, { children: "Here's a simple example that demonstrates how blocks and grids work together to process a large array:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "cpp", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "cpp", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "__global__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "void" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " processLargeArray" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " input, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " output, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N) {" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Calculate global thread index from block and thread indices" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockIdx.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockDim.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadIdx.x;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Make sure we don't access beyond array bounds" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N) {" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Each thread processes one element" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // For this example, just multiply each element by 2" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " output[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " input[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 2.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "}" }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Host code to launch the kernel" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "void" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " launchProcessing" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float*" }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: " d_input" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float*" }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: " d_output" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: " N" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") {" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Choose number of threads per block (power of 2, <= 1024)" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 256" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Calculate number of blocks needed to process N elements" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " numBlocks " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (N " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "/" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadsPerBlock;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Launch kernel with calculated grid dimensions" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " processLargeArray" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<<<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "numBlocks, threadsPerBlock" }), _jsx(_components.span, { style: { color: "#F97583" }, children: ">>>" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_input, d_output, N);" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "}" }) })] }) }) }), "\n", _jsx(_components.p, { children: "This example shows several key concepts:" }), "\n", _jsxs(_components.ol, { children: ["\n", _jsxs(_components.li, { children: ["The kernel uses ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "blockIdx.x" }) }) }) }), ", ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "blockDim.x" }) }) }) }), ", and ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "threadIdx.x" }) }) }) }), " to calculate a unique global index for each thread"] }), "\n", _jsx(_components.li, { children: "We check array bounds since the number of threads might be rounded up to the next block size" }), "\n", _jsx(_components.li, { children: "The host code calculates the number of blocks needed using ceiling division" }), "\n", _jsx(_components.li, { children: "Each thread processes exactly one element, demonstrating the one-to-one mapping between threads and data elements" }), "\n"] }), "\n", _jsx(_components.h2, { children: "Quiz" }), "\n", _jsx(QuizWizard, { questions: [{ question: "Why can't we launch a single block with 10,000 threads to handle 10,000 tasks on the GPU?", options: ["GPUs are not powerful enough to handle that many threads", "There is a hardware limit (typically 1024) on threads per block", "It would be too slow to launch that many threads", "The CPU wouldn't be able to manage that many threads"], answer: 1, explanation: "GPUs have a hardware-imposed limit (typically 1024 threads) per block, requiring larger workloads to be split across multiple blocks." }, { question: "If a GPU kernel launch uses 32 blocks with 256 threads each, how many total threads are launched?", options: ["32 threads", "256 threads", "8,192 threads", "288 threads"], answer: 2, explanation: "The total number of threads is blocks × threads per block = 32 × 256 = 8,192 threads." }, { question: "How does the GPU schedule blocks onto its Streaming Multiprocessors (SMs) when there are more blocks than available SMs?", options: ["It runs all blocks simultaneously regardless of SM count", "It schedules blocks sequentially, one after another", "It runs as many blocks as it can fit on SMs, then schedules more as blocks finish", "It splits each block across multiple SMs"], answer: 2, explanation: "The GPU schedules as many blocks as it has resources for onto its SMs, then as blocks finish, it schedules waiting blocks onto the freed SMs." }, { question: "What is the advantage of having threads grouped into blocks (as opposed to one giant batch of independent threads)?", options: ["Blocks make the code easier to write", "Blocks allow threads to share resources and cooperate within the block", "Blocks make threads run faster", "Blocks use less GPU memory"], answer: 1, explanation: "Blocks provide a way for threads to cooperate and share resources (like shared memory) within the block, while also enabling scalable scheduling across GPU resources." }, { question: "Can threads from two different blocks communicate or synchronize with each other while the kernel is running? Why or why not?", options: ["Yes, threads can always communicate across blocks", "No, threads in different blocks cannot directly communicate during kernel execution", "Yes, but only if they're on the same SM", "It depends on the GPU model"], answer: 1, explanation: "Threads in different blocks cannot directly communicate or synchronize during kernel execution. This design enables flexible scheduling of blocks across SMs." }, { question: "In the context of thread blocks, what is a typical strategy to decide the number of threads per block for a given problem?", options: ["Always use exactly 1024 threads per block", "Use a power of 2 (like 128, 256, or 512) that balances resource usage", "Use as few threads as possible", "Use a random number of threads"], answer: 1, explanation: "Common practice is to use a power of 2 (like 128, 256, or 512) that balances resource usage and occupancy, while staying well below the maximum thread limit per block." }] }), "\n", _jsx(_components.h1, { children: "Memory Management in CUDA" }), "\n", _jsxs(_components.p, { children: ["CUDA programming not only involves running code on the GPU, but also managing data transfer between the CPU (host) and GPU (device). The CPU and GPU have separate memory spaces: your computer's RAM for the CPU, and the GPU's on-board VRAM for the GPU. ", _jsx(_components.strong, { children: "You can't directly access GPU memory from the CPU or vice versa" }), " – you have to explicitly copy data across the PCIe (or NVLink) bus connecting them. This is a major difference from PyTorch's world, where moving a tensor to the GPU is one line and usually that's all you think about."] }), "\n", _jsx(_components.p, { children: "In raw CUDA C/C++, you typically do the following in your code:" }), "\n", _jsxs(_components.ol, { children: ["\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Allocate memory on the GPU" }), " for the data you need (", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "cudaMalloc" }) }) }) }), ")."] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Copy data from the host (CPU) to the device (GPU)" }), " (", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "cudaMemcpy" }) }) }) }), " with a Host-to-Device flag)."] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Launch kernels" }), " to do computation on that data (the data stays on the GPU for as many kernels as you need)."] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Copy results from the device back to the host" }), " (", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "cudaMemcpy" }) }) }) }), " with Device-to-Host flag)."] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Free the GPU memory" }), " when done (", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "cudaFree" }) }) }) }), ")."] }), "\n"] }), "\n", _jsx(ThemedImage, { light: "/blog/cuda/shared_memory_light.png", dark: "/blog/cuda/shared_memory_dark.png", alt: "CUDA Shared Memory and Thread Block Organization" }), "\n", _jsxs(_components.p, { children: ["For example, if we want to use our ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "vecAddKernel" }) }) }) }), " from before, we'd do something like:"] }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "cpp", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "cpp", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1000000" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "size_t" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " bytes " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " sizeof" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ");" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Allocate host memory and initialize" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "h_A " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ")" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: "malloc" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(bytes);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "h_B " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ")" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: "malloc" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(bytes);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "h_C " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ")" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: "malloc" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(bytes);" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// ... fill h_A and h_B with data ..." }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Allocate device memory" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d_A, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d_B, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d_C;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaMalloc" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "&" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d_A, bytes);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaMalloc" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "&" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d_B, bytes);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaMalloc" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "&" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d_C, bytes);" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Copy input arrays from host to device" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaMemcpy" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_A, h_A, bytes, cudaMemcpyHostToDevice);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaMemcpy" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_B, h_B, bytes, cudaMemcpyHostToDevice);" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Launch kernel (using, say, 256 threads per block as before)" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threads " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 256" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blocks " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (N " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threads " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "/" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threads;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "vecAddKernel" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<<<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "blocks, threads" }), _jsx(_components.span, { style: { color: "#F97583" }, children: ">>>" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_A, d_B, d_C, N);" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Copy result array back to host" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaMemcpy" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(h_C, d_C, bytes, cudaMemcpyDeviceToHost);" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "// Free device memory" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaFree" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_A);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaFree" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_B);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: "cudaFree" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(d_C);" })] })] }) }) }), "\n", _jsx(_components.h2, { children: "Quiz" }), "\n", _jsx(QuizWizard, { questions: [{ question: "What is the correct sequence of CUDA memory operations when working with GPU data?", options: ["Allocate GPU memory, copy to GPU, compute, free memory", "Copy to GPU, allocate memory, compute, copy back", "Compute on GPU, allocate memory, copy data, free memory", "Allocate CPU memory, compute, copy to GPU, free memory"], answer: 0, explanation: "As shown in the example, we must first allocate GPU memory with cudaMalloc, then copy data with cudaMemcpy (H2D), perform computation, and finally free with cudaFree." }, { question: "Why can't we directly access GPU memory from CPU code?", options: ["GPU memory is encrypted for security", "CPU and GPU have separate memory spaces connected by PCIe bus", "GPU memory is too fast for CPU to access", "Memory formats are incompatible"], answer: 1, explanation: "As explained in the text, CPU and GPU have separate memory spaces (RAM vs VRAM) connected by PCIe bus, requiring explicit data transfers." }, { question: "In the vector addition example, why do we need both h_A/h_B and d_A/d_B arrays?", options: ["To have backup copies of the data", "Because we need separate arrays in CPU and GPU memory spaces", "To improve performance with double buffering", "To handle different data types"], answer: 1, explanation: "The example shows we need host (h_) arrays in CPU memory and device (d_) arrays in GPU memory since they are separate memory spaces." }, { question: "What is the purpose of cudaMemcpy's direction flag (e.g., cudaMemcpyHostToDevice)?", options: ["To specify the source and destination of the data transfer", "To optimize the transfer speed", "To handle different data types", "To enable asynchronous transfers"], answer: 0, explanation: "As shown in the code example, cudaMemcpy's direction flag indicates whether we're copying from host to device (H2D) or device to host (D2H)." }, { question: "Why do we calculate 'bytes' as N * sizeof(float) in the memory allocation?", options: ["To ensure we allocate the correct amount of memory for our data type", "To optimize memory alignment", "To enable vectorized operations", "To match CPU memory layout"], answer: 0, explanation: "The example shows we need to calculate the total bytes based on the number of elements (N) and the size of each element (sizeof(float))." }, { question: "When should we free GPU memory with cudaFree?", options: ["After we're completely done using the GPU arrays", "Immediately after each kernel launch", "Before copying results back to CPU", "When CPU memory is low"], answer: 0, explanation: "As demonstrated in the example, we free GPU memory with cudaFree only after we've copied results back and are completely done with the GPU arrays." }] }), "\n", _jsx(_components.h1, { children: "Shared Memory and Synchronization" }), "\n", _jsx(_components.p, { children: "Unlike Python's managed memory, CUDA provides a fast, on-chip memory called shared memory, accessible by all threads in a block. Think of shared memory as a whiteboard in a team meeting room where all team members (threads) can quickly collaborate—something you might simulate in Python using shared data structures in multiprocessing, albeit at a different performance level." }), "\n", _jsxs(_components.p, { children: ["Synchronization in CUDA is achieved with ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "__syncthreads()" }) }) }) }), " (or ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "cuda.syncthreads()" }) }) }) }), " in Numba), which acts as a barrier to ensure all threads in a block reach the same point before proceeding. This prevents race conditions similar to using locks or barriers in Python multi-threading."] }), "\n", _jsx(_components.p, { children: "Here's an example using shared memory and synchronization:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "cpp", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "cpp", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "__global__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "void" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " incrementElements" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " *" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "data, " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " n) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " __shared__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tile[" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "256" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "];" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // declare shared memory array" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockIdx.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockDim.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadIdx.x;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tid " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadIdx.x;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " n) {" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Load element from global memory to shared memory" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tile[tid] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " data[idx];" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: " __syncthreads" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "();" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // ensure all loads complete" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Each thread increments its element in shared memory" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tile[tid] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: " __syncthreads" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "();" }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // ensure all threads finished updating" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Write the result back to global memory" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " data[idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tile[tid];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "}" }) })] }) }) }), "\n", _jsx(_components.h2, { children: "Quiz" }), "\n", _jsx(QuizWizard, { questions: [{ question: "What is the primary purpose of shared memory in CUDA programming?", options: ["To store thread-private variables", "To enable fast collaboration between threads in a block", "To cache global memory accesses", "To store kernel parameters"], answer: 1, explanation: "As explained, shared memory acts like a whiteboard where all threads in a block can quickly collaborate." }, { question: "What is the purpose of __syncthreads() in the incrementElements example?", options: ["To ensure all memory loads complete before processing", "To optimize memory access patterns", "To reduce thread divergence", "To improve cache hit rates"], answer: 0, explanation: "The example shows __syncthreads() being used to ensure all threads complete their loads before proceeding with computation." }, { question: "How does shared memory differ from global memory in terms of scope?", options: ["Shared memory is accessible by all threads in a block, global memory by all threads", "Shared memory is faster but smaller than global memory", "Shared memory is persistent across kernel launches", "Shared memory is managed automatically"], answer: 0, explanation: "The text explains that shared memory is accessible by all threads in a block, unlike global memory which has broader scope." }, { question: "In the incrementElements kernel, why do we declare the shared memory array inside the kernel?", options: ["Because shared memory is block-specific and allocated per kernel launch", "To optimize compiler performance", "To reduce global memory usage", "To enable automatic memory management"], answer: 0, explanation: "As shown in the example, shared memory arrays are declared inside kernels because they are specific to each block and are allocated when the kernel launches." }, { question: "Why do we need two __syncthreads() calls in the incrementElements example?", options: ["To ensure all loads complete before processing and all updates finish before writing back", "To optimize memory access patterns", "To prevent race conditions with other blocks", "To improve kernel launch speed"], answer: 0, explanation: "The example shows we need synchronization both after loading data to shared memory and after updating it, to prevent race conditions within the block." }, { question: "What happens if we forget to use __syncthreads() in the shared memory example?", options: ["Some threads might read or write data before others are ready, causing race conditions", "The kernel will fail to compile", "The GPU will automatically handle synchronization", "Performance will improve due to reduced overhead"], answer: 0, explanation: "The text implies that without __syncthreads(), threads might access shared memory before other threads have finished their operations, leading to race conditions." }] }), "\n", _jsx(_components.h1, { children: "Fused CUDA Kernels for LLMs" }), "\n", _jsxs(_components.p, { children: ["Modern large language model (LLM) workloads have driven the development of custom CUDA kernels that fuse multiple operations to reduce memory overhead and improve efficiency. ", _jsxs(CiteGroup, { children: [_jsx(Cite, { id: "fa", text: "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", url: "https://arxiv.org/abs/2205.14135" }), _jsx(Cite, { id: "qlora", text: "QLoRA: Efficient Finetuning of Quantized LLMs", url: "https://arxiv.org/abs/2305.14314" })] }), ". One of the most well-known examples is FlashAttention ", _jsx(Cite, { id: "fa", text: "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", url: "https://arxiv.org/abs/2205.14135" }), ", which optimizes Transformer self-attention by reducing memory reads and writes."] }), "\n", _jsx(_components.p, { children: "Traditional Transformer self-attention is computed as:" }), "\n", _jsx(_components.p, { children: _jsxs(_components.span, { className: "katex", children: [_jsx(_components.span, { className: "katex-mathml", children: _jsx(_components.math, { xmlns: "http://www.w3.org/1998/Math/MathML", children: _jsxs(_components.semantics, { children: [_jsxs(_components.mrow, { children: [_jsx(_components.mi, { mathvariant: "normal", children: "Attention" }), _jsx(_components.mo, { children: "⁡" }), _jsx(_components.mo, { stretchy: "false", children: "(" }), _jsx(_components.mi, { children: "Q" }), _jsx(_components.mo, { separator: "true", children: "," }), _jsx(_components.mi, { children: "K" }), _jsx(_components.mo, { separator: "true", children: "," }), _jsx(_components.mi, { children: "V" }), _jsx(_components.mo, { stretchy: "false", children: ")" }), _jsx(_components.mo, { children: "=" }), _jsx(_components.mi, { mathvariant: "normal", children: "softmax" }), _jsx(_components.mo, { children: "⁡" }), _jsxs(_components.mrow, { children: [_jsx(_components.mo, { fence: "true", children: "(" }), _jsxs(_components.mfrac, { children: [_jsxs(_components.mrow, { children: [_jsx(_components.mi, { children: "Q" }), _jsxs(_components.msup, { children: [_jsx(_components.mi, { children: "K" }), _jsx(_components.mi, { children: "T" })] })] }), _jsx(_components.msqrt, { children: _jsxs(_components.msub, { children: [_jsx(_components.mi, { children: "d" }), _jsx(_components.mi, { children: "k" })] }) })] }), _jsx(_components.mo, { fence: "true", children: ")" })] }), _jsx(_components.mo, { children: "⋅" }), _jsx(_components.mi, { children: "V" }), _jsx(_components.mo, { separator: "true", children: "," })] }), _jsx(_components.annotation, { encoding: "application/x-tex", children: "\\operatorname{Attention}(Q, K, V) = \\operatorname{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right) \\cdot V," })] }) }) }), _jsxs(_components.span, { className: "katex-html", "aria-hidden": "true", children: [_jsxs(_components.span, { className: "base", children: [_jsx(_components.span, { className: "strut", style: { height: "1em", verticalAlign: "-0.25em" } }), _jsx(_components.span, { className: "mop", children: _jsx(_components.span, { className: "mord mathrm", children: "Attention" }) }), _jsx(_components.span, { className: "mopen", children: "(" }), _jsx(_components.span, { className: "mord mathnormal", children: "Q" }), _jsx(_components.span, { className: "mpunct", children: "," }), _jsx(_components.span, { className: "mspace", style: { marginRight: "0.1667em" } }), _jsx(_components.span, { className: "mord mathnormal", style: { marginRight: "0.07153em" }, children: "K" }), _jsx(_components.span, { className: "mpunct", children: "," }), _jsx(_components.span, { className: "mspace", style: { marginRight: "0.1667em" } }), _jsx(_components.span, { className: "mord mathnormal", style: { marginRight: "0.22222em" }, children: "V" }), _jsx(_components.span, { className: "mclose", children: ")" }), _jsx(_components.span, { className: "mspace", style: { marginRight: "0.2778em" } }), _jsx(_components.span, { className: "mrel", children: "=" }), _jsx(_components.span, { className: "mspace", style: { marginRight: "0.2778em" } })] }), _jsxs(_components.span, { className: "base", children: [_jsx(_components.span, { className: "strut", style: { height: "1.8em", verticalAlign: "-0.65em" } }), _jsx(_components.span, { className: "mop", children: _jsx(_components.span, { className: "mord mathrm", children: "softmax" }) }), _jsx(_components.span, { className: "mspace", style: { marginRight: "0.1667em" } }), _jsxs(_components.span, { className: "minner", children: [_jsx(_components.span, { className: "mopen delimcenter", style: { top: "0em" }, children: _jsx(_components.span, { className: "delimsizing size2", children: "(" }) }), _jsxs(_components.span, { className: "mord", children: [_jsx(_components.span, { className: "mopen nulldelimiter" }), _jsx(_components.span, { className: "mfrac", children: _jsxs(_components.span, { className: "vlist-t vlist-t2", children: [_jsxs(_components.span, { className: "vlist-r", children: [_jsxs(_components.span, { className: "vlist", style: { height: "1.0895em" }, children: [_jsxs(_components.span, { style: { top: "-2.5864em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "3em" } }), _jsx(_components.span, { className: "sizing reset-size6 size3 mtight", children: _jsx(_components.span, { className: "mord mtight", children: _jsx(_components.span, { className: "mord sqrt mtight", children: _jsxs(_components.span, { className: "vlist-t vlist-t2", children: [_jsxs(_components.span, { className: "vlist-r", children: [_jsxs(_components.span, { className: "vlist", style: { height: "0.8622em" }, children: [_jsxs(_components.span, { className: "svg-align", style: { top: "-3em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "3em" } }), _jsx(_components.span, { className: "mord mtight", style: { paddingLeft: "0.833em" }, children: _jsxs(_components.span, { className: "mord mtight", children: [_jsx(_components.span, { className: "mord mathnormal mtight", children: "d" }), _jsx(_components.span, { className: "msupsub", children: _jsxs(_components.span, { className: "vlist-t vlist-t2", children: [_jsxs(_components.span, { className: "vlist-r", children: [_jsx(_components.span, { className: "vlist", style: { height: "0.3448em" }, children: _jsxs(_components.span, { style: { top: "-2.3488em", marginLeft: "0em", marginRight: "0.0714em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "2.5em" } }), _jsx(_components.span, { className: "sizing reset-size3 size1 mtight", children: _jsx(_components.span, { className: "mord mathnormal mtight", style: { marginRight: "0.03148em" }, children: "k" }) })] }) }), _jsx(_components.span, { className: "vlist-s", children: "" })] }), _jsx(_components.span, { className: "vlist-r", children: _jsx(_components.span, { className: "vlist", style: { height: "0.1512em" }, children: _jsx(_components.span, {}) }) })] }) })] }) })] }), _jsxs(_components.span, { style: { top: "-2.8222em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "3em" } }), _jsx(_components.span, { className: "hide-tail mtight", style: { minWidth: "0.853em", height: "1.08em" }, children: _jsx(_components.svg, { xmlns: "http://www.w3.org/2000/svg", width: "400em", height: "1.08em", viewBox: "0 0 400000 1080", preserveAspectRatio: "xMinYMin slice", children: _jsx(_components.path, { d: "M95,702\nc-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14\nc0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54\nc44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10\ns173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429\nc69,-144,104.5,-217.7,106.5,-221\nl0 -0\nc5.3,-9.3,12,-14,20,-14\nH400000v40H845.2724\ns-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7\nc-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z\nM834 80h400000v40h-400000z" }) }) })] })] }), _jsx(_components.span, { className: "vlist-s", children: "" })] }), _jsx(_components.span, { className: "vlist-r", children: _jsx(_components.span, { className: "vlist", style: { height: "0.1778em" }, children: _jsx(_components.span, {}) }) })] }) }) }) })] }), _jsxs(_components.span, { style: { top: "-3.23em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "3em" } }), _jsx(_components.span, { className: "frac-line", style: { borderBottomWidth: "0.04em" } })] }), _jsxs(_components.span, { style: { top: "-3.4461em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "3em" } }), _jsx(_components.span, { className: "sizing reset-size6 size3 mtight", children: _jsxs(_components.span, { className: "mord mtight", children: [_jsx(_components.span, { className: "mord mathnormal mtight", children: "Q" }), _jsxs(_components.span, { className: "mord mtight", children: [_jsx(_components.span, { className: "mord mathnormal mtight", style: { marginRight: "0.07153em" }, children: "K" }), _jsx(_components.span, { className: "msupsub", children: _jsx(_components.span, { className: "vlist-t", children: _jsx(_components.span, { className: "vlist-r", children: _jsx(_components.span, { className: "vlist", style: { height: "0.9191em" }, children: _jsxs(_components.span, { style: { top: "-2.931em", marginRight: "0.0714em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "2.5em" } }), _jsx(_components.span, { className: "sizing reset-size3 size1 mtight", children: _jsx(_components.span, { className: "mord mathnormal mtight", style: { marginRight: "0.13889em" }, children: "T" }) })] }) }) }) }) })] })] }) })] })] }), _jsx(_components.span, { className: "vlist-s", children: "" })] }), _jsx(_components.span, { className: "vlist-r", children: _jsx(_components.span, { className: "vlist", style: { height: "0.538em" }, children: _jsx(_components.span, {}) }) })] }) }), _jsx(_components.span, { className: "mclose nulldelimiter" })] }), _jsx(_components.span, { className: "mclose delimcenter", style: { top: "0em" }, children: _jsx(_components.span, { className: "delimsizing size2", children: ")" }) })] }), _jsx(_components.span, { className: "mspace", style: { marginRight: "0.2222em" } }), _jsx(_components.span, { className: "mbin", children: "⋅" }), _jsx(_components.span, { className: "mspace", style: { marginRight: "0.2222em" } })] }), _jsxs(_components.span, { className: "base", children: [_jsx(_components.span, { className: "strut", style: { height: "0.8778em", verticalAlign: "-0.1944em" } }), _jsx(_components.span, { className: "mord mathnormal", style: { marginRight: "0.22222em" }, children: "V" }), _jsx(_components.span, { className: "mpunct", children: "," })] })] })] }) }), "\n", _jsxs(_components.p, { children: ["where ", _jsxs(_components.span, { className: "katex", children: [_jsx(_components.span, { className: "katex-mathml", children: _jsx(_components.math, { xmlns: "http://www.w3.org/1998/Math/MathML", children: _jsxs(_components.semantics, { children: [_jsx(_components.mrow, { children: _jsx(_components.mi, { children: "Q" }) }), _jsx(_components.annotation, { encoding: "application/x-tex", children: "Q" })] }) }) }), _jsx(_components.span, { className: "katex-html", "aria-hidden": "true", children: _jsxs(_components.span, { className: "base", children: [_jsx(_components.span, { className: "strut", style: { height: "0.8778em", verticalAlign: "-0.1944em" } }), _jsx(_components.span, { className: "mord mathnormal", children: "Q" })] }) })] }), ", ", _jsxs(_components.span, { className: "katex", children: [_jsx(_components.span, { className: "katex-mathml", children: _jsx(_components.math, { xmlns: "http://www.w3.org/1998/Math/MathML", children: _jsxs(_components.semantics, { children: [_jsx(_components.mrow, { children: _jsx(_components.mi, { children: "K" }) }), _jsx(_components.annotation, { encoding: "application/x-tex", children: "K" })] }) }) }), _jsx(_components.span, { className: "katex-html", "aria-hidden": "true", children: _jsxs(_components.span, { className: "base", children: [_jsx(_components.span, { className: "strut", style: { height: "0.6833em" } }), _jsx(_components.span, { className: "mord mathnormal", style: { marginRight: "0.07153em" }, children: "K" })] }) })] }), ", and ", _jsxs(_components.span, { className: "katex", children: [_jsx(_components.span, { className: "katex-mathml", children: _jsx(_components.math, { xmlns: "http://www.w3.org/1998/Math/MathML", children: _jsxs(_components.semantics, { children: [_jsx(_components.mrow, { children: _jsx(_components.mi, { children: "V" }) }), _jsx(_components.annotation, { encoding: "application/x-tex", children: "V" })] }) }) }), _jsx(_components.span, { className: "katex-html", "aria-hidden": "true", children: _jsxs(_components.span, { className: "base", children: [_jsx(_components.span, { className: "strut", style: { height: "0.6833em" } }), _jsx(_components.span, { className: "mord mathnormal", style: { marginRight: "0.22222em" }, children: "V" })] }) })] }), " are the query, key, and value matrices, and ", _jsxs(_components.span, { className: "katex", children: [_jsx(_components.span, { className: "katex-mathml", children: _jsx(_components.math, { xmlns: "http://www.w3.org/1998/Math/MathML", children: _jsxs(_components.semantics, { children: [_jsx(_components.mrow, { children: _jsxs(_components.msub, { children: [_jsx(_components.mi, { children: "d" }), _jsx(_components.mi, { children: "k" })] }) }), _jsx(_components.annotation, { encoding: "application/x-tex", children: "d_k" })] }) }) }), _jsx(_components.span, { className: "katex-html", "aria-hidden": "true", children: _jsxs(_components.span, { className: "base", children: [_jsx(_components.span, { className: "strut", style: { height: "0.8444em", verticalAlign: "-0.15em" } }), _jsxs(_components.span, { className: "mord", children: [_jsx(_components.span, { className: "mord mathnormal", children: "d" }), _jsx(_components.span, { className: "msupsub", children: _jsxs(_components.span, { className: "vlist-t vlist-t2", children: [_jsxs(_components.span, { className: "vlist-r", children: [_jsx(_components.span, { className: "vlist", style: { height: "0.3361em" }, children: _jsxs(_components.span, { style: { top: "-2.55em", marginLeft: "0em", marginRight: "0.05em" }, children: [_jsx(_components.span, { className: "pstrut", style: { height: "2.7em" } }), _jsx(_components.span, { className: "sizing reset-size6 size3 mtight", children: _jsx(_components.span, { className: "mord mathnormal mtight", style: { marginRight: "0.03148em" }, children: "k" }) })] }) }), _jsx(_components.span, { className: "vlist-s", children: "" })] }), _jsx(_components.span, { className: "vlist-r", children: _jsx(_components.span, { className: "vlist", style: { height: "0.15em" }, children: _jsx(_components.span, {}) }) })] }) })] })] }) })] }), " is the dimension of the key vectors."] }), "\n", _jsx(_components.p, { children: "FlashAttention optimizes this by tiling the computation in GPU shared memory, reducing the number of times data is read from and written to global memory. This significantly improves efficiency, especially for long sequences." }), "\n", _jsx(_components.p, { children: "By understanding how CUDA kernels work at a low level, developers can write custom fused kernels that outperform high-level PyTorch implementations. This is particularly useful in deep learning, where memory bandwidth is often the bottleneck." }), "\n", _jsx(_components.h2, { children: "FlashAttention: PyTorch vs CUDA Implementation" }), "\n", _jsx(_components.h3, { children: "PyTorch Implementation" }), "\n", _jsx(_components.p, { children: "Here's how you might implement FlashAttention in PyTorch, using tiled computation to manage memory:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "python", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "import" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch, math" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "def" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " flash_attention_pytorch" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(Q, K, V, block_size" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "16" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "):" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " \"\"\"" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " Compute attention scores with memory-efficient block-wise operations." }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " Args:" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " Q: Query matrix [N_out x d]" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " K: Key matrix [N_inp x d]" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " V: Value matrix [N_inp x d]" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " block_size: Size of blocks for tiled computation" }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " Returns:" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " O: Output matrix [N_out x d]" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " \"\"\"" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_out, d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " Q.shape" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_inp " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " K.shape[" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "]" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Initialize output tensors" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " O " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.zeros(N_out, d, " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "device" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "Q.device)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " L " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.zeros(N_out, " }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "device" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "Q.device)" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Calculate number of blocks needed" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " T_c " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (N_inp " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "//" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size " }), _jsx(_components.span, { style: { color: "#6A737D" }, children: "# Ceiling division" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " T_r " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (N_out " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "//" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scale_factor " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " /" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " math.sqrt(d)" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Process Q and O in T_r blocks; K, V in T_c blocks" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "in" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " range" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(T_r):" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Get current block of queries" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_start " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_end " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " min" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "((i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size, N_out)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " Q_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " Q[q_start:q_end]" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Initialize block accumulators" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " O_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.zeros(q_end " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_start, d, " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "device" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "Q.device)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " L_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.zeros(q_end " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_start, " }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "device" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "Q.device)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.full((q_end " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_start, " }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "), " }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" }), _jsx(_components.span, { style: { color: "#9ECBFF" }, children: "'-inf'" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "), " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "device" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "Q.device)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " last_m " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_block.clone()" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Process K,V in blocks" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "in" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " range" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(T_c):" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " k_start " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " k_end " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " min" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "((j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " block_size, N_inp)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " K_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " K[k_start:k_end]" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " V_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " V[k_start:k_end]" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Compute attention scores for this block" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " S_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scale_factor " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (Q_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "@" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " K_block.T) " }), _jsx(_components.span, { style: { color: "#6A737D" }, children: "# [B_r x B_c]" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Update running maximum for numerical stability" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.maximum(m_block, S_block.max(" }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "dim" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "keepdim" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "True" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ").values)" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Compute attention weights with numerical stability" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " P_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.exp(S_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_block) " }), _jsx(_components.span, { style: { color: "#6A737D" }, children: "# [B_r x B_c]" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Update accumulators with scaling factor from updated maximum" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scaling_factor " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " torch.exp(last_m " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_block)" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " L_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scaling_factor " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " L_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " P_block.sum(" }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "dim" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ", " }), _jsx(_components.span, { style: { color: "#FFAB70" }, children: "keepdim" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "True" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ")" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " O_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scaling_factor " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " O_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " P_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "@" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " V_block" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " last_m " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_block.clone()" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " # Store results for this block" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " O[q_start:q_end] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " O_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "/" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " L_block " }), _jsx(_components.span, { style: { color: "#6A737D" }, children: "# Normalize with accumulated sum" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " L[q_start:q_end] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " L_block" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " return" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " O" })] })] }) }) }), "\n", _jsx(_components.p, { children: _jsx(_components.strong, { children: "Summary of PyTorch Implementation:" }) }), "\n", _jsxs(_components.ul, { children: ["\n", _jsx(_components.li, { children: "Processes matrices in blocks to save memory" }), "\n", _jsx(_components.li, { children: "Uses max-trick for numerical stability in softmax" }), "\n", _jsxs(_components.li, { children: ["Maintains running maximum (", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "last_m" }) }) }) }), ") to scale previous computations"] }), "\n", _jsx(_components.li, { children: "Accumulates results progressively to handle large sequences" }), "\n", _jsx(_components.li, { children: "Normalizes per block for better precision" }), "\n"] }), "\n", _jsx(_components.h3, { children: "CUDA Implementation" }), "\n", _jsx(_components.p, { children: "Here's the equivalent CUDA kernel that fuses these operations for maximum efficiency:" }), "\n", _jsx(_components.figure, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.pre, { style: { backgroundColor: "#24292e", color: "#e1e4e8" }, tabIndex: "0", "data-language": "cpp", "data-theme": "github-dark", children: _jsxs(_components.code, { "data-language": "cpp", "data-theme": "github-dark", style: { display: "grid" }, children: [_jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "constexpr" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 16" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "constexpr" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 128" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "extern" }), _jsx(_components.span, { style: { color: "#9ECBFF" }, children: " \"C\"" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " __global__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "void" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " flash_attention_cuda" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " output," }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // [N_out x d]" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " output_lse," }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // [N_out]" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " query," }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // [N_out x d]" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " key," }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // [N_inp x d]" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " float*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " value," }), _jsx(_components.span, { style: { color: "#6A737D" }, children: " // [N_inp x d]" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scale," })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_out," })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_inp) {" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Shared memory for block-wise computation" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " __shared__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_block[BLOCK_SIZE][HIDDEN_DIM];" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " __shared__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " k_block[BLOCK_SIZE][HIDDEN_DIM];" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " __shared__ " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " v_block[BLOCK_SIZE][HIDDEN_DIM];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Thread indices" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadIdx.x;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " ty " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " threadIdx.y;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " const" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockIdx.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tx;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Local accumulators in registers" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " -" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "INFINITY;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " l_i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " o_i[HIDDEN_DIM] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " {" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: "0.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "};" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Process input in tiles" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (N_inp " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ") " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "/" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE; " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "++" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "tile) {" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Load query block (done once per outer loop)" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "==" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " &&" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_out) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockDim.y) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d_idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " ty;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (d_idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_block[tx][d_idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " query[row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d_idx];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: " __syncthreads" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "();" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Load key and value blocks" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " ty " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_inp " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "&&" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_out) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockDim.y) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d_idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " ty;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (d_idx " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " k_block[tx][d_idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " key[(tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tx) " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d_idx];" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " v_block[tx][d_idx] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " value[(tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tx) " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d_idx];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: " __syncthreads" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "();" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Compute attention scores and update accumulators" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_out) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_prev " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_i;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Compute scores and find max for stability" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " max_score " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " -" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "INFINITY;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scores[BLOCK_SIZE];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " #pragma" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " unroll" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "&&" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_inp; " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "++" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "j) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " score " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: ";" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " #pragma" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " unroll" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM; " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "++" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " score " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " q_block[tx][d] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " k_block[j][d];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scores[j] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scale " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " score;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " max_score " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " max" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(max_score, scores[j]);" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Update running max and scale previous results" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " max" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(m_i, max_score);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scale_factor " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " exp" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(m_prev " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_i);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " l_i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scale_factor;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " #pragma" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " unroll" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM; " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "++" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " o_i[d] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " scale_factor;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Compute attention and update output" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " #pragma" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " unroll" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "&&" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tile " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " BLOCK_SIZE " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " j " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_inp; " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "++" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "j) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " p_ij " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " exp" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(scores[j] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "-" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " m_i);" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " l_i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " p_ij;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " #pragma" }), _jsx(_components.span, { style: { color: "#B392F0" }, children: " unroll" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM; " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "++" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " o_i[d] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " p_ij " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " v_block[j][d];" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#B392F0" }, children: " __syncthreads" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "();" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: " " }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: " // Write final output" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " if" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " N_out) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " float" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " inv_l " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 1.0" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "f" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " /" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " l_i;" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: " for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " (" }), _jsx(_components.span, { style: { color: "#F97583" }, children: "int" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " 0" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "; d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "<" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM; " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "++" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d) {" })] }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " output[row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " HIDDEN_DIM " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " d] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " o_i[d] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " inv_l;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: " output_lse[row] " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " l_i;" })] }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " }" }) }), "\n", _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "}" }) })] }) }) }), "\n", _jsx(_components.p, { children: _jsx(_components.strong, { children: "Summary of CUDA Implementation (Key Advantages):" }) }), "\n", _jsxs(_components.ul, { children: ["\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Faster Memory Access" }), ": Uses ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "__shared__" }) }) }) }), " memory for ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "q_block" }) }) }) }), ", ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "k_block" }) }) }) }), ", ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "v_block" }) }) }) }), " vs PyTorch's repeated ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "K_block " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " K[k_start:k_end]" })] }) }) })] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Register Usage" }), ": Stores accumulators (", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "o_i" }) }) }) }), ", ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "l_i" }) }) }) }), ", ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "m_i" }) }) }) }), ") in fast registers vs PyTorch's ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "O_block" }) }) }) }), " tensor allocations"] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Parallel Processing" }), ": Each thread computes one row via ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "row " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockIdx.x " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "*" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " BLOCK_SIZE" }), _jsx(_components.span, { style: { color: "#F97583" }, children: " +" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " tx" })] }) }) }), " vs PyTorch's sequential ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#F97583" }, children: "for" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " i " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "in" }), _jsx(_components.span, { style: { color: "#79B8FF" }, children: " range" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: "(T_r)" })] }) }) })] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Memory Coalescing" }), ": Structured memory loads with ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsxs(_components.span, { "data-line": "", children: [_jsx(_components.span, { style: { color: "#E1E4E8" }, children: "d " }), _jsx(_components.span, { style: { color: "#F97583" }, children: "+=" }), _jsx(_components.span, { style: { color: "#E1E4E8" }, children: " blockDim.y" })] }) }) }), " vs PyTorch's slice operations"] }), "\n", _jsxs(_components.li, { children: [_jsx(_components.strong, { children: "Loop Optimization" }), ": Uses ", _jsx(_components.span, { "data-rehype-pretty-code-figure": "", children: _jsx(_components.code, { "data-language": "python", "data-theme": "github-dark", style: { backgroundColor: "#24292e", color: "#e1e4e8" }, children: _jsx(_components.span, { "data-line": "", children: _jsx(_components.span, { style: { color: "#6A737D" }, children: "#pragma unroll" }) }) }) }), " to reduce loop overhead and improve performance by optimizing memory access and register usage, unlike PyTorch's slower interpreted loops."] }), "\n"] }), "\n", _jsx(_components.h2, { children: "Quiz" }), "\n", _jsx(QuizWizard, { questions: [{ question: "What is the primary reason for tracking both m_block and last_m in the PyTorch implementation, instead of just using m_block?", options: ["To enable gradient checkpointing during backpropagation", "To handle edge cases when the sequence length isn't divisible by block_size", "To correctly scale previous block computations when the maximum changes", "To optimize memory access patterns in consecutive block operations"], answer: 2, explanation: "last_m is crucial for correctly adjusting previous computations (O_block and L_block) when a new maximum attention score is found in the current block, ensuring numerical stability across block boundaries." }, { question: "In the CUDA implementation, why is the shared memory array s_block dimensioned as [BLOCK_SIZE][BLOCK_SIZE] while q_block, k_block, and v_block use [BLOCK_SIZE][HIDDEN_DIM]?", options: ["To optimize shared memory bank conflicts in matrix multiplication", "Because s_block stores intermediate attention scores between Q and K blocks", "To reduce shared memory usage since s_block is temporary", "Because attention scores form a square matrix while embeddings are rectangular"], answer: 1, explanation: "s_block stores the intermediate attention scores between Q and K blocks, which form a square matrix of size BLOCK_SIZE × BLOCK_SIZE, while the other arrays store embedding vectors of dimension HIDDEN_DIM." }, { question: "What potential issue could arise from the PyTorch implementation's use of torch.maximum() for m_block updates?", options: ["Memory fragmentation due to temporary tensor allocation", "Reduced parallelism compared to an in-place maximum operation", "Incorrect gradient computation during backpropagation", "Unnecessary synchronization points in the computation graph"], answer: 0, explanation: "torch.maximum() creates new temporary tensors, which can lead to memory fragmentation when processing many blocks, especially with large sequence lengths." }, { question: "Why does the CUDA implementation use separate loops for loading data and computing attention scores, rather than combining them?", options: ["To maintain warp execution efficiency and reduce thread divergence", "To simplify the implementation and improve code readability", "To enable compiler optimizations for memory coalescing", "To reduce register pressure in the inner loops"], answer: 0, explanation: "Separate loops maintain better warp execution efficiency by grouping similar operations together, reducing thread divergence that would occur if loading and computation were mixed." }, { question: "What is the significance of using blockDim.y in the CUDA implementation's memory loading loop?", options: ["It enables vectorized memory loads for better throughput", "It ensures proper alignment for memory coalescing", "It distributes work evenly across thread block dimensions", "It prevents shared memory bank conflicts"], answer: 2, explanation: "Using blockDim.y distributes the work of loading HIDDEN_DIM elements across the y-dimension of the thread block, enabling efficient parallel loading of the full embedding vectors." }, { question: "Why does the PyTorch implementation compute T_c and T_r using ceiling division ((N + block_size - 1) // block_size)?", options: ["To ensure all elements are processed even with non-divisible sizes", "To optimize memory alignment for tensor operations", "To match CUDA's block scheduling behavior", "To improve cache utilization in matrix operations"], answer: 0, explanation: "Ceiling division ensures that all elements are processed even when N_inp or N_out aren't perfectly divisible by block_size, by creating one extra partial block when needed." }] })] }); } function MDXContent(props = {}) { const {wrapper: MDXLayout} = { ..._provideComponents(), ...props.components }; return MDXLayout ? _jsx(MDXLayout, { ...props, children: _jsx(_createMdxContent, { ...props }) }) : _createMdxContent(props); } return { default: MDXContent }; function _missingMdxReference(id, component) { throw new Error("Expected " + (component ? "component" : "object") + " `" + id + "` to be defined: you likely forgot to import, pass, or provide it."); }

Introduction to CUDA Programming for Python Developers

Curious about PySpur?

We're open-source, Apache 2.0 licensed.