[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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ

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]:

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()

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