Repository navigation
[HiCache][Perf] fix: batch HiCache D2H submits per step for hybrid pools - #39050
Merged
ispobock merged 3 commits intoSep 17, 2026
Merged
Conversation
alphabetc1
requested review from
Ying1123,
hanming-lu,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
September 11, 2026 07:48
alphabetc1
marked this pull request as draft
September 11, 2026 07:49
alphabetc1
force-pushed
the
fix-hicache-hybrid-backup-overhead
branch
from
September 11, 2026 12:10
a091af2 to
1b09614
Compare
4 of 5 tasks
alphabetc1
marked this pull request as ready for review
September 14, 2026 14:30
alphabetc1
force-pushed
the
fix-hicache-hybrid-backup-overhead
branch
from
September 16, 2026 13:49
1bb3aec to
5297707
Compare
Collaborator
Author
|
/rerun-group radix_cache/unified_radix_tree |
Contributor
|
Results for 🚀 🚀 🚀 🚀 🚀 |
ispobock
approved these changes
Sep 16, 2026
| """ | ||
| Submit host backups queued while caching this step's finished requests. | ||
| A cache whose write-back submits inline has nothing left to flush. | ||
| """ |
Collaborator
There was a problem hiding this comment.
nit: add return None or pass to keep consistent
| # The batch that queued host backups also issues them, so a node is | ||
| # never left tree-committed but un-submitted across a scheduler step. | ||
| # A cache without queued write-back inherits the base no-op. | ||
| self.tree_cache.flush_pending_backups() |
Collaborator
There was a problem hiding this comment.
should be fine, double check disable radix cache in case
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 17, 2026
Port of upstream PR sgl-project#39050 (head 9ae5181, commits 139ee26 + 5297707 + 9ae5181, base 00a9a81). Defer per-node D2H backups so one start_writing() per scheduler step merges every node queued in that step into a single transfer op: - BasePrefixCache.flush_pending_backups() no-op default; scheduler calls it after process_batch_result; UnifiedRadixCache submits pending backups from check_hicache_events and from the blocking write_back drain; StreamingSession forwards to inner. - HybridCacheController.write(flush=False) leaves the op queued; UnifiedRadixCache._execute_kv_backup passes flush=False. - MambaPoolHost.backup_from_device_all_layer converts CPU host_indices once at the entry instead of per conv/temporal tensor in _copy_tensor_all_layers_lf_pf. Manual merge note: cherry-pick -n of 00a9a81..pr/39050 auto-merged cleanly against the fork's heavily modified files; no conflict markers. Semantic review against fork-local behavior: - Linker-mode guard preserved: flush_pending_backups() early-returns when self.linker is not None, so BackupKV actions that route to linker.offload_nodes (never cache_controller.write) do not trigger a controller submit; buffer-mode pipeline keeps its own flush_pending_writes(). - Fail-soft Mamba/SWA prefetch path and the write_back blocking drain (writing_check(write_back=True) now starts with cc.start_writing() before waiting on acks) are unchanged in intent; the merged op still carries pool_transfers via CacheOperation._merge_pool_transfers. - Mamba staged-write index fix coherent with the fork's host-pool group AND-of-flag logic: HostPoolGroup.can_use_write_back_jit = all(child flags) and MambaPoolHost._init_write_back_staging_buffers keeps its flag True so start_writing() retains CPU host indices for MLA's staged kernel; the moved conversion in backup_from_device_all_layer is the same device move the deleted per-tensor hunk did, applied once at the entry for every conv/temporal component. - Test adaptation: the PR's new ChainHostPool fake needs a dtype attribute because the fork's HostPoolGroup.__init__ reads host_pool.dtype (fork-local addition); added dtype = torch.uint8. Verification: py_compile OK on all 8 changed .py; ruff check --select=F401,F821,UP037 + ruff format --check clean with XDG_CONFIG_HOME pointed at an empty dir; scripts/lint/ check_registered_tests.py whole-tree exit 0 (no new violations; both touched test files keep their register_cpu_ci + unittest.main); test_hicache_staged_write_back_dispatch.py (20 passed, 1 skipped), test_hiradix_pp_sync_drain.py (3 passed), test_unified_radix_hicache_dispatch.py + test_unified_cache_linker_fail_soft.py + test_hicache_host_register.py (27 passed), and the fork's test_npu_mamba_async_layout.py pass. Pre-existing failures outside scope: test_mem_pool_host.py (6) and test_unified_radix_cache_unittest.py (507) fail identically on the unported HEAD. CUDA/GPU runtime validation deferred.
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 17, 2026
Port of upstream PR sgl-project#39050 (head 9ae5181, commits 139ee26 + 5297707 + 9ae5181, base 00a9a81). Defer per-node D2H backups so one start_writing() per scheduler step merges every node queued in that step into a single transfer op: - BasePrefixCache.flush_pending_backups() no-op default; scheduler calls it after process_batch_result; UnifiedRadixCache submits pending backups from check_hicache_events and from the blocking write_back drain; StreamingSession forwards to inner. - HybridCacheController.write(flush=False) leaves the op queued; UnifiedRadixCache._execute_kv_backup passes flush=False. - MambaPoolHost.backup_from_device_all_layer converts CPU host_indices once at the entry instead of per conv/temporal tensor in _copy_tensor_all_layers_lf_pf. Manual merge note: cherry-pick -n of 00a9a81..pr/39050 auto-merged cleanly against the fork's heavily modified files; no conflict markers. Semantic review against fork-local behavior: - Linker-mode guard preserved: flush_pending_backups() early-returns when self.linker is not None, so BackupKV actions that route to linker.offload_nodes (never cache_controller.write) do not trigger a controller submit; buffer-mode pipeline keeps its own flush_pending_writes(). - Fail-soft Mamba/SWA prefetch path and the write_back blocking drain (writing_check(write_back=True) now starts with cc.start_writing() before waiting on acks) are unchanged in intent; the merged op still carries pool_transfers via CacheOperation._merge_pool_transfers. - Mamba staged-write index fix coherent with the fork's host-pool group AND-of-flag logic: HostPoolGroup.can_use_write_back_jit = all(child flags) and MambaPoolHost._init_write_back_staging_buffers keeps its flag True so start_writing() retains CPU host indices for MLA's staged kernel; the moved conversion in backup_from_device_all_layer is the same device move the deleted per-tensor hunk did, applied once at the entry for every conv/temporal component. - Test adaptation: the PR's new ChainHostPool fake needs a dtype attribute because the fork's HostPoolGroup.__init__ reads host_pool.dtype (fork-local addition); added dtype = torch.uint8. Verification: py_compile OK on all 8 changed .py; ruff check --select=F401,F821,UP037 + ruff format --check clean with XDG_CONFIG_HOME pointed at an empty dir; scripts/lint/ check_registered_tests.py whole-tree exit 0 (no new violations; both touched test files keep their register_cpu_ci + unittest.main); test_hicache_staged_write_back_dispatch.py (20 passed, 1 skipped), test_hiradix_pp_sync_drain.py (3 passed), test_unified_radix_hicache_dispatch.py + test_unified_cache_linker_fail_soft.py + test_hicache_host_register.py (27 passed), and the fork's test_npu_mamba_async_layout.py pass. Pre-existing failures outside scope: test_mem_pool_host.py (6) and test_unified_radix_cache_unittest.py (507) fail identically on the unported HEAD. CUDA/GPU runtime validation deferred.
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 17, 2026
Origin: sgl-project#39701 "[HiCache] Retry failed storage backup batches and expose terminal failures" Squashed upstream commits: - c311e1c fix(hicache): retry failed storage backup batches - 3edaf3b fix(hicache): harden storage backup retry path Manual merge notes: - Applied as one squashed diff over our fork state. cache_controller.py, hiradix_cache.py, unified_radix_cache.py, metrics_collector.py and the new test landed cleanly; only hybrid_cache_controller.py conflicted. - Conflict 1 (method signature): kept our fork-local `trim_prefetch_full_head` (buffer-mode storage-prefetch pipeline, sgl-project#39283/sgl-project#39050 lineage) directly above the upstream `def _page_backup(self, operation) -> bool:` return type. Both are preserved. - Conflict 2 (sidecar write body): upstream moved the batch_set_v2 call into a nested `write_sidecars()` closure so `_retry_storage_write` can retry it. Our fork-local sgl-project#39050 `extra_info=HiCacheStorageExtraInfo(prefix_keys=...)` argument is kept INSIDE that closure, so retries still forward prefix keys. - Upstream deletes `HybridCacheController.backup_thread_func` (the base implementation now handles sidecar-only ranks and terminal-failure logging). Deletion accepted; the unused `Empty` import was dropped with it. This does not regress fork-local behavior: the base thread still calls the overridden `_page_backup`, which keeps the rank-sharded-sidecar `should_backup` path. - Fork-local `_resolve_sidecar_kv_derived_pool_transfers` keeps `transfer.keys = operation.hash_value` (unconditional) instead of upstream's `if transfer.keys is None` guard; that is pre-existing fork behavior from sgl-project#39283, untouched by this PR. - Also kept from the fork: `dp_rank` / `kv_cache_dtype` storage-config args, `prefix_keys is not None` copy semantics and `batch_get_v2(extra_info=...)`. - Test registered under test/registered/unit/mem_cache/ (valid taxonomy); 12 tests pass locally (CPU-only harness).
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 17, 2026
Origin: sgl-project#39701 "[HiCache] Retry failed storage backup batches and expose terminal failures" Squashed upstream commits: - c311e1c fix(hicache): retry failed storage backup batches - 3edaf3b fix(hicache): harden storage backup retry path Manual merge notes: - Applied as one squashed diff over our fork state. cache_controller.py, hiradix_cache.py, unified_radix_cache.py, metrics_collector.py and the new test landed cleanly; only hybrid_cache_controller.py conflicted. - Conflict 1 (method signature): kept our fork-local `trim_prefetch_full_head` (buffer-mode storage-prefetch pipeline, sgl-project#39283/sgl-project#39050 lineage) directly above the upstream `def _page_backup(self, operation) -> bool:` return type. Both are preserved. - Conflict 2 (sidecar write body): upstream moved the batch_set_v2 call into a nested `write_sidecars()` closure so `_retry_storage_write` can retry it. Our fork-local sgl-project#39050 `extra_info=HiCacheStorageExtraInfo(prefix_keys=...)` argument is kept INSIDE that closure, so retries still forward prefix keys. - Upstream deletes `HybridCacheController.backup_thread_func` (the base implementation now handles sidecar-only ranks and terminal-failure logging). Deletion accepted; the unused `Empty` import was dropped with it. This does not regress fork-local behavior: the base thread still calls the overridden `_page_backup`, which keeps the rank-sharded-sidecar `should_backup` path. - Fork-local `_resolve_sidecar_kv_derived_pool_transfers` keeps `transfer.keys = operation.hash_value` (unconditional) instead of upstream's `if transfer.keys is None` guard; that is pre-existing fork behavior from sgl-project#39283, untouched by this PR. - Also kept from the fork: `dp_rank` / `kv_cache_dtype` storage-config args, `prefix_keys is not None` copy semantics and `batch_get_v2(extra_info=...)`. - Test registered under test/registered/unit/mem_cache/ (valid taxonomy); 12 tests pass locally (CPU-only harness).
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 18, 2026
Port of upstream PR sgl-project#39050 (head 9ae5181, commits 139ee26 + 5297707 + 9ae5181, base 00a9a81). Defer per-node D2H backups so one start_writing() per scheduler step merges every node queued in that step into a single transfer op: - BasePrefixCache.flush_pending_backups() no-op default; scheduler calls it after process_batch_result; UnifiedRadixCache submits pending backups from check_hicache_events and from the blocking write_back drain; StreamingSession forwards to inner. - HybridCacheController.write(flush=False) leaves the op queued; UnifiedRadixCache._execute_kv_backup passes flush=False. - MambaPoolHost.backup_from_device_all_layer converts CPU host_indices once at the entry instead of per conv/temporal tensor in _copy_tensor_all_layers_lf_pf. Manual merge note: cherry-pick -n of 00a9a81..pr/39050 auto-merged cleanly against the fork's heavily modified files; no conflict markers. Semantic review against fork-local behavior: - Linker-mode guard preserved: flush_pending_backups() early-returns when self.linker is not None, so BackupKV actions that route to linker.offload_nodes (never cache_controller.write) do not trigger a controller submit; buffer-mode pipeline keeps its own flush_pending_writes(). - Fail-soft Mamba/SWA prefetch path and the write_back blocking drain (writing_check(write_back=True) now starts with cc.start_writing() before waiting on acks) are unchanged in intent; the merged op still carries pool_transfers via CacheOperation._merge_pool_transfers. - Mamba staged-write index fix coherent with the fork's host-pool group AND-of-flag logic: HostPoolGroup.can_use_write_back_jit = all(child flags) and MambaPoolHost._init_write_back_staging_buffers keeps its flag True so start_writing() retains CPU host indices for MLA's staged kernel; the moved conversion in backup_from_device_all_layer is the same device move the deleted per-tensor hunk did, applied once at the entry for every conv/temporal component. - Test adaptation: the PR's new ChainHostPool fake needs a dtype attribute because the fork's HostPoolGroup.__init__ reads host_pool.dtype (fork-local addition); added dtype = torch.uint8. Verification: py_compile OK on all 8 changed .py; ruff check --select=F401,F821,UP037 + ruff format --check clean with XDG_CONFIG_HOME pointed at an empty dir; scripts/lint/ check_registered_tests.py whole-tree exit 0 (no new violations; both touched test files keep their register_cpu_ci + unittest.main); test_hicache_staged_write_back_dispatch.py (20 passed, 1 skipped), test_hiradix_pp_sync_drain.py (3 passed), test_unified_radix_hicache_dispatch.py + test_unified_cache_linker_fail_soft.py + test_hicache_host_register.py (27 passed), and the fork's test_npu_mamba_async_layout.py pass. Pre-existing failures outside scope: test_mem_pool_host.py (6) and test_unified_radix_cache_unittest.py (507) fail identically on the unported HEAD. CUDA/GPU runtime validation deferred.
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 18, 2026
Origin: sgl-project#39701 "[HiCache] Retry failed storage backup batches and expose terminal failures" Squashed upstream commits: - c311e1c fix(hicache): retry failed storage backup batches - 3edaf3b fix(hicache): harden storage backup retry path Manual merge notes: - Applied as one squashed diff over our fork state. cache_controller.py, hiradix_cache.py, unified_radix_cache.py, metrics_collector.py and the new test landed cleanly; only hybrid_cache_controller.py conflicted. - Conflict 1 (method signature): kept our fork-local `trim_prefetch_full_head` (buffer-mode storage-prefetch pipeline, sgl-project#39283/sgl-project#39050 lineage) directly above the upstream `def _page_backup(self, operation) -> bool:` return type. Both are preserved. - Conflict 2 (sidecar write body): upstream moved the batch_set_v2 call into a nested `write_sidecars()` closure so `_retry_storage_write` can retry it. Our fork-local sgl-project#39050 `extra_info=HiCacheStorageExtraInfo(prefix_keys=...)` argument is kept INSIDE that closure, so retries still forward prefix keys. - Upstream deletes `HybridCacheController.backup_thread_func` (the base implementation now handles sidecar-only ranks and terminal-failure logging). Deletion accepted; the unused `Empty` import was dropped with it. This does not regress fork-local behavior: the base thread still calls the overridden `_page_backup`, which keeps the rank-sharded-sidecar `should_backup` path. - Fork-local `_resolve_sidecar_kv_derived_pool_transfers` keeps `transfer.keys = operation.hash_value` (unconditional) instead of upstream's `if transfer.keys is None` guard; that is pre-existing fork behavior from sgl-project#39283, untouched by this PR. - Also kept from the fork: `dp_rank` / `kv_cache_dtype` storage-config args, `prefix_keys is not None` copy semantics and `batch_get_v2(extra_info=...)`. - Test registered under test/registered/unit/mem_cache/ (valid taxonomy); 12 tests pass locally (CPU-only harness).
4 of 5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
With
--enable-hierarchical-cacheon a hybrid model (Inkling-Small-NVFP4: 7 full-attention layers + 35 sliding-window layers + mamba conv/temporal state), end-to-end throughput is ~5% below the same run with the host tier disabled.The cost is not the D2H bandwidth.
write_throughbacks up every new radix node on its own:HybridCacheController.write()callsstart_writing()immediately, so one finished request turns into a kernel submit per pool plus one index upload per mamba conv tensor. The transfers themselves are async on a side stream, but issuing them runs on the scheduler thread. This workload finishes up to 64 requests in the same scheduler step, so the submits arrive in bursts that the overlap scheduler cannot hide, and the GPU idles until the next step is launched.Measured on the reported configuration (4x GB300, Inkling-Small-NVFP4, TP=4,
--kv-cache-dtype mxfp8, 128 concurrent lanes x 16 turns): 12,288 backup submits over the timed window for 89.6 GiB of payload.Modifications
HybridCacheController.write()takesflush: bool = True. Withflush=Falsethe operation is left on the queue instead of being submitted immediately.UnifiedRadixCache._execute_kv_backup()passesflush=False, andcheck_hicache_events()callsstart_writing()once per scheduler step. Every node backed up during that step, across all pools and all requests, becomes a single submit.writing_check(write_back=True)flushes what is still queued before waiting for acks, so the blocking path keeps its old semantics.MambaPoolHost.backup_from_device_all_layer()moves the host index tensor to the device once in the page-first path, instead of converting it again for every conv and temporal tensor inside_copy_tensor_all_layers_lf_pf.No behaviour change other than batching: the same bytes are copied by the same kernels, and the ack accounting is unchanged (one ack per submit, carrying every node id in it).
Accuracy Tests
Not applicable: no change to kernels or model forward. Host restores were byte-compared against the device copy (payload, scales, mamba conv state) on all four ranks before and after the change, and a shape-matched L1-hit vs L2-restore probe produced identical text and identical per-token logprobs.
Speed Tests and Profiling
Same workload as above, no profiler attached, two repetitions per arm:
Backup submits over the timed window drop from 12,288 to ~700 for the same 90 GiB of payload. In a torch profiler trace covering 400 steps, the scheduler stall when 64 requests finish in one step drops from 170/153 ms to 107/98 ms (the L1 baseline is 94/82 ms), and mean
cache_finished_reqfrom 1.54 ms to 1.14 ms (L1: 0.97 ms).The residual 2% is per-node Python bookkeeping in the write-through path, which is a policy question rather than a submission-path one.
Checklist
Unit tests:
test_hicache_staged_write_back_dispatch.pygains two cases (a chain offlush=Falsewrites merges into one submit that keeps both node ids; the default still flushes per call).test/registered/unit/mem_cache/is green: 2340 passed, 1366 skipped, 430 subtests.Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): 🚫 Run #35114657772
Latest PR Test (Extra): 🚫 Run #35114656848
Latest PR Test (AMD ROCm 10): ❌ Run #35114657392