[mem_cache][12/N] refactor: stop allocators round-tripping through the KV pool - #28494
Open
alphabetc1 wants to merge 3 commits into
Open
[mem_cache][12/N] refactor: stop allocators round-tripping through the KV pool#28494alphabetc1 wants to merge 3 commits into
alphabetc1 wants to merge 3 commits into
Conversation
alphabetc1
requested review from
Ying1123,
hnyls2002,
ispobock,
merrymercy and
xiezhq-hermann
as code owners
June 17, 2026 05:37
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
alphabetc1
force-pushed
the
refactor/allocator-kvcache
branch
from
August 20, 2026 07:54
0be0ba5 to
3022f90
Compare
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
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.
Motivation
Part of #25371. Rebased onto current
mainand renumbered -- the branch had been sittingsince June and
[6/N]had meanwhile been taken by #30249.Allocators reach through the KV pool for two things they do not need it for:
get_cpu_copy/load_cpu_copypass-throughs.TokenToKVPoolAllocatorandSWATokenToKVPoolAllocatoreach define these purely to forward toself._kvcache,and
BaseTokenToKVPoolAllocatordeclares them only to raiseNotImplementedError.Retraction backup is a pool operation; routing it through the allocator adds a layer
that carries no decision.
allocator allocates
full_to_swa_index_mapping/full_to_hisparse_device_index_mappingin
__init__, hands the same object to the pool viaregister_mapping, and writesit on every alloc and free.
self._kvcache.translate_loc_from_full_to_swa(x)is thenliterally
self.full_to_swa_index_mapping[x]with an extra call frame.Modifications
get_cpu_copy/load_cpu_copyfromallocator/base.py,allocator/token.py,allocator/paged.py,allocator/swa.py. The two call sites inReq.save_kv_cache/Req.load_kv_cachego throughtoken_to_kv_pool_allocator.get_kvcache()instead.SWATokenToKVPoolAllocator.translate_loc_from_full_to_swaand the HiSparse allocators'get_last_loc_hisparse_device/free_hisparse/free_compressedindex theallocator-owned mapping directly.
full -> compressedis deliberately left as a pool call: it is genuine C4-layoutarithmetic (
(i + 1) % compress_ratio), a pool responsibility, not a mapping lookup.Comments at both sites say so, to keep the next reader from "simplifying" it too.
No behavior change: every removed indirection resolved to the same tensor index or the
same pool method.
Accuracy Test
Rebase resolution, verified on an H200 devbox against a pristine
mainbaseline on thesame box:
(
test/registered/unit/mem_cache/plusspec/test_resolve_swa_kv_pool.py;--ignoreontest_umbp_store.py, which needs themoripackage that is absent from the image.) Zerodelta, so the 1074 skips are pre-existing.
The rebase kept the original diff shape exactly (+32/-45 across 7 files). Three of the
four conflicts were context drift around the deleted methods --
mainhad addedtranslate_kv_indices_for_transfertoallocator/base.py,_copy_for_free_grouptoallocator/token.py, andPureSWATokenToKVPoolAllocatorplusswa_free_grouptoallocator/swa.py. All are preserved;PureSWATokenToKVPoolAllocatoris byte-identicalto
main, and it overridestranslate_loc_from_full_to_swawith an identity, so it isunaffected by the change here. The fourth conflict was real:
mainhad reworkedretraction backup onto
RetractionBackup(cpu_tensors=...), so the resolution keeps thatstructure and only reroutes the call through
get_kvcache().The ownership assumption was re-checked against current
mainrather than trusting theJune-era diff:
allocator/swa.pybuilds the tensor,register_mappinghands the sameobject to the pool, and the pool's
translate_loc_from_full_to_swaindexes that verytensor -- so the direct index is identical, including the
-1sentinel row. A repo-widegrep confirms no allocator-level
get_cpu_copy/load_cpu_copycaller or overridesurvives; every remaining one is on a pool.
Benchmark and Profiling Results
Not applicable -- removes call frames, changes no kernel or memory layout.
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #32693512747
Latest PR Test (Extra): ❌ Run #32693512660
Latest PR Test (AMD ROCm 7.2): ❌ Run #32693512729