From d77cbbc524ad885a2133cb2c1444996a8ac3616e Mon Sep 17 00:00:00 2001 From: Lei Zhang Date: Sun, 12 Jul 2026 12:40:18 -0700 Subject: [PATCH 1/5] feat(profile): drive Proton via the /start_profile control plane Proton could not profile serving workloads: GPU kernels run in scheduler subprocesses (mp.Process spawned in _launch_subprocesses), and those are torn down with SIGKILL on both crash and clean engine shutdown. The import-time TOKENSPEED_KERNEL_PROFILE bootstrap finalizes via atexit, which never runs under SIGKILL, so no profile was ever written; all ranks also shared one default output path. Fix by running Proton where the kernels launch, controlled through the existing profile control plane: * request_handler: new PROTON activity for /start_profile and /stop_profile. Each scheduler rank starts its own Proton session writing /-TP-[-STAGE].proton. and finalizes it on /stop_profile while the process is alive, so nothing depends on atexit or graceful exit. The existing start_step / num_steps / profile_by_stage step windows apply to Proton unchanged. * init_profile validation: reject PROTON combined with GPU or CUDA_PROFILER (CUPTI/roctracer allows one GPU profiling client per process; host-side CPU/MEM/VIZTRACER still compose), and fail cleanly when Proton is unavailable or an env-bootstrap session already owns the process. * tokenspeed_kernel.profiling: expose proton_available() and profile_config_from_env(output=...) so the control plane reuses the TOKENSPEED_KERNEL_PROFILE_* env config (data/backend/mode/ hook) with a per-rank output override. Env bootstrap behavior is unchanged and remains the right tool for single-process scripts. * bench serve: new --profile-activities flag (e.g. PROTON) sent in the /start_profile body. The profile request now uses a dedicated body dict, so profiler fields like num_steps no longer leak into every generation request payload. * Document the serving flow in the tokenspeed-kernel README; add unit tests for the new helpers and the scheduler-side PROTON dispatch, activity-conflict rejection, and step-window finalize. Signed-off-by: Lei Zhang --- python/tokenspeed/bench.py | 25 +++- .../runtime/engine/request_handler.py | 51 ++++++- test/runtime/test_request_handler_profile.py | 137 ++++++++++++++++++ tokenspeed-kernel/README.md | 12 ++ .../python/tokenspeed_kernel/profiling.py | 31 +++- tokenspeed-kernel/test/test_profiling.py | 40 +++++ 6 files changed, 288 insertions(+), 8 deletions(-) create mode 100644 test/runtime/test_request_handler_profile.py diff --git a/python/tokenspeed/bench.py b/python/tokenspeed/bench.py index c9a61f3ec..622e8e239 100755 --- a/python/tokenspeed/bench.py +++ b/python/tokenspeed/bench.py @@ -1406,6 +1406,7 @@ async def benchmark( num_warmups: int, profile: bool, profile_num_steps: int | None, + profile_activities: list[str] | None, selected_percentile_metrics: list[str], selected_percentiles: list[float], ignore_eos: bool, @@ -1485,9 +1486,13 @@ async def warmup_limited_request_func(): print("Starting profiler...") else: print(f"Starting profiler for {profile_num_steps} steps...") - extra_body = dict(extra_body or {}) + # Use a dedicated body for the /start_profile request so profiler + # fields don't leak into the generation request payloads below. + profile_body = dict(extra_body or {}) if profile_num_steps is not None: - extra_body["num_steps"] = profile_num_steps + profile_body["num_steps"] = profile_num_steps + if profile_activities is not None: + profile_body["activities"] = profile_activities profile_input = RequestFuncInput( model=model_id, model_name=model_name, @@ -1498,7 +1503,7 @@ async def warmup_limited_request_func(): logprobs=logprobs, ignore_eos=ignore_eos, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=profile_body, ) profile_output = await request_func(profile_input, session=session) if profile_output.success: @@ -1796,6 +1801,17 @@ def add_serving_cli_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--disable-tqdm", action="store_true") parser.add_argument("--profile", action="store_true") parser.add_argument("--profile-num-steps", type=int, default=None) + parser.add_argument( + "--profile-activities", + nargs="+", + choices=["CPU", "GPU", "MEM", "CUDA_PROFILER", "VIZTRACER", "PROTON"], + default=None, + help="Profiler activities for /start_profile (default: server-side " + "default, CPU and GPU via the torch profiler). PROTON drives the " + "Triton Proton profiler inside each scheduler process; it cannot be " + "combined with GPU or CUDA_PROFILER, which need the same " + "CUPTI/roctracer interface.", + ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--ignore-eos", action="store_true") parser.add_argument("--disable-ignore-eos", action="store_true") @@ -1835,6 +1851,8 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: raise ValueError("--profile-num-steps must be positive.") if not args.profile: raise ValueError("--profile-num-steps requires --profile.") + if args.profile_activities is not None and not args.profile: + raise ValueError("--profile-activities requires --profile.") if args.input_len is not None: args.random_input_len = args.input_len if args.output_len is not None: @@ -1920,6 +1938,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: num_warmups=args.num_warmups, profile=args.profile, profile_num_steps=args.profile_num_steps, + profile_activities=args.profile_activities, selected_percentile_metrics=percentile_metrics.split(","), selected_percentiles=[float(p) for p in args.metric_percentiles.split(",")], ignore_eos=args.ignore_eos, diff --git a/python/tokenspeed/runtime/engine/request_handler.py b/python/tokenspeed/runtime/engine/request_handler.py index ae2515c34..76f1ad730 100644 --- a/python/tokenspeed/runtime/engine/request_handler.py +++ b/python/tokenspeed/runtime/engine/request_handler.py @@ -27,6 +27,13 @@ import torch import zmq +from tokenspeed_kernel.profiling import ( + ProfilingState, + profile_config_from_env, + proton_available, + start_profiling, + stop_profiling, +) from viztracer import VizTracer from tokenspeed.runtime.distributed.process_group_manager import ( @@ -266,8 +273,11 @@ def handle_generate_request( ) # ------------------------------------------------------------------ - # Profiling: torch / cuda / viztracer / mem-snapshot, driven by - # /start_profile and /stop_profile control requests. + # Profiling: torch / cuda / viztracer / mem-snapshot / proton, driven + # by /start_profile and /stop_profile control requests. Proton must be + # driven from this process (not the frontend): its GPU hooks are + # per-process and the scheduler subprocess is torn down with SIGKILL, + # so an atexit-based finalize would never write the profile. # ------------------------------------------------------------------ def init_profiler(self): @@ -309,6 +319,29 @@ def init_profile( if activities is None: activities = ["CPU", "GPU"] + if "PROTON" in activities: + conflicting = sorted({"GPU", "CUDA_PROFILER"} & set(activities)) + if conflicting: + return ProfileReqOutput( + success=False, + message="PROTON cannot be combined with " + f"{', '.join(conflicting)}: CUPTI/roctracer supports only " + "one GPU profiling client per process.", + ) + if not proton_available(): + return ProfileReqOutput( + success=False, + message="Proton is not available: the installed " + "tokenspeed-triton does not provide a profiler.", + ) + if ProfilingState.get().active: + return ProfileReqOutput( + success=False, + message="A Proton session is already active in this " + "process (e.g. via TOKENSPEED_KERNEL_PROFILE); it cannot " + "be controlled through /start_profile.", + ) + self.profiler_output_dir = output_dir self.torch_profiler_with_stack = with_stack self.torch_profiler_record_shapes = record_shapes @@ -368,6 +401,15 @@ def start_profile( if "CUDA_PROFILER" in activities: torch.cuda.cudart().cudaProfilerStart() + if "PROTON" in activities: + Path(self.profiler_output_dir).mkdir(parents=True, exist_ok=True) + # Proton appends the output format extension (e.g. ".hatchet"). + proton_output = os.path.join( + self.profiler_output_dir, + f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.proton", + ) + start_profiling(profile_config_from_env(output=proton_output)) + if "VIZTRACER" in activities: Path(self.profiler_output_dir).mkdir(parents=True, exist_ok=True) self.viztracer = VizTracer( @@ -427,6 +469,11 @@ def stop_profile(self, stage: ForwardMode | None = None) -> ProfileReqOutput | N if "CUDA_PROFILER" in self.profiler_activities: torch.cuda.cudart().cudaProfilerStop() + 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() + if "VIZTRACER" in self.profiler_activities and self.viztracer is not None: self.viztracer.stop() self.viztracer.save() diff --git a/test/runtime/test_request_handler_profile.py b/test/runtime/test_request_handler_profile.py new file mode 100644 index 000000000..548795092 --- /dev/null +++ b/test/runtime/test_request_handler_profile.py @@ -0,0 +1,137 @@ +import tempfile +import unittest +from unittest import mock + +from tokenspeed_kernel import profiling + +from tokenspeed.runtime.engine import request_handler as request_handler_mod +from tokenspeed.runtime.engine.io_struct import ProfileReq, ProfileReqType +from tokenspeed.runtime.engine.request_handler import RequestHandler + + +def _make_handler() -> RequestHandler: + handler = RequestHandler.__new__(RequestHandler) + handler.forward_ct = 0 + handler.attn_tp_rank = 0 + handler.init_profiler() + return handler + + +def _start_req(output_dir: str, **kwargs) -> ProfileReq: + return ProfileReq( + type=ProfileReqType.START_PROFILE, + output_dir=output_dir, + activities=["PROTON"], + profile_id="test-profile", + **kwargs, + ) + + +class TestRequestHandlerProtonProfile(unittest.TestCase): + def setUp(self): + self.handler = _make_handler() + self.output_dir = tempfile.mkdtemp() + profiling.ProfilingState.reset() + self.addCleanup(profiling.ProfilingState.reset) + + def test_init_fails_when_proton_unavailable(self): + with mock.patch.object( + request_handler_mod, "proton_available", return_value=False + ): + result = self.handler.profile(_start_req(self.output_dir)) + + self.assertFalse(result.success) + self.assertIn("Proton is not available", result.message) + self.assertFalse(self.handler.profile_in_progress) + + def test_init_rejects_proton_mixed_with_gpu_profilers(self): + for conflicting in (["GPU"], ["CUDA_PROFILER"], ["GPU", "CUDA_PROFILER"]): + req = _start_req(self.output_dir) + req.activities = ["CPU", "PROTON", *conflicting] + + result = self.handler.profile(req) + + self.assertFalse(result.success) + for activity in conflicting: + self.assertIn(activity, result.message) + self.assertFalse(self.handler.profile_in_progress) + + def test_init_allows_proton_with_host_side_profilers(self): + req = _start_req(self.output_dir) + req.activities = ["CPU", "MEM", "VIZTRACER", "PROTON"] + + with mock.patch.object( + request_handler_mod, "proton_available", return_value=True + ): + result = self.handler.init_profile( + output_dir=self.output_dir, + start_step=None, + num_steps=None, + activities=req.activities, + with_stack=None, + record_shapes=None, + profile_by_stage=False, + profile_id="test-profile", + ) + + self.assertTrue(result.success) + + def test_init_fails_when_env_bootstrap_session_active(self): + state = profiling.ProfilingState.get() + state.enabled = True + state._session = 1 + + with mock.patch.object( + request_handler_mod, "proton_available", return_value=True + ): + result = self.handler.profile(_start_req(self.output_dir)) + + self.assertFalse(result.success) + self.assertIn("already active", result.message) + self.assertFalse(self.handler.profile_in_progress) + + def test_start_and_stop_drive_proton_session_per_rank(self): + self.handler.attn_tp_rank = 3 + with mock.patch.object( + request_handler_mod, "proton_available", return_value=True + ), mock.patch.object( + request_handler_mod, "start_profiling" + ) as start_profiling, mock.patch.object( + request_handler_mod, "stop_profiling" + ) as stop_profiling: + result = self.handler.profile(_start_req(self.output_dir)) + self.assertTrue(result.success) + self.assertTrue(self.handler.profile_in_progress) + + config = start_profiling.call_args[0][0] + self.assertEqual(config.output.rsplit("/", 1)[0], self.output_dir) + self.assertTrue(config.output.endswith("test-profile-TP-3.proton")) + stop_profiling.assert_not_called() + + result = self.handler.profile(ProfileReq(type=ProfileReqType.STOP_PROFILE)) + self.assertTrue(result.success) + stop_profiling.assert_called_once() + self.assertFalse(self.handler.profile_in_progress) + + def test_num_steps_window_finalizes_proton(self): + 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" + ) as stop_profiling: + result = self.handler.profile(_start_req(self.output_dir, num_steps=2)) + self.assertTrue(result.success) + self.assertTrue(self.handler.profile_in_progress) + + self.handler.forward_ct = 1 + self.handler._profile_batch_predicate() + stop_profiling.assert_not_called() + + self.handler.forward_ct = 2 + self.handler._profile_batch_predicate() + stop_profiling.assert_called_once() + self.assertFalse(self.handler.profile_in_progress) + + +if __name__ == "__main__": + unittest.main() diff --git a/tokenspeed-kernel/README.md b/tokenspeed-kernel/README.md index 3aaaa5b14..d9926bca5 100644 --- a/tokenspeed-kernel/README.md +++ b/tokenspeed-kernel/README.md @@ -119,6 +119,18 @@ iteration. (FLOPs / bytes) per op family, tabular reports, and Proton integration. - Runtime shape capture feeds replay and tuning workflows; `kernel_scope` scopes are visible in Proton/Chrome traces. +- End-to-end serving: pass `--profile --profile-activities PROTON` to + `tokenspeed bench serve` (or POST `/start_profile` / `/stop_profile` with + `{"activities": ["PROTON"]}`). Each scheduler process — the process where + kernels actually launch — runs its own Proton session and finalizes it on + `/stop_profile`, writing `/-TP-.proton.` + per rank. `TOKENSPEED_KERNEL_PROFILE_DATA/BACKEND/MODE/...` still select + the Proton data/backend/mode. `PROTON` composes with host-side activities + (`CPU`, `MEM`, `VIZTRACER`) but is rejected alongside `GPU` or + `CUDA_PROFILER` — CUPTI/roctracer allows one GPU profiling client per + process. The import-time `TOKENSPEED_KERNEL_PROFILE=1` + bootstrap only suits single-process scripts: serving schedulers are torn + down with SIGKILL, so its atexit finalize never writes a profile there. ### Plugins diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py b/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py index ecd17a41a..2a9028fbc 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/profiling.py @@ -42,7 +42,9 @@ "ShapeCapture", "bootstrap_profiling_from_env", "kernel_scope", + "profile_config_from_env", "profiling", + "proton_available", "shape_capture", "start_shape_capture", "start_profiling", @@ -193,9 +195,32 @@ def _is_truthy(value: str | None) -> bool: return value.strip().lower() in {"1", "true", "yes", "on"} -def _profile_config_from_env() -> ProfilingConfig: +def proton_available() -> bool: + """Report whether the Proton profiler can be used in this process. + + Returns: + True if the vendored Triton distribution provides + ``tokenspeed_triton.profiler``; False otherwise (profiling calls + become no-ops with a warning). + """ + return _HAS_PROTON + + +def profile_config_from_env(output: str | None = None) -> ProfilingConfig: + """Build a :class:`ProfilingConfig` from ``TOKENSPEED_KERNEL_PROFILE_*``. + + Args: + output: Optional Proton output prefix/path. When provided it takes + precedence over ``TOKENSPEED_KERNEL_PROFILE_OUTPUT``; callers that + profile multiple processes (e.g. one scheduler per rank) use this + to give each process a distinct output file. + + Returns: + A :class:`ProfilingConfig` with ``data``/``backend``/``mode``/``hook``/ + ``output_format`` sourced from the environment. + """ return ProfilingConfig( - output=os.environ.get(_ENV_PROFILE_OUTPUT, "profile"), + output=output or os.environ.get(_ENV_PROFILE_OUTPUT, "profile"), data=os.environ.get(_ENV_PROFILE_DATA, "tree"), backend=os.environ.get(_ENV_PROFILE_BACKEND), mode=os.environ.get(_ENV_PROFILE_MODE), @@ -316,7 +341,7 @@ def bootstrap_profiling_from_env() -> None: _BOOTSTRAPPED = True if _is_truthy(os.environ.get(_ENV_PROFILE)): - start_profiling(_profile_config_from_env()) + start_profiling(profile_config_from_env()) if _is_truthy(os.environ.get(_ENV_CAPTURE_SHAPES)): start_shape_capture() diff --git a/tokenspeed-kernel/test/test_profiling.py b/tokenspeed-kernel/test/test_profiling.py index 64bb29b51..a23737ad8 100644 --- a/tokenspeed-kernel/test/test_profiling.py +++ b/tokenspeed-kernel/test/test_profiling.py @@ -199,6 +199,46 @@ def test_bootstrap_reads_env_and_only_runs_once(monkeypatch): assert len(registrations) == 2 +def test_proton_available_reflects_import(monkeypatch): + monkeypatch.setattr(profiling, "_HAS_PROTON", True) + assert profiling.proton_available() + + monkeypatch.setattr(profiling, "_HAS_PROTON", False) + assert not profiling.proton_available() + + +def test_profile_config_from_env_defaults(monkeypatch): + for env in ( + "TOKENSPEED_KERNEL_PROFILE_OUTPUT", + "TOKENSPEED_KERNEL_PROFILE_DATA", + "TOKENSPEED_KERNEL_PROFILE_BACKEND", + "TOKENSPEED_KERNEL_PROFILE_MODE", + "TOKENSPEED_KERNEL_PROFILE_HOOK", + "TOKENSPEED_KERNEL_PROFILE_OUTPUT_FORMAT", + ): + monkeypatch.delenv(env, raising=False) + + cfg = profiling.profile_config_from_env() + assert cfg == profiling.ProfilingConfig() + + +def test_profile_config_from_env_output_override_wins(monkeypatch): + monkeypatch.setenv("TOKENSPEED_KERNEL_PROFILE_OUTPUT", "env_profile") + monkeypatch.setenv("TOKENSPEED_KERNEL_PROFILE_DATA", "trace") + monkeypatch.setenv("TOKENSPEED_KERNEL_PROFILE_BACKEND", "roctracer") + monkeypatch.setenv("TOKENSPEED_KERNEL_PROFILE_MODE", "periodic_flushing") + monkeypatch.setenv("TOKENSPEED_KERNEL_PROFILE_OUTPUT_FORMAT", "chrome_trace") + + cfg = profiling.profile_config_from_env(output="rank0/step5.proton") + + assert cfg.output == "rank0/step5.proton" + assert cfg.data == "trace" + assert cfg.backend == "roctracer" + assert cfg.mode == "periodic_flushing" + assert cfg.hook == "triton" + assert cfg.output_format == "chrome_trace" + + def test_shape_capture_records_dump_and_clear(tmp_path): output = tmp_path / "shapes.json" From 4ae5bb0f239d8cc1bb25632caf53e2f4cb769d28 Mon Sep 17 00:00:00 2001 From: Lei Zhang Date: Sun, 12 Jul 2026 21:26:23 -0700 Subject: [PATCH 2/5] Simplify README.md Signed-off-by: Lei Zhang --- tokenspeed-kernel/README.md | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tokenspeed-kernel/README.md b/tokenspeed-kernel/README.md index d9926bca5..59f9fa393 100644 --- a/tokenspeed-kernel/README.md +++ b/tokenspeed-kernel/README.md @@ -124,13 +124,8 @@ iteration. `{"activities": ["PROTON"]}`). Each scheduler process — the process where kernels actually launch — runs its own Proton session and finalizes it on `/stop_profile`, writing `/-TP-.proton.` - per rank. `TOKENSPEED_KERNEL_PROFILE_DATA/BACKEND/MODE/...` still select - the Proton data/backend/mode. `PROTON` composes with host-side activities - (`CPU`, `MEM`, `VIZTRACER`) but is rejected alongside `GPU` or - `CUDA_PROFILER` — CUPTI/roctracer allows one GPU profiling client per - process. The import-time `TOKENSPEED_KERNEL_PROFILE=1` - bootstrap only suits single-process scripts: serving schedulers are torn - down with SIGKILL, so its atexit finalize never writes a profile there. + per rank. `PROTON` composes only with host-side activities (`CPU`, `MEM`, + `VIZTRACER`). ### Plugins From 25282a3a88139d724f4c989b01ce3fd8ac810ee3 Mon Sep 17 00:00:00 2001 From: Lei Zhang Date: Sun, 12 Jul 2026 21:34:37 -0700 Subject: [PATCH 3/5] fix(profile): make per-rank profile file names unique under DP/CP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit attn_tp_rank is computed modulo the attention TP size, and the DP controller launches scheduler processes as dp_rank * attn_tp_size + attn_tp_rank, so with attention DP/CP > 1 several processes share the same TP rank. Their profile outputs under a shared output_dir used the same -TP- path and would race and overwrite each other. This affected the new Proton output as well as the pre-existing torch trace, memory snapshot, and viztracer file names. Build one process-unique rank tag per scheduler at init — TP, prefixed with DP and/or CP when those dimensions exist (e.g. DP1-CP0-TP2) — and use it for all four profiler outputs. Add unit tests for the tag shape and for two DP peers with the same TP rank producing distinct Proton output paths. Signed-off-by: Lei Zhang --- .../runtime/engine/request_handler.py | 20 +++++-- test/runtime/test_request_handler_profile.py | 59 +++++++++++++++++-- tokenspeed-kernel/README.md | 3 +- 3 files changed, 73 insertions(+), 9 deletions(-) diff --git a/python/tokenspeed/runtime/engine/request_handler.py b/python/tokenspeed/runtime/engine/request_handler.py index 76f1ad730..097924689 100644 --- a/python/tokenspeed/runtime/engine/request_handler.py +++ b/python/tokenspeed/runtime/engine/request_handler.py @@ -78,6 +78,17 @@ logger = logging.getLogger(__name__) +def _profile_rank_tag(attn_mapping) -> str: + """File-name tag identifying this scheduler process's profile outputs.""" + parts = [] + if attn_mapping.has_dp: + parts.append(f"DP{attn_mapping.dp_rank}") + if attn_mapping.has_cp: + parts.append(f"CP{attn_mapping.cp_rank}") + parts.append(f"TP{attn_mapping.tp_rank}") + return "-".join(parts) + + class RequestHandler: """ 1. Recv Reqs from ZMQ @@ -113,6 +124,7 @@ def __init__( "gloo", mapping.attn.tp_group ) self.attn_tp_src_rank = mapping.attn.tp_group[0] + self.profile_rank_tag = _profile_rank_tag(mapping.attn) self.hf_eos_token_id = hf_eos_token_id self.max_req_len = max_req_len @@ -406,7 +418,7 @@ def start_profile( # Proton appends the output format extension (e.g. ".hatchet"). proton_output = os.path.join( self.profiler_output_dir, - f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.proton", + f"{self.profile_id}-{self.profile_rank_tag}{stage_suffix}.proton", ) start_profiling(profile_config_from_env(output=proton_output)) @@ -415,7 +427,7 @@ def start_profile( self.viztracer = VizTracer( output_file=os.path.join( self.profiler_output_dir, - f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.viztracer.json", + f"{self.profile_id}-{self.profile_rank_tag}{stage_suffix}.viztracer.json", ), min_duration=int( os.environ.get("TOKENSPEED_VIZTRACER_MIN_DURATION_US", "100") @@ -453,7 +465,7 @@ def stop_profile(self, stage: ForwardMode | None = None) -> ProfileReqOutput | N self.torch_profiler.export_chrome_trace( os.path.join( self.profiler_output_dir, - f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.trace.json.gz", + f"{self.profile_id}-{self.profile_rank_tag}{stage_suffix}.trace.json.gz", ) ) torch.distributed.barrier(self.attn_tp_cpu_group) @@ -461,7 +473,7 @@ def stop_profile(self, stage: ForwardMode | None = None) -> ProfileReqOutput | N if self.profiler_activities is not None and "MEM" in self.profiler_activities: memory_profile_path = os.path.join( self.profiler_output_dir, - f"{self.profile_id}-TP-{self.attn_tp_rank}-memory{stage_suffix}.pickle", + f"{self.profile_id}-{self.profile_rank_tag}-memory{stage_suffix}.pickle", ) torch.cuda.memory._dump_snapshot(memory_profile_path) torch.cuda.memory._record_memory_history(enabled=None) diff --git a/test/runtime/test_request_handler_profile.py b/test/runtime/test_request_handler_profile.py index 548795092..aa7631df7 100644 --- a/test/runtime/test_request_handler_profile.py +++ b/test/runtime/test_request_handler_profile.py @@ -1,5 +1,6 @@ import tempfile import unittest +from types import SimpleNamespace from unittest import mock from tokenspeed_kernel import profiling @@ -9,10 +10,24 @@ from tokenspeed.runtime.engine.request_handler import RequestHandler -def _make_handler() -> RequestHandler: +def _attn_mapping( + tp_rank: int = 0, dp_rank: int | None = None, cp_rank: int | None = None +) -> SimpleNamespace: + return SimpleNamespace( + tp_rank=tp_rank, + has_dp=dp_rank is not None, + dp_rank=dp_rank or 0, + has_cp=cp_rank is not None, + cp_rank=cp_rank or 0, + ) + + +def _make_handler(attn_mapping: SimpleNamespace | None = None) -> RequestHandler: handler = RequestHandler.__new__(RequestHandler) handler.forward_ct = 0 - handler.attn_tp_rank = 0 + attn_mapping = attn_mapping or _attn_mapping() + handler.attn_tp_rank = attn_mapping.tp_rank + handler.profile_rank_tag = request_handler_mod._profile_rank_tag(attn_mapping) handler.init_profiler() return handler @@ -91,7 +106,7 @@ def test_init_fails_when_env_bootstrap_session_active(self): self.assertFalse(self.handler.profile_in_progress) def test_start_and_stop_drive_proton_session_per_rank(self): - self.handler.attn_tp_rank = 3 + self.handler = _make_handler(_attn_mapping(tp_rank=3)) with mock.patch.object( request_handler_mod, "proton_available", return_value=True ), mock.patch.object( @@ -105,7 +120,7 @@ def test_start_and_stop_drive_proton_session_per_rank(self): config = start_profiling.call_args[0][0] self.assertEqual(config.output.rsplit("/", 1)[0], self.output_dir) - self.assertTrue(config.output.endswith("test-profile-TP-3.proton")) + self.assertTrue(config.output.endswith("test-profile-TP3.proton")) stop_profiling.assert_not_called() result = self.handler.profile(ProfileReq(type=ProfileReqType.STOP_PROFILE)) @@ -113,6 +128,42 @@ def test_start_and_stop_drive_proton_session_per_rank(self): stop_profiling.assert_called_once() self.assertFalse(self.handler.profile_in_progress) + def test_rank_tag_includes_dp_cp_ranks_when_present(self): + self.assertEqual( + request_handler_mod._profile_rank_tag(_attn_mapping(tp_rank=3)), "TP3" + ) + self.assertEqual( + request_handler_mod._profile_rank_tag(_attn_mapping(tp_rank=0, dp_rank=1)), + "DP1-TP0", + ) + self.assertEqual( + request_handler_mod._profile_rank_tag( + _attn_mapping(tp_rank=2, dp_rank=1, cp_rank=0) + ), + "DP1-CP0-TP2", + ) + + def test_proton_outputs_do_not_collide_across_dp_ranks(self): + # Two DP peers share attn_tp_rank=0 but must write distinct files. + outputs = [] + for dp_rank in (0, 1): + handler = _make_handler(_attn_mapping(tp_rank=0, dp_rank=dp_rank)) + with mock.patch.object( + request_handler_mod, "proton_available", return_value=True + ), mock.patch.object( + request_handler_mod, "start_profiling" + ) as start_profiling, mock.patch.object( + request_handler_mod, "stop_profiling" + ): + result = handler.profile(_start_req(self.output_dir)) + self.assertTrue(result.success) + outputs.append(start_profiling.call_args[0][0].output) + handler.profile(ProfileReq(type=ProfileReqType.STOP_PROFILE)) + + self.assertNotEqual(outputs[0], outputs[1]) + self.assertTrue(outputs[0].endswith("test-profile-DP0-TP0.proton")) + self.assertTrue(outputs[1].endswith("test-profile-DP1-TP0.proton")) + 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 59f9fa393..79939aac4 100644 --- a/tokenspeed-kernel/README.md +++ b/tokenspeed-kernel/README.md @@ -123,7 +123,8 @@ iteration. `tokenspeed bench serve` (or POST `/start_profile` / `/stop_profile` with `{"activities": ["PROTON"]}`). Each scheduler process — the process where kernels actually launch — runs its own Proton session and finalizes it on - `/stop_profile`, writing `/-TP-.proton.` + `/stop_profile`, writing + `/[-DP][-CP]-TP.proton.` per rank. `PROTON` composes only with host-side activities (`CPU`, `MEM`, `VIZTRACER`). From b5caeda95a8f4e9ba0cf458c1dbfcd49b099bfee Mon Sep 17 00:00:00 2001 From: Lei Zhang Date: Sun, 12 Jul 2026 22:09:23 -0700 Subject: [PATCH 4/5] fix(profile): validate /start_profile before mutating profiler state init_profile set self.profile_by_stage before the PROTON activity validation, so a rejected request (e.g. activities=["PROTON", "GPU"] with profile_by_stage=true) returned failure while leaving profile_by_stage=True behind with the stage counters still None. The event loop runs _profile_batch_predicate on every batch, so the next scheduled batch entered the stage-profiling branch and crashed the scheduler with a TypeError on `None += 1`. Hoist the validation block above the first state mutation so a rejected control request leaves the handler exactly as it was, and add a regression test that replays the rejected request and drives the next-batch predicate. Signed-off-by: Lei Zhang --- python/tokenspeed/runtime/engine/request_handler.py | 6 ++++-- test/runtime/test_request_handler_profile.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/python/tokenspeed/runtime/engine/request_handler.py b/python/tokenspeed/runtime/engine/request_handler.py index 097924689..1c0bb6da8 100644 --- a/python/tokenspeed/runtime/engine/request_handler.py +++ b/python/tokenspeed/runtime/engine/request_handler.py @@ -324,13 +324,14 @@ def init_profile( message="Profiling is already in progress. Call /stop_profile first.", ) - self.profile_by_stage = profile_by_stage - if output_dir is None: output_dir = envs.TOKENSPEED_PROFILER_DIR.get() if activities is None: activities = ["CPU", "GPU"] + # All validation must precede any state mutation: the event loop runs + # _profile_batch_predicate on every batch, so a rejected request that + # left partial profiler state behind would crash the scheduler. if "PROTON" in activities: conflicting = sorted({"GPU", "CUDA_PROFILER"} & set(activities)) if conflicting: @@ -354,6 +355,7 @@ def init_profile( "be controlled through /start_profile.", ) + self.profile_by_stage = profile_by_stage self.profiler_output_dir = output_dir self.torch_profiler_with_stack = with_stack self.torch_profiler_record_shapes = record_shapes diff --git a/test/runtime/test_request_handler_profile.py b/test/runtime/test_request_handler_profile.py index aa7631df7..7e7013bce 100644 --- a/test/runtime/test_request_handler_profile.py +++ b/test/runtime/test_request_handler_profile.py @@ -8,6 +8,7 @@ from tokenspeed.runtime.engine import request_handler as request_handler_mod from tokenspeed.runtime.engine.io_struct import ProfileReq, ProfileReqType from tokenspeed.runtime.engine.request_handler import RequestHandler +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode def _attn_mapping( @@ -71,6 +72,18 @@ def test_init_rejects_proton_mixed_with_gpu_profilers(self): self.assertIn(activity, result.message) self.assertFalse(self.handler.profile_in_progress) + def test_rejected_request_leaves_profiler_state_untouched(self): + req = _start_req(self.output_dir, profile_by_stage=True, num_steps=2) + req.activities = ["PROTON", "GPU"] + + result = self.handler.profile(req) + + self.assertFalse(result.success) + self.assertFalse(self.handler.profile_by_stage) + self.handler.forward_ct = 1 + self.handler._profile_batch_predicate(ForwardMode.EXTEND) + self.assertFalse(self.handler.profile_in_progress) + def test_init_allows_proton_with_host_side_profilers(self): req = _start_req(self.output_dir) req.activities = ["CPU", "MEM", "VIZTRACER", "PROTON"] From ac4f8adcdc3a68ed945db3b9625681e1fa0ade89 Mon Sep 17 00:00:00 2001 From: Lei Zhang Date: Sun, 12 Jul 2026 13:59:34 -0700 Subject: [PATCH 5/5] feat(profile): merge Proton and VizTracer traces onto one timeline Profiling a scheduler rank with both VIZTRACER and PROTON activities yields two chrome-trace files whose event timestamps are relative to each profiler's own start, so they could not be viewed on a shared time axis. Both files carry an absolute wall-clock anchor for ts=0: VizTracer reports store viztracer_metadata.baseTimeNanoseconds, and Proton chrome traces now store a top-level baseTimeNanoseconds (added to the tokenspeed-triton fork, mirroring the torch>=2.4 chrome trace contract). That makes a fully offline merge possible. * New tokenspeed.trace_merge module and `tokenspeed merge-traces` CLI: shifts every Proton event by the delta between the two anchors and appends them to the VizTracer report, producing one Perfetto-loadable JSON where Python frames and Proton's kernel / scope lanes (including launch->kernel flow arrows) share a time axis. One file pair per scheduler rank. * Proton's synthetic pid-0 process is renamed "Trace" -> "Proton" in the merged view, and its metadata events stay timestamp-free so they cannot clobber the report's process/thread names. * Inputs missing their anchor fail with a pointed error (old tokenspeed-triton or viztracer versions). Alignment accuracy is us-to-ms grade (each profiler reconciles GPU and CPU clocks independently): good for correlating host stalls with GPU gaps, not for sub-microsecond attribution. Tested with synthetic trace pairs encoding both formats' contracts, plus an end-to-end merge of a real VizTracer report. Signed-off-by: Lei Zhang --- python/tokenspeed/cli/__main__.py | 19 ++++ python/tokenspeed/trace_merge.py | 168 ++++++++++++++++++++++++++++ test/test_trace_merge.py | 180 ++++++++++++++++++++++++++++++ tokenspeed-kernel/README.md | 7 +- 4 files changed, 373 insertions(+), 1 deletion(-) create mode 100644 python/tokenspeed/trace_merge.py create mode 100644 test/test_trace_merge.py diff --git a/python/tokenspeed/cli/__main__.py b/python/tokenspeed/cli/__main__.py index 955e76342..4dba1b569 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.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/trace_merge.py b/python/tokenspeed/trace_merge.py new file mode 100644 index 000000000..0ad31548d --- /dev/null +++ b/python/tokenspeed/trace_merge.py @@ -0,0 +1,168 @@ +# 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; the trace was written " + "by a tokenspeed-triton without absolute-anchor support. Upgrade " + "tokenspeed-triton 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/test/test_trace_merge.py b/test/test_trace_merge.py new file mode 100644 index 000000000..b347a1b94 --- /dev/null +++ b/test/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.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.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/tokenspeed-kernel/README.md b/tokenspeed-kernel/README.md index 79939aac4..080df4707 100644 --- a/tokenspeed-kernel/README.md +++ b/tokenspeed-kernel/README.md @@ -126,7 +126,12 @@ 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 + per-rank file pair: `tokenspeed merge-traces -TP0.viztracer.json + -TP0.proton.chrome_trace`. ### Plugins