diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index c55ff312b..f5be51e83 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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 diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index 8456d91ed..6c4028be6 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -305,10 +305,14 @@ class TestHiSparseUnit(unittest.TestCase): for i in range(TOP_K): if batch[b, i] < 0: continue - nd = self.device_pool.kv_buffer[layer_id][naive_locs[b, i].long()] - kd = self.device_pool.kv_buffer[layer_id][kernel_locs[b, i].long()] + naive_data = self.device_pool.kv_buffer[layer_id][ + naive_locs[b, i].long() + ] + kernel_data = self.device_pool.kv_buffer[layer_id][ + kernel_locs[b, i].long() + ] 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", )