diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 58e95ba6a..3bdc16684 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -449,14 +449,34 @@ class CudaGraphRunner: and is_ngram_supported ) + def _init_profile_context_and_memory_record(self): + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + ) + torch.cuda.memory._record_memory_history() + return profile_context + + def _post_process_after_profile(self, prof_context): + torch.cuda.memory._dump_snapshot(f"cuda_graph_runner_memory_usage.pickle") + torch.cuda.memory._record_memory_history(enabled=None) + log_message = ( + "Sorted by CUDA Time:\n" + + prof_context.key_averages(group_by_input_shape=True).table( + sort_by="cuda_time_total", row_limit=10 + ) + + "\n\nSorted by CPU Time:\n" + + prof_context.key_averages(group_by_input_shape=True).table( + sort_by="cpu_time_total", row_limit=10 + ) + + "\n\nMemory Usage is saved to cuda_graph_runner_memory_usage.pickle\n" + ) + logger.info(log_message) + def capture(self) -> None: profile_context = empty_context() if self.enable_profile_cuda_graph: - profile_context = profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=True, - ) - torch.cuda.memory._record_memory_history() + profile_context = self._init_profile_context_and_memory_record() # Trigger CUDA graph capture for specific shapes. # Capture the large shapes first so that the smaller shapes @@ -505,20 +525,7 @@ class CudaGraphRunner: save_gemlite_cache() if self.enable_profile_cuda_graph: - torch.cuda.memory._dump_snapshot(f"cuda_graph_runner_memory_usage.pickle") - torch.cuda.memory._record_memory_history(enabled=None) - log_message = ( - "Sorted by CUDA Time:\n" - + prof.key_averages(group_by_input_shape=True).table( - sort_by="cuda_time_total", row_limit=10 - ) - + "\n\nSorted by CPU Time:\n" - + prof.key_averages(group_by_input_shape=True).table( - sort_by="cpu_time_total", row_limit=10 - ) - + "\n\nMemory Usage is saved to cuda_graph_runner_memory_usage.pickle\n" - ) - logger.info(log_message) + self._post_process_after_profile(prof) def _capture_graph(self, graph, pool, stream, run_once_fn): memory_saver_adapter = TorchMemorySaverAdapter.create( diff --git a/python/sglang/srt/model_executor/npu_graph_runner.py b/python/sglang/srt/model_executor/npu_graph_runner.py index f9961a36f..c42496280 100644 --- a/python/sglang/srt/model_executor/npu_graph_runner.py +++ b/python/sglang/srt/model_executor/npu_graph_runner.py @@ -16,13 +16,22 @@ from __future__ import annotations import logging +import os import threading +from pathlib import Path from typing import TYPE_CHECKING, Optional, Union import torch from sglang.srt.configs.model_config import is_deepseek_nsa from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner +from sglang.srt.utils import is_npu + +is_npu = is_npu() + +if is_npu: + import torch_npu + from torch_npu.profiler import ProfilerActivity, profile logger = logging.getLogger(__name__) @@ -60,6 +69,34 @@ class NPUGraphRunner(CudaGraphRunner): def _cache_loc_dtype(self): return torch.int32 + def _init_profile_context_and_memory_record(self): + output_dir = os.path.join( + os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp"), "graph_capture_profile" + ) + if not Path(output_dir).exists(): + Path(output_dir).mkdir(parents=True, exist_ok=True) + logger.info( + f"Profiling starts for graph capture for NPU. Traces will be saved to: {output_dir}" + ) + experimental_config = torch_npu.profiler._ExperimentalConfig( + export_type=[torch_npu.profiler.ExportType.Text], + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + ) + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.NPU], + record_shapes=True, + profile_memory=True, + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler( + output_dir, async_mode=True + ), + experimental_config=experimental_config, + ) + return profile_context + + def _post_process_after_profile(self, prof_context): + # for NPU, profile data will be saved to disk for further analysis. + pass + def replay( self, forward_batch: ForwardBatch,