Code-Series on Jaehun's Blog

1. An Inference Engine in 1,500 Lines of CUDA: The Map

Thu, 17 Sep 2026 00:00:00 +0900

jmaczan/tiny-vllm is an inference engine for Llama 3.2 1B Instruct, written from scratch in C++ and CUDA. No PyTorch, no Hugging Face, not even a tokenizer. It opens the safetensors file itself, pushes the weights onto the GPU, and computes attention, RMSNorm, and softmax in its own kernels.

2. Weight Loading and GPU Buffer Design

Thu, 17 Sep 2026 00:00:00 +0900

In this engine, right after creating the cuBLAS handle, main() immediately does two things (src/main.cpp:557-569 ): it takes the weights out of model.safetensors and puts them on the GPU, and it preallocates the buffers that every computation from here on will use. This chapter reads both.

3. The prefill path: flowing the entire prompt at once

Thu, 17 Sep 2026 00:00:00 +0900

In Part 1 I said prefill() is a function that processes the entire prompt in one pass. All 17 tokens pass through everything from embedding to next-token selection within a single call. This part looks inside that function. We follow which kernels the tokens pass through on their way through one layer, in what order, and why that order has the shape it does, by reading each kernel’s code.

prefill is simple in code. It pops a prompt from the queue, turns it into embeddings, and runs one loop over the 16 layers.

4. The cuBLAS Transposition Trick — Why the Matrices Are Passed Reversed

Thu, 17 Sep 2026 00:00:00 +0900

In part 3, when looking at the Q/K/V projections, I wrote this sentence and moved on: “Activations and weights are row-major, but cuBLAS assumes column-major, so the call is made using only the CUBLAS_OP_T transpose flag and the lda/ldb/ldc arguments. Why it’s called this way is covered in part 4.” This part is that “why.” While part 3 looked at the order in which the kernels flow, this part pulls out the same cuBLAS call again and reads each argument one by one.

5. The decode path that produces one token — why separate kernels were made

Thu, 17 Sep 2026 00:00:00 +0900

Once prefill has streamed the whole prompt through in one pass, everything after that has to produce tokens one at a time based on the tokens produced so far. That step is decode. As we saw in part 1, prefill and decode have different operation shapes — the former is matrix × matrix, the latter vector × matrix — so this repo makes separate kernels. They’re the variants with Decode in the name, like embeddingGatherKernelDecode, ropeKernelDecode, softmaxKernelDecode, and one more pagedAttentionKernel with no prefill counterpart. This part lays those kernels next to their prefill counterparts and reads what differs and why the shape is different, and how a new token’s K/V gets appended to the cache.

6. Splitting the KV cache into 16-token blocks — reading non-contiguous blocks via the block table

Thu, 17 Sep 2026 00:00:00 +0900

Every time decode produces one token, attention re-reads all the K/V that the sequence has accumulated so far. If that K/V is held in one large contiguous buffer per sequence, sequences that finish at different times leave unused gaps and waste memory. PagedAttention, the idea that made vLLM famous, answers by splitting the KV cache into small blocks — like operating system pages — allocating them only when needed, and letting an array called the block table track “which chunk of which sequence lives in which block” (Kwon et al., SOSP 2023 ). This repository implements that idea with BLOCK_SIZE = 16-token blocks, a block table, and a single kernel, pagedAttentionKernel.

7. Overlapping multiple requests with slots and a queue — implementing continuous batching

Thu, 17 Sep 2026 00:00:00 +0900

If the GPU processes only one sequence at a time, it sits idle while that sequence finishes prefill and moves on to decode. Continuous batching is the approach of feeding the next request waiting in the queue into the spot of a request that finishes, so the GPU is never idle. This repository’s README also summarizes the principle in one paragraph — “fill slots with prompts, and when any slot finishes generating, return its result and then fill the just-emptied slot with a prompt waiting in the queue.”