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.
This commit is contained in:
@@ -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()
|
||||
@@ -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]
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user