Performance Optimization: Page Allocator Migrates from Python to C++

[lianghao208]

issue #299

[cui36]

Thanks @lianghao208!

[cui36]

Hi @lianghao208,

Thanks for the C++ migration. After running it, I want to propose two must-fixes and a design decision.

Must-fix (pushed a new commit)

1. torch_bindings.cpp is missing #include <pybind11/functional.h> (and <pybind11/stl.h> for the new vector / unordered_map returns). Without it, set_should_use_worker_ipc_callback rejects any Python callable, and KVCacheManager.__init__ always raises TypeError at kv_cache_manager.py:111. All 5 tests in test_kvcache_manager.py regress from PASSED on main to ERROR on this PR.

2. The test fixture in tests/test_kvcache_manager.py accesses Python-only attributes that don't exist on the C++ allocator: enable_page_prealloc, min_reserved_pages, reserved_page_list, num_total_pages. Replace with get_num_reserved_pages() / get_num_total_pages().

Design decision: still support kvctl limit? (proposed in a new PR: #323)

The remaining issues all stem from one question: whether the engine should still react to external writers of the shm total_size:

1. kv_cache_manager.py:230-234 comments out the resize poll.
2. C++ derives the IPC name as kvcached_engine_<pgid> while Python / kvctl use kvcached_<engine_tag>_<pgid>, the two sides aren't even looking at the same /dev/shm segment, so even if fix [sglang-integration]: kvcached's free() should be called outside free_group #1 were re-enabled, no signal would cross.

If kvctl stays a published feature (it has docs in examples/02_memory_control/, a console-script entry in setup.py, and a working test on main), all three need fixing. Cost: ~7 μs/alloc for the shm poll, but the microbench still retains ~83% of the C++ migration's speedup over main (alloc+free: main 65 μs → PR #319 23 μs → with fixes 30 μs). More optimization details in PR #323.

A future optimization that keeps kvctl working and recovers the ~7 μs would be a watcher thread that polls shm at low frequency and exposes a single atomic flag to the alloc hot path.

Note that no matter adding kvctl or not, passive elasticity (unmap-on-free, physical GPU memory floating between co-located instances via CUDA VMM) is unaffected either way. Only the active control plane is at stake here.

[lianghao208]

@cui36 Thanks for the review, they all make sense to me. I notice the PR #323 has already fixed all the correspondence problems.