From 936c10be4b7820e868028e36f7aa887a3a7a3e8f Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 21:42:52 +0800 Subject: [PATCH] [Edge][WP8] Time the steady-state talker-MTP graph call (model.mtp_graph_ms) Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ --- vllm_omni/worker/gpu_model_runner.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 560e728e6..3d689d11a 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1957,8 +1957,13 @@ class OmniGPUModelRunner(GPUModelRunner): talker_kwargs["req_infos"] = [ self.model_intermediate_buffer.setdefault(req_id, {}) for req_id in decode_req_ids ] - with current_omni_platform.set_forward_context( - None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc + from vllm_omni.edge.step_stats import StepStats + + with ( + current_omni_platform.set_forward_context( + None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc + ), + StepStats.get().timed("model.mtp_graph_ms", sync_device=getattr(self, "device", None)), ): req_embeds, code_predictor_codes = self.talker_mtp( req_input_ids,