Twelve CUDA kernels, written to saturate a GPU.
12 hand-written CUDA kernels for LLM, MoE, RAG, and graph workloads, wrapped in Rust and exposed to PyTorch as zero-copy, GPU-resident drop-in ops.
$ maturin develop --release
compiling csrc/kernels/*.cu with nvcc -O3
linking libcuda_kernels.a -> custom_cuda._native
Finished in 41.2s
$ pytest tests/ -v -m cuda
1,665 passed
$ python examples/llama_block.py
Correctness check passed - bf16 tolerance
Latency 65.401ms -> 50.836ms 1.29x
Peak memory 2268.2MB -> 1740.2MB -23.3%
The last three lines are the Llama-3-8B integration proof: a real decoder block, benchmarked, not simulated.
Hand-written CUDA kernels
RMSNorm to FP8 quantization
Correctness tests
fp32 / fp16 / bf16, real edge cases
Best single-kernel speedup
persistent-kernel Viterbi decoder
Real Llama-3-8B block
528MB VRAM saved, measured end-to-end
Language layers, zero-copy throughout
CUDA C++ → Rust/PyO3 → PyTorch
Lifecycle steps enforced per kernel
baseline → kernel → bind → test → bench → plot
Kernels reporting an honest shortfall
MatMul+Bias, Cosine Top-K, Pairwise Distance, Graph MsgPass
Peak memory bandwidth, real hardware
RTX 4070 Laptop, CUDA-event measured
Eager execution isn't slow because Python is slow.
Three structural bottlenecks recur across transformer, MoE, and retrieval workloads: none of them fixed by a faster language, only by matching kernel design to how the hardware actually moves data.
Memory bandwidth
Tens of TFLOPS of compute, only 256 GB/s of memory bandwidth. RMSNorm, SwiGLU, and RoPE do little arithmetic per byte moved: running each as a separate PyTorch op pays a full read+write of GPU memory per op, for operations that are almost pure data movement.
Intermediate allocation
Cross-entropy over a 128K-token vocabulary materializes a [batch·seq, 128,256] logits tensor, several GB in fp32, that exists only to be immediately reduced away, directly capping the batch size and context length a training step can fit in VRAM.
Kernel launch overhead
A Viterbi decoder implemented as a Python loop over timesteps launches one kernel per step. At thousands of timesteps, dispatch overhead, not compute, becomes the dominant cost, and it's essentially flat regardless of how much work each launch actually does.
Three languages, zero copies between them.
A PyTorch tensor's raw device pointer travels through Python and Rust and into a hand-written CUDA kernel, and back, without ever touching host memory or an intermediate allocation.
no CMake, no setuptools: a single `cargo build` (via `build.rs`) drives `nvcc`
An ordinary tensor already resident on the GPU: fp32, fp16, bf16, or fp8. The starting and ending point of every call.
Allocates output tensor(s) and calls the native extension. No shape or dtype logic lives here: every validation decision is pushed down into the Rust layer below.
Validates device, dtype, contiguity, and shape before any pointer reaches C++. Extracts the tensor's raw device pointer and PyTorch's current CUDA stream, then calls the extern "C" launcher: zero host-device copies, zero intermediate allocation.
Warp-shuffle reductions, register-blocked GEMM tiling, shared-memory staging, vectorized uint4/float4 loads and stores, and persistent kernels: compiled by nvcc via a Cargo build.rs script, no CMake or setuptools in the path.
1extern "C" void launch_rmsnorm_fwd(2 const void* input, const void* weight, void* output,3 int64_t rows, int64_t cols, float eps,4 cudaDataType_t dtype, cudaStream_t stream5);6 7// Rust reads this signature directly: no name mangling,8// no templates crossing the FFI boundary. Only pointers9// and integers move; the tensor's GPU memory never does.12 kernels, six domains, every result honest.
Filter by domain, or open any card for its actual optimization history: what was tried, what was measured, and whether it cleared the target that was written down before the kernel was benchmarked.
No kernel is "done" until all six steps pass.
The same lifecycle, applied uniformly to all 12 kernels: optimization techniques are applied only after profiling shows a real bottleneck, never speculatively.
every kernel below, no exceptions: a kernel isn't "done" until all six steps are complete
Baseline
PyTorch eager + torch.compile(max-autotune) reference, checked to agree with each other before either is trusted.
baselines/<kernel>.pyCUDA Kernel
A straightforward, correct first implementation: never assumed-optimal before it's measured.
csrc/kernels/<kernel>.cuRust / PyO3 Binding
Validates dtype, device, contiguity, and shape before any pointer reaches C++.
src/kernels/<kernel>.rsCorrectness Tests
Per-dtype tolerance table, non-contiguous tensors, and an edge-case battery: must pass before any benchmark counts.
tests/test_<kernel>.pyHardware Benchmark
CUDA-event timing, L2-cache flush between iterations, ≥100 measured runs, median + IQR.
benchmarks/<kernel>_bench.pyVisualization
Four standard charts per kernel, rendered non-interactively to PNG + SVG.
scripts/plot_<kernel>.py1start = torch.cuda.Event(enable_timing=True)2end = torch.cuda.Event(enable_timing=True)3torch.cuda.synchronize()4 5start.record()6kernel_under_test(*args)7end.record()8torch.cuda.synchronize()9 10elapsed_ms = start.elapsed_time(end)11# every measured iteration is preceded by a 256MB L212# cache flush: no benefiting from an artificially warm cacheReal numbers, measured on real hardware.
CUDA-event timing, L2-cache-flushed, median + IQR across ≥100 iterations: on an RTX 4070 Laptop GPU, 256 GB/s peak memory bandwidth.
Parallel Viterbi vs. per-timestep loop
persistent kernel, O(seq_len) → O(1) launches
Linear cross-entropy peak VRAM
never materializes the logits tensor
FP8 quantization vs. two-pass
80–87% of theoretical peak bandwidth
Llama-3-8B decoder block, real model
528MB VRAM saved · bf16 · seq_len=4096
Kernel 4 (Linear Cross-Entropy) reads below 1x here on purpose: its target was a ≥4x VRAM reduction, not latency, and it clears that by 12.5–31.3x. Speed and memory aren't the same axis.
Executive scorecard: all 12 kernels
Real matplotlib output, not a mockup: latency speedup vs. eager (left) and hardware saturation (right) for every kernel's headline benchmark case, including the honest low bars: MatMul+Bias at 0.1x, Cosine Top-K at 0.6x.
Parallel Viterbi: speedup vs. sequence length
Log-log scale, 5x target line barely visible at the bottom. The persistent-kernel design clears it by more than an order of magnitude at every tested sequence length, across fp32/fp16/bf16.
Does it survive contact with a real model?
Every kernel above is validated against synthetic [M, N] test tensors: necessary, not sufficient. A real Llama-3-8B-shaped decoder block, built two ways from the same weights, closes that gap.
Meta's actual released Llama-3-8B config.json values: not round numbers picked for convenience.
- ·Kernel 1: Fused RMSNorm + Residual
- ·Kernel 2: Fused SwiGLU
- ·Kernel 3: Fused RoPE
QKV / output projection GEMMs, the FFN's two big projections, and attention itself: identical PyTorch code, identical weights, in both blocks.
Numerically verified correct (bf16 tolerance) before any timing was collected. Deliberately modest relative to the per-kernel numbers above: only 3 of roughly 10 ops in the block were replaced, and they're the cheapest, most memory-bound ones already; the GEMMs and attention that dominate a real block's FLOPs are untouched in both versions.
Six optimization attempts, diagnosed before deciding.
Every rejected optimization was benchmarked, root-caused, and reverted or fixed on the evidence: not on a hunch. The two that mattered most produced their own charts.
RMSNorm+Residual: shared-memory row-staging
Hypothesized the pass-2 re-read of (x + residual) cost a full DRAM round trip, so it was staged into shared memory during pass 1 to avoid the re-read.
No measurable improvement, and a ~5% regression on several fp16/bf16 shapes: the 16–46KB shared-memory reservation per block reduced SM occupancy, and the "wasted" re-read turned out to already be an L2 hit, not a DRAM trip. Reverted to the simpler recompute version.
RoPE vectorization regressed bandwidth on nearly every shape
The first vectorized kernel made bandwidth worse than the scalar version, the opposite of the intended effect: root-caused to a fixed 256-thread block size inherited unchanged from the scalar kernel, leaving 240+ of 256 threads idle in the vectorized loop.
Fixed block size to 32 (one warp) for the vectorized path. Recovered 8–16 percentage points of bandwidth across nearly every shape (e.g. one representative shape: 67.8% → 83.1% of peak).
MatMul+Bias v1 hit a hard shared-memory-bandwidth ceiling
The naive tiled GEMM benchmarked at a flat ~1.2–1.4 TFLOPS regardless of M/K/N or dtype: 7–30x slower than eager. Cross-checked against a published benchmark of an equivalent kernel (~1.3 TFLOPS vs. cuBLAS's ~15–20) to confirm this wasn't a bug.
1D register blocking (8 output elements per thread) reached ~1.8–3.0 TFLOPS, a real 2x gain: but still 3–14x short of cuBLAS. Stopped there rather than push a rushed tensor-core rewrite under time pressure; the shortfall is reported, not hidden.
Cosine top-k v1 was 20–97x slower than eager, and got worse as candidates grew
Grid size was ceil(num_queries / 8) only. Realistic RAG shapes (few queries, huge candidate pools) meant the entire similarity scan ran on a handful of warps on a 36-SM GPU, the wrong direction entirely for a top-k kernel.
Redesigned as a two-kernel partition + merge pipeline, splitting each query's candidate pool across independent warps. Recovered to 1.7–5.1x slower than eager: still short of target, but the scaling behavior changed from catastrophic to proportional.
Pairwise distance: widening kBlockK from 8 to 16
Hypothesized the flat ~3.2–3.8 TFLOPS ceiling (regardless of embedding dimension) was the per-tile __syncthreads() overhead from kBlockK=8 failing to amortize as tile count grows.
Measured no improvement, and a regression at xlarge and small-dim cases from the doubled shared-memory footprint hurting occupancy. Reverted: confirmed the real ceiling was the same shared-memory-bandwidth bound MatMul+Bias hit, not a synchronization-overhead problem.
MoE router: 8 warps/block landed at only 1.1–2.2x, below the 3x target
Every case in this kernel executes in well under a millisecond: a launch-overhead-bound regime, not a compute-bound one, which isn't the failure mode more warps per block obviously fixes.
Doubled to 16 warps/block, halving the blocks launched per token count. The Mixtral-representative shape went from 2.2x to 3.35x, clearing the target; the rest landed at 1.4–2.3x.
RoPE speedup, after the block-size fix
The chart the launch-config bug fix actually produced: real measured bandwidth, not a projection, once the vectorized path stopped running with 240 idle threads per block.
MatMul+Bias: TFLOPS vs. cuBLAS
The honest gap: this kernel's ceiling against a tensor-core-backed vendor GEMM, reported exactly as measured rather than benchmarked against a flattering baseline.
What this library doesn't do.
- ·MatMul+Bias and Cosine Top-K remain slower than their cuBLAS-backed eager baselines: closing that gap needs tensor-core (WMMA/MMA) intrinsics, a different technique from CUDA-core tiling, and was scoped out.
- ·Pairwise Distance and Graph Message Passing meet their speedup target only in part of their tested range (modest embedding dimension; N ≥ 20,000 nodes): the shortfall at the other end is reported, not trimmed from the benchmark sweep.
- ·The integration proof (Llama-3-8B block) validates 3 of the 12 kernels composed inside a real model. The other 9 are validated in isolation against synthetic test tensors, not inside a second full model.
- ·All numbers are from one machine: an RTX 4070 Laptop GPU, 256 GB/s peak bandwidth: real and reproducible on that hardware, not vendor-reported, simulated, or re-measured on a datacenter GPU.
The signature of a bottleneck matters more than the fix.
- A flat TFLOPS or bandwidth ceiling that holds regardless of problem size is a specific, recognizable signature (shared-memory-bandwidth-bound), not just "needs more tuning": recognizing it early stopped two separate optimization attempts (MatMul+Bias, Pairwise Distance) from chasing a fix that couldn't work.
- A regression after adding vectorization is not evidence vectorization was the wrong idea: RoPE's v2 regressed because of a leftover scalar-kernel block size, not because the vectorized approach itself was flawed. Diagnosing which one it is before reverting mattered.
- The most convincing number in this project is not the biggest speedup, it's the smallest one reported honestly: MatMul+Bias at 1.8–3.0 TFLOPS, still behind cuBLAS, shipped and explained rather than quietly dropped from the results.
Twelve kernels, one benchmark harness, every number reproducible.