Repository navigation
feat(specprefill): preserve draft recurrent state at cache boundaries - #3842
Conversation
score_tokens derived cached_len and pre_lookahead_offset from cache[0].offset. On a hybrid model layer 0 is a recurrent layer whose cache holds a fixed-size summary rather than a growing sequence, so it has no offset attribute at all and the hasattr guard silently answered 0. A restored cache was therefore treated as empty. The whole prompt was re-prefilled on top of state that already held its prefix, the attention in every prefill chunk ran against that duplicated prefix, and _compute_importance then sliced keys[..., :n_prompt, :] off a cache holding cached_len + n_prompt entries -- so the importance scores, and the token selection built from them, came from the wrong keys. The mis-trimmed cache was what got stored. sparse_prefill in this same module already had the right idea: map attention layers to their cache indices and read the offset from there. _logical_cache_offset generalizes that, so no layer index is assumed and composite caches are descended into. Offsets are only compared between layers that mean the same thing by them. A bounded cache such as RotatingKVCache counts tokens ever processed rather than tokens held, and prefix_cache already refuses that comparison, so bounded layers answer only when nothing else can. Among the rest a disagreement is logged and the smallest wins, because prefilling a token twice is recoverable and skipping one is not. A cache that holds state but reports no position raises rather than answering 0, and a cache offering no emptiness probe counts as holding state. Both derivations are clamped to n_prompt. Reading the position correctly makes the exact-hit branch reachable on a hybrid model for the first time -- cached_len was previously always 0 there, so the branch was dead -- and that branch runs one token to obtain logits, leaving the cache a token longer than the prompt. n_prompt is what the old fallback meant. _compute_importance is deliberately unchanged: once cached_len is right the restored cache is the true prefix of the same sequence, so the existing slice spans exactly the prompt. The cold/warm equivalence tests cover that. The bug needed a draft cache hit to be reachable and, on hybrid GDN models, the draft cache could not produce one, so it stayed latent.
2feb2d7 to
7aec39d
Compare
|
I verified that the boundary snapshots enable disk restoration: an 8,704-token prompt restored 8,192 tokens and produced identical importance scores and selections to cold scoring. The offset-only baseline rejected that same cached prefix. I'm holding this pending the exact-hit correction in #3840, since the restored-cache path also reaches that branch. |
On an exact draft-cache hit score_tokens replayed the last prompt token on a cache that already held it, to obtain the logits the lookahead starts from. The token attended to itself twice and every recurrent layer advanced one step past the prompt; the lookahead trim shortens KV but restores neither the logits nor that state, so importance scores and the selection built from them were wrong (Qwen3.5-0.8B, 8,192-token SSD hit: 1,536 of 4,096 selected tokens differed from cold scoring). Only the N-1 state followed by a forward pass over the last token yields those logits. The draft lookup now queries tokens_to_score[:-1], so a hit can never cover the prompt end and the last token always goes through the normal prefill -- the same N-1 rule the target path applies. The replay branch is removed, and score_tokens raises on a cache that already holds the whole prompt instead of replaying; the draft workflow already falls back on that. The cache score_tokens returns is also the one stored as the prompt's, and the lookahead left it describing prompt + lookahead: only unbounded KV was sliced back. Recurrent state kept the lookahead tokens and, for a block-multiple prompt, store_cache persisted it as the final block's state. score_tokens now walks the leaf caches, descending into CacheList, and restores every leaf the slice cannot undo -- ArraysCache, a wrapped RotatingKVCache -- in place from a copy taken before the lookahead. Arrays are copied and containers rebuilt, because ArraysCache writes its state list in place and RotatingKVCache slice-assigns into its keys. The hybrid test fixture could not show either defect: its recurrent layer was stateless and its attention non-causal. Both are fixed, and the new tests compare importance and selection against cold scoring and the returned cache against a cold prefill, rather than cache length.
score_tokens holds each non-sliceable leaf cache's attributes before the lookahead and puts them back after it, so the cache it returns -- the one stored as the prompt's -- describes the prompt and not the prompt plus the lookahead tokens. The hold copied a leaf's own arrays and containers but kept nested objects by reference. SizedArraysCache, which reconstruct_cache returns for every restored recurrent layer, keeps its state on the ArraysCache it wraps, so the hold captured only that reference: decode then advanced the wrapped state in place and nothing put it back. Every restored hybrid draft cache therefore came back from scoring with recurrent state n_lookahead + 1 tokens past the prompt, and was stored that way. Measured on Qwen3.5-0.8B for an 8,193-token prompt restored at 8,192: returned-cache recurrent state differed from cold scoring by up to 19.1; with the fix, by 0.25, the drift of a different prefill chunk schedule. Nested objects are now held the same way and restored in place, with a seen-set so a reference cycle or a shared object is held once. The real-cache round-trip test covers a SizedArraysCache and fails without the recursion.
The draft prefix cache never produced a hit on a hybrid GDN model. Draft scoring called store_cache without boundary_snapshots, which the target path has passed since it was added, so every stored block carried a placeholder for each recurrent layer: the only live recurrent state sits past the trailing partial block that store_cache skips. The runtime said so on every turn -- "ArraysCache layer 0: partial prefix match detected (placeholder in last matched block)" -- and scoring re-prefilled the whole prompt every time. Capture that state and pass it, keyed by absolute token count, exactly as Scheduler does for the target model. Sliceable layers are nulled rather than carried, since store_cache re-slices those from the live cache and pinning them would hold the whole growing KV for nothing. The state for a boundary exists only while the prefill is standing on it, so the boundary is chosen on the first progress report, which is the position the prefill actually starts from. _last_reachable_boundary walks _prefill_draft's chunk loop to find the largest block-aligned position it will report. It walks rather than solving, because the last chunk is truncated to leave a token for the logits call, so the reported positions are not multiples of the step: a closed form got that wrong and silently published an earlier boundary, losing a block of reuse each time. Extracting at every block boundary and keeping the last, the other obvious approach, costs a full state extraction per block for nothing. The draft cache is allocated here rather than inside score_tokens, which otherwise hands it back only after scoring, leaving no way to read state at a boundary. A fresh cache reports offset 0, so the suffix prefilled is the whole prompt, as before. Holding that cache makes its lifetime this function's problem. Nothing may still name it, or state pulled out of it, when the buffers are returned: the reconstruct result is bound straight to draft_cache rather than through a second local, and draft_cache, the snapshots and the extracted payload are all released before the clear. Otherwise the draft KV outlives the call that exists to reclaim it, and the reclaim accounting measures a delta it did not get. Tests pin _prefill_draft's exact reporting schedule and sweep the boundary arithmetic against it, since a drift there would silently capture nothing rather than fail; a weak reference taken across the clear pins the ownership. The sliceable-type set is pinned against the scheduler's, which draft.py cannot import without a cycle.
The draft lookup now asks for all but the last prompt token, because the logits the lookahead starts from must come from a real forward pass over it. _last_reachable_boundary still let a block-aligned prompt publish its snapshot at n_to_score, which that lookup can never match: scoring the same prompt again matched the blocks below the end, found no recurrent state there, and restored nothing. The boundary is now the last block-aligned position _prefill_draft actually reports strictly below n_to_score. That position is not n_to_score - block_size: the prefill stands only on chunk ends and on the truncated last chunk at n_to_score - 1, so for 8,192 tokens in 2,048-token chunks it is 6,144. The rule still comes from walking the real loop, and the sweep against it now also asserts the prompt end is never chosen. A longer prompt extending this one still resumes at the end: store_cache publishes the returned live state as the final block's state when the store is a block multiple, and that state is the prompt's own since the lookahead is undone (8,704 tokens after a stored 8,192 restores 8,192). Measured on Qwen3.5-0.8B through a real SSD store and a fresh prefix cache: an 8,192-token prompt scored twice restores 6,144 tokens, and an 8,704-token prompt restores 8,192. Both match cold scoring exactly -- importance, selection and returned cache -- because the restore point is one the cold prefill also stands on. An end-to-end test drives the same path with a tiny random-weight Qwen3.5 and fails under the old rule.
7aec39d to
cbe542d
Compare
|
This is rebased onto #3840's latest commit and adapted to its N-1 lookup. Runtime results on Qwen3.5-0.8B through a real SSD store and a fresh prefix cache, at block size 2,048 (the value the scheduler forces for GDN) and at 256:
A new end-to-end test drives the same path with a tiny random-weight Qwen3.5 and the real prefix-cache/SSD stack, and it fails under the old rule. CI selection: 14,253 passed. Details are in the updated description. |
… times test_boundary_prediction_matches_the_real_loop_everywhere ran the real _prefill_draft for every (n, step, block) in its sweep, which took 155-177 s per CI job. The sweep now runs against a pure-Python transcription of the loop's reporting schedule, and the real loop runs only on a small matrix of the cases a transcription or a closed form gets wrong: block dividing step and not, a truncated final chunk, an exactly block-aligned prompt end, a suffix inside one chunk. A new test checks the transcription against the real loop on that matrix, and test_prefill_draft_schedule_is_exactly_this still pins the real schedule directly. Both earlier implementations still fail: the closed form over multiples of step, and the version that returned the prompt end when it was aligned. The file runs in about 4 s.
upstream's jundot#3840 and jundot#3842 added test classes to tests/test_specprefill.py at the spot where this branch adds TestPrefillOomRestoresRope; both kept. No code conflicted.
Problem
The SpecPrefill draft prefix cache never produces a hit on a hybrid GDN model.
Every scoring re-prefills the whole prompt, however much of it the cache
already holds.
Root cause
Draft scoring calls
store_cache()withoutboundary_snapshots, which thetarget path has passed since the parameter was added. Recurrent layers cannot
be re-sliced out of a live cache the way KV layers can, so without a snapshot
every stored block carries a placeholder for each recurrent layer. The only
live recurrent state sits past the trailing partial block, which
store_cacheskips. Walk-back then scans every stored block, correctly finds no usable
checkpoint, and reports a miss — every time, for every conversation. The
runtime says so directly:
Fix
Capture the recurrent state at a block boundary during the draft prefill and
pass it to
store_cachekeyed by absolute token count, exactly asSchedulerdoes for the target model.
store_cachere-slices thosefrom the live cache; pinning them in a snapshot would hold the whole growing
KV for nothing. This mirrors
scheduler.py's own boundary emission.length the caller expected. Those two should agree; if they ever did not, the
capture would silently never fire.
chunk is truncated to leave a token for the logits call, so the reported
positions are not multiples of the step. A closed form got this wrong.
keeping the last costs a full state extraction per block for nothing.
here because
score_tokenshands its own back only after scoring, too lateto read state at a boundary. So nothing may still name it, or state pulled
out of it, when the buffers are returned — otherwise the draft KV outlives
the call that exists to reclaim it, and the reclaim accounting measures a
delta it did not get.
Boundary below the prompt end
#3840 established that the logits the lookahead starts from must come from a
real forward pass over the last prompt token, so the draft lookup asks for
tokens_to_score[:-1]and a hit can never cover the prompt end. The firstrevision of this PR still let a block-aligned prompt publish its snapshot at
n_to_score. That is a boundary the lookup can never match: scoring thesame prompt again matched the blocks below the end, found only placeholders,
and restored nothing.
_last_reachable_boundary()now returns the last block-aligned position that_prefill_draftactually reports strictly belown_to_score. It stillwalks the real loop rather than solving it. The prefill stands only on chunk
ends, on the truncated last chunk at
n_to_score - 1, and onn_to_scoreitself after the logits call, so the answer is not
n_to_score - block_size.For 8,192 tokens in 2,048-token chunks, the reported positions are 2,048,
4,096, 6,144, 8,191 and 8,192, and the boundary is 6,144. There is no
special case for exact hits: the same rule applies to every prompt, and for a
non-aligned prompt it picks the same boundary as before.
A longer prompt that extends a block-aligned one still resumes at the end.
store_cachepublishes the returned live state as the final block's statewhen the store is a block multiple, and since #3840 that state is the
prompt's own.
Runtime: Qwen3.5-0.8B-MLX-4bit draft, real
BlockAwarePrefixCache+PagedSSDCacheManager, prefill step 2,048, keep 50%, greedy lookahead. Thefirst request stores; a fresh prefix-cache instance reading the same SSD
directory serves the second, which is compared with cache-free cold scoring.
Block size 2,048 is what the scheduler forces for an
ArraysCachemodel; 256is shown to exercise finer blocks.
The before column for the aligned prompt was measured at
f1c22e7d, thisPR's first commit rebased onto the updated #3840. Where the restore point is
one the cold prefill also stands on, cold and warm run the same chunks and
agree exactly on importance, selection and the returned cache. The extension
rows restore at the end of a stored prompt, where the stored prefix was
prefilled as
…2,047 + 1and cold covers the same tokens in full chunks.Their residual is that schedule drift, not cache content: two cold runs that
differ only in chunk size already differ by about 4e-4, and the SSD-restored
cache is bit-identical to an in-memory one prefilled the same way. Without
#3840's lookahead-state fix these rows showed a returned-cache Δ of about 19.
Why #3840 is a prerequisite
Two separate defects sit between a hybrid GDN draft model and a usable draft
prefix cache, and they fail differently.
#3840 is the correctness prerequisite. When a restored cache exists,
score_tokens()has to read its logical position so that only the uncachedsuffix is prefilled. On a hybrid model it read
cache[0].offset, and layer 0is recurrent with no
offset, so a restored cache was treated as empty. #3840also fixes the exact-hit replay: the lookup is now N-1, which is the contract
this PR's boundary rule is built on. And it undoes the lookahead on every
non-sliceable cache, including the
SizedArraysCachewrapper that everyrestored recurrent layer comes back as, so the cache stored after scoring
holds the prompt's own state.
This PR is the reuse enablement. Even with the position read correctly, a
hybrid GDN draft cache has no reusable recurrent boundary to restore from.
The dependency runs one way, and only for the measurements: the runtime
validation depends on #3840, because without it the cache-hit path itself is
wrong — a hit is silently discarded and the prompt re-prefilled, so there would
be nothing to measure. The regression tests here do not depend on it.
Nothing in this PR re-explains or modifies #3840's fix. The two can be reviewed
independently and merged separately; this one just has to go second.
Stacked branch. GitHub currently shows #3840's three commits along with
this PR's and the combined diff. The two commits belonging to this PR are
feat(specprefill): preserve draft recurrent state at cache boundariesandfix(specprefill): publish the draft boundary below the prompt end, touchingonly
omlx/specprefill/draft.pyandtests/test_specprefill_draft.py. Once #3840 merges I will rebase onto theupdated
mainand this diff will shrink to that commit alone.Validation
Baseline is this branch's parent — the offset fix alone — so the only
difference is the snapshot plumbing. Both arms ran back to back on the same
corpus and prompts, on a hybrid GDN pair (27B target, 0.8B draft).
partial prefix matchrejectionsThe published boundary advances with the conversation, and each hit is consumed
by #3840's suffix-only scoring, so the tokens prefilled are
n_prompt - cached.Four of the 27 scorings missed the cache; those land at 0.9x–1.0x of the fitted
baseline, so the snapshot machinery costs nothing measurable when it does not
pay off. All 27 sparse prefills completed, zero rejections, zero scoring
failures, and both arms reached the same answer on the task.
Runtime table, fit, and boundary progression
Baseline least-squares fit
-1.05 + n * 1.920e-4s, fitted onn = 13,291..34,010.
The table above is thirteen representative rows of the twenty-seven.
Published boundary, all 18 distinct values over 21 captures: 12288, 16384,
17408, 19456, 25600, 26624, 27648, 29696, 31744, 32768, 33792, 34816, 35840,
36864, 37888, 38912, 39936, 47104.
Limitations of the runtime measurement
output differs, so the comparison is the per-prompt fit above and not the
raw totals.
column is extrapolated.
a fixed-seed pair.
block_size1024, which is that machine's configured value;the upstream rule below would pick 2048 for the same model. Replaying the
measured prompt sequence through the boundary arithmetic gives the same 11
published boundaries and the same 800,768 tokens not re-prefilled at both
1024 and 2048, and 7 boundaries / 770,048 tokens at 4096 — so the block size
does not change the conclusion on this workload. Derived, not measured:
it replays the arithmetic, not the server.
How the coupling to
_prefill_draftis kept honest_last_reachable_boundary()models_prefill_draft's reporting loop. That isa coupling to another module's private function, and the failure mode if it
drifts is silent: the feature would capture nothing and the draft cache would
quietly go back to never hitting.
So the tests pin the loop itself.
test_prefill_draft_schedule_is_exactly_thisasserts the exact sequence of reported positions, driving the real
_prefill_draft; a change there breaks a test and asks a human to re-derivethe arithmetic. Asserting only the prediction is not enough — several plausible
edits to that loop leave the largest aligned position unmoved, and a weaker
version of this test passed against a deliberately perturbed loop.
test_boundary_prediction_matches_the_real_loop_everywheresweepsstep, block ∈ [2,20)andn ∈ [2,200)against the real loop, taking onlyreported positions strictly below
n, and asserts the prediction is nevernitself. An earlierclosed-form implementation was wrong 6729 times over that sweep — for instance
cached=0, m=3073, step=2048, block=1024published 2048 although 3072 wasreported and block-aligned, losing a block of reuse each time.
test_prefill_draft_holds_exactly_the_tokens_it_reportspins the other half:at every reported position the cache holds exactly that many tokens. If that
stopped holding, a checkpoint would be stored under the wrong token count and a
later restore would resume from the wrong place.
TestRepeatedBlockAlignedPromptdrives the whole path with real pieces: atiny random-weight Qwen3.5 (GDN and attention layers),
BlockAwarePrefixCacheover a
PagedSSDCacheManageron disk, and the scheduler's own stateextraction. A block-aligned 128-token prompt is stored and then scored again
through a fresh prefix-cache instance. The second request restores 96, the
boundary derived from the reported schedule 32, 64, 96, 127, 128. Importance
is
array_equalto cold, selection is identical, and the returned cacheequals a cold prefill of the prompt. Under the old rule it restores 0 and
fails.
test_restored_cache_is_released_before_the_cleartakes a weak referenceacross
sync_and_clear_cache()and fails if anything still names the restoredcache there.
_SLICEABLE_CACHE_TYPESduplicatesscheduler._KNOWN_SLICEABLE_CACHE_TYPES,because importing it would make
omlx.scheduler → omlx.specprefill.draft → omlx.scheduler. Tests have no such cycle, sotest_sliceable_cache_types_track_the_schedulerasserts the two are equal.(
CacheTypeRegistry.is_sliceableis a different question — it classifiesChunkedKVCacheas not sliceable where the scheduler deliberately does, andthe consumer of these snapshots follows the scheduler.)
On the block size, and a branch that was removed
Scheduler._enlarge_block_size_for_arrays_cacheforcespaged_cache_block_size = max(2048, prefill_step_size, qwen35_floor)for anymodel with an ArraysCache, so on a GDN hybrid the block size is always at least
the prefill step.
An earlier revision carried a fallback that re-chunked the prefill at the block
size when no boundary was reachable. Under that rule it can never fire on the
models this exists for, and where it could fire — a non-GDN model, block 256 —
every layer is sliceable, so the capture returns early and the extra chunks buy
nothing. It is removed rather than kept as a branch that only ever costs
throughput. Removing it also keeps the draft chunk schedule fixed at
prefill_step_size, which matters because GDN chunked recurrence is notbit-identical across different chunk splits.
Test gate
test_specprefill.py+test_specprefill_draft.py: 89 passed.pytest tests/ -m "not slow and not integration", which is what CI runs:14,253 passed, 356 skipped, 85 deselected, 0 failed.
ruffis clean on thetouched files.
Scope / non-goals
Only draft boundary-snapshot plumbing.
gdn_ssd_split_enabledis not enabled,the target cache is untouched, and nothing here changes the static-prefix
store, PCSR, or the SpecPrefill threshold. It does not re-fix #3840's logical
offset; it depends on it.