Skip to content

[Fix] unified-memory: Bound joint reclaim and preserve cached prefixes - #44827

Closed
ch-wan wants to merge 3 commits into
mainfrom
cheng/hot-fix/unified-joint-reclaim
Closed

ch-wan wants to merge 3 commits into
mainfrom
cheng/hot-fix/unified-joint-reclaim

Conversation

@ch-wan

@ch-wan ch-wan commented Oct 17, 2026 •

Copy link
Copy Markdown
Collaborator

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

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

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.

This continues #39294, preserving Seokwoo Song's original commits, rebased onto b62a48ee15b94ba428660ca7e021d44946b0802d.

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.
  • 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

Fresh validation of this rebased change:

  • 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 remains in #39294 and is not represented as 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.

CI States

Latest PR Test (Base): ❌ Run #38988089292
Latest PR Test (Extra): ❌ Run #38988089227
Latest PR Test (AMD ROCm 10): ❌ Run #38988089311

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants