fix: account real paged pool bytes in the prompt-cache ledger and stats#231
Merged
Conversation
The prompt-cache ledger valued a pool-backed paged entry by the layout's nominal bytes_per_block placeholder (32 B/block), so a 10.5k-token entry holding ~1.2 GiB of real KV was booked at 296 KB and the 2 GiB capacity cap never triggered. Pool stats had the same blind spot: per-sequence nominal sums missed parked retention entirely. - New dtype::size_bytes and PagedBlockPool::real_block_bytes expose the actual per-block cost (block_size x n_kv_heads x head_dim x element_size x 2) from the layer geometry captured on first write. - detach_paged records the set's real pinned-pool bytes; DetachedPagedCacheSet::nbytes prefers them (entries whose K/V never reached the pool, e.g. Int8 dense-compat, keep the nominal fallback). trim_detached_paged_to shrinks the figure by the dropped blocks. - PagedBlockPool::stats (was stats_for_sequences) reports pool-wide real bytes: reserved = allocated slab capacity, in_use = rows mapped to live blocks, covering active sequences and parked pins alike. - memory_usage_bytes no longer double-counts: parked paged sets contribute only their own dense handles and lifted sidecars, since pool_tensor_bytes already contains the pinned blocks. E2E (llama-3.2-1b, --prompt-cache-capacity-bytes 300000000): three distinct ~3.8k-token prompts book 121 MiB each (real), the third insert evicts the oldest entry (evictions_lru 1), and /v1/cache/stats paged_bytes_in_use tracks 121 -> 242 -> 242 MiB. Side finding surfaced by the real numbers: with --apc-enabled, a small partial match (e.g. the ~30-token template preamble) take()s and trims a multi-thousand-token entry down to one block, destroying its cached KV (#225's one-shot consume). Default APC-off configs are unaffected (whole-containment matches only). The non-consuming clone-and-pin adoption in #227 removes that hazard; details on the issue. Gates: core 869 lib tests (one pre-existing release-mode pack3 failure on main), server 3150, prefix-share + scheduler parity 8/8 real-model, cold clippy clean, fmt clean.
…s arm (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 #226. The prompt-cache ledger valued a pool-backed paged entry by the layout's nominal
bytes_per_blockplaceholder (32 B/block): a 10.5k-token entry holding about 1.2 GiB of real KV was booked at 296 KB, so the 2 GiB capacity cap effectively never triggered and real retention was bounded only by TTL and max_entries.dtype::size_bytesandPagedBlockPool::real_block_bytesexpose the actual per-block cost (block_size x n_kv_heads x head_dim x element_size x 2) from the layer geometry captured on first write.detach_pagedrecords the set's real pinned-pool bytes;DetachedPagedCacheSet::nbytesprefers them. Entries whose K/V never reached the pool (logical-only appends, Int8 dense-compat) keep the nominal fallback.trim_detached_paged_to(fix(server): paged backend declines every partial prefix adoption, and APC+paged loses even multiturn reuse #225) shrinks the figure by the dropped blocks.PagedBlockPool::stats(renamed fromstats_for_sequences; the per-sequence walk is gone) reports pool-wide REAL bytes:bytes_reserved= allocated slab capacity,bytes_in_use= rows mapped to live blocks. This covers active sequences and parked prompt-cache pins alike; the old per-sequence nominal sums under-reported by orders of magnitude and missed parked retention entirely.memory_usage_bytesno longer double-counts: parked paged sets contribute only their own dense handles and lifted Turbo4 sidecars, sincepool_tensor_bytesalready physically contains the pinned blocks.Validation
--prompt-cache-capacity-bytes 300000000): three distinct ~3.8k-token prompts book 121 MiB each (real),paged_bytes_in_usetracks 121 / 242 / 242 MiB, and the third insert evicts the oldest entry (evictions_lru1). Before this change the same three entries were booked at ~0.3 MB total and nothing ever evicted.pack3should_panic also failing on main), server 3150, prefix-share + scheduler real-model parity 8/8 byte-identical, cold clippy clean, fmt clean.Side finding (tracked in #227)
The real numbers made a behavioral hazard visible: with
--apc-enabled, a small partial match (for example the ~30-token chat-template preamble shared by every request) takes the stored entry (one-shottake_detached) and trims it to one block, destroying thousands of tokens of cached KV to reuse 32 (observed live:trimmed detached set from=3848 to=32). Default APC-off configurations are unaffected because matches there require whole-entry containment. The non-consuming clone-and-pin adoption planned in #227 removes the hazard; a detailed comment is on that issue.Closes #226