[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

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