Chendi.XueandHongsheng Liu
4e41b7bc23
Signed-off-by: Chendi Xue <chendi.xue@intel.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
110 lines
4.0 KiB
Python
110 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing_extensions import override
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.profiler.omni_torch_profiler import (
|
|
OmniTorchProfilerWrapper,
|
|
TorchProfilerActivity,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class XPUTorchProfilerWrapper(OmniTorchProfilerWrapper):
|
|
"""XPU-specific profiler wrapper.
|
|
|
|
Uses torch.profiler with XPU ProfilerActivity to capture Intel GPU events.
|
|
XPU events are exposed via torch.profiler.ProfilerActivity.XPU and work
|
|
with the standard torch.profiler interface.
|
|
"""
|
|
|
|
@override
|
|
def _get_default_activities(self) -> list[TorchProfilerActivity]:
|
|
"""Default to CPU + XPU profiling for Intel GPUs."""
|
|
return ["CPU", "XPU"]
|
|
|
|
@override
|
|
def _on_stop_hook(self) -> None:
|
|
"""XPU-specific stop hook that handles XPU event aggregation.
|
|
|
|
XPU profiler follows CUDA-like patterns but uses xpu_time_total metrics.
|
|
"""
|
|
rank = self.local_rank
|
|
sort_key = "self_xpu_time_total"
|
|
|
|
excel_sheets: dict[str, list[dict]] = {}
|
|
|
|
# 1) Summary op table
|
|
summary_events = self.profiler.key_averages()
|
|
excel_sheets["summary"] = self._event_list_to_rows(summary_events)
|
|
|
|
# 2) Shape-grouped op table
|
|
if self.profiler_config.torch_profiler_record_shapes:
|
|
try:
|
|
shape_events = self.profiler.key_averages(
|
|
group_by_input_shape=True,
|
|
)
|
|
excel_sheets["by_shape"] = self._event_list_to_rows(shape_events)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"[Rank %s] Failed to export shape-grouped op table: %s",
|
|
rank,
|
|
e,
|
|
)
|
|
|
|
# 3) Stack-grouped op table
|
|
if self.profiler_config.torch_profiler_with_stack:
|
|
try:
|
|
stack_events = self.profiler.key_averages(
|
|
group_by_stack_n=8,
|
|
)
|
|
excel_sheets["by_stack"] = self._event_list_to_rows(stack_events)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"[Rank %s] Failed to export stack-grouped op table: %s",
|
|
rank,
|
|
e,
|
|
)
|
|
|
|
# 4) Export stack files
|
|
try:
|
|
cpu_stack_file = self._artifact_path("stacks_cpu", ".txt")
|
|
self.profiler.export_stacks(
|
|
cpu_stack_file,
|
|
metric="self_cpu_time_total",
|
|
)
|
|
self._artifact_paths["stacks_cpu"] = cpu_stack_file
|
|
except Exception as e:
|
|
logger.warning("[Rank %s] export_stacks(cpu) failed: %s", rank, e)
|
|
|
|
# XPU stack traces
|
|
try:
|
|
xpu_stack_file = self._artifact_path("stacks_xpu", ".txt")
|
|
self.profiler.export_stacks(
|
|
xpu_stack_file,
|
|
metric="self_xpu_time_total",
|
|
)
|
|
self._artifact_paths["stacks_xpu"] = xpu_stack_file
|
|
except Exception as e:
|
|
logger.warning("[Rank %s] export_stacks(xpu) failed: %s", rank, e)
|
|
|
|
try:
|
|
self._table_path = self._write_excel_artifact("ops", excel_sheets)
|
|
except Exception as e:
|
|
logger.warning("[Rank %s] Failed to export Excel workbook: %s", rank, e)
|
|
|
|
# Print XPU time table for rank 0
|
|
if rank == 0:
|
|
try:
|
|
table = self.profiler.key_averages().table(sort_by=sort_key, row_limit=50)
|
|
logger.info("XPU Profiler Results:\n%s", table)
|
|
except Exception as e:
|
|
logger.warning("[Rank %s] Failed to generate XPU time table: %s", rank, e)
|
|
|
|
if self.dump_cpu_time_total and rank == 0:
|
|
logger.info(self.profiler.key_averages().table(sort_by="self_cpu_time_total", row_limit=50))
|