@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4453ba8d9e

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 08:45 UTC

[Core] Profiler improvements and lazy initialization (#33198)

Signed-off-by: Jaewon Lee <jaewon@meta.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>

Jaewon committed Feb 12, 2026 at 16:16 UTC 4453ba8d9ec8e35d68084a118f35ce5c48b5dae6

10 files changed +117 -37

vllm/distributed/utils.py

+40

@@ -524,3 +524,43 @@ def stateless_destroy_torch_distributed_process_group(pg: ProcessGroup) -> None:

524 """

525 pg.shutdown()

526 _unregister_process_group(pg.group_name)

527 +

528 +

529 +def get_worker_rank_suffix(global_rank: int | None = None) -> str:

530 + """Generate a descriptive rank suffix for worker identification.

531 +

532 + Returns a string like 'dp0_pp0_tp0_dcp0_ep0_rank0' including all

533 + parallel dimensions: DP, PP, TP, DCP, EP.

534 +

535 + Args:

536 + global_rank: Optional global rank to append. If not provided,

537 + only parallel dimension ranks are included.

538 +

539 + Returns:

540 + A string suffix identifying the worker's position in the

541 + distributed topology.

542 + """

543 + from vllm.distributed.parallel_state import (

544 + get_dcp_group,

545 + get_dp_group,

546 + get_ep_group,

547 + get_pp_group,

548 + get_tp_group,

549 + )

550 +

551 + try:

552 + dp_rank = get_dp_group().rank_in_group

553 + pp_rank = get_pp_group().rank_in_group

554 + tp_rank = get_tp_group().rank_in_group

555 + dcp_rank = get_dcp_group().rank_in_group

556 + ep_rank = get_ep_group().rank_in_group

557 +

558 + suffix = f"dp{dp_rank}_pp{pp_rank}_tp{tp_rank}_dcp{dcp_rank}_ep{ep_rank}"

559 + if global_rank is not None:

560 + suffix = f"{suffix}_rank{global_rank}"

561 + return suffix

562 + except Exception:

563 + # Fallback if parallel state not initialized

564 + if global_rank is not None:

565 + return f"rank{global_rank}"

566 + return ""

vllm/entrypoints/llm.py

+9 -2

@@ -1685,8 +1685,15 @@ class LLM:

1685 tokenization_kwargs=encode_kwargs,

1686 )

1687

1688 - def start_profile(self) -> None:

1689 - self.llm_engine.start_profile()

1688 + def start_profile(self, profile_prefix: str | None = None) -> None:

1689 + """Start profiling with optional custom trace prefix.

1690 +

1691 + Args:

1692 + profile_prefix: Optional prefix for the trace file names. If provided,

1693 + trace files will be named as "<prefix>_dp<X>_pp<Y>_tp<Z>".

1694 + If not provided, default naming will be used.

1695 + """

1696 + self.llm_engine.start_profile(profile_prefix)

1697

1698 def stop_profile(self) -> None:

1699 self.llm_engine.stop_profile()

vllm/v1/engine/async_llm.py

+2 -2

@@ -911,8 +911,8 @@ class AsyncLLM(EngineClient):

911 if self.errored:

912 raise self.dead_error

913

914 - async def start_profile(self) -> None:

915 - coros = [self.engine_core.profile_async(True)]

914 + async def start_profile(self, profile_prefix: str | None = None) -> None:

915 + coros = [self.engine_core.profile_async(True, profile_prefix)]

916 if self.profiler is not None:

917 coros.append(asyncio.to_thread(self.profiler.start))

918 await asyncio.gather(*coros)

vllm/v1/engine/core.py

+2 -2

@@ -568,8 +568,8 @@ class EngineCore:

568 if self.scheduler:

569 self.scheduler.shutdown()

570

571 - def profile(self, is_start: bool = True):

572 - self.model_executor.profile(is_start)

571 + def profile(self, is_start: bool = True, profile_prefix: str | None = None):

572 + self.model_executor.profile(is_start, profile_prefix)

573

574 def reset_mm_cache(self):

575 # NOTE: Since this is mainly for debugging, we don't attempt to

vllm/v1/engine/core_client.py

+12 -8

@@ -135,7 +135,7 @@ class EngineCoreClient(ABC):

135 def add_request(self, request: EngineCoreRequest) -> None:

136 raise NotImplementedError

137

138 - def profile(self, is_start: bool = True) -> None:

138 + def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:

139 raise NotImplementedError

140

141 def reset_mm_cache(self) -> None:

@@ -210,7 +210,9 @@ class EngineCoreClient(ABC):

210 async def add_request_async(self, request: EngineCoreRequest) -> None:

211 raise NotImplementedError

212

213 - async def profile_async(self, is_start: bool = True) -> None:

213 + async def profile_async(

214 + self, is_start: bool = True, profile_prefix: str | None = None

215 + ) -> None:

216 raise NotImplementedError

217

218 async def reset_mm_cache_async(self) -> None:

@@ -295,8 +297,8 @@ class InprocClient(EngineCoreClient):

297 def shutdown(self) -> None:

298 self.engine_core.shutdown()

299

298 - def profile(self, is_start: bool = True) -> None:

299 - self.engine_core.profile(is_start)

300 + def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:

301 + self.engine_core.profile(is_start, profile_prefix)

302

303 def reset_mm_cache(self) -> None:

304 self.engine_core.reset_mm_cache()

@@ -765,8 +767,8 @@ class SyncMPClient(MPClient):

767 if request_ids and not self.resources.engine_dead:

768 self._send_input(EngineCoreRequestType.ABORT, request_ids)

769

768 - def profile(self, is_start: bool = True) -> None:

769 - self.call_utility("profile", is_start)

770 + def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:

771 + self.call_utility("profile", is_start, profile_prefix)

772

773 def reset_mm_cache(self) -> None:

774 self.call_utility("reset_mm_cache")

@@ -987,8 +989,10 @@ class AsyncMPClient(MPClient):

989 """Resume the scheduler after a pause."""

990 await self.call_utility_async("resume_scheduler")

991

990 - async def profile_async(self, is_start: bool = True) -> None:

991 - await self.call_utility_async("profile", is_start)

992 + async def profile_async(

993 + self, is_start: bool = True, profile_prefix: str | None = None

994 + ) -> None:

995 + await self.call_utility_async("profile", is_start, profile_prefix)

996

997 async def reset_mm_cache_async(self) -> None:

998 await self.call_utility_async("reset_mm_cache")

vllm/v1/engine/llm_engine.py

+2 -2

@@ -326,8 +326,8 @@ class LLMEngine:

326

327 return processed_outputs.request_outputs

328

329 - def start_profile(self):

330 - self.engine_core.profile(True)

329 + def start_profile(self, profile_prefix: str | None = None):

330 + self.engine_core.profile(True, profile_prefix)

331

332 def stop_profile(self):

333 self.engine_core.profile(False)

vllm/v1/executor/abstract.py

+2 -2

@@ -238,8 +238,8 @@ class Executor(ABC):

238 def max_concurrent_batches(self) -> int:

239 return 1

240

241 - def profile(self, is_start: bool = True):

242 - self.collective_rpc("profile", args=(is_start,))

241 + def profile(self, is_start: bool = True, profile_prefix: str | None = None):

242 + self.collective_rpc("profile", args=(is_start, profile_prefix))

243

244 def save_sharded_state(

245 self,

vllm/v1/metrics/loggers.py

+2 -2

@@ -1305,8 +1305,8 @@ class StatLoggerManager:

1305 ):

1306 if engine_idx is None:

1307 engine_idx = 0

1308 - for logger in self.stat_loggers:

1309 - logger.record(

1308 + for stat_logger in self.stat_loggers:

1309 + stat_logger.record(

1310 scheduler_stats,

1311 iteration_stats,

1312 mm_cache_stats=mm_cache_stats,

vllm/v1/worker/cpu_worker.py

+1 -1

@@ -212,7 +212,7 @@ class CPUWorker(Worker):

212 )

213 return ",".join([str(x.id) for x in logical_cpu_list])

214

215 - def profile(self, is_start: bool = True):

215 + def profile(self, is_start: bool = True, profile_prefix: str | None = None):

216 if self.profiler is None:

217 raise RuntimeError("Profiler is not enabled.")

218 if is_start:

vllm/v1/worker/gpu_worker.py

+45 -16

@@ -103,20 +103,14 @@ class Worker(WorkerBase):

103 )

104

105 # Torch/CUDA profiler. Enabled and configured through profiler_config.

106 + # Profiler wrapper is created lazily in profile() when start is called,

107 + # so we have all the information needed for proper trace naming.

108 self.profiler: Any | None = None

107 - profiler_config = vllm_config.profiler_config

108 - if profiler_config.profiler == "torch":

109 - worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"

110 - self.profiler = TorchProfilerWrapper(

111 - profiler_config,

112 - worker_name=worker_name,

113 - local_rank=self.local_rank,

114 - activities=["CPU", "CUDA"],

115 - )

116 - elif profiler_config.profiler == "cuda":

117 - self.profiler = CudaProfilerWrapper(profiler_config)

118 - else:

119 - self.profiler = None

109 + self.profiler_config = vllm_config.profiler_config

110 +

111 + # Only validate profiler config is valid, don't instantiate yet

112 + if self.profiler_config.profiler not in ("torch", "cuda", None):

113 + raise ValueError(f"Unknown profiler type: {self.profiler_config.profiler}")

114

115 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER

116

@@ -677,17 +671,52 @@ class Worker(WorkerBase):

671 def take_draft_token_ids(self) -> DraftTokenIds | None:

672 return self.model_runner.take_draft_token_ids()

673

680 - def profile(self, is_start: bool = True):

681 - if self.profiler is None:

674 + def profile(self, is_start: bool = True, profile_prefix: str | None = None):

675 + # Check if profiling is enabled

676 + if self.profiler_config is None or self.profiler_config.profiler is None:

677 raise RuntimeError(

678 "Profiling is not enabled. Please set --profiler-config to enable "

679 "profiling. Example: "

680 "'--profiler-config.profiler=torch --profiler-config.torch_profiler_dir"

681 "=YOUR_DIR_PATH_TO_DUMP_TRACE'"

682 )

683 +

684 if is_start:

689 - self.profiler.start()

685 + # Generate the trace name by combining prefix with comprehensive rank suffix

686 + from vllm.distributed.utils import get_worker_rank_suffix

687 +

688 + rank_suffix = get_worker_rank_suffix(global_rank=self.rank)

689 +

690 + # Build the full trace name

691 + if profile_prefix:

692 + trace_name = f"{profile_prefix}_{rank_suffix}"

693 + else:

694 + trace_name = rank_suffix

695 +

696 + # Create the profiler wrapper only on the first start call

697 + if self.profiler is None:

698 + if self.profiler_config.profiler == "torch":

699 + self.profiler = TorchProfilerWrapper(

700 + self.profiler_config,

701 + worker_name=trace_name,

702 + local_rank=self.local_rank,

703 + activities=["CPU", "CUDA"],

704 + )

705 + logger.debug(

706 + "Starting torch profiler with trace name: %s", trace_name

707 + )

708 + elif self.profiler_config.profiler == "cuda":

709 + self.profiler = CudaProfilerWrapper(self.profiler_config)

710 + logger.debug("Starting CUDA profiler")

711 + self.profiler.start()

712 + else:

713 + # Profiler already initialized. Restart profiling but keep

714 + # the original trace name from the first initialization.

715 + self.profiler.start()

716 else:

717 + if self.profiler is None:

718 + logger.warning("Profiler was not started, nothing to stop.")

719 + return

720 self.profiler.stop()

721

722 def execute_dummy_batch(self) -> None: