From 183 ms to 19 ms per decode step

Where it started

mini-engine is an LLM inference engine I wrote from scratch in PyTorch: a paged KV cache, continuous batching, prefix caching and speculative decoding, with greedy output token-identical to Hugging Face generate. Version 0.1 was correct but slow. On an NVIDIA L4 GPU, running Qwen3-0.6B in bf16, it produced 294 output tokens per second on a batch of 64 requests. Hugging Face generate produced 1,292 and vLLM 3,256.

Attention in v0.1 was simple on purpose. Each layer gathered every sequence’s context out of the paged cache into one padded tensor and made one masked call to PyTorch’s scaled_dot_product_attention (SDPA). No custom kernels, and easy to check against a dense reference. It was also where the time went.

How I measured

benchmarks/profile_engine.py runs the throughput workload (64 requests with 256-token prompts) until every request is decoding. Then it times decode steps phase by phase (schedule, prepare, forward, sample), with a CUDA sync around each, and records a few of them with torch.profiler. The profiler’s table says which GPU operations take the time; the gap between the step time and the GPU’s busy time says how long the GPU waits on the CPU.

The rule was to fix only the largest cost in that table, then measure again. Every run was on an L4 in Colab. At the start, one decode step of 64 sequences took 183 ms: 349 tokens per second.

Four changes, one profile at a time

1. Grouped attention: 183 to 60 ms

Qwen3 uses grouped-query attention: several query heads share one key/value head. Given a boolean mask and that grouping, SDPA quietly fell back to its plain “math” path and copied K and V once for every query head. I replaced the call with two batched matrix multiplications that fold each KV head’s group of query heads into the query axis. K and V are read once per group and never copied per head.

2. Head-major KV cache: 60 to 46 ms

Now half the GPU time went to one copy: moving the gathered context into the transposed layout the matrix multiplications want. I changed the cache layout to put the KV head first, [kv_heads, blocks, block_size, head_dim], so one index over the cache slots returns the context already in the right order.

3. A Triton paged-decode kernel: 46 to 38 ms

Now the gather itself was the cost, 44% of GPU time, and it already ran at the memory bandwidth limit. A faster copy was not possible, so the fix was to stop copying. I wrote a Triton kernel that walks each sequence’s block table and reads K and V straight from the cache. One program handles one (sequence, KV head) pair and all the query heads in its group. It loops over the context in tiles with an online softmax in fp32, the FlashAttention trick: keep a running maximum and sum, and rescale them as each tile arrives. Each step reads the cache once and writes nothing extra.

4. CUDA graphs per batch bucket: 38 to 19 ms

At 38 ms the GPU was busy for only about 18 ms of each step. The rest was the CPU launching about 1,600 small kernels, one at a time. A CUDA graph records a whole sequence of launches once and replays it with a single call.

A graph records raw memory addresses, so its inputs have to live in fixed buffers. Each step copies its token ids, positions, cache slots, block tables and lengths into those buffers, then replays. Batch sizes change from step to step, so the engine captures one graph per batch-size bucket at start-up (1, 2, 4, 8, then every 16 up to 64) and pads each batch up to its bucket. Padding rows write to one spare cache block that is never given to a request, attend only to that block, and are dropped.

After this a step took 19 ms, about the time of the GPU’s own work: 3,364 tokens per second on the same profile, 9.6 times the start.

Then speculative decoding got slower

mini-engine also does n-gram speculative decoding. Each step guesses up to k next tokens by finding an earlier place where the last few tokens appeared (no draft model), and one forward pass checks all the guesses. A checking step, a verify step, carries up to k + 1 tokens per sequence. It fit neither the one-token decode kernel nor the decode graphs, so it still ran without a graph, at about 45 ms. With plain decoding now replayed from a graph, speculation made decoding slower instead of faster: 106 tokens per second against 115.

Two more pieces fixed it:

• A second Triton kernel for a few query tokens per sequence. It puts [query tokens × group heads] on the rows of one tl.dot, so the GPU’s tensor cores do the work, with a causal mask among the new tokens.
• A second set of CUDA graphs with k + 1 tokens per sequence. Sequences with fewer guesses are padded, so one graph fits any mix of guess lengths.

On the copy task, speculation now gives 4.3 times plain decoding with one request (114 to 484 tokens per second) and 3.8 times with eight (809 to 3,075).

The bug: a tensor freed after capture

The first CUDA-graph version crashed on the GPU with a device-side assert: an index out of range inside the graph. To separate my logic from CUDA, I ran the whole graph path on the CPU, with capture faked and the Triton kernels replaced by a plain dense reference. On the CPU every token matched normal decoding, so the bookkeeping was right and the problem was specific to graphs.

The cause: a graph remembers addresses, not variables. One input, logits_index, was held in a local variable. After capture its memory went back to PyTorch’s caching allocator, was given to another tensor, and every replay read whatever was there as indices. The fix was to keep every tensor a graph reads on the object, so it lives as long as the graph does. From cuda_graphs.py:

# Graphs record raw addresses: every tensor a graph reads must stay alive on self,
# or the caching allocator hands its memory to someone else after capture.
self.input_ids = torch.zeros(rows, dtype=torch.int64, device=dev)
self.positions = torch.zeros(rows, dtype=torch.int64, device=dev)

The CPU run became a permanent test, tests/test_step_graphs.py, which runs in CI on every change.

Results

The same token-id workloads went through Hugging Face generate, mini-engine and vLLM 0.32.0, each in its own process, greedy. Numbers are output tokens per second; higher is better. The raw results are in benchmarks/results/.

Qwen3-0.6B on one NVIDIA L4, bf16

• Batched decoding: 1,982 tokens per second, 6.7 times v0.1, 1.5 times Hugging Face generate, and 61% of vLLM.
• Decode step: 183 ms to 19 ms for 64 sequences.
• Speculation: with eight requests, 3,075 tokens per second. With one request, 484, ahead of vLLM on this task with or without its own speculation (171 and 135).
• Correctness held: greedy output is token-identical to Hugging Face on 8 of 8 chat prompts in fp32, with the kernels and graphs on, with and without speculation. In bf16, 4 to 5 of 8 match, because different kernels can flip the choice between two nearly tied tokens.

Limits

• One model, one GPU. Every number here is Qwen3-0.6B on one L4. The kernels were tested against the PyTorch path on several shapes, so their logic is not tied to this model, but the speed numbers are.
• Qwen3 only. The engine runs the Qwen3 architecture and nothing else yet.
• vLLM is still faster. mini-engine reaches 61% of vLLM on batched decoding and about half of it on the warm shared-prefix workload, and vLLM leads the eight-request copy task.
• Prefill is still the slow path. Prompts still go through the padded PyTorch attention, so the cold shared-prefix run is slow (370 tokens per second against vLLM’s 1,432). A paged prefill kernel is next on the roadmap.