diff --git a/docs/docs/sglang-diffusion/api/openai_api.mdx b/docs/docs/sglang-diffusion/api/openai_api.mdx index acbb6bb4b..a68aa9e82 100644 --- a/docs/docs/sglang-diffusion/api/openai_api.mdx +++ b/docs/docs/sglang-diffusion/api/openai_api.mdx @@ -70,6 +70,10 @@ The server implements an OpenAI-compatible Images API under the `/v1/images` nam **Endpoint:** `POST /v1/images/generations` +#### Request quality + +`quality` selects a model-owned sampling level when that model advertises one: use `lossless` for the reference path or `high` for a validated accelerated path. Omit it (or send OpenAI's default `auto`) to keep the runtime default. It is distinct from `output_quality`, which controls only output-file compression. The same extension is accepted by image edits and video requests. + **Python Example (b64_json response):** ```python Python diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 5c43879b2..73bb468f3 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -76,6 +76,11 @@ def _get_request_field_or_extra(request, field_name): return _get_extra_field(request, field_name) +def _runtime_sampling_quality(quality: str | None) -> str | None: + """Keep OpenAI's automatic default out of SGLang's sampling contract.""" + return None if quality in (None, "auto") else quality + + def _parse_extra_container(value: Any) -> dict[str, Any]: if isinstance(value, str): try: @@ -309,6 +314,7 @@ async def generations( use_system_prompt=_get_extra_field(request, "use_system_prompt"), use_guardrails=_get_extra_field(request, "use_guardrails"), enable_teacache=request.enable_teacache, + quality=_runtime_sampling_quality(request.quality), output_compression=request.output_compression, output_quality=request.output_quality, diffusers_kwargs=request.diffusers_kwargs, @@ -418,6 +424,7 @@ async def edits( guidance_scale: Optional[float] = Form(None), true_cfg_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), + quality: Optional[str] = Form(None), output_quality: Optional[str] = Form("default"), output_compression: Optional[int] = Form(None), enable_teacache: Optional[bool] = Form(False), @@ -478,6 +485,7 @@ async def edits( num_inference_steps=num_inference_steps, enable_teacache=enable_teacache, num_frames=num_frames, + quality=_runtime_sampling_quality(quality), output_compression=output_compression, output_quality=output_quality, enable_upscaling=enable_upscaling, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index 15e637f95..cb773f77d 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -104,6 +104,7 @@ _MULTIPART_EXTRA_FORM_FIELDS = ( "action_normalization", "condition_frame_indexes_vision", "condition_video_keep", + "quality", ) @@ -371,6 +372,7 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque "upscaling_model_path": request.upscaling_model_path, "upscaling_scale": request.upscaling_scale, "output_path": request.output_path, + "quality": _extra_value(request, "quality"), "output_compression": request.output_compression, "output_quality": request.output_quality, "perf_dump_path": request.perf_dump_path, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/fast_path_gate.py b/python/sglang/multimodal_gen/runtime/models/vaes/fast_path_gate.py new file mode 100644 index 000000000..48beb4df1 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/vaes/fast_path_gate.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared decode-scoped gate for optional VAE fast paths.""" + +from contextlib import contextmanager +from weakref import WeakKeyDictionary + +import torch.nn as nn + + +class VaeFastPathGate: + """Mutable flag shared by the wrappers installed on one VAE.""" + + __slots__ = ("enabled",) + + def __init__(self) -> None: + self.enabled = False + + +_VAE_FAST_PATH_GATES: WeakKeyDictionary[nn.Module, VaeFastPathGate] = ( + WeakKeyDictionary() +) + + +def register_vae_fast_path_gate(vae: nn.Module, gate: VaeFastPathGate) -> None: + _VAE_FAST_PATH_GATES[vae] = gate + + +@contextmanager +def use_vae_fast_path(vae: nn.Module, enabled: bool): + """Enable an installed VAE fast path for one decode and always reset it.""" + gate = _VAE_FAST_PATH_GATES.get(vae) + if gate is None: + yield + return + + previous_enabled = gate.enabled + gate.enabled = enabled + try: + yield + finally: + gate.enabled = previous_enabled diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py b/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py index 326987689..9f8d27fe3 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py @@ -8,9 +8,8 @@ decoder module family (``ResnetBlock2D`` GroupNorm+SiLU chains, All rewrites are mathematically exact re-associations of the original operators. Wrappers are installed once at VAE load and dispatch on a -request-scoped :class:`VaeFastPathGate` (published as -``_sgl_vae_fast_path_gate``): ``quality == "high"`` runs the fast paths, the -``"lossless"`` default runs the original module path bit-for-bit. +decode-scoped :class:`VaeFastPathGate`: ``quality == "high"`` runs the fast +paths, the ``"lossless"`` default runs the original module path bit-for-bit. - channels_last: run the decoder in NHWC so cuDNN convs skip the transpose kernels; parameter layout is swapped at decode entry to match the gate. @@ -28,6 +27,10 @@ import torch import torch.nn as nn import torch.nn.functional as F +from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import ( + VaeFastPathGate, + register_vae_fast_path_gate, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -43,19 +46,6 @@ except ImportError: # pragma: no cover _HAS_TRITON = False -class VaeFastPathGate: - """Mutable fast-path flag shared by every wrapper of one VAE; enabled by - ``DecodingStage`` while decoding a ``quality == "high"`` request.""" - - __slots__ = ("enabled",) - - def __init__(self) -> None: - self.enabled = False - - -GATE_ATTR = "_sgl_vae_fast_path_gate" - - # --------------------------------------------------------------------------- # Fuse A: two-pass GroupNorm(+SiLU) fusion (channels_last Triton kernel) # --------------------------------------------------------------------------- @@ -330,7 +320,7 @@ def _decoder_layout_forward(self, *args, **kwargs): memory_format=(torch.channels_last if want_cl else torch.contiguous_format) ) self._sgl_channels_last = want_cl - logger.info( + logger.debug( "%s: decoder switched to %s layout.", self._sgl_label, "channels_last (NHWC)" if want_cl else "contiguous (NCHW)", @@ -396,7 +386,7 @@ def _install_decoder_fast_paths(vae: nn.Module, label: str) -> nn.Module: m._sgl_folded_v = None m.forward = MethodType(_attn_fast_forward, m) n_norm = _install_norm_silu(decoder, ResnetBlock2D, gate) - setattr(vae, GATE_ATTR, gate) + register_vae_fast_path_gate(vae, gate) logger.info( "%s: installed quality-gated decoder fast paths " "(channels_last dispatch, %d fused upsamplers, %d fast attention " diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py b/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py index bd3934430..d9f5857f1 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py @@ -3,20 +3,19 @@ Fuses every decoder ``WanRMS_norm -> SiLU`` chain into one Triton kernel on the channels_last_3d layout. Wrappers are installed once at VAE load and -dispatch on a request-scoped :class:`VaeFastPathGate` (published as -``_sgl_vae_fast_path_gate``): ``quality == "high"`` runs the fused kernel -(not bitwise-identical to aten, hence gated), the ``"lossless"`` default -runs the original module path bit-for-bit. Install is all-or-nothing and -fail-closed. +dispatch on a decode-scoped :class:`VaeFastPathGate`: ``quality == "high"`` +runs the fused kernel (not bitwise-identical to aten, hence gated), the +``"lossless"`` default runs the original module path bit-for-bit. Install is +all-or-nothing and fail-closed. """ import torch import torch.nn as nn import torch.nn.functional as F -from sglang.multimodal_gen.runtime.models.vaes.flux2_vae_cuda_opt import ( - GATE_ATTR, +from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import ( VaeFastPathGate, + register_vae_fast_path_gate, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -133,7 +132,7 @@ def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module: n_norm = _install_norm_silu(decoder, gate) if n_norm is None: return vae - setattr(vae, GATE_ATTR, gate) + register_vae_fast_path_gate(vae, gate) logger.info( "Wan VAE: installed quality-gated fast path (%d RMSNorm+SiLU fusions).", n_norm, 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 b94e2334d..5534530af 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -20,6 +20,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im ComponentUse, ) from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE +from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import ( + use_vae_fast_path, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, @@ -312,14 +315,7 @@ class DecodingStage(PipelineStage): assert vae is not None self.vae = vae - # Request-scoped VAE fast-path gate (see flux2_vae_cuda_opt): - # quality == "high" opts this decode into the near-lossless fast - # paths; the "lossless" default keeps the bit-exact original - # module path. VAEs without installed wrappers have no gate. - gate = getattr(vae, "_sgl_vae_fast_path_gate", None) - if gate is not None: - gate.enabled = getattr(batch.sampling_params, "quality", None) == "high" - try: + with use_vae_fast_path(vae, batch.sampling_params.quality == "high"): frames = self.decode(batch.latents, server_args, vae_dtype=vae_dtype) # decode trajectory latents if needed @@ -348,9 +344,6 @@ class DecodingStage(PipelineStage): trajectory_decoded = [decoded_tensor[:, i] for i in range(T)] else: trajectory_decoded = None - finally: - if gate is not None: - gate.enabled = False frames = server_args.pipeline_config.post_decoding(frames, server_args) diff --git a/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py b/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py index 07a7a09b0..112fe2b67 100644 --- a/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py @@ -10,6 +10,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import ( _fallback_image_urls, _get_response_resize, _raise_if_image_variant_not_found, + _runtime_sampling_quality, _select_image_variant_cloud_url, _select_image_variant_path, ) @@ -40,6 +41,13 @@ def test_url_response_returns_one_item_per_output_path(): ] +def test_runtime_sampling_quality_preserves_the_openai_default(): + assert _runtime_sampling_quality(None) is None + assert _runtime_sampling_quality("auto") is None + assert _runtime_sampling_quality("lossless") == "lossless" + assert _runtime_sampling_quality("high") == "high" + + def test_image_response_includes_resize_for_every_output(): response = _build_image_response_kwargs( ["first.png", "second.png"], diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_fast_path_gate.py b/python/sglang/multimodal_gen/test/unit/test_vae_fast_path_gate.py new file mode 100644 index 000000000..55f378759 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_vae_fast_path_gate.py @@ -0,0 +1,21 @@ +import torch.nn as nn + +from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import ( + VaeFastPathGate, + register_vae_fast_path_gate, + use_vae_fast_path, +) + + +def test_vae_fast_path_gate_is_decode_scoped_and_nestable(): + vae = nn.Module() + gate = VaeFastPathGate() + register_vae_fast_path_gate(vae, gate) + + with use_vae_fast_path(vae, True): + assert gate.enabled + with use_vae_fast_path(vae, False): + assert not gate.enabled + assert gate.enabled + + assert not gate.enabled 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 cda5891e1..c12b21df3 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 @@ -22,6 +22,7 @@ def test_video_api_forwards_profiling_options(): profile=True, num_profiled_timesteps=3, profile_all_stages=False, + quality="high", ) server_args = SimpleNamespace( backend="auto", @@ -48,3 +49,4 @@ def test_video_api_forwards_profiling_options(): assert kwargs["profile"] is True assert kwargs["num_profiled_timesteps"] == 3 assert kwargs["profile_all_stages"] is False + assert kwargs["quality"] == "high"