[diffusion] Clean code (#19325)

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -639,6 +639,9 @@ def maybe_init_distributed_environment_and_model_parallel(
     if current_platform.is_cuda_alike():
         device = torch.device(f"cuda:{local_rank}")
         torch.cuda.set_device(device)
+    elif current_platform.is_npu():
+        device = torch.device(f"npu:{local_rank}")
+        torch.npu.set_device(device)
 
 
 def model_parallel_is_initialized() -> bool:

python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py

@@ -78,7 +78,7 @@ def async_a2a_communicate(
     a2a_inputs: Union[torch.Tensor, List[torch.Tensor]],
     cp_size: int,
     cp_group: ProcessGroup,
-    cp_stream: torch.cuda.Stream,
+    cp_stream: torch.get_device_module().Stream,
     local_seq_2_local_head: bool,
 ) -> Union[torch.Tensor, List[torch.Tensor]]:
     """
@@ -97,7 +97,7 @@ def async_a2a_communicate(
                 )
                 a2a_post_fns[i - 1] = post_all2all(local_seq_2_local_head, cp_size)
             if i > 1:
-                with torch.cuda.stream(cp_stream):
+                with torch.get_device_module().stream(cp_stream):
                     a2a_reqs[i - 2].wait()
                     a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2])
             if i < len(a2a_inputs):
@@ -117,10 +117,10 @@ def async_a2a_communicate(
                     a2a_inputs[i], "bs (w s) h d -> w bs s h d", w=cp_size
                 ).contiguous()
             if i > 1:
-                with torch.cuda.stream(cp_stream):
+                with torch.get_device_module().stream(cp_stream):
                     a2a_reqs[i - 2].wait()
                     a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2])
-    torch.cuda.current_stream().wait_stream(cp_stream)
+    torch.get_device_module().current_stream().wait_stream(cp_stream)
     return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs
 
 
@@ -152,7 +152,7 @@ class _SeqAllToAllQKV(torch.autograd.Function):
         k: Tensor,
         v: Tensor,
         cp_size: int,
-        cp_stream: torch.cuda.Stream,
+        cp_stream: torch.get_device_module().Stream,
         local_seq_2_local_head: bool,
     ) -> Tuple[Tensor, Tensor, Tensor]:
         ctx.group = group

python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py

@@ -33,7 +33,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.executors.sync_executor import
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -316,7 +319,7 @@ class DiffusersExecutionStage(PipelineStage):
                     return next(component.parameters()).device
                 except StopIteration:
                     pass
-        return "cuda" if torch.cuda.is_available() else "cpu"
+        return current_platform.device_type
 
     def _load_input_image(self, batch: Req) -> Image.Image | None:
         """Load input image from batch."""
@@ -565,7 +568,11 @@ class DiffusersPipeline(ComposedPipelineBase):
         """
         Determine the dtype to use for model loading.
         """
-        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        dtype = (
+            torch.bfloat16
+            if torch.get_device_module().is_bf16_supported()
+            else torch.float16
+        )
 
         if hasattr(server_args, "pipeline_config") and server_args.pipeline_config:
             dit_precision = server_args.pipeline_config.dit_precision

python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py

@@ -192,10 +192,10 @@ class LoRAPipeline(ComposedPipelineBase):
             yield []
             return
 
-        # clear CUDA cache to free up unused memory
+        # clear device cache to free up unused memory
-        if torch.cuda.is_available():
+        if torch.get_device_module().is_available():
-            torch.cuda.synchronize()
+            torch.get_device_module().synchronize()
-            torch.cuda.empty_cache()
+            torch.get_device_module().empty_cache()
 
         offload_disabled_modules = []
         for module_name in module_names:

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -315,7 +315,7 @@ class Platform:
             random.seed(seed)
             np.random.seed(seed)
             torch.manual_seed(seed)
-            torch.cuda.manual_seed_all(seed)
+            torch.get_device_module().manual_seed_all(seed)
 
     @classmethod
     def verify_model_arch(cls, model_arch: str) -> None:

python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Set, Tuple
 
 import torch
 
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -40,11 +41,13 @@ class LayerwiseOffloadManager:
         self.num_layers = num_layers
         self.pin_cpu_memory = pin_cpu_memory
         self.prefetch_size = min(max(1, prefetch_size), self.num_layers)
-        self.enabled = bool(enabled and torch.cuda.is_available())
+        self.enabled = bool(enabled and torch.get_device_module().is_available())
         if not self.enabled:
             return
-        self.device = torch.device("cuda", torch.cuda.current_device())
+        self.device = torch.device(
-        self.copy_stream = torch.cuda.Stream()
+            current_platform.device_type, torch.get_device_module().current_device()
+        )
+        self.copy_stream = torch.get_device_module().Stream()
 
         self._layer_name_re = re.compile(
             rf"(^|\.){re.escape(layers_attr_str)}\.(\d+)(\.|$)"
@@ -58,8 +61,8 @@ class LayerwiseOffloadManager:
         self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {}
         # layer indices that are already in gpu
         self._gpu_layers: Set[int] = set()
-        # layer_idx -> torch.cuda.Event for fine-grained sync, to make sure the weight is resident in pre-hook
+        # layer_idx -> torch.get_device_module().Event for fine-grained sync, to make sure the weight is resident in pre-hook
-        self._prefetch_events: Dict[int, torch.cuda.Event] = {}
+        self._prefetch_events: Dict[int, torch.get_device_module().Event] = {}
 
         self._named_parameters: Dict[str, torch.nn.Parameter] = {}
         self._named_buffers: Dict[str, torch.Tensor] = {}
@@ -144,7 +147,7 @@ class LayerwiseOffloadManager:
         for i in range(self.prefetch_size):
             self.prefetch_layer(i, non_blocking=non_blocking)
         if not non_blocking and self.copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.copy_stream)
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
 
     def get_target_with_name(self, name: str) -> torch.Tensor:
         """get the target model weight/buffer to be replaced"""
@@ -167,11 +170,11 @@ class LayerwiseOffloadManager:
             return
         if layer_idx not in self._consolidated_cpu_weights:
             return
-        self.copy_stream.wait_stream(torch.cuda.current_stream())
+        self.copy_stream.wait_stream(torch.get_device_module().current_stream())
 
         # create gpu buffer and load from CPU buffer
         gpu_buffers: Dict[torch.dtype, torch.Tensor] = {}
-        with torch.cuda.stream(self.copy_stream):
+        with torch.get_device_module().stream(self.copy_stream):
             for dtype, cpu_buffer in self._consolidated_cpu_weights[layer_idx].items():
                 gpu_buffer = torch.empty(
                     cpu_buffer.shape, dtype=dtype, device=self.device
@@ -180,7 +183,7 @@ class LayerwiseOffloadManager:
                 gpu_buffers[dtype] = gpu_buffer
 
         # record the prefetch event of this layer
-        event = torch.cuda.Event()
+        event = torch.get_device_module().Event()
         event.record(self.copy_stream)
         self._prefetch_events[layer_idx] = event
 
@@ -226,7 +229,7 @@ class LayerwiseOffloadManager:
         if not self.enabled or self.device is None:
             return
         if self.copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.copy_stream)
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
 
         for layer_idx in list(self._gpu_layers):
             self.release_layer(layer_idx)
@@ -237,7 +240,7 @@ class LayerwiseOffloadManager:
         if not self.enabled or self.device is None:
             return
         if self.copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.copy_stream)
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
 
         for layer_idx in range(self.num_layers):
             if layer_idx not in self._gpu_layers:
@@ -252,7 +255,7 @@ class LayerwiseOffloadManager:
             return
 
         if self.copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.copy_stream)
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
 
         # Collect current GPU weights and write back to CPU buffer
         for name, meta in self._weight_metadata.get(layer_idx, {}).items():
@@ -271,7 +274,7 @@ class LayerwiseOffloadManager:
         if not self.enabled or self.device is None:
             return
         if self.copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.copy_stream)
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
 
         for layer_idx in list(self._gpu_layers):
             self.sync_layer_to_cpu(layer_idx)
@@ -368,7 +371,9 @@ class LayerwiseOffloadManager:
                 if i == 0:
                     self.prepare_for_next_req(non_blocking=False)
                 if i in self._prefetch_events:
-                    torch.cuda.current_stream().wait_event(self._prefetch_events[i])
+                    torch.get_device_module().current_stream().wait_event(
+                        self._prefetch_events[i]
+                    )
 
                 # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed
                 if i % self.prefetch_size == 0:

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -120,7 +120,7 @@ def get_git_commit_hash() -> str:
 
 
 def capture_memory_snapshot() -> MemorySnapshot:
-    if not torch.cuda.is_available():
+    if not torch.get_device_module().is_available():
         return MemorySnapshot(
             allocated_mb=0.0,
             reserved_mb=0.0,
@@ -128,10 +128,10 @@ def capture_memory_snapshot() -> MemorySnapshot:
             peak_reserved_mb=0.0,
         )
 
-    allocated = torch.cuda.memory_allocated()
+    allocated = torch.get_device_module().memory_allocated()
-    reserved = torch.cuda.memory_reserved()
+    reserved = torch.get_device_module().memory_reserved()
-    peak_allocated = torch.cuda.max_memory_allocated()
+    peak_allocated = torch.get_device_module().max_memory_allocated()
-    peak_reserved = torch.cuda.max_memory_reserved()
+    peak_reserved = torch.get_device_module().max_memory_reserved()
 
     return MemorySnapshot(
         allocated_mb=allocated / (1024**2),
@@ -212,9 +212,9 @@ class StageProfiler:
             if (
                 os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
                 and self.stage_name.startswith("denoising_step_")
-                and torch.cuda.is_available()
+                and torch.get_device_module().is_available()
             ):
-                torch.cuda.synchronize()
+                torch.get_device_module().synchronize()
             self.start_time = time.perf_counter()
 
         return self
@@ -226,9 +226,9 @@ class StageProfiler:
         if (
             os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
             and self.stage_name.startswith("denoising_step_")
-            and torch.cuda.is_available()
+            and torch.get_device_module().is_available()
         ):
-            torch.cuda.synchronize()
+            torch.get_device_module().synchronize()
         execution_time_s = time.perf_counter() - self.start_time
 
         if exc_type:
@@ -254,7 +254,7 @@ class StageProfiler:
                 self.metrics.record_stage(self.stage_name, execution_time_s)
 
             # capture memory snapshot after stage if requested
-            if self.capture_memory and torch.cuda.is_available():
+            if self.capture_memory and torch.get_device_module().is_available():
                 snapshot = capture_memory_snapshot()
                 self.metrics.record_memory_snapshot(
                     f"after_{self.stage_name}", snapshot

python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py

@@ -90,9 +90,9 @@ def _torch_cleanup() -> None:
     try:
         import torch
 
-        if torch.cuda.is_available():
+        if torch.get_device_module().is_available():
-            torch.cuda.synchronize()
+            torch.get_device_module().synchronize()
-            torch.cuda.empty_cache()
+            torch.get_device_module().empty_cache()
     except Exception:
         pass