[feature] Add layerwise NVTX support (#11870)

This commit is contained in:
kyleliang-nv
2025-11-15 19:20:56 -08:00
committed by GitHub
parent 1ca205f6da
commit 597d416070
7 changed files with 718 additions and 3 deletions

View File

@@ -394,6 +394,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--enable-return-hidden-states` | Enable returning hidden states with responses. | `False` | bool flag (set to enable) |
| `--scheduler-recv-interval` | The interval to poll requests in scheduler. Can be set to >1 to reduce the overhead of this. | `1` | Type: int |
| `--numa-node` | Sets the numa node for the subprocesses. i-th element corresponds to i-th subprocess. | `None` | List[int] |
| `--enable-layerwise-nvtx-marker` | Enable layerwise NVTX profiling annotations for the model. This adds NVTX markers to every layer for detailed per-layer performance analysis with Nsight Systems. | `False` | bool flag (set to enable) |
| `--enable-attn-tp-input-scattered` | Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent. | `False` | bool flag (set to enable) |
## Debug tensor dumps

View File

@@ -122,6 +122,88 @@ You can also combine the above operations into a single command
python3 -m sglang.test.send_one --profile
```
### Profile a server with HTTP API endpoints
SGLang provides HTTP API endpoints to control profiling on a running server. This allows you to start and stop profiling programmatically, which is useful for capturing specific workload patterns.
#### Using `/start_profile` endpoint
The `/start_profile` endpoint starts profiling on the server. You can control when profiling begins and how long it runs using the following parameters:
**Basic usage:**
```bash
# Start profiling immediately for 10 steps
curl -X POST http://127.0.0.1:30000/start_profile \
-H "Content-Type: application/json" \
-d '{
"num_steps": 10
}'
```
**Parameters:**
- `output_dir` (optional): Directory where profile traces will be saved. If not specified, uses `SGLANG_TORCH_PROFILER_DIR` environment variable, or `/tmp` as the default
- `num_steps` (optional): Number of steps to profile. If not specified, profiling continues until manually stopped with `/end_profile`
- `start_step` (optional): Step number at which to start profiling (inclusive). Useful for skipping warmup iterations
- `activities` (optional): List of activities to profile, e.g., `["CPU", "GPU"]`. Default is `["CPU", "GPU"]`
- `merge_profiles` (optional): Whether to merge distributed traces. Default is `false`
**Note on step ranges:** Profiling starts at `start_step` (inclusive) and continues for `num_steps` iterations. For example, with `start_step=3` and `num_steps=10`, profiling captures steps 3, 4, 5, 6, 7, 8, 9, 10, 11, and 12 (10 steps total, starting from step 3).
**Advanced usage with `start_step`:**
```bash
# Wait 5 steps (warmup), then profile for 10 steps
curl -X POST http://127.0.0.1:30000/start_profile \
-H "Content-Type: application/json" \
-d '{
"output_dir": "/tmp/profiles",
"start_step": 5,
"num_steps": 10,
"activities": ["CPU", "GPU"]
}'
```
**Continuous profiling (manual stop):**
```bash
# Start profiling without num_steps - must manually stop with /end_profile
curl -X POST http://127.0.0.1:30000/start_profile
```
#### Using `/end_profile` endpoint
The `/end_profile` endpoint stops an ongoing profiling session and saves the trace file.
```bash
# Stop profiling and save traces
curl -X POST http://127.0.0.1:30000/end_profile
```
This is only needed when you start profiling without specifying `num_steps`. If `num_steps` is specified, profiling will automatically stop after that many steps.
#### Example workflow
```bash
# Terminal 1: Start the server
export SGLANG_TORCH_PROFILER_DIR=/tmp/profiles
python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct
# Terminal 2: Start continuous profiling
curl -X POST http://127.0.0.1:30000/start_profile \
-H "Content-Type: application/json" \
-d '{
"start_step": 3
}'
# Terminal 3: Send requests to generate load
python -m sglang.bench_serving --backend sglang --num-prompts 100
# Terminal 2: Stop profiling when done
curl -X POST http://127.0.0.1:30000/end_profile
```
### Profiler Trace Merger for Distributed Traces
SGLang now supports automatic merging of profiling traces from distributed setups with multiple parallelism types (TP, DP, PP, EP). This feature is particularly useful for analyzing performance across distributed runs.
@@ -259,6 +341,108 @@ Additionally, if you want to locate the SGLang Python source code through the cu
# some critical code
```
### Layer-wise NVTX Profiling with Nsight Systems
SGLang provides built-in layerwise NVTX annotations that can be combined with the CUDA Profiler for detailed per-layer profiling in Nsight Systems. This is particularly useful for identifying performance bottlenecks at the layer level.
#### Using `--enable-layerwise-nvtx-marker` with Nsight Systems and `/start_profile`
The `--enable-layerwise-nvtx-marker` flag automatically adds NVTX markers to every layer in your model. This is particularly powerful when combined with Nsight Systems profiling to see detailed per-layer performance.
**Method 1: Using `/start_profile` with CUDA_PROFILER (for programmatic control)**
This method allows you to control exactly when profiling starts/stops via HTTP API while Nsight Systems is running.
1. Launch the server with layerwise NVTX enabled under Nsight Systems:
```bash
# Terminal 1: Start server with nsys and capture-range option
nsys profile --trace-fork-before-exec=true \
--cuda-graph-trace=node \
--capture-range=cudaProfilerApi \
--capture-range-end=stop \
-o layerwise_profile \
python -m sglang.launch_server \
--model-path meta-llama/Llama-3.1-8B-Instruct \
--enable-layerwise-nvtx-marker \
--disable-cuda-graph
```
Note: NVTX markers are not emitted for kernel launches captured by CUDA graphs. Use `--disable-cuda-graph` to ensure all layerwise NVTX markers are emitted in the trace.
2. In another terminal, control profiling via `/start_profile` with `CUDA_PROFILER` activity:
```bash
# Terminal 2: Wait for server to be ready, then start CUDA profiling
# Wait 3 steps for warmup, then profile for 10 steps
curl -X POST http://127.0.0.1:30000/start_profile \
-H "Content-Type: application/json" \
-d '{
"start_step": 3,
"num_steps": 10,
"activities": ["CUDA_PROFILER"]
}'
```
3. Send requests to generate load:
```bash
# Terminal 3: Generate workload
python -m sglang.bench_serving --backend sglang --num-prompts 100
```
4. Profiling will automatically stop after 10 steps (due to `num_steps: 10`). If you hadn't specified `num_steps`, you would need to manually stop it:
```bash
# Terminal 2: Only needed if num_steps was not specified
curl -X POST http://127.0.0.1:30000/end_profile
```
The `--capture-range=cudaProfilerApi` option tells Nsight Systems to only capture data between `cudaProfilerStart()` and `cudaProfilerStop()` calls (triggered by `/start_profile` and `/end_profile`), reducing overhead and file size. The `start_step` parameter skips the first 3 steps to avoid capturing warmup overhead.
**Method 2: Simpler approach without `/start_profile` API**
For simpler use cases where you don't need fine-grained control over profiling start/stop, you can profile with Nsight Systems capturing the entire workload:
```bash
# Terminal 1: Start server with layerwise NVTX
# Note: --disable-cuda-graph ensures all NVTX markers are emitted
python -m sglang.launch_server \
--model-path meta-llama/Llama-3.1-8B-Instruct \
--enable-layerwise-nvtx-marker \
--disable-cuda-graph
# Terminal 2: Profile the benchmarking client
nsys profile --trace-fork-before-exec=true \
--cuda-graph-trace=node \
-o layerwise_profile \
python -m sglang.bench_serving --backend sglang --num-prompts 10
```
This approach profiles the entire client execution, including all server interactions. The layerwise NVTX markers will be visible in the Nsight Systems timeline.
**Viewing the profiling results:**
Open the generated `.qdrep` file with Nsight Systems:
```bash
nsys-ui layerwise_profile.qdrep
```
In the Nsight Systems GUI, you'll see:
- **NVTX ranges**: Each layer appears as a labeled range in the timeline with detailed information in the marker metadata
- **CUDA kernels**: All GPU kernels are shown alongside the layer annotations
- **Layer hierarchy**: The full module path (e.g., `meta-llama/Meta-Llama-3.1-8B-Instruct.model.layers.0.self_attn.qkv_proj`) helps identify specific layers. The prefix uses the full model path from `--model-path`.
- **Tensor shapes**: Input/output dimensions and parameter shapes are included in the NVTX marker data
**Benefits of layerwise NVTX profiling:**
- **Granular visibility**: See exactly which layers are taking the most time
- **Memory tracking**: Identify layers with large memory allocations
- **Bottleneck identification**: Quickly locate inefficient operations
- **Communication overhead**: In multi-GPU setups, see per-layer communication costs
- **Development debugging**: Validate that model architecture changes have the expected performance impact
## Other tips
1. You can benchmark a model using dummy weights by only providing the config.json file. This allows for quick testing of model variants without training. To do so, add `--load-format dummy` to the above commands and then you only need a correct `config.json` under the checkpoint folder.

View File

@@ -149,6 +149,7 @@ from sglang.srt.utils import (
slow_rank_detector,
xpu_has_xmx_support,
)
from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks
from sglang.srt.utils.offloader import (
create_offloader_from_server_args,
get_offloader,
@@ -772,6 +773,11 @@ class ModelRunner:
get_offloader().post_init()
# Register model for layerwise NVTX profiling if enabled
if self.server_args.enable_layerwise_nvtx_marker:
self.pyt_hooks = PytHooks()
self.pyt_hooks.register_hooks(self.model, module_prefix="model")
if self.server_args.kv_cache_dtype == "fp8_e4m3":
if self.server_args.quantization_param_path is not None:
if callable(getattr(self.model, "load_kv_cache_scales", None)):

View File

@@ -479,6 +479,7 @@ class ServerArgs:
disable_cuda_graph_padding: bool = False
enable_profile_cuda_graph: bool = False
enable_cudagraph_gc: bool = False
enable_layerwise_nvtx_marker: bool = False
enable_nccl_nvls: bool = False
enable_symm_mem: bool = False
disable_flashinfer_cutlass_moe_fp4_allgather: bool = False
@@ -3240,6 +3241,11 @@ class ServerArgs:
action="store_true",
help="Enable garbage collection during CUDA graph capture. If disabled (default), GC is frozen during capture to speed up the process.",
)
parser.add_argument(
"--enable-layerwise-nvtx-marker",
action="store_true",
help="Enable layerwise NVTX profiling annotations for the model.",
)
parser.add_argument(
"--enable-nccl-nvls",
action="store_true",

View File

@@ -0,0 +1,292 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""PyTorch hooks for layerwise NVTX profiling."""
import torch
import torch.cuda.nvtx as nvtx
class PytHooks(object):
"""This module contains all the code needed to enable forward hooks in a pytorch network.
To register the hooks for a given network, the user needs to instantiate a PytHook object.
Then call the register_hooks method.
Example:
my_hook = PytHook()
my_hook.register_hooks(my_network_model)
"""
def __init__(self):
"""Initialize module variables
Returns:
None:
Raises:
None:
"""
super().__init__()
self.module_to_name_map = {}
@staticmethod
def print_tensor(tensor_obj, prefix, tensor_list=None):
"""Descends iterators that contains Tensors and prints the Tensor
Recursive function that descends iterator type arguments until
it finds a Tensor object.
Args:
tensor_obj: Could be a Tensor or an iterator type that contains Tensors
prefix: String name to assign to the Tensor
tensor_list: List to accumulate tensor dimensions
Returns:
List of tensor dimensions
Raises:
None:
"""
if tensor_list is None:
tensor_list = []
if isinstance(tensor_obj, list) or isinstance(tensor_obj, tuple):
for ten in tensor_obj:
tensor_list = PytHooks.print_tensor(ten, prefix, tensor_list)
elif isinstance(tensor_obj, torch.Tensor):
tensor_dims = list(tensor_obj.size())
tensor_list.append(tensor_dims)
return tensor_list
def process_layer_params(self, module_obj):
"""Extract the static parameters from LLM and VLM relevant layer types
Args:
module_obj(class): Module state data structure.
Returns:
param_info(dict): Parameter meta_data for the given op.
Raises:
None
"""
param_info = {}
# Extract parameters for layers commonly used in LLMs and VLMs
if (
isinstance(module_obj, torch.nn.Conv1d)
or isinstance(module_obj, torch.nn.Conv2d)
or isinstance(module_obj, torch.nn.Conv3d)
):
conv_params = {}
conv_params["in_chan"] = module_obj.in_channels
conv_params["out_chan"] = module_obj.out_channels
conv_params["filter_dim"] = module_obj.kernel_size
conv_params["stride"] = module_obj.stride
conv_params["padding"] = module_obj.padding
conv_params["dilation"] = module_obj.dilation
conv_params["transposed"] = module_obj.transposed
conv_params["output_padding"] = module_obj.output_padding
conv_params["groups"] = module_obj.groups
conv_params["padding_mode"] = module_obj.padding_mode
param_info = conv_params
elif (
isinstance(module_obj, torch.nn.ConvTranspose1d)
or isinstance(module_obj, torch.nn.ConvTranspose2d)
or isinstance(module_obj, torch.nn.ConvTranspose3d)
):
convtranspose_params = {}
convtranspose_params["in_chan"] = module_obj.in_channels
convtranspose_params["out_chan"] = module_obj.out_channels
convtranspose_params["filter_dim"] = module_obj.kernel_size
convtranspose_params["stride"] = module_obj.stride
convtranspose_params["padding"] = module_obj.padding
convtranspose_params["dilation"] = module_obj.dilation
convtranspose_params["transposed"] = module_obj.transposed
convtranspose_params["output_padding"] = module_obj.output_padding
convtranspose_params["groups"] = module_obj.groups
convtranspose_params["padding_mode"] = module_obj.padding_mode
param_info = convtranspose_params
elif (
isinstance(module_obj, torch.nn.MaxPool1d)
or isinstance(module_obj, torch.nn.MaxPool2d)
or isinstance(module_obj, torch.nn.MaxPool3d)
):
def _handle_int_or_tuple(parameter):
if isinstance(parameter, tuple):
return list(parameter)
elif isinstance(parameter, int):
return [parameter, parameter]
pooling_params = {}
pooling_params["filter_dim"] = _handle_int_or_tuple(module_obj.kernel_size)
pooling_params["stride"] = _handle_int_or_tuple(module_obj.stride)
pooling_params["padding"] = _handle_int_or_tuple(module_obj.padding)
pooling_params["dilation"] = _handle_int_or_tuple(module_obj.dilation)
param_info = pooling_params
elif (
isinstance(module_obj, torch.nn.AvgPool1d)
or isinstance(module_obj, torch.nn.AvgPool2d)
or isinstance(module_obj, torch.nn.AvgPool3d)
):
pooling_params = {}
pooling_params["filter_dim"] = [
module_obj.kernel_size,
module_obj.kernel_size,
]
pooling_params["stride"] = [module_obj.stride, module_obj.stride]
pooling_params["padding"] = [module_obj.padding, module_obj.padding]
pooling_params["ceil_mode"] = module_obj.ceil_mode
pooling_params["count_include_pad"] = module_obj.count_include_pad
param_info = pooling_params
elif (
isinstance(module_obj, torch.nn.AdaptiveAvgPool1d)
or isinstance(module_obj, torch.nn.AdaptiveAvgPool2d)
or isinstance(module_obj, torch.nn.AdaptiveAvgPool3d)
):
pooling_params = {}
pooling_params["output_size"] = [
module_obj.output_size,
module_obj.output_size,
]
param_info = pooling_params
elif isinstance(module_obj, torch.nn.Linear):
param_info["in_features"] = module_obj.in_features
param_info["out_features"] = module_obj.out_features
elif (
isinstance(module_obj, torch.nn.BatchNorm1d)
or isinstance(module_obj, torch.nn.BatchNorm2d)
or isinstance(module_obj, torch.nn.BatchNorm3d)
):
param_info["num_features"] = module_obj.num_features
param_info["epsilon"] = module_obj.eps
param_info["momentum"] = module_obj.momentum
elif isinstance(module_obj, torch.nn.ReLU):
param_info["in_place"] = module_obj.inplace
elif isinstance(module_obj, torch.nn.Dropout):
param_info["p"] = module_obj.p
param_info["in_place"] = module_obj.inplace
elif isinstance(module_obj, torch.nn.Embedding):
param_info["num_embeddings"] = module_obj.num_embeddings
param_info["embedding_dim"] = module_obj.embedding_dim
elif isinstance(
module_obj,
(
torch.nn.Upsample,
torch.nn.UpsamplingNearest2d,
torch.nn.UpsamplingBilinear2d,
),
):
param_info["scale_factor"] = module_obj.scale_factor
return param_info
def module_fwd_hook(self, module_obj, in_tensor, out_tensor):
"""Callback function that ends the NVTX marker
Records the module name and tensor information
Called after the module executes the forward method.
Args:
module_obj: Pointer to the module object
in_tensor: Input tensor or list of tensors
out_tensor: Output tensor of the resulting forward operator
Returns:
None:
Raises:
None:
"""
nvtx.range_pop()
return
def module_fwd_pre_hook(self, module_obj, in_tensor):
"""Creates an NVTX marker with the module name in it.
This function is called before the module executes
Args:
module_obj: Module object data structure - used to get unique module name
in_tensor: Input tensor data structure
Returns:
None
Raises:
None
"""
marker_dict = {}
module_name = self.module_to_name_map.get(module_obj, "unknown")
marker_dict["Module"] = module_name
## Get trainable parameters like weights and bias
module_params = module_obj.named_parameters(recurse=False)
for idx, (param_name, param_obj) in enumerate(module_params):
if idx == 0:
marker_dict["TrainableParams"] = {}
marker_dict["TrainableParams"][param_name] = list(param_obj.size())
in_tensor_list = PytHooks.print_tensor(in_tensor, "Input")
if in_tensor_list:
marker_dict["Inputs"] = in_tensor_list
param_info = self.process_layer_params(module_obj)
if param_info:
marker_dict["StaticParams"] = param_info
nvtx.range_push("{}".format(marker_dict))
return
def register_hooks(self, network_model, module_prefix="top"):
"""User level function that activates all the hooks
The user needs to call this method from the network source code
The code descends all the modules in the network and registers their
respective hooks.
Args:
network_model: Model object for the network
module_prefix: (default: top)
Returns:
None
Raises:
Exception if a module instance is reused
"""
# Module types to skip (simple operations that don't need detailed profiling)
skip_types = (
torch.nn.Identity,
torch.nn.Dropout,
torch.nn.Dropout1d,
torch.nn.Dropout2d,
torch.nn.Dropout3d,
)
for name, module in network_model.named_modules(prefix=module_prefix):
# Skip certain module types to reduce profiling overhead
if isinstance(module, skip_types):
continue
module.register_forward_pre_hook(self.module_fwd_pre_hook)
module.register_forward_hook(self.module_fwd_hook)
if module not in self.module_to_name_map:
self.module_to_name_map[module] = name
else:
raise ValueError("Module instance {} is not unique ".format(module))
return

View File

@@ -117,7 +117,7 @@ suites = {
TestFile("test_srt_endpoint.py", 130),
TestFile("test_srt_engine.py", 450),
TestFile("test_standalone_speculative_decoding.py", 150),
TestFile("test_start_profile.py", 60),
TestFile("test_start_profile.py", 180),
TestFile("test_profile_merger.py", 60),
TestFile("test_profile_merger_http_api.py", 15),
TestFile("test_swa_unittest.py", 1),
@@ -213,8 +213,8 @@ suites = {
TestFile("lora/test_lora_openai_api.py", 30),
TestFile("openai_server/features/test_lora_openai_compatible.py", 150),
TestFile("batch_invariant/test_batch_invariant_ops.py", 10),
TestFile("test_deepseek_v3_deterministic.py", 240),
TestFile("test_cpp_radix_cache.py", 60),
TestFile("test_deepseek_v3_deterministic.py", 240),
],
"nightly-4-gpu-b200": [
TestFile("nightly/test_flashinfer_trtllm_gen_moe_backend.py", 300),

View File

@@ -1,10 +1,19 @@
"""
Usage:
python3 -m unittest test_srt_engine.TestSRTEngine.test_4_sync_async_stream_combination
# From the test/srt directory:
cd test/srt
python3 -m unittest test_start_profile.TestStartProfile
python3 -m unittest test_start_profile.TestStartProfileWithNsys
# Run specific tests:
python3 -m unittest test_start_profile.TestStartProfile.test_start_profile_1
python3 -m unittest test_start_profile.TestStartProfileWithNsys.test_start_profile_cuda_profiler
"""
import os
import shutil
import subprocess
import time
import unittest
import requests
@@ -22,6 +31,15 @@ from sglang.test.test_utils import (
OUTPUT_DIR = "./profiler_dir"
def _is_nsys_available():
"""Check if nsys (Nsight Systems) is available on the system."""
try:
result = subprocess.run(["nsys", "--version"], capture_output=True, timeout=5)
return result.returncode == 0
except (FileNotFoundError, subprocess.TimeoutExpired):
return False
class TestStartProfile(CustomTestCase):
@classmethod
@@ -112,5 +130,213 @@ class TestStartProfile(CustomTestCase):
)
class TestStartProfileWithNsys(CustomTestCase):
"""Test /start_profile with CUDA_PROFILER (requires nsys wrapper)
Each test starts its own clean server instance with nsys profiling.
"""
@classmethod
def setUpClass(cls):
if not _is_nsys_available():
raise unittest.SkipTest("nsys (Nsight Systems) is not available")
envs.SGLANG_TORCH_PROFILER_DIR.set(OUTPUT_DIR)
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
# Use a different port to avoid conflicts with other tests
cls.base_url = "http://127.0.0.1:21100"
def setUp(self):
"""Start a clean server with nsys for each test"""
# Kill any existing processes on this port
self._kill_existing_server()
# Clean up old profile files for this test
test_name = self.id().split(".")[-1] # Get test method name
self.nsys_output_file = f"nsys_profile_{test_name}"
if os.path.isdir(OUTPUT_DIR):
profile_file = os.path.join(OUTPUT_DIR, f"{self.nsys_output_file}.nsys-rep")
if os.path.exists(profile_file):
try:
os.remove(profile_file)
except OSError:
pass
# Launch server with nsys wrapper
self.process = self._popen_launch_server_with_nsys(
self.model,
self.base_url,
self.nsys_output_file,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
)
def tearDown(self):
"""Kill server and verify profile was created"""
# Kill server first to let nsys finalize the .nsys-rep file
kill_process_tree(self.process.pid)
# Also ensure nsys agent processes are killed
try:
subprocess.run(
["pkill", "-f", "nsys.*--start-agent"],
timeout=5,
stderr=subprocess.DEVNULL,
)
except (subprocess.TimeoutExpired, FileNotFoundError):
pass
# Give nsys a moment to finalize the report
time.sleep(3)
# Verify the .nsys-rep file was created
self._verify_nsys_profile_created()
def _kill_existing_server(self):
"""Kill any existing server process on our port and orphaned nsys agents"""
try:
# Kill server on our port
subprocess.run(["lsof", "-ti", ":21100"], capture_output=True, timeout=5)
subprocess.run(["pkill", "-f", "sglang.launch_server.*21100"], timeout=5)
# Kill any orphaned nsys agent processes
subprocess.run(
["pkill", "-f", "nsys.*--start-agent"],
timeout=5,
stderr=subprocess.DEVNULL, # Suppress "no process found" errors
)
time.sleep(2) # Wait for cleanup
except (subprocess.TimeoutExpired, FileNotFoundError):
pass
def _popen_launch_server_with_nsys(self, model, base_url, output_file, timeout):
"""Launch server wrapped with nsys profile -c cudaProfilerApi
Each test gets its own output file for complete isolation.
"""
_, host, port = base_url.split(":")
host = host[2:]
# Build the server launch command
command = [
"nsys",
"profile",
"-c",
"cudaProfilerApi",
"--capture-range-end",
"stop", # Stop after first cudaProfilerStop()
"-o",
os.path.join(OUTPUT_DIR, output_file),
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
]
# Create output directory if it doesn't exist
os.makedirs(OUTPUT_DIR, exist_ok=True)
# Launch the process - capture output to keep test output clean
process = subprocess.Popen(
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
# Wait for server to be ready
start_time = time.perf_counter()
elapsed = 0
with requests.Session() as session:
while elapsed < timeout:
elapsed = time.perf_counter() - start_time
return_code = process.poll()
if return_code is not None:
raise Exception(
f"Server process exited with code {return_code}. "
"Check server logs above for errors."
)
try:
response = session.get(f"{base_url}/health_generate", timeout=5)
if response.status_code == 200:
return process
except (requests.RequestException, requests.Timeout):
pass
time.sleep(5)
# Timeout reached
kill_process_tree(process.pid)
raise TimeoutError(
f"Server failed to start within {timeout} seconds. "
f"Check the server logs above for more information."
)
def _verify_nsys_profile_created(self):
"""Verify that the .nsys-rep file was created after server shutdown."""
if not os.path.isdir(OUTPUT_DIR):
raise AssertionError("Output directory does not exist.")
expected_file = f"{self.nsys_output_file}.nsys-rep"
profile_path = os.path.join(OUTPUT_DIR, expected_file)
if not os.path.exists(profile_path):
files = os.listdir(OUTPUT_DIR)
raise AssertionError(
f"Expected profile file '{expected_file}' not found. "
f"Files present: {files}"
)
def test_start_profile_cuda_profiler_with_start_step(self):
"""Test /start_profile with CUDA_PROFILER, start_step, and num_steps"""
# Use start_step to let server warm up before profiling
response = self._start_profile(
activities=["CUDA_PROFILER"], start_step=10, num_steps=3
)
self._post_request()
# Profile verification happens in tearDown()
def test_start_profile_cuda_profiler(self):
"""Test /start_profile with CUDA_PROFILER activity (no start_step)"""
# Simple num_steps test - profiling starts immediately
response = self._start_profile(activities=["CUDA_PROFILER"], num_steps=5)
self._post_request()
# Profile verification happens in tearDown()
def _start_profile(self, **kwargs):
"""Start profiling with optional parameters."""
response = requests.post(
f"{self.base_url}/start_profile",
json=kwargs if kwargs else None,
)
self.assertEqual(response.status_code, 200)
return response
def _post_request(self):
response = requests.post(
f"{self.base_url}/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
},
},
)
self.assertEqual(response.status_code, 200)
if __name__ == "__main__":
unittest.main()