Skip to content

SGLang 0.5.12 patches #679

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mohbasit wants to merge 4 commits into
base: main
Choose a base branch
from
Open

SGLang 0.5.12 patches #679

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9f5ea55
SGLang 0.5.12 patches
mohbasit May 29, 2026
e2ba3bd
SGLang 0.5.12 patches
mohbasit May 29, 2026
5441801
Merge branch 'main' into patches/sglang_patch
devalshahamd Jun 2, 2026
3009119
Merge branch 'main' into patches/sglang_patch
ajassani Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 40 additions & 13 deletions docs/Inference_analysis.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,25 +105,52 @@ Then create a container from the image.

###### SGLang Script

The build script for SGLang supports SGLang 0.5.9 and 0.5.11. It takes the path to the local TraceLens clone, the SGLang version (`--sglang-version`, default 0.5.9), and the GPU type (`--gpu-type`, default mi350). MI300 and MI350/MI355 are supported.
The build script for SGLang supports SGLang 0.5.9, 0.5.11, and 0.5.12. It takes the path to the local TraceLens clone, the SGLang version (`--sglang-version`, default 0.5.9), and the GPU type (`--gpu-type`, default mi350). MI300 and MI350/MI355 are supported.

**SGLang versions** (`--sglang-version`):

| SGLang Version | GPU Type | Base Image |
| -------------- | -------- | --------------------------------------- |
| `0.5.9` | MI300 | `lmsysorg/sglang:v0.5.9-rocm700-mi30x` |
| `0.5.9` | MI350/MI355 | `lmsysorg/sglang:v0.5.9-rocm700-mi35x` |
| `0.5.11` | MI300 | `lmsysorg/sglang:v0.5.11-rocm720-mi30x` |
| `0.5.11` | MI350/MI355 | `lmsysorg/sglang:v0.5.11-rocm720-mi35x` |
| `0.5.12` | MI300 | `lmsysorg/sglang:v0.5.12-rocm720-mi30x` |
| `0.5.12` | MI350/MI355 | `lmsysorg/sglang:v0.5.12-rocm720-mi35x` |

| Version | Patch directory | ROCm | Description |
| ------- | --------------- | ---- | ----------- |
| `0.5.9` (default) | `sglang_roofline_patches/sglang_0_5_9/` | 7.0 | Legacy stack |
| `0.5.11` | `sglang_roofline_patches/sglang_0_5_11/` | 7.2 | Previous stable |
| `0.5.12` | `sglang_roofline_patches/sglang_0_5_12/` | 7.2 | Recommended for newer stacks |

Each version subdirectory contains the same set of patch files (`cuda_graph_runner`, `fused_moe_triton_kernels`, `http_server`, `io_struct`, `kernel_shape_profiler`, `profile_utils`, `scheduler`, `scheduler_profiler_mixin`, `server_args`, `tokenizer_communicator_mixin`). Use the directory that matches your SGLang version exactly.

**Base images** (overridable with `--base-image`):

| SGLang Version | GPU Type | Base Image |
| -------------- | -------- | ---------- |
| `0.5.9` | MI300 | `lmsysorg/sglang:v0.5.9-rocm700-mi30x` |
| `0.5.9` | MI350/MI355 | `lmsysorg/sglang:v0.5.9-rocm700-mi35x` |
| `0.5.11` | MI300 | `lmsysorg/sglang:v0.5.11-rocm720-mi30x` |
| `0.5.11` | MI350/MI355 | `lmsysorg/sglang:v0.5.11-rocm720-mi35x` |
| `0.5.12` | MI300 | `lmsysorg/sglang:v0.5.12-rocm720-mi30x` |
| `0.5.12` | MI350/MI355 | `lmsysorg/sglang:v0.5.12-rocm720-mi35x` |

**Examples:**

```bash
# SGLang 0.5.12 on MI300X (recommended)
bash examples/custom_workflows/inference_analysis/build_docker_sglang.sh \
/path/to/TraceLens \
--sglang-version 0.5.12 \
--gpu-type mi300 \
-t tracelens-sglang:0.5.12-mi300

# SGLang 0.5.11 on MI355X
bash examples/custom_workflows/inference_analysis/build_docker_sglang.sh \
/path/to/TraceLens \
--sglang-version 0.5.11 \
--gpu-type mi355 \
-t tracelens-sglang:0.5.11-mi355

# SGLang 0.5.9 on MI350 (legacy default)
bash examples/custom_workflows/inference_analysis/build_docker_sglang.sh \
/path/to/TraceLens \
--gpu-type mi350 \
-t tracelens-sglang
-t tracelens-sglang:0.5.9-mi350
```

Then create a container from the image.
Expand Down Expand Up @@ -151,7 +178,7 @@ If you prefer to patch an existing environment instead of building a new image,
- Select by framework and version
- Apply: `cd /path/to/framework/../ && git apply /path/to/patchfile`
vLLM patches are in [vllm_roofline_patches](../examples/custom_workflows/inference_analysis/vllm_patches)
SGLang patches are in [sglang_roofline_patches](../examples/custom_workflows/inference_analysis/sglang_roofline_patches/)
SGLang patches are in [sglang_roofline_patches](../examples/custom_workflows/inference_analysis/sglang_roofline_patches/) — use the `sglang_0_5_9/`, `sglang_0_5_11/`, or `sglang_0_5_12/` subdirectory matching your SGLang version

##### Collection Parameters

Expand Down Expand Up @@ -196,7 +223,7 @@ The `config_vllm_v*.patch` patches add two `ProfilerConfig` flags that control g

This optional step reads the collected trace and splits it into smaller trace files or execution‑phase‑specific trace files.

Option 1: Find steady-state region of execution (highest concurrency) and separate prefill-decode and decode-only execution steps (supports vLLM v0.14–v0.20 and SGLang v0.5.9; using the patchfile is recommended). This is recommended if the tracefile is large and the user wants to extract a few representative steps automatically.
Option 1: Find steady-state region of execution (highest concurrency) and separate prefill-decode and decode-only execution steps (supports vLLM v0.14–v0.20 and SGLang v0.5.9–v0.5.12; using the patchfile is recommended). This is recommended if the tracefile is large and the user wants to extract a few representative steps automatically.

```python
python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz -o ./steady_state_analysis \
Expand Down Expand Up @@ -234,7 +261,7 @@ python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz \
--CONC 32 --OSL 1024 --R 0.8
```

Option 2: One tracefile per eager/graph execution step (supports vLLM v0.13 or higher, SGLang v0.5.9, and Atom 0.1.1). This is recommended if the user wants to perform analysis on an isolated execution step.
Option 2: One tracefile per eager/graph execution step (supports vLLM v0.13 or higher, SGLang v0.5.9–v0.5.12, and Atom 0.1.1). This is recommended if the user wants to perform analysis on an isolated execution step.

```python
python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz -o ./output --store-single-iteration
Expand Down
1 change: 1 addition & 0 deletions examples/custom_workflows/inference_analysis/build_docker_sglang.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ Base images:
0.5.12 MI355X : lmsysorg/sglang:v0.5.12-rocm720-mi35x

Examples:
$0 /path/to/TraceLens --sglang-version 0.5.12 --gpu-type mi300 -t tracelens-sglang:0.5.12-mi300
$0 /path/to/TraceLens --sglang-version 0.5.11 --gpu-type mi300 -t tracelens-sglang:0.5.11-mi300
$0 /path/to/TraceLens mi350 -t tracelens-sglang:0.5.9-mi350
EOF
Expand Down
48 changes: 26 additions & 22 deletions ...orkflows/inference_analysis/sglang_roofline_patches/sglang_0_5_12/cuda_graph_runner.patch
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py
index 05b64778a..f5717e2ad 100644
index c037b20dd..ffc70f35b 100644
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -84,6 +84,7 @@ from sglang.srt.utils import (
@@ -82,6 +82,7 @@ from sglang.srt.utils import (
require_mlp_sync,
require_mlp_tp_gather,
)
+from sglang.srt.utils.kernel_shape_profiler import disable, enable, is_enabled
from sglang.srt.utils.patch_torch import monkey_patch_torch_compile
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+from sglang.srt.utils.kernel_shape_profiler import enable, disable, is_enabled

try:
from kt_kernel import KTMoEWrapper
@@ -591,6 +592,9 @@ class CudaGraphRunner:
@@ -610,6 +611,9 @@ class CudaGraphRunner:
self.speculative_algorithm = model_runner.server_args.speculative_algorithm
self.enable_profile_cuda_graph = (
model_runner.server_args.enable_profile_cuda_graph
Expand All @@ -20,12 +20,14 @@ index 05b64778a..f5717e2ad 100644
)
self.tp_size = model_runner.server_args.tp_size
self.dp_size = model_runner.server_args.dp_size
@@ -812,9 +816,31 @@ class CudaGraphRunner:
@@ -834,9 +838,33 @@ class CudaGraphRunner:
)

def _init_profile_context_and_memory_record(self):
+ rank = get_tensor_model_parallel_rank()
+ trace_dir = os.path.join(os.environ.get("SGLANG_TORCH_PROFILER_DIR", "traces"),"capture_traces")
+ trace_dir = os.path.join(
+ os.environ.get("SGLANG_TORCH_PROFILER_DIR", "traces"), "capture_traces"
+ )
+ os.makedirs(trace_dir, exist_ok=True)
+
+ # Track which BS is currently being captured for trace file naming
Expand All @@ -52,7 +54,7 @@ index 05b64778a..f5717e2ad 100644
)
torch.cuda.memory._record_memory_history()
return profile_context
@@ -834,11 +860,14 @@ class CudaGraphRunner:
@@ -856,11 +884,14 @@ class CudaGraphRunner:
+ "\n\nMemory Usage is saved to cuda_graph_runner_memory_usage.pickle\n"
)
logger.info(log_message)
Expand All @@ -67,7 +69,7 @@ index 05b64778a..f5717e2ad 100644

def _capture_one_stream(stream_idx: Optional[int] = None):
avail_mem = get_available_gpu_memory(
@@ -907,7 +936,7 @@ class CudaGraphRunner:
@@ -929,7 +960,7 @@ class CudaGraphRunner:
_set_capture_lora_variant(None)

if self.enable_profile_cuda_graph:
Expand All @@ -76,38 +78,40 @@ index 05b64778a..f5717e2ad 100644

def _capture_graph(self, graph, pool, stream, run_once_fn):
if self.model_runner.server_args.debug_cuda_graph:
@@ -1185,15 +1214,28 @@ class CudaGraphRunner:
self.device_module.synchronize()
@@ -1186,15 +1217,30 @@ class CudaGraphRunner:
self.model_runner.tp_group.barrier()
run_once()
attn_backend.on_after_cuda_graph_warmup()
+ if self.enable_profile_cuda_graph:
+ # Activate kernel shape profiler AFTER warmup runs so that all
+ # lazily-imported modules (e.g. tilelang_kernel) are in sys.modules
+ # and auto-discovery can find them.
+ if self.enable_shape_discovery_for_cuda_graph_profile and not is_enabled():
+ if (
+ self.enable_shape_discovery_for_cuda_graph_profile
+ and not is_enabled()
+ ):
+ enable()
+ self._profiler.step()

- if get_global_graph_memory_pool() is None:
- set_global_graph_memory_pool(self.device_module.graph_pool_handle())
- # Set graph pool id globally to be able to use symmetric memory
- set_graph_pool_id(get_global_graph_memory_pool())
+
+ with torch.profiler.record_function(
+ f"capture_{num_tokens}_{self.capture_forward_mode.name}"
+ ):
+ if get_global_graph_memory_pool() is None:
+ set_global_graph_memory_pool(self.device_module.graph_pool_handle())
+ # Set graph pool id globally to be able to use symmetric memory
+ set_graph_pool_id(get_global_graph_memory_pool())
+ out = self._capture_graph(
+ graph, get_global_graph_memory_pool(), stream, run_once
+ )

- if get_global_graph_memory_pool() is None:
- set_global_graph_memory_pool(self.device_module.graph_pool_handle())
- # Set graph pool id globally to be able to use symmetric memory
- set_graph_pool_id(get_global_graph_memory_pool())
-
- out = self._capture_graph(
- graph, get_global_graph_memory_pool(), stream, run_once
- )
+ out = self._capture_graph(
+ graph, get_global_graph_memory_pool(), stream, run_once
+ )
+
+ if self.enable_profile_cuda_graph:
+ self._profiler.step()

Expand Down
Loading