[mem_cache][5/N] refactor: extract host KV cache base layer into pool_host package by alphabetc1 · Pull Request #27273 · sgl-project/sglang

alphabetc1 · 2026-06-04T11:44:33Z

Motivation

First PR of the pool_host phase of the mem_cache refactor (issue #25371). Sibling of the completed allocator series (#26675, #26676).

memory_pool_host.py is a ~2960-line module holding the host-side KV cache base class, shared allocation helpers, and every concrete host pool (MHA / MLA / Mamba / DeepSeekV4 / DSA). This PR extracts only the shared base layer into a new pool_host/ subpackage. All concrete pool classes stay in memory_pool_host.py and will be relocated one family at a time in follow-up PRs (mha, mla, mamba, deepseek_v4, nsa, group), all tracked under #25371.

Modifications

Mechanical relocation. The moved implementation bodies are unchanged; only import preambles and caller import routes are rewritten.

Symbol	From	To
`HostKVCache` (abc), `synchronized`, `HICACHE_HOST_MEMORY_RESERVE_BYTES`	`memory_pool_host.py`	`pool_host/base.py`
`HostTensorAllocator`, `get_allocator_from_storage`, `alloc_with_host_register`, `alloc_with_pin_memory`, `ALLOC_MEMORY_FUNCS`	`memory_pool_host.py`	`pool_host/common.py`
`HiSparseHostPoolMixin`	`memory_pool_host.py`	`pool_host/hisparse.py`

HiSparseHostPoolMixin gets its own pool_host/hisparse.py (rather than base.py or mla.py) because it is a concrete mixin shared by multiple host pools (MLATokenToKVPoolHost and DeepSeekV4PagedHostPool, both still in memory_pool_host.py). Keeping it separate from mla.py avoids making the DeepSeekV4/paged host path depend on an MLA-family module in follow-up moves.

pool_host/__init__.py intentionally re-exports only the package-level public API used by callers today: HostKVCache and HostTensorAllocator. Lower-level helpers stay importable from their owning submodules (pool_host.base, pool_host.common, pool_host.hisparse). The Mooncake allocator import inside get_allocator_from_storage stays lazy (function-local), matching the original.

memory_pool_host.py keeps every concrete pool class. It imports the package-level HostKVCache from pool_host; imports base-only helpers (HICACHE_HOST_MEMORY_RESERVE_BYTES, synchronized) from pool_host.base; imports host allocation helpers (ALLOC_MEMORY_FUNCS, get_allocator_from_storage) from pool_host.common; and imports HiSparseHostPoolMixin from pool_host.hisparse for its still-local implementations. ruff removed the now-orphaned abc, defaultdict, wraps, alloc_mmap, and KVCache imports from it.

Caller-site impact

12 import sites across 11 files are redirected away from memory_pool_host for moved symbols:

7 sites: bare HostKVCache imports now use sglang.srt.mem_cache.pool_host; 2 of these are inside TYPE_CHECKING blocks.
1 site: mooncake_store.py splits the not-yet-moved MLATokenToKVPoolHost import from memory_pool_host while moving HostKVCache / HostTensorAllocator to pool_host.
4 sites: ALLOC_MEMORY_FUNCS / alloc_with_pin_memory now import from pool_host.common (jit_kernel/tests/test_hicache.py, test_dsa_pool_host_unit.py, plus 2 method-local lazy imports in test_hisparse_unit.py).

A repo-wide symbol scan confirms no external caller imports the moved symbols from memory_pool_host after this PR; memory_pool_host.py itself still imports the helpers because its concrete host pools remain there for now.

Git blame preservation

git blame -C -C -C on the three relocated-code files recovers their pre-refactor ancestry from memory_pool_host.py and earlier history:

File	Traced / Total
`base.py`	171 / 176
`common.py`	86 / 93
`hisparse.py`	69 / 71

= 326 / 340 (95.9%) of relocated lines traced; the rest are rewritten import preambles. pool_host/__init__.py is a new 7-line re-export module. memory_pool_host.py is modified (not deleted), so the concrete pools' blame is preserved natively for the follow-up PRs.

Mechanical Move

Transform script (self-contained, reproduces this PR's current diff byte-for-byte from the current PR base): https://gist.github.com/alphabetc1/34850017427438e4b26fad6fb3874d2a

Current verified target: 7e7e3820c26ba0609f017b99358172a9776197de from base e1955bf57ab39a4de7f7daf9a619e4f68f952ed9.

One-click verification — run from the repo root with Python >= 3.10:

python3.10 <(curl -sL https://gist.githubusercontent.com/alphabetc1/34850017427438e4b26fad6fb3874d2a/raw/5cd0e3a845800883de7ca210eec603835a031f59/transform_pool_host_base.py)

Expected output: PASS: transform reproduces the commit exactly.

Accuracy Tests

The unit tests that exercise the moved symbols (host pools + allocation funcs) all pass on an NVIDIA H20:

python -m pytest -q \
    test/registered/unit/mem_cache/test_dsa_pool_host_unit.py \
    test/registered/unit/managers/test_hisparse_unit.py \
    python/sglang/jit_kernel/tests/test_hicache.py

21 passed, 23 warnings in 50.18s

Speed Tests and Profiling

N/A — mechanical relocation only, no hot-path change.

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.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ✅ Run #27803114008
Latest PR Test (Extra): ❌ Run #27803113902

gemini-code-assist

Code Review

This pull request refactors the host-side memory pool implementation by splitting memory_pool_host.py into a modular package under pool_host/ and updating all dependent imports. Feedback on these changes includes correcting the page alignment calculation in HostKVCache to prevent allocating an unnecessary extra page, and relaxing the type hint for the allocator parameter in alloc_with_pin_memory to avoid static type-checking issues.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-04T11:46:52Z

+        self.page_num = self.size // self.page_size + 1
+        self.size = self.page_num * self.page_size


Using self.size // self.page_size + 1 to align the host memory pool size always allocates an extra page even when self.size is already a multiple of self.page_size. It is more efficient and standard to use the round-up formula (self.size + self.page_size - 1) // self.page_size.

self.page_num = (self.size + self.page_size - 1) // self.page_size\n self.size = self.page_num * self.page_size

gemini-code-assist · 2026-06-04T11:46:52Z

+    dtype: torch.dtype,
+    device: str,
+    pin_memory: bool,
+    allocator: None,


The type hint allocator: None is overly restrictive and will cause static type checkers to complain when a HostTensorAllocator instance is passed (which happens in memory_pool_host.py). Changing it to allocator: HostTensorAllocator | None = None makes it compatible with the signature of alloc_with_host_register and type-safe.

Suggested change

allocator: None,

allocator: HostTensorAllocator | None = None,

…_host package

alphabetc1 · 2026-06-05T12:28:36Z

/tag-and-rerun-ci

xiezhq-hermann · 2026-06-10T06:48:33Z

is this file necessary? should we combine it with the base.py?

I think base.py is for abstract base classes, and utilities shared across modules should go in a separate file.

hzh0425

looks good

hzh0425 · 2026-06-15T16:41:35Z

please resolve the conflicts and rerun ci

…o refactor/mem-cache-poolhost

alphabetc1 · 2026-06-19T14:21:53Z

/rerun-group radix_cache/unified_radix_tree/hicache

github-actions · 2026-06-19T14:22:14Z

⛔ radix_cache/unified_radix_tree/hicache: Unknown test group radix_cache/unified_radix_tree/hicache.

Known groups: 4-gpu-models, 8-gpu-models, amd, ascend, attention, backends, bench_fn, chunked_prefill, constrained_decoding, core, cp, cpu, cuda_graph, debug_utils, disaggregation, dllm, dp_attn, dp_engine, ep, eplb, eval, function_call, gb300, hicache, input_embedding, jit, kernels, kv_canary, layers, lora, mla, mock_model, model_loading, models, models_e2e, moe, musa, observability, openai_server, ops, perf, pp, prefill_only, profiling, quant, radix_cache, reasoning, rl, rotary, sampling, scheduler, scripted_runtime, sessions, spec, stress, tokenizer, unit, utils, vlm, xpu

alphabetc1 · 2026-06-19T14:27:25Z

/rerun-group radix_cache hicache

github-actions · 2026-06-19T14:27:50Z

Results for /rerun-group radix_cache hicache:

🚀 1-gpu-h100 (4 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/test_cpp_radix_cache.py
cd test/ && python3 registered/radix_cache/test_swa_radix_cache_kl.py
cd test/ && python3 registered/hicache/test_hicache_spec_file_storage.py
cd test/ && python3 registered/hicache/test_hicache_variants.py

🚀 4-gpu-h100 (7 tests): ❌ View workflow run

cd test/ && python3 registered/radix_cache/test_int8_mamba_checkpoint_e2e.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py
cd test/ && python3 registered/hicache/test_hicache_storage_3fs_backend.py
cd test/ && python3 registered/hicache/test_qwen35_hicache.py

🚀 1-gpu-5090 (3 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/test_radix_attention.py
cd test/ && python3 registered/radix_cache/test_radix_cache_hit.py
cd test/ && python3 registered/hicache/test_hicache_storage.py

🚀 2-gpu-h100 (6 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py
cd test/ && python3 registered/hicache/test_hicache_spec_mooncake_storage.py
cd test/ && python3 registered/hicache/test_hicache_storage_file_backend.py
cd test/ && python3 registered/hicache/test_hicache_storage_mooncake_backend.py
cd test/ && python3 registered/hicache/test_hicache_storage_runtime_attach_detach.py

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py

alphabetc1 · 2026-06-20T03:45:53Z

View workflow run

This failed CI is irrelevant

ispobock · 2026-06-20T10:58:24Z

/rerun-test registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

github-actions · 2026-06-20T10:58:47Z

Results for /rerun-test registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

…_host package (sgl-project#27273)

alphabetc1 requested review from BBuf, DarkSharpness, HydraQYH, Ying1123, celve, hanming-lu, hnyls2002, hzh0425, ispobock, merrymercy, xiezhq-hermann, yizhang2077 and yuan-luo as code owners June 4, 2026 11:44

github-actions Bot added hicache Hierarchical Caching for SGLang jit-kernel labels Jun 4, 2026

gemini-code-assist Bot reviewed Jun 4, 2026

View reviewed changes

alphabetc1 force-pushed the refactor/mem-cache-poolhost branch from a0a9628 to fb7fcd7 Compare June 4, 2026 11:51

[mem_cache][5/N] refactor: extract host KV cache base layer into pool…

d3b6125

…_host package

alphabetc1 force-pushed the refactor/mem-cache-poolhost branch from fb7fcd7 to d3b6125 Compare June 5, 2026 11:18

Merge branch 'main' into refactor/mem-cache-poolhost

52761b9

alphabetc1 added the run-ci label Jun 5, 2026

alphabetc1 added 4 commits June 6, 2026 15:23

[mem_cache] import HostKVCache from pool_host package

5ec9b53

[mem_cache] fold host tensor allocator into pool_host base

ce38d9f

[mem_cache] move host pool allocation helpers to common

674c38d

[mem_cache] preserve pool_host package comments

7e7e382

ispobock assigned xiezhq-hermann and hzh0425 Jun 7, 2026

alphabetc1 mentioned this pull request Jun 7, 2026

[RFC][Refactor] mem_cache pool / allocator restructure #25371

Open

37 tasks

ispobock mentioned this pull request Jun 9, 2026

[Hisparse][Feature] CPU memory share across TPs #27370

Open

5 tasks

Merge branch 'main' into refactor/mem-cache-poolhost

2b0df15

xiezhq-hermann reviewed Jun 10, 2026

View reviewed changes

alphabetc1 added the bypass-fastfail label Jun 13, 2026

hzh0425 approved these changes Jun 15, 2026

View reviewed changes

hzh0425 added the run-ci-extra label Jun 15, 2026

Merge branch 'main' into refactor/mem-cache-poolhost

9f863fa

alphabetc1 force-pushed the refactor/mem-cache-poolhost branch from 1ded307 to 9f863fa Compare June 16, 2026 03:09

alphabetc1 and others added 2 commits June 16, 2026 03:25

Merge branch 'main' into refactor/mem-cache-poolhost

806f994

Merge branch 'main' into refactor/mem-cache-poolhost

8d1fcaa

alphabetc1 requested a review from huangtingwei9988 as a code owner June 16, 2026 13:49

xiezhq-hermann approved these changes Jun 18, 2026

View reviewed changes

alphabetc1 added the high priority label Jun 18, 2026

alphabetc1 and others added 4 commits June 18, 2026 13:52

Merge branch 'main' into refactor/mem-cache-poolhost

d4c9d96

Merge branch 'main' into refactor/mem-cache-poolhost

d4b9163

Merge branch 'main' into refactor/mem-cache-poolhost

9994092

Merge remote-tracking branch 'origin/refactor/mem-cache-poolhost' int…

81bf7fc

…o refactor/mem-cache-poolhost

ispobock merged commit ff1fc1f into sgl-project:main Jun 20, 2026
355 of 379 checks passed

alphabetc1 deleted the refactor/mem-cache-poolhost branch June 20, 2026 14:55

alphabetc1 mentioned this pull request Jul 6, 2026

[mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py #30249

Merged

5 tasks

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[mem_cache][5/N] refactor: extract host KV cache base layer into pool…

b21d345

…_host package (sgl-project#27273)

		self.page_num = self.size // self.page_size + 1
		self.size = self.page_num * self.page_size

	allocator: None,
	allocator: HostTensorAllocator \| None = None,

Uh oh!

Conversation

alphabetc1 commented Jun 4, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Caller-site impact

Git blame preservation

Mechanical Move

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

alphabetc1 commented Jun 5, 2026

Uh oh!

xiezhq-hermann Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

alphabetc1 Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

hzh0425 left a comment

Choose a reason for hiding this comment

Uh oh!

hzh0425 commented Jun 15, 2026

Uh oh!

alphabetc1 commented Jun 19, 2026

Uh oh!

github-actions Bot commented Jun 19, 2026

Uh oh!

alphabetc1 commented Jun 19, 2026

Uh oh!

github-actions Bot commented Jun 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alphabetc1 commented Jun 20, 2026

Uh oh!

ispobock commented Jun 20, 2026

Uh oh!

github-actions Bot commented Jun 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

alphabetc1 commented Jun 4, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Jun 19, 2026 •

edited

Loading

github-actions Bot commented Jun 20, 2026 •

edited

Loading