Project 08 — Continuous Batching

★ Build the scheduler loop at the heart of vLLM, SGLang, and TGI: decide the batch at every step.

1. What you build#

An engine whose batch membership changes every iteration: finished sequences leave immediately, waiting requests join at the next step. Then a head-to-head against Project 07 on the same workload.

This is exercise A of V.09, built out fully.

Diagram — The step loop#

flowchart TB
  ST(["step"]) --> AD["1. Admit<br/>waiting to running, while slots and KV allow"]
  AD --> PF["2. Prefill new requests<br/>emit their first token"]
  PF --> DC["3. Decode ONE token<br/>for every running request, in one batched forward"]
  DC --> RT["4. Retire finished<br/>free the slot immediately"]
  RT --> ST

  class AD queue
  class PF,DC compute
  class RT memory
  class ST neutral

2. Why it matters#

It is typically the largest single throughput win in LLM serving, and it is the rare optimization that improves latency and throughput together. It is also the most common “design it on the whiteboard” question in inference interviews.

3. Read first#

• V.09 — Continuous batching
• V.10 — PagedAttention (skim; Project 09 builds it)
• VIII.11 — vLLM architecture
• XIII.05 — Chunked prefill (for the stretch)

4. Spec#

Engine
  waiting: deque[Request]
  running: list[Request]            (≤ max_num_seqs)
  step():
     1. admit   — while room and KV memory allow
     2. prefill — newly admitted requests (separate forward, v1)
     3. decode  — ONE token for every running request, in one batched forward
     4. retire  — finished requests free their slot at once; stream tokens out

KV storage v1: slot-based
  K, V : [n_layer, max_num_seqs, n_kv_head, max_len, head_dim]   pre-allocated
  each request owns one slot index; per-slot length; attention mask from lengths

Slot-based storage wastes memory (every slot reserves max_len). That is deliberate: it keeps this project about scheduling. Project 09 replaces it with paged blocks.

You need your own model forward here — Hugging Face’s past_key_values does not let you add and remove rows cheaply. Use your Project 04/05 transformer ported to PyTorch.

5. Milestones#

1. Slot KV + batched decode. A decode step over an arbitrary subset of slots (index_select on the slot dimension), with per-slot lengths driving the mask and the position IDs.
2. Equivalence. Any request’s output is identical whether it ran alone or joined a busy batch mid-flight. This test is non-negotiable.
3. The loop. Admission → prefill → decode → retire, exactly as in V.09 §4.
4. Serve it. Plug into Project 06’s HTTP layer; the engine thread now calls step() forever.
5. Compare. Same workload as Project 07 (heavy-tailed outputs). Report throughput, TTFT p50/p95, slot utilization for static vs continuous.
6. Starve it. Shrink max_num_seqs or max_len, watch the queue grow, then implement memory-aware admission and preemption-by-recompute.

6. Starter skeleton#

type Engine struct {
	model            SlotModel
	kv               *SlotKV // one fixed KV region per slot
	freeSlots        []int
	waiting, running []*Request
}

func (e *Engine) Step() {
	// 1. admit
	var admitted []*Request
	for len(e.waiting) > 0 && len(e.freeSlots) > 0 {
		r := e.waiting[0]
		e.waiting = e.waiting[1:]
		r.Slot, e.freeSlots = e.freeSlots[len(e.freeSlots)-1], e.freeSlots[:len(e.freeSlots)-1]
		r.Admitted = time.Now()
		admitted = append(admitted, r)
	}

	// 3. one decode step for everyone ALREADY running (collected before adding the new ones)
	var dec []*Request
	for _, r := range e.running {
		if !r.Done {
			dec = append(dec, r)
		}
	}

	// 2. prefill each new request (v1: one at a time)
	for _, r := range admitted {
		logits := e.model.Prefill(r.PromptIDs, e.kv, r.Slot)
		r.Push(sample(logits, r.Params)) // first token → TTFT
		e.running = append(e.running, r)
	}

	if len(dec) > 0 {
		slots, last := make([]int, len(dec)), make([]int, len(dec))
		for i, r := range dec {
			slots[i], last[i] = r.Slot, r.OutputIDs[len(r.OutputIDs)-1]
		}
		logits := e.model.Decode(last, e.kv, slots) // (len(dec), vocab)
		for i, r := range dec {
			r.Push(sample(logits[i], r.Params))
		}
	}

	// 4. retire: the slot is reusable on the very next step
	e.running = slices.DeleteFunc(e.running, func(r *Request) bool {
		if r.Done {
			e.kv.Reset(r.Slot)
			e.freeSlots = append(e.freeSlots, r.Slot)
			r.Finish()
		}
		return r.Done
	})
}

7. What to measure#

8. Done when#

• Outputs are identical regardless of what else is in the batch.
• You have a side-by-side chart of static vs continuous on throughput and TTFT.
• You can show the ITL spike caused by a long prefill, with a timestamped trace.
• Preemption works: a deliberately undersized pool still completes every request, and you log the preemption rate.
• You can explain continuous batching to a backend engineer in 60 seconds (Checkpoint C).

9. Common pitfalls#

Cross-talk between rows. A stale mask or wrong per-slot length lets one request attend to another’s leftover KV. The equivalence test catches it; nothing else will.

Not clearing slot lengths on retire.

Sampling with shared RNG state. Outputs then depend on batch composition. Use a per-request generator.

Prefill blocking decode. In v1 it does — measure it rather than hide it.

Removing from running while iterating it.

Counting TTFT from admission instead of arrival. Queue wait is part of TTFT.

10. Stretch goals#

• Mixed batches / chunked prefill: cap prefill to N tokens per step and run it in the same step as decodes. Re-measure the ITL spike.
• Batched prefill of several new requests with padding.
• Priority admission and a fairness policy across two tenants (XII.05).
• Compare your engine’s tokens/s to vLLM on the same model and GPU. Account for the gap item by item: kernels, CUDA graphs, FlashAttention, scheduler overhead.

11. Interview questions this project answers#

1. Write the continuous batching loop in pseudocode.
2. What state does the engine keep per request?
3. Why does the gain depend on output-length variance?
4. Recompute vs swap on preemption — which, and why?
5. How does a long prompt hurt other users, and what fixes it?

12. Next#

Project 09 — KV cache manager