Mini LLM Inference Server

A small Large Language Model inference server built from first principles — not to train a model, but to understand what happens after a model receives a prompt.

prompt → tokenization → embeddings → transformer inference → logits → sampling → next token → repeat

The pipeline is built incrementally as small, self-contained components, each solving one part of inference and then assembled into the final server.

Why build this?

High-level libraries make inference look like a single line:

output = model.generate(tokens)

That line hides real systems problems this project makes explicit: how logits become probabilities, how tokens are sampled, how temperature shapes generation, how the model avoids recomputing attention (KV cache), how that cache memory is allocated, how requests are batched, and how a server trades off throughput and latency.

The goal is something small enough to read end-to-end, but complete enough to show the core mechanics behind production inference systems.

Project layout

mini-llm-inference-server/
├── mini_llm_server/        # the package
│   ├── __init__.py         # public API re-exports
│   ├── sampling.py         # sampling primitives (softmax, temperature, top-k/p, …)
│   ├── tokenizer.py        # tokenization (build_vocab, encode/decode)
│   ├── transformer.py      # tiny single-layer transformer (prefill + decode)
│   ├── kv_cache.py         # contiguous per-sequence KV cache
│   ├── paged_kv_cache.py   # paged KV blocks shared across sequences
│   ├── sequences.py        # request records and per-sequence state
│   ├── batching.py         # static batching across many sequences
│   ├── scheduler.py        # continuous batching, priority queue, preemption
│   ├── server.py           # serving API: submit, stream, collect
│   └── benchmark.py        # TTFT, ITL, throughput, percentiles
├── docs/                   # per-component deep dives, grouped by module
│   ├── sampling/           # one page per sampling primitive (+ softmax.svg)
│   ├── tokenizer/          # one page per tokenization step
│   ├── transformer/        # embeddings, projections, attention, prefill, decode
│   ├── kv_cache/           # allocating and appending to the contiguous cache
│   ├── paged_kv_cache/     # block allocator, paged attention, occupancy
│   ├── sequences/          # request lifecycle, one token at a time
│   ├── batching/           # synchronized batched decode steps
│   ├── scheduler/          # admission, preemption, continuous batching
│   ├── server/             # streaming chunks and completion responses
│   └── benchmark/          # latency and throughput metrics
├── examples/
│   ├── plot_softmax.py     # regenerates docs/sampling/softmax.svg (stdlib only)
│   └── generate_demo.py    # end-to-end: vocabulary → batched decode → text
├── requirements.txt
└── README.md

As later stages land, new modules join the package (scheduler.py, server.py), each paired with docs under a matching docs/<module>/ directory.

Quickstart

python -m pip install -r requirements.txt

A single primitive:

import numpy as np
from mini_llm_server import stable_softmax, apply_temperature

logits = np.array([2.0, 1.0, 0.1])
probs = stable_softmax(apply_temperature(logits, temperature=0.8))
print(probs)  # non-negative, sums to 1

Or the whole pipeline — tokenize, prefill, batched decode, detokenize (python -m examples.generate_demo):

import numpy as np
from mini_llm_server import (
    build_vocab, encode_prompt, decode_tokens, make_request, static_batch_generate,
)

vocab = build_vocab(["hello world"], ["<bos>", "<eos>", "<unk>"])
V, D = len(vocab["id_to_token"]), 8
rng = np.random.default_rng(0)

# A toy model with random weights — real ones would come from a checkpoint.
params = {
    "E": (rng.normal(size=(V, D)) * 0.1).astype(np.float32),
    "Wq": np.eye(D, dtype=np.float32), "Wk": np.eye(D, dtype=np.float32),
    "Wv": np.eye(D, dtype=np.float32), "Wo": np.eye(D, dtype=np.float32),
    "W_out": rng.normal(size=(D, V)).astype(np.float32),
    "max_seq_len": 32,
}

requests = [
    make_request("r0", encode_prompt("hello", vocab), 5, {}),
    make_request("r1", encode_prompt("world", vocab), 5, {}),
]
config = {"temperature": 0.9, "top_p": 0.95, "rng": np.random.default_rng(0)}

for result in static_batch_generate(params, requests, config, max_new_tokens=5):
    print(result["request_id"], result["output_ids"],
          repr(decode_tokens(result["output_ids"], vocab)))
# r0 [7, 0, 9, 6, 6] 'lrhh'
# r1 [3, 0, 10, 8, 10] ' wow'

The text is nonsense because the weights are random — the point is that every stage of a real serving path is there, from vocabulary to batched decode.

Run from the repository root so mini_llm_server is importable.

Roadmap

Progress is tracked component by component. Completed steps link to their deep dive.

Part 1 — Sampling primitives

• Numerically stable softmax
• Temperature scaling
• Top-k filtering
• Top-p filtering
• Sampling from probabilities
• Greedy token selection

Part 2 — Tokenization

• Build a vocabulary
• Encode prompts
• Decode tokens

Part 3 — Tiny Transformer with KV cache

• Embed tokens
• Linear projection
• Initialize KV cache
• Append key and value states
• Causal attention
• Model prefill
• Decode one token at a time

Part 4 — Paged KV cache

• Blocks needed
• Initialize the block allocator
• Allocate a block
• Free a block
• Append to a paged cache
• Gather KV from blocks
• Paged attention step
• Free a sequence's blocks
• KV blocks in use

Part 5 — Sequences and static batching

• Make a request
• Initialize sequence state
• Sequence decode step
• Sequence stop condition
• Generate a single sequence
• Build batch step input
• Batched decode step
• Static batch generation

Part 6 — Continuous batching and scheduling

• Has free capacity
• Continuous batch step
• Run continuous batching
• Priority queue push
• Priority queue pop
• Select admissions
• Preempt a sequence
• Schedule step

Part 7 — Serving API

• Format a stream chunk
• Submit a request
• Drive the server loop
• Collect request output
• Build a completion response

Part 8 — Benchmarking

• Time to first token
• Inter-token latency
• Aggregate throughput
• Latency percentiles
• Run the benchmark harness

Core idea

An LLM inference server is not one large function. It is a pipeline of small components with strict contracts:

correct shapes + stable numerical operations + controlled memory use + predictable scheduling = reliable inference

This project builds those contracts one at a time.

Attribution

An independent implementation created while working through the Deep-ML mini LLM inference server project. The code, explanations, and documentation here are written for learning and portfolio purposes. Platform-owned hidden tests and restricted instructional content are not included.