fix: isolate continuous batch decode state

- Match steady-state metadata to the active task IDs
- Rebuild request mappings for cached prefix pages
- Add regressions for batch refill and prefix reuse
This commit is contained in:
2026-08-09 01:01:41 +08:00
parent a33ca04f60
commit be90dfe2bd
4 changed files with 78 additions and 3 deletions
+24
View File
@@ -415,6 +415,30 @@ def test_page_pool_paged_ps64_task_extend_crosses_page():
assert len(task_cache._states["t1"].pages) >= 2
def test_page_pool_prefix_hit_populates_request_mapping():
pool = _make_paged_pool_ps64(page_size=2, max_seq_len=8, n_tokens=16)
task_cache = _make_task_cache(pool)
prompt = [11, 12, 13, 14]
assert task_cache.task_alloc("first", prompt)
task_cache.task_record_hashes("first", prompt)
task_cache.task_free("first")
assert task_cache.task_alloc("second", prompt)
second_state = task_cache._states["second"]
expected = [
page * pool.page_size + offset
for page in second_state.pages
for offset in range(pool.page_size)
]
assert second_state.cached == len(prompt)
assert (
pool.req_pool.req_to_token[second_state.req_idx, : len(prompt)].tolist()
== expected
)
def test_page_pool_paged_ps64_bind_roundtrip():
pool = _make_paged_pool_ps64(n_layers=1, n_kv_heads=2, head_dim=4)
task_cache = _make_task_cache(pool)