feat(profile): merge Proton and VizTracer traces onto one timeline

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,

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)

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 (
@@ -71,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
@@ -106,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
@@ -266,8 +285,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):
@@ -302,13 +324,38 @@ 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:
+                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.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
@@ -368,12 +415,21 @@ 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}-{self.profile_rank_tag}{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(
                 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")
@@ -411,22 +467,27 @@ 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)
 
         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)
 
         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()