From d07e8aa4a31520c0ead8b8062d6958dbc4e6ea98 Mon Sep 17 00:00:00 2001
From: Makcum888e <79456407+Makcum888e@users.noreply.github.com>
Date: Thu, 19 Feb 2026 15:33:51 +0300
Subject: [PATCH] [Diffusion] [NPU] Enable profiler on NPU (#17807)

---
 .../multimodal_gen/runtime/utils/profiler.py  | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py
index 3af3dd9bb..d01a51f86 100644
--- a/python/sglang/multimodal_gen/runtime/utils/profiler.py
+++ b/python/sglang/multimodal_gen/runtime/utils/profiler.py
@@ -3,8 +3,18 @@ import os
 
 import torch
 
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
 
+if current_platform.is_npu():
+    import torch_npu
+
+    patches = [
+        ["profiler.profile", torch_npu.profiler.profile],
+        ["profiler.schedule", torch_npu.profiler.schedule],
+    ]
+    torch_npu._apply_patches(patches)
+
 logger = init_logger(__name__)
 
 
@@ -47,12 +57,18 @@ class SGLDiffusionProfiler:
         activities = [torch.profiler.ProfilerActivity.CPU]
         if torch.cuda.is_available():
             activities.append(torch.profiler.ProfilerActivity.CUDA)
+        if current_platform.is_npu():
+            activities.append(torch_npu.profiler.ProfilerActivity.NPU)
 
         common_torch_profiler_args = dict(
             activities=activities,
             record_shapes=True,
             with_stack=True,
-            on_trace_ready=None,
+            on_trace_ready=(
+                None
+                if not current_platform.is_npu()
+                else torch_npu.profiler.tensorboard_trace_handler(self.log_dir)
+            ),
         )
         if self.full_profile:
             # profile all stages
@@ -113,6 +129,9 @@ class SGLDiffusionProfiler:
         logger.info("Stopping Profiler...")
         if torch.cuda.is_available():
             torch.cuda.synchronize()
+        if current_platform.is_npu():
+            torch.npu.synchronize()
+            export_trace = False  # set to false because our internal torch_npu.profiler will generate trace file
         self.profiler.stop()
 
         if export_trace: