CUDA 后端The CUDA backend

Last lesson watched the CPU backend make matmul fast with SIMD + multithreading (L31). But what really makes large-model inference "fly" is usually the GPU - it wins not with a few strong cores but with thousands of threads computing at once. This lesson drops into ggml's CUDA backend (ggml/src/ggml-cuda/) to see how it organizes all those threads: how one matmul is split into block-and-thread cooperation, and how on-chip fast memory keeps them fed.

A single GPU thread is feeble - lower clock than a CPU, computing one cell at a time. The GPU's strength is all in "numbers": thousands of threads at once give overwhelming throughput. So the whole craft of a GPU kernel comes down to two things: cut the work into thousands of independent little pieces (so every thread has something to do), and get the data close enough to the threads (so they do not starve waiting on VRAM). This lesson reads a real CUDA kernel skeleton, then this lesson's star - tiled matmul, then the memory hierarchy, and finally a conceptual look at flash attention.

Roadmap: first the CUDA execution model (grid / block / thread / warp), then how tiled matmul reuses data through on-chip shared memory, then the three-level memory hierarchy and "moving data costs more than computing", and finally why flash attention fuses several ops into one kernel.

The CUDA execution model: grid / block / thread

CUDA organizes thousands of threads into three levels: one kernel launch is a grid, the grid splits into many blocks, and each block holds tens to thousands of threads. One more unavoidable concept is the warp: every 32 threads form a group that executes the same instruction in perfect lockstep (this model is called SIMT, single instruction multiple threads). Threads in the same block can share a patch of on-chip __shared__ memory and can wait for each other with __syncthreads(); different blocks are largely on their own, with no ordering guaranteed between them.

First a simplest real kernel: multiply every element by a scale. It shows the most basic GPU idiom - work out "who am I" first, then compute only my share:

// real CUDA kernel: each thread computes one/few output elements (simplified from ggml-cuda/scale.cu)
__global__ void scale_f32(const float* x, float* dst, float scale, int64_t n) {
    int64_t tid    = blockIdx.x * blockDim.x + threadIdx.x;  // globally unique thread id
    int64_t stride = blockDim.x * gridDim.x;                 // total threads in one sweep
    for (int64_t i = tid; i < n; i += stride)             // grid-stride loop
        dst[i] = scale * x[i];                              // each thread on its own, no deps
}
// launch: set grid and block size
scale_f32<<<num_blocks, 256>>>(x, dst, scale, n);   // 256 = threads per block

Line by line: blockIdx.x * blockDim.x + threadIdx.x turns "which block, and which thread within it" into a globally unique thread id tid - this is the first line of almost every CUDA kernel. stride is how many elements all threads together cover in one sweep; when there is more data than threads, this grid-stride loop lets each thread hop by stride and handle a few more. The key point: all threads run the same code, landing on different data only through their different tid - that is SIMT. The <<<num_blocks, 256>>> triple-angle-bracket syntax at launch sets how many blocks to start and how many threads per block. One easily-missed but crucial point: a GPU launches far more threads than it has physical cores (tens of thousands routinely) - not waste, but the GPU's signature trick for hiding latency. When one batch of threads stalls waiting on VRAM, the hardware instantly switches to another ready batch, keeping the compute units never idle. So "more threads than you can use" is exactly the point: it fills the memory-wait time with other threads' computation - which is why a GPU favors "small tasks, massive parallelism" rather than the CPU's few strong cores.

Tiled matmul: the heart of the GPU kernel

Matmul is the heaviest op in a large model, and the key to writing it fast on a GPU is again tiling - the same idea as the CPU tiling in L31, but here we tile into shared memory. First, the flaw in the naive way: let each thread independently compute one cell of C, and it must read a whole row of A and a whole column of B from global VRAM (slow). Neighboring threads read heavily overlapping data, yet each reads its own, wasting the already-scarce VRAM bandwidth - the bottleneck is again "waiting for data", not "computing". A concrete sense: the naive version computing an N x N output reads on the order of N-cubed values from slow VRAM (every cell rescans a whole row of A and a whole column of B, with heavy duplicate reads between neighboring threads); the multiply-adds the machine is actually good at are fast, so most of the time goes to "waiting for data". That is exactly the core tension tiling resolves.

The tiled fix has a block's threads cooperate: together they first haul a small tile of A and a small tile of B from global into the block's shared __shared__ memory (on-chip, one-to-two orders of magnitude faster), wait with __syncthreads() until everyone has loaded, then every thread in the block draws from that fast memory repeatedly to compute its own cell, advancing tile by tile along K and accumulating. Below is a normalized teaching skeleton (the real mmq.cuh adds 4-bit unpacking and heavy templating on top, but the skeleton is exactly this):

// teaching tiled-matmul skeleton (real ggml-cuda/mmq.cuh adds 4-bit unpack + templates)
__global__ void matmul_tiled(const float* A, const float* B, float* C, int N) {
    __shared__ float As[T][T];                 // on-chip shared: a tile of A
    __shared__ float Bs[T][T];                 // on-chip shared: a tile of B
    int row = blockIdx.y*T + threadIdx.y;
    int col = blockIdx.x*T + threadIdx.x;
    float acc = 0.0f;                            // accumulator in a register (fastest)
    for (int k0 = 0; k0 < N; k0 += T) {         // advance along K, tile by tile
        As[threadIdx.y][threadIdx.x] = A[row*N + (k0+threadIdx.x)]; // cooperative load
        Bs[threadIdx.y][threadIdx.x] = B[(k0+threadIdx.y)*N + col];
        __syncthreads();                       // wait until the whole tile is loaded
        for (int k = 0; k < T; ++k)             // reuse the tile, all in fast memory
            acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        __syncthreads();                       // wait before overwriting next tile
    }
    C[row*N + col] = acc;                        // write each cell back just once
}

The crux is those two __syncthreads(): the first guarantees "the whole tile is on the bench" before any compute starts, the second guarantees "everyone is done" before overwriting the next tile - drop either and some thread reads half-loaded data and the result is all wrong. Once a tile is in shared memory it is reused many times by all threads in the block, amortizing the slow global accesses down to almost nothing. This carries last lesson's "load once, reuse many" tiling onto the GPU's on-chip memory: the CPU reuses cache, the GPU reuses shared memory, but the memory-saving arithmetic is identical. A rough sense of the payoff: once a T x T tile is in shared, the block's threads reuse it as rows/columns about T times - so each value hauled from global is reused roughly T-fold; with T = 32, slow VRAM traffic drops to about one thirtieth. (T is the tile width, often 16 or 32; blockDim is set to T x T so one thread owns exactly one cell.)

This skeleton is not something we invented - it is the structure of the real mmq.cuh (quantized matmul) main loop. Pull out the core few lines of the real code and compare:

// real mmq main loop (simplified from ggml-cuda/mmq.cuh mul_mat_q_process_tile)
float sum[...] = {0.0f};                          // accumulators in registers
for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
    load_tiles(x, tile_x, offset_x + kb0, ...); // load A's tile into shared
    // ... load B's tile into shared (tile_y) too ...
    __syncthreads();                          // wait for the tile to load
    vec_dot(tile_x, tile_y, sum, 0);            // reuse in-block, accumulate
    __syncthreads();                          // wait before overwriting next tile
}

The structure is identical to the teaching skeleton above: loop tiles -> load into shared -> __syncthreads -> accumulate -> __syncthreads. Almost all the real version's extra complexity lives inside the two called functions - load_tiles (a per-quant-format unpack, one for Q4_0/Q4_K/...) and vec_dot (int8 dot via dp4a, or mma on tensor cores) - while the skeleton itself is exactly the few lines you already understand.

Freezing the process into one picture is clearest - a block loads A's row-tile and B's col-tile into shared, and each thread in the block owns one cell of the output C:

The memory hierarchy: why "moving data" costs more than "computing"

We kept saying "fast memory" and "slow VRAM" above; here is the clear version. A GPU has three levels of memory, each smaller and faster as you go down: global (VRAM) is several GB, visible to all threads, but high-latency, and its bandwidth is often the whole kernel's bottleneck; shared (on-chip) is only tens of KB per block, shared just among that block's threads, one-to-two orders of magnitude faster than global; register is private to each thread and fastest - the accumulator acc in tiled matmul lives here. The entire point of tiled matmul is to move data from global up into shared and registers and reuse it.

This confirms L30's observation once more: many kernels are "memory-bound", not "compute-bound" - a GPU has so much compute it often goes unused, and what actually stalls it is the bandwidth feeding data in from VRAM. So most of GPU optimization goes into "fewer global accesses, more shared/register reuse". Grasp this and you see why, on the same card, a well-written kernel can be several times faster than a naive one - the total work is unchanged; all the savings are in movement. It is also why batching (feeding many sequences/tokens at once) almost always raises GPU utilization: bring data in once, compute on it several times, and get more value out of that one bandwidth. One more GPU-specific subtlety: when accessing global VRAM, the 32 threads of a warp should ideally read contiguous neighboring addresses, so the hardware can merge them into one large memory transaction (coalesced access); if each reads a scattered address, it degrades into dozens of tiny transactions and bandwidth utilization collapses. So a GPU kernel must not only "read global less" but also "read it tidily" - which is why a tensor's layout in VRAM matters so much for performance, and why many kernels rearrange data specifically to satisfy this rule.

Flash attention: fusing several ops into one kernel

Attention computes softmax(Q @ K^T) @ V. The naive way first computes the whole huge N x N score matrix and writes it to VRAM, reads it back for softmax, reads it back again to multiply V - just writing that N x N matrix once and reading it three times drains the bandwidth (and costs a big slab of VRAM besides). Flash attention fuses these steps into one kernel: compute in tiles, correcting with an "online softmax" as you go, never materializing the whole N x N score matrix to VRAM, writing out only the final output O.

// flash attention idea (concept pseudocode; real impl in ggml-cuda/fattn*.cu)
acc = 0;  run_max = -inf;  run_sum = 0;
for each K/V tile:                  // tiled; never materialize the full N x N scores
    s   = Q @ K_tile^T               // scores for this tile only
    m   = max(run_max, rowmax(s))    // update the running row max
    p   = exp(s - m)                 // online softmax: correct as you go
    run_sum = run_sum * exp(run_max - m) + rowsum(p)
    acc     = acc     * exp(run_max - m) + p @ V_tile
    run_max = m
O = acc / run_sum                    // normalize once; write only O, not the N x N

The trick is that "online softmax": softmax normally needs a whole row before it can fix the denominator, but flash attention keeps a running max and a running denominator (run_max / run_sum), and on each new tile rescales the accumulated result to the new max - so it can march forward tile by tile without waiting for the full row. The cost is a few extra rescale operations; the reward is never shuttling that N x N matrix in and out of VRAM - the longer the sequence, the more you save. That is the power of "op fusion": work that would otherwise take several kernels reading and writing VRAM is merged into one kernel done in a single pass, and all the savings are in the most expensive part - memory traffic. This matters most for long context: double the sequence length and that N x N matrix quadruples in area, so the naive approach's VRAM and bandwidth costs explode, whereas flash attention, never materializing it, grows only linearly with the sequence - one of the key reasons today's tens-of-thousands to hundreds-of-thousands token contexts are feasible at all.

Deeper: warp-level sum and kernel variants

Two last folds for two GPU-kernel details that are common yet easy to find baffling: how a warp sums quickly without shared memory, and why one op has so many kernels.

// sum 32 threads within a warp, no shared memory (from ggml-cuda/common.cuh)
__device__ float warp_reduce_sum(float x) {
    for (int offset = 16; offset > 0; offset >>= 1)   // 16 -> 8 -> 4 -> 2 -> 1
        x += __shfl_xor_sync(0xffffffff, x, offset, 32); // read the neighbor "offset away"
    return x;                                          // afterwards every thread holds the total
}