From b01b1a80f9fc7a744a2373aca3c7b24b27aa10e0 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:56:42 +0900 Subject: [PATCH] fix: reject a KV page pool smaller than one max_context request (#169) submit() admits requests up to max_context, but with pages < max_context / page_size such a request could never get its pages: FIFO admission blocked everything behind it and the owner loop spun. PagedKV now rejects that configuration at construction. test_paged_admission_follows_pages_not_slots used such a pool (6 pages, max_context 128); it now uses max_context 96, which its requests fit. Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/accel/pd/paged.py | 2 ++ opt_prime/tests/test_pd_batching.py | 8 +++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/opt_prime/opt_prime/accel/pd/paged.py b/opt_prime/opt_prime/accel/pd/paged.py index d9e5f0df..a9b0bf5c 100644 --- a/opt_prime/opt_prime/accel/pd/paged.py +++ b/opt_prime/opt_prime/accel/pd/paged.py @@ -166,6 +166,8 @@ def __init__(self, layers, slots, kv_heads, head_dim, dtype, device, *, pages, p """groups: `torch_models.KVGroup`s of a hybrid model; default one full-attention group of all layers""" if max_context % page_size: raise ValueError("max_context must be a multiple of page_size") + if pages < max_context // page_size: # else a max_context request is never admitted + raise ValueError("pages must hold one max_context request") self.page_size, self.pages, self.slots = page_size, pages, slots if groups is None: from opt_prime.accel.pd.torch_models import KVGroup diff --git a/opt_prime/tests/test_pd_batching.py b/opt_prime/tests/test_pd_batching.py index 047c77d3..fbe99925 100644 --- a/opt_prime/tests/test_pd_batching.py +++ b/opt_prime/tests/test_pd_batching.py @@ -133,7 +133,7 @@ def test_moe_family_and_rejections(kv): def test_paged_admission_follows_pages_not_slots(): """with paged KV a request takes only the pages it needs; admission waits for pages, not for max_context""" model = _llama() - b = ContinuousBatcher(model, kv="paged", slots=8, max_context=128, prefill_chunk=64, page_size=16, pages=6) + b = ContinuousBatcher(model, kv="paged", slots=8, max_context=96, prefill_chunk=64, page_size=16, pages=6) prompts = _prompts(5, 10, 30, seed=2) fs = [b.submit(PDRequest(f"p{i}", p, 8)) for i, p in enumerate(prompts)] # each needs 2-3 pages b.step() @@ -142,3 +142,9 @@ def test_paged_admission_follows_pages_not_slots(): for p, f in zip(prompts, fs): assert f.result() == _greedy(model, p, 8) assert len(b.kv.free_pages) == 6 and all(not o for o in b.kv.owned) + + +def test_paged_pool_smaller_than_one_request_is_rejected(): + """#169: a request admitted by max_context must fit the pool, or it blocks the FIFO queue forever""" + with pytest.raises(ValueError): + ContinuousBatcher(_llama(), kv="paged", slots=2, max_context=64, page_size=16, pages=2)