diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 123d39ae5..e1f6b2b14 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1052,7 +1052,7 @@ class HiRadixCache(RadixCache): last_node = last_node.parent return ( - torch.empty((0,), dtype=torch.int64, device=self.device), + self._empty_match_result.device_indices, last_node, ) @@ -1218,30 +1218,20 @@ class HiRadixCache(RadixCache): return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0) def match_prefix(self, params: MatchPrefixParams): - empty_value = torch.empty((0,), dtype=torch.int64, device=self.device) - - def empty_match_result(): - return MatchResult( - device_indices=empty_value, - last_device_node=self.root_node, - last_host_node=self.root_node, - host_hit_length=0, - ) - if self.disable: - return empty_match_result() + return self._empty_match_result key = params.key key, _ = key.maybe_to_bigram_view(self.is_eagle) key = key.page_aligned(self.page_size) if len(key) == 0: - return empty_match_result() + return self._empty_match_result value, last_node = self._match_prefix_helper(self.root_node, key) if value: value = torch.cat(value) else: - value = empty_value + value = self._empty_match_result.device_indices host_hit_length = 0 last_host_node = last_node diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index a707a4b24..7f9eca81f 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -335,7 +335,11 @@ class RadixCache(BasePrefixCache): self.init_metrics_collector() if self.token_to_kv_pool_allocator: - self.device = self.token_to_kv_pool_allocator.device + dev = self.token_to_kv_pool_allocator.device + if isinstance(dev, (str, torch.device)): + self.device = torch.device(dev) + else: + self.device = torch.device("cpu") else: self.device = torch.device("cpu") @@ -393,6 +397,15 @@ class RadixCache(BasePrefixCache): self.evictable_size_ = 0 self.protected_size_ = 0 self.evictable_leaves.clear() + self._empty_match_result = MatchResult( + device_indices=torch.empty( + (0,), + dtype=torch.int64, + device=self.device, + ), + last_device_node=self.root_node, + last_host_node=self.root_node, + ) self._record_all_cleared_event() def match_prefix(self, params: MatchPrefixParams) -> MatchResult: @@ -435,30 +448,19 @@ class RadixCache(BasePrefixCache): key = params.key key, _ = key.maybe_to_bigram_view(self.is_eagle) - def empty_match_result(): - return MatchResult( - device_indices=torch.empty( - (0,), - dtype=torch.int64, - device=self.device, - ), - last_device_node=self.root_node, - last_host_node=self.root_node, - ) - if self.disable or len(key) == 0: - return empty_match_result() + return self._empty_match_result key = key.page_aligned(self.page_size) if len(key) == 0: - return empty_match_result() + return self._empty_match_result value, last_node = self._match_prefix_helper(self.root_node, key) if value: value = torch.cat(value) else: - value = torch.empty((0,), dtype=torch.int64, device=self.device) + value = self._empty_match_result.device_indices return MatchResult( device_indices=value, last_device_node=last_node,