feat(runtime): scheduler-driven full_refresh bit for page-table mirror

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)

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]

tokenspeed-scheduler/bindings/python_module.cpp

@@ -73,6 +73,9 @@ void BindForwardCommonFields(Cls& cls) {
         .def_prop_ro(
             "sizes", [](const Op& op) -> const std::vector<std::int32_t>& { return op.sizes; },
             nb::rv_policy::reference_internal)
+        .def_prop_ro(
+            "full_refresh", [](const Op& op) -> const std::vector<std::int32_t>& { return op.full_refresh; },
+            nb::rv_policy::reference_internal)
         .def_prop_ro(
             "new_occupied_pages",
             [](const Op& op) {

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<std::pair<CacheBlock*, CacheBlock*>> 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<TreeNode*>(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;

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<std::vector<std::int32_t>> occupied_pages;
     std::vector<std::int32_t> begins;
     std::vector<std::int32_t> sizes;
+    // Per-request 0/1 flag mirroring ForwardOperationBase::full_refresh (1 = whole-row copy, 0 = tail delta).
+    std::vector<std::int32_t> full_refresh;
 
     std::vector<std::int32_t> input_ids;
     std::vector<std::int32_t> 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<std::int32_t>(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);