Chapter 10 Writing and Benchmarking Custom Kernels

10.1 Benchmarking and Profiling Discipline

10.1.1 Warm-Up, CUDA Synchronization, and Wall-Clock Traps

I have watched engineers spend three days optimizing a kernel and then present timing results showing it is $10 \times$ faster than PyTorch. The kernel was not faster. They forgot to call torch.cuda.synchronize(), so they measured how long the CPU took to queue the CUDA command - microseconds - instead of how long the GPU took to execute it - milliseconds. Getting benchmarking right is a prerequisite for any GPU kernel work, and getting it wrong produces numbers that are confidently, silently wrong. This topic appears at every level from junior ML engineer to staff systems engineer; the mistake is embarrassingly common and the fix takes three lines.

GPU benchmarking has exactly two mandatory steps that developers consistently miss.

Step 1: Warm-up iterations. The first time PyTorch dispatches a CUDA kernel, it performs just-in-time compilation, loads shared libraries, allocates internal caches, and primes the GPU’s state. These startup costs are one-time artifacts - not representative of steady-state performance. If you benchmark the first invocation, you measure initialization, not computation. Always discard the first 5-10 iterations before timing.

Step 2: CUDA synchronization. The GPU and CPU are independent compute devices connected by a command queue. When Python calls a CUDA operation, the CPU adds the kernel descriptor to the queue and immediately returns - it does not wait for the GPU to finish. If you call time.perf_counter() before and after a GPU operation without synchronizing, you measure how long the CPU spent dispatching the command: typically 10-100 microseconds. The actual GPU execution - the number you care about - completes asynchronously and may take 100$\times$ longer.

torch.cuda.synchronize() blocks the CPU until all queued GPU kernels have completed. Call it once immediately before starting the timer and once immediately after the timed code to drain the queue and measure true GPU execution time.

The complete benchmark pattern:

import time, torch

def benchmark(fn, n_warmup=5, n_trials=20):
    for _ in range(n_warmup):
        fn()

    times = []
    for _ in range(n_trials):
        torch.cuda.synchronize()       # drain queue; mark t=0
        t0 = time.perf_counter()
        fn()
        torch.cuda.synchronize()       # wait for GPU to finish
        t1 = time.perf_counter()
        times.append(t1 - t0)
    return sum(times) / len(times)    # mean over trials

Average over trials rather than taking the minimum. GPU clock speeds drop under thermal load; individual measurements vary by 5-10%. The mean reflects steady-state production performance better than the best-case minimum.

Without warm-up you measure initialization cost; without cuda.synchronize() you measure dispatch cost; only with both do you measure actual GPU execution time.

The GPU’s asynchronous execution model is a feature. It lets the CPU run far ahead of the GPU, filling the command queue and keeping the GPU continuously fed without waiting for each kernel to finish. This is one of the primary reasons Python - a slow interpreted language - can drive GPU computations at near-full utilization. Python is never the bottleneck because the CPU is always running one full training step ahead of the GPU.

But this same asynchrony is what makes naive wall-clock measurement wrong. When the CPU calls torch.matmul(), it takes roughly 10-50 microseconds to construct and enqueue the kernel descriptor. The actual matrix multiply on a $4096 \times 4096$ bfloat16 matrix takes roughly 40-60 milliseconds on an A100. A benchmark without synchronization reports 0.01 ms; with synchronization it reports 50 ms. The $5000 \times$ discrepancy is entirely due to missing the sync.

The JIT warm-up trap is equally insidious but less well known. PyTorch compiles CUDA code on first use, selects optimal cuBLAS kernels by running micro-benchmarks internally, and allocates persistent workspace buffers. All of this happens transparently in the background and takes 100-1000 ms per unique operation shape. The steady-state cost is zero. Benchmarking without warm-up inflates the reported cost by orders of magnitude.

A Worked Example

Benchmarking a manual (unfused) GELU vs. the fused PyTorch implementation on an A100, on a $64 , 000 , 000$- element bfloat16 tensor.

Wrong benchmark (no warmup, no sync):

t0 = time.time()
y = manual_gelu(x)   # x.shape = (64_000_000,), device='cuda'
t1 = time.time()
print(f"{(t1-t0)*1e3:.2f} ms")   # prints ~0.1 ms

Correct benchmark:

manual_time  = benchmark(lambda: manual_gelu(x))   # 8.1 ms
pytorch_time = benchmark(lambda: torch.nn.functional.gelu(x))  # 1.1 ms

The wrong benchmark reports manual_gelu at 0.1 ms, apparently $11 \times$ faster than PyTorch. The correct benchmark reports 8.1 ms - $7 . 4 \times$ slower than PyTorch. The 80$\times$ disagreement is entirely due to missing synchronization.

The manual GELU dispatches 6 separate CUDA kernels (one each for $x^{3}$, three multiplications, the tanh, and one addition). Each kernel reads and writes the full 128 MB tensor from HBM. The fused PyTorch GELU dispatches 1 kernel that reads the tensor once and writes it once. Correct benchmarking is what reveals this $7 . 4 \times$ gap; without it you would optimize in the wrong direction.

Scaling sanity check: for a 256-dim, 4-layer MLP on batch size 256, time per step is approximately 5 seconds. Scaling from 2 to 5 steps gives 2$\times$, 3$\times$, 4$\times$, 5$\times$ of single-step time - exactly linear, as expected. If scaling is non-linear when it should be linear, something is wrong: memory fragmentation, thermal throttling, or a synchronization barrier that was not there before.

What You Actually Decide in Practice

• Use time.perf_counter(), not time.time(). perf_counter has nanosecond resolution. time.time() has millisecond or worse resolution on some systems, which means all sub-millisecond GPU operations look identical.

• Set n_warmup $\geq 5$, n_trials $\geq 20$ for production benchmarks. Five warm-up iterations is enough for PyTorch’s JIT to complete on most operations. Twenty trials gives a stable mean. For kernels faster than 100 $\mu$s, increase to 100+ trials because individual measurement noise becomes significant.

• Report the mean, not the minimum. The minimum filters out thermal throttling and cache misses, producing best-case numbers that do not reflect production behavior. If you want to separate kernel execution from scheduling jitter, use CUDA events (torch.cuda.Event(enable_timing=True)) which measure inside the GPU’s own clock.

• Benchmark end-to-end, not individual kernels in isolation. A kernel that is fast in isolation may be slow when integrated because it changes memory access patterns for subsequent kernels, triggers memory reallocation, or alters the GPU’s warp scheduler state. Always benchmark the realistic workload, then drill down to identify the bottleneck.

• Never call .item(), .cpu(), or print(tensor) inside a benchmark loop. All three force a CUDA synchronize, making your GPU and CPU effectively synchronous and eliminating the asynchrony advantage. This also applies to training loops: a print(loss.item()) after each step can double step time by serializing CPU and GPU execution.

How This Shows Up in Interviews

• “How would you correctly benchmark a custom CUDA kernel’s performance?” Two mandatory steps: (1) warm-up iterations to prime JIT compilation and caches; (2) torch.cuda.synchronize() before and after the timed region to drain the GPU command queue and measure actual GPU execution time. Average over 20+ trials using time.perf_counter().

• “A colleague benchmarks their kernel and claims it is $500 \times$ faster than the reference implementation. What is the most likely explanation?” Missing cuda.synchronize(). The CPU timer is measuring how long it took to queue the CUDA command, not how long the GPU took to execute it. Typical dispatch time is 10-100 $\mu$s; the actual kernel may take 50+ ms.

• Senior follow-up: “Your benchmark shows a 10$\times$ speedup over PyTorch on a simple fused operation. What would you verify before claiming the result is real?” (1) Correctness: run an assertion that both implementations produce numerically equivalent outputs. A kernel returning garbage can appear arbitrarily fast. (2) Input size: for very small inputs, PyTorch’s dispatch overhead may dominate both paths, making the ratio meaningless. (3) Profile both: use torch.profiler to verify each dispatches what you think it does - PyTorch may dispatch to a slower kernel path for the specific shape tested. (4) Measure memory traffic: a $10 \times$ speedup on a memory-bound operation is legitimate if you reduced HBM round-trips by $10 \times$; on a compute-bound operation it would require $10 \times$ more FLOPs/second than the hardware provides, which is a red flag.

Key Takeaways

• CUDA operations are asynchronous: the CPU returns immediately after dispatch; without torch.cuda.synchronize() you measure dispatch time ($\sim$10-100 $\mu$s), not GPU execution time ($\sim$ms to seconds).

• Warm-up iterations (minimum 5) are mandatory: first iterations trigger JIT compilation and internal cuBLAS micro-benchmarking that inflates the reported cost by orders of magnitude.

• The correct benchmark pattern: warm up $\rightarrow$ sync $\rightarrow$ start timer $\rightarrow$ fn() $\rightarrow$ sync $\rightarrow$ stop timer $\rightarrow$ mean over 20+ trials.

• Use time.perf_counter(), report the mean (not min), and always verify correctness before interpreting any timing.

• Never call .item(), .cpu(), or print(tensor) inside training loops or benchmark loops: all three force an implicit CPU-GPU sync, serializing what should be asynchronous execution.

10.1.2 Reading an NVIDIA Nsight Systems Trace

A benchmark tells you your training step takes 500 ms. A profiler trace tells you 300 ms of that is the GPU idle because a print(loss.item()) call is forcing a synchronization barrier that eliminates every CPU-ahead scheduling advantage. These are different problems requiring different solutions, and you will never find the second one from benchmarks alone. Reading profiler traces fluently is the skill that separates engineers who know that their system is slow from engineers who know why and can fix it. This shows up at the senior engineer and above level in any role involving training system performance.

GPU profiling operates at two levels of resolution.

PyTorch Profiler (quick, no external tools). Instruments aten:: operations via C++ hooks and records CPU and CUDA execution times per operation:

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ]
) as prof:
    torch.cuda.synchronize()
    fn()
    torch.cuda.synchronize()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

The output table shows: operation name (Python $\rightarrow$ aten:: $\rightarrow$ CUDA kernel), self CPU time, self CUDA time, call count. For a matrix multiply it shows aten::mm on the CPU side and the specific Cutlass kernel selected for your matrix shape on the CUDA side. For a compound operation like torch.cdist, it decomposes into constituent kernels: aten::matmul (78% of GPU time), aten::pow (5%), aten::sum (3%), aten::cat (6%). This tells you exactly where to focus optimization effort.

NVIDIA Nsight Systems (detailed, requires nsys binary). Produces a .nsys-rep file for the Nsight GUI. The trace shows three horizontal lanes:

• GPU lane (“CUDA HW”): every kernel that ran on device, its start time, duration, and name.

• CPU lanes (per thread): Python and C++ execution, kernel launches, memory copies.

• Memory lane: allocation and deallocation events.

The most important feature: the GPU and CPU lanes are shown side by side on the same timeline. You can directly observe whether the GPU is idle (no kernel running) and what the CPU was doing during that gap.

NVTX annotations mark code regions that appear as colored bands in the CPU lane:

import torch.cuda.nvtx as nvtx

for step in range(n_steps):
    nvtx.range_push(f"step_{step}")
    nvtx.range_push("forward");  loss = model(x); nvtx.range_pop()
    nvtx.range_push("backward"); loss.backward();  nvtx.range_pop()
    nvtx.range_push("optimizer"); opt.step();       nvtx.range_pop()
    nvtx.range_pop()   # step

These cost microseconds and make the trace immediately interpretable.

A profiler trace tells you exactly which CUDA kernel ran, for how long, and whether the GPU was actually busy - everything benchmarks cannot tell you.

The PyTorch Profiler costs 1-5% runtime overhead and requires zero external setup. It is the right first tool for identifying which operations dominate GPU time and verifying that kernels dispatched are the kernels you intended.

Nsight Systems operates at the hardware level: it instruments the CUDA driver directly, recording every kernel launch, memory copy, and synchronization event with sub-microsecond precision. This is the right second tool when you need to understand the relationship between CPU and GPU execution - specifically, whether the GPU is idle and what the CPU is doing during that idle period.

The key asymmetry that traces reveal: different matrix shapes dispatch to completely different CUDA kernels. A $4096 \times 4096$ matmul dispatches to a Cutlass large-tile kernel; a $128 \times 128$ matmul dispatches to a different XMMA kernel. These have entirely different performance characteristics. A benchmark reports one number per implementation; a trace shows which internal path was taken for each call.

The CPU-ahead pattern is visible in any healthy Nsight trace: the CPU timeline shows the Python code for step $N + 1$ running while the GPU timeline shows step $N$ still executing. The CPU queues kernels faster than the GPU executes them. This is correct behavior. A trace where CPU and GPU steps are perfectly aligned is a warning sign: it means a synchronization barrier - usually a .item() call or a Python-level loss print - is forcing the CPU to wait for the GPU at every step. Adding one print(loss.item()) to a tight training loop can serialize what was asynchronous execution and reduce throughput by up to 50%.

A Worked Example

Profiling a 16-layer MLP (256 dimensions, batch 256) with the PyTorch Profiler.

The top-10 table by CUDA time shows (approximate numbers on A100):

• aten::mm (forward): 31% of GPU time, 16 calls per step.

• aten::mm (backward dL/dX + dL/dW): 61% of GPU time, 32 calls per step.

• aten::add: 3% GPU time (bias additions).

• aten::relu: 2% GPU time.

• CudaDeviceSynchronize: 2% (the profiler’s own barrier).

The immediate conclusion: 92% of GPU time is matrix multiplies. Optimizing the ReLU or bias operations is wasted effort. The profiler answered “where should I spend three hours?” in thirty seconds.

Now compare the Nsight Systems trace for the same model in two configurations:

Configuration A (no print): CPU timeline shows step 0 Python running for 20 ms; GPU timeline shows step 0 executing for 180 ms starting at $t = 400$ ms (after JIT startup). By step 1, the CPU is already dispatching step 2’s kernels while the GPU runs step 1. The CPU runs one full step ahead. GPU utilization is continuous - no idle gaps.

Configuration B (with print(loss.item()) after each step): Step 1 CPU duration is now 180 ms. The CPU-step boundary aligns with the GPU-step boundary. The reason: loss.item() transfers a scalar from GPU to CPU, which internally calls torch.cuda.synchronize(). The CPU blocks for 180 ms waiting for the GPU to finish, then prints, then launches step 2. The CPU-ahead advantage is gone. Step time goes from 180 ms (GPU-paced) to 360 ms (CPU-blocked): a $2 \times$ slowdown from one print statement.

What You Actually Decide in Practice

• Add NVTX annotations to every training loop before profiling. Do it once, keep it in the code permanently. The overhead is under 10 $\mu$s per annotation pair. Without them, a Nsight trace of a multi-step training loop is almost unreadable - you cannot tell where step boundaries are.

• Never call .item(), .numpy(), .cpu(), or print(tensor) inside the training loop. All four force a CUDA synchronize that serializes CPU and GPU. For loss logging, accumulate values in a GPU tensor and transfer once every 100 steps, or use a framework like wandb with asynchronous logging.

• Use PyTorch Profiler first; escalate to Nsight Systems when you need CPU/GPU timeline alignment. The Profiler answers “which operation is slow?”. Nsight Systems answers “is the GPU idle, and why?”. They are complementary. Start with the fast tool.

• When you see different CUDA kernels for different matrix shapes, benchmark both shapes separately. PyTorch dispatches to different Cutlass kernels based on matrix dimensions (and hardware). A matmul that is fast at $4096 \times 4096$ may dispatch to a less optimal path at $512 \times 512$. Profiling reveals this; benchmarking only the happy-path size hides it.

• Profile with warm-up excluded using the schedule API:

prof = torch.profiler.profile(
    schedule=torch.profiler.schedule(
        wait=2, warmup=2, active=3),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'))

This skips the first two steps (JIT), warms up for two steps, and records three steps. Without this, the profiler captures JIT compilation time which dominates and masks the real bottleneck.

How This Shows Up in Interviews

• “Training throughput is lower than expected. How do you determine whether the bottleneck is CPU or GPU?” Use Nsight Systems. Look at the GPU timeline: if there are gaps (periods where no kernel is running), the GPU is idle, and the bottleneck is CPU dispatch rate - Python is not queuing kernels fast enough. If the GPU is continuously running kernels but each takes too long, the bottleneck is GPU computation. In a healthy system, the CPU runs one full step ahead and the GPU never goes idle.

• “How would you use the PyTorch Profiler to identify which operation to optimize in an MLP?” Wrap the training step in torch.profiler.profile with CPU and CUDA activities enabled. Call key_averages().table(sort_by="cuda_time_total") to rank operations by total GPU time. The top entry (usually aten::mm) is where to focus. A 5-minute profiler run saves 5 hours of optimizing the wrong operation.

• Senior follow-up: “A profiler shows 60% of GPU time is in CudaLaunchKernel (the kernel launch overhead) rather than in actual computation kernels. What causes this and how do you fix it?” This happens when the model has many small operations instead of a few large ones. Example: an unfused GELU implemented as 6 separate element-wise operations dispatches 6 kernels; each kernel launch takes 5-20 $\mu$s. At 100 GPU kernels per step, that is 0.5-2 ms of pure launch overhead per step. In a fast model with 50 ms steps, launch overhead is 1-4% of runtime. In a very fast model with 3 ms steps, it is 17-67%. The fixes: (1) torch.compile fuses element-wise chains automatically and eliminates most of this. (2) Write a Triton or CUDA kernel that fuses the compound operation. (3) Increase batch size or hidden dimension so each kernel has more useful work to amortize the fixed launch cost.

Key Takeaways

• The PyTorch Profiler reveals which operations dominate GPU time and which CUDA kernel is dispatched per operation; use it first, as it requires no external tools.

• Nsight Systems shows the GPU and CPU timelines side by side: CPU-ahead scheduling is healthy; a CPU/GPU timeline that is perfectly aligned indicates a synchronization barrier (usually a .item() call) that serializes what should be asynchronous execution.

• A single print(loss.item()) per step can double step time by forcing a CUDA synchronize at every iteration, eliminating all CPU-ahead scheduling benefit.

• Different matrix shapes dispatch to different CUDA kernels; only a profiler reveals this; a single benchmark number hides kernel selection differences.

• Add NVTX annotations (nvtx.range_push/pop) to training loops before profiling; they cost microseconds and make traces immediately readable.

10.2 Implementing Kernels

10.2.1 A Minimal Custom CUDA Kernel (GeLU Case Study)

Your new architecture uses a custom activation function. The naive PyTorch implementation dispatches six CUDA kernels and runs at 8.1 ms. PyTorch’s built-in fused version runs at 1.1 ms. Writing a minimal custom CUDA kernel in C++ fuses all operations into a single kernel launch and cuts runtime to 1.8 ms - an 80% improvement without changing any math. Every ML infrastructure role at a frontier lab expects you to be able to sketch this pattern, explain why the fusion matters, and know where the remaining gap to PyTorch’s implementation comes from.

A CUDA kernel is a C function marked __global__ that the GPU executes on thousands of threads simultaneously. Each thread computes its global index and handles one (or a few) elements of your data. The CPU side is a thin wrapper that validates inputs, allocates the output, computes grid dimensions, and fires off the kernel.

Every CUDA kernel for element-wise operations follows the same skeleton:

// GPU side: runs on device
__global__ void gelu_kernel(const float* in, float* out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float x = in[i];
        float y = 0.5f * x * (1.0f + tanhf(
            sqrtf(2.0f / M_PI) * (x + 0.044715f * x * x * x)));
        out[i] = y;
    }
}

// CPU side: orchestrates the launch
torch::Tensor gelu_cuda(torch::Tensor x) {
    TORCH_CHECK(x.device().is_cuda(), "input must be on GPU");
    TORCH_CHECK(x.is_contiguous(), "input must be contiguous");
    auto y = torch::empty_like(x);           // not zeros_like
    int n = x.numel();
    int blocksize = 1024;
    int num_blocks = (n + blocksize - 1) / blocksize; // ceiling division
    gelu_kernel<<<num_blocks, blocksize>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), n);
    return y;
}

The GeLU approximation being computed is:

The naive PyTorch implementation of this formula dispatches six separate CUDA kernels (one each for $x^{3}$, a scale, an add, a tanh, a multiply, and a final scale). Each kernel reads and writes the full tensor to HBM. The custom kernel reads once and writes once. That is the entire reason for the speedup.

A CUDA kernel is just a C function in which each thread computes its global index i = blockIdx.x * blockDim.x + threadIdx.x, does a bounds check, and performs its element’s computation - everything else is ordinary C.

Kernel fusion eliminates redundant HBM traffic. For a 64M-element float32 tensor (256 MB), each HBM round-trip costs roughly $256$ MB$/ 2$ TB/s $\approx 128 \mu$s, plus a kernel launch overhead of $\sim$10-50 $\mu$s. Six roundtrips: $\sim$768 $\mu$s minimum bandwidth cost plus six launch overheads. The fused kernel: one read, one write, one launch - the minimum possible HBM traffic.

The contrast with torch.compile is instructive: the JIT compiler achieves 1.47 ms on the same operation (vs. our 1.8 ms) by generating a Triton kernel with vectorized loads (4 float32 values per thread instead of 1) and better inlining. For simple element-wise fusions, the compiler wins on convenience and nearly wins on performance. The custom CUDA kernel makes sense when (a) the compiler cannot discover the right fusion pattern, or (b) you need hardware-specific intrinsics not exposed through Triton.

The remaining 0.7 ms gap between our kernel (1.8 ms) and PyTorch’s (1.1 ms) comes from two sources: our kernel processes one float32 element per thread (2$\times$ HBM bytes vs. bfloat16), and it lacks the vector load instructions (float4) that allow loading 4 values per transaction and better utilizing the HBM burst width.

A Worked Example

Input: x.shape = (64_000_000,), float32, on A100. blocksize = 1024.

Grid computation.

Total threads: $62 , 500 \times 1024 = 64 , 000 , 000$. Each thread covers exactly one element.

Thread i = 1,025 (thread index 1 in block 1):

Bounds check: $1025 < 64 , 000 , 000$ - proceed. Read in[1025], compute GeLU, write out[1025].

Last thread in the last block (thread index 511 in block 62499):

First out-of-bounds thread (thread index 512 in block 62499):

All threads in block 62499 are in bounds here because $62499 \times 1024 = 63 , 998 , 976$ and $63 , 998 , 976 + 1023 = 63 , 999 , 999 < 64 , 000 , 000$. Good: no waste. If the input had 64,000,001 elements, we would need 62,501 blocks, and block 62,500 would have 1023 masked threads doing nothing.

HBM traffic: $64$M $\times 4$ bytes read $+ 64$M $\times 4$ bytes write $= 512$ MB. At 2 TB/s: $512$ MB$/ 2$ TB/s $= 0 . 26$ ms minimum. Measured: 1.8 ms - the gap is kernel launch overhead, register allocation, and the A100’s actual compute time for the transcendentals (tanhf, sqrtf).

What You Actually Decide in Practice

• Use empty_like, never zeros_like, for output allocation. zeros_like zeros every element before you write to it - an extra HBM round-trip for a tensor you are about to overwrite. For 256 MB this costs $\sim$128 $\mu$s. Use empty_like when you will write every output element.

• Always check device and contiguousness at the wrapper level. A tensor that has been transposed or sliced may be non-contiguous: its elements are not laid out consecutively in memory, so in[i] does not give you element $i$. Assert x.is_contiguous() and call x.contiguous() if needed. Failing to do this silently reads the wrong memory.

• Set CUDA_LAUNCH_BLOCKING=1 during development. Without it, CUDA errors are deferred and surface in unrelated calls, making the error message point to the wrong line. With it, errors appear immediately at the kernel launch. Remove it before benchmarking (it adds 5-15% overhead).

• blocksize = 1024 is a safe default for element-wise kernels. An A100 SM supports up to 2048 threads per SM (64 warps). With blocksize = 1024 (32 warps per block), two blocks fit per SM, giving 64 active warps - full capacity. If register pressure is high (many temporaries), reduce to 512.

• Try torch.compile before writing any C++. For the GeLU case, torch.compile delivers 1.47 ms with zero kernel code. Write a custom CUDA kernel only when the compiler misses a fusion (check the profiler), when you need hardware intrinsics (bfloat16 tensorcore ops, H100-specific warpgroup instructions), or when Triton cannot express the required operation.

How This Shows Up in Interviews

• “Walk me through how you would implement a custom activation function as a fused CUDA kernel.” Describe the two-part structure: (1) CPU wrapper - device check, contiguous check, empty_like allocation, grid computation, kernel launch. (2) GPU kernel - compute global index, bounds check, compute activation, write output. Explain that fusion eliminates multiple HBM round-trips that the naive version incurs for each sub-operation.

• “Why is the naive PyTorch GeLU 7$\times$ slower than the fused version?” Six separate CUDA kernels, each reading and writing the full tensor to HBM. For a 256 MB tensor at 2 TB/s, each round-trip costs $\sim$128 $\mu$s. Six round-trips plus six kernel launch overheads sum to $\sim$8 ms. The fused kernel: one read, one write, $\sim$0.26 ms minimum HBM cost.

• Senior follow-up: “Your custom CUDA GeLU runs at 1.8 ms but PyTorch’s runs at 1.1 ms. What accounts for the remaining 39% gap and how do you close it?” Two causes: (1) float32 vs. bfloat16 - float32 moves twice as many bytes per element; switching to bfloat16 halves HBM traffic and halves the minimum bandwidth cost. (2) Scalar vs. vectorized loads - PyTorch’s kernel loads 4 values per thread using float4 memory instructions, fully utilizing the 128-byte HBM cache line per transaction. Fix: change the kernel to advance by BLOCK_SIZE $\times$ 4, use float4 loads, and process 4 elements per thread. This brings the kernel within $\sim$10% of PyTorch’s implementation. The Triton kernel already does this automatically.

Key Takeaways

• A CUDA kernel is a C function in which each thread computes i = blockIdx.x * blockDim.x + threadIdx.x, bounds-checks against n_elements, reads in[i], computes, and writes out[i] - the entire pattern fits in ten lines of C.

• Kernel fusion eliminates HBM round-trips: naive GeLU (6 kernels) $\rightarrow$ 8.1 ms; fused custom CUDA $\rightarrow$ 1.8 ms; PyTorch fused $\rightarrow$ 1.1 ms; torch.compile $\rightarrow$ 1.47 ms.

• Use empty_like not zeros_like; assert device and contiguous; set CUDA_LAUNCH_BLOCKING=1 during development.

• torch.compile handles element-wise fusion automatically (zero kernel code, $\sim$1.47 ms for GeLU). Reserve custom CUDA/Triton for operations the compiler cannot discover or for hardware-specific intrinsics.

• The gap between hand-written and PyTorch’s implementation comes from scalar vs. vectorized loads and float32 vs. bfloat16 bandwidth: fixing both closes most of the remaining difference.

10.2.2 Triton: Block-Level Programming in Python

You need a fused softmax kernel for your inference server. Writing it in raw CUDA means tracking warp shuffle instructions, shared memory barrier placement, and hand-coded memory coalescing - easily a day of brittle C++. Triton lets you write the same kernel in Python at the thread-block level, while the compiler handles coalescing, shared memory, and thread synchronization automatically. For most ML infrastructure engineers I have worked alongside, Triton is the practical answer to “I need a custom kernel”: it delivers CUDA-class performance with Python ergonomics and a debugger that actually works. At senior-level interviews for kernelwriting roles, interviewers expect you to sketch a Triton GeLU or softmax from memory and explain the block-centric mental model that drives the design.

Triton is a Python domain-specific language first released by OpenAI in 2021. The key conceptual shift from raw CUDA: you no longer think in threads, you think in thread blocks. In CUDA, each function invocation is a single thread computing one element. In Triton, each function invocation is one program instance - one thread block - operating on a contiguous tile of elements simultaneously. The Triton compiler maps that tile to individual threads and automatically generates:

• Coalesced memory access: adjacent tile elements map to adjacent threads, which map to adjacent addresses, fully utilizing the 128-byte HBM cache line per transaction.

• Vectorized loads: the PTX output uses 4-element load instructions (ld.global.v4.f32), moving 4 float32 values per thread per transaction without any programmer annotation.

• Masked boundary handling: tl.load(ptr + offsets, mask=mask, other=0.0) suppresses outof-bounds accesses cleanly, replacing the manual if (i < n) that every CUDA kernel requires.

What Triton does not manage: scheduling across SMs. You choose the grid (number of blocks and block size), and the GPU dispatcher decides placement. This is the one degree of freedom you retain.

A complete Triton GeLU kernel:

import triton
import triton.language as tl
import torch

@triton.jit
def gelu_kernel(x_ptr, y_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)
    c = tl.sqrt(2.0 / 3.141592653589793)
    inner = c * (x + 0.044715 * x * x * x)
    y = 0.5 * x * (1.0 + tl.libdevice.tanh(inner))
    tl.store(y_ptr + offsets, y, mask=mask)

def triton_gelu(x: torch.Tensor) -> torch.Tensor:
    assert x.is_cuda and x.is_contiguous()
    y = torch.empty_like(x)
    n = x.numel()
    BLOCK_SIZE = 1024
    grid = ((n + BLOCK_SIZE - 1) // BLOCK_SIZE,)
    gelu_kernel[grid](x, y, n, BLOCK_SIZE=BLOCK_SIZE)
    return y

Compare this to the raw CUDA version from the previous subsection: the element-wise computation is identical, but tl.arange(0, BLOCK_SIZE) replaces the per-thread index arithmetic, and mask=mask replaces the explicit bounds check. The result: 1.85 ms runtime versus 1.80 ms for hand-written CUDA, with Python code that a debugger can step through.

Triton is CUDA where you write the block and the compiler writes the threads.

Against raw CUDA, Triton wins on development speed and correctness. Memory coalescing bugs are among the most common and hardest-to-diagnose sources of GPU underperformance in hand-written kernels. Triton makes them structurally impossible: the tl.arange pattern guarantees adjacent offsets for adjacent threads, which guarantees coalesced access. Shared memory management and warp barriers are generated by the compiler, not by the programmer.

Against torch.compile, Triton wins when the JIT cannot infer the right fusion pattern. For simple element-wise ops like GeLU, torch.compile already generates Triton internally and achieves 1.47 ms with no kernel code at all. The case for hand-written Triton is algorithmic novelty: online reductions over tiles, non-standard attention masking patterns, fused quantization and dequantization, or any computation where the compiler would need to understand the mathematical structure of your operation to fuse it correctly.

The trade-off Triton does not eliminate is arithmetic intensity. A GeLU is memory-bound regardless of whether it is written in Python or C++. Fusion cuts HBM round-trips from six to one - that is the large win. No amount of Triton tuning changes the roofline for an inherently bandwidth-limited kernel.

A Worked Example

Row-wise softmax over a $4096 \times 4096$ float32 matrix on an A100 (HBM bandwidth: 2 TB/s). The natural Triton design assigns one block per row, loading the entire row into an SM in a single vectorized pass.

Grid. One block per row, block size padded to the next power of two:

Kernel:

@triton.jit
def softmax_kernel(x_ptr, y_ptr, x_stride, y_stride,
                   n_cols, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(axis=0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    x = tl.load(x_ptr + row * x_stride + cols,
                 mask=mask, other=-float('inf'))
    x = x - tl.max(x, axis=0)     # numerical stability
    x_exp = tl.exp(x)
    y = x_exp / tl.sum(x_exp, axis=0)
    tl.store(y_ptr + row * y_stride + cols, y, mask=mask)

HBM traffic. Each row is read once (16 KB) and written once:

At 2 TB/s: $128$ MB$/ 2$ TB/s $= 64 \mu$s minimum. Measured: $\sim$1.9 ms. The gap comes from the sequential dependency within each row - max must complete before exp can begin, exp must complete before sum, sum before normalization - so each block has an irreducible sequential fraction that no amount of parallelism can eliminate.

Comparison across implementations:

• Manual (6-kernel PyTorch decomposition): 3.7 ms (five HBM round-trips).

• This Triton kernel: 1.9 ms (one read, one write, sequential within-block reduction).

• PyTorch built-in softmax: 1.5 ms (fused CUDA kernel with online normalization).

• torch.compile: 1.3 ms (Triton + online normalization, shape-specialized).

The 0.6 ms gap between this kernel and torch.compile closes with online (one-pass) softmax - the topic of the FlashAttention section that follows.

What You Actually Decide in Practice

• Try torch.compile before writing any Triton. For element-wise ops and standard reductions, the JIT generates a fused Triton kernel automatically. Profile with torch.profiler first: if you see a single triton_ kernel in the CUDA trace, the compiler already did the work.

• Declare BLOCK_SIZE: tl.constexpr. This tells the Triton compiler the block size is a compile-time constant, enabling loop unrolling and better register allocation. Passing BLOCK_SIZE as a runtime integer disables these optimizations silently.

• Choose block size as a power of two. Triton’s tl.arange requires power-of-two sizes. For row-wise operations, set BLOCK_SIZE to the next power of two above your column count. If a row does not fit in SRAM ($>$ 96 KB shared memory per SM on an A100), tile across the row in a loop instead.

• Always subtract the max before exponentiating in reductions. The line x = x - tl.max(x, axis=0) is not optional. Without it, logits in the hundreds produce inf from tl.exp, which propagates to NaN in the division. This is one of the easiest ways to silently break a custom attention kernel.

• Reserve raw CUDA for hardware-specific intrinsics. FlashAttention-3’s gains on H100 come from warpgroup matrix-accumulate (WGMMA) and asynchronous copy instructions that Triton’s abstraction cannot yet express. For everything else - element-wise ops, reductions, tiled attention - Triton is within 5-10% of hand-written CUDA in my experience and is far less likely to contain silent memory bugs.

How This Shows Up in Interviews

• “Sketch a fused softmax kernel in Triton.” Describe the row-per-block grid. Walk through: pid = tl.program_id(0), offset computation with tl.arange, masked load with other=-inf, in-register max subtraction for stability, tl.exp, tl.sum, normalization, masked store. Emphasize that coalescing and thread management are handled by the compiler.

• “When would you use Triton instead of torch.compile?” When the JIT cannot infer the right fusion: non-standard memory layouts, online reductions over tiles, fused quantization-dequantization, or custom attention masking patterns. torch.compile wins for standard element-wise ops and matmuls; Triton wins when you need to encode algorithmic knowledge the compiler does not have.

• Senior follow-up: “Your Triton softmax runs at 1.9 ms but torch.compile achieves 1.3 ms. What accounts for the gap and how do you close it?” The naive two-pass softmax (pass 1: find max; pass 2: exp and normalize) reads the row twice. Online softmax reduces this to a single pass: maintain running maximum $m$ and running sum $d$ across tiles, and at each new tile apply the correction $e^{m_{\text{prev}} - m_{\text{new}}}$ to the accumulated sum before adding new terms. This halves HBM reads and is the foundational technique in FlashAttention’s tiled attention kernel. Implementing it in Triton requires tracking $( m , d )$ as scalar accumulators across a loop over column tiles.

Key Takeaways

• Triton programs in blocks: pid = tl.program_id(0) gives the block index, tl.arange(0, BLOCK_SIZE) gives the tile, mask = offsets < n handles boundaries - three lines replace all thread-level index arithmetic.

• The compiler generates coalesced vectorized loads automatically, moving 4 float32 values per thread per transaction without programmer annotation.

• Performance target: Triton $\text{GeLU} \approx$ 1.85 ms; hand-written CUDA $\approx$ 1.8 ms; torch.compile $\approx$ 1.47 ms. Start with torch.compile; write Triton when the JIT cannot find the pattern.

• Always declare BLOCK_SIZE: tl.constexpr; always subtract the row max before tl.exp; choose block sizes as powers of two.

• The gap between a two-pass Triton softmax and torch.compile’s softmax is closed by online (one-pass) normalization - the technique that makes FlashAttention’s tiled attention kernel viable.

10.3 FlashAttention as a Case Study

10.3.1 Online Softmax and Tile-by-Tile Attention

When someone asks you to design an LLM serving system for 128K-token context, the first cost you should compute is how much HBM the attention operation alone consumes. Standard attention materializes an $N \times N$ score matrix in GPU global memory. At $N = 8 , 192$, 32 heads, and bf16 storage, that matrix occupies over 4 GB of HBM per layer-before activations, KV cache, or gradients. FlashAttention eliminates that $O ( N^{2} )$ allocation by computing attention entirely within SRAM-resident tiles and streaming K and V from HBM exactly once. This topic surfaces in almost every senior LLM systems interview that touches context scaling, memory efficiency, or custom kernel design; understanding it at the algorithm level, not just by name, is what distinguishes a strong candidate.

The root challenge is that softmax is a global row-wise operation. To normalize row $i$ of the attention score matrix $S = Q K^{T} / \sqrt{d}$, you need the denominator$\sum_{j} e^{S_{i j} - m_{i}}$ where $m_{i} = \max_{j} S_{i j}$. A naïve implementation therefore writes $S$ to HBM, reads it back to compute the row maxima, writes $P = \operatorname{softmax} ( S )$, and reads $P$ again to compute $P V$ -four HBM round-trips that scale as $O ( N^{2} )$ in both time and space.

The online softmax algorithm eliminates this by maintaining two running scalars per query row as K tiles stream in: a running maximum $m$ and a running normalizer $\ell$. When tile $j$ arrives with block-level statistics

The correction factor $e^{m - m}$ rescales the accumulated sum whenever the running maximum updates. The output accumulator $O$ follows the same pattern:

After processing all $T_{c} = N / B_{c}$ KV tiles, divide $O$ by $\ell$ to obtain the correctly normalized output. FlashAttention-2 structures the outer loop over query tiles of size $B_{r}$ and the inner loop over KV tiles of size $B_{c}$, dispatching one warp group per query tile so that each SM processes an independent slice of the sequence.

Online softmax converts a global normalization into a sequence of SRAM-resident corrections, enabling single-pass attention whose HBM footprint scales as $O ( N d )$ rather than $O ( N^{2} )$.

Naïve attention on an H100 is memory-bound, not compute-bound. At $N = 4 , 096$, 32 heads, and bf16, the attention score matrix alone requires $\approx 4 . 2$ GB of HBM writes and reads per forward pass per layer. At the H100’s 3.35 TB/s HBM bandwidth ceiling, those transfers consume $\approx 1 . 25$ ms per layer before any compute executes. The compute for the same attention operation at peak bf16 throughput of 1,979 TFLOP/s takes under 0.2 ms-so the bottleneck is entirely in HBM traffic.

FlashAttention keeps Q, K, V, the score tile $S$, and the output accumulator $O$ tile-resident in the $\approx$228 KB of shared memory available per SM on an H100. The trade-off is SRAM scarcity: all five tiles must fit simultaneously, which constrains the block sizes $B_{r}$ and $B_{c}$. Increasing $B_{r}$ reduces the tile-loop trip count and improves compute intensity, but reduces the number of concurrent blocks that can co-reside on the SM-so occupancy and compute intensity pull in opposite directions, and the right setting depends on the specific $d$ and hardware.

Compared with an unfused implementation, FlashAttention also eliminates multiple kernel launches-each of which would flush intermediate results to HBM. Operator fusion alone removes two of the four HBM passes; tiling removes the remaining quadratic dependence on $N$.

A Worked Example

Model: 7B-parameter decoder, $d_{\text{head}} = 128$, $H = 32$ heads, $N = 4 , 096$, bf16. Block sizes $B_{r} = B_{c} = 64$.

SRAM footprint per SM during one outer tile:

• $Q$ tile: $64 \times 128 \times 2 = 16$ KB

• $K$ tile: $64 \times 128 \times 2 = 16$ KB

• $V$ tile: $64 \times 128 \times 2 = 16$ KB

• Score tile $S$: $64 \times 64 \times 2 = 8$ KB

• Output tile $O$: $64 \times 128 \times 2 = 16$ KB

• Running stats $m$, $\ell$: $64 \times 4 \times 2 \approx 0 . 5$ KB

• Total: $\approx 73$ KB - comfortably within the 228 KB shared memory limit.

HBM traffic comparison per head, one forward pass:

• Naïve attention: write $S$ (32 MB), read $S$ for softmax (32 MB), write $P$ (32 MB), read $P$ for $P V$ (32 MB), plus Q, K, V reads and $O$ write ($\approx 4$ MB). Total $\approx 132$ MB per head.

• FlashAttention: Q read once (1 MB), K and V each read once per outer tile (1 MB each), $O$ written once (1 MB). Total $\approx 4$ MB per head.

Across 32 heads: naïve attention moves $\approx 4 . 2$ GB of HBM per layer per forward pass; FlashAttention moves $\approx 128$ MB-a 33$\times$ reduction. At 3.35 TB/s, the naïve HBM transfers cost $\approx 1 . 25$ ms per layer. FlashAttention reduces that cost to $\approx 0 . 04$ ms.

What You Actually Decide in Practice

• Block size is the primary tuning knob. For $d = 128$, start with $B_{r} = B_{c} = 64$ and benchmark $B_{r} = B_{c} = 128$. The larger block reduces the inner-loop trip count by $4 \times$ and improves arithmetic intensity, but doubles SRAM usage and reduces SM occupancy. Benchmark both; on H100 with $d = 128$ the $B_{r} = B_{c} = 128$ configuration often wins, but the gap closes for shorter sequences.

• Align head dimensions to powers of two. A model with 20 heads over a 1,600-wide hidden layer gives $d_{\text{head}} = 80$. Pad to 128 inside the kernel. Unaligned dimensions produce non-coalesced HBM loads that partially cancel the bandwidth savings.

• Do not rewrite standard FlashAttention yourself. torch.nn.functional.scaled_dot_product_attention with enable_flash_sdp(True) dispatches to a tuned FlashAttention-2 implementation. Write a custom Triton kernel only for non-standard masking patterns, hardware-specific pipelining (e.g., H100 TMA), or structurally novel attention variants.

• Causal masking adds no HBM traffic. KV tiles whose column range lies entirely above the diagonal can be skipped. Mixed tiles apply the causal mask to the score tile in SRAM before the softmax update. No additional memory movement.

• Always handle the final partial tile. Sequences not divisible by $B_{c}$ require masked loads. Forgetting boundary conditions is the most common correctness bug in custom FlashAttention kernels-the results will be wrong for the last few positions in the sequence without triggering a visible error.

How This Shows Up in Interviews

• “Why does naïve attention become slow at long context?” Strong answer: the $N \times N$ score matrix requires at minimum four HBM round-trips. At $N = 8 , 192$ and 32 heads in bf16, that is $\approx 4 . 3$ GB per layer per forward pass-fully bandwidth-limited, costing $\approx 1 . 3$ ms per layer before any compute. FlashAttention tiles K and V through SRAM, brings HBM traffic down to $O ( N d )$, and makes the operation compute-bound.

• “Walk me through one tile of the FlashAttention forward pass.” Strong answer: load $Q_{i}$, $K_{j}$, $V_{j}$ into SRAM; compute $S_{i j} = Q_{i} K_{j}^{T} / \sqrt{d}$; update running $\text{max} m^{\prime} = \max ( m , \mathrm{rowmax} ( S_{i j} ) )$; rescale running normalizer $\ell^{\prime} = \ell e^{m - m} + \mathrm{rowsum} ( e^{S_{i j} - m} )$; update output accumulator $O^{\prime} = e^{m - m} O + e^{S_{i j} - m} V_{j}$; advance to the next KV tile. After all tiles, return $O / \ell$.

• Senior follow-up: “FlashAttention-3 adds warp-specialization and uses H100 TMA. What problem does that solve and why is it H100-specific?” The remaining bottleneck in FlashAttention-2 on H100 is that the online-softmax update in registers stalls the next GEMM from issuing. H100’s WGMMA instruction and Tensor Memory Accelerator allow one warp group to issue an asynchronous GEMM while another runs the softmax update and prefetches the next KV tile via TMA DMA. A100 lacks both WG-MMA and TMA, so implementing producer-consumer pipelining requires explicit barriers that consume most of the latency benefit.

Key Takeaways

• FlashAttention’s core contribution is a data-movement strategy: compute attention entirely in SRAMresident tiles so the $N \times N$ score matrix never writes to HBM.

• The online softmax update-maintaining running $\text{max} m$ and sum $\ell$, correcting both by $e^{m_{\text{old}} - m_{\text{new}}}$ on each tile-enables single-pass normalization without materializing the full score matrix.

• HBM traffic falls from $O ( N^{2} )$ per head (naïve) to $O ( N d )$ per head (FlashAttention), a factor-of-$N / d$ reduction that scales linearly with sequence length.

• Block size $B_{r} \times B_{c}$ is the primary tuning parameter: larger blocks improve arithmetic intensity but must fit in SRAM; benchmark both 64 and 128 for your head dimension.

• In production, use scaled_dot_product_attention with the FlashAttention backend; write a custom kernel only for hardware-specific pipelining or structurally novel masking.

10.3.2 Backward-Pass Recomputation in FlashAttention

If you have trained a long-context model and watched the forward pass succeed at batch size 4 while the backward pass throws a CUDA out-of-memory error at the same configuration, the most likely culprit is the stored attention probability matrix. Standard autograd saves $P = \operatorname{softmax} ( S ) \in \mathbb{R}^{N \times N}$ per head per layer because the softmax backward requires it. At $N = 4 , 096$, 32 heads, and bf16 over 32 layers, that storage alone consumes $\approx 34$ GB-enough to OOM a single 80 GB H100 before accounting for weights, optimizer states, or KV cache. FlashAttention’s backward pass discards $P$ immediately after the forward tile completes and recomputes it tile by tile from Q, K, V during the backward pass. Understanding why that recomputation is faster than loading the stored matrix-not just why it saves memory-is the answer that separates a senior response from a mid-level one.

The softmax backward formula for attention row $i$ is:

In attention, $\partial \mathcal{L} / \partial P_{i j} = d O_{i} V_{j}^{T}$ . The row-wise contraction simplifies to a per-query-position scalar $D_{i}$:

Substituting gives the final expression for the score gradient:

$D_{i}$ is a vector of length $N$ per head (one scalar per query position) that can be computed in a single pass through $d O$ and $O$ before the tiled backward loop begins.

The key ingredient from the forward pass is not $P$ itself but the log-sum-exp statistic $L_{i} = m_{i} + \log \ell_{i}$ per query row-a vector of length $N$ per head stored in fp32. Using $L_{i}$, any tile of $P$ can be reconstructed as $P_{i j} = e^{S_{i j} - L_{i}}$, where $L_{i}$ is broadcast row-wise. The backward tile loop then proceeds entirely in SRAM:

1. Load $Q_{i}$, $K_{j}$, $V_{j}$ from HBM into SRAM.

2. Recompute $S_{i j} = Q_{i} K_{j}^{T} / \sqrt{d}$; compute $P_{i j} = e^{S_{i j} - L_{i}}$ (row-wise using stored $L_{i}$).

Storing only $L_{i} \in \mathbb{R}^{N}$ per head from the forward pass and recomputing $P$ tile by tile during the backward reduces attention activation memory from $O ( N^{2} )$ to $O ( N )$-an exchange that consistently wins on hardware where compute throughput scales much faster than HBM bandwidth.

On an H100, bf16 matmul throughput peaks at $\approx$1,979 TFLOP/s and HBM bandwidth peaks at $\approx$3.35 TB/s. Recomputing one tile of $S_{i j}$ of shape $B_{r} \times B_{c}$ with head dimension $d$ adds $2 B_{r} B_{c} d$ FLOPs and avoids reading $B_{r} B_{c} \times 2$ bytes from HBM. The arithmetic intensity of this tile recomputation is:

For $B_{r} = B_{c} = 64$ and $d = 128$, this is 32 FLOP/byte-below the H100 ridge point of $\approx$590 FLOP/byte, so the recomputed tile is memory-bound. The critical comparison, however, is not against the hardware ceiling; it is against the alternative of loading a precomputed tile from HBM. Loading $P_{i j}$ costs the same $B_{r} B_{c} \times 2$ bytes at the same bandwidth, but adds zero FLOPs-the compute units sit idle. With recomputation, those same memory cycles are accompanied by useful arithmetic, and the FLOPs finish before the next HBM load arrives. The recomputation is effectively free in terms of wall-clock time, while eliminating the $O ( N^{2} )$ activation storage.

Compared with gradient checkpointing at the layer level, FlashAttention recomputation is far more targeted. Layer-level checkpointing re-runs the full forward pass-MLP, normalization, attention alike. FlashAttention recomputes only the tile-level $Q K^{T}$ multiply and the pointwise softmax during the backward pass, a fraction of layer compute, executed at near-optimal HBM utilization.

A Worked Example

Model: 7B-parameter decoder, $d_{\text{head}} = 128$, $H = 32$ heads, $N = 4 , 096$, bf16 training with fp32 gradient accumulation. Block sizes $B_{r} = B_{c} = 64$.

Memory cost of storing $P$ per layer:

Across 32 layers: $\approx 34$ GB per training step-exceeding available HBM once weights, optimizer states, and remaining activations are included.

Memory cost of the FlashAttention approach per layer:

Across 32 layers: $\approx 16$ MB-a reduction of roughly $2 , 000 \times$.

Extra compute per layer from backward recomputation:

The backward pass recomputes $Q K^{T}$ for all $H \times ( N / B_{r} ) \times ( N / B_{c} )$ tiles:

At 1,979 TFLOP/s: $137$ GF$/ 1 , 979$ TF/s $\approx 0 . 07$ ms per layer.

HBM time saved per layer by not loading stored $P$ during backward:

Net per-layer saving: $0 . 32 - 0 . 07 = 0 . 25$ ms of backward-pass time per layer, plus the $2 , 000 \times$ activation memory reduction that makes the backward pass feasible at sequence lengths that would otherwise OOM.

What You Actually Decide in Practice

• Prefer FlashAttention recomputation over layer-level gradient checkpointing at the attention operation. Checkpointing at the layer boundary re-runs the full forward pass-MLP and normalization included. FlashAttention recomputes only the tile-level GEMM, which is already close to bandwidth-optimal and costs far less than full-layer recomputation.

• Always save $L_{i}$ in fp32 regardless of training precision. The log-sum-exp statistic is used as a subtraction offset $e^{S_{i j} - L_{i}}$ during backward. In bf16, rounding error in $L_{i}$ accumulates across tiles and produces gradient drift that is difficult to diagnose. The reference FlashAttention-2 implementation enforces fp32 storage for $L_{i}$ unconditionally.

• Precompute $D_{i}$ before the tiled backward loop. $D_{i} = \langle d O_{i} , O_{i} \rangle$ requires one pass through $d O$ and $O$, both $O ( N d )$ tensors already in HBM. Computing $D_{i}$ inside the tile loop re-reads $O$ on every iteration, multiplying HBM traffic by $T_{c} = N / B_{c}$-a $64 \times$ overhead for $N = 4 , 096$ and $B_{c} = 64$.

• Backward tile ordering affects gradient accumulation patterns. $d K_{j}$ and $d V_{j}$ accumulate contributions from all $N / B_{r}$ query tiles. FlashAttention-2 loops over query tiles in the inner loop for the $K$ and $V$ gradient pass to avoid atomics, at the cost of loading $K_{j}$ and $V_{j}$ once per outer query tile rather than once per KV tile. This increases HBM reads for K and V during the backward pass. Profile the atomic versus non-atomic variants on your hardware-the winner depends on SM count and sequence length.

• Sequence parallelism introduces an asymmetric all-reduce. With ring-attention or Ulysses-style sequence parallelism, each device holds a shard of query positions and produces partial $d K$ and $d V$ gradients. An all-reduce across devices is required before the optimizer step for $d K$ and $d V$ ; $d Q$ accumulates locally and needs no communication. Schedule the all-reduce to overlap with the next micro-batch forward pass-it is frequently on the critical path for long-context training.

How This Shows Up in Interviews

• “Why does FlashAttention save memory during training, not just inference?” Strong answer: the backward pass for standard attention requires the stored $N \times N$ probability matrix $P$. FlashAttention stores only $L_{i} \in \mathbb{R}^{N}$ per head (the log-sum-exp) and recomputes $P$ tile by tile during the backward pass-reducing attention activation memory from $O ( N^{2} )$ to $O ( N )$, a $2 , 000 \times$ reduction at $N = 4 , 096$, $H = 32$.

• “Isn’t recomputing part of the forward pass during the backward pass expensive?” Strong answer: no, because the alternative-loading the precomputed tile from HBM-costs the same bandwidth with zero compute. The recomputed tile’s FLOPs execute while the GPU waits for the next HBM load, so they overlap rather than stack. Per layer, the extra compute costs $\approx$0.07 ms while the saved HBM transfer is worth $\approx$0.32 ms-a net 0.25 ms saving per layer, plus the activation memory reduction that enables the backward pass to run at all.

• Senior follow-up: “How does FlashAttention’s backward interact with sequence parallelism, and what communication is required?” With ring attention, each device processes a contiguous shard of queries and accumulates partial $d K$ and $d V$ from remote KV tiles that are passed around the ring. After the ring completes, each device holds the full $d Q$ for its query shard (no communication needed) but only a partial $d K$ and $d V$ (an all-reduce is required). The senior answer identifies this asymmetry, notes that the all-reduce volume equals $2 \times$num_heads$\times N \times d$ bytes-the same scale as the KV weight gradients-and discusses overlapping it with the next micro-batch’s forward pass ring to hide latency.

Key Takeaways

• FlashAttention discards $P$ after each forward tile and stores only $L_{i} \in \mathbb{R}^{N}$ per head (the log-sum-exp); the $N \times N$ attention matrix is recomputed tile by tile during the backward pass.

• Attention activation memory drops from $O ( N^{2} )$ to $O ( N )$ per head-at $N = 4 , 096$ and 32 heads over 32 layers, this saves $\approx 34$ GB per training step.

• Recomputation overhead is small: extra FLOPs execute during HBM-wait cycles, adding $\approx 0 . 07$ ms per layer while saving $\approx 0 . 32$ ms of HBM transfer time.

• Always store $L_{i}$ in fp32 to prevent gradient accumulation errors from bf16 rounding in the log-sum-exp offset.

• With sequence parallelism, $d K$ and $d V$ require an all-reduce while $d Q$ does not; overlap the all-reduce with the next micro-batch’s forward ring to keep it off the critical path.