[Enhancement] add pytorch profiler ops and memory record (#2472)

Signed-off-by: bjf-frz <frz123db@gmail.com>

.gitignore

@@ -174,6 +174,7 @@ CLAUDE.md
 
 # Codex
 AGENTS.md
+.codex
 .codex/
 
 # cursor

docs/contributing/profiling.md

@@ -138,9 +138,27 @@ Single-stage diffusion serving with torch profiler:
 vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
   --omni \
   --port 8091 \
-  --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
+  --profiler-config '{
+    "profiler": "torch",
+    "torch_profiler_dir": "/tmp/vllm_profile_wan22_i2v",
+    "torch_profiler_with_stack": true,
+    "torch_profiler_with_flops": false,
+    "torch_profiler_use_gzip": true,
+    "torch_profiler_dump_cuda_time_total": true,
+    "torch_profiler_record_shapes": true,
+    "torch_profiler_with_memory": true,
+    "delay_iterations": 0,
+    "max_iterations": 0
+  }'
 ```
 
+Useful optional fields:
+
+- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files.
+- `torch_profiler_record_shapes`: export `by_shape` operator views.
+- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history.
+- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`.
+
 Single-stage diffusion serving with Nsight Systems:
 
 ```bash
@@ -177,8 +195,11 @@ For mixed-stage pipelines, use explicit `stages` and pass the same stage list to
 
 Torch profiler output:
 
-- Chrome/Perfetto traces under `torch_profiler_dir`
-- Optional aggregated CUDA-time tables under the same directory
+- Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz`
+- Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets
+- Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled
+- Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend
+- Optional aggregated CUDA-time tables under the same session directory
 
 CUDA profiler / Nsight Systems output:
 

tests/profile/test_omni_torch_profiler.py

@@ -0,0 +1,582 @@
+# tests/test_omni_torch_profiler.py
+from __future__ import annotations
+
+import gzip
+from dataclasses import dataclass
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+from openpyxl import load_workbook
+
+import vllm_omni.profiler.omni_torch_profiler as profiler_mod
+from vllm_omni.profiler.omni_torch_profiler import OmniTorchProfilerWrapper
+
+
+@pytest.fixture(autouse=True)
+def patch_worker_profiler_init(monkeypatch):
+    def fake_init(self, profiler_config):
+        self.profiler_config = profiler_config
+
+    monkeypatch.setattr(
+        profiler_mod.WorkerProfiler,
+        "__init__",
+        fake_init,
+    )
+
+
+@dataclass
+class DummyProfilerConfig:
+    torch_profiler_dir: str
+    torch_profiler_use_gzip: bool = False
+    torch_profiler_record_shapes: bool = True
+    torch_profiler_with_memory: bool = True
+    torch_profiler_with_stack: bool = True
+    torch_profiler_with_flops: bool = False
+    torch_profiler_dump_cuda_time_total: bool = False
+
+
+class FakeEvent:
+    def __init__(
+        self,
+        *,
+        name: str = "aten::mm",
+        count: int = 1,
+        input_shapes=None,
+        stack=None,
+        self_cpu_time_total: float = 10.0,
+        cpu_time_total: float = 12.0,
+        self_cuda_time_total: float = 20.0,
+        cuda_time_total: float = 25.0,
+        self_xpu_time_total: float = 0.0,
+        xpu_time_total: float = 0.0,
+        self_cpu_memory_usage: int = 128,
+        cpu_memory_usage: int = 256,
+        self_cuda_memory_usage: int = 1024,
+        cuda_memory_usage: int = 2048,
+        self_xpu_memory_usage: int = 0,
+        xpu_memory_usage: int = 0,
+        device_type: str = "CUDA",
+        node_id: int = 0,
+        overload_name: str = "",
+        is_async: bool = False,
+        is_legacy: bool = False,
+    ):
+        self.key = name
+        self.name = name
+        self.count = count
+        self.input_shapes = input_shapes if input_shapes is not None else [[2, 2], [2, 2]]
+        self.stack = stack if stack is not None else ["frame_a", "frame_b"]
+        self.self_cpu_time_total = self_cpu_time_total
+        self.cpu_time_total = cpu_time_total
+        self.self_cuda_time_total = self_cuda_time_total
+        self.cuda_time_total = cuda_time_total
+        self.self_xpu_time_total = self_xpu_time_total
+        self.xpu_time_total = xpu_time_total
+        self.self_cpu_memory_usage = self_cpu_memory_usage
+        self.cpu_memory_usage = cpu_memory_usage
+        self.self_cuda_memory_usage = self_cuda_memory_usage
+        self.cuda_memory_usage = cuda_memory_usage
+        self.self_xpu_memory_usage = self_xpu_memory_usage
+        self.xpu_memory_usage = xpu_memory_usage
+        self.device_type = device_type
+        self.node_id = node_id
+        self.overload_name = overload_name
+        self.is_async = is_async
+        self.is_legacy = is_legacy
+
+
+class FakeEventList(list):
+    def table(self, sort_by=None, row_limit=-1):
+        return f"fake_table(sort_by={sort_by}, row_limit={row_limit}, len={len(self)})"
+
+
+class FakeTorchProfiler:
+    def __init__(self, on_trace_ready=None):
+        self.started = False
+        self.stopped = False
+        self.on_trace_ready = on_trace_ready
+        self.exported_traces = []
+        self.exported_stacks = []
+
+    def start(self):
+        self.started = True
+
+    def stop(self):
+        self.stopped = True
+        if self.on_trace_ready is not None:
+            self.on_trace_ready(self)
+
+    def export_chrome_trace(self, path):
+        Path(path).write_text('{"traceEvents": []}')
+        self.exported_traces.append(path)
+
+    def export_stacks(self, path, metric):
+        Path(path).write_text(f"metric={metric}\nstack_line_1\nstack_line_2\n")
+        self.exported_stacks.append((path, metric))
+
+    def key_averages(self, group_by_input_shape=False, group_by_stack_n=0):
+        if group_by_input_shape:
+            return FakeEventList(
+                [
+                    FakeEvent(
+                        name="aten::bmm",
+                        input_shapes=[[4, 8, 16], [4, 16, 32]],
+                    )
+                ]
+            )
+        if group_by_stack_n:
+            return FakeEventList(
+                [
+                    FakeEvent(
+                        name="aten::all_reduce",
+                        stack=["python_a", "python_b", "python_c"],
+                    )
+                ]
+            )
+        return FakeEventList(
+            [
+                FakeEvent(name="aten::mm"),
+                FakeEvent(name="nccl:all_reduce"),
+            ]
+        )
+
+
+@pytest.fixture
+def fake_config(tmp_path):
+    return DummyProfilerConfig(torch_profiler_dir=str(tmp_path))
+
+
+@pytest.fixture
+def fake_profiler_factory(monkeypatch):
+    created = {}
+
+    def fake_profile(*args, **kwargs):
+        profiler = FakeTorchProfiler(on_trace_ready=kwargs.get("on_trace_ready"))
+        created["profiler"] = profiler
+        created["args"] = args
+        created["kwargs"] = kwargs
+        return profiler
+
+    monkeypatch.setattr(profiler_mod.torch.profiler, "profile", fake_profile)
+    return created
+
+
+@pytest.fixture
+def wrapper(fake_config, fake_profiler_factory):
+    return OmniTorchProfilerWrapper(
+        profiler_config=fake_config,
+        worker_name="worker0",
+        local_rank=0,
+        activities=["CPU", "CUDA"],
+    )
+
+
+def test_set_trace_filename_creates_timestamped_session_dir(wrapper, monkeypatch, tmp_path):
+    class FixedDatetime:
+        @classmethod
+        def now(cls):
+            class _Now:
+                def strftime(self, fmt):
+                    return "20260403-034200"
+
+            return _Now()
+
+    monkeypatch.setattr(profiler_mod, "datetime", FixedDatetime)
+
+    wrapper.set_trace_filename("stage_0_llm_1234567890")
+
+    session_dir = Path(wrapper._session_dir)
+    assert session_dir.exists()
+    assert session_dir.parent == tmp_path
+    assert session_dir.name == "20260403-034200_stage_0_llm_1234567890"
+
+
+def test_set_trace_filename_with_full_path_creates_timestamped_leaf(wrapper, monkeypatch, tmp_path):
+    class FixedDatetime:
+        @classmethod
+        def now(cls):
+            class _Now:
+                def strftime(self, fmt):
+                    return "20260403-111111"
+
+            return _Now()
+
+    monkeypatch.setattr(profiler_mod, "datetime", FixedDatetime)
+
+    target = tmp_path / "nested" / "stage_x"
+    wrapper.set_trace_filename(str(target))
+
+    session_dir = Path(wrapper._session_dir)
+    assert session_dir.exists()
+    assert session_dir.parent == target.parent
+    assert session_dir.name == "20260403-111111_stage_x"
+
+
+def test_on_trace_ready_exports_trace_json(wrapper):
+    wrapper.set_trace_filename("case_trace")
+
+    wrapper._on_trace_ready(wrapper.profiler)
+
+    trace_path = Path(wrapper._trace_path)
+    assert trace_path.exists()
+    assert trace_path.name == "trace_rank0.json"
+    assert trace_path.read_text() == '{"traceEvents": []}'
+
+
+def test_on_trace_ready_exports_gzip_trace(fake_config, fake_profiler_factory, monkeypatch):
+    fake_config.torch_profiler_use_gzip = True
+
+    wrapper = OmniTorchProfilerWrapper(
+        profiler_config=fake_config,
+        worker_name="worker0",
+        local_rank=0,
+        activities=["CPU", "CUDA"],
+    )
+    wrapper.set_trace_filename("case_gzip")
+
+    def fake_popen(cmd):
+        assert cmd[:2] == ["gzip", "-f"]
+        src = Path(cmd[2])
+        gz_path = src.with_suffix(src.suffix + ".gz")
+        gz_path.write_bytes(gzip.compress(src.read_bytes()))
+        src.unlink()
+
+        class DummyProc:
+            pass
+
+        return DummyProc()
+
+    monkeypatch.setattr(profiler_mod.subprocess, "Popen", fake_popen)
+
+    wrapper._on_trace_ready(wrapper.profiler)
+
+    assert wrapper._trace_path.endswith(".json.gz")
+    gz_path = Path(wrapper._trace_path)
+    assert gz_path.exists()
+    assert gzip.decompress(gz_path.read_bytes()) == b'{"traceEvents": []}'
+
+
+def test_start_enables_memory_history(wrapper, monkeypatch):
+    calls = []
+
+    monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True)
+
+    def fake_record_memory_history(*args, **kwargs):
+        calls.append((args, kwargs))
+
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_record_memory_history",
+        fake_record_memory_history,
+    )
+
+    wrapper.set_trace_filename("case_memory_start")
+    wrapper._start()
+
+    assert wrapper.profiler.started is True
+    assert wrapper._memory_history_enabled is True
+    assert len(calls) == 1
+    assert calls[0][1]["enabled"] == "all"
+    assert calls[0][1]["context"] == "all"
+    assert calls[0][1]["stacks"] == "python"
+    assert calls[0][1]["max_entries"] == 100000
+    assert calls[0][1]["clear_history"] is True
+
+
+def test_start_skips_memory_history_when_memory_disabled(fake_config, fake_profiler_factory, monkeypatch):
+    fake_config.torch_profiler_with_memory = False
+
+    wrapper = OmniTorchProfilerWrapper(
+        profiler_config=fake_config,
+        worker_name="worker0",
+        local_rank=0,
+        activities=["CPU", "CUDA"],
+    )
+
+    called = {"n": 0}
+
+    monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True)
+
+    def fake_record_memory_history(*args, **kwargs):
+        called["n"] += 1
+
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_record_memory_history",
+        fake_record_memory_history,
+    )
+
+    wrapper.set_trace_filename("case_skip_memory")
+    wrapper._start()
+
+    assert called["n"] == 0
+    assert wrapper._memory_history_enabled is False
+
+
+def test_try_dump_memory_snapshot_writes_pickle(wrapper, monkeypatch):
+    wrapper.set_trace_filename("case_snapshot")
+    wrapper._memory_history_enabled = True
+    wrapper._memory_history_backend = "CUDA"
+    wrapper._memory_history_module = profiler_mod.torch.cuda.memory
+
+    disable_calls = []
+
+    def fake_record_memory_history(*args, **kwargs):
+        disable_calls.append((args, kwargs))
+
+    def fake_dump_snapshot(path):
+        Path(path).write_bytes(b"fake pickle bytes")
+
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_record_memory_history",
+        fake_record_memory_history,
+    )
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_dump_snapshot",
+        fake_dump_snapshot,
+    )
+
+    wrapper._try_dump_memory_snapshot()
+
+    snapshot = Path(wrapper._artifact_paths["memory_snapshot"])
+    assert snapshot.exists()
+    assert snapshot.name == "memory_snapshot_rank0.pickle"
+    assert snapshot.read_bytes() == b"fake pickle bytes"
+    assert wrapper._memory_history_enabled is False
+
+    assert disable_calls[-1][1]["enabled"] is None
+
+
+def test_stop_always_dumps_memory_snapshot_on_success_path(wrapper, monkeypatch):
+    wrapper.set_trace_filename("case_stop")
+
+    record_calls = []
+    dump_calls = []
+
+    monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True)
+
+    def fake_record_memory_history(*args, **kwargs):
+        record_calls.append((args, kwargs))
+
+    def fake_dump_snapshot(path):
+        dump_calls.append(path)
+        Path(path).write_bytes(b"snapshot-bytes")
+
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_record_memory_history",
+        fake_record_memory_history,
+    )
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_dump_snapshot",
+        fake_dump_snapshot,
+    )
+
+    wrapper._start()
+    wrapper._stop()
+
+    session_dir = Path(wrapper._session_dir)
+
+    assert wrapper.profiler.started is True
+    assert wrapper.profiler.stopped is True
+    assert (session_dir / "memory_snapshot_rank0.pickle").exists()
+    assert len(dump_calls) == 1
+    assert record_calls[0][1]["enabled"] == "all"
+    assert record_calls[-1][1]["enabled"] is None
+
+
+def test_on_stop_hook_generates_stack_and_excel_artifacts(wrapper):
+    wrapper.set_trace_filename("case_artifacts")
+    wrapper._on_stop_hook()
+
+    session_dir = Path(wrapper._session_dir)
+
+    assert not (session_dir / "ops_summary_rank0.txt").exists()
+    assert not (session_dir / "ops_by_shape_rank0.txt").exists()
+    assert not (session_dir / "ops_by_stack_rank0.txt").exists()
+    assert (session_dir / "stacks_cpu_rank0.txt").exists()
+    assert (session_dir / "stacks_cuda_rank0.txt").exists()
+    assert (session_dir / "ops_rank0.xlsx").exists()
+
+
+def test_excel_contains_expected_sheets(wrapper):
+    wrapper.set_trace_filename("case_excel")
+    wrapper._on_stop_hook()
+
+    xlsx_path = Path(wrapper._session_dir) / "ops_rank0.xlsx"
+    wb = load_workbook(xlsx_path)
+
+    assert "summary" in wb.sheetnames
+    assert "by_shape" in wb.sheetnames
+    assert "by_stack" in wb.sheetnames
+
+
+def test_excel_summary_has_expected_columns(wrapper):
+    wrapper.set_trace_filename("case_excel_columns")
+    wrapper._on_stop_hook()
+
+    xlsx_path = Path(wrapper._session_dir) / "ops_rank0.xlsx"
+    wb = load_workbook(xlsx_path)
+    ws = wb["summary"]
+
+    headers = [cell.value for cell in next(ws.iter_rows(min_row=1, max_row=1))]
+    assert "name" in headers
+    assert "count" in headers
+    assert "self_cpu_time_total_us" in headers
+    assert "self_cuda_time_total_us" in headers
+    assert "self_cpu_memory_usage_bytes" in headers
+    assert "self_cuda_memory_usage_bytes" in headers
+    assert "input_shapes" in headers
+    assert "stack" in headers
+
+
+def test_get_results_returns_all_artifact_paths(wrapper, monkeypatch):
+    wrapper.set_trace_filename("case_results")
+
+    monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_record_memory_history",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        profiler_mod.torch.cuda.memory,
+        "_dump_snapshot",
+        lambda path: Path(path).write_bytes(b"snapshot"),
+    )
+
+    wrapper._start()
+    wrapper._stop()
+
+    results = wrapper.get_results()
+
+    assert "trace" in results
+    assert "table" in results
+    assert "session_dir" in results
+    assert "ops" in results
+    assert "memory_snapshot" in results
+    assert Path(results["session_dir"]).exists()
+    assert Path(results["ops"]).exists()
+    assert Path(results["table"]).exists()
+    assert Path(results["table"]).name == "ops_rank0.xlsx"
+    assert Path(results["memory_snapshot"]).exists()
+
+
+def test_start_uses_xpu_memory_history_when_available(wrapper, monkeypatch):
+    calls = []
+
+    def fake_record_memory_history(*args, **kwargs):
+        calls.append((args, kwargs))
+
+    fake_memory_module = SimpleNamespace(
+        _record_memory_history=fake_record_memory_history,
+    )
+    monkeypatch.setattr(
+        wrapper,
+        "_resolve_memory_history_backend",
+        lambda: ("XPU", fake_memory_module),
+    )
+
+    wrapper.set_trace_filename("case_xpu_memory_start")
+    wrapper._start()
+
+    assert wrapper._memory_history_enabled is True
+    assert wrapper._memory_history_backend == "XPU"
+    assert wrapper._memory_history_module is fake_memory_module
+    assert calls[0][1]["enabled"] == "all"
+
+
+def test_start_uses_npu_memory_history_when_available(wrapper, monkeypatch):
+    calls = []
+
+    def fake_record_memory_history(*args, **kwargs):
+        calls.append((args, kwargs))
+
+    fake_memory_module = SimpleNamespace(
+        _record_memory_history=fake_record_memory_history,
+    )
+    monkeypatch.setattr(
+        wrapper,
+        "_resolve_memory_history_backend",
+        lambda: ("NPU", fake_memory_module),
+    )
+
+    wrapper.set_trace_filename("case_npu_memory_start")
+    wrapper._start()
+
+    assert wrapper._memory_history_enabled is True
+    assert wrapper._memory_history_backend == "NPU"
+    assert wrapper._memory_history_module is fake_memory_module
+    assert calls[0][1]["enabled"] == "all"
+
+
+def test_start_skips_memory_history_when_backend_api_missing(wrapper, monkeypatch):
+    fake_memory_module = SimpleNamespace()
+    monkeypatch.setattr(
+        wrapper,
+        "_resolve_memory_history_backend",
+        lambda: ("XPU", fake_memory_module),
+    )
+
+    wrapper.set_trace_filename("case_missing_memory_api")
+    wrapper._start()
+
+    assert wrapper._memory_history_enabled is False
+    assert wrapper._memory_history_backend is None
+    assert wrapper._memory_history_module is None
+
+
+def test_try_dump_memory_snapshot_uses_resolved_backend_module(wrapper):
+    wrapper.set_trace_filename("case_xpu_snapshot")
+    wrapper._memory_history_enabled = True
+    wrapper._memory_history_backend = "XPU"
+
+    calls = []
+
+    def fake_record_memory_history(*args, **kwargs):
+        calls.append((args, kwargs))
+
+    def fake_dump_snapshot(path):
+        Path(path).write_bytes(b"xpu snapshot bytes")
+
+    wrapper._memory_history_module = SimpleNamespace(
+        _record_memory_history=fake_record_memory_history,
+        _dump_snapshot=fake_dump_snapshot,
+    )
+
+    wrapper._try_dump_memory_snapshot()
+
+    snapshot = Path(wrapper._artifact_paths["memory_snapshot"])
+    assert snapshot.exists()
+    assert snapshot.read_bytes() == b"xpu snapshot bytes"
+    assert calls[-1][1]["enabled"] is None
+    assert wrapper._memory_history_enabled is False
+    assert wrapper._memory_history_backend is None
+    assert wrapper._memory_history_module is None
+
+
+def test_event_list_to_rows_contains_expected_fields(wrapper):
+    rows = wrapper._event_list_to_rows(
+        [
+            FakeEvent(
+                name="aten::linear",
+                input_shapes=[[8, 16], [16, 32]],
+                stack=["f1", "f2"],
+            )
+        ]
+    )
+
+    assert len(rows) == 1
+    row = rows[0]
+    assert row["name"] == "aten::linear"
+    assert row["count"] == 1
+    assert row["self_cpu_time_total_us"] == 10.0
+    assert row["self_cuda_time_total_us"] == 20.0
+    assert row["self_cpu_memory_usage_bytes"] == 128
+    assert row["self_cuda_memory_usage_bytes"] == 1024
+    assert "[[8, 16], [16, 32]]" == row["input_shapes"]
+    assert row["stack"] == "f1\nf2"

vllm_omni/profiler/omni_torch_profiler.py

@@ -5,8 +5,8 @@ from __future__ import annotations
 
 import os
 import subprocess
-import time
-from typing import Literal
+from datetime import datetime
+from typing import Any, Literal
 
 import torch
 from typing_extensions import override
@@ -62,6 +62,13 @@ class OmniTorchProfilerWrapper(WorkerProfiler):
         self._trace_path: str | None = None
         self._table_path: str | None = None
 
+        self._activities = activities
+        self._session_dir: str | None = None
+        self._artifact_paths: dict[str, str | None] = {}
+        self._memory_history_enabled = False
+        self._memory_history_backend: str | None = None
+        self._memory_history_module = None
+
         if local_rank in (None, 0):
             logger.info_once(
                 "Omni torch profiling enabled. Traces will be saved to: %s",
@@ -72,6 +79,9 @@ class OmniTorchProfilerWrapper(WorkerProfiler):
         self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
         self.profiler = self._create_profiler(profiler_config, activities)
 
+    def _rank(self) -> int:
+        return 0 if self.local_rank is None else self.local_rank
+
     def _get_default_activities(self) -> list[TorchProfilerActivity]:
         """Get default activities for this platform.
 
@@ -106,19 +116,58 @@ class OmniTorchProfilerWrapper(WorkerProfiler):
                       Can also be a full path (e.g. from diffusion engine).
         """
         self._trace_filename = filename
+        self._session_dir = None
+        self._ensure_session_dir()
+
+    def _ensure_session_dir(self) -> str:
+        """Create one timestamped directory for this profiling run."""
+        if self._session_dir is not None:
+            os.makedirs(self._session_dir, exist_ok=True)
+            return self._session_dir
+
+        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
+        base_name = self._trace_filename or self._worker_name
+
+        if os.path.dirname(base_name):
+            parent_dir = os.path.dirname(base_name)
+            leaf_name = os.path.basename(base_name)
+            session_name = f"{ts}_{leaf_name}"
+            self._session_dir = os.path.join(parent_dir, session_name)
+        else:
+            session_name = f"{ts}_{base_name}"
+            self._session_dir = os.path.join(self._trace_dir, session_name)
+
+        os.makedirs(self._session_dir, exist_ok=True)
+        self._artifact_paths["session_dir"] = self._session_dir
+        return self._session_dir
+
+    def _artifact_path(self, stem: str, suffix: str) -> str:
+        """Build artifact path under the session directory."""
+        return os.path.join(
+            self._ensure_session_dir(),
+            f"{stem}_rank{self._rank()}{suffix}",
+        )
+
+    def _write_text_artifact(self, name: str, content: str) -> str:
+        path = self._artifact_path(name, ".txt")
+        with open(path, "w") as f:
+            f.write(content)
+        self._artifact_paths[name] = path
+        return path
+
+    def _has_cuda_like_activity(self) -> bool:
+        return any(a in self._activities for a in ("CUDA", "MUSA"))
+
+    def _get_time_sort_key(self) -> str:
+        if self._has_cuda_like_activity():
+            return "self_cuda_time_total"
+        return "self_cpu_time_total"
 
     def _on_trace_ready(self, prof) -> None:
         """Custom trace handler: export chrome trace with omni naming."""
-        rank = self.local_rank
-        filename = self._trace_filename or f"{self._worker_name}_{int(time.time())}"
-        # If filename already contains a directory, use as-is (e.g. from
-        # diffusion engine which builds full path). Otherwise join with trace_dir.
-        if os.path.dirname(filename):
-            json_file = f"{filename}_rank{rank}.json"
-        else:
-            json_file = os.path.join(self._trace_dir, f"{filename}_rank{rank}.json")
+        rank = self._rank()
 
-        os.makedirs(os.path.dirname(json_file), exist_ok=True)
+        json_file = self._artifact_path("trace", ".json")
 
         try:
             prof.export_chrome_trace(json_file)
@@ -143,18 +192,211 @@ class OmniTorchProfilerWrapper(WorkerProfiler):
             else:
                 self._trace_path = json_file
 
+            self._artifact_paths["trace"] = self._trace_path
+
         except Exception as e:
             logger.warning("[Rank %s] Failed to export trace: %s", rank, e)
 
+    def _try_enable_memory_history(self) -> None:
+        """Enable backend-specific memory history for snapshot analysis."""
+        if not self.profiler_config.torch_profiler_with_memory:
+            return
+
+        backend_name, memory_module = self._resolve_memory_history_backend()
+        if backend_name is None or memory_module is None:
+            return
+
+        record_memory_history = getattr(memory_module, "_record_memory_history", None)
+        if record_memory_history is None:
+            logger.info(
+                "[Rank %s] %s memory history is not supported on this platform",
+                self._rank(),
+                backend_name,
+            )
+            return
+
+        try:
+            record_memory_history(
+                enabled="all",
+                context="all",
+                stacks="python",
+                max_entries=100000,
+                clear_history=True,
+            )
+            self._memory_history_enabled = True
+            self._memory_history_backend = backend_name
+            self._memory_history_module = memory_module
+            logger.info("[Rank %s] %s memory history enabled", self._rank(), backend_name)
+        except Exception as e:
+            logger.warning(
+                "[Rank %s] Failed to enable %s memory history: %s",
+                self._rank(),
+                backend_name,
+                e,
+            )
+
+    def _try_dump_memory_snapshot(self) -> None:
+        """Dump a backend-specific memory snapshot into the session directory."""
+        if not self._memory_history_enabled:
+            return
+
+        try:
+            if self._memory_history_module is None or self._memory_history_backend is None:
+                return
+
+            dump_snapshot = getattr(self._memory_history_module, "_dump_snapshot", None)
+            if dump_snapshot is None:
+                logger.info(
+                    "[Rank %s] %s memory snapshot is not supported on this platform",
+                    self._rank(),
+                    self._memory_history_backend,
+                )
+                return
+
+            snapshot_file = self._artifact_path("memory_snapshot", ".pickle")
+            dump_snapshot(snapshot_file)
+            self._artifact_paths["memory_snapshot"] = snapshot_file
+            logger.info(
+                "[Rank %s] %s memory snapshot dumped to %s",
+                self._rank(),
+                self._memory_history_backend,
+                snapshot_file,
+            )
+        except Exception as e:
+            logger.warning(
+                "[Rank %s] Failed to dump %s memory snapshot: %s",
+                self._rank(),
+                self._memory_history_backend,
+                e,
+            )
+        finally:
+            try:
+                if self._memory_history_module is not None:
+                    disable_memory_history = getattr(
+                        self._memory_history_module,
+                        "_record_memory_history",
+                        None,
+                    )
+                    if disable_memory_history is not None:
+                        disable_memory_history(enabled=None)
+            except Exception:
+                pass
+            self._memory_history_enabled = False
+            self._memory_history_backend = None
+            self._memory_history_module = None
+
+    def _resolve_memory_history_backend(self) -> tuple[str | None, Any]:
+        """Resolve the memory backend that supports history/snapshot APIs."""
+        backend_specs = [
+            ("CUDA", self._has_cuda_like_activity(), getattr(torch, "cuda", None)),
+            ("NPU", "NPU" in self._activities, getattr(torch, "npu", None)),
+            ("XPU", "XPU" in self._activities, getattr(torch, "xpu", None)),
+            ("MUSA", "MUSA" in self._activities, getattr(torch, "musa", None)),
+        ]
+
+        for backend_name, enabled, device_module in backend_specs:
+            if not enabled or device_module is None:
+                continue
+
+            is_available = getattr(device_module, "is_available", None)
+            if callable(is_available) and not is_available():
+                continue
+
+            memory_module = getattr(device_module, "memory", None)
+            if memory_module is not None:
+                return backend_name, memory_module
+
+        return None, None
+
+    def _safe_get(self, obj, name: str, default=None):
+        return getattr(obj, name, default)
+
+    def _event_list_to_rows(self, event_list) -> list[dict]:
+        rows = []
+        for evt in event_list:
+            row = {
+                "name": self._safe_get(evt, "key", None) or self._safe_get(evt, "name", None),
+                "count": self._safe_get(evt, "count", None),
+                "device_type": self._safe_get(evt, "device_type", None),
+                "node_id": self._safe_get(evt, "node_id", None),
+                "self_cpu_time_total_us": self._safe_get(evt, "self_cpu_time_total", None),
+                "cpu_time_total_us": self._safe_get(evt, "cpu_time_total", None),
+                "self_cuda_time_total_us": self._safe_get(evt, "self_cuda_time_total", None),
+                "cuda_time_total_us": self._safe_get(evt, "cuda_time_total", None),
+                "self_xpu_time_total_us": self._safe_get(evt, "self_xpu_time_total", None),
+                "xpu_time_total_us": self._safe_get(evt, "xpu_time_total", None),
+                "self_cpu_memory_usage_bytes": self._safe_get(evt, "self_cpu_memory_usage", None),
+                "cpu_memory_usage_bytes": self._safe_get(evt, "cpu_memory_usage", None),
+                "self_cuda_memory_usage_bytes": self._safe_get(evt, "self_cuda_memory_usage", None),
+                "cuda_memory_usage_bytes": self._safe_get(evt, "cuda_memory_usage", None),
+                "self_xpu_memory_usage_bytes": self._safe_get(evt, "self_xpu_memory_usage", None),
+                "xpu_memory_usage_bytes": self._safe_get(evt, "xpu_memory_usage", None),
+                "input_shapes": str(self._safe_get(evt, "input_shapes", None)),
+                "stack": "\n".join(self._safe_get(evt, "stack", []) or []),
+                "overload_name": self._safe_get(evt, "overload_name", None),
+                "is_async": self._safe_get(evt, "is_async", None),
+                "is_legacy": self._safe_get(evt, "is_legacy", None),
+            }
+            rows.append(row)
+        return rows
+
+    def _write_excel_artifact(self, name: str, sheets: dict[str, list[dict]]) -> str:
+        path = self._artifact_path(name, ".xlsx")
+
+        try:
+            import pandas as pd
+        except Exception as e:
+            logger.warning(
+                "[Rank %s] pandas not available, skip Excel export: %s",
+                self._rank(),
+                e,
+            )
+            return path
+
+        with pd.ExcelWriter(path, engine="openpyxl") as writer:
+            for sheet_name, rows in sheets.items():
+                df = pd.DataFrame(rows)
+
+                safe_sheet_name = sheet_name if sheet_name else "Sheet1"
+
+                df.to_excel(
+                    writer,
+                    sheet_name=safe_sheet_name,
+                    index=False,
+                    freeze_panes=(1, 0),
+                )
+
+                ws = writer.sheets[safe_sheet_name]
+                ws.auto_filter.ref = ws.dimensions
+
+                for col_cells in ws.columns:
+                    max_len = 0
+                    col_letter = col_cells[0].column_letter
+                    for cell in col_cells[:200]:
+                        try:
+                            val = "" if cell.value is None else str(cell.value)
+                            max_len = max(max_len, len(val))
+                        except Exception:
+                            pass
+                    ws.column_dimensions[col_letter].width = min(max(max_len + 2, 12), 80)
+
+        self._artifact_paths[name] = path
+        return path
+
     @override
     def _start(self) -> None:
+        self._ensure_session_dir()
+        self._try_enable_memory_history()
         self.profiler.start()
 
     @override
     def _stop(self) -> None:
         """Stop profiler, export trace via on_trace_ready, and dump table."""
         self.profiler.stop()
-        self._on_stop_hook()
+        try:
+            self._on_stop_hook()
+        finally:
+            self._try_dump_memory_snapshot()
 
     def _on_stop_hook(self) -> None:
         """Hook called after profiler.stop().
@@ -163,6 +405,68 @@ class OmniTorchProfilerWrapper(WorkerProfiler):
         Base implementation handles CUDA time total dump.
         """
         rank = self.local_rank
+        sort_key = self._get_time_sort_key()
+
+        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)
+
+            if self._has_cuda_like_activity():
+                try:
+                    cuda_stack_file = self._artifact_path("stacks_cuda", ".txt")
+                    self.profiler.export_stacks(
+                        cuda_stack_file,
+                        metric="self_cuda_time_total",
+                    )
+                    self._artifact_paths["stacks_cuda"] = cuda_stack_file
+                except Exception as e:
+                    logger.warning("[Rank %s] export_stacks(cuda) 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)
 
         if self.profiler_config.torch_profiler_dump_cuda_time_total:
             profiler_dir = self.profiler_config.torch_profiler_dir
@@ -190,6 +494,7 @@ class OmniTorchProfilerWrapper(WorkerProfiler):
         return {
             "trace": self._trace_path,
             "table": self._table_path,
+            **self._artifact_paths,
         }