From 9866fe910b4b5ae0c677ad374252e1b2a5fb584a Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 24 Aug 2026 14:13:03 +0800 Subject: [PATCH] [diffusion] Speed up LingBot high-quality VAE decode (#36024) --- .../LingBot-World/LingBot-World-2.0.mdx | 1 + .../diffusion/LingBot-World/LingBot-World.mdx | 1 + .../configs/pipeline_configs/base.py | 3 ++ .../configs/pipeline_configs/lingbot_world.py | 1 + .../openai/realtime/realtime_adapter.py | 1 + .../runtime/pipelines_core/stages/decoding.py | 6 +++- .../multimodal_gen/runtime/utils/precision.py | 12 ++++++- .../realtime/test_lingbot_causal_denoising.py | 7 +++++ .../test/unit/test_precision_consistency.py | 21 +++++++++++++ .../test/unit/test_video_api_profiling.py | 31 +++++++++++++++++++ 10 files changed, 82 insertions(+), 2 deletions(-) diff --git a/docs/cookbook/diffusion/LingBot-World/LingBot-World-2.0.mdx b/docs/cookbook/diffusion/LingBot-World/LingBot-World-2.0.mdx index c26f33c3d..73f4c5071 100644 --- a/docs/cookbook/diffusion/LingBot-World/LingBot-World-2.0.mdx +++ b/docs/cookbook/diffusion/LingBot-World/LingBot-World-2.0.mdx @@ -84,6 +84,7 @@ Send this MessagePack map immediately after the WebSocket opens. | `num_inference_steps` | integer | No | Denoising steps per chunk. LingBot defaults to `4` when omitted. | | `guidance_scale` | number | No | Classifier-free guidance scale. Realtime LingBot commonly uses `1`. | | `negative_prompt` | string | No | Negative prompt passed to the diffusion pipeline. | +| `quality` | `"lossless"`, `"high"` | No | `lossless` keeps FP32 VAE decode. `high` uses the validated BF16 decode path for lower per-chunk latency. | | `max_chunks` | integer | No | Stop after this many chunks. Omit for a continuous session. | | `realtime_causal_sink_size` | integer | No | Number of sink frames/tokens retained in the causal attention window. | | `realtime_causal_kv_cache_num_frames` | integer | No | Number of recent frames retained in the causal KV cache window. | diff --git a/docs/cookbook/diffusion/LingBot-World/LingBot-World.mdx b/docs/cookbook/diffusion/LingBot-World/LingBot-World.mdx index 1f9cbb83b..9ffc1c326 100644 --- a/docs/cookbook/diffusion/LingBot-World/LingBot-World.mdx +++ b/docs/cookbook/diffusion/LingBot-World/LingBot-World.mdx @@ -79,6 +79,7 @@ Send this MessagePack map immediately after the WebSocket opens. | `num_inference_steps` | integer | No | Denoising steps per chunk. LingBot defaults to `4` when omitted. | | `guidance_scale` | number | No | Classifier-free guidance scale. Realtime LingBot commonly uses `1`. | | `negative_prompt` | string | No | Negative prompt passed to the diffusion pipeline. | +| `quality` | `"lossless"`, `"high"` | No | `lossless` keeps FP32 VAE decode. `high` uses the validated BF16 decode path for lower per-chunk latency. | | `max_chunks` | integer | No | Stop after this many chunks. Omit for a continuous session. | | `realtime_causal_sink_size` | integer | No | Number of sink frames/tokens retained in the causal attention window. | | `realtime_causal_kv_cache_num_frames` | integer | No | Number of recent frames retained in the causal KV cache window. | diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 8931d3d08..6bdd0b901 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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. diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py b/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py index f150dd6a3..bb9889df2 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py +++ b/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,) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py index 6b6c71ddb..55446272b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py +++ b/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, ) 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 14590065b..53cc5ed56 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/utils/precision.py b/python/sglang/multimodal_gen/runtime/utils/precision.py index 28db6d4b5..e61f5ede6 100644 --- a/python/sglang/multimodal_gen/runtime/utils/precision.py +++ b/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") diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py b/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py index 6dbd75522..890f3afb9 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py +++ b/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 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 b25b13c36..25afca467 100644 --- a/python/sglang/multimodal_gen/test/unit/test_precision_consistency.py +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_video_api_profiling.py b/python/sglang/multimodal_gen/test/unit/test_video_api_profiling.py index c12b21df3..f1631d4a5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_video_api_profiling.py +++ b/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"