CMU 11-868 L10: Accelerating Transformers on GPUs, and Where LightSeq Finds the Time

🌏 中文版

Version note: This post is based on the Spring 2026 edition of CMU 11-868 LLM Systems. Every fact was checked on 2026-09-30 against the course syllabus, the L10 slide PDF (66 pages), and the two readings. Access grade A3: slides, assignments, and starter code are all public, but the course has no public recordings, so everything below comes from the slides and papers, with page numbers.

Series: Previous: HW3: A Decoder-Only Transformer in MiniTorch | Next: HW4: Fused CUDA Kernels for Softmax and LayerNorm | Series overview

What this lecture answers

By the end of HW3 you have a trainable Transformer built on your own MiniTorch framework. It runs, but slowly. Lecture 10 asks: while the model still fits on one GPU, how much faster can single-GPU training and inference get, and where does the speed come from?

The syllabus schedules the topic over two sessions, February 16 and 18 (Part 1 and Part 2), but both use the same slide deck, which is why PDF number 11 is missing. Each session has one reading: LightSeq (inference) and LightSeq2 (training). Both come from ByteDance AI Lab, and instructor Lei Li is a co-author on each, so this lecture is the authors taking apart their own system.

Slide 6 sets the scope: this lecture only covers models smaller than GPU memory, built on the LightSeq library. Models too big for one GPU belong to the next lecture on distributed training.

Intuition: GPUs aren't slow at math, they're slow at starting work and moving data

Start with a number that surprises people. Slide 9 says a kernel launch costs about 3-5 microseconds, roughly 4,000 clock cycles. A Transformer layer has a few big matrix multiplications and a pile of small operations: bias add, dropout, residual, LayerNorm, softmax. PyTorch launches a separate kernel for each small operation by default, and each kernel reads its inputs from GPU memory and writes results back.

Slide 10 shows the cost with the simplest possible example:

• E = A + B + D in two steps (C = A + B, then E = C + D): two kernels, 4 loads and 2 stores
• One custom "add three matrices" kernel: one kernel, 3 loads and 1 store

The intermediate C never needs to touch memory. That is the core of the first technique: leave matrix multiplication to cuBLAS and fuse everything else you can.

Mechanism: LightSeq2's four techniques

Slide 8 groups the LightSeq/LightSeq2 optimizations into three categories: computational graph rewriting, rewriting dependent reductions, and memory management. The lecture then walks through four numbered techniques.

Technique 1: kernel fusion, leaving only GEMMs to cuBLAS

Slide 12 shows a full Transformer layer after re-partitioning. The GEMMs (Q/K/V projections, the two FFN linear layers, the output projection) go to cuBLAS. The operations in between, such as "bias + dropout + residual", "bias + ReLU + dropout", and "bias + reshape Q, K, V", each become one custom elementwise kernel. LayerNorm, softmax, and cross entropy become custom reduce kernels.

The slides pick a few concrete cases:

Slide 19 has a case that isn't elementwise. In an encoder-decoder model, every decoder layer projects the encoder output for cross attention. LightSeq2 concatenates all layers' weights [W1, …, WL], runs one big GEMM, and splits the result, instead of L small GEMMs.

Slides 15-16 paste code straight from LightSeq's embedding_kernels.cu. You can read the real thing on GitHub.

Technique 2: rewrite the formula, drop a synchronization

LayerNorm needs the mean μ(x) and the standard deviation σ(x). By definition, σ(x) needs μ(x) first, so the two reductions run back to back with a thread synchronization between them. Slide 21 rewrites it as:

σ(x) = √(μ(x²) − μ(x)²)

Now μ(x) and μ(x²) can be computed together, and two synchronizations become one. The slide also notes computation in FP32 and storage in FP16. Slide 22 does the same for the LayerNorm backward pass: rearrange ∇x into two sums that can run in parallel, Σ wⱼ∇yⱼ and Σ wⱼ∇yⱼxⱼ, plus elementwise work.

Softmax forward also has two reductions: a max (to prevent overflow in the exponent), then a sum of exponentials. Slide 23 says it plainly: "Two reduce: Costly!" Slide 24's answer isn't a new formula but shape-dependent tuning: rows of length ≤ 32 get one warp per row, rows between 32 and 64 get two elements per thread, and so on. Slides 25-26 show LightSeq baking parameters like block count and elements per thread into C++ templates at compile time, then picking the right instantiation at launch.

The corners of slides 22 and 24 say "You will implement LayerNorm/Softmax in hw3!". That's an older semester's numbering. In Spring 2026 these two kernels are in HW4.

Technique 3: mixed precision, with FP32 only where it's needed

Slide 27 lists the benefits of low precision: less memory for model and data, so larger batches; faster data movement between GPU memory and SMs at the same bandwidth; and up to 8x more FLOPs for FP16 than FP32. Slide 28 names the limit: forward and backward can use FP16 or FP8, but the optimizer's parameter update needs FP32. NVIDIA's APEX automates mixed precision but misses fine-grained memory optimization for LLMs.

Slide 29 shows LightSeq2's approach: parameters and gradients live in FP16 in one contiguous workspace. The trainer reads them, computes in FP32, and writes FP16 back; the FP32 copies never occupy real memory. Because every parameter sits in the same contiguous block, the whole update takes a single kernel launch instead of one per parameter tensor.

Technique 4: reuse memory based on backward-pass dependencies

Slide 30 lists each step of the self-attention backward pass and marks which intermediate tensors are never used again, so the next step can take their space. For example, once ∇out is consumed, ∇Z is written into the same buffer. The bottleneck is the gradient of the attention scores, whose size grows with the square of the sequence length; everything else fits by rotating through the same few B×L×H buffers.

The LightSeq inference paper pushes this further. Because input lengths vary, it pre-defines a maximum memory size for each kernel and lets operations with no dependency on each other share it. Section 1 of the paper reports eight times fewer memory allocations with no loss of inference speed.

Results: the numbers the slides report

Everything below is what the slides report (from the LightSeq2 experiments). Hardware and baselines differ by row:

Slide 36 also reports about 6 GB less training memory. The LightSeq2 abstract puts it as 1.4-3.5x faster than previous systems across GPUs, and a 308% training speedup on WMT14 English-German.

One thing to keep in mind when reading this table: slide 32 notes "A100 is more efficient in GEMM" to explain why the ceiling on A100 is higher than on V100. My reading: the faster matrix multiplication gets, the larger the share of time spent on non-GEMM operations, so fusing them pays off more. In other words, the speedup depends on where your bottleneck was, not on a fixed constant.

Inference: making beam search skip a sort

The second half of the deck (from p.42) turns to inference, with LightSeq (NAACL 2021) as the subject. Inference has no backward pass, but it has a bottleneck training doesn't: beam search.

Slide 45 points out that each beam search step does two things: a softmax over the whole vocabulary, then picking the top k out of k × V candidates. Sorting k × V elements is expensive. LightSeq's Hierarchical Auto Regressive Search (HARS) switches to "filter roughly, then rank precisely":

1. Split each beam's logits into k groups and take each group's max (p.46)
2. The smallest of those k maxima is a rough threshold ℛ: the k-th best is at least this large
3. Write only logits above ℛ back to memory, then sort that small set (p.47)

Slides 48-54 walk through an example with beam size 2 and vocabulary size 8: of 16 logits, only 5 need sorting. Sorting turns into a sequence of parallel max, filter, and re-rank operations.

Slide 55 adds other inference details: share tensor memory across layers, compute mostly in FP16, use float4 and half2 for bandwidth, and keep no intermediate results or gradients. Slides 61-62 report up to 14x speedup for translation inference and 6x for GPT-2 inference; the LightSeq abstract says up to 14x over TensorFlow and 1.4x over FasterTransformer.

from lightseq.training import LSTransformerEncoderLayer
config = LSTransformerEncoderLayer.get_config(
    model="bert-base",
    max_batch_tokens=4096,
    max_seq_len=512,
    fp16=True,
    local_rank=0)
ls_layer = LSTransformerEncoderLayer(config)
# replace the 1st Hugging Face layer with LightSeq2
bert_model.layer[0] = ls_layer

What this lecture doesn't cover

• FlashAttention: attention here still writes the L×L score matrix to memory; it just makes softmax fast. Fusing the whole attention without materializing the scores is the topic of this series' FlashAttention post.
• Triton: LightSeq is hand-written CUDA. The bottom of the syllabus lists an unscheduled topic, "Triton for Kernel Optimization", but it has no slides.
• Quantization: slide 64 lists it under "other approaches"; details come in the later quantization post.

How to read this lecture

1. Start with slides 9-12. Take any Transformer layer you know, count how many kernels PyTorch launches by default, and mark which ones could fuse.
2. Then read the LayerNorm and Softmax rewrites on slides 21-24. That's what HW4 has you write. Revisit the reduce kernel from HW1 and "one fewer synchronization" will mean something concrete.
3. For the papers, Section IV of LightSeq2 (The LightSeq2 System) matches the first half of the deck, and LightSeq's HARS section matches the second half.

One thing you can do tonight: run the PyTorch profiler on a forward pass of any Transformer you have, and count how many non-GEMM kernels there are and how much time they take.

Further reading

• The same ideas in Triton, and measuring before optimizing: CS336 Lecture 6: Benchmark and Profile Before Writing a Triton Kernel
• GPU memory hierarchy and the "move less data" intuition: CS336 Lecture 5: GPUs Win by Moving Data Less, Not by Making Each Thread Fast
• This course's earlier GPU programming model: L02-L04 GPU Programming and Acceleration

References

• CMU 11-868 LLM Systems (Spring 2026) course home — instructor and course description
• 11-868 Spring 2026 Syllabus — the 2/16 and 2/18 sessions share the L10 deck, one reading each
• L10 slides: Accelerating Transformer Training and Inference (PDF, 66 pages) — source of every page number here
• Wang et al., LightSeq: A High Performance Inference Library for Transformers (NAACL 2021, arXiv 2010.13887) — 4x fewer kernels, 8x fewer memory allocations, HARS, 14x/1.4x inference speedups
• Wang et al., LightSeq2: Accelerated Training for Transformer-based Models on GPUs (SC22, arXiv 2110.05722) — 1.4-3.5x training speedup, 308% on WMT14 En-De
• bytedance/lightseq (GitHub) — the CUDA kernel source the slides quote
• lightseq embedding_kernels.cu — the code on slides 15-16
• lightseq softmax_kernels.cu — the template tuning example on slides 25-26