Repository navigation
[Fix] hicache: vote the prefetch timeout once per tick, not per admission poll - #47620
xiezhq-hermann wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
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.
This comment has been minimized.
This comment has been minimized.
feaa104 to
137bbe4
Compare
3388194 to
8bf47fd
Compare
137bbe4 to
0b8be56
Compare
8bf47fd to
910a82e
Compare
0b8be56 to
7dd5f54
Compare
910a82e to
cad2377
Compare
This comment has been minimized.
This comment has been minimized.
cad2377 to
fc23921
Compare
This comment has been minimized.
This comment has been minimized.
7dd5f54 to
1d8a495
Compare
fc23921 to
c641dc1
Compare
1d8a495 to
ded5f1f
Compare
c641dc1 to
9687866
Compare
ded5f1f to
30ea114
Compare
9687866 to
75249ab
Compare
30ea114 to
ce5be85
Compare
75249ab to
5dea4ff
Compare
|
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. |
5dea4ff to
99e8cf6
Compare
|
CI note: the two |
99e8cf6 to
02bbd5b
Compare
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.
02bbd5b to
f11656d
Compare
|
Rebased on main |
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_progresscalls_can_terminate_prefetch, which on main does one_all_reduce(MAX)per call (python/sglang/srt/mem_cache/unified_radix_cache.py:2480-2490on1ad0149814). The scheduler issues that call once per waiting request with a live prefetch per scheduling pass (python/sglang/srt/managers/scheduler.py:4543onb50d439494), and the loop's exits around it (batch_is_fullfromget_num_allocatable_reqs,scheduler.py:4406-4411and4528-4540onb50d439494) 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 ownbatch_is_fulltripped, called the collective one more time, and waited in it forever while the others had moved on torecv_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 fromcheck_hicache_events, after the storage drain, under storage only (lines 4162-4164). It packs one int per live prefetch, inongoing_prefetchinsertion 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_reduceand_pp_sync)), and for every 1 calls the existing_mark_prefetch_deadline_expiredandcache_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 thedeadlinereason at 3580 and no retry at 3589). With no live prefetch no rank issues a collective. The vote runs only under policytimeout;best_effortandwait_completenever vote._can_terminate_prefetchundertimeout(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 existingis_terminated()path (line 2796); retirement (terminate_prefetch, thedeadlinereason,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 thedeadlinereason before its revoke (lines 3447-3451), as the poll would.check_prefetch_progressis now collective-free for all three policies:best_effortreturns True (line 2642),wait_completereturns False (line 2646),timeoutreads 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_countslines 3873/3875,writing_check,loading_check, the linker round), the synchronous hit query helper (line 2418) anddrain_storage_control_queues, which only the GPU test kits call.ongoing_prefetchis filled byprefetch_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, policytimeout, 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.buffer_only, PP > 1) no entry is ever kept inongoing_prefetch: PP > 0 returns at line 2448 (get_prefetch_submissionreturns an empty submission,hybrid_cache_controller.py:1245) and PP0 returns at line 2580 (submit_prefetchreturnsdecision=True,hybrid_cache_controller.py:1316), so the vote never runs there and the PP relay of_all_reduceis 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:
revokedoutcome plusdeadlinereason) 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_deadlinesto 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.expiredstate, so they agree either way.Rebase notes (2026-11-05). Rebased on main
9b6c796151without conflicts; no code change. The twobase-a-test-cpufailures 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_expiredevaluation, 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_completeand an emptyongoing_prefetchissue no collective.test_expiry_is_agreed_even_before_a_peer_starts_waiting(line 511, two-rank gloo,_two_rank_prefetch_deadlineline 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_pollsline 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_prefetchbody 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_setsline 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 withall_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)._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 main9b6c796151reads; 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 carriesprefetch_stop_policyand an emptyongoing_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_rankfixture (lines 54-55) now carriesprefetch_stop_policy(wait_complete) and an emptyongoing_prefetch, which_hicache_sync_digestreads inside the_sync_hicache_ready_countsthe cases drive directly; the three cases are unchanged (the fold is inactive underwait_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 tree02bbd5b820rebased on9b6c796151; 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.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 main9b6c796151. (TheNo acceleratorerrors quoted in the previous revision came from a mis-setCUDA_VISIBLE_DEVICESon the development box, not from the tests.)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 untouchedmainat9b6c796151(Ran 341 tests in 14.914s, OK, skipped=46).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-testto run the tests your change touches, or copy this prompt for an agent: