[Misc] SpecDecodeWorker supports profiling (vllm-project#9719)

Signed-off-by: Abatom <[email protected]> Signed-off-by: Shanshan Wang <[email protected]>
cooleel · Oct 28, 2024 · c652692 · c652692
1 parent d5668e9
commit c652692
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/vllm/spec_decode/spec_decode_worker.py b/vllm/spec_decode/spec_decode_worker.py
@@ -1038,6 +1038,14 @@ def get_cache_block_size_bytes(self):
         """
         raise NotImplementedError
 
+    def start_profile(self):
+        if isinstance(self.scorer_worker, Worker):
+            self.scorer_worker.start_profile()
+
+    def stop_profile(self):
+        if isinstance(self.scorer_worker, Worker):
+            self.scorer_worker.stop_profile()
+
 
 def split_num_cache_blocks_evenly(scorer_cache_block_size_bytes: int,
                                   proposer_cache_block_size_bytes: int,