From 0217e82a08932c5cadf92d49b3e6192cb16fed8a Mon Sep 17 00:00:00 2001 From: Makcum888e <79456407+Makcum888e@users.noreply.github.com> Date: Wed, 25 Feb 2026 21:16:03 +0300 Subject: [PATCH] [diffusion] Clean code (#19325) --- .../runtime/distributed/parallel_state.py | 3 ++ .../runtime/layers/attention/turbo_layer.py | 10 +++--- .../runtime/pipelines/diffusers_pipeline.py | 13 ++++++-- .../runtime/pipelines_core/lora_pipeline.py | 8 ++--- .../runtime/platforms/interface.py | 2 +- .../runtime/utils/layerwise_offload.py | 33 +++++++++++-------- .../runtime/utils/perf_logger.py | 20 +++++------ .../test/scripts/gen_perf_baselines.py | 6 ++-- 8 files changed, 55 insertions(+), 40 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 18edea976..147cf1444 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py index 9f94d995e..3a1642559 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py index c0d6bbfa7..847394c1e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index 035ec489b..109991460 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/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 - if torch.cuda.is_available(): - torch.cuda.synchronize() - torch.cuda.empty_cache() + # clear device cache to free up unused memory + if torch.get_device_module().is_available(): + torch.get_device_module().synchronize() + torch.get_device_module().empty_cache() offload_disabled_modules = [] for module_name in module_names: diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 1623dcfe2..8c03fb324 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py index 0bdbd841d..b6b782a01 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/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.copy_stream = torch.cuda.Stream() + self.device = torch.device( + 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 - self._prefetch_events: Dict[int, torch.cuda.Event] = {} + # 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.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: diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 3f1c0429b..d44a50d37 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/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() - reserved = torch.cuda.memory_reserved() - peak_allocated = torch.cuda.max_memory_allocated() - peak_reserved = torch.cuda.max_memory_reserved() + allocated = torch.get_device_module().memory_allocated() + reserved = torch.get_device_module().memory_reserved() + peak_allocated = torch.get_device_module().max_memory_allocated() + 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 diff --git a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py b/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py index 4a3a385d0..182223486 100644 --- a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py +++ b/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(): - torch.cuda.synchronize() - torch.cuda.empty_cache() + if torch.get_device_module().is_available(): + torch.get_device_module().synchronize() + torch.get_device_module().empty_cache() except Exception: pass