[Ascend]adapt enable-profile-cuda-graph for NPU (#12617)

Co-authored-by: Leo920320 <1530593013@qq.com>
Co-authored-by: liupengcheng <6531487+leo920320@user.noreply.gitee.com>

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(
+            profile_context = self._init_profile_context_and_memory_record()
-                activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-                record_shapes=True,
-            )
-            torch.cuda.memory._record_memory_history()
 
         # 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")
+            self._post_process_after_profile(prof)
-            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)
 
     def _capture_graph(self, graph, pool, stream, run_once_fn):
         memory_saver_adapter = TorchMemorySaverAdapter.create(

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,