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

[diffusion] CI: fix hunyuan3d JIT cache (#20773)

Browse Source 
Co-authored-by: daiweitao <dwti614707404@163.com>
This commit is contained in:
Prozac614
2026-03-20 17:51:55 +08:00
committed by GitHub
co-authored by daiweitao
parent a0a4dae67f
commit fa89d152c0
6 changed files with 198 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/csrc/render/__init__.py
+130
View File
@@ -0,0 +1,130 @@
from __future__ import annotations
import os
import shutil
import sys
from pathlib import Path
from typing import Any, Sequence
import torch
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def _get_build_directory(name: str) -> Path:
try:
from torch.utils.cpp_extension import _get_build_directory
return Path(_get_build_directory(name, False))
except (ImportError, AttributeError):
from torch.utils.cpp_extension import get_default_build_root
root = os.environ.get("TORCH_EXTENSIONS_DIR") or get_default_build_root()
if "TORCH_EXTENSIONS_DIR" not in os.environ:
cu_str = (
"cpu"
if torch.version.cuda is None
else f"cu{torch.version.cuda.replace('.', '')}"
)
py_str = (
f"py{sys.version_info.major}{sys.version_info.minor}"
f"{getattr(sys, 'abiflags', '')}"
)
root = os.path.join(root, f"{py_str}_{cu_str}")
return Path(root) / name
def _is_recoverable_load_error(
exc: BaseException, name: str, build_directory: Path
) -> bool:
message = str(exc).lower()
current = exc.__cause__ or exc.__context__
while current is not None:
message += f"\n{current}".lower()
current = current.__cause__ or current.__context__
if any(
marker in message
for marker in (
"error building extension",
"error compiling objects for extension",
"ninja",
"nvcc",
"gcc",
"g++",
"fatal error:",
"compilation terminated",
)
):
return False
if not any(
marker in message
for marker in (str(build_directory / f"{name}.so").lower(), f"{name}.so")
):
return False
return any(
marker in message
for marker in (
"undefined symbol",
"cannot open shared object file",
"no such file or directory",
"file too short",
"invalid elf header",
"wrong elf class",
"elf load command",
"dlopen",
"version `glibcxx",
)
)
def load_extension_with_recovery(
name: str,
sources: Sequence[str],
extra_cflags: Sequence[str] | None = None,
extra_cuda_cflags: Sequence[str] | None = None,
verbose: bool = False,
) -> Any:
from torch.utils.cpp_extension import load
try:
return load(
name=name,
sources=list(sources),
extra_cflags=None if extra_cflags is None else list(extra_cflags),
extra_cuda_cflags=(
None if extra_cuda_cflags is None else list(extra_cuda_cflags)
),
verbose=verbose,
)
except Exception as exc:
build_directory = _get_build_directory(name)
if not _is_recoverable_load_error(exc, name, build_directory):
raise
logger.warning(
"Detected a stale or broken JIT extension for %s at %s; clearing "
"its cache and retrying once.",
name,
build_directory,
)
sys.modules.pop(name, None)
if build_directory.exists():
shutil.rmtree(build_directory)
return load(
name=name,
sources=list(sources),
extra_cflags=None if extra_cflags is None else list(extra_cflags),
extra_cuda_cflags=(
None if extra_cuda_cflags is None else list(extra_cuda_cflags)
),
verbose=verbose,
)
__all__ = ["load_extension_with_recovery"]
python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py
+2 -3
View File
@@ -12,6 +12,7 @@ import os
from typing import List, Tuple from typing import List, Tuple
import torch import torch
from sglang.multimodal_gen.csrc.render import load_extension_with_recovery
_abs_path = os.path.dirname(os.path.abspath(__file__)) _abs_path = os.path.dirname(os.path.abspath(__file__))
_custom_rasterizer_kernel = None _custom_rasterizer_kernel = None
@@ -24,9 +25,7 @@ def _load_custom_rasterizer():
if _custom_rasterizer_kernel is not None: if _custom_rasterizer_kernel is not None:
return _custom_rasterizer_kernel return _custom_rasterizer_kernel
from torch.utils.cpp_extension import load _custom_rasterizer_kernel = load_extension_with_recovery(
_custom_rasterizer_kernel = load(
name="custom_rasterizer_kernel", name="custom_rasterizer_kernel",
sources=[ sources=[
f"{_abs_path}/rasterizer.cpp", f"{_abs_path}/rasterizer.cpp",
python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py
+2 -3
View File
@@ -12,6 +12,7 @@ import os
from typing import Tuple from typing import Tuple
import numpy as np import numpy as np
from sglang.multimodal_gen.csrc.render import load_extension_with_recovery
_abs_path = os.path.dirname(os.path.abspath(__file__)) _abs_path = os.path.dirname(os.path.abspath(__file__))
_mesh_processor_kernel = None _mesh_processor_kernel = None
@@ -24,9 +25,7 @@ def _load_mesh_processor():
if _mesh_processor_kernel is not None: if _mesh_processor_kernel is not None:
return _mesh_processor_kernel return _mesh_processor_kernel
from torch.utils.cpp_extension import load _mesh_processor_kernel = load_extension_with_recovery(
_mesh_processor_kernel = load(
name="mesh_processor_kernel", name="mesh_processor_kernel",
sources=[ sources=[
f"{_abs_path}/mesh_processor.cpp", f"{_abs_path}/mesh_processor.cpp",
python/sglang/multimodal_gen/test/server/perf_baselines.json
+60 -60
View File
@@ -1988,69 +1988,69 @@
}, },
"hunyuan3d_shape_gen": { "hunyuan3d_shape_gen": {
"stages_ms": { "stages_ms": {
"Hunyuan3DShapeBeforeDenoisingStage": 31.42, "Hunyuan3DShapeBeforeDenoisingStage": 235.65,
"Hunyuan3DShapeDenoisingStage": 3259.83, "Hunyuan3DShapeDenoisingStage": 3452.51,
"Hunyuan3DShapeExportStage": 8735.55, "Hunyuan3DShapeExportStage": 8819.6,
"Hunyuan3DShapeSaveStage": 981.64, "Hunyuan3DShapeSaveStage": 752.4,
"Hunyuan3DPaintPreprocessStage": 226071.67, "Hunyuan3DPaintPreprocessStage": 218136.45,
"Hunyuan3DPaintTexGenStage": 11083.05, "Hunyuan3DPaintTexGenStage": 10259.21,
"Hunyuan3DPaintPostprocessStage": 7469.29 "Hunyuan3DPaintPostprocessStage": 6387.55
}, },
"denoise_step_ms": { "denoise_step_ms": {
"0": 32.26, "0": 150.72,
"1": 63.34, "1": 26.65,
"2": 65.44, "2": 65.91,
"3": 65.44, "3": 68.09,
"4": 65.6, "4": 76.12,
"5": 65.81, "5": 68.16,
"6": 65.82, "6": 61.19,
"7": 65.48, "7": 68.26,
"8": 65.9, "8": 67.92,
"9": 65.77, "9": 68.26,
"10": 65.54, "10": 68.06,
"11": 65.68, "11": 68.22,
"12": 65.85, "12": 68.11,
"13": 65.77, "13": 68.19,
"14": 65.7, "14": 69.58,
"15": 65.78, "15": 66.91,
"16": 66.0, "16": 68.03,
"17": 66.15, "17": 68.36,
"18": 65.91, "18": 68.49,
"19": 66.5, "19": 67.69,
"20": 65.76, "20": 68.19,
"21": 66.08, "21": 69.1,
"22": 66.06, "22": 67.78,
"23": 66.23, "23": 68.36,
"24": 65.79, "24": 68.19,
"25": 65.58, "25": 68.26,
"26": 65.88, "26": 68.06,
"27": 65.67, "27": 68.25,
"28": 65.87, "28": 68.39,
"29": 66.09, "29": 68.26,
"30": 65.81, "30": 68.05,
"31": 65.91, "31": 68.27,
"32": 66.18, "32": 68.2,
"33": 65.93, "33": 68.19,
"34": 66.26, "34": 68.02,
"35": 66.26, "35": 68.3,
"36": 66.27, "36": 68.2,
"37": 65.57, "37": 68.48,
"38": 66.02, "38": 68.23,
"39": 66.19, "39": 68.36,
"40": 65.23, "40": 67.9,
"41": 66.11, "41": 75.76,
"42": 66.18, "42": 62.04,
"43": 65.86, "43": 66.78,
"44": 65.86, "44": 67.85,
"45": 65.92, "45": 68.11,
"46": 65.65, "46": 67.92,
"47": 65.78, "47": 68.15,
"48": 66.01, "48": 67.89,
"49": 66.08 "49": 68.3
}, },
"expected_e2e_ms": 257696.97, "expected_e2e_ms": 248171.5,
"expected_avg_denoise_ms": 65.16, "expected_avg_denoise_ms": 68.98,
"expected_median_denoise_ms": 65.86 "expected_median_denoise_ms": 68.19
}, },
"wan2_1_t2v_1.3b_frame_interp_2x": { "wan2_1_t2v_1.3b_frame_interp_2x": {
"stages_ms": { "stages_ms": {
python/sglang/multimodal_gen/test/server/test_server_common.py
+2 -2
View File
@@ -96,8 +96,8 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
if server_args.lora_path: if server_args.lora_path:
extra_args += f" --lora-path {server_args.lora_path}" extra_args += f" --lora-path {server_args.lora_path}"
# default warmup if server_args.enable_warmup:
extra_args += f" --warmup" extra_args += " --warmup"
for arg in server_args.extras: for arg in server_args.extras:
extra_args += f" {arg}" extra_args += f" {arg}"
python/sglang/multimodal_gen/test/server/testcase_configs.py
+2
View File
@@ -198,6 +198,7 @@ class DiffusionServerArgs:
dit_offload_prefetch_size: int | float | None = None dit_offload_prefetch_size: int | float | None = None
enable_cache_dit: bool = False enable_cache_dit: bool = False
text_encoder_cpu_offload: bool = False text_encoder_cpu_offload: bool = False
enable_warmup: bool = True
extras: list[str] = field(default_factory=lambda: []) extras: list[str] = field(default_factory=lambda: [])
@@ -726,6 +727,7 @@ if not current_platform.is_hip():
DiffusionServerArgs( DiffusionServerArgs(
model_path="tencent/Hunyuan3D-2", model_path="tencent/Hunyuan3D-2",
modality="3d", modality="3d",
enable_warmup=False,
), ),
HUNYUAN3D_SHAPE_sampling_params, HUNYUAN3D_SHAPE_sampling_params,
), ),