fix(perf): support multiple QNN EP context partitions

docs/commands/perf.md

@@ -153,6 +153,10 @@ traced graph's inputs don't match the provided data — for example a compiled
 context model with different input names — the trace falls back to random inputs
 and logs a warning.
 
+Op-tracing results are included in the main benchmark JSON under
+`hw_monitor.ep_proof`. The profiling CSV remains available as the raw trace
+artifact; no separate `_op_trace.json` file is written.
+
 ## Common pitfalls
 
 - **Warm-up too low on NPU.** The first several inferences on an NPU EP can be significantly slower due to kernel compilation and caching. The default of 10 warm-up iterations is usually enough for vision models, but transformer models with many operators may need `--warmup 30` or higher to reach steady-state latency.

src/winml/modelkit/commands/perf.py

@@ -2144,17 +2144,9 @@ def _io_specs_from_config(
 def _print_save_to_footer(
     console: Console,
     *,
-    trace_json: str | None,
     profiling_csv: str | None,
 ) -> None:
-    """Print save-to footer lines after the op-trace report.
-
-    Each line is rendered only when its path is supplied; if both are
-    ``None`` the helper emits nothing. The ``[dim]...[/dim]`` markup
-    softens the label so the path itself is the visual anchor.
-    """
-    if trace_json:
-        console.print(f"[dim]Op-trace JSON:[/dim] {trace_json}")
+    """Print the raw profiling artifact path after the op-trace report."""
     if profiling_csv:
         console.print(f"[dim]Profiling CSV:[/dim] {profiling_csv}")
 
@@ -3399,7 +3391,7 @@ def perf(
         # misleading JSON artifact on disk for CI consumers.
         # =================================================================
         if op_tracing:
-            from ..session.monitor.report import display_op_trace_report, write_op_trace_json
+            from ..session.monitor.report import display_op_trace_report
 
             # Both ONNX and HF inputs run through the same PerfBenchmark
             # instance, which exposes its perf context as ``_perf_ctx``.
@@ -3464,16 +3456,9 @@ def perf(
                 display_op_trace_report(trace_result, console, top_n=top_k)
             else:
                 display_op_trace_report(trace_result, console)
-            # Write the op-trace report next to the requested benchmark output
-            # file (same directory + stem, with an ``_op_trace`` suffix) so the
-            # two artifacts stay paired instead of landing under an unrelated
-            # fixed name.
-            trace_output = output.with_name(f"{output.stem}_op_trace{output.suffix}")
-            write_op_trace_json(trace_result, trace_output)
             profiling_csv = trace_result.artifacts.get("csv")
             _print_save_to_footer(
                 console,
-                trace_json=str(trace_output),
                 profiling_csv=profiling_csv,
             )
         else:

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

@@ -29,7 +29,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 from ..._env import env_flag_enabled
-from ...onnx.epcontext import select_main_epcontext_partition_name
+from ...onnx.epcontext import epcontext_partitions, select_main_epcontext_partition_name
 from ._onnx_metadata import _load_onnx_operator_data
 from .ep_monitor import WinMLEPMonitor
 from .op_metrics import (
@@ -58,6 +58,69 @@
 }
 
 
+def _metadata_int(value: object, field: str) -> int:
+    """Normalize a QNN numeric metadata value, including float strings."""
+    try:
+        return round(float(value))  # type: ignore[arg-type]
+    except (TypeError, ValueError):
+        logger.warning(
+            "QNNMonitor: could not parse %r as a number for metadata field %r; "
+            "defaulting to 0. This may corrupt cycle_to_us and duration_us values.",
+            value,
+            field,
+        )
+        return 0
+
+
+def _coalesce_partition_samples(
+    samples: list[dict[str, Any]],
+    partition_count: int,
+) -> list[dict[str, Any]]:
+    """Combine the per-partition profiling blocks emitted for each inference."""
+    if partition_count == 1:
+        return samples
+
+    combined_samples: list[dict[str, Any]] = []
+    for offset in range(0, len(samples), partition_count):
+        partition_samples = samples[offset : offset + partition_count]
+        combined: dict[str, Any] = {
+            "metadata": {
+                "hvx_threads": max(
+                    _metadata_int(sample["metadata"]["hvx_threads"], "hvx_threads")
+                    for sample in partition_samples
+                ),
+                "accel_execute_cycles": sum(
+                    _metadata_int(
+                        sample["metadata"]["accel_execute_cycles"],
+                        "accel_execute_cycles",
+                    )
+                    for sample in partition_samples
+                ),
+                "accel_execute_us": sum(
+                    _metadata_int(
+                        sample["metadata"]["accel_execute_us"],
+                        "accel_execute_us",
+                    )
+                    for sample in partition_samples
+                ),
+            },
+            "samples": [],
+        }
+        for partition_index, sample in enumerate(partition_samples):
+            metadata = sample["metadata"]
+            for op in sample["samples"]:
+                combined["samples"].append(
+                    {
+                        **op,
+                        "_partition_index": partition_index,
+                        "_accel_execute_cycles": metadata["accel_execute_cycles"],
+                        "_accel_execute_us": metadata["accel_execute_us"],
+                    }
+                )
+        combined_samples.append(combined)
+    return combined_samples
+
+
 class QNNMonitor(WinMLEPMonitor):
     """Qualcomm NPU per-op profiler via ORT's QNN EP.
 
@@ -523,65 +586,76 @@ def _parse_artifacts(self, qhas_override: Path | None = None) -> OpTraceResult:
         samples = parsed.get("samples", [])
         if self._expected_measured_samples is not None:
             expected_total = self._warmup_samples + self._expected_measured_samples
-            if len(samples) != expected_total:
+            partition_count = self._epcontext_partition_count()
+            expected_profile_blocks = expected_total * partition_count
+            if (
+                len(samples) != expected_profile_blocks
+                and partition_count == 1
+                and expected_total > 0
+                and len(samples) % expected_total == 0
+            ):
+                partition_count = len(samples) // expected_total
+                expected_profile_blocks = len(samples)
+            if len(samples) != expected_profile_blocks:
                 raise ValueError(
                     "profiling CSV sample count mismatch: "
-                    f"expected {expected_total} total samples "
+                    f"expected {expected_profile_blocks} profiling blocks "
                     f"({self._warmup_samples} warmup + "
-                    f"{self._expected_measured_samples} measured), got {len(samples)}"
+                    f"{self._expected_measured_samples} measured) across "
+                    f"{partition_count} EPContext partition(s), got {len(samples)}"
                 )
+            samples = _coalesce_partition_samples(samples, partition_count)
             samples = samples[self._warmup_samples :]
 
         artifacts: dict[str, str] = {"csv": str(csv_path)}
 
-        # Convert cycles to microseconds via the CSV-reported ratio.
-        # Use round(float(...)) rather than int() so that a float-string
-        # value like "12345.6" (legal QNN SDK output) parses correctly
-        # instead of raising ValueError → silent op-record drop.
-        def _to_int(val: object, field: str) -> int:
-            try:
-                return round(float(val))  # type: ignore[arg-type]
-            except (TypeError, ValueError):
-                logger.warning(
-                    "QNNMonitor: could not parse %r as a number for metadata field %r; "
-                    "defaulting to 0.  This may corrupt cycle_to_us and duration_us values.",
-                    val,
-                    field,
-                )
-                return 0
-
         sample_metadata: list[dict[str, int]] = []
-        operator_samples: dict[int, dict[str, Any]] = {}
+        operator_samples: dict[tuple[int, int], dict[str, Any]] = {}
         for sample in samples:
             sample_meta = sample.get("metadata", {})
-            total_cycles = _to_int(
+            total_cycles = _metadata_int(
                 sample_meta.get("accel_execute_cycles", 0) or 0,
                 "accel_execute_cycles",
             )
-            accel_us = _to_int(
+            accel_us = _metadata_int(
                 sample_meta.get("accel_execute_us", 0) or 0,
                 "accel_execute_us",
             )
-            hvx_threads = _to_int(sample_meta.get("hvx_threads", 0) or 0, "hvx_threads")
+            hvx_threads = _metadata_int(
+                sample_meta.get("hvx_threads", 0) or 0,
+                "hvx_threads",
+            )
             sample_metadata.append(
                 {
                     "hvx_threads": hvx_threads,
                     "accel_execute_cycles": total_cycles,
                     "accel_execute_us": accel_us,
                 }
             )
-            cycle_to_us = accel_us / total_cycles if total_cycles > 0 else 0.0
             for op in sample.get("samples", []):
                 op_id = op["op_id"]
+                partition_index = op.get("_partition_index", 0)
+                key = (partition_index, op_id)
                 entry = operator_samples.setdefault(
-                    op_id,
+                    key,
                     {
                         "op_path": op["op_path"],
                         "op_id": op_id,
                         "durations_us": [],
                         "percentages": [],
                     },
                 )
+                op_total_cycles = _metadata_int(
+                    op.get("_accel_execute_cycles", total_cycles),
+                    "accel_execute_cycles",
+                )
+                op_accel_us = _metadata_int(
+                    op.get("_accel_execute_us", accel_us),
+                    "accel_execute_us",
+                )
+                cycle_to_us = (
+                    op_accel_us / op_total_cycles if op_total_cycles > 0 else 0.0
+                )
                 entry["durations_us"].append(op["cycles"] * cycle_to_us)
                 entry["percentages"].append(
                     op["cycles"] / total_cycles * 100 if total_cycles > 0 else 0.0
@@ -664,6 +738,20 @@ def _metadata_mean(field: str) -> float:
             fallback_reason=fallback_reason,
         )
 
+    def _epcontext_partition_count(self) -> int:
+        """Return the number of EPContext partitions represented in each inference."""
+        if self._running_model_path is None:
+            return 1
+        try:
+            return max(1, len(epcontext_partitions(self._running_model_path)))
+        except Exception as exc:
+            logger.info(
+                "QNNMonitor: unable to inspect EPContext partitions from %s: %s",
+                self._running_model_path,
+                exc,
+            )
+            return 1
+
     def _try_qhas(
         self,
         artifacts: dict[str, str],

tests/e2e/test_perf_e2e.py

@@ -820,8 +820,9 @@ def test_op_tracing_basic_qnn_npu(self, tmp_path: Path, npu_model_arg: str):
 
         assert result.exit_code == 0, f"perf failed (exit {result.exit_code}):\n{result.output}"
         assert output_file.exists()
-        assert trace_output.exists()
-        trace = json.loads(trace_output.read_text())
+        assert not trace_output.exists()
+        output = json.loads(output_file.read_text())
+        trace = output["hw_monitor"]["ep_proof"]
         assert trace["metadata"]["device"] == "npu"
         assert trace["metadata"]["ep"] == EP_ALIASES["qnn"]
         assert trace["metadata"]["tracing_level"] == "basic"

tests/unit/commands/test_perf_optracing.py

@@ -921,7 +921,6 @@ def _invoke_text_op_trace_failure(tmp_path: Path, trace_result):
         patch("winml.modelkit.commands.perf.display_console_report") as display_report,
         patch("winml.modelkit.commands.perf.write_json_report") as write_json,
         patch("winml.modelkit.session.monitor.report.display_op_trace_report") as display_trace,
-        patch("winml.modelkit.session.monitor.report.write_op_trace_json") as write_trace,
     ):
         result = runner.invoke(
             perf,
@@ -938,7 +937,7 @@ def _invoke_text_op_trace_failure(tmp_path: Path, trace_result):
             obj={},
         )
 
-    return result, display_report, write_json, display_trace, write_trace
+    return result, display_report, write_json, display_trace
 
 
 class TestCliOpTracingDispatch:
@@ -982,7 +981,7 @@ def test_text_mode_not_run_status_exits_before_success_reports(self, tmp_path: P
             status="not_run",
         )
 
-        result, display_report, write_json, display_trace, write_trace = (
+        result, display_report, write_json, display_trace = (
             _invoke_text_op_trace_failure(tmp_path, trace)
         )
 
@@ -991,7 +990,6 @@ def test_text_mode_not_run_status_exits_before_success_reports(self, tmp_path: P
         display_report.assert_not_called()
         write_json.assert_not_called()
         display_trace.assert_not_called()
-        write_trace.assert_not_called()
 
     def test_json_mode_missing_trace_result_does_not_emit_benchmark_json(
         self, tmp_path: Path
@@ -1143,7 +1141,6 @@ def test_basic_fallback_status_exits_0_with_notice(self, tmp_path: Path):
             patch("winml.modelkit.commands.perf.display_console_report"),
             patch("winml.modelkit.commands.perf.write_json_report"),
             patch("winml.modelkit.session.monitor.report.display_op_trace_report"),
-            patch("winml.modelkit.session.monitor.report.write_op_trace_json"),
             patch("winml.modelkit.onnx.is_compiled_onnx", return_value=True),
         ):
             result = runner.invoke(
@@ -1223,7 +1220,6 @@ def test_basic_fallback_status_rejects_raw_running_model(self, tmp_path: Path):
             patch("winml.modelkit.commands.perf.display_console_report"),
             patch("winml.modelkit.commands.perf.write_json_report"),
             patch("winml.modelkit.session.monitor.report.display_op_trace_report"),
-            patch("winml.modelkit.session.monitor.report.write_op_trace_json"),
             patch("winml.modelkit.onnx.is_compiled_onnx", return_value=False),
         ):
             result = runner.invoke(
@@ -1317,7 +1313,7 @@ def _capture_write(*args, **kwargs):
 # PRD §10.5 / coreloop §8.4 mandate this test:
 #   "test_cli_op_tracing_basic_on_qnn (skip if no QNN NPU): runs
 #    wmk perf -m resnet50 --device npu --op-tracing basic, asserts CSV
-#    produced, *_op_trace.json written, at least one operator entry."
+#    produced, op trace embedded in the perf JSON, at least one operator entry."
 #
 # This is the only end-to-end proof that SC-1 holds: the headline
 # invocation produces real per-operator trace data on a QNN NPU.
@@ -1339,7 +1335,7 @@ def test_cli_op_tracing_basic_on_qnn(tmp_path):
 
     Hardware-gated. Must produce:
       * a profiling CSV under the monitor's output directory,
-      * a ``*_op_trace.json`` next to the perf JSON output,
+      * op-trace data embedded in the perf JSON output,
       * at least one operator entry, with ``status == "ok"``.
 
     A regression that silently falls back to CPU (the bug SC-1 explicitly
@@ -1378,15 +1374,11 @@ def test_cli_op_tracing_basic_on_qnn(tmp_path):
         f"perf --op-tracing basic failed (exit {result.exit_code}):\n{result.output}"
     )
 
-    # Per-op trace JSON written next to the perf output.
-    trace_files = list(tmp_path.glob("*_op_trace.json"))
-    assert trace_files, (
-        f"Expected *_op_trace.json next to {output_path}; got: {list(tmp_path.iterdir())}"
-    )
-
     import json
 
-    trace_data = json.loads(trace_files[0].read_text(encoding="utf-8"))
+    report_data = json.loads(output_path.read_text(encoding="utf-8"))
+    trace_data = report_data["hw_monitor"]["ep_proof"]
+    assert not list(tmp_path.glob("*_op_trace.json"))
     assert trace_data["status"] == "ok", (
         f"Expected status='ok' on real hardware, got {trace_data['status']!r} "
         f"with error={trace_data.get('error')!r}"