Expose can-run-cuda-graph as a regular property on the embedding result (#25711)
This commit is contained in:
@@ -586,7 +586,7 @@ class SchedulerDisaggregationPrefillMixin:
|
|||||||
self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx)
|
self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx)
|
||||||
req.time_stats.set_last_chunked_prefill_finish_time()
|
req.time_stats.set_last_chunked_prefill_finish_time()
|
||||||
|
|
||||||
can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False)
|
can_run_cuda_graph = result.can_run_cuda_graph
|
||||||
self.metrics_reporter.report_prefill_stats(
|
self.metrics_reporter.report_prefill_stats(
|
||||||
batch=batch,
|
batch=batch,
|
||||||
prefill_stats=batch.prefill_stats,
|
prefill_stats=batch.prefill_stats,
|
||||||
|
|||||||
@@ -92,7 +92,7 @@ class SchedulerDllmMixin:
|
|||||||
self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
|
self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
|
||||||
self.token_to_kv_pool_allocator.free_group_end()
|
self.token_to_kv_pool_allocator.free_group_end()
|
||||||
|
|
||||||
can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False)
|
can_run_cuda_graph = result.can_run_cuda_graph
|
||||||
self.metrics_reporter.report_prefill_stats(
|
self.metrics_reporter.report_prefill_stats(
|
||||||
batch=batch,
|
batch=batch,
|
||||||
prefill_stats=batch.prefill_stats,
|
prefill_stats=batch.prefill_stats,
|
||||||
|
|||||||
@@ -324,7 +324,7 @@ class SchedulerBatchResultProcessor:
|
|||||||
batch.reqs, batch.return_logprob, skip_stream_req
|
batch.reqs, batch.return_logprob, skip_stream_req
|
||||||
)
|
)
|
||||||
|
|
||||||
can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False)
|
can_run_cuda_graph = result.can_run_cuda_graph
|
||||||
self.metrics_reporter.report_prefill_stats(
|
self.metrics_reporter.report_prefill_stats(
|
||||||
batch=batch,
|
batch=batch,
|
||||||
prefill_stats=batch.prefill_stats,
|
prefill_stats=batch.prefill_stats,
|
||||||
|
|||||||
@@ -260,6 +260,10 @@ class EmbeddingBatchResult:
|
|||||||
pooled_hidden_states: Optional[torch.Tensor] = None
|
pooled_hidden_states: Optional[torch.Tensor] = None
|
||||||
copy_done: Optional[torch.cuda.Event] = None
|
copy_done: Optional[torch.cuda.Event] = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def can_run_cuda_graph(self) -> bool:
|
||||||
|
return False
|
||||||
|
|
||||||
def copy_to_cpu(self):
|
def copy_to_cpu(self):
|
||||||
"""Copy embeddings and pooled hidden states to CPU for overlap scheduling."""
|
"""Copy embeddings and pooled hidden states to CPU for overlap scheduling."""
|
||||||
if isinstance(self.embeddings, torch.Tensor):
|
if isinstance(self.embeddings, torch.Tensor):
|
||||||
|
|||||||
Reference in New Issue
Block a user