Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 6: Kernels, Triton, XLA — Annotated Lecture Notes & Key Takeaways

Course Information

Institution: Stanford

Original Course: Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 6: Kernels, Triton, XLA

Instructor Bio: This lecture is delivered by **Percy Liang**, Associate Professor of Computer Science at Stanford University and core faculty of the Stanford Institute for Human-Centered Artificial Intelligence (HAI). Percy Liang leads the Stanford Natural Language Processing Group, with research spanning the theoretical foundations and practical systems of language modeling, machine learning alignment, compositional semantics, and NLP evaluation. He received his PhD in Computer Science from the University of California, Berkeley and his BA in Mathematics from Harvard University. His work has been recognized with the NSF CAREER Award, Google Faculty Research Award, and multiple best paper awards at top-tier machine learning and NLP conferences. He is also widely known for creating influential benchmarks and open-source frameworks for language model research and assessment.

Course Description: This lecture dives into low-level kernel optimization and compiler technologies that maximize hardware utilization for language model workloads. It explains the role of custom compute kernels in accelerating core transformer operations, introduces the Triton programming language for writing high-performance, hardware-agnostic GPU kernels, and covers XLA (Accelerated Linear Algebra) as a graph-level compiler for optimizing end-to-end tensor computation. The session discusses when and how to apply kernel-level and compiler-level optimizations, and their impact on training throughput and inference efficiency.

If you skipped this session, the one-line version is that it's the hands-on follow-up to Monday's GPU overview, and the whole point is to get you writing your own Triton kernels, timing them, and understanding why three implementations of the exact same formula can run at very different speeds. There's one bait-and-switch to flag before you start: the title advertises XLA, and XLA is never mentioned once in the eighty-six minutes. I'll come back to that at the end, because it's the biggest hole in an otherwise very solid lecture.
 

The instructor's stated plan is refreshingly concrete. First, re-establish the hardware picture. Then talk about benchmarking and profiling — before any kernel code, on purpose. Then write kernels of increasing difficulty: element-wise, row reduction, reduction that doesn't fit, and finally matrix multiplication. The endpoint he's aiming at is FlashAttention in the assignment.
 

The Mental Model: Memory, Threads, and Why Blocks Exist

 

The lecture opens with a memory hierarchy you're meant to carry around in your head. Each GPU has somewhere between 100 and 200 streaming multiprocessors (SMs), and that number has barely moved across generations. Inside an SM sit registers — on a B200 that's roughly 65,000 of them, about 256 KB — plus L1 cache and shared memory, which are the same physical memory: you get to control the shared part, not the L1 part. Above that, one L2 cache serves the whole chip. Above that, HBM, which is huge and, at around 8 terabytes per second, still the slowest thing in the building.
 

The pattern to internalize is that capacity and speed trade off against each other. Registers and shared memory are small, local, and quick. HBM is big, distant, and slow. He keeps returning to this, and for good reason: nearly every optimization in the lecture is a way of avoiding a trip to HBM.
 

The programming model itself is only three levels deep. Threads are grouped into thread blocks (also called CTAs), and blocks form a grid; launching a kernel means launching that grid. At this level it really is as pleasant as he claims — you describe what a block and its threads should do, and correctness requires nothing more. He waves off two hardware features that complicate the picture, thread block clusters on H100/B200 and the tensor memory on B200 that sits between registers and shared memory, on the grounds that they're mostly invisible to the programmer.
 

Then he asks the question that makes the middle level make sense: why have blocks at all? For element-wise work like GELU, a flat grid of threads is perfectly natural, one element per thread, like a for-loop over your data. But softmax and matmul need threads to cooperate, and the only cooperation channel available to a lonely thread is writing to HBM and reading it back. That's far too expensive. A block, by contrast, lands on one SM, pulls a chunk of data out of HBM, lets its threads cooperate through shared memory, and writes the result back once. He calls tiling "the whole game," and this is the setup for it.
 

The pitch for Triton is that it lets you think natively at that block level rather than at the individual-thread level, and he argues that once block-thinking clicks, everything downstream gets easier.
 

The tension he names explicitly, and which organizes the rest of the hour, is that the programming model is clean and hardware-agnostic while performance is neither. You can get correct code knowing only threads, blocks, and grids. Fast code requires knowing the hardware in detail — because if you didn't care about speed, you wouldn't be writing kernels at all.
 

Where the Abstraction Leaks: Five Hardware Realities

 

He says he'll give about five examples of hardware details that decide your performance, and then delivers five-plus.
 

Warps, Lockstep, and Why Branching Hurts

 

Warps are the first thing that isn't in the clean model. Threads inside a block are bundled into groups of 32, and all 32 must execute the same instruction in lockstep every cycle. When branches force different threads down different paths, the warp runs one path, then the other — the work gets sequentialized, which is exactly what you don't want. Avoid branching where you can.
 

The flip side is the neatest trick in the whole machine. Each SM keeps many warps resident and switches between them at essentially no cost, which is not how CPU threads behave. The point is latency hiding: if one warp is stalled waiting on HBM for something like a hundred cycles, the SM immediately runs a different warp that has tensor core work to do. No idle hardware.
 

Occupancy, Register Pressure, and Thread Coarsening

 

Each thread may use at most 255 registers, and the SM's register file is fixed, so fatter threads mean fewer resident threads. He's careful not to turn this into "maximize occupancy": fewer threads each doing more work can be better, since scheduling overhead drops. The named technique is thread coarsening — have one thread chew through, say, eight elements instead of one.
 

The worked example (13:15) is worth reconstructing because it's the most concrete arithmetic in the lecture. Take blocks of 128 threads at 160 registers apiece. That's about 20,480 registers per block, and with 65,000 registers per SM you fit at most three blocks. Three blocks of four warps is twelve warps, against a ceiling of 64 — so you're running at roughly 18% occupancy, throttled not by compute but by register appetite. I liked this example, though I wish he'd then shown what you'd actually change to fix it.
 

Bank Conflicts and Coalescing

 

Shared memory is split into 32 banks, each four bytes wide, and each bank can serve one thread per cycle. Two threads hitting the same bank get serialized. The worst case is 32 threads walking down a single column of a matrix — a 32-way conflict, where your beautiful parallelism turns into a queue. The honest part: for a matmul you need rows of one operand and columns of the other, so you can't always just pick a friendly traversal order. He names swizzling as the remedy and explicitly declines to explain it.
 

Coalescing sounds similar but is a different constraint on a different resource. It's about HBM, not shared memory: the 32 requests from a warp get merged into a small number of fixed-size transactions (he quoted a range reaching about 128 bytes). When consecutive threads touch consecutive addresses, one fetch serves everyone. March down a column instead and you drag in a lot of data you'll throw away.
 

The Tail Wave

 

Blocks get scheduled onto a finite set of SMs — 148 in his example. Launch 160 blocks and you run 148, then 12, and during that second wave most of the chip sits idle. His heuristic is to pick a block count that divides the SM count.
 

A student then asked whether two blocks can share an SM to smooth that out, which is the obvious follow-up. The answer was that blocks have to stay together, that a block already using most of the tensor cores won't benefit from a neighbor, and that the real fix is resizing blocks so there's no tail. I found this the least satisfying exchange in the lecture: "it depends on the block" isn't something you can act on, and it quietly conflicts with the divide-the-SM-count heuristic from two minutes earlier.
 

His own summary is the honest one — elegant model on top, and underneath it warps, banks, coalescing, and occupancy, all of which you can only partly see through a profiler. Sometimes the scheduler just does things you don't control.
 

Measure First: Benchmarking, Profiling, and the Three GELUs

 

Before a single line of kernel code, he spends about fifteen minutes on measurement, and the reasoning is stated plainly: benchmark, profile, change, measure again, and always find the bottleneck before you start writing kernels.
 

Benchmarking gives you one end-to-end number with no breakdown, which sounds useless until you remember that end-to-end time is what you actually care about — and that collapsing to a single number lets you watch how something scales. The gotchas he lists: warm up first, because lazy compilation will otherwise pollute your first measurement; time repeatedly, because there's real variance; use CUDA events with a start and end record; synchronize, since GPU work is asynchronous and you need a barrier before reading the clock; then average, though he notes that if you're being careful you'd want a distribution or a P95.
 

The matmul scaling demo shows exactly why this matters. Time grows cubically as expected, but there's a long flat floor: up to matrices of roughly 2,000 dimensions, runtime barely moves. GPUs are built for large matmuls, and a 2×2 multiply is hopelessly inefficient no matter how you write it.
 

Profiling then tells you where time goes, and also — his emphasis — what is actually happening under high-level code. Adding two tensors in PyTorch launches one add kernel. Multiplying them launches a CUTLASS kernel whose name is informative if you can read it (29:40): SM100 means the Blackwell architecture, F32 the precision, and 64×64×16 the tile shape. Change to a 128×128 multiply and the name changes to a 32×32×16 tile. Different shapes, different kernels, chosen for you. He notes the assignment uses Nsight for this.
 

The three-GELU race is where measurement becomes an argument. GELU has a tanh approximation that's cheaper to compute, and he lines up three implementations: the naive PyTorch transcription of the formula, the built-in F.gelu, and torch.compile applied to the naive version. The naive one is slowest by a wide margin; the compiled one is much faster but still trails the built-in.
 

The profiler explains why. The naive version is a handful of separate kernels — a binary op, a unary op, an add, a tanh — because each primitive in the computation graph becomes its own kernel, and each kernel reads from HBM, computes, and writes back. Between launches, the data has to go home to HBM. The built-in is one hand-written CUDA kernel; asked why it exists, he shrugs that GELU is popular, so someone wrote one. Nothing magical. The compiled version is one fused Triton kernel (35:21), generated from the graph. Fusion means one read and one write per element instead of many.
 

Two caveats he volunteers: last year the compiled version was much closer to the built-in, and none of these numbers come from heavily tuned code. Both are worth remembering before you generalize the ranking.
 

Writing Triton Kernels: Pointers, Masks, and the PTX Underneath

 

The contrast with CUDA is drawn sharply. CUDA asks what each individual thread does: fine-grained, close to the metal, and you handle synchronization and shared-memory bookkeeping yourself — fine for element-wise work, increasingly annoying as operations get cooperative. Triton, originally from OpenAI and now fairly standard, asks what each block does: load into shared memory, compute, write to global memory. It sits between per-element CUDA and whole-tensor PyTorch. He admits the limit: if you want to exploit every new feature of the latest hardware, Triton may not give you full flexibility.
 

He also flags the mental shift that trips people up. In Triton you stop thinking functionally. You allocate the output tensor yourself and write into it; kernels don't return values, and the arguments are pointers you do arithmetic on.
 

The first kernel (43:14) is an element-wise op over an 8,000-element vector. He picks a block size of 1,024, giving eight blocks, and the launch syntax specifies the grid shape. Inside, the block asks which block it is, multiplies that id by the block size to find its starting offset, and adds a range of 0 to 1,023 to get the addresses it owns. Because the tensor length usually won't divide evenly by the block size, there's a mask: valid lanes proceed, and lanes past the end do nothing. Then a masked load, the computation, and a masked store. He stresses that every kernel in the course has that same skeleton.
 

Two student questions here are worth noting. For element-wise work, CUDA looks about the same, so the advantage only shows up later. And as for tensor cores — you don't control them; the hardware decides.
 

Then he pulls back the curtain to PTX, the intermediate assembly the compiler emits. Now you're looking at one thread rather than a block. You see loads from global memory into registers, integer registers and floating-point ones named differently, moves, multiplications by constants, and a store back to global memory at the bottom. The same code runs on every thread; what distinguishes them is the block index and the thread index handed to it. And you can see thread coarsening happening without being asked for: the compiler decided each thread should handle eight elements, since the per-thread work was thin. PTX still doesn't say which SM you're on or how warps are arranged — that stays with the hardware. He mentions that people do hand-write PTX if they think they're better than the compiler, and advises against it as a first move, though he allows that less mature accelerators sometimes need the hand-holding.
 

A student then described latency hiding back to him — a load stalls, another warp gets swapped in, the scheduler returns when the data lands — and he confirmed it. Asked why there are four warp schedulers per SM, he said he didn't know the reason. I appreciated the honesty, but it's a noticeable gap in the story he'd just finished telling.
 

Tiling: From Softmax Rows to the Matmul Finale

 

The announced progression is element-wise, then reduction over a row, then reduction where the row doesn't fit, then matmul — building toward FlashAttention.
 

One Row, One Block

 

Plain-PyTorch softmax is a counting disaster: by his tally it's five reads and three writes per element when you'd want far fewer. The Triton version assigns one block per row, which works because softmax is row-wise and blocks never need to talk to each other — there's no shared memory across blocks. He sets the block size to the number of columns rounded up to the next power of two ("for good luck"), runs one block per row, and passes strides so the kernel can find its way down the matrix. Masked lanes get filled with negative infinity rather than being skipped, which is the right identity for this operation. Then the body is almost literally the naive math: subtract the row max for stability, exponentiate, sum, divide, store.
 

The lesson he draws is the encouraging one: when your data fits inside a block, the kernel body reads almost like ordinary PyTorch. He also notes that switching to a column-wise softmax would just mean changing the strides.
 

Baby Tiling When the Row Won't Fit

 

Things stop being PyTorch-shaped when a row is 4,000 elements and your block only holds 1,024. He switches the example to row sums, which is easier to trace. Now the block owns the whole row but processes it in four tiles: each thread keeps an accumulator, loops over the tiles adding its slice, and a final reduction collapses the per-thread accumulators to a scalar. Where the accumulator physically lives — registers or shared memory — is the compiler's decision, not yours, though he says a large enough block forces it into shared memory.
 

The distinction he keeps hammering is important and easy to blur: in the element-wise kernel, the row was chopped into blocks that ran independently; here the chopping is into tiles that one block walks through itself.
 

Matmul, Arithmetic Intensity, and a Free ReLU

 

Matmul gets the full treatment (1:12:25). The naive per-output-element kernel is correct and terrible: it reads from HBM for every combination of the three dimensions, so the number of reads grows cubically while the useful computation grows at the same rate — arithmetic intensity stays flat. The idealized version loads all of A and B into shared memory, reads only quadratically, and gets intensity proportional to N. The catch is that A and B don't fit.
 

Tiling is the compromise, and his framing is memorable: globally it looks like the naive loop, locally it looks like the idealized one. Each output tile becomes one block. That block sweeps across row tiles of A and column tiles of B, loading each pair into shared memory, multiplying, and accumulating partial sums there. Only after the whole sweep does it write its tile to HBM once. Intensity rises to something on the order of the tile size — not the O(N) ideal, but respectable.
 

There's a brief stride refresher along the way: a tensor is laid out linearly, the row stride tells you how far to jump to move down a row, and a transpose simply flips the two.
 

The ReLU he bolted onto the matmul finally gets explained near the end. Since you're writing the kernel anyway, applying an element-wise activation to the accumulator before the single HBM write costs almost nothing (1:21:45). That's fusion, and it's the payoff for all the tiling work.
 

His closing summary is short: PyTorch, Triton, and PTX are levels of control; the hardware underneath is finite in SMs, banks, registers, and memory; Triton makes block-level thinking pleasant; and next time it's multi-GPU.
 

Where I Was Left Hanging

 

The XLA gap is the obvious one. It's in the lecture title and never appears; TPUs and AMD get one sentence at the 1:06 mark and then it's Nvidia for the remaining eighty-five minutes. If you came for the XLA half, you got nothing.
 

The GELU benchmark numbers also bother me. The naive time is given as "3.75" with no units stated on screen, no hardware specified, and the ranking reportedly shifted since last year. It's a fine illustration and a bad source of constants.
 

Occupancy is the argument I most wanted finished. He says you can measure it, and that higher isn't automatically better, and then leaves it there — no rule for deciding whether your 18% is a problem. Combined with the conflicting tail-wave heuristics, you come away with two plausible fixes and no way to choose between them.
 

And the last question of the day — is it better to load a whole tensor at once or grind through it element by element? — got "hard to answer in the abstract," plus an offer to talk offline. A fair punt, but it's precisely the question that would have tied the tiling argument together.
 

Where the lecture is strongest is also where it's most countable: fusion and arithmetic intensity are argued by counting reads and writes to HBM, not by appeal to intuition. That reasoning transfers. The anecdotes about warp schedulers don't.
 

So: the durable takeaway is that counting HBM round trips explains almost every design decision here, and Triton kernels are mostly a convenient way to act on that counting without descending into per-thread bookkeeping.
 

Content Disclaimer:

This article is for general reference only and does not constitute professional R&D guidance, production process advice or quality certification. All material performance data has specific test premises; readers should verify parameters against actual equipment and working conditions.