[Edge][WP2] Add initialization timeline and cold-start profiler

* vllm_omni/engine/init_timeline.py: InitTimeline phase collector; engine
  phases (stages_init, g3_device_lock, g2_spawn/ready) and worker phases
  (proc_entry, engine_core_init, init_device, load_model/load_weights,
  talker_mtp_init, attn_prewarm, cudagraph_capture, talker_mtp_capture,
  code2wav_cudagraph, predictor_setup_compile/warmup/cudagraph) recorded via
  worker_phase/timed_phase; workers append JSONL under
  VLLM_OMNI_INIT_TIMELINE_DIR (inherited at spawn) and the engine merges them.
  Enabled by VLLM_OMNI_INIT_TIMELINE=<json> or log_stats; table logged at INFO,
  partial table on orchestrator startup timeout, suggested timeouts logged.
* omni_base.py: named deploy profiles default to stage_init_timeout=900 /
  init_timeout=1500 unless given explicitly.
* gpu_model_runner._init_talker_mtp tolerates max_cudagraph_capture_size=None;
  gpu_ar_model_runner.shutdown skips CUDA graph teardown on CPU.
* benchmarks/tts/init_profile.py: runs stream_latency_bench --init-only per
  configuration (default, cold_cache, eager_stage1, eager_both,
  parallel_stage_init, edge, ...) and tabulates phase durations.
* stream_latency_bench: --parallel-stage-init / --deploy-profile flags; the
  resource poller no longer shadows Thread._stop and skips GPU sampling on CPU.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
tzhouamandClaude Fable 5.1 committed 2026-09-09 11:01:39 +08:00
1 parent f804bbd679
commit 073d31b9cd
14 files changed
+885 -24

No files matched your search

@@ -26,6 +26,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmb
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm_omni.diffusion.layers.custom_op import CustomOp
from vllm_omni.engine.init_timeline import timed_phase
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
@@ -716,6 +717,7 @@ class CodePredictorWrapper(nn.Module):
return
self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device)
@timed_phase("predictor_setup_compile")
def _setup_compile(self) -> None:
"""Lazily set up torch.compile with optional device graph capture."""
if self._compiled_model_fwd is not None:
@@ -842,6 +844,7 @@ class CodePredictorWrapper(nn.Module):
allowed = set(all_seq_lens)
return sorted(seq_len for seq_len in self._prefix_graph_seq_lens if seq_len in allowed)
@timed_phase("predictor_warmup")
def _warmup_buckets(self) -> None:
"""Warmup power-of-2 batch-size buckets to front-load Inductor compilation."""
max_bsz = self._vllm_config.scheduler_config.max_num_seqs
@@ -897,6 +900,7 @@ class CodePredictorWrapper(nn.Module):
self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids)
logger.info("code_predictor: warmup done for buckets %s", self._bucket_sizes)
@timed_phase("predictor_cudagraph")
def _capture_cuda_graphs(self) -> None:
"""Capture a CUDA graph per bucket using vLLM's global graph pool."""
from vllm.platforms import current_platform