From abf81f4114334946933b3b96f2e478265ebcd67c Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 21:32:15 +0800 Subject: [PATCH] [Edge][WP8] Model-level step timers: talker forward, sampling, talker-MTP, code-predictor sub-steps Counters model.forward_ms / model.sample_ms / model.talker_mtp_ms and predictor.substep_ms / predictor.loop_ms (VLLM_OMNI_STEP_STATS_DIR); with VLLM_OMNI_STEP_STATS_SYNC=1 the accelerator is synchronized around the model timers so async kernels are attributed (attribution runs only). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ --- vllm_omni/edge/step_stats.py | 13 ++++++- .../models/common/qwen3_code_predictor.py | 10 ++++++ vllm_omni/worker/gpu_ar_model_runner.py | 35 +++++++++++++------ 3 files changed, 46 insertions(+), 12 deletions(-) diff --git a/vllm_omni/edge/step_stats.py b/vllm_omni/edge/step_stats.py index 885256a61..09e30a3ab 100644 --- a/vllm_omni/edge/step_stats.py +++ b/vllm_omni/edge/step_stats.py @@ -28,6 +28,7 @@ from pathlib import Path from typing import Any ENV_DIR = "VLLM_OMNI_STEP_STATS_DIR" +ENV_SYNC = "VLLM_OMNI_STEP_STATS_SYNC" # "1": synchronize the accelerator around model-level timers ENV_STAGE_ID = "VLLM_OMNI_STAGE_ID" FRAME_BUDGET_MS = 80.0 # one Qwen3-TTS codec frame at 12.5 Hz _MAX_SAMPLES = 20000 @@ -81,6 +82,7 @@ class StepStats: def __init__(self, out_dir: str | None, role: str | None = None) -> None: self.out_dir = out_dir self.enabled = bool(out_dir) + self.sync = os.environ.get(ENV_SYNC, "0").strip().lower() in ("1", "true", "yes", "on") self.role = role or _default_role() self.pid = os.getpid() self._series: dict[str, _Series] = {} @@ -127,14 +129,23 @@ class StepStats: self.dump() @contextmanager - def timed(self, name: str) -> Iterator[None]: + def timed(self, name: str, sync_device: Any = None) -> Iterator[None]: + """Time a block; with ``sync_device`` (a torch device) and ``VLLM_OMNI_STEP_STATS_SYNC=1`` + the accelerator is synchronized before and after so async kernels are included.""" if not self.enabled: yield return + sync = sync_device is not None and self.sync and getattr(sync_device, "type", "cpu") == "cuda" + if sync: + import torch + + torch.accelerator.synchronize(sync_device) t0 = time.perf_counter() try: yield finally: + if sync: + torch.accelerator.synchronize(sync_device) self.add(name, (time.perf_counter() - t0) * 1000.0) def wrap(self, name: str, fn: Callable[..., Any]) -> Callable[..., Any]: diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index b22afc027..4af31a82a 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -15,6 +15,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models. from __future__ import annotations import dataclasses +import time from collections.abc import Iterable, Sequence import torch @@ -1062,7 +1063,12 @@ class CodePredictorWrapper(nn.Module): all_codes[:, 0] = layer0_code.reshape(bsz) # Autoregressive loop: predict layers 1..G-1 + from vllm_omni.edge.step_stats import StepStats + + _stats = StepStats.get() + _t_loop = time.perf_counter() if _stats.enabled else 0.0 for step in range(1, num_groups): + _t_sub = time.perf_counter() if _stats.enabled else 0.0 graph_key: int | tuple[int, int] = padded_bsz seq_len = max_seq if self._prefix_graphs_enabled: @@ -1136,6 +1142,10 @@ class CodePredictorWrapper(nn.Module): if step < num_groups - 1 or self._wrapper_config.return_proj_buf: new_embed = codec_embeds[step - 1](code) proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1) + if _stats.enabled: + _stats.add("predictor.substep_ms", (time.perf_counter() - _t_sub) * 1000.0) + if _stats.enabled: + _stats.add("predictor.loop_ms", (time.perf_counter() - _t_loop) * 1000.0) if self._wrapper_config.return_proj_buf: return all_codes, proj_buf[:bsz].clone() diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index dd362e865..18a0e9e49 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -66,6 +66,13 @@ from vllm_omni.worker.sparse_audio import resolve_sparse_mm_routing logger = init_logger(__name__) +def _step_stats(): + """Process-local per-step counters (no-op unless VLLM_OMNI_STEP_STATS_DIR is set).""" + from vllm_omni.edge.step_stats import StepStats + + return StepStats.get() + + def _to_cpu_contiguous(tensor: torch.Tensor) -> torch.Tensor: tensor = tensor.detach() if tensor.device.type == "cpu": @@ -1218,16 +1225,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex defer_finalize=defer_kv_connector_finalize, ) as kv_connector_output, ): - model_output = self._model_forward( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - sampling_metadata=self.input_batch.sampling_metadata, - logits_index=logits_indices, - sampler=self.sampler, - ) + _stats = _step_stats() + with _stats.timed("model.forward_ms", sync_device=self.device): + model_output = self._model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + sampling_metadata=self.input_batch.sampling_metadata, + logits_index=logits_indices, + sampler=self.sampler, + ) # [Omni] Map pending ropes metadata to req_ids. flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None) @@ -2051,7 +2060,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex ) with record_function_or_nullcontext("gpu_model_runner: sample"): - sampler_output = self._sample(logits, spec_decode_metadata) + with _step_stats().timed("model.sample_ms", sync_device=self.device): + sampler_output = self._sample(logits, spec_decode_metadata) self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) @@ -2128,6 +2138,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex scheduler_output.total_num_scheduled_tokens, ) + _mtp_timer = _step_stats().timed("model.talker_mtp_ms", sync_device=self.device) + _mtp_timer.__enter__() multimodal_outputs = self._run_post_sample_talker_mtp( req_ids=req_ids_output_copy, valid_sampled_token_ids=valid_sampled_token_ids, @@ -2136,6 +2148,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex sample_hidden_states=sample_hidden_states, multimodal_outputs=multimodal_outputs, ) + _mtp_timer.__exit__(None, None, None) if propose_drafts_after_bookkeeping: # ngram and other speculative decoding methods use the sampled