A from-scratch C++/CUDA inference engine for a Llama-architecture language model (SmolLM2-360M-Instruct), written to study where LLM inference time goes and how close it can get to the hardware limit. No PyTorch on the GPU: the engine has its own kernels for GEMV, attention, RoPE, RMSNorm, SwiGLU and sampling, a KV cache, INT8 weight-only quantization, batched decoding and CUDA-graph replay. cuBLAS is used only for prefill GEMMs.
$ python python/chat.py "Write a haiku about GPUs."
A GPU's power is a force,
A marvel of modern tech,
A driving force for all.
[fp16] prompt 37 tokens in 20.6 ms | 22 tokens | decode 205 tok/s (GPU time) | 0.28 s wall
GeForce GTX 1650 Ti Mobile (Turing TU117, 4 GB GDDR6, 16 SMs, no Tensor Cores), CUDA 13.1.
| FP16 weights | INT8 weight-only | |
|---|---|---|
| Weight bytes read per generated token | 690 MiB | 346 MiB |
| Decode speed, batch 1 (roofline bound) | 201 tok/s (258) | 310 tok/s (513) |
| Decode speed at 4,096 tokens of context | 149 tok/s | 198 tok/s |
| Decode speed, batch 8 (total) | 805 tok/s | 663 tok/s |
| Perplexity (HuggingFace FP32 reference: 11.420) | 11.420 | 11.468 |
| Greedy output identical to HuggingFace | 48/48 tokens | 48/48 tokens |
- Decode runs at 78% of the memory-bandwidth roofline. My GEMV kernel reads weights at 180 GB/s, 97% of the measured peak (186.5 GB/s) and 13% faster than cuBLAS.
- Split-KV ("flash-decoding") attention made long-context decode 3.6× faster (41 → 149 tok/s at 4k context). Finding and removing a register spill was a large part of it.
- cuBLAS's default algorithm choice was 7–11× too slow on this GPU. Timing every algorithm per shape, as TensorRT does with tactics, made prefill 2.1–7.1× faster.
- Correctness is checked against HuggingFace Transformers: in FP16, 100% top-1 agreement at every position of the test prompt and 99.8% over 4,088 positions of a book, identical greedy generations, identical perplexity.
Generating one token multiplies a single activation vector by every weight matrix, so each weight is read once and used for only two FLOPs. The speed limit is therefore bytes read per token ÷ memory bandwidth:
- weights: 362M parameters × 2 bytes = 690 MiB per token in FP16
- achievable read bandwidth, measured with a streaming kernel: 186.5 GB/s (the GPU's spec sheet says 192 GB/s)
- roofline: 186.5 GB/s ÷ 724 MB ≈ 258 tokens/s; measured 201 tokens/s (78%)
Weight matmuls take 4.35 ms of the 4.98 ms step (87%). The rest is attention, norms, RoPE, sampling and about 290 kernel launches.
| Kernel (batch 1, all five layer matrices) | Bandwidth | % of peak |
|---|---|---|
| One thread per output row | 11 GB/s | 6% |
| One warp per row, 2-byte loads | 171 GB/s | 92% |
| One warp per row, 16-byte loads | 180 GB/s | 97% |
cuBLAS cublasGemmEx |
159 GB/s | 85% |
The naive kernel is slow because the weights are row-major. When each thread owns
a row, the 32 threads of a warp read addresses K × 2 bytes apart, so every load
instruction touches 32 cache lines (uncoalesced). Giving each row to a whole warp
makes the loads contiguous. 16-byte vector loads (int4, 8 halves per lane)
then cut the number of load instructions 8×. The activation vector is staged in
shared memory, and the dot product accumulates in FP32.
The first attention kernel ran one thread block per query head. During decode that is only 15 blocks for a 16-SM GPU, each looping over the whole KV cache. Speed fell from 189 to 41 tok/s as the context grew to 4k tokens, far below the roofline, which only drops to 209 tok/s (the KV cache adds 40 KB per token).
The fix is the flash-decoding idea. Split the context into chunks. One block handles one KV head and one chunk, serving all 3 query heads that share that KV head (grouped-query attention), and computes a partial softmax. A second kernel merges the chunks with the log-sum-exp trick. The grid covers the maximum sequence length, so it can be captured in a CUDA graph, and blocks past the current length exit immediately.
Measured with interleaved A/B runs (each variant is a separate build, rotated across 3 rounds, median reported), because a laptop GPU's clock drifts with temperature:
| Decode attention (tokens/s) | 64 | 1,024 | 2,048 | 4,096 |
|---|---|---|---|---|
| One block per query head | 189 | 100 | 67 | 41 |
| Split-KV, first version (255 registers/thread, spilling) | 187 | 172 | 155 | 120 |
| Split-KV, final | 199 | 184 | 172 | 148 |
The first split-KV version was held back by a register spill. ptxas -v showed
255 registers per thread and 124 bytes of spills. The fully unrolled dot-product
loop made the compiler hoist all four query vectors (256 floats) into registers,
which also limited occupancy to about 2 blocks per SM. Partial unrolling plus
__launch_bounds__(128, 4) brought it to 72 registers, no spills, and gained
another 23% at 4k context.
Two things that did not matter, measured the same way:
- Reading V once per head group instead of once per query head made no difference. The repeated reads were already served from cache. An earlier measurement suggested it was slower, but that was the register spill.
- Chunk size (64, 128 or 256 positions) is within noise once registers are fixed.
Prefill (processing the prompt) multiplies many rows at once, so it uses cuBLAS
GEMMs. The first measurements were far too slow: a 16-token prompt took 107 ms,
21× a whole decode step. A standalone probe of cublasGemmEx (FP16 inputs,
FP32 accumulation) showed:
- The default heuristic runs at 0.25 TFLOP/s on this GPU.
- Algorithm 6 runs the same call at 3.07 TFLOP/s, close to the GPU's FP32 peak.
The engine now autotunes. The first time it sees a (shape, row count, output type), it times the default and all 24 algorithms, then caches the fastest. That's the same idea as TensorRT's tactic selection.
| Prompt (FP16) | cuBLAS default | Autotuned | Speedup |
|---|---|---|---|
| 16 tokens | 106.5 ms | 15.1 ms | 7.1× |
| 256 tokens | 492.5 ms | 96.7 ms | 5.1× |
| 2,048 tokens | 5,154 ms | 2,428 ms | 2.1× |
Individual GEMMs got 6.6–11.4× faster (full table in results/summary.md).
At long prompts, my simple prefill attention kernel is now the bottleneck: it
re-reads K/V for every query. A tiled FlashAttention-style kernel is the next step.
Each weight row is stored as INT8 with one FP32 scale (symmetric, per output channel, quantized at load). The GEMV dequantizes in registers right after the 16-byte load, so half as many bytes cross the memory bus.
- Speed: decode 201 → 310 tok/s (1.55×).
- Quality: perplexity 11.420 → 11.468 (+0.4%). Top-1 agreement with the FP32 reference is 95.6% per position, and greedy output is unchanged for 48 tokens.
- Why not 2×: INT8 reaches 60% of its roofline rather than 78%. The ~0.65 ms of non-matmul work per step doesn't shrink, and the INT8 GEMV runs at 165 GB/s.
A batch of sequences reads the weights once per step, so batch 8 costs only 2× the time of batch 1:
| Batch | FP16 total tok/s | INT8 total tok/s |
|---|---|---|
| 1 | 200 | 309 |
| 2 | 385 | 497 |
| 4 | 665 | 618 |
| 8 | 805 | 663 |
INT8 wins at batch 1–2 but loses from batch 4. My INT8 kernel converts each weight from int8 to float once per sequence in the inner loop, so as the batch grows the kernel becomes compute-bound instead of memory-bound. Converting once and reusing the values across the batch is the obvious fix.
The decode step is captured in a CUDA graph. The token and position live in device memory, so the same graph replays every step with no host round-trip. It made no measurable difference (200.5 vs 200.7 tok/s): a 5 ms GPU step hides the ~290 kernel launches, because the CPU enqueues them faster than the GPU runs them. Graphs matter when steps are much shorter than the launch overhead (smaller models, faster GPUs).
python/test_correctness.py compares every path against HuggingFace
Transformers (FP32, CPU) on a chat prompt:
| Check | FP16 | INT8 |
|---|---|---|
| Prefill logits, all positions (top-1 agreement) | 100% (mean |diff| 0.005) | 96.6% |
| Decode logits, one token at a time | 100% (mean |diff| 0.004) | 100% |
| Greedy generation vs reference | 48/48 tokens identical | 48/48 tokens identical |
| Batch of 4 vs batch of 1 | identical | identical |
| All four GEMV kernels | agree within 0.02 | — |
python/quality.py scores perplexity on 4,088 tokens of Alice's Adventures in
Wonderland (Project Gutenberg): 11.420 (HuggingFace), 11.420 (FP16 engine), 11.468 (INT8 engine).
Top-1 agreement with HuggingFace over those positions: 99.83% (FP16), 95.56% (INT8).
tools/convert.py safetensors (bf16) -> FP16 weight file; fuses q/k/v and gate/up projections
src/llm.cu the engine (C API, ~1,000 lines):
prefill(tokens) embed -> 32 x [RMSNorm, QKV GEMM, RoPE + KV-cache write, causal attention,
O-proj GEMM (+residual), RMSNorm, gate/up GEMM, SwiGLU, down GEMM (+residual)]
-> final norm -> LM head -> argmax (cuBLAS, autotuned)
decode step same layers for 1 token per sequence with custom GEMV kernels and
split-KV attention; reads token + position from device memory
-> captured once as a CUDA graph and replayed
python/llm.py ctypes wrapper + tokenizer + chat template
- Precision: residual stream in FP32; matmul inputs FP16; FP32 accumulation everywhere.
- Memory layout: weights are row-major
[out, in]; KV cache is[layer][sequence][kv_head][position][64]in FP16. - Reference: HuggingFace Transformers on the CPU, only for checking correctness.
- Batching: batched sequences must share a length. There is no paged KV cache or continuous batching (what vLLM / TensorRT-LLM add on top).
- Sampling and model support: greedy sampling only; the head dimension is fixed at 64.
- Prefill attention: the simple O(T²) kernel, not yet tiled.
- Profiling: Nsight Compute was not available. The driver restricts GPU
performance counters to administrators (
RmProfilingAdminOnly=1), so kernel analysis usesptxas -v, CUDA events and roofline math instead. - Hardware: one laptop GPU. Clocks vary with temperature (82–86 °C during runs), so comparisons use interleaved runs and medians.
uv venv --python 3.12 .venv
uv pip install --python .venv/bin/python -r requirements.txt --index-strategy unsafe-best-match
tools/setup.sh # download model + eval text, convert weights
make && make variants # engine + attention variants for the A/B test
cd python
python test_correctness.py # vs HuggingFace
python quality.py # perplexity
python bench.py # all performance measurements
python ablate_attention.py # interleaved attention A/B
python report.py # charts + results/summary.md
python chat.py "your question" # try itModel: SmolLM2-360M-Instruct, Apache-2.0, by Hugging Face. Evaluation text: Project Gutenberg (public domain).