diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index c54b5cbd6..a5a98981b 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -1,16 +1,11 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo -import importlib.util # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py + import logging import os -from collections.abc import Callable -from typing import TYPE_CHECKING, Any - -import diffusers -import torch -from packaging import version +from typing import TYPE_CHECKING, Any, Callable from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var @@ -58,105 +53,8 @@ if TYPE_CHECKING: SGLANG_CACHE_DIT_SECONDARY_MC: int = 3 SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER: bool = False SGLANG_CACHE_DIT_SECONDARY_TS_ORDER: int = 1 - - -def _is_hip(): - has_rocm = torch.version.hip is not None - return has_rocm - - -def _is_cuda(): - has_cuda = torch.version.cuda is not None - return has_cuda - - -def _is_musa(): - try: - if hasattr(torch, "musa") and torch.musa.is_available(): - return True - except ModuleNotFoundError: - return False - - -def _is_mps(): - return torch.backends.mps.is_available() - - -class PackagesEnvChecker: - _instance = None - - def __new__(cls): - if cls._instance is None: - cls._instance = super(PackagesEnvChecker, cls).__new__(cls) - cls._instance.initialize() - return cls._instance - - def initialize(self): - self.packages_info = { - "has_aiter": self.check_aiter(), - "diffusers_version": self.check_diffusers_version(), - } - - def check_aiter(self): - """ - Checks whether ROCm AITER library is installed - """ - try: - - logger.info("Using AITER as the attention library") - return True - except: - if _is_hip(): - logger.warning( - f'Using AMD GPUs, but library "aiter" is not installed, ' - "defaulting to other attention mechanisms" - ) - return False - - def check_flash_attn(self): - if not torch.cuda.is_available(): - return False - if _is_musa(): - logger.info( - "Flash Attention library is not supported on MUSA for the moment." - ) - return False - try: - return True - except ImportError: - logger.warning( - f'Flash Attention library "flash_attn" not found, ' - f"using pytorch attention implementation" - ) - return False - - def check_long_ctx_attn(self): - if not torch.cuda.is_available(): - return False - try: - return importlib.util.find_spec("yunchang") is not None - except ImportError: - logger.warning( - f'Ring Flash Attention library "yunchang" not found, ' - f"using pytorch attention implementation" - ) - return False - - def check_diffusers_version(self): - if version.parse( - version.parse(diffusers.__version__).base_version - ) < version.parse("0.30.0"): - raise RuntimeError( - f"Diffusers version: {version.parse(version.parse(diffusers.__version__).base_version)} is not supported," - f"please upgrade to version > 0.30.0" - ) - return version.parse(version.parse(diffusers.__version__).base_version) - - def get_packages_info(self): - return self.packages_info - - -PACKAGES_CHECKER = PackagesEnvChecker() + # model loading + SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True def get_default_cache_root() -> str: @@ -174,9 +72,58 @@ def get_default_config_root() -> str: def maybe_convert_int(value: str | None) -> int | None: - if value is None: - return None - return int(value) + return int(value) if value is not None else None + + +# helpers for environment variable definitions +def _lazy_str(key: str, default: str | None = None) -> Callable[[], str | None]: + return lambda: os.getenv(key, default) + + +def _lazy_int(key: str, default: str | int | None = None) -> Callable[[], int | None]: + def _getter(): + val = os.getenv(key) + if val is None: + return int(default) if default is not None else None + return int(val) + + return _getter + + +def _lazy_float(key: str, default: str | float) -> Callable[[], float]: + return lambda: float(os.getenv(key, str(default))) + + +def _lazy_bool(key: str, default: str = "false") -> Callable[[], bool]: + return lambda: get_bool_env_var(key, default) + + +def _lazy_bool_any(keys: list[str], default: str = "false") -> Callable[[], bool]: + def _getter(): + for key in keys: + if get_bool_env_var(key, "false"): + return True + return ( + get_bool_env_var("", default) + if not keys + else get_bool_env_var(keys[0], default) + ) + + return _getter + + +def _lazy_path( + key: str, default_func: Callable[[], str] | None = None +) -> Callable[[], str | None]: + def _getter(): + val = os.getenv(key) + if val is None: + if default_func is None: + return None + val = default_func() + return os.path.expanduser(val) + + return _getter # The begin-* and end* here are used by the documentation generator @@ -188,220 +135,188 @@ environment_variables: dict[str, Callable[[], Any]] = { # ================== Installation Time Env Vars ================== # Target device of sglang-diffusion, supporting [cuda (by default), # rocm, neuron, cpu, openvino] - "SGLANG_DIFFUSION_TARGET_DEVICE": lambda: os.getenv( + "SGLANG_DIFFUSION_TARGET_DEVICE": _lazy_str( "SGLANG_DIFFUSION_TARGET_DEVICE", "cuda" ), # Maximum number of compilation jobs to run in parallel. # By default this is the number of CPUs - "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None), + "MAX_JOBS": _lazy_str("MAX_JOBS"), # Number of threads to use for nvcc # By default this is 1. # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. - "NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None), + "NVCC_THREADS": _lazy_str("NVCC_THREADS"), # If set, sgl_diffusion will use precompiled binaries (*.so) - "SGLANG_DIFFUSION_USE_PRECOMPILED": lambda: bool( - os.environ.get("SGLANG_DIFFUSION_USE_PRECOMPILED") - ) - or bool(os.environ.get("SGLANG_DIFFUSION_PRECOMPILED_WHEEL_LOCATION")), + "SGLANG_DIFFUSION_USE_PRECOMPILED": _lazy_bool_any( + [ + "SGLANG_DIFFUSION_USE_PRECOMPILED", + "SGLANG_DIFFUSION_PRECOMPILED_WHEEL_LOCATION", + ] + ), # CMake build type # If not set, defaults to "Debug" or "RelWithDebInfo" # Available options: "Debug", "Release", "RelWithDebInfo" - "CMAKE_BUILD_TYPE": lambda: os.getenv("CMAKE_BUILD_TYPE"), + "CMAKE_BUILD_TYPE": _lazy_str("CMAKE_BUILD_TYPE"), # If set, sgl_diffusion will print verbose logs during installation - "VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))), - # Root directory for FASTVIDEO configuration files + "VERBOSE": _lazy_bool("VERBOSE"), + # Root directory for SGL-diffusion configuration files # Defaults to `~/.config/sgl_diffusion` unless `XDG_CONFIG_HOME` is set # Note that this not only affects how sgl_diffusion finds its configuration files # during runtime, but also affects how sgl_diffusion installs its configuration # files during **installation**. - "SGLANG_DIFFUSION_CONFIG_ROOT": lambda: os.path.expanduser( - os.getenv( - "SGLANG_DIFFUSION_CONFIG_ROOT", - os.path.join(get_default_config_root(), "sgl_diffusion"), - ) + "SGLANG_DIFFUSION_CONFIG_ROOT": _lazy_path( + "SGLANG_DIFFUSION_CONFIG_ROOT", + lambda: os.path.join(get_default_config_root(), "sgl_diffusion"), ), # ================== Runtime Env Vars ================== - # Root directory for FASTVIDEO cache files + # Root directory for SGL-diffusion cache files # Defaults to `~/.cache/sgl_diffusion` unless `XDG_CACHE_HOME` is set - "SGLANG_DIFFUSION_CACHE_ROOT": lambda: os.path.expanduser( - os.getenv( - "SGLANG_DIFFUSION_CACHE_ROOT", - os.path.join(get_default_cache_root(), "sgl_diffusion"), - ) + "SGLANG_DIFFUSION_CACHE_ROOT": _lazy_path( + "SGLANG_DIFFUSION_CACHE_ROOT", + lambda: os.path.join(get_default_cache_root(), "sgl_diffusion"), ), # Interval in seconds to log a warning message when the ring buffer is full - "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL": lambda: int( - os.environ.get("SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL", "60") + "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL": _lazy_int( + "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL", 60 ), # Path to the NCCL library file. It is needed because nccl>=2.19 brought # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 - "SGLANG_DIFFUSION_NCCL_SO_PATH": lambda: os.environ.get( - "SGLANG_DIFFUSION_NCCL_SO_PATH", None - ), + "SGLANG_DIFFUSION_NCCL_SO_PATH": _lazy_str("SGLANG_DIFFUSION_NCCL_SO_PATH"), # when `SGLANG_DIFFUSION_NCCL_SO_PATH` is not set, sgl_diffusion will try to find the nccl # library file in the locations specified by `LD_LIBRARY_PATH` - "LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None), + "LD_LIBRARY_PATH": _lazy_str("LD_LIBRARY_PATH"), # Internal flag to enable Dynamo fullgraph capture - "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": lambda: bool( - os.environ.get("SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0" + "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": _lazy_bool( + "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1" ), # local rank of the process in the distributed setting, used to determine # the GPU device id - "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + "LOCAL_RANK": _lazy_int("LOCAL_RANK", 0), # used to control the visible devices in the distributed setting - "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), + "CUDA_VISIBLE_DEVICES": _lazy_str("CUDA_VISIBLE_DEVICES"), # timeout for each iteration in the engine - "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": lambda: int( - os.environ.get("SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S", "60") + "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": _lazy_int( + "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S", 60 ), # Logging configuration # If set to 0, sgl_diffusion will not configure logging # If set to 1, sgl_diffusion will configure logging using the default configuration # or the configuration file specified by SGLANG_DIFFUSION_LOGGING_CONFIG_PATH - "SGLANG_DIFFUSION_CONFIGURE_LOGGING": lambda: int( - os.getenv("SGLANG_DIFFUSION_CONFIGURE_LOGGING", "1") + "SGLANG_DIFFUSION_CONFIGURE_LOGGING": _lazy_int( + "SGLANG_DIFFUSION_CONFIGURE_LOGGING", 1 ), - "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH": lambda: os.getenv( + "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH": _lazy_str( "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH" ), # this is used for configuring the default logging level - "SGLANG_DIFFUSION_LOGGING_LEVEL": lambda: os.getenv( + "SGLANG_DIFFUSION_LOGGING_LEVEL": _lazy_str( "SGLANG_DIFFUSION_LOGGING_LEVEL", "INFO" ), # if set, SGLANG_DIFFUSION_LOGGING_PREFIX will be prepended to all log messages - "SGLANG_DIFFUSION_LOGGING_PREFIX": lambda: os.getenv( - "SGLANG_DIFFUSION_LOGGING_PREFIX", "" - ), + "SGLANG_DIFFUSION_LOGGING_PREFIX": _lazy_str("SGLANG_DIFFUSION_LOGGING_PREFIX", ""), # Trace function calls # If set to 1, sgl_diffusion will trace function calls # Useful for debugging - "SGLANG_DIFFUSION_TRACE_FUNCTION": lambda: int( - os.getenv("SGLANG_DIFFUSION_TRACE_FUNCTION", "0") - ), + "SGLANG_DIFFUSION_TRACE_FUNCTION": _lazy_int("SGLANG_DIFFUSION_TRACE_FUNCTION", 0), # Path to the attention configuration file. Only used for sliding tile # attention for now. - "SGLANG_DIFFUSION_ATTENTION_CONFIG": lambda: ( - None - if os.getenv("SGLANG_DIFFUSION_ATTENTION_CONFIG", None) is None - else os.path.expanduser(os.getenv("SGLANG_DIFFUSION_ATTENTION_CONFIG", ".")) + "SGLANG_DIFFUSION_ATTENTION_CONFIG": _lazy_path( + "SGLANG_DIFFUSION_ATTENTION_CONFIG" ), # Optional override to force a specific attention backend (e.g. "aiter") - "SGLANG_DIFFUSION_ATTENTION_BACKEND": lambda: os.getenv( + "SGLANG_DIFFUSION_ATTENTION_BACKEND": _lazy_str( "SGLANG_DIFFUSION_ATTENTION_BACKEND" ), # Use dedicated multiprocess context for workers. # Both spawn and fork work - "SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD": lambda: os.getenv( + "SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD": _lazy_str( "SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD", "fork" ), # Enables torch profiler if set. Path to the directory where torch profiler # traces are saved. Note that it must be an absolute path. - "SGLANG_DIFFUSION_TORCH_PROFILER_DIR": lambda: ( - None - if os.getenv("SGLANG_DIFFUSION_TORCH_PROFILER_DIR", None) is None - else os.path.expanduser(os.getenv("SGLANG_DIFFUSION_TORCH_PROFILER_DIR", ".")) + "SGLANG_DIFFUSION_TORCH_PROFILER_DIR": _lazy_path( + "SGLANG_DIFFUSION_TORCH_PROFILER_DIR" ), # If set, sgl_diffusion will run in development mode, which will enable # some additional endpoints for developing and debugging, # e.g. `/reset_prefix_cache` - "SGLANG_DIFFUSION_SERVER_DEV_MODE": lambda: get_bool_env_var( - "SGLANG_DIFFUSION_SERVER_DEV_MODE" - ), + "SGLANG_DIFFUSION_SERVER_DEV_MODE": _lazy_bool("SGLANG_DIFFUSION_SERVER_DEV_MODE"), # If set, sgl_diffusion will enable stage logging, which will print the time # taken for each stage - "SGLANG_DIFFUSION_STAGE_LOGGING": lambda: get_bool_env_var( - "SGLANG_DIFFUSION_STAGE_LOGGING" - ), + "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"), # ================== cache-dit Env Vars ================== # Enable cache-dit acceleration for DiT inference - "SGLANG_CACHE_DIT_ENABLED": lambda: get_bool_env_var("SGLANG_CACHE_DIT_ENABLED"), + "SGLANG_CACHE_DIT_ENABLED": _lazy_bool("SGLANG_CACHE_DIT_ENABLED"), # Number of first blocks to always compute (DBCache F parameter) - "SGLANG_CACHE_DIT_FN": lambda: int(os.getenv("SGLANG_CACHE_DIT_FN", "1")), + "SGLANG_CACHE_DIT_FN": _lazy_int("SGLANG_CACHE_DIT_FN", 1), # Number of last blocks to always compute (DBCache B parameter) - "SGLANG_CACHE_DIT_BN": lambda: int(os.getenv("SGLANG_CACHE_DIT_BN", "0")), + "SGLANG_CACHE_DIT_BN": _lazy_int("SGLANG_CACHE_DIT_BN", 0), # Warmup steps before caching (DBCache W parameter) - "SGLANG_CACHE_DIT_WARMUP": lambda: int(os.getenv("SGLANG_CACHE_DIT_WARMUP", "4")), + "SGLANG_CACHE_DIT_WARMUP": _lazy_int("SGLANG_CACHE_DIT_WARMUP", 4), # Residual difference threshold (DBCache R parameter) - "SGLANG_CACHE_DIT_RDT": lambda: float(os.getenv("SGLANG_CACHE_DIT_RDT", "0.24")), + "SGLANG_CACHE_DIT_RDT": _lazy_float("SGLANG_CACHE_DIT_RDT", 0.24), # Maximum continuous cached steps (DBCache MC parameter) - "SGLANG_CACHE_DIT_MC": lambda: int(os.getenv("SGLANG_CACHE_DIT_MC", "3")), + "SGLANG_CACHE_DIT_MC": _lazy_int("SGLANG_CACHE_DIT_MC", 3), # Enable TaylorSeer calibrator - "SGLANG_CACHE_DIT_TAYLORSEER": lambda: get_bool_env_var( - "SGLANG_CACHE_DIT_TAYLORSEER", default="false" - ), + "SGLANG_CACHE_DIT_TAYLORSEER": _lazy_bool("SGLANG_CACHE_DIT_TAYLORSEER", "false"), # TaylorSeer order (1 or 2) - "SGLANG_CACHE_DIT_TS_ORDER": lambda: int( - os.getenv("SGLANG_CACHE_DIT_TS_ORDER", "1") - ), + "SGLANG_CACHE_DIT_TS_ORDER": _lazy_int("SGLANG_CACHE_DIT_TS_ORDER", 1), # SCM preset: none, slow, medium, fast, ultra - "SGLANG_CACHE_DIT_SCM_PRESET": lambda: os.getenv( - "SGLANG_CACHE_DIT_SCM_PRESET", "none" - ), + "SGLANG_CACHE_DIT_SCM_PRESET": _lazy_str("SGLANG_CACHE_DIT_SCM_PRESET", "none"), # SCM custom compute bins (e.g., "8,3,3,2,2") - "SGLANG_CACHE_DIT_SCM_COMPUTE_BINS": lambda: os.getenv( - "SGLANG_CACHE_DIT_SCM_COMPUTE_BINS", None - ), + "SGLANG_CACHE_DIT_SCM_COMPUTE_BINS": _lazy_str("SGLANG_CACHE_DIT_SCM_COMPUTE_BINS"), # SCM custom cache bins (e.g., "1,2,2,2,3") - "SGLANG_CACHE_DIT_SCM_CACHE_BINS": lambda: os.getenv( - "SGLANG_CACHE_DIT_SCM_CACHE_BINS", None - ), + "SGLANG_CACHE_DIT_SCM_CACHE_BINS": _lazy_str("SGLANG_CACHE_DIT_SCM_CACHE_BINS"), # SCM policy: dynamic or static - "SGLANG_CACHE_DIT_SCM_POLICY": lambda: os.getenv( - "SGLANG_CACHE_DIT_SCM_POLICY", "dynamic" - ), - # ================== cache-dit Secondary Transformer Env Vars ================== - # For dual-transformer models like Wan2.2 (high-noise + low-noise experts) - # These parameters configure the secondary transformer (transformer_2) - # If not set, they inherit from the primary transformer settings - # Number of first blocks to always compute for secondary transformer - "SGLANG_CACHE_DIT_SECONDARY_FN": lambda: int( - os.getenv( - "SGLANG_CACHE_DIT_SECONDARY_FN", os.getenv("SGLANG_CACHE_DIT_FN", "1") - ) - ), - # Number of last blocks to always compute for secondary transformer - "SGLANG_CACHE_DIT_SECONDARY_BN": lambda: int( - os.getenv( - "SGLANG_CACHE_DIT_SECONDARY_BN", os.getenv("SGLANG_CACHE_DIT_BN", "0") - ) - ), - # Warmup steps before caching for secondary transformer - "SGLANG_CACHE_DIT_SECONDARY_WARMUP": lambda: int( - os.getenv( - "SGLANG_CACHE_DIT_SECONDARY_WARMUP", - os.getenv("SGLANG_CACHE_DIT_WARMUP", "4"), - ) - ), - # Residual difference threshold for secondary transformer - "SGLANG_CACHE_DIT_SECONDARY_RDT": lambda: float( - os.getenv( - "SGLANG_CACHE_DIT_SECONDARY_RDT", os.getenv("SGLANG_CACHE_DIT_RDT", "0.24") - ) - ), - # Maximum continuous cached steps for secondary transformer - "SGLANG_CACHE_DIT_SECONDARY_MC": lambda: int( - os.getenv( - "SGLANG_CACHE_DIT_SECONDARY_MC", os.getenv("SGLANG_CACHE_DIT_MC", "3") - ) - ), - # Enable TaylorSeer for secondary transformer - "SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER": lambda: get_bool_env_var( - "SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER", - default=os.getenv("SGLANG_CACHE_DIT_TAYLORSEER", "false"), - ), - # TaylorSeer order for secondary transformer - "SGLANG_CACHE_DIT_SECONDARY_TS_ORDER": lambda: int( - os.getenv( - "SGLANG_CACHE_DIT_SECONDARY_TS_ORDER", - os.getenv("SGLANG_CACHE_DIT_TS_ORDER", "1"), - ) + "SGLANG_CACHE_DIT_SCM_POLICY": _lazy_str("SGLANG_CACHE_DIT_SCM_POLICY", "dynamic"), + # model loading + "SGLANG_USE_RUNAI_MODEL_STREAMER": _lazy_bool( + "SGLANG_USE_RUNAI_MODEL_STREAMER", "true" ), } +# Add cache-dit Secondary Transformer Env Vars via programmatic generation to reduce duplication +_CACHE_DIT_SECONDARY_CONFIGS = [ + ("FN", int, "1"), + ("BN", int, "0"), + ("WARMUP", int, "4"), + ("RDT", float, "0.24"), + ("MC", int, "3"), + ("TS_ORDER", int, "1"), +] + + +def _create_secondary_getter(suffix, type_func, default_val): + primary_key = f"SGLANG_CACHE_DIT_{suffix}" + secondary_key = f"SGLANG_CACHE_DIT_SECONDARY_{suffix}" + + def _getter(): + val = os.getenv(secondary_key) + if val is not None: + return type_func(val) + return type_func(os.getenv(primary_key, str(default_val))) + + return secondary_key, _getter + + +for suffix, type_func, default_val in _CACHE_DIT_SECONDARY_CONFIGS: + key, getter = _create_secondary_getter(suffix, type_func, default_val) + environment_variables[key] = getter + + +# Special handling for boolean secondary var (TaylorSeer) +def _secondary_taylorseer_getter(): + return get_bool_env_var( + "SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER", + default=os.getenv("SGLANG_CACHE_DIT_TAYLORSEER", "false"), + ) + + +environment_variables["SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER"] = ( + _secondary_taylorseer_getter +) + # end-env-vars-definition - - def __getattr__(name: str): # lazy evaluation of environment variables if name in environment_variables: @@ -411,27 +326,3 @@ def __getattr__(name: str): def __dir__(): return list(environment_variables.keys()) - - -def get_torch_distributed_backend() -> str: - if torch.cuda.is_available(): - return "nccl" - elif _is_musa(): - return "mccl" - elif _is_mps(): - return "gloo" - else: - raise NotImplementedError( - "No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available" - ) - - -def get_device(local_rank: int) -> torch.device: - if torch.cuda.is_available(): - return torch.device("cuda", local_rank) - elif _is_musa(): - return torch.device("musa", local_rank) - elif _is_mps(): - return torch.device("mps") - else: - return torch.device("cpu") diff --git a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py index cab45f2a5..d9915fd8c 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py @@ -23,6 +23,7 @@ from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_ from sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator import ( CpuCommunicator, ) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import ( init_logger, suppress_stdout, @@ -44,7 +45,6 @@ _group_name_counter: dict[str, int] = {} def get_local_torch_device() -> torch.device: """Return the torch device for the current rank.""" - from sglang.multimodal_gen.runtime.platforms import current_platform return ( torch.device(f"cuda:{envs.LOCAL_RANK}") @@ -846,7 +846,7 @@ class PipelineGroupCoordinator(GroupCoordinator): assert self.cpu_group is not None assert self.device_group is not None - self.device = envs.get_device(local_rank) + self.device = current_platform.get_device(local_rank) self.recv_buffer_set: bool = False self.recv_tasks_queue: List[Tuple[str, int]] = [] diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 683b75027..57d418c8d 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -341,7 +341,9 @@ def initialize_model_parallel( """ if backend is None: - backend = envs.get_torch_distributed_backend() + from sglang.multimodal_gen.runtime.platforms import current_platform + + backend = current_platform.get_torch_distributed_backend_str() # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() diff --git a/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/python/sglang/multimodal_gen/runtime/layers/custom_op.py index 32cfe3aba..f4cd030d6 100644 --- a/python/sglang/multimodal_gen/runtime/layers/custom_op.py +++ b/python/sglang/multimodal_gen/runtime/layers/custom_op.py @@ -8,23 +8,11 @@ from typing import Any import torch.nn as nn -from sglang.multimodal_gen.runtime.utils.common import ( - is_cpu, - is_cuda, - is_hip, - is_npu, - is_xpu, -) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) -_is_cuda = is_cuda() -_is_hip = is_hip() -_is_cpu = is_cpu() -_is_npu = is_npu() -_is_xpu = is_xpu() - class CustomOp(nn.Module): """ @@ -70,13 +58,13 @@ class CustomOp(nn.Module): return self.forward_native(*args, **kwargs) def dispatch_forward(self) -> Callable: - if _is_cuda: + if current_platform.is_cuda(): return self.forward_cuda - elif _is_hip: + elif current_platform.is_hip(): return self.forward_hip - elif _is_npu: + elif current_platform.is_npu(): return self.forward_npu - elif _is_xpu: + elif current_platform.is_xpu(): return self.forward_xpu else: return self.forward_native diff --git a/python/sglang/multimodal_gen/runtime/platforms/__init__.py b/python/sglang/multimodal_gen/runtime/platforms/__init__.py index c87fb3aa9..ee515bb24 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/__init__.py +++ b/python/sglang/multimodal_gen/runtime/platforms/__init__.py @@ -139,8 +139,7 @@ def resolve_current_platform_cls_qualname() -> str: _current_platform: Platform | None = None _init_trace: str = "" -if TYPE_CHECKING: - current_platform: Platform +current_platform: Platform def __getattr__(name: str): @@ -150,12 +149,6 @@ def __getattr__(name: str): # Platform` so that they can inherit `Platform` class. Therefore, # we cannot resolve `current_platform` during the import of # `sglang.multimodal_gen.runtime.platforms`. - # 2. when users use out-of-tree platform plugins, they might run - # `import sgl_diffusion`, some sgl_diffusion internal code might access - # `current_platform` during the import, and we need to make sure - # `current_platform` is only resolved after the plugins are loaded - # (we have tests for this, if any developer violate this, they will - # see the test failures). global _current_platform if _current_platform is None: platform_cls_qualname = resolve_current_platform_cls_qualname() diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index f8284dc9f..136ea29b3 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -20,7 +20,6 @@ from sglang.multimodal_gen.runtime.platforms.interface import ( Platform, PlatformEnum, ) -from sglang.multimodal_gen.runtime.utils.common import is_blackwell, is_sm120 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import import_pynvml @@ -213,7 +212,7 @@ class CudaPlatformBase(Platform): elif selected_backend in [ AttentionBackendEnum.FA, ]: - if is_blackwell(): + if cls.is_blackwell(): from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( set_fa_ver, ) @@ -224,14 +223,14 @@ class CudaPlatformBase(Platform): raise ValueError(f"Invalid attention backend for {cls.device_name}") else: - if is_blackwell(): + if cls.is_blackwell(): from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( set_fa_ver, ) set_fa_ver(4) target_backend = AttentionBackendEnum.FA - if is_sm120(): + if cls.is_sm120(): try: from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401 SageAttention3Backend, diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 19f80f85c..82cff8624 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -6,6 +6,7 @@ from __future__ import annotations import enum import random +from functools import lru_cache from typing import TYPE_CHECKING, NamedTuple import numpy as np @@ -90,28 +91,79 @@ class Platform: supported_quantization: list[str] = [] + @lru_cache(maxsize=1) def is_cuda(self) -> bool: - return self._enum == PlatformEnum.CUDA + return self.is_cuda_static() + @lru_cache(maxsize=1) def is_rocm(self) -> bool: - return self._enum == PlatformEnum.ROCM + return self.is_rocm_static() + @lru_cache(maxsize=1) def is_tpu(self) -> bool: return self._enum == PlatformEnum.TPU + @lru_cache(maxsize=1) def is_cpu(self) -> bool: return self._enum == PlatformEnum.CPU + @classmethod + @lru_cache(maxsize=1) + def is_blackwell(cls): + if not cls.is_cuda_static(): + return False + return torch.cuda.get_device_capability()[0] == 10 + + @classmethod + @lru_cache(maxsize=1) + def is_sm120(cls): + if not cls.is_cuda_static(): + return False + return torch.cuda.get_device_capability()[0] == 12 + + @classmethod + def is_cuda_static(cls) -> bool: + return getattr(cls, "_enum", None) == PlatformEnum.CUDA + + @classmethod + def is_rocm_static(cls) -> bool: + return getattr(cls, "_enum", None) == PlatformEnum.ROCM + + @lru_cache(maxsize=1) + def is_hpu(self) -> bool: + return hasattr(torch, "hpu") and torch.hpu.is_available() + + @lru_cache(maxsize=1) + def is_xpu(self) -> bool: + return hasattr(torch, "xpu") and torch.xpu.is_available() + + @lru_cache(maxsize=1) + def is_npu(self) -> bool: + return hasattr(torch, "npu") and torch.npu.is_available() + def is_out_of_tree(self) -> bool: return self._enum == PlatformEnum.OOT + @lru_cache(maxsize=1) def is_cuda_alike(self) -> bool: """Stateless version of :func:`torch.cuda.is_available`.""" return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM) + @lru_cache(maxsize=1) def is_mps(self) -> bool: return self._enum == PlatformEnum.MPS + @lru_cache(maxsize=1) + def is_musa(self): + try: + return hasattr(torch, "musa") and torch.musa.is_available() + except ModuleNotFoundError: + return False + + @lru_cache(maxsize=1) + def is_hip(self) -> bool: + return self.is_rocm() + @classmethod def get_attn_backend_cls_str( cls, @@ -168,6 +220,30 @@ class Platform: """Get the total memory of a device in bytes.""" raise NotImplementedError + @lru_cache(maxsize=1) + def get_device(self, local_rank: int) -> torch.device: + if self.is_cuda() or self.is_rocm(): + return torch.device("cuda", local_rank) + elif self.is_musa(): + return torch.device("musa", local_rank) + elif self.is_mps(): + return torch.device("mps") + else: + return torch.device("cpu") + + @lru_cache(maxsize=1) + def get_torch_distributed_backend_str(self) -> str: + if self.is_cuda_alike(): + return "nccl" + elif self.is_musa(): + return "mccl" + elif self.is_mps(): + return "gloo" + else: + raise NotImplementedError( + "No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available" + ) + @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: """ diff --git a/python/sglang/multimodal_gen/runtime/utils/common.py b/python/sglang/multimodal_gen/runtime/utils/common.py index 156800491..134609b2a 100644 --- a/python/sglang/multimodal_gen/runtime/utils/common.py +++ b/python/sglang/multimodal_gen/runtime/utils/common.py @@ -1,6 +1,5 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo -import importlib import ipaddress import logging import os @@ -241,48 +240,6 @@ def get_zmq_socket( # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip -@lru_cache(maxsize=1) -def is_hip() -> bool: - return torch.version.hip is not None - - -@lru_cache(maxsize=1) -def is_cuda(): - return torch.cuda.is_available() and torch.version.cuda - - -@lru_cache(maxsize=1) -def is_cuda_alike(): - return is_cuda() or is_hip() - - -@lru_cache(maxsize=1) -def is_blackwell(): - if not is_cuda(): - return False - return torch.cuda.get_device_capability()[0] == 10 - - -@lru_cache(maxsize=1) -def is_sm120(): - if not is_cuda(): - return False - return torch.cuda.get_device_capability()[0] == 12 - - -@lru_cache(maxsize=1) -def is_hpu() -> bool: - return hasattr(torch, "hpu") and torch.hpu.is_available() - - -@lru_cache(maxsize=1) -def is_xpu() -> bool: - return hasattr(torch, "xpu") and torch.xpu.is_available() - - -@lru_cache(maxsize=1) -def is_npu() -> bool: - return hasattr(torch, "npu") and torch.npu.is_available() @lru_cache(maxsize=1) @@ -295,11 +252,6 @@ def is_host_cpu_x86() -> bool: ) -@lru_cache(maxsize=1) -def is_cpu() -> bool: - return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" and is_host_cpu_x86() - - # cuda @@ -309,16 +261,6 @@ def set_cuda_arch(): os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}" -def is_flashinfer_available(): - """ - Check whether flashinfer is available. - As of Oct. 6, 2024, it is only available on NVIDIA GPUs. - """ - # if not get_bool_env_var("SGLANG_IS_FLASHINFER_AVAILABLE", default="true"): - # return False - return importlib.util.find_spec("flashinfer") is not None and is_cuda() - - # env var managements _warned_bool_env_var_keys = set() diff --git a/python/sglang/multimodal_gen/runtime/utils/distributed.py b/python/sglang/multimodal_gen/runtime/utils/distributed.py index e63a046b8..80afa0721 100644 --- a/python/sglang/multimodal_gen/runtime/utils/distributed.py +++ b/python/sglang/multimodal_gen/runtime/utils/distributed.py @@ -21,6 +21,7 @@ def broadcast_pyobj( The `rank` here refer to the source rank on global process group (regardless of dist_group argument). """ + device = torch.device( current_platform.device_type if not force_cpu_device else "cpu" ) diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index a0af3baf6..e9e1c3de9 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -18,8 +18,6 @@ from typing import Any, cast import sglang.multimodal_gen.envs as envs -SGLANG_DIFFUSION_CONFIGURE_LOGGING = envs.SGLANG_DIFFUSION_CONFIGURE_LOGGING -SGLANG_DIFFUSION_LOGGING_CONFIG_PATH = envs.SGLANG_DIFFUSION_LOGGING_CONFIG_PATH SGLANG_DIFFUSION_LOGGING_LEVEL = envs.SGLANG_DIFFUSION_LOGGING_LEVEL SGLANG_DIFFUSION_LOGGING_PREFIX = envs.SGLANG_DIFFUSION_LOGGING_PREFIX diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 2beed4250..3b59ea3ef 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -16,7 +16,7 @@ import pytest import requests from openai import OpenAI -from sglang.multimodal_gen.runtime.utils.common import is_hip +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord from sglang.multimodal_gen.test.server.conftest import _GLOBAL_PERF_RESULTS @@ -52,7 +52,11 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: # Skip ring attention tests on AMD/ROCm - Ring Attention requires Flash Attention # which is not available on AMD. Use Ulysses parallelism instead. - if is_hip() and server_args.ring_degree is not None and server_args.ring_degree > 1: + if ( + current_platform.is_hip() + and server_args.ring_degree is not None + and server_args.ring_degree > 1 + ): pytest.skip( f"Skipping {case.id}: Ring Attention (ring_degree={server_args.ring_degree}) " "requires Flash Attention which is not available on AMD/ROCm" diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 5f8f789ae..d02d82f36 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -21,7 +21,8 @@ import pytest from openai import Client, OpenAI from sglang.multimodal_gen.benchmarks.compare_perf import calculate_upper_bound -from sglang.multimodal_gen.runtime.utils.common import is_hip, kill_process_tree +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.utils.common import kill_process_tree from sglang.multimodal_gen.runtime.utils.logging_utils import ( globally_suppress_loggers, init_logger, @@ -150,7 +151,7 @@ class ServerContext: # ROCm/AMD: Extra cleanup to ensure GPU memory is released between tests # This is needed because ROCm memory release can be slower than CUDA - if is_hip(): + if current_platform.is_hip(): self._cleanup_rocm_gpu_memory() # Clean up downloaded models if HF cache is not persistent # This prevents disk exhaustion in CI when cache is not mounted @@ -318,7 +319,7 @@ class ServerManager: """Start the diffusion server and wait for readiness.""" # ROCm/AMD: Wait for GPU memory to be clear before starting # This prevents OOM when running sequential tests on ROCm - if is_hip(): + if current_platform.is_hip(): self._wait_for_rocm_gpu_memory_clear() log_dir, perf_log_path = prepare_perf_log() @@ -551,7 +552,7 @@ class PerformanceValidator: For AMD GPUs, uses 100% higher tolerance and issues warning instead of assertion. """ # Check if running on AMD GPU - is_amd = is_hip() + is_amd = current_platform.is_hip() if is_amd: # Use 100% higher tolerance for AMD (2x the expected value) @@ -753,7 +754,7 @@ def get_generate_fn( job_completed = False is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1" # Check if running on AMD GPU - use longer timeout - is_amd = is_hip() + is_amd = current_platform.is_hip() if is_baseline_generation_mode: timeout = 3600.0 elif is_amd: