chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Kalyan Kumar
2026-03-05 12:52:56 +05:30
committed by GitHub
parent 2bdd89a6cd
commit c1df359b44
5 changed files with 98 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
python/sglang/bench_one_batch.py
View File

@@ -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)

29
python/sglang/bench_serving.py
View File

@@ -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)

3
python/sglang/profiler.py
View File

@@ -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()

2
python/sglang/srt/managers/scheduler_profiler_mixin.py
View File

@@ -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
]

23
python/sglang/test/bench_one_batch_server_internal.py
View File

@@ -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,