Skip to content

02. Model-Serving Platform — Low-Level Design

~20 min read · Part 3 of 4 (Overview → HLD → LLD → Q&A)

The HLD named the boxes. This file opens the four that carry the design's weight — the continuous-batching scheduler, the paged KV-cache manager, the autoscaler, and the version-rollout controller — and pins down the data, the algorithms, and the concurrency corners where an inference platform actually breaks.

Data models

The control plane's source of truth is small and read-hot. Routing is data the router reads, not code it deploys, which is what makes rollback a sub-second config edit.

CREATE TABLE model_version (
    model        VARCHAR(64)  NOT NULL,       -- "chat-model"
    version      VARCHAR(32)  NOT NULL,       -- "v43"
    artifact_uri VARCHAR(256) NOT NULL,       -- s3://.../v43/  (~26 GB)
    state        SMALLINT     NOT NULL,       -- 0=loading 1=ready 2=canary 3=stable 4=retired
    sla_tpot_ms  INT          NOT NULL DEFAULT 30,
    created_at   TIMESTAMP    NOT NULL,
    PRIMARY KEY (model, version)
);

CREATE TABLE routing_rule (                   -- the live traffic split
    model        VARCHAR(64)  NOT NULL,
    version      VARCHAR(32)  NOT NULL,
    traffic_pct  SMALLINT     NOT NULL,        -- must sum to 100 per model
    PRIMARY KEY (model, version)
);

Two deliberate choices. First, traffic_pct is a per-version row, not a single "canary %" field, so a model can run three versions at once (v41 retiring, v42 stable, v43 canary) and the router just samples the distribution. Second, the SLA (sla_tpot_ms) is stored per version, because the Rollout Judge compares a canary against a fixed contract, not merely against the incumbent — a new version that is faster than v42 but still over the absolute ceiling should not pass.

Inside a replica, the scheduler tracks each in-flight request and the KV-cache manager tracks GPU memory as pages. These are in-memory structures, not tables, but their shape is the crux of the design:

@dataclass
class Sequence:
    req_id:      str
    tenant_id:   str
    prompt_ids:  list[int]        # 500 tokens in the scenario
    output_ids:  list[int]        # grows one per decode step
    max_tokens:  int              # e.g. 200
    kv_pages:    list[int]        # page indices owned in the KV pool
    stage:       Stage            # WAITING | PREFILL | DECODE | DONE
    arrived_at:  float

# KV cache as a paged pool (PagedAttention-style):
PAGE_TOKENS = 16                  # tokens of K/V per page
# page_size_bytes = 0.8 MB/token × 16 = ~12.8 MB per page
# usable HBM ~13 GB / 12.8 MB  ≈ ~1,020 pages per GPU

The KV cache is deliberately not one contiguous buffer per request. A request's attention state grows token by token to an unknown final length, and pre-reserving its maximum wastes memory on requests that stop early. Paging it into 16-token blocks means a request holds only ceil(current_len / 16) pages and the pool stays near-fully-packed, which is the difference between ~23 and ~45 concurrent sequences per GPU.

Component internals

Component 1 — Continuous-batching scheduler

Responsibility: every step, choose the set of sequences that run the next forward pass together, admitting new arrivals and evicting finished ones without ever waiting for a batch window to fill.

The classic mistake it avoids is static batching — collect N requests, run them to completion together, then take the next N. Static batching wastes the GPU whenever sequences in a batch finish at different lengths (short ones idle waiting for the longest), and it makes a late arrival wait for the whole batch to drain. Continuous batching instead treats the batch as a rolling set edited every step.

class Scheduler:
    def step(self) -> None:
        # 1. Admit new arrivals into the running set, up to caps.
        while self.waiting and self._can_admit():
            seq = self.waiting.peek()
            need = ceil(len(seq.prompt_ids) / PAGE_TOKENS)
            if not self.kv.can_allocate(need):       # KV pressure → stop admitting
                break
            if self._tenant_over_quota(seq.tenant_id): # fairness cap
                self.waiting.rotate(); continue
            self.kv.allocate(seq, need)
            seq.stage = PREFILL
            self.running.append(self.waiting.pop())

        # 2. Run one fused forward pass: prefill new prompts + decode the rest.
        batch = [s for s in self.running if s.stage in (PREFILL, DECODE)]
        tokens = self.model.forward(batch)          # the GPU step

        # 3. Append generated tokens; retire or preempt.
        for seq, tok in zip(batch, tokens):
            seq.output_ids.append(tok)
            seq.stage = DECODE
            if tok == EOS or len(seq.output_ids) >= seq.max_tokens:
                self.kv.free(seq); seq.stage = DONE
                self.running.remove(seq); self._stream_finish(seq)
            elif len(seq.output_ids) % PAGE_TOKENS == 0:
                self.kv.grow(seq, 1)                 # next page as it grows
            else:
                self._stream_token(seq, tok)

    def _can_admit(self) -> bool:
        # Stop admitting before TPOT breaks: batch size caps at the point
        # where the measured step time still clears the SLA.
        return len(self.running) < self.max_batch and self.kv.free_pages() > SAFETY

Two caps govern it. max_batch is set where the measured step time still clears the TPOT SLA — in the scenario, ~45. And kv.can_allocate gates admission on free pages, because admitting a sequence you cannot grow leads to mid-generation preemption, which is far more expensive than making it wait in the queue a moment longer.

Component 2 — Paged KV-cache manager

Responsibility: hand out and reclaim GPU memory for attention state in fixed pages, keeping the pool packed and giving the scheduler an honest free-page count.

class KVCacheManager:
    def __init__(self, total_pages: int):
        self.free: deque[int] = deque(range(total_pages))   # ~1,020 on our GPU
        self.table: dict[str, list[int]] = {}               # req_id -> page ids

    def can_allocate(self, n: int) -> bool:  return len(self.free) >= n
    def allocate(self, seq, n):              seq.kv_pages = [self.free.popleft() for _ in range(n)]; self.table[seq.req_id] = seq.kv_pages
    def grow(self, seq, n):                  seq.kv_pages += [self.free.popleft() for _ in range(n)]
    def free(self, seq):                     self.free.extend(seq.kv_pages); del self.table[seq.req_id]

    def preempt(self, seq):
        # Under pressure, reclaim a victim's pages; it re-prefills later.
        self.free.extend(seq.kv_pages); seq.kv_pages = []; seq.stage = WAITING

The non-obvious part is preempt. When the pool is full and a high-priority request must be admitted, the manager evicts a victim's pages back to the free list and the scheduler moves that victim back to WAITING — it will re-prefill its prompt-plus-generated-so-far when memory frees. Preemption is correct but expensive (recompute), so the scheduler's job is to avoid needing it by not over-admitting. An optional refinement is prefix caching: pages holding a shared prompt prefix (a common system prompt across requests) are kept and reused rather than recomputed, and cold prefixes can be offloaded to CPU RAM instead of dropped.

Component 3 — Rollout controller + judge

Responsibility: drive a version from 1% to 100% through automated gates, and reverse instantly on regression.

class RolloutController:
    def advance(self, dep: Deployment):
        step = dep.steps[dep.cursor]                 # e.g. {pct: 5}
        self.autoscaler.ensure_replicas(dep.version, for_pct=step.pct)
        self.registry.set_traffic(dep.model, dep.version, step.pct)  # config edit
        self.bake_until = now() + dep.bake_window     # e.g. +10 min

    def judge(self, dep: Deployment) -> Verdict:
        cur = metrics.window(dep.version, last="5m")
        base = metrics.window(dep.incumbent, last="5m")
        if cur.error_rate > base.error_rate + 0.005:        return ROLLBACK
        if cur.p99_tpot_ms > dep.sla_tpot_ms:               return ROLLBACK
        if cur.p99_tpot_ms > base.p99_tpot_ms * 1.20:       return ROLLBACK
        if now() >= self.bake_until:                        return ADVANCE
        return HOLD

The judge checks three things, not one: an absolute SLA ceiling, a relative regression bound against the incumbent, and an error-rate delta. A version can be within the absolute SLA yet 20% slower than what it replaces, and a serious platform treats that as a regression worth catching before 100%.

Core algorithm — surviving the canary regression at 1% (threaded scenario)

Walk the full scenario through the machinery. The platform serves 1,000 req/s on ~150 GPUs running v42, holding p99 TPOT at ~22 ms against a 30 ms SLA. An operator deploys v43 as a canary.

  1. Warm the canary. The autoscaler launches enough v43 replicas for the first step. Each pulls ~26 GB from the artifact store and loads it — call it 3 minutes of cold start, done while v42 still serves 100%. v43 replicas pass health checks and register.
  2. Route 1%. The controller sets routing_rule for v43 to traffic_pct = 1. The router now sends ~10 of the 1,000 req/s to v43; the other 990 stay on v42. This is the blast-radius cap: whatever is wrong with v43 can hurt at most 10 req/s.
  3. Bake and measure. Over a 10-minute window the judge collects v43's telemetry. Suppose v43 ships a subtly larger tokenizer vocabulary that inflates its per-step time: its p99 TPOT measures 41 ms against the 30 ms SLA and v42's 22 ms.
  4. Judge trips. On the next evaluation, cur.p99_tpot_ms (41) > sla_tpot_ms (30)ROLLBACK. No human needed.
  5. Reset traffic. The controller writes v43's traffic_pct back to 0. Because routing is config the router re-reads within seconds, the ~10 req/s on v43 immediately shift back to v42. Total exposure: ~10 req/s for the bake window, zero after. The 990 req/s on v42 never moved, and p99 across the platform never left 22 ms for 99% of traffic.
  6. Contrast the naive path. Had the platform flipped v43 to 100% in one deploy, all 1,000 req/s would have jumped to 41 ms TPOT — a platform-wide SLA breach — and rollback would have meant a redeploy, minutes long, while every request suffered. Canarying at 1% turned a company-wide incident into a 10-req/s footnote in a dashboard. That gap is the entire value of the rollout machinery.

Sequence diagram — an inference request through continuous batching

Client        Router        Scheduler       KVCache        Model(GPU)
  │  POST infer  │              │               │               │
  ├─────────────▶│  pick v42    │               │               │
  │              │  replica     │               │               │
  │              ├─────────────▶│ enqueue(req)  │               │
  │              │              ├─ can_allocate?▶│  (500 tok →   │
  │              │              │◀── yes ────────┤   32 pages)   │
  │              │              ├─ allocate ────▶│               │
  │              │              │  stage=PREFILL │               │
  │              │              ├──────── forward(batch) ───────▶│  prefill
  │◀─────────────┼──────────────┤◀── first token ───────────────┤  (TTFT)
  │  token 0     │              │               │               │
  │              │              │  ── loop per decode step ──    │
  │              │              ├──────── forward(batch) ───────▶│  decode
  │◀─────────────┼──────────────┤◀── next token ────────────────┤  (TPOT)
  │  token n     │              ├─ grow(seq) ──▶│ (every 16 tok) │
  │              │              │      ...       │               │
  │              │              │  EOS at 200    │               │
  │              │              ├─ free(seq) ──▶│  (32 pages back)│
  │◀─────────────┼──────────────┤ finish_reason │               │

The request is admitted, prefilled once (first token, TTFT), then advances one token per decode step alongside every other in-flight sequence (each token is one TPOT), growing its KV pages as it lengthens, and finally frees all its pages the instant it stops — which is what lets a queued request take its slot on the very next step.

Concurrency and edge cases

  • Version pin during rollout. A request with X-Version-Pin: v42 must land on v42 even while v43 is at 25% — the router honors the pin before sampling the split, so a client that has qualified a specific version is never silently upgraded mid-rollout.
  • Retry idempotency after a replica death. A dropped in-flight generation is retried against another replica, restarting from the prompt. Because generation is not idempotent token-for-token (sampling is random), the retry may produce different output; the contract is "a valid completion," not "the same completion." Requests carry a client-supplied request_id so a retry that races the original does not get double-counted in usage.
  • KV-cache race between admit and grow. The scheduler must not admit a new sequence (consuming free pages) if doing so would starve an in-flight sequence that needs to grow on the same step. Admission checks free pages after reserving growth headroom for the running set, so a decode step never fails to allocate the one page a live sequence needs — starving a half-finished generation is worse than delaying a new one.
  • Preemption thrashing under overload. If arrivals exceed capacity, naive preemption evicts and re-prefills the same sequences repeatedly, and throughput collapses toward zero. The guard is upstream: the gateway sheds with 429 once queue depth crosses the point where admitted work already saturates the fleet, so the scheduler is never asked to run more than it can without preempting.
  • Routing-config read-your-writes. A rollback edit must be visible to every router promptly, or some routers keep sending traffic to a failing version. Routers watch the config store (or long-poll) rather than caching on a fixed TTL, so a rollback propagates in seconds; the last-known-good copy is only a fallback for when the store itself is unreachable.
  • Tenant fairness under a burst. A tenant firing 500 concurrent long generations is bounded by a per-tenant cap on batch slots and KV pages in _tenant_over_quota, so its excess requests rotate to the back of the waiting queue instead of evicting a neighbor's in-flight sequences. Isolation is enforced in the scheduler's admission step, not by giving the tenant its own GPUs.