[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:
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
|
||||
|
||||
Reference in new issue
Block a user