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

[raikonenfnu]

Summary

TL;DR: req_to_page is a runtime mirror of the scheduler-owned logical page row. Updating that mirror with only the appended tail is unsafe when a scheduler transition also republishes, adopts, or rebuilds earlier non-tail pages. In those cases, the old mirror can keep stale physical page IDs and can even alias a recycled page into multiple logical slots.

This PR adds a scheduler-produced full_refresh bit. When the bit is set, runtime refreshes that request's whole req_to_page row from forward_op.occupied_pages starting at logical page 0. When the bit is not set, runtime keeps the cheaper tail-delta update.

The bit is set in tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp for transitions that can change non-tail page identity:

• applyEventAndGenerateOp(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.

• applyEventAndGenerateOp(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.

• applyEventAndGenerateOp(ScheduleDecodeEvent) when the request came 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 looks like a decode-tail update.

• applyEventAndGenerateOp(ScheduleDecodeFromRetractedEvent): resuming a retracted request re-matches and rebuilds the row from scratch; prefix pages are re-derived and can land on different physical pages, plus a fresh tail.

Pure Decoding -> Decoding remains tail-only and keeps full_refresh=false.

The runtime side consumes full_refresh before its zero-size early return, so a full-refresh row is honored even if no new tail page was allocated. The existing capacity clamp is preserved for rows that exceed req_to_page width.

The Kimi K2.5 MXFP4 AIME smoke threshold remains 0.75; local validation showed the scheduler-driven mirror fix did not raise that limit-4 score to 1.0, so the PR should not turn that smoke into a perfect-score gate.

Test Plan

• pytest -q test/runtime/test_update_block_table_capacity_clamp.py
• pre-commit run --all-files

[raikonenfnu]

A fully concrete example (real sentence, real tokens)

Let's make it painfully literal. Use a tiny page_size = 4 (4 tokens per page) so the boundaries are easy to see.

System prompt (shared by everyone):

"You are a helpful assistant."

Tokenized (pretend each word/punct is one token):

pos:    0     1     2     3        4          5
tok:  "You" "are"  "a" "helpful" "assistant"  "."

Request A = system prompt + "Hi":

pos:    0     1     2     3        4          5      6
tok:  "You" "are"  "a" "helpful" "assistant"  "."   "Hi"

With page_size = 4:

logical page 0 = positions 0..3   -> "You are a helpful"
logical page 1 = positions 4..7   -> "assistant . Hi <next>"

Timeline

Step 1 — Request A runs first (cold). It computes its own KV into private pages:

A's mirror:   logical 0 -> P0     ("You are a helpful")
              logical 1 -> P1     ("assistant . Hi ...")

At this point page 0's data ("You are a helpful") is written into physical box P0. Append-only holds: those bytes never change.

Step 2 — scheduler publishes the shared prefix. Because "You are a helpful" is a clean, full page that other requests will reuse, the radix tree adopts a canonical copy of it:

radix tree:   node["You are a helpful"] -> physical page CACHED

Step 3 — Request B arrives. Same system prompt, different question:

"You are a helpful assistant. What is 2+2?"

B's logical page 0 is also "You are a helpful" — identical tokens → identical hash → prefix hit. B does not recompute or allocate for page 0; it just borrows CACHED:

B's mirror:   logical 0 -> CACHED   (borrowed, free)
              logical 1 -> P5       (B's own tail: "assistant . What is")
              ...

Step 4 — Request A gets retracted then resumed (memory pressure kicked it out, now it's back). On resume, A re-matches the tree. The scheduler sees page 0 is already the canonical CACHED and says: drop your private duplicate P0, point at CACHED.

A's logical 0 sticky note flips:   P0  ->  CACHED
A's private page P0 is freed  -> returned to the pool

This is the non-tail repoint: logical page 0 (NOT the tail) changed which physical box it points to, from P0 to CACHED. The tokens are still "You are a helpful" — only the physical copy changed.

Step 5 — A appends its next token (say it generates "there" at position 7, opening a new page 2). The allocator needs a fresh physical page and hands back... P0, because we just freed it in step 4:

A's authoritative row now:   logical 0 -> CACHED
                             logical 1 -> P1
                             logical 2 -> P0      (recycled!)

The bug, in tokens

The scheduler knows the correct row is [CACHED, P1, P0]. But its delta message only mentions the appended tail: "begin=2, wrote page P0." A mirror that trusts only the tail keeps its old page-0 note:

STALE mirror (append-only):   logical 0 -> P0     <- never updated from step 4
                              logical 1 -> P1
                              logical 2 -> P0     <- the recycled tail

Now logical 0 and logical 2 both point at physical box P0. So:

tokens 0..3   ("You are a helpful")   -> physical P0
tokens 8..11  (A's newest tokens)     -> physical P0   <- SAME BOX

When the attention kernel reads A's newest tokens' KV, it instead reads "You are a helpful"'s KV. Request A's generation silently drifts — even though sampling is greedy and the prompt never changed. That's exactly the deterministic-divergence symptom in the RFC.

The fix:

When required(based on full_refresh flag of the request), refresh/copy entire row ([CACHED, P1, P0]) into the mirror from logical page 0, not just the tail.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 404f73dae5

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 91ae189b32

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2dd87c97f8

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".