[diffusion] chore: enable channels-last 3D VAE convs by default (#23200)

This commit is contained in:
Xiaoyu Zhang
2026-05-04 22:59:31 +08:00
committed by GitHub
parent 05aed5e1d5
commit 4b6d44641b
3 changed files with 14 additions and 2 deletions
+3 -2
View File
@@ -56,7 +56,8 @@ if TYPE_CHECKING:
# model loading
SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True
SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND: str | None = None
SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: bool = False
SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: bool = True
SGLANG_USE_CUDA_HUNYUANVIDEO_GROUP_NORM_SILU: bool = False
SGLANG_USE_ROCM_VAE: bool = False
SGLANG_USE_ROCM_CUDNN_BENCHMARK: bool = False
@@ -248,7 +249,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
# taken for each stage
"SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"),
"SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D": _lazy_bool(
"SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", "false"
"SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", "true"
),
# ================== cache-dit Env Vars ==================
# Enable cache-dit acceleration for DiT inference
@@ -297,6 +297,9 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = {
),
},
"mova_360p_ring1_uly2": {
ComponentType.VAE: ComponentSkip(
"Representative MOVA VAE accuracy is covered by mova_360p_tp2; ring/ulysses topology does not exercise a distinct VAE component"
),
ComponentType.TRANSFORMER: ComponentSkip(
"HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
),
@@ -752,6 +752,14 @@ def _run_staged_native_component_accuracy_case(
engine_cls.clear_memory()
ref = ref.to(device=device, dtype=torch.bfloat16).eval()
if component == ComponentType.VAE:
from sglang.multimodal_gen import envs
from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import (
_convert_conv3d_weights_to_channels_last_3d,
)
if torch.cuda.is_available() and envs.SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D:
_convert_conv3d_weights_to_channels_last_3d(ref)
ref_call = profile.prepare_reference_call(ref, inputs)
ref_autocast = (
torch.autocast(