Repository navigation
fix(sglang): split --hicache-size between host pools as native SGLang does - #600
Conversation
|
/ci run |
kvcached's default layout interleaves all layers (and, for MHA, K with
V) in one buffer, so a layer's MHA or MLA view has a token stride that
spans every layer, and a Mamba state view has a slot stride that spans
every layer's states. SGLang HiCache backs MHA and MLA pools up through
per-layer pointer tables that assume dense rows, and copies Mamba states
assuming dense slots. The elastic MHA pool leaves the pointer tables
unset, so the first backup fails; the elastic MLA pool publishes them,
so MLA backups copy rows of other layers and tokens without any error.
Seen on SGLang 0.5.20 with --enable-hierarchical-cache on an A100:
Qwen2.5-7B and Qwen3.5-4B stop at the first prefill with
File ".../pool_host/mha.py", line 393, in _resolve_device_transfer_buffers
device_pool.k_data_ptrs,
AttributeError: 'ElasticMHATokenToKVPool' object has no attribute 'k_data_ptrs'
and DeepSeek-V2-Lite serves all 30 prompts, but every one that is loaded
back from the host diverges from native SGLang at its first token, with a
token outside native's top 5. With KVCACHED_CONTIGUOUS_LAYOUT=false,
where each Mamba layer has a buffer of its own, Qwen3.5 stops at the
first copy of a Mamba state to the host with
RuntimeError: ... transfer_mamba: item_size must be 16-byte aligned (uint4)
The host pools now copy through SGLang's one-layer JIT HiCache kernel,
which reads each view's own stride: MHA and MLA backups one layer at a
time (their loads already used that kernel), and Mamba states in both
directions. The kernel copies fixed-size rows and is compiled for each
row size, so a Mamba state is copied as rows of at most 4 KiB of
kvcached's flat state buffer rather than as one megabyte-sized row.
Host layouts, MTP draft pools, layer sharding and decode context
parallelism that these copies do not cover are refused at startup. MHA
and MLA pools with dense views (KVCACHED_CONTIGUOUS_LAYOUT=false) keep
SGLang's copies. The Mamba host pool sizes itself from one tensor of all
layers, so with per-layer buffers it is shown the layers stacked, and the
states are copied one layer at a time.
SGLang splits --hicache-size between a hybrid model's device pools in proportion to get_kv_size_bytes(). kvcached's elastic MHA pools report the size of their buffers, which span the whole virtual reservation, about a GPU per pool, so the MHA host pool takes almost all of the host memory. Seen on SGLang 0.5.20 with Qwen3.5-4B, --enable-hierarchical-cache and --hicache-size 14 (with the strided HiCache copies applied): the Mamba host pool gets 42 slots (2.16 GB), where native SGLang gives it 129. With sync scheduling (--disable-overlap-schedule) each prompt keeps a Mamba state per prefill chunk, and 16 prompts sent twice did not hit the host cache once on the second pass. The split now counts an elastic pool's logical capacity: its reported bytes scaled by its token capacity over its buffers' length. Native pools and Mamba pools are counted as before. The same 16 prompts then hit 4096 tokens each, the prefill chunk boundary that native SGLang also hits in this mode.
b9aca45 to
ceadcb0
Compare
939da65 to
51950c2
Compare
RixinLiu
changed the title
rishabhsinha17
left a comment
rishabhsinha17
left a comment
There was a problem hiding this comment.
LGTM. the split counted the virtual reservation, which starved the Mamba host pool, and the fix counts the capacity SGLang sized its tree with.
i verified which pool the split sees and the number it gets. build_hybrid_mamba_stack hands kvcache.full_kv_pool itself to _split_hicache_size (hybrid_pool_assembler.py:1485). scaling by (size + page_size) over k_buffer[0].shape[0] equals get_kv_size_bytes_phy() in both layouts, caps at the native total for a request larger than the reservation, returns 0 for a zero-layer pool with no division, and leaves Mamba pools alone. pool.size is only set in init, so a shrink does not move the split.
51950c2 is the same delta as 939da65 rebased onto ceadcb0. the basic gpu run on 939da65 ended in a cancel 28 seconds in, right after the image pulls and before any test ran, so the new head still needs a /ci run.
cpu manifest at 51950c2 1518 passed, failure set identical to main. 11 tests in the file pass. lint and mypy clean.
shipiyouniao
left a comment
shipiyouniao
left a comment
There was a problem hiding this comment.
LGTM for the sizing change. Using logical pool capacity rather than the virtual reservation matches the native host-budget calculation, while native and Mamba pools keep their existing accounting. The direct-backend issue in #599 still needs fixing in the base.
- Preserve logical pool sizing on the latest HiCache copy fixes. - Cover pool capacity bounds, budget conservation and native fallback.
shipiyouniao
left a comment
shipiyouniao
left a comment
There was a problem hiding this comment.
LGTM at 0288c80. The baseline sync preserves the sizing change; current-head CPU, lint/type checks and private Codex review pass.
#599 is merged; this PR is based on main, including the direct-backend index-device fix.
Problem
SGLang splits
--hicache-sizebetween a hybrid model's device pools in proportion toget_kv_size_bytes(). kvcached's elastic MHA pools report the size of their buffers, which span the whole virtual reservation (about a GPU per pool), while Mamba pools report their slots. The MHA host pool therefore takes almost all of the host memory.On an A100 with SGLang 0.5.20, Qwen3.5-4B,
--enable-hierarchical-cache --hicache-size 14:Native SGLang gives the same model a 7.40 GB MHA host pool and 129 Mamba slots. With sync scheduling (--disable-overlap-schedule), 16 prompts sent twice did not hit the host cache once on the second pass. A separate run that sent only 4 prompts, with everything else the same, hit 4096 tokens for each prompt on the second pass, so the 42 Mamba slots are too few for 16 prompts.
Change
The split counts an elastic pool's logical capacity: its reported bytes scaled by its token capacity over its buffers' length. Native pools and Mamba pools are counted as before, and
get_kv_size_bytes()itself is unchanged for its other callers.Validation
0288c809,tests/test_sglang_hicache.pyhas 24 cases, including MHA/MLA capacity bounds, zero-layer pools, budget conservation and native fallback. The CPU suite passes on Python 3.10–3.13: 1588 passed, 15 skipped per version. Pre-commit, mypy for all four versions, and GPU-free C++ tests pass. These are constructed regressions; no new hybrid serving run was performed.KVCACHED_MAX_CACHED_TOKENS=0, 16 prompts of 4.5k-7.5k tokens sent twice):