Repository navigation
fix: preserve both private draft pages in shifted prefix checkpoints - #165
Merged
lukealonso merged 1 commit intoNov 5, 2026
Merged
Conversation
A draft page table shifted by one physical page can have a complete draft page that is still private at the target prefix endpoint. Snapshotting only the partial page loses that history after restore, including at consumed=257 with 256-token pages. Preserve the preceding boundary page and the current partial page, and validate conflicts across both parts before copying. Expand GLM continuation storage and its planning costs accordingly. Qwen's existing custom checkpoint path already preserves both boundary pages; other shared paged recipes do not declare trailing draft pages. Update regression fixtures to use the native shifted table and shared prefix pages, and repair the existing single-rank maintenance fixture's completion-output wrapper. Validation on RTX PRO 4500 SM120: the original GPU reproducer lost both draft planes without a transfer error; the fix restores exact bytes. Transfer, maintenance and pool tests: 36 passed, 2 tensor-parallel tests skipped. Related CPU tests: 19 passed, 1 GPU test skipped. Full GLM serving and TP2/TP4 were not run. Existing exported packages remain loadable; re-export to obtain the fix. Design-Signed-Off-By: Codex (GPT-6); design-review=6aeb Signed-off-by: Maxrio9 <70799635+Maxrio9@users.noreply.github.com>
original-el8
added
area:caching
This was referenced Nov 12, 2026
yatesdr
added a commit
that referenced
this pull request
Nov 12, 2026
GLM-5.3 cores with MTP drafts, exported from main, serve their first request and then fail every prefix-cache restore with "prepared model reported device error 72". Each GLM state row saves the multi-token prediction draft cache as a two-page region, because #165 preserves both private draft pages across a shifted page boundary. PreparedStateTransfer groups these trailing transfers by a byte size that _copy_pages uses as PAGE_BYTES in three places: the stride between cache pages, the offset of the second saved page, and the copy length. All three must be one cache page. #165 grouped by cache.shape[1], but merge 2b38bf9 resolved a conflict on the adjacent line by restoring the older region.nbytes key. That key is two pages, so each snapshot copy now reads the wrong cache page, starts one full region too far, and writes twice the page length. In the GLM row layout this overwrites last_hidden and consumed_tokens. Snapshot copies are not validated, so publication succeeds. The next restore finds a saved consumed count that differs from its endpoint, fails validation, and maintenance rejects the command on every rank. PreparedStateTransfer groups trailing pages by cache.shape[1] again. That is equal to region.nbytes // 2, which the preceding check requires. A CPU test builds the prepared transfer under FakeTensorMode and checks that the trailing copy specializes on one 128-byte page rather than the 256-byte region. Prefix-cache restores of GLM-5.3 MTP state succeed again, and snapshot copies no longer write beyond the draft region. Only GLM-5.3 passes trailing pages; other recipes are unaffected. Cores exported from main before this change keep the oversized copy in their captured graphs and must be re-exported. Evidence, on 2x RTX PRO 6000 (SM120), GLM-5.3-Flash TP2, concurrency 8, prefill chunk 4096, 5 MTP drafts, speculative capacity 48: - Exported at f7c468c (before the merge): 3.1K tokens, then 3.6K, 4.1K, 6K and 8K reusing the cache, all succeed. Exported at 5fb495d: the first request succeeds and every later request fails with error 72. - Exported at 0572fd1 with this change and the separate block-FP8 tracing fix, which main needs to export GLM-5.3: two passes of 3.1K, 3.6K, 4.1K, 6K, 8K and 128K prompts succeed. The second pass reuses the complete prompt each time (128K time to first token 0.16 s against 15.8 s cold). No device errors. A greedy factual prompt returns the expected answer. - tests/recipes/test_glm53_state_transfer_cuda.py and tests/components/test_state_transfer.py on one SM120 GPU: 19 passed with this change; 14 failed without it. Another 19 tests in those files need the native qualification Cargo feature, which the test build lacked, and did not run either way. - The new CPU test fails without this change and passes with it. Not tested: SM103/SM121, multi-host TP, and concurrent requests. Design-Signed-Off-By: yatesdr; design-review=ed22
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replaces #16: rebased onto main
6fa01a49and the commit now carries the DCO sign-off CONTRIBUTING.md requires; the change is otherwise unchanged. Re-validated on this tree on an RTX PRO 4500 (SM120):tests/recipes/test_glm53_state_transfer_cuda.pyandtest_glm53_maintenance_cuda.pywithFLARE_RUN_CUDA_TESTS=1and thequalificationCargo feature offlare-heliopausegive 32 passed, 2 skipped (the two tensor-parallel cases). Control: current main sources with this PR's tests fail attest_mtp_trailing_partial_page_restores_at_target_page_boundaries[257](stopped at the first failure,-x), so the defect is still present on main.Problem (demonstrated defect)
draft_page_tableswithpage_offset = 1, so Helios packspages[1 + i]and draftlogical page j lives in physical request page j + 1.
floor(E / P)physical pages, and the draft cursor lags by one token.PreparedStateTransfersaved only draft logical page(E-1)//P, and only when(E-1) % P != 0. With E = k·P + r:A restore at those endpoints silently loses draft K/V history. Target outputs stay exact, but drafter acceptance
drops, which breaks the complete-speculative-state guarantee. The existing GLM test built
draft_tables = tables.clone()without the shift, so it couldn't catch this. The qwen38 (Flash-Next) custom checkpoint path alreadysaves both boundary pages (
DRAFT_PART 0/1).Reproduced on GPU with the GLM pool (256-token pages): snapshot at endpoint 257, then restore, lost
draft.latent_pageanddraft.index_pagebytes with no transfer error.Change
across both parts, and apply the same predicates as qwen38.
1, 2, 4, 255, 256, 257, 258, 511, 512, 513 plus cross-part alias cases.
mismatch.
Validation
RTX PRO 4500 SM120:
Not run: full GLM serving, TP2/TP4.
Compatibility
re-exported.
Disclosure
Written by an AI coding agent (Codex, GPT-6) with a staged design-review acknowledgement, and validated on my hardware (RTX PRO 4500 Blackwell, SM120, driver 615.78.08 open module, CUDA 13.4). Claude Code reviewed the storage-contract, CPU harness and reporting aspects and prepared this PR; it did not author or review kernel code.
🤖 Generated with Claude Code
https://claude.ai/code/session_01BJmfhfkwQM1xjdQVaUKwns