Triton and hand-rolled kernels

Kernels are how you stop paying for memory traffic you do not need. FlashAttention is the well-known title for fused attention that keeps the T x T matrix from landing in HBM. You will not re-implement FlashAttention this week. You will write a Triton kernel so the words "program id", "block", and "mask" are not theatre.

Coding: vector_add

import torch
import triton
import triton.language as tl

@triton.jit
def vector_add(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    n = out.numel()
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
    vector_add[grid](x, y, out, n, BLOCK_SIZE=1024)
    return out

Run it on CUDA if you have it. If you only have a Mac or CPU this week, still type the kernel and write the launch story: grid of blocks, each block covers BLOCK_SIZE elements, mask handles the tail. Then time x + y in PyTorch versus the kernel on a machine that can run Triton. Report those timings only.

Why attention fusion matters

Naive attention writes the score matrix to HBM, reads it for softmax, writes probabilities, reads them for the value mix. That is a memory novel. FlashAttention tiles the computation so softmax happens in SRAM and the T x T tensor never becomes a full resident. Your TinyGPT from week 3 is the naive replica. This week you know why production serving does not keep that matrix.

Assignment

• Land vector_add and a unit check against x + y.
• Write a 20-line note: SRAM vs HBM, and why a fused kernel can beat a correct PyTorch graph.
• Optional stretch: a fused multiply-add. Do not claim FlashAttention parity.

Opinion checkpoint

Next: Cluster orchestration.