Skip to content

[Fix] hicache: vote the prefetch timeout once per tick, not per admission poll - #47620

Open
xiezhq-hermann wants to merge 1 commit into
mainfrom
hicache/scheduler-thread-transfers
Open

xiezhq-hermann wants to merge 1 commit into
mainfrom
hicache/scheduler-thread-transfers

Conversation

@xiezhq-hermann

@xiezhq-hermann xiezhq-hermann commented Oct 31, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

Under --hicache-storage-prefetch-policy timeout (the default) the per-request deadline verdict is agreed by a collective issued from inside the admission loop: check_prefetch_progress calls _can_terminate_prefetch, which on main does one _all_reduce(MAX) per call (python/sglang/srt/mem_cache/unified_radix_cache.py:2480-2490 on 1ad0149814). The scheduler issues that call once per waiting request with a live prefetch per scheduling pass (python/sglang/srt/managers/scheduler.py:4543 on b50d439494), and the loop's exits around it (batch_is_full from get_num_allocatable_reqs, scheduler.py:4406-4411 and 4528-4540 on b50d439494) are evaluated by each rank on its own. That is the shape reported in #30760 (TP=4, no PP): one rank examined one more candidate than its peers before its own batch_is_full tripped, called the collective one more time, and waited in it forever while the others had moved on to recv_requests. #44614 centralized the per-request budget (_StoragePrefetchWait, started at the first blocked poll, _prefetch_deadline_expired, _mark_prefetch_deadline_expired) but left the vote in the poll; #36425 had moved the vote into _can_terminate_prefetch; #33029 and #44340 (both closed, not merged) proposed a once-per-tick vote, and this PR reuses that part of #44340.

This PR removes the per-poll collective that #30760 reported and makes the remaining rank-consistency premise (every rank polls the same requests) fail loud in the same tick on every rank instead of hanging. Even when the counts match, the per-poll vote costs one gloo all_reduce per waiting request per pass (prior measurements: roughly 60-80 us at TP2 and 0.6-0.9 ms at TP8), paid for every request parked behind a slow storage lookup.

Modifications

All in python/sglang/srt/mem_cache/unified_radix_cache.py; no new flags, no new collective.

  • _vote_prefetch_deadlines (lines 2657-2685) runs once per scheduler tick from check_hicache_events, after the storage drain, under storage only (lines 4162-4164). It packs one int per live prefetch, in ongoing_prefetch insertion order, evaluated on PP rank 0 with the existing _prefetch_deadline_expired (a wait that has not started yet contributes 0), does ONE _all_reduce(MAX) (line 2681: TP MAX on PP0, relayed along PP by _pp_sync, lines 425-459 (_all_reduce and _pp_sync)), and for every 1 calls the existing _mark_prefetch_deadline_expired and cache_controller.terminate_prefetch (lines 2684-2685), so every rank records the same verdict in the same tick and the transfer worker and the hit drain stop spending on the cut span at once (the drain already handles a terminated operation: is_terminated() at lines 3447 and 3575, with the deadline reason at 3580 and no retry at 3589). With no live prefetch no rank issues a collective. The vote runs only under policy timeout; best_effort and wait_complete never vote.
  • _can_terminate_prefetch under timeout (lines 2647-2654) returns the recorded verdict (wait.expired) and issues no collective. It keeps @rank_consensus(same_results=True): the verdict is rank-agreed by the vote. Since the vote also terminates, a poll on an expired operation normally takes the existing is_terminated() path (line 2796); retirement (terminate_prefetch, the deadline reason, revoke_pending_prefetch / _handle_prefetch_result, lines 2803-2812) is unchanged and still happens at the poll, or at the lockstep drain when the hit or the final ack lands first (lines 3575-3592 and 3613-3632). A parked hit cut by the vote is finished with the deadline reason before its revoke (lines 3447-3451), as the poll would.
  • check_prefetch_progress is now collective-free for all three policies: best_effort returns True (line 2642), wait_complete returns False (line 2646), timeout reads the verdict (line 2654). The only collective reachable from an admission poll is the PP ticket path's own MAX in _check_pp_prefetch_progress (line 2721), left as is. The remaining _all_reduce* call sites in the file are the tick's own syncs (_sync_hicache_ready_counts lines 3873/3875, writing_check, loading_check, the linker round), the synchronous hit query helper (line 2418) and drain_storage_control_queues, which only the GPU test kits call.
  • Rank consistency of the vote's shape, made fail-loud: ongoing_prefetch is filled by prefetch_from_storage (@rank_consensus, line 2425) and emptied by the lockstep drain and by the admission polls. The polls retire per rank, so the number and length of the packed vote would depend on the premise that every rank polls the same requests ([Bug] HiCache prefetch all_reduce deadlock with TP=4, no PP — mismatched call count in check_prefetch_progress #30760's maintainer position: get_num_allocatable_reqs, hence the poll sequence, is identical across TP ranks). _hicache_sync_digest (lines 2687-2705) folds the live prefetch set (rids and attempt ids in insertion order, crc32) into the digest the ready-count MIN sync already carries as [digest, -digest] (line 3847; asserts at 3855-3858 for the shm path and 3878-3881 for the gloo path, which also covers PP stages through the relay). A divergent set therefore raises on every rank at the start of the next tick, before any vote is issued, instead of one rank blocking in a vote its peer never issues or the two voting on tensors of different length. The row width is unchanged, so the shm fast path (8-slot table) is unaffected; the fold is active only when the vote would read the dict (storage on, policy timeout, non-empty). Main is exposed to the same premise one step earlier (the poll's own collective) and hangs; other policies keep main's behavior.
  • In PP ticket mode (buffer_only, PP > 1) no entry is ever kept in ongoing_prefetch: PP > 0 returns at line 2448 (get_prefetch_submission returns an empty submission, hybrid_cache_controller.py:1245) and PP0 returns at line 2580 (submit_prefetch returns decision=True, hybrid_cache_controller.py:1316), so the vote never runs there and the PP relay of _all_reduce is only exercised in cache-mode PP.

Behavior changes to be aware of: a request whose budget ran out is cut on every rank at the next tick's vote instead of at its own poll, so the cut can land up to one tick later than before; between the vote and the poll the operation is terminated but still listed, so a hit or final ack landing in that window is settled by the drain (same outcome accounting path as a controller-terminated query: revoked outcome plus deadline reason) rather than by the poll.

Companion note for downstream forks that subclass the tick (prod_inference's Meta subclass runs its own packed buffer-mode vote and per-poll agreement before calling super().check_hicache_events()): after this change that subclass should either override _vote_prefetch_deadlines to a no-op in buffer mode or drop its own vote in favor of this one, to avoid two packed collectives per tick; both feed the same _mark_prefetch_deadline_expired / wait.expired state, so they agree either way.

Rebase notes (2026-11-05). Rebased on main 9b6c796151 without conflicts; no code change. The two base-a-test-cpu failures on the previous CI run (test_decode_queue_cleanup.py, test_prefill_cuda_graph_dcp.py) were main's and are fixed upstream by #47758 and #47736.

Accuracy Tests

No model-output change: the policy's verdict is the same _prefetch_deadline_expired evaluation, agreed by the same MAX, read by the same poll.

CPU unit tests, test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py, TestStoragePrefetchTimeout (line 456):

  • test_polls_never_reduce_and_the_tick_votes_once (line 457): five live prefetches, four polled in one pass issue 0 collectives; the tick issues exactly one MAX over [0,0,0,0,0]; at a time long past the budget the four polls still return False and issue nothing, the next tick votes [1,1,1,1,0] (the never-polled request's wait never started), the four voted operations are terminated by the tick and the parked one is not, every polled request is then cut by its poll with no retry armed; best_effort, wait_complete and an empty ongoing_prefetch issue no collective.
  • test_expiry_is_agreed_even_before_a_peer_starts_waiting (line 511, two-rank gloo, _two_rank_prefetch_deadline line 262): only rank 0 started the wait and only its clock ran out; rank 0 polls (returns False, no collective) while rank 1 does not poll; the tick marks the operation expired on both ranks with one collective each; the poll on both returns True; neither rank re-issues the prefetch. Adapted from main's version, which asserted that the poll itself returned True on both ranks: that assertion pinned the per-poll vote this PR removes, so the agreement is now asserted after the tick, and the polls are asserted not to reduce.
  • test_uneven_polls_before_a_verdict_issue_no_collective (line 522, two-rank gloo, _two_rank_uneven_polls line 299): [Bug] HiCache prefetch all_reduce deadlock with TP=4, no PP — mismatched call count in check_prefetch_progress #30760's trigger before any verdict; both ranks start two waits, then rank 0 polls two requests and rank 1 one; no collective is issued by the polls, the tick votes [1,1] once on each rank, and both ranks retire both requests with the same verdicts. Against the old per-poll vote (main's _can_terminate_prefetch body patched in) it fails at its first uneven poll, which votes and cuts by itself.
  • test_divergent_live_prefetch_sets_raise_on_every_rank_in_the_same_tick (line 532, two-rank gloo, _two_rank_divergent_live_sets line 336): [Bug] HiCache prefetch all_reduce deadlock with TP=4, no PP — mismatched call count in check_prefetch_progress #30760's shape after a verdict; after an agreed [1,1] vote rank 0 polls two requests and rank 1 one, so the live sets differ by one entry; the next tick's ready-count sync raises the digest assertion on both ranks (checked with all_gather_object), and the sequence of collectives is identical on both (MIN, MAX, MIN: no vote was issued with mismatched lengths). With the fold removed from _hicache_sync_digest (scratch negative control) the same worker fails differently: rank 0 issues no vote and the peer waits alone in its one-element MAX until the gloo timeout (Timed out waiting 30000ms for send operation to complete).
  • The CPU fixtures these cases run the tick on (_arm_tick, line 216) also pre-set storage-worker state (storage_stop_event, the worker thread handles, storage_workers_alive) and backup-retry state (_host_capacity_backup_retries, _backup_write_specs) that nothing on main 9b6c796151 reads; the sibling PRs of this stack (storage-worker liveness, backup retry) read them from the tick, so the lines are inert here and keep the fixtures valid with those PRs applied.
  • test/registered/unit/mem_cache/test_hicache_pp_sync_drain.py: test_pp_relays_the_deadline_vote_after_the_counts (line 90) runs the tick on a cache-mode PP0 and a PP1 stage, each with one live prefetch: PP0 reads its own clock, issues the counts (8-slot MIN) and then the vote [1] (MAX), marks and terminates; PP1 contributes [0], never evaluates _prefetch_deadline_expired, receives [1] second on the same path and marks and terminates the same way. The fixture (lines 32-35) now carries prefetch_stop_policy and an empty ongoing_prefetch; the two existing cases are unchanged, and the buffer-mode follower's _all_reduce.assert_not_called() now also pins that an empty vote issues nothing.
  • test/registered/unit/managers/scheduler_components/test_shm_tag_exchange.py: the _hicache_tp_rank fixture (lines 54-55) now carries prefetch_stop_policy (wait_complete) and an empty ongoing_prefetch, which _hicache_sync_digest reads inside the _sync_hicache_ready_counts the cases drive directly; the three cases are unchanged (the fold is inactive under wait_complete, so the 8/9/10-value rows take the same shm or gloo path as before).

Results on the development box (python -m unittest, on the committed tree 02bbd5b820 rebased on 9b6c796151; CPU unless stated):

  • test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py: Ran 19 tests in 41.326s, OK (main: 16).
  • test/registered/unit/mem_cache/test_hicache_pp_sync_drain.py: Ran 3 tests in 0.350s, OK (main: 2).
  • test/registered/unit/managers/scheduler_components/test_shm_tag_exchange.py: Ran 3 tests in 21.222s, OK.
  • GPU, one NVIDIA B200 (CUDA 13.0): test/registered/unit/mem_cache/unified_radix_cache/test_unified_radix_cache_storage.py: Ran 281 tests in 39.880s, OK (skipped=57); identical on untouched main 9b6c796151. (The No accelerator errors quoted in the previous revision came from a mis-set CUDA_VISIBLE_DEVICES on the development box, not from the tests.)
  • GPU, same B200: test/registered/unit/mem_cache/unified_radix_cache/test_unified_radix_cache_hicache.py (the whole file): Ran 341 tests in 17.219s, OK (skipped=46); identical on untouched main at 9b6c796151 (Ran 341 tests in 14.914s, OK, skipped=46).
  • pre-commit clean on the changed files.

Speed Tests and Profiling

Collectives per scheduling pass with N waiting requests under timeout: main 1 (ready-count sync) + N; this PR 2 (ready-count sync + one packed vote), 1 when no prefetch is live. The digest fold adds a crc32 over the live rids per tick on every rank (no extra slot, no extra collective). Not re-measured on GPU for this revision; the per-request gloo cost quoted in Motivation (60-80 us at TP2, 0.6-0.9 ms at TP8) is a prior measurement, not one from this revision.

Not run: GPU serving, PP > 1 end to end, PD disaggregation, mooncake/hf3fs/nixl backends; the GPU unit suite above (one NVIDIA B200, CUDA 13.0) is the only GPU evidence. The repository CI re-runs on push.

Checklist


CI States

Latest PR Test (Base): 🚫 Run #42136241276
Latest PR Test (Extra): ❌ Run #42136241343
Latest PR Test (AMD ROCm 10): ⏳ Run #42136241257

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

Tip: disable this comment in your organization's Code Review settings.

@github-actions

This comment has been minimized.

@xiezhq-hermann xiezhq-hermann added the run-ci CI: run the baseline test suite on this PR label Oct 31, 2026
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-agreement branch from feaa104 to 137bbe4 Compare October 31, 2026 21:29
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 3388194 to 8bf47fd Compare October 31, 2026 21:29
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-agreement branch from 137bbe4 to 0b8be56 Compare October 31, 2026 21:48
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 8bf47fd to 910a82e Compare October 31, 2026 21:48
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-agreement branch from 0b8be56 to 7dd5f54 Compare October 31, 2026 23:02
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 910a82e to cad2377 Compare October 31, 2026 23:02
@github-actions

This comment has been minimized.

@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from cad2377 to fc23921 Compare October 31, 2026 23:21
@github-actions

This comment has been minimized.

@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-agreement branch from 7dd5f54 to 1d8a495 Compare October 31, 2026 23:36
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from fc23921 to c641dc1 Compare October 31, 2026 23:36
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-agreement branch from 1d8a495 to ded5f1f Compare November 1, 2026 01:54
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from c641dc1 to 9687866 Compare November 1, 2026 01:54
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-agreement branch from ded5f1f to 30ea114 Compare November 1, 2026 04:46
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 9687866 to 75249ab Compare November 1, 2026 04:46
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/storage-workers-fail-loud branch from 30ea114 to ce5be85 Compare November 1, 2026 19:44
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 75249ab to 5dea4ff Compare November 1, 2026 19:46
@xiezhq-hermann xiezhq-hermann changed the title [Refactor] hicache: agree storage transfer completion on the scheduler thread [Fix] hicache: vote the prefetch timeout once per tick, not per admission poll Nov 1, 2026
@xiezhq-hermann

Copy link
Copy Markdown
Collaborator Author

Scope change. This PR no longer carries the scheduler-thread agreement engine. After reviewing the design we are keeping the worker-thread agreement that cache mode has run on in production, and treating backend call bounds as the backend's responsibility (a backend returns or raises within its own bound and must not touch the caller's buffers afterwards). The four PRs of the former series are repurposed, each now independent on main, to the concrete gaps found while auditing that design. New title and body describe the new content; the previous commits are replaced. The branch name is historical.

@xiezhq-hermann xiezhq-hermann reopened this Nov 1, 2026
@xiezhq-hermann
xiezhq-hermann removed this pull request from stack #47679 November 1, 2026 19:52
@xiezhq-hermann
xiezhq-hermann changed the base branch from hicache/storage-workers-fail-loud to main November 1, 2026 19:53
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 5dea4ff to 99e8cf6 Compare November 1, 2026 20:05
@xiezhq-hermann

Copy link
Copy Markdown
Collaborator Author

CI note: the two base-a-test-cpu shards failing here are test/registered/unit/disaggregation/test_decode_queue_cleanup.py (SimpleNamespace has no tree_cache) and test/registered/unit/model_executor/test_prefill_cuda_graph_dcp.py (SimpleNamespace has no skips_prefill_metadata_on_idle_dp_rank). Both fail identically on an untouched checkout of main at b50d439494, the base of this PR; neither file is touched here. Lint and every other shard pass.

@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 99e8cf6 to 02bbd5b Compare November 6, 2026 00:09
Under the default timeout prefetch policy the per-request deadline verdict
was agreed by a MAX collective issued from inside the admission loop, once
per waiting request per scheduling pass. The loop's exits are evaluated per
rank, so one rank could enter a collective its peers never issued (#30760),
and even matched counts cost one gloo all_reduce per waiting request.

check_hicache_events now votes once per tick: one packed MAX over every
live prefetch in ongoing_prefetch order, evaluated on PP rank 0 with the
existing per-request budget and relayed along PP by _all_reduce. For every
agreed 1 the tick marks the deadline expired and terminates the operation,
so the transfer worker and the hit drain stop spending on it; retirement
stays with the admission poll or the drain. _can_terminate_prefetch under
timeout reads the recorded verdict and issues no collective, so
check_prefetch_progress is collective-free for every policy.

Because the polls retire per rank, the vote's shape would depend on the
poll sequence. The ready-count MIN sync's piggybacked digest now folds in
the live prefetch set, so a divergent set raises on every rank in the same
tick instead of hanging in a vote of a different length; the row width is
unchanged, so the shm fast path is unaffected.

Tests: the poll-only agreement case in test_storage_prefetch_lifecycle.py
is adapted to agree through the tick; new cases pin zero collectives from
polls and one from the tick, uneven polls before a verdict, the same-tick
raise on divergent live sets after a verdict, and the PP relay of the vote.
@xiezhq-hermann
xiezhq-hermann force-pushed the hicache/scheduler-thread-transfers branch from 02bbd5b to f11656d Compare November 6, 2026 00:21
@xiezhq-hermann

Copy link
Copy Markdown
Collaborator Author

Rebased on main 9b6c796151 (head f11656d); the description is updated (line references, test evidence incl. the GPU unit suites on a B200, and a "Rebase notes" paragraph in Modifications). The two base-a-test-cpu failures of the previous run were main's and are fixed by #47758 / #47736. CI re-triggered on push.

This branch has not been deployed

No deployments
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 unified-radix-cache

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant