diff --git a/python/tokenspeed/runtime/execution/cache_loc_kernel.py b/python/tokenspeed/runtime/execution/cache_loc_kernel.py index 733082409..aee338d2c 100644 --- a/python/tokenspeed/runtime/execution/cache_loc_kernel.py +++ b/python/tokenspeed/runtime/execution/cache_loc_kernel.py @@ -395,20 +395,32 @@ def flatten_and_to_device(data, dtype=torch.int32): tensor = torch.tensor(flat, dtype=dtype, device="cpu", pin_memory=True) return tensor.to(device, non_blocking=True) - # sizes[i] is the number of newly allocated pages for request i. - if all(n == 0 for n in forward_op.sizes): - return - max_pages = req_to_page.shape[1] - # Clamp a request that would overflow req_to_page instead of crashing the - # engine. Happens when MTP accept-rate collapse keeps a request alive past - # context_len; its KV drops but it will be finished shortly. sizes = list(forward_op.sizes) begins = list(forward_op.begins) # new_occupied_pages is a list-of-lists [batch, size_i] of page ids; # take a shallow copy so we can trim the offending request's row. new_occupied_pages = [list(row) for row in forward_op.new_occupied_pages] request_ids = list(forward_op.request_ids) + # full_refresh[i]==1 means the op rebuilt/adopted the prefix (non-tail) pages, + # so copy the scheduler's full page row from logical page 0 instead of the + # begin/size tail delta (a tail-only copy would leave stale/aliased entries). + full_refresh = forward_op.full_refresh + for i, refresh in enumerate(full_refresh): + if refresh: + row = list(forward_op.occupied_pages[i]) + begins[i] = 0 + sizes[i] = len(row) + new_occupied_pages[i] = row + + # sizes[i] is the number of pages to copy for request i (tail delta, or the + # whole row for a full-refresh request). + if all(n == 0 for n in sizes): + return + + # Clamp a request that would overflow req_to_page instead of crashing the + # engine. Happens when MTP accept-rate collapse keeps a request alive past + # context_len; its KV drops but it will be finished shortly. for i, (begin, size) in enumerate(zip(begins, sizes)): if begin + size > max_pages: clamped = max(0, max_pages - begin) diff --git a/test/runtime/test_update_block_table_capacity_clamp.py b/test/runtime/test_update_block_table_capacity_clamp.py index cc70b4a92..15972aa29 100644 --- a/test/runtime/test_update_block_table_capacity_clamp.py +++ b/test/runtime/test_update_block_table_capacity_clamp.py @@ -39,6 +39,8 @@ def _make_forward_op( new_occupied_pages: list[list[int]] | None = None, request_ids: list[str] | None = None, request_pool_indices: list[int] | None = None, + full_refresh: list[int] | None = None, + occupied_pages: list[list[int]] | None = None, ) -> SimpleNamespace: """Build a minimal forward_op stand-in with just the fields the function reads.""" if new_occupied_pages is None: @@ -47,12 +49,20 @@ def _make_forward_op( request_ids = [f"req-{i}" for i in range(len(begins))] if request_pool_indices is None: request_pool_indices = list(range(len(begins))) + if full_refresh is None: + full_refresh = [0 for _ in begins] + if occupied_pages is None: + # For tail-only rows occupied_pages is unused; default it to the tail + # delta so the field is always present and well-shaped. + occupied_pages = [list(row) for row in new_occupied_pages] return SimpleNamespace( begins=list(begins), sizes=list(sizes), new_occupied_pages=new_occupied_pages, request_ids=request_ids, request_pool_indices=request_pool_indices, + full_refresh=list(full_refresh), + occupied_pages=occupied_pages, ) @@ -221,3 +231,201 @@ def emit(self, record: logging.LogRecord) -> None: msgs = [r.getMessage() for r in captured_records] assert any("my-bad-req" in m for m in msgs), msgs assert any("page copy would exceed req_to_page capacity" in m for m in msgs), msgs + + +def test_update_block_table_full_refresh_replaces_prefix_page(monkeypatch): + """full_refresh must copy repointed non-tail pages, not only append the tail. + + If a prefix page changes from private P0 to canonical CACHED and the freed + P0 is immediately reused as the new tail, a tail-only mirror update would + leave logical page 0 and logical page 2 pointing at the same physical page. + """ + from tokenspeed.runtime.execution import cache_loc_kernel + + req_to_page = torch.zeros(8, 513, dtype=torch.int32) + # req[0]: old mirror had [10, 11]. The scheduler authoritative row is + # [99, 11, 10]: prefix repointed to cached page 99, freed page 10 reused + # as the new tail. Copying only begin=2,new=[10] would leave [10, 11, 10]. + # req[1]: tail-only -> begin=200, only the 2-page delta. + forward_op = _make_forward_op( + begins=[2, 200], + sizes=[1, 2], + new_occupied_pages=[[10], [50, 51]], + occupied_pages=[[99, 11, 10], [40, 41, 42]], + full_refresh=[1, 0], + ) + + captured: dict = {} + + def fake_update_req_to_page( + req_to_page, + req_pool_indices, + new_occupied_pages, + new_occupied_pages_num, + pages_copy_starts, + ): + captured["num"] = new_occupied_pages_num.tolist() + captured["starts"] = pages_copy_starts.tolist() + captured["pages"] = new_occupied_pages.tolist() + + monkeypatch.setattr(cache_loc_kernel, "update_req_to_page", fake_update_req_to_page) + cache_loc_kernel.update_block_table( + forward_op, device="cpu", req_to_page=req_to_page + ) + + # req[0] refreshed the whole 3-page row from start 0; req[1] kept its tail + # delta. The flattened full-refresh row must be [99, 11, 10], not the + # tail-only aliasing update [10]. + assert captured["num"] == [3, 2] + assert captured["starts"] == [0, 200] + assert captured["pages"] == [99, 11, 10, 50, 51] + + +def test_update_block_table_full_refresh_respects_clamp(monkeypatch): + """A full_refresh row longer than max_pages must still be clamped, not crash.""" + from tokenspeed.runtime.execution import cache_loc_kernel + + req_to_page = torch.zeros(8, 4, dtype=torch.int32) + # Whole row has 6 pages but req_to_page only holds 4 → clamp to 4 at start 0. + forward_op = _make_forward_op( + begins=[3], + sizes=[1], + new_occupied_pages=[[999]], + occupied_pages=[[10, 11, 12, 13, 14, 15]], + full_refresh=[1], + ) + + captured: dict = {} + + def fake_update_req_to_page( + req_to_page, + req_pool_indices, + new_occupied_pages, + new_occupied_pages_num, + pages_copy_starts, + ): + captured["num"] = new_occupied_pages_num.tolist() + captured["starts"] = pages_copy_starts.tolist() + captured["pages"] = new_occupied_pages.tolist() + + monkeypatch.setattr(cache_loc_kernel, "update_req_to_page", fake_update_req_to_page) + cache_loc_kernel.update_block_table( + forward_op, device="cpu", req_to_page=req_to_page + ) + + # Clamped to the first 4 ids. Un-clamped this would be num=[6], + # pages=[10,11,12,13,14,15] -> a 2-column OOB write past the width-4 row. + assert captured["starts"] == [0] + assert captured["num"] == [4] + assert captured["pages"] == [10, 11, 12, 13] + + +def test_update_block_table_full_refresh_overrides_zero_size(monkeypatch): + """full_refresh=1 must copy the row even when the op's tail size is 0 + (early-return must be computed on the effective, post-branch sizes).""" + from tokenspeed.runtime.execution import cache_loc_kernel + + req_to_page = torch.zeros(8, 513, dtype=torch.int32) + forward_op = _make_forward_op( + begins=[5], + sizes=[0], + new_occupied_pages=[[]], + occupied_pages=[[10, 11, 12]], + full_refresh=[1], + ) + + captured: dict = {} + + def fake_update_req_to_page( + req_to_page, + req_pool_indices, + new_occupied_pages, + new_occupied_pages_num, + pages_copy_starts, + ): + captured["num"] = new_occupied_pages_num.tolist() + captured["starts"] = pages_copy_starts.tolist() + captured["pages"] = new_occupied_pages.tolist() + + monkeypatch.setattr(cache_loc_kernel, "update_req_to_page", fake_update_req_to_page) + cache_loc_kernel.update_block_table( + forward_op, device="cpu", req_to_page=req_to_page + ) + + assert captured["num"] == [3] + assert captured["starts"] == [0] + assert captured["pages"] == [10, 11, 12] + + +def test_update_block_table_full_refresh_decode_tail_repoints_prefix(monkeypatch): + """PrefillDone->Decode can append a decode tail while also repointing + completed prefix pages; full_refresh must make that a whole-row copy.""" + from tokenspeed.runtime.execution import cache_loc_kernel + + req_to_page = torch.zeros(8, 513, dtype=torch.int32) + forward_op = _make_forward_op( + begins=[3], + sizes=[1], + new_occupied_pages=[[99]], + occupied_pages=[[20, 21, 22, 99]], + full_refresh=[1], + ) + + captured: dict = {} + + def fake_update_req_to_page( + req_to_page, + req_pool_indices, + new_occupied_pages, + new_occupied_pages_num, + pages_copy_starts, + ): + captured["num"] = new_occupied_pages_num.tolist() + captured["starts"] = pages_copy_starts.tolist() + captured["pages"] = new_occupied_pages.tolist() + + monkeypatch.setattr(cache_loc_kernel, "update_req_to_page", fake_update_req_to_page) + cache_loc_kernel.update_block_table( + forward_op, device="cpu", req_to_page=req_to_page + ) + + assert captured["num"] == [4] + assert captured["starts"] == [0] + assert captured["pages"] == [20, 21, 22, 99] + + +def test_update_block_table_full_refresh_prefill_chunk_repoints_prefix(monkeypatch): + """Continuation prefill can publish prior chunk pages while adding a new + tail chunk; full_refresh must refresh the whole scheduler row.""" + from tokenspeed.runtime.execution import cache_loc_kernel + + req_to_page = torch.zeros(8, 513, dtype=torch.int32) + forward_op = _make_forward_op( + begins=[2], + sizes=[2], + new_occupied_pages=[[80, 81]], + occupied_pages=[[30, 31, 80, 81]], + full_refresh=[1], + ) + + captured: dict = {} + + def fake_update_req_to_page( + req_to_page, + req_pool_indices, + new_occupied_pages, + new_occupied_pages_num, + pages_copy_starts, + ): + captured["num"] = new_occupied_pages_num.tolist() + captured["starts"] = pages_copy_starts.tolist() + captured["pages"] = new_occupied_pages.tolist() + + monkeypatch.setattr(cache_loc_kernel, "update_req_to_page", fake_update_req_to_page) + cache_loc_kernel.update_block_table( + forward_op, device="cpu", req_to_page=req_to_page + ) + + assert captured["num"] == [4] + assert captured["starts"] == [0] + assert captured["pages"] == [30, 31, 80, 81] diff --git a/tokenspeed-scheduler/bindings/python_module.cpp b/tokenspeed-scheduler/bindings/python_module.cpp index c97d25513..abc60edaf 100644 --- a/tokenspeed-scheduler/bindings/python_module.cpp +++ b/tokenspeed-scheduler/bindings/python_module.cpp @@ -73,6 +73,9 @@ void BindForwardCommonFields(Cls& cls) { .def_prop_ro( "sizes", [](const Op& op) -> const std::vector& { return op.sizes; }, nb::rv_policy::reference_internal) + .def_prop_ro( + "full_refresh", [](const Op& op) -> const std::vector& { return op.full_refresh; }, + nb::rv_policy::reference_internal) .def_prop_ro( "new_occupied_pages", [](const Op& op) { diff --git a/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp b/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp index bf643e331..70d80edd2 100644 --- a/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp +++ b/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp @@ -715,6 +715,8 @@ PrefillOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::Sched auto match = event.GetMatchResult(); #endif auto op = applyPrefillEvent(request, event, FlatGroupIds()); + // First chunk adopts/rebuilds the prefix, so refresh the whole mirror row. + op.full_refresh = true; #if TOKENSPEED_FLAT_KVCACHE // Host-loaded pages ride the same LoadBackOperation channel as radix loadbacks. std::vector> load_pairs = event.TakeFlatLoadPairs(); @@ -757,6 +759,9 @@ PrefillOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::Sched PrefillOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::SchedulePrefillEvent event) { auto op = applyPrefillEvent(request, event, FlatGroupIds()); + // Continuation prefill may publish/adopt completed prior-chunk pages before + // scheduling the next tail chunk, so earlier logical pages can be repointed. + op.full_refresh = true; #if !TOKENSPEED_FLAT_KVCACHE if (hybrid_prefix_cache_) { hybrid_prefix_cache_->CommitChunk(op.request_id, const_cast(request->GetDeviceNode())); @@ -810,6 +815,11 @@ DecodeOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::Schedu #endif auto op = applyDecodeEvent(request, std::move(event), config_.decode_input_tokens, FlatGroupIds()); + if (came_from_prefill_done) { + // PrefillDone -> Decode may publish/adopt completed prefill pages before + // reserving the decode tail, so earlier logical pages can be repointed. + op.full_refresh = true; + } if (need_bootstrap_token) { op.decode_input_id = bootstrap_token; } @@ -853,6 +863,7 @@ DecodeOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::Schedu .occupied_pages = std::move(all_pages), .begin = 0, .size = sz, + .full_refresh = true, }}; op.decode_input_id = request->GetLastToken(); op.hist_token_len = request->TokenSize() - 1; diff --git a/tokenspeed-scheduler/csrc/scheduler/operations/forward.h b/tokenspeed-scheduler/csrc/scheduler/operations/forward.h index 293cdcc35..8ed7cffa3 100644 --- a/tokenspeed-scheduler/csrc/scheduler/operations/forward.h +++ b/tokenspeed-scheduler/csrc/scheduler/operations/forward.h @@ -41,6 +41,8 @@ struct ForwardOperationBase { std::int32_t begin; // Number of newly allocated pages (starting at occupied_pages[begin]). std::int32_t size; + // True when this op rebuilt the prefix: refresh the whole mirror row, not just the tail delta. + bool full_refresh{false}; std::int32_t prefill_length; @@ -91,6 +93,8 @@ struct FlatForwardOperation { std::vector> occupied_pages; std::vector begins; std::vector sizes; + // Per-request 0/1 flag mirroring ForwardOperationBase::full_refresh (1 = whole-row copy, 0 = tail delta). + std::vector full_refresh; std::vector input_ids; std::vector shifted_input_ids; @@ -133,6 +137,7 @@ struct FlatForwardOperation { occupied_pages.push_back(std::move(inner.occupied_pages)); begins.push_back(inner.begin); sizes.push_back(inner.size); + full_refresh.push_back(static_cast(inner.full_refresh)); mamba_working_indices.push_back(inner.mamba_working_idx); mamba_checkpoint_dst_indices.push_back(inner.mamba_checkpoint_dst_idx); mamba_cow_src_indices.push_back(inner.mamba_cow_src_idx);