Skip to content

[HiCache][Perf] fix: batch HiCache D2H submits per step for hybrid pools - #39050

Merged
ispobock merged 3 commits into
sgl-project:mainfrom
alphabetc1:fix-hicache-hybrid-backup-overhead
Sep 17, 2026
Merged

ispobock merged 3 commits into
sgl-project:mainfrom
alphabetc1:fix-hicache-hybrid-backup-overhead

Conversation

@alphabetc1

@alphabetc1 alphabetc1 commented Sep 11, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

With --enable-hierarchical-cache on 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_through backs up every new radix node on its own: HybridCacheController.write() calls start_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() takes flush: bool = True. With flush=False the operation is left on the queue instead of being submitted immediately.
  • UnifiedRadixCache._execute_kv_backup() passes flush=False, and check_hicache_events() calls start_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:

build L1 tok/s L2 tok/s gap
before 12,181 11,564 5.1%
after 12,211 / 12,159 11,913 / 11,952 2.0%

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_req from 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.py gains two cases (a chain of flush=False writes 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

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. 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
  4. 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 #35114657772
Latest PR Test (Extra): 🚫 Run #35114656848
Latest PR Test (AMD ROCm 10): ❌ Run #35114657392

@github-actions github-actions Bot added hicache Hierarchical Caching for SGLang unified-radix-cache labels Sep 11, 2026
@alphabetc1
alphabetc1 marked this pull request as draft September 11, 2026 07:49
@alphabetc1
alphabetc1 force-pushed the fix-hicache-hybrid-backup-overhead branch from a091af2 to 1b09614 Compare September 11, 2026 12:10
@alphabetc1 alphabetc1 changed the title fix: batch HiCache D2H submits per step and pipeline the ready-count all_reduce fix: batch HiCache D2H submits per step for hybrid pools Sep 11, 2026
@alphabetc1
alphabetc1 marked this pull request as ready for review September 14, 2026 14:30
@alphabetc1 alphabetc1 added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) labels Sep 16, 2026
@alphabetc1
alphabetc1 force-pushed the fix-hicache-hybrid-backup-overhead branch from 1bb3aec to 5297707 Compare September 16, 2026 13:49
@alphabetc1

Copy link
Copy Markdown
Collaborator Author

/rerun-group radix_cache/unified_radix_tree

@github-actions

github-actions Bot commented Sep 16, 2026 •

Copy link
Copy Markdown
Contributor

Results for /rerun-group radix_cache/unified_radix_tree:

🚀 4-gpu-h100 (5 tests): ❌ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/linker/test_unified_cache_linker_kl_dsv4.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

🚀 8-gpu-h200 (4 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/linker/test_unified_cache_linker_kl_glm52.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4_pp.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_glm52.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mimo.py

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

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py

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

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py

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

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py

"""
Submit host backups queued while caching this step's finished requests.
A cache whose write-back submits inline has nothing left to flush.
"""

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).
@alphabetc1 alphabetc1 changed the title fix: batch HiCache D2H submits per step for hybrid pools [HiCache][Perf] fix: batch HiCache D2H submits per step for hybrid pools Sep 17, 2026
@ispobock
ispobock merged commit 6ca866e into sgl-project:main Sep 17, 2026
334 of 371 checks passed
@alphabetc1
alphabetc1 deleted the fix-hicache-hybrid-backup-overhead branch September 17, 2026 09:57
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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

hicache Hierarchical Caching for SGLang run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) unified-radix-cache

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants