[feature] Add layerwise NVTX support (#11870)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)):
|
||||
|
||||
@@ -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",
|
||||
|
||||
292
python/sglang/srt/utils/nvtx_pytorch_hooks.py
Normal file
292
python/sglang/srt/utils/nvtx_pytorch_hooks.py
Normal 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
|
||||
@@ -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),
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user