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)
|
||||
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(
|
||||
batch=batch,
|
||||
prefill_stats=batch.prefill_stats,
|
||||
|
||||
@@ -92,7 +92,7 @@ class SchedulerDllmMixin:
|
||||
self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
|
||||
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(
|
||||
batch=batch,
|
||||
prefill_stats=batch.prefill_stats,
|
||||
|
||||
@@ -324,7 +324,7 @@ class SchedulerBatchResultProcessor:
|
||||
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(
|
||||
batch=batch,
|
||||
prefill_stats=batch.prefill_stats,
|
||||
|
||||
@@ -260,6 +260,10 @@ class EmbeddingBatchResult:
|
||||
pooled_hidden_states: Optional[torch.Tensor] = None
|
||||
copy_done: Optional[torch.cuda.Event] = None
|
||||
|
||||
@property
|
||||
def can_run_cuda_graph(self) -> bool:
|
||||
return False
|
||||
|
||||
def copy_to_cpu(self):
|
||||
"""Copy embeddings and pooled hidden states to CPU for overlap scheduling."""
|
||||
if isinstance(self.embeddings, torch.Tensor):
|
||||
|
||||
Reference in New Issue
Block a user