chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/__init__.py

@@ -0,0 +1 @@
+"""Diffusion breakable CUDA graph runtime helpers."""

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/__init__.py

@@ -0,0 +1 @@
+"""Model-specific prompt padders for diffusion breakable CUDA graph."""

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/ideogram.py

@@ -0,0 +1,131 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0
+# ==============================================================================
+"""Ideogram-4 breakable CUDA graph (BCG) prompt padding."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
+    prompt_padding as bcg_utils,
+)
+from sglang.multimodal_gen.runtime.layers.attention import DynamicVarlenMaskMeta
+
+_SEQUENCE_PADDING_INDICATOR = -1
+_OUTPUT_IMAGE_INDICATOR = 2
+_LLM_TOKEN_INDICATOR = 3
+_DYNAMIC_MASK_META_ATTR = "_sglang_bcg_ideogram_attn_mask_meta"
+
+
+def is_ideogram_transformer(current_model: Any, call_kwargs: dict) -> bool:
+    return (
+        bcg_utils.transformer_class_name_matches(current_model, "ideogram")
+        and "llm_features" in call_kwargs
+        and "x" in call_kwargs
+        and "indicator" in call_kwargs
+        and "position_ids" in call_kwargs
+    )
+
+
+def _unwrap_model(current_model: Any) -> Any:
+    for attr in ("module", "_orig_mod"):
+        wrapped = getattr(current_model, attr, None)
+        if wrapped is not None:
+            current_model = wrapped
+    return current_model
+
+
+def _dynamic_mask_meta(current_model: Any) -> DynamicVarlenMaskMeta:
+    model = _unwrap_model(current_model)
+    meta = getattr(model, _DYNAMIC_MASK_META_ATTR, None)
+    if not isinstance(meta, DynamicVarlenMaskMeta):
+        meta = DynamicVarlenMaskMeta()
+        setattr(model, _DYNAMIC_MASK_META_ATTR, meta)
+    return meta
+
+
+def _first_indicator(call_kwargs: dict) -> torch.Tensor | None:
+    indicator = bcg_utils.first_tensor(call_kwargs.get("indicator"))
+    if not torch.is_tensor(indicator) or indicator.dim() < 2:
+        return None
+    return indicator
+
+
+def _text_and_image_lengths(indicator: torch.Tensor) -> tuple[int, int] | None:
+    row = indicator[0]
+    if not torch.any(row == _LLM_TOKEN_INDICATOR):
+        return None
+    image_positions = (row == _OUTPUT_IMAGE_INDICATOR).nonzero(as_tuple=False)
+    if image_positions.numel() == 0:
+        return None
+    text_seq = int(image_positions[0].item())
+    if text_seq <= 0:
+        return None
+    image_seq = int(row.numel()) - text_seq
+    if image_seq <= 0:
+        return None
+    return text_seq, image_seq
+
+
+def _pad_total_dim(obj: Any, *, source: int, target: int, value: float = 0) -> Any:
+    return bcg_utils.pad_nested_dim(
+        obj, dim=1, source=source, target=target, value=value
+    )
+
+
+def pad_ideogram_prompt_kwargs(
+    call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
+) -> dict:
+    indicator = _first_indicator(call_kwargs)
+    if indicator is None:
+        return call_kwargs
+
+    lengths = _text_and_image_lengths(indicator)
+    if lengths is None:
+        return call_kwargs
+    text_seq, image_seq = lengths
+
+    bucket = bcg_utils.select_text_bucket(text_seq, buckets)
+    if bucket is None:
+        return call_kwargs
+
+    source_total = text_seq + image_seq
+    target_total = bucket + image_seq
+    out = dict(call_kwargs)
+
+    if source_total < target_total:
+        for key in ("llm_features", "x"):
+            if key in out and out[key] is not None:
+                out[key] = _pad_total_dim(
+                    out[key], source=source_total, target=target_total
+                )
+        if out.get("position_ids") is not None:
+            out["position_ids"] = _pad_total_dim(
+                out["position_ids"], source=source_total, target=target_total
+            )
+        if out.get("segment_ids") is not None:
+            out["segment_ids"] = _pad_total_dim(
+                out["segment_ids"],
+                source=source_total,
+                target=target_total,
+                value=_SEQUENCE_PADDING_INDICATOR,
+            )
+        if out.get("indicator") is not None:
+            out["indicator"] = _pad_total_dim(
+                out["indicator"], source=source_total, target=target_total
+            )
+        if out.get("attn_mask") is not None:
+            out["attn_mask"] = _pad_total_dim(
+                out["attn_mask"], source=source_total, target=target_total
+            )
+
+    if out.get("attn_mask") is not None:
+        out["attn_mask_meta"] = _dynamic_mask_meta(current_model)
+
+    return out
+
+
+bcg_utils.register_prompt_padder(is_ideogram_transformer, pad_ideogram_prompt_kwargs)

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/qwen_image.py

@@ -0,0 +1,102 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Qwen-Image breakable CUDA graph (BCG) prompt padding.
+
+Qwen-Image / Qwen-Image-Edit carry text length on dim 1 of
+``encoder_hidden_states`` and a separate ``freqs_cis`` text-rope cache plus
+``txt_seq_lens``; they may not pass an explicit prompt mask, so this padder
+synthesizes one. Registered with the base denoising stage's padder registry.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
+    prompt_padding as bcg_utils,
+)
+
+
+def is_qwen_transformer(current_model: Any, call_kwargs: dict) -> bool:
+    return (
+        bcg_utils.transformer_class_name_matches(current_model, "qwen")
+        and "txt_seq_lens" in call_kwargs
+        and "freqs_cis" in call_kwargs
+    )
+
+
+def pad_qwen_prompt_kwargs(
+    call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
+) -> dict:
+    ehs = call_kwargs.get("encoder_hidden_states")
+    ehs_tensor = bcg_utils.first_tensor(ehs)
+    if not torch.is_tensor(ehs_tensor) or ehs_tensor.dim() < 2:
+        return call_kwargs
+
+    seq = ehs_tensor.shape[1]
+    bucket = bcg_utils.select_text_bucket(seq, buckets)
+    if bucket is None:
+        return call_kwargs
+
+    out = dict(call_kwargs)
+    if seq < bucket:
+        out["encoder_hidden_states"] = bcg_utils.pad_nested_dim(
+            ehs, dim=1, source=seq, target=bucket
+        )
+        if (
+            "encoder_hidden_states_2" in out
+            and out["encoder_hidden_states_2"] is not None
+        ):
+            out["encoder_hidden_states_2"] = bcg_utils.pad_nested_dim(
+                out["encoder_hidden_states_2"], dim=1, source=seq, target=bucket
+            )
+
+    mask = out.get("encoder_hidden_states_mask")
+    if mask is None:
+        mask = torch.ones(
+            ehs_tensor.shape[:2],
+            device=ehs_tensor.device,
+            dtype=torch.bool,
+        )
+    if mask is not None:
+        out["encoder_hidden_states_mask"] = bcg_utils.pad_nested_dim(
+            mask, dim=1, source=seq, target=bucket
+        )
+
+    if "encoder_attention_mask" in out and out["encoder_attention_mask"] is not None:
+        out["encoder_attention_mask"] = bcg_utils.pad_nested_dim(
+            out["encoder_attention_mask"], dim=1, source=seq, target=bucket
+        )
+
+    freqs_cis = out.get("freqs_cis")
+    if isinstance(freqs_cis, tuple) and len(freqs_cis) == 2:
+        img_cache, txt_cache = freqs_cis
+        txt_cache = bcg_utils.pad_nested_dim(
+            txt_cache, dim=0, source=seq, target=bucket
+        )
+        out["freqs_cis"] = (img_cache, txt_cache)
+    elif isinstance(freqs_cis, list) and len(freqs_cis) == 2:
+        img_cache, txt_cache = freqs_cis
+        txt_cache = bcg_utils.pad_nested_dim(
+            txt_cache, dim=0, source=seq, target=bucket
+        )
+        out["freqs_cis"] = [img_cache, txt_cache]
+
+    out["txt_seq_lens"] = bcg_utils.bucket_txt_seq_lens(out.get("txt_seq_lens"), bucket)
+    return out
+
+
+bcg_utils.register_prompt_padder(is_qwen_transformer, pad_qwen_prompt_kwargs)

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/zimage.py

@@ -0,0 +1,169 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0
+# ==============================================================================
+"""Z-Image breakable CUDA graph (BCG) prompt padding."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
+    prompt_padding as bcg_utils,
+)
+
+
+def is_zimage_transformer(current_model: Any, call_kwargs: dict) -> bool:
+    return (
+        bcg_utils.transformer_class_name_matches(current_model, "zimage")
+        and "encoder_hidden_states" in call_kwargs
+        and "freqs_cis" in call_kwargs
+    )
+
+
+def _first_caption_tensor(encoder_hidden_states: Any) -> torch.Tensor | None:
+    tensor = bcg_utils.first_tensor(encoder_hidden_states)
+    if not torch.is_tensor(tensor):
+        return None
+    if tensor.dim() == 2:
+        return tensor
+    if tensor.dim() == 3:
+        return tensor[0]
+    return None
+
+
+def _caption_seq_len(tensor: torch.Tensor) -> int:
+    if tensor.dim() == 2:
+        return int(tensor.shape[0])
+    if tensor.dim() == 3:
+        return int(tensor.shape[1])
+    raise ValueError("Z-Image caption tensor must have rank 2 or 3")
+
+
+def _pad_caption(obj: Any, *, target: int) -> Any:
+    if torch.is_tensor(obj):
+        if obj.dim() == 2:
+            return bcg_utils.pad_tensor_dim(obj, 0, target)
+        if obj.dim() == 3:
+            return bcg_utils.pad_tensor_dim(obj, 1, target)
+        return obj
+    if isinstance(obj, list):
+        return [_pad_caption(item, target=target) for item in obj]
+    if isinstance(obj, tuple):
+        return tuple(_pad_caption(item, target=target) for item in obj)
+    return obj
+
+
+def _unwrap_model(current_model: Any) -> Any:
+    for attr in ("module", "_orig_mod"):
+        wrapped = getattr(current_model, attr, None)
+        if wrapped is not None:
+            current_model = wrapped
+    return current_model
+
+
+def _build_caption_freqs(current_model: Any, *, target: int, device: torch.device):
+    rotary_emb = getattr(_unwrap_model(current_model), "rotary_emb", None)
+    if rotary_emb is None:
+        return None
+
+    axes = [
+        torch.arange(1, target + 1, dtype=torch.int32, device=device),
+        torch.zeros(target, dtype=torch.int32, device=device),
+        torch.zeros(target, dtype=torch.int32, device=device),
+    ]
+    cap_pos_ids = torch.stack(axes, dim=-1)
+    return rotary_emb(cap_pos_ids)
+
+
+def _pad_caption_freqs(freqs_cis: Any, current_model: Any, *, target: int) -> Any:
+    if not isinstance(freqs_cis, (tuple, list)) or len(freqs_cis) != 2:
+        return freqs_cis
+
+    cap_cache, image_cache = freqs_cis
+    cap_tensor = bcg_utils.first_tensor(cap_cache)
+    if torch.is_tensor(cap_tensor) and cap_tensor.dim() >= 1:
+        cap_freqs = _build_caption_freqs(
+            current_model, target=target, device=cap_tensor.device
+        )
+        if cap_freqs is not None:
+            cap_cache = cap_freqs
+
+    if isinstance(freqs_cis, tuple):
+        return (cap_cache, image_cache)
+    return [cap_cache, image_cache]
+
+
+def _caption_mask(
+    call_kwargs: dict, *, caption: torch.Tensor, seq: int, bucket: int
+) -> torch.Tensor:
+    mask = bcg_utils.first_tensor(call_kwargs.get("encoder_hidden_states_mask"))
+    if not torch.is_tensor(mask):
+        mask = bcg_utils.first_tensor(call_kwargs.get("encoder_attention_mask"))
+    if torch.is_tensor(mask):
+        if mask.dim() == 1:
+            mask = mask[:seq].unsqueeze(0)
+        elif mask.dim() >= 2:
+            mask = mask[:, :seq]
+        mask = mask.to(device=caption.device, dtype=torch.bool)
+    else:
+        batch = int(caption.shape[0]) if caption.dim() == 3 else 1
+        mask = torch.ones((batch, seq), device=caption.device, dtype=torch.bool)
+    return bcg_utils.pad_tensor_dim(mask, 1, bucket)
+
+
+def pad_zimage_prompt_kwargs(
+    call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
+) -> dict:
+    caption = _first_caption_tensor(call_kwargs.get("encoder_hidden_states"))
+    if caption is None:
+        return call_kwargs
+
+    seq = _caption_seq_len(caption)
+    cap_freq = None
+    freqs_cis = call_kwargs.get("freqs_cis")
+    if isinstance(freqs_cis, (tuple, list)) and len(freqs_cis) == 2:
+        cap_freq = bcg_utils.first_tensor(freqs_cis[0])
+    cap_freq_len = int(cap_freq.shape[0]) if torch.is_tensor(cap_freq) else seq
+
+    bucket = bcg_utils.select_text_bucket(max(seq, cap_freq_len), buckets)
+    if bucket is None:
+        return call_kwargs
+
+    out = {
+        key: value
+        for key, value in call_kwargs.items()
+        if key
+        in {
+            "hidden_states",
+            "timestep",
+            "guidance",
+            "encoder_hidden_states",
+            "encoder_attention_mask",
+            "encoder_hidden_states_mask",
+            "freqs_cis",
+            "image_seq_len_target",
+            "patch_size",
+            "f_patch_size",
+        }
+    }
+
+    if seq < bucket:
+        out["encoder_hidden_states"] = _pad_caption(
+            out["encoder_hidden_states"], target=bucket
+        )
+
+    caption_mask = _caption_mask(call_kwargs, caption=caption, seq=seq, bucket=bucket)
+    out["encoder_hidden_states_mask"] = caption_mask
+    out["caption_valid_lens"] = caption_mask.sum(dim=1).to(dtype=torch.long)
+    out["_use_caption_valid_mask"] = True
+    if out.get("encoder_attention_mask") is not None:
+        out["encoder_attention_mask"] = out["encoder_hidden_states_mask"]
+    out["freqs_cis"] = _pad_caption_freqs(
+        out.get("freqs_cis"), current_model, target=bucket
+    )
+    return out
+
+
+bcg_utils.register_prompt_padder(is_zimage_transformer, pad_zimage_prompt_kwargs)

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py

@@ -0,0 +1,306 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for breakable CUDA graph (BCG) prompt padding.
+
+These helpers bucket prompt-conditioning inputs by sequence length so diffusion
+DiT forward calls with different prompt lengths can reuse captured CUDA graphs.
+Model-specific padders can register custom handling under
+``breakable_cuda_graph.model_padders``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Prompt-conditioning kwarg keys, grouped by which dim carries the text length.
+PROMPT_MASK_KEYS = (
+    "encoder_attention_mask",
+    "encoder_hidden_states_mask",
+    "attention_mask",
+    "text_mask",
+    "prompt_attention_mask",
+    "negative_attention_mask",
+    "prompt_embeds_mask",
+    "negative_prompt_embeds_mask",
+)
+TEXT_DIM1_KEYS = (
+    "encoder_hidden_states",
+    "encoder_hidden_states_2",
+    "encoder_attention_mask",
+    "encoder_hidden_states_mask",
+    "attention_mask",
+    "text_mask",
+    "text_ids",
+    "text_pos_ids",
+    "txt_ids",
+    "prompt_embeds",
+    "negative_prompt_embeds",
+    "prompt_attention_mask",
+    "negative_attention_mask",
+    "prompt_embeds_mask",
+    "negative_prompt_embeds_mask",
+    "audio_encoder_hidden_states",
+    "audio_encoder_attention_mask",
+)
+TEXT_DIM0_KEYS = (
+    "txt_freqs_cis",
+    "text_freqs_cis",
+)
+TEXT_SEQ_LEN_KEYS = (
+    "txt_seq_lens",
+    "text_seq_lens",
+)
+
+
+def first_tensor(obj: Any) -> torch.Tensor | None:
+    """First tensor leaf found by depth-first traversal (dicts in sorted-key
+    order), or ``None``."""
+    if torch.is_tensor(obj):
+        return obj
+    if isinstance(obj, (list, tuple)):
+        for item in obj:
+            tensor = first_tensor(item)
+            if tensor is not None:
+                return tensor
+    if isinstance(obj, dict):
+        for key in sorted(obj):
+            tensor = first_tensor(obj[key])
+            if tensor is not None:
+                return tensor
+    return None
+
+
+def select_text_bucket(seq: int, buckets: tuple[int, ...]) -> int | None:
+    """Smallest bucket that fits ``seq``; ``None`` (and a warning) when ``seq``
+    exceeds the largest bucket so the caller runs that length eagerly."""
+    for bucket in buckets:
+        if seq <= bucket:
+            return bucket
+    logger.warning(
+        "[Diffusion BCG] text length %d exceeds max bucket %d; not padding "
+        "(this length captures its own graph). Raise --bcg-text-buckets.",
+        seq,
+        buckets[-1],
+    )
+    return None
+
+
+def pad_tensor_dim(tensor: Any, dim: int, target: int, value: float = 0) -> Any:
+    if not torch.is_tensor(tensor) or tensor.dim() <= dim:
+        return tensor
+    seq = tensor.shape[dim]
+    if seq >= target:
+        return tensor
+    pad = [0, 0] * tensor.dim()
+    pad_index = 2 * (tensor.dim() - dim - 1) + 1
+    pad[pad_index] = target - seq
+    return torch.nn.functional.pad(tensor, tuple(pad), value=value)
+
+
+def pad_nested_dim(
+    obj: Any,
+    *,
+    dim: int,
+    source: int,
+    target: int,
+    value: float = 0,
+) -> Any:
+    if torch.is_tensor(obj):
+        if obj.dim() > dim and obj.shape[dim] == source:
+            return pad_tensor_dim(obj, dim, target, value)
+        return obj
+    if isinstance(obj, list):
+        return [
+            pad_nested_dim(item, dim=dim, source=source, target=target, value=value)
+            for item in obj
+        ]
+    if isinstance(obj, tuple):
+        return tuple(
+            pad_nested_dim(item, dim=dim, source=source, target=target, value=value)
+            for item in obj
+        )
+    return obj
+
+
+def bucket_txt_seq_lens(txt_seq_lens: Any, bucket: int) -> Any:
+    if txt_seq_lens is None:
+        return txt_seq_lens
+    if torch.is_tensor(txt_seq_lens):
+        return torch.full_like(txt_seq_lens, bucket)
+    if isinstance(txt_seq_lens, list):
+        return [bucket_txt_seq_lens(seq_len, bucket) for seq_len in txt_seq_lens]
+    if isinstance(txt_seq_lens, tuple):
+        return tuple(bucket_txt_seq_lens(seq_len, bucket) for seq_len in txt_seq_lens)
+    if isinstance(txt_seq_lens, int):
+        return bucket
+    return txt_seq_lens
+
+
+def prompt_seq_and_dim(call_kwargs: dict) -> tuple[int, int] | None:
+    """Return ``(text_seq_len, seq_dim)`` inferred from the prompt embeddings or
+    a prompt mask, or ``None`` when no text conditioning is present."""
+    ehs_tensor = first_tensor(call_kwargs.get("encoder_hidden_states"))
+    if torch.is_tensor(ehs_tensor) and ehs_tensor.dim() >= 2:
+        if ehs_tensor.dim() == 2:
+            return int(ehs_tensor.shape[0]), 0
+        return int(ehs_tensor.shape[1]), 1
+
+    for key in PROMPT_MASK_KEYS:
+        tensor = first_tensor(call_kwargs.get(key))
+        if torch.is_tensor(tensor) and tensor.dim() >= 2:
+            if tensor.shape[0] == 1:
+                return int(tensor.shape[1]), 1
+            return int(tensor.shape[0]), 0
+    return None
+
+
+def pad_nested_text_dim(
+    obj: Any,
+    *,
+    source: int,
+    target: int,
+    preferred_dim: int,
+) -> Any:
+    if torch.is_tensor(obj):
+        if obj.dim() > preferred_dim and obj.shape[preferred_dim] == source:
+            return pad_tensor_dim(obj, preferred_dim, target)
+        for dim in (1, 0):
+            if dim != preferred_dim and obj.dim() > dim and obj.shape[dim] == source:
+                return pad_tensor_dim(obj, dim, target)
+        return obj
+    if isinstance(obj, list):
+        return [
+            pad_nested_text_dim(
+                item, source=source, target=target, preferred_dim=preferred_dim
+            )
+            for item in obj
+        ]
+    if isinstance(obj, tuple):
+        return tuple(
+            pad_nested_text_dim(
+                item, source=source, target=target, preferred_dim=preferred_dim
+            )
+            for item in obj
+        )
+    if isinstance(obj, dict):
+        return {
+            key: pad_nested_text_dim(
+                value, source=source, target=target, preferred_dim=preferred_dim
+            )
+            for key, value in obj.items()
+        }
+    return obj
+
+
+def bucket_text_seq_lens(obj: Any, *, target: int) -> Any:
+    if isinstance(obj, int) and not isinstance(obj, bool):
+        return target
+    if isinstance(obj, list):
+        return [bucket_text_seq_lens(item, target=target) for item in obj]
+    if isinstance(obj, tuple):
+        return tuple(bucket_text_seq_lens(item, target=target) for item in obj)
+    return obj
+
+
+def pad_masked_prompt_kwargs(call_kwargs: dict, buckets: tuple[int, ...]) -> dict:
+    """Generic, model-agnostic prompt padding for models that pass a prompt
+    attention mask alongside their text embeddings."""
+    seq_and_dim = prompt_seq_and_dim(call_kwargs)
+    if seq_and_dim is None:
+        return call_kwargs
+    seq, seq_dim = seq_and_dim
+    has_mask = any(
+        first_tensor(call_kwargs.get(key)) is not None for key in PROMPT_MASK_KEYS
+    )
+    if not has_mask:
+        return call_kwargs
+    bucket = select_text_bucket(seq, buckets)
+    if bucket is None or seq == bucket:
+        return call_kwargs
+
+    out = dict(call_kwargs)
+    for key in TEXT_DIM1_KEYS:
+        if key in out and out[key] is not None:
+            out[key] = pad_nested_text_dim(
+                out[key], source=seq, target=bucket, preferred_dim=seq_dim
+            )
+    for key in TEXT_DIM0_KEYS:
+        if key in out and out[key] is not None:
+            out[key] = pad_nested_dim(out[key], dim=0, source=seq, target=bucket)
+    for key in TEXT_SEQ_LEN_KEYS:
+        if key in out and out[key] is not None:
+            out[key] = bucket_text_seq_lens(out[key], target=bucket)
+    return out
+
+
+def transformer_class_name_matches(current_model: Any, needle: str) -> bool:
+    """True when ``current_model`` (or its ``module`` / ``_orig_mod`` wrapper)
+    is a transformer whose qualified class name contains ``needle``."""
+    candidates = [current_model]
+    for attr in ("module", "_orig_mod"):
+        wrapped = getattr(current_model, attr, None)
+        if wrapped is not None:
+            candidates.append(wrapped)
+    for candidate in candidates:
+        cls = type(candidate)
+        name = f"{cls.__module__}.{cls.__qualname__}".lower()
+        if needle in name:
+            return True
+    return False
+
+
+# --- Model-specific prompt-padder registry ------------------------------- #
+# Each model that needs custom prompt padding registers a (predicate, padder)
+# pair from its own module in ``model_specific_stages`` so the base denoising
+# stage stays model-agnostic. ``padder(call_kwargs, current_model, buckets)``
+# returns the padded kwargs.
+PromptPadder = Callable[[dict, Any, tuple], dict]
+_PROMPT_PADDERS: list[tuple[Callable[[Any, dict], bool], PromptPadder]] = []
+
+
+def register_prompt_padder(
+    predicate: Callable[[Any, dict], bool], padder: PromptPadder
+) -> None:
+    _PROMPT_PADDERS.append((predicate, padder))
+
+
+def select_prompt_padder(current_model: Any, call_kwargs: dict) -> PromptPadder | None:
+    """Return the registered model-specific padder for ``current_model``, or
+    ``None`` to fall back to :func:`pad_masked_prompt_kwargs`."""
+    _ensure_model_padders_registered()
+    for predicate, padder in _PROMPT_PADDERS:
+        if predicate(current_model, call_kwargs):
+            return padder
+    return None
+
+
+_model_padders_registered = False
+
+
+def _ensure_model_padders_registered() -> None:
+    """Import the model-specific padder modules once so they register."""
+    global _model_padders_registered
+    if _model_padders_registered:
+        return
+    _model_padders_registered = True
+    from sglang.multimodal_gen.runtime.breakable_cuda_graph.model_padders import (  # noqa: F401
+        ideogram,
+        qwen_image,
+        zimage,
+    )

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/replay_token.py

@@ -0,0 +1,31 @@
+"""Replay-token tracking for diffusion BCG replays.
+
+The SRT BCG core does not stamp replays; the diffusion runner sets a fresh
+token around each graph replay so replay-local caches (e.g. varlen attention
+mask metadata in ``DynamicVarlenMaskMeta``) can be rebuilt once per replay
+while still being reused across the break points of that same replay.
+``get_current_replay_token`` returns ``None`` outside a replay (including
+during capture).
+"""
+
+import itertools
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+_current_replay_token_var: ContextVar[int | None] = ContextVar(
+    "mm_bcg_replay_token", default=None
+)
+_replay_token_counter = itertools.count(1)
+
+
+def get_current_replay_token() -> int | None:
+    return _current_replay_token_var.get()
+
+
+@contextmanager
+def replay_token_scope():
+    token = _current_replay_token_var.set(next(_replay_token_counter))
+    try:
+        yield
+    finally:
+        _current_replay_token_var.reset(token)

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py

@@ -0,0 +1,472 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Breakable CUDA graph (BCG) runner for diffusion DiT transformers.
+
+A runner wraps a callable ``nn.Module`` and turns it into an *eager runner* that
+transparently proxies every attribute to the wrapped module and, when called,
+replays a previously captured graph for the input signature — or runs the
+module eagerly when no graph was captured for that signature. Capture is an
+explicit, idempotent ``capture()`` call (driven at warmup) so that serving never
+triggers a fresh capture.
+
+This file is intentionally local to ``multimodal_gen``: diffusion reuses the
+low-level SRT BCG primitives, but the capture/replay runner owns diffusion DiT
+signature handling, static tensor buffers, prompt-bucket warmup, and fallback
+behavior.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph.replay_token import (
+    replay_token_scope,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    enable_breakable_cuda_graph,
+)
+
+# Log under the multimodal_gen namespace so the diffusion server's logging
+# config surfaces the "[Diffusion BCG] captured ..." lines.
+logger = logging.getLogger(__name__)
+
+
+def _env_int(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        logger.warning("[BCG] ignoring invalid integer %s=%r", name, raw)
+        return default
+
+
+def _env_float(name: str, default: float) -> float:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        return float(raw)
+    except ValueError:
+        logger.warning("[BCG] ignoring invalid float %s=%r", name, raw)
+        return default
+
+
+def _map_tensors(obj, fn):
+    """Rebuild ``obj`` applying ``fn`` to every tensor leaf, recursing into
+    list/tuple/dict containers; everything else passes through unchanged."""
+    if torch.is_tensor(obj):
+        return fn(obj)
+    if isinstance(obj, tuple):
+        return tuple(_map_tensors(o, fn) for o in obj)
+    if isinstance(obj, list):
+        return [_map_tensors(o, fn) for o in obj]
+    if isinstance(obj, dict):
+        return {k: _map_tensors(v, fn) for k, v in obj.items()}
+    return obj
+
+
+def _flatten_tensors(obj, out: list):
+    """Depth-first collect every tensor leaf into ``out`` (deterministic order:
+    dicts traversed in sorted-key order to match across calls)."""
+    if torch.is_tensor(obj):
+        out.append(obj)
+    elif isinstance(obj, (list, tuple)):
+        for o in obj:
+            _flatten_tensors(o, out)
+    elif isinstance(obj, dict):
+        for k in sorted(obj):
+            _flatten_tensors(obj[k], out)
+
+
+def _flatten_kwargs(kwargs: dict[str, Any]) -> list[torch.Tensor]:
+    out: list[torch.Tensor] = []
+    for name in sorted(kwargs):
+        _flatten_tensors(kwargs[name], out)
+    return out
+
+
+def _signature_leaf(obj: Any) -> Any:
+    if torch.is_tensor(obj):
+        return ("tensor", tuple(obj.shape), str(obj.dtype))
+    if isinstance(obj, tuple):
+        return ("tuple", tuple(_signature_leaf(o) for o in obj))
+    if isinstance(obj, list):
+        return ("list", tuple(_signature_leaf(o) for o in obj))
+    if isinstance(obj, dict):
+        return (
+            "dict",
+            tuple((k, _signature_leaf(obj[k])) for k in sorted(obj)),
+        )
+    if obj is None or isinstance(obj, (bool, int, float, str)):
+        return ("const", obj)
+    return ("object", type(obj).__module__, type(obj).__qualname__, id(obj))
+
+
+def _signature_kwargs(kwargs: dict[str, Any]) -> tuple:
+    return tuple((name, _signature_leaf(kwargs[name])) for name in sorted(kwargs))
+
+
+def _signature_summary_leaf(sig: Any, *, depth: int = 0) -> Any:
+    if not isinstance(sig, tuple) or not sig:
+        return sig
+
+    tag = sig[0]
+    if tag == "tensor":
+        return sig
+    if tag == "const":
+        value = sig[1]
+        if isinstance(value, str) and len(value) > 64:
+            value = value[:61] + "..."
+        return (tag, value)
+    if tag == "object":
+        return sig[:3]
+    if depth >= 2:
+        return (tag, "...")
+    if tag in ("tuple", "list"):
+        items = sig[1]
+        preview = tuple(
+            _signature_summary_leaf(item, depth=depth + 1) for item in items[:4]
+        )
+        if len(items) > 4:
+            preview += (("...", len(items) - 4),)
+        return (tag, len(items), preview)
+    if tag == "dict":
+        items = sig[1]
+        preview = tuple(
+            (key, _signature_summary_leaf(value, depth=depth + 1))
+            for key, value in items[:4]
+        )
+        if len(items) > 4:
+            preview += (("...", len(items) - 4),)
+        return (tag, len(items), preview)
+    return sig
+
+
+def _signature_summary(key: tuple) -> tuple:
+    return tuple((name, _signature_summary_leaf(value)) for name, value in key[:16]) + (
+        (("...", len(key) - 16),) if len(key) > 16 else ()
+    )
+
+
+def _clone_output(out: Any) -> Any:
+    if torch.is_tensor(out):
+        return out.clone()
+    if isinstance(out, tuple):
+        return tuple(_clone_output(o) for o in out)
+    if isinstance(out, list):
+        return [_clone_output(o) for o in out]
+    return out
+
+
+@dataclass
+class _CaptureEntry:
+    graph: BreakableCUDAGraph
+    # full captured kwargs with persistent static buffers at every tensor leaf
+    static_kwargs: dict[str, Any]
+    # the same static buffers, flattened in _flatten_kwargs order (replay copies
+    # live tensors into these positionally)
+    static_leaves: list[torch.Tensor]
+    output: Any
+    num_segments: int
+
+
+class _CaptureRejected(RuntimeError):
+    pass
+
+
+class BaseBreakableCudaGraphRunner:
+    """Eager runner around ``transformer`` with an explicit capture/replay API.
+
+    The capture/replay contract:
+
+    * :meth:`capture` captures a BCG graph for the given input signature, once
+      (idempotent). It is intended to be driven at warmup so that every
+      signature served later is already captured.
+    * :meth:`replay` copies live inputs into the captured static buffers and
+      replays the graph, returning a clone of the captured output.
+    * :meth:`__call__` is the *eager runner*: it replays when a graph exists for
+      the signature and otherwise runs ``transformer`` eagerly. It never
+      captures, so serving never pays a capture cost.
+
+    Any attribute not defined on the runner is proxied to ``transformer`` so the
+    runner can stand in for the wrapped module ("other functions directly
+    pass").
+    """
+
+    def __init__(
+        self,
+        transformer: nn.Module,
+        device: torch.device,
+        pool=None,
+    ) -> None:
+        self.transformer = transformer
+        self.device = device
+        self.device_module = torch.get_device_module(device)
+        # One shared mempool across all captured graphs/segments so per-block
+        # intermediates can be reclaimed and weak-ref'd safely.
+        self._pool = (
+            pool if pool is not None else self.device_module.graph_pool_handle()
+        )
+        self._capture_stream = self.device_module.Stream(device=device)
+        self.entries: dict[tuple, _CaptureEntry] = {}
+        # Signatures we have given up capturing (capture raised); run eager.
+        self._blocked: set[tuple] = set()
+        self._disabled_reason: str | None = None
+        self.max_entries = max(0, _env_int("SGLANG_DIFFUSION_BCG_MAX_ENTRIES", 32))
+        self.max_segments = max(0, _env_int("SGLANG_DIFFUSION_BCG_MAX_SEGMENTS", 128))
+
+    def __getattr__(self, name: str) -> Any:
+        # Only reached for attributes the runner itself does not define; proxy
+        # them to the wrapped transformer so callers can treat the runner as a
+        # transparent stand-in. Use __dict__ to avoid recursing through
+        # __getattr__ before ``transformer`` is assigned in __init__.
+        try:
+            transformer = self.__dict__["transformer"]
+        except KeyError as e:  # pragma: no cover - during/ before __init__
+            raise AttributeError(name) from e
+        return getattr(transformer, name)
+
+    # ------------------------------------------------------------------ #
+    # Public capture / replay API
+    # ------------------------------------------------------------------ #
+    @torch.no_grad()
+    def capture(self, **kwargs) -> bool:
+        """Capture a graph for ``kwargs``'s signature if not already captured.
+
+        Idempotent: returns ``True`` when a graph is available for the
+        signature afterwards (already captured or newly captured), ``False``
+        when capture is disabled/blocked or failed (the caller then runs eager).
+        """
+        if self._disabled_reason is not None:
+            return False
+        key = self._signature(kwargs)
+        if key in self._blocked:
+            return False
+        if key in self.entries:
+            return True
+        try:
+            entry = self._capture(kwargs, key)
+        except Exception as e:  # noqa: BLE001 — never break generation on capture
+            logger.warning(
+                "[Diffusion BCG] capture failed for signature %s (%s); "
+                "this signature will run eager.",
+                _signature_summary(key),
+                e,
+            )
+            self._blocked.add(key)
+            return False
+        self.entries[key] = entry
+        self._evict_entries_if_needed()
+        return True
+
+    def _should_capture_on_call(self, key: tuple) -> bool:
+        """Whether ``__call__`` may lazily capture an unseen signature.
+
+        Base runners only ever capture through the explicit :meth:`capture`
+        API, so this returns ``False``: serving never records a fresh graph.
+        Subclasses gate lazy capture on a warmup window (see the diffusion
+        runner) so warmup can capture by simply driving the forward as usual.
+        """
+        return False
+
+    @torch.no_grad()
+    def __call__(self, **kwargs) -> Any:
+        """Eager runner: replay a captured graph, else run ``transformer``.
+
+        While serving this never captures, so no new graph is recorded once
+        warmup is over. During the warmup window subclasses opt into lazy
+        capture via :meth:`_should_capture_on_call`.
+        """
+        if self._disabled_reason is not None:
+            return self.transformer(**kwargs)
+        key = self._signature(kwargs)
+        entry = self.entries.get(key)
+        if entry is None:
+            if not self._should_capture_on_call(key):
+                return self.transformer(**kwargs)
+            if not self.capture(**kwargs):
+                return self.transformer(**kwargs)
+            entry = self.entries[key]
+        return self.replay(entry, kwargs)
+
+    def replay(self, entry: _CaptureEntry, kwargs: dict[str, Any]) -> Any:
+        live_leaves = _flatten_kwargs(kwargs)
+        if len(live_leaves) != len(entry.static_leaves):
+            # Structure changed under a matching shape key — should not happen;
+            # fall back to eager rather than copy mismatched buffers.
+            return self.transformer(**kwargs)
+        for buf, live in zip(entry.static_leaves, live_leaves):
+            buf.copy_(live, non_blocking=True)
+        with replay_token_scope():
+            entry.graph.replay()
+        # Clone so the caller can hold the result across the next replay / the
+        # other CFG branch (which shares this static output buffer when shapes
+        # match). The clone is one cheap DtoD copy relative to the full DiT.
+        return _clone_output(entry.output)
+
+    # ------------------------------------------------------------------ #
+    # Internals
+    # ------------------------------------------------------------------ #
+    def _signature(self, kwargs: dict[str, Any]) -> tuple:
+        """Capture key for tensor leaves and non-tensor control values.
+
+        Tensor leaves are keyed by shape+dtype so their values can change per
+        replay. Non-tensor leaves are baked into the captured Python control
+        flow, so simple constants must be part of the key as well. Mutable
+        objects are keyed by identity to avoid replaying a graph whose eager
+        break points still reference a previous request's state object.
+        """
+        return _signature_kwargs(kwargs)
+
+    def _empty_cache(self) -> None:
+        empty_cache = getattr(self.device_module, "empty_cache", None)
+        if callable(empty_cache):
+            empty_cache()
+
+    @staticmethod
+    def _drop_entry(entry: _CaptureEntry) -> None:
+        entry.graph._break_fns.clear()
+        entry.graph._segments.clear()
+        entry.static_kwargs.clear()
+        entry.static_leaves.clear()
+        entry.output = None
+
+    def reset(self, *, disabled_reason: str | None = None) -> None:
+        for entry in self.entries.values():
+            self._drop_entry(entry)
+        self.entries.clear()
+        self._blocked.clear()
+        self._pool = None
+        self._empty_cache()
+        if disabled_reason is not None:
+            self._disabled_reason = disabled_reason
+
+    def _capture_limit_reason(self, entry: _CaptureEntry) -> str | None:
+        if self.max_segments and entry.num_segments > self.max_segments:
+            return (
+                f"captured {entry.num_segments} segments, above "
+                f"SGLANG_DIFFUSION_BCG_MAX_SEGMENTS={self.max_segments}"
+            )
+        return None
+
+    def _evict_entries_if_needed(self) -> None:
+        if not self.max_entries:
+            return
+        while len(self.entries) > self.max_entries:
+            evicted_key = next(iter(self.entries))
+            entry = self.entries.pop(evicted_key)
+            self._drop_entry(entry)
+            logger.info(
+                "[Diffusion BCG] evicted oldest capture for signature %s "
+                "(SGLANG_DIFFUSION_BCG_MAX_ENTRIES=%d)",
+                _signature_summary(evicted_key),
+                self.max_entries,
+            )
+        self._empty_cache()
+
+    def _capture(self, kwargs: dict[str, Any], key: tuple) -> _CaptureEntry:
+        if self._pool is None:
+            self._pool = self.device_module.graph_pool_handle()
+
+        # Persistent static buffers at every tensor leaf; bake non-tensors.
+        def _to_static(t: torch.Tensor) -> torch.Tensor:
+            # Static buffers live on the capture device. A CPU input (e.g. a
+            # scalar timestep/sigma or an index tensor built on the host)
+            # would otherwise force a CPU->CUDA copy inside the captured
+            # region, which is illegal; place its buffer on the device so the
+            # only host->device copy happens here, before capture, and replay
+            # is device-to-device.
+            if t.device.type == "cpu":
+                buf = torch.empty(t.shape, dtype=t.dtype, device=self.device)
+            else:
+                buf = torch.empty_like(t)
+            buf.copy_(t)
+            return buf
+
+        static_kwargs = {
+            name: _map_tensors(v, _to_static) for name, v in kwargs.items()
+        }
+        static_leaves = _flatten_kwargs(static_kwargs)
+
+        # Warm up on the capture stream so cuBLAS/cuDNN/Triton workspaces and
+        # any lazy JIT are materialized before capture (mirrors the LLM runner
+        # and torch.cuda.make_graphed_callables).
+        self.device_module.synchronize()
+        with self.device_module.stream(self._capture_stream):
+            for _ in range(2):
+                self.transformer(**static_kwargs)
+        self._capture_stream.synchronize()
+        self.device_module.synchronize()
+
+        graph = BreakableCUDAGraph()
+        with enable_breakable_cuda_graph():
+            with BreakableCUDAGraphCapture(
+                cuda_graph=graph, pool=self._pool, stream=self._capture_stream
+            ):
+                output = self.transformer(**static_kwargs)
+        self.device_module.synchronize()
+
+        logger.info(
+            "[Diffusion BCG] captured %d segment(s), %d tensor input(s) for "
+            "signature %s",
+            len(graph._segments),
+            len(static_leaves),
+            _signature_summary(key),
+        )
+        entry = _CaptureEntry(
+            graph=graph,
+            static_kwargs=static_kwargs,
+            static_leaves=static_leaves,
+            output=output,
+            num_segments=len(graph._segments),
+        )
+        limit_reason = self._capture_limit_reason(entry)
+        if limit_reason is not None:
+            self._drop_entry(entry)
+            self.reset(disabled_reason=limit_reason)
+            raise _CaptureRejected(
+                f"{limit_reason}; disabling this BCG runner and using eager"
+            )
+        return entry
+
+
+class DiffusionBreakableCudaGraphRunner(BaseBreakableCudaGraphRunner):
+    """Capture/replay a diffusion DiT ``transformer`` with BCG.
+
+    Unknown attributes proxy to the wrapped transformer, so the runner can
+    stand in for the module while only intercepting ``forward`` calls.
+    """
+
+    def _should_capture_on_call(self, key) -> bool:
+        try:
+            from sglang.multimodal_gen.runtime.managers.forward_context import (
+                get_forward_context,
+            )
+
+            forward_batch = get_forward_context().forward_batch
+        except Exception:
+            return False
+        return bool(getattr(forward_batch, "is_warmup", False))

python/sglang/multimodal_gen/runtime/cache/__init__.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Cache acceleration module for SGLang-diffusion
+
+This module provides various caching strategies to accelerate
+diffusion transformer (DiT) inference:
+
+- TeaCache: Temporal similarity-based caching for diffusion models
+- cache-dit integration: Block-level caching with DBCache and TaylorSeer
+
+"""
+
+from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
+    CacheDitConfig,
+    enable_cache_on_dual_transformer,
+    enable_cache_on_transformer,
+    get_scm_mask,
+)
+from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheContext, TeaCacheMixin
+
+__all__ = [
+    # TeaCache (always available)
+    "TeaCacheContext",
+    "TeaCacheMixin",
+    # cache-dit integration (lazy-loaded, requires cache-dit package)
+    "CacheDitConfig",
+    "enable_cache_on_transformer",
+    "enable_cache_on_dual_transformer",
+    "get_scm_mask",
+]

python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py

@@ -0,0 +1,683 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+cache-dit integration module for SGLang DiT pipelines.
+
+This module provides helper functions to enable cache-dit acceleration
+on transformer modules in SGLang's modular pipeline architecture.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+import torch.distributed as dist
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_ring_parallel_world_size,
+    get_tp_world_size,
+    get_ulysses_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+import cache_dit
+from cache_dit import (
+    BlockAdapter,
+    DBCacheConfig,
+    ForwardPattern,
+    ParamsModifier,
+    TaylorSeerCalibratorConfig,
+    steps_mask,
+)
+from cache_dit.caching.block_adapters import BlockAdapterRegister
+from cache_dit.parallelism import ParallelismBackend, ParallelismConfig
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_dit_group
+
+_original_similarity = None
+
+
+def _patch_cache_dit_similarity():
+    from cache_dit.caching.cache_contexts import cache_manager
+
+    global _original_similarity
+    if _original_similarity is not None:
+        return
+
+    _original_similarity = cache_manager.CachedContextManager.similarity
+
+    def patched_similarity(self, t1, t2, *, threshold, parallelized=False, prefix="Fn"):
+        if not parallelized:
+            return _original_similarity(
+                self,
+                t1,
+                t2,
+                threshold=threshold,
+                parallelized=parallelized,
+                prefix=prefix,
+            )
+
+        sp_group = getattr(self, "_sglang_sp_group", None)
+        tp_group = getattr(self, "_sglang_tp_group", None)
+        tp_sp_group = getattr(self, "_sglang_tp_sp_group", None)
+        target_group = tp_sp_group or sp_group or tp_group
+
+        if target_group is None:
+            return _original_similarity(
+                self,
+                t1,
+                t2,
+                threshold=threshold,
+                parallelized=parallelized,
+                prefix=prefix,
+            )
+
+        # Adapted from https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/caching/cache_contexts/cache_manager.py#L495-L523
+        condition_thresh = self.get_important_condition_threshold()
+        if condition_thresh > 0.0:
+            raw_diff = (t1 - t2).abs()
+            token_m_df = raw_diff.mean(dim=-1)
+            token_m_t1 = t1.abs().mean(dim=-1)
+            token_diff = token_m_df / token_m_t1
+            condition = token_diff > condition_thresh
+            if condition.sum() > 0:
+                condition = condition.unsqueeze(-1).expand_as(raw_diff)
+                mean_diff = raw_diff[condition].mean()
+                mean_t1 = t1[condition].abs().mean()
+            else:
+                mean_diff = (t1 - t2).abs().mean()
+                mean_t1 = t1.abs().mean()
+        else:
+            mean_diff = (t1 - t2).abs().mean()
+            mean_t1 = t1.abs().mean()
+
+        dist.all_reduce(mean_diff, op=dist.ReduceOp.AVG, group=target_group)
+        dist.all_reduce(mean_t1, op=dist.ReduceOp.AVG, group=target_group)
+
+        diff = (mean_diff / mean_t1).item()
+        self.add_residual_diff(diff)
+        return diff < threshold
+
+    cache_manager.CachedContextManager.similarity = patched_similarity
+
+
+def _build_parallelism_config(
+    sp_group: Optional[torch.distributed.ProcessGroup],
+    tp_group: Optional[torch.distributed.ProcessGroup],
+):
+    if sp_group is None and tp_group is None:
+        return None
+
+    ulysses_size = None
+    ring_size = None
+    if sp_group is not None:
+        ulysses_size = get_ulysses_parallel_world_size()
+        ring_size = get_ring_parallel_world_size()
+
+    tp_size = None
+    if tp_group is not None:
+        tp_size = get_tp_world_size()
+
+    return ParallelismConfig(
+        backend=ParallelismBackend.AUTO,
+        ulysses_size=ulysses_size,
+        ring_size=ring_size,
+        tp_size=tp_size,
+    )
+
+
+def _mark_transformer_parallelized(transformer, config, sp_group, tp_group):
+    if config is None:
+        return
+
+    transformer._is_parallelized = True
+    transformer._parallelism_config = config
+
+
+def get_scm_mask(
+    preset: str,
+    num_inference_steps: int,
+    compute_bins: Optional[List[int]] = None,
+    cache_bins: Optional[List[int]] = None,
+) -> Optional[List[int]]:
+    """
+    Get SCM mask using cache-dit's steps_mask().
+
+    This is a thin wrapper that delegates to cache-dit's built-in
+    steps_mask() function which handles all presets and scaling logic.
+
+    Args:
+        preset: Preset name ("none", "slow", "medium", "fast", "ultra").
+        compute_bins: Custom compute bins (overrides preset).
+        cache_bins: Custom cache bins (overrides preset).
+
+    Returns:
+        SCM mask list (1=compute, 0=cache), or None if disabled.
+    """
+    if preset == "none" and not (compute_bins and cache_bins):
+        return None
+
+    # Use cache-dit's steps_mask() directly
+    mask = steps_mask(
+        compute_bins=compute_bins,
+        cache_bins=cache_bins,
+        total_steps=num_inference_steps,
+        mask_policy=preset if preset != "none" else "medium",
+    )
+
+    compute_count = sum(mask)
+    cache_count = len(mask) - compute_count
+    logger.info(
+        "SCM: generated mask with %d compute steps, %d cache steps (preset=%s)",
+        compute_count,
+        cache_count,
+        preset,
+    )
+
+    return mask
+
+
+@dataclass
+class CacheDitConfig:
+    """Configuration for cache-dit integration.
+
+    Attributes:
+        enabled: Whether to enable cache-dit acceleration.
+        Fn_compute_blocks: Number of first blocks to always compute (DBCache F).
+        Bn_compute_blocks: Number of last blocks to always compute (DBCache B).
+        max_warmup_steps: Number of warmup steps before caching starts (DBCache W).
+        residual_diff_threshold: Threshold for residual difference (DBCache R).
+        max_continuous_cached_steps: Maximum consecutive cached steps (DBCache MC).
+        enable_taylorseer: Whether to enable TaylorSeer calibrator.
+        taylorseer_order: Order of Taylor expansion (1 or 2).
+        num_inference_steps: Total number of inference steps (required for transformer-only mode).
+        steps_computation_mask: Binary mask for step-level caching (1=compute, 0=cache).
+            Generated by get_scm_mask() (wrapper around cache_dit.steps_mask()).
+        steps_computation_policy: Caching policy for SCM ("dynamic" or "static").
+    """
+
+    enabled: bool = False
+    Fn_compute_blocks: int = 1
+    Bn_compute_blocks: int = 0
+    # Use 4 as default warmup steps instead of 8 in cache-dit, thus making
+    # DBCache work for few steps distilled models, e.g., Z-Image w/ 8-steps.
+    max_warmup_steps: int = 4
+    # Use a relatively higher residual diff threshold (namely, 0.24) as default
+    # to allow more aggressive caching due to we have already applied max continuous
+    # cached steps limit, otherwise, we should use a lower threshold here like 0.12.
+    residual_diff_threshold: float = 0.24
+    max_continuous_cached_steps: int = 3
+    # TaylorSeer is not suitable for few steps distilled models, so, we choose
+    # to disable it by default. Reference:
+    # - From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers,
+    #   https://arxiv.org/pdf/2503.06923
+    # - FoCa: Forecast then Calibrate: Feature Caching as ODE for Efficient
+    #   Diffusion Transformers, https://arxiv.org/pdf/2508.16211
+    enable_taylorseer: bool = False
+    taylorseer_order: int = 1
+    num_inference_steps: Optional[int] = None
+    # SCM fields (generated by _maybe_enable_cache_dit from env configuration)
+    steps_computation_mask: Optional[List[int]] = None
+    steps_computation_policy: str = "dynamic"
+
+
+@dataclass(frozen=True)
+class DualTransformerBlockAdapterSpec:
+    """BlockAdapter metadata for dual-transformer DiT pipelines.
+
+    This describes the cache-dit-facing structure of a pair of transformers.
+    The denoising loop semantics live in DenoisingStage; this spec only covers
+    how cache-dit should find blocks and interpret each block's forward.
+    """
+
+    blocks_attr: tuple[str, str]
+    blocks_name: Optional[List[str]]
+    forward_pattern: List[ForwardPattern]
+    check_forward_pattern: bool
+    check_num_outputs: bool
+    has_separate_cfg: bool
+
+
+DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS: dict[str, DualTransformerBlockAdapterSpec] = {
+    "wan2.2": DualTransformerBlockAdapterSpec(
+        blocks_attr=("blocks", "blocks"),
+        blocks_name=None,
+        forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+        check_forward_pattern=True,
+        check_num_outputs=False,
+        has_separate_cfg=True,
+    ),
+    "ideogram4": DualTransformerBlockAdapterSpec(
+        blocks_attr=("layers", "layers"),
+        blocks_name=["layers", "layers"],
+        forward_pattern=[ForwardPattern.Pattern_3, ForwardPattern.Pattern_3],
+        check_forward_pattern=False,
+        check_num_outputs=False,
+        has_separate_cfg=False,
+    ),
+}
+
+
+# Custom BlockAdapter for DiT models absent from cache-dit's BlockAdapterRegister.
+# Value: (blocks attr, forward_pattern). forward_pattern must
+# match the block's forward signature (see cache_dit.ForwardPattern; e.g., ERNIE
+# uses Pattern_3). has_separate_cfg follows the run (passed by
+# enable_cache_on_transformer); cache-dit auto-resolves the remaining
+# fields.
+_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern]] = {
+    "ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3),
+    "Krea2Transformer2DModel": ("transformer_blocks", ForwardPattern.Pattern_3),
+}
+
+
+def _build_custom_block_adapter(
+    transformer: torch.nn.Module,
+    has_separate_cfg: bool = False,
+) -> Optional[BlockAdapter]:
+    """Build a manual BlockAdapter for a model absent from cache-dit's registry,
+    or None if the class is unknown."""
+    spec = _CUSTOM_BLOCK_ADAPTER_SPECS.get(transformer.__class__.__name__)
+    if spec is None:
+        return None
+    blocks_attr, forward_pattern = spec
+    blocks = getattr(transformer, blocks_attr, None)
+    if blocks is None:
+        raise ValueError(
+            f"Transformer {transformer.__class__.__name__} has no attribute "
+            f"{blocks_attr!r} for cache-dit blocks."
+        )
+    return BlockAdapter(
+        transformer=transformer,
+        blocks=blocks,
+        forward_pattern=forward_pattern,
+        has_separate_cfg=has_separate_cfg,
+    )
+
+
+def enable_cache_on_transformer(
+    transformer: torch.nn.Module,
+    config: CacheDitConfig,
+    model_name: str = "transformer",
+    sp_group: Optional[torch.distributed.ProcessGroup] = None,
+    tp_group: Optional[torch.distributed.ProcessGroup] = None,
+    has_separate_cfg: bool = False,
+) -> torch.nn.Module:
+    """Enable cache-dit on a transformer module, by wrapping the module with cache-dit
+
+    This function enables cache-dit acceleration using the BlockAdapterRegister
+    for pre-registered models
+
+    Args:
+        model_name: Name of the model for logging purposes.
+        sp_group: Sequence parallel process group (for Ulysses/Ring).
+        tp_group: Tensor parallel process group.
+        has_separate_cfg: Whether the run issues separate conditional/unconditional
+            passes per step (CFG). Used by custom adapters (ERNIE, Krea-2); a
+            mismatch only disables caching, never corrupts output.
+
+    """
+    if not config.enabled:
+        return transformer
+
+    if config.num_inference_steps is None:
+        raise ValueError(
+            "num_inference_steps is required for transformer-only mode. "
+            "Please provide it in CacheDitConfig."
+        )
+
+    # Prefer the standard path (transformer pre-registered in cache-dit). For
+    # models absent from the registry, fall back to a manual BlockAdapter (see
+    # _build_custom_block_adapter).
+    custom_adapter = None
+    if not BlockAdapterRegister.is_supported(transformer):
+        custom_adapter = _build_custom_block_adapter(
+            transformer, has_separate_cfg=has_separate_cfg
+        )
+        if custom_adapter is None:
+            transformer_cls_name = transformer.__class__.__name__
+            raise ValueError(
+                f"{transformer_cls_name} is not officially supported by cache-dit. "
+                "Supported cache-dit DiT families include Flux, QwenImage, HunyuanDiT, "
+                "HunyuanVideo, Wan, CogVideoX, Mochi, and others. "
+                "Please ensure your transformer belongs to one of these families or "
+                "define a custom BlockAdapter."
+            )
+
+    # Build cache config (including SCM fields if provided)
+    cache_config = DBCacheConfig(
+        num_inference_steps=config.num_inference_steps,
+        Fn_compute_blocks=config.Fn_compute_blocks,
+        Bn_compute_blocks=config.Bn_compute_blocks,
+        max_warmup_steps=config.max_warmup_steps,
+        residual_diff_threshold=config.residual_diff_threshold,
+        max_continuous_cached_steps=config.max_continuous_cached_steps,
+        # SCM fields
+        steps_computation_mask=config.steps_computation_mask,
+        steps_computation_policy=config.steps_computation_policy,
+    )
+
+    # Build calibrator config if TaylorSeer is enabled
+    calibrator_config = None
+    if config.enable_taylorseer:
+        calibrator_config = TaylorSeerCalibratorConfig(
+            taylorseer_order=config.taylorseer_order,
+        )
+
+    # Enable cache-dit on the transformer
+    logger.info(
+        "Enabling cache-dit on %s with config: Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, "
+        "TaylorSeer=%s (order=%d), steps=%d",
+        model_name,
+        config.Fn_compute_blocks,
+        config.Bn_compute_blocks,
+        config.max_warmup_steps,
+        config.residual_diff_threshold,
+        config.max_continuous_cached_steps,
+        config.enable_taylorseer,
+        config.taylorseer_order,
+        config.num_inference_steps,
+    )
+
+    # Log SCM configuration if enabled
+    if config.steps_computation_mask:
+        compute_steps = sum(config.steps_computation_mask)
+        cache_steps = len(config.steps_computation_mask) - compute_steps
+        logger.info(
+            "SCM enabled: %d compute steps, %d cache steps, policy=%s",
+            compute_steps,
+            cache_steps,
+            config.steps_computation_policy,
+        )
+
+    parallelism_config = _build_parallelism_config(sp_group, tp_group)
+    if parallelism_config is not None:
+        _patch_cache_dit_similarity()
+
+    _mark_transformer_parallelized(transformer, parallelism_config, sp_group, tp_group)
+
+    # Custom path: pass a pre-built BlockAdapter, bypassing the registry.
+    # Standard path: let enable_cache discover the registered adapter.
+    target = transformer
+    if custom_adapter is not None:
+        target = custom_adapter
+        logger.info(
+            "Enabling cache-dit on %s via custom BlockAdapter (%s).",
+            model_name,
+            custom_adapter.forward_pattern,
+        )
+    cache_dit.enable_cache(
+        target,
+        cache_config=cache_config,
+        calibrator_config=calibrator_config,
+        parallelism_config=None,
+    )
+
+    if parallelism_config is not None:
+        context_manager = getattr(transformer, "_context_manager", None)
+        if context_manager is not None:
+            context_manager._sglang_sp_group = sp_group
+            context_manager._sglang_tp_group = tp_group
+            # In mixed TP + SP (Ulysses/Ring) mode, cache-dit decisions must be consistent
+            # across the full TP×SP model-parallel slice. Prefer using SGLang's DIT group
+            # as a conservative superset group; fallback to None.
+            tp_sp_group = None
+            if sp_group is not None and tp_group is not None:
+                tp_sp_group = get_dit_group()
+
+            context_manager._sglang_tp_sp_group = tp_sp_group
+
+    return transformer
+
+
+def enable_cache_on_dual_transformer(
+    transformer: torch.nn.Module,
+    transformer_2: torch.nn.Module,
+    primary_config: CacheDitConfig,
+    secondary_config: CacheDitConfig,
+    model_name: str = "wan2.2",
+    sp_group: Optional[torch.distributed.ProcessGroup] = None,
+    tp_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> tuple[torch.nn.Module, torch.nn.Module]:
+    """Enable cache-dit on dual transformers using BlockAdapter.
+
+    For models with two transformers, cache-dit requires enabling cache on both
+    simultaneously via BlockAdapter. The two transformers may be split by denoising
+    range, or run as paired conditional/unconditional branches. This cannot be done
+    by calling enable_cache separately on each transformer.
+
+    Args:
+        primary_config: CacheDitConfig for primary transformer.
+        secondary_config: CacheDitConfig for secondary transformer.
+        sp_group: Sequence parallel process group (for Ulysses/Ring).
+        tp_group: Tensor parallel process group.
+    """
+    adapter_spec = DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS.get(model_name)
+    if adapter_spec is None:
+        raise ValueError(
+            f"Dual-transformer cache-dit is only supported for "
+            f"{sorted(DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS)}, got {model_name}."
+        )
+
+    if not primary_config.enabled:
+        return transformer, transformer_2
+
+    if primary_config.num_inference_steps is None:
+        raise ValueError(
+            "num_inference_steps is required for dual-transformer mode. "
+            "Please provide it in CacheDitConfig."
+        )
+
+    # Build DBCacheConfig for primary transformer
+    primary_cache_config = DBCacheConfig(
+        num_inference_steps=primary_config.num_inference_steps,
+        Fn_compute_blocks=primary_config.Fn_compute_blocks,
+        Bn_compute_blocks=primary_config.Bn_compute_blocks,
+        max_warmup_steps=primary_config.max_warmup_steps,
+        residual_diff_threshold=primary_config.residual_diff_threshold,
+        max_continuous_cached_steps=primary_config.max_continuous_cached_steps,
+        steps_computation_mask=primary_config.steps_computation_mask,
+        steps_computation_policy=primary_config.steps_computation_policy,
+    )
+
+    # Build DBCacheConfig for secondary transformer
+    secondary_cache_config = DBCacheConfig(
+        num_inference_steps=secondary_config.num_inference_steps,
+        Fn_compute_blocks=secondary_config.Fn_compute_blocks,
+        Bn_compute_blocks=secondary_config.Bn_compute_blocks,
+        max_warmup_steps=secondary_config.max_warmup_steps,
+        residual_diff_threshold=secondary_config.residual_diff_threshold,
+        max_continuous_cached_steps=secondary_config.max_continuous_cached_steps,
+        steps_computation_mask=secondary_config.steps_computation_mask,
+        steps_computation_policy=secondary_config.steps_computation_policy,
+    )
+
+    # Build calibrator configs if TaylorSeer is enabled
+    primary_calibrator = None
+    if primary_config.enable_taylorseer:
+        primary_calibrator = TaylorSeerCalibratorConfig(
+            taylorseer_order=primary_config.taylorseer_order,
+        )
+
+    secondary_calibrator = None
+    if secondary_config.enable_taylorseer:
+        secondary_calibrator = TaylorSeerCalibratorConfig(
+            taylorseer_order=secondary_config.taylorseer_order,
+        )
+
+    # Build ParamsModifier for each transformer
+    primary_modifier = ParamsModifier(
+        cache_config=primary_cache_config,
+        calibrator_config=primary_calibrator,
+    )
+    secondary_modifier = ParamsModifier(
+        cache_config=secondary_cache_config,
+        calibrator_config=secondary_calibrator,
+    )
+
+    # Log configuration
+    logger.info(
+        "Enabling cache-dit on %s dual transformers with BlockAdapter",
+        model_name,
+    )
+    logger.info(
+        "  Primary (transformer): Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, TaylorSeer=%s",
+        primary_config.Fn_compute_blocks,
+        primary_config.Bn_compute_blocks,
+        primary_config.max_warmup_steps,
+        primary_config.residual_diff_threshold,
+        primary_config.max_continuous_cached_steps,
+        primary_config.enable_taylorseer,
+    )
+    logger.info(
+        "  Secondary transformer: Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, TaylorSeer=%s",
+        secondary_config.Fn_compute_blocks,
+        secondary_config.Bn_compute_blocks,
+        secondary_config.max_warmup_steps,
+        secondary_config.residual_diff_threshold,
+        secondary_config.max_continuous_cached_steps,
+        secondary_config.enable_taylorseer,
+    )
+
+    # Log SCM configuration if enabled
+    if primary_config.steps_computation_mask:
+        compute_steps = sum(primary_config.steps_computation_mask)
+        cache_steps = len(primary_config.steps_computation_mask) - compute_steps
+        logger.info(
+            "  SCM enabled for primary transformer: %d compute steps, %d cache steps, policy=%s",
+            compute_steps,
+            cache_steps,
+            primary_config.steps_computation_policy,
+        )
+    if secondary_config.steps_computation_mask:
+        compute_steps = sum(secondary_config.steps_computation_mask)
+        cache_steps = len(secondary_config.steps_computation_mask) - compute_steps
+        logger.info(
+            "  SCM enabled for secondary transformer: %d compute steps, %d cache steps, policy=%s",
+            compute_steps,
+            cache_steps,
+            secondary_config.steps_computation_policy,
+        )
+
+    parallelism_config = _build_parallelism_config(sp_group, tp_group)
+    if parallelism_config is not None:
+        _patch_cache_dit_similarity()
+
+    _mark_transformer_parallelized(transformer, parallelism_config, sp_group, tp_group)
+    _mark_transformer_parallelized(
+        transformer_2, parallelism_config, sp_group, tp_group
+    )
+
+    transformer_blocks_attr, transformer_2_blocks_attr = adapter_spec.blocks_attr
+    transformer_blocks = getattr(transformer, transformer_blocks_attr, None)
+    transformer_2_blocks = getattr(transformer_2, transformer_2_blocks_attr, None)
+    if transformer_blocks is None or transformer_2_blocks is None:
+        raise ValueError(
+            f"Dual transformers for {model_name} must expose cache-dit block "
+            f"attributes {adapter_spec.blocks_attr}. "
+            f"transformer has {transformer_blocks_attr}: "
+            f"{transformer_blocks is not None}, secondary transformer has "
+            f"{transformer_2_blocks_attr}: {transformer_2_blocks is not None}"
+        )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=[transformer, transformer_2],
+            blocks=[transformer_blocks, transformer_2_blocks],
+            blocks_name=adapter_spec.blocks_name,
+            forward_pattern=adapter_spec.forward_pattern,
+            params_modifiers=[primary_modifier, secondary_modifier],
+            check_forward_pattern=adapter_spec.check_forward_pattern,
+            check_num_outputs=adapter_spec.check_num_outputs,
+            has_separate_cfg=adapter_spec.has_separate_cfg,
+        ),
+        parallelism_config=None,
+    )
+
+    if parallelism_config is not None:
+        for t in [transformer, transformer_2]:
+            context_manager = getattr(t, "_context_manager", None)
+            if context_manager is not None:
+                context_manager._sglang_sp_group = sp_group
+                context_manager._sglang_tp_group = tp_group
+                tp_sp_group = None
+                if sp_group is not None and tp_group is not None:
+                    try:
+                        tp_sp_group = get_dit_group()
+                    except Exception:
+                        tp_sp_group = None
+                context_manager._sglang_tp_sp_group = tp_sp_group
+
+    return transformer, transformer_2
+
+
+def refresh_context_on_transformer(
+    transformer: torch.nn.Module,
+    num_inference_steps: int,
+    scm_preset: str | None = None,
+    verbose: bool = False,
+) -> None:
+    """Refresh cache-dit context for transformer."""
+    steps_computation_mask = None
+    if scm_preset is not None:
+        steps_computation_mask = cache_dit.steps_mask(
+            mask_policy=scm_preset, total_steps=num_inference_steps
+        )
+    cache_dit.refresh_context(
+        transformer,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_inference_steps,
+            steps_computation_mask=steps_computation_mask,
+            steps_computation_policy=scm_preset,
+        ),
+        verbose=verbose,
+    )
+    logger.debug(f"cache-dit refreshed on transformer (steps={num_inference_steps})")
+
+
+def refresh_context_on_dual_transformer(
+    transformer: torch.nn.Module,
+    transformer_2: torch.nn.Module,
+    num_high_noise_steps: int,
+    num_low_noise_steps: int,
+    scm_preset: str | None = None,
+    verbose: bool = False,
+    steps_computation_mask: Optional[List[int]] = None,
+    steps_computation_mask_2: Optional[List[int]] = None,
+    steps_computation_policy: str | None = None,
+) -> None:
+    """Refresh cache-dit context for dual transformers."""
+    high_noise_steps_computation_mask = steps_computation_mask
+    low_noise_steps_computation_mask = steps_computation_mask_2
+    if high_noise_steps_computation_mask is None and scm_preset is not None:
+        high_noise_steps_computation_mask = cache_dit.steps_mask(
+            mask_policy=scm_preset, total_steps=num_high_noise_steps
+        )
+    if low_noise_steps_computation_mask is None and scm_preset is not None:
+        low_noise_steps_computation_mask = cache_dit.steps_mask(
+            mask_policy=scm_preset, total_steps=num_low_noise_steps
+        )
+    policy = (
+        steps_computation_policy if steps_computation_policy is not None else scm_preset
+    )
+    cache_dit.refresh_context(
+        transformer,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_high_noise_steps,
+            steps_computation_mask=high_noise_steps_computation_mask,
+            steps_computation_policy=policy,
+        ),
+        verbose=verbose,
+    )
+    cache_dit.refresh_context(
+        transformer_2,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_low_noise_steps,
+            steps_computation_mask=low_noise_steps_computation_mask,
+            steps_computation_policy=policy,
+        ),
+        verbose=verbose,
+    )
+    logger.debug(
+        f"cache-dit refreshed on dual transformers (steps={num_high_noise_steps}, {num_low_noise_steps})"
+    )

python/sglang/multimodal_gen/runtime/cache/teacache.py

@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+TeaCache: Temporal similarity-based caching for diffusion models.
+
+TeaCache accelerates diffusion inference by selectively skipping redundant
+computation when consecutive diffusion steps are similar enough. This is
+achieved by tracking the L1 distance between modulated inputs across timesteps.
+
+Key concepts:
+- Modulated input: The input to transformer blocks after timestep conditioning
+- L1 distance: Measures how different consecutive timesteps are
+- Threshold: When accumulated L1 distance exceeds threshold, force computation
+- CFG support: Separate caches for positive and negative branches
+
+References:
+- TeaCache: Accelerating Diffusion Models with Temporal Similarity
+  https://arxiv.org/abs/2411.14324
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.configs.models import DiTConfig
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams
+
+
+@dataclass
+class TeaCacheContext:
+    """Common context extracted for TeaCache skip decision.
+
+    This context is populated from the forward_batch and forward_context
+    during each denoising step, providing all information needed to make
+    cache decisions.
+
+    Attributes:
+        current_timestep: Current denoising timestep index (0-indexed).
+        num_inference_steps: Total number of inference steps.
+        do_cfg: Whether classifier-free guidance is enabled.
+        is_cfg_negative: True if currently processing negative CFG branch.
+        teacache_thresh: Threshold for accumulated L1 distance.
+        coefficients: Polynomial coefficients for L1 rescaling.
+        teacache_params: Full TeaCacheParams for model-specific access.
+    """
+
+    current_timestep: int
+    num_inference_steps: int
+    do_cfg: bool
+    is_cfg_negative: bool  # For CFG branch selection
+    teacache_thresh: float
+    coefficients: list[float]
+    teacache_params: "TeaCacheParams"  # Full params for model-specific access
+
+
+class TeaCacheMixin:
+    """
+    Mixin class providing TeaCache optimization functionality.
+
+    TeaCache accelerates diffusion inference by selectively skipping redundant
+    computation when consecutive diffusion steps are similar enough.
+
+    This mixin should be inherited by DiT model classes that want to support
+    TeaCache optimization. It provides:
+    - State management for tracking L1 distances
+    - CFG-aware caching (separate caches for positive/negative branches)
+    - Decision logic for when to compute vs. use cache
+
+    Example usage in a DiT model:
+        class MyDiT(TeaCacheMixin, BaseDiT):
+            def __init__(self, config, **kwargs):
+                super().__init__(config, **kwargs)
+                self._init_teacache_state()
+
+            def forward(self, hidden_states, timestep, ...):
+                ctx = self._get_teacache_context()
+                if ctx is not None:
+                    # Compute modulated input (model-specific, e.g., after timestep embedding)
+                    modulated_input = self._compute_modulated_input(hidden_states, timestep)
+                    is_boundary = (ctx.current_timestep == 0 or
+                                   ctx.current_timestep >= ctx.num_inference_steps - 1)
+
+                    should_calc = self._compute_teacache_decision(
+                        modulated_inp=modulated_input,
+                        is_boundary_step=is_boundary,
+                        coefficients=ctx.coefficients,
+                        teacache_thresh=ctx.teacache_thresh,
+                    )
+
+                    if not should_calc:
+                        # Use cached residual (must implement retrieve_cached_states)
+                        return self.retrieve_cached_states(hidden_states)
+
+                # Normal forward pass...
+                output = self._transformer_forward(hidden_states, timestep, ...)
+
+                # Cache states for next step
+                if ctx is not None:
+                    self.maybe_cache_states(output, hidden_states)
+
+                return output
+
+    Subclass implementation notes:
+        - `_compute_modulated_input()`: Model-specific method to compute the input
+          after timestep conditioning (used for L1 distance calculation)
+        - `retrieve_cached_states()`: Must be overridden to return cached output
+        - `maybe_cache_states()`: Override to store states for cache retrieval
+
+    Attributes:
+        cnt: Counter for tracking steps.
+        enable_teacache: Whether TeaCache is enabled.
+        previous_modulated_input: Cached modulated input for positive branch.
+        previous_residual: Cached residual for positive branch.
+        accumulated_rel_l1_distance: Accumulated L1 distance for positive branch.
+        is_cfg_negative: Whether currently processing negative CFG branch.
+        _supports_cfg_cache: Whether this model supports CFG cache separation.
+
+    CFG-specific attributes (only when _supports_cfg_cache is True):
+        previous_modulated_input_negative: Cached input for negative branch.
+        previous_residual_negative: Cached residual for negative branch.
+        accumulated_rel_l1_distance_negative: L1 distance for negative branch.
+    """
+
+    # Models that support CFG cache separation (wan/hunyuan/zimage)
+    # Models not in this set (flux/qwen) auto-disable TeaCache when CFG is enabled
+    _CFG_SUPPORTED_PREFIXES: set[str] = {"wan", "hunyuan", "zimage"}
+    config: DiTConfig
+
+    def _init_teacache_state(self) -> None:
+        """Initialize TeaCache state. Call this in subclass __init__."""
+        # Common TeaCache state
+        self.cnt = 0
+        self.enable_teacache = True
+        # Flag indicating if this model supports CFG cache separation
+        self._supports_cfg_cache = (
+            self.config.prefix.lower() in self._CFG_SUPPORTED_PREFIXES
+        )
+
+        # Always initialize positive cache fields (used in all modes)
+        self.previous_modulated_input: torch.Tensor | None = None
+        self.previous_residual: torch.Tensor | None = None
+        self.accumulated_rel_l1_distance: float = 0.0
+
+        self.is_cfg_negative = False
+        # CFG-specific fields initialized to None (created when CFG is used)
+        # These are only used when _supports_cfg_cache is True AND do_cfg is True
+        if self._supports_cfg_cache:
+            self.previous_modulated_input_negative: torch.Tensor | None = None
+            self.previous_residual_negative: torch.Tensor | None = None
+            self.accumulated_rel_l1_distance_negative: float = 0.0
+
+    def reset_teacache_state(self) -> None:
+        """Reset all TeaCache state at the start of each generation task."""
+        self.cnt = 0
+
+        # Primary cache fields (always present)
+        self.previous_modulated_input = None
+        self.previous_residual = None
+        self.accumulated_rel_l1_distance = 0.0
+        self.is_cfg_negative = False
+        self.enable_teacache = True
+        # CFG negative cache fields (always reset, may be unused)
+        if self._supports_cfg_cache:
+            self.previous_modulated_input_negative = None
+            self.previous_residual_negative = None
+            self.accumulated_rel_l1_distance_negative = 0.0
+
+    def _compute_l1_and_decide(
+        self,
+        modulated_inp: torch.Tensor,
+        coefficients: list[float],
+        teacache_thresh: float,
+    ) -> tuple[float, bool]:
+        """
+        Compute L1 distance and decide whether to calculate or use cache.
+
+        Args:
+            modulated_inp: Current timestep's modulated input.
+            coefficients: Polynomial coefficients for L1 rescaling.
+            teacache_thresh: Threshold for cache decision.
+
+        Returns:
+            Tuple of (new_accumulated_distance, should_calc).
+        """
+        prev_modulated_inp = (
+            self.previous_modulated_input_negative
+            if self.is_cfg_negative
+            else self.previous_modulated_input
+        )
+
+        # Defensive check: if previous input is not set, force calculation
+        if prev_modulated_inp is None:
+            return 0.0, True
+
+        # Compute relative L1 distance
+        diff = modulated_inp - prev_modulated_inp
+        rel_l1 = (diff.abs().mean() / prev_modulated_inp.abs().mean()).cpu().item()
+
+        # Apply polynomial rescaling
+        rescale_func = np.poly1d(coefficients)
+
+        accumulated_rel_l1_distance = (
+            self.accumulated_rel_l1_distance_negative
+            if self.is_cfg_negative
+            else self.accumulated_rel_l1_distance
+        )
+        accumulated_rel_l1_distance = accumulated_rel_l1_distance + rescale_func(rel_l1)
+
+        if accumulated_rel_l1_distance >= teacache_thresh:
+            # Threshold exceeded: force compute and reset accumulator
+            return 0.0, True
+        # Cache hit: keep accumulated distance
+        return accumulated_rel_l1_distance, False
+
+    def _compute_teacache_decision(
+        self,
+        modulated_inp: torch.Tensor,
+        is_boundary_step: bool,
+        coefficients: list[float],
+        teacache_thresh: float,
+    ) -> bool:
+        """
+        Compute cache decision for TeaCache.
+
+        Args:
+            modulated_inp: Current timestep's modulated input.
+            is_boundary_step: True for boundary timesteps that always compute.
+            coefficients: Polynomial coefficients for L1 rescaling.
+            teacache_thresh: Threshold for cache decision.
+
+        Returns:
+            True if forward computation is needed, False to use cache.
+        """
+        if not self.enable_teacache:
+            return True
+
+        if is_boundary_step:
+            new_accum, should_calc = 0.0, True
+        else:
+            new_accum, should_calc = self._compute_l1_and_decide(
+                modulated_inp=modulated_inp,
+                coefficients=coefficients,
+                teacache_thresh=teacache_thresh,
+            )
+
+        # Advance baseline and accumulator for the active branch
+        if not self.is_cfg_negative:
+            self.previous_modulated_input = modulated_inp.clone()
+            self.accumulated_rel_l1_distance = new_accum
+        elif self._supports_cfg_cache:
+            self.previous_modulated_input_negative = modulated_inp.clone()
+            self.accumulated_rel_l1_distance_negative = new_accum
+
+        return should_calc
+
+    def _get_teacache_context(self) -> TeaCacheContext | None:
+        """
+        Check TeaCache preconditions and extract common context.
+
+        Returns:
+            TeaCacheContext if TeaCache is enabled and properly configured,
+            None if should skip TeaCache logic entirely.
+        """
+        from sglang.multimodal_gen.runtime.managers.forward_context import (
+            get_forward_context,
+        )
+
+        forward_context = get_forward_context()
+        forward_batch = forward_context.forward_batch
+
+        # Early return checks
+        if (
+            forward_batch is None
+            or not forward_batch.enable_teacache
+            or forward_batch.teacache_params is None
+        ):
+            return None
+
+        teacache_params = forward_batch.teacache_params
+
+        # Extract common values
+        current_timestep = forward_context.current_timestep
+        num_inference_steps = forward_batch.num_inference_steps
+        do_cfg = forward_batch.do_classifier_free_guidance
+        is_cfg_negative = forward_batch.is_cfg_negative
+
+        # Reset at first timestep
+        if current_timestep == 0 and not self.is_cfg_negative:
+            self.reset_teacache_state()
+
+        return TeaCacheContext(
+            current_timestep=current_timestep,
+            num_inference_steps=num_inference_steps,
+            do_cfg=do_cfg,
+            is_cfg_negative=is_cfg_negative,
+            teacache_thresh=teacache_params.teacache_thresh,
+            coefficients=teacache_params.get_coefficients(),
+            teacache_params=teacache_params,
+        )
+
+    def maybe_cache_states(
+        self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor
+    ) -> None:
+        """Cache states for later retrieval. Override in subclass if needed."""
+        pass
+
+    def should_skip_forward_for_cached_states(self, **kwargs: dict[str, Any]) -> bool:
+        """Check if forward can be skipped using cached states."""
+        return False
+
+    def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Retrieve cached states. Must be implemented by subclass."""
+        raise NotImplementedError("retrieve_cached_states is not implemented")

python/sglang/multimodal_gen/runtime/disaggregation/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Disaggregation support for diffusion pipelines."""

python/sglang/multimodal_gen/runtime/disaggregation/disagg_args.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Compatibility shim for disaggregated diffusion argument helpers."""
+
+from __future__ import annotations
+
+import argparse
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.server_args.disagg import DisaggServerArgsMixin
+
+# Keep the historical disagg_args import path working.
+DISAGG_RESULT_PORT_OFFSETS = DisaggServerArgsMixin.DISAGG_RESULT_PORT_OFFSETS
+DisaggArgsMixin = DisaggServerArgsMixin
+
+
+def add_disagg_cli_args(parser: argparse.ArgumentParser) -> None:
+    """Register disaggregated-diffusion CLI args through ServerArgs."""
+
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+    ServerArgs.add_disagg_cli_args(parser)
+
+
+def convert_disagg_role_string(kwargs: dict) -> None:
+    """Convert ``disagg_role`` from string to ``RoleType`` enum in-place."""
+
+    if "disagg_role" in kwargs and isinstance(kwargs["disagg_role"], str):
+        kwargs["disagg_role"] = RoleType.from_string(kwargs["disagg_role"])

python/sglang/multimodal_gen/runtime/disaggregation/dispatch_policy.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dispatch policies for multi-instance disaggregated diffusion pipelines."""
+
+import abc
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+class DispatchPolicy(abc.ABC):
+    def __init__(self, num_instances: int):
+        if num_instances < 1:
+            raise ValueError(f"num_instances must be >= 1, got {num_instances}")
+        self._num_instances = num_instances
+
+    @property
+    def num_instances(self) -> int:
+        return self._num_instances
+
+    @abc.abstractmethod
+    def select(self, active_counts: list[int] | None = None) -> int: ...
+
+    def select_with_capacity(self, free_slots: list[int]) -> int | None:
+        """Select an instance that has free capacity, or None if all full."""
+        if not any(s > 0 for s in free_slots):
+            return None
+        return self.select(active_counts=None)
+
+    def record_completion(self, instance_id: int) -> None:
+        pass
+
+
+class RoundRobin(DispatchPolicy):
+    def __init__(self, num_instances: int):
+        super().__init__(num_instances)
+        self._lock = threading.Lock()
+        self._next = 0
+
+    def select(self, active_counts: list[int] | None = None) -> int:
+        with self._lock:
+            chosen = self._next
+            self._next = (self._next + 1) % self._num_instances
+        return chosen
+
+    def select_with_capacity(self, free_slots: list[int]) -> int | None:
+        with self._lock:
+            for _ in range(self._num_instances):
+                idx = self._next
+                self._next = (self._next + 1) % self._num_instances
+                if free_slots[idx] > 0:
+                    return idx
+            return None
+
+
+class MaxFreeSlotsFirst(DispatchPolicy):
+    """Dispatch to the instance with the most free slots."""
+
+    def __init__(self, num_instances: int, max_slots_per_instance: int = 1):
+        super().__init__(num_instances)
+        self._max_slots = max_slots_per_instance
+        self._lock = threading.Lock()
+        self._tiebreak = 0
+
+    def select(self, active_counts: list[int] | None = None) -> int:
+        with self._lock:
+            if active_counts is None or len(active_counts) != self._num_instances:
+                chosen = self._tiebreak % self._num_instances
+                self._tiebreak += 1
+                return chosen
+
+            best_id = 0
+            best_free = self._max_slots - active_counts[0]
+            for i in range(1, self._num_instances):
+                free = self._max_slots - active_counts[i]
+                if free > best_free:
+                    best_free = free
+                    best_id = i
+                elif free == best_free:
+                    if i == (self._tiebreak % self._num_instances):
+                        best_id = i
+
+            self._tiebreak += 1
+
+            if best_free <= 0:
+                logger.warning(
+                    "All %d instances are at capacity (%d slots each), "
+                    "dispatching to instance %d anyway",
+                    self._num_instances,
+                    self._max_slots,
+                    best_id,
+                )
+
+            return best_id
+
+    def select_with_capacity(self, free_slots: list[int]) -> int | None:
+        with self._lock:
+            best_id = -1
+            best_free = 0
+            for i in range(self._num_instances):
+                if free_slots[i] > best_free:
+                    best_free = free_slots[i]
+                    best_id = i
+                elif free_slots[i] == best_free and best_free > 0:
+                    if i == (self._tiebreak % self._num_instances):
+                        best_id = i
+
+            self._tiebreak += 1
+
+            if best_id < 0:
+                return None
+            return best_id
+
+
+class PoolDispatcher:
+    """Wraps three independent dispatch policies for encoder/denoiser/decoder pools."""
+
+    def __init__(
+        self,
+        num_encoders: int,
+        num_denoisers: int,
+        num_decoders: int,
+        policy_name: str = "round_robin",
+        **kwargs,
+    ):
+        self.encoder_policy = create_dispatch_policy(
+            policy_name, num_encoders, **kwargs
+        )
+        self.denoiser_policy = create_dispatch_policy(
+            policy_name, num_denoisers, **kwargs
+        )
+        self.decoder_policy = create_dispatch_policy(
+            policy_name, num_decoders, **kwargs
+        )
+
+    def select_encoder(self, active_counts: list[int] | None = None) -> int:
+        return self.encoder_policy.select(active_counts)
+
+    def select_denoiser(self, active_counts: list[int] | None = None) -> int:
+        return self.denoiser_policy.select(active_counts)
+
+    def select_decoder(self, active_counts: list[int] | None = None) -> int:
+        return self.decoder_policy.select(active_counts)
+
+    def select_encoder_with_capacity(self, free_slots: list[int]) -> int | None:
+        return self.encoder_policy.select_with_capacity(free_slots)
+
+    def select_denoiser_with_capacity(self, free_slots: list[int]) -> int | None:
+        return self.denoiser_policy.select_with_capacity(free_slots)
+
+    def select_decoder_with_capacity(self, free_slots: list[int]) -> int | None:
+        return self.decoder_policy.select_with_capacity(free_slots)
+
+
+def create_dispatch_policy(name: str, num_instances: int, **kwargs) -> DispatchPolicy:
+    policies = {
+        "round_robin": RoundRobin,
+        "max_free_slots": MaxFreeSlotsFirst,
+    }
+    cls = policies.get(name)
+    if cls is None:
+        raise ValueError(
+            f"Unknown dispatch policy '{name}'. Available: {list(policies.keys())}"
+        )
+    return cls(num_instances=num_instances, **kwargs)

python/sglang/multimodal_gen/runtime/disaggregation/metrics.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Observability metrics for disaggregated diffusion pipelines."""
+
+import threading
+import time
+from dataclasses import dataclass
+
+
+@dataclass
+class _RequestTiming:
+    start_time: float
+    stage_start: float = 0.0
+
+
+@dataclass
+class RoleStats:
+    role: str
+    requests_completed: int = 0
+    requests_failed: int = 0
+    requests_in_flight: int = 0
+    requests_timed_out: int = 0
+    queue_depth: int = 0
+    last_latency_s: float = 0.0
+    avg_latency_s: float = 0.0
+    max_latency_s: float = 0.0
+    throughput_rps: float = 0.0
+    uptime_s: float = 0.0
+
+    def to_dict(self) -> dict:
+        return {
+            "role": self.role,
+            "requests_completed": self.requests_completed,
+            "requests_failed": self.requests_failed,
+            "requests_in_flight": self.requests_in_flight,
+            "requests_timed_out": self.requests_timed_out,
+            "queue_depth": self.queue_depth,
+            "last_latency_s": round(self.last_latency_s, 4),
+            "avg_latency_s": round(self.avg_latency_s, 4),
+            "max_latency_s": round(self.max_latency_s, 4),
+            "throughput_rps": round(self.throughput_rps, 4),
+            "uptime_s": round(self.uptime_s, 1),
+        }
+
+
+class DisaggMetrics:
+    """Thread-safe metrics collector for a single disagg role."""
+
+    def __init__(self, role: str):
+        self._role = role
+        self._lock = threading.Lock()
+        self._start_time = time.monotonic()
+
+        self._completed = 0
+        self._failed = 0
+        self._timed_out = 0
+
+        self._in_flight: dict[str, _RequestTiming] = {}
+
+        self._last_latency = 0.0
+        self._max_latency = 0.0
+        self._total_latency = 0.0
+
+        self._completion_times: list[float] = []
+        self._throughput_window_s = 60.0
+
+        self._queue_depth = 0
+
+    @property
+    def role(self) -> str:
+        return self._role
+
+    def record_request_start(self, request_id: str) -> None:
+        with self._lock:
+            self._in_flight[request_id] = _RequestTiming(start_time=time.monotonic())
+
+    def record_request_complete(self, request_id: str) -> None:
+        now = time.monotonic()
+        with self._lock:
+            timing = self._in_flight.pop(request_id, None)
+            if timing is not None:
+                latency = now - timing.start_time
+                self._last_latency = latency
+                self._max_latency = max(self._max_latency, latency)
+                self._total_latency += latency
+
+            self._completed += 1
+            self._completion_times.append(now)
+            self._prune_completion_times(now)
+
+    def record_request_failed(self, request_id: str) -> None:
+        with self._lock:
+            self._in_flight.pop(request_id, None)
+            self._failed += 1
+
+    def record_request_timeout(self, request_id: str) -> None:
+        with self._lock:
+            self._in_flight.pop(request_id, None)
+            self._timed_out += 1
+
+    def update_queue_depth(self, depth: int) -> None:
+        with self._lock:
+            self._queue_depth = depth
+
+    def snapshot(self) -> RoleStats:
+        now = time.monotonic()
+        with self._lock:
+            self._prune_completion_times(now)
+            total = self._completed + self._failed
+            avg_latency = self._total_latency / total if total > 0 else 0.0
+            rps = (
+                len(self._completion_times) / self._throughput_window_s
+                if self._completion_times
+                else 0.0
+            )
+
+            return RoleStats(
+                role=self._role,
+                requests_completed=self._completed,
+                requests_failed=self._failed,
+                requests_in_flight=len(self._in_flight),
+                requests_timed_out=self._timed_out,
+                queue_depth=self._queue_depth,
+                last_latency_s=self._last_latency,
+                avg_latency_s=avg_latency,
+                max_latency_s=self._max_latency,
+                throughput_rps=rps,
+                uptime_s=now - self._start_time,
+            )
+
+    def _prune_completion_times(self, now: float) -> None:
+        cutoff = now - self._throughput_window_s
+        while self._completion_times and self._completion_times[0] < cutoff:
+            self._completion_times.pop(0)

python/sglang/multimodal_gen/runtime/disaggregation/request_state.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Request state machine for disaggregated diffusion pipelines."""
+
+import enum
+import logging
+import threading
+import time
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+
+class RequestState(enum.Enum):
+    """Lifecycle states for a disagg pipeline request.
+
+    *_WAITING: request queued, awaiting a free buffer slot.
+    *_RUNNING: request dispatched to a specific instance.
+    """
+
+    PENDING = "pending"
+    ENCODER_WAITING = "encoder_waiting"
+    ENCODER_RUNNING = "encoder_running"
+    ENCODER_DONE = "encoder_done"
+    DENOISING_WAITING = "denoising_waiting"
+    DENOISING_RUNNING = "denoising_running"
+    DENOISING_DONE = "denoising_done"
+    DECODER_WAITING = "decoder_waiting"
+    DECODER_RUNNING = "decoder_running"
+    DONE = "done"
+    FAILED = "failed"
+    TIMED_OUT = "timed_out"
+
+
+_TERMINAL_STATES = {RequestState.DONE, RequestState.FAILED, RequestState.TIMED_OUT}
+_ACTIVE_STATES = set(RequestState) - _TERMINAL_STATES
+
+# Normal (non-failure) transitions.  FAILED and TIMED_OUT are handled
+# separately in transition() — any active state can reach them.
+_VALID_TRANSITIONS: dict[RequestState, set[RequestState]] = {
+    RequestState.PENDING: {RequestState.ENCODER_WAITING, RequestState.ENCODER_RUNNING},
+    RequestState.ENCODER_WAITING: {RequestState.ENCODER_RUNNING},
+    RequestState.ENCODER_RUNNING: {RequestState.ENCODER_DONE},
+    RequestState.ENCODER_DONE: {
+        RequestState.DENOISING_WAITING,
+        RequestState.DENOISING_RUNNING,
+    },
+    RequestState.DENOISING_WAITING: {RequestState.DENOISING_RUNNING},
+    RequestState.DENOISING_RUNNING: {RequestState.DENOISING_DONE},
+    RequestState.DENOISING_DONE: {
+        RequestState.DECODER_WAITING,
+        RequestState.DECODER_RUNNING,
+    },
+    RequestState.DECODER_WAITING: {RequestState.DECODER_RUNNING},
+    RequestState.DECODER_RUNNING: {RequestState.DONE},
+}
+
+
+@dataclass
+class RequestRecord:
+    request_id: str
+    state: RequestState = RequestState.PENDING
+    submit_time: float = field(default_factory=time.monotonic)
+    last_transition_time: float = field(default_factory=time.monotonic)
+    encoder_instance: int | None = None
+    denoiser_instance: int | None = None
+    decoder_instance: int | None = None
+    error: str | None = None
+
+    def elapsed_s(self) -> float:
+        return time.monotonic() - self.submit_time
+
+    def is_terminal(self) -> bool:
+        return self.state in _TERMINAL_STATES
+
+
+class RequestTracker:
+    """Thread-safe tracker for request state machines."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._requests: dict[str, RequestRecord] = {}
+
+    def submit(self, request_id: str) -> RequestRecord:
+        with self._lock:
+            if request_id in self._requests:
+                raise ValueError(f"Duplicate request_id: {request_id}")
+            record = RequestRecord(request_id=request_id)
+            self._requests[request_id] = record
+            return record
+
+    def transition(
+        self,
+        request_id: str,
+        new_state: RequestState,
+        *,
+        error: str | None = None,
+        encoder_instance: int | None = None,
+        denoiser_instance: int | None = None,
+        decoder_instance: int | None = None,
+    ) -> RequestRecord:
+        with self._lock:
+            record = self._requests.get(request_id)
+            if record is None:
+                raise ValueError(f"Unknown request_id: {request_id}")
+
+            old_state = record.state
+
+            if new_state in _TERMINAL_STATES and new_state != RequestState.DONE:
+                # FAILED / TIMED_OUT: allowed from any active state
+                if old_state not in _ACTIVE_STATES:
+                    raise ValueError(
+                        f"Cannot transition {request_id} from terminal state "
+                        f"{old_state.value} to {new_state.value}"
+                    )
+            elif new_state not in _VALID_TRANSITIONS.get(old_state, set()):
+                raise ValueError(
+                    f"Invalid transition for {request_id}: "
+                    f"{old_state.value} -> {new_state.value}"
+                )
+
+            record.state = new_state
+            record.last_transition_time = time.monotonic()
+            if error is not None:
+                record.error = error
+            if encoder_instance is not None:
+                record.encoder_instance = encoder_instance
+            if denoiser_instance is not None:
+                record.denoiser_instance = denoiser_instance
+            if decoder_instance is not None:
+                record.decoder_instance = decoder_instance
+
+            logger.debug(
+                "Request %s: %s -> %s", request_id, old_state.value, new_state.value
+            )
+            return record
+
+    def get(self, request_id: str) -> RequestRecord | None:
+        with self._lock:
+            return self._requests.get(request_id)
+
+    def remove(self, request_id: str) -> RequestRecord | None:
+        with self._lock:
+            return self._requests.pop(request_id, None)
+
+    def find_timed_out(self, timeout_s: float) -> list[str]:
+        now = time.monotonic()
+        with self._lock:
+            return [
+                r.request_id
+                for r in self._requests.values()
+                if r.state in _ACTIVE_STATES and (now - r.submit_time) > timeout_s
+            ]
+
+    def snapshot(self) -> dict:
+        with self._lock:
+            state_counts = {}
+            for r in self._requests.values():
+                state_counts[r.state.value] = state_counts.get(r.state.value, 0) + 1
+            return {
+                "total": len(self._requests),
+                "active": sum(
+                    1 for r in self._requests.values() if not r.is_terminal()
+                ),
+                "by_state": state_counts,
+            }

python/sglang/multimodal_gen/runtime/disaggregation/roles.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Role definitions for diffusion pipeline disaggregation."""
+
+from enum import Enum
+
+_ROLE_ALIASES = {"denoising": "denoiser"}
+
+
+class RoleType(str, Enum):
+    MONOLITHIC = "monolithic"
+    ENCODER = "encoder"
+    DENOISER = "denoiser"
+    DECODER = "decoder"
+    SERVER = "server"  # Head node (no GPU, routes requests)
+
+    @classmethod
+    def from_string(cls, value: str) -> "RoleType":
+        v = _ROLE_ALIASES.get(value.lower(), value.lower())
+        try:
+            return cls(v)
+        except ValueError:
+            raise ValueError(
+                f"Invalid role: {value}. Must be one of: {', '.join([r.value for r in cls])}"
+            ) from None
+
+    @classmethod
+    def choices(cls) -> list[str]:
+        return [role.value for role in cls] + sorted(_ROLE_ALIASES)
+
+
+def get_module_role(module_name: str) -> "RoleType | None":
+    """Classify a module name to its primary role. Returns None for shared modules."""
+    encoder_prefixes = (
+        "text_encoder",
+        "tokenizer",
+        "image_encoder",
+        "image_processor",
+        "processor",
+        "connectors",
+        "vision_language_encoder",
+    )
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in encoder_prefixes
+    ):
+        return RoleType.ENCODER
+
+    if module_name in {"hy3dshape_conditioner", "hy3dshape_image_processor"}:
+        return RoleType.ENCODER
+
+    denoising_prefixes = (
+        "transformer",
+        "unconditional_transformer",
+        "video_dit",
+        "audio_dit",
+        "dual_tower_bridge",
+    )
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in denoising_prefixes
+    ):
+        return RoleType.DENOISER
+
+    if module_name == "hy3dshape_model":
+        return RoleType.DENOISER
+
+    decoder_prefixes = ("vae", "audio_vae", "video_vae", "vocoder")
+    if any(
+        module_name == p or module_name.startswith(p + "_") for p in decoder_prefixes
+    ):
+        return RoleType.DECODER
+
+    if module_name == "hy3dshape_vae":
+        return RoleType.DECODER
+
+    return None
+
+
+def filter_modules_for_role(
+    module_names: list[str],
+    role: "RoleType",
+    *,
+    extra_allowed_modules: set[str] | None = None,
+) -> list[str]:
+    """Filter module names to only those needed by the given role."""
+    if role in (RoleType.MONOLITHIC, RoleType.SERVER):
+        return module_names
+
+    extra_allowed_modules = extra_allowed_modules or set()
+    filtered = []
+    for name in module_names:
+        module_role = get_module_role(name)
+
+        if module_role is None:
+            filtered.append(name)
+        elif module_role == role:
+            filtered.append(name)
+        elif name in extra_allowed_modules:
+            filtered.append(name)
+
+    return filtered

python/sglang/multimodal_gen/runtime/disaggregation/transport/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Transport layer for disaggregated diffusion pipelines."""

python/sglang/multimodal_gen/runtime/disaggregation/transport/allocator.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Buddy-system memory allocator for TransferTensorBuffer."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from dataclasses import dataclass
+from functools import lru_cache
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Block:
+    offset: int  # byte offset from pool start
+    size: int
+    allocated: bool = False
+    request_id: str | None = None
+
+
+class BuddyAllocator:
+    """Power-of-2 buddy-system allocator for pinned memory."""
+
+    def __init__(self, pool_size: int, min_block_size: int = 1 << 20):
+        if min_block_size <= 0 or (min_block_size & (min_block_size - 1)) != 0:
+            raise ValueError(
+                f"min_block_size must be a power of 2, got {min_block_size}"
+            )
+
+        self._min_block_size = min_block_size
+        self._pool_size = self._next_power_of_2(max(pool_size, min_block_size))
+        self._lock = threading.Lock()
+
+        # Free lists indexed by order: order 0 = min_block_size, order 1 = 2*min_block_size, ...
+        self._max_order = self._size_to_order(self._pool_size)
+        self._free_lists: list[list[int]] = [[] for _ in range(self._max_order + 1)]
+
+        self._blocks: dict[int, Block] = {}
+
+        root = Block(offset=0, size=self._pool_size)
+        self._blocks[0] = root
+        self._free_lists[self._max_order].append(0)
+
+        self._allocated_bytes = 0
+        self._num_allocations = 0
+
+    @property
+    def pool_size(self) -> int:
+        return self._pool_size
+
+    def allocate(self, size: int, request_id: str | None = None) -> int | None:
+        """Allocate a block of at least `size` bytes. Returns offset or None."""
+        if size <= 0:
+            raise ValueError(f"Allocation size must be positive, got {size}")
+
+        alloc_size = max(self._next_power_of_2(size), self._min_block_size)
+        target_order = self._size_to_order(alloc_size)
+
+        if target_order > self._max_order:
+            logger.warning(
+                "Requested size %d exceeds pool size %d", size, self._pool_size
+            )
+            return None
+
+        with self._lock:
+            return self._allocate_locked(target_order, request_id)
+
+    def free(self, offset: int) -> bool:
+        """Free the block at the given offset and coalesce with buddy if possible."""
+        with self._lock:
+            return self._free_locked(offset)
+
+    def get_block_info(self, offset: int) -> Block | None:
+        with self._lock:
+            return self._blocks.get(offset)
+
+    def get_stats(self) -> dict:
+        with self._lock:
+            free_blocks_by_order = {}
+            for order, offsets in enumerate(self._free_lists):
+                if offsets:
+                    block_size = self._min_block_size << order
+                    free_blocks_by_order[block_size] = len(offsets)
+
+            return {
+                "pool_size": self._pool_size,
+                "min_block_size": self._min_block_size,
+                "allocated_bytes": self._allocated_bytes,
+                "free_bytes": self._pool_size - self._allocated_bytes,
+                "num_allocations": self._num_allocations,
+                "num_blocks": len(self._blocks),
+                "free_blocks_by_size": free_blocks_by_order,
+            }
+
+    def count_free_slots(self, slot_size: int) -> int:
+        """Count how many allocations of the given size can fit."""
+        if slot_size <= 0:
+            return 0
+        alloc_size = max(self._next_power_of_2(slot_size), self._min_block_size)
+
+        with self._lock:
+            count = 0
+            for order in range(self._size_to_order(alloc_size), self._max_order + 1):
+                for _ in self._free_lists[order]:
+                    block_size = self._min_block_size << order
+                    count += block_size // alloc_size
+            return count
+
+    # --- Internal (caller must hold self._lock) ---
+
+    def _allocate_locked(self, target_order: int, request_id: str | None) -> int | None:
+        found_order = -1
+        for order in range(target_order, self._max_order + 1):
+            if self._free_lists[order]:
+                found_order = order
+                break
+
+        if found_order < 0:
+            return None
+
+        offset = self._free_lists[found_order].pop(0)
+        block = self._blocks[offset]
+
+        # Split down to target_order
+        while found_order > target_order:
+            found_order -= 1
+            buddy_size = self._min_block_size << found_order
+            buddy_offset = offset + buddy_size
+
+            buddy = Block(offset=buddy_offset, size=buddy_size)
+            self._blocks[buddy_offset] = buddy
+            self._free_lists[found_order].append(buddy_offset)
+
+            block.size = buddy_size
+
+        block.allocated = True
+        block.request_id = request_id
+        self._allocated_bytes += block.size
+        self._num_allocations += 1
+
+        return offset
+
+    def _free_locked(self, offset: int) -> bool:
+        block = self._blocks.get(offset)
+        if block is None or not block.allocated:
+            return False
+
+        block.allocated = False
+        block.request_id = None
+        self._allocated_bytes -= block.size
+        self._num_allocations -= 1
+
+        self._coalesce(block)
+        return True
+
+    def _coalesce(self, block: Block) -> None:
+        """Recursively merge with buddy if both are free."""
+        while block.size < self._pool_size:
+            buddy_offset = block.offset ^ block.size
+            buddy = self._blocks.get(buddy_offset)
+
+            if buddy is None or buddy.allocated or buddy.size != block.size:
+                break
+
+            order = self._size_to_order(buddy.size)
+            self._free_lists[order].remove(buddy_offset)
+
+            if buddy_offset < block.offset:
+                del self._blocks[block.offset]
+                buddy.size *= 2
+                block = buddy
+            else:
+                del self._blocks[buddy_offset]
+                block.size *= 2
+
+        order = self._size_to_order(block.size)
+        self._free_lists[order].append(block.offset)
+
+    def _size_to_order(self, size: int) -> int:
+        order = 0
+        s = self._min_block_size
+        while s < size:
+            s <<= 1
+            order += 1
+        return order
+
+    @staticmethod
+    @lru_cache(maxsize=256)
+    def _next_power_of_2(n: int) -> int:
+        if n <= 0:
+            return 1
+        n -= 1
+        n |= n >> 1
+        n |= n >> 2
+        n |= n >> 4
+        n |= n >> 8
+        n |= n >> 16
+        n |= n >> 32
+        return n + 1

python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py

@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+"""TransferTensorBuffer: memory staging area for disaggregated tensor transfer."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.transport.allocator import (
+    BuddyAllocator,
+)
+from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
+    str_to_dtype,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SlotHandle:
+    request_id: str
+    offset: int  # byte offset in the pool
+    size: int  # allocated size in bytes
+    tensor_views: dict[str, torch.Tensor | list[torch.Tensor]] = field(
+        default_factory=dict
+    )
+
+
+class TransferTensorBuffer:
+    """Memory pool for staging tensor payloads between roles.
+
+    Wraps a contiguous block of memory (CPU pinned or GPU) with a BuddyAllocator.
+    """
+
+    def __init__(
+        self,
+        pool_size: int,
+        min_block_size: int = 1 << 20,
+        role_name: str = "unknown",
+        device: str = "cpu",
+    ):
+        self._role_name = role_name
+        self._device = device
+        self._allocator = BuddyAllocator(pool_size, min_block_size)
+        actual_size = self._allocator.pool_size
+
+        if device == "cpu":
+            self._pool = torch.empty(actual_size, dtype=torch.uint8, pin_memory=True)
+        else:
+            self._pool = torch.empty(actual_size, dtype=torch.uint8, device=device)
+        self._pool_ptr = self._pool.data_ptr()
+
+        pool_location = "pinned CPU" if device == "cpu" else f"GPU ({device})"
+        logger.info(
+            "TransferTensorBuffer[%s]: allocated %d MiB %s memory "
+            "(min_block=%d KiB)",
+            role_name,
+            actual_size >> 20,
+            pool_location,
+            min_block_size >> 10,
+        )
+
+    @property
+    def pool_size(self) -> int:
+        return self._allocator.pool_size
+
+    @property
+    def device(self) -> str:
+        return self._device
+
+    @property
+    def pool_data_ptr(self) -> int:
+        return self._pool_ptr
+
+    def allocate(self, size: int, request_id: str) -> SlotHandle | None:
+        """Allocate a slot. Returns None if pool is full."""
+        offset = self._allocator.allocate(size, request_id=request_id)
+        if offset is None:
+            logger.warning(
+                "TransferTensorBuffer[%s]: allocation failed for %s (%d bytes). "
+                "Pool stats: %s",
+                self._role_name,
+                request_id,
+                size,
+                self._allocator.get_stats(),
+            )
+            return None
+
+        block = self._allocator.get_block_info(offset)
+        return SlotHandle(
+            request_id=request_id,
+            offset=offset,
+            size=block.size if block else size,
+        )
+
+    def free(self, handle: SlotHandle) -> bool:
+        return self._allocator.free(handle.offset)
+
+    def write_tensor(
+        self,
+        handle: SlotHandle,
+        name: str,
+        tensor: torch.Tensor,
+        byte_offset: int = 0,
+        stream: torch.Stream | None = None,
+    ) -> int:
+        """Copy a tensor into the pool slot. Returns bytes written."""
+        src_tensor = tensor.contiguous()
+        nbytes = src_tensor.numel() * src_tensor.element_size()
+
+        if byte_offset + nbytes > handle.size:
+            raise ValueError(
+                f"Write exceeds slot: offset={byte_offset}, nbytes={nbytes}, "
+                f"slot_size={handle.size}"
+            )
+
+        dst = self._pool[
+            handle.offset + byte_offset : handle.offset + byte_offset + nbytes
+        ]
+        src_bytes = src_tensor.view(torch.uint8).reshape(-1)
+
+        if stream is not None:
+            with torch.get_device_module().stream(stream):
+                dst.copy_(src_bytes, non_blocking=True)
+        else:
+            dst.copy_(src_bytes, non_blocking=True)
+
+        return nbytes
+
+    def read_tensor(
+        self,
+        handle: SlotHandle,
+        shape: list[int],
+        dtype: torch.dtype,
+        byte_offset: int = 0,
+        device: torch.device | str = "cpu",
+        stream: torch.Stream | None = None,
+    ) -> torch.Tensor:
+        """Read a tensor from the pool slot. Returns a clone on target device."""
+        nbytes = 1
+        for s in shape:
+            nbytes *= s
+        nbytes *= torch.tensor([], dtype=dtype).element_size()
+
+        raw = self._pool[
+            handle.offset + byte_offset : handle.offset + byte_offset + nbytes
+        ]
+        src = raw.view(dtype).reshape(shape)
+
+        pool_dev = str(self._pool.device)
+        target_dev = str(device)
+
+        same_device = pool_dev == target_dev
+
+        if same_device:
+            # Clone to decouple tensor lifetime from pool slot
+            if stream is not None:
+                with torch.get_device_module().stream(stream):
+                    return src.clone()
+            return src.clone()
+
+        if stream is not None:
+            with torch.get_device_module().stream(stream):
+                return src.to(device, non_blocking=True)
+        return src.to(device, non_blocking=True)
+
+    def write_tensors_from_gpu(
+        self,
+        handle: SlotHandle,
+        tensors: dict[str, torch.Tensor | list[torch.Tensor] | None],
+        stream: torch.Stream | None = None,
+    ) -> dict[str, list[dict]]:
+        """Batch-write GPU tensors into a slot. Returns a manifest for later reads."""
+        manifest: dict[str, list[dict]] = {}
+        byte_offset = 0
+
+        # Ensure copy stream sees all prior compute kernels
+        if stream is not None:
+            stream.wait_stream(torch.get_device_module().current_stream())
+
+        for name, value in tensors.items():
+            if value is None:
+                continue
+
+            entries = []
+            if isinstance(value, torch.Tensor):
+                nbytes = self.write_tensor(handle, name, value, byte_offset, stream)
+                entries.append(
+                    {
+                        "offset": byte_offset,
+                        "shape": list(value.shape),
+                        "dtype": str(value.dtype).replace("torch.", ""),
+                    }
+                )
+                byte_offset += nbytes
+                byte_offset = (byte_offset + 511) & ~511  # align to 512B
+
+            elif isinstance(value, list):
+                for i, t in enumerate(value):
+                    if t is None:
+                        continue
+                    nbytes = self.write_tensor(
+                        handle, f"{name}[{i}]", t, byte_offset, stream
+                    )
+                    entries.append(
+                        {
+                            "offset": byte_offset,
+                            "shape": list(t.shape),
+                            "dtype": str(t.dtype).replace("torch.", ""),
+                            "list_index": i,
+                        }
+                    )
+                    byte_offset += nbytes
+                    byte_offset = (byte_offset + 511) & ~511
+
+            if entries:
+                manifest[name] = entries
+
+        return manifest
+
+    def read_tensors_from_manifest(
+        self,
+        handle: SlotHandle,
+        manifest: dict[str, list[dict]],
+        device: torch.device | str = "cpu",
+        stream: torch.Stream | None = None,
+    ) -> dict[str, torch.Tensor | list[torch.Tensor]]:
+        """Batch-read tensors from a slot using a manifest."""
+        result: dict[str, torch.Tensor | list[torch.Tensor]] = {}
+
+        for name, entries in manifest.items():
+            if not entries:
+                continue
+            has_list_index = any("list_index" in e for e in entries)
+
+            if has_list_index:
+                max_idx = max(e.get("list_index", 0) for e in entries) + 1
+                tensors = [None] * max_idx
+                for entry in entries:
+                    t = self.read_tensor(
+                        handle,
+                        entry["shape"],
+                        str_to_dtype(entry["dtype"]),
+                        entry["offset"],
+                        device,
+                        stream,
+                    )
+                    tensors[entry["list_index"]] = t
+                result[name] = tensors
+            else:
+                entry = entries[0]
+                result[name] = self.read_tensor(
+                    handle,
+                    entry["shape"],
+                    str_to_dtype(entry["dtype"]),
+                    entry["offset"],
+                    device,
+                    stream,
+                )
+
+        return result
+
+    def free_slots_count(self, typical_request_size: int) -> int:
+        """Estimate how many requests of typical size can still be buffered."""
+        return self._allocator.count_free_slots(typical_request_size)
+
+    def get_stats(self) -> dict:
+        alloc_stats = self._allocator.get_stats()
+        alloc_stats["role"] = self._role_name
+        return alloc_stats

python/sglang/multimodal_gen/runtime/disaggregation/transport/codec.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Zero-copy tensor codec for ZMQ multipart messages.
+
+Frame 0: JSON metadata (tensor descriptors + scalar fields)
+Frame 1-N: Raw tensor data buffers (one per tensor)
+"""
+
+import ctypes
+import json
+import logging
+from dataclasses import dataclass
+
+import torch
+import zmq
+
+logger = logging.getLogger(__name__)
+
+_DTYPE_TO_STR = {
+    torch.float16: "float16",
+    torch.float32: "float32",
+    torch.float64: "float64",
+    torch.bfloat16: "bfloat16",
+    torch.int8: "int8",
+    torch.int16: "int16",
+    torch.int32: "int32",
+    torch.int64: "int64",
+    torch.uint8: "uint8",
+    torch.bool: "bool",
+}
+_STR_TO_DTYPE = {v: k for k, v in _DTYPE_TO_STR.items()}
+
+
+def dtype_to_str(dtype: torch.dtype) -> str:
+    s = _DTYPE_TO_STR.get(dtype)
+    if s is None:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+    return s
+
+
+def str_to_dtype(s: str) -> torch.dtype:
+    d = _STR_TO_DTYPE.get(s)
+    if d is None:
+        raise ValueError(f"Unknown dtype string: {s}")
+    return d
+
+
+class TensorWrapper:
+    """Expose a CPU-contiguous tensor's data buffer for zero-copy ZMQ send."""
+
+    def __init__(self, tensor: torch.Tensor):
+        if tensor.is_cuda or tensor.is_npu:
+            tensor = tensor.cpu()
+        if not tensor.is_contiguous():
+            tensor = tensor.contiguous()
+        self.tensor = tensor
+        data_ptr = tensor.data_ptr()
+        total_bytes = tensor.numel() * tensor.element_size()
+        self._c_buf = (ctypes.c_char * total_bytes).from_address(data_ptr)
+        self._view = memoryview(self._c_buf)
+
+
+@dataclass
+class TensorDescriptor:
+    field_name: str
+    shape: list[int]
+    dtype: str
+    list_index: int = -1  # -1 means not part of a list
+
+    def to_dict(self) -> dict:
+        return {
+            "field_name": self.field_name,
+            "shape": self.shape,
+            "dtype": self.dtype,
+            "list_index": self.list_index,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "TensorDescriptor":
+        return cls(
+            field_name=d["field_name"],
+            shape=d["shape"],
+            dtype=d["dtype"],
+            list_index=d.get("list_index", -1),
+        )
+
+
+def pack_tensors(
+    tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+    scalar_fields: dict | None = None,
+) -> tuple[bytes, list[TensorWrapper]]:
+    """Pack tensor fields into metadata + buffer list for send_multipart."""
+    descriptors = []
+    buffers = []
+
+    for field_name, value in tensor_fields.items():
+        if value is None:
+            continue
+
+        if isinstance(value, torch.Tensor):
+            wrapper = TensorWrapper(value)
+            descriptors.append(
+                TensorDescriptor(
+                    field_name=field_name,
+                    shape=list(value.shape),
+                    dtype=dtype_to_str(value.dtype),
+                )
+            )
+            buffers.append(wrapper)
+
+        elif isinstance(value, list):
+            for i, t in enumerate(value):
+                if t is None:
+                    continue
+                if not isinstance(t, torch.Tensor):
+                    raise TypeError(
+                        f"Expected Tensor in list for field '{field_name}', "
+                        f"got {type(t)}"
+                    )
+                wrapper = TensorWrapper(t)
+                descriptors.append(
+                    TensorDescriptor(
+                        field_name=field_name,
+                        shape=list(t.shape),
+                        dtype=dtype_to_str(t.dtype),
+                        list_index=i,
+                    )
+                )
+                buffers.append(wrapper)
+
+    metadata = {
+        "tensor_descriptors": [d.to_dict() for d in descriptors],
+        "scalar_fields": scalar_fields or {},
+    }
+    metadata_bytes = json.dumps(metadata, separators=(",", ":")).encode("utf-8")
+    return metadata_bytes, buffers
+
+
+def send_tensors(
+    socket: zmq.Socket,
+    tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+    scalar_fields: dict | None = None,
+    flags: int = 0,
+) -> None:
+    """Send tensors over ZMQ using multipart with zero-copy."""
+    metadata_bytes, buffers = pack_tensors(tensor_fields, scalar_fields)
+    parts: list = [metadata_bytes]
+    parts.extend(w._view if isinstance(w, TensorWrapper) else w for w in buffers)
+    socket.send_multipart(parts, flags=flags, copy=True)
+
+
+def unpack_tensors(
+    parts: list,
+    device: str | torch.device = "cpu",
+) -> tuple[dict[str, torch.Tensor | list[torch.Tensor]], dict]:
+    """Unpack multipart message frames into tensor fields and scalar fields."""
+    metadata_frame = parts[0]
+    metadata_bytes = (
+        bytes(metadata_frame.buffer)
+        if hasattr(metadata_frame, "buffer")
+        else bytes(metadata_frame)
+    )
+    metadata = json.loads(metadata_bytes)
+
+    descriptors = [
+        TensorDescriptor.from_dict(d) for d in metadata["tensor_descriptors"]
+    ]
+    scalar_fields = metadata.get("scalar_fields", {})
+
+    if len(parts) - 1 != len(descriptors):
+        raise ValueError(
+            f"Expected {len(descriptors)} tensor frames, got {len(parts) - 1}"
+        )
+
+    tensor_fields: dict[str, torch.Tensor | list[torch.Tensor]] = {}
+    list_sizes: dict[str, int] = {}
+    for desc in descriptors:
+        if desc.list_index >= 0:
+            current_max = list_sizes.get(desc.field_name, 0)
+            list_sizes[desc.field_name] = max(current_max, desc.list_index + 1)
+
+    for field_name, size in list_sizes.items():
+        tensor_fields[field_name] = [None] * size
+
+    for i, desc in enumerate(descriptors):
+        frame = parts[i + 1]
+        buf = frame.buffer if hasattr(frame, "buffer") else bytes(frame)
+        dtype = str_to_dtype(desc.dtype)
+        # clone() to own the memory (decouple from ZMQ buffer lifetime)
+        tensor = torch.frombuffer(buf, dtype=dtype).reshape(desc.shape).clone()
+        if device != "cpu" and device != torch.device("cpu"):
+            tensor = tensor.to(device)
+
+        if desc.list_index >= 0:
+            tensor_fields[desc.field_name][desc.list_index] = tensor
+        else:
+            tensor_fields[desc.field_name] = tensor
+
+    return tensor_fields, scalar_fields

python/sglang/multimodal_gen/runtime/disaggregation/transport/engine.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Transfer engine abstraction for tensor transfer between role instances."""
+
+import logging
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger(__name__)
+
+_MOONCAKE_AVAILABLE = None
+
+
+def _check_mooncake() -> bool:
+    global _MOONCAKE_AVAILABLE
+    if _MOONCAKE_AVAILABLE is None:
+        try:
+            from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (  # noqa: F401
+                MooncakeTransferEngine as _MTE,
+            )
+
+            _MOONCAKE_AVAILABLE = True
+        except ImportError:
+            _MOONCAKE_AVAILABLE = False
+    return _MOONCAKE_AVAILABLE
+
+
+class BaseTransferEngine(ABC):
+    """Abstract transfer engine for data movement between roles."""
+
+    @property
+    def supports_gpu_direct(self) -> bool:
+        return False
+
+    @property
+    @abstractmethod
+    def session_id(self) -> str: ...
+
+    @abstractmethod
+    def register_buffer(self, ptr: int, length: int) -> None: ...
+
+    @abstractmethod
+    def deregister_buffer(self, ptr: int) -> None: ...
+
+    @abstractmethod
+    def transfer_sync(
+        self, dst_session_id: str, src_addr: int, dst_addr: int, length: int
+    ) -> int:
+        """Returns 0 on success, negative on failure."""
+
+    @abstractmethod
+    def batch_transfer_sync(
+        self,
+        dst_session_id: str,
+        src_addrs: list[int],
+        dst_addrs: list[int],
+        lengths: list[int],
+    ) -> int: ...
+
+
+class MooncakeDiffusionEngine(BaseTransferEngine):
+    """Production engine backed by MooncakeTransferEngine (RDMA)."""
+
+    @property
+    def supports_gpu_direct(self) -> bool:
+        return True
+
+    def __init__(
+        self,
+        hostname: str,
+        gpu_id: int = 0,
+        ib_device: str | None = None,
+    ):
+        from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+            MooncakeTransferEngine,
+        )
+
+        self._engine = MooncakeTransferEngine(
+            hostname=hostname,
+            gpu_id=gpu_id,
+            ib_device=ib_device,
+        )
+        logger.info(
+            "MooncakeDiffusionEngine initialized: session_id=%s",
+            self._engine.session_id,
+        )
+
+    @property
+    def session_id(self) -> str:
+        return self._engine.session_id
+
+    def register_buffer(self, ptr: int, length: int) -> None:
+        self._engine.register(ptr, length)
+
+    def deregister_buffer(self, ptr: int) -> None:
+        self._engine.deregister(ptr)
+
+    def transfer_sync(
+        self, dst_session_id: str, src_addr: int, dst_addr: int, length: int
+    ) -> int:
+        return self._engine.transfer_sync(dst_session_id, src_addr, dst_addr, length)
+
+    def batch_transfer_sync(
+        self,
+        dst_session_id: str,
+        src_addrs: list[int],
+        dst_addrs: list[int],
+        lengths: list[int],
+    ) -> int:
+        return self._engine.batch_transfer_sync(
+            dst_session_id, src_addrs, dst_addrs, lengths
+        )
+
+
+def create_transfer_engine(
+    hostname: str = "127.0.0.1",
+    gpu_id: int = 0,
+    ib_device: str | None = None,
+) -> BaseTransferEngine:
+    """Factory: returns MooncakeDiffusionEngine if mooncake is available."""
+    if not _check_mooncake():
+        raise RuntimeError(
+            "Mooncake transfer engine is required for disaggregated diffusion "
+            "but is not installed. Please install mooncake first."
+        )
+    return MooncakeDiffusionEngine(
+        hostname=hostname, gpu_id=gpu_id, ib_device=ib_device
+    )

python/sglang/multimodal_gen/runtime/disaggregation/transport/manager.py

@@ -0,0 +1,391 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-instance transfer manager for disaggregated diffusion roles."""
+
+import logging
+import threading
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.transport.buffer import (
+    SlotHandle,
+    TransferTensorBuffer,
+)
+from sglang.multimodal_gen.runtime.disaggregation.transport.engine import (
+    BaseTransferEngine,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StagedTransfer:
+    request_id: str
+    slot: SlotHandle
+    manifest: dict
+    scalar_fields: dict = field(default_factory=dict)
+
+
+@dataclass
+class PendingReceive:
+    request_id: str
+    slot: SlotHandle
+
+
+class DiffusionTransferManager:
+    """Manages tensor transfers for a single role instance.
+
+    Owns a TransferTensorBuffer (memory pool) and a BaseTransferEngine (RDMA or mock).
+    """
+
+    def __init__(
+        self,
+        engine: BaseTransferEngine,
+        buffer: TransferTensorBuffer,
+    ):
+        self._engine = engine
+        self._buffer = buffer
+        self._lock = threading.Lock()
+
+        self._engine.register_buffer(self._buffer.pool_data_ptr, self._buffer.pool_size)
+
+        self._staged: dict[str, StagedTransfer] = {}
+        self._pending_receives: dict[str, PendingReceive] = {}
+
+        logger.info(
+            "DiffusionTransferManager initialized: session=%s, pool=%d bytes",
+            self._engine.session_id,
+            self._buffer.pool_size,
+        )
+
+    @property
+    def session_id(self) -> str:
+        return self._engine.session_id
+
+    @property
+    def pool_data_ptr(self) -> int:
+        return self._buffer.pool_data_ptr
+
+    @property
+    def pool_size(self) -> int:
+        return self._buffer.pool_size
+
+    def stage_tensors(
+        self,
+        request_id: str,
+        tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+        scalar_fields: dict | None = None,
+        stream: torch.Stream | None = None,
+    ) -> StagedTransfer | None:
+        """Stage GPU tensors into the local TransferBuffer. Returns None on allocation failure."""
+        total_size = 0
+        for name, t in tensor_fields.items():
+            if t is None:
+                continue
+            if isinstance(t, list):
+                for ti in t:
+                    total_size += ti.nelement() * ti.element_size()
+            else:
+                total_size += t.nelement() * t.element_size()
+
+        if total_size == 0:
+            staged = StagedTransfer(
+                request_id=request_id,
+                slot=None,
+                manifest={},
+                scalar_fields=scalar_fields or {},
+            )
+            with self._lock:
+                self._staged[request_id] = staged
+            return staged
+
+        slot = self._buffer.allocate(total_size, request_id)
+        if slot is None:
+            logger.warning(
+                "TransferManager: failed to allocate %d bytes for %s",
+                total_size,
+                request_id,
+            )
+            return None
+
+        manifest = self._buffer.write_tensors_from_gpu(slot, tensor_fields, stream)
+
+        if stream is not None:
+            stream.synchronize()
+        elif torch.get_device_module().is_available():
+            torch.get_device_module().synchronize()
+
+        staged = StagedTransfer(
+            request_id=request_id,
+            slot=slot,
+            manifest=manifest,
+            scalar_fields=scalar_fields or {},
+        )
+        with self._lock:
+            self._staged[request_id] = staged
+
+        logger.debug(
+            "TransferManager: staged %s (%d bytes, offset=%d)",
+            request_id,
+            total_size,
+            slot.offset,
+        )
+        return staged
+
+    def stage_tensors_async(
+        self,
+        request_id: str,
+        tensor_fields: dict[str, torch.Tensor | list[torch.Tensor] | None],
+        scalar_fields: dict | None = None,
+        stream: torch.Stream | None = None,
+    ) -> tuple[StagedTransfer | None, torch.Event | None]:
+        """Stage GPU tensors, returning a CUDA event instead of blocking.
+
+        Caller MUST wait on the event before reading buffer data.
+        """
+        total_size = 0
+        for name, t in tensor_fields.items():
+            if t is None:
+                continue
+            if isinstance(t, list):
+                for ti in t:
+                    total_size += ti.nelement() * ti.element_size()
+            else:
+                total_size += t.nelement() * t.element_size()
+
+        if total_size == 0:
+            staged = StagedTransfer(
+                request_id=request_id,
+                slot=None,
+                manifest={},
+                scalar_fields=scalar_fields or {},
+            )
+            with self._lock:
+                self._staged[request_id] = staged
+            return staged, None
+
+        slot = self._buffer.allocate(total_size, request_id)
+        if slot is None:
+            logger.warning(
+                "TransferManager: failed to allocate %d bytes for %s",
+                total_size,
+                request_id,
+            )
+            return None, None
+
+        manifest = self._buffer.write_tensors_from_gpu(slot, tensor_fields, stream)
+
+        d2h_event = None
+        if stream is not None:
+            d2h_event = torch.get_device_module().Event()
+            d2h_event.record(stream)
+        elif torch.get_device_module().is_available():
+            d2h_event = torch.get_device_module().Event()
+            d2h_event.record(torch.get_device_module().current_stream())
+
+        staged = StagedTransfer(
+            request_id=request_id,
+            slot=slot,
+            manifest=manifest,
+            scalar_fields=scalar_fields or {},
+        )
+        with self._lock:
+            self._staged[request_id] = staged
+
+        logger.debug(
+            "TransferManager: staged_async %s (%d bytes, offset=%d)",
+            request_id,
+            total_size,
+            slot.offset,
+        )
+        return staged, d2h_event
+
+    def load_tensors_async(
+        self,
+        request_id: str,
+        manifest: dict,
+        device: torch.device | str = current_platform.device_type,
+        stream: torch.Stream | None = None,
+    ) -> tuple[
+        dict[str, torch.Tensor | list[torch.Tensor]],
+        torch.get_device_module().Event | None,
+    ]:
+        """Load tensors from receive slot to GPU, returning a CUDA event.
+
+        Caller MUST wait on the event before using the returned tensors.
+        """
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+
+        if pending is None:
+            raise ValueError(
+                f"TransferManager: no pending receive slot for {request_id}"
+            )
+
+        tensors = self._buffer.read_tensors_from_manifest(
+            pending.slot, manifest, device=device, stream=stream
+        )
+
+        load_event = None
+        if stream is not None:
+            load_event = torch.get_device_module().Event()
+            load_event.record(stream)
+        elif torch.get_device_module().is_available():
+            load_event = torch.get_device_module().Event()
+            load_event.record(torch.get_device_module().current_stream())
+
+        logger.debug(
+            "TransferManager: loaded_async %d tensor fields for %s to %s",
+            len(tensors),
+            request_id,
+            device,
+        )
+        return tensors, load_event
+
+    def push_to_peer(
+        self,
+        request_id: str,
+        dest_session_id: str,
+        dest_addr: int,
+        transfer_size: int,
+    ) -> bool:
+        """Push staged data to a remote peer's buffer via RDMA. Returns True on success."""
+        with self._lock:
+            staged = self._staged.get(request_id)
+
+        if staged is None:
+            logger.error("TransferManager: no staged transfer for %s", request_id)
+            return False
+
+        if staged.slot is None:
+            return True
+
+        src_addr = self._buffer.pool_data_ptr + staged.slot.offset
+        ret = self._engine.transfer_sync(
+            dest_session_id, src_addr, dest_addr, transfer_size
+        )
+
+        if ret == 0:
+            logger.debug(
+                "TransferManager: pushed %s (%d bytes) to %s",
+                request_id,
+                transfer_size,
+                dest_session_id,
+            )
+        else:
+            logger.error(
+                "TransferManager: RDMA push failed for %s (ret=%d)",
+                request_id,
+                ret,
+            )
+
+        return ret == 0
+
+    def free_staged(self, request_id: str) -> None:
+        with self._lock:
+            staged = self._staged.pop(request_id, None)
+
+        if staged and staged.slot is not None:
+            self._buffer.free(staged.slot)
+            logger.debug("TransferManager: freed staged slot for %s", request_id)
+
+    def allocate_receive_slot(
+        self, request_id: str, size: int
+    ) -> PendingReceive | None:
+        """Allocate a local buffer slot to receive incoming data."""
+        slot = self._buffer.allocate(size, request_id)
+        if slot is None:
+            logger.warning(
+                "TransferManager: failed to allocate receive slot (%d bytes) for %s",
+                size,
+                request_id,
+            )
+            return None
+
+        pending = PendingReceive(request_id=request_id, slot=slot)
+        with self._lock:
+            self._pending_receives[request_id] = pending
+
+        logger.debug(
+            "TransferManager: allocated receive slot for %s (offset=%d, size=%d)",
+            request_id,
+            slot.offset,
+            slot.size,
+        )
+        return pending
+
+    def load_tensors(
+        self,
+        request_id: str,
+        manifest: dict,
+        device: torch.device | str = current_platform.device_type,
+        stream: torch.Stream | None = None,
+    ) -> dict[str, torch.Tensor | list[torch.Tensor]]:
+        """Load tensors from a receive slot into GPU memory."""
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+
+        if pending is None:
+            raise ValueError(
+                f"TransferManager: no pending receive slot for {request_id}"
+            )
+
+        tensors = self._buffer.read_tensors_from_manifest(
+            pending.slot, manifest, device=device, stream=stream
+        )
+
+        if stream is not None:
+            stream.synchronize()
+        elif torch.get_device_module().is_available():
+            torch.get_device_module().synchronize()
+
+        logger.debug(
+            "TransferManager: loaded %d tensor fields for %s to %s",
+            len(tensors),
+            request_id,
+            device,
+        )
+        return tensors
+
+    def register_prealloc_as_receive(
+        self, request_id: str, slot: "SlotHandle"
+    ) -> "PendingReceive":
+        """Register a pre-allocated slot as a pending receive (fast path)."""
+        pending = PendingReceive(request_id=request_id, slot=slot)
+        with self._lock:
+            self._pending_receives[request_id] = pending
+        return pending
+
+    def free_receive_slot(self, request_id: str) -> None:
+        with self._lock:
+            pending = self._pending_receives.pop(request_id, None)
+
+        if pending:
+            self._buffer.free(pending.slot)
+            logger.debug("TransferManager: freed receive slot for %s", request_id)
+
+    def get_receive_slot_addr(self, request_id: str) -> int | None:
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+        if pending is None:
+            return None
+        return self._buffer.pool_data_ptr + pending.slot.offset
+
+    def get_receive_slot_offset(self, request_id: str) -> int | None:
+        with self._lock:
+            pending = self._pending_receives.get(request_id)
+        if pending is None:
+            return None
+        return pending.slot.offset
+
+    def get_staged_info(self, request_id: str) -> StagedTransfer | None:
+        with self._lock:
+            return self._staged.get(request_id)
+
+    def free_slots_count(self, typical_size: int = 64 * 1024 * 1024) -> int:
+        return self._buffer.free_slots_count(typical_size)
+
+    def cleanup(self) -> None:
+        self._engine.deregister_buffer(self._buffer.pool_data_ptr)
+        logger.info("DiffusionTransferManager cleaned up")

python/sglang/multimodal_gen/runtime/disaggregation/transport/protocol.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Transfer protocol messages for disaggregated diffusion.
+
+All messages are sent as ZMQ multipart with a b"__transfer__" discriminator
+in frame[0] and JSON payload in frame[1].
+"""
+
+import json
+import logging
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+TRANSFER_MAGIC = b"__transfer__"
+
+
+class TransferMsgType:
+    # Instance → DiffusionServer
+    STAGED = "transfer_staged"
+    ALLOCATED = "transfer_allocated"
+    PUSHED = "transfer_pushed"
+    DONE = "transfer_done"
+
+    # DiffusionServer → Instance
+    ALLOC = "transfer_alloc"
+    PUSH = "transfer_push"
+    READY = "transfer_ready"
+
+    # Registration
+    REGISTER = "transfer_register"
+    REGISTER_ACK = "transfer_register_ack"
+
+
+@dataclass
+class TransferStagedMsg:
+    msg_type: str = TransferMsgType.STAGED
+    request_id: str = ""
+    data_size: int = 0
+    manifest: dict = None
+    session_id: str = ""
+    pool_ptr: int = 0
+    slot_offset: int = 0
+
+    def __post_init__(self):
+        if self.manifest is None:
+            self.manifest = {}
+
+
+@dataclass
+class TransferAllocMsg:
+    msg_type: str = TransferMsgType.ALLOC
+    request_id: str = ""
+    data_size: int = 0
+    source_role: str = ""
+
+
+@dataclass
+class TransferAllocatedMsg:
+    msg_type: str = TransferMsgType.ALLOCATED
+    request_id: str = ""
+    session_id: str = ""
+    pool_ptr: int = 0
+    slot_offset: int = 0
+    slot_size: int = 0
+
+
+@dataclass
+class TransferPushMsg:
+    msg_type: str = TransferMsgType.PUSH
+    request_id: str = ""
+    dest_session_id: str = ""
+    dest_addr: int = 0
+    transfer_size: int = 0
+
+
+@dataclass
+class TransferPushedMsg:
+    msg_type: str = TransferMsgType.PUSHED
+    request_id: str = ""
+
+
+@dataclass
+class TransferReadyMsg:
+    msg_type: str = TransferMsgType.READY
+    request_id: str = ""
+    manifest: dict = None
+    slot_offset: int = 0
+    scalar_fields: dict = None
+
+    def __post_init__(self):
+        if self.manifest is None:
+            self.manifest = {}
+        if self.scalar_fields is None:
+            self.scalar_fields = {}
+
+
+@dataclass
+class TransferDoneMsg:
+    msg_type: str = TransferMsgType.DONE
+    request_id: str = ""
+    error: str | None = None
+
+
+@dataclass
+class TransferRegisterMsg:
+    msg_type: str = TransferMsgType.REGISTER
+    role: str = ""
+    session_id: str = ""
+    pool_ptr: int = 0
+    pool_size: int = 0
+    # The instance's own work endpoint (e.g. tcp://host:port). Used by the
+    # DiffusionServer to key peer info by URL index (i.e. the same index used
+    # to build the PUSH work-socket list), so the control plane and the RDMA
+    # data plane cannot drift when instances register in a different order
+    # than --*-urls.
+    work_endpoint: str = ""
+    # Pre-allocated receive slots: [{"offset": int, "size": int, "slot_id": int, "addr": int}]
+    preallocated_slots: list = field(default_factory=list)
+
+
+def encode_transfer_msg(msg: Any) -> list[bytes]:
+    """Encode as [TRANSFER_MAGIC, json_payload_bytes]."""
+    if hasattr(msg, "__dataclass_fields__"):
+        d = asdict(msg)
+    elif isinstance(msg, dict):
+        d = msg
+    else:
+        raise TypeError(f"Cannot encode transfer message: {type(msg)}")
+
+    return [TRANSFER_MAGIC, json.dumps(d, separators=(",", ":")).encode("utf-8")]
+
+
+def decode_transfer_msg(frames: list[bytes]) -> dict:
+    if len(frames) < 2 or frames[0] != TRANSFER_MAGIC:
+        raise ValueError(f"Not a transfer message: frame[0]={frames[0]!r}")
+    return json.loads(frames[1])
+
+
+def is_transfer_message(frames: list) -> bool:
+    return len(frames) >= 2 and (
+        frames[0] == TRANSFER_MAGIC
+        or (isinstance(frames[0], memoryview) and bytes(frames[0]) == TRANSFER_MAGIC)
+        or (hasattr(frames[0], "bytes") and frames[0].bytes == TRANSFER_MAGIC)
+    )

python/sglang/multimodal_gen/runtime/distributed/__init__.py

@@ -0,0 +1,63 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from sglang.multimodal_gen.runtime.distributed.communication_op import *
+from sglang.multimodal_gen.runtime.distributed.group_coordinator import (
+    get_local_torch_device,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    cleanup_dist_env_and_memory,
+    get_decode_parallel_group_coordinator,
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+    get_dp_group,
+    get_dp_rank,
+    get_dp_world_size,
+    get_sp_group,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+    get_tp_group,
+    get_tp_rank,
+    get_tp_world_size,
+    get_world_group,
+    get_world_rank,
+    get_world_size,
+    init_distributed_environment,
+    initialize_model_parallel,
+    maybe_init_distributed_environment_and_model_parallel,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.distributed.utils import *
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+__all__ = [
+    # Initialization
+    "init_distributed_environment",
+    "initialize_model_parallel",
+    "cleanup_dist_env_and_memory",
+    "model_parallel_is_initialized",
+    "maybe_init_distributed_environment_and_model_parallel",
+    # World group
+    "get_world_group",
+    "get_world_rank",
+    "get_world_size",
+    # Data parallel group
+    "get_dp_group",
+    "get_dp_rank",
+    "get_dp_world_size",
+    # Sequence parallel group
+    "get_sp_group",
+    "get_sp_parallel_rank",
+    "get_sp_world_size",
+    # Tensor parallel group
+    "get_tp_group",
+    "get_tp_rank",
+    "get_tp_world_size",
+    # Decode parallel group
+    "get_decode_parallel_group_coordinator",
+    "get_decode_parallel_rank",
+    "get_decode_parallel_world_size",
+    # Get torch device
+    "get_local_torch_device",
+]

python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
+    _apply_cfg_postprocess,
+    _unwrap,
+    _wrap,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    cfg_model_parallel_all_gather,
+    cfg_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
+        CFGBranch,
+        CFGPolicy,
+    )
+
+# Tracks (n_branches, cfg_world_size, cfg_rank) tuples already logged so the
+# dispatch table is printed once per unique configuration, not once per step.
+_logged_dispatch_keys: set[tuple[int, int, int]] = set()
+
+
+def _run(
+    predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
+    bid: int,
+    branches,
+) -> tuple[torch.Tensor, ...]:
+    branch = branches[bid]
+    device = get_local_torch_device()
+    local_branch = dataclasses.replace(
+        branch,
+        kwargs={
+            k: v.to(device) if isinstance(v, torch.Tensor) else v
+            for k, v in branch.kwargs.items()
+        },
+    )
+    raw = predict_fn(local_branch)
+    return _wrap(raw)
+
+
+def run_cfg_parallel(
+    policy: CFGPolicy,
+    predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
+) -> list[torch.Tensor | tuple[torch.Tensor, ...]]:
+    """Dispatch CFG branches across ranks, all-gather results, return in branch order.
+
+    ``predict_fn`` is a closure capturing all step-varying state
+    (latent_model_input, timestep, model, etc.).  It is called with each
+    assigned ``CFGBranch`` and must return the raw ``_predict_noise`` output.
+
+    Idle ranks (cfg_world_size > n_branches) run branch 0 as a dummy forward
+    to obtain tensor shapes for the all-gather.
+
+    Returns a list indexed to match ``policy.branches``, identical on every rank.
+    """
+
+    cfg_rank = get_classifier_free_guidance_rank()
+    cfg_world_size = get_classifier_free_guidance_world_size()
+    branches = policy.branches
+    n_branches = len(branches)
+    assignments = dispatch_branches(n_branches, cfg_world_size)
+    branches_assigned_to_local_rank = assignments[cfg_rank]
+    max_num_branches_per_rank = max(len(a) for a in assignments)
+
+    if cfg_world_size > n_branches:
+        logger.warning_once(
+            "cfg_parallel_size=%d > n_branches=%d; %d GPU(s) will be idle for CFG",
+            cfg_world_size,
+            n_branches,
+            cfg_world_size - n_branches,
+        )
+
+    dispatch_key = (n_branches, cfg_world_size, cfg_rank)
+    if dispatch_key not in _logged_dispatch_keys:
+        _logged_dispatch_keys.add(dispatch_key)
+        branch_names = (
+            [branches[i].name for i in branches_assigned_to_local_rank]
+            if branches_assigned_to_local_rank
+            else ["(idle)"]
+        )
+        logger.info(
+            "CFG parallel dispatch: rank %d/%d -> [%s]",
+            cfg_rank,
+            cfg_world_size,
+            ", ".join(branch_names),
+        )
+
+    # perform the forward for local branches
+    predicts_from_local_branches: list[tuple[torch.Tensor, ...]] = [
+        _run(predict_fn, bid, branches) for bid in branches_assigned_to_local_rank
+    ]
+
+    if not predicts_from_local_branches:  # idle rank: run branch 0 for tensor shapes
+        predicts_from_local_branches.append(_run(predict_fn, 0, branches))
+
+    # pad the predicts to the length of max_num_branches_per_rank, to prepare for the all-gather later
+    ref = predicts_from_local_branches[0]
+    while len(predicts_from_local_branches) < max_num_branches_per_rank:
+        # TODO: cache this zero
+        predicts_from_local_branches.append(tuple(torch.zeros_like(t) for t in ref))
+
+    # All-gather each slot and output element with separate_tensors=True.
+    # all_slots[slot][elem] = list[Tensor] indexed by CFG rank; no reshape.
+    all_slots: list[list[list[torch.Tensor]]] = [
+        [
+            cfg_model_parallel_all_gather(p, dim=0, separate_tensors=True)
+            for p in slot_pred
+        ]
+        for slot_pred in predicts_from_local_branches
+    ]
+
+    # reorder the results in branch order: branch bid -> owner rank, slot.
+    n_elems = len(ref)
+    final: list[torch.Tensor | tuple[torch.Tensor, ...]] = []
+    for bid in range(n_branches):
+        owner = bid % cfg_world_size
+        slot = bid // cfg_world_size
+        elems = tuple(all_slots[slot][ei][owner] for ei in range(n_elems))
+        final.append(_unwrap(elems))
+    return final
+
+
+def run_two_branch_cfg_parallel(
+    policy: CFGPolicy,
+    predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
+    cfg_scale: float,
+    batch,
+    pipeline_config,
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    """Run standard two-pass CFG with the old all-reduce combine.
+
+    This keeps the existing WAN baselines: it avoids gathering both branch
+    predictions, and it preserves the bf16 arithmetic order used before the
+    multi-branch CFG dispatcher was added.
+    """
+
+    cfg_rank = get_classifier_free_guidance_rank()
+    pred_t = _run(predict_fn, cfg_rank, policy.branches)
+
+    if cfg_rank == 0:
+        partial = tuple(cfg_scale * p for p in pred_t)
+        cond_t = pred_t
+    else:
+        partial = tuple((1 - cfg_scale) * p for p in pred_t)
+        cond_t = tuple(torch.empty_like(p) for p in pred_t)
+
+    results = [cfg_model_parallel_all_reduce(p) for p in partial]
+    cond_t = tuple(get_cfg_group().broadcast(p, src=0) for p in cond_t)
+    results[0] = _apply_cfg_postprocess(results[0], cond_t[0], batch, pipeline_config)
+    return _unwrap(tuple(results))
+
+
+def dispatch_branches(n_branches: int, n_ranks: int) -> list[list[int]]:
+    """Assign branches to ranks in Round-robin fashion
+
+    Returns a list of length ``n_ranks`` where element ``r`` contains the
+    branch indices assigned to rank ``r``.  Branch ``i`` goes to rank
+    ``i % n_ranks``.
+
+    Example: 4 passes, 2 GPUs:
+        rank 0 -> [0, 2],  rank 1 -> [1, 3]
+    """
+    assignments: list[list[int]] = [[] for _ in range(n_ranks)]
+    for i in range(n_branches):
+        assignments[i % n_ranks].append(i)
+    return assignments

python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+
+
+@dataclass
+class CFGBranch:
+    """Immutable specification of one CFG branch forward pass.
+
+    Built once before the denoising loop; read-only across all steps.
+    """
+
+    name: str
+    is_conditional: bool
+    kwargs: dict[str, Any]
+
+    def configure_batch(self, batch: Req) -> None:
+        """Set batch state before this branch's forward pass.
+
+        Override for richer per-branch context (e.g. a branch index instead of
+        a single boolean) when a model needs more than two guidance modes.
+        """
+        batch.is_cfg_negative = not self.is_conditional
+
+
+@dataclass
+class CFGPolicy:
+    """Owns the CFG branches for one generation run and combines their predictions.
+
+    Built once before the denoising loop via ``build()``, then used read-only
+    across all steps.  Subclass and override ``build()`` / ``combine()`` for
+    custom CFG schemes (N-branch, multi-output, etc.).
+
+    The default implementation handles standard 2-branch CFG.  With a single
+    branch (CFG disabled) ``combine()`` returns the prediction unchanged.
+    """
+
+    branches: list[CFGBranch] = field(default_factory=list)
+
+    def build(
+        self,
+        batch: Req,
+        image_kwargs: dict[str, Any],
+        pos_cond_kwargs: dict[str, Any],
+        neg_cond_kwargs: dict[str, Any],
+    ) -> CFGPolicy:
+        """Return a new policy with branches populated.
+
+        Called once before the denoising loop.  The returned policy is
+        immutable for the lifetime of the run.  Override to declare N branches.
+        """
+        branches = [CFGBranch("conditional", True, {**image_kwargs, **pos_cond_kwargs})]
+        if batch.do_classifier_free_guidance:
+            branches.append(
+                CFGBranch("unconditional", False, {**image_kwargs, **neg_cond_kwargs})
+            )
+        return dataclasses.replace(self, branches=branches)
+
+    def combine(
+        self,
+        predictions: list[torch.Tensor | tuple[torch.Tensor, ...]],
+        batch: Req,
+        cfg_scale: float,
+        pipeline_config: Any,
+        *,
+        cfg_parallel: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Combine branch predictions into the final noise estimate.
+
+        Default: standard 2-branch CFG formula applied element-wise, followed
+        by normalization / rescale / model-specific postprocess.
+        Single-branch (CFG disabled): returns the prediction unchanged.
+        Override for N-branch or multi-output models.
+        """
+        if len(predictions) == 1:
+            return predictions[0]
+        pos_t = _wrap(predictions[0])
+        neg_t = _wrap(predictions[1])
+        if cfg_parallel:
+            # Match the old CFG-parallel calculation: multiply the positive
+            # prediction by cfg_scale and the negative prediction by
+            # (1 - cfg_scale) before adding them. The serial CFG formula is
+            # mathematically equivalent, but bf16 rounding changes WAN outputs.
+            results = [
+                cfg_scale * p + (1 - cfg_scale) * n for p, n in zip(pos_t, neg_t)
+            ]
+        else:
+            results = [n + cfg_scale * (p - n) for p, n in zip(pos_t, neg_t)]
+        results[0] = _apply_cfg_postprocess(
+            results[0], pos_t[0], batch, pipeline_config
+        )
+        return _unwrap(tuple(results))
+
+
+# Helpers used by CFGPolicy and run_cfg_parallel.
+
+
+def _wrap(
+    pred: torch.Tensor | tuple[torch.Tensor, ...],
+) -> tuple[torch.Tensor, ...]:
+    return pred if isinstance(pred, tuple) else (pred,)
+
+
+def _unwrap(
+    pred: tuple[torch.Tensor, ...],
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    return pred[0] if len(pred) == 1 else pred
+
+
+def _apply_cfg_postprocess(
+    noise_pred: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    batch: Req,
+    pipeline_config: Any,
+) -> torch.Tensor:
+    if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
+        noise_pred = _apply_cfg_normalization(
+            noise_pred, noise_pred_cond, float(batch.cfg_normalization)
+        )
+    if batch.guidance_rescale > 0.0:
+        noise_pred = _rescale_noise_cfg(
+            noise_pred, noise_pred_cond, guidance_rescale=batch.guidance_rescale
+        )
+    return pipeline_config.postprocess_cfg_noise(batch, noise_pred, noise_pred_cond)
+
+
+def _apply_cfg_normalization(
+    noise_pred: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    cfg_normalization: float,
+) -> torch.Tensor:
+    cond_f = noise_pred_cond.float()
+    pred_f = noise_pred.float()
+    ori_norm = torch.linalg.vector_norm(cond_f)
+    new_norm = torch.linalg.vector_norm(pred_f)
+    max_norm = ori_norm * cfg_normalization
+    if new_norm > max_norm:
+        noise_pred = noise_pred * (max_norm / new_norm)
+    return noise_pred
+
+
+def _rescale_noise_cfg(
+    noise_cfg: torch.Tensor,
+    noise_pred_text: torch.Tensor,
+    guidance_rescale: float = 0.0,
+) -> torch.Tensor:
+    std_text = noise_pred_text.std(
+        dim=list(range(1, noise_pred_text.ndim)), keepdim=True
+    )
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    return guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg

python/sglang/multimodal_gen/runtime/distributed/communication_op.py

@@ -0,0 +1,67 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py
+
+import torch
+import torch.distributed as dist
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_cfg_group,
+    get_sp_group,
+    get_tp_group,
+)
+
+
+def tensor_model_parallel_all_reduce(
+    input_: torch.Tensor, tp_group: dist.ProcessGroup = None
+) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    tp_group = tp_group or get_tp_group()
+    return tp_group.all_reduce(input_)
+
+
+def tensor_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1, tp_group: dist.ProcessGroup = None
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    tp_group = tp_group or get_tp_group()
+    return tp_group.all_gather(input_, dim)
+
+
+# TODO: remove model, make it sequence_parallel
+def sequence_model_parallel_all_to_all_4D(
+    input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1
+) -> torch.Tensor:
+    """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group."""
+    return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim)
+
+
+def sequence_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_sp_group().all_gather(input_, dim)
+
+
+def sequence_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    return get_sp_group().all_reduce(input_)
+
+
+def cfg_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1, separate_tensors: bool = False
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_cfg_group().all_gather(input_, dim, separate_tensors)
+
+
+def cfg_model_parallel_all_reduce(
+    input_: torch.Tensor,
+    op: torch._C._distributed_c10d.ReduceOp = torch._C._distributed_c10d.ReduceOp.SUM,
+) -> torch.Tensor:
+    """All-reduce the input tensor across CFG parallel group."""
+    if not input_.is_contiguous():
+        input_ = input_.contiguous()
+    return get_cfg_group().all_reduce(input_, op=op)

python/sglang/multimodal_gen/runtime/distributed/device_communicators/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py

@@ -0,0 +1,306 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py
+
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+from torch.distributed import ProcessGroup, ReduceOp
+
+
+class DistributedAutograd:
+    """Collection of autograd functions for distributed operations.
+
+    This class provides custom autograd functions for distributed operations like all_reduce,
+    all_gather, and all_to_all. Each operation is implemented as a static inner class with
+    proper forward and backward implementations.
+    """
+
+    class AllReduce(torch.autograd.Function):
+        """Differentiable all_reduce operation.
+
+        The gradient of all_reduce is another all_reduce operation since the operation
+        combines values from all ranks equally.
+        """
+
+        @staticmethod
+        def forward(
+            ctx: Any,
+            group: ProcessGroup,
+            input_: Tensor,
+            op: dist.ReduceOp | None = None,
+        ) -> Tensor:
+            ctx.group = group
+            ctx.op = op
+            output = input_.clone()
+            dist.all_reduce(output, group=group, op=op)
+            return output
+
+        @staticmethod
+        def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None]:
+            grad_output = grad_output.clone()
+            dist.all_reduce(grad_output, group=ctx.group, op=ctx.op)
+            return None, grad_output, None
+
+    class AllGather(torch.autograd.Function):
+        """Differentiable all_gather operation.
+
+        The operation gathers tensors from all ranks and concatenates them along a specified dimension.
+        The backward pass uses reduce_scatter to efficiently distribute gradients back to source ranks.
+        """
+
+        @staticmethod
+        def forward(
+            ctx: Any, group: ProcessGroup, input_: Tensor, world_size: int, dim: int
+        ) -> Tensor:
+            ctx.group = group
+            ctx.world_size = world_size
+            ctx.dim = dim
+            ctx.input_shape = input_.shape
+
+            input_size = input_.size()
+            output_size = (input_size[0] * world_size,) + input_size[1:]
+            output_tensor = torch.empty(
+                output_size, dtype=input_.dtype, device=input_.device
+            )
+
+            dist.all_gather_into_tensor(output_tensor, input_, group=group)
+
+            output_tensor = output_tensor.reshape((world_size,) + input_size)
+            output_tensor = output_tensor.movedim(0, dim)
+            output_tensor = output_tensor.reshape(
+                input_size[:dim]
+                + (world_size * input_size[dim],)
+                + input_size[dim + 1 :]
+            )
+            return output_tensor
+
+        @staticmethod
+        def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None, None]:
+            # Split the gradient tensor along the gathered dimension
+            dim_size = grad_output.size(ctx.dim) // ctx.world_size
+            grad_chunks = grad_output.reshape(
+                grad_output.shape[: ctx.dim]
+                + (ctx.world_size, dim_size)
+                + grad_output.shape[ctx.dim + 1 :]
+            )
+            grad_chunks = grad_chunks.movedim(ctx.dim, 0)
+
+            # Each rank only needs its corresponding gradient
+            grad_input = torch.empty(
+                ctx.input_shape, dtype=grad_output.dtype, device=grad_output.device
+            )
+            dist.reduce_scatter_tensor(
+                grad_input, grad_chunks.contiguous(), group=ctx.group
+            )
+
+            return None, grad_input, None, None
+
+    class AllToAll4D(torch.autograd.Function):
+        """Differentiable all_to_all operation specialized for 4D tensors.
+
+        This operation is particularly useful for attention operations where we need to
+        redistribute data across ranks for efficient parallel processing.
+
+        The operation supports two modes:
+        1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads
+        2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions
+        """
+
+        @staticmethod
+        def forward(
+            ctx: Any,
+            group: ProcessGroup,
+            input_: Tensor,
+            world_size: int,
+            scatter_dim: int,
+            gather_dim: int,
+        ) -> Tensor:
+            ctx.group = group
+            ctx.world_size = world_size
+            ctx.scatter_dim = scatter_dim
+            ctx.gather_dim = gather_dim
+
+            if world_size == 1:
+                return input_
+
+            assert (
+                input_.dim() == 4
+            ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
+
+            if scatter_dim == 2 and gather_dim == 1:
+                bs, shard_seqlen, hn, hd = input_.shape
+                assert hn % world_size == 0, (
+                    f"head dimension ({hn}) must be divisible by sequence "
+                    f"parallel world size ({world_size})"
+                )
+                seqlen = shard_seqlen * world_size
+                shard_hn = hn // world_size
+
+                input_ = input_.transpose(0, 2).contiguous()  # hn, shard_seqlen, bs, hd
+                output = torch.empty_like(input_)
+
+                dist.all_to_all_single(
+                    output, input_, group=group
+                )  # hn, shard_seqlen, bs, hd
+
+                output = torch.cat(
+                    output.split(shard_hn), dim=1
+                )  # sharded hn, seqlen, bs, hd
+
+                output = output.transpose(
+                    0, 2
+                ).contiguous()  # bs, seqlen, sharded_hn, hd
+
+                return output
+            elif scatter_dim == 1 and gather_dim == 2:
+                bs, seqlen, shard_hn, hd = input_.shape
+                assert seqlen % world_size == 0, (
+                    f"sequence dimension ({seqlen}) must be divisible by sequence "
+                    f"parallel world size ({world_size})"
+                )
+                hn = shard_hn * world_size
+                shard_seqlen = seqlen // world_size
+
+                input_ = input_.transpose(0, 2).contiguous()  # shard_hn, seqlen, bs, hd
+
+                input_ = (
+                    input_.reshape(shard_hn, world_size, shard_seqlen, bs, hd)
+                    .transpose(0, 1)
+                    .reshape(shard_hn * world_size, shard_seqlen, bs, hd)
+                    .contiguous()
+                )
+
+                output = torch.empty_like(input_)
+
+                dist.all_to_all_single(output, input_, group=group)
+
+                output = output.transpose(
+                    0, 2
+                ).contiguous()  # bs, seqlen, sharded_hn, hd
+
+                return output
+            else:
+                raise RuntimeError(
+                    f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. "
+                    f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported."
+                )
+
+        @staticmethod
+        def backward(
+            ctx: Any, grad_output: Tensor
+        ) -> tuple[None, Tensor, None, None, None]:
+            if ctx.world_size == 1:
+                return None, grad_output, None, None, None
+
+            # For backward pass, we swap scatter_dim and gather_dim
+            output = DistributedAutograd.AllToAll4D.apply(
+                ctx.group, grad_output, ctx.world_size, ctx.gather_dim, ctx.scatter_dim
+            )
+            return None, output, None, None, None
+
+
+class DeviceCommunicatorBase:
+    """
+    Base class for device-specific communicator with autograd support.
+    It can use the `cpu_group` to initialize the communicator.
+    If the device has PyTorch integration (PyTorch can recognize its
+    communication backend), the `device_group` will also be given.
+    """
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        self.device = device or torch.device("cpu")
+        self.cpu_group = cpu_group
+        self.device_group = device_group
+        self.unique_name = unique_name
+        self.rank = dist.get_rank(cpu_group)
+        self.world_size = dist.get_world_size(cpu_group)
+        self.ranks = dist.get_process_group_ranks(cpu_group)
+        self.global_rank = dist.get_rank()
+        self.global_world_size = dist.get_world_size()
+        self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank)
+
+    def all_reduce(
+        self, input_: torch.Tensor, op: dist.ReduceOp | None = ReduceOp.SUM
+    ) -> torch.Tensor:
+        """Performs an all_reduce operation with gradient support."""
+        return DistributedAutograd.AllReduce.apply(self.device_group, input_, op)
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        """Performs an all_gather operation with gradient support."""
+        if dim < 0:
+            dim += input_.dim()
+        return DistributedAutograd.AllGather.apply(
+            self.device_group, input_, self.world_size, dim
+        )
+
+    def all_to_all_4D(
+        self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1
+    ) -> torch.Tensor:
+        """Performs a 4D all-to-all operation with gradient support."""
+        return DistributedAutograd.AllToAll4D.apply(
+            self.device_group, input_, self.world_size, scatter_dim, gather_dim
+        )
+
+    def gather(
+        self, input_: torch.Tensor, dst: int = 0, dim: int = -1
+    ) -> torch.Tensor | None:
+        """
+        NOTE: We assume that the input tensor is on the same device across
+        all the ranks.
+        NOTE: `dst` is the local rank of the destination rank.
+        """
+        world_size = self.world_size
+        assert (
+            -input_.dim() <= dim < input_.dim()
+        ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Allocate output tensor.
+        if self.rank_in_group == dst:
+            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+        else:
+            gather_list = None
+        # Gather.
+        torch.distributed.gather(
+            input_, gather_list, dst=self.ranks[dst], group=self.device_group
+        )
+        if self.rank_in_group == dst:
+            output_tensor = torch.cat(gather_list, dim=dim)
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        """Sends a tensor to the destination rank in a non-blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+        torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(
+        self, size: torch.Size, dtype: torch.dtype, src: int | None = None
+    ) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def destroy(self) -> None:
+        pass

python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py

@@ -0,0 +1,162 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/cpu_communicator.py
+
+import os
+
+import torch
+from torch.distributed import ProcessGroup
+
+from .base_device_communicator import DeviceCommunicatorBase
+
+
+class CpuCommunicator(DeviceCommunicatorBase):
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        from sglang.multimodal_gen.runtime.platforms import current_platform
+        from sglang.multimodal_gen.runtime.platforms.interface import CpuArchEnum
+
+        super().__init__(cpu_group, device, device_group, unique_name)
+        self.dist_module = torch.distributed
+
+        if (
+            (current_platform.get_cpu_architecture() == CpuArchEnum.X86)
+            and hasattr(torch.ops._C, "init_shm_manager")
+            and unique_name.startswith("tp")
+        ):
+            self.dist_module = _CPUSHMDistributed(self)
+
+    def all_reduce(
+        self,
+        input_: torch.Tensor,
+        op: torch.distributed.ReduceOp | None = torch.distributed.ReduceOp.SUM,
+    ) -> torch.Tensor:
+        self.dist_module.all_reduce(input_, group=self.device_group, op=op)
+        return input_
+
+    def gather(
+        self, input_: torch.Tensor, dst: int = 0, dim: int = -1
+    ) -> torch.Tensor | None:
+        """
+        NOTE: We assume that the input tensor is on the same device across
+        all the ranks.
+        NOTE: `dst` is the local rank of the destination rank.
+        """
+        world_size = self.world_size
+        assert (
+            -input_.dim() <= dim < input_.dim()
+        ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Allocate output tensor.
+        if self.rank_in_group == dst:
+            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+        else:
+            gather_list = None
+
+        # Gather.
+        self.dist_module.gather(
+            input_, gather_list, dst=self.ranks[dst], group=self.device_group
+        )
+
+        if self.rank_in_group == dst:
+            output_tensor = torch.cat(gather_list, dim=dim)
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+        input_size = input_.size()
+        # NOTE: we have to use concat-style all-gather here,
+        # stack-style all-gather has compatibility issues with
+        # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
+        output_size = (input_size[0] * self.world_size,) + input_size[1:]
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            output_size, dtype=input_.dtype, device=input_.device
+        )
+        # All-gather.
+        self.dist_module.all_gather_into_tensor(
+            output_tensor, input_, group=self.device_group
+        )
+
+        # Reshape
+        output_tensor = output_tensor.reshape((self.world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim]
+            + (self.world_size * input_size[dim],)
+            + input_size[dim + 1 :]
+        )
+        return output_tensor
+
+
+class _CPUSHMDistributed:
+
+    def __init__(self, communicator: CpuCommunicator):
+        instance_identifier = os.environ["VLLM_DIST_IDENT"]
+        unique_name = communicator.unique_name
+        instance_identifier = f"{instance_identifier}-{unique_name}"
+        self.communicator = communicator
+
+        group_ranks = [str(rank) for rank in self.communicator.ranks]
+        shm_group_identifier = f"[{'-'.join(group_ranks)}]"
+        self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm"
+
+        self.handle = self._init_cpu_shm()
+
+    def _init_cpu_shm(self) -> int:
+        handle = torch.ops._C.init_shm_manager(
+            self.group_name,
+            self.communicator.world_size,
+            self.communicator.rank,
+        )
+        torch.distributed.barrier(self.communicator.device_group)
+        torch.ops._C.join_shm_manager(
+            handle,
+            self.group_name,
+        )
+        torch.distributed.barrier(self.communicator.device_group)
+
+        return int(handle)
+
+    def all_reduce(
+        self, input: torch.Tensor, group: ProcessGroup | None = None
+    ) -> None:
+        torch.ops._C.shm_allreduce(self.handle, input)
+
+    def gather(
+        self,
+        input: torch.Tensor,
+        gather_list: list[torch.Tensor] | None,
+        dst: int = -1,
+        group: ProcessGroup | None = None,
+    ) -> None:
+        # Note: different from the torch gather, here we use local dst rank.
+        torch.ops._C.shm_gather(
+            self.handle,
+            input,
+            gather_list,
+            torch.distributed.get_group_rank(group, dst),
+        )
+
+    def all_gather_into_tensor(
+        self,
+        output: torch.Tensor,
+        input: torch.Tensor,
+        group: ProcessGroup | None = None,
+    ) -> None:
+        torch.ops._C.shm_all_gather(self.handle, input, output)

python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py

@@ -0,0 +1,80 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py
+
+import torch
+from torch.distributed import ProcessGroup
+
+from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator import (
+    DeviceCommunicatorBase,
+)
+
+
+class CudaCommunicator(DeviceCommunicatorBase):
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        super().__init__(cpu_group, device, device_group, unique_name)
+
+        from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl import (
+            PyNcclCommunicator,
+        )
+
+        self.pynccl_comm: PyNcclCommunicator | None = None
+        if self.world_size > 1:
+            self.pynccl_comm = PyNcclCommunicator(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+    def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None):
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        out = pynccl_comm.all_reduce(input_, op=op)
+        if out is None:
+            # fall back to the default all-reduce using PyTorch.
+            # this usually happens during testing.
+            # when we run the model, allreduce only happens for the TP
+            # group, where we always have either custom allreduce or pynccl.
+            out = input_.clone()
+            torch.distributed.all_reduce(out, group=self.device_group, op=op)
+        return out
+
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        """Sends a tensor to the destination rank in a non-blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.send(tensor, dst)
+        else:
+            torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(
+        self, size: torch.Size, dtype: torch.dtype, src: int | None = None
+    ) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.recv(tensor, src)
+        else:
+            torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def destroy(self) -> None:
+        if self.pynccl_comm is not None:
+            self.pynccl_comm = None

python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py

@@ -0,0 +1,259 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py
+
+# ===================== import region =====================
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup, ReduceOp
+
+from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl_wrapper import (
+    NCCLLibrary,
+    buffer_type,
+    cudaStream_t,
+    ncclComm_t,
+    ncclDataTypeEnum,
+    ncclRedOpTypeEnum,
+    ncclUniqueId,
+)
+from sglang.multimodal_gen.runtime.distributed.utils import StatelessProcessGroup
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import current_stream
+
+logger = init_logger(__name__)
+
+
+class PyNcclCommunicator:
+
+    def __init__(
+        self,
+        group: ProcessGroup | StatelessProcessGroup,
+        device: int | str | torch.device,
+        library_path: str | None = None,
+    ):
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the PyNcclCommunicator to. If None,
+                it will be bind to f"cuda:{local_rank}".
+            library_path: the path to the NCCL library. If None, it will
+                use the default library path.
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device.
+        """
+        if not isinstance(group, StatelessProcessGroup):
+            assert dist.is_initialized()
+            assert (
+                dist.get_backend(group) != dist.Backend.NCCL
+            ), "PyNcclCommunicator should be attached to a non-NCCL group."
+            # note: this rank is the rank in the group
+            self.rank = dist.get_rank(group)
+            self.world_size = dist.get_world_size(group)
+        else:
+            self.rank = group.rank
+            self.world_size = group.world_size
+
+        self.group = group
+
+        # if world_size == 1, no need to create communicator
+        if self.world_size == 1:
+            self.available = False
+            self.disabled = True
+            return
+        try:
+            self.nccl = NCCLLibrary(library_path)
+        except Exception:
+            # disable because of missing NCCL library
+            # e.g. in a non-GPU environment
+            self.available = False
+            self.disabled = True
+            return
+
+        self.available = True
+        self.disabled = False
+
+        logger.info("sglang-diffusion is using nccl==%s", self.nccl.ncclGetVersion())
+
+        if self.rank == 0:
+            # get the unique id from NCCL
+            self.unique_id = self.nccl.ncclGetUniqueId()
+        else:
+            # construct an empty unique id
+            self.unique_id = ncclUniqueId()
+
+        if not isinstance(group, StatelessProcessGroup):
+            tensor = torch.ByteTensor(list(self.unique_id.internal))
+            ranks = dist.get_process_group_ranks(group)
+            # arg `src` in `broadcast` is the global rank
+            dist.broadcast(tensor, src=ranks[0], group=group)
+            byte_list = tensor.tolist()
+            for i, byte in enumerate(byte_list):
+                self.unique_id.internal[i] = byte
+        else:
+            self.unique_id = group.broadcast_obj(self.unique_id, src=0)
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+        # nccl communicator and stream will use this device
+        # `torch.cuda.device` is a context manager that changes the
+        # current cuda device to the specified one
+        with torch.cuda.device(device):
+            self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                self.world_size, self.unique_id, self.rank
+            )
+
+            stream = current_stream()
+            # A small all_reduce for warmup.
+            data = torch.zeros(1, device=device)
+            self.all_reduce(data)
+            if stream is not None:
+                stream.synchronize()
+            del data
+
+    def all_reduce(
+        self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None
+    ) -> torch.Tensor:
+        if self.disabled:
+            return None
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert in_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {in_tensor.device}"
+        )
+
+        out_tensor = torch.empty_like(in_tensor)
+
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllReduce(
+            buffer_type(in_tensor.data_ptr()),
+            buffer_type(out_tensor.data_ptr()),
+            in_tensor.numel(),
+            ncclDataTypeEnum.from_torch(in_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+        return out_tensor
+
+    def all_gather(
+        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllGather(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            input_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def reduce_scatter(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclReduceScatter(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            output_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def send(self, tensor: torch.Tensor, dst: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclSend(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            dst,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def recv(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclRecv(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        if src == self.rank:
+            sendbuff = buffer_type(tensor.data_ptr())
+            # NCCL requires the sender also to have a receive buffer
+            recvbuff = buffer_type(tensor.data_ptr())
+        else:
+            sendbuff = buffer_type()
+            recvbuff = buffer_type(tensor.data_ptr())
+        self.nccl.ncclBroadcast(
+            sendbuff,
+            recvbuff,
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )

python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py

@@ -0,0 +1,451 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py
+
+# This file is a pure Python wrapper for the NCCL library.
+# The main purpose is to use NCCL combined with CUDA graph.
+# Before writing this script, we tried the following approach:
+# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
+#  often gets stuck when initializing the NCCL communicator.
+# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
+#  contains many other potential cuda APIs, that are not allowed during
+#  capturing the CUDA graph. For further details, please check
+# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
+#
+# Another rejected idea is to write a C/C++ binding for NCCL. It is usually
+# doable, but we often encounter issues related with nccl versions, and need
+# to switch between different versions of NCCL. See
+# https://github.com/NVIDIA/nccl/issues/1234 for more details.
+# A C/C++ binding is not flexible enough to handle this. It requires
+# recompilation of the code every time we want to switch between different
+# versions. This current implementation, with a **pure** Python wrapper, is
+# more flexible. We can easily switch between different versions of NCCL by
+# changing the environment variable `SGLANG_DIFFUSION_NCCL_SO_PATH`, or the `so_file`
+# variable in the code.
+
+# TODO(will): support SGLANG_DIFFUSION_NCCL_SO_PATH
+
+import ctypes
+import platform
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch.distributed import ReduceOp
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import find_nccl_library
+
+logger = init_logger(__name__)
+
+# === export types and functions from nccl to Python ===
+# for the original nccl definition, please check
+# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
+
+ncclResult_t = ctypes.c_int
+ncclComm_t = ctypes.c_void_p
+
+
+class ncclUniqueId(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * 128)]
+
+
+cudaStream_t = ctypes.c_void_p
+buffer_type = ctypes.c_void_p
+
+ncclDataType_t = ctypes.c_int
+
+
+class ncclDataTypeEnum:
+    ncclInt8 = 0
+    ncclChar = 0
+    ncclUint8 = 1
+    ncclInt32 = 2
+    ncclInt = 2
+    ncclUint32 = 3
+    ncclInt64 = 4
+    ncclUint64 = 5
+    ncclFloat16 = 6
+    ncclHalf = 6
+    ncclFloat32 = 7
+    ncclFloat = 7
+    ncclFloat64 = 8
+    ncclDouble = 8
+    ncclBfloat16 = 9
+    ncclNumTypes = 10
+
+    @classmethod
+    def from_torch(cls, dtype: torch.dtype) -> int:
+        if dtype == torch.int8:
+            return cls.ncclInt8
+        if dtype == torch.uint8:
+            return cls.ncclUint8
+        if dtype == torch.int32:
+            return cls.ncclInt32
+        if dtype == torch.int64:
+            return cls.ncclInt64
+        if dtype == torch.float16:
+            return cls.ncclFloat16
+        if dtype == torch.float32:
+            return cls.ncclFloat32
+        if dtype == torch.float64:
+            return cls.ncclFloat64
+        if dtype == torch.bfloat16:
+            return cls.ncclBfloat16
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+ncclRedOp_t = ctypes.c_int
+
+
+class ncclRedOpTypeEnum:
+    ncclSum = 0
+    ncclProd = 1
+    ncclMax = 2
+    ncclMin = 3
+    ncclAvg = 4
+    ncclNumOps = 5
+
+    @classmethod
+    def from_torch(cls, op: ReduceOp) -> int:
+        if op == ReduceOp.SUM:
+            return cls.ncclSum
+        if op == ReduceOp.PRODUCT:
+            return cls.ncclProd
+        if op == ReduceOp.MAX:
+            return cls.ncclMax
+        if op == ReduceOp.MIN:
+            return cls.ncclMin
+        if op == ReduceOp.AVG:
+            return cls.ncclAvg
+        raise ValueError(f"Unsupported op: {op}")
+
+
+@dataclass
+class Function:
+    name: str
+    restype: Any
+    argtypes: list[Any]
+
+
+class NCCLLibrary:
+    exported_functions = [
+        # const char* ncclGetErrorString(ncclResult_t result)
+        Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
+        # ncclResult_t  ncclGetVersion(int *version);
+        Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]),
+        # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
+        Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]),
+        # ncclResult_t  ncclCommInitRank(
+        #   ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+        # note that ncclComm_t is a pointer type, so the first argument
+        # is a pointer to a pointer
+        Function(
+            "ncclCommInitRank",
+            ncclResult_t,
+            [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int],
+        ),
+        # ncclResult_t  ncclAllReduce(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllReduce",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclAllGather(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllGather",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclReduceScatter(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclReduceScatter",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclSend(
+        #   const void* sendbuff, size_t count, ncclDataType_t datatype,
+        #   int dest, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclSend",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclRecv(
+        #   void* recvbuff, size_t count, ncclDataType_t datatype,
+        #   int src, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclRecv",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t ncclBroadcast(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, int root, ncclComm_t comm,
+        #   cudaStream_t stream);
+        Function(
+            "ncclBroadcast",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # be cautious! this is a collective call, it will block until all
+        # processes in the communicator have called this function.
+        # because Python object destruction can happen in random order,
+        # it is better not to call it at all.
+        # ncclResult_t  ncclCommDestroy(ncclComm_t comm);
+        Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
+    ]
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}
+
+    def __init__(self, so_file: str | None = None):
+
+        so_file = so_file or find_nccl_library()
+
+        try:
+            if so_file not in NCCLLibrary.path_to_dict_mapping:
+                lib = ctypes.CDLL(so_file)
+                NCCLLibrary.path_to_library_cache[so_file] = lib
+            self.lib = NCCLLibrary.path_to_library_cache[so_file]
+        except Exception as e:
+            logger.error(
+                "Failed to load NCCL library from %s ."
+                "It is expected if you are not running on NVIDIA/AMD/MTHREADS GPUs."
+                "Otherwise, the nccl library might not exist, be corrupted "
+                "or it does not support the current platform %s."
+                "If you already have the library, please set the "
+                "environment variable SGLANG_DIFFUSION_NCCL_SO_PATH"
+                " to point to the correct nccl library path.",
+                so_file,
+                platform.platform(),
+            )
+            raise e
+
+        if so_file not in NCCLLibrary.path_to_dict_mapping:
+            _funcs: dict[str, Any] = {}
+            for func in NCCLLibrary.exported_functions:
+                f = getattr(self.lib, func.name)
+                f.restype = func.restype
+                f.argtypes = func.argtypes
+                _funcs[func.name] = f
+            NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
+        self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
+
+    def ncclGetErrorString(self, result: ncclResult_t) -> str:
+        return str(self._funcs["ncclGetErrorString"](result).decode("utf-8"))
+
+    def NCCL_CHECK(self, result: ncclResult_t) -> None:
+        if result != 0:
+            error_str = self.ncclGetErrorString(result)
+            raise RuntimeError(f"NCCL error: {error_str}")
+
+    def ncclGetVersion(self) -> str:
+        version = ctypes.c_int()
+        self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
+        version_str = str(version.value)
+        # something like 21903 --> "2.19.3"
+        major = version_str[0].lstrip("0")
+        minor = version_str[1:3].lstrip("0")
+        patch = version_str[3:].lstrip("0")
+        return f"{major}.{minor}.{patch}"
+
+    def ncclGetUniqueId(self) -> ncclUniqueId:
+        unique_id = ncclUniqueId()
+        self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
+        return unique_id
+
+    def ncclCommInitRank(
+        self, world_size: int, unique_id: ncclUniqueId, rank: int
+    ) -> ncclComm_t:
+        comm = ncclComm_t()
+        self.NCCL_CHECK(
+            self._funcs["ncclCommInitRank"](
+                ctypes.byref(comm), world_size, unique_id, rank
+            )
+        )
+        return comm
+
+    def ncclAllReduce(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllReduce"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclReduceScatter(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclReduceScatter"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclAllGather(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # which is an aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllGather"](
+                sendbuff, recvbuff, count, datatype, comm, stream
+            )
+        )
+
+    def ncclSend(
+        self,
+        sendbuff: buffer_type,
+        count: int,
+        datatype: int,
+        dest: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream)
+        )
+
+    def ncclRecv(
+        self,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        src: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream)
+        )
+
+    def ncclBroadcast(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        root: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclBroadcast"](
+                sendbuff, recvbuff, count, datatype, root, comm, stream
+            )
+        )
+
+    def ncclCommDestroy(self, comm: ncclComm_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
+
+
+__all__ = [
+    "NCCLLibrary",
+    "ncclDataTypeEnum",
+    "ncclRedOpTypeEnum",
+    "ncclUniqueId",
+    "ncclComm_t",
+    "cudaStream_t",
+    "buffer_type",
+]

python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py

@@ -0,0 +1,91 @@
+# Reference: https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
+
+
+import torch
+
+
+class Singleton:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(Singleton, cls).__new__(cls, *args, **kwargs)
+        return cls._instance
+
+
+class ProcessGroupSingleton(Singleton):
+    def __init__(self):
+        self.ULYSSES_PG = None
+        self.RING_PG = None
+
+
+PROCESS_GROUP = ProcessGroupSingleton()
+
+
+def set_seq_parallel_pg_by_sp_groups(
+    sp_ulysses_degree,
+    sp_ring_degree,
+    rank: int,
+    sp_groups: list[list[int]],
+    use_ulysses_low: bool = True,
+):
+    """Create Ulysses/Ring process groups inside each SP group.
+
+    This is required when TP>1, because SP groups are not necessarily made of
+    consecutive global ranks (e.g., tp-sp order makes SP ranks strided).
+
+    Args:
+        sp_ulysses_degree: ulysses degree inside SP.
+        sp_ring_degree: ring degree inside SP.
+        rank: global rank of current process.
+        sp_groups: list of global-rank lists for each SP group.
+        use_ulysses_low: keep the same semantics as the original function.
+    """
+    sp_degree = sp_ring_degree * sp_ulysses_degree
+    assert sp_degree > 0
+    assert all(
+        len(g) == sp_degree for g in sp_groups
+    ), f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}"
+
+    ulyssess_pg = None
+    ring_pg = None
+
+    num_ulysses_pgs = sp_ring_degree
+    num_ring_pgs = sp_ulysses_degree
+
+    def _map_indices_to_ranks(ranks: list[int], indices: list[int]) -> list[int]:
+        return [ranks[i] for i in indices]
+
+    # Important: call torch.distributed.new_group in the same order on all ranks.
+    for sp_ranks in sp_groups:
+        if use_ulysses_low:
+            for i in range(num_ulysses_pgs):
+                idx = list(range(i * sp_ulysses_degree, (i + 1) * sp_ulysses_degree))
+                ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ulysses_ranks)
+                if rank in ulysses_ranks:
+                    ulyssess_pg = group
+
+            for i in range(num_ring_pgs):
+                idx = list(range(i, sp_degree, num_ring_pgs))
+                ring_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ring_ranks)
+                if rank in ring_ranks:
+                    ring_pg = group
+        else:
+            for i in range(num_ring_pgs):
+                idx = list(range(i * sp_ring_degree, (i + 1) * sp_ring_degree))
+                ring_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ring_ranks)
+                if rank in ring_ranks:
+                    ring_pg = group
+
+            for i in range(num_ulysses_pgs):
+                idx = list(range(i, sp_degree, num_ulysses_pgs))
+                ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ulysses_ranks)
+                if rank in ulysses_ranks:
+                    ulyssess_pg = group
+
+    PROCESS_GROUP.ULYSSES_PG = ulyssess_pg
+    PROCESS_GROUP.RING_PG = ring_pg

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -0,0 +1,912 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py
+# Copyright 2023 The vLLM team.
+# Adapted from
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# Adapted from
+# Copyright 2024 xDiT team.
+# Adapted from
+# https://github.com/vllm-project/vllm/blob/main/vllm/distributed/parallel_state.py
+# Copyright 2023 The vLLM team.
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""sglang-diffusion distributed state.
+
+It takes over the control of the distributed environment from PyTorch.
+The typical workflow is:
+
+- call `init_distributed_environment` to initialize the distributed environment.
+- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to
+ initialize the model parallel groups.
+
+- any code dealing with the distributed stuff
+
+- call `destroy_model_parallel` to destroy the model parallel groups.
+- call `destroy_distributed_environment` to destroy the distributed environment.
+
+If you only need to use the distributed environment without model parallelism,
+ you can skip the model parallel initialization and destruction steps.
+"""
+
+import contextlib
+import datetime
+import os
+import weakref
+from collections import namedtuple
+from collections.abc import Callable
+from contextlib import contextmanager
+from multiprocessing import shared_memory
+from typing import Any, List, Optional
+from unittest.mock import patch
+
+import torch
+import torch.distributed
+from torch.distributed import ProcessGroup
+
+import sglang.multimodal_gen.envs as envs
+from sglang.multimodal_gen.runtime.distributed.utils import StatelessProcessGroup
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+from ..utils.distributed import RankGenerator
+from .group_coordinator import (
+    GroupCoordinator,
+    PipelineGroupCoordinator,
+    SequenceParallelGroupCoordinator,
+    get_local_torch_device,
+)
+
+logger = init_logger(__name__)
+
+_WORLD: GroupCoordinator | None = None
+_TP: GroupCoordinator | None = None
+_SP: SequenceParallelGroupCoordinator | None = None
+_PP: PipelineGroupCoordinator | None = None
+_CFG: GroupCoordinator | None = None
+_DP: GroupCoordinator | None = None
+_VAE_DECODE: GroupCoordinator | None = None
+_DIT: ProcessGroup | None = None
+_VAE: ProcessGroup | None = None
+_VAE_DECODE_PARALLEL_AXES = "tp-sp-pp-cfg"
+
+TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
+
+
+def _split_tensor_dict(
+    tensor_dict: dict[str, torch.Tensor | Any],
+) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
+    """Split the tensor dictionary into two parts:
+    1. A list of (key, value) pairs. If the value is a tensor, it is replaced
+         by its metadata.
+    2. A list of tensors.
+    """
+    metadata_list: list[tuple[str, Any]] = []
+    tensor_list: list[torch.Tensor] = []
+    for key, value in tensor_dict.items():
+        if isinstance(value, torch.Tensor):
+            # Note: we cannot use `value.device` here,
+            # because it contains not only the device type but also the device
+            # index (e.g. "cuda:0"). We only need the device type.
+            # receiving side will set the device index.
+            device = value.device.type
+            metadata_list.append(
+                (key, TensorMetadata(device, value.dtype, value.size()))
+            )
+            tensor_list.append(value)
+        else:
+            metadata_list.append((key, value))
+    return metadata_list, tensor_list
+
+
+_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
+
+
+def _register_group(group: "GroupCoordinator") -> None:
+    _groups[group.unique_name] = weakref.ref(group)
+
+
+def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
+    assert group_name in _groups, f"Group {group_name} is not found."
+    group = _groups[group_name]()
+    if group is None:
+        raise ValueError(f"Group {group_name} is destroyed.")
+    return group._all_reduce_out_place(tensor)
+
+
+def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
+    return torch.empty_like(tensor)
+
+
+def get_world_group() -> GroupCoordinator:
+    assert _WORLD is not None, "world group is not initialized"
+    return _WORLD
+
+
+def world_group_is_initialized() -> bool:
+    return _WORLD is not None
+
+
+def init_world_group(
+    ranks: list[int], local_rank: int, backend: str
+) -> GroupCoordinator:
+    return GroupCoordinator(
+        group_ranks=[ranks],
+        local_rank=local_rank,
+        torch_distributed_backend=backend,
+        use_device_communicator=True,
+        group_name="world",
+    )
+
+
+def _sync_srt_world_group() -> None:
+    import sglang.srt.distributed.parallel_state as srt_parallel_state
+
+    if srt_parallel_state._WORLD is None:
+        srt_parallel_state._WORLD = _WORLD
+
+
+def _clear_srt_world_group() -> None:
+    import sglang.srt.distributed.parallel_state as srt_parallel_state
+
+    if srt_parallel_state._WORLD is _WORLD:
+        srt_parallel_state._WORLD = None
+
+
+def init_parallel_group_coordinator(
+    group_ranks: List[List[int]],
+    local_rank: int,
+    backend: str,
+    parallel_mode: str,
+    **kwargs,
+) -> GroupCoordinator:
+    """Return a group coordinator for the given parallel mode."""
+    assert parallel_mode in [
+        "data",
+        "pipeline",
+        "tensor",
+        "sequence",
+        "classifier_free_guidance",
+        "vae_decode",
+    ], f"parallel_mode {parallel_mode} is not supported"
+    if parallel_mode == "pipeline":
+        return PipelineGroupCoordinator(
+            group_ranks=group_ranks,
+            local_rank=local_rank,
+            torch_distributed_backend=backend,
+            group_name="pp_group",
+        )
+    elif parallel_mode == "sequence":
+        return SequenceParallelGroupCoordinator(
+            group_ranks=group_ranks,
+            local_rank=local_rank,
+            torch_distributed_backend=backend,
+            group_name="sp_group",
+            **kwargs,
+        )
+    else:
+        return GroupCoordinator(
+            group_ranks=group_ranks,
+            local_rank=local_rank,
+            torch_distributed_backend=backend,
+            use_device_communicator=parallel_mode != "tensor",
+            use_srt_custom_allreduce=parallel_mode == "tensor",
+            group_name=(
+                "tp_group"
+                if parallel_mode == "tensor"
+                else (
+                    "vae_decode_group" if parallel_mode == "vae_decode" else "cfg_group"
+                )
+            ),
+        )
+
+
+def _get_vae_decode_group_ranks(
+    rank_generator: RankGenerator,
+) -> list[list[int]]:
+    # VAE decode happens after each DP replica owns a different request result.
+    # Decode can shard one request across TP/SP/PP/CFG ranks, but must not cross DP.
+    return rank_generator.get_ranks(_VAE_DECODE_PARALLEL_AXES)
+
+
+def get_tp_group() -> GroupCoordinator:
+    assert _TP is not None, "tensor model parallel group is not initialized"
+    return _TP
+
+
+def init_distributed_environment(
+    world_size: int = 1,
+    rank: int = 0,
+    distributed_init_method: str = "env://",
+    local_rank: int = 0,
+    backend: str | None = None,
+    device_id: torch.device | None = None,
+    timeout: int | None = None,
+):
+    # Determine the appropriate backend based on the platform
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    if backend is None:
+        backend = current_platform.get_torch_distributed_backend_str()
+        logger.info(
+            "Using %s backend for %s platform", backend, current_platform.device_name
+        )
+
+    logger.debug(
+        "world_size=%d rank=%d local_rank=%d "
+        "distributed_init_method=%s backend=%s timeout=%s",
+        world_size,
+        rank,
+        local_rank,
+        distributed_init_method,
+        backend,
+        timeout,
+    )
+    if not torch.distributed.is_initialized():
+        assert distributed_init_method is not None, (
+            "distributed_init_method must be provided when initializing "
+            "distributed environment"
+        )
+
+        # For MPS, MUSA, and XPU, don't pass device_id as it doesn't support device indices
+        extra_args = (
+            {}
+            if (
+                current_platform.is_mps()
+                or current_platform.is_musa()
+                or current_platform.is_npu()
+                or current_platform.is_cpu()
+                or current_platform.is_xpu()
+            )
+            else dict(device_id=device_id)
+        )
+
+        if timeout is not None:
+
+            extra_args["timeout"] = datetime.timedelta(seconds=timeout)
+            logger.info(f"Setting distributed timeout to {timeout} seconds")
+
+        torch.distributed.init_process_group(
+            backend=backend,
+            init_method=distributed_init_method,
+            world_size=world_size,
+            rank=rank,
+            **extra_args,
+        )
+
+    # set the local rank
+    # local_rank is not available in torch ProcessGroup,
+    # see https://github.com/pytorch/pytorch/issues/122816
+    if local_rank == -1:
+        # local rank not set, this usually happens in single-node
+        # setting, where we can use rank as local rank
+        if distributed_init_method == "env://":
+            local_rank = envs.LOCAL_RANK
+        else:
+            local_rank = rank
+    global _WORLD
+    if _WORLD is None:
+        ranks = list(range(torch.distributed.get_world_size()))
+        _WORLD = init_world_group(ranks, local_rank, backend)
+    else:
+        assert (
+            _WORLD.world_size == torch.distributed.get_world_size()
+        ), "world group already initialized with a different world size"
+    _sync_srt_world_group()
+
+
+def get_sp_group() -> SequenceParallelGroupCoordinator:
+    assert _SP is not None, "sequence parallel group is not initialized"
+    return _SP
+
+
+def get_dp_group() -> GroupCoordinator:
+    assert _DP is not None, "data parallel group is not initialized"
+    return _DP
+
+
+# xDiT
+def initialize_model_parallel(
+    data_parallel_size: int = 1,
+    classifier_free_guidance_degree: int = 1,
+    sequence_parallel_degree: Optional[int] = None,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    tensor_parallel_degree: int = 1,
+    pipeline_parallel_degree: int = 1,
+    vae_parallel_size: int = 0,
+    backend: Optional[str] = None,
+) -> None:
+    """
+    Initialize model parallel groups.
+
+    Arguments:
+        data_parallel_size: number of data parallelism groups.
+        classifier_free_guidance_degree: number of GPUs used for Classifier Free Guidance (CFG)
+        sequence_parallel_degree: number of GPUs used for sequence parallelism. sequence_parallel_degree = ulysses_degree * ring_degree
+        ulysses_degree: number of GPUs used for ulysses sequence parallelism.
+        ring_degree: number of GPUs used for ring sequence parallelism.
+        tensor_parallel_degree: number of GPUs used for tensor parallelism.
+        pipeline_parallel_degree: number of GPUs used for pipeline parallelism.
+        backend: distributed backend of pytorch collective comm.
+
+    Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
+    use 2 groups to parallelize the batch dim(dp), 2 groups to parallelize
+    split batch caused by CFG, and 2 GPUs to parallelize sequence.
+
+    dp_degree (2) * cfg_degree (2) * sp_degree (2) * pp_degree (2) = 16.
+
+    The present function will create 8 data-parallel groups,
+    8 CFG group, 8 pipeline-parallel group, and
+    8 sequence-parallel groups:
+        8 data-parallel groups:
+            [g0, g8], [g1, g9], [g2, g10], [g3, g11],
+            [g4, g12], [g5, g13], [g6, g14], [g7, g15]
+        8 CFG-parallel groups:
+            [g0, g4], [g1, g5], [g2, g6], [g3, g7],
+            [g8, g12], [g9, g13], [g10, g14], [g11, g15]
+        8 sequence-parallel groups:
+            [g0, g1], [g2, g3], [g4, g5], [g6, g7],
+            [g8, g9], [g10, g11], [g12, g13], [g14, g15]
+        8 pipeline-parallel groups:
+            [g0, g2], [g4, g6], [g8, g10], [g12, g14],
+            [g1, g3], [g5, g7], [g9, g11], [g13, g15]
+    Note that for efficiency, the caller should make sure adjacent ranks
+    are on the same DGX box. For example if we are using 2 DGX-1 boxes
+    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
+    ranks 8 to 15 belong to the second box.
+    """
+
+    if backend is None:
+        from sglang.multimodal_gen.runtime.platforms import current_platform
+
+        backend = current_platform.get_torch_distributed_backend_str()
+    # Get world size and rank. Ensure some consistencies.
+    assert torch.distributed.is_initialized()
+    world_size: int = torch.distributed.get_world_size()
+    backend = backend or torch.distributed.get_backend(get_world_group().device_group)
+
+    dit_parallel_size = (
+        data_parallel_size
+        * classifier_free_guidance_degree
+        * sequence_parallel_degree
+        * pipeline_parallel_degree
+        * tensor_parallel_degree
+    )
+
+    if world_size < dit_parallel_size:
+        raise RuntimeError(
+            f"world_size ({world_size}) is less than "
+            f"tensor_parallel_degree ({tensor_parallel_degree}) x "
+            f"pipeline_parallel_degree ({pipeline_parallel_degree}) x"
+            f"sequence_parallel_degree ({sequence_parallel_degree}) x"
+            f"classifier_free_guidance_degree "
+            f"({classifier_free_guidance_degree}) x"
+            f"data_parallel_degree ({data_parallel_size})"
+        )
+
+    rank_generator: RankGenerator = RankGenerator(
+        tensor_parallel_degree,
+        sequence_parallel_degree,
+        pipeline_parallel_degree,
+        classifier_free_guidance_degree,
+        data_parallel_size,
+        "tp-sp-pp-cfg-dp",
+    )
+    global _DP
+    assert _DP is None, "data parallel group is already initialized"
+    _DP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("dp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="data",
+    )
+
+    global _CFG
+    assert _CFG is None, "classifier_free_guidance group is already initialized"
+    _CFG = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("cfg"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="classifier_free_guidance",
+    )
+    global _PP
+    assert _PP is None, "pipeline model parallel group is already initialized"
+    _PP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("pp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="pipeline",
+    )
+
+    global _SP
+    assert _SP is None, "sequence parallel group is already initialized"
+
+    try:
+        from .parallel_groups import PROCESS_GROUP as _YC_PROCESS_GROUP
+        from .parallel_groups import (
+            set_seq_parallel_pg_by_sp_groups as _set_seq_parallel_pg_by_sp_groups,
+        )
+    except ImportError:
+        _set_seq_parallel_pg_by_sp_groups = None
+
+        class _DummyProcessGroup:
+            ULYSSES_PG = torch.distributed.group.WORLD
+            RING_PG = torch.distributed.group.WORLD
+
+        PROCESS_GROUP = _DummyProcessGroup()
+    else:
+        # Build SGLang Diffusion SP sub-groups based on the true SP groups. This is
+        # critical when TP>1, because SP groups may be strided in global ranks
+        # (e.g., tp-sp order).
+        sp_groups = rank_generator.get_ranks("sp")
+        _set_seq_parallel_pg_by_sp_groups(
+            sp_ulysses_degree=ulysses_degree,
+            sp_ring_degree=ring_degree,
+            rank=get_world_group().rank,
+            sp_groups=sp_groups,
+        )
+        PROCESS_GROUP = _YC_PROCESS_GROUP
+
+    _SP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("sp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="sequence",
+        ulysses_group=PROCESS_GROUP.ULYSSES_PG,
+        ring_group=PROCESS_GROUP.RING_PG,
+    )
+
+    global _TP
+    assert _TP is None, "Tensor parallel group is already initialized"
+    _TP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("tp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="tensor",
+    )
+
+    global _VAE_DECODE
+    assert _VAE_DECODE is None, "VAE decode parallel group is already initialized"
+    _VAE_DECODE = init_parallel_group_coordinator(
+        group_ranks=_get_vae_decode_group_ranks(rank_generator),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="vae_decode",
+    )
+
+    if vae_parallel_size > 0:
+        init_vae_group(dit_parallel_size, vae_parallel_size, backend)
+    init_dit_group(dit_parallel_size, backend)
+
+
+def get_sp_world_size() -> int:
+    """Return world size for the sequence model parallel group."""
+    return get_sp_group().world_size
+
+
+def get_sp_parallel_rank() -> int:
+    """Return my rank for the sequence model parallel group."""
+    return get_sp_group().rank_in_group
+
+
+def get_world_size() -> int:
+    """Return world size for the world group."""
+    return get_world_group().world_size
+
+
+def get_world_rank() -> int:
+    """Return my rank for the world group."""
+    return get_world_group().rank
+
+
+def get_dp_world_size() -> int:
+    """Return world size for the data parallel group."""
+    return get_dp_group().world_size
+
+
+def get_dp_rank() -> int:
+    """Return my rank for the data parallel group."""
+    return get_dp_group().rank_in_group
+
+
+def maybe_init_distributed_environment_and_model_parallel(
+    tp_size: int,
+    sp_size: int,
+    cfg_degree: int = 1,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    dp_size: int = 1,
+    distributed_init_method: str = "env://",
+    dist_timeout: int | None = None,
+):
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    if _WORLD is not None and model_parallel_is_initialized():
+        # make sure the tp and sp sizes are correct
+        assert (
+            get_tp_world_size() == tp_size
+        ), f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}"
+        assert (
+            get_sp_world_size() == sp_size
+        ), f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}"
+        return
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    rank = int(os.environ.get("RANK", 0))
+    device = get_local_torch_device()
+    logger.info(
+        "Initializing distributed environment with world_size=%d, device=%s, timeout=%s",
+        world_size,
+        device,
+        dist_timeout,
+        main_process_only=False,
+    )
+
+    init_distributed_environment(
+        world_size=world_size,
+        rank=rank,
+        local_rank=local_rank,
+        distributed_init_method=distributed_init_method,
+        device_id=device,
+        backend=current_platform.get_torch_distributed_backend_str(),
+        timeout=dist_timeout,
+    )
+    initialize_model_parallel(
+        data_parallel_size=dp_size,
+        classifier_free_guidance_degree=cfg_degree,
+        tensor_parallel_degree=tp_size,
+        ulysses_degree=ulysses_degree,
+        ring_degree=ring_degree,
+        sequence_parallel_degree=sp_size,
+    )
+
+    # Only set CUDA device if we're on a CUDA platform
+    if current_platform.is_cuda_alike():
+        device = torch.device(f"cuda:{local_rank}")
+        torch.cuda.set_device(device)
+    elif current_platform.is_npu():
+        device = torch.device(f"npu:{local_rank}")
+        torch.npu.set_device(device)
+
+
+def model_parallel_is_initialized() -> bool:
+    """Check if model parallel groups are initialized."""
+    return (
+        _DP is not None
+        and _CFG is not None
+        and _SP is not None
+        and _PP is not None
+        and _TP is not None
+        and _VAE_DECODE is not None
+    )
+
+
+_TP_STATE_PATCHED = False
+
+
+@contextmanager
+def patch_tensor_parallel_group(tp_group: GroupCoordinator):
+    """Patch the tp group temporarily until this function ends.
+
+    This method is for draft workers of speculative decoding to run draft model
+    with different tp degree from that of target model workers.
+
+    """
+    global _TP_STATE_PATCHED
+    assert not _TP_STATE_PATCHED, "Should not call when it's already patched"
+
+    _TP_STATE_PATCHED = True
+    old_tp_group = get_tp_group()
+    global _TP
+    _TP = tp_group
+    try:
+        yield
+    finally:
+        # restore the original state
+        _TP_STATE_PATCHED = False
+        _TP = old_tp_group
+
+
+def get_tp_world_size() -> int:
+    """Return world size for the tensor model parallel group."""
+    return get_tp_group().world_size
+
+
+def get_tp_rank() -> int:
+    """Return my rank for the tensor model parallel group."""
+    return get_tp_group().rank_in_group
+
+
+def destroy_distributed_environment() -> None:
+    global _WORLD
+    _clear_srt_world_group()
+    if _WORLD:
+        _WORLD.destroy()
+    _WORLD = None
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
+
+
+def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
+    destroy_model_parallel()
+    destroy_distributed_environment()
+    with contextlib.suppress(AssertionError):
+        torch.distributed.destroy_process_group()
+    if shutdown_ray:
+        import ray  # Lazy import Ray
+
+        ray.shutdown()
+
+
+def is_the_same_node_as(
+    pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0
+) -> list[int]:
+    """
+    This is a collective operation that returns if each rank is in the same node
+    as the source rank. It tests if processes are attached to the same
+    memory system (shared access to shared memory).
+    """
+    if isinstance(pg, ProcessGroup):
+        assert (
+            torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL
+        ), "in_the_same_node_as should be tested with a non-NCCL group."
+        # local rank inside the group
+        rank = torch.distributed.get_rank(group=pg)
+        world_size = torch.distributed.get_world_size(group=pg)
+
+        # global ranks of the processes in the group
+        ranks = torch.distributed.get_process_group_ranks(pg)
+    else:
+        rank = pg.rank
+        world_size = pg.world_size
+        ranks = list(range(world_size))
+
+    # local tensor in each process to store the result
+    is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32)
+
+    magic_message = b"magic_message"
+    shm = None
+
+    try:
+        with contextlib.suppress(OSError):
+            if rank == source_rank:
+                # create a shared memory segment
+                shm = shared_memory.SharedMemory(create=True, size=128)
+                shm.buf[: len(magic_message)] = magic_message
+                if isinstance(pg, ProcessGroup):
+                    torch.distributed.broadcast_object_list(
+                        [shm.name], src=ranks[source_rank], group=pg
+                    )
+                else:
+                    pg.broadcast_obj(shm.name, src=source_rank)
+                is_in_the_same_node[rank] = 1
+            else:
+                # try to open the shared memory segment
+                if isinstance(pg, ProcessGroup):
+                    recv = [None]
+                    torch.distributed.broadcast_object_list(
+                        recv, src=ranks[source_rank], group=pg
+                    )
+                    name = recv[0]
+                else:
+                    name = pg.broadcast_obj(None, src=source_rank)
+                # fix to https://stackoverflow.com/q/62748654/9191338
+                # Python incorrectly tracks shared memory even if it is not
+                # created by the process. The following patch is a workaround.
+                with patch(
+                    "multiprocessing.resource_tracker.register",
+                    lambda *args, **kwargs: None,
+                ):
+                    shm = shared_memory.SharedMemory(name=name)
+                if shm.buf[: len(magic_message)] == magic_message:
+                    is_in_the_same_node[rank] = 1
+    except Exception as e:
+        logger.error("Error ignored in is_in_the_same_node: %s", e)
+    finally:
+        if shm:
+            shm.close()
+
+    if isinstance(pg, ProcessGroup):
+        torch.distributed.barrier(group=pg)
+    else:
+        pg.barrier()
+
+    # clean up the shared memory segment
+    with contextlib.suppress(OSError):
+        if rank == source_rank and shm:
+            shm.unlink()
+
+    if isinstance(pg, ProcessGroup):
+        torch.distributed.all_reduce(is_in_the_same_node, group=pg)
+        aggregated_data = is_in_the_same_node
+    else:
+        aggregated_data = torch.zeros_like(is_in_the_same_node)
+        for i in range(world_size):
+            rank_data = pg.broadcast_obj(is_in_the_same_node, src=i)
+            aggregated_data += rank_data
+
+    return [x == 1 for x in aggregated_data.tolist()]
+
+
+def get_tensor_model_parallel_world_size() -> int:
+    """Return world size for the tensor model parallel group."""
+    return get_tp_world_size()
+
+
+def get_tensor_model_parallel_rank() -> int:
+    """Return my rank for the tensor model parallel group."""
+    return get_tp_rank()
+
+
+def get_sequence_parallel_world_size() -> int:
+    """Return world size for the sequence parallel group."""
+    return get_sp_world_size()
+
+
+def get_sequence_parallel_rank() -> int:
+    """Return my rank for the sequence parallel group."""
+    return get_sp_parallel_rank()
+
+
+def get_ulysses_parallel_world_size() -> int:
+    return get_sp_group().ulysses_world_size
+
+
+def get_ulysses_parallel_rank() -> int:
+    return get_sp_group().ulysses_rank
+
+
+def get_ring_parallel_world_size() -> int:
+    return get_sp_group().ring_world_size
+
+
+def get_ring_parallel_rank() -> int:
+    return get_sp_group().ring_rank
+
+
+# PP
+def get_pp_group() -> PipelineGroupCoordinator:
+    assert _PP is not None, "pipeline model parallel group is not initialized"
+    return _PP
+
+
+def get_pipeline_parallel_world_size() -> int:
+    """Return world size for the pipeline model parallel group."""
+    return get_pp_group().world_size
+
+
+def get_pipeline_parallel_rank() -> int:
+    """Return my rank for the pipeline model parallel group."""
+    return get_pp_group().rank_in_group
+
+
+def is_pipeline_first_stage() -> bool:
+    """Return True if in the first pipeline model parallel stage, False otherwise."""
+    return get_pipeline_parallel_rank() == 0
+
+
+def is_pipeline_last_stage() -> bool:
+    """Return True if in the last pipeline model parallel stage, False otherwise."""
+    return get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1)
+
+
+# CFG
+def get_cfg_group() -> GroupCoordinator:
+    assert (
+        _CFG is not None
+    ), "classifier_free_guidance parallel group is not initialized"
+    return _CFG
+
+
+def get_classifier_free_guidance_world_size() -> int:
+    """Return world size for the classifier_free_guidance parallel group."""
+    return get_cfg_group().world_size
+
+
+def get_classifier_free_guidance_rank() -> int:
+    """Return my rank for the classifier_free_guidance parallel group."""
+    return get_cfg_group().rank_in_group
+
+
+def get_data_parallel_world_size() -> int:
+    """Return world size for the data parallel group."""
+    return get_dp_world_size()
+
+
+def get_data_parallel_rank() -> int:
+    """Return my rank for the data parallel group."""
+    return get_dp_rank()
+
+
+def is_dp_last_group() -> bool:
+    """Return True if in the last data parallel group, False otherwise."""
+    return (
+        get_sequence_parallel_rank() == (get_sequence_parallel_world_size() - 1)
+        and get_classifier_free_guidance_rank()
+        == (get_classifier_free_guidance_world_size() - 1)
+        and get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1)
+    )
+
+
+def get_dit_world_size() -> int:
+    """Return world size for the DiT model (excluding VAE)."""
+    return (
+        get_data_parallel_world_size()
+        * get_classifier_free_guidance_world_size()
+        * get_sequence_parallel_world_size()
+        * get_pipeline_parallel_world_size()
+        * get_tensor_model_parallel_world_size()
+    )
+
+
+def get_vae_parallel_group() -> ProcessGroup:
+    assert _VAE is not None, "VAE parallel group is not initialized"
+    return _VAE
+
+
+def get_vae_parallel_world_size() -> int:
+    """Return world size for the VAE parallel group."""
+    return torch.distributed.get_world_size(group=get_vae_parallel_group())
+
+
+def get_vae_parallel_rank() -> int:
+    """Return my rank for the VAE parallel group."""
+    return torch.distributed.get_rank(group=get_vae_parallel_group())
+
+
+def get_decode_parallel_group_coordinator() -> GroupCoordinator:
+    assert _VAE_DECODE is not None, "VAE decode parallel group is not initialized"
+    return _VAE_DECODE
+
+
+def get_decode_parallel_world_size() -> int:
+    return get_decode_parallel_group_coordinator().world_size
+
+
+def get_decode_parallel_rank() -> int:
+    return get_decode_parallel_group_coordinator().rank_in_group
+
+
+def init_dit_group(
+    dit_parallel_size: int,
+    backend: str,
+) -> None:
+    global _DIT
+    assert _DIT is None, "DIT group is already initialized"
+    _DIT = torch.distributed.new_group(
+        ranks=list(range(dit_parallel_size)), backend=backend
+    )
+
+
+def get_dit_group() -> ProcessGroup:
+    assert _DIT is not None, "DIT group is not initialized"
+    return _DIT
+
+
+def init_vae_group(
+    dit_parallel_size: int,
+    vae_parallel_size: int,
+    backend: str,
+):
+    # Initialize VAE group first
+    global _VAE
+    assert _VAE is None, "VAE parallel group is already initialized"
+    vae_ranks = list(range(dit_parallel_size, dit_parallel_size + vae_parallel_size))
+    _VAE = torch.distributed.new_group(ranks=vae_ranks, backend=backend)
+
+
+def destroy_model_parallel() -> None:
+    """Set the groups to none and destroy them."""
+    global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE
+
+    for group in (_TP, _SP, _DP, _CFG, _PP, _VAE_DECODE):
+        if group is not None:
+            group.destroy()
+
+    for group in (_DIT, _VAE):
+        if group is not None:
+            torch.distributed.destroy_process_group(group)
+
+    _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE = (None,) * 8

python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unified sequence-parallel shard / pad / gather helpers.
+
+Layout invariant: padding always sits at the end of the LAST rank's local
+chunk, so the ulysses-gathered sequence carries one contiguous pad block at its
+global tail. `tail_attn_meta` then lets attention skip that block for free
+(the pad becomes its own varlen segment - no repacking, no mask compute).
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_ring_parallel_world_size,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+)
+
+# Text shorter than this stays replicated instead of SP-sharded (see
+# plan_text_strategy). 0 = always shard when legal; H100 bench showed sharding
+# wins from trivial lengths on, so the knob exists only as an escape hatch.
+_TEXT_SHARD_MIN = int(os.environ.get("SGLANG_SP_TEXT_SHARD_MIN", "0"))
+
+
+@dataclass(frozen=True)
+class SpShard:
+    """Facts of one tail-padded even shard, shared by tensors of that stream."""
+
+    orig_len: int  # real tokens (global)
+    local_len: int  # per-rank chunk length (equal on every rank)
+    num_pad: int  # pad tokens, all at the last rank's local tail
+    sp_size: int
+    sp_rank: int
+
+    @property
+    def local_pad(self) -> int:
+        """Pad rows inside THIS rank's chunk (tail rows of the last rank)."""
+        return self.num_pad if self.sp_rank == self.sp_size - 1 else 0
+
+    @property
+    def local_real_len(self) -> int:
+        return self.local_len - self.local_pad
+
+
+def build_shard_plan(seq_len: int) -> SpShard:
+    """Shard math only; tensors are sliced separately via `shard_like`."""
+    sp_size = get_sp_world_size()
+    if sp_size <= 1:
+        return SpShard(seq_len, seq_len, 0, 1, 0)
+    local_len = (seq_len + sp_size - 1) // sp_size
+    return SpShard(
+        orig_len=seq_len,
+        local_len=local_len,
+        num_pad=local_len * sp_size - seq_len,
+        sp_size=sp_size,
+        sp_rank=get_sp_parallel_rank(),
+    )
+
+
+def shard_like(
+    x: torch.Tensor, shard: SpShard, dim: int = 1, pad_mode: str = "zeros"
+) -> torch.Tensor:
+    """Apply a planned shard to one tensor (RoPE caches use the same plan as
+    hidden states so their chunks stay aligned)."""
+    if shard.sp_size <= 1:
+        return x
+    if shard.num_pad > 0:
+        if pad_mode == "repeat_last":
+            pad = x.narrow(dim, x.shape[dim] - 1, 1)
+            pad = pad.expand(
+                *[shard.num_pad if i == dim else -1 for i in range(x.dim())]
+            )
+            x = torch.cat([x, pad], dim=dim)
+        else:
+            # F.pad pads dims last-to-first: (left, right) pairs from dim -1.
+            pads = [0, 0] * (x.dim() - 1 - dim) + [0, shard.num_pad]
+            x = F.pad(x, pads)
+    return x.narrow(dim, shard.sp_rank * shard.local_len, shard.local_len)
+
+
+def shard_seq(
+    x: torch.Tensor, dim: int = 1, pad_mode: str = "zeros"
+) -> tuple[torch.Tensor, SpShard]:
+    """
+    mode:
+        zeroes: pad with zeroes at tail
+        repeat_last: repeat the last token, only for rotary embedding
+    """
+    shard = build_shard_plan(x.shape[dim])
+    return shard_like(x, shard, dim=dim, pad_mode=pad_mode), shard
+
+
+def gather_seq(local: torch.Tensor, orig_len: int, dim: int = 1) -> torch.Tensor:
+    """All-gather an SP-sharded sequence and trim the tail padding"""
+    if get_sp_world_size() <= 1:
+        return local
+    full = sequence_model_parallel_all_gather(local.contiguous(), dim=dim)
+    if full.shape[dim] > orig_len:
+        full = full.narrow(dim, 0, orig_len)
+    return full
+
+
+def shard_seq_prefix(
+    x: torch.Tensor, prefix_len: int, shard: SpShard, dim: int = 0
+) -> torch.Tensor:
+    """Shard only the leading ``prefix_len`` rows (e.g. the text segment of a
+    joint RoPE cache) with an existing plan; the remainder is kept as-is."""
+    rest = x.shape[dim] - prefix_len
+    return torch.cat(
+        [
+            shard_like(x.narrow(dim, 0, prefix_len), shard, dim=dim),
+            x.narrow(dim, prefix_len, rest),
+        ],
+        dim=dim,
+    )
+
+
+def join_seqs(
+    prefix: torch.Tensor, body: torch.Tensor, local_pad: int, dim: int = 1
+) -> torch.Tensor:
+    """Concatenate local sharded ``[prefix (txt tokens, padding tokens), body (img tokens)]`` for joint attention, while relocating the
+     prefix's ``local_pad`` tail rows behind the body.
+
+    Why leave the padding at tail: the shard pads the *text* chunk, but the local joint layout is
+    [text, image].
+
+    In naive implementation, after the ulysses all-to-all, that pad would sit mid-sequence (of last rank)
+    ([... txt_last, PAD, img_last]), which required further mem copy (for the padding tokens), inefficient in this case
+
+    With the pad relocated behind the image, the padding forms one global-tail block that the zero-copy varlen
+    path (tail_attn_meta, implemented in USPAttention.forward) skips for free
+    """
+    if local_pad > 0:
+        real = prefix.shape[dim] - local_pad
+        return torch.cat(
+            [
+                # txt tokens
+                prefix.narrow(dim, 0, real),
+                body,
+                # leave the padding at global-tail
+                prefix.narrow(dim, real, local_pad),
+            ],
+            dim=dim,
+        )
+    return torch.cat([prefix, body], dim=dim)
+
+
+def split_seqs(
+    joint: torch.Tensor, prefix_len: int, local_pad: int, dim: int = 1
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Inverse of ``join_seqs``: recover ``(prefix, body)`` from the joint output, with the pad rows rejoining the prefix tail so the residual text
+    stream keeps its per-rank shape.
+
+     ([... txt_last, PAD, img_last]) -> prefix (txt + pad), body (img)
+    """
+    total = joint.shape[dim]
+    if local_pad > 0:
+        real = prefix_len - local_pad
+        body_end = total - local_pad
+        prefix = torch.cat(
+            [joint.narrow(dim, 0, real), joint.narrow(dim, body_end, local_pad)],
+            dim=dim,
+        )
+        return prefix, joint.narrow(dim, real, body_end - real)
+    return (
+        joint.narrow(dim, 0, prefix_len),
+        joint.narrow(dim, prefix_len, total - prefix_len),
+    )
+
+
+def should_shard_text(txt_len: int) -> bool:
+    """True when the joint-attention text stream should be SP-sharded here
+    (see plan_text_strategy for the policy)."""
+    return get_sp_world_size() > 1 and plan_text_strategy(txt_len) == "shard"
+
+
+def tail_attn_meta(
+    shard: SpShard,
+    batch_size: int,
+    device: torch.device,
+    image_seq_len: int = 0,
+) -> dict | None:
+    """Per-request attention meta for a tail-padded shard: `cu_seqlens_tail`
+    splits each batch row into [valid | pad] varlen segments over the gathered
+    layout, so USPAttention runs varlen FA on the padded q/k/v with zero
+    repacking. Built once per request, reused by every block."""
+    if shard.sp_size <= 1 or shard.num_pad == 0:
+        return None
+    seq = shard.sp_size * (shard.local_len + image_seq_len)
+    valid = seq - shard.num_pad
+    row = torch.tensor([valid, shard.num_pad], dtype=torch.int32, device=device)
+    seglens = row.repeat(batch_size)
+    cu_seqlens = torch.zeros(2 * batch_size + 1, dtype=torch.int32, device=device)
+    cu_seqlens[1:] = torch.cumsum(seglens, dim=0)
+    return {
+        "pad_start": valid,
+        "pad_end": seq,
+        "local_pad": shard.local_pad,
+        "cu_seqlens_tail": cu_seqlens,
+        "max_seqlen_tail": max(valid, shard.num_pad),
+    }
+
+
+def plan_text_strategy(txt_len: int) -> str:
+    """Choose "shard" or "replicate" for the joint-attention text stream.
+
+    Prefer "shard" by default. for small sequence (shorter than SGLANG_SP_TEXT_SHARD_MIN), choose "replicate" for better performance
+
+    """
+    sp_size = get_sp_world_size()
+    if sp_size <= 1:
+        return "replicate"
+    if txt_len % sp_size != 0 and get_ring_parallel_world_size() > 1:
+        return "replicate"
+    if txt_len < _TEXT_SHARD_MIN:
+        return "replicate"
+    return "shard"

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

@@ -0,0 +1,196 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py
+
+# Copyright 2023 The vLLM team.
+# Adapted from
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+import dataclasses
+import pickle
+import time
+from collections import deque
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+from torch.distributed import TCPStore
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def ensure_divisibility(numerator, denominator) -> None:
+    """Ensure that numerator is divisible by the denominator."""
+    assert numerator % denominator == 0, "{} is not divisible by {}".format(
+        numerator, denominator
+    )
+
+
+def divide(numerator: int, denominator: int) -> int:
+    """Ensure that numerator is divisible by the denominator and return
+    the division value."""
+    ensure_divisibility(numerator, denominator)
+    return numerator // denominator
+
+
+def split_tensor_along_last_dim(
+    tensor: torch.Tensor,
+    num_partitions: int,
+    contiguous_split_chunks: bool = False,
+) -> Sequence[torch.Tensor]:
+    """Split a tensor along its last dimension.
+
+    Arguments:
+        tensor: input tensor.
+        num_partitions: number of partitions to split the tensor
+        contiguous_split_chunks: If True, make each chunk contiguous
+                                 in memory.
+
+    Returns:
+        A list of Tensors
+    """
+    # Get the size and dimension.
+    last_dim = tensor.dim() - 1
+    last_dim_size = divide(tensor.size()[last_dim], num_partitions)
+    # Split.
+    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
+    # NOTE: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tuple(tensor_list)
+
+
+@dataclasses.dataclass
+class StatelessProcessGroup:
+    """A dataclass to hold a metadata store, and the rank, world_size of the
+    group. Only use it to communicate metadata between processes.
+    For data-plane communication, create NCCL-related objects.
+    """
+
+    rank: int
+    world_size: int
+    store: torch._C._distributed_c10d.Store
+    data_expiration_seconds: int = 3600  # 1 hour
+
+    # dst rank -> counter
+    send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict)
+    # src rank -> counter
+    recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)
+    broadcast_send_counter: int = 0
+    broadcast_recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)
+
+    # A deque to store the data entries, with key and timestamp.
+    entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque)
+
+    def __post_init__(self):
+        assert self.rank < self.world_size
+        self.send_dst_counter = {i: 0 for i in range(self.world_size)}
+        self.recv_src_counter = {i: 0 for i in range(self.world_size)}
+        self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)}
+
+    def send_obj(self, obj: Any, dst: int):
+        """Send an object to a destination rank."""
+        self.expire_data()
+        key = f"send_to/{dst}/{self.send_dst_counter[dst]}"
+        self.store.set(key, pickle.dumps(obj))
+        self.send_dst_counter[dst] += 1
+        self.entries.append((key, time.perf_counter()))
+
+    def expire_data(self) -> None:
+        """Expire data that is older than `data_expiration_seconds` seconds."""
+        while self.entries:
+            # check the oldest entry
+            key, timestamp = self.entries[0]
+            if time.perf_counter() - timestamp > self.data_expiration_seconds:
+                self.store.delete_key(key)
+                self.entries.popleft()
+            else:
+                break
+
+    def recv_obj(self, src: int) -> Any:
+        """Receive an object from a source rank."""
+        obj = pickle.loads(
+            self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}")
+        )
+        self.recv_src_counter[src] += 1
+        return obj
+
+    def broadcast_obj(self, obj: Any | None, src: int) -> Any:
+        """Broadcast an object from a source rank to all other ranks.
+        It does not clean up after all ranks have received the object.
+        Use it for limited times, e.g., for initialization.
+        """
+        if self.rank == src:
+            self.expire_data()
+            key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}"
+            self.store.set(key, pickle.dumps(obj))
+            self.broadcast_send_counter += 1
+            self.entries.append((key, time.perf_counter()))
+            return obj
+        else:
+            key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}"
+            recv_obj = pickle.loads(self.store.get(key))
+            self.broadcast_recv_src_counter[src] += 1
+            return recv_obj
+
+    def all_gather_obj(self, obj: Any) -> list[Any]:
+        """All gather an object from all ranks."""
+        gathered_objs = []
+        for i in range(self.world_size):
+            if i == self.rank:
+                gathered_objs.append(obj)
+                self.broadcast_obj(obj, src=self.rank)
+            else:
+                recv_obj = self.broadcast_obj(None, src=i)
+                gathered_objs.append(recv_obj)
+        return gathered_objs
+
+    def barrier(self):
+        """A barrier to synchronize all ranks."""
+        for i in range(self.world_size):
+            if i == self.rank:
+                self.broadcast_obj(None, src=self.rank)
+            else:
+                self.broadcast_obj(None, src=i)
+
+    @staticmethod
+    def create(
+        host: str,
+        port: int,
+        rank: int,
+        world_size: int,
+        data_expiration_seconds: int = 3600,
+    ) -> "StatelessProcessGroup":
+        """A replacement for `torch.distributed.init_process_group` that does not
+        pollute the global state.
+
+        If we have process A and process B called `torch.distributed.init_process_group`
+        to form a group, and then we want to form another group with process A, B, C,
+        D, it is not possible in PyTorch, because process A and process B have already
+        formed a group, and process C and process D cannot join that group. This
+        function is a workaround for this issue.
+
+        `torch.distributed.init_process_group` is a global call, while this function
+        is a stateless call. It will return a `StatelessProcessGroup` object that can be
+        used for exchanging metadata. With this function, process A and process B
+        can call `StatelessProcessGroup.create` to form a group, and then process A, B,
+        C, and D can call `StatelessProcessGroup.create` to form another group.
+        """  # noqa
+        store = TCPStore(
+            host_name=host,
+            port=port,
+            world_size=world_size,
+            is_master=(rank == 0),
+        )
+
+        return StatelessProcessGroup(
+            rank=rank,
+            world_size=world_size,
+            store=store,
+            data_expiration_seconds=data_expiration_seconds,
+        )

python/sglang/multimodal_gen/runtime/entrypoints/__init__.py

@@ -0,0 +1,4 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+from sglang.multimodal_gen.runtime.utils.logging_utils import globally_suppress_loggers
+
+globally_suppress_loggers()

python/sglang/multimodal_gen/runtime/entrypoints/cli/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/entrypoints/cli/cli_types.py

@@ -0,0 +1,30 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/types.py
+
+import argparse
+
+from sglang.multimodal_gen.utils import FlexibleArgumentParser
+
+
+class CLISubcommand:
+    """Base class for CLI subcommands"""
+
+    name: str
+
+    def cmd(
+        self, args: argparse.Namespace, unknown_args: list[str] | None = None
+    ) -> None:
+        """Execute the command with the given arguments"""
+        raise NotImplementedError
+
+    def validate(self, args: argparse.Namespace) -> None:
+        """Validate the arguments for this command"""
+        pass
+
+    def subparser_init(
+        self, subparsers: argparse._SubParsersAction
+    ) -> FlexibleArgumentParser:
+        """Initialize the subparser for this command"""
+        raise NotImplementedError

python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py

@@ -0,0 +1,248 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
+
+import argparse
+import dataclasses
+import json
+import os
+from typing import cast
+
+from sglang.multimodal_gen import DiffGenerator
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    SamplingParams,
+    generate_request_id,
+)
+from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand
+from sglang.multimodal_gen.runtime.entrypoints.cli.utils import (
+    RaiseNotImplementedAction,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import GenerationResult
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.perf_logger import (
+    MemorySnapshot,
+    PerformanceLogger,
+    RequestMetrics,
+)
+from sglang.multimodal_gen.utils import FlexibleArgumentParser
+
+logger = init_logger(__name__)
+
+
+def _resolve_cli_sampling_params_cls(server_args: ServerArgs) -> type[SamplingParams]:
+    pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
+    if pipeline_class_name:
+        from sglang.multimodal_gen.registry import get_pipeline_config_classes
+
+        config_classes = get_pipeline_config_classes(pipeline_class_name)
+        if config_classes is not None:
+            _, sampling_params_cls = config_classes
+            return sampling_params_cls
+
+    try:
+        from sglang.multimodal_gen.registry import get_model_info
+
+        model_info = get_model_info(
+            server_args.model_path,
+            backend=server_args.backend,
+            model_id=server_args.model_id,
+        )
+        if model_info is not None:
+            return model_info.sampling_param_cls
+    except Exception as exc:
+        logger.debug("Falling back to base SamplingParams for CLI parsing: %s", exc)
+
+    return SamplingParams
+
+
+def add_multimodal_gen_generate_args(parser: argparse.ArgumentParser):
+    """Add the arguments for the generate command."""
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="",
+        required=False,
+        help="Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional.",
+    )
+    parser.add_argument(
+        "--perf-dump-path",
+        type=str,
+        default=None,
+        required=False,
+        help="Path to dump the performance metrics (JSON) for the run.",
+    )
+    parser.add_argument(
+        "--output-file-path",
+        type=str,
+        default=None,
+        required=False,
+        help="Convenience alias that sets both --output-path and --output-file-name.",
+    )
+
+    parser = ServerArgs.add_cli_args(parser)
+    parser = SamplingParams.add_cli_args(parser)
+
+    parser.add_argument(
+        "--text-encoder-configs",
+        action=RaiseNotImplementedAction,
+        help="JSON array of text encoder configurations (NOT YET IMPLEMENTED)",
+    )
+
+    return parser
+
+
+def _apply_output_file_path_override(
+    args: argparse.Namespace, sampling_params_kwargs: dict
+):
+    output_file_path = args.output_file_path
+    if not output_file_path:
+        return
+
+    output_path = os.path.dirname(output_file_path) or "."
+    sampling_params_kwargs["output_path"] = output_path
+    sampling_params_kwargs["output_file_name"] = os.path.basename(output_file_path)
+
+
+def maybe_dump_performance(
+    args: argparse.Namespace,
+    server_args,
+    prompt: str,
+    results: GenerationResult | list[GenerationResult] | None,
+):
+    """dump performance if necessary"""
+    if not (args.perf_dump_path and results):
+        return
+
+    if isinstance(results, list):
+        result = results[0] if results else None
+    else:
+        result = results
+
+    metrics_dict = result.metrics
+    if not (args.perf_dump_path and metrics_dict):
+        return
+
+    metrics = RequestMetrics(request_id=metrics_dict.get("request_id"))
+    metrics.stages = metrics_dict.get("stages", {})
+    metrics.steps = metrics_dict.get("steps", [])
+    metrics.total_duration_ms = metrics_dict.get("total_duration_ms", 0)
+
+    # restore memory snapshots from serialized dict
+    memory_snapshots_dict = metrics_dict.get("memory_snapshots", {})
+    for checkpoint_name, snapshot_dict in memory_snapshots_dict.items():
+        snapshot = MemorySnapshot(
+            allocated_mb=snapshot_dict.get("allocated_mb", 0.0),
+            reserved_mb=snapshot_dict.get("reserved_mb", 0.0),
+            peak_allocated_mb=snapshot_dict.get("peak_allocated_mb", 0.0),
+            peak_reserved_mb=snapshot_dict.get("peak_reserved_mb", 0.0),
+        )
+        metrics.memory_snapshots[checkpoint_name] = snapshot
+
+    PerformanceLogger.dump_benchmark_report(
+        file_path=args.perf_dump_path,
+        metrics=metrics,
+        meta={
+            "prompt": prompt,
+            "model": server_args.model_path,
+        },
+        tag="cli_generate",
+    )
+
+
+def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None):
+    """The entry point for the generate command."""
+    args.request_id = "mocked_fake_id_for_offline_generate"
+
+    server_args = ServerArgs.from_cli_args(args, unknown_args)
+    sampling_params_cls = _resolve_cli_sampling_params_cls(server_args)
+
+    sampling_params_kwargs = {}
+    config_file = getattr(args, "config", None)
+    # respect config file by overriding args with args parsed from it
+    if config_file:
+        config_args = ServerArgs.load_config_file(config_file) or {}
+        sampling_param_fields = {
+            field.name for field in dataclasses.fields(sampling_params_cls)
+        }
+        sampling_params_kwargs.update(
+            {
+                key: value
+                for key, value in config_args.items()
+                if key in sampling_param_fields and value is not None
+            }
+        )
+
+    sampling_params_kwargs.update(sampling_params_cls.get_cli_args(args))
+    _apply_output_file_path_override(args, sampling_params_kwargs)
+    sampling_params_kwargs["request_id"] = generate_request_id()
+
+    # Handle diffusers-specific kwargs passed via CLI
+    if hasattr(args, "diffusers_kwargs") and args.diffusers_kwargs:
+        try:
+            sampling_params_kwargs["diffusers_kwargs"] = json.loads(
+                args.diffusers_kwargs
+            )
+            logger.info(
+                "Parsed diffusers_kwargs: %s",
+                sampling_params_kwargs["diffusers_kwargs"],
+            )
+        except json.JSONDecodeError as e:
+            logger.error("Failed to parse --diffusers-kwargs as JSON: %s", e)
+            raise ValueError(
+                f"--diffusers-kwargs must be valid JSON. Got: {args.diffusers_kwargs}"
+            ) from e
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=server_args.model_path, server_args=server_args, local_mode=True
+    )
+
+    results = generator.generate(sampling_params_kwargs=sampling_params_kwargs)
+
+    prompt = sampling_params_kwargs.get("prompt")
+    maybe_dump_performance(args, server_args, prompt, results)
+
+
+class GenerateSubcommand(CLISubcommand):
+    """The `generate` subcommand for the sglang-diffusion CLI"""
+
+    def __init__(self) -> None:
+        self.name = "generate"
+        super().__init__()
+        self.init_arg_names = self._get_init_arg_names()
+        self.generation_arg_names = self._get_generation_arg_names()
+
+    def _get_init_arg_names(self) -> list[str]:
+        """Get names of arguments for DiffGenerator initialization"""
+        return ["num_gpus", "tp_size", "sp_size", "model_path"]
+
+    def _get_generation_arg_names(self) -> list[str]:
+        """Get names of arguments for generate_video method"""
+        return [field.name for field in dataclasses.fields(SamplingParams)]
+
+    def cmd(
+        self, args: argparse.Namespace, unknown_args: list[str] | None = None
+    ) -> None:
+        generate_cmd(args, unknown_args)
+
+    def validate(self, args: argparse.Namespace) -> None:
+        """Validate the arguments for this command"""
+        if args.num_gpus is not None and args.num_gpus <= 0:
+            raise ValueError("Number of gpus must be positive")
+
+        if args.config and not os.path.exists(args.config):
+            raise ValueError(f"Config file not found: {args.config}")
+
+    def subparser_init(
+        self, subparsers: argparse._SubParsersAction
+    ) -> FlexibleArgumentParser:
+        generate_parser = subparsers.add_parser(
+            "generate",
+            help="Run inference on a model",
+            usage="sglang generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]",
+        )
+
+        generate_parser = add_multimodal_gen_generate_args(generate_parser)
+
+        return cast(FlexibleArgumentParser, generate_parser)

python/sglang/multimodal_gen/runtime/entrypoints/cli/main.py

@@ -0,0 +1,44 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py
+
+from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand
+from sglang.multimodal_gen.runtime.entrypoints.cli.generate import GenerateSubcommand
+from sglang.multimodal_gen.runtime.entrypoints.cli.serve import ServeSubcommand
+from sglang.multimodal_gen.utils import FlexibleArgumentParser
+
+
+def generate_cmd_init() -> list[CLISubcommand]:
+    return [GenerateSubcommand(), ServeSubcommand()]
+
+
+def cmd_init() -> list[CLISubcommand]:
+    """Initialize all commands from separate modules"""
+    commands = []
+    commands.extend(generate_cmd_init())
+    return commands
+
+
+def main() -> None:
+    parser = FlexibleArgumentParser(description="sglang-diffusion CLI")
+    parser.add_argument("-v", "--version", action="version", version="0.1.0")
+
+    subparsers = parser.add_subparsers(required=False, dest="subparser")
+
+    cmds = {}
+    for cmd in cmd_init():
+        cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
+        cmds[cmd.name] = cmd
+    args, unknown_args = parser.parse_known_args()
+    if args.subparser in cmds:
+        cmds[args.subparser].validate(args)
+
+    if hasattr(args, "dispatch_function"):
+        args.dispatch_function(args, unknown_args=unknown_args)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py

@@ -0,0 +1,75 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+from typing import cast
+
+from sglang.multimodal_gen.apps.webui import run_sgl_diffusion_webui
+from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand
+from sglang.multimodal_gen.runtime.launch_server import (
+    dispatch_launch,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.utils import FlexibleArgumentParser
+
+
+def add_multimodal_gen_serve_args(parser: argparse.ArgumentParser):
+    """Add the arguments for the serve command."""
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="",
+        required=False,
+        help="Read CLI options from a config JSON or YAML file.",
+    )
+    return ServerArgs.add_cli_args(parser)
+
+
+def execute_serve_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None):
+    """The entry point for the serve command."""
+    # use server-based warmup for production
+    server_args = ServerArgs.from_cli_args(
+        args, unknown_args, default_args={"warmup_mode": "server"}
+    )
+
+    dispatch_launch(server_args)
+
+    if server_args.webui:
+        run_sgl_diffusion_webui(server_args)
+
+
+class ServeSubcommand(CLISubcommand):
+    """The `serve` subcommand for the sglang-diffusion CLI"""
+
+    def __init__(self) -> None:
+        self.name = "serve"
+        super().__init__()
+
+    def cmd(
+        self, args: argparse.Namespace, unknown_args: list[str] | None = None
+    ) -> None:
+        execute_serve_cmd(args, unknown_args)
+
+    def validate(self, args: argparse.Namespace) -> None:
+        """Validate the arguments for this command"""
+        if args.config and not os.path.exists(args.config):
+            raise ValueError(f"Config file not found: {args.config}")
+
+    def subparser_init(
+        self, subparsers: argparse._SubParsersAction
+    ) -> FlexibleArgumentParser:
+        serve_parser = subparsers.add_parser(
+            "serve",
+            help="Launch the server and start FastAPI listener.",
+            usage="sglang serve --model-path MODEL_PATH_OR_ID [OPTIONS]",
+        )
+
+        serve_parser = add_multimodal_gen_serve_args(serve_parser)
+
+        return cast(FlexibleArgumentParser, serve_parser)
+
+
+def cmd_init() -> list[CLISubcommand]:
+    return [ServeSubcommand()]

python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py

@@ -0,0 +1,75 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+import shlex
+import subprocess
+import sys
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class RaiseNotImplementedAction(argparse.Action):
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        raise NotImplementedError(f"The {option_string} option is not yet implemented")
+
+
+def launch_distributed(
+    num_gpus: int, args: list[str], master_port: int | None = None
+) -> int:
+    """
+    Launch a distributed job with the given arguments
+
+    Args:
+        num_gpus: Number of GPUs to use
+        args: Arguments to pass to v1_sgl_diffusion_inference.py (defaults to sys.argv[1:])
+        master_port: Port for the master process (default: random)
+    """
+
+    current_env = os.environ.copy()
+    python_executable = sys.executable
+    project_root = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "../../../..")
+    )
+    main_script = os.path.join(
+        project_root, "sgl_diffusion/sample/v1_sgl_diffusion_inference.py"
+    )
+
+    cmd = [
+        python_executable,
+        "-m",
+        "torch.distributed.run",
+        f"--nproc_per_node={num_gpus}",
+    ]
+
+    if master_port is not None:
+        cmd.append(f"--master_port={master_port}")
+
+    cmd.append(main_script)
+    cmd.extend(args)
+
+    logger.info("Running inference with %d GPU(s)", num_gpus)
+    logger.info("Launching command: %s", shlex.join(cmd))
+
+    current_env["PYTHONIOENCODING"] = "utf-8"
+    process = subprocess.Popen(
+        cmd,
+        env=current_env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+        bufsize=1,
+        encoding="utf-8",
+        errors="replace",
+    )
+
+    if process.stdout:
+        for line in iter(process.stdout.readline, ""):
+            print(line.strip())
+
+    return process.wait()

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -0,0 +1,713 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+"""
+DiffGenerator module for sglang-diffusion.
+
+This module provides a consolidated interface for generating images/videos using
+diffusion models.
+"""
+
+import dataclasses
+import multiprocessing as mp
+import os
+import time
+from contextlib import ExitStack
+from typing import Any, List, Union
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    GenerationResult,
+    ListLorasReq,
+    MergeLoraWeightsReq,
+    SetLoraReq,
+    ShutdownReq,
+    UnmergeLoraWeightsReq,
+    expand_request_outputs,
+    format_lora_message,
+    prepare_request,
+    save_outputs,
+)
+from sglang.multimodal_gen.runtime.launch_server import launch_server
+from sglang.multimodal_gen.runtime.pipelines_core import Req
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
+from sglang.multimodal_gen.runtime.server_warmup import (
+    run_sync_client_warmup,
+    should_run_explicit_client_warmup,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    GREEN,
+    RESET,
+    init_logger,
+    log_batch_completion,
+    log_generation_timer,
+)
+from sglang.multimodal_gen.runtime.utils.trace_wrapper import (
+    init_diffusion_tracing,
+    trace_req,
+)
+
+logger = init_logger(__name__)
+
+try:
+    # Set the start method to 'spawn' to avoid CUDA errors in forked processes.
+    # This must be done at the top level of the module, before any CUDA context
+    # or other processes are initialized.
+    mp.set_start_method("spawn", force=True)
+except RuntimeError:
+    # The start method can only be set once per program execution.
+    pass
+
+
+class DiffGenerator:
+    """
+    A unified class for generating images/videos using diffusion models.
+
+    This class provides a simple interface for image/video generation with rich
+    customization options, similar to popular frameworks like HF Diffusers.
+    """
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+    ):
+        """
+        Initialize the generator.
+
+        Args:
+            server_args: The inference arguments
+        """
+        self.server_args = server_args
+        self.port_args = PortArgs.from_server_args(server_args)
+
+        # The executor is now a client to the Scheduler service
+        self.local_scheduler_process: list[mp.Process] | None = None
+        self.owns_scheduler_client: bool = False
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        local_mode: bool = True,
+        **kwargs,
+    ) -> "DiffGenerator":
+        """
+        Create a DiffGenerator from a pretrained model.
+
+        Priority level: Default pipeline config < User's pipeline config < User's kwargs
+        """
+        # If users also provide some kwargs, it will override the ServerArgs and PipelineConfig.
+
+        if (server_args := kwargs.get("server_args", None)) is not None:
+            if isinstance(server_args, ServerArgs):
+                pass
+            elif isinstance(server_args, dict):
+                server_args = ServerArgs.from_kwargs(**server_args)
+        else:
+            server_args = ServerArgs.from_kwargs(**kwargs)
+
+        return cls.from_server_args(server_args, local_mode=local_mode)
+
+    @classmethod
+    def from_server_args(
+        cls, server_args: ServerArgs, local_mode: bool = True
+    ) -> "DiffGenerator":
+        """
+        Create a DiffGenerator with the specified arguments.
+
+        Args:
+            server_args: The inference arguments
+
+        Returns:
+            The created DiffGenerator
+        """
+        instance = cls(
+            server_args=server_args,
+        )
+        init_diffusion_tracing(server_args, "DiffGenerator")
+
+        logger.info(f"Local mode: {local_mode}")
+        if local_mode:
+            instance.local_scheduler_process = instance._start_local_server_if_needed()
+            instance.owns_scheduler_client = True
+            instance._run_client_warmup_if_needed()
+        else:
+            # In remote mode, we just need to connect and check.
+            sync_scheduler_client.initialize(server_args)
+            instance._check_remote_scheduler()
+            instance.owns_scheduler_client = True
+        return instance
+
+    def _start_local_server_if_needed(
+        self,
+    ) -> list[mp.Process]:
+        """Check if a local server is running; if not, start it and return the process handles."""
+        # First, we need a client to test the server. Initialize it temporarily.
+        sync_scheduler_client.initialize(self.server_args)
+
+        processes = launch_server(self.server_args, launch_http_server=False)
+
+        return processes
+
+    def _run_client_warmup_if_needed(self) -> None:
+        if not should_run_explicit_client_warmup(self.server_args):
+            return
+
+        run_sync_client_warmup(self.server_args, sync_scheduler_client.forward)
+
+    def _check_remote_scheduler(self):
+        """Check if the remote scheduler is accessible."""
+        if not sync_scheduler_client.ping():
+            raise ConnectionError(
+                f"Could not connect to remote scheduler at "
+                f"{self.server_args.scheduler_endpoint} with `local mode` as False. "
+                "Please ensure the server is running."
+            )
+        logger.info(
+            f"Successfully connected to remote scheduler at "
+            f"{self.server_args.scheduler_endpoint}."
+        )
+
+    @staticmethod
+    def _resolve_image_paths_per_prompt(
+        prompts: list[str], image_paths: str | list[str] | None
+    ) -> list[str | list[str] | None]:
+        if len(prompts) <= 1:
+            return [image_paths]
+
+        if not isinstance(image_paths, list) or len(image_paths) <= 1:
+            return [image_paths for _ in prompts]
+
+        if len(image_paths) != len(prompts):
+            raise ValueError(
+                "When using multiple prompts with multiple input images, "
+                "provide either one shared image or exactly one image per prompt."
+            )
+
+        return [[image_path] for image_path in image_paths]
+
+    def generate(
+        self,
+        sampling_params_kwargs: dict | None = None,
+        external_trace_header: dict[str, str] | None = None,
+    ) -> GenerationResult | list[GenerationResult] | None:
+        """Generate image(s)/video(s) based on the given prompt(s).
+
+        Returns a single GenerationResult for a single prompt, a list for
+        multiple prompts, or None when every request failed.
+        """
+        # 1. prepare requests
+        prompts = self._resolve_prompts(
+            sampling_params_kwargs.get("prompt"),
+            sampling_params_kwargs.get("prompt_path"),
+        )
+        user_output_file_name = sampling_params_kwargs.get("output_file_name")
+
+        if len(prompts) > 1 and user_output_file_name is not None:
+            raise ValueError(
+                "Cannot use multiple prompts with a fixed output_file_name. "
+                "Either remove --output-file-name or use a single prompt."
+            )
+
+        sampling_params_orig = SamplingParams.from_user_sampling_params_args(
+            self.server_args.model_path,
+            server_args=self.server_args,
+            **sampling_params_kwargs,
+        )
+
+        request_groups: list[list[Req]] = []
+        image_paths_per_prompt = self._resolve_image_paths_per_prompt(
+            prompts, sampling_params_orig.image_path
+        )
+
+        for i, p in enumerate(prompts):
+            sampling_params = dataclasses.replace(
+                sampling_params_orig,
+                prompt=p,
+                output_file_name=user_output_file_name,
+                image_path=image_paths_per_prompt[i],
+            )
+            # `dataclasses.replace` drops non-field attrs; restore
+            # `_explicit_fields` so InputValidationStage honors user-supplied
+            # width/height, and mark the keys overridden above as explicit.
+            sampling_params._explicit_fields = getattr(
+                sampling_params_orig, "_explicit_fields", set()
+            ) | {"prompt", "output_file_name", "image_path"}
+            sampling_params._set_output_file_name()
+            req = prepare_request(
+                server_args=self.server_args,
+                sampling_params=sampling_params,
+                external_trace_header=external_trace_header,
+            )
+            request_groups.append(
+                expand_request_outputs(
+                    req,
+                    num_prompts=len(prompts),
+                    prompt_index=i,
+                )
+            )
+
+        results: list[GenerationResult] = []
+        total_start_time = time.perf_counter()
+        global_output_index = 0
+
+        for requests in request_groups:
+            try:
+                timer_prompt = [req.prompt for req in requests]
+                logger.info("Processing %d grouped request(s)", len(requests))
+                with ExitStack() as stack:
+                    for req in requests:
+                        stack.enter_context(trace_req(req.trace_ctx))
+                    timer = stack.enter_context(
+                        log_generation_timer(logger, timer_prompt)
+                    )
+                    output_batch = self._send_to_scheduler_and_wait_for_response(
+                        requests
+                    )
+                    if output_batch.error:
+                        raise Exception(f"{output_batch.error}")
+
+                    if (
+                        output_batch.output is None
+                        and output_batch.output_file_paths is None
+                    ):
+                        logger.error("Received empty output from scheduler")
+                        continue
+
+                    if requests[0].save_output and requests[0].return_file_paths_only:
+                        output_file_paths = output_batch.output_file_paths or []
+                        self._validate_output_count(
+                            len(output_file_paths), len(requests)
+                        )
+                        for idx, path in enumerate(output_file_paths):
+                            req = requests[idx]
+                            results.append(
+                                GenerationResult(
+                                    **self._result_common(
+                                        req, output_batch, timer.duration, idx
+                                    ),
+                                    prompt_index=global_output_index + idx,
+                                    output_file_path=path,
+                                )
+                            )
+                    elif requests[0].data_type == DataType.MESH:
+                        output_file_paths = output_batch.output_file_paths or []
+                        self._validate_output_count(
+                            len(output_file_paths), len(requests)
+                        )
+                        for idx, sample in enumerate(output_file_paths):
+                            req = requests[idx]
+                            results.append(
+                                GenerationResult(
+                                    **self._result_common(
+                                        req, output_batch, timer.duration, idx
+                                    ),
+                                    prompt_index=global_output_index + idx,
+                                    output_file_path=sample,
+                                )
+                            )
+                    else:
+                        self._validate_output_count(
+                            len(output_batch.output), len(requests)
+                        )
+                        samples_out: list[Any] = []
+                        audios_out: list[Any] = []
+                        frames_out: list[Any] = []
+                        save_outputs(
+                            output_batch.output,
+                            requests[0].data_type,
+                            requests[0].fps,
+                            requests[0].save_output,
+                            lambda idx: requests[idx].output_file_path(1, 0),
+                            audio=output_batch.audio,
+                            audio_sample_rate=output_batch.audio_sample_rate,
+                            samples_out=samples_out,
+                            audios_out=audios_out,
+                            frames_out=frames_out,
+                            output_compression=requests[0].output_compression,
+                            enable_frame_interpolation=requests[
+                                0
+                            ].enable_frame_interpolation,
+                            frame_interpolation_exp=requests[0].frame_interpolation_exp,
+                            frame_interpolation_scale=requests[
+                                0
+                            ].frame_interpolation_scale,
+                            frame_interpolation_model_path=requests[
+                                0
+                            ].frame_interpolation_model_path,
+                            enable_upscaling=requests[0].enable_upscaling,
+                            upscaling_model_path=requests[0].upscaling_model_path,
+                            upscaling_scale=requests[0].upscaling_scale,
+                        )
+
+                        for idx in range(len(samples_out)):
+                            req = requests[idx]
+                            results.append(
+                                GenerationResult(
+                                    **self._result_common(
+                                        req, output_batch, timer.duration, idx
+                                    ),
+                                    samples=samples_out[idx],
+                                    frames=frames_out[idx],
+                                    audio=audios_out[idx],
+                                    prompt_index=global_output_index + idx,
+                                    output_file_path=req.output_file_path(1, 0),
+                                )
+                            )
+            except Exception as e:
+                logger.error("Generation failed: %s", e, exc_info=True)
+            finally:
+                global_output_index += len(requests)
+
+        total_gen_time = time.perf_counter() - total_start_time
+        if self.server_args.batching_max_size > 1:
+            log_batch_completion(
+                logger,
+                len(results),
+                total_gen_time,
+            )
+        self._log_summary(results)
+
+        if not results:
+            return None
+        return results[0] if len(results) == 1 else results
+
+    def generate_action(
+        self,
+        sampling_params_kwargs: dict | None = None,
+        external_trace_header: dict[str, str] | None = None,
+    ) -> dict[str, Any]:
+        sampling_params_kwargs = sampling_params_kwargs or {}
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.server_args.model_path,
+            server_args=self.server_args,
+            **sampling_params_kwargs,
+        )
+        if sampling_params.data_type != DataType.ACTION:
+            raise ValueError(
+                f"generate_action requires an ACTION pipeline, got {sampling_params.data_type}"
+            )
+
+        req = prepare_request(
+            server_args=self.server_args,
+            sampling_params=sampling_params,
+            external_trace_header=external_trace_header,
+        )
+        output_batch = self._send_to_scheduler_and_wait_for_response(req)
+        if output_batch.error:
+            raise RuntimeError(output_batch.error)
+        if output_batch.output is None:
+            raise RuntimeError("action policy returned no output")
+        return output_batch.output[0]
+
+    def _resolve_prompts(
+        self,
+        prompt: str | list[str] | None,
+        prompt_path: str | None = None,
+    ) -> list[str]:
+        """Collect prompts from the argument or from a prompt file."""
+        path = prompt_path or self.server_args.prompt_file_path
+        if path is not None:
+            if not os.path.exists(path):
+                raise FileNotFoundError(f"Prompt text file not found: {path}")
+            with open(path, encoding="utf-8") as f:
+                prompts = [line.strip() for line in f if line.strip()]
+            if not prompts:
+                raise ValueError(f"No prompts found in file: {path}")
+            logger.info("Found %d prompts in %s", len(prompts), path)
+            return prompts
+
+        if prompt is None:
+            return [" "]
+        if isinstance(prompt, str):
+            return [prompt]
+        return list(prompt)
+
+    def _log_summary(self, results: list[GenerationResult]) -> None:
+        if not results:
+            return
+        if self.server_args.warmup:
+            total_duration_ms = results[0].metrics.get("total_duration_ms", 0)
+            logger.info(
+                f"Warmed-up request processed in {GREEN}%.2f{RESET} seconds (with warmup excluded)",
+                total_duration_ms / 1000.0,
+            )
+
+        peak_memories = [r.peak_memory_mb for r in results if r.peak_memory_mb]
+        if peak_memories:
+            logger.info(
+                f"Memory usage - Max peak: {max(peak_memories):.2f} MB, "
+                f"Avg peak: {sum(peak_memories) / len(peak_memories):.2f} MB"
+            )
+
+    @staticmethod
+    def _result_common(
+        req: Req,
+        output_batch: OutputBatch,
+        generation_time: float,
+        output_index: int | None = None,
+    ) -> dict[str, Any]:
+        metrics = output_batch.metrics
+        if (
+            output_index is not None
+            and output_batch.metrics_list is not None
+            and output_index < len(output_batch.metrics_list)
+        ):
+            metrics = output_batch.metrics_list[output_index]
+        if req.data_type == DataType.ACTION:
+            size = ("action",)
+        else:
+            size = (req.height, req.width, req.num_frames)
+        return dict(
+            prompt=req.prompt,
+            size=size,
+            generation_time=generation_time,
+            peak_memory_mb=output_batch.peak_memory_mb,
+            metrics=metrics.to_dict() if metrics else {},
+            action=output_batch.action_pred,
+            trajectory_latents=output_batch.trajectory_latents,
+            trajectory_timesteps=output_batch.trajectory_timesteps,
+            rollout_trajectory_data=output_batch.rollout_trajectory_data,
+            trajectory_decoded=output_batch.trajectory_decoded,
+        )
+
+    @staticmethod
+    def _validate_output_count(output_count: int, request_count: int) -> None:
+        if output_count != request_count:
+            raise RuntimeError(
+                f"Expected {request_count} outputs, got {output_count} from scheduler"
+            )
+
+    def _send_to_scheduler_and_wait_for_response(self, batch: list[Req]) -> OutputBatch:
+        """
+        Sends a request to the scheduler and waits for a response.
+        """
+        return sync_scheduler_client.forward(batch)
+
+    # LoRA
+    def _send_lora_request(self, req: Any, success_msg: str, failure_msg: str):
+        response = sync_scheduler_client.forward(req)
+        if response.error is None:
+            logger.info(success_msg)
+            return response
+        else:
+            error_msg = response.error
+            raise RuntimeError(f"{failure_msg}: {error_msg}")
+
+    def set_lora(
+        self,
+        lora_nickname: Union[str, List[str]],
+        lora_path: Union[str, None, List[Union[str, None]]] = None,
+        target: Union[str, List[str]] = "all",
+        strength: Union[float, List[float]] = 1.0,
+        merge_mode: str | None = None,
+    ) -> None:
+        """
+        Set LoRA adapter(s) for the specified transformer(s).
+        Supports both single LoRA (backward compatible) and multiple LoRA adapters.
+
+        Args:
+            lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings.
+            lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None.
+            target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
+                Valid values:
+                - "all": Apply to all transformers (default)
+                - "transformer": Apply only to the primary transformer (high noise for Wan2.2)
+                - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
+                - "critic": Apply only to the critic model
+            strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
+            merge_mode: Optional LoRA merge mode: "auto", "merge", or "dynamic".
+        """
+        req = SetLoraReq(
+            lora_nickname=lora_nickname,
+            lora_path=lora_path,
+            target=target,
+            strength=strength,
+            merge_mode=merge_mode,
+        )
+        nickname_str, target_str, strength_str = format_lora_message(
+            lora_nickname, target, strength
+        )
+
+        self._send_lora_request(
+            req,
+            f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})",
+            "Failed to set LoRA adapter",
+        )
+
+    def unmerge_lora_weights(self, target: str = "all") -> None:
+        """
+        Unmerge LoRA weights from the base model.
+
+        Args:
+            target: Which transformer(s) to unmerge.
+        """
+        req = UnmergeLoraWeightsReq(target=target)
+        self._send_lora_request(
+            req,
+            f"Successfully unmerged LoRA weights (target: {target})",
+            "Failed to unmerge LoRA weights",
+        )
+
+    def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None:
+        """
+        Merge LoRA weights into the base model.
+
+        Args:
+            target: Which transformer(s) to merge.
+            strength: LoRA strength for merge, default 1.0.
+        """
+        req = MergeLoraWeightsReq(target=target, strength=strength)
+        self._send_lora_request(
+            req,
+            f"Successfully merged LoRA weights (target: {target}, strength: {strength})",
+            "Failed to merge LoRA weights",
+        )
+
+    def list_loras(self) -> dict:
+        """List loaded LoRA adapters and current application status per module."""
+        output = self._send_lora_request(
+            req=ListLorasReq(),
+            success_msg="Successfully listed LoRA adapters",
+            failure_msg="Failed to list LoRA adapters",
+        )
+        # _send_lora_request already raises on error, so output.error is always None here
+        return output.output or {}
+
+    def _ensure_lora_state(
+        self,
+        lora_path: str | None,
+        lora_nickname: str | None = None,
+        merge_lora: bool = True,
+    ) -> None:
+        """
+        Ensure the LoRA state matches the desired configuration.
+
+        Note: This method does not cache client-side state. The server handles
+        idempotent operations, so redundant calls are safe but may have minor overhead.
+        """
+        if lora_path is None:
+            # Unmerge all LoRA weights when no lora_path is provided
+            self.unmerge_lora_weights()
+            return
+
+        lora_nickname = lora_nickname or self.server_args.lora_nickname
+
+        # Set the LoRA adapter (server handles idempotent logic)
+        self.set_lora(lora_nickname, lora_path)
+
+        # Merge or unmerge based on the merge_lora flag
+        if merge_lora:
+            self.merge_lora_weights()
+        else:
+            self.unmerge_lora_weights()
+
+    def generate_with_lora(
+        self,
+        prompt: str | list[str] | None = None,
+        sampling_params: SamplingParams | None = None,
+        *,
+        lora_path: str | None = None,
+        lora_nickname: str | None = None,
+        merge_lora: bool = True,
+        **kwargs,
+    ):
+        self._ensure_lora_state(
+            lora_path=lora_path, lora_nickname=lora_nickname, merge_lora=merge_lora
+        )
+        return self.generate(
+            sampling_params_kwargs=dict(
+                prompt=prompt,
+                sampling_params=sampling_params,
+                **kwargs,
+            )
+        )
+
+    def shutdown(self):
+        """
+        Shutdown the generator.
+        If in local mode, it also shuts down the scheduler server.
+        """
+        # sends the shutdown command to the server
+        if self.local_scheduler_process and self.owns_scheduler_client:
+            try:
+                sync_scheduler_client.forward(ShutdownReq(), timeout_ms=5000)
+            except Exception:
+                pass
+
+        if self.local_scheduler_process:
+            for process in self.local_scheduler_process:
+                process.join(timeout=10)
+                if process.is_alive():
+                    logger.warning(
+                        f"Local worker {process.name} did not terminate gracefully, forcing."
+                    )
+                    process.terminate()
+                    process.join(timeout=1)
+                    if process.is_alive():
+                        process.kill()
+                        process.join(timeout=1)
+            self.local_scheduler_process = None
+
+        if self.owns_scheduler_client:
+            sync_scheduler_client.close()
+            self.owns_scheduler_client = False
+
+    def _force_shutdown_local_processes(self) -> None:
+        local_scheduler_process = getattr(self, "local_scheduler_process", None)
+        log = globals().get("logger")
+        if local_scheduler_process:
+            for process in local_scheduler_process:
+                if process.is_alive():
+                    if log is not None:
+                        log.warning(
+                            f"Local worker {process.name} did not terminate gracefully, forcing."
+                        )
+                    process.terminate()
+            for process in local_scheduler_process:
+                process.join(timeout=1)
+                if process.is_alive():
+                    if log is not None:
+                        log.warning(
+                            f"Local worker {process.name} did not terminate after terminate(), killing."
+                        )
+                    process.kill()
+                    process.join(timeout=1)
+            self.local_scheduler_process = None
+
+        if getattr(self, "owns_scheduler_client", False):
+            try:
+                client = globals().get("sync_scheduler_client")
+                if client is not None:
+                    client.close()
+            finally:
+                self.owns_scheduler_client = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.shutdown()
+
+    def __del__(self):
+        owns_scheduler_client = bool(getattr(self, "owns_scheduler_client", False))
+        local_scheduler_process = getattr(self, "local_scheduler_process", None)
+        log = globals().get("logger")
+        if owns_scheduler_client:
+            if log is not None:
+                log.warning(
+                    "Generator was garbage collected without being shut down. "
+                    "Forcing local server and client cleanup."
+                )
+            self._force_shutdown_local_processes()
+        elif local_scheduler_process:
+            if log is not None:
+                log.warning(
+                    "Generator was garbage collected without being shut down. "
+                    "Forcing local server cleanup."
+                )
+            self._force_shutdown_local_processes()

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -0,0 +1,413 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import asyncio
+import base64
+import os
+import signal
+import uuid
+from contextlib import asynccontextmanager, suppress
+from typing import TYPE_CHECKING
+
+import httpx
+import torch
+from fastapi import APIRouter, FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    VertexGenerateReqInput,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime import (
+    realtime_video_api,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params
+from sglang.multimodal_gen.runtime.entrypoints.post_training import (
+    rollout_api,
+    weights_api,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    prepare_request,
+    save_outputs,
+)
+from sglang.multimodal_gen.runtime.entrypoints.vla import api as vla_api
+from sglang.multimodal_gen.runtime.entrypoints.vla import openpi
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
+from sglang.multimodal_gen.runtime.server_warmup import (
+    run_async_client_warmup,
+    should_run_synthetic_server_warmup,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils.json_response import orjson_response
+from sglang.version import __version__
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+
+logger = init_logger(__name__)
+
+VERTEX_ROUTE = os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate")
+SERVER_WARMUP_BYPASS_PATHS = (
+    "/health",
+    "/health_generate",
+    "/model_info",
+    "/server_info",
+)
+
+
+async def _wait_until_http_ready(server_args: ServerArgs) -> None:
+    """for server warmup"""
+    health_url = f"{server_args.url()}/health"
+    # Probe the local server directly: a loopback readiness check must never be
+    # routed through an HTTP proxy. trust_env=False also avoids crashing startup
+    # on a malformed proxy env var, since httpx parses *_PROXY/NO_PROXY when the
+    # client is constructed (raising httpx.InvalidURL before any request). See #28493.
+    async with httpx.AsyncClient(trust_env=False) as client:
+        for _ in range(120):
+            try:
+                response = await client.get(health_url, timeout=5.0)
+                if response.status_code == 200:
+                    return
+            except httpx.HTTPError:
+                pass
+            await asyncio.sleep(1.0)
+    raise RuntimeError(f"HTTP server did not become ready at {health_url}")
+
+
+async def _run_server_warmup_after_http_ready(
+    server_args: ServerArgs, warmup_done: asyncio.Event
+) -> None:
+    try:
+        if not should_run_synthetic_server_warmup(server_args):
+            warmup_done.set()
+            return
+
+        await _wait_until_http_ready(server_args)
+
+        await run_async_client_warmup(
+            server_args,
+            async_scheduler_client.forward,
+            fail_open=server_args.warmup_resolutions is None,
+        )
+        logger.info("The server is fired up and ready to roll!")
+        warmup_done.set()
+    except asyncio.CancelledError:
+        raise
+    except Exception as e:
+        logger.error("Server warmup failed; aborting startup: %s", e, exc_info=True)
+        os.kill(os.getpid(), signal.SIGTERM)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    from sglang.multimodal_gen.runtime.scheduler_client import (
+        async_scheduler_client,
+        run_zeromq_broker,
+    )
+
+    # 1. Initialize the singleton client that connects to the backend Scheduler
+    server_args = app.state.server_args
+    async_scheduler_client.initialize(server_args)
+    warmup_done = asyncio.Event()
+    app.state.server_warmup_done = warmup_done
+
+    # 2. Start the ZMQ Broker in the background to handle offline requests
+    broker_task = asyncio.create_task(run_zeromq_broker(server_args))
+    warmup_task = None
+    if server_args.server_warmup:
+        warmup_task = asyncio.create_task(
+            _run_server_warmup_after_http_ready(server_args, warmup_done)
+        )
+    else:
+        warmup_done.set()
+
+    try:
+        yield
+    finally:
+        if warmup_task is not None and not warmup_task.done():
+            warmup_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await warmup_task
+
+        # On shutdown
+        logger.info("FastAPI app is shutting down...")
+        broker_task.cancel()
+        async_scheduler_client.close()
+
+
+# Health router
+health_router = APIRouter()
+
+
+@health_router.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+@health_router.get("/models", deprecated=True)
+async def get_models(request: Request):
+    """
+    Get information about the model served by this server.
+
+    .. deprecated::
+        Use /v1/models instead for OpenAI-compatible model discovery.
+        This endpoint will be removed in a future version.
+    """
+    from sglang.multimodal_gen.registry import get_model_info
+
+    server_args: ServerArgs = request.app.state.server_args
+    model_info = get_model_info(server_args.model_path, model_id=server_args.model_id)
+
+    response = {
+        "model_path": server_args.model_path,
+        "num_gpus": server_args.num_gpus,
+        "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,
+    }
+
+    if model_info:
+        response["pipeline_name"] = model_info.pipeline_cls.pipeline_name
+        response["pipeline_class"] = model_info.pipeline_cls.__name__
+
+    return response
+
+
+@health_router.get("/server_info")
+async def server_info_endpoint(request: Request):
+    """Get server information.
+
+    Returns fields compatible with the LLM engine's /server_info so that
+    the model gateway can discover diffusion workers.
+    """
+    server_args: ServerArgs = request.app.state.server_args
+
+    return {
+        "model_path": server_args.model_path,
+        "served_model_name": server_args.model_id or server_args.model_path,
+        "tp_size": server_args.tp_size,
+        "dp_size": server_args.dp_size,
+        "version": __version__,
+    }
+
+
+@health_router.get("/model_info")
+async def model_info_endpoint(request: Request):
+    """Get model information.
+
+    Returns fields compatible with the LLM engine's /model_info so that
+    the model gateway can detect capabilities for diffusion workers.
+    """
+    from sglang.multimodal_gen.registry import get_model_info
+
+    server_args: ServerArgs = request.app.state.server_args
+    task_type = server_args.pipeline_config.task_type
+
+    try:
+        registry_info = get_model_info(
+            server_args.model_path,
+            backend=server_args.backend,
+            model_id=server_args.model_id,
+        )
+    except Exception:
+        logger.warning("Failed to resolve model info from registry", exc_info=True)
+        registry_info = None
+
+    return {
+        # Fields consumed by the model gateway for worker discovery
+        "model_path": server_args.model_path,
+        "is_generation": True,
+        "model_type": "diffusion",
+        "architectures": (
+            [registry_info.pipeline_cls.__name__] if registry_info else None
+        ),
+        # Fields matching the LLM engine's /model_info shape
+        "has_image_understanding": task_type.accepts_image_input(),
+        "has_audio_understanding": False,
+        # Diffusion-specific fields
+        "task_type": task_type.name,
+        "is_image_gen": task_type.is_image_gen(),
+    }
+
+
+@health_router.get("/health_generate")
+async def health_generate():
+    # TODO : health generate endpoint
+    return {"status": "ok"}
+
+
+@health_router.get("/stats")
+async def stats_endpoint(request: Request):
+    """Get runtime statistics including disagg pipeline metrics.
+
+    Returns queue depth, request counts, latency, throughput, etc.
+    Sends a GetDisaggStatsReq to the scheduler via ZMQ and returns the result.
+    """
+    from sglang.multimodal_gen.runtime.entrypoints.utils import GetDisaggStatsReq
+
+    server_args: ServerArgs = request.app.state.server_args
+    response: dict = {
+        "status": "ok",
+        "model_path": server_args.model_path,
+    }
+
+    # Query the scheduler for disagg metrics
+    try:
+        stats_response = await async_scheduler_client.forward(GetDisaggStatsReq())
+        if hasattr(stats_response, "output") and stats_response.output is not None:
+            response["disagg"] = stats_response.output
+    except Exception as e:
+        response["disagg"] = {"error": str(e)}
+
+    return response
+
+
+def make_serializable(obj):
+    """Recursively converts Tensors to None for JSON serialization."""
+    if isinstance(obj, torch.Tensor):
+        return None
+    if isinstance(obj, dict):
+        return {k: make_serializable(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [make_serializable(v) for v in obj]
+    return obj
+
+
+def encode_video_to_base64(file_path: str):
+    if not os.path.exists(file_path):
+        return None
+    with open(file_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+async def forward_to_scheduler(
+    req_obj: "Req",
+    sp: SamplingParams,
+):
+    """Forwards request to scheduler and processes the result."""
+    try:
+        response = await async_scheduler_client.forward(req_obj)
+        if response.output is None and response.output_file_paths is None:
+            raise RuntimeError("Model generation returned no output.")
+
+        if response.output_file_paths:
+            output_file_path = response.output_file_paths[0]
+        else:
+            output_file_path = sp.output_file_path()
+            save_outputs(
+                [response.output[0]],
+                sp.data_type,
+                sp.fps,
+                True,
+                lambda _idx: output_file_path,
+                audio=response.audio,
+                audio_sample_rate=response.audio_sample_rate,
+                enable_frame_interpolation=sp.enable_frame_interpolation,
+                frame_interpolation_exp=sp.frame_interpolation_exp,
+                frame_interpolation_scale=sp.frame_interpolation_scale,
+                frame_interpolation_model_path=sp.frame_interpolation_model_path,
+                enable_upscaling=sp.enable_upscaling,
+                upscaling_model_path=sp.upscaling_model_path,
+                upscaling_scale=sp.upscaling_scale,
+            )
+
+        if hasattr(response, "model_dump"):
+            data = response.model_dump()
+        else:
+            data = response if isinstance(response, dict) else vars(response)
+
+        if output_file_path:
+            logger.info("Processing output file: %s", output_file_path)
+            b64_video = encode_video_to_base64(output_file_path)
+
+            if b64_video:
+                data["output"] = b64_video
+                data.pop("video_data", None)
+                data.pop("video_tensor", None)
+
+        return make_serializable(data)
+
+    except Exception as e:
+        logger.error("Error during generation: %s", e, exc_info=True)
+        return {"error": str(e)}
+
+
+vertex_router = APIRouter()
+
+
+@vertex_router.post(VERTEX_ROUTE)
+async def vertex_generate(vertex_req: VertexGenerateReqInput):
+    if not vertex_req.instances:
+        return orjson_response({"predictions": []})
+
+    server_args = get_global_server_args()
+    params = vertex_req.parameters or {}
+
+    futures = []
+
+    for inst in vertex_req.instances:
+        rid = f"vertex_{uuid.uuid4()}"
+
+        sp = build_sampling_params(
+            rid,
+            prompt=inst.get("prompt") or inst.get("text"),
+            image_path=inst.get("image") or inst.get("image_url"),
+            num_frames=params.get("num_frames"),
+            fps=params.get("fps"),
+            width=params.get("width"),
+            height=params.get("height"),
+            guidance_scale=params.get("guidance_scale"),
+            save_output=params.get("save_output"),
+        )
+
+        backend_req = prepare_request(server_args, sampling_params=sp)
+        futures.append(forward_to_scheduler(backend_req, sp))
+
+    results = await asyncio.gather(*futures)
+
+    return orjson_response({"predictions": results})
+
+
+def create_app(server_args: ServerArgs):
+    """
+    Create and configure the FastAPI application instance.
+    """
+    app = FastAPI(lifespan=lifespan)
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    @app.middleware("http")
+    async def wait_for_server_warmup(request: Request, call_next):
+        warmup_done = getattr(request.app.state, "server_warmup_done", None)
+        if (
+            warmup_done is not None
+            and not warmup_done.is_set()
+            and request.url.path not in SERVER_WARMUP_BYPASS_PATHS
+        ):
+            await warmup_done.wait()
+        return await call_next(request)
+
+    app.include_router(health_router)
+    app.include_router(vertex_router)
+
+    from sglang.multimodal_gen.runtime.entrypoints.openai import common_api, mesh_api
+
+    app.include_router(common_api.router)
+    app.include_router(image_api.router)
+    app.include_router(video_api.router)
+    app.include_router(realtime_video_api.router)
+    if server_args.pipeline_config.task_type.is_action_gen():
+        app.include_router(vla_api.router)
+        app.include_router(openpi.router)
+    app.include_router(mesh_api.router)
+    app.include_router(weights_api.router)
+    app.include_router(rollout_api.router)
+
+    app.state.server_args = server_args
+    return app

python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py

@@ -0,0 +1,252 @@
+import time
+from typing import Any, List, Optional, Union
+
+from fastapi import APIRouter, Body, HTTPException
+from pydantic import BaseModel, Field
+
+from sglang.multimodal_gen.registry import get_model_info
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    ListLorasReq,
+    MergeLoraWeightsReq,
+    SetLoraReq,
+    UnmergeLoraWeightsReq,
+    format_lora_message,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils.json_response import orjson_response
+
+router = APIRouter(prefix="/v1")
+logger = init_logger(__name__)
+
+
+class ModelCard(BaseModel):
+    """Model cards."""
+
+    id: str
+    object: str = "model"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    owned_by: str = "sglang"
+    root: Optional[str] = None
+    parent: Optional[str] = None
+    max_model_len: Optional[int] = None
+
+
+class DiffusionModelCard(ModelCard):
+    """Extended ModelCard with diffusion-specific fields."""
+
+    num_gpus: Optional[int] = None
+    task_type: Optional[str] = None
+    dit_precision: Optional[str] = None
+    vae_precision: Optional[str] = None
+    pipeline_name: Optional[str] = None
+    pipeline_class: Optional[str] = None
+
+
+async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str):
+    try:
+        output: OutputBatch = await async_scheduler_client.forward(req)
+        if output.error is None:
+            return {"status": "ok", "message": success_msg}
+        else:
+            error_msg = output.error
+            raise HTTPException(status_code=500, detail=f"{failure_msg}: {error_msg}")
+    except Exception as e:
+        if isinstance(e, HTTPException):
+            raise
+        logger.error(f"Error during '{failure_msg}': {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.post("/set_lora")
+async def set_lora(
+    lora_nickname: Union[str, List[str]] = Body(..., embed=True),
+    lora_path: Optional[Union[str, List[Optional[str]]]] = Body(None, embed=True),
+    target: Union[str, List[str]] = Body("all", embed=True),
+    strength: Union[float, List[float]] = Body(1.0, embed=True),
+    merge_mode: Optional[str] = Body(None, embed=True),
+):
+    """
+    Set LoRA adapter(s) for the specified transformer(s).
+    Supports both single LoRA (backward compatible) and multiple LoRA adapters.
+
+    Args:
+        lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings.
+        lora_path: Path(s) to the LoRA adapter(s) (local path or HF repo id).
+            Can be a string, None, or a list of strings/None. Must match the length of lora_nickname.
+        target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
+            If a list, must match the length of lora_nickname. Valid values:
+            - "all": Apply to all transformers (default)
+            - "transformer": Apply only to the primary transformer (high noise for Wan2.2)
+            - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
+            - "critic": Apply only to the critic model
+        strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
+            If a list, must match the length of lora_nickname. Values < 1.0 reduce the effect,
+            values > 1.0 amplify the effect.
+        merge_mode: Optional LoRA merge mode: "auto", "merge", or "dynamic".
+    """
+    req = SetLoraReq(
+        lora_nickname=lora_nickname,
+        lora_path=lora_path,
+        target=target,
+        strength=strength,
+        merge_mode=merge_mode,
+    )
+    nickname_str, target_str, strength_str = format_lora_message(
+        lora_nickname, target, strength
+    )
+
+    return await _handle_lora_request(
+        req,
+        f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})",
+        "Failed to set LoRA adapter",
+    )
+
+
+@router.post("/merge_lora_weights")
+async def merge_lora_weights(
+    target: str = Body("all", embed=True),
+    strength: float = Body(1.0, embed=True),
+):
+    """
+    Merge LoRA weights into the base model.
+
+    Args:
+        target: Which transformer(s) to merge. One of "all", "transformer",
+                "transformer_2", "critic".
+        strength: LoRA strength for merge, default 1.0. Values < 1.0 reduce the effect,
+            values > 1.0 amplify the effect.
+    """
+    req = MergeLoraWeightsReq(target=target, strength=strength)
+    return await _handle_lora_request(
+        req,
+        f"Successfully merged LoRA weights (target: {target}, strength: {strength})",
+        "Failed to merge LoRA weights",
+    )
+
+
+@router.post("/unmerge_lora_weights")
+async def unmerge_lora_weights(
+    target: str = Body("all", embed=True),
+):
+    """
+    Unmerge LoRA weights from the base model.
+
+    Args:
+        target: Which transformer(s) to unmerge. One of "all", "transformer",
+                "transformer_2", "critic".
+    """
+    req = UnmergeLoraWeightsReq(target=target)
+    return await _handle_lora_request(
+        req,
+        f"Successfully unmerged LoRA weights (target: {target})",
+        "Failed to unmerge LoRA weights",
+    )
+
+
+@router.get("/model_info")
+async def model_info():
+    """Get the model information."""
+    server_args = get_global_server_args()
+    if not server_args:
+        raise HTTPException(status_code=500, detail="Server args not initialized")
+
+    result = {
+        "model_path": server_args.model_path,
+    }
+    return result
+
+
+@router.get("/list_loras")
+async def list_loras():
+    """List loaded LoRA adapters and current application status per module."""
+    try:
+        req = ListLorasReq()
+        output: OutputBatch = await async_scheduler_client.forward(req)
+        if output.error is None:
+            return output.output or {}
+        else:
+            raise HTTPException(status_code=500, detail=output.error)
+    except Exception as e:
+        if isinstance(e, HTTPException):
+            raise
+        logger.error(f"Error during 'list_loras': {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get("/models")
+async def available_models():
+    """Show available models. OpenAI-compatible endpoint with extended diffusion info."""
+    server_args = get_global_server_args()
+    if not server_args:
+        raise HTTPException(status_code=500, detail="Server args not initialized")
+
+    model_info = get_model_info(
+        server_args.model_path,
+        backend=server_args.backend,
+        model_id=server_args.model_id,
+    )
+
+    card_kwargs = {
+        "id": server_args.model_path,
+        "root": server_args.model_path,
+        # Extended diffusion-specific fields
+        "num_gpus": server_args.num_gpus,
+        "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,
+    }
+
+    if model_info:
+        card_kwargs["pipeline_name"] = model_info.pipeline_cls.pipeline_name
+        card_kwargs["pipeline_class"] = model_info.pipeline_cls.__name__
+
+    model_card = DiffusionModelCard(**card_kwargs)
+
+    # Return dict directly to preserve extended fields (ModelList strips them)
+    return {"object": "list", "data": [model_card.model_dump()]}
+
+
+@router.get("/models/{model:path}")
+async def retrieve_model(model: str):
+    """Retrieve a model instance. OpenAI-compatible endpoint with extended diffusion info."""
+    server_args = get_global_server_args()
+    if not server_args:
+        raise HTTPException(status_code=500, detail="Server args not initialized")
+
+    if model != server_args.model_path:
+        return orjson_response(
+            {
+                "error": {
+                    "message": f"The model '{model}' does not exist",
+                    "type": "invalid_request_error",
+                    "param": "model",
+                    "code": "model_not_found",
+                }
+            },
+            status_code=404,
+        )
+
+    model_info = get_model_info(
+        server_args.model_path,
+        backend=server_args.backend,
+        model_id=server_args.model_id,
+    )
+
+    card_kwargs = {
+        "id": model,
+        "root": model,
+        "num_gpus": server_args.num_gpus,
+        "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,
+    }
+
+    if model_info:
+        card_kwargs["pipeline_name"] = model_info.pipeline_cls.pipeline_name
+        card_kwargs["pipeline_class"] = model_info.pipeline_cls.__name__
+
+    # Return dict to preserve extended fields
+    return DiffusionModelCard(**card_kwargs).model_dump()

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -0,0 +1,447 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import base64
+import contextlib
+import json
+import os
+import time
+from typing import Any, List, Optional
+
+from fastapi import (
+    APIRouter,
+    File,
+    Form,
+    HTTPException,
+    Path,
+    Query,
+    Request,
+    UploadFile,
+)
+from fastapi.responses import FileResponse
+
+from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    ImageGenerationsRequest,
+    ImageResponse,
+    ImageResponseData,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage
+from sglang.multimodal_gen.runtime.entrypoints.openai.stores import IMAGE_STORE
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
+    add_common_data_to_response,
+    build_sampling_params,
+    choose_output_image_ext,
+    flatten_extra_params,
+    merge_image_input_list,
+    process_generation_batch,
+    save_image_to_path,
+    temp_dir_if_disabled,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.srt.observability.trace import extract_trace_headers
+
+router = APIRouter(prefix="/v1/images", tags=["images"])
+
+
+def _get_extra_field(request, field_name):
+    """Get a field from model_extra, with fallback to nested extra_body dict."""
+    extra = request.model_extra or {}
+    value = extra.get(field_name)
+    if value is not None:
+        return value
+    if field_name == "use_guardrails" and extra.get("guardrails") is not None:
+        return extra["guardrails"]
+
+    for container_name in ("extra_body", "extra_json", "extra_args", "extra_params"):
+        value = _parse_extra_container(extra.get(container_name)).get(field_name)
+        if value is not None:
+            return value
+
+    return value
+
+
+def _parse_extra_container(value: Any) -> dict[str, Any]:
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except Exception:
+            return {}
+    if isinstance(value, dict):
+        return flatten_extra_params(dict(value))
+    return {}
+
+
+def _read_b64_for_paths(paths: list[str]) -> list[str]:
+    """Read and base64-encode each file. Must be called before cloud upload deletes them."""
+    result = []
+    for path in paths:
+        with open(path, "rb") as f:
+            result.append(base64.b64encode(f.read()).decode("utf-8"))
+    return result
+
+
+def _build_image_response_kwargs(
+    save_file_path_list: list[str],
+    resp_format: str,
+    prompt: str,
+    request_id: str,
+    result: OutputBatch,
+    *,
+    b64_list: list[str] | None = None,
+    cloud_url: str | None = None,
+    fallback_url: str | None = None,
+    is_persistent: bool = True,
+) -> dict:
+    """Build ImageResponse data list.
+
+    For b64_json: uses pre-read b64_list (call _read_b64_for_paths first).
+    For url: uses cloud_url or fallback_url.
+    file_path is omitted when is_persistent=False to avoid exposing stale temp paths.
+    """
+    ret = None
+    if resp_format == "b64_json":
+        if not b64_list:
+            raise ValueError("b64_list required for b64_json response_format")
+        data = [
+            ImageResponseData(
+                b64_json=b64,
+                revised_prompt=prompt,
+                file_path=os.path.abspath(path) if is_persistent else None,
+            )
+            for b64, path in zip(b64_list, save_file_path_list)
+        ]
+        ret = {"data": data}
+    elif resp_format == "url":
+        url = cloud_url or fallback_url
+        if not url:
+            raise HTTPException(
+                status_code=400,
+                detail="response_format='url' requires cloud storage to be configured.",
+            )
+        ret = {
+            "data": [
+                ImageResponseData(
+                    url=url,
+                    revised_prompt=prompt,
+                    file_path=(
+                        os.path.abspath(save_file_path_list[0])
+                        if is_persistent
+                        else None
+                    ),
+                )
+            ],
+        }
+    else:
+        raise HTTPException(
+            status_code=400, detail=f"response_format={resp_format} is not supported"
+        )
+
+    ret = add_common_data_to_response(ret, request_id=request_id, result=result)
+
+    return ret
+
+
+@router.post("/generations", response_model=ImageResponse)
+async def generations(
+    request: ImageGenerationsRequest,
+    raw_request: Request,
+):
+    request_id = generate_request_id()
+    server_args = get_global_server_args()
+    is_cosmos3 = "cosmos3" in (server_args.model_path or "").lower()
+    ext = (
+        "png"
+        if is_cosmos3 and request.output_format is None
+        else choose_output_image_ext(request.output_format, request.background)
+    )
+
+    with temp_dir_if_disabled(server_args.output_path) as output_dir:
+        sampling = build_sampling_params(
+            request_id,
+            prompt=request.prompt,
+            size=request.size,
+            width=request.width,
+            height=request.height,
+            num_outputs_per_prompt=max(1, min(int(request.n or 1), 10)),
+            output_file_name=f"{request_id}.{ext}",
+            output_path=output_dir,
+            num_frames=1,
+            seed=request.seed,
+            generator_device=request.generator_device,
+            num_inference_steps=request.num_inference_steps,
+            guidance_scale=request.guidance_scale,
+            true_cfg_scale=request.true_cfg_scale,
+            negative_prompt=request.negative_prompt,
+            max_sequence_length=(
+                request.max_sequence_length
+                if request.max_sequence_length is not None
+                else _get_extra_field(request, "max_sequence_length")
+            ),
+            flow_shift=(
+                request.flow_shift
+                if request.flow_shift is not None
+                else _get_extra_field(request, "flow_shift")
+            ),
+            use_duration_template=_get_extra_field(request, "use_duration_template"),
+            use_resolution_template=_get_extra_field(
+                request, "use_resolution_template"
+            ),
+            use_system_prompt=_get_extra_field(request, "use_system_prompt"),
+            use_guardrails=_get_extra_field(request, "use_guardrails"),
+            enable_teacache=request.enable_teacache,
+            output_compression=request.output_compression,
+            output_quality=request.output_quality,
+            diffusers_kwargs=request.diffusers_kwargs,
+            enable_upscaling=request.enable_upscaling,
+            upscaling_model_path=request.upscaling_model_path,
+            upscaling_scale=request.upscaling_scale,
+            perf_dump_path=request.perf_dump_path,
+            use_pe=_get_extra_field(request, "use_pe"),
+            preset=_get_extra_field(request, "preset"),
+            progressive_mode=(
+                request.progressive_mode
+                if request.progressive_mode is not None
+                else _get_extra_field(request, "progressive_mode")
+            ),
+            progressive_levels=(
+                request.progressive_levels
+                if request.progressive_levels is not None
+                else _get_extra_field(request, "progressive_levels")
+            ),
+            progressive_delta=(
+                request.progressive_delta
+                if request.progressive_delta is not None
+                else _get_extra_field(request, "progressive_delta")
+            ),
+        )
+        trace_headers = extract_trace_headers(raw_request.headers)
+        batch = prepare_request(
+            server_args=server_args,
+            sampling_params=sampling,
+            external_trace_header=trace_headers,
+        )
+        # Add diffusers_kwargs if provided
+        if request.diffusers_kwargs:
+            batch.extra["diffusers_kwargs"] = request.diffusers_kwargs
+
+        save_file_path_list, result = await process_generation_batch(
+            async_scheduler_client, batch
+        )
+        save_file_path = save_file_path_list[0]
+        resp_format = (request.response_format or "b64_json").lower()
+        if (
+            is_cosmos3
+            and "response_format" not in request.model_fields_set
+            and request.response_format == "url"
+        ):
+            resp_format = "b64_json"
+
+        # read b64 before cloud upload may delete the local file
+        b64_list = (
+            _read_b64_for_paths(save_file_path_list)
+            if resp_format == "b64_json"
+            else None
+        )
+
+        cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
+
+        is_persistent = server_args.output_path is not None
+        await IMAGE_STORE.upsert(
+            request_id,
+            {
+                "id": request_id,
+                "created_at": int(time.time()),
+                "file_path": None if cloud_url or not is_persistent else save_file_path,
+                "url": cloud_url,
+            },
+        )
+
+        response_kwargs = _build_image_response_kwargs(
+            save_file_path_list,
+            resp_format,
+            request.prompt,
+            request_id,
+            result,
+            b64_list=b64_list,
+            cloud_url=cloud_url,
+            fallback_url=f"/v1/images/{request_id}/content" if is_persistent else None,
+            is_persistent=is_persistent,
+        )
+
+    return ImageResponse(**response_kwargs)
+
+
+@router.post("/edits", response_model=ImageResponse)
+async def edits(
+    raw_request: Request,
+    image: Optional[List[UploadFile]] = File(None),
+    image_array: Optional[List[UploadFile]] = File(None, alias="image[]"),
+    url: Optional[List[str]] = Form(None),
+    url_array: Optional[List[str]] = Form(None, alias="url[]"),
+    prompt: str = Form(...),
+    mask: Optional[UploadFile] = File(None),
+    model: Optional[str] = Form(None),
+    n: Optional[int] = Form(1),
+    response_format: Optional[str] = Form(None),
+    size: Optional[str] = Form(None),
+    output_format: Optional[str] = Form(None),
+    background: Optional[str] = Form("auto"),
+    seed: Optional[int] = Form(None),
+    generator_device: Optional[str] = Form("cuda"),
+    user: Optional[str] = Form(None),
+    negative_prompt: Optional[str] = Form(None),
+    guidance_scale: Optional[float] = Form(None),
+    true_cfg_scale: Optional[float] = Form(None),
+    num_inference_steps: Optional[int] = Form(None),
+    output_quality: Optional[str] = Form("default"),
+    output_compression: Optional[int] = Form(None),
+    enable_teacache: Optional[bool] = Form(False),
+    enable_upscaling: Optional[bool] = Form(False),
+    upscaling_model_path: Optional[str] = Form(None),
+    upscaling_scale: Optional[int] = Form(4),
+    num_frames: int = Form(1),
+):
+    request_id = generate_request_id()
+    server_args = get_global_server_args()
+    # Resolve images from either `image` or `image[]` (OpenAI SDK sends `image[]` when list is provided)
+    images = image or image_array
+    urls = url or url_array
+
+    if (not images or len(images) == 0) and (not urls or len(urls) == 0):
+        raise HTTPException(
+            status_code=422, detail="Field 'image' or 'url' is required"
+        )
+
+    image_list = merge_image_input_list(images, urls)
+
+    with contextlib.ExitStack() as stack:
+        uploads_dir = stack.enter_context(
+            temp_dir_if_disabled(server_args.input_save_path)
+        )
+        output_dir = stack.enter_context(temp_dir_if_disabled(server_args.output_path))
+
+        input_paths = []
+        try:
+            for idx, img in enumerate(image_list):
+                filename = img.filename if hasattr(img, "filename") else f"image_{idx}"
+                input_path = await save_image_to_path(
+                    img,
+                    os.path.join(uploads_dir, f"{request_id}_{idx}_{filename}"),
+                    prefer_remote_source=server_args.input_save_path is None,
+                )
+                input_paths.append(input_path)
+        except Exception as e:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Failed to process image source: {str(e)}",
+            )
+
+        ext = choose_output_image_ext(output_format, background)
+        sampling = build_sampling_params(
+            request_id,
+            prompt=prompt,
+            size=size,
+            num_outputs_per_prompt=max(1, min(int(n or 1), 10)),
+            output_file_name=f"{request_id}.{ext}",
+            output_path=output_dir,
+            image_path=input_paths,
+            seed=seed,
+            generator_device=generator_device,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            true_cfg_scale=true_cfg_scale,
+            num_inference_steps=num_inference_steps,
+            enable_teacache=enable_teacache,
+            num_frames=num_frames,
+            output_compression=output_compression,
+            output_quality=output_quality,
+            enable_upscaling=enable_upscaling,
+            upscaling_model_path=upscaling_model_path,
+            upscaling_scale=upscaling_scale,
+        )
+        trace_headers = extract_trace_headers(raw_request.headers)
+        batch = prepare_request(
+            server_args=server_args,
+            sampling_params=sampling,
+            external_trace_header=trace_headers,
+        )
+        save_file_path_list, result = await process_generation_batch(
+            async_scheduler_client, batch
+        )
+        save_file_path = save_file_path_list[0]
+        resp_format = (response_format or "b64_json").lower()
+
+        # read b64 before cloud upload may delete the local file
+        b64_list = (
+            _read_b64_for_paths(save_file_path_list)
+            if resp_format == "b64_json"
+            else None
+        )
+
+        cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
+
+        is_persistent = server_args.output_path is not None
+        is_input_persistent = server_args.input_save_path is not None
+        await IMAGE_STORE.upsert(
+            request_id,
+            {
+                "id": request_id,
+                "created_at": int(time.time()),
+                "file_path": None if cloud_url or not is_persistent else save_file_path,
+                "url": cloud_url,
+                "input_image_paths": input_paths if is_input_persistent else None,
+                "num_input_images": len(input_paths),
+            },
+        )
+
+        response_kwargs = _build_image_response_kwargs(
+            save_file_path_list,
+            resp_format,
+            prompt,
+            request_id,
+            result,
+            b64_list=b64_list,
+            cloud_url=cloud_url,
+            fallback_url=f"/v1/images/{request_id}/content" if is_persistent else None,
+            is_persistent=is_persistent,
+        )
+
+    return ImageResponse(**response_kwargs)
+
+
+@router.get("/{image_id}/content")
+async def download_image_content(
+    image_id: str = Path(...), variant: Optional[str] = Query(None)
+):
+    item = await IMAGE_STORE.get(image_id)
+    if not item:
+        raise HTTPException(status_code=404, detail="Image not found")
+
+    if item.get("url"):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Image has been uploaded to cloud storage. Please use the cloud URL: {item.get('url')}",
+        )
+
+    file_path = item.get("file_path")
+    if not file_path:
+        raise HTTPException(
+            status_code=404,
+            detail="Image was not persisted on disk (output_path is disabled). Use b64_json response_format or configure cloud storage.",
+        )
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="Image is still being generated")
+
+    ext = os.path.splitext(file_path)[1].lower()
+    media_type = "image/jpeg"
+    if ext == ".png":
+        media_type = "image/png"
+    elif ext == ".webp":
+        media_type = "image/webp"
+
+    return FileResponse(
+        path=file_path, media_type=media_type, filename=os.path.basename(file_path)
+    )

python/sglang/multimodal_gen/runtime/entrypoints/openai/mesh_api.py

@@ -0,0 +1,296 @@
+import asyncio
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from fastapi import (
+    APIRouter,
+    File,
+    Form,
+    HTTPException,
+    Path,
+    Query,
+    Request,
+    UploadFile,
+)
+from fastapi.responses import FileResponse
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    SamplingParams,
+    generate_request_id,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    MeshGenerationsRequest,
+    MeshListResponse,
+    MeshResponse,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage
+from sglang.multimodal_gen.runtime.entrypoints.openai.stores import MESH_STORE
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
+    add_common_data_to_response,
+    merge_image_input_list,
+    process_generation_batch,
+    save_image_to_path,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+router = APIRouter(prefix="/v1/meshes", tags=["meshes"])
+
+
+def _normalize_format(fmt: Optional[str]) -> str:
+    fmt = (fmt or "glb").lower()
+    return fmt if fmt in ("glb", "obj") else "glb"
+
+
+def _build_sampling_params_from_request(
+    request_id: str, req: MeshGenerationsRequest, image_path: Optional[str] = None
+) -> SamplingParams:
+    ext = _normalize_format(req.output_format)
+
+    server_args = get_global_server_args()
+    sampling_kwargs: Dict[str, Any] = {
+        "request_id": request_id,
+        "prompt": req.prompt,
+        "num_frames": 1,
+        "image_path": [image_path] if image_path else None,
+        "save_output": True,
+        "output_file_name": f"{request_id}.{ext}",
+        "seed": req.seed,
+        "generator_device": req.generator_device,
+    }
+    if req.num_inference_steps is not None:
+        sampling_kwargs["num_inference_steps"] = req.num_inference_steps
+    if req.guidance_scale is not None:
+        sampling_kwargs["guidance_scale"] = req.guidance_scale
+    if req.negative_prompt is not None:
+        sampling_kwargs["negative_prompt"] = req.negative_prompt
+
+    return SamplingParams.from_user_sampling_params_args(
+        model_path=server_args.model_path,
+        server_args=server_args,
+        **sampling_kwargs,
+    )
+
+
+def _mesh_job_from_sampling(
+    request_id: str, req: MeshGenerationsRequest, sampling: SamplingParams
+) -> Dict[str, Any]:
+    return {
+        "id": request_id,
+        "object": "mesh",
+        "model": req.model or "",
+        "status": "queued",
+        "progress": 0,
+        "created_at": int(time.time()),
+        "format": _normalize_format(req.output_format),
+        "file_path": os.path.abspath(sampling.output_file_path()),
+    }
+
+
+async def _dispatch_job_async(job_id: str, batch: Req) -> None:
+    from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+
+    try:
+        save_file_path_list, result = await process_generation_batch(
+            async_scheduler_client, batch
+        )
+        save_file_path = save_file_path_list[0]
+
+        file_size = None
+        if os.path.exists(save_file_path):
+            file_size = os.path.getsize(save_file_path)
+
+        cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
+
+        update_fields: Dict[str, Any] = {
+            "status": "completed",
+            "progress": 100,
+            "completed_at": int(time.time()),
+            "url": cloud_url,
+            "file_path": save_file_path if not cloud_url else None,
+            "file_size_bytes": file_size,
+        }
+        update_fields = add_common_data_to_response(
+            update_fields, request_id=job_id, result=result
+        )
+        await MESH_STORE.update_fields(job_id, update_fields)
+    except Exception as e:
+        logger.error(f"{e}")
+        await MESH_STORE.update_fields(
+            job_id, {"status": "failed", "error": {"message": str(e)}}
+        )
+
+
+@router.post("", response_model=MeshResponse)
+async def create_mesh(
+    request: Request,
+    image: Optional[List[UploadFile]] = File(None),
+    image_array: Optional[List[UploadFile]] = File(None, alias="image[]"),
+    url: Optional[List[str]] = Form(None),
+    url_array: Optional[List[str]] = Form(None, alias="url[]"),
+    prompt: Optional[str] = Form("generate 3d mesh"),
+    model: Optional[str] = Form(None),
+    seed: Optional[int] = Form(None),
+    generator_device: Optional[str] = Form("cuda"),
+    guidance_scale: Optional[float] = Form(None),
+    num_inference_steps: Optional[int] = Form(None),
+    negative_prompt: Optional[str] = Form(None),
+    output_format: Optional[str] = Form("glb"),
+):
+    content_type = request.headers.get("content-type", "").lower()
+    request_id = generate_request_id()
+    server_args = get_global_server_args()
+
+    input_path = None
+
+    if "multipart/form-data" in content_type:
+        images = image or image_array
+        urls = url or url_array
+        image_list = merge_image_input_list(images, urls)
+
+        if not image_list:
+            raise HTTPException(
+                status_code=422,
+                detail="Field 'image' or 'url' is required for mesh generation",
+            )
+
+        uploads_dir = os.path.join("outputs", "uploads")
+        os.makedirs(uploads_dir, exist_ok=True)
+        img = image_list[0]
+        filename = img.filename if hasattr(img, "filename") else "input_image"
+        try:
+            input_path = await save_image_to_path(
+                img, os.path.join(uploads_dir, f"{request_id}_{filename}")
+            )
+        except Exception as e:
+            raise HTTPException(
+                status_code=400, detail=f"Failed to process image source: {str(e)}"
+            )
+
+        req = MeshGenerationsRequest(
+            prompt=prompt or "generate 3d mesh",
+            model=model,
+            seed=seed,
+            generator_device=generator_device,
+            num_inference_steps=num_inference_steps,
+            negative_prompt=negative_prompt,
+            output_format=output_format,
+            **(
+                {"guidance_scale": guidance_scale} if guidance_scale is not None else {}
+            ),
+        )
+    else:
+        try:
+            body = await request.json()
+        except Exception:
+            body = {}
+        try:
+            payload: Dict[str, Any] = dict(body or {})
+
+            if payload.get("input_image"):
+                img_src = payload.pop("input_image")
+                uploads_dir = os.path.join("outputs", "uploads")
+                os.makedirs(uploads_dir, exist_ok=True)
+                input_path = await save_image_to_path(
+                    img_src,
+                    os.path.join(uploads_dir, f"{request_id}_input_image"),
+                )
+
+            req = MeshGenerationsRequest(**payload)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Invalid request body: {e}")
+
+    if not input_path:
+        raise HTTPException(
+            status_code=422,
+            detail="An input image is required for mesh generation",
+        )
+
+    sampling_params = _build_sampling_params_from_request(request_id, req, input_path)
+    job = _mesh_job_from_sampling(request_id, req, sampling_params)
+    await MESH_STORE.upsert(request_id, job)
+
+    batch = prepare_request(
+        server_args=server_args,
+        sampling_params=sampling_params,
+    )
+
+    asyncio.create_task(_dispatch_job_async(request_id, batch))
+    return MeshResponse(**job)
+
+
+@router.get("", response_model=MeshListResponse)
+async def list_meshes(
+    after: Optional[str] = Query(None),
+    limit: Optional[int] = Query(None, ge=1, le=100),
+    order: Optional[str] = Query("desc"),
+):
+    order = (order or "desc").lower()
+    if order not in ("asc", "desc"):
+        order = "desc"
+    jobs = await MESH_STORE.list_values()
+
+    reverse = order != "asc"
+    jobs.sort(key=lambda j: j.get("created_at", 0), reverse=reverse)
+
+    if after is not None:
+        try:
+            idx = next(i for i, j in enumerate(jobs) if j["id"] == after)
+            jobs = jobs[idx + 1 :]
+        except StopIteration:
+            jobs = []
+
+    if limit is not None:
+        jobs = jobs[:limit]
+    items = [MeshResponse(**j) for j in jobs]
+    return MeshListResponse(data=items)
+
+
+@router.get("/{mesh_id}", response_model=MeshResponse)
+async def retrieve_mesh(mesh_id: str = Path(...)):
+    job = await MESH_STORE.get(mesh_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Mesh not found")
+    return MeshResponse(**job)
+
+
+@router.delete("/{mesh_id}", response_model=MeshResponse)
+async def delete_mesh(mesh_id: str = Path(...)):
+    job = await MESH_STORE.pop(mesh_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Mesh not found")
+    job["status"] = "deleted"
+    return MeshResponse(**job)
+
+
+@router.get("/{mesh_id}/content")
+async def download_mesh_content(
+    mesh_id: str = Path(...), variant: Optional[str] = Query(None)
+):
+    job = await MESH_STORE.get(mesh_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Mesh not found")
+
+    if job.get("url"):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Mesh has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}",
+        )
+
+    file_path = job.get("file_path")
+    if not file_path or not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="Generation is still in-progress")
+
+    ext = os.path.splitext(file_path)[1].lower()
+    media_type = {
+        ".glb": "model/gltf-binary",
+        ".obj": "text/plain",
+    }.get(ext, "application/octet-stream")
+
+    return FileResponse(
+        path=file_path, media_type=media_type, filename=os.path.basename(file_path)
+    )

python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py

@@ -0,0 +1,227 @@
+import time
+import uuid
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+# Image API protocol models
+class ImageResponseData(BaseModel):
+    b64_json: Optional[str] = None
+    url: Optional[str] = None
+    revised_prompt: Optional[str] = None
+    file_path: Optional[str] = None
+
+
+class ImageResponse(BaseModel):
+    id: str
+    created: int = Field(default_factory=lambda: int(time.time()))
+    data: List[ImageResponseData]
+    peak_memory_mb: Optional[float] = None
+    inference_time_s: Optional[float] = None
+
+
+class ImageGenerationsRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    prompt: str
+    model: Optional[str] = None
+    n: Optional[int] = 1
+    quality: Optional[str] = "auto"
+    response_format: Optional[str] = "url"  # url | b64_json
+    size: Optional[str] = "1024x1024"  # e.g., 1024x1024
+    style: Optional[str] = "vivid"
+    background: Optional[str] = "auto"  # transparent | opaque | auto
+    output_format: Optional[str] = None  # png | jpeg | webp
+    user: Optional[str] = None
+    # SGLang extensions
+    width: Optional[int] = None
+    height: Optional[int] = None
+    num_inference_steps: Optional[int] = None
+    guidance_scale: Optional[float] = None
+    true_cfg_scale: Optional[float] = (
+        None  # for CFG vs guidance distillation (e.g., QwenImage)
+    )
+    seed: Optional[Union[int, List[int]]] = None
+    generator_device: Optional[str] = "cuda"
+    negative_prompt: Optional[str] = None
+    output_quality: Optional[str] = "default"
+    output_compression: Optional[int] = None
+    enable_teacache: Optional[bool] = False
+    max_sequence_length: Optional[int] = None
+    flow_shift: Optional[float] = None
+    # Upscaling
+    enable_upscaling: Optional[bool] = False
+    upscaling_model_path: Optional[str] = None
+    upscaling_scale: Optional[int] = 4
+    diffusers_kwargs: Optional[Dict[str, Any]] = None  # kwargs for diffusers backend
+    # Performance profiling
+    perf_dump_path: Optional[str] = None
+    # Progressive resolution generation
+    progressive_mode: Optional[str] = None
+    progressive_levels: Optional[int] = None
+    progressive_delta: Optional[float] = None
+
+
+# Video API protocol models
+class VideoResponse(BaseModel):
+    id: str
+    object: str = "video"
+    model: str = "sora-2"
+    status: str = "queued"
+    progress: int = 0
+    created_at: int = Field(default_factory=lambda: int(time.time()))
+    size: str = ""
+    seconds: str = "4"
+    quality: str = "standard"
+    url: Optional[str] = None
+    remixed_from_video_id: Optional[str] = None
+    completed_at: Optional[int] = None
+    expires_at: Optional[int] = None
+    error: Optional[Dict[str, Any]] = None
+    file_path: Optional[str] = None
+    file_paths: Optional[List[str]] = None
+    num_outputs: Optional[int] = None
+    peak_memory_mb: Optional[float] = None
+    inference_time_s: Optional[float] = None
+    action: Optional[Dict[str, Any]] = None
+
+
+class VideoGenerationsRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    prompt: str
+    input_reference: Optional[str] = None
+    reference_url: Optional[str] = None
+    video_path: Optional[str] = None
+    video_url: Optional[str] = None
+    model: Optional[str] = None
+    n: Optional[int] = 1
+    num_outputs_per_prompt: Optional[int] = None
+    seconds: Optional[int] = 4
+    size: Optional[str] = ""
+    fps: Optional[int] = None
+    num_frames: Optional[int] = None
+    seed: Optional[Union[int, List[int]]] = None
+    generator_device: Optional[str] = "cuda"
+    # SGLang extensions
+    width: Optional[int] = None
+    height: Optional[int] = None
+    num_inference_steps: Optional[int] = None
+    guidance_scale: Optional[float] = None
+    guidance_scale_2: Optional[float] = None
+    true_cfg_scale: Optional[float] = (
+        None  # for CFG vs guidance distillation (e.g., QwenImage)
+    )
+    negative_prompt: Optional[str] = None
+    max_sequence_length: Optional[int] = None
+    flow_shift: Optional[float] = None
+    enable_teacache: Optional[bool] = False
+    # Frame interpolation
+    enable_frame_interpolation: Optional[bool] = False
+    frame_interpolation_exp: Optional[int] = 1  # 1=2×, 2=4×
+    frame_interpolation_scale: Optional[float] = 1.0
+    frame_interpolation_model_path: Optional[str] = None
+    # Upscaling
+    enable_upscaling: Optional[bool] = False
+    upscaling_model_path: Optional[str] = None
+    upscaling_scale: Optional[int] = 4
+    output_quality: Optional[str] = "default"
+    output_compression: Optional[int] = None
+    output_path: Optional[str] = None
+    diffusers_kwargs: Optional[Dict[str, Any]] = None  # kwargs for diffusers backend
+    # Performance profiling
+    perf_dump_path: Optional[str] = None
+
+
+class VideoListResponse(BaseModel):
+    data: List[VideoResponse]
+    object: str = "list"
+
+
+class VideoRemixRequest(BaseModel):
+    prompt: str
+
+
+class RealtimeVideoGenerationsRequest(VideoGenerationsRequest):
+    type: Literal["init"]
+    # WebSocket does not support multipart/form-data image uploads
+    first_frame: Optional[bytes | str] = None
+    condition_inputs: Optional[Dict[str, Any]] = None
+    max_chunks: Optional[int] = Field(default=None, ge=1)
+    seed: Optional[int] = 42
+    guidance_scale: Optional[float] = 1.0
+    size: Optional[str] = "832x480"
+    profile: Optional[bool] = False
+    num_profiled_timesteps: Optional[int] = None
+    profile_all_stages: Optional[bool] = False
+    realtime_output_format: Optional[Literal["raw", "webp", "jpeg"]] = None
+    realtime_preview_max_width: Optional[int] = None
+    realtime_output_pacing: Optional[bool] = False
+    realtime_causal_sink_size: Optional[int] = None
+    realtime_causal_kv_cache_num_frames: Optional[int] = None
+
+
+class RealtimeEvent(BaseModel):
+    type: Literal["event"]
+    kind: str
+    payload: Any = None
+    event_id: Optional[int] = None
+
+
+# Mesh API protocol models
+class MeshResponse(BaseModel):
+    id: str
+    object: str = "mesh"
+    model: str = ""
+    status: str = "queued"
+    progress: int = 0
+    created_at: int = Field(default_factory=lambda: int(time.time()))
+    format: str = "glb"
+    url: Optional[str] = None
+    completed_at: Optional[int] = None
+    expires_at: Optional[int] = None
+    error: Optional[Dict[str, Any]] = None
+    file_path: Optional[str] = None
+    file_size_bytes: Optional[int] = None
+    peak_memory_mb: Optional[float] = None
+    inference_time_s: Optional[float] = None
+
+
+class MeshGenerationsRequest(BaseModel):
+    prompt: str = "generate 3d mesh"
+    input_image: Optional[str] = None
+    model: Optional[str] = None
+    seed: Optional[Union[int, List[int]]] = None
+    generator_device: Optional[str] = "cuda"
+    num_inference_steps: Optional[int] = None
+    guidance_scale: Optional[float] = None
+    negative_prompt: Optional[str] = None
+    output_format: Optional[str] = "glb"
+
+
+class MeshListResponse(BaseModel):
+    data: List[MeshResponse]
+    object: str = "list"
+
+
+@dataclass
+class BaseReq(ABC):
+    rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True)
+    http_worker_ipc: Optional[str] = field(default=None, kw_only=True)
+
+    def regenerate_rid(self):
+        """Generate a new request ID and return it."""
+        if isinstance(self.rid, list):
+            self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))]
+        else:
+            self.rid = uuid.uuid4().hex
+        return self.rid
+
+
+@dataclass
+class VertexGenerateReqInput(BaseReq):
+    instances: List[dict]
+    parameters: Optional[dict] = None

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

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py

@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    RealtimeEvent,
+    RealtimeVideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import (
+    BaseRealtimeModelAdapter,
+    RealtimeChunkInputs,
+    build_realtime_sampling_params,
+    save_realtime_first_frame,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.constants import (
+    LINGBOT_CAMERA_ACTIONS_CONDITION,
+    LINGBOT_PROMPT_UPDATED_CONDITION,
+)
+from sglang.multimodal_gen.runtime.realtime.control_signals import (
+    ControlSignalQueue,
+    ParsedControlEventPayload,
+    parse_control_event_payload,
+)
+from sglang.multimodal_gen.runtime.realtime.states import (
+    RealtimeCameraControlState,
+)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import (
+        GenerateSession,
+        RealtimeChunkContext,
+    )
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+LINGBOT_REALTIME_DEFAULT_NUM_INFERENCE_STEPS = 4
+LINGBOT_REALTIME_MIN_CONDITION_CHUNKS = 2
+COMPOSITE_INPUT_EVENT_KIND = "composite_input"
+
+
+class LingBotWorldRealtimeState(RealtimeCameraControlState):
+    def __init__(self):
+        super().__init__(
+            min_pulse_items=1,
+            script_maxlen=512,
+            max_transitions=512,
+        )
+        self.prompt_queue = ControlSignalQueue(max_events={"prompt": 1})
+
+    def clear(self) -> None:
+        super().clear()
+        self.prompt_queue.clear()
+
+    def receive_prompt(self, prompt: str, *, event_id: int | None = None) -> None:
+        self.prompt_queue.push("prompt", prompt, event_id=event_id)
+
+    def parse_camera_control_event_payload(
+        self,
+        payload: Any,
+        *,
+        event_id: int | None,
+    ) -> ParsedControlEventPayload:
+        return parse_control_event_payload(
+            payload,
+            event_id=event_id,
+            kind="camera_actions",
+            normalize_state_payload=self._normalize_state_actions,
+            validate_script_payload=LingBotWorldRealtimeAdapter._validate_camera_actions,
+        )
+
+    def receive_parsed_camera_control_event_payload(
+        self,
+        parsed: ParsedControlEventPayload,
+        *,
+        event_id: int | None,
+    ) -> str:
+        if parsed.mode == "state":
+            transitions = parsed.payload
+            self.receive_camera_state_transitions(transitions)
+            return f"kind=camera_actions, mode=state, transitions={len(transitions)}"
+
+        camera_actions = parsed.payload
+        self.receive_camera_action_script(camera_actions, event_id=event_id)
+        return f"kind=camera_actions, mode=script, frames={len(camera_actions)}"
+
+    def receive_camera_control_event_payload(
+        self,
+        payload: Any,
+        *,
+        event_id: int | None,
+    ) -> str:
+        parsed = self.parse_camera_control_event_payload(payload, event_id=event_id)
+        return self.receive_parsed_camera_control_event_payload(
+            parsed, event_id=event_id
+        )
+
+    def sample_prompt(self) -> str:
+        prompt = self.prompt_queue.pop_latest("prompt")
+        if not isinstance(prompt, str):
+            raise ValueError("prompt event payload must be a string")
+        self.latest_sampled_event_id = self.prompt_queue.last_sampled_seq_id("prompt")
+        return prompt
+
+    def has_prompt(self) -> bool:
+        return self.prompt_queue.has_events("prompt")
+
+
+class LingBotWorldRealtimeAdapter(BaseRealtimeModelAdapter):
+    def create_state(self) -> LingBotWorldRealtimeState:
+        return LingBotWorldRealtimeState()
+
+    def _state(self, session: GenerateSession) -> LingBotWorldRealtimeState:
+        state = session.adapter_state
+        if not isinstance(state, LingBotWorldRealtimeState):
+            raise TypeError("LingBot realtime adapter state is not initialized")
+        return state
+
+    async def on_init(
+        self,
+        session: GenerateSession,
+        request: RealtimeVideoGenerationsRequest,
+    ) -> None:
+        condition_inputs = request.condition_inputs or {}
+        camera_actions = condition_inputs.get(LINGBOT_CAMERA_ACTIONS_CONDITION)
+        if camera_actions is not None:
+            state = self._state(session)
+            state.receive_camera_action_script(
+                self._validate_camera_actions(camera_actions)
+            )
+
+        await save_realtime_first_frame(session, request)
+
+    @staticmethod
+    def _validate_camera_actions(payload: Any) -> list[list[str]]:
+        if not isinstance(payload, list):
+            raise ValueError("camera_actions event payload must be list[list[str]]")
+        normalized = []
+        for frame_actions in payload:
+            if not isinstance(frame_actions, list):
+                raise ValueError("camera_actions event payload must be list[list[str]]")
+            normalized.append(list(frame_actions))
+        return normalized
+
+    def ingest_event(
+        self,
+        session: GenerateSession,
+        event: RealtimeEvent,
+    ) -> str:
+        state = self._state(session)
+        if event.kind == "camera_actions":
+            return self._ingest_camera_actions(state, event.payload, event.event_id)
+        elif event.kind == "prompt":
+            return self._ingest_prompt(state, event.payload, event.event_id)
+        elif event.kind == COMPOSITE_INPUT_EVENT_KIND:
+            return self._ingest_composite_input(state, event.payload, event.event_id)
+        raise ValueError(f"unsupported event kind: {event.kind}")
+
+    def _ingest_camera_actions(
+        self,
+        state: LingBotWorldRealtimeState,
+        payload: Any,
+        event_id: int | None,
+    ) -> str:
+        return state.receive_camera_control_event_payload(
+            payload,
+            event_id=event_id,
+        )
+
+    def _ingest_prompt(
+        self,
+        state: LingBotWorldRealtimeState,
+        payload: Any,
+        event_id: int | None,
+    ) -> str:
+        prompt = self._validate_prompt_payload(payload)
+        state.receive_prompt(prompt, event_id=event_id)
+        return f"kind=prompt, prompt_len={len(prompt)}"
+
+    @staticmethod
+    def _validate_prompt_payload(payload: Any) -> str:
+        if not isinstance(payload, str) or not payload:
+            raise ValueError("prompt event payload must be a non-empty string")
+        return payload
+
+    def _ingest_composite_input(
+        self,
+        state: LingBotWorldRealtimeState,
+        payload: Any,
+        event_id: int | None,
+    ) -> str:
+        if not isinstance(payload, dict):
+            raise ValueError("composite_input event payload must be a map")
+        input_types = payload.get("input_types")
+        if not isinstance(input_types, list) or not input_types:
+            raise ValueError(
+                "composite_input event payload requires non-empty input_types"
+            )
+
+        parsed_inputs = []
+        for input_type in input_types:
+            if not isinstance(input_type, str) or not input_type:
+                raise ValueError(
+                    "composite_input input_types must contain non-empty strings"
+                )
+            if input_type not in payload:
+                raise ValueError(f"composite_input event payload requires {input_type}")
+            parsed_inputs.append(
+                (
+                    input_type,
+                    self._parse_composite_input_item(
+                        state,
+                        input_type,
+                        payload[input_type],
+                        event_id,
+                    ),
+                )
+            )
+
+        input_logs = []
+        for input_type, parsed_payload in parsed_inputs:
+            input_logs.append(
+                self._ingest_parsed_composite_input_item(
+                    state,
+                    input_type,
+                    parsed_payload,
+                    event_id,
+                )
+            )
+        return f"kind=composite_input, inputs={input_logs}"
+
+    def _parse_composite_input_item(
+        self,
+        state: LingBotWorldRealtimeState,
+        input_type: str,
+        payload: Any,
+        event_id: int | None,
+    ) -> Any:
+        if input_type == "camera_actions":
+            return state.parse_camera_control_event_payload(
+                payload,
+                event_id=event_id,
+            )
+        if input_type == "prompt":
+            return self._validate_prompt_payload(payload)
+        raise ValueError(f"unsupported composite_input type: {input_type}")
+
+    def _ingest_parsed_composite_input_item(
+        self,
+        state: LingBotWorldRealtimeState,
+        input_type: str,
+        parsed_payload: Any,
+        event_id: int | None,
+    ) -> str:
+        if input_type == "camera_actions":
+            return state.receive_parsed_camera_control_event_payload(
+                parsed_payload,
+                event_id=event_id,
+            )
+        if input_type == "prompt":
+            state.receive_prompt(parsed_payload, event_id=event_id)
+            return f"kind=prompt, prompt_len={len(parsed_payload)}"
+        raise ValueError(f"unsupported composite_input type: {input_type}")
+
+    def sample_chunk_inputs(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+        chunk_size: int,
+    ) -> RealtimeChunkInputs:
+        """Samples user inputs (conditions) for the current RealtimeChunk from RealtimeStates"""
+        state = self._state(session)
+        request = session.request
+        if request is None:
+            raise ValueError("realtime request is not initialized")
+
+        prompt_updated = False
+        if chunk.index == 0:
+            prompt = request.prompt
+        elif state.has_prompt():
+            prompt = state.sample_prompt()
+            request.prompt = prompt
+            prompt_updated = True
+        else:
+            prompt = request.prompt
+
+        condition_inputs = {}
+        if prompt_updated:
+            condition_inputs[LINGBOT_PROMPT_UPDATED_CONDITION] = True
+        camera_actions = state.sample_camera_actions(chunk_size)
+        if camera_actions is not None:
+            condition_inputs[LINGBOT_CAMERA_ACTIONS_CONDITION] = camera_actions
+        return RealtimeChunkInputs(prompt=prompt, condition_inputs=condition_inputs)
+
+    def build_sampling_params(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+        chunk_inputs: RealtimeChunkInputs,
+        chunk_size: int,
+    ):
+        request = session.request
+        if request is None:
+            raise ValueError("realtime request is not initialized")
+
+        num_frames = self._condition_num_frames(
+            request=request,
+            server_args=server_args,
+            chunk_size=chunk_size,
+        )
+
+        return build_realtime_sampling_params(
+            chunk.request_id,
+            request=request,
+            chunk_inputs=chunk_inputs,
+            num_frames=num_frames,
+            num_inference_steps=(
+                request.num_inference_steps
+                or LINGBOT_REALTIME_DEFAULT_NUM_INFERENCE_STEPS
+            ),
+            chunk_size=chunk_size,
+        )
+
+    @staticmethod
+    def _condition_num_frames(
+        *,
+        request: RealtimeVideoGenerationsRequest,
+        server_args: ServerArgs | None,
+        chunk_size: int,
+    ) -> int:
+        if server_args is None:
+            return int(request.num_frames or 0)
+
+        # encode one extra blank condition chunk so repeat-last never reuses
+        # the first-frame image mask on later realtime chunks
+        temporal_ratio = int(
+            server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        )
+        required_latent_frames = chunk_size * LINGBOT_REALTIME_MIN_CONDITION_CHUNKS
+        required_num_frames = (required_latent_frames - 1) * temporal_ratio + 1
+        return max(int(request.num_frames or 0), required_num_frames)
+
+    def get_realtime_event_id(self, session: GenerateSession) -> int | None:
+        return self._state(session).latest_sampled_event_id
+
+    def clear_state(self, session: GenerateSession) -> None:
+        state = session.adapter_state
+        if isinstance(state, LingBotWorldRealtimeState):
+            state.clear()

python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py

@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    RealtimeEvent,
+    RealtimeVideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import (
+    BaseRealtimeModelAdapter,
+    RealtimeChunkInputs,
+    build_realtime_sampling_params,
+    save_realtime_first_frame,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base import (
+    normalize_sana_wm_camera_actions,
+    parse_sana_wm_action_string,
+    snap_sana_wm_num_frames,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.self_forcing import (
+    SanaWMSelfForcingSampler,
+)
+from sglang.multimodal_gen.runtime.realtime.states import (
+    RealtimeCameraControlState,
+)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import (
+        GenerateSession,
+        RealtimeChunkContext,
+    )
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
+        OutputBatch,
+    )
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+SANA_WM_DEFAULT_SIZE = "1280x704"
+SANA_WM_DEFAULT_NUM_FRAMES = 1081
+SANA_WM_DEFAULT_FPS = 16
+SANA_WM_DEFAULT_STEPS = 4
+SANA_WM_DEFAULT_GUIDANCE = 1.0
+SANA_WM_CONTROL_PULSE_FRAMES = 8
+
+
+def _normalize_sana_wm_state_actions(actions: list[Any]) -> list[str]:
+    return [str(action).lower() for action in actions]
+
+
+class SanaWMRealtimeAdapterState(RealtimeCameraControlState):
+    def __init__(self):
+        super().__init__(
+            min_pulse_items=SANA_WM_CONTROL_PULSE_FRAMES,
+            script_maxlen=2048,
+            max_transitions=512,
+            normalize_state_actions=_normalize_sana_wm_state_actions,
+        )
+        self.base_condition_inputs: dict[str, Any] = {}
+
+    def clear(self) -> None:
+        super().clear()
+        self.base_condition_inputs.clear()
+
+    def receive_camera_control_event_payload(
+        self,
+        payload: Any,
+        *,
+        event_id: int | None,
+    ) -> str:
+        return super().receive_camera_control_event_payload(
+            payload,
+            event_id=event_id,
+            validate_camera_actions=SanaWMRealtimeAdapter._validate_camera_actions,
+        )
+
+
+class SanaWMRealtimeAdapter(BaseRealtimeModelAdapter):
+    def create_state(self) -> SanaWMRealtimeAdapterState:
+        return SanaWMRealtimeAdapterState()
+
+    def _state(self, session: GenerateSession) -> SanaWMRealtimeAdapterState:
+        state = session.adapter_state
+        if not isinstance(state, SanaWMRealtimeAdapterState):
+            raise TypeError("SANA-WM realtime adapter state is not initialized")
+        return state
+
+    @staticmethod
+    def _validate_camera_actions(payload: Any) -> list[list[str]]:
+        return normalize_sana_wm_camera_actions(
+            payload, error_label="camera_actions event payload"
+        )
+
+    @staticmethod
+    def _raw_frame_count(result: OutputBatch) -> int | None:
+        if result.raw_frame_batches is None:
+            return None
+        return sum(len(frames) for frames in result.raw_frame_batches)
+
+    async def on_init(
+        self,
+        session: GenerateSession,
+        request: RealtimeVideoGenerationsRequest,
+    ) -> None:
+        request.size = request.size or SANA_WM_DEFAULT_SIZE
+        if request.num_frames is not None:
+            request.num_frames = int(request.num_frames)
+        else:
+            # Open-ended session: keep num_frames unset so prepare_next_request
+            # samples uniform action chunks (no front-loaded segmentation), and
+            # flag the stage explicitly via condition_inputs —
+            # build_sampling_params strips None fields, so the per-chunk batch
+            # would otherwise carry the SamplingParams default num_frames.
+            request.condition_inputs = {
+                **(request.condition_inputs or {}),
+                "sana_wm_open_ended": True,
+            }
+        request.fps = int(request.fps or SANA_WM_DEFAULT_FPS)
+        request.num_inference_steps = int(
+            request.num_inference_steps or SANA_WM_DEFAULT_STEPS
+        )
+        request.guidance_scale = float(
+            request.guidance_scale or SANA_WM_DEFAULT_GUIDANCE
+        )
+        if request.negative_prompt is None:
+            request.negative_prompt = ""
+        if request.generator_device is None:
+            request.generator_device = "cuda"
+
+        state = self._state(session)
+        condition_inputs = dict(request.condition_inputs or {})
+        camera_actions = condition_inputs.pop("camera_actions", None)
+        action = condition_inputs.pop("action", None)
+        if camera_actions is not None and action is not None:
+            raise ValueError("pass only one of camera_actions or action")
+        if camera_actions is not None:
+            state.receive_camera_control_event_payload(camera_actions, event_id=None)
+        if action is not None:
+            if not isinstance(action, str) or not action:
+                raise ValueError("action condition input must be a non-empty string")
+            state.receive_camera_action_script(
+                parse_sana_wm_action_string(action), event_id=None
+            )
+        state.base_condition_inputs = condition_inputs
+
+        await save_realtime_first_frame(
+            session,
+            request,
+            required_error="SANA-WM realtime requires first_frame",
+            cache_remote_urls=True,
+        )
+
+    def ingest_event(
+        self,
+        session: GenerateSession,
+        event: RealtimeEvent,
+    ) -> str:
+        state = self._state(session)
+        if event.kind == "camera_actions":
+            return state.receive_camera_control_event_payload(
+                event.payload,
+                event_id=event.event_id,
+            )
+        if event.kind == "action":
+            if not isinstance(event.payload, str) or not event.payload:
+                raise ValueError("action event payload must be a non-empty string")
+            camera_actions = parse_sana_wm_action_string(event.payload)
+            state.receive_camera_action_script(camera_actions, event_id=event.event_id)
+            return f"kind=action, frames={len(camera_actions)}"
+        raise ValueError(f"unsupported event kind: {event.kind}")
+
+    def sample_chunk_inputs(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+        chunk_size: int,
+    ) -> RealtimeChunkInputs:
+        action_chunk_size = self._action_chunk_size(
+            session,
+            server_args,
+            chunk,
+            chunk_size,
+        )
+        state = self._state(session)
+        request = session.request
+        if request is None:
+            raise ValueError("realtime request is not initialized")
+
+        condition_inputs = dict(state.base_condition_inputs) if chunk.index == 0 else {}
+        camera_actions = state.sample_camera_actions(action_chunk_size)
+        if camera_actions is not None:
+            condition_inputs["camera_actions"] = camera_actions
+        return RealtimeChunkInputs(
+            prompt=request.prompt,
+            condition_inputs=condition_inputs,
+        )
+
+    def build_sampling_params(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+        chunk_inputs: RealtimeChunkInputs,
+        chunk_size: int,
+    ):
+        request = session.request
+        if request is None:
+            raise ValueError("realtime request is not initialized")
+
+        return build_realtime_sampling_params(
+            chunk.request_id,
+            request=request,
+            chunk_inputs=chunk_inputs,
+            num_frames=request.num_frames,
+            num_inference_steps=request.num_inference_steps,
+            chunk_size=chunk_size,
+        )
+
+    def _action_chunk_size(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+        chunk_size: int,
+    ) -> int:
+        temporal_compression = int(
+            server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        )
+        # Match action sampling to the latent span used by the batch path. Chunk
+        # 0 may carry a front-loaded remainder, so a fixed nfpb*tc action count
+        # would read static-padded camera poses and drift from batch output.
+        action_chunk_size = chunk_size * temporal_compression
+        req_num_frames = (
+            session.request.num_frames if session.request is not None else None
+        )
+        if req_num_frames is not None:
+            snapped = snap_sana_wm_num_frames(
+                int(req_num_frames), stride=temporal_compression
+            )
+            latent_t = (snapped - 1) // temporal_compression + 1
+            segments = SanaWMSelfForcingSampler.create_autoregressive_segments(
+                latent_t, chunk_size
+            )
+            idx = int(chunk.index)
+            if 0 <= idx and idx + 1 < len(segments):
+                action_chunk_size = (
+                    segments[idx + 1] - segments[idx]
+                ) * temporal_compression
+        return action_chunk_size
+
+    def get_realtime_event_id(self, session: GenerateSession) -> int | None:
+        return self._state(session).latest_sampled_event_id
+
+    def on_chunk_complete(self, session: GenerateSession, result: OutputBatch) -> None:
+        if session.request is not None and self._raw_frame_count(result) == 0:
+            session.request.max_chunks = session.generate_chunk_cnt + 1
+        session.generate_chunk_completed()
+
+    def clear_state(self, session: GenerateSession) -> None:
+        state = session.adapter_state
+        if isinstance(state, SanaWMRealtimeAdapterState):
+            state.clear()

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

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    RealtimeVideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.realtime.session import (
+    RealtimeSession,
+)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import (
+        BaseRealtimeModelAdapter,
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class RealtimeChunkContext:
+    session_id: str
+    index: int
+    request_id: str
+
+
+class GenerateSession:
+    """A realtime generation session"""
+
+    def __init__(self):
+        self.id = uuid4().hex
+        self.request: RealtimeVideoGenerationsRequest | None = None
+        self.input_temp_dir: str | None = None
+        self.generate_chunk_cnt = 0
+        self.current_chunk: RealtimeChunkContext | None = None
+        self.realtime_session = RealtimeSession()
+        self.adapter: BaseRealtimeModelAdapter | None = None
+        self.adapter_state: Any = None
+        self.output_pace_next_send_at: float | None = None
+        self.output_pace_last_event_id: int | None = None
+
+    def set_adapter(self, adapter: BaseRealtimeModelAdapter):
+        self.adapter = adapter
+        self.adapter_state = adapter.create_state()
+
+    def set_request(self, request: RealtimeVideoGenerationsRequest):
+        self.request = request
+
+    def dispose(self):
+        if self.adapter is not None:
+            self.adapter.dispose(self)
+        self.request = None
+        self.input_temp_dir = None
+        self.generate_chunk_cnt = 0
+        self.current_chunk = None
+        self.adapter = None
+        self.adapter_state = None
+        self.output_pace_next_send_at = None
+        self.output_pace_last_event_id = None
+        self.realtime_session.dispose()
+
+    def new_chunk(self) -> RealtimeChunkContext:
+        if self.current_chunk is not None:
+            raise RuntimeError("previous realtime chunk is still active")
+        chunk = RealtimeChunkContext(
+            session_id=self.id,
+            index=self.generate_chunk_cnt,
+            request_id=f"{self.id}_{uuid4().hex}",
+        )
+        self.current_chunk = chunk
+        return chunk
+
+    def generate_chunk_completed(self):
+        self.generate_chunk_cnt += 1
+        self.current_chunk = None
+
+    def reached_max_chunks(self) -> bool:
+        return (
+            self.request is not None
+            and self.request.max_chunks is not None
+            and self.generate_chunk_cnt >= self.request.max_chunks
+        )

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

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import hashlib
+import os
+import tempfile
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from fastapi import WebSocket
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    RealtimeEvent,
+    RealtimeVideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_output_adapter import (
+    RawRGBRealtimeOutputAdapter,
+    RealtimeFrameSendStats,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
+    build_sampling_params,
+    save_image_to_path,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    prepare_request,
+)
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import (
+        GenerateSession,
+        RealtimeChunkContext,
+    )
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
+        OutputBatch,
+        Req,
+    )
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+@dataclass(slots=True)
+class RealtimeChunkInputs:
+    """Sampled from realtime control state, consumed by the Req"""
+
+    prompt: str
+    condition_inputs: dict[str, Any] = field(default_factory=dict)
+
+
+async def save_realtime_first_frame(
+    session: GenerateSession,
+    request: RealtimeVideoGenerationsRequest,
+    *,
+    required_error: str | None = None,
+    cache_remote_urls: bool = False,
+) -> None:
+    first_frame = request.first_frame
+    if first_frame is None:
+        if required_error is not None:
+            raise ValueError(required_error)
+        return
+
+    server_args = get_global_server_args()
+    if server_args.input_save_path is not None:
+        uploads_dir = server_args.input_save_path
+        os.makedirs(uploads_dir, exist_ok=True)
+    else:
+        if session.input_temp_dir is None:
+            session.input_temp_dir = tempfile.mkdtemp(prefix="sglang_input_")
+        uploads_dir = session.input_temp_dir
+
+    if (
+        cache_remote_urls
+        and isinstance(first_frame, str)
+        and first_frame.lower().startswith(("http://", "https://"))
+    ):
+        suffix = os.path.splitext(first_frame.split("?", 1)[0])[1]
+        digest = hashlib.sha256(first_frame.encode("utf-8")).hexdigest()[:16]
+        target_path = os.path.join(uploads_dir, f"realtime_ref_{digest}{suffix}")
+        if os.path.exists(target_path):
+            request.first_frame = target_path
+            return
+    else:
+        target_path = os.path.join(uploads_dir, f"{session.id}_first_frame")
+
+    request.first_frame = await save_image_to_path(first_frame, target_path)
+
+
+def build_realtime_sampling_params(
+    request_id: str,
+    *,
+    request: RealtimeVideoGenerationsRequest,
+    chunk_inputs: RealtimeChunkInputs,
+    num_frames: int | None,
+    num_inference_steps: int | None,
+    chunk_size: int,
+):
+    return build_sampling_params(
+        request_id,
+        prompt=chunk_inputs.prompt,
+        size=request.size,
+        num_frames=num_frames,
+        fps=request.fps,
+        image_path=request.first_frame,
+        output_file_name=request_id,
+        save_output=False,
+        seed=request.seed,
+        generator_device=request.generator_device,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=request.guidance_scale,
+        guidance_scale_2=request.guidance_scale_2,
+        negative_prompt=request.negative_prompt,
+        enable_teacache=request.enable_teacache,
+        enable_frame_interpolation=request.enable_frame_interpolation,
+        frame_interpolation_exp=request.frame_interpolation_exp,
+        frame_interpolation_scale=request.frame_interpolation_scale,
+        frame_interpolation_model_path=request.frame_interpolation_model_path,
+        enable_upscaling=request.enable_upscaling,
+        upscaling_model_path=request.upscaling_model_path,
+        upscaling_scale=request.upscaling_scale,
+        diffusers_kwargs=request.diffusers_kwargs,
+        profile=request.profile,
+        num_profiled_timesteps=request.num_profiled_timesteps,
+        profile_all_stages=request.profile_all_stages,
+        perf_dump_path=request.perf_dump_path,
+        output_path=request.output_path,
+        output_compression=request.output_compression,
+        output_quality=request.output_quality,
+        condition_inputs=chunk_inputs.condition_inputs,
+        realtime_chunk_size=chunk_size,
+    )
+
+
+class BaseRealtimeModelAdapter:
+    def __init__(self):
+        self.output_adapter = RawRGBRealtimeOutputAdapter()
+
+    async def on_init(
+        self,
+        session: GenerateSession,
+        request: RealtimeVideoGenerationsRequest,
+    ) -> None:
+        raise NotImplementedError
+
+    def create_state(self) -> Any:
+        """create a state for managing runtime states"""
+        raise NotImplementedError
+
+    def ingest_event(
+        self,
+        session: GenerateSession,
+        event: RealtimeEvent,
+    ) -> str:
+        """
+        Ingest a realtime endpoint event and install it into the model's realtime control queues
+        """
+        raise NotImplementedError
+
+    async def wait_for_next_chunk(self, session: GenerateSession) -> None:
+        del session
+
+    def get_chunk_size(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+    ) -> int:
+        del session, chunk
+        arch_config = server_args.pipeline_config.dit_config.arch_config
+        return int(getattr(arch_config, "num_frames_per_block", 3))
+
+    def sample_chunk_inputs(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+        chunk_size: int,
+    ) -> RealtimeChunkInputs:
+        raise NotImplementedError
+
+    def build_sampling_params(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+        chunk_inputs: RealtimeChunkInputs,
+        chunk_size: int,
+    ):
+        raise NotImplementedError
+
+    def get_realtime_event_id(self, session: GenerateSession) -> int | None:
+        del session
+        return None
+
+    def prepare_next_request(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+    ) -> Req:
+        chunk_size = self.get_chunk_size(session, server_args, chunk)
+        chunk_inputs = self.sample_chunk_inputs(
+            session,
+            server_args,
+            chunk,
+            chunk_size,
+        )
+        sampling_params = self.build_sampling_params(
+            session,
+            server_args,
+            chunk,
+            chunk_inputs,
+            chunk_size,
+        )
+        batch = prepare_request(
+            server_args=server_args,
+            sampling_params=sampling_params,
+        )
+        self.apply_realtime_request_fields(
+            batch,
+            session,
+            chunk,
+            event_id=self.get_realtime_event_id(session),
+        )
+        return batch
+
+    def apply_realtime_request_fields(
+        self,
+        batch: Req,
+        session: GenerateSession,
+        chunk: RealtimeChunkContext,
+        *,
+        event_id: int | None,
+    ) -> None:
+        batch.realtime_session_id = session.id
+        batch.return_raw_frames = True
+        batch.block_idx = chunk.index
+        batch.realtime_event_id = event_id
+        if session.request is None:
+            return
+        batch.realtime_output_format = session.request.realtime_output_format
+        batch.realtime_preview_max_width = session.request.realtime_preview_max_width
+        batch.realtime_output_pacing = bool(session.request.realtime_output_pacing)
+        batch.realtime_causal_sink_size = session.request.realtime_causal_sink_size
+        batch.realtime_causal_kv_cache_num_frames = (
+            session.request.realtime_causal_kv_cache_num_frames
+        )
+
+    async def send_output(
+        self,
+        ws: WebSocket,
+        session: GenerateSession,
+        result: OutputBatch,
+        batch: Req,
+    ) -> RealtimeFrameSendStats:
+        """send the generate output (usually frames) back via websocket"""
+        return await self.output_adapter.send(ws, session, result, batch)
+
+    def on_chunk_complete(self, session: GenerateSession, result: OutputBatch) -> None:
+        del result
+        session.generate_chunk_completed()
+
+    def clear_state(self, session: GenerateSession) -> None:
+        del session
+
+    def dispose(self, session: GenerateSession) -> None:
+        self.clear_state(session)
+        self.output_adapter.reset()

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

@@ -0,0 +1,610 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+import io
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypedDict
+
+import msgspec.msgpack
+from fastapi import WebSocket
+from PIL import Image
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.timer import (
+    RealtimeStageTimer,
+)
+from sglang.multimodal_gen.runtime.utils.realtime_video import (
+    JPEG_FRAME_CONTENT_TYPE,
+    RAW_RGB_CHANNELS,
+    RAW_RGB_CONTENT_TYPE,
+    WEBP_FRAME_CONTENT_TYPE,
+)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import (
+        GenerateSession,
+    )
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
+        OutputBatch,
+        Req,
+    )
+
+
+class RealtimeFrameBatchHeader(TypedDict, total=False):
+    type: str
+    request_id: str
+    chunk_index: int
+    content_type: str
+    num_frames: int
+    total_size: int
+    format: str
+    width: int
+    height: int
+    channels: int
+    bytes_per_frame: int
+    raw_size: int
+    encoding: str
+    delta_reference: str
+    payload_lengths: list[int]
+    event_id: int
+    frame_batch_index: int
+    num_frame_batches: int
+    is_final_frame_batch: bool
+
+
+class RealtimeFrameBatchMessage(RealtimeFrameBatchHeader, total=False):
+    payload: bytes
+
+
+class RealtimeFrameSendStats(TypedDict):
+    header_pack_ms: float
+    header_write_ms: float
+    raw_payload_build_ms: float
+    raw_write_ms: float
+    ws_write_ms: float
+    pace_wait_ms: float
+    raw_bytes: int
+    ws_payload_bytes: int
+    num_frames: int
+    num_batches: int
+    frame_shape: tuple[int, int, int] | None
+    content_type: str
+
+
+def empty_frame_send_stats(content_type: str = "") -> RealtimeFrameSendStats:
+    return {
+        "header_pack_ms": 0.0,
+        "header_write_ms": 0.0,
+        "raw_payload_build_ms": 0.0,
+        "raw_write_ms": 0.0,
+        "ws_write_ms": 0.0,
+        "pace_wait_ms": 0.0,
+        "raw_bytes": 0,
+        "ws_payload_bytes": 0,
+        "num_frames": 0,
+        "num_batches": 0,
+        "frame_shape": None,
+        "content_type": content_type,
+    }
+
+
+def _raw_rgb_frame_metadata(batch: Req) -> dict[str, int | str]:
+    frame_width = batch.width
+    frame_height = batch.height
+    if frame_width is None or frame_height is None:
+        return {}
+
+    frame_width = int(frame_width)
+    frame_height = int(frame_height)
+    if batch.enable_upscaling:
+        upscaling_scale = int(batch.upscaling_scale or 1)
+        frame_width *= upscaling_scale
+        frame_height *= upscaling_scale
+
+    return {
+        "format": "rgb24",
+        "width": frame_width,
+        "height": frame_height,
+        "channels": RAW_RGB_CHANNELS,
+        "bytes_per_frame": frame_width * frame_height * RAW_RGB_CHANNELS,
+    }
+
+
+def _frame_shape_from_metadata(
+    metadata: dict[str, int | str] | None,
+) -> tuple[int, int, int] | None:
+    if not metadata:
+        return None
+    return (
+        int(metadata["height"]),
+        int(metadata["width"]),
+        int(metadata["channels"]),
+    )
+
+
+RAW_RGB_FRAMES_PER_WS_MESSAGE = 16
+ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE = 6
+FRAME_BATCH_PACK_OFFLOAD_BYTES = 64 * 1024
+WEBP_DEFAULT_QUALITY = 90
+JPEG_DEFAULT_QUALITY = 95
+JPEG_SUBSAMPLING = 0
+RAW_LOSSLESS_OUTPUT_FORMAT = "raw"
+ENCODED_PREVIEW_FORMATS = {"webp", "jpeg"}
+
+
+@dataclass(frozen=True)
+class _TransportPayload:
+    content_type: str
+    payload: bytes
+    metadata: dict[str, int | str | bool | list[int]]
+
+
+def _split_frame_batch(
+    frames: list[bytes],
+    frames_per_message: int = RAW_RGB_FRAMES_PER_WS_MESSAGE,
+) -> list[list[bytes]]:
+    if not frames:
+        return [frames]
+    return [
+        frames[i : i + frames_per_message]
+        for i in range(0, len(frames), frames_per_message)
+    ]
+
+
+def _encode_rgb_frame_to_webp(
+    frame: bytes,
+    *,
+    width: int,
+    height: int,
+    quality: int,
+    preview_max_width: int | None,
+) -> bytes:
+    buffer = io.BytesIO()
+    image = _resize_preview_image(
+        Image.frombytes("RGB", (width, height), frame),
+        preview_max_width=preview_max_width,
+    )
+    image.save(
+        buffer,
+        format="WEBP",
+        quality=quality,
+        method=0,
+    )
+    return buffer.getvalue()
+
+
+def _encode_rgb_frame_to_jpeg(
+    frame: bytes,
+    *,
+    width: int,
+    height: int,
+    quality: int,
+    preview_max_width: int | None,
+) -> bytes:
+    buffer = io.BytesIO()
+    image = _resize_preview_image(
+        Image.frombytes("RGB", (width, height), frame),
+        preview_max_width=preview_max_width,
+    )
+    image.save(
+        buffer,
+        format="JPEG",
+        quality=quality,
+        subsampling=JPEG_SUBSAMPLING,
+    )
+    return buffer.getvalue()
+
+
+def _preview_dimensions(
+    *,
+    width: int,
+    height: int,
+    preview_max_width: int | None,
+) -> tuple[int, int]:
+    if (
+        preview_max_width is None
+        or preview_max_width <= 0
+        or width <= preview_max_width
+    ):
+        return width, height
+    preview_width = int(preview_max_width)
+    preview_height = max(1, round(height * preview_width / width))
+    return preview_width, preview_height
+
+
+def _resize_preview_image(
+    image: Image.Image,
+    *,
+    preview_max_width: int | None,
+) -> Image.Image:
+    width, height = image.size
+    preview_width, preview_height = _preview_dimensions(
+        width=width,
+        height=height,
+        preview_max_width=preview_max_width,
+    )
+    if (preview_width, preview_height) == image.size:
+        return image
+    return image.resize((preview_width, preview_height), Image.Resampling.BICUBIC)
+
+
+def _pack_frame_batch_message(
+    header: RealtimeFrameBatchHeader,
+    payload: bytes,
+) -> bytes:
+    message: RealtimeFrameBatchMessage = {
+        **header,
+        "type": "frame_batch",
+        "payload": payload,
+    }
+    return msgspec.msgpack.encode(message)
+
+
+def _pack_frame_batch_header(header: RealtimeFrameBatchHeader) -> bytes:
+    return msgspec.msgpack.encode(header)
+
+
+def _build_transport_payload(
+    transport_frames: list[bytes],
+    *,
+    content_type: str,
+    metadata: dict[str, int | str],
+    output_format: str | None,
+    transport_quality: int | None,
+    preview_max_width: int | None,
+) -> _TransportPayload:
+    payload_content_type = content_type
+    payload_metadata: dict[str, int | str | bool | list[int]] = {}
+    raw_payload = b""
+
+    if (
+        output_format in ENCODED_PREVIEW_FORMATS
+        and content_type == RAW_RGB_CONTENT_TYPE
+        and transport_frames
+    ):
+        if output_format == "webp":
+            encoded_frames = [
+                _encode_rgb_frame_to_webp(
+                    frame,
+                    width=int(metadata["width"]),
+                    height=int(metadata["height"]),
+                    quality=int(transport_quality or WEBP_DEFAULT_QUALITY),
+                    preview_max_width=preview_max_width,
+                )
+                for frame in transport_frames
+            ]
+            payload_content_type = WEBP_FRAME_CONTENT_TYPE
+        else:
+            encoded_frames = [
+                _encode_rgb_frame_to_jpeg(
+                    frame,
+                    width=int(metadata["width"]),
+                    height=int(metadata["height"]),
+                    quality=int(transport_quality or JPEG_DEFAULT_QUALITY),
+                    preview_max_width=preview_max_width,
+                )
+                for frame in transport_frames
+            ]
+            payload_content_type = JPEG_FRAME_CONTENT_TYPE
+        raw_payload = b"".join(encoded_frames)
+        preview_width, preview_height = _preview_dimensions(
+            width=int(metadata["width"]),
+            height=int(metadata["height"]),
+            preview_max_width=preview_max_width,
+        )
+        payload_metadata = {
+            "format": output_format,
+            "encoding": output_format,
+            "source_width": int(metadata["width"]),
+            "source_height": int(metadata["height"]),
+            "preview_width": preview_width,
+            "preview_height": preview_height,
+            "width": preview_width,
+            "height": preview_height,
+            "payload_lengths": [len(frame) for frame in encoded_frames],
+        }
+    elif content_type == RAW_RGB_CONTENT_TYPE and transport_frames:
+        raw_payload = b"".join(transport_frames)
+        payload_metadata = {
+            "raw_size": len(raw_payload),
+            "encoding": RAW_LOSSLESS_OUTPUT_FORMAT,
+        }
+    else:
+        raw_payload = b"".join(transport_frames)
+
+    return _TransportPayload(
+        content_type=payload_content_type,
+        payload=raw_payload,
+        metadata=payload_metadata,
+    )
+
+
+def _should_build_payload_off_loop(
+    *,
+    content_type: str,
+    output_format: str | None,
+    transport_frames: list[bytes],
+) -> bool:
+    if content_type != RAW_RGB_CONTENT_TYPE or not transport_frames:
+        return False
+    return output_format in ENCODED_PREVIEW_FORMATS or output_format is None
+
+
+def _is_encoded_preview_transport(
+    *,
+    content_type: str,
+    output_format: str | None,
+) -> bool:
+    return (
+        output_format in ENCODED_PREVIEW_FORMATS
+        and content_type == RAW_RGB_CONTENT_TYPE
+    )
+
+
+async def _build_encoded_preview_payloads(
+    split_batches: list[list[bytes]],
+    *,
+    content_type: str,
+    metadata: dict[str, int | str],
+    output_format: str,
+    transport_quality: int | None,
+    preview_max_width: int | None,
+    event_id: int | None,
+) -> list[_TransportPayload]:
+    return list(
+        await asyncio.gather(
+            *(
+                _build_encoded_preview_payload(
+                    transport_frames,
+                    metadata=metadata,
+                    output_format=output_format,
+                    transport_quality=transport_quality,
+                    preview_max_width=preview_max_width,
+                )
+                for transport_frames in split_batches
+            )
+        )
+    )
+
+
+async def _build_encoded_preview_payload(
+    transport_frames: list[bytes],
+    *,
+    metadata: dict[str, int | str],
+    output_format: str,
+    transport_quality: int | None,
+    preview_max_width: int | None,
+) -> _TransportPayload:
+    width = int(metadata["width"])
+    height = int(metadata["height"])
+    if output_format == "webp":
+        encoded_frames = list(
+            await asyncio.gather(
+                *(
+                    asyncio.to_thread(
+                        _encode_rgb_frame_to_webp,
+                        frame,
+                        width=width,
+                        height=height,
+                        quality=int(transport_quality or WEBP_DEFAULT_QUALITY),
+                        preview_max_width=preview_max_width,
+                    )
+                    for frame in transport_frames
+                )
+            )
+        )
+        payload_content_type = WEBP_FRAME_CONTENT_TYPE
+    else:
+        encoded_frames = list(
+            await asyncio.gather(
+                *(
+                    asyncio.to_thread(
+                        _encode_rgb_frame_to_jpeg,
+                        frame,
+                        width=width,
+                        height=height,
+                        quality=int(transport_quality or JPEG_DEFAULT_QUALITY),
+                        preview_max_width=preview_max_width,
+                    )
+                    for frame in transport_frames
+                )
+            )
+        )
+        payload_content_type = JPEG_FRAME_CONTENT_TYPE
+
+    preview_width, preview_height = _preview_dimensions(
+        width=width,
+        height=height,
+        preview_max_width=preview_max_width,
+    )
+    return _TransportPayload(
+        content_type=payload_content_type,
+        payload=b"".join(encoded_frames),
+        metadata={
+            "format": output_format,
+            "encoding": output_format,
+            "source_width": width,
+            "source_height": height,
+            "preview_width": preview_width,
+            "preview_height": preview_height,
+            "width": preview_width,
+            "height": preview_height,
+            "payload_lengths": [len(frame) for frame in encoded_frames],
+        },
+    )
+
+
+class RawRGBRealtimeOutputAdapter:
+    """send raw RGB over WebSocket using lossless transport"""
+
+    def __init__(self) -> None:
+        pass
+
+    def reset(self) -> None:
+        pass
+
+    async def send(
+        self,
+        ws: WebSocket,
+        session: GenerateSession,
+        result: OutputBatch,
+        batch: Req,
+    ) -> RealtimeFrameSendStats:
+        """send frames through ws"""
+        content_type = result.raw_frame_content_type
+        if result.raw_frame_batches is None:
+            return empty_frame_send_stats(content_type)
+        if batch.block_idx == 0:
+            self.reset()
+
+        frame_metadata = (
+            result.raw_frame_metadata or _raw_rgb_frame_metadata(batch)
+            if content_type == RAW_RGB_CONTENT_TYPE
+            else {}
+        )
+        output_format = getattr(batch, "realtime_output_format", None)
+        preview_max_width = getattr(batch, "realtime_preview_max_width", None)
+        stats = await self._send_frame_batches(
+            ws,
+            result.raw_frame_batches,
+            content_type=content_type,
+            chunk_index_start=batch.block_idx,
+            request_id=batch.request_id,
+            event_id=getattr(batch, "realtime_event_id", None),
+            frame_metadata=frame_metadata,
+            output_format=output_format,
+            transport_quality=getattr(batch, "output_compression", None),
+            preview_max_width=preview_max_width,
+        )
+        stats["frame_shape"] = _frame_shape_from_metadata(frame_metadata)
+        return stats
+
+    async def _send_frame_batches(
+        self,
+        ws: WebSocket,
+        frame_batches: list[list[bytes]],
+        *,
+        content_type: str,
+        chunk_index_start: int,
+        request_id: str,
+        event_id: int | None = None,
+        frame_metadata: dict[str, int | str] | None = None,
+        output_format: str | None = None,
+        transport_quality: int | None = None,
+        preview_max_width: int | None = None,
+    ) -> RealtimeFrameSendStats:
+        chunk_index = chunk_index_start
+        metadata = frame_metadata or {}
+        stats = empty_frame_send_stats(content_type)
+        for frames in frame_batches:
+            split_batches = (
+                _split_frame_batch(frames, ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE)
+                if _is_encoded_preview_transport(
+                    content_type=content_type,
+                    output_format=output_format,
+                )
+                else (
+                    _split_frame_batch(frames)
+                    if content_type == RAW_RGB_CONTENT_TYPE
+                    else [frames]
+                )
+            )
+            num_frame_batches = len(split_batches)
+            encoded_preview_payloads: list[_TransportPayload] | None = None
+            if _is_encoded_preview_transport(
+                content_type=content_type,
+                output_format=output_format,
+            ):
+                timer = RealtimeStageTimer()
+                encoded_preview_payloads = await _build_encoded_preview_payloads(
+                    split_batches,
+                    content_type=content_type,
+                    metadata=metadata,
+                    output_format=output_format,
+                    transport_quality=transport_quality,
+                    preview_max_width=preview_max_width,
+                    event_id=event_id,
+                )
+                stats["raw_payload_build_ms"] += timer.mark_ms()
+            for frame_batch_index, transport_frames in enumerate(split_batches):
+                timer = RealtimeStageTimer()
+                transport_metadata = metadata
+                if encoded_preview_payloads is not None:
+                    transport_payload = encoded_preview_payloads[frame_batch_index]
+                else:
+                    if _should_build_payload_off_loop(
+                        content_type=content_type,
+                        output_format=output_format,
+                        transport_frames=transport_frames,
+                    ):
+                        transport_payload = await asyncio.to_thread(
+                            _build_transport_payload,
+                            transport_frames,
+                            content_type=content_type,
+                            metadata=metadata,
+                            output_format=output_format,
+                            transport_quality=transport_quality,
+                            preview_max_width=preview_max_width,
+                        )
+                    else:
+                        transport_payload = _build_transport_payload(
+                            transport_frames,
+                            content_type=content_type,
+                            metadata=metadata,
+                            output_format=output_format,
+                            transport_quality=transport_quality,
+                            preview_max_width=preview_max_width,
+                        )
+                    stats["raw_payload_build_ms"] += timer.mark_ms()
+
+                header: RealtimeFrameBatchHeader = {
+                    "type": "frame_batch_header",
+                    "request_id": request_id,
+                    "chunk_index": chunk_index,
+                    "content_type": transport_payload.content_type,
+                    "num_frames": len(transport_frames),
+                    "total_size": len(transport_payload.payload),
+                    "frame_batch_index": frame_batch_index,
+                    "num_frame_batches": num_frame_batches,
+                    "is_final_frame_batch": frame_batch_index == num_frame_batches - 1,
+                }
+                if event_id is not None:
+                    header["event_id"] = event_id
+                header.update(transport_metadata)
+                header.update(transport_payload.metadata)
+
+                if len(transport_payload.payload) >= FRAME_BATCH_PACK_OFFLOAD_BYTES:
+                    header_payload = _pack_frame_batch_header(header)
+                    stats["header_pack_ms"] += timer.mark_ms()
+
+                    await ws.send_bytes(header_payload)
+                    stats["header_write_ms"] += timer.mark_ms()
+
+                    await ws.send_bytes(transport_payload.payload)
+                    stats["raw_write_ms"] += timer.mark_ms()
+
+                    stats["ws_payload_bytes"] += len(header_payload) + len(
+                        transport_payload.payload
+                    )
+                else:
+                    message_payload = _pack_frame_batch_message(
+                        header,
+                        transport_payload.payload,
+                    )
+                    stats["header_pack_ms"] += timer.mark_ms()
+
+                    stats["header_write_ms"] += timer.mark_ms()
+                    await ws.send_bytes(message_payload)
+                    stats["raw_write_ms"] += timer.mark_ms()
+
+                    stats["ws_payload_bytes"] += len(message_payload)
+
+                stats["raw_bytes"] += sum(len(frame) for frame in transport_frames)
+                stats["num_frames"] += len(transport_frames)
+                stats["num_batches"] += 1
+                stats["content_type"] = transport_payload.content_type
+            chunk_index += 1
+
+        stats["ws_write_ms"] = stats["header_write_ms"] + stats["raw_write_ms"]
+        return stats

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

@@ -0,0 +1,500 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import asyncio
+import shutil
+import time
+from typing import TYPE_CHECKING
+
+import msgspec.msgpack
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    RealtimeEvent,
+    RealtimeVideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import (
+    GenerateSession,
+    RealtimeChunkContext,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_output_adapter import (
+    RealtimeFrameSendStats,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import (
+    get_realtime_model_adapter,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.timer import (
+    RealtimeStageTimer,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
+    process_generation_batch,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    ReleaseRealtimeSessionReq,
+)
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+
+logger = init_logger(__name__)
+router = APIRouter(prefix="/v1/realtime_video", tags=["realtime"])
+_ACTIVE_SESSION_IDS: set[str] = set()
+_ACTIVE_SESSION_WAIT_SECONDS = 1.0
+_ACTIVE_SESSION_WAIT_INTERVAL_SECONDS = 0.1
+
+
+def _transport_ms(value: float) -> int:
+    return max(0, int(value + 0.5))
+
+
+async def _wait_for_active_session_slot(
+    *,
+    timeout_s: float = _ACTIVE_SESSION_WAIT_SECONDS,
+    interval_s: float = _ACTIVE_SESSION_WAIT_INTERVAL_SECONDS,
+) -> bool:
+    deadline = time.monotonic() + timeout_s
+    while _ACTIVE_SESSION_IDS and time.monotonic() < deadline:
+        await asyncio.sleep(interval_s)
+    return not _ACTIVE_SESSION_IDS
+
+
+def _log_realtime_chunk_timing(
+    session: GenerateSession,
+    chunk: RealtimeChunkContext,
+    batch: "Req",
+    request_prepare_ms: float,
+    scheduler_forward_ms: float,
+    chunk_total_ms: float,
+    send_stats: RealtimeFrameSendStats,
+) -> None:
+    logger.info(
+        "realtime chunk timing: session_id=%s request_id=%s "
+        "chunk_idx=%s event_id=%s condition_kinds=%s "
+        "request_prepare=%.2fms scheduler_forward=%.2fms "
+        "output_pace=%.2fms "
+        "header_pack=%.2fms "
+        "header_write=%.2fms raw_payload_build=%.2fms raw_write=%.2fms "
+        "ws_write=%.2fms chunk_total=%.2fms batches=%d frames=%d "
+        "frame_shape=%s raw_bytes=%d ws_payload_bytes=%d content_type=%s",
+        session.id,
+        chunk.request_id,
+        batch.block_idx,
+        getattr(batch, "realtime_event_id", None),
+        sorted(batch.condition_inputs) if batch.condition_inputs else [],
+        request_prepare_ms,
+        scheduler_forward_ms,
+        send_stats["pace_wait_ms"],
+        send_stats["header_pack_ms"],
+        send_stats["header_write_ms"],
+        send_stats["raw_payload_build_ms"],
+        send_stats["raw_write_ms"],
+        send_stats["ws_write_ms"],
+        chunk_total_ms,
+        send_stats["num_batches"],
+        send_stats["num_frames"],
+        send_stats["frame_shape"],
+        send_stats["raw_bytes"],
+        send_stats["ws_payload_bytes"],
+        send_stats["content_type"],
+    )
+
+
+async def _send_realtime_chunk_stats(
+    ws: WebSocket,
+    session: GenerateSession,
+    chunk: RealtimeChunkContext,
+    batch: "Req",
+    request_prepare_ms: float,
+    scheduler_forward_ms: float,
+    chunk_total_ms: float,
+    send_stats: RealtimeFrameSendStats,
+) -> None:
+    await ws.send_bytes(
+        msgspec.msgpack.encode(
+            {
+                "type": "chunk_stats",
+                "session_id": session.id,
+                "request_id": chunk.request_id,
+                "chunk_index": batch.block_idx,
+                "event_id": getattr(batch, "realtime_event_id", None),
+                "request_prepare_ms": _transport_ms(request_prepare_ms),
+                "scheduler_forward_ms": _transport_ms(scheduler_forward_ms),
+                "pace_wait_ms": _transport_ms(send_stats["pace_wait_ms"]),
+                "header_write_ms": _transport_ms(send_stats["header_write_ms"]),
+                "raw_payload_build_ms": _transport_ms(
+                    send_stats["raw_payload_build_ms"]
+                ),
+                "raw_write_ms": _transport_ms(send_stats["raw_write_ms"]),
+                "ws_write_ms": _transport_ms(send_stats["ws_write_ms"]),
+                "chunk_total_ms": _transport_ms(chunk_total_ms),
+                "num_batches": send_stats["num_batches"],
+                "num_frames": send_stats["num_frames"],
+                "raw_bytes": send_stats["raw_bytes"],
+                "ws_payload_bytes": send_stats["ws_payload_bytes"],
+                "content_type": send_stats["content_type"],
+            }
+        )
+    )
+
+
+async def _generate_loop(ws: WebSocket, session: GenerateSession):
+    adapter = session.adapter
+    if adapter is None:
+        raise ValueError("realtime adapter is not initialized")
+
+    pending_send_task = None
+    while not session.reached_max_chunks():
+        try:
+            if pending_send_task is not None and pending_send_task.done():
+                await pending_send_task
+                pending_send_task = None
+
+            # send to scheduler and generate video chunk
+            server_args = get_global_server_args()
+
+            await adapter.wait_for_next_chunk(session)
+
+            timer = RealtimeStageTimer()
+            chunk_started = time.perf_counter()
+
+            chunk = session.new_chunk()
+            batch = adapter.prepare_next_request(
+                session,
+                server_args,
+                chunk,
+            )
+            if batch.condition_inputs:
+                logger.debug(
+                    "consume realtime conditions, session_id=%s, block_idx=%s, kinds=%s",
+                    session.id,
+                    batch.block_idx,
+                    sorted(batch.condition_inputs),
+                )
+            request_prepare_ms = timer.mark_ms()
+
+            _, result = await process_generation_batch(async_scheduler_client, batch)
+            scheduler_forward_ms = timer.mark_ms()
+
+            # finish
+            adapter.on_chunk_complete(session, result)
+            if pending_send_task is not None:
+                await pending_send_task
+            if getattr(batch, "realtime_output_pacing", False):
+                await _send_output_and_log(
+                    ws,
+                    session,
+                    chunk,
+                    batch,
+                    result,
+                    request_prepare_ms,
+                    scheduler_forward_ms,
+                    chunk_started,
+                )
+                pending_send_task = None
+            else:
+                pending_send_task = asyncio.create_task(
+                    _send_output_and_log(
+                        ws,
+                        session,
+                        chunk,
+                        batch,
+                        result,
+                        request_prepare_ms,
+                        scheduler_forward_ms,
+                        chunk_started,
+                    )
+                )
+
+        except asyncio.CancelledError:
+            if pending_send_task is not None:
+                pending_send_task.cancel()
+                await _await_realtime_task(pending_send_task)
+            logger.info("generation completed, session_id=%s", session.id)
+            break
+        except WebSocketDisconnect:
+            if pending_send_task is not None:
+                pending_send_task.cancel()
+                await _await_realtime_task(pending_send_task)
+            logger.info(
+                "client disconnected during generation, session_id=%s", session.id
+            )
+            break
+        except Exception as e:
+            if pending_send_task is not None:
+                pending_send_task.cancel()
+                await _await_realtime_task(pending_send_task)
+            err_msg = str(e).splitlines()[0]
+            logger.error("error during generate loop: %s", err_msg)
+            try:
+                await write_error_msg(f"error during generate loop: {err_msg}", ws)
+            except Exception as send_error:
+                logger.error(
+                    "error during sending complete msg: %s",
+                    send_error,
+                )
+            break
+    else:
+        if pending_send_task is not None:
+            await pending_send_task
+        logger.info(
+            "generation reached max chunks, session_id=%s, max_chunks=%s",
+            session.id,
+            session.request.max_chunks if session.request is not None else None,
+        )
+
+
+async def _send_output_and_log(
+    ws: WebSocket,
+    session: GenerateSession,
+    chunk: RealtimeChunkContext,
+    batch: "Req",
+    result,
+    request_prepare_ms: float,
+    scheduler_forward_ms: float,
+    chunk_started: float,
+) -> RealtimeFrameSendStats:
+    if session.adapter is None:
+        raise ValueError("realtime adapter is not initialized")
+    pace_wait_ms = await _wait_for_realtime_output_slot(session, batch, result)
+    send_stats = await session.adapter.send_output(
+        ws,
+        session,
+        result,
+        batch,
+    )
+    send_stats["pace_wait_ms"] = pace_wait_ms
+    chunk_total_ms = (time.perf_counter() - chunk_started) * 1000
+    _log_realtime_chunk_timing(
+        session,
+        chunk,
+        batch,
+        request_prepare_ms,
+        scheduler_forward_ms,
+        chunk_total_ms,
+        send_stats,
+    )
+    await _send_realtime_chunk_stats(
+        ws,
+        session,
+        chunk,
+        batch,
+        request_prepare_ms,
+        scheduler_forward_ms,
+        chunk_total_ms,
+        send_stats,
+    )
+    return send_stats
+
+
+def _result_num_frames(result) -> int:
+    if result.raw_frame_batches is None:
+        return 0
+    return sum(len(frames) for frames in result.raw_frame_batches)
+
+
+def _output_pacing_fps(batch: "Req") -> float:
+    fps = float(batch.fps or 0)
+    if batch.enable_frame_interpolation:
+        fps *= 2 ** int(batch.frame_interpolation_exp or 1)
+    return fps
+
+
+async def _wait_for_realtime_output_slot(
+    session: GenerateSession,
+    batch: "Req",
+    result,
+) -> float:
+    if not getattr(batch, "realtime_output_pacing", False):
+        return 0.0
+
+    frame_count = _result_num_frames(result)
+    output_fps = _output_pacing_fps(batch)
+    if frame_count <= 0 or output_fps <= 0:
+        return 0.0
+
+    now = time.perf_counter()
+    next_send_at = session.output_pace_next_send_at
+    if next_send_at is None:
+        next_send_at = now
+    if (
+        batch.realtime_event_id is not None
+        and batch.realtime_event_id != session.output_pace_last_event_id
+    ):
+        next_send_at = min(next_send_at, now)
+        session.output_pace_last_event_id = batch.realtime_event_id
+
+    wait_s = max(0.0, next_send_at - now)
+    if wait_s > 0:
+        await asyncio.sleep(wait_s)
+
+    send_started_at = time.perf_counter()
+    session.output_pace_next_send_at = (
+        max(next_send_at, send_started_at) + frame_count / output_fps
+    )
+    return wait_s * 1000
+
+
+async def _await_realtime_task(task: asyncio.Task | None) -> None:
+    if task is None:
+        return
+    try:
+        await task
+    except (asyncio.CancelledError, WebSocketDisconnect):
+        pass
+    except Exception as e:
+        logger.debug("realtime task exited with error: %s", e)
+
+
+async def _listen_events(ws: WebSocket, session: GenerateSession):
+    """listen for user events: usually condition inputs"""
+    async for message in ws.iter_bytes():
+        data = None
+        try:
+            data = msgspec.msgpack.decode(message)
+            if not isinstance(data, dict):
+                raise ValueError("realtime event must be a map")
+            realtime_event = RealtimeEvent.model_validate(data)
+            if session.adapter is None:
+                raise ValueError("realtime adapter is not initialized")
+            event_log = session.adapter.ingest_event(session, realtime_event)
+            logger.info(
+                "receive realtime event, session_id=%s, event_id=%s, %s",
+                session.id,
+                realtime_event.event_id,
+                event_log,
+            )
+        except Exception as e:
+            event_kind = data.get("kind") if isinstance(data, dict) else None
+            logger.warning("invalid event, kind=%s, error=%s", event_kind, e)
+            await write_error_msg("invalid event", ws)
+            continue
+
+
+async def _listen_generate_request(ws: WebSocket, session: GenerateSession):
+    while True:
+        try:
+            data = msgspec.msgpack.decode(await ws.receive_bytes())
+            if not isinstance(data, dict):
+                raise ValueError("generate request must be a map")
+
+            realtime_req = RealtimeVideoGenerationsRequest.model_validate(data)
+            adapter = get_realtime_model_adapter(get_global_server_args())
+            session.set_adapter(adapter)
+            await adapter.on_init(session, realtime_req)
+
+            # Keep session state update atomic with validated request.
+            session.set_request(realtime_req)
+            break
+        except WebSocketDisconnect:
+            raise
+        except Exception as e:
+            logger.warning(
+                "invalid generate request, session_id=%s, error=%s",
+                session.id,
+                e,
+            )
+            await write_error_msg("invalid generate request", ws)
+            continue
+
+
+async def _cleanup_realtime_session(
+    session: GenerateSession,
+    generate_task: asyncio.Task | None,
+    listen_task: asyncio.Task | None,
+) -> None:
+    logger.info("terminating session, session_id=%s", session.id)
+    for task in (generate_task, listen_task):
+        if task and not task.done():
+            task.cancel()
+    for task in (generate_task, listen_task):
+        if task is None:
+            continue
+        await _await_realtime_task(task)
+    try:
+        await async_scheduler_client.forward(
+            ReleaseRealtimeSessionReq(session_id=session.id)
+        )
+    except Exception as e:
+        logger.warning(
+            "failed to release realtime session on scheduler, session_id=%s, error=%s",
+            session.id,
+            e,
+        )
+    if session.input_temp_dir is not None:
+        shutil.rmtree(session.input_temp_dir, ignore_errors=True)
+    session.dispose()
+
+
+async def _close_realtime_websocket(
+    websocket: WebSocket,
+    *,
+    code: int,
+    reason: str,
+) -> None:
+    try:
+        await websocket.close(code=code, reason=reason)
+    except (RuntimeError, WebSocketDisconnect):
+        pass
+
+
+async def _wait_for_server_warmup(websocket: WebSocket) -> None:
+    warmup_done = getattr(websocket.app.state, "server_warmup_done", None)
+    if warmup_done is not None and not warmup_done.is_set():
+        await warmup_done.wait()
+
+
+@router.websocket("/generate")
+async def generate(websocket: WebSocket):
+    """endpoint for creating a new realtime session"""
+    await websocket.accept()
+    await _wait_for_server_warmup(websocket)
+    if _ACTIVE_SESSION_IDS and not await _wait_for_active_session_slot():
+        logger.warning(
+            "reject realtime session because another session is active: %s",
+            sorted(_ACTIVE_SESSION_IDS),
+        )
+        try:
+            await write_error_msg(
+                "another realtime session is already active", websocket
+            )
+        finally:
+            await websocket.close(code=1008)
+        return
+
+    session = GenerateSession()
+    _ACTIVE_SESSION_IDS.add(session.id)
+    generate_task = None
+    listen_task = None
+    try:
+        # receive new generate request
+        await _listen_generate_request(websocket, session)
+
+        # continuously generate video chunk
+        generate_task = asyncio.create_task(_generate_loop(websocket, session))
+        # continuously listen for user events
+        listen_task = asyncio.create_task(_listen_events(websocket, session))
+
+        wait_tasks = [generate_task, listen_task]
+        await asyncio.wait(wait_tasks, return_when=asyncio.FIRST_COMPLETED)
+        if generate_task.done() and session.reached_max_chunks():
+            await _close_realtime_websocket(
+                websocket,
+                code=1000,
+                reason="generation complete",
+            )
+
+    except WebSocketDisconnect:
+        logger.info("client disconnected, session_id=%s", session.id)
+    finally:
+        try:
+            await _cleanup_realtime_session(session, generate_task, listen_task)
+        finally:
+            _ACTIVE_SESSION_IDS.discard(session.id)
+
+
+async def write_error_msg(error_msg: str, websocket: WebSocket):
+    await websocket.send_bytes(
+        msgspec.msgpack.encode({"type": "error", "content": error_msg})
+    )

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

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import (
+    BaseRealtimeModelAdapter,
+)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+_REALTIME_ADAPTER_REGISTRY: dict[type, type[BaseRealtimeModelAdapter]] = {}
+_BUILTIN_ADAPTERS_REGISTERED = False
+
+
+def register_realtime_model_adapter(
+    pipeline_config_cls: type,
+    adapter_cls: type[BaseRealtimeModelAdapter],
+) -> None:
+    _REALTIME_ADAPTER_REGISTRY[pipeline_config_cls] = adapter_cls
+
+
+def _register_builtin_realtime_model_adapters() -> None:
+    global _BUILTIN_ADAPTERS_REGISTERED
+    if _BUILTIN_ADAPTERS_REGISTERED:
+        return
+
+    from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import (
+        LingBotWorldCausalDMDConfig,
+    )
+    from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import (
+        SanaWMRealtimeConfig,
+    )
+    from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters.lingbot_world_realtime_adapter import (
+        LingBotWorldRealtimeAdapter,
+    )
+    from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters.sana_wm_realtime_adapter import (
+        SanaWMRealtimeAdapter,
+    )
+
+    register_realtime_model_adapter(
+        LingBotWorldCausalDMDConfig,
+        LingBotWorldRealtimeAdapter,
+    )
+    register_realtime_model_adapter(
+        SanaWMRealtimeConfig,
+        SanaWMRealtimeAdapter,
+    )
+    _BUILTIN_ADAPTERS_REGISTERED = True
+
+
+def get_realtime_model_adapter(
+    server_args: ServerArgs,
+) -> BaseRealtimeModelAdapter:
+    _register_builtin_realtime_model_adapters()
+
+    pipeline_config = server_args.pipeline_config
+    for config_cls in type(pipeline_config).__mro__:
+        adapter_cls = _REALTIME_ADAPTER_REGISTRY.get(config_cls)
+        if adapter_cls is not None:
+            return adapter_cls()
+
+    raise ValueError(
+        "Realtime video is not supported for pipeline config "
+        f"{type(pipeline_config).__name__}; no realtime adapter is registered."
+    )

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

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import time
+
+
+class RealtimeStageTimer:
+    __slots__ = ("_last", "_start")
+
+    def __init__(self):
+        now = time.perf_counter()
+        self._start = now
+        self._last = now
+
+    def mark_ms(self) -> float:
+        now = time.perf_counter()
+        elapsed_ms = (now - self._last) * 1000.0
+        self._last = now
+        return elapsed_ms
+
+    def total_ms(self) -> float:
+        return (time.perf_counter() - self._start) * 1000.0

python/sglang/multimodal_gen/runtime/entrypoints/openai/storage.py

@@ -0,0 +1,109 @@
+import asyncio
+import os
+from typing import Optional
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class CloudStorage:
+    def __init__(self):
+        self.enabled = os.getenv("SGLANG_CLOUD_STORAGE_TYPE", "").lower() == "s3"
+        if not self.enabled:
+            return
+
+        try:
+            import boto3
+        except ImportError:
+            logger.error(
+                "boto3 is not installed. Please install it with `pip install boto3` to use cloud storage."
+            )
+            self.enabled = False
+            return
+
+        self.bucket_name = os.getenv("SGLANG_S3_BUCKET_NAME")
+        if not self.bucket_name:
+            self.enabled = False
+            return
+
+        endpoint_url = os.getenv("SGLANG_S3_ENDPOINT_URL") or None
+        region_name = os.getenv("SGLANG_S3_REGION_NAME") or None
+
+        self.client = boto3.client(
+            "s3",
+            aws_access_key_id=os.getenv("SGLANG_S3_ACCESS_KEY_ID"),
+            aws_secret_access_key=os.getenv("SGLANG_S3_SECRET_ACCESS_KEY"),
+            endpoint_url=endpoint_url,
+            region_name=region_name,
+        )
+        self.endpoint_url = endpoint_url
+        self.region_name = region_name
+
+    def is_enabled(self) -> bool:
+        return self.enabled
+
+    async def upload_file(self, local_path: str, destination_key: str) -> Optional[str]:
+        if not self.is_enabled():
+            return None
+
+        def _sync_upload():
+            """Synchronous part of the upload to run in a thread."""
+            ext = os.path.splitext(local_path)[1].lower()
+            content_type = {
+                ".png": "image/png",
+                ".jpg": "image/jpeg",
+                ".jpeg": "image/jpeg",
+                ".webp": "image/webp",
+                ".mp4": "video/mp4",
+                ".glb": "model/gltf-binary",
+                ".obj": "text/plain",
+            }.get(ext, "application/octet-stream")
+
+            # Use the client created once in __init__
+            self.client.upload_file(
+                local_path,
+                self.bucket_name,
+                destination_key,
+                ExtraArgs={"ContentType": content_type},
+            )
+
+        try:
+            # Offload the blocking I/O call to a thread executor
+            await asyncio.get_running_loop().run_in_executor(None, _sync_upload)
+        except Exception as e:
+            # If upload fails, log the error and return None for fallback
+            logger.error(f"Upload failed for {destination_key}: {e}")
+            return None
+
+        # Simplified URL generation with a default region
+        if self.endpoint_url:
+            url = (
+                f"{self.endpoint_url.rstrip('/')}/{self.bucket_name}/{destination_key}"
+            )
+        else:
+            region = self.region_name or "us-east-1"
+            url = f"https://{self.bucket_name}.s3.{region}.amazonaws.com/{destination_key}"
+
+        logger.info(f"Uploaded {local_path} to {url}")
+        return url
+
+    async def upload_and_cleanup(self, file_path: str) -> Optional[str]:
+        """Helper to upload a file and delete the local copy if successful."""
+        if not self.is_enabled():
+            return None
+
+        key = os.path.basename(file_path)
+        url = await self.upload_file(file_path, key)
+
+        if url:
+            try:
+                # pass if removal fails
+                os.remove(file_path)
+            except OSError as e:
+                logger.warning(f"Failed to remove temporary file {file_path}: {e}")
+        return url
+
+
+# Global instance
+cloud_storage = CloudStorage()

python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py

@@ -0,0 +1,48 @@
+import asyncio
+from typing import Any, Dict, List, Optional
+
+
+class AsyncDictStore:
+    """A small async-safe in-memory key-value store for dict items.
+
+    This encapsulates the usual pattern of a module-level dict guarded by
+    an asyncio.Lock and provides simple CRUD methods that are safe to call
+    concurrently from FastAPI request handlers and background tasks.
+    """
+
+    def __init__(self) -> None:
+        self._items: Dict[str, Dict[str, Any]] = {}
+        self._lock = asyncio.Lock()
+
+    async def upsert(self, key: str, value: Dict[str, Any]) -> None:
+        async with self._lock:
+            self._items[key] = value
+
+    async def update_fields(
+        self, key: str, updates: Dict[str, Any]
+    ) -> Optional[Dict[str, Any]]:
+        async with self._lock:
+            item = self._items.get(key)
+            if item is None:
+                return None
+            item.update(updates)
+            return item
+
+    async def get(self, key: str) -> Optional[Dict[str, Any]]:
+        async with self._lock:
+            return self._items.get(key)
+
+    async def pop(self, key: str) -> Optional[Dict[str, Any]]:
+        async with self._lock:
+            return self._items.pop(key, None)
+
+    async def list_values(self) -> List[Dict[str, Any]]:
+        async with self._lock:
+            return list(self._items.values())
+
+
+# Global stores shared by OpenAI entrypoints
+# [request_id, dict]
+VIDEO_STORE = AsyncDictStore()
+IMAGE_STORE = AsyncDictStore()
+MESH_STORE = AsyncDictStore()

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -0,0 +1,452 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+import asyncio
+import inspect
+import json
+import os
+import shutil
+import tempfile
+import time
+from contextlib import contextmanager
+from typing import Any, Generator, List, Optional, Union
+
+import httpx
+from fastapi import HTTPException, UploadFile
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    ListLorasReq,
+    MergeLoraWeightsReq,
+    SetLoraReq,
+    ShutdownReq,
+    UnmergeLoraWeightsReq,
+    format_lora_message,
+    save_outputs,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.common import parse_size
+from sglang.multimodal_gen.runtime.utils.image_io import save_base64_image_to_path
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    init_logger,
+    log_batch_completion,
+    log_generation_timer,
+)
+from sglang.multimodal_gen.runtime.utils.trace_wrapper import trace_req
+
+# re-export LoRA protocol types for backward compatibility
+__all__ = [
+    "SetLoraReq",
+    "MergeLoraWeightsReq",
+    "UnmergeLoraWeightsReq",
+    "ListLorasReq",
+    "ShutdownReq",
+    "format_lora_message",
+]
+
+logger = init_logger(__name__)
+
+OUTPUT_QUALITY_MAPPER = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
+DEFAULT_FPS = 24
+DEFAULT_VIDEO_SECONDS = 4
+
+
+def _bad_request(message: str) -> HTTPException:
+    return HTTPException(status_code=400, detail=message)
+
+
+def _parse_size_or_raise(size: str) -> tuple[int, int]:
+    width, height = parse_size(size)
+    if width is None or height is None or width <= 0 or height <= 0:
+        raise _bad_request("size must be formatted as positive WIDTHxHEIGHT")
+    return width, height
+
+
+def _validate_positive_int(kwargs: dict[str, Any], name: str) -> None:
+    value = kwargs.get(name)
+    if value is not None and int(value) <= 0:
+        raise _bad_request(f"{name} must be positive")
+
+
+def flatten_extra_params(payload: Any) -> dict[str, Any]:
+    """Promote vLLM-Omni-style extra_params into regular request fields."""
+    if not isinstance(payload, dict):
+        return {}
+
+    extra_params = payload.pop("extra_params", None)
+    if isinstance(extra_params, str):
+        try:
+            extra_params = json.loads(extra_params)
+        except Exception:
+            extra_params = None
+    if not isinstance(extra_params, dict):
+        if "guardrails" in payload:
+            payload.setdefault("use_guardrails", payload["guardrails"])
+        return payload
+
+    for key, value in extra_params.items():
+        payload.setdefault(key, value)
+    if "guardrails" in extra_params:
+        payload.setdefault("use_guardrails", extra_params["guardrails"])
+
+    return payload
+
+
+@contextmanager
+def temp_dir_if_disabled(
+    configured_path: str | None,
+) -> Generator[str, None, None]:
+    """Yield *configured_path* when it is set, otherwise create a temporary
+    directory that is automatically removed when the context exits."""
+    if configured_path is not None:
+        os.makedirs(configured_path, exist_ok=True)
+        yield configured_path
+    else:
+        tmp = tempfile.mkdtemp(prefix="sglang_")
+        try:
+            yield tmp
+        finally:
+            shutil.rmtree(tmp, ignore_errors=True)
+
+
+def choose_output_image_ext(
+    output_format: Optional[str], background: Optional[str]
+) -> str:
+    fmt = (output_format or "").lower()
+    if fmt in {"png", "webp", "jpeg", "jpg"}:
+        return "jpg" if fmt == "jpeg" else fmt
+    if (background or "auto").lower() == "transparent":
+        return "png"
+    return "jpg"
+
+
+def build_sampling_params(request_id: str, **kwargs) -> SamplingParams:
+    """Build SamplingParams from request parameters.
+
+    Handles size parsing, output_quality resolution, and None filtering before
+    delegating to SamplingParams.from_user_sampling_params_args. Callers pass
+    only the parameters they have; None values are stripped automatically so
+    that SamplingParams defaults apply.
+    """
+    server_args = get_global_server_args()
+
+    # pop HTTP-layer params that aren't SamplingParams fields
+    output_quality = kwargs.pop("output_quality", None)
+
+    has_explicit_compression = kwargs.get("output_compression") is not None
+
+    # parse "WxH" size string if provided
+    size = kwargs.pop("size", None)
+    if size:
+        w, h = _parse_size_or_raise(size)
+        # treat None dimensions as unset so parsed size can fill them
+        if kwargs.get("width") is None:
+            kwargs["width"] = w
+        if kwargs.get("height") is None:
+            kwargs["height"] = h
+
+    for name in (
+        "width",
+        "height",
+        "num_frames",
+        "num_inference_steps",
+        "num_outputs_per_prompt",
+    ):
+        _validate_positive_int(kwargs, name)
+
+    # filter out None values to let SamplingParams defaults apply
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+    kwargs.setdefault("save_output", True)
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        model_path=server_args.model_path,
+        server_args=server_args,
+        request_id=request_id,
+        **kwargs,
+    )
+
+    # resolve output_quality → output_compression with the correct data_type.
+    # SamplingParams.__post_init__ may have resolved with the wrong data_type
+    # (default VIDEO) before _adjust() set the correct one.
+    if not has_explicit_compression and output_quality is not None:
+        resolved = adjust_output_quality(output_quality, sampling_params.data_type)
+        if resolved is not None:
+            sampling_params.output_compression = resolved
+
+    return sampling_params
+
+
+async def save_image_to_path(
+    image: Union[UploadFile, bytes, str],
+    target_path: str,
+    *,
+    prefer_remote_source: bool = False,
+) -> str:
+    input_path = await _maybe_url_image(
+        image, target_path, prefer_remote_source=prefer_remote_source
+    )
+    if input_path is None:
+        input_path = await _save_upload_to_path(image, target_path)
+    return input_path
+
+
+# Helpers
+async def _save_upload_to_path(
+    upload: Union[UploadFile, bytes], target_path: str
+) -> str:
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+    if isinstance(upload, bytes):
+        content = upload
+    elif isinstance(upload, (bytearray, memoryview)):
+        content = bytes(upload)
+    else:
+        read = getattr(upload, "read", None)
+        if not callable(read):
+            raise TypeError(f"Unsupported image upload type: {type(upload).__name__}")
+        content = read()
+        if inspect.isawaitable(content):
+            content = await content
+        if isinstance(content, (bytearray, memoryview)):
+            content = bytes(content)
+        if not isinstance(content, bytes):
+            raise TypeError(
+                f"Image upload read() returned {type(content).__name__}, expected bytes"
+            )
+    with open(target_path, "wb") as f:
+        f.write(content)
+    return target_path
+
+
+async def _maybe_url_image(
+    img_url: str,
+    target_path: str,
+    *,
+    prefer_remote_source: bool = False,
+) -> str | None:
+    if not isinstance(img_url, str):
+        return None
+
+    if img_url.lower().startswith(("http://", "https://")):
+        # Only bypass persistence when the caller explicitly disables input saves.
+        # Otherwise keep the prefetch outside the measured server stages.
+        if prefer_remote_source:
+            return img_url
+        # download image from URL and persist on disk
+        input_path = await _save_url_image_to_path(img_url, target_path)
+        return input_path
+    elif img_url.startswith("data:image"):
+        if prefer_remote_source:
+            return img_url
+        # encode image base64 url and persist on disk
+        input_path = save_base64_image_to_path(img_url, target_path)
+        return input_path
+    else:
+        raise ValueError("Unsupported image url format")
+
+
+async def _save_url_image_to_path(image_url: str, target_path: str) -> str:
+    """Download image from URL and save to target path."""
+
+    def _is_retryable_download_error(error: Exception) -> bool:
+        if isinstance(error, httpx.HTTPStatusError):
+            status_code = error.response.status_code
+            # Retry on rate limit and transient server-side failures.
+            return status_code == 429 or 500 <= status_code < 600
+        # Retry on transient network/protocol issues.
+        return isinstance(
+            error,
+            (
+                httpx.TimeoutException,
+                httpx.NetworkError,
+                httpx.RemoteProtocolError,
+            ),
+        )
+
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+    max_attempts = 3
+    backoff_seconds = 0.2
+    last_error: Exception | None = None
+
+    try:
+        async with httpx.AsyncClient(follow_redirects=True) as client:
+            for attempt in range(1, max_attempts + 1):
+                try:
+                    response = await client.get(image_url, timeout=10.0)
+                    response.raise_for_status()
+
+                    # Determine file extension from content type or URL after downloading
+                    if not os.path.splitext(target_path)[1]:
+                        content_type = response.headers.get("content-type", "").lower()
+
+                        url_path = image_url.split("?")[0]
+                        _, url_ext = os.path.splitext(url_path)
+                        url_ext = url_ext.lower()
+
+                        if url_ext in {
+                            ".jpg",
+                            ".jpeg",
+                            ".png",
+                            ".webp",
+                            ".gif",
+                            ".bmp",
+                        }:
+                            ext = ".jpg" if url_ext == ".jpeg" else url_ext
+                        elif content_type.startswith("image/"):
+                            if "jpeg" in content_type or "jpg" in content_type:
+                                ext = ".jpg"
+                            elif "png" in content_type:
+                                ext = ".png"
+                            elif "webp" in content_type:
+                                ext = ".webp"
+                            else:
+                                ext = ".jpg"  # Default to jpg
+                        elif content_type == "application/octet-stream":
+                            # for octet-stream, if we couldn't get it from URL, default to jpg
+                            ext = ".jpg"
+                        else:
+                            raise ValueError(
+                                f"URL does not point to an image. Content-Type: {content_type}"
+                            )
+                        target_path = f"{target_path}{ext}"
+
+                    with open(target_path, "wb") as f:
+                        f.write(response.content)
+
+                    return target_path
+                except Exception as e:
+                    last_error = e
+                    if attempt == max_attempts or not _is_retryable_download_error(e):
+                        raise
+                    wait_s = backoff_seconds * (2 ** (attempt - 1))
+                    logger.warning(
+                        "Retrying image download (%s/%s) for %s after %.1fs due to: %s",
+                        attempt,
+                        max_attempts,
+                        image_url,
+                        wait_s,
+                        e,
+                    )
+                    await asyncio.sleep(wait_s)
+    except Exception as e:
+        final_error = last_error or e
+        raise Exception(
+            f"Failed to download image from URL {image_url}: {str(final_error)}"
+        )
+
+
+async def process_generation_batch(
+    scheduler_client: AsyncSchedulerClient,
+    batch,
+) -> tuple[list[str], OutputBatch]:
+    total_start_time = time.perf_counter()
+    with trace_req(batch.trace_ctx), log_generation_timer(logger, batch.prompt):
+        result = await scheduler_client.forward([batch])
+
+        if (
+            result.output is None
+            and result.output_file_paths is None
+            and result.raw_frame_batches is None
+        ):
+            error_msg = result.error or "Unknown error"
+            raise RuntimeError(
+                f"Model generation returned no output. Error from scheduler: {error_msg}"
+            )
+
+        save_file_path_list = []
+        if result.output_file_paths:
+            save_file_path_list = result.output_file_paths
+        elif result.output is not None:
+            num_outputs = len(result.output)
+            save_file_path_list = save_outputs(
+                result.output,
+                batch.data_type,
+                batch.fps,
+                batch.save_output,
+                lambda idx: str(batch.output_file_path(num_outputs, idx)),
+                audio=result.audio,
+                audio_sample_rate=result.audio_sample_rate,
+                output_compression=batch.output_compression,
+                enable_frame_interpolation=batch.enable_frame_interpolation,
+                frame_interpolation_exp=batch.frame_interpolation_exp,
+                frame_interpolation_scale=batch.frame_interpolation_scale,
+                frame_interpolation_model_path=batch.frame_interpolation_model_path,
+                enable_upscaling=batch.enable_upscaling,
+                upscaling_model_path=batch.upscaling_model_path,
+                upscaling_scale=batch.upscaling_scale,
+            )
+
+    total_time = time.perf_counter() - total_start_time
+    if get_global_server_args().batching_max_size > 1:
+        log_batch_completion(
+            logger,
+            len(save_file_path_list),
+            total_time,
+        )
+
+    if result.peak_memory_mb and result.peak_memory_mb > 0:
+        logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB")
+
+    return save_file_path_list, result
+
+
+def merge_image_input_list(*inputs: Union[List, Any, None]) -> List:
+    """
+    Merge multiple image input sources into a single list.
+
+    This function handles both single items and lists of items, merging them
+    into a single flattened list. Useful for processing images, URLs, or other
+    multimedia inputs that can come as either single items or lists.
+
+    Args:
+        *inputs: Variable number of inputs, each can be None, single item, or list
+
+    Returns:
+        List: Flattened list of all non-None inputs
+
+    Example:
+        >>> merge_image_input_list(["img1", "img2"], "img3", None)
+        ["img1", "img2", "img3"]
+    """
+    result = []
+    for input_item in inputs:
+        if input_item is not None:
+            if isinstance(input_item, list):
+                result.extend(input_item)
+            else:
+                result.append(input_item)
+    return result
+
+
+def add_common_data_to_response(
+    response: dict, request_id: str, result: OutputBatch
+) -> dict:
+    if result.peak_memory_mb and result.peak_memory_mb > 0:
+        response["peak_memory_mb"] = result.peak_memory_mb
+
+    if result.metrics and result.metrics.total_duration_s > 0:
+        response["inference_time_s"] = result.metrics.total_duration_s
+
+    response["id"] = request_id
+
+    if result.action_pred is not None:
+        t = result.action_pred
+        response["action"] = {
+            "data": t.tolist(),
+            "shape": list(t.shape),
+            "dtype": str(t.dtype).replace("torch.", ""),
+            "raw_action_dim": result.action_raw_action_dim,
+            "action_mode": result.action_mode,
+            "domain_id": result.action_domain_id,
+        }
+
+    return response
+
+
+def adjust_output_quality(output_quality: str, data_type: DataType = None) -> int:
+    if output_quality == "default":
+        return 50 if data_type == DataType.VIDEO else 75
+    return OUTPUT_QUALITY_MAPPER.get(output_quality, None)

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -0,0 +1,741 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import asyncio
+import json
+import os
+import shutil
+import tempfile
+import time
+from typing import Any, Dict, Optional
+
+from fastapi import (
+    APIRouter,
+    File,
+    Form,
+    HTTPException,
+    Path,
+    Query,
+    Request,
+    UploadFile,
+)
+from fastapi.responses import FileResponse
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    SamplingParams,
+    generate_request_id,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+    VideoGenerationsRequest,
+    VideoListResponse,
+    VideoResponse,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage
+from sglang.multimodal_gen.runtime.entrypoints.openai.stores import VIDEO_STORE
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
+    DEFAULT_FPS,
+    DEFAULT_VIDEO_SECONDS,
+    add_common_data_to_response,
+    build_sampling_params,
+    flatten_extra_params,
+    merge_image_input_list,
+    process_generation_batch,
+    save_image_to_path,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.observability.trace import extract_trace_headers
+
+logger = init_logger(__name__)
+router = APIRouter(prefix="/v1/videos", tags=["videos"])
+
+_VIDEO_EXTENSIONS = {
+    ".avi",
+    ".gif",
+    ".m4v",
+    ".mkv",
+    ".mov",
+    ".mp4",
+    ".mpeg",
+    ".mpg",
+    ".webm",
+}
+
+
+def _extra_value(request: VideoGenerationsRequest, name: str) -> Any:
+    return (request.model_extra or {}).get(name)
+
+
+def _request_value(request: VideoGenerationsRequest, name: str) -> Any:
+    value = getattr(request, name, None)
+    if value is not None:
+        return value
+    return _extra_value(request, name)
+
+
+def _parse_form_extra_value(value: Any) -> Any:
+    if not isinstance(value, str):
+        return value
+    try:
+        return json.loads(value)
+    except Exception:
+        return value
+
+
+def _is_probably_video_source(source: Any) -> bool:
+    content_type = (getattr(source, "content_type", "") or "").lower()
+    if content_type.startswith("video/"):
+        return True
+
+    if isinstance(source, str):
+        if source.lower().startswith("data:video"):
+            return True
+        source_name = source
+    else:
+        source_name = getattr(source, "filename", None)
+
+    if not source_name:
+        return False
+    source_name = str(source_name).split("?", 1)[0].split("#", 1)[0]
+    return os.path.splitext(source_name)[1].lower() in _VIDEO_EXTENSIONS
+
+
+def _is_cosmos3_server(server_args) -> bool:
+    from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
+
+    return isinstance(server_args.pipeline_config, Cosmos3Config)
+
+
+def _normalize_optional_string(value: Any) -> Any:
+    if isinstance(value, str) and not value.strip():
+        return None
+    return value
+
+
+def _coerce_optional_int_list(value: Any) -> list[int] | None:
+    value = _parse_form_extra_value(value)
+    if value is None:
+        return None
+    if isinstance(value, str) and not value.strip():
+        return None
+    if isinstance(value, (list, tuple)):
+        return [int(item) for item in value]
+    return [int(value)]
+
+
+def _resolve_video_path(req: VideoGenerationsRequest) -> str | None:
+    video_path = _request_value(req, "video_path") or _request_value(req, "video_url")
+    if video_path:
+        return str(video_path)
+
+    input_reference = _request_value(req, "input_reference")
+    if _is_probably_video_source(input_reference):
+        return str(input_reference)
+
+    reference_url = _request_value(req, "reference_url")
+    if _is_probably_video_source(reference_url):
+        return str(reference_url)
+
+    return None
+
+
+def _resolve_image_path(
+    req: VideoGenerationsRequest, video_path: str | None
+) -> str | None:
+    image_path = _request_value(req, "input_reference")
+    if video_path and image_path == video_path:
+        return None
+    if _is_probably_video_source(image_path):
+        return None
+    return image_path
+
+
+def _resolve_sound_duration(
+    req: VideoGenerationsRequest, *, num_frames: int, fps: int
+) -> float | None:
+    generate_sound = _request_value(req, "generate_sound")
+    sound_duration = _request_value(req, "sound_duration")
+
+    if generate_sound is False:
+        return 0.0
+    if sound_duration is not None:
+        return float(sound_duration)
+    if generate_sound is True:
+        return float(num_frames) / float(fps)
+    return None
+
+
+def _cosmos3_sampling_param_kwargs(
+    req: VideoGenerationsRequest, *, num_frames: int, fps: int
+) -> Dict[str, Any]:
+    """Map HTTP/API aliases to Cosmos3SamplingParams field names."""
+    kwargs: Dict[str, Any] = {}
+
+    sound_duration = _resolve_sound_duration(req, num_frames=num_frames, fps=fps)
+    if sound_duration is not None:
+        kwargs["sound_duration"] = sound_duration
+
+    condition_frame_indexes = _request_value(req, "condition_frame_indexes")
+    if condition_frame_indexes is None:
+        condition_frame_indexes = _request_value(req, "condition_frame_indexes_vision")
+    condition_frame_indexes = _coerce_optional_int_list(condition_frame_indexes)
+    if condition_frame_indexes is not None:
+        kwargs["condition_frame_indexes"] = condition_frame_indexes
+
+    for name in (
+        "condition_video_keep",
+        "action_mode",
+        "domain_id",
+        "domain_name",
+        "raw_action_dim",
+        "action_fps",
+        "action",
+        "action_view_point",
+        "action_normalization",
+    ):
+        value = _parse_form_extra_value(_request_value(req, name))
+        value = _normalize_optional_string(value)
+        if value is not None:
+            kwargs[name] = value
+
+    return kwargs
+
+
+def _build_video_sampling_params(request_id: str, request: VideoGenerationsRequest):
+    """Resolve video-specific defaults (fps, seconds → num_frames) then
+    delegate to the shared build_sampling_params."""
+    server_args = get_global_server_args()
+    seconds = request.seconds if request.seconds is not None else DEFAULT_VIDEO_SECONDS
+    fps = request.fps if request.fps is not None else DEFAULT_FPS
+    num_frames = request.num_frames if request.num_frames is not None else fps * seconds
+    num_outputs = request.num_outputs_per_prompt
+    if num_outputs is None:
+        num_outputs = request.n or 1
+    video_path = _resolve_video_path(request)
+    image_path = _resolve_image_path(request, video_path)
+    cosmos3_kwargs = {}
+    if _is_cosmos3_server(server_args):
+        cosmos3_kwargs = _cosmos3_sampling_param_kwargs(
+            request, num_frames=num_frames, fps=fps
+        )
+        if server_args.pipeline_config.action_stats_path is not None:
+            cosmos3_kwargs["action_stats_path"] = (
+                server_args.pipeline_config.action_stats_path
+            )
+
+    return build_sampling_params(
+        request_id,
+        prompt=request.prompt,
+        num_outputs_per_prompt=max(1, min(int(num_outputs), 10)),
+        size=request.size,
+        width=request.width,
+        height=request.height,
+        num_frames=num_frames,
+        fps=fps,
+        image_path=image_path,
+        video_path=video_path,
+        output_file_name=request_id,
+        seed=request.seed,
+        generator_device=request.generator_device,
+        num_inference_steps=request.num_inference_steps,
+        guidance_scale=request.guidance_scale,
+        guidance_scale_2=request.guidance_scale_2,
+        negative_prompt=request.negative_prompt,
+        max_sequence_length=request.max_sequence_length,
+        flow_shift=request.flow_shift,
+        use_duration_template=_extra_value(request, "use_duration_template"),
+        use_resolution_template=_extra_value(request, "use_resolution_template"),
+        use_system_prompt=_extra_value(request, "use_system_prompt"),
+        use_guardrails=_extra_value(request, "use_guardrails"),
+        enable_teacache=request.enable_teacache,
+        enable_frame_interpolation=request.enable_frame_interpolation,
+        frame_interpolation_exp=request.frame_interpolation_exp,
+        frame_interpolation_scale=request.frame_interpolation_scale,
+        frame_interpolation_model_path=request.frame_interpolation_model_path,
+        enable_upscaling=request.enable_upscaling,
+        upscaling_model_path=request.upscaling_model_path,
+        upscaling_scale=request.upscaling_scale,
+        output_path=request.output_path,
+        output_compression=request.output_compression,
+        output_quality=request.output_quality,
+        perf_dump_path=request.perf_dump_path,
+        diffusers_kwargs=request.diffusers_kwargs,
+        **cosmos3_kwargs,
+    )
+
+
+# extract metadata which http_server needs to know
+def _video_job_from_sampling(
+    request_id: str, req: VideoGenerationsRequest, sampling: SamplingParams
+) -> Dict[str, Any]:
+    size_str = f"{sampling.width}x{sampling.height}"
+    seconds = int(round((sampling.num_frames or 0) / float(sampling.fps or 24)))
+    return {
+        "id": request_id,
+        "object": "video",
+        "model": req.model or "sora-2",
+        "status": "queued",
+        "progress": 0,
+        "created_at": int(time.time()),
+        "size": size_str,
+        "seconds": str(seconds),
+        "quality": "standard",
+        "file_path": os.path.abspath(sampling.output_file_path()),
+    }
+
+
+async def _save_first_input_image(
+    image_sources,
+    request_id: str,
+    uploads_dir: str,
+    *,
+    prefer_remote_source: bool = False,
+) -> str | None:
+    """Save the first input image from a list of sources and return its path."""
+    image_list = merge_image_input_list(image_sources)
+    if not image_list:
+        return None
+    image = image_list[0]
+
+    os.makedirs(uploads_dir, exist_ok=True)
+
+    filename = image.filename if hasattr(image, "filename") else "url_image"
+    target_path = os.path.join(uploads_dir, f"{request_id}_{filename}")
+    return await save_image_to_path(
+        image, target_path, prefer_remote_source=prefer_remote_source
+    )
+
+
+async def _dispatch_job_async(
+    job_id: str,
+    batch: Req,
+    *,
+    temp_dirs: list[str] | None = None,
+    output_persistent: bool = True,
+) -> None:
+    from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+
+    try:
+        save_file_path_list, result = await process_generation_batch(
+            async_scheduler_client, batch
+        )
+        save_file_path = save_file_path_list[0]
+
+        cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
+
+        persistent_path = (
+            save_file_path if not cloud_url and output_persistent else None
+        )
+        update_fields = {
+            "status": "completed",
+            "progress": 100,
+            "completed_at": int(time.time()),
+            "url": cloud_url,
+            "file_path": persistent_path,
+            "file_paths": (
+                [os.path.abspath(path) for path in save_file_path_list]
+                if output_persistent
+                else None
+            ),
+            "num_outputs": len(save_file_path_list),
+        }
+        update_fields = add_common_data_to_response(
+            update_fields, request_id=job_id, result=result
+        )
+        await VIDEO_STORE.update_fields(job_id, update_fields)
+    except Exception as e:
+        logger.error(f"{e}")
+        await VIDEO_STORE.update_fields(
+            job_id, {"status": "failed", "error": {"message": str(e)}}
+        )
+    finally:
+        for td in temp_dirs or []:
+            shutil.rmtree(td, ignore_errors=True)
+
+
+# TODO: support image to video generation
+@router.post("", response_model=VideoResponse)
+async def create_video(
+    request: Request,
+    # multipart/form-data fields (optional; used only when content-type is multipart)
+    prompt: Optional[str] = Form(None),
+    input_reference: Optional[UploadFile] = File(None),
+    reference_url: Optional[str] = Form(None),
+    video_reference: Optional[UploadFile] = File(None),
+    video_url: Optional[str] = Form(None),
+    video_path: Optional[str] = Form(None),
+    model: Optional[str] = Form(None),
+    n: Optional[int] = Form(1),
+    num_outputs_per_prompt: Optional[int] = Form(None),
+    seconds: Optional[int] = Form(None),
+    size: Optional[str] = Form(None),
+    fps: Optional[int] = Form(None),
+    num_frames: Optional[int] = Form(None),
+    seed: Optional[int] = Form(None),
+    generator_device: Optional[str] = Form("cuda"),
+    negative_prompt: Optional[str] = Form(None),
+    guidance_scale: Optional[float] = Form(None),
+    num_inference_steps: Optional[int] = Form(None),
+    max_sequence_length: Optional[int] = Form(None),
+    flow_shift: Optional[float] = Form(None),
+    enable_teacache: Optional[bool] = Form(None),
+    enable_frame_interpolation: Optional[bool] = Form(None),
+    frame_interpolation_exp: Optional[int] = Form(None),
+    frame_interpolation_scale: Optional[float] = Form(None),
+    frame_interpolation_model_path: Optional[str] = Form(None),
+    enable_upscaling: Optional[bool] = Form(None),
+    upscaling_model_path: Optional[str] = Form(None),
+    upscaling_scale: Optional[int] = Form(None),
+    output_quality: Optional[str] = Form(None),
+    output_compression: Optional[int] = Form(None),
+    output_path: Optional[str] = Form(None),
+    extra_params: Optional[str] = Form(None),
+    extra_body: Optional[str] = Form(None),
+):
+    content_type = request.headers.get("content-type", "").lower()
+    request_id = generate_request_id()
+
+    server_args = get_global_server_args()
+    task_type = server_args.pipeline_config.task_type
+
+    # Resolve input upload directory (may be a temp dir when saving is disabled)
+    temp_dirs: list[str] = []
+    if server_args.input_save_path is not None:
+        uploads_dir = server_args.input_save_path
+        os.makedirs(uploads_dir, exist_ok=True)
+    else:
+        uploads_dir = tempfile.mkdtemp(prefix="sglang_input_")
+        temp_dirs.append(uploads_dir)
+
+    # Resolve output directory
+    effective_output_path = server_args.output_path
+    output_persistent = True
+    if "multipart/form-data" not in content_type:
+        # JSON body may carry a per-request output_path; checked after parsing below
+        pass
+
+    if "multipart/form-data" in content_type:
+        if not prompt:
+            raise HTTPException(status_code=400, detail="prompt is required")
+
+        video_input_path = None
+        image_sources = merge_image_input_list(input_reference, reference_url)
+        if video_reference is not None:
+            video_input_path = await _save_first_input_image(
+                video_reference,
+                request_id,
+                uploads_dir,
+                prefer_remote_source=server_args.input_save_path is None,
+            )
+        elif video_path or video_url:
+            video_input_path = video_path or video_url
+        elif input_reference is not None and _is_probably_video_source(input_reference):
+            video_input_path = await _save_first_input_image(
+                input_reference,
+                request_id,
+                uploads_dir,
+                prefer_remote_source=server_args.input_save_path is None,
+            )
+            image_sources = merge_image_input_list(reference_url)
+        elif reference_url and _is_probably_video_source(reference_url):
+            video_input_path = reference_url
+            image_sources = merge_image_input_list(input_reference)
+
+        # Validate image input based on model task type
+        if task_type.requires_image_input() and not image_sources:
+            raise HTTPException(
+                status_code=400,
+                detail="input_reference or reference_url is required for image-to-video generation",
+            )
+        input_path = None
+        if image_sources:
+            try:
+                input_path = await _save_first_input_image(
+                    image_sources,
+                    request_id,
+                    uploads_dir,
+                    prefer_remote_source=server_args.input_save_path is None,
+                )
+            except Exception as e:
+                raise HTTPException(
+                    status_code=400, detail=f"Failed to process image source: {str(e)}"
+                )
+
+        # Parse extra_body JSON (if provided in multipart form) to get fps/num_frames overrides
+        extra_from_form: Dict[str, Any] = {}
+        if extra_body:
+            try:
+                extra_from_form = flatten_extra_params(json.loads(extra_body))
+            except Exception:
+                extra_from_form = {}
+        if extra_params:
+            try:
+                extra_from_form.update(
+                    flatten_extra_params({"extra_params": json.loads(extra_params)})
+                )
+            except Exception:
+                pass
+
+        def form_value(name: str, value: Any) -> Any:
+            selected = value if value is not None else extra_from_form.get(name)
+            return _parse_form_extra_value(selected)
+
+        raw_form = await request.form()
+        for key in (
+            "use_duration_template",
+            "use_resolution_template",
+            "use_system_prompt",
+            "use_guardrails",
+            "guardrails",
+            "video_path",
+            "video_url",
+            "generate_sound",
+            "sound_duration",
+            "condition_frame_indexes",
+            "action_mode",
+            "domain_id",
+            "domain_name",
+            "raw_action_dim",
+            "action_fps",
+            "action",
+            "action_view_point",
+            "action_normalization",
+            "condition_frame_indexes_vision",
+            "condition_video_keep",
+        ):
+            if key in raw_form and key not in extra_from_form:
+                extra_from_form[key] = _parse_form_extra_value(raw_form[key])
+        flatten_extra_params(extra_from_form)
+
+        request_field_names = set(VideoGenerationsRequest.model_fields)
+        extra_request_fields = {
+            key: value
+            for key, value in extra_from_form.items()
+            if key not in request_field_names
+        }
+        fps_val = form_value("fps", fps)
+        num_frames_val = form_value("num_frames", num_frames)
+
+        req = VideoGenerationsRequest(
+            prompt=prompt,
+            input_reference=input_path,
+            video_path=form_value("video_path", video_input_path),
+            video_url=form_value("video_url", video_url),
+            model=form_value("model", model),
+            n=form_value("n", n),
+            num_outputs_per_prompt=form_value(
+                "num_outputs_per_prompt", num_outputs_per_prompt
+            ),
+            seconds=form_value("seconds", seconds) or 4,
+            size=form_value("size", size),
+            fps=fps_val,
+            num_frames=num_frames_val,
+            seed=form_value("seed", seed),
+            generator_device=form_value("generator_device", generator_device),
+            negative_prompt=form_value("negative_prompt", negative_prompt),
+            num_inference_steps=form_value("num_inference_steps", num_inference_steps),
+            guidance_scale=form_value("guidance_scale", guidance_scale),
+            max_sequence_length=form_value("max_sequence_length", max_sequence_length),
+            flow_shift=form_value("flow_shift", flow_shift),
+            enable_teacache=form_value("enable_teacache", enable_teacache),
+            enable_frame_interpolation=form_value(
+                "enable_frame_interpolation", enable_frame_interpolation
+            ),
+            frame_interpolation_exp=form_value(
+                "frame_interpolation_exp", frame_interpolation_exp
+            ),
+            frame_interpolation_scale=form_value(
+                "frame_interpolation_scale", frame_interpolation_scale
+            ),
+            frame_interpolation_model_path=form_value(
+                "frame_interpolation_model_path", frame_interpolation_model_path
+            ),
+            enable_upscaling=form_value("enable_upscaling", enable_upscaling),
+            upscaling_model_path=form_value(
+                "upscaling_model_path", upscaling_model_path
+            ),
+            upscaling_scale=form_value("upscaling_scale", upscaling_scale),
+            output_compression=form_value("output_compression", output_compression),
+            output_quality=form_value("output_quality", output_quality),
+            output_path=form_value("output_path", output_path),
+            diffusers_kwargs=form_value("diffusers_kwargs", None),
+            **extra_request_fields,
+        )
+    else:
+        try:
+            body = await request.json()
+        except Exception:
+            body = {}
+        try:
+            # If client uses extra_body, merge it into the top-level payload
+            payload: Dict[str, Any] = dict(body or {})
+            extra = payload.pop("extra_body", None)
+            if isinstance(extra, str):
+                extra = json.loads(extra)
+            if isinstance(extra, dict):
+                payload.update(flatten_extra_params(extra))
+            # openai may turn extra_body to extra_json
+            extra_json = payload.pop("extra_json", None)
+            if isinstance(extra_json, str):
+                extra_json = json.loads(extra_json)
+            if isinstance(extra_json, dict):
+                payload.update(flatten_extra_params(extra_json))
+            flatten_extra_params(payload)
+            # Validate image input based on model task type
+            if payload.get("video_url") and not payload.get("video_path"):
+                payload["video_path"] = payload["video_url"]
+            if _is_probably_video_source(payload.get("reference_url")):
+                payload.setdefault("video_path", payload.get("reference_url"))
+            if _is_probably_video_source(payload.get("input_reference")):
+                payload.setdefault("video_path", payload.get("input_reference"))
+
+            has_image_input = (
+                payload.get("reference_url")
+                and not _is_probably_video_source(payload.get("reference_url"))
+            ) or (
+                payload.get("input_reference")
+                and not _is_probably_video_source(payload.get("input_reference"))
+            )
+            if task_type.requires_image_input() and not has_image_input:
+                raise HTTPException(
+                    status_code=400,
+                    detail="input_reference or reference_url is required for image-to-video generation",
+                )
+            # for non-multipart/form-data type
+            if payload.get("reference_url") and not _is_probably_video_source(
+                payload.get("reference_url")
+            ):
+                try:
+                    input_path = await _save_first_input_image(
+                        payload.get("reference_url"),
+                        request_id,
+                        uploads_dir,
+                        prefer_remote_source=server_args.input_save_path is None,
+                    )
+                except Exception as e:
+                    raise HTTPException(
+                        status_code=400,
+                        detail=f"Failed to process image source: {str(e)}",
+                    )
+                payload["input_reference"] = input_path
+            req = VideoGenerationsRequest(**payload)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Invalid request body: {e}")
+
+    # Resolve per-request output_path override
+    effective_output_path = req.output_path or server_args.output_path
+    if effective_output_path is None:
+        output_tmp = tempfile.mkdtemp(prefix="sglang_output_")
+        temp_dirs.append(output_tmp)
+        effective_output_path = output_tmp
+        output_persistent = False
+
+    # Inject resolved output_path so _build_video_sampling_params picks it up
+    req.output_path = effective_output_path
+
+    logger.debug(f"Server received from create_video endpoint: req={req}")
+
+    try:
+        sampling_params = _build_video_sampling_params(request_id, req)
+    except (ValueError, TypeError) as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    job = _video_job_from_sampling(request_id, req, sampling_params)
+    await VIDEO_STORE.upsert(request_id, job)
+
+    # Build Req for scheduler
+    trace_headers = extract_trace_headers(request.headers)
+    batch = prepare_request(
+        server_args=server_args,
+        sampling_params=sampling_params,
+        external_trace_header=trace_headers,
+    )
+    # Add diffusers_kwargs if provided
+    if req.diffusers_kwargs:
+        batch.extra["diffusers_kwargs"] = req.diffusers_kwargs
+        if "max_sequence_length" in req.diffusers_kwargs:
+            batch.max_sequence_length = req.diffusers_kwargs["max_sequence_length"]
+        if "flow_shift" in req.diffusers_kwargs:
+            batch.flow_shift = req.diffusers_kwargs["flow_shift"]
+    # Enqueue the job asynchronously and return immediately
+    asyncio.create_task(
+        _dispatch_job_async(
+            request_id,
+            batch,
+            temp_dirs=temp_dirs or None,
+            output_persistent=output_persistent,
+        )
+    )
+    return VideoResponse(**job)
+
+
+@router.get("", response_model=VideoListResponse)
+async def list_videos(
+    after: Optional[str] = Query(None),
+    limit: Optional[int] = Query(None, ge=1, le=100),
+    order: Optional[str] = Query("desc"),
+):
+    # Normalize order
+    order = (order or "desc").lower()
+    if order not in ("asc", "desc"):
+        order = "desc"
+    jobs = await VIDEO_STORE.list_values()
+
+    reverse = order != "asc"
+    jobs.sort(key=lambda j: j.get("created_at", 0), reverse=reverse)
+
+    if after is not None:
+        try:
+            idx = next(i for i, j in enumerate(jobs) if j["id"] == after)
+            jobs = jobs[idx + 1 :]
+        except StopIteration:
+            jobs = []
+
+    if limit is not None:
+        jobs = jobs[:limit]
+    items = [VideoResponse(**j) for j in jobs]
+    return VideoListResponse(data=items)
+
+
+@router.get("/{video_id}", response_model=VideoResponse)
+async def retrieve_video(video_id: str = Path(...)):
+    job = await VIDEO_STORE.get(video_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Video not found")
+    return VideoResponse(**job)
+
+
+# TODO: support aborting a job.
+@router.delete("/{video_id}", response_model=VideoResponse)
+async def delete_video(video_id: str = Path(...)):
+    job = await VIDEO_STORE.pop(video_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Video not found")
+    # Mark as deleted in response semantics
+    job["status"] = "deleted"
+    return VideoResponse(**job)
+
+
+@router.get("/{video_id}/content")
+async def download_video_content(
+    video_id: str = Path(...), variant: Optional[str] = Query(None)
+):
+    job = await VIDEO_STORE.get(video_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Video not found")
+
+    if job.get("url"):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Video has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}",
+        )
+
+    file_path = job.get("file_path")
+    if not file_path or not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="Generation is still in-progress")
+
+    media_type = "video/mp4"  # default variant
+    return FileResponse(
+        path=file_path, media_type=media_type, filename=os.path.basename(file_path)
+    )

python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py

@@ -0,0 +1,110 @@
+"""Request/response data structures for post-training APIs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+
+@dataclass
+class UpdateWeightFromDiskReqInput:
+    """Request to update model weights from disk for diffusion models."""
+
+    model_path: str
+    flush_cache: bool = True
+    target_modules: list[str] | None = None
+
+
+@dataclass
+class UpdateWeightFromTensorReqInput:
+    """Request to update model weights from tensor payloads for diffusion models."""
+
+    serialized_named_tensors: list[str | bytes]
+    load_format: str | None = None
+    target_modules: list[str] | None = None
+
+
+@dataclass
+class UpdateWeightFromTensorCheckerReqInput:
+    """Request to verify live module weights against expected SHA-256 values."""
+
+    target_module: str
+    expected_named_tensors_sha256: dict[str, str]
+
+
+@dataclass
+class GetWeightsChecksumReqInput:
+    """Compute SHA-256 checksum of loaded module weights for verification."""
+
+    module_names: list[str] | None = None
+
+
+@dataclass
+class ReleaseMemoryOccupationReqInput:
+    """Request to release (sleep) GPU memory occupation for the diffusion engine."""
+
+    pass
+
+
+@dataclass
+class ResumeMemoryOccupationReqInput:
+    """Request to resume (wake) GPU memory occupation for the diffusion engine."""
+
+    pass
+
+
+class RolloutRequest(BaseModel):
+    prompt: str
+    negative_prompt: Optional[str] = None
+    seed: Optional[int] = None
+    generator_device: str = "cuda"
+
+    width: Optional[int] = None
+    height: Optional[int] = None
+    num_inference_steps: Optional[int] = None
+    num_outputs_per_prompt: Optional[int] = None
+
+    guidance_scale: Optional[float] = None
+    true_cfg_scale: Optional[float] = None
+
+    # video-specific (ignored by image pipelines)
+    num_frames: Optional[int] = None
+    fps: Optional[int] = None
+
+    rollout: bool = True
+    rollout_sde_type: str = "sde"
+    rollout_noise_level: float = 0.7
+    rollout_log_prob_no_const: bool = False
+    rollout_debug_mode: bool = True
+
+    rollout_return_denoising_env: bool = False
+    rollout_return_dit_trajectory: bool = False
+
+    # 0-indexed denoising-loop step filters. None = all steps.
+    rollout_sde_step_indices: Optional[list[int]] = None
+    rollout_return_step_indices: Optional[list[int]] = None
+
+    image_path: Optional[list[str]] = None
+
+    # suppress verbose per-request logging (also gates peak_memory_mb collection)
+    suppress_logs: bool = False
+
+    extra_sampling_params: Optional[dict[str, Any]] = None
+
+
+class RolloutResponse(BaseModel):
+    request_id: str
+    prompt: str
+    seed: int
+
+    generated_output: Any = None
+
+    rollout_log_probs: Optional[dict[str, Any]] = None
+    rollout_debug_tensors: Optional[dict[str, Any]] = None
+    denoising_env: Optional[dict[str, Any]] = None
+    dit_trajectory: Optional[dict[str, Any]] = None
+
+    inference_time_s: Optional[float] = None
+    peak_memory_mb: Optional[float] = None

python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py

@@ -0,0 +1,329 @@
+"""Rollout HTTP API (``POST /rollout/generate``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import ORJSONResponse
+
+from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params
+from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import (
+    RolloutRequest,
+    RolloutResponse,
+)
+from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import (
+    _maybe_serialize,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
+    RolloutDebugTensors,
+    RolloutDenoisingEnv,
+    RolloutDitTrajectory,
+    RolloutTrajectoryData,
+)
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+router = APIRouter(prefix="/rollout", tags=["rollout"])
+
+
+def _extract_single_sample_tensor(
+    obj: Any, sample_idx: int, batch_size: int, *, current_key: str | None = None
+) -> Any:
+    if isinstance(obj, torch.Tensor):
+        if obj.dim() >= 1 and obj.shape[0] == batch_size:
+            return obj[sample_idx].contiguous()
+        return obj
+    if isinstance(obj, dict):
+        return {
+            k: _extract_single_sample_tensor(v, sample_idx, batch_size, current_key=k)
+            for k, v in obj.items()
+        }
+    if isinstance(obj, list):
+        if current_key == "img_shapes" and len(obj) == batch_size:
+            return [obj[sample_idx]]
+        return [
+            _extract_single_sample_tensor(
+                v, sample_idx, batch_size, current_key=current_key
+            )
+            for v in obj
+        ]
+    if isinstance(obj, tuple):
+        return tuple(
+            _extract_single_sample_tensor(
+                v, sample_idx, batch_size, current_key=current_key
+            )
+            for v in obj
+        )
+    return obj
+
+
+def _slice_rollout_trajectory_for_sample(
+    rtd: RolloutTrajectoryData | None,
+    sample_idx: int,
+    batch_size: int,
+) -> RolloutTrajectoryData | None:
+    if rtd is None:
+        return None
+    log_probs = rtd.rollout_log_probs
+    if (
+        isinstance(log_probs, torch.Tensor)
+        and log_probs.dim() >= 1
+        and log_probs.shape[0] == batch_size
+    ):
+        log_probs = log_probs[sample_idx].contiguous()
+    debug_tensors = None
+    if rtd.rollout_debug_tensors:
+        rd = rtd.rollout_debug_tensors
+        debug_tensors = RolloutDebugTensors(
+            rollout_variance_noises=_extract_single_sample_tensor(
+                rd.rollout_variance_noises, sample_idx, batch_size
+            ),
+            rollout_prev_sample_means=_extract_single_sample_tensor(
+                rd.rollout_prev_sample_means, sample_idx, batch_size
+            ),
+            rollout_noise_std_devs=_extract_single_sample_tensor(
+                rd.rollout_noise_std_devs, sample_idx, batch_size
+            ),
+            rollout_model_outputs=_extract_single_sample_tensor(
+                rd.rollout_model_outputs, sample_idx, batch_size
+            ),
+        )
+    denoising_env = None
+    if rtd.denoising_env:
+        env = rtd.denoising_env
+        denoising_env = RolloutDenoisingEnv(
+            image_kwargs=(
+                _extract_single_sample_tensor(env.image_kwargs, sample_idx, batch_size)
+                if env.image_kwargs
+                else None
+            ),
+            pos_cond_kwargs=(
+                _extract_single_sample_tensor(
+                    env.pos_cond_kwargs, sample_idx, batch_size
+                )
+                if env.pos_cond_kwargs
+                else None
+            ),
+            neg_cond_kwargs=(
+                _extract_single_sample_tensor(
+                    env.neg_cond_kwargs, sample_idx, batch_size
+                )
+                if env.neg_cond_kwargs
+                else None
+            ),
+            guidance=(
+                _extract_single_sample_tensor(env.guidance, sample_idx, batch_size)
+                if env.guidance is not None
+                else None
+            ),
+        )
+    dit_trajectory = None
+    if rtd.dit_trajectory:
+        dit = rtd.dit_trajectory
+        dit_trajectory = RolloutDitTrajectory(
+            latents=_extract_single_sample_tensor(dit.latents, sample_idx, batch_size),
+            timesteps=dit.timesteps,
+        )
+    return RolloutTrajectoryData(
+        rollout_log_probs=log_probs,
+        rollout_debug_tensors=debug_tensors,
+        denoising_env=denoising_env,
+        dit_trajectory=dit_trajectory,
+    )
+
+
+def _serialize_rollout_trajectory(
+    rtd: RolloutTrajectoryData | None,
+    *,
+    serialized_dit_timesteps: dict | None = None,
+) -> tuple[dict | None, dict | None, dict | None, dict | None]:
+    """Return order: rollout_log_probs, rollout_debug_tensors, denoising_env, dit_trajectory."""
+    if rtd is None:
+        return None, None, None, None
+    serialized_log_probs = _maybe_serialize(rtd.rollout_log_probs)
+    serialized_debug_tensors = None
+    if rtd.rollout_debug_tensors:
+        rd = rtd.rollout_debug_tensors
+        serialized_debug_tensors = {
+            "rollout_variance_noises": _maybe_serialize(rd.rollout_variance_noises),
+            "rollout_prev_sample_means": _maybe_serialize(rd.rollout_prev_sample_means),
+            "rollout_noise_std_devs": _maybe_serialize(rd.rollout_noise_std_devs),
+            "rollout_model_outputs": _maybe_serialize(rd.rollout_model_outputs),
+        }
+    serialized_denoising_env = None
+    if rtd.denoising_env:
+        env = rtd.denoising_env
+        serialized_denoising_env = {
+            "image_kwargs": (
+                _maybe_serialize(env.image_kwargs) if env.image_kwargs else None
+            ),
+            "pos_cond_kwargs": (
+                _maybe_serialize(env.pos_cond_kwargs) if env.pos_cond_kwargs else None
+            ),
+            "neg_cond_kwargs": (
+                _maybe_serialize(env.neg_cond_kwargs) if env.neg_cond_kwargs else None
+            ),
+            "guidance": (
+                _maybe_serialize(env.guidance) if env.guidance is not None else None
+            ),
+        }
+    serialized_dit_trajectory = None
+    if rtd.dit_trajectory:
+        dit = rtd.dit_trajectory
+        serialized_dit_trajectory = {
+            "latents": (
+                _maybe_serialize(dit.latents) if dit.latents is not None else None
+            ),
+            "timesteps": serialized_dit_timesteps,
+        }
+    return (
+        serialized_log_probs,
+        serialized_debug_tensors,
+        serialized_denoising_env,
+        serialized_dit_trajectory,
+    )
+
+
+def _build_response(
+    request_id: str, prompt: str, seed: int, rollout: bool, result: OutputBatch
+) -> list[RolloutResponse]:
+    """
+    rollout: bool - set to False when evaluating the model
+    """
+    batch_size = result.output.shape[0]
+    inference_time_s = (
+        result.metrics.total_duration_s
+        if result.metrics and result.metrics.total_duration_s > 0
+        else None
+    )
+    peak_memory_mb = result.peak_memory_mb if result.peak_memory_mb > 0 else None
+    rollout_trajectory_data = result.rollout_trajectory_data
+    if rollout:
+        assert (
+            rollout_trajectory_data is not None
+        ), "rollout_trajectory_data must be present when rollout=True"
+
+    serialized_dit_timesteps = None
+    if rollout and rollout_trajectory_data and rollout_trajectory_data.dit_trajectory:
+        serialized_dit_timesteps = _maybe_serialize(
+            rollout_trajectory_data.dit_trajectory.timesteps
+        )
+
+    responses: list[RolloutResponse] = []
+    for sample_idx in range(batch_size):
+        out_i = result.output[sample_idx]
+        if isinstance(out_i, torch.Tensor):
+            out_i = out_i.contiguous()
+        serialized_generated_output = _maybe_serialize(out_i)
+        if not rollout:
+            responses.append(
+                RolloutResponse(
+                    request_id=request_id,
+                    prompt=prompt,
+                    seed=seed,
+                    generated_output=serialized_generated_output,
+                    inference_time_s=inference_time_s,
+                    peak_memory_mb=peak_memory_mb,
+                )
+            )
+            continue
+        per_sample_trajectory = _slice_rollout_trajectory_for_sample(
+            result.rollout_trajectory_data, sample_idx, batch_size
+        )
+        (
+            serialized_log_probs,
+            serialized_debug_tensors,
+            serialized_denoising_env,
+            serialized_dit_trajectory,
+        ) = _serialize_rollout_trajectory(
+            per_sample_trajectory,
+            serialized_dit_timesteps=serialized_dit_timesteps,
+        )
+        responses.append(
+            RolloutResponse(
+                request_id=request_id,
+                prompt=prompt,
+                seed=seed,
+                generated_output=serialized_generated_output,
+                rollout_log_probs=serialized_log_probs,
+                rollout_debug_tensors=serialized_debug_tensors,
+                denoising_env=serialized_denoising_env,
+                dit_trajectory=serialized_dit_trajectory,
+                inference_time_s=inference_time_s,
+                peak_memory_mb=peak_memory_mb,
+            )
+        )
+    return responses
+
+
+def _build_sampling_kwargs(request: RolloutRequest) -> dict:
+    sampling_kwargs: dict = dict(
+        prompt=request.prompt,
+        negative_prompt=request.negative_prompt,
+        seed=request.seed,
+        generator_device=request.generator_device,
+        width=request.width,
+        height=request.height,
+        num_inference_steps=request.num_inference_steps,
+        num_outputs_per_prompt=request.num_outputs_per_prompt,
+        guidance_scale=request.guidance_scale,
+        true_cfg_scale=request.true_cfg_scale,
+        num_frames=request.num_frames,
+        fps=request.fps,
+        image_path=request.image_path,
+        rollout=request.rollout,
+        rollout_sde_type=request.rollout_sde_type,
+        rollout_noise_level=request.rollout_noise_level,
+        rollout_log_prob_no_const=request.rollout_log_prob_no_const,
+        rollout_debug_mode=request.rollout_debug_mode,
+        rollout_return_denoising_env=request.rollout_return_denoising_env,
+        rollout_return_dit_trajectory=request.rollout_return_dit_trajectory,
+        rollout_sde_step_indices=request.rollout_sde_step_indices,
+        rollout_return_step_indices=request.rollout_return_step_indices,
+        suppress_logs=request.suppress_logs,
+        save_output=False,
+        return_trajectory_latents=False,
+        return_trajectory_decoded=False,
+    )
+    if request.extra_sampling_params:
+        sampling_kwargs.update(request.extra_sampling_params)
+        sampling_kwargs["rollout"] = request.rollout
+    return {k: v for k, v in sampling_kwargs.items() if v is not None}
+
+
+@router.post("/generate", response_model=list[RolloutResponse])
+async def rollout_generate(request: RolloutRequest):
+    request_id = generate_request_id()
+    server_args = get_global_server_args()
+    sampling_kwargs = _build_sampling_kwargs(request)
+    try:
+        sampling_params = build_sampling_params(request_id, **sampling_kwargs)
+    except Exception as exc:
+        raise HTTPException(
+            status_code=400, detail=f"Invalid sampling params: {exc}"
+        ) from exc
+    pipeline_request = prepare_request(
+        server_args=server_args, sampling_params=sampling_params
+    )
+    try:
+        output_batch: OutputBatch = await async_scheduler_client.forward(
+            pipeline_request
+        )
+    except Exception as exc:
+        logger.error("Rollout generation failed: %s", exc, exc_info=True)
+        raise HTTPException(
+            status_code=500, detail=f"Generation failed: {exc}"
+        ) from exc
+    if output_batch.error:
+        raise HTTPException(status_code=500, detail=output_batch.error)
+    rollout_responses = _build_response(
+        request_id, request.prompt, request.seed, request.rollout, output_batch
+    )
+    return ORJSONResponse(content=[r.model_dump() for r in rollout_responses])

python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py

@@ -0,0 +1,48 @@
+"""Tensor serialization for post-training / rollout HTTP responses."""
+
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+import numpy as np
+import torch
+from safetensors.torch import load, save
+
+
+def tensor_to_base64(t: torch.Tensor) -> str:
+    t = t.detach().contiguous().cpu()
+    raw = save({"t": t})
+    return base64.b64encode(raw).decode("ascii")
+
+
+def base64_to_tensor(s: str) -> torch.Tensor:
+    raw = base64.b64decode(s)
+    return load(raw)["t"]
+
+
+def _maybe_serialize(obj: Any) -> Any:
+    if isinstance(obj, torch.Tensor):
+        return {
+            "__tensor__": True,
+            "data": tensor_to_base64(obj),
+            "shape": list(obj.shape),
+            "dtype": str(obj.dtype),
+        }
+    if isinstance(obj, np.ndarray):
+        return _maybe_serialize(torch.from_numpy(obj))
+    if isinstance(obj, dict):
+        return {k: _maybe_serialize(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_maybe_serialize(v) for v in obj]
+    return obj
+
+
+def _maybe_deserialize(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        if obj.get("__tensor__"):
+            return base64_to_tensor(obj["data"])
+        return {k: _maybe_deserialize(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_maybe_deserialize(v) for v in obj]
+    return obj

python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py

@@ -0,0 +1,199 @@
+"""Weight update API for the diffusion engine."""
+
+from fastapi import APIRouter, Request
+
+from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import (
+    GetWeightsChecksumReqInput,
+    ReleaseMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqInput,
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightFromTensorCheckerReqInput,
+    UpdateWeightFromTensorReqInput,
+)
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.srt.utils.json_response import orjson_response
+
+router = APIRouter()
+
+
+@router.post("/update_weights_from_disk")
+async def update_weights_from_disk(request: Request):
+    """Update model weights from disk inplace without restarting the server."""
+    body = await request.json()
+    model_path = body.get("model_path")
+    if not model_path:
+        return orjson_response(
+            {"success": False, "message": "model_path is required"},
+            status_code=400,
+        )
+
+    req = UpdateWeightFromDiskReqInput(
+        model_path=model_path,
+        flush_cache=body.get("flush_cache", True),
+        target_modules=body.get("target_modules"),
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response(
+            {"success": False, "message": str(e)},
+            status_code=500,
+        )
+
+    if response.output is None:
+        return orjson_response(
+            {
+                "success": False,
+                "message": response.error or "Unknown status",
+            },
+            status_code=500,
+        )
+
+    result = response.output
+    return orjson_response(
+        result,
+        status_code=200 if result["success"] else 400,
+    )
+
+
+@router.post("/update_weights_from_tensor")
+async def update_weights_from_tensor(request: Request):
+    """Update model weights from serialized tensor payloads."""
+    body = await request.json()
+    serialized_named_tensors = body.get("serialized_named_tensors")
+    if not serialized_named_tensors:
+        return orjson_response(
+            {"success": False, "message": "serialized_named_tensors is required"},
+            status_code=400,
+        )
+
+    req = UpdateWeightFromTensorReqInput(
+        serialized_named_tensors=serialized_named_tensors,
+        load_format=body.get("load_format"),
+        target_modules=body.get("target_modules"),
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response(
+            {"success": False, "message": str(e)},
+            status_code=500,
+        )
+
+    result = response.output
+    return orjson_response(
+        result,
+        status_code=200 if result["success"] else 400,
+    )
+
+
+@router.post("/update_weights_from_tensor_checker")
+async def update_weights_from_tensor_checker(request: Request):
+    """Verify live module weights against expected SHA-256 values."""
+    body = await request.json()
+    target_module = body.get("target_module")
+    if not target_module:
+        return orjson_response(
+            {"success": False, "message": "target_module is required"},
+            status_code=400,
+        )
+
+    expected_named_tensors_sha256 = body.get("expected_named_tensors_sha256")
+    if (
+        not isinstance(expected_named_tensors_sha256, dict)
+        or not expected_named_tensors_sha256
+    ):
+        return orjson_response(
+            {
+                "success": False,
+                "message": "expected_named_tensors_sha256 is required",
+            },
+            status_code=400,
+        )
+
+    req = UpdateWeightFromTensorCheckerReqInput(
+        target_module=target_module,
+        expected_named_tensors_sha256=expected_named_tensors_sha256,
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response(
+            {"success": False, "message": str(e)},
+            status_code=500,
+        )
+
+    result = response.output
+    success = result.get("success", False)
+    message = result.get("message", "Unknown status")
+    return orjson_response(
+        {"success": success, "message": message},
+        status_code=200 if success else 400,
+    )
+
+
+@router.post("/get_weights_checksum")
+async def get_weights_checksum(request: Request):
+    """Return SHA-256 checksum of each requested module's weights."""
+    body = await request.json()
+    req = GetWeightsChecksumReqInput(
+        module_names=body.get("module_names"),
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response({"error": str(e)}, status_code=500)
+
+    return orjson_response(response.output, status_code=200)
+
+
+@router.post("/release_memory_occupation")
+async def release_memory_occupation():
+    """Release GPU memory occupation (sleep the engine)."""
+    try:
+        response = await async_scheduler_client.forward(
+            ReleaseMemoryOccupationReqInput()
+        )
+    except Exception as e:
+        return orjson_response({"success": False, "message": str(e)}, status_code=500)
+
+    if response.output is None:
+        return orjson_response(
+            {
+                "success": False,
+                "message": response.error or "Unknown status",
+            },
+            status_code=500,
+        )
+
+    payload = response.output
+    success = bool(payload["success"])
+    return orjson_response(payload, status_code=200 if success else 400)
+
+
+@router.post("/resume_memory_occupation")
+async def resume_memory_occupation():
+    """Resume GPU memory occupation (wake the engine)."""
+    try:
+        response = await async_scheduler_client.forward(
+            ResumeMemoryOccupationReqInput()
+        )
+    except Exception as e:
+        return orjson_response({"success": False, "message": str(e)}, status_code=500)
+
+    if response.output is None:
+        return orjson_response(
+            {
+                "success": False,
+                "message": response.error or "Unknown status",
+            },
+            status_code=500,
+        )
+
+    payload = response.output
+    success = bool(payload["success"])
+    return orjson_response(payload, status_code=200 if success else 400)

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

@@ -0,0 +1,755 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+"""
+DiffGenerator module for sglang-diffusion.
+
+This module provides a consolidated interface for generating videos using
+diffusion models.
+"""
+
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from copy import copy
+from dataclasses import dataclass, field
+from typing import Any, Callable, List, Optional, Sequence, Union
+
+import imageio
+import numpy as np
+import torch
+from PIL import Image
+
+try:
+    import scipy.io.wavfile as scipy_wavfile
+except ImportError:  # pragma: no cover
+    scipy_wavfile = None
+
+try:
+    import imageio_ffmpeg as _imageio_ffmpeg
+except ImportError:  # pragma: no cover
+    _imageio_ffmpeg = None
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
+from sglang.srt.observability.trace import TraceReqContext
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class SetLoraReq:
+    lora_nickname: Union[str, List[str]]
+    lora_path: Optional[Union[str, List[Optional[str]]]] = None
+    target: Union[str, List[str]] = "all"
+    strength: Union[float, List[float]] = 1.0
+    merge_mode: Optional[str] = None
+
+
+@dataclass
+class MergeLoraWeightsReq:
+    target: str = "all"
+    strength: float = 1.0
+
+
+@dataclass
+class UnmergeLoraWeightsReq:
+    target: str = "all"
+
+
+@dataclass
+class ListLorasReq:
+    pass
+
+
+@dataclass
+class ShutdownReq:
+    pass
+
+
+@dataclass
+class ReleaseRealtimeSessionReq:
+    session_id: str
+
+
+@dataclass
+class GetDisaggStatsReq:
+    """Request to get disagg pipeline metrics from the scheduler."""
+
+    pass
+
+
+def format_lora_message(
+    lora_nickname: Union[str, List[str]],
+    target: Union[str, List[str]],
+    strength: Union[float, List[float]],
+) -> tuple[str, str, str]:
+    """Format success message for single or multiple LoRAs."""
+    if isinstance(lora_nickname, list):
+        nickname_str = ", ".join(lora_nickname)
+        target_str = ", ".join(target) if isinstance(target, list) else target
+        strength_str = (
+            ", ".join(f"{s:.2f}" for s in strength)
+            if isinstance(strength, list)
+            else f"{strength:.2f}"
+        )
+    else:
+        nickname_str = lora_nickname
+        target_str = target if isinstance(target, str) else ", ".join(target)
+        strength_str = (
+            f"{strength:.2f}"
+            if isinstance(strength, (int, float))
+            else ", ".join(f"{s:.2f}" for s in strength)
+        )
+    return nickname_str, target_str, strength_str
+
+
+@dataclass
+class GenerationResult:
+    """Result of a single generation request from DiffGenerator."""
+
+    samples: Any = None
+    frames: Any = None
+    audio: Any = None
+    action: Any = None  # [T, raw_action_dim] predicted action (policy/inverse_dynamics)
+    prompt: str | None = None
+    size: tuple | None = None  # (height, width, num_frames)
+    generation_time: float = 0.0
+    peak_memory_mb: float = 0.0
+    metrics: dict = field(default_factory=dict)
+    trajectory_latents: Any = None
+    trajectory_timesteps: Any = None
+    rollout_trajectory_data: Any = None
+    trajectory_decoded: Any = None
+    prompt_index: int = 0
+    output_file_path: str | None = None
+
+
+@dataclass
+class MaterializedOutput:
+    sample: Any
+    frames: list[Any]
+    audio: Any = None
+    fps: int = 0
+
+
+def normalize_output_seeds(
+    seed: int | list[int],
+    *,
+    num_outputs_per_prompt: int,
+    num_prompts: int = 1,
+    prompt_index: int = 0,
+) -> list[int]:
+    """
+    return a list of seed with size equal to `num_outputs_per_prompt`
+    """
+    if num_outputs_per_prompt <= 0:
+        raise ValueError(
+            f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}"
+        )
+
+    if isinstance(seed, list):
+        seeds = [int(item) for item in seed]
+        total_outputs = num_outputs_per_prompt * num_prompts
+        if len(seeds) == num_outputs_per_prompt:
+            return seeds
+        if len(seeds) == total_outputs:
+            start = prompt_index * num_outputs_per_prompt
+            return seeds[start : start + num_outputs_per_prompt]
+        raise ValueError(
+            "seed list length must match num_outputs_per_prompt "
+            f"({num_outputs_per_prompt}) or total outputs ({total_outputs}), "
+            f"got {len(seeds)}"
+        )
+
+    base_seed = int(seed)
+    return [base_seed + i for i in range(num_outputs_per_prompt)]
+
+
+def _with_output_index_suffix(output_file_name: str, output_index: int) -> str:
+    base, ext = os.path.splitext(output_file_name)
+    return f"{base}_{output_index}{ext}"
+
+
+def _copy_trace_ctx_for_output(req: Req, request_id: str | None, output_index: int):
+    trace_ctx = req.trace_ctx
+    if output_index == 0 or not trace_ctx.tracing_enable:
+        return trace_ctx
+
+    output_trace_ctx = TraceReqContext(
+        rid=request_id,
+        module_name=trace_ctx.module_name,
+        external_trace_header=trace_ctx.external_trace_header,
+    )
+    output_trace_ctx.trace_req_start()
+    return output_trace_ctx
+
+
+def _copy_req_for_output(
+    req: Req,
+    *,
+    request_id: str | None,
+    output_index: int,
+) -> Req:
+    """Create a lightweight per-output ``Req`` without deep-copying tensors."""
+    output_req = copy(req)
+    output_req.sampling_params = copy(req.sampling_params)
+    output_req.extra = dict(req.extra)
+    output_req.condition_inputs = dict(req.condition_inputs)
+    output_req.trace_ctx = _copy_trace_ctx_for_output(req, request_id, output_index)
+    return output_req
+
+
+def expand_request_outputs(
+    req: Req,
+    *,
+    num_prompts: int = 1,
+    prompt_index: int = 0,
+) -> list[Req]:
+    """
+    Expand a req to a list with size equal to `num_prompts`
+    """
+    num_outputs = int(req.num_outputs_per_prompt)
+    # each req must has different seed
+    seeds = normalize_output_seeds(
+        req.seed,
+        num_outputs_per_prompt=num_outputs,
+        num_prompts=num_prompts,
+        prompt_index=prompt_index,
+    )
+
+    if num_outputs == 1:
+        req.seed = seeds[0]
+        req.seeds = None
+        req.generator = None
+        return [req]
+
+    expanded: list[Req] = []
+    for output_index, seed in enumerate(seeds):
+        output_request_id = (
+            f"{req.request_id}:{output_index}" if req.request_id is not None else None
+        )
+        output_req = _copy_req_for_output(
+            req, request_id=output_request_id, output_index=output_index
+        )
+        output_req.seed = seed
+        output_req.num_outputs_per_prompt = 1
+        output_req.seeds = None
+        output_req.generator = None
+        output_req.extra["parent_request_id"] = req.request_id
+        output_req.extra["output_index"] = output_index
+
+        if output_request_id is not None:
+            output_req.request_id = output_request_id
+
+        if req.output_file_name:
+            output_req.output_file_name = _with_output_index_suffix(
+                req.output_file_name, output_index
+            )
+        output_req.validate()
+        expanded.append(output_req)
+
+    return expanded
+
+
+def _normalize_audio_to_numpy(audio: Any) -> np.ndarray | None:
+    """Convert audio (torch / numpy) into a float32 numpy array in [-1, 1], best-effort."""
+    if audio is None:
+        return None
+    if isinstance(audio, torch.Tensor):
+        audio_np = audio.detach().float().clamp(-1.0, 1.0).cpu().numpy()
+    elif isinstance(audio, np.ndarray):
+        audio_np = audio.astype(np.float32, copy=False)
+        audio_np = np.clip(audio_np, -1.0, 1.0)
+    else:
+        return None
+
+    # 1. Squeeze leading singleton dimensions (Batch, etc.)
+    while audio_np.ndim > 1 and audio_np.shape[0] == 1:
+        audio_np = audio_np.squeeze(0)
+
+    # 2. Handle (C, L) -> (L, C)
+    if audio_np.ndim == 2 and audio_np.shape[0] < audio_np.shape[1]:
+        audio_np = audio_np.transpose(1, 0)
+
+    # 3. Final safety check: if still 2D and channels (dim 1) is huge, something is wrong
+    if audio_np.ndim == 2 and audio_np.shape[1] > 256 and audio_np.shape[0] == 1:
+        audio_np = audio_np.flatten()
+
+    return audio_np
+
+
+def _pick_audio_sample_rate(
+    *,
+    audio_np: np.ndarray,
+    audio_sample_rate: Optional[int],
+    fps: int,
+    num_frames: int,
+) -> int:
+    """Pick a plausible sample rate, falling back to inferring from video duration."""
+    selected_sr = int(audio_sample_rate) if audio_sample_rate is not None else None
+    if selected_sr is None or not (8000 <= selected_sr <= 192000):
+        selected_sr = 24000
+        try:
+            duration_s = float(num_frames) / float(fps) if fps else 0.0
+            if duration_s > 0:
+                audio_len = (
+                    int(audio_np.shape[0])
+                    if audio_np.ndim == 2
+                    else int(audio_np.shape[-1])
+                )
+                inferred_sr = int(round(float(audio_len) / duration_s))
+                if 8000 <= inferred_sr <= 192000:
+                    selected_sr = inferred_sr
+        except Exception:
+            pass
+    return selected_sr
+
+
+def _resolve_ffmpeg_exe() -> str:
+    ffmpeg_exe = "ffmpeg"
+    ffmpeg_on_path = shutil.which("ffmpeg")
+    if ffmpeg_on_path:
+        ffmpeg_exe = ffmpeg_on_path
+    try:
+        if _imageio_ffmpeg is not None:
+            ffmpeg_exe = _imageio_ffmpeg.get_ffmpeg_exe()
+    except Exception:
+        pass
+
+    ffmpeg_ok = False
+    if ffmpeg_exe:
+        if os.path.isabs(ffmpeg_exe):
+            ffmpeg_ok = os.path.exists(ffmpeg_exe)
+        else:
+            ffmpeg_ok = shutil.which(ffmpeg_exe) is not None
+    if not ffmpeg_ok:
+        raise RuntimeError("ffmpeg not found")
+    return ffmpeg_exe
+
+
+def _mux_audio_np_into_mp4(
+    *,
+    save_file_path: str,
+    audio_np: np.ndarray,
+    sample_rate: int,
+    ffmpeg_exe: str,
+) -> None:
+    merged_path = save_file_path.rsplit(".", 1)[0] + ".tmp_mux.mp4"
+    tmp_wav_path = None
+    try:
+        if scipy_wavfile is None:
+            raise RuntimeError(
+                "scipy is required to mux audio into mp4 (pip install scipy)"
+            )
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            tmp_wav_path = f.name
+        scipy_wavfile.write(tmp_wav_path, sample_rate, audio_np)
+        subprocess.run(
+            [
+                ffmpeg_exe,
+                "-y",
+                "-i",
+                save_file_path,
+                "-i",
+                tmp_wav_path,
+                "-c:v",
+                "copy",
+                "-c:a",
+                "aac",
+                "-strict",
+                "experimental",
+                merged_path,
+            ],
+            check=True,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        os.replace(merged_path, save_file_path)
+    finally:
+        if tmp_wav_path:
+            try:
+                os.remove(tmp_wav_path)
+            except OSError:
+                pass
+        if os.path.exists(merged_path):
+            try:
+                os.remove(merged_path)
+            except OSError:
+                pass
+
+
+def _maybe_mux_audio_into_mp4(
+    *,
+    save_file_path: str,
+    audio: Any,
+    frames: list,
+    fps: int,
+    audio_sample_rate: Optional[int],
+) -> None:
+    """Best-effort mux audio into an already-written mp4 at save_file_path.
+
+    Any failure should keep the silent video and only log a warning.
+    """
+    audio_np = _normalize_audio_to_numpy(audio)
+    if audio_np is None:
+        return
+    selected_sr = _pick_audio_sample_rate(
+        audio_np=audio_np,
+        audio_sample_rate=audio_sample_rate,
+        fps=fps,
+        num_frames=len(frames),
+    )
+
+    try:
+        ffmpeg_exe = _resolve_ffmpeg_exe()
+        _mux_audio_np_into_mp4(
+            save_file_path=save_file_path,
+            audio_np=audio_np,
+            sample_rate=selected_sr,
+            ffmpeg_exe=ffmpeg_exe,
+        )
+    except Exception as e:
+        logger.warning(
+            "Failed to mux audio into mp4 (saved silent video): %s",
+            str(e),
+        )
+
+
+def prepare_request(
+    server_args: ServerArgs,
+    sampling_params: SamplingParams,
+    external_trace_header: dict[str, str] | None = None,
+) -> Req:
+    """
+    Create a Req object with sampling_params as a parameter.
+    """
+    req = Req(
+        sampling_params=sampling_params,
+        VSA_sparsity=server_args.attention_backend_config.VSA_sparsity,
+    )
+    sampling_params.apply_request_extra(req)
+    if getattr(sampling_params, "max_sequence_length", None) is not None:
+        req.max_sequence_length = sampling_params.max_sequence_length
+
+    diffusers_kwargs = getattr(sampling_params, "diffusers_kwargs", None)
+    if diffusers_kwargs and "max_sequence_length" in diffusers_kwargs:
+        req.max_sequence_length = diffusers_kwargs["max_sequence_length"]
+
+    if not isinstance(req.prompt, str):
+        raise TypeError(f"`prompt` must be a string, but got {type(req.prompt)}")
+
+    req_width = getattr(req, "width", None)
+    req_height = getattr(req, "height", None)
+    if (req_width is not None and req_width <= 0) or (
+        req_height is not None and req_height <= 0
+    ):
+        raise ValueError(
+            f"Height and width must be positive, got height={req_height}, width={req_width}"
+        )
+
+    if server_args.enable_trace:
+        trace_ctx = TraceReqContext(
+            rid=sampling_params.request_id,
+            module_name="diffusion",
+            external_trace_header=external_trace_header,
+        )
+        trace_ctx.trace_req_start()
+        req.trace_ctx = trace_ctx
+
+    return req
+
+
+def attach_audio_to_video_sample(
+    sample: Any,
+    audio: Any,
+    output_idx: int,
+) -> Any:
+    """Attach per-sample audio for video outputs when available."""
+    audio = select_output_audio(audio, output_idx)
+    if audio is None:
+        return sample
+    if not (isinstance(sample, (tuple, list)) and len(sample) == 2):
+        return (sample, audio)
+    return sample
+
+
+def select_output_audio(audio: Any, output_idx: int) -> Any:
+    if isinstance(audio, torch.Tensor) and audio.ndim >= 2:
+        return audio[output_idx] if audio.shape[0] > output_idx else None
+    if isinstance(audio, np.ndarray) and audio.ndim >= 2:
+        return audio[output_idx] if audio.shape[0] > output_idx else None
+    return audio
+
+
+def _split_sample_audio(sample: Any) -> tuple[Any, Any]:
+    if isinstance(sample, (tuple, list)) and len(sample) == 2:
+        return sample[0], sample[1]
+    return sample, None
+
+
+def _sample_to_uint8_frames(sample: Any) -> list[Any]:
+    """return numpy frames in THCW format"""
+    if isinstance(sample, torch.Tensor):
+        # sample is raw tensor
+        if sample.dim() == 3:
+            sample = sample.unsqueeze(1)
+        sample = (sample * 255).clamp(0, 255).to(torch.uint8)
+        videos = sample.permute(1, 2, 3, 0).contiguous().cpu().numpy()
+        return list(videos)
+
+    if not isinstance(sample, np.ndarray):
+        raise TypeError(f"Unsupported sample type: {type(sample)}")
+
+    # sample is numpy frames
+    arr = sample
+    if arr.ndim == 3:
+        if arr.shape[-1] in (1, 3, 4):
+            arr = arr[None, ...]
+        else:
+            arr = arr[..., None]
+    if arr.ndim != 4:
+        raise ValueError(f"Unexpected numpy sample shape: {tuple(arr.shape)}")
+
+    if arr.shape[-1] not in (1, 3, 4) and arr.shape[0] in (1, 3, 4):
+        t = torch.from_numpy(arr)
+        if t.dim() == 3:
+            t = t.unsqueeze(1)
+        t = (t * 255).clamp(0, 255).to(torch.uint8)
+        videos = t.permute(1, 2, 3, 0).contiguous().cpu().numpy()
+        return list(videos)
+
+    if arr.dtype != np.uint8:
+        arr = (np.clip(arr, 0.0, 1.0) * 255.0).astype(np.uint8)
+    return list(arr)
+
+
+def materialize_output_sample(
+    sample: Any,
+    data_type: DataType,
+    fps: int,
+    *,
+    enable_frame_interpolation: bool = False,
+    frame_interpolation_exp: int = 1,
+    frame_interpolation_scale: float = 1.0,
+    frame_interpolation_model_path: Optional[str] = None,
+    enable_upscaling: bool = False,
+    upscaling_model_path: Optional[str] = None,
+    upscaling_scale: int = 4,
+) -> MaterializedOutput:
+    """materialize samples, apply postprocessing if applicable"""
+    sample_without_audio, audio = _split_sample_audio(sample)
+    frames = _sample_to_uint8_frames(sample_without_audio)
+
+    # frames are uint8 numpy arrays in THWC format at this point
+    if enable_frame_interpolation and data_type == DataType.VIDEO and len(frames) > 1:
+        from sglang.multimodal_gen.runtime.postprocess import (
+            interpolate_video_frames,
+        )
+
+        frames, multiplier = interpolate_video_frames(
+            frames,
+            exp=frame_interpolation_exp,
+            scale=frame_interpolation_scale,
+            model_path=frame_interpolation_model_path,
+        )
+        fps = fps * multiplier
+
+    if enable_upscaling and frames:
+        from sglang.multimodal_gen.runtime.postprocess import upscale_frames
+
+        frames = upscale_frames(
+            frames,
+            model_path=upscaling_model_path,
+            scale=upscaling_scale,
+        )
+
+    return MaterializedOutput(sample=sample, frames=frames, audio=audio, fps=fps)
+
+
+def save_materialized_output(
+    materialized: MaterializedOutput,
+    data_type: DataType,
+    save_file_path: Optional[str],
+    *,
+    save_output: bool = True,
+    audio_sample_rate: Optional[int] = None,
+    output_compression: Optional[int] = None,
+) -> None:
+    if not save_output:
+        return
+    if not save_file_path:
+        logger.info("No output path provided, output not saved")
+        return
+
+    os.makedirs(os.path.dirname(save_file_path), exist_ok=True)
+    if data_type == DataType.VIDEO:
+        quality = output_compression / 10 if output_compression is not None else 5
+        imageio.mimsave(
+            save_file_path,
+            materialized.frames,
+            fps=materialized.fps,
+            format=data_type.get_default_extension(),
+            codec="libx264",
+            quality=quality,
+        )
+
+        _maybe_mux_audio_into_mp4(
+            save_file_path=save_file_path,
+            audio=materialized.audio,
+            frames=materialized.frames,
+            fps=materialized.fps,
+            audio_sample_rate=audio_sample_rate,
+        )
+    else:
+        quality = output_compression if output_compression is not None else 75
+        if len(materialized.frames) > 1:
+            for i, image in enumerate(materialized.frames):
+                parts = save_file_path.rsplit(".", 1)
+                if len(parts) == 2:
+                    indexed_path = f"{parts[0]}_{i}.{parts[1]}"
+                else:
+                    indexed_path = f"{save_file_path}_{i}"
+                _save_image_frame(indexed_path, image, quality, output_compression)
+        else:
+            _save_image_frame(
+                save_file_path, materialized.frames[0], quality, output_compression
+            )
+    logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}")
+
+
+def _save_image_frame(
+    path: str, frame: np.ndarray, quality: int | None, output_compression: int | None
+) -> None:
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".png":
+        compress_level = 1
+        if output_compression is not None and output_compression != 75:
+            compress_level = max(0, min(9, round(output_compression / 100 * 9)))
+        if frame.ndim == 3 and frame.shape[-1] == 1:
+            frame = frame[..., 0]
+        Image.fromarray(frame).save(path, format="PNG", compress_level=compress_level)
+    else:
+        imageio.imwrite(path, frame, quality=quality)
+
+
+def save_outputs(
+    outputs: Sequence[Any],
+    data_type: DataType,
+    fps: int,
+    save_output: bool,
+    build_output_path: Callable[[int], str],
+    *,
+    audio: Any = None,
+    audio_sample_rate: Optional[int] = None,
+    samples_out: Optional[list[Any]] = None,
+    audios_out: Optional[list[Any]] = None,
+    frames_out: Optional[list[Any]] = None,
+    output_compression: Optional[int] = None,
+    enable_frame_interpolation: bool = False,
+    frame_interpolation_exp: int = 1,
+    frame_interpolation_scale: float = 1.0,
+    frame_interpolation_model_path: Optional[str] = None,
+    enable_upscaling: bool = False,
+    upscaling_model_path: Optional[str] = None,
+    upscaling_scale: int = 4,
+) -> list[str]:
+    output_paths: list[str] = []
+    for idx, sample in enumerate(outputs):
+        save_file_path = build_output_path(idx)
+        if data_type == DataType.ACTION:
+            if samples_out is not None:
+                samples_out.append(sample)
+            if audios_out is not None:
+                audios_out.append(None)
+            if frames_out is not None:
+                frames_out.append([])
+            if save_output and save_file_path:
+                os.makedirs(os.path.dirname(save_file_path) or ".", exist_ok=True)
+                with open(save_file_path, "w", encoding="utf-8") as f:
+                    json.dump(sample, f, ensure_ascii=False)
+                logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}")
+            output_paths.append(save_file_path)
+            continue
+
+        if data_type == DataType.VIDEO:
+            sample = attach_audio_to_video_sample(sample, audio, idx)
+
+        frames = post_process_sample(
+            sample,
+            data_type,
+            fps,
+            save_output,
+            save_file_path,
+            audio_sample_rate=audio_sample_rate,
+            output_compression=output_compression,
+            enable_frame_interpolation=enable_frame_interpolation,
+            frame_interpolation_exp=frame_interpolation_exp,
+            frame_interpolation_scale=frame_interpolation_scale,
+            frame_interpolation_model_path=frame_interpolation_model_path,
+            enable_upscaling=enable_upscaling,
+            upscaling_model_path=upscaling_model_path,
+            upscaling_scale=upscaling_scale,
+        )
+
+        if samples_out is not None:
+            samples_out.append(sample)
+        if audios_out is not None:
+            if data_type == DataType.VIDEO:
+                audios_out.append(select_output_audio(audio, idx))
+            else:
+                audios_out.append(audio)
+        if frames_out is not None:
+            frames_out.append(frames)
+        output_paths.append(save_file_path)
+    return output_paths
+
+
+def post_process_sample(
+    sample: Any,
+    data_type: DataType,
+    fps: int,
+    save_output: bool = True,
+    save_file_path: Optional[str] = None,
+    audio_sample_rate: Optional[int] = None,
+    output_compression: Optional[int] = None,
+    enable_frame_interpolation: bool = False,
+    frame_interpolation_exp: int = 1,
+    frame_interpolation_scale: float = 1.0,
+    frame_interpolation_model_path: Optional[str] = None,
+    enable_upscaling: bool = False,
+    upscaling_model_path: Optional[str] = None,
+    upscaling_scale: int = 4,
+) -> list[Any]:
+    """materialize frames and save outputs (optional)"""
+    if data_type == DataType.ACTION:
+        return []
+
+    materialized = materialize_output_sample(
+        sample,
+        data_type,
+        fps,
+        enable_frame_interpolation=enable_frame_interpolation,
+        frame_interpolation_exp=frame_interpolation_exp,
+        frame_interpolation_scale=frame_interpolation_scale,
+        frame_interpolation_model_path=frame_interpolation_model_path,
+        enable_upscaling=enable_upscaling,
+        upscaling_model_path=upscaling_model_path,
+        upscaling_scale=upscaling_scale,
+    )
+    save_materialized_output(
+        materialized,
+        data_type,
+        save_file_path,
+        save_output=save_output,
+        audio_sample_rate=audio_sample_rate,
+        output_compression=output_compression,
+    )
+    return materialized.frames

python/sglang/multimodal_gen/runtime/entrypoints/vla/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

python/sglang/multimodal_gen/runtime/entrypoints/vla/api.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from fastapi import APIRouter, HTTPException, Request, Response, WebSocket
+
+from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
+    action_generation_response,
+    action_metadata,
+    action_raw_response,
+    infer_action,
+    pack_msgpack,
+    unpack_msgpack,
+)
+from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import (
+    run_action_msgpack_ws,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.srt.utils.json_response import orjson_response
+
+router = APIRouter(prefix="/v1/actions", tags=["actions"])
+
+
+def _wants_msgpack(request: Request) -> bool:
+    content_type = request.headers.get("content-type", "").lower()
+    accept = request.headers.get("accept", "").lower()
+    return "msgpack" in content_type or "msgpack" in accept
+
+
+def _response_format(payload: dict) -> str:
+    runtime = payload.get("runtime") or {}
+    response_format = str(runtime.get("response_format", "envelope")).lower()
+    if response_format not in ("envelope", "raw"):
+        raise ValueError("runtime.response_format must be 'envelope' or 'raw'")
+    return response_format
+
+
+def _prefer_numpy_output(payload: dict) -> None:
+    runtime = payload.setdefault("runtime", {})
+    runtime.setdefault("output_format", "numpy")
+
+
+@router.post("/generations")
+async def create_action_generation(request: Request):
+    server_args: ServerArgs = request.app.state.server_args
+    try:
+        if "msgpack" in request.headers.get("content-type", "").lower():
+            payload = unpack_msgpack(await request.body())
+        else:
+            payload = await request.json()
+        wants_msgpack = _wants_msgpack(request)
+        if wants_msgpack:
+            _prefer_numpy_output(payload)
+        output = await infer_action(payload, server_args)
+        if _response_format(payload) == "raw":
+            response = action_raw_response(output, preserve_numpy=wants_msgpack)
+        else:
+            response = action_generation_response(
+                output,
+                server_args,
+                preserve_numpy=wants_msgpack,
+            )
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    if wants_msgpack:
+        return Response(
+            content=pack_msgpack(response), media_type="application/msgpack"
+        )
+    return orjson_response(response)
+
+
+@router.get("/metadata")
+async def action_metadata_endpoint(request: Request):
+    return orjson_response(action_metadata(request.app.state.server_args))
+
+
+@router.websocket("/realtime")
+async def action_realtime_ws(websocket: WebSocket):
+    server_args: ServerArgs = websocket.app.state.server_args
+    await run_action_msgpack_ws(
+        websocket,
+        server_args,
+        prepare_payload=_prefer_numpy_output,
+        build_response=lambda output: action_generation_response(
+            output,
+            server_args,
+            preserve_numpy=True,
+        ),
+    )

python/sglang/multimodal_gen/runtime/entrypoints/vla/openpi.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, WebSocket
+
+from sglang.multimodal_gen.runtime.entrypoints.vla.ws_utils import (
+    run_action_msgpack_ws,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+router = APIRouter()
+
+
+def _prefer_numpy_output(observation: dict[str, Any]) -> None:
+    observation.setdefault("output_format", "numpy")
+
+
+@router.websocket("/openpi/policy")
+async def openpi_policy_ws(websocket: WebSocket):
+    server_args: ServerArgs = websocket.app.state.server_args
+    await run_action_msgpack_ws(
+        websocket,
+        server_args,
+        prepare_payload=_prefer_numpy_output,
+        build_response=lambda output: output,
+    )

python/sglang/multimodal_gen/runtime/entrypoints/vla/protocol.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import base64
+import dataclasses
+import io
+import time
+import uuid
+from functools import lru_cache
+from typing import Any
+
+import numpy as np
+from PIL import Image
+
+from sglang.multimodal_gen.configs.sample.vla import VLASamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+def pack_numpy_payload(obj):
+    if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {obj.dtype}")
+    if isinstance(obj, np.ndarray):
+        return {
+            b"__ndarray__": True,
+            b"data": obj.tobytes(),
+            b"dtype": obj.dtype.str,
+            b"shape": obj.shape,
+        }
+    if isinstance(obj, np.generic):
+        return {
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
+        }
+    return obj
+
+
+def unpack_numpy_payload(obj):
+    ndarray_marker = obj.get("__ndarray__") or obj.get(b"__ndarray__")
+    npgeneric_marker = obj.get("__npgeneric__") or obj.get(b"__npgeneric__")
+    data = obj.get("data", obj.get(b"data"))
+    dtype = obj.get("dtype", obj.get(b"dtype"))
+    shape = obj.get("shape", obj.get(b"shape"))
+    if ndarray_marker:
+        return np.ndarray(
+            buffer=data,
+            dtype=np.dtype(dtype),
+            shape=shape,
+        )
+    if npgeneric_marker:
+        return np.dtype(dtype).type(data)
+    return obj
+
+
+def pack_msgpack(payload: Any) -> bytes:
+    import msgpack
+
+    return msgpack.packb(payload, default=pack_numpy_payload, use_bin_type=True)
+
+
+def unpack_msgpack(payload: bytes) -> Any:
+    import msgpack
+
+    return msgpack.unpackb(payload, object_hook=unpack_numpy_payload, raw=False)
+
+
+def _decode_b64_image(payload: dict[str, Any]) -> Image.Image:
+    data = payload.get("b64_json") or payload.get("base64")
+    if not data:
+        raise ValueError("image payload requires b64_json")
+    if isinstance(data, str) and "," in data and data.startswith("data:"):
+        data = data.split(",", 1)[1]
+    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")
+
+
+def _decode_tensor_payload(payload: dict[str, Any]) -> Any:
+    values = payload.get("values")
+    if values is None:
+        values = payload.get("data")
+    if values is None:
+        return payload
+    dtype = payload.get("dtype")
+    array = np.asarray(values, dtype=np.dtype(dtype) if dtype else None)
+    shape = payload.get("shape")
+    if shape is not None:
+        array = array.reshape(tuple(shape))
+    return array
+
+
+def _normalize_image_value(value: Any) -> Any:
+    if not isinstance(value, dict):
+        return value
+    if "b64_json" in value or "base64" in value:
+        return _decode_b64_image(value)
+    if "values" in value or "data" in value:
+        return _decode_tensor_payload(value)
+    return value
+
+
+def _normalize_observation(observation: dict[str, Any]) -> dict[str, Any]:
+    normalized = dict(observation)
+    images = normalized.get("images")
+    if isinstance(images, dict):
+        normalized["images"] = {
+            name: _normalize_image_value(value) for name, value in images.items()
+        }
+    state = normalized.get("state")
+    if isinstance(state, dict):
+        normalized["state"] = _decode_tensor_payload(state)
+    observation_state = normalized.get("observation.state")
+    if isinstance(observation_state, dict):
+        normalized["observation.state"] = _decode_tensor_payload(observation_state)
+    noise = normalized.get("noise")
+    if isinstance(noise, dict):
+        normalized["noise"] = _decode_tensor_payload(noise)
+    observation_noise = normalized.get("observation.noise")
+    if isinstance(observation_noise, dict):
+        normalized["observation.noise"] = _decode_tensor_payload(observation_noise)
+    return normalized
+
+
+def images_from_observation(
+    observation: dict[str, Any],
+    pipeline_config: Any,
+) -> dict[str, Any]:
+    if isinstance(observation.get("images"), dict):
+        images = dict(observation["images"])
+    else:
+        images = {}
+    for key in pipeline_config.image_keys:
+        if key in observation:
+            images[key] = observation[key]
+        full_key = f"observation.images.{key}"
+        if full_key in observation:
+            images[key] = observation[full_key]
+    return {name: _normalize_image_value(value) for name, value in images.items()}
+
+
+def action_metadata(server_args: ServerArgs) -> dict[str, Any]:
+    pipeline_config = server_args.pipeline_config
+    policy_family = getattr(
+        pipeline_config,
+        "policy_family",
+        type(pipeline_config).__name__.removesuffix("PipelineConfig").lower(),
+    )
+    return {
+        "object": "action.metadata",
+        "model": server_args.model_id or server_args.model_path,
+        "model_path": server_args.model_path,
+        "policy_family": policy_family,
+        "input": {
+            "image_keys": list(pipeline_config.image_keys),
+            "image_size": list(pipeline_config.image_size),
+            "state_dim": pipeline_config.state_dim,
+        },
+        "output": {
+            "action_type": "continuous",
+            "action_horizon": pipeline_config.action_horizon,
+            "action_dim": pipeline_config.output_action_dim,
+            "padded_action_dim": pipeline_config.action_dim,
+            "dtype": "float32",
+        },
+        "runtime": {
+            "materialize_dtype": pipeline_config.materialize_dtype,
+            "enable_autocast": pipeline_config.enable_autocast,
+            "parallelism": {
+                "num_gpus": server_args.num_gpus,
+                "tp_size": server_args.tp_size,
+                "sp_degree": server_args.sp_degree,
+                "ulysses_degree": server_args.ulysses_degree,
+                "ring_degree": server_args.ring_degree,
+                "prefix_strategy": pipeline_config.prefix_parallel_strategy,
+                "action_strategy": pipeline_config.action_parallel_strategy,
+                "layout_version": pipeline_config.parallel_layout_version,
+            },
+        },
+        "defaults": {
+            "num_inference_steps": pipeline_config.default_num_inference_steps,
+            "prefix_cache": (
+                "auto" if pipeline_config.enable_global_prefix_cache else False
+            ),
+            "cuda_graph": "auto" if pipeline_config.enable_action_cuda_graph else False,
+        },
+        "capabilities": {
+            "exact_prefix_cache": True,
+            "cuda_graph": pipeline_config.enable_action_cuda_graph,
+            "realtime_websocket": True,
+            "openpi_websocket": True,
+            "batch_inputs": False,
+            "multiple_candidates": False,
+        },
+    }
+
+
+def _runtime_bool(value: Any, default: bool) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, str):
+        value = value.lower()
+        if value == "auto":
+            return default
+        if value in ("true", "1", "yes"):
+            return True
+        if value in ("false", "0", "no"):
+            return False
+    return bool(value)
+
+
+def _action_request_to_observation(payload: dict[str, Any]) -> dict[str, Any]:
+    if "input" not in payload:
+        return _normalize_observation(payload)
+
+    input_payload = payload.get("input") or {}
+    observation = dict(input_payload.get("observation") or {})
+    if "task" in input_payload:
+        observation["prompt"] = input_payload["task"]
+    elif "prompt" in input_payload:
+        observation["prompt"] = input_payload["prompt"]
+    if "images" in input_payload:
+        observation["images"] = input_payload["images"]
+    if "state" in input_payload:
+        observation["state"] = input_payload["state"]
+    if "noise" in input_payload:
+        observation["noise"] = input_payload["noise"]
+    return _normalize_observation(observation)
+
+
+@lru_cache(maxsize=32)
+def _resolve_action_sampling_params_cls_cached(
+    model_path: str,
+    backend: str | None,
+    model_id: str | None,
+    pipeline_class_name: str | None,
+) -> type[VLASamplingParams]:
+    if pipeline_class_name:
+        from sglang.multimodal_gen.registry import get_pipeline_config_classes
+
+        config_classes = get_pipeline_config_classes(pipeline_class_name)
+        if config_classes is not None:
+            _, sampling_params_cls = config_classes
+            if issubclass(sampling_params_cls, VLASamplingParams):
+                return sampling_params_cls
+
+    from sglang.multimodal_gen.registry import get_model_info
+
+    model_info = get_model_info(
+        model_path,
+        backend=backend,
+        model_id=model_id,
+    )
+    sampling_params_cls = model_info.sampling_param_cls
+    if not issubclass(sampling_params_cls, VLASamplingParams):
+        raise ValueError(
+            f"Action endpoint requires VLASamplingParams, got {sampling_params_cls.__name__}"
+        )
+    return sampling_params_cls
+
+
+def _resolve_action_sampling_params_cls(
+    server_args: ServerArgs,
+) -> type[VLASamplingParams]:
+    return _resolve_action_sampling_params_cls_cached(
+        server_args.model_path,
+        getattr(server_args, "backend", None),
+        getattr(server_args, "model_id", None),
+        getattr(server_args, "pipeline_class_name", None),
+    )
+
+
+@lru_cache(maxsize=32)
+def _sampling_params_field_names(
+    sampling_params_cls: type[VLASamplingParams],
+) -> frozenset[str]:
+    return frozenset(field.name for field in dataclasses.fields(sampling_params_cls))
+
+
+def build_action_sampling_params(
+    payload: dict[str, Any],
+    server_args: ServerArgs,
+) -> VLASamplingParams:
+    pipeline_config = server_args.pipeline_config
+    observation = _action_request_to_observation(payload)
+    parameters = dict(payload.get("parameters") or {})
+    runtime = dict(payload.get("runtime") or {})
+    if "return_timing" in payload and "return_timing" not in runtime:
+        runtime["return_timing"] = payload["return_timing"]
+    images = images_from_observation(observation, pipeline_config)
+    state = observation.get("state")
+    if state is None:
+        state = observation.get("observation.state")
+    noise = observation.get("noise")
+    if noise is None:
+        noise = observation.get("observation.noise")
+    prompt = observation.get("prompt") or observation.get("task") or ""
+    prefix_cache = runtime.get("prefix_cache")
+    if prefix_cache is None:
+        prefix_cache = observation.get("enable_prefix_cache")
+    if prefix_cache is None:
+        prefix_cache = observation.get("enable_pi_prefix_cache")
+    cuda_graph = runtime.get("cuda_graph")
+    if cuda_graph is None:
+        cuda_graph = observation.get("enable_cuda_graph")
+    if cuda_graph is None:
+        cuda_graph = observation.get("enable_pi_cuda_graph")
+    output_format = str(
+        runtime.get(
+            "output_format",
+            parameters.get(
+                "output_format",
+                observation.get("output_format", "list"),
+            ),
+        )
+    ).lower()
+    if output_format not in ("list", "numpy"):
+        raise ValueError("output_format must be 'list' or 'numpy'")
+
+    sampling_params_cls = _resolve_action_sampling_params_cls(server_args)
+    sampling_kwargs = {
+        "request_id": payload.get("request_id") or payload.get("id"),
+        "prompt": prompt,
+        "images": images,
+        "image_masks": observation.get("image_masks"),
+        "camera_order": observation.get("camera_order"),
+        "state": state,
+        "noise": noise,
+        "observation": observation,
+        "action_horizon": int(
+            parameters.get(
+                "action_horizon",
+                observation.get("action_horizon", pipeline_config.action_horizon),
+            )
+        ),
+        "action_dim": int(
+            parameters.get(
+                "action_dim",
+                observation.get("action_dim", pipeline_config.action_dim),
+            )
+        ),
+        "num_inference_steps": int(
+            parameters.get(
+                "num_inference_steps",
+                observation.get(
+                    "num_inference_steps",
+                    pipeline_config.default_num_inference_steps,
+                ),
+            )
+        ),
+        "output_format": output_format,
+        "return_timing": _runtime_bool(runtime.get("return_timing"), True),
+        "enable_prefix_cache": _runtime_bool(prefix_cache, True),
+        "enable_cuda_graph": _runtime_bool(cuda_graph, True),
+    }
+    supported_fields = _sampling_params_field_names(sampling_params_cls)
+    sp = sampling_params_cls(
+        **{
+            name: value
+            for name, value in sampling_kwargs.items()
+            if name in supported_fields
+        }
+    )
+    sp._adjust(server_args)
+    return sp
+
+
+async def infer_action(
+    payload: dict[str, Any],
+    server_args: ServerArgs,
+) -> dict[str, Any]:
+    sp = build_action_sampling_params(payload, server_args)
+    req = prepare_request(server_args, sp)
+    response = await async_scheduler_client.forward(req)
+    if getattr(response, "error", None):
+        raise RuntimeError(response.error)
+    if response.output is None:
+        raise RuntimeError("action policy returned no output")
+    return response.output[0]
+
+
+def action_generation_response(
+    output: dict[str, Any],
+    server_args: ServerArgs,
+    *,
+    preserve_numpy: bool = False,
+) -> dict[str, Any]:
+    actions = output["actions"]
+    if isinstance(actions, np.ndarray):
+        action_shape = list(actions.shape)
+        action_values = actions if preserve_numpy else actions.tolist()
+    else:
+        horizon = len(actions) if isinstance(actions, list) else 0
+        action_dim = len(actions[0]) if horizon and isinstance(actions[0], list) else 0
+        action_shape = [horizon, action_dim]
+        action_values = actions
+    response = {
+        "id": output.get("request_id") or f"act_{uuid.uuid4().hex}",
+        "object": "action.generation",
+        "created": int(time.time()),
+        "model": server_args.model_id or server_args.model_path,
+        "data": [
+            {
+                "index": 0,
+                "input_index": 0,
+                "candidate_index": 0,
+                "action": {
+                    "type": "continuous",
+                    "dtype": "float32",
+                    "shape": action_shape,
+                    "values": action_values,
+                },
+            }
+        ],
+        "usage": {
+            "action_horizon": action_shape[0] if action_shape else 0,
+            "action_dim": action_shape[1] if len(action_shape) > 1 else 0,
+            "denoise_steps": output.get("parameters", {}).get(
+                "num_inference_steps",
+                server_args.pipeline_config.default_num_inference_steps,
+            ),
+            "prefix_cache_hit": bool(output.get("cache", {}).get("hit", False)),
+        },
+    }
+    if "timings" in output:
+        response["timings"] = output["timings"]
+    if "cache" in output:
+        response["cache"] = output["cache"]
+    if "parallel" in output:
+        response["parallel"] = output["parallel"]
+    return response
+
+
+def action_raw_response(
+    output: dict[str, Any],
+    *,
+    preserve_numpy: bool = False,
+) -> dict[str, Any]:
+    response = dict(output)
+    actions = response.get("actions")
+    if isinstance(actions, np.ndarray) and not preserve_numpy:
+        response["actions"] = actions.tolist()
+    return response

python/sglang/multimodal_gen/runtime/entrypoints/vla/ws_utils.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import time
+import traceback
+from collections.abc import Callable
+from typing import Any
+
+from fastapi import WebSocket, WebSocketDisconnect
+
+from sglang.multimodal_gen.runtime.entrypoints.vla.protocol import (
+    action_metadata,
+    infer_action,
+    pack_msgpack,
+    unpack_msgpack,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+async def run_action_msgpack_ws(
+    websocket: WebSocket,
+    server_args: ServerArgs,
+    *,
+    prepare_payload: Callable[[dict[str, Any]], None],
+    build_response: Callable[[dict[str, Any]], dict[str, Any]],
+) -> None:
+    await websocket.accept()
+    await websocket.send_bytes(pack_msgpack(action_metadata(server_args)))
+
+    prev_total_time = None
+    while True:
+        try:
+            start_time = time.monotonic()
+            payload = unpack_msgpack(await websocket.receive_bytes())
+            prepare_payload(payload)
+            infer_start = time.monotonic()
+            output = await infer_action(payload, server_args)
+            response = build_response(output)
+            response.setdefault("server_timing", {})["infer_ms"] = (
+                time.monotonic() - infer_start
+            ) * 1000
+            if prev_total_time is not None:
+                response["server_timing"]["prev_total_ms"] = prev_total_time * 1000
+            await websocket.send_bytes(pack_msgpack(response))
+            prev_total_time = time.monotonic() - start_time
+        except WebSocketDisconnect:
+            break
+        except Exception:
+            try:
+                await websocket.send_bytes(
+                    pack_msgpack({"error": traceback.format_exc()})
+                )
+            except Exception:
+                pass
+            await websocket.close(code=1011, reason="Internal server error")
+            raise

python/sglang/multimodal_gen/runtime/ipc_array.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Helpers for transferring large numpy arrays between local scheduler processes."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+_MIN_FILE_REF_BYTES = 32 << 20
+
+
+@dataclass
+class NumpyArrayFileRef:
+    path: str
+
+    def materialize(self) -> np.ndarray:
+        try:
+            return np.load(self.path, allow_pickle=False)
+        finally:
+            try:
+                os.unlink(self.path)
+            except FileNotFoundError:
+                pass
+
+
+def is_local_endpoint(endpoint: str) -> bool:
+    return endpoint.startswith(
+        ("tcp://127.0.0.1:", "tcp://localhost:", "ipc://", "inproc://")
+    )
+
+
+def spill_large_arrays_to_file_refs(value: Any) -> Any:
+    directory = _array_ipc_dir()
+    if directory is None:
+        return value
+    return _spill_large_arrays_to_file_refs(value, directory)
+
+
+def _spill_large_arrays_to_file_refs(value: Any, directory: str) -> Any:
+    if isinstance(value, np.ndarray) and value.nbytes >= _MIN_FILE_REF_BYTES:
+        # only spill if the array size is above the threshold. if not, it's not worth it
+        return _spill_array(value, directory)
+    if isinstance(value, list):
+        return [_spill_large_arrays_to_file_refs(item, directory) for item in value]
+    if isinstance(value, tuple):
+        return tuple(
+            _spill_large_arrays_to_file_refs(item, directory) for item in value
+        )
+    return value
+
+
+def materialize_file_refs(value: Any) -> Any:
+    if isinstance(value, NumpyArrayFileRef):
+        return value.materialize()
+    if isinstance(value, list):
+        return [materialize_file_refs(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(materialize_file_refs(item) for item in value)
+    return value
+
+
+def _spill_array(array: np.ndarray, directory: str) -> NumpyArrayFileRef:
+    if not array.flags.c_contiguous:
+        array = np.ascontiguousarray(array)
+
+    fd, path = tempfile.mkstemp(
+        prefix="sgldiffusion-array-",
+        suffix=".npy",
+        dir=directory,
+    )
+    try:
+        with os.fdopen(fd, "wb") as f:
+            np.save(f, array, allow_pickle=False)
+    except Exception:
+        try:
+            os.unlink(path)
+        except FileNotFoundError:
+            pass
+        raise
+    return NumpyArrayFileRef(path=path)
+
+
+def _array_ipc_dir() -> str | None:
+    shm_path = Path("/dev/shm")
+    if shm_path.is_dir() and os.access(shm_path, os.W_OK):
+        return str(shm_path)
+    return None

python/sglang/multimodal_gen/runtime/launch_server.py

@@ -0,0 +1,788 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import dataclasses
+import multiprocessing as mp
+import os
+import signal
+import sys
+import threading
+import time
+
+import psutil
+import uvicorn
+
+from sglang.multimodal_gen.runtime.disaggregation.orchestrator import (
+    DiffusionServer,
+)
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app
+from sglang.multimodal_gen.runtime.entrypoints.utils import ShutdownReq
+from sglang.multimodal_gen.runtime.managers.gpu_worker import run_scheduler_process
+from sglang.multimodal_gen.runtime.scheduler_client import SchedulerClient
+from sglang.multimodal_gen.runtime.server_args import (
+    ServerArgs,
+    prepare_server_args,
+    set_global_server_args,
+)
+from sglang.multimodal_gen.runtime.utils.common import is_port_available
+from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger
+from sglang.multimodal_gen.runtime.utils.trace_wrapper import init_diffusion_tracing
+from sglang.multimodal_gen.utils import kill_itself_when_parent_died
+
+_SCHEDULER_SHUTDOWN_TIMEOUT_MS = 5000
+_WORKER_JOIN_TIMEOUT_S = 10
+_WORKER_TERMINATE_TIMEOUT_S = 1
+_WORKER_KILL_TIMEOUT_S = 1
+
+
+def _find_available_port(
+    start: int = 10000, avoid: set[int] | None = None, max_attempts: int = 100
+) -> int:
+    """Find an available port starting from *start*, skipping ports in *avoid*."""
+    if avoid is None:
+        avoid = set()
+    port = max(1024, min(start, 65535))
+    for _ in range(max_attempts):
+        if port not in avoid and is_port_available(port):
+            return port
+        port += 1
+        if port > 65535:
+            port = 1024
+    raise RuntimeError(
+        f"No available port found after {max_attempts} attempts (start={start})"
+    )
+
+
+def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
+    """Kill the process and all its child processes."""
+    # Remove sigchld handler to avoid spammy logs.
+    if threading.current_thread() is threading.main_thread():
+        signal.signal(signal.SIGCHLD, signal.SIG_DFL)
+
+    if parent_pid is None:
+        parent_pid = os.getpid()
+        include_parent = False
+
+    try:
+        itself = psutil.Process(parent_pid)
+    except psutil.NoSuchProcess:
+        return
+
+    children = itself.children(recursive=True)
+    for child in children:
+        if child.pid == skip_pid:
+            continue
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+    if include_parent:
+        try:
+            if parent_pid == os.getpid():
+                itself.kill()
+                sys.exit(0)
+
+            itself.kill()
+
+            # Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes),
+            # so we send an additional signal to kill them.
+            itself.send_signal(signal.SIGQUIT)
+        except psutil.NoSuchProcess:
+            pass
+
+
+def _process_names(processes) -> str:
+    return ", ".join(getattr(p, "name", repr(p)) for p in processes)
+
+
+def _join_processes_with_deadline(processes, timeout_s: float) -> None:
+    deadline = time.monotonic() + timeout_s
+    for process in processes:
+        remaining_s = max(0.0, deadline - time.monotonic())
+        process.join(timeout=remaining_s)
+
+
+def _terminate_alive_processes(processes, timeout_s: float) -> list:
+    alive = [p for p in processes if p.is_alive()]
+    if not alive:
+        return []
+
+    logger.warning(
+        "Worker process(es) did not exit in time; terminating: %s",
+        _process_names(alive),
+    )
+    for process in alive:
+        process.terminate()
+    _join_processes_with_deadline(alive, timeout_s)
+    return [p for p in alive if p.is_alive()]
+
+
+def _kill_alive_processes(processes, timeout_s: float) -> None:
+    alive = [p for p in processes if p.is_alive()]
+    if not alive:
+        return
+
+    logger.warning(
+        "Worker process(es) did not terminate in time; killing: %s",
+        _process_names(alive),
+    )
+    for process in alive:
+        process.kill()
+    _join_processes_with_deadline(alive, timeout_s)
+
+
+def _run_http_server_process(server_args: ServerArgs) -> None:
+    kill_itself_when_parent_died()
+    launch_http_server_only(server_args)
+
+
+def _request_monolithic_scheduler_shutdown(server_args: ServerArgs) -> None:
+    if server_args.disagg_role != RoleType.MONOLITHIC:
+        return
+
+    client = SchedulerClient()
+    try:
+        client.initialize(server_args)
+        client.forward(ShutdownReq(), timeout_ms=_SCHEDULER_SHUTDOWN_TIMEOUT_MS)
+    except Exception as e:
+        logger.warning("Failed to request graceful scheduler shutdown: %s", e)
+    finally:
+        client.close()
+
+
+def shutdown_scheduler_processes(
+    server_args: ServerArgs | None,
+    processes: list,
+    *,
+    request_shutdown: bool = True,
+) -> None:
+    if not processes:
+        return
+
+    if request_shutdown and server_args is not None:
+        _request_monolithic_scheduler_shutdown(server_args)
+
+    _join_processes_with_deadline(processes, _WORKER_JOIN_TIMEOUT_S)
+    alive = _terminate_alive_processes(processes, _WORKER_TERMINATE_TIMEOUT_S)
+    _kill_alive_processes(alive, _WORKER_KILL_TIMEOUT_S)
+
+
+def launch_server(server_args: ServerArgs, launch_http_server: bool = True):
+    """
+    Args:
+        launch_http_server: False for offline local mode
+    """
+    configure_logger(server_args)
+
+    # Start a new server with multiple worker processes
+    logger.info("Starting server...")
+
+    num_gpus = server_args.num_gpus
+    processes = []
+
+    # Pipes for master to talk to slaves
+    task_pipes_to_slaves_w = []
+    task_pipes_to_slaves_r = []
+    for _ in range(num_gpus - 1):
+        r, w = mp.Pipe(duplex=False)
+        task_pipes_to_slaves_r.append(r)
+        task_pipes_to_slaves_w.append(w)
+
+    # Pipes for slaves to talk to master
+    result_pipes_from_slaves_w = []
+    result_pipes_from_slaves_r = []
+    for _ in range(num_gpus - 1):
+        r, w = mp.Pipe(duplex=False)
+        result_pipes_from_slaves_r.append(r)
+        result_pipes_from_slaves_w.append(w)
+
+    # Launch all worker processes
+    master_port = server_args.master_port
+    scheduler_pipe_readers = []
+    scheduler_pipe_writers = []
+
+    for i in range(num_gpus):
+        reader, writer = mp.Pipe(duplex=False)
+        scheduler_pipe_writers.append(writer)
+        if i == 0:  # Master worker
+            process = mp.Process(
+                target=run_scheduler_process,
+                args=(
+                    i,  # local_rank
+                    i,  # rank
+                    master_port,
+                    server_args,
+                    writer,
+                    None,  # No task pipe to read from master
+                    None,  # No result pipe to write to master
+                    task_pipes_to_slaves_w,
+                    result_pipes_from_slaves_r,
+                ),
+                name=f"sglang-diffusionWorker-{i}",
+                daemon=True,
+            )
+        else:  # Slave workers
+            process = mp.Process(
+                target=run_scheduler_process,
+                args=(
+                    i,  # local_rank
+                    i,  # rank
+                    master_port,
+                    server_args,
+                    writer,
+                    None,  # No task pipe to read from master
+                    None,  # No result pipe to write to master
+                    task_pipes_to_slaves_r[i - 1],
+                    result_pipes_from_slaves_w[i - 1],
+                ),
+                name=f"sglang-diffusionWorker-{i}",
+                daemon=True,
+            )
+        scheduler_pipe_readers.append(reader)
+        process.start()
+        processes.append(process)
+
+    # Wait for all workers to be ready
+    scheduler_infos = []
+    for writer in scheduler_pipe_writers:
+        writer.close()
+
+    # Close unused pipe ends in parent process
+    for p in task_pipes_to_slaves_w:
+        p.close()
+    for p in task_pipes_to_slaves_r:
+        p.close()
+    for p in result_pipes_from_slaves_w:
+        p.close()
+    for p in result_pipes_from_slaves_r:
+        p.close()
+
+    for i, reader in enumerate(scheduler_pipe_readers):
+        try:
+            data = reader.recv()
+        except EOFError:
+            logger.error(
+                f"Rank {i} scheduler is dead. Please check if there are relevant logs."
+            )
+            processes[i].join()
+            logger.error(f"Exit code: {processes[i].exitcode}")
+            raise
+
+        if data["status"] != "ready":
+            raise RuntimeError(
+                "Initialization failed. Please see the error messages above."
+            )
+        scheduler_infos.append(data)
+        reader.close()
+
+    logger.debug("All workers are ready")
+
+    if launch_http_server:
+        if server_args.pipeline_config.task_type.is_action_gen():
+            logger.info(
+                "VLA pipeline ready: model=%s; per-request details are "
+                "debug-only (use --log-level debug).",
+                server_args.model_id or server_args.model_path,
+            )
+        logger.info("Starting FastAPI server.")
+        if server_args.webui:
+            logger.info("Launch FastAPI server in another process because of webui.")
+            http_server_process = mp.Process(
+                target=_run_http_server_process,
+                args=(server_args,),
+                name="sglang-diffusion-webui",
+                daemon=True,
+            )
+            http_server_process.start()
+        else:
+            try:
+                launch_http_server_only(server_args)
+            finally:
+                shutdown_scheduler_processes(server_args, processes)
+
+    return processes
+
+
+def launch_pool_disagg_server(
+    server_args: ServerArgs,
+    encoder_gpus: list[list[int]],
+    denoiser_gpus: list[list[int]],
+    decoder_gpus: list[list[int]],
+    launch_http_server: bool = True,
+):
+    """Launch a pool-based disaggregated server with N:M:K independent role instances.
+
+    DiffusionServer orchestrates the full pipeline, dispatching at every
+    role transition (Encoder → Denoiser → Decoder).
+
+    Args:
+        server_args: Base server configuration
+        encoder_gpus: List of GPU ID lists, one per encoder instance.
+            e.g., [[0], [2]] for 2 encoder instances on GPUs 0 and 2.
+        denoiser_gpus: List of GPU ID lists, one per denoiser instance.
+            e.g., [[1], [3]] for 2 denoiser instances.
+        decoder_gpus: List of GPU ID lists, one per decoder instance.
+            e.g., [[0], [2]] for 2 decoder instances (can share with encoder).
+        launch_http_server: Whether to launch the HTTP server.
+
+    Example:
+        launch_pool_disagg_server(server_args,
+            encoder_gpus=[[0], [2]],
+            denoiser_gpus=[[1], [3]],
+            decoder_gpus=[[0], [2]],
+        )
+    """
+    configure_logger(server_args)
+
+    num_encoders = len(encoder_gpus)
+    num_denoisers = len(denoiser_gpus)
+    num_decoders = len(decoder_gpus)
+    logger.info(
+        "Starting pool disagg server: %d encoder(s), %d denoiser(s), %d decoder(s)...",
+        num_encoders,
+        num_denoisers,
+        num_decoders,
+    )
+
+    host = server_args.host or "127.0.0.1"
+
+    def find_port(start):
+        return _find_available_port(start)
+
+    # Allocate endpoints
+    port_cursor = server_args.scheduler_port + 3000
+
+    # Per-instance work endpoints (instance binds PULL, DS connects PUSH)
+    encoder_work_endpoints = []
+    for i in range(num_encoders):
+        p = find_port(port_cursor)
+        encoder_work_endpoints.append(f"tcp://{host}:{p}")
+        port_cursor = p + 1
+
+    denoiser_work_endpoints = []
+    for i in range(num_denoisers):
+        p = find_port(port_cursor)
+        denoiser_work_endpoints.append(f"tcp://{host}:{p}")
+        port_cursor = p + 1
+
+    decoder_work_endpoints = []
+    for i in range(num_decoders):
+        p = find_port(port_cursor)
+        decoder_work_endpoints.append(f"tcp://{host}:{p}")
+        port_cursor = p + 1
+
+    # Per-role-type result endpoints (DS binds PULL, instances connect PUSH)
+    # Use deterministic convention: scheduler_port + {1,2,3}
+    base_port = server_args.scheduler_port
+    encoder_result_ep = f"tcp://{host}:{base_port + 1}"
+    denoiser_result_ep = f"tcp://{host}:{base_port + 2}"
+    decoder_result_ep = f"tcp://{host}:{base_port + 3}"
+
+    logger.info(
+        "Pool endpoints allocated: %d work + 3 result endpoints",
+        num_encoders + num_denoisers + num_decoders,
+    )
+
+    # Launch all role instances
+    all_processes = []
+
+    role_configs = [
+        (RoleType.ENCODER, encoder_gpus, encoder_work_endpoints, encoder_result_ep),
+        (
+            RoleType.DENOISER,
+            denoiser_gpus,
+            denoiser_work_endpoints,
+            denoiser_result_ep,
+        ),
+        (RoleType.DECODER, decoder_gpus, decoder_work_endpoints, decoder_result_ep),
+    ]
+
+    for role_type, gpu_lists, work_eps, result_ep in role_configs:
+        for inst_idx, gpu_ids in enumerate(gpu_lists):
+            num_role_gpus = len(gpu_ids)
+
+            # Per-role parallelism: use explicit overrides if set, else None (auto-derive)
+            role_par = server_args.get_role_parallelism(role_type)
+
+            role_overrides = {
+                "disagg_role": role_type,
+                "disagg_mode": True,
+                "pool_work_endpoint": work_eps[inst_idx],
+                "pool_result_endpoint": result_ep,
+                "num_gpus": num_role_gpus,
+                "warmup": role_type == RoleType.ENCODER,
+                "server_warmup": False,
+                "scheduler_port": find_port(port_cursor),
+                "master_port": find_port(port_cursor + 100),
+                # Per-role parallelism (None = auto-derive from num_gpus)
+                "tp_size": role_par["tp_size"],
+                "sp_degree": role_par["sp_degree"],
+                "ulysses_degree": role_par["ulysses_degree"],
+                "ring_degree": role_par["ring_degree"],
+            }
+            port_cursor = role_overrides["master_port"] + 100
+
+            base_dict = {
+                f.name: getattr(server_args, f.name)
+                for f in dataclasses.fields(server_args)
+            }
+            base_dict.update(role_overrides)
+            base_dict.pop("pipeline_config", None)
+            role_args = ServerArgs.from_kwargs(**base_dict)
+
+            pool_ctx = mp.get_context("spawn")
+            inst_readers = []
+
+            # Spawn all ranks first — NCCL init blocks until all ranks connect
+            for rank_idx in range(num_role_gpus):
+                reader, writer = pool_ctx.Pipe(duplex=False)
+                gpu_id = gpu_ids[rank_idx]
+
+                process = pool_ctx.Process(
+                    target=_run_disagg_role_process,
+                    args=(gpu_id, rank_idx, rank_idx, role_args, writer, [], []),
+                    name=f"sglang-pool-{role_type.value}-{inst_idx}-r{rank_idx}",
+                    daemon=True,
+                )
+                process.start()
+                all_processes.append(process)
+                inst_readers.append(reader)
+
+            # Wait for all ranks to be ready (after all are spawned)
+            for rank_idx, reader in enumerate(inst_readers):
+                try:
+                    data = reader.recv()
+                except EOFError:
+                    logger.error(
+                        "Pool %s[%d] rank %d is dead.",
+                        role_type.value,
+                        inst_idx,
+                        rank_idx,
+                    )
+                    raise
+                if data.get("status") != "ready":
+                    raise RuntimeError(
+                        f"Pool {role_type.value}[{inst_idx}] rank {rank_idx} "
+                        "failed to initialize."
+                    )
+                reader.close()
+
+            logger.info(
+                "Pool %s[%d] ready on GPU(s) %s (work=%s)",
+                role_type.value.upper(),
+                inst_idx,
+                gpu_ids,
+                work_eps[inst_idx],
+            )
+
+    logger.info("All pool role instances ready")
+
+    # Start DiffusionServer
+    frontend_endpoint = f"tcp://{host}:{server_args.scheduler_port}"
+
+    diffusion_server = DiffusionServer(
+        frontend_endpoint=frontend_endpoint,
+        encoder_work_endpoints=encoder_work_endpoints,
+        denoiser_work_endpoints=denoiser_work_endpoints,
+        decoder_work_endpoints=decoder_work_endpoints,
+        encoder_result_endpoint=encoder_result_ep,
+        denoiser_result_endpoint=denoiser_result_ep,
+        decoder_result_endpoint=decoder_result_ep,
+        dispatch_policy_name=server_args.disagg_dispatch_policy,
+        timeout_s=float(server_args.disagg_timeout),
+    )
+    diffusion_server.start()
+
+    if not diffusion_server.wait_ready(timeout=30.0):
+        raise RuntimeError("DiffusionServer failed to bind sockets within 30 seconds")
+
+    if launch_http_server:
+        logger.info(
+            "Starting FastAPI server (connected to DiffusionServer at port %d).",
+            server_args.scheduler_port,
+        )
+        try:
+            launch_http_server_only(server_args)
+        finally:
+            diffusion_server.stop()
+            shutdown_scheduler_processes(
+                server_args, all_processes, request_shutdown=False
+            )
+
+    return all_processes
+
+
+def _run_disagg_role_process(
+    gpu_id: int,
+    _local_rank: int,
+    rank: int,
+    server_args: ServerArgs,
+    pipe_writer: mp.connection.Connection,
+    task_pipes: list,
+    result_pipes: list,
+):
+    """Entry point for a disagg role process.
+
+    Uses the physical GPU index (gpu_id) as local_rank so that
+    torch.cuda.set_device(local_rank) selects the correct GPU.
+    This avoids relying on CUDA_VISIBLE_DEVICES remapping, which
+    may not work if CUDA was pre-initialized in the parent process.
+    """
+    run_scheduler_process(
+        local_rank=gpu_id,
+        rank=rank,
+        master_port=server_args.master_port,
+        server_args=server_args,
+        pipe_writer=pipe_writer,
+        task_pipe_r=None,
+        result_pipe_w=None,
+        task_pipes_to_slaves=task_pipes,
+        result_pipes_from_slaves=result_pipes,
+    )
+
+
+def launch_http_server_only(server_args):
+    init_diffusion_tracing(server_args, "DiffHTTPServer")
+
+    # set for endpoints to access global_server_args
+    set_global_server_args(server_args)
+    app = create_app(server_args)
+    uvicorn.run(
+        app,
+        use_colors=True,
+        log_level=server_args.log_level,
+        host=server_args.host,
+        port=server_args.port,
+        reload=False,
+        ws_per_message_deflate=False,
+    )
+
+
+def parse_url_string(url_str: str) -> list[str]:
+    """Parse a semicolon-separated URL string into a list.
+
+    Example: "tcp://10.0.0.1:35000;tcp://10.0.0.2:35000" -> ["tcp://...", "tcp://..."]
+    """
+    return [u.strip() for u in url_str.split(";") if u.strip()]
+
+
+def launch_disagg_server(server_args: ServerArgs):
+    """Launch DiffusionServer head node + HTTP server (--disagg-role server).
+
+    No GPU workers are spawned. Connects to remote role instances
+    specified by --encoder-urls, --denoiser-urls, --decoder-urls.
+
+    Result endpoints use deterministic convention:
+        encoder result: scheduler_port + 1
+        denoiser result: scheduler_port + 2
+        decoder result: scheduler_port + 3
+    """
+    configure_logger(server_args)
+
+    for name, val in [
+        ("--encoder-urls", server_args.encoder_urls),
+        ("--denoiser-urls", server_args.denoiser_urls),
+        ("--decoder-urls", server_args.decoder_urls),
+    ]:
+        if val is None:
+            raise ValueError(f"{name} is required for --disagg-role server")
+
+    host = server_args.host or "127.0.0.1"
+    base_port = server_args.scheduler_port
+
+    encoder_work_endpoints = parse_url_string(server_args.encoder_urls)
+    denoiser_work_endpoints = parse_url_string(server_args.denoiser_urls)
+    decoder_work_endpoints = parse_url_string(server_args.decoder_urls)
+
+    encoder_result_ep = f"tcp://{host}:{base_port + 1}"
+    denoiser_result_ep = f"tcp://{host}:{base_port + 2}"
+    decoder_result_ep = f"tcp://{host}:{base_port + 3}"
+
+    frontend_endpoint = f"tcp://{host}:{base_port}"
+
+    logger.info(
+        "Starting DiffusionServer: %d encoder(s), %d denoiser(s), %d decoder(s)",
+        len(encoder_work_endpoints),
+        len(denoiser_work_endpoints),
+        len(decoder_work_endpoints),
+    )
+    logger.info("  Frontend: %s", frontend_endpoint)
+    logger.info("  Encoder work endpoints: %s", encoder_work_endpoints)
+    logger.info("  Denoiser work endpoints: %s", denoiser_work_endpoints)
+    logger.info("  Decoder work endpoints: %s", decoder_work_endpoints)
+    logger.info(
+        "  Result endpoints: encoder=%s, denoiser=%s, decoder=%s",
+        encoder_result_ep,
+        denoiser_result_ep,
+        decoder_result_ep,
+    )
+
+    diffusion_server = DiffusionServer(
+        frontend_endpoint=frontend_endpoint,
+        encoder_work_endpoints=encoder_work_endpoints,
+        denoiser_work_endpoints=denoiser_work_endpoints,
+        decoder_work_endpoints=decoder_work_endpoints,
+        encoder_result_endpoint=encoder_result_ep,
+        denoiser_result_endpoint=denoiser_result_ep,
+        decoder_result_endpoint=decoder_result_ep,
+        dispatch_policy_name=server_args.disagg_dispatch_policy,
+        timeout_s=float(server_args.disagg_timeout),
+    )
+    diffusion_server.start()
+
+    if not diffusion_server.wait_ready(timeout=30.0):
+        raise RuntimeError("DiffusionServer failed to bind sockets within 30 seconds")
+
+    logger.info(
+        "Starting HTTP server (connected to DiffusionServer at port %d).",
+        base_port,
+    )
+    try:
+        launch_http_server_only(server_args)
+    finally:
+        diffusion_server.stop()
+
+
+def launch_disagg_role(server_args: ServerArgs):
+    """Launch a standalone disaggregated role instance (--disagg-role encoder/denoising/decoder).
+
+    The instance:
+    1. Binds its work PULL socket on tcp://0.0.0.0:{scheduler_port}
+    2. Connects its result PUSH socket to the DiffusionServer head node
+       (derived from --disagg-server-addr + role offset)
+    3. Spawns GPU worker processes for the assigned role.
+    """
+    configure_logger(server_args)
+
+    role_type = server_args.disagg_role
+    if server_args.disagg_server_addr is None:
+        raise ValueError(
+            "--disagg-server-addr is required for --disagg-role " f"{role_type.value}"
+        )
+
+    # Derive endpoints
+    work_endpoint = server_args.derive_pool_work_endpoint()
+    result_endpoint = server_args.derive_pool_result_endpoint()
+
+    logger.info(
+        "Starting disagg role: %s, num_gpus=%d",
+        role_type.value,
+        server_args.num_gpus,
+    )
+    logger.info("  Work endpoint (bind): %s", work_endpoint)
+    logger.info("  Result endpoint (connect): %s", result_endpoint)
+    logger.info(
+        "  P2P: hostname=%s, ib_device=%s, pool_size=%d",
+        server_args.disagg_p2p_hostname,
+        server_args.disagg_ib_device,
+        server_args.disagg_transfer_pool_size,
+    )
+
+    # Build role-specific ServerArgs
+    # Use a different port for the scheduler's internal ROUTER socket to avoid
+    # conflicting with the pool work PULL socket (both bind on scheduler_port).
+    internal_scheduler_port = _find_available_port(
+        start=server_args.scheduler_port + 100, avoid={server_args.scheduler_port}
+    )
+
+    role_par = server_args.get_role_parallelism(role_type)
+    role_overrides = {
+        "disagg_role": role_type,
+        "disagg_mode": True,
+        "pool_work_endpoint": work_endpoint,
+        "pool_result_endpoint": result_endpoint,
+        "warmup": role_type == RoleType.ENCODER,
+        "server_warmup": False,
+        "scheduler_port": internal_scheduler_port,
+        # Per-role parallelism (None = auto-derive from num_gpus)
+        "tp_size": role_par["tp_size"],
+        "sp_degree": role_par["sp_degree"],
+        "ulysses_degree": role_par["ulysses_degree"],
+        "ring_degree": role_par["ring_degree"],
+    }
+
+    base_dict = {
+        f.name: getattr(server_args, f.name) for f in dataclasses.fields(server_args)
+    }
+    base_dict.update(role_overrides)
+    base_dict.pop("pipeline_config", None)
+    role_args = ServerArgs.from_kwargs(**base_dict)
+
+    # Spawn GPU worker processes
+    # NOTE: All ranks must be spawned before waiting for ready signals,
+    # because NCCL init_process_group blocks until all ranks connect.
+    num_gpus = server_args.num_gpus
+    base_gpu_id = server_args.base_gpu_id
+    pool_ctx = mp.get_context("spawn")
+    processes = []
+    readers = []
+
+    for rank_idx in range(num_gpus):
+        reader, writer = pool_ctx.Pipe(duplex=False)
+        gpu_id = base_gpu_id + rank_idx
+
+        process = pool_ctx.Process(
+            target=_run_disagg_role_process,
+            args=(gpu_id, rank_idx, rank_idx, role_args, writer, [], []),
+            name=f"sglang-{role_type.value}-r{rank_idx}",
+            daemon=True,
+        )
+        process.start()
+        processes.append(process)
+        readers.append(reader)
+
+    # Wait for all ranks to be ready (after all are spawned)
+    for rank_idx, reader in enumerate(readers):
+        try:
+            data = reader.recv()
+        except EOFError:
+            logger.error(
+                "Role %s rank %d is dead.",
+                role_type.value,
+                rank_idx,
+            )
+            raise
+        if data.get("status") != "ready":
+            raise RuntimeError(
+                f"Role {role_type.value} rank {rank_idx} failed to initialize."
+            )
+        reader.close()
+
+    logger.info(
+        "Role %s ready (%d GPU(s), work=%s)",
+        role_type.value.upper(),
+        num_gpus,
+        work_endpoint,
+    )
+
+    # Block until interrupted
+    try:
+        for p in processes:
+            p.join()
+    except KeyboardInterrupt:
+        logger.info("Role %s shutting down.", role_type.value)
+    finally:
+        shutdown_scheduler_processes(role_args, processes, request_shutdown=False)
+
+
+def dispatch_launch(server_args: ServerArgs):
+    """Route to the correct launch function based on --disagg-role."""
+    role = server_args.disagg_role
+    if role == RoleType.MONOLITHIC:
+        launch_server(server_args)
+    elif role == RoleType.SERVER:
+        launch_disagg_server(server_args)
+    elif role in (RoleType.ENCODER, RoleType.DENOISER, RoleType.DECODER):
+        launch_disagg_role(server_args)
+    else:
+        raise ValueError(f"Unknown disagg_role: {role}")
+
+
+if __name__ == "__main__":
+    server_args = prepare_server_args(sys.argv[1:])
+
+    try:
+        dispatch_launch(server_args)
+    finally:
+        kill_process_tree(os.getpid(), include_parent=False)

python/sglang/multimodal_gen/runtime/layers/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/layers/activation.py

@@ -0,0 +1,178 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/activation.py
+"""Custom activation functions."""
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_cuda = current_platform.is_cuda()
+_is_hip = current_platform.is_hip()
+_is_npu = current_platform.is_npu()
+_is_xpu = current_platform.is_xpu()
+
+if _is_cuda:
+    from sglang.jit_kernel.activation import silu_and_mul
+elif _is_hip or _is_xpu:
+    from sgl_kernel import silu_and_mul
+
+
+if _is_npu:
+    import torch_npu
+# TODO (will): remove this dependency
+from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
+
+
+@CustomOp.register("silu_and_mul")
+class SiluAndMul(CustomOp):
+    """An activation function for SwiGLU.
+
+    The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
+
+    Shapes:
+        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
+        return: (num_tokens, d) or (batch_size, seq_len, d)
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        silu_and_mul(x, out)
+        return out
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        d = x.shape[-1] // 2
+        return F.silu(x[..., :d]) * x[..., d:]
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        out = torch_npu.npu_swiglu(x)
+        return out
+
+    def forward_musa(self, x: torch.Tensor) -> torch.Tensor:
+        return nn.SwishGLU()(x)
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        silu_and_mul(x, out)
+        return out
+
+
+@CustomOp.register("gelu_and_mul")
+class GeluAndMul(CustomOp):
+    """An activation function for GeGLU.
+
+    The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
+
+    Shapes:
+        x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
+        return: (batch_size, seq_len, d) or (num_tokens, d)
+    """
+
+    def __init__(self, approximate: str = "none"):
+        super().__init__()
+        self.approximate = approximate
+        if approximate not in ("none", "tanh"):
+            raise ValueError(f"Unknown approximate mode: {approximate}")
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        y_npu, _ = torch_npu.npu_geglu(
+            x,
+            dim=-1,
+            approximate=1 if self.approximate == "tanh" else 0,
+            activate_left=True,
+        )
+        return y_npu
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        d = x.shape[-1] // 2
+        return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:]
+
+    def extra_repr(self) -> str:
+        return f"approximate={repr(self.approximate)}"
+
+
+@CustomOp.register("gelu_new")
+class NewGELU(CustomOp):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        c = math.sqrt(2.0 / math.pi)
+        return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+@CustomOp.register("quick_gelu")
+class QuickGELU(CustomOp):
+    # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90
+    def __init__(self):
+        super().__init__()
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        return x * torch.sigmoid(1.702 * x)
+
+
+_ACTIVATION_REGISTRY = {
+    "gelu": nn.GELU,
+    "gelu_new": NewGELU,
+    "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"),
+    "relu": nn.ReLU,
+    "silu": nn.SiLU,
+    "quick_gelu": QuickGELU,
+}
+
+
+def get_act_fn(act_fn_name: str) -> nn.Module:
+    """Get an activation function by name."""
+    act_fn_name = act_fn_name.lower()
+    if act_fn_name not in _ACTIVATION_REGISTRY:
+        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
+
+    return _ACTIVATION_REGISTRY[act_fn_name]()
+
+
+_ACTIVATION_AND_MUL_REGISTRY = {
+    "gelu": GeluAndMul,
+    "silu": SiluAndMul,
+}
+
+
+def get_act_and_mul_fn(act_fn_name: str) -> nn.Module:
+    """Get an activation-and-mul (i.e. SiluAndMul) function by name."""
+    act_fn_name = act_fn_name.lower()
+    if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY:
+        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
+
+    return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]()

python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py

@@ -0,0 +1,414 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+from collections import defaultdict
+from typing import Any
+
+import numpy as np
+
+from sglang.multimodal_gen.utils import dict_to_3d_list
+
+
+def configure_sta(
+    mode: str = "STA_searching",
+    layer_num: int = 40,
+    time_step_num: int = 50,
+    head_num: int = 40,
+    **kwargs,
+) -> list[list[list[Any]]]:
+    """
+    Configure Sliding Tile Attention (STA) parameters based on the specified mode.
+
+    Parameters:
+    ----------
+    mode : str
+        The STA mode to use. Options are:
+        - 'STA_searching': Generate a set of mask candidates for initial search
+        - 'STA_tuning': Select best mask strategy based on previously saved results
+        - 'STA_inference': Load and use a previously tuned mask strategy
+    layer_num: int, number of layers
+    time_step_num: int, number of timesteps
+    head_num: int, number of heads
+
+    **kwargs : dict
+        Mode-specific parameters:
+
+        For 'STA_searching':
+        - mask_candidates: list of str, optional, mask candidates to use
+        - mask_selected: list of int, optional, indices of selected masks
+
+        For 'STA_tuning':
+        - mask_search_files_path: str, required, path to mask search results
+        - mask_candidates: list of str, optional, mask candidates to use
+        - mask_selected: list of int, optional, indices of selected masks
+        - skip_time_steps: int, optional, number of time steps to use full attention (default 12)
+        - save_dir: str, optional, directory to save mask strategy (default "mask_candidates")
+
+        For 'STA_inference':
+        - load_path: str, optional, path to load mask strategy (default "mask_candidates/mask_strategy.json")
+    """
+    valid_modes = ["STA_searching", "STA_tuning", "STA_inference", "STA_tuning_cfg"]
+    if mode not in valid_modes:
+        raise ValueError(f"Mode must be one of {valid_modes}, got {mode}")
+
+    if mode == "STA_searching":
+        # Get parameters with defaults
+        mask_candidates: list[str] | None = kwargs.get("mask_candidates")
+        if mask_candidates is None:
+            raise ValueError("mask_candidates is required for STA_searching mode")
+        mask_selected: list[int] = kwargs.get(
+            "mask_selected", list(range(len(mask_candidates)))
+        )
+
+        # Parse selected masks
+        selected_masks: list[list[int]] = []
+        for index in mask_selected:
+            mask = mask_candidates[index]
+            masks_list = [int(x) for x in mask.split(",")]
+            selected_masks.append(masks_list)
+
+        # Create 3D mask structure with fixed dimensions (t=50, l=60)
+        masks_3d: list[list[list[list[int]]]] = []
+        for i in range(time_step_num):  # Fixed t dimension = 50
+            row = []
+            for j in range(layer_num):  # Fixed l dimension = 60
+                row.append(selected_masks)  # Add all masks at each position
+            masks_3d.append(row)
+
+        return masks_3d
+
+    elif mode == "STA_tuning":
+        # Get required parameters
+        mask_search_files_path: str | None = kwargs.get("mask_search_files_path")
+        if not mask_search_files_path:
+            raise ValueError("mask_search_files_path is required for STA_tuning mode")
+
+        # Get optional parameters with defaults
+        mask_candidates_tuning: list[str] | None = kwargs.get("mask_candidates")
+        if mask_candidates_tuning is None:
+            raise ValueError("mask_candidates is required for STA_tuning mode")
+        mask_selected_tuning: list[int] = kwargs.get(
+            "mask_selected", list(range(len(mask_candidates_tuning)))
+        )
+        skip_time_steps_tuning: int | None = kwargs.get("skip_time_steps")
+        save_dir_tuning: str | None = kwargs.get("save_dir", "mask_candidates")
+
+        # Parse selected masks
+        selected_masks_tuning: list[list[int]] = []
+        for index in mask_selected_tuning:
+            mask = mask_candidates_tuning[index]
+            masks_list = [int(x) for x in mask.split(",")]
+            selected_masks_tuning.append(masks_list)
+
+        # Read JSON results
+        results = read_specific_json_files(mask_search_files_path)
+        averaged_results = average_head_losses(results, selected_masks_tuning)
+
+        # Add full attention mask for specific cases
+        full_attention_mask_tuning: list[int] | None = kwargs.get("full_attention_mask")
+        if full_attention_mask_tuning is not None:
+            selected_masks_tuning.append(full_attention_mask_tuning)
+
+        # Select best mask strategy
+        timesteps_tuning: int = kwargs.get("timesteps", time_step_num)
+        if skip_time_steps_tuning is None:
+            skip_time_steps_tuning = 12
+        mask_strategy, sparsity, strategy_counts = select_best_mask_strategy(
+            averaged_results,
+            selected_masks_tuning,
+            skip_time_steps_tuning,
+            timesteps_tuning,
+            head_num,
+        )
+
+        # Save mask strategy
+        if save_dir_tuning is not None:
+            os.makedirs(save_dir_tuning, exist_ok=True)
+            file_path = os.path.join(
+                save_dir_tuning, f"mask_strategy_s{skip_time_steps_tuning}.json"
+            )
+            with open(file_path, "w") as f:
+                json.dump(mask_strategy, f, indent=4)
+            print(f"Successfully saved mask_strategy to {file_path}")
+
+        # Print sparsity and strategy counts for information
+        print(f"Overall sparsity: {sparsity:.4f}")
+        print("\nStrategy usage counts:")
+        total_heads = time_step_num * layer_num * head_num  # Fixed dimensions
+        for strategy, count in strategy_counts.items():
+            print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)")
+
+        # Convert dictionary to 3D list with fixed dimensions
+        mask_strategy_3d = dict_to_3d_list(
+            mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num
+        )
+
+        return mask_strategy_3d
+    elif mode == "STA_tuning_cfg":
+        # Get required parameters for both positive and negative paths
+        mask_search_files_path_pos: str | None = kwargs.get(
+            "mask_search_files_path_pos"
+        )
+        mask_search_files_path_neg: str | None = kwargs.get(
+            "mask_search_files_path_neg"
+        )
+        save_dir_cfg: str | None = kwargs.get("save_dir")
+
+        if (
+            not mask_search_files_path_pos
+            or not mask_search_files_path_neg
+            or not save_dir_cfg
+        ):
+            raise ValueError(
+                "mask_search_files_path_pos, mask_search_files_path_neg, and save_dir are required for STA_tuning_cfg mode"
+            )
+
+        # Get optional parameters with defaults
+        mask_candidates_cfg: list[str] | None = kwargs.get("mask_candidates")
+        if mask_candidates_cfg is None:
+            raise ValueError("mask_candidates is required for STA_tuning_cfg mode")
+        mask_selected_cfg: list[int] = kwargs.get(
+            "mask_selected", list(range(len(mask_candidates_cfg)))
+        )
+        skip_time_steps_cfg: int | None = kwargs.get("skip_time_steps")
+
+        # Parse selected masks
+        selected_masks_cfg: list[list[int]] = []
+        for index in mask_selected_cfg:
+            mask = mask_candidates_cfg[index]
+            masks_list = [int(x) for x in mask.split(",")]
+            selected_masks_cfg.append(masks_list)
+
+        # Read JSON results for both positive and negative paths
+        pos_results = read_specific_json_files(mask_search_files_path_pos)
+        neg_results = read_specific_json_files(mask_search_files_path_neg)
+        # Combine positive and negative results into one list
+        combined_results = pos_results + neg_results
+
+        # Average the combined results
+        averaged_results = average_head_losses(combined_results, selected_masks_cfg)
+
+        # Add full attention mask for specific cases
+        full_attention_mask_cfg: list[int] | None = kwargs.get("full_attention_mask")
+        if full_attention_mask_cfg is not None:
+            selected_masks_cfg.append(full_attention_mask_cfg)
+
+        timesteps_cfg: int = kwargs.get("timesteps", time_step_num)
+        if skip_time_steps_cfg is None:
+            skip_time_steps_cfg = 12
+        # Select best mask strategy using combined results
+        mask_strategy, sparsity, strategy_counts = select_best_mask_strategy(
+            averaged_results,
+            selected_masks_cfg,
+            skip_time_steps_cfg,
+            timesteps_cfg,
+            head_num,
+        )
+
+        # Save mask strategy
+        os.makedirs(save_dir_cfg, exist_ok=True)
+        file_path = os.path.join(
+            save_dir_cfg, f"mask_strategy_s{skip_time_steps_cfg}.json"
+        )
+        with open(file_path, "w") as f:
+            json.dump(mask_strategy, f, indent=4)
+        print(f"Successfully saved mask_strategy to {file_path}")
+
+        # Print sparsity and strategy counts for information
+        print(f"Overall sparsity: {sparsity:.4f}")
+        print("\nStrategy usage counts:")
+        total_heads = time_step_num * layer_num * head_num  # Fixed dimensions
+        for strategy, count in strategy_counts.items():
+            print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)")
+
+        # Convert dictionary to 3D list with fixed dimensions
+        mask_strategy_3d = dict_to_3d_list(
+            mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num
+        )
+
+        return mask_strategy_3d
+
+    else:  # STA_inference
+        # Get parameters with defaults
+        load_path: str | None = kwargs.get(
+            "load_path", "mask_candidates/mask_strategy.json"
+        )
+        if load_path is None:
+            raise ValueError("load_path is required for STA_inference mode")
+
+        # Load previously saved mask strategy
+        with open(load_path) as f:
+            mask_strategy = json.load(f)
+
+        # Convert dictionary to 3D list with fixed dimensions
+        mask_strategy_3d = dict_to_3d_list(
+            mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num
+        )
+
+        return mask_strategy_3d
+
+
+# Helper functions
+
+
+def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]:
+    """Read and parse JSON files containing mask search results."""
+    json_contents: list[dict[str, Any]] = []
+
+    # List files only in the current directory (no walk)
+    files = os.listdir(folder_path)
+    # Filter files
+    matching_files = [f for f in files if "mask" in f and f.endswith(".json")]
+    print(f"Found {len(matching_files)} matching files: {matching_files}")
+
+    for file_name in matching_files:
+        file_path = os.path.join(folder_path, file_name)
+        with open(file_path) as file:
+            data = json.load(file)
+            json_contents.append(data)
+
+    return json_contents
+
+
+def average_head_losses(
+    results: list[dict[str, Any]], selected_masks: list[list[int]]
+) -> dict[str, dict[str, np.ndarray]]:
+    """Average losses across all prompts for each mask strategy."""
+    # Initialize a dictionary to store the averaged results
+    averaged_losses: dict[str, dict[str, np.ndarray]] = {}
+    loss_type = "L2_loss"
+    # Get all loss types (e.g., 'L2_loss')
+    averaged_losses[loss_type] = {}
+
+    for mask in selected_masks:
+        mask_str = str(mask)
+        data_shape = np.array(results[0][loss_type][mask_str]).shape
+        accumulated_data = np.zeros(data_shape)
+
+        # Sum across all prompts
+        for prompt_result in results:
+            accumulated_data += np.array(prompt_result[loss_type][mask_str])
+
+        # Average by dividing by number of prompts
+        averaged_data = accumulated_data / len(results)
+        averaged_losses[loss_type][mask_str] = averaged_data
+
+    return averaged_losses
+
+
+def select_best_mask_strategy(
+    averaged_results: dict[str, dict[str, np.ndarray]],
+    selected_masks: list[list[int]],
+    skip_time_steps: int = 12,
+    timesteps: int = 50,
+    head_num: int = 40,
+) -> tuple[dict[str, list[int]], float, dict[str, int]]:
+    """Select the best mask strategy for each head based on loss minimization."""
+    best_mask_strategy: dict[str, list[int]] = {}
+    loss_type = "L2_loss"
+    # Get the shape of time steps and layers
+    layers = len(averaged_results[loss_type][str(selected_masks[0])][0])
+
+    # Counter for sparsity calculation
+    total_tokens = 0  # total number of masked tokens
+    total_length = 0  # total sequence length
+
+    strategy_counts: dict[str, int] = {str(strategy): 0 for strategy in selected_masks}
+    full_attn_strategy = selected_masks[-1]  # Last strategy is full attention
+    print(f"Strategy {full_attn_strategy}, skip first {skip_time_steps} steps ")
+
+    for t in range(timesteps):
+        for layer_idx in range(layers):
+            for h in range(head_num):
+                if t < skip_time_steps:  # First steps use full attention
+                    strategy = full_attn_strategy
+                else:
+                    # Get losses for this head across all strategies
+                    head_losses = []
+                    for strategy in selected_masks[:-1]:  # Exclude full attention
+                        head_losses.append(
+                            averaged_results[loss_type][str(strategy)][t][layer_idx][h]
+                        )
+
+                    # Find which strategy gives minimum loss
+                    best_strategy_idx = np.argmin(head_losses)
+                    strategy = selected_masks[best_strategy_idx]
+
+                best_mask_strategy[f"{t}_{layer_idx}_{h}"] = strategy
+
+                # Calculate sparsity
+                nums = strategy  # strategy is already a list of numbers
+                total_tokens += (
+                    nums[0] * nums[1] * nums[2]
+                )  # masked tokens for chosen strategy
+                total_length += (
+                    full_attn_strategy[0]
+                    * full_attn_strategy[1]
+                    * full_attn_strategy[2]
+                )
+
+                # Count strategy usage
+                strategy_counts[str(strategy)] += 1
+
+    overall_sparsity = 1 - total_tokens / total_length
+
+    return best_mask_strategy, overall_sparsity, strategy_counts
+
+
+def save_mask_search_results(
+    mask_search_final_result: list[dict[str, list[float]]],
+    prompt: str,
+    mask_strategies: list[str],
+    output_dir: str = "output/mask_search_result/",
+) -> str | None:
+    if not mask_search_final_result:
+        print("No mask search results to save")
+        return None
+
+    # Create result dictionary with defaultdict for nested lists
+    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
+        "L2_loss": defaultdict(list),
+        "L1_loss": defaultdict(list),
+    }
+
+    mask_selected = list(range(len(mask_strategies)))
+    selected_masks: list[list[int]] = []
+    for index in mask_selected:
+        mask = mask_strategies[index]
+        masks_list = [int(x) for x in mask.split(",")]
+        selected_masks.append(masks_list)
+
+    # Process each mask strategy
+    for i, mask_strategy in enumerate(selected_masks):
+        mask_strategy_str = str(mask_strategy)
+        # Process L2 loss
+        step_results: list[list[float]] = []
+        for step_data in mask_search_final_result:
+            if isinstance(step_data, dict) and "L2_loss" in step_data:
+                layer_losses = [float(loss) for loss in step_data["L2_loss"]]
+                step_results.append(layer_losses)
+        mask_search_dict["L2_loss"][mask_strategy_str] = step_results
+
+        step_results = []
+        for step_data in mask_search_final_result:
+            if isinstance(step_data, dict) and "L1_loss" in step_data:
+                layer_losses = [float(loss) for loss in step_data["L1_loss"]]
+                step_results.append(layer_losses)
+        mask_search_dict["L1_loss"][mask_strategy_str] = step_results
+
+    # Create the output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Create a filename based on the first 20 characters of the prompt
+    filename = prompt[:50].replace(" ", "_")
+    filepath = os.path.join(output_dir, f"mask_search_{filename}.json")
+
+    # Save the results to a JSON file
+    with open(filepath, "w") as f:
+        json.dump(mask_search_dict, f, indent=4)
+
+    print(f"Successfully saved mask research results to {filepath}")
+
+    return filepath

python/sglang/multimodal_gen/runtime/layers/attention/__init__.py

@@ -0,0 +1,38 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.layer import (
+    DynamicVarlenMaskMeta,
+    LocalAttention,
+    UlyssesAttention,
+    UlyssesAttention_VSA,
+    USPAttention,
+    build_varlen_mask_meta,
+    build_varlen_mask_meta_from_lengths,
+    build_varlen_mask_meta_from_ranges,
+)
+from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import MinimalA2AAttnOp
+
+__all__ = [
+    "USPAttention",
+    "LocalAttention",
+    "DynamicVarlenMaskMeta",
+    "UlyssesAttention",
+    "UlyssesAttention_VSA",
+    "MinimalA2AAttnOp",
+    "AttentionBackend",
+    "AttentionMetadata",
+    "AttentionMetadataBuilder",
+    # "AttentionState",
+    "get_attn_backend",
+    "build_varlen_mask_meta",
+    "build_varlen_mask_meta_from_lengths",
+    "build_varlen_mask_meta_from_ranges",
+]

python/sglang/multimodal_gen/runtime/layers/attention/backends/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py

@@ -0,0 +1,207 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+
+import aiter
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER_GFX95
+
+logger = logging.getLogger(__name__)
+
+_use_fp8_attn = os.environ.get("SGLANG_DIFFUSION_AITER_FP8_ATTN", "0") == "1"
+_fp8_dtype = torch.float8_e4m3fn
+
+# fmha_fwd_hd128_fp8_gfx950 ASM kernel. Support full MHA with q/k/v head_dim == 128 -- e.g., Wan 2.2 self- and cross-attention.
+_FMHA_FP8_HEAD_DIM = 128
+
+
+if _use_fp8_attn:
+    logger.info("DiT FP8 attention enabled via SGLANG_DIFFUSION_AITER_FP8_ATTN=1")
+
+
+def _can_use_fmha_fp8_prefill(
+    q_head_dim: int,
+    k_head_dim: int,
+    v_head_dim: int,
+    num_heads: int,
+    num_kv_heads: int,
+) -> bool:
+    """True if MHA q/k/v head_dim==128 on a gfx950-class arch."""
+    if not USE_AITER_GFX95:
+        return False
+    if num_kv_heads != num_heads:
+        return False
+    return q_head_dim == k_head_dim == v_head_dim == _FMHA_FP8_HEAD_DIM
+
+
+def _fmha_fp8_prefill_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: float,
+    is_causal: bool,
+    q_scale: torch.Tensor,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> torch.Tensor:
+    """
+    FP8 FMHA prefill via aiter.flash_attn_fp8_pertensor_func.
+
+    Expects q, k, v as (batch, seqlen, nheads, 128) FP8, contiguous.
+    """
+
+    def _ensure_fp8_descale(scale: torch.Tensor) -> torch.Tensor:
+        """Per-tensor descale as shape (1,) float32 for flash_attn_fp8_pertensor_func."""
+        return scale.to(dtype=torch.float32).reshape(1).contiguous()
+
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    q_descale = _ensure_fp8_descale(q_scale)
+    k_descale = _ensure_fp8_descale(k_scale)
+    v_descale = _ensure_fp8_descale(v_scale)
+
+    return aiter.flash_attn_fp8_pertensor_func(
+        q,
+        k,
+        v,
+        q_descale,
+        k_descale,
+        v_descale,
+        causal=is_causal,
+        softmax_scale=softmax_scale,
+        window_size=(-1, -1, 0),
+    )
+
+
+class AITerBackend(AttentionBackend):
+    """
+    Backend for AITemplate attention implementation.
+    """
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.AITER
+
+    @staticmethod
+    def get_impl_cls() -> type["AITerImpl"]:
+        return AITerImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        # AITer backend does not require special metadata.
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError("AITer backend does not have a metadata builder.")
+
+
+class AITerImpl(AttentionImpl):
+    """
+    Implementation of attention using AITemplate.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        dropout_p: float = 0.0,
+        **extra_impl_args,
+    ) -> None:
+        if num_kv_heads is not None and num_kv_heads != num_heads:
+            raise NotImplementedError(
+                "AITer backend does not support Grouped Query Attention yet."
+            )
+        self.causal = causal
+        self.dropout_p = dropout_p
+        self.softmax_scale = softmax_scale
+
+    @torch.compiler.disable
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        """
+        Performs attention using one of:
+          - _fmha_fp8_prefill_attention (FP8, SGLANG_DIFFUSION_AITER_FP8_ATTN=1 when eligible)
+          - flash_attn_func (BF16, default or FP8 fallback for unsupported shapes)
+
+        Args:
+            query: Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            key: Key tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            value: Value tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            attn_metadata: Metadata for the attention operation (unused).
+
+        Returns:
+            Output tensor of shape [batch_size, seq_len, num_heads, head_dim]
+        """
+        if _use_fp8_attn:
+            if query.dtype != _fp8_dtype:
+                q_fp8, q_scale = aiter.per_tensor_quant(query, quant_dtype=_fp8_dtype)
+                k_fp8, k_scale = aiter.per_tensor_quant(key, quant_dtype=_fp8_dtype)
+                v_fp8, v_scale = aiter.per_tensor_quant(value, quant_dtype=_fp8_dtype)
+            else:
+                q_fp8, k_fp8, v_fp8 = query, key, value
+                one = torch.tensor(1.0, dtype=torch.float32, device=query.device)
+                q_scale = k_scale = v_scale = one
+
+            d_q = q_fp8.shape[-1]
+            d_k = k_fp8.shape[-1]
+            d_v = v_fp8.shape[-1]
+            h_q = q_fp8.shape[2]
+            h_kv = k_fp8.shape[2]
+
+            if _can_use_fmha_fp8_prefill(d_q, d_k, d_v, h_q, h_kv):
+                return _fmha_fp8_prefill_attention(
+                    q_fp8,
+                    k_fp8,
+                    v_fp8,
+                    softmax_scale=self.softmax_scale,
+                    is_causal=self.causal,
+                    q_scale=q_scale,
+                    k_scale=k_scale,
+                    v_scale=v_scale,
+                )
+
+            logger.warning_once(
+                "FP8 FMHA prefill unsupported for this shape (need gfx950-class AITER, "
+                "full MHA, q/k/v head_dim=%d; got q=%d, k=%d, v=%d, num_heads=%d, "
+                "num_kv_heads=%d). Falling back to BF16.",
+                _FMHA_FP8_HEAD_DIM,
+                d_q,
+                d_k,
+                d_v,
+                h_q,
+                h_kv,
+            )
+
+        # BF16 path
+        output, _ = aiter.flash_attn_func(
+            query,
+            key,
+            value,
+            dropout_p=self.dropout_p,
+            causal=self.causal,
+            return_attn_probs=False,
+            return_lse=True,
+        )
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+class AITERSageBackend(AttentionBackend):
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.AITER_SAGE
+
+    @staticmethod
+    def get_impl_cls() -> type["AITERSageImpl"]:
+        return AITERSageImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        # AITER Sage backend does not require special metadata.
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError(
+            "AITER Sage backend does not have a metadata builder."
+        )
+
+
+class AITERSageImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        dropout_p: float = 0.0,
+        **extra_impl_args,
+    ) -> None:
+
+        try:
+            from aiter.ops.triton.attention.fav3_sage import fav3_sage_wrapper_func
+
+            self.aiter_sage_attn_fn = fav3_sage_wrapper_func
+        except ImportError:
+            raise ImportError(
+                "AITER Sage attention is not available, please update AITER version."
+            )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        """
+        Performs attention using aiter sage backend.
+
+        Args:
+            query: Query tensor of shape [batch_size, seq_len, head_num, head_dim]
+            key: Key tensor of shape [batch_size, seq_len, head_num, head_dim]
+            value: Value tensor of shape [batch_size, seq_len, head_num, head_dim]
+            attn_metadata: Metadata for the attention operation (unused).
+
+        Returns:
+            Output tensor of shape [batch_size, seq_len, head_num, head_dim]
+        """
+
+        output = self.aiter_sage_attn_fn(query, key, value)
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py

@@ -0,0 +1,104 @@
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class AscendFAMetadata:
+    pass
+
+
+class AscendFAMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> AttentionMetadata:
+        return AscendFAMetadata()
+
+
+class AscendFABackend(AttentionBackend):
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["AscendFAImpl"]:
+        return AscendFAImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return AscendFAMetadataBuilder
+
+
+class AscendFAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        return_softmax_lse: bool = False,
+    ) -> torch.Tensor:
+        mask = None
+        num_heads, num_key_value_heads = query.shape[2], key.shape[2]
+        if self.causal:
+            seq_len = query.shape[1]
+            mask = torch.triu(
+                torch.ones(seq_len, seq_len, device=query.device), diagonal=1
+            ).bool()
+        # transpose to bs, heads, seq_len, head_dim
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        output, lse = torch.ops.npu.npu_fused_infer_attention_score(
+            query,
+            key,
+            value,
+            num_heads=num_heads,
+            num_key_value_heads=num_key_value_heads,
+            scale=self.softmax_scale,
+            input_layout="BNSD",
+            softmax_lse_flag=return_softmax_lse,
+            atten_mask=mask,
+        )
+        output = output.transpose(1, 2)
+        if return_softmax_lse:
+            return output, lse
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py

@@ -0,0 +1,179 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    pass
+
+import torch
+
+from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+class AttentionBackend(ABC):
+    """Abstract class for attention backends."""
+
+    # For some attention backends, we allocate an output tensor before
+    # calling the custom op. When piecewise cudagraph is enabled, this
+    # makes sure the output tensor is allocated inside the cudagraph.
+    accept_output_buffer: bool = False
+
+    @staticmethod
+    @abstractmethod
+    def get_enum() -> AttentionBackendEnum:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> type["AttentionImpl"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    # @staticmethod
+    # @abstractmethod
+    # def get_state_cls() -> Type["AttentionState"]:
+    #     raise NotImplementedError
+
+    # @classmethod
+    # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
+    #     return cls.get_metadata_cls()(*args, **kwargs)
+
+    @staticmethod
+    @abstractmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return None
+
+
+@dataclass
+class AttentionMetadata:
+    """Attention metadata for prefill and decode batched together."""
+
+    # Current step of diffusion process
+    current_timestep: int
+
+    def asdict_zerocopy(self, skip_fields: set[str] | None = None) -> dict[str, Any]:
+        """Similar to dataclasses.asdict, but avoids deepcopying."""
+        if skip_fields is None:
+            skip_fields = set()
+        # Note that if we add dataclasses as fields, they will need
+        # similar handling.
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self)
+            if field.name not in skip_fields
+        }
+
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+
+class AttentionMetadataBuilder(ABC, Generic[T]):
+    """Abstract class for attention metadata builders."""
+
+    @abstractmethod
+    def __init__(self) -> None:
+        """Create the builder, remember some configuration and parameters."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare(self) -> None:
+        """Prepare for one batch."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def build(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> AttentionMetadata:
+        """Build attention metadata with on-device tensors."""
+        raise NotImplementedError
+
+
+class AttentionLayer(Protocol):
+
+    _k_scale: torch.Tensor
+    _v_scale: torch.Tensor
+    _k_scale_float: float
+    _v_scale_float: float
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor: ...
+
+
+class AttentionImpl(ABC, Generic[T]):
+
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        raise NotImplementedError
+
+    def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor:
+        """Preprocess QKV tensor before performing attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom preprocessing
+        like reshaping, tiling, scaling, or other transformations.
+
+        Called AFTER all_to_all for distributed attention
+
+        """
+        return qkv
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        """Postprocess the output tensor after the attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom postprocessing
+        like untiling, scaling, or other transformations.
+
+        Called BEFORE all_to_all for distributed attention
+
+        """
+
+        return output
+
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+
+def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
+    return wrap_method_with_debug_kernel_once(
+        attn_impl,
+        "forward",
+        op_name=f"diffusion.attn_impl.{attn_impl.__class__.__name__}.forward",
+    )

python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py

@@ -0,0 +1,279 @@
+from dataclasses import dataclass
+from typing import Any
+
+import attentions  # noqa: F401
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
+    LaserAttentionBackend,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+BSA_BLOCK_SIZE = 128
+
+
+class BlockSparseAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.BLOCK_SPARSE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["BlockSparseAttentionImpl"]:
+        return BlockSparseAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["BlockSparseAttentionMetadata"]:
+        return BlockSparseAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["BlockSparseAttentionMetadataBuilder"]:
+        return BlockSparseAttentionMetadataBuilder
+
+
+@dataclass
+class BlockSparseAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    skip_first_steps: int
+    sparsity: float
+    block_frame_stride: int
+
+
+class BlockSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        skip_first_steps: int,
+        sparsity: float,
+        raw_latent_shape: list[int],
+        patch_size: tuple[int, int, int],
+        **kwargs: dict[str, Any],
+    ) -> BlockSparseAttentionMetadata:
+        """
+        Builds BlockSparseAttention metadata.
+
+        Args:
+            current_timestep: The current diffusion timestep.
+            skip_first_steps: Number of initial timesteps to skip before applying
+                sparsity. Must be non‑negative.
+            sparsity: Fraction of tokens to drop (block‑wise) in the block sparse
+                attention mechanism. Must be in the range [0.0, 1.0).
+            raw_latent_shape: Shape of the latent tensor before patching.
+            patch_size: Patch size as (T, height, width). Only the height
+                and width components are used to divide the latent dimensions.
+            **kwargs: Additional keyword arguments (ignored, but accepted for
+                compatibility with base class or calling conventions).
+
+        Returns:
+            BlockSparseAttentionMetadata
+        Note:
+            The `block_frame_stride` is needed to set the first blocks to be non‑sparse.
+        """
+        if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
+            raise ValueError(
+                (
+                    "Invalid attention metadata values."
+                    f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
+                    f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
+                )
+            )
+
+        if sparsity == 0.0:
+            logger.warning(
+                (
+                    "Sparsity is set to 0.0, which means no tokens will be dropped."
+                    "For better performance use Laser Attention or increase sparsity."
+                )
+            )
+
+        if len(raw_latent_shape) >= 5:
+            latent_height, latent_width = raw_latent_shape[3:5]
+        else:
+            latent_height, latent_width = raw_latent_shape[-2:]
+
+        latent_height //= patch_size[1]
+        latent_width //= patch_size[2]
+
+        frame_stride = latent_height * latent_width
+        block_frame_stride = (frame_stride + BSA_BLOCK_SIZE - 1) // BSA_BLOCK_SIZE
+
+        return BlockSparseAttentionMetadata(
+            current_timestep=current_timestep,
+            skip_first_steps=skip_first_steps,
+            sparsity=sparsity,
+            block_frame_stride=block_frame_stride,
+        )
+
+
+class BlockSparseAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.block_size = BSA_BLOCK_SIZE
+        self.stride = 8
+        self.default_tokens = 214748647
+
+        self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _get_estimate_mask(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        sparsity: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.sparse_block_estimate(
+            query=query,
+            key=key,
+            actual_seq_lengths=None,
+            actual_seq_lengths_kv=None,
+            input_layout="BNSD",
+            stride=self.stride,
+            sparse_size=self.block_size,
+            num_heads=query.shape[1],
+            num_key_value_heads=key.shape[1],
+            scale_value=self.softmax_scale / self.stride,
+            threshold=1.0,
+            causal=self.causal,
+            keep_sink=True,
+            keep_recent=True,
+            row_sparse=1.0 - sparsity,
+        )
+
+    def _block_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        smask: torch.Tensor,
+        sct: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.ops.attentions.block_sparse_attention(
+            query=query,
+            key=key,
+            value=value,
+            sparse_mask=smask,
+            sparse_count_table=sct,
+            input_layout="BNSD",
+            sparse_size=self.block_size,
+            num_heads=query.shape[1],
+            num_key_value_heads=key.shape[1],
+            scale_value=self.softmax_scale,
+            causal=self.causal,
+            inner_precise=1,
+            pre_tokens=self.default_tokens,
+            next_tokens=self.default_tokens,
+            actual_seq_lengths=None,
+            actual_seq_lengths_kv=None,
+        )
+
+    def _get_smask(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        block_frame_stride: int,
+        sparsity: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        smask, sct = self._get_estimate_mask(
+            query,
+            key,
+            sparsity,
+        )
+
+        seq_len = smask.shape[2]
+
+        # Set the first blocks to be non-sparse to ensure the quality of the first few steps
+        smask[:, :, :block_frame_stride, :seq_len] = 1
+        smask[:, :, :seq_len, :block_frame_stride] = 1
+        smask = smask.to(torch.int8)
+        sct = smask.sum(dim=-1, dtype=torch.int32)
+        return smask, sct
+
+    def _adaptive_block_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        block_frame_stride: int,
+        sparsity: float,
+    ) -> torch.Tensor:
+        # TODO Currently implementation for BSND input layout has quality issues
+        # When the implementation is improved, transposes can be removed
+        q = query.permute(0, 2, 1, 3).contiguous()
+        k = key.permute(0, 2, 1, 3).contiguous()
+        v = value.permute(0, 2, 1, 3).contiguous()
+
+        smask, sct = self._get_smask(
+            q,
+            k,
+            block_frame_stride,
+            sparsity,
+        )
+        output = self._block_sparse_attention(q, k, v, smask, sct)
+        output = output.permute(0, 2, 1, 3).contiguous()
+
+        return output
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
+            output = self.laser_attn_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            output = self._adaptive_block_sparse_attention(
+                query,
+                key,
+                value,
+                attn_metadata.block_frame_stride,
+                attn_metadata.sparsity,
+            )
+
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py

@@ -0,0 +1,445 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
+from sglang.multimodal_gen.runtime.layers.utils import register_custom_op
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+)
+
+
+def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+# -----------------------------
+# Fake implementations for schema / tracing
+# custom op schema requires FIXED return structure.
+# We provide TWO ops:
+# 1) out-only op: always returns Tensor
+# 2) out+lse op: always returns Tuple[Tensor, Tensor]
+# -----------------------------
+def flash_attn_varlen_func_fake_out(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> torch.Tensor:
+    assert ver == 4, "only support flash attention v4"
+    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+    head_dim_v = v.shape[-1]
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (
+            batch_size + 1,
+        ), "cu_seqlens_q must have shape (batch_size + 1,)"
+        assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
+        assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
+
+    assert q.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+
+    q_batch_seqlen_shape = (
+        (batch_size, seqlen_q) if cu_seqlens_q is None else (q.shape[0],)
+    )
+    out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
+    return out
+
+
+def flash_attn_varlen_func_fake_out_lse(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = True,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert ver == 4, "only support flash attention v4"
+    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+        total_q = batch_size * seqlen_q
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+        total_q = q.shape[0]
+    head_dim_v = v.shape[-1]
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (
+            batch_size + 1,
+        ), "cu_seqlens_q must have shape (batch_size + 1,)"
+        assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
+        assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
+
+    assert q.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+
+    q_batch_seqlen_shape = (
+        (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,)
+    )
+    lse_shape = (
+        (batch_size, num_head, seqlen_q)
+        if cu_seqlens_q is None
+        else (num_head, total_q)
+    )
+
+    out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
+    lse = q.new_empty(lse_shape, dtype=torch.float32)
+    return out, lse
+
+
+# -----------------------------
+# Registered custom ops
+# NOTE: fixed return schemas to avoid:
+# "Object of type 'Tensor' is not an instance of 'sequence'"
+# -----------------------------
+@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out)
+def flash_attn_varlen_func_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> torch.Tensor:
+    if window_size is None:
+        window_size = [-1, -1]
+    if return_softmax_lse:
+        raise ValueError(
+            "flash_attn_varlen_func_op is out-only op; return_softmax_lse must be False. "
+            "Use flash_attn_varlen_func_op_lse for (out, lse)."
+        )
+    return flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        qv=qv,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        window_size=tuple(window_size),
+        attention_chunk=attention_chunk,
+        softcap=softcap,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        return_softmax_lse=False,
+        sinks=sinks,
+        ver=ver,
+    )
+
+
+@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out_lse)
+def flash_attn_varlen_func_op_lse(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = True,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if window_size is None:
+        window_size = [-1, -1]
+    if not return_softmax_lse:
+        raise ValueError(
+            "flash_attn_varlen_func_op_lse is out+lse op; return_softmax_lse must be True. "
+            "Use flash_attn_varlen_func_op for out-only."
+        )
+    return flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        qv=qv,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        window_size=tuple(window_size),
+        attention_chunk=attention_chunk,
+        softcap=softcap,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        return_softmax_lse=True,
+        sinks=sinks,
+        ver=ver,
+    )
+
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+
+fa_ver = 3
+
+
+def set_fa_ver(ver: int) -> None:
+    global fa_ver
+    fa_ver = ver
+
+
+@dataclass
+class FlashAttentionMetadata:
+    # Sequence lengths for the forward batch
+    # Maximum sequence length for query
+    max_seqlen_q: int = 1
+    # Maximum sequence length for key
+    max_seqlen_k: int = 0
+    # Cumulative sequence lengths for query
+    cu_seqlens_q: torch.Tensor = None
+    # Cumulative sequence lengths for key
+    cu_seqlens_k: torch.Tensor = None
+
+
+class FlashAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore
+        self,
+        raw_latent_shape=list,
+        **kwargs: dict[str, Any],
+    ) -> FlashAttentionMetadata:
+        # TODO: put empty values here to be set at first-run, since the q_len calculation can be complicated
+        return FlashAttentionMetadata(max_seqlen_q=None, max_seqlen_k=None)
+
+
+class FlashAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+
+class FlashAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.attention_metadata = FlashAttentionMetadata()
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+        *,
+        return_softmax_lse: bool = False,
+    ):
+        if attn_metadata is not None:
+            if attn_metadata.max_seqlen_q is None:
+                attn_metadata.max_seqlen_q = query.shape[1]
+            if attn_metadata.max_seqlen_k is None:
+                attn_metadata.max_seqlen_k = key.shape[1]
+            max_seqlen_q = attn_metadata.max_seqlen_q
+            max_seqlen_k = attn_metadata.max_seqlen_k
+        else:
+            max_seqlen_q = query.shape[1]
+            max_seqlen_k = key.shape[1]
+
+        # FA version selection:
+        # - fa_ver == 3: call python function (can return Tensor or (Tensor, Tensor) depending on flag)
+        # - fa_ver == 4: call custom ops with FIXED return schema
+        if fa_ver == 3:
+            flash_attn_op = flash_attn_varlen_func
+            output = flash_attn_op(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=None,
+                cu_seqlens_k=None,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                return_softmax_lse=return_softmax_lse,
+                ver=fa_ver,
+            )
+            return output
+
+        if fa_ver == 4:
+            if return_softmax_lse:
+                out_tensor, softmax_lse = flash_attn_varlen_func_op_lse(
+                    q=query,
+                    k=key,
+                    v=value,
+                    cu_seqlens_q=None,
+                    cu_seqlens_k=None,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    softmax_scale=self.softmax_scale,
+                    causal=self.causal,
+                    return_softmax_lse=True,
+                    ver=fa_ver,
+                )
+                return out_tensor, softmax_lse
+            out_tensor = flash_attn_varlen_func_op(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=None,
+                cu_seqlens_k=None,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                return_softmax_lse=False,
+                ver=fa_ver,
+            )
+            return out_tensor
+
+        raise ValueError(f"flash attention version {fa_ver} is not supported.")

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py

@@ -0,0 +1,79 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
+    flash_attn_func,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class FlashAttention2Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA2
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttention2Impl"]:
+        return FlashAttention2Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+
+class FlashAttention2Impl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ):
+        output = flash_attn_func(
+            q=query,  # type: ignore[no-untyped-call]
+            k=key,
+            v=value,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            max_seqlen_q=None,
+            max_seqlen_k=None,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+        )
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py

@@ -0,0 +1,191 @@
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPABackend
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+# Import to use torch.ops.attentions, install package with sgl_kernel_npu
+try:
+    import attentions  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        (
+            "The required 'attentions' package is not installed."
+            "The package can be installed with sgl_kernel_npu"
+        )
+    ) from e
+
+logger = init_logger(__name__)
+
+
+class LaserAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.LASER_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["LaserAttentionImpl"]:
+        return LaserAttentionImpl
+
+
+class LaserAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.softmax_scale = softmax_scale
+
+        # After preprocess input layout should be BNSD.
+        self.seqlen_base = 256
+        self.seqlen_index = 2
+        self.dim_index = 3
+        self.dim_base = 128
+        self.max_token = 2**31 - 1
+        self.seq_len_pad_base = 256
+
+        # the laser attention operator has issues with small seq_len
+        self.min_seqlen = 2048
+        self.sdpa_impl = SDPABackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _pad(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        """
+        Pad the input tensor along the sequence length and head dimension.
+        to multiples of base values. self.seqlen_index and self.dim_index should be positive integers.
+        """
+
+        seq_len = input_tensor.size(self.seqlen_index)
+        head_dim = input_tensor.size(self.dim_index)
+
+        pad_seq = 0
+        if seq_len % self.seqlen_base != 0:
+            pad_seq = ((seq_len // self.seqlen_base) + 1) * self.seqlen_base - seq_len
+
+        pad_dim = 0
+        if head_dim % self.dim_base != 0:
+            pad_dim = ((head_dim // self.dim_base) + 1) * self.dim_base - head_dim
+
+        if pad_seq == 0 and pad_dim == 0:
+            return input_tensor
+
+        pad_list = [0] * (2 * input_tensor.ndim)
+
+        pad_list[len(pad_list) - 2 * self.seqlen_index - 1] = pad_seq
+        pad_list[len(pad_list) - 2 * self.dim_index - 1] = pad_dim
+
+        return torch.nn.functional.pad(input_tensor, pad_list)
+
+    def _la_preprocess_input(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Currently BSND input layout is not supported
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+
+        if q.dtype != torch.float16:
+            q = q.to(torch.float16)
+            k = k.to(torch.float16)
+            v = v.to(torch.float16)
+
+        q = self._pad(q)
+        k = self._pad(k)
+        v = self._pad(v)
+
+        return q, k, v
+
+    def _la_postprocess_output(
+        self,
+        attention_out: torch.Tensor,
+        dtype: torch.dtype,
+        qseqlen: int,
+        head_dim: int,
+    ) -> torch.Tensor:
+        if dtype != attention_out.dtype:
+            attention_out = attention_out.to(dtype)
+
+        attention_out = attention_out[:, :, :qseqlen, :head_dim]
+        attention_out = attention_out.transpose(1, 2).contiguous()
+        return attention_out
+
+    def _laser_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        head_num: int,
+        pre_tokens: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.la(
+            query=query,
+            key=key,
+            value=value,
+            atten_mask=None,
+            alibi_mask=None,
+            drop_mask=None,
+            scale_value=self.softmax_scale,
+            head_num=head_num,
+            input_layout="BNSD",
+            keep_prob=1.0,
+            pre_tokens=pre_tokens,
+            next_tokens=1,
+            is_highPrecision=True,
+        )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        q_seqlen, head_dim = query.shape[1], query.shape[3]
+        kv_seqlen = key.shape[1]
+
+        if q_seqlen < self.min_seqlen or kv_seqlen != q_seqlen:
+            output = self.sdpa_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            pre_tokens = self.max_token
+            if kv_seqlen % self.seq_len_pad_base != 0:
+                pre_tokens = (
+                    kv_seqlen // self.seq_len_pad_base + 1
+                ) * self.seq_len_pad_base - kv_seqlen
+
+            q, k, v = self._la_preprocess_input(query, key, value)
+            _, la_output = self._laser_attention(q, k, v, q.shape[1], pre_tokens)
+            output = self._la_postprocess_output(
+                la_output, query.dtype, q_seqlen, head_dim
+            )
+
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py

@@ -0,0 +1,414 @@
+import math
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import attentions  # noqa: F401
+import torch
+from einops import rearrange
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
+    LaserAttentionBackend,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class RainFusionAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.RAIN_FUSION_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["RainFusionAttentionImpl"]:
+        return RainFusionAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["RainFusionAttentionMetadata"]:
+        return RainFusionAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["RainFusionAttentionMetadataBuilder"]:
+        return RainFusionAttentionMetadataBuilder
+
+
+@dataclass
+class RainFusionAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    skip_first_steps: int
+    sparsity: float
+    latent_shape: list[int]
+
+
+class RainFusionAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        skip_first_steps: int,
+        sparsity: float,
+        raw_latent_shape: list[int],
+        patch_size: tuple[int, int, int],
+        **kwargs: dict[str, Any],
+    ) -> RainFusionAttentionMetadata:
+        if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
+            raise ValueError(
+                (
+                    "Invalid attention metadata values."
+                    f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
+                    f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
+                )
+            )
+
+        if sparsity == 0.0:
+            logger.warning(
+                (
+                    "Sparsity is set to 0.0, which means no tokens will be dropped."
+                    "For better performance use Laser Attention or increase sparsity."
+                )
+            )
+
+        latent_shape = raw_latent_shape[-3:]
+        latent_shape = [latent_shape[i] // patch_size[i] for i in range(3)]
+
+        return RainFusionAttentionMetadata(
+            current_timestep=current_timestep,
+            skip_first_steps=skip_first_steps,
+            sparsity=sparsity,
+            latent_shape=latent_shape,
+        )
+
+
+class RainFusionAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.block_size = 128
+        self.inner_precise = 0
+
+        self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _avgpool(
+        self, input_tensor: torch.Tensor, pool_size: int = 128
+    ) -> torch.Tensor:
+        batch, seqlen, heads, dim = input_tensor.shape
+        x = input_tensor.permute(0, 2, 3, 1).reshape(batch * heads, dim, seqlen)
+
+        pooled = torch.nn.functional.avg_pool1d(
+            x, kernel_size=pool_size, stride=pool_size, ceil_mode=True
+        )
+        out = pooled.reshape(batch, heads, dim, -1).permute(0, 3, 1, 2).contiguous()
+
+        return out
+
+    def _get_mask_index(self, mask: torch.Tensor) -> torch.Tensor:
+        batch_size, num_heads, seq_len, _ = mask.shape
+
+        mask_reshaped = mask.reshape(-1, seq_len)
+        row_indices = torch.arange(
+            seq_len, device=mask.device, dtype=torch.float32
+        ).unsqueeze(0)
+
+        sorted_vals = torch.where(mask_reshaped, row_indices, seq_len)
+        sorted_vals, _ = torch.sort(sorted_vals, dim=-1)
+        valid_count = mask_reshaped.sum(dim=-1, keepdim=True)
+        keep_mask = row_indices < valid_count
+        result = torch.where(keep_mask, sorted_vals, -1)
+
+        pos_matrix = result.reshape(batch_size, num_heads, seq_len, seq_len).to(
+            torch.int64
+        )
+        return pos_matrix
+
+    def _get_blockwise_mask(
+        self,
+        qkv_pool: torch.Tensor,
+        sparsity: float,
+        scale: float,
+        pool_size: int,
+        latent_shape: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        first_frame_len = latent_shape[1] * latent_shape[2]
+
+        query_pool, key_pool, value_pool = torch.chunk(qkv_pool, 3, dim=0)
+        attn_scores = (
+            query_pool.permute(0, 2, 1, 3) @ key_pool.permute(0, 2, 3, 1) * scale
+        )
+
+        keep_len = math.ceil(attn_scores.shape[-1] * (1 - sparsity))
+
+        topk_values, _ = torch.topk(attn_scores, k=keep_len, dim=-1)
+        mask = attn_scores >= topk_values[..., -1:]
+
+        firstframe_block_num = (first_frame_len + pool_size - 1) // pool_size
+        if firstframe_block_num > 0:
+            mask[:, :, :firstframe_block_num, :] = True
+            mask[:, :, :, :firstframe_block_num] = True
+
+        select_idx = self._get_mask_index(mask)
+        select_idx = select_idx[0].transpose(0, 1)
+        select_num_idx = mask[0].transpose(0, 1).sum(dim=-1)
+        return select_idx, select_num_idx
+
+    def _rearrange_with_remaining(
+        self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
+    ) -> torch.Tensor:
+        """
+        b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d
+        or
+        b n (f hn hb wn wb) d -> b n (f hn wn hb wb) d
+        """
+        tq, hq, wq = latent_shape
+        first_frame_len, frame_num = hq * wq, tq
+
+        b, s, n, d = tensor.shape
+
+        if (hq % 8 != 0) or (wq % 8 != 0):
+            tensor_first = tensor[:, :first_frame_len, :, :]
+            tensor = tensor[:, first_frame_len:, :, :]
+            tensor_hwt = rearrange(
+                tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
+            )
+            if hq % 8 != 0:
+                tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
+                tensor_h_r = tensor_h_r.reshape(b, frame_num - 1, -1, n, d)
+            if wq % 8 != 0:
+                tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
+                tensor_w_r = tensor_w_r.reshape(b, frame_num - 1, -1, n, d)
+            tensor_hwt = rearrange(
+                tensor_hwt,
+                "b f (hn hb) (wn wb) n d -> b f (hn wn hb wb) n d",
+                f=frame_num - 1,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+            if hq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
+            if wq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=2)
+            tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
+            tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
+        else:
+            tensor_hwt = rearrange(
+                tensor,
+                "b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d",
+                f=frame_num,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+
+        return tensor_hwt
+
+    def _inv_rearrange_with_remaining(
+        self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
+    ) -> torch.Tensor:
+        tq, hq, wq = latent_shape
+        first_frame_len, frame_num = hq * wq, tq
+
+        b, s, n, d = tensor.shape
+
+        if (hq % 8 != 0) or (wq % 8 != 0):
+            tensor_first = tensor[:, :first_frame_len, :, :]
+            tensor = tensor[:, first_frame_len:, :, :]
+            tensor_hwt = rearrange(
+                tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
+            )
+            if hq % 8 != 0:
+                tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
+            if wq % 8 != 0:
+                tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
+            tensor_hwt = tensor_hwt.reshape(b, frame_num - 1, -1, n, d)
+            tensor_hwt = rearrange(
+                tensor_hwt,
+                "b f (hn wn hb wb) n d -> b f (hn hb) (wn wb) n d",
+                f=frame_num - 1,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+            if wq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=3)
+            if hq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
+            tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
+            tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
+        else:
+            tensor_hwt = rearrange(
+                tensor,
+                "b (f hn wn hb wb) n h -> b (f hn hb wn wb) n h",
+                f=frame_num,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+
+        return tensor_hwt
+
+    def _do_tensor_rearrange_pooling(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        pool_size: int,
+        latent_shape: tuple[int, int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Tensor block rearrangement + pooling operation
+        """
+        tensor = torch.cat((query, key, value), dim=0)
+
+        tensor = self._rearrange_with_remaining(tensor, latent_shape)
+        tensor_pool = self._avgpool(tensor, pool_size)
+
+        query_, key_, value_ = torch.chunk(tensor, 3, dim=0)
+        return query_, key_, value_, tensor_pool
+
+    def _rain_fusion_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        select_idx: torch.Tensor,
+        select_num_idx: torch.Tensor,
+        blockshape: List[int],
+        scale: float = 1.0,
+        head_num: int = 1,
+        input_layout: str = "TND",
+        actual_seq_lengths=Optional[torch.Tensor],
+        actual_seq_lengths_kv=Optional[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.rainfusionattention(
+            query=query,
+            key=key,
+            value=value,
+            select_idx=select_idx,
+            select_num_idx=select_num_idx,
+            blockshape=blockshape,
+            attn_mask=None,
+            actual_seq_qlen=actual_seq_lengths,
+            actual_seq_kvlen=actual_seq_lengths_kv,
+            block_table=None,
+            q_input_layout=input_layout,
+            kv_input_layout=input_layout,
+            head_num=head_num,
+            mask_type=0,
+            scale=scale,
+            inner_precise=self.inner_precise,
+            block_size=0,
+        )
+
+    def _rain_fusion_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        latent_shape: tuple[int, int, int],
+        sparsity: float,
+    ):
+        q, k, v, qkv_pool = self._do_tensor_rearrange_pooling(
+            query, key, value, self.block_size, latent_shape
+        )
+
+        select_idx, select_num_idx = self._get_blockwise_mask(
+            qkv_pool,
+            sparsity,
+            self.softmax_scale,
+            self.block_size,
+            latent_shape,
+        )
+
+        batch_size, seqlen_q, head_num, head_dim = q.shape
+        seqlen_kv = k.shape[1]
+
+        layout = "TND"
+        q = q.reshape(-1, head_num, head_dim)
+        k = k.reshape(-1, head_num, head_dim)
+        v = v.reshape(-1, head_num, head_dim)
+
+        actual_seq_lengths = [seqlen_q] * batch_size
+        actual_seq_lengths_kv = [seqlen_kv] * batch_size
+
+        out, _ = self._rain_fusion_attention(
+            q,
+            k,
+            v,
+            scale=self.softmax_scale,
+            head_num=head_num,
+            input_layout=layout,
+            select_idx=select_idx,
+            select_num_idx=select_num_idx,
+            blockshape=[self.block_size, self.block_size],
+            actual_seq_lengths=actual_seq_lengths,
+            actual_seq_lengths_kv=actual_seq_lengths_kv,
+        )
+
+        out = out.reshape(batch_size, seqlen_q, head_num, head_dim)
+        out = self._inv_rearrange_with_remaining(out, latent_shape)
+        return out
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
+            output = self.laser_attn_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            output = self._rain_fusion_sparse_attention(
+                query,
+                key,
+                value,
+                attn_metadata.latent_shape,
+                attn_metadata.sparsity,
+            )
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py

@@ -0,0 +1,74 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+from sageattention import sageattn
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (  # FlashAttentionMetadata,
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttentionImpl"]:
+        return SageAttentionImpl
+
+
+class SageAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        *,
+        return_softmax_lse: bool = False,
+    ) -> torch.Tensor:
+        output = sageattn(
+            query,
+            key,
+            value,
+            # since input is (batch_size, seq_len, head_num, head_dim)
+            tensor_layout="NHD",
+            is_causal=self.causal,
+            sm_scale=self.softmax_scale,
+            return_lse=return_softmax_lse,
+        )
+        if return_softmax_lse:
+            output, softmax_lse = output
+            return output, softmax_lse
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py

@@ -0,0 +1,92 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn.functional as F
+from sageattn3 import sageattn3_blackwell
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttention3Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_ATTN_3
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+
+class SageAttention3Impl(AttentionImpl):
+    _warned_gqa_fallback_global: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        # SageAttention3's Blackwell kernel assumes MHA (Hq == Hkv). For GQA/MQA
+        # (Hq != Hkv), fall back to torch SDPA which supports GQA.
+        if key.shape[1] != query.shape[1]:
+            if query.shape[1] % key.shape[1] != 0:
+                raise ValueError(
+                    "GQA/MQA requires query heads to be a multiple of KV heads, "
+                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
+                )
+            if not type(self)._warned_gqa_fallback_global:
+                logger.warning(
+                    "SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA."
+                )
+                type(self)._warned_gqa_fallback_global = True
+            output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                is_causal=self.causal,
+                dropout_p=self.dropout,
+                scale=self.softmax_scale,
+                enable_gqa=True,
+            )
+        else:
+            output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
+        output = output.transpose(1, 2)
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py

@@ -0,0 +1,95 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from contextlib import nullcontext
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (  # FlashAttentionMetadata,
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
+    SDPBackend.CUDNN_ATTENTION,
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+]
+
+
+class SDPABackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.TORCH_SDPA
+
+    @staticmethod
+    def get_impl_cls() -> type["SDPAImpl"]:
+        return SDPAImpl
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SDPAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+        self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        # transpose to bs, heads, seq_len, head_dim
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        attn_kwargs = {
+            "attn_mask": None,
+            "dropout_p": self.dropout,
+            "is_causal": self.causal,
+            "scale": self.softmax_scale,
+        }
+        if query.shape[1] != key.shape[1]:
+            attn_kwargs["enable_gqa"] = True
+        sdpa_context = (
+            sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+            if self.allow_cudnn_sdp and query.device.type == "cuda"
+            else nullcontext()
+        )
+        with sdpa_context:
+            output = torch.nn.functional.scaled_dot_product_attention(
+                query, key, value, **attn_kwargs
+            )
+        output = output.transpose(1, 2)
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py

@@ -0,0 +1,316 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import json
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from einops import rearrange
+
+from sglang.multimodal_gen.runtime.distributed import get_sp_group
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import (
+    ForwardContext,
+    get_forward_context,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import dict_to_3d_list
+
+try:
+    from st_attn import sliding_tile_attention
+
+    st_attn_backend_available = True
+except Exception:
+    st_attn_backend_available = False
+
+logger = init_logger(__name__)
+
+
+class RangeDict(dict):
+
+    def __getitem__(self, item: int) -> str:
+        for key in self.keys():
+            if isinstance(key, tuple):
+                low, high = key
+                if low <= item <= high:
+                    return str(super().__getitem__(key))
+            elif key == item:
+                return str(super().__getitem__(key))
+        raise KeyError(f"seq_len {item} not supported for STA")
+
+
+class SlidingTileAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # TODO(will-refactor): check this
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SLIDING_TILE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
+        return SlidingTileAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
+        return SlidingTileAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
+        return SlidingTileAttentionMetadataBuilder
+
+
+@dataclass
+class SlidingTileAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    STA_param: list[
+        list[Any]
+    ]  # each timestep with one metadata, shape [num_layers, num_heads]
+
+
+class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        STA_param: list[list[Any]],
+        current_timestep: int,
+        **kwargs: dict[str, Any],
+    ) -> SlidingTileAttentionMetadata:
+        param = STA_param
+        if param is None:
+            return SlidingTileAttentionMetadata(
+                current_timestep=current_timestep, STA_param=[]
+            )
+        return SlidingTileAttentionMetadata(
+            current_timestep=current_timestep, STA_param=param[current_timestep]
+        )
+
+
+class SlidingTileAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        if not st_attn_backend_available:
+            raise ValueError("st attn not supported")
+        # TODO(will-refactor): for now this is the mask strategy, but maybe we should
+        # have a more general config for STA?
+        mask_strategy_file_path = (
+            get_global_server_args().attention_backend_config.mask_strategy_file_path
+        )
+        if mask_strategy_file_path is None:
+            raise ValueError("SGLANG_DIFFUSION_ATTENTION_CONFIG is not set")
+
+        # TODO(kevin): get mask strategy for different STA modes
+        with open(mask_strategy_file_path) as f:
+            mask_strategy = json.load(f)
+        self.mask_strategy = dict_to_3d_list(mask_strategy)
+
+        self.prefix = prefix
+        sp_group = get_sp_group()
+        self.sp_size = sp_group.world_size
+        # STA config
+        self.STA_base_tile_size = [6, 8, 8]
+        self.dit_seq_shape_mapping = RangeDict(
+            {
+                (115200, 115456): "30x48x80",
+                82944: "36x48x48",
+                69120: "18x48x80",
+            }
+        )
+        self.full_window_mapping = {
+            "30x48x80": [5, 6, 10],
+            "36x48x48": [6, 6, 6],
+            "18x48x80": [3, 6, 10],
+        }
+
+    def tile(self, x: torch.Tensor) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2],
+        )
+
+    def untile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(
+            x,
+            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2],
+        )
+        return x
+
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        img_sequence_length = qkv.shape[1]
+        self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length]
+        self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str]
+        self.dit_seq_shape_int = list(map(int, self.dit_seq_shape_str.split("x")))
+        self.img_seq_length = (
+            self.dit_seq_shape_int[0]
+            * self.dit_seq_shape_int[1]
+            * self.dit_seq_shape_int[2]
+        )
+        return self.tile(qkv)
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.untile(output)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+        if self.mask_strategy is None:
+            raise ValueError("mask_strategy cannot be None for SlidingTileAttention")
+        if self.mask_strategy[0] is None:
+            raise ValueError("mask_strategy[0] cannot be None for SlidingTileAttention")
+
+        timestep = attn_metadata.current_timestep
+        forward_context: ForwardContext = get_forward_context()
+        forward_batch = forward_context.forward_batch
+        if forward_batch is None:
+            raise ValueError("forward_batch cannot be None")
+        # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl'
+        layer_idx = int(self.prefix.split(".")[-3])
+        if attn_metadata.STA_param is None or len(attn_metadata.STA_param) <= layer_idx:
+            raise ValueError("Invalid STA_param")
+        STA_param = attn_metadata.STA_param[layer_idx]
+
+        text_length = q.shape[1] - self.img_seq_length
+        has_text = text_length > 0
+
+        query = q.transpose(1, 2).contiguous()
+        key = k.transpose(1, 2).contiguous()
+        value = v.transpose(1, 2).contiguous()
+
+        head_num = query.size(1)
+        sp_group = get_sp_group()
+        current_rank = sp_group.rank_in_group
+        start_head = current_rank * head_num
+
+        # searching or tuning mode
+        if len(STA_param) < head_num * sp_group.world_size:
+            sparse_attn_hidden_states_all = []
+            full_mask_window = STA_param[-1]
+            for window_size in STA_param[:-1]:
+                sparse_hidden_states = sliding_tile_attention(
+                    query,
+                    key,
+                    value,
+                    [window_size] * head_num,
+                    text_length,
+                    has_text,
+                    self.dit_seq_shape_str,
+                ).transpose(1, 2)
+                sparse_attn_hidden_states_all.append(sparse_hidden_states)
+
+            hidden_states = sliding_tile_attention(
+                query,
+                key,
+                value,
+                [full_mask_window] * head_num,
+                text_length,
+                has_text,
+                self.dit_seq_shape_str,
+            ).transpose(1, 2)
+
+            attn_L2_loss = []
+            attn_L1_loss = []
+            # average loss across all heads
+            for sparse_attn_hidden_states in sparse_attn_hidden_states_all:
+                # L2 loss
+                attn_L2_loss_ = (
+                    torch.mean(
+                        (sparse_attn_hidden_states.float() - hidden_states.float())
+                        ** 2,
+                        dim=[0, 1, 3],
+                    )
+                    .cpu()
+                    .numpy()
+                )
+                attn_L2_loss_ = [round(float(x), 6) for x in attn_L2_loss_]
+                attn_L2_loss.append(attn_L2_loss_)
+                # L1 loss
+                attn_L1_loss_ = (
+                    torch.mean(
+                        torch.abs(
+                            sparse_attn_hidden_states.float() - hidden_states.float()
+                        ),
+                        dim=[0, 1, 3],
+                    )
+                    .cpu()
+                    .numpy()
+                )
+                attn_L1_loss_ = [round(float(x), 6) for x in attn_L1_loss_]
+                attn_L1_loss.append(attn_L1_loss_)
+
+            layer_loss_save = {"L2_loss": attn_L2_loss, "L1_loss": attn_L1_loss}
+
+            if forward_batch.is_cfg_negative:
+                if forward_batch.mask_search_final_result_neg is not None:
+                    forward_batch.mask_search_final_result_neg[timestep].append(
+                        layer_loss_save
+                    )
+            else:
+                if forward_batch.mask_search_final_result_pos is not None:
+                    forward_batch.mask_search_final_result_pos[timestep].append(
+                        layer_loss_save
+                    )
+        else:
+            windows = [STA_param[head_idx + start_head] for head_idx in range(head_num)]
+
+            hidden_states = sliding_tile_attention(
+                query,
+                key,
+                value,
+                windows,
+                text_length,
+                has_text,
+                self.dit_seq_shape_str,
+            ).transpose(1, 2)
+
+        return hidden_states

python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py

@@ -0,0 +1,695 @@
+"""
+Copyright (c) 2025 by SLA team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+
+This implementation is adapted from: from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py and https://github.com/thu-ml/SLA/blob/main/SageSLA/core.py
+Citation (please cite if you use this code):
+
+@article{zhang2025sla,
+  title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention},
+  author={Jintao Zhang and Haoxu Wang and Kai Jiang and Shuo Yang and Kaiwen Zheng and Haocheng Xi and Ziteng Wang and Hongzhou Zhu and Min Zhao and Ion Stoica and Joseph E. Gonzalez and Jun Zhu and Jianfei Chen},
+  journal={arXiv preprint arXiv:2509.24006},
+  year={2025}
+}
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# ==================================SLA Functions===================================
+def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
+    arg_k = k - torch.mean(
+        k, dim=-2, keepdim=True
+    )  # smooth-k technique in SageAttention
+    pooled_qblocks = mean_pool(q, BLKQ)
+    pooled_kblocks = mean_pool(arg_k, BLKK)
+    pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
+
+    K = pooled_score.shape[-1]
+    topk = min(K, int(topk_ratio * K))
+    lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
+
+    sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
+    sparse_map.scatter_(-1, lut, 1)
+    return sparse_map, lut, topk
+
+
+def mean_pool(x, BLK):
+    assert x.is_contiguous()
+
+    B, H, L, D = x.shape
+    L_BLOCKS = (L + BLK - 1) // BLK
+    x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
+
+    grid = (L_BLOCKS, B * H)
+    compress_kernel[grid](x, x_mean, L, D, BLK)
+    return x_mean
+
+
+@triton.jit
+def compress_kernel(
+    X,
+    XM,
+    L: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_L: tl.constexpr,
+):
+    idx_l = tl.program_id(0)
+    idx_bh = tl.program_id(1)
+
+    offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
+    offs_d = tl.arange(0, D)
+
+    x_offset = idx_bh * L * D
+    xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
+    x = tl.load(
+        X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L
+    )
+
+    nx = min(BLOCK_L, L - idx_l * BLOCK_L)
+    x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
+    tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))
+
+
+@triton.jit
+def _attn_fwd(
+    Q,
+    K,
+    V,
+    qk_scale: tl.constexpr,
+    topk: tl.constexpr,
+    LUT,
+    LSE,
+    OS,
+    L: tl.constexpr,
+    M_BLOCKS: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    idx_m = tl.program_id(0).to(tl.int64)
+    idx_bh = tl.program_id(1).to(tl.int64)
+
+    qkv_offset = idx_bh * L * D
+    lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk
+    lse_offset = idx_bh * L
+    offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, D)
+
+    Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None]
+    V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :]
+    OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    LUT_ptr = LUT + lut_offset
+    LSE_ptrs = LSE + lse_offset + offs_m
+
+    m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
+
+    q = tl.load(Q_ptrs, mask=offs_m[:, None] < L)
+    for block_idx in tl.range(topk):
+        idx_n = tl.load(LUT_ptr + block_idx)
+        n_mask = offs_n < L - idx_n * BLOCK_N
+
+        k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
+        qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634)  # = 1 / ln(2)
+        if L - idx_n * BLOCK_N < BLOCK_N:
+            qk = tl.where(n_mask[None, :], qk, float("-inf"))
+
+        v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
+        local_m = tl.max(qk, 1)
+        new_m = tl.maximum(m_i, local_m)
+        qk = qk - new_m[:, None]
+
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        alpha = tl.math.exp2(m_i - new_m)
+        o_s = o_s * alpha[:, None]
+        o_s += tl.dot(p.to(v.dtype), v)
+
+        l_i = l_i * alpha + l_ij
+        m_i = new_m
+
+    o_s = o_s / l_i[:, None]
+    tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
+
+    m_i += tl.math.log2(l_i)
+    tl.store(LSE_ptrs, m_i, mask=offs_m < L)
+
+
+def _get_cuda_arch(device_index: int) -> str:
+    """Get CUDA architecture string for the given device."""
+    major, minor = torch.cuda.get_device_capability(device_index)
+    return f"sm{major}{minor}"
+
+
+# ==================================SLA Class===================================
+class SparseLinearAttentionBackend(AttentionBackend):
+    """Sparse Linear Attention Backend for efficient attention computation."""
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SLA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SparseLinearAttentionImpl"]:
+        return SparseLinearAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SparseLinearAttentionMetadata"]:
+        return SparseLinearAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SparseLinearAttentionMetadataBuilder"]:
+        return SparseLinearAttentionMetadataBuilder
+
+
+@dataclass
+class SparseLinearAttentionMetadata(AttentionMetadata):
+    """Metadata for Sparse Linear Attention computation."""
+
+    # Basic attention parameters
+    current_timestep: int
+
+    # Sparse attention configuration
+    topk_ratio: float = 0.1
+
+
+class SparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder):
+    """Builder for SparseLinearAttentionMetadata."""
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        topk_ratio: float = 0.1,
+        **kwargs: dict[str, Any],
+    ) -> SparseLinearAttentionMetadata:
+        return SparseLinearAttentionMetadata(
+            current_timestep=current_timestep,
+            topk_ratio=topk_ratio,
+        )
+
+
+class SparseLinearAttentionImpl(AttentionImpl, nn.Module):
+    """Implementation of sparse linear attention for the backend."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool = False,
+        softmax_scale: float | None = None,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        # SLA-specific parameters - matched to TurboDiffusion defaults
+        topk_ratio: float = 0.1,  # TurboDiffusion uses topk=0.1
+        feature_map: str = "softmax",
+        BLKQ: int = 128,  # TurboDiffusion uses BLKQ=128
+        BLKK: int = 64,  # TurboDiffusion uses BLKK=64
+        use_bf16: bool = True,
+        **extra_impl_args,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        # SLA-specific config
+        self.topk_ratio = topk_ratio
+        self.BLKQ = BLKQ
+        self.BLKK = BLKK
+        self.dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+        # Learnable linear projection for combining sparse + linear attention
+        self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32)
+
+        # Feature map for linear attention
+        # Type annotation for callables
+        self.feature_map_q: Callable[[torch.Tensor], torch.Tensor]
+        self.feature_map_k: Callable[[torch.Tensor], torch.Tensor]
+        if feature_map == "elu":
+            self.feature_map_q = lambda x: F.elu(x) + 1
+            self.feature_map_k = lambda x: F.elu(x) + 1
+        elif feature_map == "relu":
+            self.feature_map_q = F.relu
+            self.feature_map_k = F.relu
+        elif feature_map == "softmax":
+            self.feature_map_q = lambda x: F.softmax(x, dim=-1)
+            self.feature_map_k = lambda x: F.softmax(x, dim=-1)
+        else:
+            raise ValueError(f"Unknown feature map: {feature_map}")
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        """Initialize projection weights to zero for residual-like behavior."""
+        with torch.no_grad():
+            nn.init.zeros_(self.proj_l.weight)
+            nn.init.zeros_(self.proj_l.bias)  # type: ignore[arg-type]
+
+    def _calc_linear_attention_with_torch(self, q, k, v):
+        kv = torch.matmul(k.transpose(-1, -2), v)
+        k_sum = torch.sum(k, dim=-2, keepdim=True)
+        return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseLinearAttentionMetadata = None,
+    ) -> torch.Tensor:
+        """Forward pass for sparse linear attention.
+
+        Args:
+            query: query tensor of shape (B, H, L, D)
+            key: key tensor of shape (B, H, L, D)
+            value: value tensor of shape (B, H, L, D)
+            attn_metadata: attention metadata containing configuration
+        Returns:
+            output tensor of shape (B, H, L, D)
+        """
+        dtype = query.dtype
+
+        # Transpose for computation
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+
+        # Get sparse attention map
+        sparse_map, lut, real_topk = get_block_map(
+            query, key, topk_ratio=self.topk_ratio, BLKQ=self.BLKQ, BLKK=self.BLKK
+        )
+
+        # Convert to computation dtype
+        query = query.to(self.dtype)
+        key = key.to(self.dtype)
+        value = value.to(self.dtype)
+
+        # Sparse attention computation
+        o_s = _attention.apply(
+            query, key, value, sparse_map, lut, real_topk, self.BLKQ, self.BLKK
+        )
+
+        # Apply feature maps
+        query = self.feature_map_q(query).to(self.dtype)  # c_q
+        key = self.feature_map_k(key).to(self.dtype)  # c_k
+        # Linear attention computation
+        o_l = self._calc_linear_attention_with_torch(query, key, value)
+
+        # Apply projection and combine results
+        with torch.amp.autocast("cuda", dtype=self.dtype):
+            o_l = self.proj_l(o_l)
+
+        # Combine sparse and linear attention
+        output = (o_s + o_l).to(dtype).transpose(1, 2)
+
+        return output
+
+
+class _attention(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None):
+        assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
+        assert k_block_id.is_contiguous() and lut.is_contiguous()
+
+        # We recommend the following two settings
+        assert BLOCK_M == 64 or BLOCK_M == 128
+        assert BLOCK_N == 64
+
+        B, H, L, D = q.shape
+        if qk_scale is None:
+            qk_scale = D**-0.5
+
+        M_BLOCKS = triton.cdiv(L, BLOCK_M)
+
+        o_s = torch.empty_like(v)
+        lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32)
+
+        grid = (M_BLOCKS, B * H)
+        _attn_fwd[grid](
+            q,
+            k,
+            v,
+            qk_scale,
+            topk,
+            lut,
+            lse,
+            o_s,
+            L,
+            M_BLOCKS,
+            D,
+            BLOCK_M,
+            BLOCK_N,
+            num_warps=4 if q.shape[-1] == 64 else 8,
+            num_stages=3,
+        )
+
+        ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s)
+        ctx.qk_scale = qk_scale
+        ctx.topk = topk
+        ctx.BLOCK_M = BLOCK_M
+        ctx.BLOCK_N = BLOCK_N
+        return o_s
+
+
+# ==================================SageSLA Class===================================
+SAGESLA_ENABLED = True
+try:
+    import spas_sage_attn._fused as fused
+    import spas_sage_attn._qattn as qattn
+    from spas_sage_attn.utils import block_map_lut_triton, get_vanilla_qk_quant
+except ImportError:
+    SAGESLA_ENABLED = False
+
+SAGE2PP_ENABLED = True
+try:
+    from spas_sage_attn._qattn import (
+        qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold,
+    )
+except ImportError:
+    SAGE2PP_ENABLED = False
+
+
+class SageSparseLinearAttentionBackend(AttentionBackend):
+    """Quantized Sparse-Linear Attention backend using SageAttention kernels."""
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_SLA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SageSparseLinearAttentionImpl"]:
+        return SageSparseLinearAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SageSparseLinearAttentionMetadata"]:
+        return SageSparseLinearAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SageSparseLinearAttentionMetadataBuilder"]:
+        return SageSparseLinearAttentionMetadataBuilder
+
+
+@dataclass
+class SageSparseLinearAttentionMetadata(AttentionMetadata):
+    """Metadata for Sage Sparse Linear Attention computation."""
+
+    # Basic attention parameters
+    current_timestep: int
+
+    # Sparse attention configuration
+    topk_ratio: float = 0.1
+
+
+class SageSparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder):
+    """Builder for SageSparseLinearAttentionMetadata."""
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        topk_ratio: float = 0.1,
+        **kwargs: dict[str, Any],
+    ) -> SageSparseLinearAttentionMetadata:
+        return SageSparseLinearAttentionMetadata(
+            current_timestep=current_timestep,
+            topk_ratio=topk_ratio,
+        )
+
+
+class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool = False,
+        softmax_scale: float | None = None,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        topk_ratio: float = 0.5,
+        feature_map: str = "softmax",
+        use_bf16: bool = True,
+        **extra_impl_args,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        assert (
+            SAGESLA_ENABLED
+        ), "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA."
+
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.softmax_scale = softmax_scale if softmax_scale else head_size**-0.5
+        self.causal = causal
+        self.prefix = prefix
+
+        self.topk_ratio = topk_ratio
+        self.dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+        # Learnable linear projection for combining sparse + linear attention
+        self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32)
+
+        # Feature map for linear attention
+        # Type annotation for callables
+        self.feature_map_q: Callable[[torch.Tensor], torch.Tensor]
+        self.feature_map_k: Callable[[torch.Tensor], torch.Tensor]
+        if feature_map == "elu":
+            self.feature_map_q = lambda x: F.elu(x) + 1
+            self.feature_map_k = lambda x: F.elu(x) + 1
+        elif feature_map == "relu":
+            self.feature_map_q = F.relu
+            self.feature_map_k = F.relu
+        elif feature_map == "softmax":
+            self.feature_map_q = lambda x: F.softmax(x, dim=-1)
+            self.feature_map_k = lambda x: F.softmax(x, dim=-1)
+        else:
+            raise ValueError(f"Unknown feature map: {feature_map}")
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        """Initialize projection weights to zero for residual-like behavior."""
+        with torch.no_grad():
+            nn.init.zeros_(self.proj_l.weight)
+            nn.init.zeros_(self.proj_l.bias)  # type: ignore[arg-type]
+
+    def _calc_linear_attention_with_torch(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ):
+        kv = torch.matmul(k.transpose(-1, -2), v)
+        k_sum = torch.sum(k, dim=-2, keepdim=True)
+        return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Forward pass for Sage Sparse Linear attention with quantized kernels.
+        Args:
+            query: query tensor of shape (B, L, H, D)
+            key: key tensor of shape (B, L, H, D)
+            value: value tensor of shape (B, L, H, D)
+            attn_metadata: attention metadata containing configuration
+        Returns:
+            output tensor of shape (B, L, H, D)
+        """
+        dtype = query.dtype
+
+        # Transpose from (B, L, H, D) to SLA format (B, H, L, D)
+        q = query.transpose(1, 2).contiguous()
+        k = key.transpose(1, 2).contiguous()
+        v = value.transpose(1, 2).contiguous()
+
+        # Determine block sizes based on GPU architecture
+        arch = _get_cuda_arch(q.device.index)
+
+        if arch == "sm90":
+            BLKQ = 64
+            BLKK = 128
+        else:
+            BLKQ = 128
+            BLKK = 64
+        # Compute block-sparse attention pattern
+        sparse_map, lut, real_topk = get_block_map(
+            q, k, topk_ratio=self.topk_ratio, BLKQ=BLKQ, BLKK=BLKK
+        )
+
+        # Convert to compute dtype
+        q = q.to(self.dtype)
+        k = k.to(self.dtype)
+        v = v.to(self.dtype)
+
+        ########## SPARGE BEGIN ##########
+        km = k.mean(dim=-2, keepdim=True)
+        headdim = q.size(-1)
+        assert headdim in [
+            64,
+            128,
+        ], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale."
+
+        # Quantize Q, K to INT8
+        q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK)
+        lut, valid_block_num = block_map_lut_triton(sparse_map)
+        scale = 1.0 / (headdim**0.5)
+
+        o_s = torch.empty_like(q)
+
+        if arch in ("sm80", "sm86", "sm87"):
+            pvthreshold = torch.full(
+                (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device
+            )
+            v_fp16 = v.to(torch.float16)
+            qattn.qk_int8_sv_f16_accum_f16_block_sparse_attn_inst_buf_with_pv_threshold(
+                q_int8,
+                k_int8,
+                v_fp16,
+                o_s,
+                lut,
+                valid_block_num,
+                pvthreshold,
+                q_scale,
+                k_scale,
+                1,
+                False,
+                1,
+                scale,
+                0,
+            )
+        else:
+            b, h_kv, kv_len, head_dim = v.shape
+            padded_len = (kv_len + 127) // 128 * 128
+            v_transposed_permutted = torch.empty(
+                (b, h_kv, head_dim, padded_len), dtype=v.dtype, device=v.device
+            )
+            fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1)
+            v_fp8 = torch.empty(
+                v_transposed_permutted.shape, dtype=torch.float8_e4m3fn, device=v.device
+            )
+            v_scale = torch.empty(
+                (b, h_kv, head_dim), dtype=torch.float32, device=v.device
+            )
+            fused.scale_fuse_quant_cuda(
+                v_transposed_permutted, v_fp8, v_scale, kv_len, 2.25, 1
+            )
+
+            if arch == "sm90":
+                qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90(
+                    q_int8,
+                    k_int8,
+                    v_fp8,
+                    o_s,
+                    lut,
+                    valid_block_num,
+                    q_scale,
+                    k_scale,
+                    v_scale,
+                    1,
+                    False,
+                    1,
+                    scale,
+                )
+            else:
+                pvthreshold = torch.full(
+                    (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device
+                )
+                if SAGE2PP_ENABLED:
+                    qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
+                        q_int8,
+                        k_int8,
+                        v_fp8,
+                        o_s,
+                        lut,
+                        valid_block_num,
+                        pvthreshold,
+                        q_scale,
+                        k_scale,
+                        v_scale,
+                        1,
+                        False,
+                        1,
+                        scale,
+                        0,
+                    )
+                else:
+                    qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
+                        q_int8,
+                        k_int8,
+                        v_fp8,
+                        o_s,
+                        lut,
+                        valid_block_num,
+                        pvthreshold,
+                        q_scale,
+                        k_scale,
+                        v_scale,
+                        1,
+                        False,
+                        1,
+                        scale,
+                        0,
+                    )
+
+        ########## SPARGE END ##########
+
+        # Linear attention with feature maps
+        q_linear = self.feature_map_q(q).to(self.dtype)
+        k_linear = self.feature_map_k(k).to(self.dtype)
+        o_l = self._calc_linear_attention_with_torch(q_linear, k_linear, v)
+
+        # Project linear attention output and combine
+        with torch.amp.autocast("cuda", dtype=self.dtype):
+            o_l = self.proj_l(o_l)
+
+        # Combine sparse and linear outputs
+        output = (o_s + o_l).to(dtype).transpose(1, 2)
+
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py

@@ -0,0 +1,562 @@
+"""
+Sparse Video Gen 2 (SAP) attention backend.
+
+This is a baseline integration that wires the backend into the
+attention framework.
+
+Adapted from https://github.com/svg-project/Sparse-VideoGen/blob/main/svg/models/wan/attention.py
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+try:
+    from svg.kernels.triton.permute import (
+        apply_inverse_permutation_triton,
+        permute_tensor_by_labels_triton,
+    )
+    from svg.kmeans_utils import (
+        batch_kmeans_Euclid,
+        dynamic_block_sparse_fwd_flashinfer,
+        identify_dynamic_map,
+    )
+
+    svg2_available = True
+except ImportError:
+    svg2_available = False
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SparseVideoGen2AttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SparseVideoGen2AttentionImpl"]:
+        return SparseVideoGen2AttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SparseVideoGen2AttentionMetadata"]:
+        return SparseVideoGen2AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SparseVideoGen2AttentionMetadataBuilder"]:
+        return SparseVideoGen2AttentionMetadataBuilder
+
+
+@dataclass
+class Svg2LayerCache:
+    # centroids for kmeans clustering
+    q_centroids: torch.Tensor | None = None
+    k_centroids: torch.Tensor | None = None
+    centroids_initialized: bool = False
+
+
+@dataclass
+class Svg2Cache:
+    layers: dict[int, Svg2LayerCache] = field(default_factory=dict)
+
+    def get_layer(self, layer_idx: int) -> Svg2LayerCache:
+        layer_cache = self.layers.get(layer_idx)
+        if layer_cache is None:
+            layer_cache = Svg2LayerCache()
+            self.layers[layer_idx] = layer_cache
+        return layer_cache
+
+
+@dataclass
+class SparseVideoGen2AttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    num_q_centroids: int
+    num_k_centroids: int
+    top_p_kmeans: float
+    min_kc_ratio: float
+    kmeans_iter_init: int
+    kmeans_iter_step: int
+    zero_step_kmeans_init: bool
+    first_layers_fp: float
+    first_times_fp: float
+    context_length: int
+    num_frame: int
+    frame_size: int
+    cache: Svg2Cache
+    prompt_length: int | None = None
+    max_seqlen_q: int | None = None
+    max_seqlen_k: int | None = None
+
+
+def _require_kwarg(kwargs: dict[str, Any], name: str) -> Any:
+    if name not in kwargs:
+        raise ValueError(
+            f"Missing required argument for SparseVideoGen2Attention: {name}"
+        )
+    return kwargs[name]
+
+
+class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore[override]
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, ...],
+        patch_size: tuple[int, int, int],
+        cache: Svg2Cache,
+        num_q_centroids: int,
+        num_k_centroids: int,
+        top_p_kmeans: float,
+        min_kc_ratio: float,
+        kmeans_iter_init: int,
+        kmeans_iter_step: int,
+        zero_step_kmeans_init: bool,
+        first_layers_fp: float,
+        first_times_fp: float,
+        context_length: int = 0,
+        prompt_length: int | None = None,
+        **kwargs: dict[str, Any],
+    ) -> SparseVideoGen2AttentionMetadata:
+        raw_shape = tuple(raw_latent_shape)
+        if len(raw_shape) == 5:
+            t, h, w = raw_shape[2:5]
+        elif len(raw_shape) == 3:
+            t, h, w = raw_shape
+        else:
+            raise ValueError(
+                "raw_latent_shape must be (T, H, W) or (B, C, T, H, W) for SAP attention"
+            )
+        pt, ph, pw = patch_size
+        if t % pt != 0 or h % ph != 0 or w % pw != 0:
+            raise ValueError(
+                "raw_latent_shape must be divisible by patch_size for SAP attention"
+            )
+
+        num_frame = t // pt
+        frame_size = (h // ph) * (w // pw)
+
+        return SparseVideoGen2AttentionMetadata(
+            current_timestep=current_timestep,
+            num_q_centroids=num_q_centroids,
+            num_k_centroids=num_k_centroids,
+            top_p_kmeans=top_p_kmeans,
+            min_kc_ratio=min_kc_ratio,
+            kmeans_iter_init=kmeans_iter_init,
+            kmeans_iter_step=kmeans_iter_step,
+            zero_step_kmeans_init=zero_step_kmeans_init,
+            first_layers_fp=first_layers_fp,
+            first_times_fp=first_times_fp,
+            context_length=context_length,
+            prompt_length=prompt_length,
+            num_frame=num_frame,
+            frame_size=frame_size,
+            cache=cache,
+        )
+
+
+class SparseVideoGen2AttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        if causal:
+            raise ValueError(
+                "Sparse Video Gen 2 attention does not support causal attention"
+            )
+        if not svg2_available:
+            raise ImportError(
+                "Sparse Video Gen 2 attention backend requires svg package to be installed"
+                "Please install it by following the instructions at "
+                "https://github.com/svg-project/Sparse-VideoGen"
+            )
+        self.prefix = prefix
+        self.layer_idx = self._get_layer_idx(prefix)
+
+    def _get_layer_idx(self, prefix: str) -> int:
+        parts = prefix.split(".")
+        if len(parts) < 3:
+            raise ValueError(
+                f"Invalid prefix for SparseVideoGen2AttentionImpl: {prefix}"
+            )
+        return int(parts[-3])
+
+    def kmeans_init(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+        qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid(
+            query.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_q_centroids,
+            max_iters=attn_metadata.kmeans_iter_init,
+        )
+        klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid(
+            key.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_k_centroids,
+            max_iters=attn_metadata.kmeans_iter_init,
+        )
+
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        layer_cache.q_centroids = qcentroids
+        layer_cache.k_centroids = kcentroids
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def kmeans_step(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid(
+            query.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_q_centroids,
+            max_iters=attn_metadata.kmeans_iter_step,
+            init_centroids=layer_cache.q_centroids,
+        )
+        klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid(
+            key.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_k_centroids,
+            max_iters=attn_metadata.kmeans_iter_step,
+            init_centroids=layer_cache.k_centroids,
+        )
+
+        layer_cache.q_centroids = qcentroids
+        layer_cache.k_centroids = kcentroids
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def kmeans_clustering(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        if not layer_cache.centroids_initialized:
+            (
+                qlabels,
+                qcentroids,
+                qcluster_sizes,
+                qiter,
+                klabels,
+                kcentroids,
+                kcluster_sizes,
+                kiter,
+            ) = self.kmeans_init(query, key, attn_metadata)
+            layer_cache.centroids_initialized = True
+            logger.debug(
+                "Centroids initialized at layer %s (init iters: %s).",
+                self.layer_idx,
+                attn_metadata.kmeans_iter_init,
+            )
+        else:
+            (
+                qlabels,
+                qcentroids,
+                qcluster_sizes,
+                qiter,
+                klabels,
+                kcentroids,
+                kcluster_sizes,
+                kiter,
+            ) = self.kmeans_step(query, key, attn_metadata)
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def semantic_aware_permutation(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+
+        # 1. Kmeans clustering
+        (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        ) = self.kmeans_clustering(query, key, attn_metadata)
+
+        # 2. Identify dynamic map
+        q_cluster_sizes = qcluster_sizes.view(
+            cfg, num_heads, attn_metadata.num_q_centroids
+        )
+        k_cluster_sizes = kcluster_sizes.view(
+            cfg, num_heads, attn_metadata.num_k_centroids
+        )
+
+        dynamic_map = identify_dynamic_map(
+            qcentroids.view(cfg, num_heads, attn_metadata.num_q_centroids, dim),
+            kcentroids.view(cfg, num_heads, attn_metadata.num_k_centroids, dim),
+            q_cluster_sizes,
+            k_cluster_sizes,
+            attn_metadata.top_p_kmeans,
+            attn_metadata.min_kc_ratio,
+        )
+
+        # 3. Permute the query, key, value
+        q_permuted, q_sorted_indices = permute_tensor_by_labels_triton(
+            query, qlabels, dim=2
+        )
+        k_permuted, k_sorted_indices = permute_tensor_by_labels_triton(
+            key, klabels, dim=2
+        )
+        v_permuted, v_sorted_indices = permute_tensor_by_labels_triton(
+            value, klabels, dim=2, sorted_indices=k_sorted_indices
+        )
+
+        return (
+            q_permuted,
+            k_permuted,
+            v_permuted,
+            dynamic_map,
+            q_cluster_sizes,
+            k_cluster_sizes,
+            q_sorted_indices,
+        )
+
+    def _hunyuan_dynamic_map_post_processing(
+        self,
+        q_perm: torch.Tensor,
+        k_perm: torch.Tensor,
+        v_perm: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        dyn_map: torch.Tensor,
+        qc_sz_s: torch.Tensor,
+        kc_sz_s: torch.Tensor,
+        q_sorted_indices: torch.Tensor,
+        video_length: int,
+        context_length: int,
+        prompt_length: int,
+        unprompt_length: int,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        # Place the permuted video tokens back and keep text tokens at the tail.
+        query[:, :, :-context_length, :] = q_perm
+        key[:, :, :-context_length, :] = k_perm
+        value[:, :, :-context_length, :] = v_perm
+
+        # Add prompt/unprompt clusters to the dynamic map.
+        dyn_map = F.pad(dyn_map, (0, 2, 0, 2), value=0)
+        dyn_map[:, :, -2, :-1] = True
+        dyn_map[:, :, :-1, -2] = True
+        dyn_map[:, :, -1, -1] = True
+
+        qc_sz_s = F.pad(qc_sz_s, (0, 2), value=0)
+        qc_sz_s[:, :, -2] = prompt_length
+        qc_sz_s[:, :, -1] = unprompt_length
+        kc_sz_s = F.pad(kc_sz_s, (0, 2), value=0)
+        kc_sz_s[:, :, -2] = prompt_length
+        kc_sz_s[:, :, -1] = unprompt_length
+
+        q_sorted_indices = F.pad(q_sorted_indices, (0, context_length), value=0)
+        q_sorted_indices[:, video_length:] = torch.arange(
+            video_length,
+            video_length + context_length,
+            device=q_sorted_indices.device,
+        )
+        return query, key, value, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ) -> torch.Tensor:
+        torch.backends.cuda.preferred_linalg_library(backend="magma")
+        res = None
+        # bshd -> bhsd
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+        batch_size, num_heads, seq_len, dim = query.size()
+
+        context_length, num_frame, frame_size = (
+            attn_metadata.context_length,
+            attn_metadata.num_frame,
+            attn_metadata.frame_size,
+        )
+        prompt_length = attn_metadata.prompt_length
+        if prompt_length is None:
+            prompt_length = context_length
+
+        assert (
+            seq_len == context_length + num_frame * frame_size
+        ), f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}"
+
+        # Determine if we use Full Attention to calculate
+        full_attention_flag = False
+
+        if self.layer_idx < attn_metadata.first_layers_fp:
+            full_attention_flag = True
+        if attn_metadata.current_timestep > attn_metadata.first_times_fp:
+            full_attention_flag = True
+
+        if full_attention_flag:
+            if attn_metadata.zero_step_kmeans_init:
+                video_length = attn_metadata.num_frame * attn_metadata.frame_size
+                query_video = query[:, :, :video_length, :].contiguous()
+                key_video = key[:, :, :video_length, :].contiguous()
+                self.kmeans_clustering(query_video, key_video, attn_metadata)
+
+            with sdpa_kernel(
+                SDPBackend.CUDNN_ATTENTION
+            ):  # not sure why we need to force cudnn here, but it's faster than flash attention
+                output_hidden_states = torch.nn.functional.scaled_dot_product_attention(
+                    query, key, value, dropout_p=0.0, is_causal=False
+                )
+
+            res = output_hidden_states.reshape(
+                batch_size, num_heads, seq_len, dim
+            ).transpose(1, 2)
+        else:
+            if context_length > 0:
+                video_length = num_frame * frame_size
+                unprompt_length = max(context_length - prompt_length, 0)
+                query_video = query[:, :, :video_length, :].contiguous()
+                key_video = key[:, :, :video_length, :].contiguous()
+                value_video = value[:, :, :video_length, :].contiguous()
+
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self.semantic_aware_permutation(
+                    query_video, key_video, value_video, attn_metadata
+                )
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self._hunyuan_dynamic_map_post_processing(
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    query,
+                    key,
+                    value,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                    video_length,
+                    context_length,
+                    prompt_length,
+                    unprompt_length,
+                )
+            else:
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self.semantic_aware_permutation(query, key, value, attn_metadata)
+
+            output_permuted = dynamic_block_sparse_fwd_flashinfer(
+                q_perm, k_perm, v_perm, dyn_map, qc_sz_s, kc_sz_s, is_cpu=False
+            )
+
+            attn_output = apply_inverse_permutation_triton(
+                output_permuted, q_sorted_indices, dim=2
+            )
+
+            res = attn_output.reshape(batch_size, num_heads, seq_len, dim).transpose(
+                1, 2
+            )
+
+        torch.backends.cuda.preferred_linalg_library(
+            backend="default"
+        )  # reset to default
+        return res.contiguous()