Project 14 — Distributed Inference

Cross the machine boundary: split prefill from decode, ship the KV cache over a network, and learn when that is cheaper than recomputing it.

1. What you build#

A two-role serving system:

                 ┌──────────────┐   KV blocks    ┌──────────────┐
client → router →│ PREFILL node │ ─────────────→ │ DECODE node  │→ stream → client
                 └──────────────┘   (network)    └──────────────┘

A prefill worker processes the prompt and produces the KV cache; a transfer layer moves it; a decode worker (your Project 08 engine) admits the request with KV already populated and generates. A router picks a worker of each role. You then compare it against the co-located baseline and work out the break-even.

Diagram — Request flow, with the fallback#

sequenceDiagram
    participant C as Client
    participant R as Router
    participant P as Prefill node
    participant D as Decode node
    C->>R: POST /v1/completions
    R->>P: prompt + target decode node
    P->>P: prefill, sample first token
    alt transfer OK
        P->>D: header + KV blocks + first token
        D->>D: allocate blocks, admit to running set
    else transfer failed or decode node full
        R->>D: prompt only
        D->>D: local prefill
    end
    loop until EOS
        D-->>R: token
        R-->>C: token
    end

2. Why it matters#

Prefill is compute-bound and bursty; decode is memory-bound and steady (V.03). Co-locating them means each interferes with the other and neither can be scaled or tuned separately. Disaggregation is how large deployments address this, and by 2026 it is a mainstream configuration in the major engines rather than a research idea. Its feasibility rests on one number: how fast KV can move versus how fast it can be recomputed. This project makes you measure that number.

3. Read first#

• IX.08 — Interconnects
• IX.09 — Communication cost math
• IX.10 — Multi-node
• IX.11 — Distributed KV cache
• XIII.06 — Disaggregation
• XIII.07 — KV offloading

4. Spec#

Router        POST /v1/completions
              1. choose prefill worker (least queued prompt tokens)
              2. choose decode worker  (least running sequences / most free KV blocks)
              3. tell prefill worker where to send KV ; proxy the token stream from decode

Prefill node  run prefill ; sample first token ; send {request meta, first token, KV blocks}

Transfer      v1  TCP, length-prefixed frames, raw tensor bytes
              v2  zero-copy send (memoryview / sendfile), pinned buffers, optional FP16→INT8
              framing: header{request_id, n_layer, n_tokens, dtype, block_size} + payload
              integrity: checksum per frame

Decode node   allocate blocks (Project 09) ; write KV ; admit into running set ; decode

Failure       prefill dies → router retries on another prefill node
              transfer fails or decode node is full → fall back to local prefill on decode node

5. Milestones#

1. Do the arithmetic first. For your model: KV bytes per token × prompt length = payload. Divide by your link bandwidth → transfer time. Divide prompt length by your measured prefill tokens/s → recompute time. Predict which wins at 100, 1k, 8k tokens on 1 GbE, 10 GbE, 100 GbE.
2. KV export/import in one process. Export after prefill, import into a fresh engine, continue decoding. Tokens identical to the uninterrupted run.
3. Across processes over localhost TCP.
4. Across machines (two VMs, or tc netem to impose bandwidth and latency on loopback).
5. Router with role-aware selection and stream proxying.
6. Baseline. Same hardware, both workers running co-located Project 08 engines.
7. Interference experiment. Steady decode traffic + bursts of 8k-token prompts. Compare decode ITL p99 for co-located vs disaggregated.
8. Failure handling and fallback.

6. Starter skeleton#

HDR = struct.Struct("!16sIIIH")     # request_id, n_layer, n_tokens, payload_bytes, dtype_code

def send_kv(sock, rid, k, v):       # k, v: [n_layer, n_kv_head, T, head_dim], contiguous
    kb, vb = k.cpu().numpy(), v.cpu().numpy()
    sock.sendall(HDR.pack(rid, k.shape[0], k.shape[2], kb.nbytes + vb.nbytes, DTYPES[k.dtype]))
    sock.sendall(memoryview(kb)); sock.sendall(memoryview(vb))      # no pickle, no copies

def recv_kv(sock, cfg):
    rid, L, T, nbytes, dt = HDR.unpack(recv_exact(sock, HDR.size))
    buf = bytearray(nbytes); view = memoryview(buf); got = 0
    while got < nbytes:
        got += sock.recv_into(view[got:], nbytes - got)
    half = nbytes // 2
    shape = (L, cfg.n_kv_head, T, cfg.head_dim)
    k = torch.frombuffer(view[:half], dtype=NP2T[dt]).view(shape)
    v = torch.frombuffer(view[half:], dtype=NP2T[dt]).view(shape)
    return rid, k, v

def should_transfer(n_tokens, kv_bytes_per_tok, link_bps, prefill_tps, overhead_s=0.002):
    return n_tokens * kv_bytes_per_tok / link_bps + overhead_s < n_tokens / prefill_tps

7. What to measure#

8. Done when#

• Decode from transferred KV is token-identical to local.
• You have predicted and measured break-even curves on one chart.
• You showed the ITL-under-burst improvement, or showed that your link is too slow for it and explained which number would have to change.
• Killing the prefill worker mid-request does not fail the request.
• You can state in one sentence when disaggregation is the wrong choice.

9. Common pitfalls#

Pickling tensors. Slow and copies several times. Send raw bytes.

GPU → pageable host → socket → pageable host → GPU. Four copies. The data path, not the wire, becomes the bottleneck — which is why real systems use RDMA and GPU-direct transports.

Ignoring what a fast prefill node costs. Disaggregation only pays if each role is separately saturated.

Model/config mismatch between roles. Different dtype, block size, or quantization makes KV garbage with no error. Put a schema version and checksum in the header.

Not transferring the first sampled token and sampler state.

Assuming a datacenter network. On 1 GbE, recompute wins by a wide margin. That is a valid result — report it.

10. Stretch goals#

• Pipeline parallelism across the two nodes (layers 0-5 on A, 6-11 on B) and its bytes-per-hop cost: batch × d_model × bytes per token. Compare with P/D disaggregation.
• Compress KV in transit (FP16 → INT8, or lossless) and re-plot the break-even.
• Shared KV store: both decode nodes pull prefixes from a Redis/disk tier instead of recomputing — the idea behind LMCache and peer-to-peer KV sharing.
• Run nccl-tests across two real nodes and compare against your TCP path.
• Stream KV layer by layer so decode can start layer 0 before the last layer arrives.

11. Interview questions this project answers#

1. Why separate prefill and decode onto different machines?
2. What determines whether KV transfer beats recomputation? Do the arithmetic.
3. How would you size the prefill-to-decode ratio?
4. What goes wrong if the two roles run slightly different model builds?
5. Why do disaggregated designs depend on RDMA-class networking?

12. Next#

Project 15 — Mini inference platform