From 1de52e7751025b7e7686efb7864586291e43c2de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A7=91=E8=8B=B1?= Date: Sun, 27 Oct 2024 12:18:03 +0800 Subject: [PATCH] [Misc] SpecDecodeWorker supports profiling (#9719) Signed-off-by: Abatom Signed-off-by: qishuai --- vllm/spec_decode/spec_decode_worker.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/vllm/spec_decode/spec_decode_worker.py b/vllm/spec_decode/spec_decode_worker.py index 316db43502d3b..9f7ef2f8d851c 100644 --- 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,