jzzliu
73ff04d77f
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.
114 lines
4.0 KiB
Python
114 lines
4.0 KiB
Python
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)
|