Vraj Patel BATCHER ← Writing

Aug 2026 · updated Oct 2026 · 6 min read

Building a paged inference engine from scratch, and measuring it honestly

A mini LLM inference server, built without looking at vLLM's source: a paged KV cache, a continuous batching scheduler, and a FlashAttention-2 kernel in Triton that reads K and V through a block table. It now serves GPT-2 small, and I benchmarked it against Hugging Face and vLLM 0.21.0. vLLM wins by a lot, and most of this post is about why.

Updated Oct 2026. The first version of this post benchmarked a small model I trained myself, where the paged engine lost to a naive loop at every concurrency level. That sweep ran before I found two kernel bugs, so I have withdrawn it. Everything below uses GPT-2 small after the fixes. The old numbers are at the end, clearly marked.

Key numbers · RTX 5080 Laptop

82 → 1,030tok/s, 1 → 32 concurrent requests, GPT-2 small
3.5xHugging Face with a KV cache, at 32 requests
~23xvLLM 0.21.0's lead over mine at 32 requests
~0.005Prefill log-prob gap vs Hugging Face, after the fixes

01What this is

I built this to understand how LLM serving works under the hood, not just how to train a model. The design follows the PagedAttention paper. Working out the mechanisms myself, without vLLM's source, is the only way I'd know for certain what my version does and doesn't handle.

The transformer underneath is my own decoder-only build, adapted from an earlier project (cpp-gpu-inference/en-de-transformer). A serving layer for the engine (FastAPI, Docker, CI/CD to AWS ECR) lives in paged-engine-serving.

Hardware
RTX 5080 Laptop, 16 GB, WSL2
Stack
PyTorch 2.11.0 (CUDA 12.8 build), Triton 3.6.0
Model
GPT-2 small (124M), fp32 weights, fp16 KV pool
Workload
Every request generates exactly 64 new tokens; all systems get the same pre-tokenized prompts

02The engine

  • Paged KV cache. Keys and values live in fixed-size blocks instead of one contiguous buffer per sequence, so many requests can share GPU memory without fragmentation.
  • Continuous batching scheduler. Requests are admitted roughly in order, but the scheduler can look a few requests ahead and admit whichever ones fit. If too many requests get skipped in a row, it falls back to blocking mode until memory frees up. New requests join a batch that is already running instead of waiting for it to finish.
  • FlashAttention-2 kernel (Triton). Reads K and V straight from the paged block cache through a block table.
  • INT8 KV cache option. Stores keys and values in int8 (symmetric, one scale per block per head) instead of fp16, which roughly halves KV memory. I wrote this before the kernel fixes below and have not re-validated it since, so treat it as unverified. All GPT-2 results use the fp16 KV path.
Block allocator, scheduler phases and the kernel rewrite details

The first piece is BlockManager (core/block_manager.py): a fixed-size pool of KV cache blocks with allocate, release_blocks, num_free_blocks, and can_allocate. I tested it with interleaved-sequence correctness tests: allocate and free sequences in overlapping, out-of-order patterns and check that the free-block count and block ownership stay consistent throughout, not just at the end.

Scheduler.schedule() runs four phases every call: _free_finished releases blocks for sequences that finished last step, _allocate_decode grows the block allocation for sequences still generating, _admit_waiting decides which queued requests get to start, and _build_output assembles the batch for this step. Whether a sequence is doing prefill or decode isn't stored as a flag anywhere. It's derived on the fly by comparing token counts.

I had a standalone FA-2 kernel in Triton from an earlier project. Paging meant it could no longer assume each sequence's K/V lived in one contiguous buffer, so I rewrote the block-reading logic: new kernel parameters block_table, num_blocks_per_seq, q_len_per_seq, kv_len_per_seq, and block_size, with the inner loop walking a sequence's block table to find where its K/V actually live in the pool.

The model runner batches whatever the scheduler returns into one flat packed tensor for the step, prefill and decode sequences together, rather than a padded rectangular batch. One forward pass, then each sequence's last-token logits get pulled out with vectorized indexing.

INT8 KV cache design

Granularity is per-block-per-head, symmetric quantization with no zero-point. The cache is int8 from the very first token, never staged in fp16 and quantized later. On every new token write, the block's existing entries get dequantized with the stored scale, the new token's K/V gets appended, a new max-abs scale is recomputed over the whole set, and everything gets requantized and written back. On the read side, the kernel dequantizes immediately before each of the two tl.dot calls. The identical-output check I ran on it (same 4 prompts, FP16 build vs INT8 build) predates the kernel fixes, which is why I now call this path unverified.

03Correctness, and the bugs a real model found

The first version of the kernel was tested against PyTorch scaled_dot_product_attention on three sequences of different lengths stored in scrambled blocks. That test only covered prefill, where the query length equals the key length.

Running a real model exposed two bugs it missed:

  • Stride order in the kernel launcher. After I changed the cache layout, the launcher still passed the head stride and the slot stride in the old order. Head 0 read the right memory and every other head read the wrong memory. Prefill logits were badly wrong, even for a one token prompt.
  • Causal mask during decode. The mask compared key positions against the query's index inside the padded query tensor. For a single decode token that index is 0, so the token could only attend to the first key. The mask now uses the query's real position, kv_len - q_len + offset.

To catch this kind of thing, the repo now has a GPT-2 small parity harness that loads the Hugging Face weights into the paged model and compares against HF, down to per-layer hidden states, which is how the stride bug was found.

After the fixes, prefill matches HF to within about 0.005 in log-probability at prompt lengths from 1 to 21 tokens (fp32 weights), and the relative error of each layer's hidden state stays around 1e-4.

Three earlier bugs, from the first build details

EOS off-by-one. Sequences stop and free their blocks as soon as they emit an EOS token, instead of always running to max_len. Implementing that surfaced an off-by-one: the batch index was only incremented inside the non-EOS branch, so a sequence hitting EOS mid-step would leave the index pointing at the wrong row for the rest of that step. Fixed by incrementing the index unconditionally, outside the branch.

num_blocks was set to 10000 (block_size=64, 6 decoder layers), which pre-allocates a large fixed KV cache pool, about 3.75 GB, at model construction time regardless of how many sequences are actually active. That's a real fixed cost of block-based paging worth knowing about, and it skewed early memory comparisons until I accounted for it.

Both models were staying resident on GPU at the same time during testing, naive and paged loaded together, which confounded early memory results. Fixed by explicitly freeing the naive model between test phases.

04Throughput

Numbers are the median of 3 runs for the HF loops and my engine, and the median of 5 for vLLM. The two HF baselines handle requests one at a time, so part of my engine's lead over them is just batching. vLLM batches too, which makes it the fair comparison.

Against Hugging Face: slower at 1 request, 3.5x faster at 32

    Batching is doing real work. My engine matches HF with a KV cache at 4 concurrent requests and reaches 3.5x its throughput at 32. It is slower at a single request.

    Against vLLM: about 12x behind at 1 request, about 23x at 32

      Same GPU, same prompts, same token counts, same scale. On a shared axis, my engine's best result is a sliver next to vLLM's.
      Throughput in tokens per second · GPT-2 small · 64 new tokens per request · RTX 5080 Laptop
      Concurrent requestsHF, no cacheHF, KV cachePaged (this repo)vLLM 0.21.0
      125329782949
      42552942973,663
      82522974936,667
      1624829575313,058
      322512951,03023,247

      05Where mine loses, and why

      • vLLM is about 12x faster than my engine at 1 request and about 23x faster at 32.
      • Dividing wall time by the 64 steps gives roughly 12 ms per step for my engine at 1 request and 31 ms at 32. vLLM takes about 1 ms at 1 request and 1.4 ms at 32. So my engine has a large fixed cost per step plus a cost of about 0.6 ms for each extra sequence, and vLLM has almost neither.
      • Scheduling, batch assembly and the final .tolist() account for only 1 to 2 percent of step time, so the cost is inside the model call. My guess is the per-layer host work in the attention wrapper (rebuilding the block table, padding and unpadding the query in Python, a few GPU syncs per layer). I have not profiled it, so that is a hypothesis.

      Things that make the vLLM comparison imperfect: vLLM runs fp16 weights while mine are fp32, and it uses CUDA graphs, fused kernels and asynchronous scheduling. It also ran in a separate environment (vLLM 0.21.0 built for CUDA 12.9, FlashAttention 2 backend, FlashInfer sampler disabled, pin_memory=False because of WSL). Same GPU, same prompts, same token counts.

      Approximate time per decode step (wall time / 64 steps)
      System1 request32 requests
      Paged (this repo)~12 ms~31 ms
      vLLM 0.21.0~1 ms~1.4 ms

      06The memory ceiling: inconclusive

      I also tried to measure a memory ceiling, the max concurrent sequences before OOM, but under WSL2 the reported GPU memory was physically impossible for this card, so that measurement is inconclusive and not included.

      Why I'm not reporting a max-concurrency number as fact

      The peak memory figure reported at the last successful point (n=64) was approximately 32 GB, on a 16 GB GPU. That isn't physically possible, so the number is wrong somewhere, not a real measurement of paged's memory ceiling.

      Traced it to WSL2: CUDA allocations there can spill into a shared virtual memory pool backed by system RAM instead of failing cleanly at the real VRAM limit, which means torch.cuda.max_memory_allocated() isn't trustworthy for physical GPU capacity in this environment.

      I could chase this down with an nvidia-smi-based measurement instead. I decided to report it honestly as inconclusive rather than publish a number I already know is wrong.

      07Open items

      Stated plainly, because a list of finished things with no unfinished things next to it isn't credible.

      • Cut the per-step host overhead. Compute the block table and batch metadata once per step instead of once per layer, vectorize the query padding, and try CUDA graphs. This is the main thing holding throughput back.
      • Re-test the INT8 KV path against the same GPT-2 harness.
      • Re-validate the live demo and core/model_runner.py. Both predate the kernel fixes; the GPT-2 harness uses its own runner in gpt2/paged_run.py.
      • Re-measure the memory ceiling with an nvidia-smi-based check, ideally outside WSL2.
      • Not started: CPU offloading or recomputation when memory runs out (the scheduler waits), speculative decoding, beam search. Greedy decoding only.

      08Reproduce it

      Everything is in github.com/VrajPatel105/paged-inference-engine. Run everything from the repo root; vLLM needs its own environment because it pins a different PyTorch build.

      File map
      SectionFile
      Block allocatorcore/block_manager.py
      Schedulercore/scheduler.py, core/sequence.py
      Paged FlashAttention-2 kernelkernels/flash_attention.py
      GPT-2 paged runnergpt2/paged_run.py
      Parity vs Hugging Facetests/test_gpt2_parity.py, tests/test_gpt2_engine_parity.py, gpt2/diagnose.py
      Benchmarksbench/benchmark_gpt2.py, bench/benchmark_vllm.py
      Commands
      StepCommand
      Save HF reference outputs (once)python -m gpt2.reference
      Prefill paritypytest -s tests/test_gpt2_parity.py
      Engine paritypytest -s tests/test_gpt2_engine_parity.py
      HF vs pagedpython -m bench.benchmark_gpt2
      vLLM (separate env)python bench/benchmark_vllm.py

      09Appendix

      Earlier sweep (withdrawn: ran before the kernel fixes) Aug 2026

      The first benchmark used a small model I trained myself, against a naive baseline: no KV cache reuse, no batching, one request processed fully at a time. It showed the paged engine below the naive loop at every concurrency level. That sweep ran before the two kernel fixes in section 03, so these numbers come from a broken kernel and do not describe the engine. They are kept here only so the change is visible.

      Withdrawn · aggregate throughput, tokens/sec · max_len 50 tokens · RTX 5080 Laptop
      ConcurrencyNaive tok/sPaged tok/s
      149195
      4542152
      8550166
      16542173
      32543176

      One methodology lesson from that run still holds: warmup runs before timing starts mattered more than I expected. Without them, one-time CUDA and Triton kernel compilation gets counted as request latency. Paged TTFT measured ~1100ms before adding warmup, and ~69ms after, a compile-cost artifact, not anything close to the real per-request cost.