Skip to content
All tags

#compiler

5 posts

TPU, JAX, and Pallas: One Attention Kernel on TPU, from XLA Fusions to Splash Attention | CMU 11-868 Lectures 12–13

Across two lectures and more than 200 slides, Google's Srinath Mandalapu traces one attention computation from Python down to TPU VLIW instructions. L12 covers the JAX ecosystem, the memory and compute units of TPU Ironwood, and how XLA compiles attention into three fused kernels. L13 covers what XLA cannot do: using Pallas to control movement between HBM and VMEM yourself, writing FlashAttention, then adding block sparsity to get Splash Attention. The ideas match the GPU version. The difference is that on TPU the compiler does most of the scheduling, and Pallas is how you take loops and block sizes back into your own hands.

Performance Optimization Beyond the Experts — Halide's Algorithm/Schedule Split, Autoschedulers, and LLM Agents | Stanford CS149 Lecture 13

L13 asks what to do when there are too few people who can write fast code. The slides offer three answers. First, raise the level of abstraction: Halide splits what to compute (the algorithm) from how to compute it (the schedule), so one line of schedule turns the same blur into a tiled, vectorized, multi-core version. Second, intelligent search: because the schedule space is well defined, search plus a learned cost model can generate schedules automatically. Third, the emerging option of LLM agents: have a model write CUDA, run it, read the profiler, reflect, and revise, and let it improve itself with a database of examples or prompt optimization. The last slide leaves you with a question: is the real value in DSL design or in the LLM agent?

Programming Specialized Hardware — ThunderKittens Tames H100 Asynchrony, Dataflow Replaces It with Metapipelines | Stanford CS149 Lecture 11

L11 asks what programmers pay once hardware specializes for AI. On the H100, saturating Tensor Cores means 16×16 tiles, TMA moving data asynchronously, and producer and consumer warps running as a pipeline. That's hard to write, which is why DSLs like ThunderKittens exist. The other route is a dataflow architecture (SambaNova SN40L): describe the computation with parallel patterns such as map, reduce, and zip, and let the compiler handle tiling, metapipelining, and placement. The slides say this can fuse an entire Llama 3.1 8B decoder layer into one kernel.

From C to x86-64, Reading Disassembly for the First Time | Stanford CS107 Lecture 14

CS107 Lecture 14 dissects the ten x86-64 instructions for sum_array: addresses and machine bytes appear on the left, AT&T assembly on the right, and the reader's job is to recover C-level effects from opcodes, operands, registers, and control flow—not to write assembly from scratch.

TypeScript 7: A Native Go Rewrite Makes Type Checking Fast, but Migration Goes Beyond tsc

TypeScript 7 rewrites the compiler and language service in Go. The order-of-magnitude gains come from native execution and shared-memory parallelism, while old APIs and compiler options define the migration cost.