I built an Inference Engine that beats vLLM

The field of serving LLMs, optimising it for inference and other related stuff is Inference Engineering. And this field has been doing numbers recently because inference can be costly. Granted that training a frontier LLM can cost hundreds of millions, it is still a one-time cost while inference depends entirely on the scale. There is no fixed upper bound for how costly it can get.

There are many vectors on which Inference can be optimised

• cost
• scale
• latency
• accuracy

Before 2025 I mostly worked on side projects related to PyTorch, SLLM training/fine tuning etc but last yeear I had a little Agentic AI time skip arc and lost touch of the other side of AI.

So:

And I started building an inference engine from scratch called Helios (after the greek sun god).

GitHub: github.com/sidmanale643/helios

The goal was simple:

Keep reading the book, keep implementing to build a lightweight but powerful inference engine in python with extremely readable code. The current scope is a small model that fits on a single GPU

Model

I used a Qwen-3 4B model because it would fit in a free Tesla T4 and I was well aware of the internals. I didnt want to spend time on learning new architectures because I was too eager about inference :D

Weights are obviously downloaded from HuggingFace.

Specs:

dtype

FP16 for older generations like Turing and BF16 whenever supported

API

Helios exposes a OpenAI Style API for inference

KV Caching

is an inference technique where instead of recalculating the K and V vectors for each new token, we simply cache all the seen K and V vectors across all layers and heads during prefill and reuse them at decode time

For in depth guide, checkout: KV Cache From First Principles

Prefill vs Decode

Prefill is the stage in inference before the first token is even predicted. This step is essentially where the Q, K and V matrices are projected, KVs are cached etc.

Decode is the stage at which and after the tokens are predicted auto regressively by reading in the KV caches, writing the slot for the new token and so on.

Helios runs both in a single tick. For every request that I decoding, a token is taken and the remaining budget is allocated for chunked prefill and the bydget is controlled by HELIOS_PREFILL_CHUNK_SIZE.

If the tick contains only decode tokens all of them are batched and run together by reusing existing page tables.

Earlier Helios had different workers for prefill and decode but after they were bundled into a single forward the peak aggregate throughput increased by upto 23%.

torch.sdpa and Flash Attention

torch.nn.functional.scaled_dot_product_attention (SDPA) is PyTorch's attention API. Instead of manually doing softmax(QKᵀ / √d) × V, you give PyTorch the Q, K and V tensors and it picks an appropriate and efficient implementation for the current hardware like Flash Attention on supported GPUs

This helps us optimise one the most expensive operation in an LLM without reinventing the wheel.

In Helios, we project the Q,K and V matrices then apply RoPE to Q and K, then send them through SDPA with the right causal mask. For Qwen3 Helios enables GQA so multiple query heads can share the smaller set of KV heads without creating unnecessary copies.

Scheduler

The scheduler is the queue and worker that decides which requests may use the GPU and when.

Helios keeps a bounded FIFO wait list and a background worker thread. HTTP requests become jobs. Each batch is one Engine tick.

Admission is FIFO, capped by HELIOS_MAX_BATCH_SIZE. A request that cannot fit the KV budget is rejected. A request that fits the budget but not current free memory stays at the head of the queue. Finished or cancelled jobs leave, and later jobs can start.

Continuous batching

Requests do not wait for each other to finish. Requests arrive, stay in process or get finished and leave as soon as finished, independently.

Each tick first takes one decode token for every request that is already generating. Remaining token budget goes to unfinished prompts.

If a prompt has waited at least 100 milliseconds, the oldest waiting prompt is taken instead so long prefills still move. Selected tokens run in one packed model forward, with independent positions and page tables. Decode-only ticks reuse decode page-table buffers. A request keeps its KV pages across ticks until it hits EOS or max_tokens, then its result is returned and its reservation is released.

Paged attention

Paged attention stores the KV cache in fixed size pages from a shared pool, instead of one large contiguous cache per request.

Helios uses 256-token pages in one GPU pool. Each request has a page table of the pages it owns. New tokens allocate pages from the free page list. Prefill and decode write K and V into those pages, then use PyTorch varlen_attn as the attention mechanism using the page table.

When the last request (or cache) that holds a page drops it, the page goes back to the pool. Admission still reserves the request’s max length so decode has space.

Prefix caching

KV Caching is per request and the caches are evicted immediately after. Prefix caching helps us apply KV caching across requests for use cases like for eg agents where the same blocks of tokens can be sent repeatedly across tens of requests.

Helios hashes complete 256-token prompt blocks in order along with the previous blocks hash. Once a request arrives it looks up the longest matching chain, restores those pages into the request cache, and prefills only the unmatched later blocks.

Partial last blocks are not stored. When a request finishes, completed prompt pages can stay in the prefix cache with a sliding TTL. Requests that share a cached prefix share the same physical pages. Eviction frees unused blocks first, so some parent blocks might still stay. Shorter prefixes than 256 tokens are not cached.

At all steps of development and making architectural changes I wanted to compare how the performance was affected

I created a small 30-request benchmark covering three workload types:

• prefill heavy
• decode heavy
• balanced

This was an extremely extremely important step because this acted as my compass during development . It helped me determine if progress improved or regressed, which parts of the engine are bottlenecks like torch.compile caused TTFT to increase by 100x, scheduler was starving the prefill and many others.

Metrics measured

Latency

• TTFT (time_to_first_token_seconds) — p50 / p95
• E2E (end_to_end_seconds) — p50 / p95

Throughput

• Decode tok/s (decode_tokens_per_second) — p50 / p95
• Prefill tok/s (prefill_tokens_per_second) — p50 / p95
• Batch output tok/s (output_tokens_per_second)
• Requests/s (requests_per_second)

Batch / concurrency

• Peak and average client concurrency
• Drain time and drain fraction

Work done

• Prompt tokens, output tokens
• Completion ratio and max-token hit rate
• Cache hit rate / restored tokens

Helios vs vLLM

Since Helios does not support torch.compile and CUDA graphs, I wanted the benchmarking setup to be comparable so I ran vLLM in --eager mode. Eager model disables most of the CUDA/triton graphs and torch.compile but still keeps other features like chunked prefill, continuous batching, paged attention etc etc.

Both servers used BF16 on an NVIDIA L4 GPU.

vLLM delivered much lower request latency. Its median TTFT was 590 ms, compared with 740 ms for Helios, and its p95 TTFT was 1.79 s versus 4.02 s.

The gap is most visible for long prompts, even with the graphs disabled vLLM is much more efficient at scheduling requests and optimising prefil/decode cycles.

vLLM's p50 decode speed was also higher than Helios'.

But even so, Helios had a higher throughput even without any CUDA graphs and compilations, which was extremely encouraging.

Eager mode does not turn vLLM into plain PyTorch. It still uses its optimized attention backend, paged KV-cache management, continuous batching, and chunked-prefill scheduler.

• Mixture-of-Experts support: add support for MoE models with efficient expert routing
• torch.compile: test where compilation actually improves latency and throughput
• CUDA Graphs: reduce CPU launch overhead during decode by capturing repeatable GPU workloads
• Newer architectures: move beyond Qwen3-4B and support larger models with architectures such as MLA and DSA
• KV cache offloading: move inactive KV blocks between GPU and CPU memory to support longer contexts and more concurrent requests
• Parallelism: add tensor and pipeline parallelism to move Helios beyond single-GPU serving

Also I am going to start a blog series explaining all the architectural choices and implementations.