diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 1e90ca3a8..5edae2a08 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -224,6 +224,7 @@ class PipelineConfig: # VAE configuration vae_config: VAEConfig = field(default_factory=VAEConfig) vae_precision: str = "fp32" + vae_decode_precision: str | None = None vae_tiling: bool = True vae_slicing: bool = False vae_sp: bool = True @@ -798,6 +799,17 @@ class PipelineConfig: choices=["fp32", "fp16", "bf16"], help="Precision for VAE", ) + parser.add_argument( + f"--{prefix_with_dot}vae-decode-precision", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}vae_decode_precision", + default=PipelineConfig.vae_decode_precision, + choices=["fp32", "fp16", "bf16"], + help=( + "Optional decode-only VAE precision override. " + "Defaults to --vae-precision when unset." + ), + ) parser.add_argument( f"--{prefix_with_dot}vae-tiling", action=StoreBoolean, diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py index 730f17b97..dcd23479f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py @@ -18,6 +18,7 @@ class LongLive2T2VConfig(Wan2_2_TI2V_5B_Config): is_causal: bool = True task_type: ModelTaskType = ModelTaskType.TI2V vae_precision: str = "bf16" + vae_decode_precision: str = "bf16" flow_shift: float | None = 5.0 dmd_denoising_steps: list[int] | None = field( diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py index f8daf9bfc..893851e69 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py @@ -89,6 +89,7 @@ class WanT2V480PConfig(PipelineConfig): # Precision for each component precision: str = "bf16" vae_precision: str = "fp32" + vae_decode_precision: str = "bf16" text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",)) def __post_init__(self): @@ -239,6 +240,7 @@ class Wan2_2_TI2V_5B_Config(WanT2V480PConfig, WanI2VCommonConfig): flow_shift: float | None = 5.0 task_type: ModelTaskType = ModelTaskType.TI2V expand_timesteps: bool = True + vae_decode_precision: str = "fp32" # ti2v, 5B vae_stride = (4, 16, 16) @@ -269,6 +271,7 @@ class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): flow_shift: float | None = 12.0 boundary_ratio: float | None = 0.875 + vae_decode_precision: str = "fp32" def __post_init__(self) -> None: self.dit_config.boundary_ratio = self.boundary_ratio @@ -279,6 +282,7 @@ class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): class Wan2_2_I2V_A14B_Config(WanI2V720PConfig): flow_shift: float | None = 5.0 boundary_ratio: float | None = 0.900 + vae_decode_precision: str = "fp32" def __post_init__(self) -> None: super().__post_init__() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 874b4e3f5..58d590eed 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -165,6 +165,7 @@ async def get_models(request: Request): "task_type": server_args.pipeline_config.task_type.name, "dit_precision": server_args.pipeline_config.dit_precision, "vae_precision": server_args.pipeline_config.vae_precision, + "vae_decode_precision": server_args.pipeline_config.vae_decode_precision, } if model_info: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index 60559c218..d739ebe0f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -34,7 +34,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.precision import ( autocast_context, autocast_enabled, - resolve_precision, + resolve_decode_precision, temporary_module_dtype, ) from sglang.multimodal_gen.runtime.utils.torch_compile import ( @@ -117,9 +117,7 @@ class DecodingStage(PipelineStage): def component_uses( self, server_args: ServerArgs, stage_name: str | None = None ) -> list[ComponentUse]: - vae_dtype = resolve_precision( - server_args, self.component_name, precision_attr="vae_precision" - ) + vae_dtype = resolve_decode_precision(server_args, self.component_name) stage_name = self._component_stage_name(stage_name) return [ ComponentUse( @@ -204,7 +202,9 @@ class DecodingStage(PipelineStage): latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents) server_args: Configuration containing: - disable_autocast: Whether to disable automatic mixed precision (default: False) - - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") + - pipeline_config.vae_decode_precision: optional decode-only + VAE precision ("fp32", "fp16", "bf16") + - pipeline_config.vae_precision: fallback VAE precision - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency Returns: @@ -212,10 +212,8 @@ class DecodingStage(PipelineStage): normalized to [0, 1] range and moved to CPU as float32 """ latents = latents.to(get_local_torch_device()) - # Setup VAE precision from user policy. - vae_dtype = resolve_precision( - server_args, self.component_name, precision_attr="vae_precision" - ) + # The caller resolves the decode-only override before component use so + # residency and execution agree on the target dtype. vae_autocast_enabled = autocast_enabled(vae_dtype, server_args.disable_autocast) # scale and shift @@ -281,9 +279,7 @@ class DecodingStage(PipelineStage): # load vae if not already loaded (used for memory constrained devices) self.load_model() - vae_dtype = resolve_precision( - server_args, self.component_name, precision_attr="vae_precision" - ) + vae_dtype = resolve_decode_precision(server_args, self.component_name) with self.use_declared_component( component_name=self.component_name, module=self.vae, diff --git a/python/sglang/multimodal_gen/runtime/utils/precision.py b/python/sglang/multimodal_gen/runtime/utils/precision.py index cfb8b894e..238dacdf5 100644 --- a/python/sglang/multimodal_gen/runtime/utils/precision.py +++ b/python/sglang/multimodal_gen/runtime/utils/precision.py @@ -29,6 +29,25 @@ def resolve_precision( return precision_to_dtype(precision, field_name or precision_attr) +def resolve_decode_precision(server_args, component_name: str = "vae") -> torch.dtype: + pipeline_config = server_args.pipeline_config + if component_name in ("audio_vae", "vocoder"): + return resolve_precision( + server_args, + component_name, + precision_attr="audio_vae_precision", + ) + + decode_precision = getattr(pipeline_config, "vae_decode_precision", None) + if decode_precision is not None: + return precision_to_dtype(decode_precision, "vae_decode_precision") + return resolve_precision( + server_args, + component_name, + precision_attr="vae_precision", + ) + + def resolve_component_precision(server_args, module_name: str) -> Optional[torch.dtype]: pipeline_config = getattr(server_args, "pipeline_config", None) if pipeline_config is None: diff --git a/python/sglang/multimodal_gen/test/unit/test_decoding_stage_parallelism.py b/python/sglang/multimodal_gen/test/unit/test_decoding_stage_parallelism.py index 1ca909efd..edc376c0e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_decoding_stage_parallelism.py +++ b/python/sglang/multimodal_gen/test/unit/test_decoding_stage_parallelism.py @@ -2,6 +2,7 @@ import unittest from types import SimpleNamespace from unittest.mock import patch +import torch import torch.nn as nn from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( @@ -19,6 +20,19 @@ class FakeVAE(nn.Module): class TestDecodingStageParallelism(unittest.TestCase): + def test_component_use_honors_decode_precision_override(self): + stage = DecodingStage(FakeVAE()) + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_precision="fp32", + vae_decode_precision="bf16", + ) + ) + + [component_use] = stage.component_uses(server_args) + + self.assertEqual(component_use.target_dtype, torch.bfloat16) + def test_cfg_parallel_uses_replicated_decode_when_decode_group_has_multiple_ranks( self, ): diff --git a/python/sglang/multimodal_gen/test/unit/test_longlive2_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_longlive2_pipeline_config.py index 3646cbce1..59578dfa4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_longlive2_pipeline_config.py +++ b/python/sglang/multimodal_gen/test/unit/test_longlive2_pipeline_config.py @@ -17,6 +17,10 @@ class TestLongLive2AdjustNumFrames(unittest.TestCase): def test_rounds_to_causal_block_aligned_latents(self): self.assertEqual(self.config.adjust_num_frames(65), 61) + def test_preserves_bf16_vae_decode_precision(self): + self.assertEqual(self.config.vae_precision, "bf16") + self.assertEqual(self.config.vae_decode_precision, "bf16") + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_precision_consistency.py b/python/sglang/multimodal_gen/test/unit/test_precision_consistency.py index d52f824c2..b25b13c36 100644 --- a/python/sglang/multimodal_gen/test/unit/test_precision_consistency.py +++ b/python/sglang/multimodal_gen/test/unit/test_precision_consistency.py @@ -70,6 +70,7 @@ autocast_enabled = precision.autocast_enabled get_module_dtype = precision.get_module_dtype precision_to_dtype = precision.precision_to_dtype resolve_component_precision = precision.resolve_component_precision +resolve_decode_precision = precision.resolve_decode_precision resolve_precision = precision.resolve_precision temporary_module_dtype = precision.temporary_module_dtype @@ -104,6 +105,7 @@ class TestDiffusionPrecisionConsistency(unittest.TestCase): def _server_args(self, **overrides): config = { "vae_precision": "fp16", + "vae_decode_precision": None, "audio_vae_precision": "bf16", "dit_precision": "fp32", "image_encoder_precision": "fp16", @@ -128,6 +130,18 @@ class TestDiffusionPrecisionConsistency(unittest.TestCase): with self.assertRaisesRegex(ValueError, "Unsupported custom_precision"): precision_to_dtype("fp8", "custom_precision") + def test_decode_precision_override_and_fallback(self): + self.assertEqual( + resolve_decode_precision(self._server_args()), + torch.float16, + ) + self.assertEqual( + resolve_decode_precision(self._server_args(vae_decode_precision="bf16")), + torch.bfloat16, + ) + with self.assertRaisesRegex(ValueError, "Unsupported vae_decode_precision"): + resolve_decode_precision(self._server_args(vae_decode_precision="fp8")) + def test_component_precision_mapping(self): server_args = self._server_args() expected = { diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index b4f326198..5b2335312 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -688,6 +688,28 @@ class TestWarmupImageIsModelValid(unittest.TestCase): class TestOffloadDefaults(unittest.TestCase): + def test_wan_decode_precision_defaults(self): + for pipeline_config in ( + WanT2V480PConfig(), + WanI2V480PConfig(), + ): + with self.subTest(pipeline_config=pipeline_config.__class__.__name__): + self.assertEqual(pipeline_config.vae_precision, "fp32") + self.assertEqual(pipeline_config.vae_decode_precision, "bf16") + + for pipeline_config in ( + FastWan2_2_TI2V_5B_Config(), + Wan2_2_T2V_A14B_Config(), + Wan2_2_I2V_A14B_Config(), + ): + with self.subTest(pipeline_config=pipeline_config.__class__.__name__): + self.assertEqual(pipeline_config.vae_precision, "fp32") + self.assertEqual(pipeline_config.vae_decode_precision, "fp32") + + generic_config = PipelineConfig() + self.assertEqual(generic_config.vae_precision, "fp32") + self.assertIsNone(generic_config.vae_decode_precision) + def _from_dict_with_pipeline_config( self, pipeline_config,