pico-vllm

11:42 AM

Wanted to actually understand vLLM instead of just using it. Not "read the paper and nod along" understand — build-it-with-my-own-hands understand. So: pico-vllm. A from-scratch, single-GPU inference engine implementing PagedAttention-style KV-cache management and continuous batching, on a 4050 with 6GB VRAM.

About pico-vllm A minimal serving engine that reimplements the two ideas that make vLLM fast — paged KV-cache memory management and iteration-level continuous batching — verified against a real Qwen2.5-0.5B-Instruct model, not a toy.

Repo: github.com/grvwrk/pico-vllm

Scoped it down hard before writing anything. No Triton, no custom CUDA kernels, no tensor parallelism. Gather the scattered KV blocks with plain PyTorch indexing, run scaled_dot_product_attention on the result. The interesting engineering here is memory management and scheduling, not kernel writing — decided early that pretending otherwise would just mean shipping nothing.

Reading before writing code

• PagedAttention / vLLM paper — Efficient Memory Management for LLM Serving with PagedAttention (Kwon et al., SOSP 2023)
• Orca — A Distributed Serving System for Transformer-Based Generative Models (OSDI 2022), where the actual term "iteration-level scheduling" comes from
• nano-vllm (GeeeekExplorer) as a reference implementation — read it after getting my own block manager working, not before, so I wouldn't unconsciously copy the structure and lose the ability to explain my own decisions

Day 1

Started at the very bottom, before touching any GPU code at all: how do you even divide memory into fixed-size blocks. Sounds trivial. Wasn't, because I kept mixing units — total_capacity in bytes vs. tokens, block_size overloaded to mean two different things in the same function. Had to force myself into a rule: the block manager thinks purely in block indices and token counts, full stop. Bytes only exist at the one point where the real tensor gets allocated, nowhere else.

free_block_pool = set(range(num_block))
allocated_blocks = {}  # seq_id -> [block indices], in order

A set for the free pool — any free block is as good as any other, so O(1) grab/return beats anything fancier. An ordered list per sequence, not a set, because blocks have to stay in order: block 0 holds tokens 0–15, block 1 holds 16–31, and later code depends on reading them back in that sequence. allocate_block raises instead of blocking when the pool's empty — decided early that "what happens when memory runs out" isn't this class's job. I didn't have a scheduler yet, so that decision had nowhere to live regardless.

Wrote three isolated tests for it: no cross-sequence collisions, correct exhaustion behavior, correct reuse after freeing. All passed clean. Good, boring, exactly what a foundational piece should be.

Then wired it to a real GPU tensor. Argued with myself over the shape for a while before landing on:

(num_layers, num_block, num_kv_heads, num_tokens_per_block, head_dim)

head_dim last so writing one token is one contiguous memory chunk, not a strided mess across the tensor. num_layers/num_block first because that's the actual access pattern — key_cache[layer, block] should be a clean slice.

Wanted to verify this wasn't just "shapes that compile" but actually correct, so I ran a real sentence through Qwen2.5-0.5B-Instruct and compared against its own past_key_values. Immediately hit AttributeError: 'DynamicCache' object has no attribute 'key_cache' — HF changed the cache internals since whatever version I half-remembered from tutorials. Had to dir() my way down through DynamicCache → .layers[0] → .keys like an actual detective, no shortcuts.

That same test caught something real, not just an API mismatch: this model uses grouped-query attention. 14 query heads, only 2 actual KV heads. My cache was sized using the wrong head count until I printed model.eval() and compared k_proj's output size against q_proj's directly. Sizing the cache by KV heads instead of query heads is exactly what real vLLM does — and I only actually understood why by getting it wrong first, watching the test fail, and tracing it back.

Also spent an annoying stretch chasing a torchvision::nms does not exist error that had nothing to do with any of this — a broken torch/torchvision version pairing pulled in by an unrelated import chain. Fixed, moved on.

Day 2

Built Sequence — the per-sequence bookkeeping layer sitting between the block manager and the cache: which blocks does this sequence own, in order, and how many real tokens has it written so far.

Almost gave Sequence a direct reference to BlockManager so it could self-allocate whenever it ran out of room. Talked myself out of it deliberately: deciding who gets memory when the system's under contention is a scheduling decision, not something one sequence should unilaterally grab for itself. It also would've made Sequence untestable without a live BlockManager sitting next to it in every single test. Went with "raise a ValueError and let the caller allocate and retry" instead — same pattern the block manager already used.

Found a genuinely embarrassing edge case while doing this: seq_len % block_size == 0 is True both when a block has just filled up and when seq_len == 0, i.e. nothing's been written at all. Without an added seq_len > 0 guard, the very first token a sequence ever receives would incorrectly demand a brand new block it didn't need. Small bug. The kind that's completely invisible until you actually trace through the numbers by hand.

Wrote three tests here too: normal appends within one block, the boundary-crossing case including the raised-error path, and a check that written data is actually retrievable afterward — not just that the counters look right.

Then the gather step: reconstructing a sequence's real KV data across however many blocks it owns, with the padding trimmed off. Realized something useful before writing any code — only the last block a sequence owns can ever be partially full, because a new block only ever gets requested once the current one is entirely done. So the simplest correct approach is concatenate everything first, slice once at the very end, rather than trimming block-by-block inside the loop.

Attention itself needed grouped-query head expansion — repeat_interleave(7, dim=0), specifically not .repeat(). repeat_interleave gives [kvhead0×7, kvhead1×7]; plain .repeat gives [kvhead0, kvhead1, kvhead0, kvhead1, ...]. Only one of those actually lines each query head up with the correct KV head it's supposed to attend through — got this wrong on paper first, walked through it with real numbers, then wrote the code.

Learned PyTorch forward hooks for the first time specifically to verify this against real attention output — hooked q_proj to grab the raw query, and o_proj's input (not output) as ground truth, since that's the real attention result before the final projection.

First run failed. Not randomly wrong — consistently wrong, drifting further off the longer the sequence got. That pattern was the clue: diagnosed it as RoPE, specifically that my captured query was pre-rotation while HF's cached keys (pulled from past_key_values) are already post-rotation. Rotated key attending against an unrotated query. Applied the same rotary embedding to my query before comparing, reran, passed clean.

Later

Wired the full 24-layer forward pass next — every linear layer, every layernorm, the MLP blocks, all using the model's real pretrained weights untouched. The only thing swapped out is attention, where paged_attention slots in.

First full generation test diverged at token 3 out of 5 — not garbage, not zero, just slightly off. Rather than assume bug, checked the actual pattern: tokens 0–2 matched exactly, and the divergence didn't compound wildly after that. That shape pointed at floating-point drift, not a logic error. Rewrote the test to compare per-step logits closeness against HF re-run on the identical prefix my own pipeline had produced, instead of requiring two independently-generated greedy trajectories to agree token-for-token — a much fairer bar once floating point noise is in play. In bfloat16: small, flat, non-compounding drift. Reran the exact same test in float32: clean match, tight tolerance. Confirmed it was precision, not a bug, by actually changing the one variable and watching the result change — not by just deciding to believe it.

Built the naive batching baseline deliberately, not as an afterthought — needed something demonstrably wasteful to compare the real scheduler against, and matched generation lengths across sequences wouldn't have shown anything, so the test prompts had to be mismatched on purpose (three short, one long).

Then the actual scheduler: WAITING / RUNNING / FINISHED states per sequence, admission and eviction re-evaluated on every single iteration, not once per batch. A finished sequence frees its blocks the instant it's done, rather than sitting there occupying a slot while the rest of the batch grinds on.

Hit a real crash while stress-testing this: set the block pool artificially low to force the scheduler's waiting queue to actually do something, and a running sequence needing a new block with the pool fully dry just crashed with an unhandled RuntimeError. Not a bug in the scheduling logic exactly — more an honest missing feature (no preemption path exists yet). Fixed the test scenario rather than pretending the gap wasn't real, and wrote it down as a known limitation instead of a solved problem.

The numbers that made this worth writing about:

• Naive vs. continuous, same four prompts, same hardware, same model: 135 wasted forward passes eliminated down to zero, and 2.47 → 6.93 tokens/sec — a 2.8x throughput improvement from eviction efficiency alone.
• Squeezed the block pool down to 3 blocks for 4 requests specifically to force admission to matter. It did — one request sat queued for 8 full decode steps and only got admitted once an earlier sequence finished and freed its block. Naive batching structurally can't do this; it has no concept of a waiting queue and would've failed outright trying to allocate for everyone up front.

After the core was solid, kept pushing: refactored the single-model forward pass into a proper ModelRunner/ModelRegistry pattern (mirroring how real vLLM supports 200+ architectures — one adapter class per model family, everything above it architecture-agnostic), wrapped the scheduler in an async SchedulerService so concurrent HTTP requests actually get continuously batched together instead of each spinning up isolated work, and added SSE streaming plus an OpenAI-compatible /v1/chat/completions endpoint on top.

Looking Ahead

No preemption yet — if a running sequence needs a block and the pool's fully dry, it still just raises instead of evicting someone else gracefully. Documented, not solved. Also want to eventually swap the gather-then-SDPA attention for a real fused kernel, purely to measure the actual delta with my own hands now that the memory/scheduling logic underneath it is something I've verified layer by layer and actually trust.

More updates soon.