[seokwoosong]

This PR is part of a stack (oldest at bottom):

• [Fix] unified-memory: Avoid gather scratch during Mamba and Conv compaction #44097
• -> [Fix] unified-memory: Bound joint reclaim and preserve cached prefixes #39294 (merged)

CI is run only on the top PR, #44097, which contains both changes.

Merged into main as ee30dc4d40993045273f0a90260db7efa302c214. #44097 has been rebased onto this merge and now targets main directly.

Current head: 1adc6a6229fecca1b58cb56df6732c395ea65334.

Motivation

Joint FULL/SWA allocations can evict more cached prefixes than needed. FULL/SWA/Mamba allocations can fail despite recoverable capacity, or discard the entire cache before rejecting an impossible demand.

Seokwoo Song's original commits are preserved, rebased onto ae28306a74fc9ba56daab9ccb06dcf45659f58af.

Modifications

• Stop allocation-driven eviction at a read-only sufficient-reclaim predicate; prepare and verify the physical layout at bounded phase boundaries.
• Reject demands proven impossible by byte or ID bounds before eviction, and bound token/state reclaim quotas.
• Compute exact unsharded paged-extend demand for the unified SWA allocators while preserving conservative demand elsewhere.
• Update the dLLM KV-reuse test allocator to implement the paged-demand interface with the conservative default, fixing the CPU CI failure.
• Adapt the closed-move-gate regression to the current physical-page copy interface. Preserve main's SWA-preferred reclaim planning for storage-backed HiCache and decode disaggregation.

Accuracy Tests

Latest main rebase: ae28306a74fc9ba56daab9ccb06dcf45659f58af. Both PRs replayed without conflicts, and all eight patches are unchanged by git range-diff. The exact combined top commit passed 87 tests / 211 subtests on H200 (Python 3.12 / Torch 2.13.0+cu130), covering dLLM KV reuse, prefill budget, joint reclaim, SWA page estimation, Mamba/Conv compaction and views (CPU and CUDA), and adjacent HiCache ordering/regressions. All changed-file pre-commit checks and git diff --check passed.

Earlier CI follow-up, before this main rebase: reproduced the get_extend_allocation_demand failure in test_dllm_fdfo_kv_reuse.py before the correction. After updating the fake allocator in #39294, 66 tests / 162 subtests passed across the dLLM KV-reuse, prefill-budget, joint-eviction, tri-pool reclaim, and SWA paged-extend suites on the H200 host (Python 3.12 / Torch 2.13.0+cu130). All six dLLM tests passed. Changed-file pre-commit checks passed.

Earlier validation before this main rebase:

• H200 host, Python 3.12 / Torch 2.13.0+cu130: 220 tests / 1,050 subtests passed across 14 allocator/cache files using Python TreeCore. These are CPU allocator tests.
• Freshly built current-main Rust TreeCore on the H200 host: 38 tests / 155 subtests passed in the focused reclaim suite; session cases use the supported Python fallback. A separate probe asserted that the ordinary cache selected RustUnifiedTreeCore.
• With the compaction fix from [Fix] unified-memory: Avoid gather scratch during Mamba and Conv compaction #44097 applied in a separate combined snapshot on B300, Torch 2.14.1+cu130: 61 tests / 201 subtests passed, including the CUDA compaction cases.
• Changed-file pre-commit checks and git diff --check passed.

Controlled real-allocator comparisons against unchanged main production code, with the same test fixtures, reproduced these results in both eager and lazy modes:

Case	Main	This change
FULL/SWA, 6-token demand	Succeeds; evicts 5 cached tokens	Succeeds; evicts 3
FULL/SWA/Mamba, 4-token demand	Fails	Succeeds; evicts 1 cached token
FULL/SWA/Mamba, 200-token impossible demand	Fails; evicts all 96 cached tokens	Fails; preserves all 96

The selected regression cases fail on main and pass on the candidate. Review checked quota semantics, deferred frees, page rounding, locked prefixes, changing move gates, pending reuse, and the compatibility of the rebase.

The FLOAT geometry predicates are sufficient certificates, not a complete feasibility solver. No fresh model-serving, full HiCache, or multi-GPU result is claimed. Historical serving evidence is retained below and is not validation of this head.

Speed Tests and Profiling

The controlled allocator comparisons above validate cache preservation and successful recovery, not serving throughput. No new serving-speed claim.

Checklist

• Format code with pre-commit.
• Add and validate regression tests.
• Follow SGLang code style.
• Documentation: not applicable; no new user-facing API or configuration.
• Fresh model accuracy / serving performance benchmarks: not run for this rebase.

Historical description and validation before the current rebase

Motivation

In unified-memory configurations, KV + SWA allocation can evict more prefix cache than necessary. Mamba + KV + SWA allocation can also fail despite recoverable joint capacity, or discard cached entries before rejecting an impossible request. This change bounds joint reclaim for these allocators to preserve reusable cache and recover capacity for feasible allocations.

Modifications

• Stop the victim walk when the allocator’s sufficient reclaim predicate is satisfied, then prepare and check the actual allocation.
• Keep the per-victim feasibility predicate read-only and bound preparation to reclaim phase boundaries.
• Use the exact unsharded paged-extend demand where supported, retaining the conservative fallback for other allocators and sharded layouts.
• Reject demands proven impossible by byte/ID bounds before eviction, preserving cached entries in those cases; bound remaining recovery attempts. This is not a complete feasibility solver for every FLOAT layout.

Accuracy Tests

The focused three-file CPU suite passed 54 tests / 162 subtests (23.73 seconds). Both lazy modes reproduced the allocator results below: two-pool demand 6 reclaimed 5 → 3 cached tokens; three-pool demand 4 changed failure → success; impossible three-pool demand 200 preserved all 96 cached tokens instead of evicting them. Baseline production code was unchanged; only test fixtures and regression files were supplied for comparison. These checks used Python TreeCore.

The current revision 3af4a9dd2bdb6cc0c5e7f27e3a09cbd5438fc647 additionally clarifies the two evict_for_alloc parameter contracts in docstrings only; executable ASTs are unchanged relative to the tested revision. Changed-file pre-commit checks passed. All-files pre-commit also passed for the preceding API migration. GPU, Rust TreeCore, and serving results below remain historical results from the explicitly listed older revisions and were not rerun after this rebase.

Validation comparing upstream (e17935820a47) with this fix (fb1d221b849d):

Check	Result
CPU, 14 files	281 tests / 1,628 subtests passed
Rust TreeCore focused tests	38 tests / 155 subtests passed
GPU allocator payload tests, RTX 5090	3 passed in both unit and serving environments; real KV and Mamba moves exercised
Formatting	All-files pre-commit passed

Actual allocator case	Before	After
Two-pool demand 6	5 cached tokens reclaimed; allocation succeeds	3 reclaimed; allocation succeeds
Three-pool demand 4	Allocation fails	Allocation succeeds, reclaiming 1 cached token
Impossible three-pool demand 200	Rejects after deleting all 96 cached tokens	Rejects while preserving all 96

These are controlled allocator cases, not serving throughput claims. The real-GPU pending-reuse race branch was not reached; no coverage is claimed for that branch. Request-owned SWA WINDOW_FREE events were also not observed in the serving workloads, including the supplemental longer-prefix lane.

Serving comparison	Result
Mellum (KV + SWA)	48/48 primary correctness sequences and 3,840/3,840 paired measured primary sequences match; mapping/concurrent output differences require the qualifications below
Reduced Inkling (KV + SWA + Mamba)	48/48 primary correctness sequences and 3,840/3,840 paired measured primary sequences match
Qwen3.5-4B (KV + Mamba)	Both versions show the same two concurrent text variants; no universal token-ID parity or causal explanation is claimed

Supplemental C2 collection is complete. A separate 128-output-token smoke test matched all 29 paired request sequences (Mellum 13, Inkling 16), with 58 requests, zero errors and zero reported retractions. This is limited longer-output coverage, not a throughput or quality benchmark. Output variability in other workloads is retained below.

Serving output comparisons, targeted diagnostics, and limitations

The primary Mellum and reduced Inkling correctness comparisons matched all 96 token-ID sequences. All 7,680 paired measured request sequences also matched across the 16 primary performance runs; the supplemental results and qualifications are reported separately below. The completed pressure mapping compared 636 requests across six paired configurations: 635 token-ID sequences matched. The one Mellum mismatch occurred at the 4K cap when a revisit retained 896 rather than 512 prefix tokens. It reproduced with caching enabled; all seven requests matched with caching disabled. A focused 48-request diagnostic reproduced the split and controlled retained-prefix depth: both versions produced identical 32-token outputs and full-vocabulary logits when matched at 512 tokens, and also when matched at 896 tokens. Changing depth from 512 to 896 reproduced the same output change on the baseline alone. Retained FULL/SWA KV matched seed bytes at prefix matching and immediately before each consuming forward. This localizes the observed reproducer to retained-prefix/compute-path sensitivity; the synchronizing diagnostic does not exclude unobserved races or establish output equivalence for other workloads.

Across two supplemental Mellum C2 pairs, 254/260 measured token-ID sequences and 46/52 warmup sequences matched. Both revisions produced all observed variants for the same input and retained-prefix depth, including the initially candidate-only warmup variant. All six measured differences correlated with request forward-entry order; that association alone is not an exact batch or kernel-level causal proof. A separate 232-request CPU-trace diagnostic found identical outputs in all three actual batch-sequence groups shared by both revisions. A fourth sequence group appeared only on the baseline, and all four attempted same-batch admission conditions were not realized. The diagnostic added 12 short warmups per server absent from the original C2 runs. CPU metadata and observed output prefixes do not prove actual GPU inputs, KV payload integrity, the final sampling decision, or race freedom. These descriptive C2 results are kept separate from primary performance measurements.

In fresh Qwen performance runs, one concurrent request produced the same two observed text variants in both versions; no candidate-only text variant was observed. Earlier token-ID triage also found within-candidate variation, and a candidate repeat matched both baseline runs. The allocator/scheduling/numerical cause remains unresolved. Performance text hashes are not complete token-ID parity coverage.

Speed Tests and Profiling

Fresh RTX 5090 / WSL measurements, using four Before/After pairs per workload. Historical serving numbers are excluded. Changes below are medians of paired changes; CPU is scheduler process time.

Model / workload	Cached tokens per revisit, Before → After	TTFT p50 change	Output tokens/s change	CPU change
Mellum (KV + SWA) / eviction threshold	512 → 512	-0.54%	+0.45%	-0.42%
Reduced Inkling (KV + SWA + Mamba) / eviction threshold	0 → 512	-25.57% (-10.85 ms)	+0.60%	-0.58%
Qwen3.5-4B (KV + Mamba) / serial control	Unchanged	-0.24%	+0.90%	-0.87%
Qwen3.5-4B (KV + Mamba) / C4 control	Unchanged	-0.42%	+1.05%	-0.99%

Mellum and Inkling use the preselected 8K-cap, oldest-first conditions at their baseline-calibrated eviction thresholds (9 and 12 cold requests). Inkling's revisit TTFT decreased in all four pairs (-23.35% to -26.49%). Whole-workload throughput changes were small, and cache reuse did not increase in every condition. A separate diagnostic replay at Inkling's threshold observed scheduled prefill work for the two revisits decrease from 2,048 to 1,024 tokens; this is the scheduler's extend-token count, not measured GPU FLOPs. None of the predeclared performance alerts triggered; this is not statistical equivalence or a general speedup claim. Qwen's output-variability limitation above remains.

All Mellum and Inkling primary conditions

Changes are median paired percentages; throughput ranges retain every pair, including outliers. Each condition has 20 measured episodes plus two warm episodes per server, with two hot-prefix revisits per episode.

Model / cold requests	Cached tokens per revisit, Before → After	Tokens/s change (range)	Revisit TTFT p50 change	CPU change
Mellum (KV + SWA) / 7	896 → 896	+0.85% (-0.38% to +14.91%)	-0.94%	-0.80%
Mellum (KV + SWA) / 9	512 → 512	+0.45% (+0.14% to +12.67%)	-0.54%	-0.42%
Mellum (KV + SWA) / 11	0 → 0	-0.55% (-5.24% to -0.05%)	+2.79%	+0.55%
Mellum (KV + SWA) / 5 (2K hot prefixes)	1536 → 1536	-0.47% (-0.83% to +8.67%)	-0.78%	+0.46%
Reduced Inkling (KV + SWA + Mamba) / 10	512 → 512	-0.34% (-1.38% to +0.46%)	+0.21%	+0.29%
Reduced Inkling (KV + SWA + Mamba) / 12	0 → 512	+0.60% (-1.09% to +1.09%)	-25.57%	-0.58%
Reduced Inkling (KV + SWA + Mamba) / 14	0 → 0	-0.05% (-0.73% to +0.51%)	+0.58%	+0.03%

Mellum's highest short-prefix pressure had lower throughput in all four pairs (median -0.55%, range -5.24% to -0.05%); its revisit TTFT p50 increased by a median 1.99 ms. This stays below the predeclared alert thresholds but is retained rather than described as zero regression. Some Mellum pairs had large positive throughput outliers; they are not the claimed benefit. Primary Mellum cache-hit totals were unchanged; separate mapping shows retention improvements at other capacities.

The 16 primary servers completed 15,360 measured requests and 1,728 warmup requests, with no request errors or reported retractions. All paired measured token-ID sequences matched. The additional 2K-prefix Mellum lane was specified before outcome inspection and retains its separate supplemental status.

Supplemental concurrent hot-prefix revisits

Two AB/BA pairs per model, using fixed seed 3001 at nominal cap 8192 and the baseline-calibrated threshold. Seeding and cold pressure remain serial; only the two hot revisits overlap (C2). Each server has 10 measured and two warm episodes, with no additional short warmups. These are descriptive two-pair results, separate from the four-pair primary measurements.

Model	Cached tokens per revisit, Before → After	Revisit TTFT p50, pair 1 / pair 2	Tokens/s, pair 1 / pair 2	CPU, pair 1 / pair 2
Mellum (KV + SWA)	704 → 704	+0.30% / -2.40%	-0.86% / -0.76%	+0.85% / +0.74%
Reduced Inkling (KV + SWA + Mamba)	0 → 512	-22.01% / -24.38%	+0.82% / +0.59%	-0.73% / -0.53%

All eight servers completed 1,392 requests (1,160 measured and 232 warmup), with zero request errors or reported retractions. Inkling matched 320/320 measured and 64/64 warmup paired token-ID sequences. Mellum matched 254/260 measured and 46/52 warmup sequences; the output-variability investigation and its limitations are described above. The small Mellum throughput decreases in both pairs are retained; no equivalence or broad throughput improvement is claimed.

Mellum / Inkling setup and workload

Before: upstream e17935820a47; After: fb1d221b849d. RTX 5090 (32 GB), WSL, Python 3.12.3, Torch 2.14.1, Triton 3.8.0, sglang-kernel 0.4.9. TP1, Python TreeCore, HiCache off, CUDA graph capture disabled, page size 128, chunked prefill 512, context length 4096, max running requests 4. Greedy 32-token outputs with EOS ignored; fixed token-ID inputs, not a language-quality benchmark.

Model	Pool layout	Attention	Static memory fraction	Actual shared arena at nominal cap 8192
Mellum2-12B-A2.5B-Base (KV + SWA) (271755e48ab6)	FULL + SWA	FA4	0.90	398,196,736 bytes
Inkling (KV + SWA + Mamba) (85b071f87d9b), local 8-layer / hidden-size-768 configuration	FULL + SWA (FLOAT) + Mamba	FA4	0.65	105,906,176 bytes

Inkling uses the same local SM120 compatibility override in both arms (inkling_moe.is_sm121 = lambda: True), Triton MoE/Mamba/linear attention, 128 Mamba slots, Mamba/full ratio 0.1, and extra_buffer tracking every 128 tokens. This is the stated local model configuration, not a claim about the full original model. Per-pool address limits share one arena and are not additive reservations; nominal token caps are not physical capacities. Both arms' actual allocator geometry was checked.

Each serial episode flushes cache, seeds two independent 1024-token hot prefixes, sends 512-token distinct cold requests, then revisits both hot prefixes oldest-first. Baseline-only calibration (seed 1001) finds the first cold request that triggers eviction; primary conditions use N-2, N, N+2 at cap 8192. Mellum's separate 2048-token hot-prefix condition has its own threshold of 5. Calibration does not use candidate outcomes. Four pairs use seeds 3001–3004 and AB/BA/BA/AB order, with the same seeded condition permutation in both arms. Each server has 12 short warmup requests and two full warm episodes per condition, excluded from measured aggregates. Flush/checkpoint work is outside timed episodes; client timing includes request setup/bookkeeping. Detailed diagnostic hooks are absent from performance runs.

Inputs are independently generated Python-style assignments. For each seed/case, random.Random(seed * 1000 + case) supplies an identifier, then 700 lines using names from record,item,cache,count,entry,result,value,index,table,request and integers below 100000. Each line has the form name_j = value # Store this independent record and preserve its value. The model tokenizer's encode(..., add_special_tokens=True) is truncated to the required length. Hot cases are 0/1, cold cases begin at 100; the maximum shared prefix is below one 128-token page. The long lane extends the same hot inputs and keeps their first 1024 IDs unchanged.

Alerts require a median throughput loss >5% or CPU increase >10%, or specified latency increases >5% and >2 ms, with at least three of four pairs adverse. Revisit TTFT p50 and phase E2E p50 are primary latency checks; p95 checks are exploratory. A triggered condition receives one bounded confirmation block; none triggered here. These thresholds do not establish equivalence or exclude smaller overhead.

Pressure mapping: all cache-reuse conditions

Each row is one paired diagnostic episode, with two hot-prefix revisits. Values are total cached input tokens across those two requests; these are not repeated throughput measurements. Oldest/recent specifies which hot prefix was revisited first.

Model / nominal token cap / hot-prefix length	Cold requests	Revisit order	Before cached tokens	After cached tokens
Mellum (KV + SWA) / 8192 / 1024	7	oldest	1,792	1,792
Mellum (KV + SWA) / 8192 / 1024	7	recent	1,792	1,792
Mellum (KV + SWA) / 8192 / 1024	8	oldest	1,792	1,792
Mellum (KV + SWA) / 8192 / 1024	8	recent	1,792	1,792
Mellum (KV + SWA) / 8192 / 1024	9	oldest	1,024	1,024
Mellum (KV + SWA) / 8192 / 1024	9	recent	1,408	1,408
Mellum (KV + SWA) / 8192 / 1024	11	oldest	0	0
Mellum (KV + SWA) / 8192 / 1024	11	recent	512	512
Mellum (KV + SWA) / 4096 / 1024	1	oldest	1,792	1,792
Mellum (KV + SWA) / 4096 / 1024	1	recent	1,792	1,792
Mellum (KV + SWA) / 4096 / 1024	2	oldest	1,792	1,792
Mellum (KV + SWA) / 4096 / 1024	2	recent	1,792	1,792
Mellum (KV + SWA) / 4096 / 1024	3	oldest	512	1,024
Mellum (KV + SWA) / 4096 / 1024	3	recent	512	1,408
Mellum (KV + SWA) / 4096 / 1024	5	oldest	0	0
Mellum (KV + SWA) / 4096 / 1024	5	recent	512	512
Mellum (KV + SWA) / 16384 / 1024	21	oldest	1,792	1,792
Mellum (KV + SWA) / 16384 / 1024	21	recent	1,792	1,792
Mellum (KV + SWA) / 16384 / 1024	22	oldest	1,792	1,792
Mellum (KV + SWA) / 16384 / 1024	22	recent	1,792	1,792
Mellum (KV + SWA) / 16384 / 1024	23	oldest	0	1,024
Mellum (KV + SWA) / 16384 / 1024	23	recent	896	1,408
Mellum (KV + SWA) / 16384 / 1024	25	oldest	0	0
Mellum (KV + SWA) / 16384 / 1024	25	recent	0	512
Mellum (KV + SWA) / 8192 / 2048	3	oldest	3,840	3,840
Mellum (KV + SWA) / 8192 / 2048	3	recent	3,840	3,840
Mellum (KV + SWA) / 8192 / 2048	4	oldest	3,840	3,840
Mellum (KV + SWA) / 8192 / 2048	4	recent	3,840	3,840
Mellum (KV + SWA) / 8192 / 2048	5	oldest	3,072	3,072
Mellum (KV + SWA) / 8192 / 2048	5	recent	3,456	3,456
Mellum (KV + SWA) / 8192 / 2048	7	oldest	512	1,024
Mellum (KV + SWA) / 8192 / 2048	7	recent	2,432	2,432
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	10	oldest	1,024	1,024
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	10	recent	1,024	1,024
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	11	oldest	512	1,024
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	11	recent	512	1,024
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	12	oldest	0	1,024
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	12	recent	0	512
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	14	oldest	0	0
Reduced Inkling (KV + SWA + Mamba) / 8192 / 1024	14	recent	0	512
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	3	oldest	1,024	1,024
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	3	recent	1,024	1,024
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	4	oldest	512	1,024
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	4	recent	512	1,024
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	5	oldest	0	1,024
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	5	recent	0	1,024
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	7	oldest	0	0
Reduced Inkling (KV + SWA + Mamba) / 4096 / 1024	7	recent	0	512

The sole token-ID mismatch was Mellum, cap 4096, cold 3, recent-first (one of seven requests in that episode); its causal validation is separate. All other mapping request sequences matched. Request-owned out-of-window SWA frees were not reached in these workloads, including the longer Mellum prompts.

Qwen setup, repetition, and limitations

RTX 5090 under WSL, Qwen3.5-4B snapshot 851bf6e806ef, Triton attention/Mamba/linear attention, TP1, unified memory with FULL+Mamba, page size 128, chunked prefill 512, graph capture disabled, HiCache off. Python TreeCore; Torch 2.14.1, Triton 3.8.0, sglang-kernel 0.4.9. Nominal max-total-tokens 8192, context 4096, static fraction 0.65, Mamba cache slots 128, Mamba/full ratio 0.1, extra_buffer tracking at 128 tokens. Actual allocator geometry: FULL 1,636 pages/virtual IDs at 4,194,304 bytes/page; Mamba 133 pages/virtual IDs at 51,511,296 bytes/page, sharing a 6,861,881,344-byte arena. The nominal token cap is not its physical capacity. Fixed token-ID traffic, greedy 32-token output; this is not a language-quality benchmark.

Four paired comparisons in AB/BA/BA/AB order, with 1,472 measured requests/server: 64 serial cycles × 15 requests and 16 concurrent cycles × 32 requests. Each server also ran 12 short warmup requests and two full warm cycles per workload. Four separate A/A runs are retained and excluded from A/B aggregates. Total: 11,776 paired measured requests plus 5,888 A/A requests, with 0 request errors and 0 reported retractions. All 12 servers had identical allocator geometry and clean reset/accounting checks.

Serial throughput changes ranged from -0.35% to +1.48%; concurrent changes from -0.35% to +3.04%. Reset/checkpoint work is outside each timed cycle. Client timing includes client setup and request bookkeeping; CPU is scheduler process time between cycle boundaries, not allocator-only CPU. Performance runs have no detailed hot-path diagnostic hooks. A small failed CPU import probe occurred during baseline A/A1; A/A variation is not treated as a calibrated noise floor. Conditional additional performance repeats were not triggered. Diagnostic replays are separate.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. (N/A: internal allocator bugfix with no new user-facing API or configuration.)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (Paired serving output-consistency and performance results are reported above, including observed differences and limitations; no language-quality benchmark claim.)
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #39029680033
Latest PR Test (Extra): ❌ Run #39029679972
Latest PR Test (AMD ROCm 10): ❌ Run #39029680058

[seokwoosong]

Follow-up to this PR: #33091

Fixes a case where FULL/SWA capacity checks pass individually, but the combined allocation still cannot fit in the shared pool.

[ch-wan]

[bug][P1] allocation_ready is a mutating operation passed as a loop predicate; it runs once per evicted node. The callback here is recovery.prepare_token_allocation(num_tokens), which is not a query: it calls _flush_deferred_free_group(...) unconditionally, and then _relieve_for_alloc(self, num_tokens) -- the urgent-compaction ladder that calls flush_for_allocation() on every chain member and relocates pages.

Follow where it is consumed. UnifiedRadixCache._evict_device defines target_reached(component_type), whose first statement is if allocation_ready is not None and allocation_ready(): return True, and target_reached(ct) is evaluated at the bottom of the per-node while True: walk (unified_radix_cache.py:907) as well as per component at :874. So a deep eviction round performs a deferred-free flush per evicted node, and a full two-sided compaction per node for as long as the byte bound looks satisfiable. That is O(nodes) compactions for one allocation, on the path that runs under memory pressure.

Suggestion: Split the capability in two: a pure token_allocation_ready(num_tokens) for the loop predicate (it already exists and is side-effect free), and call the mutating prepare_token_allocation only at the boundaries -- once before the walk and once after it -- rather than from target_reached. If mid-walk reclaim is genuinely needed, gate it on progress (e.g. only after N nodes, or only when live bytes dropped since the last attempt) so it cannot run per node.

[seokwoosong]

Addressed in fb1d221. The per-victim callback is now read-only, while deferred-free processing and compaction run at bounded allocator phase boundaries. Regression tests cover predicate purity and preparation counts.

[ch-wan]

[bug] The second pass is quota'd at the entire cache, so an infeasible allocation evicts everything. The fallback _evict is issued with num_tokens=self.full_evictable_size(), swa_num_tokens=self.swa_evictable_size(), mamba_num=self.mamba_evictable_size() -- i.e. every evictable entry -- and stops only when allocation_ready() turns true or the cache is exhausted. Before this change the walk was bounded by the per-component shortfall.

The stopping condition cannot be reached for a demand the pool can never serve: _token_allocation_byte_shortfall is required_bytes + live_bytes + sink_bytes - total_bytes, so once required_bytes alone exceeds total_bytes - sink_bytes, the shortfall stays positive even at zero live bytes, prepare_token_allocation keeps returning False, and the walk runs to exhaustion. The result is a fully dropped prefix cache and a failed allocation, where the old code would have evicted only the shortfall and failed. Since evict_from_tree_cache is also called speculatively from capacity probes, an over-large num_tokens reaching this path is not exotic.

Suggestion: Bound the second pass -- stop as soon as a round frees nothing (the contract the deleted tri-pool loop used), or refuse to start it when _token_allocation_byte_shortfall(num_tokens) exceeds what the whole cache could release. Add a case asserting that a request larger than the pool leaves the cache intact.

[seokwoosong]

Removed the whole-cache fallback and bounded recovery by token/state reclaim quotas. Requests proven impossible by byte or ID bounds are rejected before eviction, with regression tests verifying that cached prefixes and payloads remain intact.

[ch-wan]

[suggestion] This fixes the same bug as #36729, which rewrites the same three files with an incompatible design. The defect is real and well-demonstrated -- on main, full_available_size()=4, swa_available_size()=4 and a 4-token request take neither branch of if full_available_size < num_tokens or swa_available_size < num_tokens, so nothing is evicted although joint availability is 3. But #36729 (yonghao/ump-shared-capacity) is fixing exactly that, and the two changes collide head-on:

• both add a capability protocol to allocator/base.py;
• this PR inserts a 14-line recovery block into the body of evict_from_tree_cache, which Use a shared byte budget for unified hybrid-SWA memory #36729 deletes wholesale and replaces with a one-line delegation to allocator.evict_to_free_tokens(...);
• this PR deletes UnifiedMambaSWATokenToKVPoolAllocator.evict_to_free_tokens, which Use a shared byte budget for unified hybrid-SWA memory #36729 makes the central dispatch point (it adds two such overrides).

A textual merge will resolve to silently dropping one of the two designs. The approaches also differ in kind: #36729 computes a minimal joint reclaim plan up front (reclaim_plan(n, n, ...) over the byte model) and realizes compaction once via ensure_capacity, where this PR re-probes a mutating readiness predicate per evicted node. The former bounds both the eviction and the compaction work by construction.

Suggestion: Settle this with the #36729 author before either lands. If #36729's reclaim_plan already covers the joint case -- it is byte-joint by construction -- the most valuable part of this PR is the regression suite: test_unified_joint_allocation_eviction.py builds a real shared pool at FULL 4 / SWA 4 / joint 3, which is precisely the case #36729 has no test for. Porting those 13 cases onto #36729 would preserve the work and avoid the conflict.

[seokwoosong]

Reworked this on top of the merged #36729 design, retaining allocator-owned dispatch and the joint reclaim planner. The original two-pool case is already fixed; this revision addresses the remaining over-reclaim and three-pool recovery/cache-preservation cases, with regression coverage.

[ch-wan]

[suggestion] The branch is a no-op; both arms do the same thing. BasePrefixCache.evict_for_alloc declares allocation_ready: Optional[Callable[[], bool]] = None, so self.inner.evict_for_alloc(params, allocation_ready=None) is identical to self.inner.evict_for_alloc(params). All three overrides in the tree accept the keyword (I checked), so there is no wrapped cache that would reject it.

Suggestion: Forward unconditionally: return self.inner.evict_for_alloc(params, allocation_ready=allocation_ready).

[seokwoosong]

The forwarding override is no longer present in the rebased upstream structure, so the redundant branch is absent.

[ch-wan]

[suggestion] A predicate named *_ready that mutates the allocator will be misread. Everything at this call site reads as a query -- if allocation_ready is not None and allocation_ready(): return True inside a function called target_reached -- but the callable threaded in from common.py flushes deferred frees and can compact both sub-pools. The next person to add a target_reached() call, or to hoist it out of the loop, has no local signal that they are moving reclaim work.

Suggestion: Rename the parameter to something that states the side effect (try_recover_allocation), or -- better, and consistent with the P1 comment on common.py -- pass the pure token_allocation_ready here and keep the mutating call at the boundaries.

[seokwoosong]

The callback is now read-only and named allocation_reclaim_satisfied. It indicates sufficient reclaim, while any required layout preparation remains the allocator’s responsibility after the walk.

[seokwoosong]

@ch-wan following up on the review here.

I'm planning to rework this PR on top of #36729's allocator-owned design: keep its dispatch/reclaim planner as the base, and apply the remaining fixes for joint allocation that this PR tries to fix.

[seokwoosong]

@ch-wan

Reworked this PR on top of the merged #36729, removing the overlapping changes and focusing on the remaining over-reclaim and three-pool allocation issues. Updated the regression tests and added fresh Before/After validation results to the description.