Cache empty MatchResult in RadixCache (#24470)
This commit is contained in:
@@ -1052,7 +1052,7 @@ class HiRadixCache(RadixCache):
|
|||||||
last_node = last_node.parent
|
last_node = last_node.parent
|
||||||
|
|
||||||
return (
|
return (
|
||||||
torch.empty((0,), dtype=torch.int64, device=self.device),
|
self._empty_match_result.device_indices,
|
||||||
last_node,
|
last_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1218,30 +1218,20 @@ class HiRadixCache(RadixCache):
|
|||||||
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
|
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
|
||||||
|
|
||||||
def match_prefix(self, params: MatchPrefixParams):
|
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:
|
if self.disable:
|
||||||
return empty_match_result()
|
return self._empty_match_result
|
||||||
|
|
||||||
key = params.key
|
key = params.key
|
||||||
key, _ = key.maybe_to_bigram_view(self.is_eagle)
|
key, _ = key.maybe_to_bigram_view(self.is_eagle)
|
||||||
key = key.page_aligned(self.page_size)
|
key = key.page_aligned(self.page_size)
|
||||||
if len(key) == 0:
|
if len(key) == 0:
|
||||||
return empty_match_result()
|
return self._empty_match_result
|
||||||
|
|
||||||
value, last_node = self._match_prefix_helper(self.root_node, key)
|
value, last_node = self._match_prefix_helper(self.root_node, key)
|
||||||
if value:
|
if value:
|
||||||
value = torch.cat(value)
|
value = torch.cat(value)
|
||||||
else:
|
else:
|
||||||
value = empty_value
|
value = self._empty_match_result.device_indices
|
||||||
|
|
||||||
host_hit_length = 0
|
host_hit_length = 0
|
||||||
last_host_node = last_node
|
last_host_node = last_node
|
||||||
|
|||||||
@@ -335,7 +335,11 @@ class RadixCache(BasePrefixCache):
|
|||||||
self.init_metrics_collector()
|
self.init_metrics_collector()
|
||||||
|
|
||||||
if self.token_to_kv_pool_allocator:
|
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:
|
else:
|
||||||
self.device = torch.device("cpu")
|
self.device = torch.device("cpu")
|
||||||
|
|
||||||
@@ -393,6 +397,15 @@ class RadixCache(BasePrefixCache):
|
|||||||
self.evictable_size_ = 0
|
self.evictable_size_ = 0
|
||||||
self.protected_size_ = 0
|
self.protected_size_ = 0
|
||||||
self.evictable_leaves.clear()
|
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()
|
self._record_all_cleared_event()
|
||||||
|
|
||||||
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
||||||
@@ -435,30 +448,19 @@ class RadixCache(BasePrefixCache):
|
|||||||
key = params.key
|
key = params.key
|
||||||
key, _ = key.maybe_to_bigram_view(self.is_eagle)
|
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:
|
if self.disable or len(key) == 0:
|
||||||
return empty_match_result()
|
return self._empty_match_result
|
||||||
|
|
||||||
key = key.page_aligned(self.page_size)
|
key = key.page_aligned(self.page_size)
|
||||||
|
|
||||||
if len(key) == 0:
|
if len(key) == 0:
|
||||||
return empty_match_result()
|
return self._empty_match_result
|
||||||
|
|
||||||
value, last_node = self._match_prefix_helper(self.root_node, key)
|
value, last_node = self._match_prefix_helper(self.root_node, key)
|
||||||
if value:
|
if value:
|
||||||
value = torch.cat(value)
|
value = torch.cat(value)
|
||||||
else:
|
else:
|
||||||
value = torch.empty((0,), dtype=torch.int64, device=self.device)
|
value = self._empty_match_result.device_indices
|
||||||
return MatchResult(
|
return MatchResult(
|
||||||
device_indices=value,
|
device_indices=value,
|
||||||
last_device_node=last_node,
|
last_device_node=last_node,
|
||||||
|
|||||||
Reference in New Issue
Block a user