Add XPU profiler activity support in benchmark code (#12981)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/bench_one_batch.py

@@ -57,7 +57,7 @@ import multiprocessing
 import os
 import time
 from types import SimpleNamespace
-from typing import Tuple
+from typing import Optional, Tuple
 
 import numpy as np
 import torch
@@ -79,8 +79,6 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
     configure_logger,
     get_bool_env_var,
-    is_cuda_alike,
-    is_xpu,
     kill_process_tree,
     maybe_reindex_device_id,
     require_mlp_sync,
@@ -90,15 +88,6 @@ from sglang.srt.utils import (
 )
 from sglang.srt.utils.hf_transformers_utils import get_tokenizer
 
-profile_activities = [torch.profiler.ProfilerActivity.CPU] + [
-    profiler_activity
-    for available, profiler_activity in [
-        (is_cuda_alike(), torch.profiler.ProfilerActivity.CUDA),
-        (is_xpu(), torch.profiler.ProfilerActivity.XPU),
-    ]
-    if available
-]
-
 
 def start_profile(profile_activities, profile_record_shapes=False, rank_print=print):
     """
@@ -118,6 +107,8 @@ def start_profile(profile_activities, profile_record_shapes=False, rank_print=pr
             activities.append(torch.profiler.ProfilerActivity.CPU)
         if "GPU" in profile_activities:
             activities.append(torch.profiler.ProfilerActivity.CUDA)
+        if "XPU" in profile_activities:
+            activities.append(torch.profiler.ProfilerActivity.XPU)
         if activities:
             profiler = torch.profiler.profile(
                 activities=activities,
@@ -179,6 +170,8 @@ class BenchArgs:
     profile_activities: Tuple[str] = ("CPU", "GPU")
     profile_stage: str = "all"
     profile_filename_prefix: str = "profile"
+    profile_start_step: Optional[int] = None
+    profile_steps: Optional[int] = None
 
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
@@ -217,8 +210,8 @@ class BenchArgs:
             type=str,
             nargs="+",
             default=["CPU", "GPU"],
-            choices=["CPU", "GPU", "CUDA_PROFILER"],
+            choices=["CPU", "GPU", "CUDA_PROFILER", "XPU"],
-            help="Profiler activities: CPU, GPU, CUDA_PROFILER. If CPU/GPU, use torch profiler. If CUDA_PROFILER, use CUDA profiler.",
+            help="Profiler activities: CPU, GPU, XPU, CUDA_PROFILER. If CPU/GPU/XPU, use torch profiler. If CUDA_PROFILER, use CUDA profiler.",
         )
         parser.add_argument(
             "--profile-stage",
@@ -234,14 +227,32 @@ class BenchArgs:
             help="Prefix of the profiling file names. The full profiling result file(s) be "
             '"[profile_filename_prefix]_batch[batch_size]_input[input_len]_output[output_len].trace.json.gz"',
         )
+        parser.add_argument(
+            "--profile-start-step",
+            type=int,
+            default=None,
+            help="Decode step at which to start profiling (0-indexed). If not specified, defaults to output_len // 2.",
+        )
+        parser.add_argument(
+            "--profile-steps",
+            type=int,
+            default=None,
+            help="Number of decode steps to profile starting from profile-start-step. If not specified, profiles only one step.",
+        )
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace):
         # use the default value's type to cast the args into correct types.
         attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)]
-        return cls(
+        result = {}
-            **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs}
+        for attr, attr_type in attrs:
-        )
+            value = getattr(args, attr)
+            # Handle None values - don't try to cast them
+            if value is None or attr_type == type(None):
+                result[attr] = value
+            else:
+                result[attr] = attr_type(value)
+        return cls(**result)
 
 
 def load_model(server_args, port_args, gpu_id, tp_rank):
@@ -525,6 +536,8 @@ def latency_test_run_once(
     profile_filename_prefix,
     profile_stage,
     tp_rank,
+    profile_start_step=None,
+    profile_steps=None,
 ):
     max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len)
     if batch_size > max_batch_size:
@@ -582,12 +595,17 @@ def latency_test_run_once(
     measurement_results["prefill_throughput"] = throughput
 
     decode_latencies = []
-    profile_step_of_interest = output_len // 2
+    # Determine profiling start step and end step
+    profile_start = (
+        profile_start_step if profile_start_step is not None else (output_len // 2)
+    )
+    profile_end = profile_start + (profile_steps if profile_steps is not None else 1)
     enable_profile_decode = profile and profile_stage in ["all", "decode"]
+    profiler = None
     for i in range(output_len - 1):
         synchronize(device)
-        profiler = None
+        # Start profiler at the specified step
-        if enable_profile_decode and i == profile_step_of_interest:
+        if enable_profile_decode and i == profile_start:
             profiler = start_profile(
                 profile_activities,
                 profile_record_shapes=profile_record_shapes,
@@ -599,7 +617,8 @@ def latency_test_run_once(
         synchronize(device)
         latency = time.perf_counter() - tic
 
-        if enable_profile_decode and i == profile_step_of_interest:
+        # Stop profiler after the specified number of steps
+        if enable_profile_decode and profiler is not None and i >= profile_end - 1:
             trace_filename = _create_torch_profiler_filename(
                 profile_filename_prefix, batch_size, input_len, output_len, "decode"
             )
@@ -611,6 +630,7 @@ def latency_test_run_once(
                 trace_filename=trace_filename,
                 stage="decode",
             )
+            profiler = None
 
         tot_latency += latency
         throughput = batch_size / latency
@@ -686,6 +706,8 @@ def latency_test(
         profile_filename_prefix="",
         profile_stage="all",
         tp_rank=tp_rank,
+        profile_start_step=None,
+        profile_steps=None,
     )
 
     rank_print("Benchmark ...")
@@ -736,6 +758,8 @@ def latency_test(
             bench_args.profile_filename_prefix,
             bench_args.profile_stage,
             tp_rank,
+            bench_args.profile_start_step,
+            bench_args.profile_steps,
         )
         if ret is not None:
             result_list.append(ret)

python/sglang/bench_serving.py

@@ -730,6 +730,14 @@ async def async_request_profile(api_url: str) -> RequestFuncOutput:
                 # stop_profile doesn't need any parameters
                 body = {}
             print(f"async_request_profile {api_url=} {body=}")
+            # Add optional profiling parameters if provided
+            if (
+                hasattr(args, "profile_start_step")
+                and args.profile_start_step is not None
+            ):
+                body["start_step"] = str(args.profile_start_step)
+            if hasattr(args, "profile_steps") and args.profile_steps is not None:
+                body["num_steps"] = str(args.profile_steps)
             async with session.post(url=api_url, json=body) as response:
                 if response.status == 200:
                     output.success = True
@@ -1312,8 +1320,10 @@ async def benchmark(
     if is_multi_turn:
         outputs = [x for output in outputs for x in output]
 
-    # Stop profiler
+    # Stop profiler (only if profile_steps was not provided, as it auto-stops)
-    if profile:
+    if profile and not (
+        hasattr(args, "profile_steps") and args.profile_steps is not None
+    ):
         if pd_separated:
             if pd_profile_urls:
                 await _call_profile_pd(pd_profile_urls, "stop")
@@ -2016,7 +2026,20 @@ if __name__ == "__main__":
         type=str,
         nargs="+",
         default=["CPU", "GPU"],
-        choices=["CPU", "GPU", "CUDA_PROFILER"],
+        choices=["CPU", "GPU", "CUDA_PROFILER", "XPU"],
+        help="Profiler activities to capture: CPU, GPU, XPU, CUDA_PROFILER.",
+    )
+    parser.add_argument(
+        "--profile-start-step",
+        type=int,
+        default=None,
+        help="Start profiling after this many forward steps. Useful for warmup.",
+    )
+    parser.add_argument(
+        "--profile-steps",
+        type=int,
+        default=None,
+        help="Number of steps to profile. If specified, profiling stops automatically after this many steps.",
     )
     parser.add_argument("--profile-num-steps", type=int, default=None)
     parser.add_argument("--profile-by-stage", action="store_true", default=False)

python/sglang/profiler.py

@@ -26,6 +26,7 @@ def run_profile(
     profile_by_stage: bool = False,
     merge_profiles: bool = False,
     profile_prefix: Optional[str] = None,
+    start_step: Optional[int] = None,
 ) -> str:
     if output_dir is None:
         output_dir = PROFILER_DIR
@@ -57,6 +58,8 @@ def run_profile(
         "merge_profiles": merge_profiles,
         "profile_prefix": profile_prefix,
     }
+    if start_step is not None:
+        json_data["start_step"] = str(start_step)
 
     response = requests.post(url=url + "/start_profile", json=json_data)
     response.raise_for_status()

python/sglang/srt/managers/scheduler_profiler_mixin.py

@@ -154,6 +154,8 @@ class SchedulerProfilerMixin:
             "CPU": torch.profiler.ProfilerActivity.CPU,
             "GPU": torch.profiler.ProfilerActivity.CUDA,
         }
+        if hasattr(torch.profiler.ProfilerActivity, "XPU"):
+            activity_map["XPU"] = torch.profiler.ProfilerActivity.XPU
         torchprof_activities = [
             activity_map[a] for a in activities if a in activity_map
         ]

python/sglang/test/bench_one_batch_server_internal.py

@@ -94,6 +94,8 @@ class BenchArgs:
     skip_warmup: bool = False
     show_report: bool = False
     profile: bool = False
+    profile_activities: Tuple[str] = ("CPU", "GPU")
+    profile_start_step: Optional[int] = None
     profile_steps: int = 5
     profile_by_stage: bool = False
     profile_prefix: Optional[str] = None
@@ -141,6 +143,20 @@ class BenchArgs:
         parser.add_argument("--skip-warmup", action="store_true")
         parser.add_argument("--show-report", action="store_true")
         parser.add_argument("--profile", action="store_true")
+        parser.add_argument(
+            "--profile-activities",
+            type=str,
+            nargs="+",
+            default=("CPU", "GPU"),
+            choices=["CPU", "GPU", "XPU"],
+            help="Profiler activities: CPU, GPU, XPU. use torch profiler.",
+        )
+        parser.add_argument(
+            "--profile-start-step",
+            type=int,
+            default=BenchArgs.profile_start_step,
+            help="Start profiling after this many forward steps. Useful for warmup.",
+        )
         parser.add_argument(
             "--profile-steps", type=int, default=BenchArgs.profile_steps
         )
@@ -394,6 +410,8 @@ def run_one_case(
     result_filename: str,
     tokenizer: PreTrainedTokenizer | AutoProcessor,
     profile: bool = False,
+    profile_activities: Tuple[str] = ("CPU", "GPU"),
+    profile_start_step: Optional[int] = None,
     profile_steps: int = BenchArgs.profile_steps,
     profile_by_stage: bool = False,
     profile_prefix: Optional[str] = BenchArgs.profile_prefix,
@@ -523,10 +541,11 @@ def run_one_case(
         profile_link: str = run_profile(
             url=url,
             num_steps=profile_steps,
-            activities=["CPU", "GPU"],
+            profile_activities=profile_activities,
             output_dir=profile_output_dir,
             profile_by_stage=profile_by_stage,
             profile_prefix=profile_prefix,
+            start_step=profile_start_step,
         )
 
     # Get metrics before the request (for cache hit rate calculation)
@@ -918,6 +937,8 @@ def run_benchmark_internal(
                             parallel_batch=bench_args.parallel_batch,
                             cache_hit_rate=bench_args.cache_hit_rate,
                             profile=bench_args.profile,
+                            profile_activities=bench_args.profile_activities,
+                            profile_start_step=bench_args.profile_start_step,
                             profile_steps=bench_args.profile_steps,
                             profile_by_stage=bench_args.profile_by_stage,
                             profile_prefix=profile_prefix,