Making small PyTorch ops faster on a laptop GPU

First, check what your GPU is actually running at

My earlier kernel numbers all came from a GPU stuck in performance state P5, with memory at 810 MHz, even in the laptop's Performance profile. I assumed it was a laptop thing. Then I ran a six-second copy that saturates memory and watched the clocks: still P5, still 810 MHz, and a "SW Power Cap" flag active while the GPU drew 46 W of an 80 W limit. Locking the memory clock with nvidia-smi was accepted and ignored.

The culprit was nvidia-powerd, the Dynamic Boost daemon. Stopping it took the same copy from 35.8 GB/s to 356 GB/s. Idle still drops back to P8, so nothing is pinned high. It needs no BIOS change, firmware or install, and a reboot brings the daemon back.

Shell · check, then stop the daemon for this boot

systemctl is-active nvidia-powerd
nvidia-smi --query-gpu=pstate,clocks.mem --format=csv
sudo systemctl stop nvidia-powerd     # undo: sudo systemctl start nvidia-powerd

If you benchmark on a laptop with an NVIDIA GPU, it is worth checking this before trusting any number. All results below are in the faster P0 state, and every experiment checks that the daemon is inactive before it times anything.

Where the time goes in a small op

Inside a CUDA graph, an empty kernel costs about 0.64 µs per launch on this GPU. Add one round trip to DRAM and it is 1.17 µs. My hand-tuned GEMV kernels for shapes up to 256×256 were already within 0 to 5% of that floor, so there was little left to win inside a single kernel. The waiting happens between kernels: each op starts, asks DRAM for its weights, and waits.

Prefetching the next layer while this one runs

While layer i runs, a tiny kernel on a second CUDA stream asks L2 to fetch the weights that layer i+1 will read. When layer i+1 starts, its weights are already close. The consumer kernels stay unchanged (cuBLAS through F.linear, cuDNN, autograd), and the prefetch never writes memory, so every output stays bit-identical. I checked that on every prefetch result in this post.

CUDA C++ · the whole prefetch kernel

extern "C" __global__ void l2_hint(const char* p, long long bytes, int stride, unsigned* sink) {
    long long step = (long long)gridDim.x * blockDim.x * stride;
    for (long long off = (long long)(blockIdx.x * blockDim.x + threadIdx.x) * stride; off < bytes; off += step)
        asm volatile("prefetch.global.L2 [%0];" :: "l"(p + off));
}

Python · run op i on the main stream, prefetch op i+1 on a side stream

for i, op in enumerate(self.ops):
    j = i + self.distance
    if j < len(self.ops) and self.plans[j]:
        self.side.wait_stream(main)          # starts when op i starts
        with torch.cuda.stream(self.side):
            for t in self.reads[j]:
                prefetch(t, **self.plans[j])
    y = op(y)
main.wait_stream(self.side)

To be sure the effect came from L2 and not from some scheduling side effect, I ran a decoy: the same prefetch traffic aimed at the wrong layer. The decoy always made things slower. Then I mapped where the real thing helps.

Three things decide the outcome. The current and next layer have to fit in L2 together. The prefetch has to finish before the current op does, and one 256-thread block issues roughly 140 GB/s of prefetches. Every extra block competes with the op you are trying to speed up. Near the L2 limit the best setting jumps around: at one shape, one block lost 14%, two blocks gained 15% and four blocks lost 9.5%.

Measuring instead of predicting

My first policy was a formula built from those three constraints. On shapes it had never seen it saved 5.5% on average, short of the 8% target I wrote down before running it. It never hurt anything, but it left most of the gain on the table.

What worked was measuring. When the CUDA graph is captured, the wrapper times a handful of candidate plans for each distinct layer shape on throwaway weights, keeps the best one only if its 95% confidence interval clears +1%, and otherwise runs plain PyTorch. That takes 0.2 to 0.6 seconds per shape.

It works in training too

In the backward pass the order flips: computing the gradient of layer i reads W_i again, so the prefetch should fetch W_{i-1}. With small batches that saved 8.84% of forward plus backward time on held-out shapes, with gradients that match plain PyTorch bit for bit.

You can install this on an unmodified model with hooks. They only act inside the model's own forward and backward, so calling a submodule on its own behaves exactly like plain PyTorch.

Python · prefetch hooks on an existing model, then capture a training step

from api import PrefetchHooks

def step(model, x):
    for p in model.parameters():
        p.grad = None
    model(x).float().square().mean().backward()

# warm up on a side stream first (a PyTorch rule for capturing backward)
hooks = PrefetchHooks(model, x, backward=True, step=step)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
    step(model, x)
graph.replay()
hooks.remove()   # back to plain PyTorch

Beyond language models

None of this is specific to transformers. The prefetch works on any chain of layers that read weights, so I tried it on MLP blocks, convolutions, a ViT, text encoders, a recommendation MLP and a reinforcement learning policy. Some of it worked well and some of it did nothing, and the split follows the same rule as before.

Python · any list of modules, tuned when it is first called

from api import PrefetchSequential

blocks = [nn.Sequential(nn.LayerNorm(768), nn.Linear(768, 768), nn.GELU()) for _ in range(72)]
model = PrefetchSequential([b.cuda().bfloat16() for b in blocks], example_input)

graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
    y = model(static_x)       # same output bits as nn.Sequential(*blocks)

MLP blocks with a LayerNorm, a linear layer and a GELU saved 8 to 12% at batch 4 and 16, and a generic nn.Sequential of such blocks saved 7.6% at width 512 and batch 4, and 10% at width 768 and batch 8. Small-batch training of MLP stacks saved 8.84%. Each of those layers weighs 1 to 1.5 MiB.

The convolutions and the end-to-end models barely moved. Batch-1 convolutions spend 14 to 35 µs inside cuDNN for at most 1.2 MiB of weights, so there is little fetch latency to hide. The text encoders and the ViT carry 3.4 to 6 MiB per layer, and the FP32 1024×1024 layers in the recommendation and RL models weigh 4 MiB, so two consecutive layers never fit in 4 MiB of L2. The ViT also runs 197 tokens per layer, which keeps the GPU busy already. Going finer, with hooks on every submodule, forced me to switch off PyTorch's fused attention fast path in the encoders, which costs 4 to 5% by itself.

On a GPU with 4 MiB of L2, cross-op prefetch pays for chains of layers up to roughly 2.6 MiB each, at small batch, whatever the model family.

A full decode step is a different story

Next I tried prefetch on whole decoder layers with the dimensions of SmolLM2-135M and Qwen2.5-0.5B. It saved 1.3%. A profile explained why: the large MLP projections already run at about 60% of peak memory bandwidth, and attention plus norms and elementwise ops took a third of the step. In the smallest model, attention alone was 45% of the gap between the step time and the time needed just to read the weights and KV cache.

So I wrote a decode attention kernel that handles one new token in a single launch: it appends the new key and value to the cache, computes the softmax in FP32, splits the cache across blocks, and lets the last block to finish merge the partial results. PyTorch's path for the same work is two copies plus a two-kernel flash attention. My first version was 2.4 times slower than PyTorch, because each warp waited on memory once per token. Keeping more loads in flight fixed that.

CUDA C++ · the last block for each head merges the partial softmax results

if (tid == 0) ticket = atomicAdd(tickets + head, 1u);
__syncthreads();
if (ticket != (unsigned)(splits - 1)) return;   // not the last block for this head
__threadfence();
// merge (max, sum, acc[64]) from every split, write the output,
// then reset the ticket so the same CUDA graph can be replayed
if (tid == 0) tickets[head] = 0u;

v4 loses with long caches and many heads, where PyTorch's flash path uses tensor cores and mine does plain FP32 math. The same trick as before handles that: time both at capture and fall back to PyTorch when mine does not win. Against the strongest PyTorch setup I could build (torch.compile for the elementwise ops, SDPA attention, the whole step in one CUDA graph), that choice plus the prefetch saved 12.66% on 12 model and context combinations it had not seen, up to 30.6% for SmolLM2-135M with 192 cached tokens. Repeating it with the GPU cooled down before each configuration gave 12.36%.

My hand-written fused RMSNorm and SiLU kernels beat eager PyTorch by 5.4% and lost to torch.compile by up to 5.3%, so I left elementwise fusion to the compiler.

I also compared against FlashInfer 0.7.2, built locally with a conda copy of nvcc 12.8. It beat PyTorch's path everywhere it ran, and my kernel beat it in 8 of 9 cases up to 1280 tokens. It does not support a GQA group of 7, which is what Qwen2.5-0.5B uses. I added it as a third option: it got picked in 3 of 12 configurations and never made anything slower.

Try it on your GPU

All of it lives in klab, my small CUDA kernel lab, under research/l2-prefetch. The experiments run on klab's own pieces: its NVRTC kernel loader, its CUDA graph timing harness and its GPU lock. Confirmatory tests had their success criteria fixed and hashed before they ran, and the ones that failed are reported here as results too.

Shell · run the GPU tests on your machine

git clone https://github.com/waybarrios/klab && cd klab
python -m unittest research/l2-prefetch/tests/test_prefetch.py \
                   research/l2-prefetch/tests/test_decode.py

The open question I care most about is what happens on GPUs with a much larger L2. Ada and Hopper cards have tens of megabytes, so layers that fail here might fit there. If you run the tests or the experiments on another GPU, or find a shape where this makes things worse, please open an issue with the shape, the precision and what you saw. Results that disagree with mine are just as welcome.

Still in progress

This is a snapshot of ongoing work. A stricter repeat of the decode-step test, waiting for the GPU to cool to 60 °C before every measurement, is running as I write this. Next on my list: larger models, where reading fewer bytes matters more than hiding latency, head sizes other than 64 in the attention kernel, and wiring all of this into the klab command line.

I will keep posting updates to the klab repository as they land, and this post will change when the numbers do. If something here looks wrong to you, I would like to hear about it.