Repository navigation
Support unified memory page-envelope transfers in PD - #39477
Merged
Merged
Conversation
Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
Co-authored-by: yhzhuang <yhzhuang@fb.com>
Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
Minimize both component reclaim quotas under the joint capacity predicate so a blocked compaction path cannot turn one allocation shortfall into an all-SWA eviction. Original prod_inference commit: 56c7082a92bc0c9024585387460e64b065652918 Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
Co-authored-by: yhzhuang <yhzhuang@fb.com>
# Conflicts: # python/sglang/srt/mem_cache/common.py
# Conflicts: # python/sglang/srt/mem_cache/common.py
# Conflicts: # python/sglang/srt/mem_cache/common.py
…pacity # Conflicts: # python/sglang/srt/managers/schedule_policy.py # python/sglang/srt/managers/scheduler_components/invariant_checker.py # python/sglang/srt/mem_cache/kv_cache_configurator.py # python/sglang/srt/mem_cache/multi_ended_allocator.py # python/sglang/srt/model_executor/pool_configurator.py
# Conflicts: # python/sglang/srt/disaggregation/base/conn.py # python/sglang/srt/mem_cache/kv_cache_configurator.py
ch-wan
reviewed
Sep 16, 2026
ZYHowell
commented
Sep 17, 2026
ZYHowell
commented
Sep 17, 2026
ZYHowell
commented
Sep 17, 2026
`DecodePreallocQueue` carried both capacity algorithms and picked between them with `isinstance(allocator, UnifiedSWATokenToKVPoolAllocator)`: a per-side token comparison for pools whose sides own their own buffers, and a shared-byte reservation for the unified layout. The scheduler had to reconstruct the allocator's own accounting to ask the second question -- adding `full_available_size() + evictable - budget` back into the demand -- and naming one class meant the answer depended on which sibling a layout happened to subclass. Following `check_decode_capacity`, the pool now answers and the scheduler only calls. Three methods, each with the separate-buffer behaviour as the default and the shared-envelope behaviour as an override: - `prealloc_fits` prices a (full, swa) demand against the scheduler's budgets. The default compares per side and never reads `tree_cache`; the override reads what the tree could reclaim and prices the whole ask in bytes. - `reclaim_for_prealloc` frees room and names the shortfall it could not meet. Defaults to the sliding-window evict on `SWATokenToKVPoolAllocator`, since only SWA-family allocators reach it. - `has_shared_byte_envelope` states the layout fact behind both consequences: the sides are priced together, and the answer is post-reclaim, so admitting on it still owes the reclaim. `swa_capacity_and_available` needed no override at all -- the base already returned exactly the pair the static branch was recomputing. Behaviour is unchanged for every layout, including the tri-pool: its `can_reserve` refuses an evictable allowance, so it keeps the inherited per-side default it already took under the old type test. One branch remains in `_check_if_req_exceed_kv_capacity`, marked in place: its two sides bound a different length, so folding them would change which requests are refused. Scheduler tests that admit requests now bind the real separate-buffer implementations onto their allocator double, or a bare `MagicMock` returns a truthy `Mock` and the decision under test stops being made anywhere. Verified on cheng-wan-h200-8gpu: 266 passed / 2 skipped / 1403 subtests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…tion Two bugs from review of the previous commit, plus its two comment-style slips. **Tri-pool PD admission could crash the decode node.** `full_available_size()` and `swa_available_size()` both take `schedulable_available_size()`, which credits the peer's drainable holes, so each side reports capacity backed by the *same* shared gap. Comparing each against its own token budget therefore double-counts those bytes: a pair that fits neither together was admitted, `alloc_extend_swa_tail` then priced them jointly through `_fits_page_demand` and returned None, and `_pre_alloc` asserts on that (`KV cache is full! Bug in memory estimation`) rather than refusing the request. The previous commit's note that the tri-pool "keeps the per-side default it already took" described the control flow correctly but did not make per-side admission safe, now that tail allocation is asymmetric on one envelope. `UnifiedMambaSWATokenToKVPoolAllocator.prealloc_fits` now prices the pair on the float chain's own page grid before applying the budgets. The budgets still bind -- they carry decode headroom the allocator cannot see. **Hybrid-Mamba (non-SWA) decode retraction raised NotImplementedError.** Unified PD is forced onto `cpu_tensor`, and `Req.offload_kv_cache` calls the allocator's `get_cpu_copy`; `UnifiedMambaTokenToKVPoolAllocator` inherited the raising base. It now translates virtual FULL ids to physical before the pool, the way `UnifiedSWAKVPool` already does for the SWA layout. `UnifiedMLATokenToKVPool` grew the physical->kernel hop its own class docstring states, since the MLA parent indexes `kv_buffer` by kernel-facing ids -- delegating without it would have read the wrong rows silently. `has_shared_byte_envelope` bundled two contracts that the tri-pool answers differently: it is one buffer, but its `prealloc_fits` prices the pool as it stands rather than post-reclaim. Split into `prealloc_fits_assumes_reclaim` (owed reclaim) and `prealloc_ceiling_fits` (an ask that can never fit, or None when the pool has no ceiling of its own), which also lets `_check_if_req_exceed_kv_capacity` drop its last layout test. Dropped two comments that narrated the refactor rather than stating a fact the code does not -- that rationale belongs in the PR, per the repo comment rules. Verified on cheng-wan-h200-8gpu: 278 passed / 2 skipped / 1403 subtests. The new tri-pool case fails without the fix (admits an infeasible pair). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous regression case asked for `2 * available_size()` on both sides with artificial budgets. It did fail without the fix, but through the budget conjunct rather than the joint grid -- it was not the shape the production bug takes. Scanning the whole (full, swa) space on the tri-pool fixture: at page_size 1 there is no pair that each side can host alone yet the grid refuses, in any of five pool states. The slack only opens on a paged grid -- at page_size 4 there are 16 such pairs, and asking for exactly each side's own `available_size` is one of them. The case now pins that, and asserts the per-side precondition alongside the joint refusal so the double-count is visible in the test rather than implied. Verified red without the fix. Also trimmed the `prealloc_fits` docstring to the two facts a reader cannot recover here, dropping the `_pre_alloc` assert story -- same rule as the two comments removed in the previous commit. Verified on cheng-wan-h200-8gpu: 309 passed / 2 skipped / 1443 subtests across 18 files, including the MLA/envelope/gate tests that cover the pools the previous commit touched and the 1-GPU SWA tail allocation test. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commit fixed two halves of that path by reading the code and ran neither. Both now have a case, and both were verified red without their fix. `UnifiedMLATokenToKVPool.get_cpu_copy` / `load_cpu_copy` round-trip through PHYSICAL ids at page_size 1 and 4, checked against the kernel-id formula the class docstring states. Without the rewrite the parent indexes `kv_buffer` at a different row, so the restore returns other tokens' KV -- silently, which is why a round-trip assert is the guard rather than a raise. `UnifiedMambaTokenToKVPoolAllocator.get_cpu_copy` / `load_cpu_copy` hand the pool physical ids. The case pins that the v2p table is not identity for the allocated run first, so a delegate that passed `req_to_token`'s virtual ids straight through -- the shape the fix replaced -- fails it rather than coincidentally agreeing. Both live in `test_unified_mla_views.py`, which already had the CPU unified MLA + mamba fixture and the `_kernel_id` helper. The two call sites need a stated `dcp_enabled` (and `attn_dcp_size` for the composite), so they take `get_parallel().override(...)` rather than publishing a whole ServerArgs. Verified on cheng-wan-h200-8gpu: 311 passed / 2 skipped / 1445 subtests across 18 files. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Only conflict is the import block in `test_decode_radix_lock_ref.py`, where main added `CustomTestCase` next to the allocator-double helper this branch introduced; both imports are kept. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Main tightened the registered-test layout: kernel tests live under
test/registered/kernels/{ops,benchmark}/<group>/. This one landed at
test/registered/kernel/mem_cache/ before that, and the sync makes it a
lint failure.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
ch-wan
added
bypass-fastfail
run-ci-extra
Collaborator
|
/tag-and-rerun-ci |
Bring in the unified-memory rerun group and current main updates.
Collaborator
Author
|
/rerun-group unified-memory |
Contributor
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 |
Use the existing separate-buffer allocator double and publish the runtime configuration needed by the preallocation capacity checks. Preserve the HiSparse host-backed capacity assertions.
Collaborator
Author
|
/rerun-group unified-memory |
Contributor
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 |
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
4 participants
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.
Original PR: #36730 — previous reviews and comments.
Motivation
Unified memory pools expose virtual token IDs while disaggregated transfer backends operate on physical buffers. PD transfers need an explicit page-envelope contract and physical index translation, including distinct target and independent-draft index vectors when their layouts differ. Decode preallocation should charge SWA only for newly allocated tail pages, including the Mamba/SWA/FULL tri-pool layout.
Targets
main, with merged #36729 supplied by the existing base. This review update preserves the current base and propagates through #39478 and #39479.Modifications
UnifiedSWAAllocatorBase. Price only newly bound SWA pages, derive their IDs from the allocation, and preserve an already-bound partial prefix page. Remove the duplicate two-pool override and its page-alignment assertions.self.token_to_kv_pool_allocatorconsistently in the decode reservation helpers.merge_and_sort_free()calls from unified paged allocation: that method only updatesfree_pages/release_pages, while these kernels allocate fromfree_virtual_idsand already reject requests with insufficient virtual pages. Physical free-page sorting remains in compaction; the removed calls cannot replenish or reorder the virtual IDs used by these kernels.Accuracy Tests
Focused allocator tests verify tail binding, physical translations, prefix preservation and release accounting. Full model accuracy and multi-node PD serving were not rerun for this review update.
Speed Tests and Profiling
No serving benchmark or profiling run; no new performance result is claimed.
Test plan
Using Python 3.12 in
/data/venvs/pd-kl-tier5with local source onPYTHONPATH:git diff --checkpassed.Checklist
Review and Merge Process
The original PR linked at the top contains prior reviews and comments. Review and CI status for this replacement PR must be evaluated separately.
CI States
Latest PR Test (Base): ❌ Run #35379101664
Latest PR Test (Extra): 🚫 Run #35379101178
Latest PR Test (AMD ROCm 10): ❌ Run #35379101701