How I Built an LLM Inference Engine From Scratch in CUDA (235 tokens/s on a 6 GB Laptop GPU)

My laptop's RTX 3050 now generates 235 tokens per second with TinyLlama-1.1B, using an inference engine where every GPU kernel is written by hand. That is 77% of the fastest this GPU could possibly go, and it took three rounds of measure, find the worst part, rewrite.

The project is tinyinfer: about 2,200 lines of C++ and CUDA that run Llama-family models with no PyTorch and no cuBLAS. This post is the story of those three rounds: what the benchmarks showed each time, what I changed, and the one "optimization" that made things slower.

TL;DR

When a model generates text one token at a time, every weight in the network is read from GPU memory exactly once per token, and each weight is used for only one multiply-add. The GPU's arithmetic units sit mostly idle; the real question is how fast the memory can feed them.

That gives a hard ceiling you can calculate before writing any code:

tokens/s  ≤  memory bandwidth / bytes of weights per token

My RTX 3050 6 GB Laptop GPU has a 96-bit GDDR6 bus at 168 GB/s. TinyLlama-1.1B stored in 16-bit floats reads 2.07 GB per token, a ceiling of about 81 tokens/s. Stored in 4 bits (plus one fp16 scale per 64 weights) it reads 0.55 GB, a ceiling of about 305 tokens/s.

This framing shaped everything after it. "Is it fast?" became "what fraction of 168 GB/s does each kernel reach?", which is a number I could measure for every kernel and every matrix shape.

The first version did the obvious things well. Each decoder layer runs RMSNorm, three projections for queries, keys and values, rotary position embeddings, attention over the KV cache, an output projection, and a gated MLP. Almost all the bytes live in the matrix-vector products (GEMV), so that kernel got the most care:

Result: 200.7 tokens/s, 66% of the ceiling. The kernel benchmark showed where the rest went:

Attention at 2,048 tokens of context reached only 17 GB/s. Two clear targets.

The worst matrices were short and wide. The down projection has 2,048 rows of 5,632 inputs, so each block of 8 rows copied 22 KB of input vector into shared memory to read only 22 KB of 4-bit weights. Half of the block's traffic went to staging, and each lane had only a couple of loads in flight at a time.

The fix was to give each warp R rows instead of one. A lane then issues R independent 16-byte loads per step, which keeps more memory requests in flight, and the staged input vector is shared across 8 × R rows. The right R depends on the shape, so instead of guessing, the engine times R = 1, 2 and 4 on every distinct matrix shape when the model loads and keeps the fastest. That takes a few milliseconds.

The autotuner's choices were not what I would have guessed: R = 2 for the down projection, R = 1 for most of the others.

TinyLlama uses grouped-query attention: 32 query heads but only 4 key/value heads, so 8 query heads share each K/V head. The v0.1 kernel ran one block per query head, which meant every K/V row was read from memory 8 times. It also used only 32 blocks on a 20-SM GPU.

First attempt (v0.2): split-K. One block per KV head and 128-step chunk of the context, serving all 8 grouped heads, with each warp keeping a running softmax (the FlashAttention trick) and a small kernel merging the chunks. K/V is now read once, and long contexts spread over many blocks. At 2,048 tokens: 125 µs → 58 µs. But at 128 tokens it got slower, 10 µs → 26 µs.

The cause: each timestep needed one warp-wide reduction per query head, 8 in a row, each a chain of 5 dependent shuffles. The kernel was waiting on latency, not memory.

Second attempt (v0.3): a shared-memory tile, one lane per timestep. The block copies its 64-step K/V tile into shared memory once. Each warp takes one query head, and each lane scores one timestep by walking that key row on its own; the rows are padded by one element so 32 lanes reading 32 rows hit 32 different banks. The softmax over the whole tile then needs just two warp reductions (max and sum) instead of eight per timestep.

Time per layer, TinyLlama shape. At long context v0.3 is 3.6× faster than v0.1; at very short context the simple kernel still edges it out, which is a cost of a few microseconds per layer.

The query, key and value projections read the same input, and so do the gate and up projections. Concatenating their weight matrices when the model loads turns five matrix-vector launches per layer into two. The small 256-row key and value matrices, which used the GPU poorly on their own, disappear into a 2,560-row matrix. Together with the new attention, decoding went from 212.3 to 235.2 tokens/s, 77% of the ceiling, and still averaged 224.7 tokens/s over a 1,900-token run.

The same benchmark also caught a regression. For models without grouped heads, such as Llama-2-7B where every query head has its own K/V head, there is nothing to share, so staging the tile in shared memory is pure overhead: 269 µs against 219 µs for the simple kernel at 2,048 tokens. The fix was one line, using the new kernel only when heads are grouped. I would not have found it without benchmarking a model shape I wasn't optimizing for.

TinyLlama-1.1B, 4-bit, RTX 3050 6 GB Laptop GPU, greedy decoding of 256 tokens.

Fast and wrong is easy, so correctness was checked in layers, and every optimization had to pass all of them before its numbers counted:

transformers library, converts it, and compares the logits at all 64 positions: about 8e-4 relative error with 100% argmax agreement in 16-bit, 8-bit and 4-bit. The remaining difference comes from storing the KV cache in fp16.tinyinfer --verify 16 on TinyLlama: worst relative error 2.5e-4, no argmax mismatches. This also paid off on something that looked like a bug. Early on, the model answered "hii" with a confident paragraph about a Quran app. Instead of guessing, I compared my prompt tokens with HuggingFace's official chat template: identical. The engine was fine; that's just a 1.1B model at temperature 0.7.

The biggest remaining gap is prompt processing. Prompts go through the one-token decode path, so a 1,000-token prompt costs as much as generating 1,000 tokens. Processing the prompt as one batch turns those matrix-vector products into matrix-matrix products, where tensor cores apply and the work becomes compute-bound instead of bandwidth-bound. That is the next version.

Can you run an LLM on a 6 GB laptop GPU?

Yes. With 4-bit weights, TinyLlama-1.1B is a 0.68 GB file and Llama-2-7B is 3.8 GB, so both fit in 6 GB of VRAM (the 7B model with a 1,024-token context). tinyinfer runs TinyLlama at 235 tokens/s on an RTX 3050 Laptop GPU.

Why is LLM token generation memory-bound?

Generating one token reads every weight once and uses it for a single multiply-add, so the GPU spends its time waiting on memory, not computing. The speed limit is memory bandwidth divided by the bytes of weights read per token.

How is tinyinfer different from llama.cpp or vLLM?

It is a small, readable engine (about 2,200 lines) built to show how fast hand-written CUDA kernels can get on a consumer GPU, with every kernel tested against a CPU reference. llama.cpp and vLLM support far more models and features; tinyinfer is a learning-friendly, from-scratch implementation focused on single-user decoding speed.

Which GPUs does it support?

NVIDIA GPUs from Turing (sm_75) and newer, on Windows and Linux, with CUDA 12 or 13.

The code, the benchmark tables and the tests are at github.com/Aadithya-kl/tinyinfer. If you have an NVIDIA GPU, run tinyinfer --bench 256 and send me the number.

Originally published on Medium.