feat: add --memory flag to measure RAM and VRAM per phase

src/winml/modelkit/commands/_live_chart.py

@@ -253,10 +253,10 @@ def _render_status(
         # CPU-only mode drops the adapter cell + device-memory cell since we
         # have no live values to populate them with.
         cpu_cell = f"CPU: {cpu_pct:.1f}%"
-        ram_cell = f"Mem: {ram_mb:.0f} MB"
+        ram_cell = f"RAM: {ram_mb:.0f} MB"
         if self._show_adapter:
             adapter_cell = f"{self._adapter_label}: {mean_util:.1f}% avg ({current_util:.1f}% now)"
-            mem_cell = f"Device Mem: {memory_local_mb:.0f}/{memory_shared_mb:.0f} MB (local/shared)"
+            mem_cell = f"VRAM: {memory_local_mb:.0f}/{memory_shared_mb:.0f} MB (local/shared)"
             row2 = f"  {adapter_cell:<30}| {cpu_cell:<12}|  {ram_cell}  |  {mem_cell}"
         else:
             row2 = f"  {cpu_cell:<12}|  {ram_cell}"

src/winml/modelkit/commands/perf.py

@@ -84,6 +84,7 @@ class BenchmarkConfig:
     skip_build: bool = True
     allow_unsupported_nodes: bool = False
     monitor: bool = False
+    memory: bool = True
     ep: EPNameOrAlias | None = None
     ep_options: dict[str, str] | None = None
     shape_config: dict | None = None
@@ -140,6 +141,9 @@ class BenchmarkResult:
     # Hardware monitor metrics (from HWMonitor.to_dict())
     hw_monitor: dict[str, Any] | None = None
 
+    # Memory profile dict (rss deltas from memory_tracker)
+    memory_profile: dict[str, float] | None = None
+
     def to_dict(self) -> dict[str, Any]:
         """Convert to dictionary for JSON serialization."""
         result: dict[str, Any] = {
@@ -183,6 +187,8 @@ def to_dict(self) -> dict[str, Any]:
         }
         if self.hw_monitor:
             result["hw_monitor"] = self.hw_monitor
+        if self.memory_profile:
+            result["memory"] = self.memory_profile
         return result
 
 
@@ -295,6 +301,7 @@ def __init__(self, config: BenchmarkConfig) -> None:
         self.config = config
         self._model: WinMLPreTrainedModel | WinMLCompositeModel | None = None
         self._inputs: dict[str, np.ndarray] | None = None
+        self._memory: dict[str, float] | None = None
 
     @property
     def _is_composite(self) -> bool:
@@ -345,7 +352,7 @@ def run(self) -> BenchmarkResult | dict[str, BenchmarkResult]:
             ORT session, so each sub-model is benchmarked individually rather
             than timing the aggregate ``forward()`` pass.
         """
-        # [1] Load model
+        # [1] Load model (build pipeline: optimize, cache, etc.)
         logger.info("Loading model: %s", self.config.model_id)
         self._load_model()
         assert self._model is not None
@@ -381,15 +388,39 @@ def _run_single(self) -> BenchmarkResult:
         Returns:
             BenchmarkResult with timing statistics
         """
+        import gc
+
         assert self._model is not None
 
+        # Initialize memory tracking variables
+        adapter_luid: str | None = None
+        rss_baseline = rss_after_compile = 0.0
+        vram_local_baseline = vram_shared_baseline = 0.0
+        vram_local_compile = vram_shared_compile = 0.0
+
+        # Memory: baseline right before compile() — excludes all Python lib
+        # imports, EP DLLs, and build pipeline overhead. Measures only ORT
+        # session compilation (model weights loaded into memory).
+        if self.config.memory:
+            from ..session.monitor.memory_tracker import get_rss_mb, get_vram_mb
+
+            adapter_luid = self._resolve_adapter_luid()
+            gc.collect()
+            rss_baseline = get_rss_mb()
+            vram_local_baseline, vram_shared_baseline = get_vram_mb(adapter_luid)
+
         # [2] Generate inputs
         logger.info("Generating benchmark inputs")
         self._generate_inputs()
 
         # Compile session early so model.device is resolved for display
         self._single._session.compile()
 
+        if self.config.memory:
+            gc.collect()
+            rss_after_compile = get_rss_mb()
+            vram_local_compile, vram_shared_compile = get_vram_mb(adapter_luid)
+
         # Print model info before benchmark starts
         _print_model_info(
             self._single.io_config,
@@ -407,6 +438,30 @@ def _run_single(self) -> BenchmarkResult:
         )
         stats = self._run_benchmark()
 
+        if self.config.memory:
+            rss_after_inference = get_rss_mb()
+            vram_local_infer, vram_shared_infer = get_vram_mb(adapter_luid)
+            self._memory = {
+                "rss_baseline_mb": round(rss_baseline, 2),
+                "rss_after_compile_mb": round(rss_after_compile, 2),
+                "rss_after_inference_mb": round(rss_after_inference, 2),
+                "rss_model_load_delta_mb": round(rss_after_compile - rss_baseline, 2),
+                "rss_inference_delta_mb": round(rss_after_inference - rss_after_compile, 2),
+                "rss_total_delta_mb": round(rss_after_inference - rss_baseline, 2),
+                "vram_local_after_inference_mb": round(vram_local_infer, 2),
+                "vram_shared_after_inference_mb": round(vram_shared_infer, 2),
+                "vram_local_model_load_delta_mb": round(
+                    vram_local_compile - vram_local_baseline, 2
+                ),
+                "vram_local_inference_delta_mb": round(vram_local_infer - vram_local_compile, 2),
+                "vram_local_total_delta_mb": round(vram_local_infer - vram_local_baseline, 2),
+                "vram_shared_model_load_delta_mb": round(
+                    vram_shared_compile - vram_shared_baseline, 2
+                ),
+                "vram_shared_inference_delta_mb": round(vram_shared_infer - vram_shared_compile, 2),
+                "vram_shared_total_delta_mb": round(vram_shared_infer - vram_shared_baseline, 2),
+            }
+
         # [4] Collect results
         logger.info("Collecting results")
         return self._collect_results(stats)
@@ -474,6 +529,31 @@ def _generate_inputs(self) -> None:
             batch_size=self.config.batch_size,
         )
 
+    def _resolve_adapter_luid(self) -> str | None:
+        """Resolve adapter LUID for VRAM queries."""
+        import sys
+
+        if sys.platform != "win32":
+            return None
+
+        assert self._model is not None
+        device = self._single.device or self.config.device
+        if device == "cpu":
+            return None
+
+        try:
+            from ..sysinfo.pdh_adapters import resolve_adapter_luid
+
+            ep_name = self._single.ep_name
+            for kind in [device] if device in ("npu", "gpu") else ["npu", "gpu"]:
+                luid = resolve_adapter_luid(kind, ep_name=ep_name)
+                if luid:
+                    return luid
+            return None
+        except Exception:
+            logger.debug("Could not resolve adapter LUID", exc_info=True)
+            return None
+
     def _run_benchmark(self) -> PerfStats:
         """Execute benchmark iterations with timing."""
         if self.config.monitor:
@@ -607,6 +687,8 @@ def _collect_results(self, stats: PerfStats) -> BenchmarkResult:
             running_model_path=str(self._single.running_model_path),
             # Hardware monitor metrics (only present when --monitor is used)
             hw_monitor=getattr(self, "_hw_metrics", None),
+            # Memory profile (only present when --memory is used)
+            memory_profile=self._memory,
         )
 
 
@@ -951,7 +1033,6 @@ def display_console_report(result: BenchmarkResult, console: Console) -> None:
         console.print("[bold]Hardware (during benchmark)[/bold]")
         cpu = result.hw_monitor.get("cpu", {})
         ram = result.hw_monitor.get("ram", {})
-        dev_mem = result.hw_monitor.get("device_memory", {})
         # to_dict() emits both "npu" (always) and "gpu" (when monitoring GPU).
         # device_kind is None when CPU/RAM-only — drop the adapter line entirely
         # rather than printing zeroed "NPU: 0.0% avg".
@@ -962,15 +1043,35 @@ def display_console_report(result: BenchmarkResult, console: Console) -> None:
                 f"  {device_kind.upper()}: {adapter.get('mean_pct', 0):.1f}% avg, "
                 f"{adapter.get('peak_pct', 0):.1f}% peak  |  "
                 f"CPU: {cpu.get('mean_pct', 0):.1f}% avg  |  "
-                f"Mem: {ram.get('used_mb', 0):.0f} MB"
+                f"RAM: {ram.get('used_mb', 0):.0f} MB"
             )
+        else:
             console.print(
-                f"  Device Mem: {dev_mem.get('local_peak_mb', 0):.0f}/"
-                f"{dev_mem.get('shared_peak_mb', 0):.0f} MB (local/shared)"
+                f"  CPU: {cpu.get('mean_pct', 0):.1f}% avg  |  RAM: {ram.get('used_mb', 0):.0f} MB"
             )
-        else:
+
+    # Memory section (only when --memory is enabled)
+    if result.memory_profile:
+        mem = result.memory_profile
+        console.print()
+        console.print("[bold]Memory:[/bold]")
+        console.print(
+            f"  RAM:  {mem['rss_after_inference_mb']:.1f} MB -> "
+            f"model load: {mem['rss_model_load_delta_mb']:+.1f} MB  |  "
+            f"inference: {mem['rss_inference_delta_mb']:+.1f} MB  |  "
+            f"total: {mem['rss_total_delta_mb']:+.1f} MB"
+        )
+        vram_local = mem.get("vram_local_after_inference_mb", 0)
+        vram_shared = mem.get("vram_shared_after_inference_mb", 0)
+        if vram_local > 0 or vram_shared > 0:
             console.print(
-                f"  CPU: {cpu.get('mean_pct', 0):.1f}% avg  |  Mem: {ram.get('used_mb', 0):.0f} MB"
+                f"  VRAM: {vram_local:.1f}/{vram_shared:.1f} MB (local/shared) -> "
+                f"model load: {mem['vram_local_model_load_delta_mb']:+.1f}/"
+                f"{mem['vram_shared_model_load_delta_mb']:+.1f} MB  |  "
+                f"inference: {mem['vram_local_inference_delta_mb']:+.1f}/"
+                f"{mem['vram_shared_inference_delta_mb']:+.1f} MB  |  "
+                f"total: {mem['vram_local_total_delta_mb']:+.1f}/"
+                f"{mem['vram_shared_total_delta_mb']:+.1f} MB"
             )
 
     console.print()
@@ -1278,6 +1379,12 @@ def _run_simple_loop(
     show_default=True,
     help="Show live hardware utilization chart for the benchmarked device (NPU, GPU, or CPU)",
 )
+@click.option(
+    "--memory/--no-memory",
+    default=True,
+    show_default=True,
+    help="Measure process and device memory at each benchmark phase",
+)
 @click.option(
     "--op-tracing",
     "op_tracing",
@@ -1311,6 +1418,7 @@ def perf(
     allow_unsupported_nodes: bool,
     module_class: str | None,
     monitor: bool,
+    memory: bool,
     op_tracing: str | None,
     output_format: cli_utils.OutputFormat,
     verbose: int,
@@ -1459,6 +1567,7 @@ def perf(
         no_compile=no_compile,
         allow_unsupported_nodes=allow_unsupported_nodes,
         monitor=monitor,
+        memory=memory,
         ep=ep,
         ep_options=ep_provider_options,
         shape_config=shape_config,

src/winml/modelkit/session/monitor/live_display.py

@@ -218,7 +218,7 @@ def _render_status(
         cpu_now = cpu_samples[-1] if cpu_samples else 0.0
         ram_mb = self._hw.ram_used_mb
         cpu_cell = f"CPU: {cpu_now:.1f}%"
-        ram_cell = f"Mem: {ram_mb:.0f} MB"
+        ram_cell = f"RAM: {ram_mb:.0f} MB"
 
         if not self._show_adapter:
             return f"  {cpu_cell:<12}| {ram_cell}"
@@ -229,7 +229,7 @@ def _render_status(
         mem_shared = self._hw.peak_memory_shared_mb
 
         adapter_cell = f"{self._adapter_label}: {adapter_mean:.1f}% avg ({adapter_now:.1f}% now)"
-        mem_cell = f"Device Mem: {mem_local:.0f}/{mem_shared:.0f} MB"
+        mem_cell = f"VRAM: {mem_local:.0f}/{mem_shared:.0f} MB"
 
         return f"  {adapter_cell:<32}| {cpu_cell:<12}| {ram_cell:<16}| {mem_cell}"
 

src/winml/modelkit/session/monitor/memory_tracker.py

@@ -0,0 +1,54 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Process memory helper for perf benchmarking."""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+
+import psutil
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_rss_mb() -> float:
+    """Return current RSS in MB for this process."""
+    return psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
+
+
+def get_vram_mb(adapter_luid: str | None) -> tuple[float, float]:
+    """Return current VRAM usage as (local_mb, shared_mb) via PDH.
+
+    Returns (0.0, 0.0) on non-Windows, if no adapter_luid, or on failure.
+    """
+    if sys.platform != "win32" or not adapter_luid:
+        return 0.0, 0.0
+
+    try:
+        from ._pdh import PdhQuery
+
+        pid = os.getpid()
+        q = PdhQuery()
+        q.open()
+        q.add_counter(
+            "local",
+            rf"\GPU Process Memory(pid_{pid}_luid_{adapter_luid}_phys_0)\Local Usage",
+        )
+        q.add_counter(
+            "shared",
+            rf"\GPU Process Memory(pid_{pid}_luid_{adapter_luid}_phys_0)\Shared Usage",
+        )
+        # Memory counters are absolute (not rate-based), single collect suffices.
+        values = q.collect()
+        q.close()
+        local = (values.get("local") or 0) / (1024 * 1024)
+        shared = (values.get("shared") or 0) / (1024 * 1024)
+        return local, shared
+    except Exception:
+        logger.debug("VRAM query failed", exc_info=True)
+        return 0.0, 0.0

tests/unit/commands/test_perf_cli.py

@@ -608,6 +608,7 @@ def test_format_text_shows_console_report(
         mock_result.samples_per_sec = 100.0
         mock_result.batches_per_sec = 100.0
         mock_result.hw_monitor = None
+        mock_result.memory_profile = None
         mock_instance = MagicMock()
         mock_instance.run.return_value = mock_result
         mock_benchmark_class.return_value = mock_instance

tests/unit/session/monitor/test_memory_tracker.py

@@ -0,0 +1,25 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests for the memory_tracker module."""
+
+from __future__ import annotations
+
+from winml.modelkit.session.monitor.memory_tracker import get_rss_mb
+
+
+class TestGetRssMb:
+    """Test process RSS retrieval."""
+
+    def test_returns_positive_float(self) -> None:
+        rss = get_rss_mb()
+        assert isinstance(rss, float)
+        assert rss > 0
+
+    def test_increases_after_allocation(self) -> None:
+        before = get_rss_mb()
+        _data = [bytearray(1024 * 1024) for _ in range(10)]  # ~10 MB
+        after = get_rss_mb()
+        assert after >= before
+        assert _data is not None

tests/unit/session/test_ep_monitor.py

@@ -1516,11 +1516,11 @@ def test_cpu_only_status_omits_adapter_cell(self):
         )
         # Adapter line / device-memory line gone; CPU + Mem remain.
         assert "CPU: 12.5%" in status
-        assert "Mem: 8000 MB" in status
+        assert "RAM: 8000 MB" in status
         assert "NPU:" not in status
         assert "GPU:" not in status
         assert "Adapter:" not in status
-        assert "Device Mem:" not in status
+        assert "VRAM:" not in status
 
     def test_cpu_only_chart_legend_omits_adapter_swatch(self):
         """The chart legend should advertise only CPU when no adapter polled."""
@@ -1554,4 +1554,4 @@ def test_adapter_kind_status_keeps_adapter_cell(self):
             ram_mb=8000.0,
         )
         assert "GPU: 42.0% avg" in status
-        assert "Device Mem:" in status
+        assert "VRAM:" in status