Module 01

Why Triton Exists & the Programming Model

The framing question of every Triton interview: "why is Triton's model different from CUDA?" Get the block-vs-SIMT distinction right and you've cleared the conceptual bar.

🎯

What interviewers probe here

Whether you understand Triton's block/tile programming model vs CUDA's per-thread SIMT — and can name what the compiler does for you. The #1 separator: candidates who think a Triton "program" is a CUDA thread.

Two Philosophies

🧠

CUDA: Scalar Program, Blocked Threads

You write per-thread scalar code (acc += A[m,k]*B[k,n]) and manually manage coalescing, shared memory, and __syncthreads.

📦

Triton: Blocked Program, Scalar Threads

You write tile-level code (acc += A[m:m+BM, :] @ B[:, n:n+BN]); the compiler maps threads underneath. There is no threadIdx.

You are one program (tl.program_id) that owns a whole tile — not one thread among many.

What the Compiler Does For You

In CUDA you hand-tune these; in Triton the compiler handles them via block-level data-flow analysis. Naming three of these is a standard interview answer:

Coalescing Automatic memory coalescing & thread swizzling — you never arrange threads by hand.

Shared memory Automatic shared-memory allocation + synchronization — no manual __shared__ / __syncthreads.

Vectorization + tensor cores Automatic vectorization, tensor-core instruction selection, prefetching, async-copy scheduling.

Compilation path: @triton.jit Python → Triton IR → TritonGPU IR → LLVM → PTX → SASS, launched over a grid of program instances.

The Framing Round

1. What is a Triton "program" analogous to in CUDA?

2. Which of these do you NOT manage by hand in Triton?

⚠

Common follow-ups & gotchas

• "Block" in Triton = a program's tile of data, not a CUDA thread block. Same word, different meaning.
• You cannot control individual threads — only tile shapes, num_warps, num_stages.
• "Does Triton always beat cuBLAS?" — No; its win is near-peak performance with far less code for custom fused ops.

Module 02

Your First Kernel — Vector Add

The warm-up. Every Triton kernel follows this skeleton: compute offsets, mask the boundary, load, compute, store. Get the masking right and you won't segfault in a live coding round.

🎯

What interviewers probe here

Rarely the actual bar — but the boundary mask is where candidates crash live. If n_elements isn't a multiple of BLOCK_SIZE, the last program runs out of bounds without it.

The Universal Skeleton

1

program_id → offsets

→

2

mask boundary

→

3

tl.load

→

4

compute

→

5

tl.store

The Complete Kernel

Triton

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr,
    n_elements, BLOCK_SIZE: tl.constexpr):
  pid = tl.program_id(axis=0)
  start = pid * BLOCK_SIZE
  offs = start + tl.arange(0, BLOCK_SIZE)
  mask = offs < n_elements
  x = tl.load(x_ptr + offs, mask=mask)
  y = tl.load(y_ptr + offs, mask=mask)
  tl.store(out_ptr + offs, x + y, mask=mask)

grid = lambda m: (triton.cdiv(n, m['BLOCK_SIZE']),)
add_kernel[grid](x, y, out, n, BLOCK_SIZE=1024)

Line by line

Tensors pass as pointers to their first element.

BLOCK_SIZE is constexpr — compile-time, usable as a shape.

program_id — which tile this program owns.

tl.arange builds the offset vector.

mask guards the ragged last block.

tl.load/store honor the mask — masked lanes skipped.

grid lambda gets meta so cdiv sees BLOCK_SIZE. Launch is async.

Spot the Bug: Illegal Memory Access

This kernel crashes with "illegal memory access" when n = 1000, BLOCK_SIZE = 1024. Click the line at fault.

offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

x = tl.load(x_ptr + offs) # no mask!

tl.store(out_ptr + offs, x + 1.0)

Conceptual Round

1. Why must BLOCK_SIZE be tl.constexpr?

2. Does add_kernel[grid](...) block until the GPU finishes?

⚠

Common gotchas

• Forgetting the mask when size isn't a multiple of the block.
• Passing BLOCK_SIZE positionally — meta-params must be keyword args.
• This uses plain pointer arithmetic (x_ptr + offs), not tl.make_block_ptr — block pointers appear in matmul/attention. Don't conflate them.

Module 03

The Memory Model & Masking

Where candidates break in live coding. Triton manipulates pointers into tensors, not tensors — strides, broadcasting, and the right other= value for a masked reduction are the details that separate strong from weak.

🎯

What interviewers probe here

Whether you build addresses from strides (never assume contiguity) and pick the correct other= fill for masked lanes — -inf for a max, 0.0 for a sum. This is a top boundary-handling separator.

Pointers, Strides, Broadcasting

An address is base + i*stride_i + j*stride_j. Build 2D pointer blocks by broadcasting offs_i[:, None] against offs_j[None, :]. Strides come from PyTorch's .stride() and are in elements, not bytes.

Strided 2D tile load

offs_m = pid_m*BM + tl.arange(0, BM)
offs_n = pid_n*BN + tl.arange(0, BN)
ptrs = base + offs_m[:, None]*stride_m \
            + offs_n[None, :]*stride_n
mask = (offs_m[:, None] < M) & \
       (offs_n[None, :] < N)
tile = tl.load(ptrs, mask=mask, other=0.0)

Key ideas

Broadcasting [:, None] × [None, :] builds the 2D address grid.

2D mask guards both dimensions at once.

other= is the value masked-off lanes return — choose it for your reduction.

Spot the Bug: The Wrong Fill Value

This kernel computes a row max but returns wrong answers for rows with all-negative values. Click the mistake.

offs = tl.arange(0, BLOCK_SIZE)

row = tl.load(p + offs, mask=offs<n, other=0.0) # max reduction!

m = tl.max(row, axis=0)

💡

The rule

Pick the neutral element for your reduction: other=-float('inf') for max, other=0.0 for sum, other=float('inf') for min. The fill must never affect the result.

Conceptual Round

1. How do you correctly index a transposed/strided tensor in Triton?

2. Triton strides are measured in...?

⚠

Common gotchas

• Wrong other= for the reduction (the bug above).
• Broadcasting mistakes — [:, None] vs [None, :] transposed.
• Forgetting the mask on tl.store → out-of-bounds writes.
• Assuming contiguity instead of threading strides through.

Module 04

Fused Softmax

The canonical Triton coding exercise. "Write a fused softmax kernel" tests masking, numerical stability, and single-pass reduction all at once — and sets up Flash Attention.

🎯

What interviewers probe here

Two things: why subtract the row max (numerical stability + shift-invariance), and why fusion is faster (softmax is memory-bound; separate ops round-trip DRAM per op).

One Pass, On-Chip

Load a whole row into SRAM, then do max → subtract → exp → sum → divide on-chip in one pass, and write once. The naive PyTorch softmax round-trips DRAM per op (~5MN read + 3MN write); the fused kernel touches DRAM ~once in / once out (~2MN) → roughly 4× less traffic.

Triton

row = tl.load(row_ptr + cols,
    mask=cols < n_cols, other=-float('inf'))
row = row - tl.max(row, axis=0)   # stability
num = tl.exp(row)
den = tl.sum(num, axis=0)
tl.store(out_ptr + cols, num / den,
    mask=cols < n_cols)

Why each line

other=-inf so padded lanes never win the max.

subtract max — prevents exp() overflow; softmax is shift-invariant so the answer is unchanged.

exp then sum — the reduction, in SRAM.

one store — the whole op touched DRAM twice total.

Fused vs Unfused — Watch the DRAM Traffic

Step through what happens to a row under each approach. This is the narration interviewers want when they ask "why is fusion faster?"

0 / 4 messages

Conceptual Round

1. Why subtract the row max before exp, and why is it still correct?

2. Why is a fused softmax faster than exp/sum/div as three ops?

⚠

Common gotchas

• One-block-per-row requires the row to fit in SRAM — fails for huge n_cols.
• Forgetting other=-inf (Module 3 bug).
• Reducing over the wrong axis.
• BLOCK_SIZE = next_power_of_2(n_cols) — reductions want power-of-two width, extra lanes masked.

Module 05

Autotuning & Performance

Autotuning itself is engineering practice, but the concepts get tested: "what does num_stages do?" and "when does Triton recompile?" Don't confuse pipelining depth with warp count.

🎯

What interviewers probe here

The meaning of the tuning knobs — especially the classic mix-up: num_stages = software-pipelining depth, num_warps = threads per program (÷32). And what triggers recompilation.

The Autotune Decorator

Triton

@triton.autotune(
  configs=[
    triton.Config({'BLOCK_SIZE': 128}, num_warps=4),
    triton.Config({'BLOCK_SIZE': 1024}, num_warps=8)],
  key=['n_elements'])
@triton.jit
def kernel(...): ...

The knobs

configs — candidate tile shapes + warp/stage counts; Triton benchmarks all and caches the winner.

num_warps — threads per program = 32 × num_warps.

key — arg names whose value change re-triggers benchmarking (e.g. shapes).

num_stages Software-pipelining depth — how many iterations of loads overlap with compute. More = more latency hiding, more SMEM pressure. Not warps.

reset_to_zero / restore_value Every config runs during tuning — reset accumulator outputs, or restore mutated inputs, or results corrupt.

Spot the Bug: Wrong Results During Tuning

A kernel accumulates into an output buffer in place. It's correct without autotuning, but produces garbage once @triton.autotune is added. Click the fix it's missing.

@triton.autotune(configs=[cfg_a, cfg_b],

key=['M','N']) # no reset_to_zero for the accumulator

@triton.jit
def kernel(out_ptr, ...): # does out += partial

Conceptual Round

1. What does num_stages control?

2. When does Triton re-run autotuning?

⚠

Common gotchas

• Missing a shape dim from key → stale config for new shapes.
• First call is slow (tuning cost); TRITON_PRINT_AUTOTUNING=1 shows the winner.
• num_stages too high overflows shared memory → compile failure.
• triton.heuristics must come before triton.autotune, not after.

Module 06

Blocked Matrix Multiply

The block matmul: a 2D grid of C-tiles, a K-loop with tl.dot, an fp32 accumulator, and L2-friendly grouped ordering. The concepts here are the bridge to Flash Attention.

🎯

What interviewers probe here

Why the accumulator is fp32 even with fp16 in/out, and what GROUP_SIZE_M buys you (L2 reuse) — the kind of "why is this arranged this way" question that separates candidates.

Each Program Owns One C-Tile

Triton (K-loop core)

acc = tl.zeros((BM, BN), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BK)):
  a = tl.load(a_ptrs, mask=k_mask, other=0.0)
  b = tl.load(b_ptrs, mask=k_mask, other=0.0)
  acc = tl.dot(a, b, acc)   # fp32 accumulate
  a_ptrs += BK * stride_ak
  b_ptrs += BK * stride_bk
c = acc.to(tl.float16)      # cast once at the end

Key facts

acc is fp32 — summing many fp16 products loses precision fast; tensor cores natively accumulate in fp32.

tl.dot(a, b, acc) = block matmul added into acc. Use input_precision= (not the deprecated allow_tf32).

K-loop mask handles K not divisible by BK.

cast once at the end.

Match the Matmul Concept to Its Purpose

Drag each piece onto what it's for.

fp32 accumulator

GROUP_SIZE_M ordering

K-loop mask

tl.dot

Preserve precision when summing many fp16 products

Drop here

Reorder tile traversal for better L2 cache reuse

Drop here

Handle K not divisible by the block size

Drop here

Block matmul mapped onto tensor cores

Drop here

💡

Grouped ordering, concretely

GROUP_SIZE_M finishes a group of row-tiles before moving across columns, so the same A/B tiles stay hot in L2. Triton's 9×9 example: row-major needs 90 tile-loads for the first 9 outputs; grouped needs 54 — a real ~10% gain. It changes only order, never the math.

Conceptual Round

1. Why accumulate in fp32 when inputs and output are fp16?

2. Does GROUP_SIZE_M change the result?

⚠

Common gotchas

• Off-by-one in the pid → (pid_m, pid_n) decomposition.
• Forgetting the K-loop mask when K isn't divisible by BLOCK_K.
• Applying a fused activation after the fp16 cast instead of on the fp32 accumulator.
• Using deprecated allow_tf32 — use input_precision="tf32". For the CUDA-level version of this, see GEMM from Scratch.

Module 07

Flash Attention in Triton

The highest-value module in the course. "Explain online softmax and how Flash Attention avoids the N×N matrix" is a staple ML-breadth question — and the #1 misconception is thinking it reduces FLOPs. It doesn't.

🎯

What interviewers probe here

That Flash Attention trades memory for recompute: same (or more) FLOPs, far less HBM traffic, O(N) memory instead of O(N²). Candidates who say "it's faster because fewer operations" fail. The bottleneck was memory movement, not math.

Never Materialize the N×N Score Matrix

Standard attention computes the full T×T score matrix QKᵀ, softmaxes it, then multiplies by V — that matrix is O(N²) in HBM. Flash Attention keeps a Q-tile in SRAM and streams over K/V tiles, maintaining running statistics so it never stores the scores.

m_i (running max) Init -inf. Tracks the max score per query row across tiles seen so far.

l_i (running denom) The softmax denominator, accumulated tile by tile.

acc (output) The running weighted sum of V. Rescaled when the max grows.

The Rescaling Trick — Why It's Exact

The subtle part: when a new K/V tile raises the running max, all prior contributions must be retroactively rescaled. Step through why this equals the true full-row softmax.

0 / 4 messages

The Signature ML-Breadth Round

1. Why is Flash Attention faster?

2. Why multiply acc and l_i by α = exp(m_old − m_new) when the max grows?

⚠

Common gotchas

• Initializing m_i to 0 instead of -inf.
• Forgetting to rescale acc (or l_i) → wrong normalization.
• Doing the final /l_i inside the loop instead of once in the epilogue.
• The modern Triton tutorial uses exp2 (base-2, faster) with the scale folded by 1/ln2, and TensorDescriptor loads — teach the algorithm with the portable block-pointer form first.

🔗

Where this leads

Online softmax is the bridge to the LLM-inference world — KV cache, prefill/decode, PagedAttention. See vLLM & Distributed Kernels for the serving side.

Module 08

Integration with PyTorch & torch.compile

Mostly engineering practice — but the autograd question ("how do you make your kernel differentiable?") is a common senior follow-up, and knowing TRITON_INTERPRET=1 shows real hands-on experience.

🎯

What interviewers probe here

Framed as methodology, not a quiz: "you wrote a fast kernel — how do you make it differentiable, composable with torch.compile, and how do you debug it?" Naming register_autograd and TRITON_INTERPRET=1 is the signal.

Three Levels of Integration

1

Call from @torch.compile (PyTorch ≥ 2.3)

Inductor traces through the grid launch and fuses/optimizes around your kernel. Simplest path.

2

Composable op: triton_op + wrap_triton (≥ 2.6)

Makes a raw kernel a first-class PyTorch op — composable with tensor subclasses, FlopCounter, CPU fallback, AOTInductor, and torch.compile can trace into it (unlike opaque custom_op).

3

Autograd: prefer register_autograd

Over torch.autograd.Function (which has composability footguns under torch.compile). Provide backward + setup_context; backward must be PyTorch ops.

Making a Kernel Differentiable

torch.library (≥ 2.6)

from torch.library import triton_op, wrap_triton

@triton_op("mylib::mysin", mutates_args={})
def mysin(x): 
  out = torch.empty_like(x)
  wrap_triton(sin_kernel)[(x.numel(),)](x, out, ...)
  return out

def backward(ctx, grad):
  (x,) = ctx.saved_tensors
  return grad * x.cos()
def setup_context(ctx, inputs, output):
  (x,) = inputs; ctx.save_for_backward(x)
mysin.register_autograd(backward, setup_context=setup_context)

Why this way

triton_op makes it a real op subsystems can see.

wrap_triton at the launch site.

backward in PyTorch ops (or its own triton_op).

setup_context saves what backward needs.

register_autograd — not autograd.Function, which breaks under compile.

Conceptual Round

1. Why wrap a kernel in triton_op vs just calling it in a compiled function?

2. Your kernel gives wrong results and you can't attach a GPU debugger. Fastest way to inspect tensors?

🏆

You've completed Triton Kernel Mastery

Programming model → first kernel → memory & masking → fused softmax → autotuning → blocked matmul → Flash Attention → PyTorch integration. The interview through-line: Triton abstracts away threads so you reason at the tile level about memory movement — and memory movement is what makes kernels fast.

⚠

Version-awareness (say these to sound current)

• tl.dot: use input_precision=, not deprecated allow_tf32.
• triton_op/wrap_triton need PyTorch ≥ 2.6; basic kernels in torch.compile ≥ 2.3.
• Prefer register_autograd over torch.autograd.Function for compile-safety.
• For the "is my kernel memory- or compute-bound" question, apply the roofline model.