chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] Clean code (#19325)

Browse Source
This commit is contained in:
Makcum888e
2026-02-25 21:16:03 +03:00
committed by GitHub
parent 2fb239450e
commit 0217e82a08
8 changed files with 55 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
View File

@@ -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:

10
python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py
View File

@@ -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

13
python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
View File

@@ -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

8
python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
View File

@@ -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:

2
python/sglang/multimodal_gen/runtime/platforms/interface.py
View File

@@ -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:

33
python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py
View File

@@ -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:

20
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
View File

@@ -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

6
python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py
View File

@@ -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