CS336 — LLM Systems

From-scratch implementations built for Stanford CS336 (Language Modeling from Scratch), covering GPU kernel programming and distributed training systems for large language models.

Contents

• Flash Attention (Triton)
• Fully Sharded Data Parallel (FSDP)

---

Flash Attention (Triton)

A from-scratch implementation of Flash Attention as a custom Triton GPU kernel. Includes a PyTorch autograd reference implementation, a Triton forward/backward kernel, and a Modal-based benchmarking script comparing against naive attention and PyTorch's fused scaled_dot_product_attention.

Contents

• flash_attention.py

  • MyFlashAttnAutogradFunctionClass — reference PyTorch implementation of flash attention (online softmax, single-pass forward + analytic backward) used to validate correctness of the Triton kernel.
  • TritonFlashAttention — torch.autograd.Function wrapping the Triton kernel, usable as a drop-in replacement for standard attention.
  • flash_fwd_kernel — the Triton JIT kernel implementing the tiled forward pass: block-wise QKᵀ, online softmax (running max/sum), causal masking, and accumulation of the output block, following the FlashAttention algorithm.

• benchmark_attention.py

  • Modal app (attention-benchmark) that runs on an H100 and times forward and backward passes across a grid of d_model and sequence lengths, comparing:
    • naive — explicit QKᵀ → softmax → PV, materializing the full attention matrix (the O(N²) memory baseline)
    • torch_sdpa — torch.nn.functional.scaled_dot_product_attention (PyTorch's fused/dispatched kernel, for reference)
    • triton_flash — the custom Triton kernel above
  • Reports forward/backward latency, peak memory, and speedup/memory reduction relative to the naive baseline.

How it works

Standard attention is memory-bound: it materializes an (seq_len, seq_len) score matrix in HBM, which the GPU has to write and re-read for softmax and the final matmul with V. Flash attention avoids this by tiling Q, K, V into blocks that fit in SRAM and computing the output incrementally with an online (running) softmax — keeping track of a running max m and running normalizer l per query, and rescaling the accumulated output O as new key/value tiles are processed. The score matrix itself is never fully materialized, which turns an O(N²) memory-traffic problem into roughly O(N²d / SRAM_size), and lets attention scale to much longer sequences before running out of memory.

Setup

pip install modal
modal setup   # one-time auth, opens a browser to link your Modal account

Local dependencies for running/testing outside Modal:

pip install torch triton einops numpy

Running the benchmark

From the directory containing flash_attention.py and benchmark_attention.py (so the local module import resolves):

modal run benchmark_attention.py

This builds a Modal image (torch, triton, einops, numpy), ships flash_attention.py into it, and runs the benchmark remotely on an H100. Output streams live to your terminal: per-config timings as each implementation finishes, followed by a summary table of forward/backward speedup and memory reduction versus the naive baseline.

Note: the naive implementation is expected to run out of memory at long sequence lengths — that's the intended demonstration of flash attention's memory advantage, not a bug.

Benchmark results

Measured on Modal (H100), batch_size=8, comparing Triton flash attention against naive PyTorch attention (full QKᵀ materialization):

d_model	seq_len	fwd speedup	bwd speedup	peak mem reduction
16	256	6.10x	0.63x	3.21x
16	1024	4.89x	0.76x	12.18x
16	4096	9.73x	0.76x	18.61x
32	256	0.27x	0.65x	2.96x
32	1024	4.62x	0.72x	10.40x
32	4096	1.62x	0.73x	17.39x
64	256	0.61x	0.62x	2.58x
64	1024	11.77x	0.59x	8.13x
64	4096	2.22x	0.68x	15.39x

Takeaways:

• Memory reduction is the clearest, most consistent win — up to ~18.6x peak memory reduction over naive attention at seq_len=4096, growing with sequence length as expected, since naive attention's O(N²) score matrix is exactly what flash attention avoids materializing.
• Forward speedup is real but noisy, especially at small sizes (e.g. d_model=32, seq_len=256 shows a 0.27x "speedup" — actually a slowdown, most likely Triton kernel launch overhead dominating at tiny problem sizes where there isn't much memory pressure to relieve yet). Speedup trends upward and stabilizes at larger d_model/seq_len.
• Backward speedup is consistently below 1x (0.59x–0.76x) — the Triton kernel's backward pass is currently slower than naive PyTorch. This matches a known limitation of this implementation: the backward pass recomputes attention weights in plain PyTorch from the saved logsumexp L rather than using a fused Triton backward kernel, so it doesn't get the same memory-bound advantage the forward pass does. A fused Triton backward kernel (following FlashAttention-2's tiled backward approach) would be the natural next step to close this gap.

Status / known limitations

• Q_TILE_SIZE / K_TILE_SIZE are currently set conservatively (16); larger tiles (64–128) tuned to the target GPU's SRAM would better reflect production flash-attention kernel performance.
• Backward pass for TritonFlashAttention currently recomputes attention weights in PyTorch (not a fused Triton backward kernel) — benchmark data above confirms this makes backward consistently slower than naive PyTorch (0.59x–0.76x). A natural next step would be a fused Triton backward kernel following the same tiling approach as the forward pass.

---

Fully Sharded Data Parallel (FSDP) — from scratch

A from-scratch PyTorch implementation of Fully Sharded Data Parallel training. Shards nn.Linear and nn.Embedding parameters across ranks (ZeRO-3 style) and overlaps communication with compute via forward/backward hooks, exposed as a transparent nn.Module wrapper.

Contents

• _ShardedLayer — per-layer sharding and communication state.

  • Splits each wrapped module's parameters into a flat buffer, pads to be divisible by world_size, and keeps only this rank's contiguous shard as the "true" parameter (local_shard) that the optimizer sees and updates.
  • start_gather / finish_gather_and_install — async all_gather_into_tensor to reconstruct the full parameter on-device right before it's needed, then splices the flat buffer back into the module's named tensors.
  • start_reduce_scatter / finish_reduce_scatter — async reduce_scatter_tensor (sum, then divided by world_size) to combine gradients across ranks and leave each rank holding only the gradient for its own shard — no full gradient is ever materialized on any single rank.
  • free — drops the gathered full parameter after use so peak memory stays at shard size, not full-model size.

• FSDP — the user-facing wrapper.

  • Walks the wrapped module's submodules, replaces every nn.Linear / nn.Embedding's parameters with a _ShardedLayer, and registers the local FP32 shards as this module's own parameters (so AdamW(fsdp_model.parameters(), ...) works with no special-casing).
  • Registers forward pre/post hooks and full backward pre/post hooks per layer to gather parameters just before they're needed and free them immediately after, with 2-layer-ahead prefetching in both directions (i+2 on the forward pass, i-2 on the backward pass) so the next layer's all-gather overlaps with the current layer's compute.
  • finish_gradient_synchronization() — blocks until every in-flight reduce-scatter (and any stray all-gather) has completed, to be called after loss.backward() and before the optimizer step.

How it works

Standard data parallelism (DDP) replicates full model parameters, gradients, and optimizer state on every GPU — memory cost scales with model size regardless of how many GPUs you have. FSDP instead shards parameters, gradients, and optimizer state across ranks, so each GPU only permanently holds 1/world_size of the model.

The trick to making this fast rather than just memory-efficient is overlap: right before a layer's forward (or backward) needs its full parameters, an async all_gather reconstructs them from all ranks' shards — and while that layer computes, the next layer's gather is already kicked off in the background via the hook-based prefetch. After compute, the full parameter is freed immediately, so at any instant only ~2 layers' worth of full parameters are resident, not the whole model. Gradients are synchronized with reduce_scatter instead of all_reduce: each rank ends up with the (mean) gradient for only its own shard, rather than every rank redundantly holding the full gradient — the same sharding principle applied to the backward pass.

Design choices / precision

• FP32 master shards: local_shard (what the optimizer updates) is always FP32, while the gathered, compute-time parameter can be cast to a lower compute_dtype (e.g. bf16) if passed to FSDP(...). This mirrors the mixed-precision "FP32 master weights" pattern used in production FSDP/ZeRO implementations.
• Gradient reduction: reduce-scatter sums gradients across ranks and divides by world_size, matching DDP's mean-gradient semantics.

Status / known limitations

• Fixed prefetch depth: gather-ahead distance is hardcoded to 2 layers in both directions, not adaptive to layer size or available memory.
• Hard sync point: finish_gradient_synchronization() blocks on every pending reduce-scatter; there's no overlap between the optimizer step and communication (an optimization real FSDP implementations pursue further).
• Only shards nn.Linear and nn.Embedding — other parameterized modules (e.g. nn.LayerNorm) are left unsharded and fully replicated.

---