fix: adopt partially matched paged prefixes instead of declining#230
Merged
Conversation
try_adopt_cached_prefix declined every paged adoption whose matched_len fell short of the stored entry, so the default backend (auto -> paged under batching) never reused divergent-sibling prefixes, and with --apc-enabled the APC block clamp made even strict multiturn continuation miss (measured: multiturn cached 11067 -> 0). The paged arm now floors matched_len to the pool block boundary and trims the detached set before adopting: the new CachePool::trim_detached_paged_to releases both references of every dropped tail block (allocation + detach pin, mirroring release_detached_paged), clamps the dense-handle offsets and bookkeeping lengths so adopt restores RoPE at the trimmed length, and rejects Turbo4 sidecar layouts and sliding-window states. Flooring to the pool block boundary means no partially filled tail block survives, so the suffix re-prefill always starts on a fresh block and needs no copy-on-write. Whole-entry matches skip the trim and stay bit-exact with the old path. Validation: - New real-model gates paged_partial_prefix_share_parity (qwen3 + llama3): store with APC on clamps a 52-token entry match to 48, the scheduler mirror floors to 32, trims, adopts, and decodes byte-identically to a cold run. - New unit tests: trim releases the dropped block fully and the trimmed set adopts + gathers byte-identically; misaligned and oversized targets are rejected. - HTTP end-to-end (llama-3.2-1b, --apc-enabled, default paged): burst cached 0 -> 3648 tokens, seqshare 0 -> 25536, multiturn 0 -> 10976, matching the dense backend; seqshare footprint peak 1412 MiB vs dense 1507 MiB; wall times within noise. - 868 core lib tests (one pre-existing release-mode pack3 failure on main, unrelated), 3150 server lib tests, cold clippy clean.
…binding (review follow-up)
sjswerdloff
added a commit
to sjswerdloff/mlxcel
that referenced
this pull request
Jul 3, 2026
…q_len Two tests introduced upstream in lablup#230 (95df7cb) and lablup#231 (d402468) panic at `paged_detach.rs:1069`: debug_assert!(handle.offset >= target_tokens as i32); handle.offset = target_tokens as i32; The assertion's comment claims a freshly detached pool-backed handle "always satisfies offset == seq_len > target_tokens". This is false. For a paged sequence, K/V data lives in the shared block pool. The dense-side handles (`sequence.caches[]`) are metadata placeholders whose `offset` field is NEVER advanced during writes — `write_prefill` touches the paged state, not the dense-cache counter. When `detach_paged` builds the set via `sequence.caches.iter_mut().map(|c| c.clone_handle())`, each handle's offset is 0. `trim_detached_paged_to(&mut set, 8)` then hits `debug_assert!(0 >= 8)` and panics. The panic was masking, not caused by, the design. The line immediately below the assertion — `handle.offset = target_tokens as i32` — is the correct behavior: re-anchor the dense-side metadata offset to the trimmed length so a later adopt resumes RoPE at the prefix boundary. This mirrors `clone_detached_paged_prefix`'s `pool_backed_handle_clone(target_tokens as i32)` (line 885). Verified empirically: the two tests fail identically at 95df7cb (their birth commit), i.e. they never passed on upstream main. Not a fork regression. Fix: rewrite the comment to describe actual behavior and tighten the assertion to catch what would genuinely be corruption (0 < offset < target_tokens — a partially-advanced handle) while allowing both legitimate states: * Pool-backed handle: offset == 0 (dense K/V lives in the pool) * Dense-backed handle (Int8 dense-compat): offset >= seq_len Tests: cache::paged_detach::tests::trim_detached_paged_to_enables_partial_prefix_adoption cache::paged_detach::tests::detached_paged_set_accounts_real_pool_bytes Both now pass. Full paged_detach::tests suite: 35/35 green. Full cache::paged: 87/87 green. Broader `cache` selector: 449/449 green. Note: running `cargo test cache::` (all 925 cache tests) now segfaults at aggregate teardown *after* all tests complete. Pre-fix this suite reported clean 413/2 because the two panics leaked resources in a way that MLX FFI teardown tolerated; the fix un-masks a downstream order-dependent FFI teardown issue. Filed as follow-up. Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>
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
1 participant
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.
Summary
Implements #225.
try_adopt_cached_prefixdeclined every paged adoption whosematched_lenfell short of the stored entry, so the default backend (auto resolves to paged under batching) never reused divergent-sibling prefixes, and with--apc-enabledthe APC block clamp made even strict multiturn continuation miss (measured: multiturn cached tokens 11067 with APC off, 0 with APC on).matched_lento the pool block boundary and trims the detached set before adopting. Flooring means no partially filled tail block survives, so the suffix re-prefill always starts on a fresh block and needs no copy-on-write against a shared tail. Whole-entry matches skip the trim and stay bit-exact with the old path.CachePool::trim_detached_paged_to: releases BOTH references of every dropped tail block (the block-table allocation and the detach pin, mirroringrelease_detached_paged), clamps the dense-handle offsets and bookkeeping lengths so the subsequent adopt restores RoPE at the trimmed length, and rejects Turbo4 sidecar layouts and sliding-window (logical_start > 0) states. Release failures are collected and reported after the trim finishes so the set's bookkeeping never goes inconsistent halfway; the caller then releases the remaining pins and falls back to a cold prefill.min_prefix_tokensand reported as the adopted length in logs, observability, and the scheduler return value.Validation
paged_partial_prefix_share_parity(qwen3 + llama3): a store with APC enabled clamps a 52-token entry match to 48 tokens, the scheduler mirror floors to 32, trims, adopts (B keeps A's first physical block, drops the second), re-prefills the rest, and decodes byte-identically to a cold run.--apc-enabled, default paged backend, scripts/bench_memory_footprint.py): burst cached tokens 0 to 3648, seqshare 0 to 25536, multiturn 0 to 10976, matching the dense backend's numbers; seqshare footprint peak 1412 MiB vs 1507 MiB on dense; wall times within run noise.paged_prefix_share_parity,paged_scheduler_parity,paged_real_model_parity.pack3should-panic failure also present on main, unrelated), cold clippy clean, fmt clean.Concurrent siblings are still limited by the one-shot
take_detached(only the first adopter wins); that is #227.Closes #225