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

[diffusion] refactor: centralize hardware platform detection and streamline environment variable management (#15842)

Browse Source
This commit is contained in:
Mick
2025-12-26 22:16:18 +08:00
committed by GitHub
parent cf34d0ab32
commit 8dc6f0fc4d
12 changed files with 260 additions and 365 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/envs.py
+152 -261
View File
@@ -1,16 +1,11 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import importlib.util
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py
import logging import logging
import os import os
from collections.abc import Callable from typing import TYPE_CHECKING, Any, Callable
from typing import TYPE_CHECKING, Any
import diffusers
import torch
from packaging import version
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var 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_MC: int = 3
SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER: bool = False SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER: bool = False
SGLANG_CACHE_DIT_SECONDARY_TS_ORDER: int = 1 SGLANG_CACHE_DIT_SECONDARY_TS_ORDER: int = 1
# model loading
SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True
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()
def get_default_cache_root() -> str: 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: def maybe_convert_int(value: str | None) -> int | None:
if value is None: 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 return None
return int(value) val = default_func()
return os.path.expanduser(val)
return _getter
# The begin-* and end* here are used by the documentation generator # 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 ================== # ================== Installation Time Env Vars ==================
# Target device of sglang-diffusion, supporting [cuda (by default), # Target device of sglang-diffusion, supporting [cuda (by default),
# rocm, neuron, cpu, openvino] # rocm, neuron, cpu, openvino]
"SGLANG_DIFFUSION_TARGET_DEVICE": lambda: os.getenv( "SGLANG_DIFFUSION_TARGET_DEVICE": _lazy_str(
"SGLANG_DIFFUSION_TARGET_DEVICE", "cuda" "SGLANG_DIFFUSION_TARGET_DEVICE", "cuda"
), ),
# Maximum number of compilation jobs to run in parallel. # Maximum number of compilation jobs to run in parallel.
# By default this is the number of CPUs # 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 # Number of threads to use for nvcc
# By default this is 1. # By default this is 1.
# If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. # 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) # If set, sgl_diffusion will use precompiled binaries (*.so)
"SGLANG_DIFFUSION_USE_PRECOMPILED": lambda: bool( "SGLANG_DIFFUSION_USE_PRECOMPILED": _lazy_bool_any(
os.environ.get("SGLANG_DIFFUSION_USE_PRECOMPILED") [
) "SGLANG_DIFFUSION_USE_PRECOMPILED",
or bool(os.environ.get("SGLANG_DIFFUSION_PRECOMPILED_WHEEL_LOCATION")), "SGLANG_DIFFUSION_PRECOMPILED_WHEEL_LOCATION",
]
),
# CMake build type # CMake build type
# If not set, defaults to "Debug" or "RelWithDebInfo" # If not set, defaults to "Debug" or "RelWithDebInfo"
# Available options: "Debug", "Release", "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 # If set, sgl_diffusion will print verbose logs during installation
"VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))), "VERBOSE": _lazy_bool("VERBOSE"),
# Root directory for FASTVIDEO configuration files # Root directory for SGL-diffusion configuration files
# Defaults to `~/.config/sgl_diffusion` unless `XDG_CONFIG_HOME` is set # Defaults to `~/.config/sgl_diffusion` unless `XDG_CONFIG_HOME` is set
# Note that this not only affects how sgl_diffusion finds its configuration files # Note that this not only affects how sgl_diffusion finds its configuration files
# during runtime, but also affects how sgl_diffusion installs its configuration # during runtime, but also affects how sgl_diffusion installs its configuration
# files during **installation**. # files during **installation**.
"SGLANG_DIFFUSION_CONFIG_ROOT": lambda: os.path.expanduser( "SGLANG_DIFFUSION_CONFIG_ROOT": _lazy_path(
os.getenv(
"SGLANG_DIFFUSION_CONFIG_ROOT", "SGLANG_DIFFUSION_CONFIG_ROOT",
os.path.join(get_default_config_root(), "sgl_diffusion"), lambda: os.path.join(get_default_config_root(), "sgl_diffusion"),
)
), ),
# ================== Runtime Env Vars ================== # ================== 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 # Defaults to `~/.cache/sgl_diffusion` unless `XDG_CACHE_HOME` is set
"SGLANG_DIFFUSION_CACHE_ROOT": lambda: os.path.expanduser( "SGLANG_DIFFUSION_CACHE_ROOT": _lazy_path(
os.getenv(
"SGLANG_DIFFUSION_CACHE_ROOT", "SGLANG_DIFFUSION_CACHE_ROOT",
os.path.join(get_default_cache_root(), "sgl_diffusion"), lambda: os.path.join(get_default_cache_root(), "sgl_diffusion"),
)
), ),
# Interval in seconds to log a warning message when the ring buffer is full # Interval in seconds to log a warning message when the ring buffer is full
"SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL": lambda: int( "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL": _lazy_int(
os.environ.get("SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL", "60") "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL", 60
), ),
# Path to the NCCL library file. It is needed because nccl>=2.19 brought # 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 # 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": _lazy_str("SGLANG_DIFFUSION_NCCL_SO_PATH"),
"SGLANG_DIFFUSION_NCCL_SO_PATH", None
),
# when `SGLANG_DIFFUSION_NCCL_SO_PATH` is not set, sgl_diffusion will try to find the nccl # 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` # 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 # Internal flag to enable Dynamo fullgraph capture
"SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": lambda: bool( "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": _lazy_bool(
os.environ.get("SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0" "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1"
), ),
# local rank of the process in the distributed setting, used to determine # local rank of the process in the distributed setting, used to determine
# the GPU device id # 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 # 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 # timeout for each iteration in the engine
"SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": lambda: int( "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": _lazy_int(
os.environ.get("SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S", "60") "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S", 60
), ),
# Logging configuration # Logging configuration
# If set to 0, sgl_diffusion will not configure logging # If set to 0, sgl_diffusion will not configure logging
# If set to 1, sgl_diffusion will configure logging using the default configuration # If set to 1, sgl_diffusion will configure logging using the default configuration
# or the configuration file specified by SGLANG_DIFFUSION_LOGGING_CONFIG_PATH # or the configuration file specified by SGLANG_DIFFUSION_LOGGING_CONFIG_PATH
"SGLANG_DIFFUSION_CONFIGURE_LOGGING": lambda: int( "SGLANG_DIFFUSION_CONFIGURE_LOGGING": _lazy_int(
os.getenv("SGLANG_DIFFUSION_CONFIGURE_LOGGING", "1") "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" "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH"
), ),
# this is used for configuring the default logging level # 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" "SGLANG_DIFFUSION_LOGGING_LEVEL", "INFO"
), ),
# if set, SGLANG_DIFFUSION_LOGGING_PREFIX will be prepended to all log messages # if set, SGLANG_DIFFUSION_LOGGING_PREFIX will be prepended to all log messages
"SGLANG_DIFFUSION_LOGGING_PREFIX": lambda: os.getenv( "SGLANG_DIFFUSION_LOGGING_PREFIX": _lazy_str("SGLANG_DIFFUSION_LOGGING_PREFIX", ""),
"SGLANG_DIFFUSION_LOGGING_PREFIX", ""
),
# Trace function calls # Trace function calls
# If set to 1, sgl_diffusion will trace function calls # If set to 1, sgl_diffusion will trace function calls
# Useful for debugging # Useful for debugging
"SGLANG_DIFFUSION_TRACE_FUNCTION": lambda: int( "SGLANG_DIFFUSION_TRACE_FUNCTION": _lazy_int("SGLANG_DIFFUSION_TRACE_FUNCTION", 0),
os.getenv("SGLANG_DIFFUSION_TRACE_FUNCTION", "0")
),
# Path to the attention configuration file. Only used for sliding tile # Path to the attention configuration file. Only used for sliding tile
# attention for now. # attention for now.
"SGLANG_DIFFUSION_ATTENTION_CONFIG": lambda: ( "SGLANG_DIFFUSION_ATTENTION_CONFIG": _lazy_path(
None "SGLANG_DIFFUSION_ATTENTION_CONFIG"
if os.getenv("SGLANG_DIFFUSION_ATTENTION_CONFIG", None) is None
else os.path.expanduser(os.getenv("SGLANG_DIFFUSION_ATTENTION_CONFIG", "."))
), ),
# Optional override to force a specific attention backend (e.g. "aiter") # 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" "SGLANG_DIFFUSION_ATTENTION_BACKEND"
), ),
# Use dedicated multiprocess context for workers. # Use dedicated multiprocess context for workers.
# Both spawn and fork work # 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" "SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD", "fork"
), ),
# Enables torch profiler if set. Path to the directory where torch profiler # Enables torch profiler if set. Path to the directory where torch profiler
# traces are saved. Note that it must be an absolute path. # traces are saved. Note that it must be an absolute path.
"SGLANG_DIFFUSION_TORCH_PROFILER_DIR": lambda: ( "SGLANG_DIFFUSION_TORCH_PROFILER_DIR": _lazy_path(
None "SGLANG_DIFFUSION_TORCH_PROFILER_DIR"
if os.getenv("SGLANG_DIFFUSION_TORCH_PROFILER_DIR", None) is None
else os.path.expanduser(os.getenv("SGLANG_DIFFUSION_TORCH_PROFILER_DIR", "."))
), ),
# If set, sgl_diffusion will run in development mode, which will enable # If set, sgl_diffusion will run in development mode, which will enable
# some additional endpoints for developing and debugging, # some additional endpoints for developing and debugging,
# e.g. `/reset_prefix_cache` # e.g. `/reset_prefix_cache`
"SGLANG_DIFFUSION_SERVER_DEV_MODE": lambda: get_bool_env_var( "SGLANG_DIFFUSION_SERVER_DEV_MODE": _lazy_bool("SGLANG_DIFFUSION_SERVER_DEV_MODE"),
"SGLANG_DIFFUSION_SERVER_DEV_MODE"
),
# If set, sgl_diffusion will enable stage logging, which will print the time # If set, sgl_diffusion will enable stage logging, which will print the time
# taken for each stage # taken for each stage
"SGLANG_DIFFUSION_STAGE_LOGGING": lambda: get_bool_env_var( "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"),
"SGLANG_DIFFUSION_STAGE_LOGGING"
),
# ================== cache-dit Env Vars ================== # ================== cache-dit Env Vars ==================
# Enable cache-dit acceleration for DiT inference # 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) # 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) # 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) # 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) # 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) # 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 # Enable TaylorSeer calibrator
"SGLANG_CACHE_DIT_TAYLORSEER": lambda: get_bool_env_var( "SGLANG_CACHE_DIT_TAYLORSEER": _lazy_bool("SGLANG_CACHE_DIT_TAYLORSEER", "false"),
"SGLANG_CACHE_DIT_TAYLORSEER", default="false"
),
# TaylorSeer order (1 or 2) # TaylorSeer order (1 or 2)
"SGLANG_CACHE_DIT_TS_ORDER": lambda: int( "SGLANG_CACHE_DIT_TS_ORDER": _lazy_int("SGLANG_CACHE_DIT_TS_ORDER", 1),
os.getenv("SGLANG_CACHE_DIT_TS_ORDER", "1")
),
# SCM preset: none, slow, medium, fast, ultra # SCM preset: none, slow, medium, fast, ultra
"SGLANG_CACHE_DIT_SCM_PRESET": lambda: os.getenv( "SGLANG_CACHE_DIT_SCM_PRESET": _lazy_str("SGLANG_CACHE_DIT_SCM_PRESET", "none"),
"SGLANG_CACHE_DIT_SCM_PRESET", "none"
),
# SCM custom compute bins (e.g., "8,3,3,2,2") # 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": _lazy_str("SGLANG_CACHE_DIT_SCM_COMPUTE_BINS"),
"SGLANG_CACHE_DIT_SCM_COMPUTE_BINS", None
),
# SCM custom cache bins (e.g., "1,2,2,2,3") # 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": _lazy_str("SGLANG_CACHE_DIT_SCM_CACHE_BINS"),
"SGLANG_CACHE_DIT_SCM_CACHE_BINS", None
),
# SCM policy: dynamic or static # SCM policy: dynamic or static
"SGLANG_CACHE_DIT_SCM_POLICY": lambda: os.getenv( "SGLANG_CACHE_DIT_SCM_POLICY": _lazy_str("SGLANG_CACHE_DIT_SCM_POLICY", "dynamic"),
"SGLANG_CACHE_DIT_SCM_POLICY", "dynamic" # model loading
), "SGLANG_USE_RUNAI_MODEL_STREAMER": _lazy_bool(
# ================== cache-dit Secondary Transformer Env Vars ================== "SGLANG_USE_RUNAI_MODEL_STREAMER", "true"
# 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"),
)
), ),
} }
# 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 # end-env-vars-definition
def __getattr__(name: str): def __getattr__(name: str):
# lazy evaluation of environment variables # lazy evaluation of environment variables
if name in environment_variables: if name in environment_variables:
@@ -411,27 +326,3 @@ def __getattr__(name: str):
def __dir__(): def __dir__():
return list(environment_variables.keys()) 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")
python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py
+2 -2
View File
@@ -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 ( from sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator import (
CpuCommunicator, CpuCommunicator,
) )
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import ( from sglang.multimodal_gen.runtime.utils.logging_utils import (
init_logger, init_logger,
suppress_stdout, suppress_stdout,
@@ -44,7 +45,6 @@ _group_name_counter: dict[str, int] = {}
def get_local_torch_device() -> torch.device: def get_local_torch_device() -> torch.device:
"""Return the torch device for the current rank.""" """Return the torch device for the current rank."""
from sglang.multimodal_gen.runtime.platforms import current_platform
return ( return (
torch.device(f"cuda:{envs.LOCAL_RANK}") torch.device(f"cuda:{envs.LOCAL_RANK}")
@@ -846,7 +846,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
assert self.cpu_group is not None assert self.cpu_group is not None
assert self.device_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_buffer_set: bool = False
self.recv_tasks_queue: List[Tuple[str, int]] = [] self.recv_tasks_queue: List[Tuple[str, int]] = []
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
+3 -1
View File
@@ -341,7 +341,9 @@ def initialize_model_parallel(
""" """
if backend is None: 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. # Get world size and rank. Ensure some consistencies.
assert torch.distributed.is_initialized() assert torch.distributed.is_initialized()
world_size: int = torch.distributed.get_world_size() world_size: int = torch.distributed.get_world_size()
python/sglang/multimodal_gen/runtime/layers/custom_op.py
+5 -17
View File
@@ -8,23 +8,11 @@ from typing import Any
import torch.nn as nn import torch.nn as nn
from sglang.multimodal_gen.runtime.utils.common import ( from sglang.multimodal_gen.runtime.platforms import current_platform
is_cpu,
is_cuda,
is_hip,
is_npu,
is_xpu,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) 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): class CustomOp(nn.Module):
""" """
@@ -70,13 +58,13 @@ class CustomOp(nn.Module):
return self.forward_native(*args, **kwargs) return self.forward_native(*args, **kwargs)
def dispatch_forward(self) -> Callable: def dispatch_forward(self) -> Callable:
if _is_cuda: if current_platform.is_cuda():
return self.forward_cuda return self.forward_cuda
elif _is_hip: elif current_platform.is_hip():
return self.forward_hip return self.forward_hip
elif _is_npu: elif current_platform.is_npu():
return self.forward_npu return self.forward_npu
elif _is_xpu: elif current_platform.is_xpu():
return self.forward_xpu return self.forward_xpu
else: else:
return self.forward_native return self.forward_native
python/sglang/multimodal_gen/runtime/platforms/__init__.py
-7
View File
@@ -139,7 +139,6 @@ def resolve_current_platform_cls_qualname() -> str:
_current_platform: Platform | None = None _current_platform: Platform | None = None
_init_trace: str = "" _init_trace: str = ""
if TYPE_CHECKING:
current_platform: Platform current_platform: Platform
@@ -150,12 +149,6 @@ def __getattr__(name: str):
# Platform` so that they can inherit `Platform` class. Therefore, # Platform` so that they can inherit `Platform` class. Therefore,
# we cannot resolve `current_platform` during the import of # we cannot resolve `current_platform` during the import of
# `sglang.multimodal_gen.runtime.platforms`. # `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 global _current_platform
if _current_platform is None: if _current_platform is None:
platform_cls_qualname = resolve_current_platform_cls_qualname() platform_cls_qualname = resolve_current_platform_cls_qualname()
python/sglang/multimodal_gen/runtime/platforms/cuda.py
+3 -4
View File
@@ -20,7 +20,6 @@ from sglang.multimodal_gen.runtime.platforms.interface import (
Platform, Platform,
PlatformEnum, 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.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import import_pynvml from sglang.multimodal_gen.utils import import_pynvml
@@ -213,7 +212,7 @@ class CudaPlatformBase(Platform):
elif selected_backend in [ elif selected_backend in [
AttentionBackendEnum.FA, AttentionBackendEnum.FA,
]: ]:
if is_blackwell(): if cls.is_blackwell():
from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
set_fa_ver, set_fa_ver,
) )
@@ -224,14 +223,14 @@ class CudaPlatformBase(Platform):
raise ValueError(f"Invalid attention backend for {cls.device_name}") raise ValueError(f"Invalid attention backend for {cls.device_name}")
else: else:
if is_blackwell(): if cls.is_blackwell():
from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
set_fa_ver, set_fa_ver,
) )
set_fa_ver(4) set_fa_ver(4)
target_backend = AttentionBackendEnum.FA target_backend = AttentionBackendEnum.FA
if is_sm120(): if cls.is_sm120():
try: try:
from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401 from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401
SageAttention3Backend, SageAttention3Backend,
python/sglang/multimodal_gen/runtime/platforms/interface.py
+78 -2
View File
@@ -6,6 +6,7 @@ from __future__ import annotations
import enum import enum
import random import random
from functools import lru_cache
from typing import TYPE_CHECKING, NamedTuple from typing import TYPE_CHECKING, NamedTuple
import numpy as np import numpy as np
@@ -90,28 +91,79 @@ class Platform:
supported_quantization: list[str] = [] supported_quantization: list[str] = []
@lru_cache(maxsize=1)
def is_cuda(self) -> bool: def is_cuda(self) -> bool:
return self._enum == PlatformEnum.CUDA return self.is_cuda_static()
@lru_cache(maxsize=1)
def is_rocm(self) -> bool: def is_rocm(self) -> bool:
return self._enum == PlatformEnum.ROCM return self.is_rocm_static()
@lru_cache(maxsize=1)
def is_tpu(self) -> bool: def is_tpu(self) -> bool:
return self._enum == PlatformEnum.TPU return self._enum == PlatformEnum.TPU
@lru_cache(maxsize=1)
def is_cpu(self) -> bool: def is_cpu(self) -> bool:
return self._enum == PlatformEnum.CPU 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: def is_out_of_tree(self) -> bool:
return self._enum == PlatformEnum.OOT return self._enum == PlatformEnum.OOT
@lru_cache(maxsize=1)
def is_cuda_alike(self) -> bool: def is_cuda_alike(self) -> bool:
"""Stateless version of :func:`torch.cuda.is_available`.""" """Stateless version of :func:`torch.cuda.is_available`."""
return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM) return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
@lru_cache(maxsize=1)
def is_mps(self) -> bool: def is_mps(self) -> bool:
return self._enum == PlatformEnum.MPS 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 @classmethod
def get_attn_backend_cls_str( def get_attn_backend_cls_str(
cls, cls,
@@ -168,6 +220,30 @@ class Platform:
"""Get the total memory of a device in bytes.""" """Get the total memory of a device in bytes."""
raise NotImplementedError 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 @classmethod
def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
""" """
python/sglang/multimodal_gen/runtime/utils/common.py
-58
View File
@@ -1,6 +1,5 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import importlib
import ipaddress import ipaddress
import logging import logging
import os import os
@@ -241,48 +240,6 @@ def get_zmq_socket(
# https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip # 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) @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 # cuda
@@ -309,16 +261,6 @@ def set_cuda_arch():
os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}" 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 # env var managements
_warned_bool_env_var_keys = set() _warned_bool_env_var_keys = set()
python/sglang/multimodal_gen/runtime/utils/distributed.py
+1
View File
@@ -21,6 +21,7 @@ def broadcast_pyobj(
The `rank` here refer to the source rank on global process group (regardless The `rank` here refer to the source rank on global process group (regardless
of dist_group argument). of dist_group argument).
""" """
device = torch.device( device = torch.device(
current_platform.device_type if not force_cpu_device else "cpu" current_platform.device_type if not force_cpu_device else "cpu"
) )
python/sglang/multimodal_gen/runtime/utils/logging_utils.py
-2
View File
@@ -18,8 +18,6 @@ from typing import Any, cast
import sglang.multimodal_gen.envs as envs 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_LEVEL = envs.SGLANG_DIFFUSION_LOGGING_LEVEL
SGLANG_DIFFUSION_LOGGING_PREFIX = envs.SGLANG_DIFFUSION_LOGGING_PREFIX SGLANG_DIFFUSION_LOGGING_PREFIX = envs.SGLANG_DIFFUSION_LOGGING_PREFIX
python/sglang/multimodal_gen/test/server/test_server_common.py
+6 -2
View File
@@ -16,7 +16,7 @@ import pytest
import requests import requests
from openai import OpenAI 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.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.server.conftest import _GLOBAL_PERF_RESULTS 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 # Skip ring attention tests on AMD/ROCm - Ring Attention requires Flash Attention
# which is not available on AMD. Use Ulysses parallelism instead. # 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( pytest.skip(
f"Skipping {case.id}: Ring Attention (ring_degree={server_args.ring_degree}) " f"Skipping {case.id}: Ring Attention (ring_degree={server_args.ring_degree}) "
"requires Flash Attention which is not available on AMD/ROCm" "requires Flash Attention which is not available on AMD/ROCm"
python/sglang/multimodal_gen/test/server/test_server_utils.py
+6 -5
View File
@@ -21,7 +21,8 @@ import pytest
from openai import Client, OpenAI from openai import Client, OpenAI
from sglang.multimodal_gen.benchmarks.compare_perf import calculate_upper_bound 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 ( from sglang.multimodal_gen.runtime.utils.logging_utils import (
globally_suppress_loggers, globally_suppress_loggers,
init_logger, init_logger,
@@ -150,7 +151,7 @@ class ServerContext:
# ROCm/AMD: Extra cleanup to ensure GPU memory is released between tests # ROCm/AMD: Extra cleanup to ensure GPU memory is released between tests
# This is needed because ROCm memory release can be slower than CUDA # 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() self._cleanup_rocm_gpu_memory()
# Clean up downloaded models if HF cache is not persistent # Clean up downloaded models if HF cache is not persistent
# This prevents disk exhaustion in CI when cache is not mounted # 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.""" """Start the diffusion server and wait for readiness."""
# ROCm/AMD: Wait for GPU memory to be clear before starting # ROCm/AMD: Wait for GPU memory to be clear before starting
# This prevents OOM when running sequential tests on ROCm # This prevents OOM when running sequential tests on ROCm
if is_hip(): if current_platform.is_hip():
self._wait_for_rocm_gpu_memory_clear() self._wait_for_rocm_gpu_memory_clear()
log_dir, perf_log_path = prepare_perf_log() 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. For AMD GPUs, uses 100% higher tolerance and issues warning instead of assertion.
""" """
# Check if running on AMD GPU # Check if running on AMD GPU
is_amd = is_hip() is_amd = current_platform.is_hip()
if is_amd: if is_amd:
# Use 100% higher tolerance for AMD (2x the expected value) # Use 100% higher tolerance for AMD (2x the expected value)
@@ -753,7 +754,7 @@ def get_generate_fn(
job_completed = False job_completed = False
is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1" is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
# Check if running on AMD GPU - use longer timeout # Check if running on AMD GPU - use longer timeout
is_amd = is_hip() is_amd = current_platform.is_hip()
if is_baseline_generation_mode: if is_baseline_generation_mode:
timeout = 3600.0 timeout = 3600.0
elif is_amd: elif is_amd: