Expose can-run-cuda-graph as a regular property on the embedding result (#25711)

This commit is contained in:
fzyzcjy
2026-05-19 09:18:00 +08:00
committed by GitHub
parent a740f8de33
commit fb7e49d4eb
4 changed files with 7 additions and 3 deletions
+1 -1
View File
@@ -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,
+1 -1
View File
@@ -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,
+4
View File
@@ -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):