diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index ca4bad804..8cf0aee1b 100644 --- a/python/sglang/bench_one_batch.py +++ b/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"], - help="Profiler activities: CPU, GPU, CUDA_PROFILER. If CPU/GPU, use torch profiler. If CUDA_PROFILER, use CUDA profiler.", + choices=["CPU", "GPU", "CUDA_PROFILER", "XPU"], + 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( - **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} - ) + result = {} + 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 - if enable_profile_decode and i == profile_step_of_interest: + # Start profiler at the specified step + 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) diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index 12b086fd8..dec3109c2 100644 --- a/python/sglang/bench_serving.py +++ b/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 - if profile: + # Stop profiler (only if profile_steps was not provided, as it auto-stops) + 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) diff --git a/python/sglang/profiler.py b/python/sglang/profiler.py index 0cbca82e1..ebc7a100e 100644 --- a/python/sglang/profiler.py +++ b/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() diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py index 7d08f12b3..c02ed7997 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/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 ] diff --git a/python/sglang/test/bench_one_batch_server_internal.py b/python/sglang/test/bench_one_batch_server_internal.py index 5aea6a7da..aa864d952 100644 --- a/python/sglang/test/bench_one_batch_server_internal.py +++ b/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,