From fa89d152c030664855d51bc26dc31b6cab41f801 Mon Sep 17 00:00:00 2001 From: Prozac614 Date: Fri, 20 Mar 2026 17:51:55 +0800 Subject: [PATCH] [diffusion] CI: fix hunyuan3d JIT cache (#20773) Co-authored-by: daiweitao --- .../multimodal_gen/csrc/render/__init__.py | 130 ++++++++++++++++++ .../render/hunyuan3d_rasterizer/__init__.py | 5 +- .../csrc/render/mesh_processor/__init__.py | 5 +- .../test/server/perf_baselines.json | 120 ++++++++-------- .../test/server/test_server_common.py | 4 +- .../test/server/testcase_configs.py | 2 + 6 files changed, 198 insertions(+), 68 deletions(-) create mode 100644 python/sglang/multimodal_gen/csrc/render/__init__.py diff --git a/python/sglang/multimodal_gen/csrc/render/__init__.py b/python/sglang/multimodal_gen/csrc/render/__init__.py new file mode 100644 index 000000000..d7277542c --- /dev/null +++ b/python/sglang/multimodal_gen/csrc/render/__init__.py @@ -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"] diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py index 0d778479e..6d410130b 100644 --- a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py +++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py @@ -12,6 +12,7 @@ import os from typing import List, Tuple import torch +from sglang.multimodal_gen.csrc.render import load_extension_with_recovery _abs_path = os.path.dirname(os.path.abspath(__file__)) _custom_rasterizer_kernel = None @@ -24,9 +25,7 @@ def _load_custom_rasterizer(): if _custom_rasterizer_kernel is not None: return _custom_rasterizer_kernel - from torch.utils.cpp_extension import load - - _custom_rasterizer_kernel = load( + _custom_rasterizer_kernel = load_extension_with_recovery( name="custom_rasterizer_kernel", sources=[ f"{_abs_path}/rasterizer.cpp", diff --git a/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py b/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py index 21fc15b96..03fb86410 100644 --- a/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py +++ b/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py @@ -12,6 +12,7 @@ import os from typing import Tuple import numpy as np +from sglang.multimodal_gen.csrc.render import load_extension_with_recovery _abs_path = os.path.dirname(os.path.abspath(__file__)) _mesh_processor_kernel = None @@ -24,9 +25,7 @@ def _load_mesh_processor(): if _mesh_processor_kernel is not None: return _mesh_processor_kernel - from torch.utils.cpp_extension import load - - _mesh_processor_kernel = load( + _mesh_processor_kernel = load_extension_with_recovery( name="mesh_processor_kernel", sources=[ f"{_abs_path}/mesh_processor.cpp", diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index a0b5b5576..4d5c3161d 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -1988,69 +1988,69 @@ }, "hunyuan3d_shape_gen": { "stages_ms": { - "Hunyuan3DShapeBeforeDenoisingStage": 31.42, - "Hunyuan3DShapeDenoisingStage": 3259.83, - "Hunyuan3DShapeExportStage": 8735.55, - "Hunyuan3DShapeSaveStage": 981.64, - "Hunyuan3DPaintPreprocessStage": 226071.67, - "Hunyuan3DPaintTexGenStage": 11083.05, - "Hunyuan3DPaintPostprocessStage": 7469.29 + "Hunyuan3DShapeBeforeDenoisingStage": 235.65, + "Hunyuan3DShapeDenoisingStage": 3452.51, + "Hunyuan3DShapeExportStage": 8819.6, + "Hunyuan3DShapeSaveStage": 752.4, + "Hunyuan3DPaintPreprocessStage": 218136.45, + "Hunyuan3DPaintTexGenStage": 10259.21, + "Hunyuan3DPaintPostprocessStage": 6387.55 }, "denoise_step_ms": { - "0": 32.26, - "1": 63.34, - "2": 65.44, - "3": 65.44, - "4": 65.6, - "5": 65.81, - "6": 65.82, - "7": 65.48, - "8": 65.9, - "9": 65.77, - "10": 65.54, - "11": 65.68, - "12": 65.85, - "13": 65.77, - "14": 65.7, - "15": 65.78, - "16": 66.0, - "17": 66.15, - "18": 65.91, - "19": 66.5, - "20": 65.76, - "21": 66.08, - "22": 66.06, - "23": 66.23, - "24": 65.79, - "25": 65.58, - "26": 65.88, - "27": 65.67, - "28": 65.87, - "29": 66.09, - "30": 65.81, - "31": 65.91, - "32": 66.18, - "33": 65.93, - "34": 66.26, - "35": 66.26, - "36": 66.27, - "37": 65.57, - "38": 66.02, - "39": 66.19, - "40": 65.23, - "41": 66.11, - "42": 66.18, - "43": 65.86, - "44": 65.86, - "45": 65.92, - "46": 65.65, - "47": 65.78, - "48": 66.01, - "49": 66.08 + "0": 150.72, + "1": 26.65, + "2": 65.91, + "3": 68.09, + "4": 76.12, + "5": 68.16, + "6": 61.19, + "7": 68.26, + "8": 67.92, + "9": 68.26, + "10": 68.06, + "11": 68.22, + "12": 68.11, + "13": 68.19, + "14": 69.58, + "15": 66.91, + "16": 68.03, + "17": 68.36, + "18": 68.49, + "19": 67.69, + "20": 68.19, + "21": 69.1, + "22": 67.78, + "23": 68.36, + "24": 68.19, + "25": 68.26, + "26": 68.06, + "27": 68.25, + "28": 68.39, + "29": 68.26, + "30": 68.05, + "31": 68.27, + "32": 68.2, + "33": 68.19, + "34": 68.02, + "35": 68.3, + "36": 68.2, + "37": 68.48, + "38": 68.23, + "39": 68.36, + "40": 67.9, + "41": 75.76, + "42": 62.04, + "43": 66.78, + "44": 67.85, + "45": 68.11, + "46": 67.92, + "47": 68.15, + "48": 67.89, + "49": 68.3 }, - "expected_e2e_ms": 257696.97, - "expected_avg_denoise_ms": 65.16, - "expected_median_denoise_ms": 65.86 + "expected_e2e_ms": 248171.5, + "expected_avg_denoise_ms": 68.98, + "expected_median_denoise_ms": 68.19 }, "wan2_1_t2v_1.3b_frame_interp_2x": { "stages_ms": { 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 690c5ba62..aa19cce01 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -96,8 +96,8 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: if server_args.lora_path: extra_args += f" --lora-path {server_args.lora_path}" - # default warmup - extra_args += f" --warmup" + if server_args.enable_warmup: + extra_args += " --warmup" for arg in server_args.extras: extra_args += f" {arg}" diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index f213b9f2a..11324977f 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -198,6 +198,7 @@ class DiffusionServerArgs: dit_offload_prefetch_size: int | float | None = None enable_cache_dit: bool = False text_encoder_cpu_offload: bool = False + enable_warmup: bool = True extras: list[str] = field(default_factory=lambda: []) @@ -726,6 +727,7 @@ if not current_platform.is_hip(): DiffusionServerArgs( model_path="tencent/Hunyuan3D-2", modality="3d", + enable_warmup=False, ), HUNYUAN3D_SHAPE_sampling_params, ),