diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 84e476214..a64344844 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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, diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 6d35864ab..ee1b8d076 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -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, diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index baf09cfb7..6a96fb4ef 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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, diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 31c5b375c..33db3942b 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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):