diff --git a/python/tokenspeed/cli/__main__.py b/python/tokenspeed/cli/__main__.py index 955e76342..cefa51a87 100644 --- a/python/tokenspeed/cli/__main__.py +++ b/python/tokenspeed/cli/__main__.py @@ -42,6 +42,12 @@ def _env(args: argparse.Namespace) -> None: env_main() +def _merge_traces(args: argparse.Namespace) -> None: + from tokenspeed.cli.trace_merge import main as merge_traces_main + + merge_traces_main(args.merge_args) + + def _version(args: argparse.Namespace) -> None: from tokenspeed.version import __version__ @@ -77,6 +83,14 @@ def main() -> None: ) env_parser.set_defaults(func=_env) + merge_traces_parser = subparsers.add_parser( + "merge-traces", + add_help=False, + help="Merge a Proton chrome trace into a VizTracer report on a " + "shared timeline.", + ) + merge_traces_parser.set_defaults(func=_merge_traces, merge_args=[]) + version_parser = subparsers.add_parser( "version", help="Print the TokenSpeed version.", @@ -94,6 +108,11 @@ def main() -> None: args.func(args) return + if args.func is _merge_traces: + args.merge_args = extra_args + args.func(args) + return + if args.func is _serve: raw = list(sys.argv[2:]) args.func(args, raw) diff --git a/python/tokenspeed/cli/trace_merge.py b/python/tokenspeed/cli/trace_merge.py new file mode 100644 index 000000000..1ea2475de --- /dev/null +++ b/python/tokenspeed/cli/trace_merge.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Merge a Proton chrome trace into a VizTracer report on one timeline. + +Both profilers can run in the same scheduler process (``/start_profile`` +with ``{"activities": ["VIZTRACER", "PROTON"]}``) but write separate files +with timestamps relative to their own start. Each file also records the +absolute wall-clock time its ``ts=0`` refers to — ``viztracer_metadata. +baseTimeNanoseconds`` in VizTracer reports, top-level +``baseTimeNanoseconds`` in Proton chrome traces — so the two timelines can +be shifted onto a shared axis after the fact. + +Usage (one file pair per scheduler rank): + + tokenspeed merge-traces -TP0.viztracer.json \ + -TP0.proton.chrome_trace -o -TP0-merged.json + +Open the merged report in vizviewer or https://ui.perfetto.dev — Python +frames and Proton's kernel/scope lanes share one time axis. Alignment +accuracy is microsecond-to-millisecond grade (the profilers reconcile GPU +and CPU clocks independently); use it to correlate host activity with GPU +gaps, not for sub-microsecond attribution. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +__all__ = ["merge_proton_viztracer", "main"] + +# Proton labels its single synthetic process (pid 0) "Trace"; rename it so +# the lanes are recognizable next to the VizTracer process in a merged view. +_PROTON_PROCESS_NAME = "Proton" + + +def _load_json(path: str | Path, kind: str) -> dict[str, Any]: + try: + with open(path, encoding="utf-8") as f: + return json.load(f) + except json.JSONDecodeError as exc: + raise ValueError(f"{kind} file {path} is not valid JSON: {exc}") from exc + + +def merge_proton_viztracer( + viztracer_path: str | Path, + proton_path: str | Path, + output_path: str | Path, +) -> int: + """Merge a Proton chrome trace into a VizTracer report. + + Shifts every Proton event by the difference between the two files' + ``baseTimeNanoseconds`` anchors and appends them to the VizTracer + report's ``traceEvents``, producing one chrome-trace JSON on a shared + time axis. + + Args: + viztracer_path: VizTracer report (``.json``) saved with a viztracer + version that records ``viztracer_metadata.baseTimeNanoseconds``. + proton_path: Proton trace (``data="trace"`` session dumped as + ``chrome_trace``) with a top-level ``baseTimeNanoseconds``. + output_path: Where to write the merged chrome-trace JSON. + + Returns: + The number of Proton events merged into the report. + + Raises: + ValueError: If either input lacks its absolute time anchor or is + not valid JSON. + """ + viztracer_json = _load_json(viztracer_path, "VizTracer report") + proton_json = _load_json(proton_path, "Proton trace") + + viztracer_base_ns = viztracer_json.get("viztracer_metadata", {}).get( + "baseTimeNanoseconds" + ) + if viztracer_base_ns is None: + raise ValueError( + f"{viztracer_path} has no viztracer_metadata.baseTimeNanoseconds; " + "re-record with a viztracer version that stores the report's " + "absolute time base." + ) + + proton_base_ns = proton_json.get("baseTimeNanoseconds") + if proton_base_ns is None: + raise ValueError( + f"{proton_path} has no baseTimeNanoseconds; Upgrade " + "tokenspeed-proton and re-record." + ) + + offset_us = (proton_base_ns - viztracer_base_ns) / 1000.0 + + proton_events = proton_json.get("traceEvents", []) + for event in proton_events: + if event.get("ph") == "M": + # Metadata events carry no meaningful timestamp; drop any so + # they cannot clobber the report's process/thread names. + event.pop("ts", None) + if ( + event.get("name") == "process_name" + and event.get("args", {}).get("name") == "Trace" + ): + event["args"]["name"] = _PROTON_PROCESS_NAME + elif "ts" in event: + event["ts"] += offset_us + + viztracer_json.setdefault("traceEvents", []).extend(proton_events) + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", encoding="utf-8") as f: + json.dump(viztracer_json, f) + return len(proton_events) + + +def main(argv: list[str] | None = None) -> None: + """CLI entry point for ``tokenspeed merge-traces``. + + Args: + argv: Argument list to parse; defaults to ``sys.argv[1:]``. + """ + parser = argparse.ArgumentParser( + prog="tokenspeed merge-traces", + description="Merge a Proton chrome trace into a VizTracer report " + "on a shared timeline (one file pair per scheduler rank).", + ) + parser.add_argument("viztracer_json", help="VizTracer report (.json)") + parser.add_argument("proton_trace", help="Proton chrome trace (.chrome_trace)") + parser.add_argument( + "-o", + "--output", + default=None, + help="Output path (default: -merged.json)", + ) + args = parser.parse_args(argv) + + output = args.output + if output is None: + viztracer_path = Path(args.viztracer_json) + output = viztracer_path.with_name(f"{viztracer_path.stem}-merged.json") + + merged = merge_proton_viztracer(args.viztracer_json, args.proton_trace, output) + print(f"Merged {merged} Proton events into {output}") + + +if __name__ == "__main__": + main() diff --git a/python/tokenspeed/runtime/engine/request_handler.py b/python/tokenspeed/runtime/engine/request_handler.py index eb4f09a51..f651f3923 100644 --- a/python/tokenspeed/runtime/engine/request_handler.py +++ b/python/tokenspeed/runtime/engine/request_handler.py @@ -375,6 +375,12 @@ def init_profile( message="Proton is not available: the installed " "tokenspeed-triton does not provide a profiler.", ) + if torch.version.hip and "HIP_VISIBLE_DEVICES" in os.environ: + return ProfileReqOutput( + success=False, + message="Proton on AMD requires ROCR_VISIBLE_DEVICES; " + "unset HIP_VISIBLE_DEVICES before calling /start_profile.", + ) if ProfilingState.get().active: return ProfileReqOutput( success=False, @@ -450,7 +456,19 @@ def start_profile( self.profiler_output_dir, f"{self.profile_id}-{self.profile_rank_tag}{stage_suffix}.proton", ) - start_profiling(profile_config_from_env(output=proton_output)) + try: + start_profiling(profile_config_from_env(output=proton_output)) + except Exception as exc: + logger.exception("Failed to start Proton profiling") + if self.torch_profiler is not None: + self.torch_profiler.stop() + self.torch_profiler = None + if "MEM" in activities: + torch.cuda.memory._record_memory_history(enabled=None) + return ProfileReqOutput( + success=False, + message=f"Failed to start Proton profiling: {exc}", + ) if "VIZTRACER" in activities: Path(self.profiler_output_dir).mkdir(parents=True, exist_ok=True) @@ -511,11 +529,18 @@ def stop_profile(self, stage: ForwardMode | None = None) -> ProfileReqOutput | N if "CUDA_PROFILER" in self.profiler_activities: torch.cuda.cudart().cudaProfilerStop() + proton_error: Exception | None = None if "PROTON" in self.profiler_activities: # Finalizes the session and writes the profile now, while this # process is still alive (shutdown is SIGKILL; no atexit). - stop_profiling() - torch.distributed.barrier(self.attn_tp_cpu_group) + try: + stop_profiling() + except Exception as exc: + logger.exception("Failed to finalize Proton profiling") + proton_error = exc + finally: + # Do not reply until every TP peer has finished writing. + torch.distributed.barrier(self.attn_tp_cpu_group) if "VIZTRACER" in self.profiler_activities and self.viztracer is not None: self.viztracer.stop() @@ -531,6 +556,11 @@ def stop_profile(self, stage: ForwardMode | None = None) -> ProfileReqOutput | N self.profile_in_progress = False self.profiler_start_forward_ct = None + if proton_error is not None: + return ProfileReqOutput( + success=False, + message=f"Failed to finalize Proton profiling: {proton_error}", + ) return ProfileReqOutput(success=True, message="Succeeded.") def _profile_batch_predicate(self, forward_mode=None): diff --git a/test/cli/test_trace_merge.py b/test/cli/test_trace_merge.py new file mode 100644 index 000000000..5919261c4 --- /dev/null +++ b/test/cli/test_trace_merge.py @@ -0,0 +1,180 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Tests for :mod:`tokenspeed.cli.trace_merge` using synthetic trace files. + +The synthetic files encode the on-disk contracts this tool depends on: +VizTracer reports store their absolute time base in +``viztracer_metadata.baseTimeNanoseconds``; Proton chrome traces store it +top-level as ``baseTimeNanoseconds`` with relative-microsecond event +timestamps, metadata (``ph == "M"``) events without timestamps, and a +synthetic pid-0 process named "Trace". +""" + +from __future__ import annotations + +import json + +import pytest + +from tokenspeed.cli.trace_merge import main, merge_proton_viztracer + +VIZTRACER_BASE_NS = 1_000_000_000_000_000 +# Proton session started 2.5 ms after the viztracer clock origin. +PROTON_BASE_NS = VIZTRACER_BASE_NS + 2_500_000 +EXPECTED_OFFSET_US = 2_500.0 + + +def _viztracer_report() -> dict: + return { + "displayTimeUnit": "us", + "traceEvents": [ + { + "ph": "X", + "ts": 100.0, + "dur": 50.0, + "pid": 4242, + "tid": 4242, + "name": "forward", + "cat": "fee", + }, + ], + "viztracer_metadata": { + "version": "1.1.1", + "baseTimeNanoseconds": VIZTRACER_BASE_NS, + }, + } + + +def _proton_trace() -> dict: + return { + "displayTimeUnit": "us", + "baseTimeNanoseconds": PROTON_BASE_NS, + "traceEvents": [ + { + "ph": "M", + "pid": 0, + "tid": 0, + "name": "process_name", + "args": {"name": "Trace"}, + }, + { + "ph": "M", + "pid": 0, + "tid": 100, + "name": "thread_name", + "args": {"name": "GPU Stream 7"}, + }, + { + "ph": "X", + "ts": 10.0, + "dur": 5.0, + "pid": 0, + "tid": 100, + "name": "gemm.mm[triton_mm]", + "cat": "kernel", + }, + { + "ph": "s", + "ts": 8.0, + "pid": 0, + "tid": 1, + "id": 1, + "cat": "flow", + "name": "launch->kernel", + }, + { + "ph": "f", + "ts": 10.0, + "pid": 0, + "tid": 100, + "id": 1, + "cat": "flow", + "name": "launch->kernel", + }, + ], + } + + +def _write(tmp_path, name: str, payload: dict) -> str: + path = tmp_path / name + path.write_text(json.dumps(payload)) + return str(path) + + +def test_merge_shifts_proton_events_onto_viztracer_axis(tmp_path): + viztracer_path = _write(tmp_path, "run.viztracer.json", _viztracer_report()) + proton_path = _write(tmp_path, "run.proton.chrome_trace", _proton_trace()) + output_path = tmp_path / "merged.json" + + merged_count = merge_proton_viztracer(viztracer_path, proton_path, output_path) + assert merged_count == 5 + + merged = json.loads(output_path.read_text()) + events = merged["traceEvents"] + assert len(events) == 6 + + # VizTracer events are untouched; its metadata anchor is preserved. + assert events[0] == _viztracer_report()["traceEvents"][0] + assert merged["viztracer_metadata"]["baseTimeNanoseconds"] == VIZTRACER_BASE_NS + + by_name = {e["name"]: e for e in events[1:]} + # Proton's synthetic process is renamed for the merged view, and + # metadata events stay timestamp-free. + assert by_name["process_name"]["args"]["name"] == "Proton" + assert "ts" not in by_name["process_name"] + assert by_name["thread_name"]["args"]["name"] == "GPU Stream 7" + + # Timed events (complete and flow) are shifted by the anchor delta. + assert by_name["gemm.mm[triton_mm]"]["ts"] == 10.0 + EXPECTED_OFFSET_US + assert by_name["gemm.mm[triton_mm]"]["dur"] == 5.0 + flow_ts = sorted(e["ts"] for e in events if e.get("cat") == "flow") + assert flow_ts == [8.0 + EXPECTED_OFFSET_US, 10.0 + EXPECTED_OFFSET_US] + + +def test_merge_rejects_viztracer_report_without_anchor(tmp_path): + report = _viztracer_report() + del report["viztracer_metadata"]["baseTimeNanoseconds"] + viztracer_path = _write(tmp_path, "run.viztracer.json", report) + proton_path = _write(tmp_path, "run.proton.chrome_trace", _proton_trace()) + + with pytest.raises(ValueError, match="viztracer_metadata.baseTimeNanoseconds"): + merge_proton_viztracer(viztracer_path, proton_path, tmp_path / "out.json") + + +def test_merge_rejects_proton_trace_without_anchor(tmp_path): + trace = _proton_trace() + del trace["baseTimeNanoseconds"] + viztracer_path = _write(tmp_path, "run.viztracer.json", _viztracer_report()) + proton_path = _write(tmp_path, "run.proton.chrome_trace", trace) + + with pytest.raises(ValueError, match="baseTimeNanoseconds"): + merge_proton_viztracer(viztracer_path, proton_path, tmp_path / "out.json") + + +def test_cli_defaults_output_next_to_viztracer_report(tmp_path, capsys): + viztracer_path = _write(tmp_path, "run.viztracer.json", _viztracer_report()) + proton_path = _write(tmp_path, "run.proton.chrome_trace", _proton_trace()) + + main([viztracer_path, proton_path]) + + default_output = tmp_path / "run.viztracer-merged.json" + assert default_output.exists() + assert "Merged 5 Proton events" in capsys.readouterr().out diff --git a/test/runtime/test_request_handler_profile.py b/test/runtime/test_request_handler_profile.py index 6f05181d2..28ee629ac 100644 --- a/test/runtime/test_request_handler_profile.py +++ b/test/runtime/test_request_handler_profile.py @@ -126,6 +126,36 @@ def test_init_fails_when_env_bootstrap_session_active(self): self.assertIn("already active", result.message) self.assertFalse(self.handler.profile_in_progress) + def test_init_rejects_hip_visible_devices_for_amd_proton(self): + with ( + mock.patch.object( + request_handler_mod, "proton_available", return_value=True + ), + mock.patch.object(request_handler_mod.torch.version, "hip", "7.2"), + mock.patch.dict( + request_handler_mod.os.environ, {"HIP_VISIBLE_DEVICES": "0"} + ), + ): + result = self.handler.profile(_start_req(self.output_dir)) + + self.assertFalse(result.success) + self.assertIn("ROCR_VISIBLE_DEVICES", result.message) + self.assertFalse(self.handler.profile_in_progress) + + def test_start_returns_failure_when_proton_cannot_initialize(self): + with mock.patch.object( + request_handler_mod, "proton_available", return_value=True + ), mock.patch.object( + request_handler_mod, + "start_profiling", + side_effect=RuntimeError("rocprofiler unavailable"), + ): + result = self.handler.profile(_start_req(self.output_dir)) + + self.assertFalse(result.success) + self.assertIn("Failed to start Proton profiling", result.message) + self.assertFalse(self.handler.profile_in_progress) + def test_start_and_stop_drive_proton_session_per_rank(self): self.handler = _make_handler(_attn_mapping(tp_rank=3)) with mock.patch.object( @@ -187,7 +217,7 @@ def test_proton_outputs_do_not_collide_across_dp_ranks(self): def test_stop_profile_barriers_tp_peers_after_proton_finalize(self): # Only attn-TP rank 0 replies to /stop_profile; the reply must wait - # until every TP peer has finalized its proton file. + # until every TP peer has finalized its Proton file. self.handler.attn_tp_cpu_group = object() with mock.patch.object( @@ -204,6 +234,24 @@ def test_stop_profile_barriers_tp_peers_after_proton_finalize(self): self.assertTrue(result.success) self.barrier.assert_called_once_with(self.handler.attn_tp_cpu_group) + def test_stop_profile_reports_proton_finalize_failure(self): + self.handler.attn_tp_cpu_group = object() + + with mock.patch.object( + request_handler_mod, "proton_available", return_value=True + ), mock.patch.object(request_handler_mod, "start_profiling"), mock.patch.object( + request_handler_mod, + "stop_profiling", + side_effect=RuntimeError("finalize failed"), + ): + self.handler.profile(_start_req(self.output_dir)) + result = self.handler.profile(ProfileReq(type=ProfileReqType.STOP_PROFILE)) + + self.assertFalse(result.success) + self.assertIn("Failed to finalize Proton profiling", result.message) + self.assertFalse(self.handler.profile_in_progress) + self.barrier.assert_called_once_with(self.handler.attn_tp_cpu_group) + def test_num_steps_window_finalizes_proton(self): with mock.patch.object( request_handler_mod, "proton_available", return_value=True diff --git a/tokenspeed-kernel/README.md b/tokenspeed-kernel/README.md index 79939aac4..e58e3362b 100644 --- a/tokenspeed-kernel/README.md +++ b/tokenspeed-kernel/README.md @@ -126,7 +126,11 @@ iteration. `/stop_profile`, writing `/[-DP][-CP]-TP.proton.` per rank. `PROTON` composes only with host-side activities (`CPU`, `MEM`, - `VIZTRACER`). + `VIZTRACER`). To see Python activity and Proton's kernel lanes on one + Perfetto timeline, profile with `VIZTRACER` + `PROTON` + (`TOKENSPEED_KERNEL_PROFILE_DATA=trace`, + `TOKENSPEED_KERNEL_PROFILE_OUTPUT_FORMAT=chrome_trace`), then merge the + traces with `tokenspeed merge-traces`. ### Plugins diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py b/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py index 2a9028fbc..4053f0e18 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py @@ -28,6 +28,7 @@ import warnings from contextlib import contextmanager from dataclasses import asdict, dataclass +from numbers import Real from pathlib import Path from typing import Any @@ -189,6 +190,19 @@ def __exit__(self, *args: object) -> None: _BOOTSTRAPPED = False +def _proton_metrics(metrics: dict[str, object]) -> dict[str, object]: + """Keep only the metric types accepted by Proton's Python API.""" + supported: dict[str, object] = {} + for name, value in metrics.items(): + if hasattr(value, "data_ptr") or isinstance(value, Real): + supported[name] = value + elif isinstance(value, (list, tuple)) and all( + isinstance(element, Real) for element in value + ): + supported[name] = list(value) + return supported + + def _is_truthy(value: str | None) -> bool: if value is None: return False @@ -261,10 +275,13 @@ def stop_profiling() -> None: return output_format = state._config.output_format if state._config is not None else "" - proton.finalize(state._session, output_format) - state._session = None - state._config = None - state.enabled = False + try: + proton.finalize(state._session, output_format) + finally: + # Keep wrapper state recoverable even when report serialization fails. + state._session = None + state._config = None + state.enabled = False def start_shape_capture() -> None: @@ -306,13 +323,27 @@ def kernel_scope( kernel_name: str = "", **metrics: object, ): + """Return a Proton scope for one kernel launch. + + Proton accepts numeric or tensor metrics only, so ``dtype`` is retained by + :class:`ShapeCapture` but is not emitted as a Proton scope metric. + + Args: + family: Kernel family name. + mode: Kernel operation name. + dtype: Kernel data type recorded by shape capture. + kernel_name: Selected kernel implementation name. + **metrics: Numeric kernel-shape metrics for Proton. + + Returns: + A Proton scope when profiling is active, otherwise a no-op scope. + """ state = ProfilingState.get() if not state.active: return _NOOP_SCOPE name = f"{family}.{mode}[{kernel_name}]" if kernel_name else f"{family}.{mode}" - scope_metrics = {"dtype": str(dtype)} - scope_metrics.update(metrics) + scope_metrics = _proton_metrics(metrics) return proton.scope(name, metrics=scope_metrics) diff --git a/tokenspeed-kernel/test/test_profiling.py b/tokenspeed-kernel/test/test_profiling.py index a23737ad8..7e00e3c2f 100644 --- a/tokenspeed-kernel/test/test_profiling.py +++ b/tokenspeed-kernel/test/test_profiling.py @@ -112,6 +112,29 @@ def test_start_and_stop_calls_proton(monkeypatch): assert fake.finalize_calls == [((123, "chrome_trace"), {})] +def test_stop_clears_state_when_proton_finalize_fails(monkeypatch): + fake = _FakeProton() + original_finalize = fake.finalize + + def fail_finalize(*args, **kwargs): + _ = args, kwargs + raise RuntimeError("write failed") + + monkeypatch.setattr(profiling, "_HAS_PROTON", True) + monkeypatch.setattr(profiling, "proton", fake) + monkeypatch.setattr(fake, "finalize", fail_finalize) + + profiling.start_profiling() + with pytest.raises(RuntimeError, match="write failed"): + profiling.stop_profiling() + + assert not profiling.ProfilingState.get().active + + monkeypatch.setattr(fake, "finalize", original_finalize) + assert profiling.start_profiling() == 123 + profiling.stop_profiling() + + def test_start_warns_when_proton_missing(monkeypatch): monkeypatch.setattr(profiling, "_HAS_PROTON", False) monkeypatch.setattr(profiling, "proton", None) @@ -159,7 +182,6 @@ def test_kernel_scope_uses_proton_scope_when_active(monkeypatch): ( "gemm.mm[triton_mm_fp8_scaled]", { - "dtype": "torch.float16", "M": 32, "N": 64, "K": 128, @@ -169,6 +191,36 @@ def test_kernel_scope_uses_proton_scope_when_active(monkeypatch): assert fake.scope_log == ["enter", "exit"] +def test_kernel_scope_filters_unsupported_proton_metrics(monkeypatch): + fake = _FakeProton() + monkeypatch.setattr(profiling, "_HAS_PROTON", True) + monkeypatch.setattr(profiling, "proton", fake) + profiling.start_profiling() + + with profiling.kernel_scope( + "quantization", + "mxfp4", + torch.float16, + kernel_name="triton", + shape=(2, 3), + scale_size=32, + scale_layout="ue8m0", + has_global_scale=False, + ): + pass + + assert fake.scope_calls == [ + ( + "quantization.mxfp4[triton]", + { + "shape": [2, 3], + "scale_size": 32, + "has_global_scale": False, + }, + ) + ] + + def test_bootstrap_reads_env_and_only_runs_once(monkeypatch): fake = _FakeProton() registrations: list[object] = []