I built an inference engine from scratch. Here is what I learnt.

LLMs feel almost suspiciously simple from the outside.

You send some text to an API. A few seconds later, words start appearing on the screen. It is easy to think the interesting part is the model. Load it, call generate(), send the output back.

That is true until you try to serve it across more than one machine.

Then a slightly annoying question appears: which machine should generate the next token?

The answer is not “whichever one is free.” It is “the machine that already has the state needed to produce that next token.” That one constraint shaped nearly every decision in this project.

I built a distributed inference engine to understand that constraint from first principles. The project uses a TypeScript coordinator and Rust workers, loads GGUF models through llama.cpp, streams tokens with SSE, and keeps conversational state pinned to the worker that owns it.

This is not me trying to rebuild vLLM in a weekend. It is a deliberately small system for learning the systems problems hiding behind “generate text.”

First, what even is an inference engine?

#

An inference engine is the runtime layer between an application and an LLM. It accepts a prompt, prepares the model state, generates tokens, streams them back, and cleans up afterwards.

For one model on one laptop, that can be one process. For multiple workers, it becomes a distributed system with some awkward properties:

• model weights are large and live on each worker
• a conversation creates state as it runs
• every generated token depends on the ones before it
• memory, especially KV-cache memory, becomes the resource that decides whether a request is safe to accept

My initial goal was intentionally narrower than “distributed model parallelism.” I did not split one model across machines. Each worker runs a full model replica. The coordinator chooses a worker for a request, and that worker owns the inference session.

That gives horizontal capacity for independent conversations while keeping the data plane local.

flowchart TB
    C[Client] -->|POST /coordinator/infer| CO[Coordinator
TypeScript / Express]
    CO -->|select healthy worker| S[Scheduler]
    S --> CO
    CO -->|prefill + decode| W1[Rust worker 1
model + KV cache]
    CO -->|prefill + decode| W2[Rust worker 2
model + KV cache]
    CO -->|SSE tokens| C
    W1 -->|heartbeat + load| CO
    W2 -->|heartbeat + load| CO

The important thing in that diagram is what the coordinator does not own: model weights, token-level execution state, and the KV cache. Those belong to workers.

The thing that changed the design: prefill is not decode

#

Before writing the services, I had to understand the inference loop well enough to draw boundaries around it.

There are two phases.

Prefill processes the prompt. It tokenizes the input, runs it through the model, and creates the attention state required for future generation. That state is called the KV cache.

Decode generates one token at a time. Each new token depends on the cache and on the tokens generated before it. This is the part users see streaming back.

sequenceDiagram
    participant Client
    participant Coordinator
    participant Worker
    Client->>Coordinator: prompt + conversation_id
    Coordinator->>Worker: prefill(session_id, prompt)
    Note right of Worker: build / extend KV cache
    Worker-->>Coordinator: prefill complete
    Coordinator->>Worker: decode(session_id)
    loop one token at a time
        Worker-->>Coordinator: SSE token(seq)
        Coordinator-->>Client: SSE token(seq)
    end

The KV cache is not some optional optimisation I could bolt on later. In systems terms, it is the in-memory continuation state of a request. If Worker A created it, moving decode to Worker B means B has no idea where that conversation is. It either needs the cache moved over the network or it has to replay the prompt.

Both options are expensive. So I locked in an invariant early:

Decode always runs on the worker that owns the session’s KV cache.

That one invariant gave me sticky sessions, worker-local cache lifecycle, and very honest failure behaviour. It also explains why “stateless workers” are not enough for this job.

Starting with boundaries, not clever code

#

I began with the boring documents: architecture, state ownership, request lifecycle, invariants, and failure modes. It sounds slow, but it prevented the codebase from becoming a coordinator that secretly knew everything.

The three main pieces are simple.

I also wrote down a few rules that are more valuable than they first look:

1. KV cache never moves across workers.
2. A conversation stays pinned to its worker while its session is alive.
3. A worker failure affects that worker’s sessions, not everyone else.
4. The scheduler makes routing decisions but does not mutate inference state.
5. A slow client must not be allowed to slow model execution forever.

These are not implementation details. They are guardrails. When an edge case shows up, I can ask which rule it would violate instead of improvising state transfers until the demo works.

Building the request path

#

With those boundaries in place, the first version was intentionally mocked. The worker accepted prefill, stored fake state keyed by session ID, and accepted decode calls. Only after the contracts existed did I connect a real GGUF model.

The Rust worker now uses llama_cpp to load a model, create an inference session, advance that session with the prompt during prefill, and generate completion tokens during decode. The coordinator exposes POST /coordinator/infer, then talks to worker endpoints for prefill and decode.

New conversations go through this path:

flowchart TD
    A[New request] --> B{System can admit it?}
    B -- no --> R[503: reject early]
    B -- yes --> C[Find healthy workers]
    C --> D[Filter workers at capacity]
    D --> E[Score available workers]
    E --> F[Prefill on selected worker]
    F -- worker rejects/fails --> G{Another worker left?}
    G -- yes --> E
    G -- no --> H[Return failure]
    F -- success --> I[Register conversation -> worker + session]
    I --> J[Decode and stream tokens]

For existing conversations, the path is deliberately different. The coordinator first looks up the conversation ID. If the worker and session still exist, it sends a continue prefill request to that exact worker, extending the same context instead of starting over elsewhere.

If the session is gone or full, the API returns a clear reset-required response. Pretending that the conversation could seamlessly continue would be worse than admitting that the state is gone.

Scheduling: load balancing, but memory-aware

#

At first I thought scheduling would be the centrepiece: complex heuristics, lots of scores, maybe a clever algorithm.

The useful version is much more boring.

Workers send heartbeats containing their identity, URL, liveness, active-session count, and KV-cache usage. The coordinator marks a worker alive, stale, or dead based on the age of those heartbeats. Before routing a new request it does an inexpensive admission-control check, estimates the cache requirement, and refuses work if the system is already constrained.

For candidates that remain, the scheduler uses a weighted load score:

score = 0.6 × session utilisation + 0.4 × KV-cache utilisation

The worker with the lowest score wins.

It is intentionally a pure function over soft state. If the scheduler restarts, it loses only its recent view. Heartbeats rebuild that view; it never owned a model session in the first place.

That separation matters because the system’s scarce resource is usually not the coordinator’s CPU. It is the memory occupied by live model sessions. A worker can look free in terms of requests while being unable to take another long context safely.

Failures were a feature, not a footnote

#

The first honest question I asked was: what happens when a worker disappears mid-request?

There is no magical answer. There are only trade-offs.

If a worker fails during prefill, the coordinator can try a different healthy worker. The session has not become useful yet, so redoing the prompt is acceptable.

If it fails during decode, the KV cache vanished with it. The stream fails and the conversation is invalidated. A future request must start a fresh session. In theory I could re-prefill the full conversation automatically; in practice, making that transparent needs a durable source of chat history and a conscious product decision about replay cost. I chose to surface the truth.

stateDiagram-v2
    [*] --> Prefill
    Prefill --> Decode: session created
    Prefill --> RetryElsewhere: worker failure
    RetryElsewhere --> Prefill
    Decode --> Complete: final token
    Decode --> SessionLost: worker failure
    SessionLost --> [*]
    Complete --> Reusable: keep session for next turn
    Reusable --> Prefill: continue prompt
    Reusable --> Expired: TTL / eviction
    Expired --> [*]

This approach is less flashy than claiming fault tolerance everywhere, but it gives each failure a defined blast radius. Other workers, their models, and their caches remain untouched.

Streaming is its own distributed-systems problem

#

Getting tokens from a worker to a client sounds easy until the client is slow.

The worker generates a numbered stream of SSE events. The coordinator reads that stream, checks sequence numbers, puts tokens into a bounded buffer, and writes them to the client with a deadline. It also records buffer occupancy, write latency, dropped tokens, and stream lifecycle events.

flowchart LR
    W[Worker generates token] --> CH[Bounded channel]
    CH --> WS[Worker SSE stream]
    WS --> R[Coordinator reader]
    R --> B[Bounded buffer
64 tokens]
    B --> WR[Deadline-bound writer]
    WR --> CL[Client]
    B -. overflow .-> DROP[Drop oldest + log]
    WR -. timeout .-> END[End client stream]

The decision here was important: protect the worker before protecting a stalled client. A slow browser should not make a model process accumulate an unbounded amount of output in memory. The coordinator can end that client stream while keeping the worker and other sessions healthy.

There is a second subtlety: a normal completed stream, a client disconnect, and a hard decode failure are not the same thing. A normal completion can leave the session alive for the next turn. A hard decode failure tears it down. Making those paths explicit stopped the conversation registry from becoming a pile of optimistic assumptions.

The implementation choices

#

The split was practical rather than ideological.

• TypeScript + Express at the edge made the HTTP API, request validation, orchestration, and SSE forwarding quick to iterate on.
• Rust + Axum on workers provided a compact place for the model process, sessions, cache limits, heartbeats, and streaming.
• HTTP and SSE kept the protocol inspectable while I was learning. I can curl the health endpoints and watch a stream without a special client.
• GGUF + llama.cpp via Rust bindings made it possible to use small, quantized local models instead of making the project dependent on a GPU-only serving stack.
• Docker Compose gives the coordinator and worker a portable demo path, while the same URLs and heartbeat contracts can support more workers.

There are deliberately approximate pieces too. For admission control, cache size is currently estimated from prompt length rather than measured from the backend’s exact allocation. That is good enough to learn the control flow, but real production capacity planning would use model-aware token counts and actual memory telemetry.

What I would change next

#

The system works as an educational inference platform, but “works” should not be confused with “finished.”

First, I would replace rough cache estimates with tokenizer- and model-specific accounting. Then I would add proper cancellation so a disconnected client can cancel generation work, not only stop receiving it. I would also make the coordinator highly available or persist the minimal session metadata needed for recovery.

For higher throughput, the next serious piece is continuous batching: admitting and scheduling prefill/decode work in small batches while preserving per-session ordering. That is much more interesting than blindly adding more workers.

I would still avoid tensor parallelism in this project’s first version. Splitting one model across devices introduces another layer of communication and synchronisation. It solves a different problem: fitting or accelerating a single model that cannot comfortably run on one worker. Here I wanted to understand request-level distribution and state locality first.

The biggest lesson

#

I started this project thinking distributed inference was mainly about sending requests to multiple machines.

It is not.

It is about deciding where state is allowed to live, making that ownership obvious, and being truthful when that state disappears. The model call is one part of the system. The harder part is everything that lets that call survive real clients, full memory, slow networks, stale workers, and the next turn in the conversation.

That is what made building this fun: every “simple” decision had a systems consequence hiding behind it.

The code and the design notes are open source at ishanjain1502/distributed-inference-engine. If you are building something similar, I would start with the invariants. They are much cheaper to change in a Markdown file than after your KV cache has quietly become everyone’s problem.