6. Inference Serving

Continuous Batching Scheduler

seniorai-mldistributed-systems

Naive ("static") batching waits for every request in a batch to fully finish before admitting any new ones -- if one request needs 500 tokens and another needs 5, the GPU sits there holding a mostly-idle slot for the short request's remaining 495 steps. vLLM's core innovation, continuous batching, admits a new request the INSTANT a slot frees up, one generation step at a time, so the batch stays as full as possible continuously instead of draining and refilling in lockstep.

Implement ContinuousBatchScheduler.step() -- one generation step across the whole system:

  1. Admit requests from the waiting queue (FIFO) to fill any free slots in active, up to max_batch_size.
  2. Process one token for every currently active request (decrement its remaining token count).
  3. Remove any request whose remaining count reached 0 -- it's done, which frees its slot for step 1 on the NEXT call to step().

Return the sorted list of request ids that were active (processed) THIS step, before removal.

Example -- the property that matters

sched = ContinuousBatchScheduler(max_batch_size=2)
sched.add_request("A", 3)  # needs 3 tokens
sched.add_request("B", 1)  # needs 1 token -- finishes fast
sched.add_request("C", 2)  # needs 2 tokens

sched.step()  # -> ["A", "B"]   -- C still waiting, batch is full
sched.step()  # -> ["A", "C"]   -- B finished last step, C admitted IMMEDIATELY

C doesn't wait for A (the longest-running request) to finish -- it gets B's slot the moment B is done. That immediacy is the whole point.

The missing link: a free slot isn't the same as free memory

Counting slots is a simplification. Each active request's KV cache is a real chunk of GPU memory, and different requests need different amounts (a request with a long prompt costs more than a short one). A scheduler that only checks "is there a free slot" can admit a request that fits the SLOT budget but blows the MEMORY budget -- which in a real system means an out-of-memory crash, not a slow response.

Add a memory budget: __init__(self, max_batch_size, max_memory=None), and add_request(self, request_id, num_tokens, memory_cost=1) (defaults chosen so passing neither preserves the original slot-only behavior exactly). Admission in step() must now check BOTH conditions: room in active (max_batch_size) AND that admitting wouldn't push total active memory over max_memory. If the request at the FRONT of the waiting queue doesn't currently fit in the remaining memory budget, stop admitting for this step entirely -- don't skip ahead to a smaller request further back just because it happens to fit (strict FIFO, even at the cost of some wasted capacity).

Example

sched = ContinuousBatchScheduler(max_batch_size=5, max_memory=10)
sched.add_request("A", num_tokens=3, memory_cost=8)
sched.add_request("B", num_tokens=3, memory_cost=8)
sched.step()  # -> ["A"]   -- B doesn't fit (8 + 8 > 10), even though a slot is free

Constraints

  • Admission (step 1) happens BEFORE token processing (step 2) within the same step() call -- a newly-admitted request gets processed in the very step it was admitted.
  • active never exceeds max_batch_size, and (when max_memory is set) total active memory never exceeds max_memory.
  • A request whose own memory_cost exceeds max_memory can never be admitted at all, and (since this is strict FIFO) will block every request behind it forever -- a real system would reject such a request outright at submission time; that rejection policy is out of scope here.

Real-world grounding

Real schedulers make this decision even more involved on top of the memory accounting above -- prioritization/fairness across requests, and prefill vs. decode phases needing very different amounts of compute -- but "admit into any free slot immediately, without waiting for the whole batch to drain, subject to a real resource budget" is the core idea underneath all of it.

hints
scheduler.pysign in to run
Loading...
test result

This is a preview. Sign in to run this exercise against the real tests.