[diffusion] Speed up LingBot high-quality VAE decode (#36024)

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -225,6 +225,9 @@ class PipelineConfig:
     vae_config: VAEConfig = field(default_factory=VAEConfig)
     vae_precision: str = "fp32"
     vae_decode_precision: str | None = None
+    # Optional request-scoped override. The loader keeps the reference decode
+    # dtype resident so lossless requests never consume pre-rounded weights.
+    vae_decode_precision_high: str | None = None
     vae_tiling: bool = True
     # Bounds the attention grid the diffusion decoder's stages see, which is
     # what makes a full-length decode tractable.

python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py

@@ -282,6 +282,7 @@ class LingBotWorldI2VConfig(Wan2_2_I2V_A14B_Config):
     dit_config: DiTConfig = field(default_factory=LingBotWorldVideoConfig)
     flow_shift: float | None = 10.0
     boundary_ratio: float | None = 0.947
+    vae_decode_precision_high: str = "bf16"
     text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
     preprocess_text_funcs: tuple[Callable[[str], str] | None, ...] = field(
         default_factory=lambda: (lingbot_prompt_clean,)

python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py

@@ -126,6 +126,7 @@ def build_realtime_sampling_params(
         output_path=request.output_path,
         output_compression=request.output_compression,
         output_quality=request.output_quality,
+        quality=getattr(request, "quality", None),
         condition_inputs=chunk_inputs.condition_inputs,
         realtime_chunk_size=chunk_size,
     )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -321,7 +321,11 @@ class DecodingStage(PipelineStage):
         # load vae if not already loaded (used for memory constrained devices)
         self.load_model()
 
-        vae_dtype = resolve_decode_precision(server_args, self.component_name)
+        vae_dtype = resolve_decode_precision(
+            server_args,
+            self.component_name,
+            quality=batch.sampling_params.quality,
+        )
         with self.use_declared_component(
             component_name=self.component_name,
             module=self.vae,

python/sglang/multimodal_gen/runtime/utils/precision.py

@@ -29,7 +29,12 @@ 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:
+def resolve_decode_precision(
+    server_args,
+    component_name: str = "vae",
+    *,
+    quality: str | None = None,
+) -> torch.dtype:
     pipeline_config = server_args.pipeline_config
     if component_name in ("audio_vae", "vocoder"):
         return resolve_precision(
@@ -38,6 +43,11 @@ def resolve_decode_precision(server_args, component_name: str = "vae") -> torch.
             precision_attr="audio_vae_precision",
         )
 
+    if quality == "high":
+        high_precision = getattr(pipeline_config, "vae_decode_precision_high", None)
+        if high_precision is not None:
+            return precision_to_dtype(high_precision, "vae_decode_precision_high")
+
     decode_precision = getattr(pipeline_config, "vae_decode_precision", None)
     if decode_precision is not None:
         return precision_to_dtype(decode_precision, "vae_decode_precision")

python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py

@@ -6,6 +6,7 @@ import torch
 
 from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import (
     LingBotWorldCausalDMDConfig,
+    LingBotWorldV2CausalDMDConfig,
 )
 from sglang.multimodal_gen.configs.quantization.qvg_kv import QVGKVQuantArgs
 from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
@@ -37,6 +38,12 @@ from sglang.multimodal_gen.runtime.realtime.states import RealtimeCausalDiTState
 LINGBOT_INTERACTIVE_KV_WINDOW_ENV = "SGLANG_LINGBOT_ENABLE_INTERACTIVE_KV_WINDOW"
 
 
+def test_lingbot_quality_high_uses_bf16_vae_decode_only():
+    for config in (LingBotWorldCausalDMDConfig(), LingBotWorldV2CausalDMDConfig()):
+        assert config.vae_decode_precision == "fp32"
+        assert config.vae_decode_precision_high == "bf16"
+
+
 def test_lingbot_denoising_stage_does_not_own_realtime_cache_refs():
     stage = LingBotWorldCausalDMDDenoisingStage.__new__(
         LingBotWorldCausalDMDDenoisingStage

python/sglang/multimodal_gen/test/unit/test_precision_consistency.py

@@ -106,6 +106,7 @@ class TestDiffusionPrecisionConsistency(unittest.TestCase):
         config = {
             "vae_precision": "fp16",
             "vae_decode_precision": None,
+            "vae_decode_precision_high": None,
             "audio_vae_precision": "bf16",
             "dit_precision": "fp32",
             "image_encoder_precision": "fp16",
@@ -139,8 +140,28 @@ class TestDiffusionPrecisionConsistency(unittest.TestCase):
             resolve_decode_precision(self._server_args(vae_decode_precision="bf16")),
             torch.bfloat16,
         )
+        self.assertEqual(
+            resolve_decode_precision(
+                self._server_args(vae_decode_precision_high="bf16"),
+                quality="high",
+            ),
+            torch.bfloat16,
+        )
+        self.assertEqual(
+            resolve_decode_precision(
+                self._server_args(vae_decode_precision_high="bf16"),
+                quality="lossless",
+            ),
+            torch.float16,
+        )
         with self.assertRaisesRegex(ValueError, "Unsupported vae_decode_precision"):
             resolve_decode_precision(self._server_args(vae_decode_precision="fp8"))
+        with self.assertRaisesRegex(
+            ValueError, "Unsupported vae_decode_precision_high"
+        ):
+            resolve_decode_precision(
+                self._server_args(vae_decode_precision_high="fp8"), quality="high"
+            )
 
     def test_component_precision_mapping(self):
         server_args = self._server_args()

python/sglang/multimodal_gen/test/unit/test_video_api_profiling.py

@@ -2,8 +2,13 @@ from types import SimpleNamespace
 from unittest.mock import patch
 
 from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    RealtimeVideoGenerationsRequest,
     VideoGenerationsRequest,
 )
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import (
+    RealtimeChunkInputs,
+    build_realtime_sampling_params,
+)
 from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import (
     _build_video_sampling_params,
 )
@@ -50,3 +55,29 @@ def test_video_api_forwards_profiling_options():
     assert kwargs["num_profiled_timesteps"] == 3
     assert kwargs["profile_all_stages"] is False
     assert kwargs["quality"] == "high"
+
+
+def test_realtime_video_api_forwards_sampling_quality():
+    request = RealtimeVideoGenerationsRequest(
+        type="init",
+        prompt="profile this realtime request",
+        first_frame="cat.png",
+        quality="high",
+    )
+    chunk_inputs = RealtimeChunkInputs(prompt=request.prompt)
+
+    with patch(
+        "sglang.multimodal_gen.runtime.entrypoints.openai.realtime."
+        "realtime_adapter.build_sampling_params",
+        side_effect=lambda request_id, **kwargs: kwargs,
+    ):
+        kwargs = build_realtime_sampling_params(
+            "realtime-profile-request",
+            request=request,
+            chunk_inputs=chunk_inputs,
+            num_frames=9,
+            num_inference_steps=4,
+            chunk_size=9,
+        )
+
+    assert kwargs["quality"] == "high"