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

@@ -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()

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]

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)