GPU Programming Model, Simulated · AI infra interview lesson

Why It Matters

Thirty-two threads reading adjacent floats touch different sectors than widely spaced reads. CPU models make access patterns visible; actual speed still requires GPU measurements.

FP32 vector addition performs one add for 8 bytes read and 4 bytes written: 1/12 FLOP per byte. Its Roofline ceiling must use the matching FP32 execution units, not FP16 Tensor Core throughput. An aligned warp reading consecutive FP32 values needs four 32-byte sectors; stride 32 may need 32 sectors. Eight times as many sectors does not imply an eightfold kernel slowdown: caches, reuse, instructions and other bottlenecks affect timing. These CPU models estimate accesses and resources, providing hypotheses for GPU measurements.

Walk the Shapes First

N = 10 elements, blockDim = 4, plus a warp shrunk to 8 lanes:

Rules read from the table: the thread count is rounded up to whole blocks, so idx < N must mask the extra threads; each warp-reduce step halves the active lanes and the result lands only in lane 0; global access is charged per sector, shared memory per number of distinct words on the same bank.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: launch simulation, bounds check, grid-stride loop

New idea: a launch runs the same function for every (blockIdx, threadIdx), and a thread finds its data through idx = blockIdx * blockDim + threadIdx. A grid-stride loop lets a fixed number of threads cover an array of any length: idx += blockDim * gridDim.

import torch, math
torch.manual_seed(0)

def simulate_launch(grid_dim, block_dim, kernel, *args):
    # a launch without __syncthreads: every (block, thread) runs the same function
    for b in range(grid_dim):
        for t in range(block_dim):
            kernel(b, t, block_dim, grid_dim, *args)

def vector_add_kernel(blockIdx, threadIdx, blockDim, gridDim, a, b, c, n):
    idx = blockIdx * blockDim + threadIdx
    if idx < n:                       # the last block has idle threads
        c[idx] = a[idx] + b[idx]

Continue with a grid-stride loop

The launch may have fewer threads than elements. After processing idx, a thread advances by the total grid size. Together the threads cover the array while still checking the tail boundary.

def vector_add_grid_stride(blockIdx, threadIdx, blockDim, gridDim, a, b, c, n):
    idx = blockIdx * blockDim + threadIdx
    stride = blockDim * gridDim       # number of threads in the whole grid
    while idx < n:
        c[idx] = a[idx] + b[idx]
        idx += stride

n = 1000
a, b = torch.randn(n), torch.randn(n)
block = 256
grid = (n + block - 1) // block       # ceil_div: 4 blocks, 1024 threads, 24 idle
c = torch.zeros(n)
simulate_launch(grid, block, vector_add_kernel, a, b, c, n)
assert torch.allclose(c, a + b)

c2 = torch.zeros(n)
simulate_launch(2, 128, vector_add_grid_stride, a, b, c2, n)   # 256 threads cover 1000 elements
assert torch.allclose(c2, a + b)

# vectorized view of the same index math: [grid, block] index grid + mask
idx = torch.arange(grid)[:, None] * block + torch.arange(block)[None, :]
mask = idx < n
assert idx.shape == (4, 256) and int(mask.sum()) == n and int((~mask).sum()) == 24

The 2D case is the same: row = blockIdx.y * blockDim.y + threadIdx.y, col = blockIdx.x * blockDim.x + threadIdx.x, each with its own bounds check. x is the fastest-varying dimension, so map threadIdx.x to contiguous columns in memory; Step 3 shows why.

Step 2: warp shuffle, warp reduce, divergence count of block reduce

New idea: the 32 lanes of a warp can read each other's registers directly. __shfl_down_sync(v, offset) gives lane i the value of lane i+offset, and out-of-range lanes keep their own value; __shfl_xor_sync(v, mask) gives lane i the value of lane i ^ mask. Reduction formulas:

• shfl_down: v += shfl_down(v, o) for o = 16, 8, 4, 2, 1; the result is only in lane 0
• xor butterfly: v += shfl_xor(v, o), also 5 steps; every lane gets the total

A warp is a [num_warps, 32] tensor here, one row per warp.

import torch
torch.manual_seed(0)
WARP = 32

def shfl_down(vals, offset):
    # vals: [num_warps, 32]; lane i reads lane i+offset; out-of-range lanes keep their own value
    lane = torch.arange(WARP)
    src = lane + offset
    out = vals.clone()
    ok = src < WARP
    out[:, ok] = vals[:, src[ok]]
    return out

def shfl_xor(vals, mask):
    lane = torch.arange(WARP)
    return vals[:, lane ^ mask]

def warp_reduce_sum(vals):
    for offset in (16, 8, 4, 2, 1):
        vals = vals + shfl_down(vals, offset)
    return vals[:, 0]                 # only lane 0 holds the full sum

def warp_allreduce_sum(vals):
    for mask in (16, 8, 4, 2, 1):
        vals = vals + shfl_xor(vals, mask)
    return vals                       # every lane holds the full sum

x = torch.randn(5, WARP)
assert torch.allclose(warp_reduce_sum(x), x.sum(1), atol=1e-5)
assert torch.allclose(warp_allreduce_sum(x), x.sum(1, keepdim=True).expand(-1, WARP), atol=1e-5)
# after the shfl_down loop, lanes other than 0 hold partial sums only
v = x.clone()
for offset in (16, 8, 4, 2, 1):
    v = v + shfl_down(v, offset)
assert torch.allclose(v[:, 0], x.sum(1), atol=1e-5) and not torch.allclose(v[:, 1], x.sum(1))

Inside a block the reduction is a tree in shared memory. AIInfraGuide's Reduce chapter optimizes it version by version: V0 picks active threads with tid % (2s) == 0 (interleaved addressing, heavy divergence); V1 uses index = 2*s*tid, removing divergence but adding bank conflicts; V2 switches to tid < s, halving from s = blockDim/2 (sequential addressing); V3 does one add while loading, so fewer threads sit idle; V4 unrolls the last warp; V5 fully unrolls the loop with templates; V6 replaces the last levels with warp shuffle; V7 adds float4 vectorized loads and a grid-stride loop.

The code simulates V0 and V2 and counts, per round, the warps in which only some lanes are active (those warps diverge). When N is not a multiple of the block size, out-of-range threads load the identity: 0 for sum, -inf for max.

import torch, math
torch.manual_seed(0)

def divergent_warps(active):
    w = active.view(-1, 32)
    return int((w.any(1) & ~w.all(1)).sum())

def block_reduce(x, block=256, mode="sequential", op=torch.add, identity=0.0):
    n = x.numel()
    grid = (n + block - 1) // block
    smem = torch.full((grid * block,), identity, dtype=x.dtype)
    smem[:n] = x                                  # out-of-range threads load the identity
    smem = smem.view(grid, block)
    tid = torch.arange(block)
    strides = [1 << i for i in range(int(math.log2(block)))]   # 1, 2, ..., block/2
    if mode == "sequential":
        strides = strides[::-1]                   # block/2, ..., 1
    div = 0
    for s in strides:
        active = (tid % (2 * s) == 0) if mode == "interleaved" else (tid < s)
        div += divergent_warps(active)
        dst = tid[active]
        smem[:, dst] = op(smem[:, dst], smem[:, dst + s])   # one __syncthreads() per round
    return smem[:, 0], div                        # one partial result per block

x = torch.randn(1000)
p0, div0 = block_reduce(x, mode="interleaved")
p2, div2 = block_reduce(x, mode="sequential")
assert p0.shape == (4,) and torch.allclose(p0.sum(), x.sum(), atol=1e-4)
assert torch.allclose(p2.sum(), x.sum(), atol=1e-4)
assert (div0, div2) == (47, 5)                    # V0 diverges in almost every round, V2 only when s < 32

# max-reduce must pad with -inf, not 0
neg = -torch.rand(1000) - 1.0                     # all values in [-2, -1)
pm, _ = block_reduce(neg, op=torch.maximum, identity=float("-inf"))
assert pm.max() == neg.max()
pz, _ = block_reduce(neg, op=torch.maximum, identity=0.0)
assert pz.max() == 0.0                            # padding with 0 gives a wrong answer

In V2 only the 5 rounds s = 16, 8, 4, 2, 1 leave warp 0 partially active. Those are exactly the rounds V4 unrolls and V6 replaces with shuffles.

Step 3: sector counting, bank conflicts, padding and swizzle

New idea: global memory moves data in 32 B sectors, and one warp access costs the number of sectors it covers: sectors = len(unique(byte_add… (Pro lesson)

Separate global access from shared-memory conflicts

Sector counts describe global-memory transactions. Shared-memory bank conflicts concern different addresses mapping to one bank. (Pro lesson)

Step 4: an occupancy calculator

New idea: the number of blocks resident on one SM is the minimum of several limits: max blocks per SM, max warps, total registers, total sha… (Pro lesson)

Step 5: tiled GEMM (block tiling + thread tiling)

New idea: C[M,N] = A[M,K] @ B[K,N] . In the naive kernel every output element reads 2K numbers from global memory, 2*M*N*K in total. (Pro lesson)

Reuse inputs across multiple outputs per thread

Block tiling shares input across threads; thread tiling gives one thread multiple output accumulators. (Pro lesson)

In Real Systems

Triton : store_kvcache_kernel in nano-vllm's nanovllm/layers/attention.py is a Triton kernel. (Pro lesson)

Common Pitfalls

No bounds check → writing c[idx] = a[idx] + b[idx] without if (idx < n) → the extra threads of the last block write out of bounds; (Pro lesson)

Say It in the Interview

A CUDA launch runs the same function on every (blockIdx, threadIdx), so each thread computes a global index, guards it with idx < n , and a… (Pro lesson)

Hands-on Notebook

The notebook has 7 exercises. Write each implementation first, then run its tests; (Pro lesson)