Dataflow Architectures for AI

A PyTorch model passes through several compiler stages before it runs on an accelerator, and each stage makes decisions that change what the hardware executes. In this course you will follow a model through that stack and connect each compiler decision to the kernels that run on the device. You will study the PyTorch 2 compilation path through TorchDynamo, AOTAutograd, and TorchInductor, including graph breaks, guards, and recompilation. You will then examine the intermediate representations and compiler passes that matter most for accelerators: decomposition, fusion, tiling, layout, scheduling, and code generation. The course compares three compiler stacks on the same classes of problems: TorchInductor with Triton on GPUs, XLA with HLO on TPUs, and the Spyre compilation stack with SuperDSC and the emerging KTIR representation. Worked examples show how a feed-forward layer is mapped across Spyre cores and how one attention graph is handled by XLA and torch-spyre. In the lab, you will benchmark eager and compiled execution, inspect a captured graph, repair a graph break, and connect compiler artifacts to the code your run produced. A 18-question graded quiz at the end covers every module of the course.

By the end of this course, you will be able to explain how dataflow execution works, trace dependencies through a computational graph, and compare how different accelerators schedule and execute the same AI workload. You will begin with the classical dataflow model: operations as nodes, dependencies as edges, and execution triggered when inputs are ready. From there, you will compare static and dynamic scheduling and examine how compile-time planning introduces constraints such as shape specialization. You will use torch.fx to inspect model graphs and learn how guards capture assumptions about shapes and program state. The course then compares four implementations of dataflow: TPU systolic arrays, Groq's deterministic tensor streaming, the Cerebras wafer-scale fabric, and IBM Spyre's 32-core programmable dataflow design. Triton is used as a contrasting example of tile-level programming on GPUs. A final case study follows one transformer attention block through TPU and Spyre compilation flows. In the lab, you will trace a transformer block, visualize the graph, annotate FLOPs and memory footprint, and predict the critical path on two hardware profiles. You should be comfortable with Python, PyTorch, and transformer architectures.

When a compiled model returns a wrong answer, the error often appears far from the operation that caused it. This course teaches a systematic method for debugging compiled AI workloads: find the stage where a failure first appears, and separate logic errors from expected numerical differences. You will learn why familiar eager-debugging techniques become less reliable after graph capture, fusion, asynchronous execution, and compiler scheduling. An eight-layer model organizes failures from Python source down to hardware execution. You will use differential debugging against a trusted baseline to narrow the search. The course covers backend ablation, compiler IR, provenance, compile logs, kernel replay, and forward comparison across PyTorch/Inductor, JAX/XLA, and torch-spyre. You will also study failure modes that become especially important on statically planned architectures, including padding errors, scratchpad spills, host fallbacks, relayouts, and deadlocks. A worked example follows a real torch-spyre tensor-copy bug from symptom to root cause. In the lab, you will diagnose three planted failures: an injected NaN, a recompile storm, and a silent numerical divergence. For each one, you will name the layer, pick a comparison criterion and record the evidence. A 20-question graded quiz at the end covers every module of the course.

In LLM inference, moving data often takes longer than computing on it. This course covers how memory capacity, bandwidth, placement and data movement affect inference on different accelerator architectures. You will use arithmetic intensity and the ridge point to explain why low-batch decode is often bandwidth-bound and how the KV cache changes that balance as context grows. You will compare prefill and decode across multiple accelerator designs and study how data placement differs between hardware-managed caches and compiler-managed local memories. The course covers tiling, layout, alignment, software-managed scratchpads, and FlashAttention from a memory-systems perspective. The second half focuses on the KV cache: how to size it, how MHA and GQA change its footprint, and how serving systems manage it with continuous batching, PagedAttention, prefix caching, chunked prefill, speculative decoding, and tiered storage. You will also examine how these techniques change on a compiled-schedule architecture such as Spyre. In the lab, you will measure decode time and memory use on a Colab GPU across batch sizes and context lengths, compare the results with analytical predictions, and explain the gap between the model and the measurement. A 22-question graded quiz at the end covers every module of the course.

This course shows how to predict how fast an LLM inference step should run on given hardware, check that prediction against measurements, and use the difference to find practical optimizations. You will start by bounding the time of an operation from its FLOPs and bytes, then use the roofline model and ridge point to distinguish compute-limited from memory-limited execution. You will learn to count traffic at different memory levels and reason about how compiler placement decisions change off-chip movement on scratchpad-based architectures. The course then moves from analytical models to measurement. You will use PyTorch Profiler to read execution traces and separate time dominated by compute, memory movement, synchronization, and software overhead. A worked example takes a slow decode loop from diagnosis to fix. You will also study benchmarking methodology for compiled code, compiler cost models, shape bucketing, warmup precompilation, and batched decode. In the lab, you will measure achieved compute and bandwidth on a GPU, place prefill and decode on a roofline, run a batch sweep, report time to first token and time per output token, and predict how a scratchpad target would change off-chip traffic. A 20-question graded quiz at the end covers every module of the course.