2 · Just enough GPU to read this code

This chapter explains the handful of hardware facts the kernels are built around. You don't need the whole CUDA programming guide: just how work is split into threads, how bytes reach those threads, and how a tensor core multiplies. Each idea is tied to the line of submission.py that depends on it.

2.1 The machine

2.2 Grid, CTA, warp, lane

A kernel launch creates a grid of thread blocks. NVIDIA's documents also call a thread block a CTA (cooperative thread array); this guide uses "CTA" because the code does.

Each CTA is a group of threads that:

• runs on one SM;
• shares that SM's shared memory;
• can synchronize with __syncthreads().

The threads of a CTA come in warps of 32. A warp executes one instruction at a time for all 32 of its threads ("lanes"). The SM keeps many warps resident and switches between them every cycle. When one warp is waiting for memory, another warp issues instructions. That is how a GPU hides memory latency: lots of independent warps, each with lots of loads in flight.

In our GEMV kernels:

For example, qkv_gate at m = 1 launches 2304/16 = 144 CTAs of 17 warps each.

2.3 The memory hierarchy

Our GEMVs read weights straight from HBM into registers. They never stage weights through shared memory, because every weight byte is used by exactly one warp, exactly once. Shared memory appears only at the end, to add up the warps' partial sums.

2.4 Feeding a memory-bound kernel

To keep 4.8 TB/s flowing, the GPU must have enough bytes "in the air" at every moment. This is Little's law: bytes in flight = bandwidth × latency.

4.8 TB/s × ~0.7 µs  ≈  3.4 MB in flight  ≈  25 KB per SM, all the time

A thread's load takes about a microsecond to come back. So each warp must issue several loads before it needs the first result, and each SM must host enough warps.

Three habits in our code serve this goal:

• Wide loads. Each lane loads 16 bytes at once (uint4, four 32-bit words). 32 lanes × 16 B = 512 bytes per warp instruction.
• Coalescing. When the 32 lanes of a warp read 512 consecutive bytes, the memory system serves them as one big transaction. Scattered addresses would cost many transactions. Chapter 5 shows how we reorder the weights so that every load is perfectly contiguous.
• Loads in flight. The template parameters U (plain GEMV) and UG (compressed GEMV) say how many K blocks a warp loads before using them. PIPE chooses whether the next batch is loaded while the current one is computed. More warps per CTA (NW) means more loads in flight per SM.

The load instruction itself is special:

ld.global.nc.L1::no_allocate.v4.u32

2.5 Tensor cores and mma.sync.m16n8k16

A tensor core multiplies small matrix tiles in one instruction. The variant our GEMVs use is:

mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32   D = A · B + C

Together, the 32 lanes of a warp hold the three tiles. Each lane holds a few elements, called its fragment. That one warp instruction performs 16 × 8 × 16 = 2,048 multiply-adds. Which lane holds which element is fixed by the hardware, with lane = 4·g + t:

How our GEMV uses it ("swap A and B").

• The 16 rows of the A tile are 16 weight rows, and its 16 columns are 16 values of k.
• The 8 columns of the B tile are up to 8 tokens.
• The accumulator therefore holds the outputs of 16 rows × 8 tokens.
• At m = 1, seven of the eight token columns are wasted, but that costs nothing that matters: the kernel is waiting for memory anyway.
• At m = 16 or 32, the kernel runs the mma for 2 or 4 blocks of 8 tokens. That count is NB, computed by nb_of(m).

Why a tensor core for a matrix-vector product? Instruction count. Without it, every 16 bytes of weights would need a dozen conversion and fused multiply-add instructions per lane. Those instructions would compete with the loads for issue slots. One mma per 16 bytes keeps the SM free to keep memory busy.

2.6 Reductions inside a CTA and across CTAs

The K dimension is long (up to 8192), so it is split among the warps of a CTA. Each warp ends with partial sums for the same 16 rows × tokens, which must be added up.

• Inside a CTA:
• every warp writes its partial to shared memory;
• __syncthreads() waits until all have written;
• a few threads add them up.
• Across CTAs ("split-K"):
• When a matrix has few rows, one CTA per 16-row block would leave SMs idle. The router has 256 rows, which is only 16 CTAs for 132 SMs.
• So the K dimension is also split across S CTAs, which must combine their partials through global memory.
• Chapter 6 (gemv_kernel) and chapter 9 (Triton) show two ways to do that in one launch. Lab 5 simulates both.

2.7 Templates: one compiled kernel per configuration

The kernels are C++ templates, such as gemv_kernel<NW, U, NB, DUAL, PIPE>, because:

• Arrays like float acc[NB][4] stay in registers only if their size and every index are known at compile time. A runtime index would push them to slow local memory.
• #pragma unroll loops with compile-time trip counts become straight-line code, with no loop overhead and with loads that can be scheduled early.

The price is that every configuration is a separate compiled kernel. submission.py therefore lists exactly the combinations it uses (the C3(...) and Z3(...) lines, chapter 6). Asking for any other combination at run time raises "unsupported gemv config".

2.8 PTX, asm volatile and SASS

• PTX is NVIDIA's virtual instruction set.
• asm volatile("...") lets CUDA C++ emit an exact PTX instruction. We use it for the streaming load, mma.sync, and the byte-shuffling prmt and lop3 (chapter 7).
• The driver compiles PTX into SASS, the real machine code, which you can inspect with cuobjdump -sass.

2.9 CUDA graphs and launch cost

Launching a kernel from Python costs several microseconds of CPU time. A CUDA graph records the launches once and replays them, so the judge measures GPU time only. Two consequences for the code:

1. kernel_fn must not wait for the GPU or read GPU values on the CPU: no .item(), no printing tensors. The graph cannot contain that.
2. Every kernel still pays roughly 1–2 µs of fixed cost inside the GPU: start-up, ramp-up, and the tail where the last CTAs finish. This is a big deal for the tiny router (floor 1.8 µs, ours 5.4 µs).

cd ~/projects/okc/muse_bf16_proj/learn/labs
python3 lab2_fragments.py

Part 1 of the lab prints the three ownership maps above and checks an emulated mma against numpy. Part 2 belongs to chapter 5; you can run it now or later.

Check yourself