diff --git a/docs/Inference_analysis.md b/docs/Inference_analysis.md index 1ae4c4539..47bdfa87f 100644 --- a/docs/Inference_analysis.md +++ b/docs/Inference_analysis.md @@ -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. @@ -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 @@ -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 \ @@ -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 diff --git a/examples/custom_workflows/inference_analysis/build_docker_sglang.sh b/examples/custom_workflows/inference_analysis/build_docker_sglang.sh index 4c0b98e77..9e0a5b7cc 100755 --- a/examples/custom_workflows/inference_analysis/build_docker_sglang.sh +++ b/examples/custom_workflows/inference_analysis/build_docker_sglang.sh @@ -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 diff --git a/examples/custom_workflows/inference_analysis/sglang_roofline_patches/sglang_0_5_12/cuda_graph_runner.patch b/examples/custom_workflows/inference_analysis/sglang_roofline_patches/sglang_0_5_12/cuda_graph_runner.patch index 959648016..b5553f950 100644 --- a/examples/custom_workflows/inference_analysis/sglang_roofline_patches/sglang_0_5_12/cuda_graph_runner.patch +++ b/examples/custom_workflows/inference_analysis/sglang_roofline_patches/sglang_0_5_12/cuda_graph_runner.patch @@ -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 @@ -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 @@ -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) @@ -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: @@ -76,8 +78,7 @@ 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() @@ -85,14 +86,13 @@ index 05b64778a..f5717e2ad 100644 + # 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}" + ): @@ -100,14 +100,18 @@ index 05b64778a..f5717e2ad 100644 + 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()