[HiCache & PD]Fixed detailed cache hit breakdown in PD scenarios. (#21764)
This commit is contained in:
@@ -959,6 +959,9 @@ class DecodeTransferQueue:
|
|||||||
# Case 3: Success - commit the transfer
|
# Case 3: Success - commit the transfer
|
||||||
decode_req.req.output_ids.append(output_id[0].item())
|
decode_req.req.output_ids.append(output_id[0].item())
|
||||||
decode_req.req.cached_tokens = cached_tokens[0].item()
|
decode_req.req.cached_tokens = cached_tokens[0].item()
|
||||||
|
decode_req.req.cached_tokens_device = cached_tokens[1].item()
|
||||||
|
decode_req.req.cached_tokens_host = cached_tokens[2].item()
|
||||||
|
decode_req.req.cached_tokens_storage = cached_tokens[3].item()
|
||||||
if not self.spec_algorithm.is_none():
|
if not self.spec_algorithm.is_none():
|
||||||
decode_req.req.output_topk_p = output_topk_p
|
decode_req.req.output_topk_p = output_topk_p
|
||||||
decode_req.req.output_topk_index = output_topk_index
|
decode_req.req.output_topk_index = output_topk_index
|
||||||
|
|||||||
@@ -251,6 +251,9 @@ class MetadataBuffers:
|
|||||||
|
|
||||||
self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
|
self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
|
||||||
self.cached_tokens[req.metadata_buffer_index][0] = req.cached_tokens
|
self.cached_tokens[req.metadata_buffer_index][0] = req.cached_tokens
|
||||||
|
self.cached_tokens[req.metadata_buffer_index][1] = req.cached_tokens_device
|
||||||
|
self.cached_tokens[req.metadata_buffer_index][2] = req.cached_tokens_host
|
||||||
|
self.cached_tokens[req.metadata_buffer_index][3] = req.cached_tokens_storage
|
||||||
if req.return_logprob:
|
if req.return_logprob:
|
||||||
if req.output_token_logprobs_val: # not none or empty list
|
if req.output_token_logprobs_val: # not none or empty list
|
||||||
self.output_token_logprobs_val[req.metadata_buffer_index][0] = (
|
self.output_token_logprobs_val[req.metadata_buffer_index][0] = (
|
||||||
|
|||||||
@@ -55,15 +55,10 @@ class SchedulerOutputProcessorMixin:
|
|||||||
"""Get detailed cache breakdown for a request, if available.
|
"""Get detailed cache breakdown for a request, if available.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
- None if HiCache is not enabled
|
- None if no cached tokens at all
|
||||||
- {"device": X, "host": Y} if HiCache enabled but L3 storage is not
|
- {"device": X, "host": Y} without storage breakdown
|
||||||
- {"device": X, "host": Y, "storage": Z, "storage_backend": "..."} if L3 enabled
|
- {"device": X, "host": Y, "storage": Z} with storage breakdown
|
||||||
"""
|
"""
|
||||||
# Only show details if HiCache is enabled
|
|
||||||
if not getattr(self, "enable_hierarchical_cache", False):
|
|
||||||
return None
|
|
||||||
|
|
||||||
# Only show if there are any cached tokens
|
|
||||||
if (
|
if (
|
||||||
req.cached_tokens_device > 0
|
req.cached_tokens_device > 0
|
||||||
or req.cached_tokens_host > 0
|
or req.cached_tokens_host > 0
|
||||||
@@ -78,6 +73,13 @@ class SchedulerOutputProcessorMixin:
|
|||||||
details["storage"] = req.cached_tokens_storage
|
details["storage"] = req.cached_tokens_storage
|
||||||
details["storage_backend"] = self._get_storage_backend_type()
|
details["storage_backend"] = self._get_storage_backend_type()
|
||||||
return details
|
return details
|
||||||
|
|
||||||
|
if req.cached_tokens > 0:
|
||||||
|
return {
|
||||||
|
"device": req.cached_tokens,
|
||||||
|
"host": 0,
|
||||||
|
}
|
||||||
|
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def process_batch_result_prebuilt(self: Scheduler, batch: ScheduleBatch):
|
def process_batch_result_prebuilt(self: Scheduler, batch: ScheduleBatch):
|
||||||
|
|||||||
Reference in New Issue
Block a user