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

Signed-off-by: Abatom <[email protected]>
KuntaiDu · Oct 27, 2024 · 2dee2c4 · 2dee2c4
1 parent 69ddbbb
commit 2dee2c4
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,