Triton, CUTLASS, and PyTorch Operators

Suppose a PyTorch profile shows that layer normalization and a following residual addition take 180 microseconds for a particular tensor shape. A developer can express a fused replacement as ordinary PyTorch operations, ask torch.compile to generate code, write a Triton kernel, assemble a CUTLASS operation, or write CUDA C++. Every route can compute the same numbers. The routes differ in how directly they expose thread assignment, memory layout, synchronization, matrix instructions, automatic differentiation, and integration with the rest of PyTorch.

Chapter 11 follows one optimized operation from a Python call to GPU execution. The first lesson compares implementation options using evidence from the missing capability, not a blanket ranking of languages. The second lesson launches a small Triton vector-add program and connects program identifiers, blocked offsets, masks, and compile-time parameters to the output array. The third lesson registers an operation with PyTorch so fake tensors, autograd, autocast, and torch.compile can reason about the operation without executing unsafe guesses.

Correct output from one contiguous FP16 tensor does not make a production operator complete. Real framework calls can contain odd sizes, transposed strides, non-contiguous views, BF16 values, gradient recording, graph capture, and devices that lack a specialized instruction. An operator test suite must make those cases observable. The fast implementation also needs an eligibility predicate: a Boolean decision that sends only supported inputs to the optimized kernel and sends every other valid input to a correct fallback.

PyTorch code begins on a central processing unit (CPU) and launches the selected work on a graphics processing unit (GPU). CUTLASS commonly builds matrix multiplication (GEMM) kernels. Triton and compiler tools can print an intermediate representation (IR), which is the compiler-owned program form between source code and target instructions. The lesson connects each printed artifact to one measured launch.

The final lesson constructs the evidence expected for a maintained or upstream contribution. A benchmark table covers representative shapes, including shapes with no improvement. Correctness tests state tolerances and edge cases. Compiler logs show whether the operator composes with graph capture. The resulting lesson treats kernel code as one part of a software interface that must remain correct as models, devices, and framework versions change.