Skip to content

[mem_cache][12/N] refactor: stop allocators round-tripping through the KV pool - #28494

Open
alphabetc1 wants to merge 3 commits into
sgl-project:mainfrom
alphabetc1:refactor/allocator-kvcache
Open

[mem_cache][12/N] refactor: stop allocators round-tripping through the KV pool#28494
alphabetc1 wants to merge 3 commits into
sgl-project:mainfrom
alphabetc1:refactor/allocator-kvcache

Conversation

@alphabetc1

@alphabetc1 alphabetc1 commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Part of #25371. Rebased onto current main and renumbered -- the branch had been sitting
since 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:

  1. get_cpu_copy / load_cpu_copy pass-throughs. TokenToKVPoolAllocator and
    SWATokenToKVPoolAllocator each define these purely to forward to self._kvcache,
    and BaseTokenToKVPoolAllocator declares them only to raise NotImplementedError.
    Retraction backup is a pool operation; routing it through the allocator adds a layer
    that carries no decision.
  2. Round-trips through the pool to read a tensor the allocator itself owns. The
    allocator allocates full_to_swa_index_mapping / full_to_hisparse_device_index_mapping
    in __init__, hands the same object to the pool via register_mapping, and writes
    it on every alloc and free. self._kvcache.translate_loc_from_full_to_swa(x) is then
    literally self.full_to_swa_index_mapping[x] with an extra call frame.

Modifications

  • Drop get_cpu_copy / load_cpu_copy from allocator/base.py, allocator/token.py,
    allocator/paged.py, allocator/swa.py. The two call sites in Req.save_kv_cache /
    Req.load_kv_cache go through token_to_kv_pool_allocator.get_kvcache() instead.
  • SWATokenToKVPoolAllocator.translate_loc_from_full_to_swa and the HiSparse allocators'
    get_last_loc_hisparse_device / free_hisparse / free_compressed index the
    allocator-owned mapping directly.
  • full -> compressed is deliberately left as a pool call: it is genuine C4-layout
    arithmetic ((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 main baseline on the
same box:

      rebased branch:   1746 passed, 1074 skipped, 201 subtests passed
  pristine main:        1746 passed, 1074 skipped, 201 subtests passed

(test/registered/unit/mem_cache/ plus spec/test_resolve_swa_kv_pool.py; --ignore on
test_umbp_store.py, which needs the mori package that is absent from the image.) Zero
delta, 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 -- main had added
translate_kv_indices_for_transfer to allocator/base.py, _copy_for_free_group to
allocator/token.py, and PureSWATokenToKVPoolAllocator plus swa_free_group to
allocator/swa.py. All are preserved; PureSWATokenToKVPoolAllocator is byte-identical
to main, and it overrides translate_loc_from_full_to_swa with an identity, so it is
unaffected by the change here. The fourth conflict was real: main had reworked
retraction backup onto RetractionBackup(cpu_tensors=...), so the resolution keeps that
structure and only reroutes the call through get_kvcache().

The ownership assumption was re-checked against current main rather than trusting the
June-era diff: allocator/swa.py builds the tensor, register_mapping hands the same
object to the pool, and the pool's translate_loc_from_full_to_swa indexes that very
tensor -- so the direct index is identical, including the -1 sentinel row. A repo-wide
grep confirms no allocator-level get_cpu_copy / load_cpu_copy caller or override
survives; every remaining one is on a pool.

Benchmark and Profiling Results

Not applicable -- removes call frames, changes no kernel or memory layout.

Checklist

  • Format your code according to the Code Formatting with Pre-Commit.
  • Add unit tests as outlined in the Running Unit Tests.
  • Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
  • Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling.
  • For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
  • Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

🤖 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

@gemini-code-assist

Copy link
Copy Markdown
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
alphabetc1 force-pushed the refactor/allocator-kvcache branch from 0be0ba5 to 3022f90 Compare August 20, 2026 07:54
@alphabetc1 alphabetc1 changed the title [mem_cache][6/N] refactor: decouple allocator runtime from KV pool [mem_cache][12/N] refactor: stop allocators round-tripping through the KV pool Aug 20, 2026
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.

1 participant