From 6fa10b8a2b8ee5eb52106411ee1fccd5a7c782f6 Mon Sep 17 00:00:00 2001
From: Stanley Winata <stanley.winata@amd.com>
Date: Sun, 12 Jul 2026 05:14:29 +0000
Subject: [PATCH 1/4] feat(runtime): scheduler-driven full_refresh bit for
 page-table mirror

TL;DR: the req_to_page mirror was updated with only the appended tail, but some scheduler steps also repoint earlier non-tail pages, leaving stale or aliased row entries. Add a scheduler-produced full_refresh bit so the runtime can refresh the entire row when the authoritative scheduler row may have changed before the tail.

The bit is set in scheduler/operations/forward.cpp at transitions that can rebuild, publish, adopt, or recover prefix pages:

- SchedulePrefillFirstChunkEvent: the first chunk adopts/rebuilds the prefix from a fresh match, so the mirror must refresh the whole row rather than only the appended tail.

- SchedulePrefillEvent: continuation prefill can publish/adopt completed prior-chunk pages before scheduling the next tail chunk, so earlier logical pages may be repointed even though the op also has a tail delta.

- ScheduleDecodeEvent from PrefillDone: the transition can publish/adopt completed prefill pages before reserving the decode tail, so earlier logical pages may be repointed while the op still has a decode tail delta.

- ScheduleDecodeFromRetractedEvent: resuming a retracted request re-matches and rebuilds the row from scratch; prefix pages can land on different physical pages plus a fresh tail.

Pure decode continuation stays tail-only. Runtime update_block_table consumes full_refresh by replacing begin/size/new_occupied_pages with the scheduler's occupied_pages row from logical page 0, while preserving clamp behavior for oversized rows.

Add focused tests for full-row vs tail selection, continuation-prefill prefix repoint, decode-tail plus prefix repoint, clamp on an oversized full-refresh row, and full_refresh overriding a zero tail size.

Validation:

- python -m pytest -q test/runtime/test_update_block_table_capacity_clamp.py

- pre-commit run --all-files

Signed-off-by: Stanley Winata <stanley.winata@amd.com>
---
 .../runtime/execution/cache_loc_kernel.py     |  26 ++-
 ...2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml |   2 +-
 .../test_update_block_table_capacity_clamp.py | 202 ++++++++++++++++++
 .../bindings/python_module.cpp                |   3 +
 .../csrc/scheduler/operations/forward.cpp     |  11 +
 .../csrc/scheduler/operations/forward.h       |   5 +
 6 files changed, 241 insertions(+), 8 deletions(-)

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/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml b/test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml
index 66bb71f0d..a9222166d 100644
--- a/test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml
+++ b/test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml
@@ -51,4 +51,4 @@ eval:
     --generation-config '{"do_sample":false,"temperature":0.0,"max_tokens":8192}'
 report:
   github_step_summary: true
-score_threshold: 0.75
+score_threshold: 1.0
diff --git a/test/runtime/test_update_block_table_capacity_clamp.py b/test/runtime/test_update_block_table_capacity_clamp.py
index cc70b4a92..b741f8bbc 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,195 @@ 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_copies_whole_row(monkeypatch):
+    """A full_refresh=1 request must copy the scheduler's full occupied_pages
+    row from logical page 0, ignoring the begin/size tail delta; a full_refresh=0
+    request in the same batch keeps the tail-only delta."""
+    from tokenspeed.runtime.execution import cache_loc_kernel
+
+    req_to_page = torch.zeros(8, 513, dtype=torch.int32)
+    # req[0]: full refresh -> whole row from 0, delta fields ignored.
+    # req[1]: tail-only -> begin=200, only the 2-page delta.
+    forward_op = _make_forward_op(
+        begins=[100, 200],
+        sizes=[1, 2],
+        new_occupied_pages=[[999], [50, 51]],
+        occupied_pages=[[10, 11, 12, 13], [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 4-page row from start 0; req[1] kept its tail delta.
+    assert captured["num"] == [4, 2]
+    assert captured["starts"] == [0, 200]
+    # Flattened: whole row of req[0] then the tail delta of req[1].
+    assert captured["pages"] == [10, 11, 12, 13, 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<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) {
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<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;
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<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);

From 3a67b89e3cb4e3e61a9685a8d98fe5f5d3211c32 Mon Sep 17 00:00:00 2001
From: Stanley Winata <stanley.winata@amd.com>
Date: Sun, 12 Jul 2026 06:35:54 +0000
Subject: [PATCH 2/4] fix(runtime): refresh page-table mirror from scheduler
 rows

Always copy updated req_to_page rows from forward_op.occupied_pages starting at logical page 0 instead of applying begin/size tail deltas. The scheduler owns the authoritative page row, and several prefix-cache transitions can repoint non-tail pages while still reporting a tail update; refreshing the whole row prevents stale or aliased mirror entries.

Keep the existing capacity clamp for rows that exceed req_to_page width, but apply it to the authoritative row payload. Update the page-table tests to cover row-wide refresh, oversized-row clamping, and the zero-update early return.

Validation:

- python -m pytest -q test/runtime/test_update_block_table_capacity_clamp.py

- pre-commit run --all-files

Signed-off-by: Stanley Winata <stanley.winata@amd.com>
---
 .../runtime/execution/cache_loc_kernel.py     | 51 +++++-------
 .../test_update_block_table_capacity_clamp.py | 82 ++++++++-----------
 2 files changed, 56 insertions(+), 77 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/cache_loc_kernel.py b/python/tokenspeed/runtime/execution/cache_loc_kernel.py
index aee338d2c..67d2adb61 100644
--- a/python/tokenspeed/runtime/execution/cache_loc_kernel.py
+++ b/python/tokenspeed/runtime/execution/cache_loc_kernel.py
@@ -395,33 +395,22 @@ 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)
 
-    max_pages = req_to_page.shape[1]
-    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):
+    # If no row reports a page-table update, the existing device mirror is current.
+    if all(n == 0 for n in forward_op.sizes):
         return
 
+    max_pages = req_to_page.shape[1]
+    # The scheduler owns the authoritative logical page row. Refresh every
+    # updated row from occupied_pages so non-tail page repoints cannot leave
+    # stale or aliased entries in the Python/device mirror.
+    pages_copy_starts = [0 for _ in forward_op.begins]
+    pages_to_copy = [list(row) for row in forward_op.occupied_pages]
+    pages_copy_sizes = [len(row) for row in pages_to_copy]
+    request_ids = list(forward_op.request_ids)
     # 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)):
+    for i, (begin, size) in enumerate(zip(pages_copy_starts, pages_copy_sizes)):
         if begin + size > max_pages:
             clamped = max(0, max_pages - begin)
             logger.warning(
@@ -437,21 +426,21 @@ def flatten_and_to_device(data, dtype=torch.int32):
                 max_pages,
                 clamped,
             )
-            sizes[i] = clamped
-            # Keep new_occupied_pages[i] consistent with the clamped size so
+            pages_copy_sizes[i] = clamped
+            # Keep pages_to_copy[i] consistent with the clamped size so
             # the kernel's cumsum-based offsets stay aligned across the batch.
-            new_occupied_pages[i] = new_occupied_pages[i][:clamped]
+            pages_to_copy[i] = pages_to_copy[i][:clamped]
 
-    new_occupied_pages_num = flatten_and_to_device(sizes, dtype=torch.int32)
-    pages_copy_starts = flatten_and_to_device(begins, dtype=torch.int32)
-    new_occupied_pages_t = flatten_and_to_device(new_occupied_pages, dtype=torch.int32)
+    pages_to_copy_num = flatten_and_to_device(pages_copy_sizes, dtype=torch.int32)
+    pages_copy_starts_t = flatten_and_to_device(pages_copy_starts, dtype=torch.int32)
+    pages_to_copy_t = flatten_and_to_device(pages_to_copy, dtype=torch.int32)
     request_pool_indices = flatten_and_to_device(
         forward_op.request_pool_indices, dtype=torch.int64
     )
     update_req_to_page(
         req_to_page=req_to_page,
         req_pool_indices=request_pool_indices,
-        new_occupied_pages=new_occupied_pages_t,
-        new_occupied_pages_num=new_occupied_pages_num,
-        pages_copy_starts=pages_copy_starts,
+        new_occupied_pages=pages_to_copy_t,
+        new_occupied_pages_num=pages_to_copy_num,
+        pages_copy_starts=pages_copy_starts_t,
     )
diff --git a/test/runtime/test_update_block_table_capacity_clamp.py b/test/runtime/test_update_block_table_capacity_clamp.py
index b741f8bbc..14ce84f94 100644
--- a/test/runtime/test_update_block_table_capacity_clamp.py
+++ b/test/runtime/test_update_block_table_capacity_clamp.py
@@ -52,8 +52,8 @@ def _make_forward_op(
     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.
+        # Keep a well-shaped authoritative row for tests that do not care
+        # about row-wide refresh semantics.
         occupied_pages = [list(row) for row in new_occupied_pages]
     return SimpleNamespace(
         begins=list(begins),
@@ -74,12 +74,13 @@ def test_update_block_table_does_not_raise_on_overflow(monkeypatch):
     """
     from tokenspeed.runtime.execution import cache_loc_kernel
 
-    # max_pages=513 (the value from the real crash). req[1] is the offender:
-    # begin=513 + size=1 = 514 > 513.
+    # max_pages=513 (the value from the real crash). req[1]'s authoritative
+    # scheduler row is the offender: len(row)=514 > 513.
     req_to_page = torch.zeros(8, 513, dtype=torch.int32)
     forward_op = _make_forward_op(
         begins=[400, 513, 100],
         sizes=[2, 1, 3],
+        occupied_pages=[list(range(2)), list(range(514)), list(range(3))],
     )
 
     captured: dict = {}
@@ -102,14 +103,14 @@ def fake_update_req_to_page(
         forward_op, device="cpu", req_to_page=req_to_page
     )
 
-    # The offender (req[1]) got clamped to 0, others unchanged.
-    assert captured["num"] == [2, 0, 3]
-    # begins are unchanged.
-    assert captured["starts"] == [400, 513, 100]
-    # And the flattened pages array dropped the offending request's entry,
+    # The offender (req[1]) got clamped to max row width, others unchanged.
+    assert captured["num"] == [2, 513, 3]
+    # Whole-row refresh always starts at logical page 0.
+    assert captured["starts"] == [0, 0, 0]
+    # And the flattened pages array trimmed the offending request's row,
     # so the cumsum-based offsets the kernel uses stay consistent.
-    # req[0] contributes 2 pages, req[1] contributes 0, req[2] contributes 3 → 5 total.
-    assert len(captured["pages"]) == 5
+    # req[0] contributes 2 pages, req[1] contributes 513, req[2] contributes 3.
+    assert len(captured["pages"]) == 518
 
 
 def test_update_block_table_passthrough_when_no_overflow(monkeypatch):
@@ -142,16 +143,17 @@ def fake_update_req_to_page(
     assert captured["num"] == [1, 2, 1]
 
 
-def test_update_block_table_clamp_partial_overflow(monkeypatch):
-    """If begin < max_pages and begin+size > max_pages, clamp to (max - begin)."""
+def test_update_block_table_clamp_oversized_row(monkeypatch):
+    """If the authoritative scheduler row is wider than req_to_page, clamp it."""
     from tokenspeed.runtime.execution import cache_loc_kernel
 
-    req_to_page = torch.zeros(8, 513, dtype=torch.int32)
-    # begin=512 + size=4 = 516 > 513, but begin < 513 → clamp to size=1.
+    req_to_page = torch.zeros(8, 3, dtype=torch.int32)
+    # Whole-row refresh starts at 0, so a 4-page row clamps to width 3.
     forward_op = _make_forward_op(
         begins=[512],
         sizes=[4],
         new_occupied_pages=[[700, 701, 702, 703]],
+        occupied_pages=[[700, 701, 702, 703]],
     )
 
     captured: dict = {}
@@ -171,9 +173,9 @@ def fake_update_req_to_page(
         forward_op, device="cpu", req_to_page=req_to_page
     )
 
-    assert captured["num"] == [1]
-    # Only the first page from new_occupied_pages survives (the one that fits).
-    assert captured["pages"] == [700]
+    assert captured["num"] == [3]
+    # Only the first three pages from the scheduler row fit.
+    assert captured["pages"] == [700, 701, 702]
 
 
 def test_update_block_table_zero_total_returns_early(monkeypatch):
@@ -218,6 +220,7 @@ def emit(self, record: logging.LogRecord) -> None:
             begins=[513],
             sizes=[1],
             request_ids=["my-bad-req"],
+            occupied_pages=[list(range(514))],
         )
         with mock.patch.object(
             cache_loc_kernel, "update_req_to_page", lambda **kw: None
@@ -233,15 +236,13 @@ def emit(self, record: logging.LogRecord) -> None:
     assert any("page copy would exceed req_to_page capacity" in m for m in msgs), msgs
 
 
-def test_update_block_table_full_refresh_copies_whole_row(monkeypatch):
-    """A full_refresh=1 request must copy the scheduler's full occupied_pages
-    row from logical page 0, ignoring the begin/size tail delta; a full_refresh=0
-    request in the same batch keeps the tail-only delta."""
+def test_update_block_table_copies_whole_row_for_updated_requests(monkeypatch):
+    """Every updated request copies the scheduler's full occupied_pages row
+    from logical page 0, ignoring the begin/size tail delta."""
     from tokenspeed.runtime.execution import cache_loc_kernel
 
     req_to_page = torch.zeros(8, 513, dtype=torch.int32)
-    # req[0]: full refresh -> whole row from 0, delta fields ignored.
-    # req[1]: tail-only -> begin=200, only the 2-page delta.
+    # Both rows refresh from occupied_pages, regardless of the full_refresh bit.
     forward_op = _make_forward_op(
         begins=[100, 200],
         sizes=[1, 2],
@@ -268,11 +269,9 @@ def fake_update_req_to_page(
         forward_op, device="cpu", req_to_page=req_to_page
     )
 
-    # req[0] refreshed the whole 4-page row from start 0; req[1] kept its tail delta.
-    assert captured["num"] == [4, 2]
-    assert captured["starts"] == [0, 200]
-    # Flattened: whole row of req[0] then the tail delta of req[1].
-    assert captured["pages"] == [10, 11, 12, 13, 50, 51]
+    assert captured["num"] == [4, 3]
+    assert captured["starts"] == [0, 0]
+    assert captured["pages"] == [10, 11, 12, 13, 40, 41, 42]
 
 
 def test_update_block_table_full_refresh_respects_clamp(monkeypatch):
@@ -314,9 +313,10 @@ def fake_update_req_to_page(
     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)."""
+def test_update_block_table_zero_size_returns_early_even_with_occupied_pages(
+    monkeypatch,
+):
+    """If the scheduler reports no updated rows, keep the existing mirror."""
     from tokenspeed.runtime.execution import cache_loc_kernel
 
     req_to_page = torch.zeros(8, 513, dtype=torch.int32)
@@ -328,27 +328,17 @@ def test_update_block_table_full_refresh_overrides_zero_size(monkeypatch):
         full_refresh=[1],
     )
 
-    captured: dict = {}
+    called = {"v": False}
 
-    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()
+    def fake_update_req_to_page(**kwargs):
+        called["v"] = True
 
     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]
+    assert called["v"] is False
 
 
 def test_update_block_table_full_refresh_decode_tail_repoints_prefix(monkeypatch):

From 048e5b1d6658302943f6587e60cc13a5685fcb0d Mon Sep 17 00:00:00 2001
From: Stanley Winata <stanley.winata@amd.com>
Date: Sun, 12 Jul 2026 07:13:29 +0000
Subject: [PATCH 3/4] Revert "fix(runtime): refresh page-table mirror from
 scheduler rows"

This reverts commit 2dd87c97f80aec71e3f9c5f37dd269a4267728e1.

The always-refresh experiment did not improve local Kimi K2.5 MXFP4 AIME25 accuracy: both scheduler-driven full_refresh and row-wide refresh scored 0.75 on limit=4. Keep the scheduler-produced full_refresh mechanism as the targeted fix and avoid the broader per-step row refresh until we have a stronger signal.

Restore the Kimi K2.5 MXFP4 AIME smoke threshold to 0.75, matching the observed stable score while this mirror fix is validated separately.

Validation:

- python -m pytest -q test/runtime/test_update_block_table_capacity_clamp.py

- pre-commit run --all-files

Signed-off-by: Stanley Winata <stanley.winata@amd.com>
---
 .../runtime/execution/cache_loc_kernel.py     | 51 +++++++-----
 ...2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml |  2 +-
 .../test_update_block_table_capacity_clamp.py | 82 +++++++++++--------
 3 files changed, 78 insertions(+), 57 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/cache_loc_kernel.py b/python/tokenspeed/runtime/execution/cache_loc_kernel.py
index 67d2adb61..aee338d2c 100644
--- a/python/tokenspeed/runtime/execution/cache_loc_kernel.py
+++ b/python/tokenspeed/runtime/execution/cache_loc_kernel.py
@@ -395,22 +395,33 @@ 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)
 
-    # If no row reports a page-table update, the existing device mirror is current.
-    if all(n == 0 for n in forward_op.sizes):
-        return
-
     max_pages = req_to_page.shape[1]
-    # The scheduler owns the authoritative logical page row. Refresh every
-    # updated row from occupied_pages so non-tail page repoints cannot leave
-    # stale or aliased entries in the Python/device mirror.
-    pages_copy_starts = [0 for _ in forward_op.begins]
-    pages_to_copy = [list(row) for row in forward_op.occupied_pages]
-    pages_copy_sizes = [len(row) for row in pages_to_copy]
+    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(pages_copy_starts, pages_copy_sizes)):
+    for i, (begin, size) in enumerate(zip(begins, sizes)):
         if begin + size > max_pages:
             clamped = max(0, max_pages - begin)
             logger.warning(
@@ -426,21 +437,21 @@ def flatten_and_to_device(data, dtype=torch.int32):
                 max_pages,
                 clamped,
             )
-            pages_copy_sizes[i] = clamped
-            # Keep pages_to_copy[i] consistent with the clamped size so
+            sizes[i] = clamped
+            # Keep new_occupied_pages[i] consistent with the clamped size so
             # the kernel's cumsum-based offsets stay aligned across the batch.
-            pages_to_copy[i] = pages_to_copy[i][:clamped]
+            new_occupied_pages[i] = new_occupied_pages[i][:clamped]
 
-    pages_to_copy_num = flatten_and_to_device(pages_copy_sizes, dtype=torch.int32)
-    pages_copy_starts_t = flatten_and_to_device(pages_copy_starts, dtype=torch.int32)
-    pages_to_copy_t = flatten_and_to_device(pages_to_copy, dtype=torch.int32)
+    new_occupied_pages_num = flatten_and_to_device(sizes, dtype=torch.int32)
+    pages_copy_starts = flatten_and_to_device(begins, dtype=torch.int32)
+    new_occupied_pages_t = flatten_and_to_device(new_occupied_pages, dtype=torch.int32)
     request_pool_indices = flatten_and_to_device(
         forward_op.request_pool_indices, dtype=torch.int64
     )
     update_req_to_page(
         req_to_page=req_to_page,
         req_pool_indices=request_pool_indices,
-        new_occupied_pages=pages_to_copy_t,
-        new_occupied_pages_num=pages_to_copy_num,
-        pages_copy_starts=pages_copy_starts_t,
+        new_occupied_pages=new_occupied_pages_t,
+        new_occupied_pages_num=new_occupied_pages_num,
+        pages_copy_starts=pages_copy_starts,
     )
diff --git a/test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml b/test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml
index a9222166d..66bb71f0d 100644
--- a/test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml
+++ b/test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml
@@ -51,4 +51,4 @@ eval:
     --generation-config '{"do_sample":false,"temperature":0.0,"max_tokens":8192}'
 report:
   github_step_summary: true
-score_threshold: 1.0
+score_threshold: 0.75
diff --git a/test/runtime/test_update_block_table_capacity_clamp.py b/test/runtime/test_update_block_table_capacity_clamp.py
index 14ce84f94..b741f8bbc 100644
--- a/test/runtime/test_update_block_table_capacity_clamp.py
+++ b/test/runtime/test_update_block_table_capacity_clamp.py
@@ -52,8 +52,8 @@ def _make_forward_op(
     if full_refresh is None:
         full_refresh = [0 for _ in begins]
     if occupied_pages is None:
-        # Keep a well-shaped authoritative row for tests that do not care
-        # about row-wide refresh semantics.
+        # 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),
@@ -74,13 +74,12 @@ def test_update_block_table_does_not_raise_on_overflow(monkeypatch):
     """
     from tokenspeed.runtime.execution import cache_loc_kernel
 
-    # max_pages=513 (the value from the real crash). req[1]'s authoritative
-    # scheduler row is the offender: len(row)=514 > 513.
+    # max_pages=513 (the value from the real crash). req[1] is the offender:
+    # begin=513 + size=1 = 514 > 513.
     req_to_page = torch.zeros(8, 513, dtype=torch.int32)
     forward_op = _make_forward_op(
         begins=[400, 513, 100],
         sizes=[2, 1, 3],
-        occupied_pages=[list(range(2)), list(range(514)), list(range(3))],
     )
 
     captured: dict = {}
@@ -103,14 +102,14 @@ def fake_update_req_to_page(
         forward_op, device="cpu", req_to_page=req_to_page
     )
 
-    # The offender (req[1]) got clamped to max row width, others unchanged.
-    assert captured["num"] == [2, 513, 3]
-    # Whole-row refresh always starts at logical page 0.
-    assert captured["starts"] == [0, 0, 0]
-    # And the flattened pages array trimmed the offending request's row,
+    # The offender (req[1]) got clamped to 0, others unchanged.
+    assert captured["num"] == [2, 0, 3]
+    # begins are unchanged.
+    assert captured["starts"] == [400, 513, 100]
+    # And the flattened pages array dropped the offending request's entry,
     # so the cumsum-based offsets the kernel uses stay consistent.
-    # req[0] contributes 2 pages, req[1] contributes 513, req[2] contributes 3.
-    assert len(captured["pages"]) == 518
+    # req[0] contributes 2 pages, req[1] contributes 0, req[2] contributes 3 → 5 total.
+    assert len(captured["pages"]) == 5
 
 
 def test_update_block_table_passthrough_when_no_overflow(monkeypatch):
@@ -143,17 +142,16 @@ def fake_update_req_to_page(
     assert captured["num"] == [1, 2, 1]
 
 
-def test_update_block_table_clamp_oversized_row(monkeypatch):
-    """If the authoritative scheduler row is wider than req_to_page, clamp it."""
+def test_update_block_table_clamp_partial_overflow(monkeypatch):
+    """If begin < max_pages and begin+size > max_pages, clamp to (max - begin)."""
     from tokenspeed.runtime.execution import cache_loc_kernel
 
-    req_to_page = torch.zeros(8, 3, dtype=torch.int32)
-    # Whole-row refresh starts at 0, so a 4-page row clamps to width 3.
+    req_to_page = torch.zeros(8, 513, dtype=torch.int32)
+    # begin=512 + size=4 = 516 > 513, but begin < 513 → clamp to size=1.
     forward_op = _make_forward_op(
         begins=[512],
         sizes=[4],
         new_occupied_pages=[[700, 701, 702, 703]],
-        occupied_pages=[[700, 701, 702, 703]],
     )
 
     captured: dict = {}
@@ -173,9 +171,9 @@ def fake_update_req_to_page(
         forward_op, device="cpu", req_to_page=req_to_page
     )
 
-    assert captured["num"] == [3]
-    # Only the first three pages from the scheduler row fit.
-    assert captured["pages"] == [700, 701, 702]
+    assert captured["num"] == [1]
+    # Only the first page from new_occupied_pages survives (the one that fits).
+    assert captured["pages"] == [700]
 
 
 def test_update_block_table_zero_total_returns_early(monkeypatch):
@@ -220,7 +218,6 @@ def emit(self, record: logging.LogRecord) -> None:
             begins=[513],
             sizes=[1],
             request_ids=["my-bad-req"],
-            occupied_pages=[list(range(514))],
         )
         with mock.patch.object(
             cache_loc_kernel, "update_req_to_page", lambda **kw: None
@@ -236,13 +233,15 @@ def emit(self, record: logging.LogRecord) -> None:
     assert any("page copy would exceed req_to_page capacity" in m for m in msgs), msgs
 
 
-def test_update_block_table_copies_whole_row_for_updated_requests(monkeypatch):
-    """Every updated request copies the scheduler's full occupied_pages row
-    from logical page 0, ignoring the begin/size tail delta."""
+def test_update_block_table_full_refresh_copies_whole_row(monkeypatch):
+    """A full_refresh=1 request must copy the scheduler's full occupied_pages
+    row from logical page 0, ignoring the begin/size tail delta; a full_refresh=0
+    request in the same batch keeps the tail-only delta."""
     from tokenspeed.runtime.execution import cache_loc_kernel
 
     req_to_page = torch.zeros(8, 513, dtype=torch.int32)
-    # Both rows refresh from occupied_pages, regardless of the full_refresh bit.
+    # req[0]: full refresh -> whole row from 0, delta fields ignored.
+    # req[1]: tail-only -> begin=200, only the 2-page delta.
     forward_op = _make_forward_op(
         begins=[100, 200],
         sizes=[1, 2],
@@ -269,9 +268,11 @@ def fake_update_req_to_page(
         forward_op, device="cpu", req_to_page=req_to_page
     )
 
-    assert captured["num"] == [4, 3]
-    assert captured["starts"] == [0, 0]
-    assert captured["pages"] == [10, 11, 12, 13, 40, 41, 42]
+    # req[0] refreshed the whole 4-page row from start 0; req[1] kept its tail delta.
+    assert captured["num"] == [4, 2]
+    assert captured["starts"] == [0, 200]
+    # Flattened: whole row of req[0] then the tail delta of req[1].
+    assert captured["pages"] == [10, 11, 12, 13, 50, 51]
 
 
 def test_update_block_table_full_refresh_respects_clamp(monkeypatch):
@@ -313,10 +314,9 @@ def fake_update_req_to_page(
     assert captured["pages"] == [10, 11, 12, 13]
 
 
-def test_update_block_table_zero_size_returns_early_even_with_occupied_pages(
-    monkeypatch,
-):
-    """If the scheduler reports no updated rows, keep the existing mirror."""
+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)
@@ -328,17 +328,27 @@ def test_update_block_table_zero_size_returns_early_even_with_occupied_pages(
         full_refresh=[1],
     )
 
-    called = {"v": False}
+    captured: dict = {}
 
-    def fake_update_req_to_page(**kwargs):
-        called["v"] = True
+    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 called["v"] is False
+    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):

From 959056de523befb9cfac2f18774387e4eb364e4d Mon Sep 17 00:00:00 2001
From: Stanley Winata <stanley.winata@amd.com>
Date: Sun, 12 Jul 2026 16:11:14 +0000
Subject: [PATCH 4/4] test(runtime): clarify full-refresh page alias regression

Update the full-refresh block-table test to encode the concrete corruption mode: a non-tail prefix page is repointed to a canonical cached page while the freed physical page is reused as the new tail. A tail-only mirror update would leave logical page 0 and logical page 2 aliasing the same physical page.

Keep the test CPU-only by asserting the update_block_table contract passed to the page-table kernel: full_refresh rows copy the scheduler's occupied_pages from logical page 0, while ordinary rows still use their tail delta.

Validation:

- python -m pytest -q test/runtime/test_update_block_table_capacity_clamp.py

- pre-commit run --all-files

Signed-off-by: Stanley Winata <stanley.winata@amd.com>
---
 .../test_update_block_table_capacity_clamp.py | 30 +++++++++++--------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/test/runtime/test_update_block_table_capacity_clamp.py b/test/runtime/test_update_block_table_capacity_clamp.py
index b741f8bbc..15972aa29 100644
--- a/test/runtime/test_update_block_table_capacity_clamp.py
+++ b/test/runtime/test_update_block_table_capacity_clamp.py
@@ -233,20 +233,25 @@ def emit(self, record: logging.LogRecord) -> None:
     assert any("page copy would exceed req_to_page capacity" in m for m in msgs), msgs
 
 
-def test_update_block_table_full_refresh_copies_whole_row(monkeypatch):
-    """A full_refresh=1 request must copy the scheduler's full occupied_pages
-    row from logical page 0, ignoring the begin/size tail delta; a full_refresh=0
-    request in the same batch keeps the tail-only delta."""
+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]: full refresh -> whole row from 0, delta fields ignored.
+    # 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=[100, 200],
+        begins=[2, 200],
         sizes=[1, 2],
-        new_occupied_pages=[[999], [50, 51]],
-        occupied_pages=[[10, 11, 12, 13], [40, 41, 42]],
+        new_occupied_pages=[[10], [50, 51]],
+        occupied_pages=[[99, 11, 10], [40, 41, 42]],
         full_refresh=[1, 0],
     )
 
@@ -268,11 +273,12 @@ def fake_update_req_to_page(
         forward_op, device="cpu", req_to_page=req_to_page
     )
 
-    # req[0] refreshed the whole 4-page row from start 0; req[1] kept its tail delta.
-    assert captured["num"] == [4, 2]
+    # 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]
-    # Flattened: whole row of req[0] then the tail delta of req[1].
-    assert captured["pages"] == [10, 11, 12, 13, 50, 51]
+    assert captured["pages"] == [99, 11, 10, 50, 51]
 
 
 def test_update_block_table_full_refresh_respects_clamp(monkeypatch):