[RuneFang]

Fix HiCache TP _all_reduce_attn_groups deadlock in _get_new_batch_prefill_raw

TL;DR

HiRadixCache.check_prefetch_progress() issues a TP-group collective
(_all_reduce_attn_groups) but is called from inside the
for req in self.waiting_queue loop of
Scheduler._get_new_batch_prefill_raw(). That loop's trip count is a function
of per-rank local state (get_num_allocatable_reqs() → available_size()
of the KV-cache token pool), so different TP ranks can call the collective a
different number of times in the same scheduler tick. Once the NCCL collective
sequence desyncs, the whole TP group deadlocks (see the exact stack in
#30760).

Fix: pre-compute prefetch progress once per tick for the entire (already
broadcast) waiting_queue, in a new HiRadixCache.bulk_check_prefetch_progress
helper. The scheduler loop body then only does a local dict.get() and never
enters a collective, so its per-rank trip count is irrelevant to collective
ordering.

Closes #30760.

---

Root cause (with code pointers on main)

1. Scheduler._get_new_batch_prefill_raw()
   (scheduler.py:~3030) walks
   the waiting queue and short-circuits on batch_is_full:

   for req in self.waiting_queue:
       ...
       if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs):
           running_batch.batch_is_full = True
       if running_batch.batch_is_full:
           break
       if self.enable_hicache_storage:
           prefetch_done = self.tree_cache.check_prefetch_progress(req.rid)
           ...

2. get_num_allocatable_reqs() returns
   req_to_token_pool.available_size() // …, which is not synchronized
   across TP ranks — slot allocation and release happen locally per rank, so
   the value can drift by one for a single tick.

3. check_prefetch_progress → can_terminate_prefetch unconditionally calls
   _all_reduce_attn_groups (see hiradix_cache.py:can_terminate_prefetch).
   So the number of collectives fired per tick is
   min(len(waiting_queue), get_num_allocatable_reqs_for_this_rank).

4. If rank 0 sees one more allocatable slot than rank 1 in a given tick,
   rank 0 executes the loop body one more time and fires the collective one
   more time. NCCL blocks rank 0 waiting for the peers to join the
   (N+1)-th collective; the peers meanwhile move on to the next scheduler
   tick and call broadcast_pyobj / forward. Permanent hang.

Observed py-spy stack on the hung rank (matches the issue verbatim):

all_reduce                (torch/distributed/distributed_c10d.py)
_all_reduce_attn_groups   (hiradix_cache.py)
check_prefetch_progress   (hiradix_cache.py)          <-- inside for-loop
_get_new_batch_prefill_raw(scheduler.py)
get_new_batch_prefill     (scheduler.py)
get_next_batch_to_run     (scheduler.py)
event_loop_overlap        (scheduler.py)

---

Fix

python/sglang/srt/mem_cache/hiradix_cache.py

New public method on HiRadixCache:

def bulk_check_prefetch_progress(self, req_ids: list) -> dict:
    if not self.enable_storage:
        return {}
    result = {}
    for rid in req_ids:
        result[rid] = self.check_prefetch_progress(rid)
    return result

The collective count is now exactly len(req_ids) on every rank because
waiting_queue is already TP-broadcast, so req_ids is identical everywhere.
Semantics of check_prefetch_progress (and its side effects — popping
ongoing_prefetch, updating prefetch_loaded_tokens_by_reqid, etc.) are
completely unchanged.

python/sglang/srt/managers/scheduler.py::_get_new_batch_prefill_raw

Call bulk_check_prefetch_progress before the loop, then convert the
loop's collective call into a pure dict lookup:

prefetch_progress_map = None
if self.enable_hicache_storage and hasattr(
    self.tree_cache, "bulk_check_prefetch_progress"
):
    prefetch_progress_map = self.tree_cache.bulk_check_prefetch_progress(
        [req.rid for req in self.waiting_queue]
    )

for req in self.waiting_queue:
    ...
    if prefetch_progress_map is not None:
        prefetch_done = prefetch_progress_map.get(req.rid, True)
        if not prefetch_done:
            continue
        loaded_tokens = self.tree_cache.pop_prefetch_loaded_tokens(req.rid)
        if loaded_tokens > 0:
            req.storage_hit_length = loaded_tokens
    elif self.enable_hicache_storage:
        # legacy per-req path, preserved for tree_cache impls that have not
        # yet opted into the bulk API (see "Impact" below).
        prefetch_done = self.tree_cache.check_prefetch_progress(req.rid)
        ...

Why this is correct

• Same total number of check_prefetch_progress calls per tick as before
  (one per req in waiting_queue), just moved out of the loop.
• Same side effects in the same order (iterated over the same list).
• Same prefetch_done boolean returned to the loop body.
• Fewer collectives per tick in one edge case: previously the loop broke
  early on batch_is_full and skipped check_prefetch_progress for tail
  reqs; now we always call it for the whole queue. This is intentional — it
  is the invariant that fixes the deadlock. The extra calls do not change
  batch composition and are bounded by len(waiting_queue).

---

Impact analysis (why this is safe)

Surface	Result
Public API compatibility	✅ Additive only
Non-HiCache path (enable_hicache_storage=False)	✅ Byte-identical
HiCache without a storage backend	✅ Byte-identical (early return)
HiCache + storage backend (target of the fix)	✅ Behavior equivalent + no hang
pop_prefetch_loaded_tokens order / storage_hit_length	✅ Same order, same values
Non-HiRadixCache tree_cache implementations (UnifiedRadixCache, HiMambaRadixCache, …)	✅ Falls back to legacy per-req path via hasattr(...) guard
Extra collectives per tick	➖ At most len(waiting_queue); same big-O as before
Latency impact in HiCache path	➖ No measurable regression; the extra calls only affect ticks where the loop previously broke early on batch_is_full

Reproduction

The scenario described in #30760 is a per-tick race between per-rank
available_size() drift and a collective inside check_prefetch_progress.
Waiting for the race to occur naturally requires long-running production
traffic, which is why the original issue has no minimal repro. To make the
race deterministic and reviewable, a companion branch ships a
zero-source-change reproducer:

Repro branch: https://github.com/RuneFang/sglang/tree/repo%2330760

That branch contains only the repro_30760/ directory (launch script,
usercustomize.py injector, evidence logs and py-spy stacks). It does not
touch any source under python/sglang/. Applying this PR on top of the
repro branch turns the reproducer's verdict from DEADLOCKED to
ALL REQUESTS SUCCEEDED.

Key idea of the injector (why it faithfully reproduces #30760)

The injector monkey-patches a single method and gates the extra collective on
the call site, not on any timing:

# Applied in every spawned TP scheduler worker via PYTHONPATH/usercustomize.
_orig_check = HiRadixCache.check_prefetch_progress

def _patched_check(self, req_id):
    if torch.distributed.get_rank() == 0 and not _called_from_bulk():
        # Rank 0 issues ONE extra collective when check_prefetch_progress is
        # invoked directly from the per-loop call site in the scheduler --
        # this is exactly the "one more candidate request" observation from
        # the issue. When the call originates from bulk_check_prefetch_progress
        # (the fixed path), _called_from_bulk() returns True and this branch
        # is skipped, so every rank stays lock-stepped.
        probe = torch.tensor([0, 0], dtype=torch.int)
        self._all_reduce_attn_groups(probe, torch.distributed.ReduceOp.MAX)
    return _orig_check(self, req_id)

HiRadixCache.check_prefetch_progress = _patched_check

The stack-based gate directly encodes the causal chain from the root cause:

Code path (call site of check_prefetch_progress)	Pre-fix behavior	Post-fix behavior
Inside _get_new_batch_prefill_raw's per-req loop	Rank 0 fires an extra collective → deadlock	(no longer used — the fix removes this call site)
Inside bulk_check_prefetch_progress (new)	(does not exist)	Extra collective is suppressed → all ranks lock-stepped

Batch composition is not perturbed (no get_num_allocatable_reqs patching, no
is_empty override), which is why the fixed run does not exhibit any
unrelated forward-time hangs.

Evidence (captured on the repro branch)

Pre-fix, rank 0 (py-spy dump — matches the stack reported in #30760):

all_reduce                        (torch/distributed/distributed_c10d.py:3075)
_all_reduce_attn_groups           (hiradix_cache.py:217)
check_prefetch_progress           (hiradix_cache.py, via injector)
_get_new_batch_prefill_raw        (scheduler.py:3056)
get_new_batch_prefill             (scheduler.py:2920)
get_next_batch_to_run             (scheduler.py:2859)
event_loop_overlap                (scheduler.py:1636)
run_event_loop                    (scheduler.py:1560)
run_scheduler_process             (scheduler.py:4761)

Pre-fix, other TP ranks: already advanced one tick ahead into
event_loop_overlap (running the next forward/broadcast_pyobj), so the
NCCL collective sequence is desynchronized -> permanent hang.

Post-fix, all ranks: every request in the 12-request burst returns 200 OK
with sub-second wall time; /health responsive throughout the run; no
heartbeat loss in the server log.

---

FAQ: Why not just synchronize available_size() across ranks?

available_size() is by design a per-rank local counter
(len(free_slots)). Turning it into a globally consistent value means
injecting a new collective into the scheduler's per-tick inner loop and
polluting the semantics of a deliberately local counter.

It would also only paper over the symptom — the root cause is the pattern
"a per-rank-state-gated loop containing a collective", which exists
elsewhere too (e.g. _try_hicache_queue_load_back).

This PR adds zero new collectives: it just moves the collective that
already lives inside check_prefetch_progress into a loop whose input
(the broadcast waiting_queue) is already cross-rank consistent,
removing the coupling structurally.

---

CI States

Latest PR Test (Base): ❌ Run #30612117737
Latest PR Test (Extra): ❌ Run #30612117623

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[github-actions]

Thanks @RuneFang. Closing this because it has had no updates in 91 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.