[mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py by alphabetc1 · Pull Request #30249 · sgl-project/sglang

alphabetc1 · 2026-07-06T10:29:51Z

Motivation

Part of the pool_host phase of the mem_cache refactor (issue #25371), continuing directly from #27273 ([5/N]), which extracted the shared host base layer into the pool_host/ subpackage.

That PR left every concrete host pool (MHA / MLA / Mamba / DeepSeekV4 / DSA) in memory_pool_host.py, to be relocated one family at a time. This PR relocates the MHA family — the first concrete pool family to move out of the ~3600-line memory_pool_host.py into its own pool_host/mha.py module.

Modifications

Mechanical relocation. The moved implementation bodies are byte-for-byte unchanged; only the new module's import preamble and caller import routes are rewritten.

Symbol	From	To
`MHATokenToKVPoolHost`, `MHATokenToKOnlyPoolHost`, `AsymmetricMHATokenToKVPoolHost`, `get_mha_host_pool_cls`	`memory_pool_host.py`	`pool_host/mha.py`
`_WRITE_BACK_STAGING_PAGE_CHUNK`	`memory_pool_host.py`	`pool_host/base.py`

The whole MHA family moves together because it is tightly coupled: AsymmetricMHATokenToKVPoolHost subclasses MHATokenToKVPoolHost, and get_mha_host_pool_cls dispatches between the two based on the device pool's K/V dims. Splitting them would create cross-module inheritance for no benefit.

_WRITE_BACK_STAGING_PAGE_CHUNK is not MHA-specific — it is also read by the MLA / Mamba / DeepSeekV4 / DSA host pools that stay in memory_pool_host.py. Leaving it there would force pool_host/mha.py to import back up from the aggregator module (an inverted layering that becomes a hard import cycle the moment memory_pool_host.py ever imports mha). Relocating it into pool_host/base.py — which already owns the sibling shared constant HICACHE_HOST_MEMORY_RESERVE_BYTES — keeps the dependency direction pointing down into the package, and both modules import it from there (no duplication).

Consistent with #27273, pool_host/__init__.py is not touched: the MHA symbols are imported directly from pool_host.mha at every call site, with no package-level re-export. ruff removed the now-orphaned MHA-only imports from memory_pool_host.py (MHATokenToKOnlyPool, MHATokenToKVPool, and several sgl_kernel / jit_kernel transfer helpers used only by the moved code).

Caller-site impact

12 import sites are redirected from memory_pool_host to pool_host.mha for the moved symbols; a repo-wide symbol scan confirms no other caller imports them from memory_pool_host after this PR.

Runtime (5):

kv_cache_builder.py — get_mha_host_pool_cls (method-local lazy import)
hiradix_cache.py — get_mha_host_pool_cls
hybrid_cache/hybrid_pool_assembler.py — MHATokenToKOnlyPoolHost, get_mha_host_pool_cls
disaggregation/decode_kvcache_offload_manager.py — get_mha_host_pool_cls
storage/aibrix_kvcache/unit_test.py — MHATokenToKVPoolHost

Tests + benchmark (7):

test_mem_pool_host.py, jit/test_hicache.py, jit/test_kvcacheio_asymmetric.py, test_minimax_sparse_pool_host_unit.py, benchmark/hf3fs/bench_zerocopy.py — import redirects only.
test_asymmetric_mha_pool_host_unit.py — import redirect plus 4 mock.patch targets retargeted to pool_host.mha (the Asymmetric pool's kernels are now resolved in the mha module namespace).
test_hicache_staged_write_back_dispatch.py — import redirect plus a new MHA_POOL_HOST_MODULE constant; the MHA test's 4 mock.patch targets move to it, while the MLA test's patches stay on MEMORY_POOL_HOST_MODULE (that pool is still in memory_pool_host.py).

Git blame preservation

git blame -C -C -C on pool_host/mha.py recovers the moved code's pre-refactor ancestry:

File	Traced / Total
`pool_host/mha.py`	1227 / 1233 (99.5%)

The 6 untraced lines are the rewritten import preamble. base.py gains 2 lines (the relocated constant + a blank). memory_pool_host.py is modified (not deleted), so the remaining pools' blame is preserved natively for the follow-up PRs.

Mechanical Move

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

Current verified target: 5c6b17c285e5aebb70e0b4770af4f4e3475f9581 from base 24c42c90be82c24433ab10e65db24e52a48d8084.

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

python3 <(curl -sL https://gist.githubusercontent.com/alphabetc1/bc0d1350cf7d5af01f50a27a1b59b313/raw/f7274ac087a4aefe8f56530308b41a154fb6e138/transform_move_mha_host_pool.py)

Expected output: PASS: transform reproduces the commit exactly.

Accuracy Tests

Mechanical relocation with byte-for-byte identical implementation bodies, so no model-output change is possible. Correctness of the move is established by the reproducible transform above.

The unit tests that exercise the moved symbols are dispatched via CI (I don't have an on-box CUDA runtime to attach local results):

/rerun-test test/registered/unit/mem_cache/test_mem_pool_host.py test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py test/registered/jit/test_hicache.py test/registered/jit/test_kvcacheio_asymmetric.py test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py

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 #28785703297
Latest PR Test (Extra): ✅ Run #28785703187

…ha.py

gemini-code-assist

Code Review

This pull request refactors the codebase by moving MHA-related host memory pool classes (including MHATokenToKVPoolHost, MHATokenToKOnlyPoolHost, AsymmetricMHATokenToKVPoolHost, and get_mha_host_pool_cls) from memory_pool_host.py to a new dedicated module under pool_host/mha.py, updating all corresponding imports across benchmarks, tests, and core components. The review feedback correctly points out that the import of _WRITE_BACK_STAGING_PAGE_CHUNK in memory_pool_host.py is now redundant and should be removed since the class using it was relocated.

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-07-06T10:31:58Z


 from sglang.srt.mem_cache.pool_host import HostKVCache
 from sglang.srt.mem_cache.pool_host.base import (
+    _WRITE_BACK_STAGING_PAGE_CHUNK,


The import of _WRITE_BACK_STAGING_PAGE_CHUNK is no longer needed in memory_pool_host.py because MHATokenToKVPoolHost (the only class that was using it) has been moved to pool_host/mha.py. This import can be safely removed.

alphabetc1 · 2026-07-06T10:35:58Z

/rerun-test test/registered/unit/mem_cache/test_mem_pool_host.py test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py test/registered/jit/test_hicache.py test/registered/jit/test_kvcacheio_asymmetric.py test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py

github-actions · 2026-07-06T10:36:28Z

Results for /rerun-test test/registered/unit/mem_cache/test_mem_pool_host.py test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py test/registered/jit/test_hicache.py test/registered/jit/test_kvcacheio_asymmetric.py test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py:

🚀 ubuntu-latest (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/mem_cache/test_mem_pool_host.py
cd test/ && python3 registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py
cd test/ && python3 registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py

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

cd test/ && python3 registered/jit/test_hicache.py
cd test/ && python3 registered/jit/test_kvcacheio_asymmetric.py

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

cd test/ && python3 registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py

alphabetc1 · 2026-07-07T02:43:26Z

The failing jit-kernel-b200-test (test_ltx2_qknorm_split_rope.py, bit-exact torch.equal) is a pre-existing B200 issue from #29708, unrelated to this PR — it only runs here because we touch test/registered/jit/ files. The same cases fail identically on an unrelated PR, e.g. https://github.com/sgl-project/sglang/actions/runs/28795054955/job/85383471009

hzh0425 · 2026-07-08T12:12:12Z

The diffusion test looks unrelated

hzh0425 · 2026-07-08T12:12:35Z

All other tests have passed

…gl-project#30249)

[mem_cache][6/N] refactor: move MHA host-pool family into pool_host/m…

5c6b17c

…ha.py

alphabetc1 requested review from ByronHsu, HaiShaw, ShangmingCai, Ying1123, hanming-lu, hnyls2002, hzh0425, ispobock, merrymercy, sogalin, xiezhq-hermann and yizhang2077 as code owners July 6, 2026 10:29

github-actions Bot added the hicache Hierarchical Caching for SGLang label Jul 6, 2026

gemini-code-assist Bot reviewed Jul 6, 2026

View reviewed changes

alphabetc1 added run-ci bypass-fastfail run-ci-extra labels Jul 6, 2026

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

c63c63c

ispobock assigned xiezhq-hermann, hzh0425 and ispobock Jul 7, 2026

hzh0425 approved these changes Jul 8, 2026

View reviewed changes

hzh0425 merged commit 108a183 into sgl-project:main Jul 8, 2026
430 of 464 checks passed

amd-bot mentioned this pull request Jul 9, 2026

[CI Monitor] Daily Report - 2026-07-09 bingxche/sglang-ci-bot#130

Open

alphabetc1 deleted the refactor/mem-cache-poolhost-mha branch July 9, 2026 02:31

michaelzhang-ai mentioned this pull request Jul 10, 2026

[AMD] Cherry-pick AMD fixes into release/v0.5.15 #30722

Closed

3 tasks

This was referenced Jul 14, 2026

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

Open

[mem_cache][7/N] refactor: move MLATokenToKVPoolHost to pool_host.mla #30616

Merged

[mem_cache][8/N] refactor: move MambaPoolHost to pool_host.mamba #31180

Merged

alphabetc1 mentioned this pull request Aug 20, 2026

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

Open

6 tasks

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

[mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py (s…

b221a46

…gl-project#30249)

Navigation Menu

Uh oh!

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

[mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py#30249
hzh0425 merged 2 commits into
sgl-project:mainfrom
alphabetc1:refactor/mem-cache-poolhost-mha

alphabetc1 commented Jul 6, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 6, 2026

Uh oh!

alphabetc1 commented Jul 6, 2026

Uh oh!

github-actions Bot commented Jul 6, 2026 •

edited

Loading

Uh oh!

alphabetc1 commented Jul 7, 2026

Uh oh!

hzh0425 commented Jul 8, 2026

Uh oh!

hzh0425 commented Jul 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

alphabetc1 commented Jul 6, 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 Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

alphabetc1 commented Jul 6, 2026

Uh oh!

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

Uh oh!

Uh oh!

alphabetc1 commented Jul 7, 2026

Uh oh!

hzh0425 commented Jul 8, 2026

Uh oh!

hzh0425 commented Jul 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

alphabetc1 commented Jul 6, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Jul 6, 2026 •

edited

Loading