GPU Kernels & Compilers: A Compressed Interview Course

What they actually test START HERE

A kernel/compiler interview is not a LeetCode screen. The interviewer is probing whether you can reason about where the time and the bytes go, and whether you can move fluidly between four layers of the stack without losing the thread.

The five questions behind every question

1. Where are the FLOPs and where are the bytes? Is this op compute-bound, memory-bound, or communication-bound? (Module 04.)
2. What does the memory hierarchy force? Tiling, fusion, recompute-vs-store — all fall out of HBM << SRAM << registers.
3. What does the hardware demand? Tensor-core shapes, alignment, coalescing, occupancy, precision/scaling.
4. What can the compiler do for you, and where does it stop? Fusion boundaries, why some patterns need a hand-written kernel.
5. How do you know it's right and fast? Numerical checks vs a reference, autotuning, regression tests, profiling.

Linear algebra you actually need FOUNDATION

You don't need spectral theory for these interviews. You need to be fluent in shapes, reductions, and the fact that almost everything reduces to matmul.

The four core operations

Op	Shapes	FLOPs	Arithmetic intensity
Dot product	(K)·(K) → scalar	2K	~1 (memory-bound)
Matrix–vector (GEMV)	(M,K)·(K) → (M)	2MK	~1–2 (memory-bound)
Matrix–matrix (GEMM)	(M,K)·(K,N) → (M,N)	2MNK	scales with tile size (compute-bound)
Outer product	(M)⊗(N) → (M,N)	MN	<1 (very memory-bound)

Why “everything is a matmul”

A linear layer is \(Y = XW^\top + b\). A convolution lowers to a matmul (im2col / implicit GEMM). Attention is two matmuls with a softmax between them. This is why a fast, well-scheduled GEMM is the single most valuable kernel on the planet — and why tensor cores exist to do nothing else.

Matrix multiplication & tiling CORE KERNEL

A naive GEMM is correct and slow: each output element re-reads a full row of \(A\) and column of \(B\) from HBM. The whole game is reuse — load a block once into fast memory, do as much math as possible on it, then move on.

The memory hierarchy that forces tiling

Level	~Latency	~Bandwidth	Size
Registers	~1 cycle	~10s TB/s	~256 KB / SM
Shared mem / SRAM	~20 cycles	~10s TB/s	~100–228 KB / SM
L2 cache	~200 cycles	several TB/s	tens of MB
HBM (global)	~400–600 cycles	~1.5–8 TB/s	tens of GB

The jump from HBM to SRAM is ~10–100×. Tiling exists to pay the HBM cost once per block and amortize it over \(O(\text{tile})\) FLOPs.

A real tiled GEMM in Triton

Triton lets you write the tile schedule in Python while it handles vectorization, shared-memory staging, and tensor-core selection. This is the canonical kernel to be able to sketch.

import triton
import triton.language as tl

@triton.jit
def matmul_kernel(A, B, C, M, N, K,
                  sam, sak, sbk, sbn, scm, scn,
                  BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BM + tl.arange(0, BM)         # rows of C this block owns
    offs_n = pid_n * BN + tl.arange(0, BN)         # cols of C this block owns
    offs_k = tl.arange(0, BK)

    a_ptrs = A + offs_m[:, None] * sam + offs_k[None, :] * sak
    b_ptrs = B + offs_k[:, None] * sbk + offs_n[None, :] * sbn

    acc = tl.zeros((BM, BN), dtype=tl.float32)     # accumulate in FP32
    for k in range(0, K, BK):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k, other=0.0)
        acc += tl.dot(a, b)                         # ← maps to tensor-core MMA
        a_ptrs += BK * sak
        b_ptrs += BK * sbk

    c = acc.to(C.dtype.element_ty)
    c_ptrs = C + offs_m[:, None] * scm + offs_n[None, :] * scn
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)

The whole kernel is: pick your output tile → loop over K loading A/B tiles → tl.dot accumulates in FP32 → store once. tl.dot is what the backend turns into mma/wgmma/tcgen05.

Tensor cores: the MMA the tl.dot becomes

Tensor cores execute a small fixed-shape matrix-multiply-accumulate (MMA) per instruction. You don't get to pick arbitrary shapes — the hardware shape constrains your tiling.

Generation	Instruction	Issued by	Representative shape (m·n·k)
Ampere	mma	1 warp	m16 n8 k16 (or k8)
Hopper	wgmma	warpgroup (4 warps)	m64 nN k16, N∈{16…256}
Blackwell	tcgen05.mma	single thread → tensor memory	up to m256 n256 k16 across 2 SMs

GPU execution model FOUNDATION

You can't reason about a kernel without the cost model underneath it. The GPU is a throughput machine: it hides latency with massive parallelism rather than caches and out-of-order tricks.

The five things that govern speed

• Occupancy — enough resident warps per SM to hide HBM latency. Bounded by registers/thread and shared memory/block. Too many registers → fewer warps → stalls show through.
• Coalescing — the 32 lanes of a warp should touch contiguous global addresses so the hardware merges them into a few 128-byte transactions. Strided/scattered access multiplies your byte traffic.
• Shared-memory bank conflicts — SMEM has 32 banks; if lanes hit the same bank with different addresses, accesses serialize. Padding / swizzling fixes it.
• Divergence — an if that splits a warp executes both sides masked. Branchy per-lane code wastes throughput.
• Latency hiding — the GPU tolerates ~500-cycle HBM latency by having other warps ready to issue. This is why “make it more parallel” often beats “make each thread faster.”

Roofline & the three bounds MENTAL MODEL

The roofline is the single most useful diagram in this whole field, and it's the lens behind half the interview questions. It answers: is my kernel limited by compute, by memory bandwidth, or by communication?

The three bounds, by example

Bound	Limited by	Typical ops	The fix
Compute	peak FLOP/s (tensor cores)	large square GEMM, conv	better tiling, higher MMA utilization, lower precision
Memory	HBM bandwidth	softmax, LayerNorm, bias+act, GEMV, attention's I/O	fuse to cut round-trips; recompute instead of store
Communication	NVLink / network BW & latency	all-reduce, all-to-all (MoE), pipeline/tensor parallel	overlap comm with compute; reduce message volume; topology-aware collectives

Fused element-wise kernels FUSION 101

Elementwise and reduction ops are memory-bound (Module 04). The lever is fusion: do multiple ops in one pass so intermediate tensors never touch HBM.

Fused bias + activation (epilogue)

The pattern: compute the GEMM tile, then apply bias and the activation before the single store. As a standalone elementwise kernel it looks like this; in practice you splice it into the matmul epilogue.

@triton.jit
def bias_gelu_kernel(X, Bias, Y, M, N, sxm, sxn, BM: tl.constexpr, BN: tl.constexpr):
    pid_m, pid_n = tl.program_id(0), tl.program_id(1)
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    ptr = X + offs_m[:, None] * sxm + offs_n[None, :] * sxn

    x = tl.load(ptr, mask=mask, other=0.0).to(tl.float32)
    b = tl.load(Bias + offs_n, mask=offs_n < N, other=0.0).to(tl.float32)
    x = x + b[None, :]
    # tanh approximation of GELU
    y = 0.5 * x * (1.0 + tl.math.tanh(0.7978845608 * (x + 0.044715 * x * x * x)))
    tl.store(Y + offs_m[:, None] * sxm + offs_n[None, :] * sxn,
             y.to(Y.dtype.element_ty), mask=mask)

Note FP32 math on BF16/FP16 storage — compute in high precision, store in low. That theme repeats everywhere.

Row reduction — the other primitive

Softmax, LayerNorm, RMSNorm, and loss functions are all row reductions followed by an elementwise pass. The skill is doing the reduction in one block with a tree/warp reduction, keeping the row resident.

@triton.jit
def row_sum_kernel(X, Out, N, BN: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BN)
    x = tl.load(X + row * N + cols, mask=cols < N, other=0.0)
    tl.store(Out + row, tl.sum(x, axis=0))   # tl.sum lowers to a warp/tree reduction

Softmax, LayerNorm, RMSNorm FUSED FORWARD

These three are the canonical fused-reduction kernels and they show up constantly. Know the math, the numerical-stability trick, and the one-pass kernel for each.

Numerically-stable softmax

@triton.jit
def softmax_kernel(X, Y, N, BN: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BN)
    mask = cols < N
    x = tl.load(X + row * N + cols, mask=mask, other=-float("inf")).to(tl.float32)
    x = x - tl.max(x, axis=0)              # stability shift
    num = tl.exp(x)
    y = num / tl.sum(num, axis=0)
    tl.store(Y + row * N + cols, y.to(Y.dtype.element_ty), mask=mask)

Online (streaming) softmax — the FlashAttention seed

When the row doesn't fit in SRAM, you process it in blocks while maintaining a running max \(m\) and running denominator \(\ell\), rescaling the accumulator when the max grows. This recurrence is exactly what makes FlashAttention possible (Module 09).

LayerNorm vs RMSNorm

	LayerNorm	RMSNorm
Formula	\(y=\dfrac{x-\mu}{\sqrt{\sigma^2+\epsilon}}\,\gamma+\beta\)	\(y=\dfrac{x}{\sqrt{\frac1d\sum_i x_i^2+\epsilon}}\,\gamma\)
Stats needed	mean and variance (2 reductions)	mean-square (1 reduction)
Recenters?	yes (subtracts μ)	no
Params	γ and β	γ only
Why used	classic, stable	cheaper, ~same quality — default in LLaMA-style models

@triton.jit
def rmsnorm_kernel(X, W, Y, N, eps, BN: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BN); mask = cols < N
    x = tl.load(X + row * N + cols, mask=mask, other=0.0).to(tl.float32)
    ms = tl.sum(x * x, axis=0) / N                 # mean of squares (1 pass)
    inv = 1.0 / tl.sqrt(ms + eps)
    w = tl.load(W + cols, mask=mask, other=0.0).to(tl.float32)
    y = x * inv * w
    tl.store(Y + row * N + cols, y.to(Y.dtype.element_ty), mask=mask)

RMSNorm drops the mean subtraction and β — one reduction instead of two, no second statistic to carry. That's the whole reason it's popular.

Low precision & block scaling NUMERICS

Lower precision means more FLOP/s and less HBM traffic — it moves you up and right on the roofline. The catch is dynamic range and quantization error. This module is the numerics an interviewer will push hardest on for a 2025-era kernel role.

Float anatomy & the format zoo

A float is sign · exponent · mantissa. Exponent bits buy range; mantissa bits buy precision. Every low-precision format is a different bet on that trade.

Format	Bits (S/E/M)	Buys	Used for
FP32	1 / 8 / 23	reference precision	accumulation, master weights
TF32	1 / 8 / 10	FP32 range, less mantissa	tensor-core matmul inputs (Ampere+)
BF16	1 / 8 / 7	FP32 range, low precision	training (no loss scaling needed)
FP16	1 / 5 / 10	precision, narrow range	inference; training needs loss scaling
FP8 E4M3	1 / 4 / 3	more precision	fwd activations & weights
FP8 E5M2	1 / 5 / 2	more range	gradients (backward)
FP4 E2M1	1 / 2 / 1	4 bits → 16 values total	weights/activations with block scaling

FP16/BF16 matmul from first principles

The non-negotiable rule: inputs low precision, accumulate in FP32. Tensor cores are built for exactly this — they take FP16/BF16 (or FP8) operands and accumulate into an FP32 register tile. Summing thousands of products in BF16 would lose the small terms; FP32 accumulation preserves them.

In PTX the accumulator type is explicit — the .f32 in the MMA opcode is the FP32 register tile:

// Ampere MMA: BF16 inputs, FP32 accumulate
mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32  {d0,d1,d2,d3}, {a..}, {b..}, {c0,c1,c2,c3};
//                         shape            acc  A    B    acc

FP8, FP4, and why FP4 requires scaling

FP4 E2M1 has just 16 representable values spanning roughly ±6. You cannot represent a realistic weight distribution with that directly — so you scale. Each small block of values shares a scale factor that maps the block's range into FP4's tiny range; you store the 4-bit values plus the per-block scale.

MXFP4 vs NVFP4 — know the difference cold

	MXFP4 (OCP microscaling)	NVFP4 (Blackwell)
Element	FP4 E2M1	FP4 E2M1
Block size	32 values	16 values
Scale format	E8M0 (power-of-two only)	FP8 E4M3 (fractional) + per-tensor FP32
Levels of scaling	one (per-block)	two (per-block + per-tensor)
Error	higher (coarser block, pow-2 scale)	lower (finer block, fractional scale)
Backing	open standard (AMD, Intel, NV, …)	NVIDIA Blackwell-native

NVFP4's two wins: a smaller block (16) narrows the dynamic range each scale must cover, and a fractional E4M3 scale (not just powers of two) fits that range more tightly. Both reduce quantization error; reports put MXFP4 at needing materially more training tokens to match NVFP4 loss.

Tensor-core constraints & block-scaled matmul

Low-precision MMAs impose hard rules your kernel must satisfy:

• Format pairs are fixed — the MMA opcode names operand and accumulator types; you can't feed arbitrary combinations.
• K must align to the MMA K and to the scaling block — e.g. FP8/FP4 paths want K as a multiple of the block (16/32) so each MMA step consumes whole blocks with their scale.
• Scales have a prescribed layout — Blackwell's block-scaled MMA reads the scale tensor in a specific interleaved order; getting that layout wrong is a top cause of “correct in NumPy, garbage on device.”

# Block-scaled matmul, conceptual inner loop (FP4 weights + per-block scales)
acc = tl.zeros((BM, BN), tl.float32)
for k in range(0, K, BLOCK):                  # BLOCK = 16 (NVFP4) or 32 (MXFP4)
    a4 = load_fp4(A, ...)                      # packed 4-bit
    b4 = load_fp4(B, ...)
    sa = load_scale(Ascale, ...)              # one scale per block
    sb = load_scale(Bscale, ...)
    a = dequant(a4) * sa                       # → higher precision
    b = dequant(b4) * sb
    acc += tl.dot(a, b)                        # FP32 accumulate
# On Blackwell this whole block+scale path is a single hardware-scaled MMA.

Matmul in practice SHAPES · TUNING · TESTS

The same GEMM code performs wildly differently across shapes. Recognizing the shape regime — and tuning, then locking it down with tests — is the practical core of the job.

Three shape regimes

Shape	Example	Regime	Strategy
Square	4096³	compute-bound	large tiles, max tensor-core util
Tall-skinny	(65536, 128)·(128, 256)	K too small to fill tile	split-K / stream-K to expose parallelism
MLP prefill	(B·S, d)·(d, 4d)	compute-bound (large B·S)	fuse bias+act epilogue (Module 05)
Decode (M=1)	(1, d)·(d, 4d)	memory-bound (GEMV)	weight-only quant; batch requests

Autotuning

The best BM/BN/BK, num_warps, and pipeline num_stages depend on shape, dtype, and architecture. You don't guess — you search, then cache the winner per shape key.

configs = [
    triton.Config({'BM':128,'BN':256,'BK':64}, num_warps=8, num_stages=3),
    triton.Config({'BM':128,'BN':128,'BK':64}, num_warps=8, num_stages=4),
    triton.Config({'BM':64, 'BN':128,'BK':64}, num_warps=4, num_stages=4),
    triton.Config({'BM':64, 'BN':64, 'BK':32}, num_warps=4, num_stages=3),
]
@triton.autotune(configs=configs, key=['M','N','K'])   # re-tune when shape changes
@triton.jit
def matmul_kernel(...):
    ...

Numerical checks & regression testing

A kernel isn't done until it's provably correct against a reference and guarded against regressions. Tolerances must match the dtype — BF16 results are not bit-exact to FP32.

import torch

def check(M, N, K, dtype):
    a = torch.randn(M, K, device='cuda', dtype=dtype)
    b = torch.randn(K, N, device='cuda', dtype=dtype)
    ref = (a.float() @ b.float())                  # FP32 reference
    out = triton_matmul(a, b).float()
    # tolerance scales with K (accumulated rounding) and with dtype
    rtol = {torch.float16: 1e-2, torch.bfloat16: 2e-2, torch.float32: 1e-4}[dtype]
    torch.testing.assert_close(out, ref, rtol=rtol, atol=rtol * K**0.5)

# regression matrix: shapes × dtypes, run in CI; also assert TFLOP/s ≥ baseline
for dt in (torch.float16, torch.bfloat16):
    for (M,N,K) in [(4096,4096,4096), (65536,256,128), (1,4096,4096)]:
        check(M, N, K, dt)

Attention & FlashAttention THE BIG ONE

Attention is two matmuls with a softmax between them. The whole performance story is that the middle matrix is \(S\times S\) — quadratic in sequence length — and you must never write it to HBM.

FlashAttention: tile + online softmax, never materialize S×S

FlashAttention keeps a block of \(Q\) in SRAM and streams blocks of \(K,V\) past it, updating the output with the online-softmax recurrence from Module 06. The \(S\times S\) matrix exists only one tile at a time, in SRAM.

FlashAttention forward kernel (Triton-style)

@triton.jit
def flash_attn_fwd(Q, K, V, O, scale, S, D,
                   BQ: tl.constexpr, BK: tl.constexpr, CAUSAL: tl.constexpr):
    qb = tl.program_id(0)                              # this program owns one Q block
    offs_q = qb * BQ + tl.arange(0, BQ)
    offs_d = tl.arange(0, D)
    q = tl.load(Q + offs_q[:, None]*D + offs_d[None, :])   # BQ×D stays resident

    m = tl.full((BQ,), -float("inf"), tl.float32)     # running row max
    l = tl.zeros((BQ,), tl.float32)                   # running denominator
    acc = tl.zeros((BQ, D), tl.float32)               # running output

    kv_end = (qb + 1) * BQ if CAUSAL else S           # causal: skip future blocks
    for kb in range(0, kv_end, BK):
        offs_k = kb + tl.arange(0, BK)
        k = tl.load(K + offs_k[:, None]*D + offs_d[None, :])
        v = tl.load(V + offs_k[:, None]*D + offs_d[None, :])

        s = tl.dot(q, tl.trans(k)) * scale            # BQ×BK scores (in SRAM only)
        if CAUSAL:                                    # mask the diagonal block
            s = tl.where(offs_q[:, None] >= offs_k[None, :], s, -float("inf"))

        m_new = tl.maximum(m, tl.max(s, axis=1))      # online softmax update
        p = tl.exp(s - m_new[:, None])
        alpha = tl.exp(m - m_new)                     # correction for old acc/denom
        l = l * alpha + tl.sum(p, axis=1)
        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
        m = m_new

    acc = acc / l[:, None]                            # finalize
    tl.store(O + offs_q[:, None]*D + offs_d[None, :], acc.to(O.dtype.element_ty))

The three carried values m, l, acc and the alpha rescale are the entire trick. Drop the FP32 accumulators or the rescale and you get NaNs or wrong sums.

Causal masking, head dim, and seq-len sweeps

• Causal masking: token \(i\) attends only to \(j\le i\). Whole K/V blocks above the diagonal are skipped (the kv_end bound) — roughly half the work — and only the diagonal block needs the per-element triangular mask. Forgetting to skip the upper blocks leaves an easy 2× on the table.
• Head dimension \(d\): usually 64 or 128. It sets the SRAM tile footprint (Q,K,V blocks are \(B\!\cdot\!d\)); larger \(d\) raises arithmetic intensity but shrinks how many rows fit per block. Small \(d\) attention is more memory-bound.
• Sequence-length sweep: standard attention's memory grows \(O(S^2)\) and OOMs; FlashAttention's memory is \(O(S)\) while compute stays \(O(S^2 d)\). So as \(S\) grows the gap widens — the canonical benchmark plot you should be able to describe.

Mixture of Experts & distributed SCALE-OUT

MoE replaces a dense FFN with many experts and routes each token to only a few. You get more parameters at constant FLOPs — but you buy a communication problem, and that's where the kernel/systems interview goes.

Routing & top-k gating

logits = x @ Wr                      # [T, E]  token→expert scores
probs  = logits.softmax(-1)
topv, topi = probs.topk(k, dim=-1)   # [T, k] weights and expert ids
topv = topv / topv.sum(-1, keepdim=True)   # renormalize the kept gates
# dispatch: sort/permute tokens so each expert sees a contiguous run
order = topi.view(-1).argsort()

Dispatch / combine, grouped GEMM, all-to-all

• Dispatch: permute tokens so each expert gets a contiguous block; when experts live on other GPUs (expert parallelism), this permute is an all-to-all — every device sends each other device its tokens.
• Grouped GEMM: experts share weight shapes but receive different token counts (ragged M). One kernel does a batch of variable-M GEMMs without launching one kernel per expert and without padding waste.
• Combine: a second all-to-all routes outputs back to the originating device, then the weighted sum applies the gate weights.

Load balancing, expert parallelism, overlap

• Load balancing: routing can pile tokens onto a few popular experts → idle GPUs + stragglers. Fixes: an auxiliary load-balancing loss that pushes the router toward uniform usage, a capacity factor (max tokens/expert, drop or reroute overflow), and noisy/again gating.
• Expert parallelism (EP): shard experts across GPUs. Each token may need an expert on another device → the all-to-alls. EP is usually combined with data/tensor/pipeline parallelism.
• Overlap: the killer optimization. While expert \(e\)'s GEMM runs, prefetch/communicate the next group's tokens — overlap all-to-all with grouped-GEMM compute so the comm-bound dispatch hides behind compute.

Bottleneck taxonomy & a microbenchmark

Stage	Bound	Why	Lever
Router top-k	cheap / memory	tiny GEMM + sort	rarely the bottleneck
Dispatch all-to-all	communication	cross-device token shuffle	overlap; fewer bytes; topology-aware
Expert grouped GEMM	compute (prefill) / memory (decode)	big GEMMs vs M=1 GEMV	tile/MMA util; weight-only quant at decode
Combine all-to-all	communication	route outputs back	overlap with next layer

# Dispatch microbenchmark: is your MoE comm- or compute-bound?
import torch, torch.distributed as dist, time
def bench_all_to_all(T, d, iters=50):
    x = torch.randn(T, d, device='cuda', dtype=torch.bfloat16)
    out = torch.empty_like(x)
    torch.cuda.synchronize(); t0 = time.time()
    for _ in range(iters):
        dist.all_to_all_single(out, x)        # the dispatch primitive
    torch.cuda.synchronize()
    bytes_moved = T * d * 2 * iters           # BF16 = 2B/elt
    gbps = bytes_moved / (time.time() - t0) / 1e9
    return gbps                                # compare to NVLink/network peak
# If gbps ≈ link peak → comm-bound (overlap it). If experts' TFLOP/s ≈ tensor-core
# peak → compute-bound. Whichever saturates first is your bottleneck.

Compiler: XLA / StableHLO / PTX / MLIR EMPHASIS

This is the axis to win on. The story is one pipeline: a framework graph is progressively lowered through stable, then internal, then hardware IRs, and the compiler's superpower — and its limit — is fusion.

StableHLO — the portability contract

Purpose: StableHLO is a stable, versioned operation set (an MLIR dialect) that sits between ML frameworks (JAX, PyTorch, TensorFlow) and ML compilers (XLA, IREE). It's the interop layer — a frontend can serialize StableHLO with backward-compatibility guarantees and hand it to any compiler. Think of it as “the bytecode of ML graphs.”

// StableHLO for  y = exp(a + b)  — note explicit tensor types & shapes
func.func @f(%a: tensor<4x8xf32>, %b: tensor<4x8xf32>) -> tensor<4x8xf32> {
  %0 = stablehlo.add %a, %b : tensor<4x8xf32>
  %1 = stablehlo.exponential %0 : tensor<4x8xf32>
  return %1 : tensor<4x8xf32>
}

HLO & fusion — XLA's single most important pass

XLA lowers StableHLO into its internal HLO and runs optimization passes. The one that matters most is fusion: merge multiple ops into one kernel so intermediate tensors stay in registers/SRAM instead of round-tripping HBM. Since most non-GEMM ML ops are memory-bound (Module 04), fusion is where XLA earns its speedups.

Lowering boundaries — where the compiler stops (and you write a kernel)

This is the question that separates a compiler engineer from a framework user. Fusion and codegen stop at boundaries like:

• Library calls: big GEMMs/convs go to cuBLAS/cuDNN as custom calls — XLA schedules them but won't fuse through them. An epilogue may fuse onto the GEMM, but two library GEMMs won't merge.
• Incompatible iteration spaces: fusing a reduction's consumer with an unrelated elementwise op can force recompute or block tiling that isn't profitable; the cost model declines.
• Opaque / data-dependent ops: dynamic shapes, sorts, scatter with unknown extents, control flow — hard to fuse and tile statically.
• Algorithms the op set can't express efficiently: FlashAttention's online-softmax tiling isn't something XLA derives from softmax(QKᵀ)V primitives — so it ships as a custom call (a hand-written Triton/CUDA kernel). Same for many fused MoE paths.

MLIR, PTX, and where Triton sits

• MLIR is the compiler infrastructure: multiple dialects (StableHLO, linalg, affine, gpu, nvvm, llvm) and passes that progressively lower from high-level to hardware. StableHLO is one dialect; lowering walks down to the LLVM dialect.
• PTX is NVIDIA's virtual ISA — forward-compatible assembly. ptxas compiles PTX → SASS, the actual per-architecture machine code. LLVM/NVVM emit PTX.
• Triton is a kernel DSL with its own MLIR-based pipeline (Triton IR → Triton GPU IR → LLVM → PTX). It's how you write the custom-call kernels XLA can't generate — your FlashAttention, your fused MoE GEMM.

Training & gradients context CONTEXT

Even for a kernel role you must speak training fluently — half the kernels you'll write are backward passes, and the memory tricks (recompute, checkpointing) come straight from the training loop.

Backward is just more matmuls and reductions

Autodiff applies the chain rule op-by-op (reverse-mode / VJP). The key fact for kernels: the backward of a linear layer is two GEMMs.

The memory ideas you already met

• Activation checkpointing / recompute: don't store every activation for backward — recompute it from a checkpoint. Trades FLOPs for memory. FlashAttention's backward does exactly this: it recomputes the attention tiles rather than storing the S×S probabilities.
• Mixed-precision training: keep FP32 master weights, do compute in BF16, accumulate in FP32. FP16 additionally needs loss scaling (Module 07) so small gradients don't underflow.
• Gradient checking = numerical regression test: compare the analytic gradient to a finite-difference estimate \(\frac{\mathcal L(\theta+\epsilon)-\mathcal L(\theta-\epsilon)}{2\epsilon}\) within tolerance — the training-side analog of the kernel numerical checks in Module 08.

Interview drills + cheat sheet REVIEW

Rapid-fire — answer out loud

One-glance cheat sheet

MMA shapes	Issuer	FP formats (S/E/M)
Ampere mma m16n8k16	1 warp	BF16 1/8/7 · FP16 1/5/10
Hopper wgmma m64nNk16	warpgroup (4)	FP8 E4M3 1/4/3 · E5M2 1/5/2
Blackwell tcgen05 ≤256·256·16	single thread + tensor mem	FP4 E2M1 1/2/1 (block-scaled)

---

Further reading (primary sources)

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — arXiv:2205.14135.
• OpenXLA — XLA architecture & StableHLO docs (openxla.org).
• NVIDIA — “Introducing NVFP4 for Efficient and Accurate Low-Precision Inference” (developer.nvidia.com).
• OCP Microscaling (MX) formats specification — for MXFP4/MXFP8.
• NVIDIA PTX ISA & CUTLASS tutorials — for mma/wgmma/tcgen05 and tensor-memory GEMMs.

Built as a compressed study aid. Math via KaTeX, code via highlight.js. Verify exact hardware peak numbers against the spec sheet for your target part before quoting them in an interview.