[Bug Fix] Ensure prefill_info_table is populated before honoring disagg_prefill_dp_rank (#22990)
Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
This commit is contained in:
@@ -421,6 +421,7 @@ class DecodePreallocQueue:
|
||||
|
||||
# Fast path: cache-only lookup, no network calls
|
||||
prefill_dp_rank = self._resolve_prefill_dp_rank(req)
|
||||
logger.debug(f"prefill_dp_rank: {prefill_dp_rank}")
|
||||
if prefill_dp_rank is not None:
|
||||
decode_req.kv_receiver.init(prefill_dp_rank)
|
||||
return
|
||||
@@ -428,13 +429,14 @@ class DecodePreallocQueue:
|
||||
self.pending_reqs.append(decode_req)
|
||||
|
||||
def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]:
|
||||
if req.disagg_prefill_dp_rank is not None:
|
||||
return req.disagg_prefill_dp_rank
|
||||
|
||||
prefill_info = self.kv_manager.prefill_info_table.get(_bootstrap_addr(req))
|
||||
# If None, it will go to the slow path and resolve prefill_info by _ensure_prefill_info then cache it
|
||||
if prefill_info is None:
|
||||
return None
|
||||
|
||||
if req.disagg_prefill_dp_rank is not None:
|
||||
return req.disagg_prefill_dp_rank
|
||||
|
||||
if prefill_info.dp_size == 1:
|
||||
return 0
|
||||
|
||||
|
||||
Reference in New Issue
Block a user