Repository navigation
Conversation
|
@ispobock @xiezhq-hermann this is a reload perf fix. When an evicted prefix comes back from the host KV tier it runs the per-token kernel (~4.3 GB/s) even when the host bytes are contiguous, so this adds a bulk/staging fast-path that moves the slab in one copy. Measured 1.49x on a 128k reload (GLM-5.2-FP8, 0.694s vs 1.036s), byte-identical to the per-token path with a clean fallback for the non-contiguous cases. Same as my other HiCache PR (#33883), it's stuck behind run-ci and I can't set that label from my account, so a |
|
@ispobock @xiezhq-hermann this one is mergeable and just needs run-ci to fire the suite. It's the bulk/staging reload fast-path: 1.49x on a 128k evicted-prefix reload (0.694s vs 1.036s) with byte-identical output. Could you add run-ci and take a review pass? |
5005989 to
69354a2
Compare
|
@xiezhq-hermann @ispobock rebased and mergeable here, but its blocked on run-ci which I cant trigger myself. Could one of you add run-ci and review the HiCache staging reload path? |
|
@ispobock @xiezhq-hermann following up on the HiCache staging-reload fast-path. Rebased and mergeable, blocked on run-ci which I cant add myself. Could you add it and review when you can? |
When an evicted prefix is reloaded from the host KV tier, the transfer runs the per-token kernel (`transfer_kv_per_layer_mla`) at ~4.3 GB/s no matter how the bytes are laid out, because it pays a per-token indexing cost for every token. The host pool makes that worse: it hands out free slot ids first-fit and appends freed ids on free, so after a few fill/evict cycles a re-admitted prefix lands on a scrambled set of host offsets. So the reload is a fine-grained scatter/gather. Two changes turn that reload into a large contiguous copy: 1. The host slot allocator hands out a contiguous run whenever one is free, so a prefix's KV occupies a contiguous host region. Capacity and the success/failure set are identical to the stock allocator (it only returns None when there isn't enough total free, same as before), and when the pool is too fragmented for one run it falls back to gathering the largest runs first. So the slots handed out are always valid, distinct and in range; only the packing changes. 2. A reload fast-path in `load_to_device_per_layer`. When the host run for a layer is one ascending unit-stride run, it moves the slab in one shot instead of the per-token gather: if the device target slots are also one run, a single `copy_`; if the device slots are scattered (the device pool allocator does not pack them), it copies the contiguous host slab into a contiguous device staging buffer once and places the rows at their real slots with an on-device `index_copy_`. This covers both halves of a DSA reload, the MLA KV anchor and the DSA indexer, which ride the same prefix. Any non-contiguous host case, dtype/shape mismatch, or a non-contiguous destination falls back to the exact stock kernel, so a reload is never wrong and the path with nothing to gain is unchanged. Correctness is byte-identical to the per-token kernel. A CPU check reproduces the kernel's per-index gather over the whole tensor for the KV anchor, the DSA indexer page shape, and the scattered-device staging tier, and confirms the fast-path rejects dtype/shape/non-contiguous inputs so they fall back. The reload lands the data at the exact device slots the radix cache recorded, so the model reads the right KV. Measured served on GLM-5.2-FP8 (tp8, host L2 tier, `--hicache-io-backend kernel --hicache-mem-layout layer_first`, KV in bf16 so the baseline reload is the jit indexed kernel): a 128k evicted-prefix reload is 0.694 s with the fast-path vs 1.036 s on stock (median of 3 each), a 1.49x speedup, staging tier firing on every KV and indexer layer with no fallback. Both sides return the same reloaded answer (the planted needle) with cache_frac 0.9996, so the reload is byte-correct. The clock isn't pinnable on this box so treat it as first-pass; the ~5% run-to-run spread is well inside the 1.49x gap. The isolated per-layer copy is much faster (~57 vs ~4-5 GB/s); the served win is smaller because the reload also pays the fixed resident-hit and scheduling floor, which the copy doesn't touch. Signed-off-by: adlashab <adlashab@amd.com>
69354a2 to
485e540
Compare
|
@hzh0425 @huangtingwei9988 rebased onto current main (head 485e540) so it clears the CI min base. This is the HiCache staging reload fast path, it bulk-copies contiguous host runs instead of page by page: 1.49x on a 128k evicted-prefix reload (0.694s vs 1.036s) with byte-identical output. I don't have run-ci access, could one of you add the label and review the pool_host change? |
|
@huangtingwei9988 can you help take a look as this primarily touches changes you priorly did |
…t head take_contiguous_run can hand out a run from further down free_slots, so an in-use slot leaked to the head was no longer checked and test_double_alloc stopped raising. alloc() now also checks the head slots that first-fit would have handed out, so a leak there trips the assert again even when the contiguous run comes from somewhere else in the list. Also add a unit test for the contiguous pick after a sorted merge. Signed-off-by: adlashab <adlashab@amd.com>
|
@huangtingwei9988 @xiezhq-hermann thanks for the merge from main and the run-ci. The base-a red in test_mem_pool_host.py is from this PR. test_double_alloc puts a used slot at the head of free_slots, and the contiguous run lookup in alloc() picks the free run behind it, so the double-alloc assert never fired (the retry then got None because the pool was used up by then). I pushed f344806 on top of your merge. alloc() now also checks the head slots that first-fit would have handed out, so a leaked slot there trips the assert again, and I added a small test for the contiguous pick. Locally test_mem_pool_host.py was 21/22 before and is 23/23 now, and the whole unit/mem_cache dir fails the same set as main in my env. CI on that commit stopped at the min-base gate (60878c3 went into main after your merge), so I merged main once more in ca5086f, no conflicts and main didnt touch the PR files. The other reds look unrelated to me. AMD tracing, kda and minimax_indexer_cp fail the same way on other PRs today, the mi35x kda/mxfp8 jobs are kernel tests that never build the host pool, a few mi300 jobs died with "runner lost communication", and the NPU HiCache failure is in pool_host/mha.py which this PR doesnt touch (all NPU base-b jobs are red on other PRs too). |
Signed-off-by: adlashab <adlashab@amd.com>
pool_host/base.py: import conflict only, main dropped the KVCache import, kept take_contiguous_run. alloc() and _merge_release_slots() are unchanged on main. pool_host/mla.py merged without a conflict but the layer_first kernel reload still passed device_pool.kv_buffer[device_layer_id] to bulk_reload. After sgl-project#43711 that name only exists on the kernel_ascend path, so point bulk_reload at the captured self.device_kv_buffers[layer_id] like the jit and transfer_kv_per_layer_mla calls next to it (this also picks the right buffer for packed MTP draft layers). Signed-off-by: adlashab <adlashab@amd.com>
|
@xiezhq-hermann @huangtingwei9988 the branch went conflicting after the pool_host rework on main, so I merged main into it again in b74d982 (plain merge, no rebase). The only textual conflict was the import block in pool_host/base.py, main dropped the KVCache import. alloc() and _merge_release_slots() didnt change on main, so the contiguous-run pick and the head check for double-alloc carry over as they were. The part that mattered was mla.py. It merged without a conflict, but our bulk_reload call still passed device_pool.kv_buffer[device_layer_id], and after #43711 device_layer_id only exists on the kernel_ascend path, so every layer_first kernel reload would raise UnboundLocalError. I pointed it at the captured self.device_kv_buffers[layer_id], same as the jit and transfer_kv_per_layer_mla calls right below it, which also gives the right buffer for MTP draft layers. dsa.py still indexes device_pool.index_k_with_scale_buffer on main so nothing to change there. Net diff vs main is the same as before apart from that one line. Ran it on MI355X in rocm/sgl-dev:v0.5.21-rocm720-mi35x-20261002, main d88d5c3 vs the merge (CPU tests):
pre-commit passes on the changed files. |
|
@xiezhq-hermann @huangtingwei9988 CI on b74d982 finished. base-a and the GPU base-b shards are all green, but 5 of the base-c shards (b200 x3, h20, h200) got cancelled when the run was stopped at 01:51 UTC on 10-13, so base-c never completed. Could one of you rerun those? I went through the other reds, none of them come from this diff:
After base-c it only needs a review of the pool_host change. |
This branch has not been deployed
4 participants
When an evicted prefix is reloaded from the host KV tier, the transfer runs the per-token kernel (
transfer_kv_per_layer_mla) at ~4.3 GB/s no matter how the bytes are laid out, because it pays a per-token indexing cost for every token. The host pool makes that worse: it hands out free slot ids first-fit and appends freed ids on free, so after a few fill/evict cycles a re-admitted prefix lands on a scrambled set of host offsets. So the reload is a fine-grained scatter/gather.Two changes turn that reload into a large contiguous copy:
The host slot allocator hands out a contiguous run whenever one is free, so a prefix's KV occupies a contiguous host region. Capacity and the success/failure set are identical to the stock allocator (it only returns None when there isn't enough total free, same as before), and when the pool is too fragmented for one run it falls back to gathering the largest runs first. So the slots handed out are always valid, distinct and in range; only the packing changes.
A reload fast-path in
load_to_device_per_layer. When the host run for a layer is one ascending unit-stride run, it moves the slab in one shot instead of the per-token gather: if the device target slots are also one run, a singlecopy_; if the device slots are scattered (the device pool allocator does not pack them), it copies the contiguous host slab into a contiguous device staging buffer once and places the rows at their real slots with an on-deviceindex_copy_. This covers both halves of a DSA reload, the MLA KV anchor and the DSA indexer, which ride the same prefix. Any non-contiguous host case, dtype/shape mismatch, or a non-contiguous destination falls back to the exact stock kernel, so a reload is never wrong and the path with nothing to gain is unchanged.Correctness is byte-identical to the per-token kernel. A CPU check reproduces the kernel's per-index gather over the whole tensor for the KV anchor, the DSA indexer page shape, and the scattered-device staging tier, and confirms the fast-path rejects dtype/shape/non-contiguous inputs so they fall back. The reload lands the data at the exact device slots the radix cache recorded, so the model reads the right KV.
Measured served on GLM-5.2-FP8 (tp8, host L2 tier,
--hicache-io-backend kernel --hicache-mem-layout layer_first, KV in bf16 so the baseline reload is the jit indexed kernel): a 128k evicted-prefix reload is 0.694 s with the fast-path vs 1.036 s on stock (median of 3 each), a 1.49x speedup, staging tier firing on every KV and indexer layer with no fallback. Both sides return the same reloaded answer (the planted needle) with cache_frac 0.9996, so the reload is byte-correct. The clock isn't pinnable on this box so treat it as first-pass; the ~5% run-to-run spread is well inside the 1.49x gap. The isolated per-layer copy is much faster (~57 vs ~4-5 GB/s); the served win is smaller because the reload also pays the fixed resident-hit and scheduling floor, which the copy doesn't touch.CI States
Latest PR Test (Base): 🚫 Run #38279147817
Latest PR Test (Extra): ❌ Run #38279147543
Latest PR Test (AMD ROCm 10): ❌ Run #38279147381