From 73ff04d77f21665d98fae6b66b3b48748c39e6dc Mon Sep 17 00:00:00 2001 From: jzzliu Date: Sat, 7 Nov 2026 20:27:21 -0700 Subject: [PATCH] Add continuous-batching engine with variable-length KV cache Requests of different lengths share one left-padded cache with explicit per-row RoPE positions and a key mask. Requests join mid-flight and leave on EOS/length, and unused padding columns are trimmed. Greedy output is verified identical to solo generation. --- scripts/bench_batching.py | 51 ++++++++++ src/microgpt/batching.py | 195 ++++++++++++++++++++++++++++++++++++++ tests/test_batching.py | 113 ++++++++++++++++++++++ 3 files changed, 359 insertions(+) create mode 100644 scripts/bench_batching.py create mode 100644 src/microgpt/batching.py create mode 100644 tests/test_batching.py diff --git a/scripts/bench_batching.py b/scripts/bench_batching.py new file mode 100644 index 0000000..5d1133a --- /dev/null +++ b/scripts/bench_batching.py @@ -0,0 +1,51 @@ +"""Throughput of the continuous-batching engine vs serving requests one at a time. + + python scripts/bench_batching.py +""" + +from __future__ import annotations + +import time + +import torch + +from microgpt.batching import BatchEngine +from microgpt.model import GPT, GPTConfig + + +def main() -> None: + torch.manual_seed(0) + cfg = GPTConfig(vocab_size=96, block_size=256, n_layer=6, n_head=6, n_embd=192, n_kv_head=2) + model = GPT(cfg).eval() + g = torch.Generator().manual_seed(0) + reqs = [] + for _ in range(16): + length = int(torch.randint(8, 48, (1,), generator=g)) + reqs.append((torch.randint(0, 96, (length,), generator=g).tolist(), 64)) + total = sum(n for _, n in reqs) + print(f"{len(reqs)} requests, prompts 8-47 tokens, 64 new tokens each ({total} total)\n") + print("| strategy | wall (s) | tokens/s | speedup | mean batch |") + print("|:--|--:|--:|--:|--:|") + + model.generate(torch.tensor([reqs[0][0]]), 8, temperature=0) # warm up + t0 = time.perf_counter() + for prompt, n in reqs: + model.generate(torch.tensor([prompt]), n, temperature=0) + base = time.perf_counter() - t0 + print(f"| one at a time | {base:.1f} | {total / base:.0f} | 1.0x | 1.0 |") + + for max_batch in (2, 4, 8, 16): + eng = BatchEngine(model, max_batch=max_batch) + for prompt, n in reqs: + eng.submit(prompt, n) + t0 = time.perf_counter() + eng.run() + dt = time.perf_counter() - t0 + print( + f"| continuous, max_batch={max_batch} | {dt:.1f} | {total / dt:.0f} " + f"| {base / dt:.1f}x | {eng.mean_batch_size:.1f} |" + ) + + +if __name__ == "__main__": + main() diff --git a/src/microgpt/batching.py b/src/microgpt/batching.py new file mode 100644 index 0000000..06da2a8 --- /dev/null +++ b/src/microgpt/batching.py @@ -0,0 +1,195 @@ +"""Continuous batching: serve many variable-length requests from one shared KV cache. + +How it works +------------ +* Every active sequence is one row of a single batched cache (B, n_kv, S, hd). +* Rows have different lengths, so shorter rows are **left-padded** to the common width S and a + boolean `mask` (B, S) marks which columns are real. Padding is never attended to. +* RoPE positions are passed explicitly per row (`positions`), so padding does not shift a + sequence's positions and cached keys keep the rotation they were stored with. +* New requests are prefilled on their own and then spliced into the batch (padding either side + as needed) while others are mid-decode. Finished requests are dropped from the batch and any + columns that became padding for *every* remaining row are trimmed away. + +Because positions and masking are exact, greedy output for a request is identical to running it +alone with `GPT.generate` (this is what the tests check). With sampling, a shared RNG means +results depend on batch composition, as in any real batched server. + +Limitation: a sequence ends when it reaches `block_size` tokens (no sliding window here). +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from typing import Deque, Dict, List, Optional, Tuple + +import torch + +from microgpt.model import GPT, KVCache +from microgpt.sampling import sample_next + + +@dataclass +class _Active: + rid: int + max_new: int + eos_id: Optional[int] + out: List[int] # generated tokens so far + cache_len: int # tokens already in this row's cache == position of the pending token + pending: int # last sampled token, not yet fed through the model + + +def _pad_left(t: torch.Tensor, n: int) -> torch.Tensor: + """Prepend n zero columns along the sequence axis (dim 2) of a (B, H, S, hd) tensor.""" + if n == 0: + return t + pad = torch.zeros(t.size(0), t.size(1), n, t.size(3), dtype=t.dtype, device=t.device) + return torch.cat((pad, t), dim=2) + + +def _pad_mask_left(m: torch.Tensor, n: int) -> torch.Tensor: + if n == 0: + return m + return torch.cat((torch.zeros(m.size(0), n, dtype=torch.bool, device=m.device), m), dim=1) + + +class BatchEngine: + def __init__( + self, + model: GPT, + max_batch: int = 8, + temperature: float = 0.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + generator: Optional[torch.Generator] = None, + ) -> None: + self.model = model.eval() + self.max_batch = max_batch + self.sampling = dict(temperature=temperature, top_k=top_k, top_p=top_p, generator=generator) + self.waiting: Deque[Tuple[int, List[int], int, Optional[int]]] = deque() + self.active: List[_Active] = [] + self.cache: Optional[KVCache] = None + self.mask: Optional[torch.Tensor] = None # (B, S) bool, True = real token + self.results: Dict[int, List[int]] = {} + self._next_id = 0 + self.stats = {"decode_steps": 0, "generated": 0, "batch_rows": 0, "prefills": 0} + + # ------------------------------------------------------------------ public API + def submit(self, prompt: List[int], max_new_tokens: int, eos_id: Optional[int] = None) -> int: + block = self.model.cfg.block_size + if not 1 <= len(prompt) <= block: + raise ValueError(f"prompt length must be in [1, {block}], got {len(prompt)}") + if max_new_tokens < 1: + raise ValueError("max_new_tokens must be >= 1") + rid = self._next_id + self._next_id += 1 + self.waiting.append((rid, list(prompt), max_new_tokens, eos_id)) + return rid + + def step(self) -> bool: + """Admit waiting requests, then advance every active row by one token. + + Returns True while there is still work left.""" + self._admit() + if self.active: + self._decode_step() + return bool(self.waiting or self.active) + + def run(self) -> Dict[int, List[int]]: + while self.step(): + pass + return self.results + + @property + def mean_batch_size(self) -> float: + steps = self.stats["decode_steps"] + return self.stats["batch_rows"] / steps if steps else 0.0 + + # ------------------------------------------------------------------ internals + def _sample(self, logits: torch.Tensor) -> List[int]: + return sample_next(logits, **self.sampling).flatten().tolist() + + def _finished(self, a: _Active) -> bool: + hit_eos = a.eos_id is not None and a.out[-1] == a.eos_id + # The pending token needs a free position (< block_size) to be fed next. + return len(a.out) >= a.max_new or hit_eos or a.cache_len >= self.model.cfg.block_size + + def _admit(self) -> None: + while self.waiting and len(self.active) < self.max_batch: + rid, prompt, max_new, eos_id = self.waiting.popleft() + with torch.no_grad(): + logits, _, cache = self.model(torch.tensor([prompt])) + first = self._sample(logits[:, -1])[0] + self.stats["prefills"] += 1 + self.stats["generated"] += 1 + row = _Active(rid, max_new, eos_id, [first], len(prompt), first) + if self._finished(row): + self.results[rid] = row.out + else: + self._splice(cache, row) + + def _splice(self, cache: KVCache, row: _Active) -> None: + """Add a freshly prefilled single-row cache to the batch, aligning widths.""" + length = row.cache_len + if self.cache is None: + self.cache = cache + self.mask = torch.ones(1, length, dtype=torch.bool) + else: + assert self.mask is not None + width = self.mask.size(1) + n = max(width, length) + self.cache = [ + ( + torch.cat((_pad_left(k0, n - width), _pad_left(k1, n - length)), dim=0), + torch.cat((_pad_left(v0, n - width), _pad_left(v1, n - length)), dim=0), + ) + for (k0, v0), (k1, v1) in zip(self.cache, cache) + ] + self.mask = torch.cat( + ( + _pad_mask_left(self.mask, n - width), + _pad_mask_left(torch.ones(1, length, dtype=torch.bool), n - length), + ), + dim=0, + ) + self.active.append(row) + + @torch.no_grad() + def _decode_step(self) -> None: + assert self.cache is not None and self.mask is not None + rows = len(self.active) + idx = torch.tensor([[a.pending] for a in self.active]) + positions = torch.tensor([[a.cache_len] for a in self.active]) + mask = torch.cat((self.mask, torch.ones(rows, 1, dtype=torch.bool)), dim=1) + logits, _, self.cache = self.model( + idx, cache=self.cache, positions=positions, key_mask=mask + ) + self.mask = mask + for a, tok in zip(self.active, self._sample(logits[:, -1])): + a.cache_len += 1 + a.out.append(tok) + a.pending = tok + self.stats["decode_steps"] += 1 + self.stats["batch_rows"] += rows + self.stats["generated"] += rows + self._retire() + + def _retire(self) -> None: + keep = [i for i, a in enumerate(self.active) if not self._finished(a)] + for a in self.active: + if self._finished(a): + self.results[a.rid] = a.out + if len(keep) == len(self.active): + return + if not keep: + self.active, self.cache, self.mask = [], None, None + return + assert self.cache is not None and self.mask is not None + sel = torch.tensor(keep) + self.active = [self.active[i] for i in keep] + mask = self.mask[sel] + # Drop leading columns that are padding for every surviving row. + first = int(mask.any(dim=0).nonzero()[0]) + self.mask = mask[:, first:] + self.cache = [(k[sel][:, :, first:], v[sel][:, :, first:]) for k, v in self.cache] diff --git a/tests/test_batching.py b/tests/test_batching.py new file mode 100644 index 0000000..8d51584 --- /dev/null +++ b/tests/test_batching.py @@ -0,0 +1,113 @@ +import pytest +import torch +from helpers import make_model + +from microgpt.batching import BatchEngine +from microgpt.model import GPTConfig + +CFG = GPTConfig(vocab_size=50, block_size=32, n_layer=2, n_head=4, n_embd=32, n_kv_head=2) + + +@pytest.fixture(scope="module") +def model(): + return make_model(CFG, seed=3) + + +def alone(model, prompt, n): + out = model.generate(torch.tensor([prompt]), n, temperature=0) + return out[0, len(prompt) :].tolist() + + +def prompt_of(length, seed): + g = torch.Generator().manual_seed(seed) + return torch.randint(0, CFG.vocab_size, (length,), generator=g).tolist() + + +def test_left_padded_forward_matches_unpadded(model): + """The core mechanic: padding + explicit positions + key mask is invisible to real tokens.""" + prompts = [prompt_of(n, n) for n in (3, 9, 6)] + width = max(map(len, prompts)) + idx = torch.zeros(3, width, dtype=torch.long) + key_mask = torch.zeros(3, width, dtype=torch.bool) + positions = torch.zeros(3, width, dtype=torch.long) + for r, p in enumerate(prompts): + idx[r, width - len(p) :] = torch.tensor(p) + key_mask[r, width - len(p) :] = True + positions[r, width - len(p) :] = torch.arange(len(p)) + with torch.no_grad(): + batched = model(idx, positions=positions, key_mask=key_mask)[0] + for r, p in enumerate(prompts): + solo = model(torch.tensor([p]))[0][0] + assert torch.allclose(batched[r, width - len(p) :], solo, atol=1e-5) + assert torch.isfinite(batched).all() # padded rows must not produce NaN + + +def test_engine_matches_individual_generation(model): + """5 requests, batch of 3: forces queueing and mid-flight admission.""" + reqs = [(prompt_of(n, 10 + n), m) for n, m in [(3, 12), (8, 5), (5, 14), (1, 9), (7, 10)]] + eng = BatchEngine(model, max_batch=3) + ids = [eng.submit(p, m) for p, m in reqs] + results = eng.run() + for rid, (p, m) in zip(ids, reqs): + assert results[rid] == alone(model, p, m) + assert eng.cache is None and not eng.active # fully drained + assert eng.mean_batch_size > 1.5 # it actually batched + + +def test_requests_joining_mid_flight(model): + a, b, c = prompt_of(4, 1), prompt_of(10, 2), prompt_of(2, 3) + eng = BatchEngine(model, max_batch=4) + ra = eng.submit(a, 15) + for _ in range(4): + eng.step() + rb = eng.submit(b, 8) # longer prompt than the current batch width + eng.step() + rc = eng.submit(c, 6) # shorter prompt than the current batch width + results = eng.run() + assert results[ra] == alone(model, a, 15) + assert results[rb] == alone(model, b, 8) + assert results[rc] == alone(model, c, 6) + + +def test_eos_stops_a_sequence_without_disturbing_others(model): + p, q = prompt_of(5, 21), prompt_of(6, 22) + free_run = alone(model, p, 10) + eos = free_run[3] + cut = free_run.index(eos) + 1 # first occurrence, inclusive + eng = BatchEngine(model, max_batch=2) + rp = eng.submit(p, 10, eos_id=eos) + rq = eng.submit(q, 10) + results = eng.run() + assert results[rp] == free_run[:cut] + assert results[rq] == alone(model, q, 10) + + +def test_sequence_stops_at_block_size(model): + p = prompt_of(30, 5) # 2 free positions + the first token from the prefill + eng = BatchEngine(model) + rid = eng.submit(p, 10) + out = eng.run()[rid] + assert len(out) == 3 + assert out == alone(model, p, 3) + + +def test_finished_rows_release_cache_columns(model): + long_p, short_p = prompt_of(20, 7), prompt_of(2, 8) + eng = BatchEngine(model, max_batch=2) + eng.submit(long_p, 5) # leaves after 4 decode steps (first token comes from the prefill) + eng.submit(short_p, 10) + widths = [] + while eng.step(): + if eng.mask is not None: + widths.append(eng.mask.size(1)) + assert max(widths) >= 20 and widths[-1] < 20 # width shrank after the long row left + + +def test_submit_validation(model): + eng = BatchEngine(model) + with pytest.raises(ValueError): + eng.submit([], 3) + with pytest.raises(ValueError): + eng.submit([1] * (CFG.block_size + 1), 3) + with pytest.raises(ValueError): + eng.submit([1], 0)