minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Bug Fix] Ensure prefill_info_table is populated before honoring disagg_prefill_dp_rank (#22990)

Browse Source 
Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
This commit is contained in:
Byron Hsu
2026-04-17 11:10:31 +08:00
committed by GitHub
co-authored by Byron Hsu
parent 04a53955b9
commit cf9845f8e3
2 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+5 -3
View File
@@ -421,6 +421,7 @@ class DecodePreallocQueue:
# Fast path: cache-only lookup, no network calls # Fast path: cache-only lookup, no network calls
prefill_dp_rank = self._resolve_prefill_dp_rank(req) 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: if prefill_dp_rank is not None:
decode_req.kv_receiver.init(prefill_dp_rank) decode_req.kv_receiver.init(prefill_dp_rank)
return return
@@ -428,13 +429,14 @@ class DecodePreallocQueue:
self.pending_reqs.append(decode_req) self.pending_reqs.append(decode_req)
def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]: 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)) 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: if prefill_info is None:
return None return None
if req.disagg_prefill_dp_rank is not None:
return req.disagg_prefill_dp_rank
if prefill_info.dp_size == 1: if prefill_info.dp_size == 1:
return 0 return 0
test/registered/unit/managers/test_hisparse_unit.py
+7 -3
View File
@@ -305,10 +305,14 @@ class TestHiSparseUnit(unittest.TestCase):
for i in range(TOP_K): for i in range(TOP_K):
if batch[b, i] < 0: if batch[b, i] < 0:
continue continue
nd = self.device_pool.kv_buffer[layer_id][naive_locs[b, i].long()] naive_data = self.device_pool.kv_buffer[layer_id][
kd = self.device_pool.kv_buffer[layer_id][kernel_locs[b, i].long()] naive_locs[b, i].long()
]
kernel_data = self.device_pool.kv_buffer[layer_id][
kernel_locs[b, i].long()
]
self.assertTrue( self.assertTrue(
torch.allclose(nd.float(), kd.float(), atol=1e-2), torch.allclose(naive_data.float(), kernel_data.float(), atol=1e-2),
f"{msg}layer {layer_id}, b{b} idx {i}: naive != kernel", f"{msg}layer {layer_id}, b{b} idx {i}: naive != kernel",
) )