GPU Kernel Optimization

A fused Triton kernel for RMSNorm, written after profiling a real LLM's inference to find where time actually goes, not by guessing. Runs on a single consumer GPU, no cluster, no cloud API, no framework doing the kernel work for you.

Results

Model: Qwen2.5 0.5B Instruct, loaded directly in PyTorch.

Profiling (forward pass, single token decode, CUDA event timing since this GPU's CUPTI support isn't there yet, see Architecture): Attention 55.1% of wall time, RMSNorm 19.1%, MLP 11.2%. RMSNorm is the target, not attention, it is memory bound and fusable, attention here is mostly matmul PyTorch already delegates to a tuned vendor library.

Kernel correctness: matches the model's actual RMSNorm layer within float16 rounding tolerance across decode and prefill shapes (scripts/validate_rmsnorm.py, tests/).

Isolated kernel benchmark (warmup runs, CUDA event timing, 200 repeats, median): 2.7x to 3.3x faster than PyTorch's native RMSNorm across every shape tested (scripts/benchmark_rmsnorm.py).

End to end: patched into a real model.generate() call, 50 new tokens, 10 timed runs after warmup: 1.114x speedup, 10.2% of total wall clock removed (scripts/benchmark_end_to_end.py). That lines up with RMSNorm's measured 19% share of forward pass time once tokenization and sampling overhead outside the forward pass is accounted for.

Architecture

• Kernel (src/rmsnorm_kernel.py): one Triton program per row, computes mean square, rsqrt, and the weight multiply in a single kernel launch, instead of PyTorch's several separate elementwise and reduction kernels for the same sequence of operations.
• Profiling (scripts/profile_model.py): torch.profiler's CUDA activity tracking depends on CUPTI, which fails to initialize on this GPU (CUPTI_ERROR_INVALID_DEVICE), most likely because this architecture is new enough that CUPTI does not recognize it yet on this driver. Profiling instead uses torch.cuda.Event timing wrapped around each submodule type through forward hooks, real CUDA driver timestamps, not dependent on CUPTI at all.
• Validation before benchmarking: correctness is checked against the model's own RMSNorm layer before any speed claim, and the kernel is benchmarked in isolation before being integrated, so an end to end number is never the first thing measured.

Running it

python -m venv .venv
.venv\Scripts\activate
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
pip install -r requirements.txt
python scripts/profile_model.py
python scripts/validate_rmsnorm.py
python scripts/benchmark_rmsnorm.py
python scripts/benchmark_end_to_end.py
pytest tests/

The nightly PyTorch build is not a preference, this GPU's architecture (Blackwell) is new enough that the stable release does not support it yet.