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

• tinyinfer is an open-source LLM inference engine written from scratch in C++ and CUDA, with no PyTorch and no cuBLAS.
• It runs Llama-family models (TinyLlama, Llama-2, Vicuna) at 235 tokens/s on an RTX 3050 6 GB laptop GPU, 77% of the hardware's memory-bandwidth limit.
• The three biggest wins: 4-bit weights, autotuned multi-row GEMV kernels, and GQA-aware flash-decoding attention (3.6× faster at long context).
• Output matches the HuggingFace reference within 0.1%.

The speed limit is memory, not math

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.

v0.1: a straightforward engine at 66%

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:

• One warp per output row, each lane loading 16 bytes at a time so a warp reads 512 contiguous bytes.
• The input vector in shared memory, padded with one float every 32, so the lanes' strided reads never hit the same memory bank.
• 4-bit weights unpacked in registers with an exponent trick: OR the 4-bit value into the mantissa of 2²³ and subtract, which is one integer op and one float add instead of a slower int-to-float conversion.
• The whole decode step captured as one CUDA graph, so about 250 kernel launches per token cost one launch.

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.

Round 1: let each warp do more rows, and let the GPU choose

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.

Round 2: attention when eight heads share one K/V

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.

Round 3: fusion, and an optimization that made things slower

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.

How I knew it was still correct

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

1. A CPU reference engine. A plain C++ version of the whole model, written to be obviously correct rather than fast.
2. The CPU engine against HuggingFace. A test builds a random Llama in the 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.
3. Every GPU kernel against the CPU reference. Matrix-vector in all three formats, with every rows-per-warp setting, on shapes that include odd ones like 333 × 11008; attention with group sizes 1, 3, 4 and 8, head sizes 64 to 256, and contexts that cross chunk boundaries.
4. The whole GPU engine against the CPU engine on the real model. tinyinfer --verify 16 on TinyLlama: worst relative error 2.5e-4, no argmax mismatches.
5. The tokenizer and chat template against the originals. Token ids match TinyLlama's own tokenizer and chat template exactly.

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.

What I'd tell someone starting, and what's next

• Compute the ceiling first. Bandwidth divided by bytes per token tells you how far off you are, and per-kernel percentages tell you where.
• Benchmark every shape, not just the one you're tuning. The R = 4 setting I expected to win usually didn't, and the attention regression only showed up on a model I wasn't targeting.
• Build the reference before the fast version. Every one of these optimizations was a rewrite of a kernel, and the tests are why I could make them in an afternoon without fear.
• Latency is a bottleneck too. The v0.2 attention read the right bytes but stalled on chains of shuffles. Fewer dependent steps beat fewer bytes there.

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.

FAQ

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.