ember

A mini-vLLM — an LLM inference server built from scratch to understand the two ideas that make LLM serving fast.

Modern serving engines (vLLM, TGI, SGLang) get their throughput from two ideas that have nothing to do with the model architecture and everything to do with systems:

1. A paged KV cache — the attention cache is split into fixed-size blocks and each sequence keeps a block table mapping its logical positions to physical blocks. KV memory needn't be contiguous, so it packs tightly and fragments by at most one block per sequence instead of one max-length reservation per sequence.
2. Continuous batching — the scheduler decides batch membership every decode step, not once per request. A finished sequence frees its blocks immediately; a queued request is admitted into the very next step.

ember implements both, from scratch, in ~1.5k readable lines of Python — including a from-scratch GPT-2 that loads the real weights, so it emits real text. No vllm, no transformers modelling code; the only thing borrowed is the weight file.

from ember import LLM, SamplingParams

llm = LLM("gpt2")                                  # weights auto-downloaded once
outs = llm.generate(["The capital of France is"],  # served with continuous batching
                    SamplingParams(temperature=0.0, max_tokens=12))
print(outs[0].text)
#  the capital of the French Republic, and the capital of the

---

The core idea, in one picture

A sequence's logical token positions are mapped through a block table to non-contiguous physical blocks in one shared KV pool. Position p lives at block_table[p // block_size], offset p % block_size.

   sequence A (block_size = 4)                shared KV pool (physical blocks)
   logical:  [ 0 1 2 3 | 4 5 6 7 | 8 ]        ┌────┬────┬────┬────┬────┬────┐
   block_table = [ 2, 0, 5 ] ────────┐        │ B0 │ B1 │ B2 │ B3 │ B4 │ B5 │
                                     └──────►  │ A₁ │ .. │ A₀ │ .. │ .. │ A₂ │
   sequence B                                 └────┴────┴────┴────┴────┴────┘
   block_table = [ 1, 3 ] ───────────────────────────►  B1, B3

When a sequence finishes, its blocks return to the free list and any waiting sequence reuses them on the next step. That is the whole reason a paged engine fits far more concurrent sequences into the same memory than static per-sequence buffers.

Request lifecycle

flowchart LR
  R([request]) --> W[waiting queue]
  W -->|admit if blocks + budget allow| PF[prefill: whole prompt in one step]
  PF --> RUN[running]
  RUN -->|1 token / step| DEC[decode]
  DEC --> RUN
  DEC -->|pool full| PRE[preempt youngest -> recompute]
  PRE --> W
  RUN -->|EOS / stop / max_tokens| DONE([finish -> free blocks])

Loading

Every tick, the scheduler runs one prefill step or one decode step across the whole batch; finished sequences are retired and their KV reclaimed before the next tick.

---

What's implemented

Area	Detail
Paged KV cache	Block allocator + per-sequence block tables; O(1) allocate/free; fragmentation bounded to one block/sequence.
Continuous batching	Iteration-level scheduler with prefill-priority admission, per-step token/sequence budgets.
Preemption	Recompute-based: when the pool is full the youngest sequences are evicted (KV freed, tokens kept) and re-prefilled later.
Model	GPT-2 (124M–1.5B) from scratch in PyTorch — learned pos-emb, pre-LN blocks, tied LM head — loads official HF weights.
Attention	Paged scatter/gather into a shared pool; supports grouped-query attention (n_kv_head < n_head).
Sampling	Greedy, temperature, top-k, nucleus (top-p); per-request seeded RNG.
Serving	OpenAI-compatible POST /v1/completions, streaming (SSE) and non-streaming, on an async engine.
Offline	LLM.generate([...]) batch API and an ember CLI.

Correctness

Greedy decoding is a deterministic function of (weights, prompt), so paging and batching — which are pure implementation details — must never change the output. The test suite pins exactly that:

• block-size invariance — identical tokens for block_size ∈ {1, 2, 3, 7, 16, 64};
• batch invariance — a prompt generated alone == generated alongside others;
• preemption invariance — a squeezed KV pool that forces eviction/recompute still yields identical tokens.

An off-by-one in the slot math or the scheduler bookkeeping breaks one of these immediately.

pytest -q        # 13 passed

---

Quickstart

python -m venv .venv && source .venv/bin/activate
pip install -e .                          # installs deps AND the `ember` / `ember-serve` commands

# offline batch (downloads gpt2 weights on first run)
python examples/offline_batch.py

# serve an OpenAI-compatible endpoint
ember-serve --model gpt2                  # http://127.0.0.1:8000

pip install -e . is what creates the ember and ember-serve console scripts (from [project.scripts]). Installing only requirements.txt gets the libraries but not those commands. With the venv active they're on your PATH; otherwise call them by path, e.g. .venv/bin/ember-serve.

curl -N localhost:8000/v1/completions -H 'Content-Type: application/json' \
  -d '{"prompt":"Once upon a time","max_tokens":40,"temperature":0.8,"stream":true}'

Runs on CUDA, Apple-Silicon MPS, or CPU (auto-detected).

Serving your own model

Besides the GPT-2 presets, ember serves any model in its GPT-2 weight layout — a local export dir or a Hugging Face repo with config.json + model.safetensors. For example, scribe is a ~30M GPT trained from scratch on TinyStories; ember runs it directly:

from ember import LLM, SamplingParams
llm = LLM("mbsdeepak/scribe")                       # or a local ./export dir
print(llm.generate(["Once upon a time"], SamplingParams(max_tokens=60))[0].text)

ember-serve --model mbsdeepak/scribe                # same, as a streaming API

Layout

ember/
  cache/block_manager.py   # the paged cache: allocator + block tables
  core/
    scheduler.py           # continuous batching: admit / decode / preempt
    model_runner.py        # scheduler decisions  ->  batched tensors; owns the KV pool
    engine.py              # add_request + step() loop
    sampler.py             # greedy / temp / top-k / top-p
    sequence.py            # the unit of work
  model/
    gpt2.py                # GPT-2 from scratch, with paged attention
    weights.py             # load official GPT-2 weights (HF Conv1D -> Linear)
  entrypoints/
    llm.py                 # offline LLM().generate([...])
    async_engine.py        # async loop that streams tokens
    api_server.py          # OpenAI-compatible FastAPI app

Deliberately out of scope (for now)

This is a teaching-grade engine: correctness and readability first. The attention scatter/gather is done with indexing + scaled_dot_product_attention per sequence rather than a fused kernel, so it is not fast — the point is that the data structures and scheduling are exactly what a production engine uses. Natural next steps, roughly in order:

• a fused paged-attention kernel (Triton) so throughput is real;
• prefix caching — share KV blocks across requests with a common prefix (copy-on-write block tables);
• chunked prefill and a proper prefill/decode mix per step;
• continuous-batching throughput benchmarks vs. static batching;
• paged KV quantization (fp8/int8).

Why this exists

I already built the layers around the model — a gateway (bulkhead), context engineering (loom), evals (gauntlet), tracing (sonar) and the agent runtime (cogs). ember is the missing middle: the thing that actually serves the tokens. Same bias as the rest — build it to understand it, minimal dependencies, code you can read end-to-end.

License

MIT