[diffusion] model: support SANA-WM with streaming support (#27531)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: sjmshsh <88866917+sjmshsh@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py

@@ -0,0 +1,253 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Inference-side helpers for the bidirectional fused GDN path.
+
+Precision knob: env var ``FUSED_GDN_PRECISION`` or ``PRECISION_OVERRIDE``:
+  0=IEEE fp32 dots, 1=TF32, 2=bf16 TC + fp32 state [default], 3=bf16 TC + bf16 state.
+"""
+
+# ruff: noqa: E501
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# =====================================================================
+#  GPU-adaptive kernel config
+# =====================================================================
+
+
+def _get_kernel_config() -> dict:
+    """Return optimal kernel parameters for the current GPU.
+
+    STATE_FP32 (fp32 state_prev) needs ~128KB SRAM (H100 228KB), vs ~96KB for
+    bf16 state_prev (fits GB10's 101KB).
+    """
+    if not torch.cuda.is_available():
+        return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 4, "STATE_FP32": False}
+    smem = torch.cuda.get_device_properties(0).shared_memory_per_multiprocessor
+    state_fp32 = smem >= 150 * 1024  # H100 (228KB) yes, GB10 (101KB) no
+    return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 8, "STATE_FP32": state_fp32}
+
+
+_KCFG = None
+
+
+def _kcfg():
+    global _KCFG
+    if _KCFG is None:
+        _KCFG = _get_kernel_config()
+    return _KCFG
+
+
+# precision=0 → IEEE fp32 dots + fp32 state  (DOT_PRECISION=2, STATE_FP32=1)
+# precision=1 → TF32  dots   + fp32 state    (DOT_PRECISION=1, STATE_FP32=1)
+# precision=2 → bf16  dots   + fp32 state    (DOT_PRECISION=0, STATE_FP32=1) [default]
+# precision=3 → bf16  dots   + bf16 state    (DOT_PRECISION=0, STATE_FP32=0)
+def _precision_params(precision: int) -> tuple:
+    if precision == 0:
+        return 2, True
+    elif precision == 1:
+        return 1, True
+    elif precision == 3:
+        return 0, False
+    else:  # default
+        return 0, True
+
+
+_env_prec = os.environ.get("FUSED_GDN_PRECISION", None)
+PRECISION_OVERRIDE: int | None = int(_env_prec) if _env_prec is not None else None
+
+
+def _resolve_launch_config() -> tuple:
+    """Returns (prec, dot_prec, state_fp32, num_warps).
+
+    Uses ``PRECISION_OVERRIDE`` when set, else ``_kcfg()`` (per-GPU SRAM).
+    num_warps clamped to 4 when dots run on fp32 operands (more registers).
+    """
+    cfg = _kcfg()
+    prec = PRECISION_OVERRIDE if PRECISION_OVERRIDE is not None else 2
+    dot_prec, state_fp32 = _precision_params(prec)
+    if PRECISION_OVERRIDE is None:
+        state_fp32 = cfg["STATE_FP32"]
+    nw = cfg["num_warps"]
+    if dot_prec >= 1:
+        nw = min(nw, 4)
+    return prec, dot_prec, state_fp32, nw
+
+
+def prepare_rope_tables(
+    rotary_emb, N: int, D: int, device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Complex rotary_emb `(1, 1, N, D//2)` → expanded (N, D) cos/sin tables.
+
+    Encodes the interleaved-pair rotation
+        y[2i]   = x[2i]*cos[i] - x[2i+1]*sin[i]
+        y[2i+1] = x[2i]*sin[i] + x[2i+1]*cos[i]
+    as  y[d] = x[d]*cos_exp[d] + x[d^1]*sin_exp[d]
+    where sin_exp[2i] = -sin[i], sin_exp[2i+1] = +sin[i].
+
+    Returns (cos_exp, sin_exp) both (N, D) float32, contiguous.
+    """
+    if rotary_emb is None:
+        return (
+            torch.ones(N, D, device=device, dtype=torch.float32),
+            torch.zeros(N, D, device=device, dtype=torch.float32),
+        )
+    freqs = rotary_emb.squeeze(0).squeeze(0)  # (N, D//2) complex
+    cos_half = freqs.real.float()
+    sin_half = freqs.imag.float()
+    rope_cos = cos_half.repeat_interleave(2, dim=-1)
+    rope_sin = torch.stack([-sin_half, sin_half], dim=-1).reshape(N, D)
+    return rope_cos.contiguous(), rope_sin.contiguous()
+
+
+def _precompute_inv_rms(
+    qkv: torch.Tensor, idx: int, C: int, eps: float = 1e-5
+) -> torch.Tensor:
+    """Compute 1/RMS for one component of QKV over the full C = H*D channel dim.
+
+    qkv: (B, N, 3, H, D); idx: 0=Q, 1=K, 2=V; C: H*D. Returns (B, N) float32.
+    """
+    raw = qkv[:, :, idx].float()  # (B, N, H, D)
+    sq_sum = (raw * raw).sum(dim=(-2, -1))  # (B, N)
+    return torch.rsqrt(sq_sum / C + eps)
+
+
+# =====================================================================
+#  Fused single-pass Q+K inverse-RMS Triton kernel
+# =====================================================================
+# Single Triton launch that reads each `(b, n)` row of `qkv` once and emits
+# both `q_inv_rms[b, n]` and `k_inv_rms[b, n]`. Replaces two separate PyTorch
+# scans (cast→square→sum→rsqrt) over `qkv[:, :, 0]` and `qkv[:, :, 1]`.
+#
+# Layout assumed: `qkv` is (B, N, 3, H, D) contiguous, so the C = H*D channels
+# for a given (b, n, qkv_idx) live in a contiguous memory span.
+
+
+@triton.jit
+def _fused_qk_inv_rms_kernel(
+    qkv_ptr,  # *T_in     (B, N, 3, H, D), contiguous
+    q_inv_rms_ptr,  # *float32  (B, N)
+    k_inv_rms_ptr,  # *float32  (B, N)
+    N: tl.constexpr,
+    C: tl.constexpr,  # H * D
+    eps,
+    BLOCK_C: tl.constexpr,
+):
+    bn_id = tl.program_id(0)
+    qkv_row_stride = 3 * C
+    row_base = bn_id * qkv_row_stride
+    q_base = row_base
+    k_base = row_base + C
+
+    offs = tl.arange(0, BLOCK_C)
+    mask = offs < C
+
+    q_vals = tl.load(qkv_ptr + q_base + offs, mask=mask, other=0.0).to(tl.float32)
+    k_vals = tl.load(qkv_ptr + k_base + offs, mask=mask, other=0.0).to(tl.float32)
+
+    q_sq = tl.sum(q_vals * q_vals, axis=0)
+    k_sq = tl.sum(k_vals * k_vals, axis=0)
+
+    inv_c = 1.0 / C
+    q_inv = tl.rsqrt(q_sq * inv_c + eps)
+    k_inv = tl.rsqrt(k_sq * inv_c + eps)
+
+    tl.store(q_inv_rms_ptr + bn_id, q_inv)
+    tl.store(k_inv_rms_ptr + bn_id, k_inv)
+
+
+def fused_qk_inv_rms(
+    qkv: torch.Tensor,
+    eps: float = 1e-5,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Single-pass Triton fused Q+K inverse-RMS.
+
+    Replaces two ``_precompute_inv_rms`` scans with one launch that reads each
+    ``(b, n)`` row of ``qkv`` exactly once.
+    qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32.
+    """
+    assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)"
+    assert (
+        qkv.dim() == 5 and qkv.shape[2] == 3
+    ), f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}"
+    B, N, _, H, D = qkv.shape
+    C = H * D
+    q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device)
+    k_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device)
+    BLOCK_C = triton.next_power_of_2(C)
+    _fused_qk_inv_rms_kernel[(B * N,)](
+        qkv,
+        q_inv_rms,
+        k_inv_rms,
+        N=N,
+        C=C,
+        eps=eps,
+        BLOCK_C=BLOCK_C,
+    )
+    return q_inv_rms, k_inv_rms
+
+
+# =====================================================================
+#  Bidirectional GDN entry point (delegates to chunkwise)
+# =====================================================================
+
+
+def fused_bigdn_func(
+    qkv: torch.Tensor,  # (B, N, 3, H, D)
+    q_inv_rms: torch.Tensor,  # (B, N) float32
+    k_inv_rms: torch.Tensor,  # (B, N) float32
+    q_norm_weight: torch.Tensor,  # (C,) float32
+    k_norm_weight: torch.Tensor,  # (C,) float32
+    rope_cos: torch.Tensor,  # (N, D) float32
+    rope_sin: torch.Tensor,  # (N, D) float32
+    beta: torch.Tensor,  # (B, H, F, S)
+    decay: torch.Tensor,  # (B, H, F)
+    F: int,
+    S: int,
+    k_scale: float,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Bidirectional fused GDN. Returns ``(B, N, H, D)``.
+
+    Thin entry point kept for call-site stability; delegates to
+    :func:`fused_bigdn_bidi_chunkwise` from ``fused_gdn_chunkwise``.
+    """
+    from sglang.jit_kernel.diffusion.triton.sana_wm_gdn_chunkwise import (
+        fused_bigdn_bidi_chunkwise,
+    )
+
+    return fused_bigdn_bidi_chunkwise(
+        qkv,
+        q_inv_rms,
+        k_inv_rms,
+        q_norm_weight,
+        k_norm_weight,
+        rope_cos,
+        rope_sin,
+        beta,
+        decay,
+        F=F,
+        S=S,
+        k_scale=k_scale,
+        eps=eps,
+    )

python/sglang/multimodal_gen/apps/realtime_webui/app.js

@@ -431,10 +431,7 @@ async function decodeFrameBatch(header, data) {
 }
 
 function isWorkerDecodableContentType(contentType) {
-  return (
-    isWorkerDecodableRawContentType(contentType) ||
-    isEncodedPreviewContentType(contentType)
-  );
+  return isWorkerDecodableRawContentType(contentType);
 }
 
 function isWorkerDecodableRawContentType(contentType) {
@@ -1238,7 +1235,7 @@ async function payloadToArrayBuffer(data) {
   return data.arrayBuffer();
 }
 
-function drawFrame(image) {
+function drawFrame(image, { close = true, markRendered = true } = {}) {
   const sourceWidth = image.width;
   const sourceHeight = image.height;
   let drawSource = image;
@@ -1258,9 +1255,9 @@ function drawFrame(image) {
   ctx.imageSmoothingEnabled = true;
   ctx.imageSmoothingQuality = "high";
   ctx.drawImage(drawSource, 0, 0, sourceWidth, sourceHeight);
-  renderedPreviewFrames += 1;
+  if (markRendered) renderedPreviewFrames += 1;
   setPreviewState("live");
-  if (!(image instanceof ImageData)) image.close?.();
+  if (close && !(image instanceof ImageData)) image.close?.();
 }
 
 function renderLoop(now) {
@@ -1528,11 +1525,11 @@ function receive(data, epoch) {
       const payload = message.payload;
       delete message.payload;
       enqueueDecodeBatch(message, payload, epoch);
-      setStatus("Live", "live");
+      if (!renderedPreviewFrames) setStatus("Receiving", "live");
       return;
     }
     pendingHeader = message;
-    if (pendingHeader) setStatus("Live", "live");
+    if (pendingHeader && !renderedPreviewFrames) setStatus("Receiving", "live");
     return;
   }
   const header = pendingHeader;
@@ -1557,6 +1554,9 @@ async function decodeAndEnqueueFrameBatch(header, data, epoch) {
     return;
   }
   const now = performance.now();
+  if (!renderedPreviewFrames && decodedFrames.length) {
+    drawFrame(decodedFrames[0].image, { close: false, markRendered: false });
+  }
   // record source frames before preview playback can hold or drop for latency
   recordDecodedFrameBatch(decodedFrames);
   const enqueueResult = playbackController.enqueueDecodedFrames(header, decodedFrames, now);

python/sglang/multimodal_gen/apps/realtime_webui/index.html

@@ -163,6 +163,6 @@
       </section>
     </main>
   <script src="./playback_controller.js?v=realtime-playback-v13"></script>
-  <script src="./app.js?v=realtime-record-v73"></script>
+  <script src="./app.js?v=realtime-record-v75"></script>
 </body>
 </html>

python/sglang/multimodal_gen/configs/models/dits/sana_wm.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks
+
+
+@dataclass
+class SanaWMArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_blocks_or_transformer_blocks]
+    )
+
+    # --- Core dims (upstream: depth=20, hidden=2240, heads=20, linear_head_dim=112) ---
+    patch_size: int = 1
+    in_channels: int = 128  # LTX-2 VAE latent channels
+    out_channels: int = 128
+    num_layers: int = 20
+
+    # Patch embedder uses (1, patch_size, patch_size) — temporal patch is always 1.
+    patch_size_t: int = 1
+
+    num_attention_heads: int = 20
+    attention_head_dim: int = 112  # = linear_head_dim
+    linear_head_dim: int = 112
+
+    # --- Cross-attention (text conditioning) ---
+    # In upstream, cross-attn uses num_heads (=20) with head_dim = hidden/num_heads = 112.
+    num_cross_attention_heads: int = 20
+    cross_attention_head_dim: int = 112
+    cross_attention_dim: int = 2240  # query dim used inside MultiHeadCrossAttention
+    cross_norm: bool = True
+
+    # Gemma-2-2b-it hidden size (input to y_embedder.y_proj).
+    caption_channels: int = 2304
+    model_max_length: int = 300
+    y_norm: bool = True
+    y_norm_scale_factor: float = 0.01
+    y_norm_eps: float = 1e-5
+
+    mlp_ratio: float = 3.0
+    qk_norm: bool = True
+    norm_eps: float = 1e-6
+    timestep_norm_scale_factor: float = 1.0
+
+    # --- Hybrid GDN/Softmax attention ---
+    # softmax_every_n=4 => blocks where (i+1)%4 == 0 use softmax main branch,
+    # i.e. block indices {3, 7, 11, 15, 19}.
+    softmax_every_n: int = 4
+
+    # --- GDN ShortConvolution params ---
+    conv_kernel_size: int = 4
+    k_conv_only: bool = True
+    chunk_gdn_chunk_size: int = 21
+    update_rule: str = "torch_chunk"  # main branch update rule
+    cam_update_rule: str = "torch_chunk"  # camera branch update rule
+    # main GDN scan backend: "auto" uses the SANA-WM Triton fast path on
+    # supported CUDA inference runs, otherwise falls back to the torch scan.
+    gdn_backend: str = "auto"
+
+    # --- Camera conditioning ---
+    cam_attn_compress: int = 1  # cam_dim == in_dim
+    init_cam_from_base: bool = True
+    use_chunk_plucker_post_attn: bool = True
+    use_chunk_plucker_input: bool = False
+    chunk_plucker_channels: int = 48  # 8 orig frames × 6D Plücker
+    chunk_plucker_post_attn_blocks: int = 20
+
+    chunk_split_strategy: str = "first_chunk_plus_one"
+    chunk_size: int = 10
+    # Upstream currently forwards chunk metadata through the softmax blocks but
+    # does not apply a chunk-causal mask there. Keep this disabled by default
+    # for checkpoint-output parity; it can be enabled for experiments.
+    use_chunked_softmax_attention: bool = False
+
+    # --- Temporal FFN (GLUMBConvTemp) ---
+    ffn_type: str = "GLUMBConvTemp"
+    t_kernel_size: int = 3
+    mlp_acts: tuple = field(default_factory=lambda: ("silu", "silu", None))
+
+    # --- Position embedding ---
+    pos_embed_type: str = "wan_rope"
+
+    # --- VAE coupling (LTX-2) ---
+    vae_temporal_stride: int = 8  # original-frames per latent frame
+    vae_spatial_stride: int = 32  # pixels per latent token (per spatial axis)
+
+    sample_size: int = 32  # legacy, unused
+    guidance_embeds: bool = False
+    class_dropout_prob: float = 0.0
+
+    # the released checkpoints store raw upstream parameter names; streaming
+    # also keeps an unused all-zero pos_embed while the native model uses RoPE
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            "^pos_embed$": "",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class SanaWMConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=SanaWMArchConfig)
+    prefix: str = "SanaWM"

python/sglang/multimodal_gen/configs/models/dits/sana_wm_refiner.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks
+
+
+@dataclass
+class SanaWMRefinerArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_blocks_or_transformer_blocks]
+    )
+
+    # Core dims
+    in_channels: int = 128
+    out_channels: int = 128
+    patch_size: int = 1
+    patch_size_t: int = 1
+    num_layers: int = 28
+    num_attention_heads: int = 32
+    attention_head_dim: int = 64
+    cross_attention_dim: int = 4096
+    caption_channels: int = 4096
+
+    qk_norm: bool = True
+    norm_eps: float = 1e-6
+    apply_gated_attention: bool = False
+
+    timestep_scale_multiplier: float = 1000.0
+    rope_type: str = "interleaved"
+
+    # RoPE coord generation
+    sampling_rate: int = 16000
+    hop_length: int = 160
+    scale_factors: tuple = (8, 32, 32)
+    base_num_frames: int = 20
+    base_height: int = 2048
+    base_width: int = 2048
+    causal_offset: int = 1
+
+    # Map Diffusers-style param keys to sglang's LTX-2 primitive naming.
+    # The refiner reuses LTX2Attention / LTX2FeedForward, so it inherits the
+    # same naming differences vs Diffusers:
+    #   * ff.net.0.proj / ff.net.2 -> proj_in / proj_out  (LTX2FeedForward)
+    #   * norm_q / norm_k          -> q_norm / k_norm     (LTX2Attention)
+    # Keep this aligned with LTX2ArchConfig.param_names_mapping in ltx_2.py.
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^(transformer_blocks\.\d+\.ff)\.net\.0\.proj\.(.*)$": r"\1.proj_in.\2",
+            r"^(transformer_blocks\.\d+\.ff)\.net\.2\.(.*)$": r"\1.proj_out.\2",
+            r"(.*)\.norm_q\.(.*)$": r"\1.q_norm.\2",
+            r"(.*)\.norm_k\.(.*)$": r"\1.k_norm.\2",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class SanaWMRefinerConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=SanaWMRefinerArchConfig)
+    prefix: str = "SanaWMRefiner"

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

@@ -317,6 +317,12 @@ class PipelineConfig:
     def prepare_calculated_size(self, image):
         return self.calculate_condition_image_size(image, image.width, image.height)
 
+    def preprocess_realtime_condition_image(self, batch, _vae_image_processor) -> bool:
+        """Realtime hook: optionally preprocess the first-frame condition image
+        in-place. Return True if handled (skip the standard path), False to fall
+        back to the normal condition-image preprocessing. Default: not handled."""
+        return False
+
     def prepare_image_processor_kwargs(self, batch, neg=False):
         return {}
 
@@ -744,6 +750,31 @@ class PipelineConfig:
             help="Override the selected pipeline config's resolution setting. Only applies to pipelines that define a resolution field.",
         )
 
+        # SANA-WM streaming knobs. default=None so they only apply to pipeline
+        # configs that define these fields (e.g. SanaWMPipelineConfig); other
+        # configs are left untouched by update_config_from_args.
+        parser.add_argument(
+            f"--{prefix_with_dot}streaming",
+            action=StoreBoolean,
+            dest=f"{prefix_with_dot.replace('-', '_')}streaming",
+            default=None,
+            help="SANA-WM: enable chunk-causal streaming (forward_long) generation.",
+        )
+        parser.add_argument(
+            f"--{prefix_with_dot}refiner-chunked",
+            action=StoreBoolean,
+            dest=f"{prefix_with_dot.replace('-', '_')}refiner_chunked",
+            default=None,
+            help="SANA-WM: chunk-wise streaming refiner (vs whole-clip dense refiner).",
+        )
+        parser.add_argument(
+            f"--{prefix_with_dot}num-frame-per-block",
+            type=int,
+            dest=f"{prefix_with_dot.replace('-', '_')}num_frame_per_block",
+            default=None,
+            help="SANA-WM: latent frames per streaming chunk (default 3).",
+        )
+
         # DiT configuration
         parser.add_argument(
             f"--{prefix_with_dot}dit-precision",
@@ -941,6 +972,30 @@ class PipelineConfig:
                 )
             # 1.5. Adjust pipeline config for fine-tuned VAE if needed
             pipeline_config_cls = model_info.pipeline_config_cls
+            # If an explicit pipeline_class_name refines the model-default config
+            # (e.g. SanaWMRealtimePipeline -> SanaWMRealtimeConfig, a subclass of
+            # the model-resolved SanaWMPipelineConfig), prefer the pipeline's own
+            # config so realtime-only wiring (the /v1/realtime_video adapter) is
+            # selected. Only applies when the explicit config strictly subclasses
+            # the model default, so non-realtime pipelines are unaffected.
+            if pipeline_class_name:
+                explicit_config_classes = get_pipeline_config_classes(
+                    pipeline_class_name
+                )
+                if explicit_config_classes is not None:
+                    explicit_config_cls = explicit_config_classes[0]
+                    if (
+                        isinstance(explicit_config_cls, type)
+                        and isinstance(pipeline_config_cls, type)
+                        and explicit_config_cls is not pipeline_config_cls
+                        and issubclass(explicit_config_cls, pipeline_config_cls)
+                    ):
+                        logger.info(
+                            f"Refining pipeline config {pipeline_config_cls.__name__} "
+                            f"-> {explicit_config_cls.__name__} for explicit "
+                            f"pipeline_class_name={pipeline_class_name}"
+                        )
+                        pipeline_config_cls = explicit_config_cls
         vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path")
         if vae_path is None:
             component_paths = kwargs.get(

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

@@ -24,6 +24,11 @@ from sglang.multimodal_gen.runtime.distributed import (
     get_sp_world_size,
 )
 
+# Distilled 3-step stage-2 sigma schedule. Single source of truth shared by
+# LTX2TwoStagePipeline and the SANA-WM refiner stages (matches NVlabs
+# `inference_sana_wm.py` / the LTX-2 distilled refiner).
+STAGE_2_DISTILLED_SIGMA_VALUES: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0)
+
 
 def pack_text_embeds(
     text_hidden_states: torch.Tensor,

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

@@ -24,3 +24,4 @@ class ModelDeploymentConfig:
     fsdp_auto_min_available_memory_gb: float | None = None
     fsdp_auto_requires_cfg: bool = True
     fsdp_auto_requires_default_parallelism: bool = True
+    auto_enable_cfg_parallel: bool = True

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

@@ -0,0 +1,335 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig
+from sglang.multimodal_gen.configs.models.dits.sana_wm import SanaWMConfig
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig
+from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config
+from sglang.multimodal_gen.configs.models.vaes.ltx_video import LTXVideoVAEConfig
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ModelTaskType,
+    PipelineConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def sana_wm_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
+    """Extract Gemma-2 last hidden state as text conditioning (same as SANA T2I)."""
+    return outputs.last_hidden_state
+
+
+SANA_WM_CHI_PROMPT: tuple[str, ...] = (
+    'Given a user prompt, generate an "Enhanced prompt" that provides detailed '
+    "visual descriptions suitable for image generation. Evaluate the level of "
+    "detail in the user prompt:",
+    "- If the prompt is simple, focus on adding specifics about colors, shapes, "
+    "sizes, textures, and spatial relationships to create vivid and concrete scenes.",
+    "- If the prompt is already detailed, refine and enhance the existing details "
+    "slightly without overcomplicating.",
+    "Here are examples of how to transform or refine prompts:",
+    "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled "
+    "up in a round shape, sleeping peacefully on a warm sunny windowsill, "
+    "surrounded by pots of blooming red flowers.",
+    "- User Prompt: A busy city street -> Enhanced: A bustling city street scene "
+    "at dusk, featuring glowing street lamps, a diverse crowd of people in "
+    "colorful clothing, and a double-decker bus passing by towering glass skyscrapers.",
+    "Please generate only the enhanced description for the prompt below and "
+    "avoid including any additional commentary or evaluations:",
+    "User Prompt: ",
+)
+
+
+@dataclass
+class SanaWMPipelineConfig(PipelineConfig):
+    """Pipeline config for the SANA-WM TI2V world model (text + first-frame image -> video,
+    optional 6-DoF camera trajectory, optional Stage-2 LTX-2 refiner)."""
+
+    task_type: ModelTaskType = ModelTaskType.TI2V
+
+    # SanaWMBeforeDenoisingStage._splice_first_frame handles condition-image
+    # resize + VAE-encode itself, so bypass the framework's generic TI2V
+    # preprocessing in InputValidationStage. Without this, the framework path
+    # reads `vae_config.arch_config.scale_factor_spatial` which LTXVideoVAEArchConfig
+    # does not expose (it uses `spatial_compression_ratio` instead). LTX-2 sets
+    # the same flag for the same reason -- both are TI2V on LTXVideoVAEConfig.
+    skip_input_image_preprocess: bool = True
+
+    # --- Guidance ---
+    # SANA-WM uses standard CFG via guidance_scale; no embedded guidance token.
+    should_use_guidance: bool = False
+
+    enable_autocast: bool = False
+
+    # --- Streaming self-forcing (S1c) ---
+    # When ``streaming`` is set, the pipeline uses the autoregressive
+    # SanaWMStreamingDenoisingStage (forward_long, chunk-by-chunk) instead of the
+    # one-shot bidirectional denoise. ``num_frame_per_block`` is the streaming
+    # chunk size in LATENT frames (distinct from the DiT's intra-attention
+    # ``arch.chunk_size``). ``denoising_step_list`` must end in 0.
+    streaming: bool = False
+    num_frame_per_block: int = 3
+    num_cached_blocks: int = 2
+    sink_token: bool = True
+    denoising_step_list: tuple = (1000, 960, 889, 727, 0)
+    streaming_cfg_scale: float = 1.0
+    # Streaming refiner (S2b): chunked LTX-2 sink/current refiner.
+    sink_size: int = 1
+    refiner_block_size: int = 3
+    refiner_kv_max_frames: int = 11
+    refiner_seed: int = 42
+    # True -> chunked streaming refiner (low-latency, causal); False -> whole-clip
+    # dense refiner (global context, max quality, non-streaming).
+    refiner_chunked: bool = True
+
+    # --- DiT ---
+    dit_config: DiTConfig = field(default_factory=SanaWMConfig)
+
+    # --- VAE: LTX-2 (128ch, 8× temporal, 32× spatial) ---
+    vae_config: VAEConfig = field(default_factory=LTXVideoVAEConfig)
+    vae_precision: str = "bf16"
+    # Match NVlabs SANA-WM inference: long videos must use LTX-2 spatial
+    # tiling plus framewise temporal decode to avoid oversized Conv3d pads.
+    vae_tiling: bool = True
+    vae_sp: bool = False  # no VAE SP for now
+    vae_framewise_encoding: bool = True
+    vae_framewise_decoding: bool = True
+    vae_tile_sample_min_num_frames: int = 96
+    vae_tile_sample_stride_num_frames: int = 64
+
+    # Load both encoder and decoder (need encoder for first-frame conditioning)
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.vae_config.use_tiling = self.vae_tiling
+        self.vae_config.use_temporal_tiling = self.vae_framewise_decoding
+        self.vae_config.tile_sample_min_num_frames = self.vae_tile_sample_min_num_frames
+        self.vae_config.tile_sample_stride_num_frames = (
+            self.vae_tile_sample_stride_num_frames
+        )
+        self.vae_config.blend_num_frames = (
+            self.vae_tile_sample_min_num_frames - self.vae_tile_sample_stride_num_frames
+        )
+
+    # --- Text encoder: Gemma-2-2b-it (single encoder, same as SANA T2I) ---
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Gemma2Config(),)
+    )
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
+    text_encoder_extra_args: list[dict] = field(
+        default_factory=lambda: [
+            {
+                # Match NVlabs SANA-WM prompt encoding: positive and negative
+                # branches must have the same token dimension for CFG concat.
+                "padding": "max_length",
+                "return_attention_mask": True,
+            }
+        ]
+    )
+    chi_prompt: tuple[str, ...] = SANA_WM_CHI_PROMPT
+    preprocess_text_funcs: tuple[Callable | None, ...] = field(
+        default_factory=lambda: (None,)
+    )
+    postprocess_text_funcs: tuple[Callable, ...] = field(
+        default_factory=lambda: (sana_wm_postprocess_text,)
+    )
+
+    # --- Scheduler ---
+    # linear_flow training schedule from the released config.yaml.
+    flow_shift: float = 9.95
+    # Official NVlabs/Sana inference resolves inference_flow_shift first and
+    # falls back to flow_shift only when it is absent.
+    inference_flow_shift: float | None = 9.8
+
+    # --- Video shape ---
+    # VAE strides: (temporal, spatial_h, spatial_w)
+    vae_stride: tuple = (8, 32, 32)  # LTX-2 VAE temporal=8, spatial=32
+
+    # --- Camera conditioning ---
+    camera_conditioning: bool = True  # set False to disable camera branch
+
+    # --- Deployment ---
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            auto_dit_layerwise_offload=True,
+            # Conservative auto-FSDP gate for the 720p world-model path. Users
+            # can still force FSDP explicitly on smaller cards.
+            fsdp_auto_min_available_memory_gb=60,
+        )
+
+    # --- Latent shape ---
+    def prepare_latent_shape(self, batch, batch_size: int, num_frames: int):
+        """
+        Returns 5D latent shape: (B, 128, T_latent, H_sp, W_sp).
+        T_latent = ceil((num_frames - 1) / temporal_stride) + 1
+        """
+        t_stride = self.vae_stride[0]
+        h_stride = self.vae_stride[1]
+        w_stride = self.vae_stride[2] if len(self.vae_stride) > 2 else h_stride
+
+        if batch.height % h_stride != 0 or batch.width % w_stride != 0:
+            raise ValueError(
+                "SANA-WM height/width must be divisible by the LTX-2 spatial "
+                f"stride ({h_stride}, {w_stride}); got "
+                f"height={batch.height}, width={batch.width}."
+            )
+
+        T_latent = (num_frames - 1) // t_stride + 1
+        H_sp = batch.height // h_stride
+        W_sp = batch.width // w_stride
+        z_dim = self.vae_config.arch_config.latent_channels  # 128
+
+        return (batch_size, z_dim, T_latent, H_sp, W_sp)
+
+    def adjust_num_frames(self, num_frames: int) -> int:
+        """Ensure (num_frames - 1) is divisible by VAE temporal stride."""
+        t_stride = self.vae_stride[0]
+        if (num_frames - 1) % t_stride != 0:
+            adjusted = ((num_frames - 1) // t_stride) * t_stride + 1
+            logger.warning(
+                f"num_frames - 1 must be divisible by temporal stride {t_stride}. "
+                f"Rounding {num_frames} → {adjusted}."
+            )
+            return adjusted
+        return num_frames
+
+    # --- Text embedding accessors ---
+    def get_pos_prompt_embeds(self, batch):
+        return batch.prompt_embeds[0]
+
+    def get_neg_prompt_embeds(self, batch):
+        return batch.negative_prompt_embeds[0]
+
+    # --- Conditioning kwargs for DenoisingStage ---
+    def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        """Build positive conditioning kwargs passed to SanaWMTransformer3DModel.forward.
+
+        The DiT forward signature consumes:
+          * encoder_hidden_states  -- Gemma-2 embeddings (set by DenoisingStage)
+          * timestep               -- diffusion step (set by DenoisingStage)
+          * encoder_attention_mask -- text padding mask
+          * camera_conditions      -- (B, T_lat, 20) latent-frame UCPE raymap
+          * chunk_plucker          -- (B, 48, T_lat, H, W) packed Plücker raymap
+        """
+        out = {}
+
+        m = batch.prompt_attention_mask
+        if isinstance(m, (list, tuple)):
+            out["encoder_attention_mask"] = m[0] if m else None
+        elif m is not None:
+            out["encoder_attention_mask"] = m
+
+        # Camera conditioning (built by SanaWMBeforeDenoisingStage)
+        if hasattr(batch, "extra") and batch.extra:
+            cc = batch.extra.get("camera_conditions", None)
+            cp = batch.extra.get("chunk_plucker", None)
+            if cc is not None:
+                out["camera_conditions"] = cc
+            if cp is not None:
+                out["chunk_plucker"] = cp
+
+        return out
+
+    def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        """Build negative conditioning kwargs for CFG.
+
+        Camera/plucker are structural video conditions, not text conditions.
+        NVlabs SANA-WM duplicates them for both CFG branches and only swaps
+        text embeddings/masks.
+        """
+        out = {}
+        m = batch.negative_attention_mask
+        if isinstance(m, (list, tuple)):
+            out["encoder_attention_mask"] = m[0] if m else None
+        elif m is not None:
+            out["encoder_attention_mask"] = m
+        if hasattr(batch, "extra") and batch.extra:
+            cc = batch.extra.get("camera_conditions", None)
+            cp = batch.extra.get("chunk_plucker", None)
+            if cc is not None:
+                out["camera_conditions"] = cc
+            if cp is not None:
+                out["chunk_plucker"] = cp
+        return out
+
+    # --- Post-processing ---
+    def post_denoising_loop(self, latents: torch.Tensor, batch) -> torch.Tensor:
+        """No token un-packing needed; 5D latents are already spatial."""
+        return latents
+
+    def shard_latents_for_sp(self, batch, latents):
+        # SANA-WM uses frame-wise GDN recurrent scan and a temporal depth-wise
+        # conv (GLUMBConvTemp, t_kernel=3) that both span across frames. Splitting
+        # the latent along T would truncate the GDN hidden state and drop the
+        # GLUMBConvTemp halo at rank boundaries, producing silent wrong outputs.
+        # Camera/Plücker tensors are also indexed in lockstep with T and would
+        # need matching shards. Disable SP until a halo-exchange-aware impl lands.
+        return latents, False
+
+    def gather_latents_for_sp(self, latents):
+        return latents
+
+    def get_decode_scale_and_shift(self, device, dtype, vae):
+        """Invert the LTX-2 latent normalization used before denoising.
+
+        SANA-WM uses the LTX-2 VAE. Upstream encodes as
+        ``(z - latents_mean) * scaling_factor / latents_std`` and decodes by
+        applying the inverse transform.
+        """
+        latents_mean = getattr(vae, "latents_mean", None)
+        latents_std = getattr(vae, "latents_std", None)
+
+        scaling_factor = (
+            getattr(getattr(vae, "config", None), "scaling_factor", None)
+            or getattr(vae, "scaling_factor", None)
+            or getattr(self.vae_config.arch_config, "scaling_factor", None)
+            or 1.0
+        )
+        if isinstance(scaling_factor, (int, float)) and float(scaling_factor) == 0.0:
+            scaling_factor = 1.0
+
+        if isinstance(latents_mean, torch.Tensor) and isinstance(
+            latents_std, torch.Tensor
+        ):
+            latents_mean = latents_mean.to(device=device, dtype=dtype).view(
+                1, -1, 1, 1, 1
+            )
+            latents_std = latents_std.to(device=device, dtype=dtype).view(
+                1, -1, 1, 1, 1
+            )
+            sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view(
+                1, 1, 1, 1, 1
+            )
+            return sf / latents_std, latents_mean
+
+        sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view(
+            1, 1, 1, 1, 1
+        )
+        return sf, None
+
+
+class SanaWMRealtimeConfig(SanaWMPipelineConfig):
+    """Realtime alias of the SANA-WM pipeline config.
+
+    Same numerics/fields as SanaWMPipelineConfig (our correct streaming pipeline);
+    exists so the realtime-serving registry can key adapters on a realtime config
+    class (matching the upstream/LingBot-World pattern) without renaming the base.
+    """
+
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            auto_dit_layerwise_offload=True,
+            auto_disable_component_offload_min_available_memory_gb=120,
+            auto_disable_component_offload_components=("dit",),
+            auto_enable_cfg_parallel=False,
+        )

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -702,7 +702,7 @@ class SamplingParams:
                 raise
 
         user_kwargs = dict(kwargs)
-        user_kwargs.pop("diffusers_kwargs", None)
+        diffusers_kwargs = user_kwargs.pop("diffusers_kwargs", None)
 
         user_sampling_params = type(sampling_params)(*args, **user_kwargs)
         # TODO: refactor
@@ -710,6 +710,8 @@ class SamplingParams:
             user_sampling_params, explicit_fields=set(user_kwargs.keys())
         )
         sampling_params._explicit_fields = set(user_kwargs.keys())
+        if diffusers_kwargs is not None:
+            sampling_params.diffusers_kwargs = diffusers_kwargs
         sampling_params._adjust(server_args)
 
         sampling_params._validate_with_pipeline_config(server_args.pipeline_config)
@@ -945,6 +947,36 @@ class SamplingParams:
                 '--image-path "img1.png" "img2.png"'
             ),
         )
+        add_argument(
+            "--action",
+            type=str,
+            help=(
+                "SANA-WM WASD/IJKL action DSL, e.g. "
+                "'w-80,jw-40,w-40,lw-60,w-100'. Model-specific fields are "
+                "ignored by other pipelines."
+            ),
+        )
+        add_argument(
+            "--translation-speed",
+            "--translation_speed",
+            type=float,
+            dest="translation_speed",
+            help="SANA-WM action DSL per-frame translation speed.",
+        )
+        add_argument(
+            "--rotation-speed-deg",
+            "--rotation_speed_deg",
+            type=float,
+            dest="rotation_speed_deg",
+            help="SANA-WM action DSL per-frame rotation speed in degrees.",
+        )
+        add_argument(
+            "--pitch-limit-deg",
+            "--pitch_limit_deg",
+            type=float,
+            dest="pitch_limit_deg",
+            help="SANA-WM action DSL absolute pitch clamp in degrees.",
+        )
         add_argument(
             "--moba-config-path",
             type=str,

python/sglang/multimodal_gen/configs/sample/sana_wm.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sampling parameters for SANA-WM TI2V world model generation."""
+
+from dataclasses import dataclass
+from typing import Any, Optional, Sequence, Union
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+
+# Type alias for the camera tensor inputs. Accept torch.Tensor, numpy arrays,
+# or nested Python lists — coerced to torch.Tensor downstream in the stage.
+CameraTensorLike = Union[Any, Sequence[Sequence[Sequence[float]]]]
+
+
+@dataclass
+class SanaWMSamplingParams(SamplingParams):
+    """Default sampling parameters for SANA-WM 720p (704×1280) 16fps video.
+
+    Frame counts must satisfy (num_frames - 1) % 8 == 0.
+
+    Optional camera conditioning:
+        camera_to_world: (T, 4, 4) extrinsics, one per output frame.
+        intrinsics:      (T, 3, 3) pinhole intrinsics, one per output frame.
+        action:          WASD/IJKL action DSL (e.g. "w-80,jw-40,w-40"), rolled
+                         out to camera_to_world before the camera branch.
+    Omitted camera_to_world -> static identity camera. Omitted intrinsics ->
+    centered heuristic; pass explicit intrinsics for closest NVlabs parity.
+    """
+
+    data_type: DataType = DataType.VIDEO
+
+    # Resolution: 720p landscape (LTX-2 VAE requires multiples of 32)
+    height: int = 704
+    width: int = 1280
+
+    # 49 = (49-1)/8 = 6 latent frames → ~3 seconds at 16fps
+    num_frames: int = 49
+
+    # SANA-WM is trained at 16fps (override base default of 24).
+    fps: int = 16
+
+    num_inference_steps: int = 20
+
+    guidance_scale: float = 4.5
+
+    # NVlabs' SANA-WM inference defaults to an empty negative prompt.
+    negative_prompt: str = ""
+
+    # --- Camera trajectory (6-DoF) — optional ---
+    camera_to_world: Optional[CameraTensorLike] = None
+    intrinsics: Optional[CameraTensorLike] = None
+    action: Optional[str] = None
+    translation_speed: float = (
+        0.04  # match official streaming (STREAMING_TRANSLATION_SPEED)
+    )
+    rotation_speed_deg: float = 1.2
+    pitch_limit_deg: float = 85.0
+
+    def build_request_extra(self) -> dict[str, Any]:
+        extra = super().build_request_extra()
+        if self.action is not None and self.camera_to_world is not None:
+            raise ValueError(
+                "SANA-WM accepts either action or camera_to_world, not both."
+            )
+        if self.camera_to_world is not None:
+            extra["camera_to_world"] = self.camera_to_world
+        if self.intrinsics is not None:
+            extra["intrinsics"] = self.intrinsics
+        if self.action is not None:
+            extra["action"] = self.action
+            extra["translation_speed"] = self.translation_speed
+            extra["rotation_speed_deg"] = self.rotation_speed_deg
+            extra["pitch_limit_deg"] = self.pitch_limit_deg
+        return extra

python/sglang/multimodal_gen/registry.py

@@ -77,6 +77,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImagePipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import (
     StableDiffusion3PipelineConfig,
 )
@@ -131,6 +132,7 @@ from sglang.multimodal_gen.configs.sample.qwenimage import (
     QwenImageSamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams
+from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams
 from sglang.multimodal_gen.configs.sample.stablediffusion3 import (
     StableDiffusion3SamplingParams,
 )
@@ -977,6 +979,20 @@ def _register_configs():
         ],
     )
 
+    # SANA-WM (register BEFORE generic SANA T2I to prevent "sana" detector false-match)
+    register_configs(
+        sampling_param_cls=SanaWMSamplingParams,
+        pipeline_config_cls=SanaWMPipelineConfig,
+        hf_model_paths=[
+            "Efficient-Large-Model/SANA-WM_bidirectional",
+            "Efficient-Large-Model/SANA-WM_streaming",
+        ],
+        model_detectors=[
+            # Match "sana-wm" or "sana_wm" but NOT plain T2I "sana" checkpoints.
+            lambda hf_id: ("sana-wm" in hf_id.lower() or "sana_wm" in hf_id.lower()),
+        ],
+    )
+
     # Cosmos3 — single checkpoint serves T2V, I2V, and T2I. Mode is dispatched
     # per-request inside the pipeline from ``num_frames`` and ``image_path``.
     # Both Nano (8B) and Super (32B) share the same pipeline; arch dimensions
@@ -1005,7 +1021,13 @@ def _register_configs():
             "Efficient-Large-Model/Sana_1600M_512px_diffusers",
             "Efficient-Large-Model/Sana_600M_512px_diffusers",
         ],
-        model_detectors=[lambda hf_id: "sana" in hf_id.lower()],
+        model_detectors=[
+            lambda hf_id: (
+                "sana" in hf_id.lower()
+                and "sana-wm" not in hf_id.lower()
+                and "sana_wm" not in hf_id.lower()
+            )
+        ],
     )
 
     # FireRed-Image-Edit

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

@@ -73,6 +73,22 @@ async def _wait_until_http_ready(server_args: ServerArgs) -> None:
     raise RuntimeError(f"HTTP server did not become ready at {health_url}")
 
 
+def _is_realtime_serving(server_args: ServerArgs) -> bool:
+    """A realtime pipeline establishes per-session state over the WebSocket, so
+    the synthetic server-warmup request (which has no session) cannot run — it
+    would fail in the realtime stage and abort startup. Detect it via the
+    realtime-adapter registry and skip server warmup."""
+    try:
+        from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import (
+            get_realtime_model_adapter,
+        )
+
+        get_realtime_model_adapter(server_args)
+        return True
+    except Exception:
+        return False
+
+
 async def _run_server_warmup_after_http_ready(
     server_args: ServerArgs, warmup_done: asyncio.Event
 ) -> None:
@@ -81,6 +97,7 @@ async def _run_server_warmup_after_http_ready(
             not server_args.warmup
             or not server_args.server_warmup
             or server_args.warmup_resolutions is not None
+            or _is_realtime_serving(server_args)
         ):
             warmup_done.set()
             return

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

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 import os
 import tempfile
-from collections import deque
 from typing import TYPE_CHECKING, Any
 
 from fastapi import WebSocket
@@ -28,12 +27,13 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
 from sglang.multimodal_gen.runtime.entrypoints.utils import (
     prepare_request,
 )
+from sglang.multimodal_gen.runtime.realtime.camera_controls import (
+    RealtimeCameraControlState,
+)
 from sglang.multimodal_gen.runtime.realtime.condition_events import (
     ConditionEvent,
     ConditionEventQueue,
     ControlSignal,
-    ControlStateSamplingQueue,
-    ControlStateTransition,
 )
 from sglang.multimodal_gen.runtime.server_args import get_global_server_args
 
@@ -53,22 +53,18 @@ LINGBOT_REALTIME_DEFAULT_NUM_INFERENCE_STEPS = 4
 LINGBOT_REALTIME_MIN_CONDITION_CHUNKS = 2
 
 
-class LingBotWorldRealtimeState:
+class LingBotWorldRealtimeState(RealtimeCameraControlState):
     def __init__(self):
-        self.events = ConditionEventQueue(max_events={"prompt": 1})
-        self.camera_state = ControlStateSamplingQueue(
-            default_item=[],
+        super().__init__(
             min_pulse_items=1,
+            script_maxlen=512,
             max_transitions=512,
         )
-        self.camera_script_queue: deque[ControlSignal] = deque(maxlen=512)
-        self.latest_sampled_event_id: int | None = None
+        self.events = ConditionEventQueue(max_events={"prompt": 1})
 
     def clear(self) -> None:
+        super().clear()
         self.events.clear()
-        self.camera_state.clear()
-        self.camera_script_queue.clear()
-        self.latest_sampled_event_id = None
 
     def receive_prompt(self, prompt: str, *, event_id: int | None = None) -> None:
         self.events.push(
@@ -82,125 +78,17 @@ class LingBotWorldRealtimeState:
             )
         )
 
-    def receive_camera_script(
-        self,
-        camera_actions: list[list[str]],
-        *,
-        event_id: int | None = None,
-    ) -> None:
-        self.camera_script_queue.clear()
-        self.camera_state.clear()
-        for actions in camera_actions:
-            self.camera_script_queue.append(
-                ControlSignal(
-                    kind="camera_actions",
-                    payload=list(actions),
-                    seq_id=event_id,
-                )
-            )
-
-    def receive_camera_state_transitions(
-        self,
-        transitions: list[ControlStateTransition],
-    ) -> None:
-        self.camera_script_queue.clear()
-        self.camera_state.push_many(transitions)
-
-    def receive_camera_actions(
-        self,
-        camera_actions: list[list[str]],
-        *,
-        event_id: int | None = None,
-    ) -> None:
-        self.receive_camera_script(camera_actions, event_id=event_id)
-
-    def receive_camera_state(
-        self,
-        actions: list[str],
-        *,
-        event_id: int | None = None,
-        timestamp_ms: int | None = None,
-    ) -> None:
-        self.receive_camera_state_transitions(
-            [
-                ControlStateTransition(
-                    payload=list(actions),
-                    seq_id=event_id,
-                    timestamp_ms=timestamp_ms,
-                )
-            ]
-        )
-
-    def _sample_camera_script(self, chunk_size: int) -> list[list[str]]:
-        chunk: list[list[str]] = []
-        latest_event_id = self.latest_sampled_event_id
-        while self.camera_script_queue and len(chunk) < chunk_size:
-            signal = self.camera_script_queue.popleft()
-            chunk.append(list(signal.payload))
-            latest_event_id = signal.seq_id
-        while len(chunk) < chunk_size:
-            chunk.append([])
-        self.latest_sampled_event_id = latest_event_id
-        return chunk
-
-    def _camera_state_transition(
-        self,
-        actions: list[str],
-        *,
-        event_id: int | None,
-        timestamp_ms: int | None,
-    ) -> ControlStateTransition:
-        return ControlStateTransition(
-            payload=list(actions),
-            seq_id=event_id,
-            timestamp_ms=timestamp_ms,
-        )
-
-    def _camera_transitions_from_event_payload(
-        self,
-        payload: dict[str, Any],
-        *,
-        event_id: int | None,
-    ) -> list[ControlStateTransition]:
-        transitions = payload.get("transitions")
-        if not isinstance(transitions, list):
-            raise ValueError("camera_actions state payload requires transitions")
-        result = []
-        for transition in transitions:
-            if not isinstance(transition, dict):
-                raise ValueError("camera_actions transition must be a map")
-            actions = transition.get("actions")
-            if not isinstance(actions, list):
-                raise ValueError("camera_actions transition actions must be a list")
-            timestamp_ms = transition.get("client_ts_ms")
-            if timestamp_ms is not None:
-                timestamp_ms = int(timestamp_ms)
-            result.append(
-                self._camera_state_transition(
-                    list(actions),
-                    event_id=event_id,
-                    timestamp_ms=timestamp_ms,
-                )
-            )
-        return result
-
     def receive_camera_event_payload(
         self,
         payload: Any,
         *,
         event_id: int | None,
     ) -> str:
-        if isinstance(payload, dict) and payload.get("mode") == "state":
-            transitions = self._camera_transitions_from_event_payload(
-                payload,
-                event_id=event_id,
-            )
-            self.receive_camera_state_transitions(transitions)
-            return f"kind=camera_actions, mode=state, transitions={len(transitions)}"
-
-        camera_actions = LingBotWorldRealtimeAdapter._validate_camera_actions(payload)
-        self.receive_camera_script(camera_actions, event_id=event_id)
-        return f"kind=camera_actions, mode=script, frames={len(camera_actions)}"
+        return super().receive_camera_event_payload(
+            payload,
+            event_id=event_id,
+            validate_camera_actions=LingBotWorldRealtimeAdapter._validate_camera_actions,
+        )
 
     def sample_prompt(self) -> str:
         prompt = self.events.pop_latest("prompt")
@@ -209,20 +97,6 @@ class LingBotWorldRealtimeState:
         self.latest_sampled_event_id = self.events.last_sampled_seq_id("prompt")
         return prompt
 
-    def sample_camera_actions(self, chunk_size: int) -> list[list[str]] | None:
-        """samples a sequence of camera actions for the chunk with chunk_size frames
-
-        Args:
-            chunk_size: number of frames
-        """
-        if self.camera_script_queue:
-            return self._sample_camera_script(chunk_size)
-        action_list = self.camera_state.sample_chunk(chunk_size)
-        if action_list is None:
-            return None
-        self.latest_sampled_event_id = self.camera_state.latest_sampled_seq_id()
-        return [list(actions) for actions in action_list]
-
     def has_prompt(self) -> bool:
         return self.events.has_events("prompt")
 

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

@@ -0,0 +1,359 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import hashlib
+import os
+import tempfile
+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_adapter import (
+    RealtimeChunkInputs,
+    RealtimeModelAdapter,
+)
+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.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.camera_controls import (
+    RealtimeCameraControlState,
+)
+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
+
+
+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_event_payload(
+        self,
+        payload: Any,
+        *,
+        event_id: int | None,
+    ) -> str:
+        return super().receive_camera_event_payload(
+            payload,
+            event_id=event_id,
+            validate_camera_actions=SanaWMRealtimeAdapter._validate_camera_actions,
+        )
+
+
+class SanaWMRealtimeAdapter(RealtimeModelAdapter):
+    name = "sana_wm_realtime"
+
+    def __init__(self):
+        self.output_adapter = RawRGBRealtimeOutputAdapter()
+
+    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:
+        if request.first_frame is None:
+            raise ValueError("SANA-WM realtime requires first_frame")
+
+        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_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_script(
+                parse_sana_wm_action_string(action), event_id=None
+            )
+        state.base_condition_inputs = condition_inputs
+
+        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 isinstance(
+            request.first_frame, str
+        ) and request.first_frame.lower().startswith(("http://", "https://")):
+            suffix = os.path.splitext(request.first_frame.split("?", 1)[0])[1]
+            digest = hashlib.sha256(request.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")
+        image_path = await save_image_to_path(request.first_frame, target_path)
+        request.first_frame = image_path
+
+    async def wait_for_next_chunk(self, session: GenerateSession) -> None:
+        del session
+
+    def ingest_event(
+        self,
+        session: GenerateSession,
+        event: RealtimeEvent,
+    ) -> str:
+        state = self._state(session)
+        if event.kind == "camera_actions":
+            return state.receive_camera_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_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,
+        chunk: RealtimeChunkContext,
+        action_chunk_size: int,
+    ) -> RealtimeChunkInputs:
+        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,
+        chunk: RealtimeChunkContext,
+        chunk_inputs: RealtimeChunkInputs,
+        chunk_size: int,
+    ):
+        request = session.request
+        if request is None:
+            raise ValueError("realtime request is not initialized")
+
+        return build_sampling_params(
+            chunk.request_id,
+            prompt=chunk_inputs.prompt,
+            size=request.size,
+            num_frames=request.num_frames,
+            fps=request.fps,
+            image_path=request.first_frame,
+            output_file_name=chunk.request_id,
+            save_output=False,
+            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,
+            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,
+        )
+
+    def prepare_next_request(
+        self,
+        session: GenerateSession,
+        server_args: ServerArgs,
+        chunk: RealtimeChunkContext,
+    ) -> Req:
+        arch_config = server_args.pipeline_config.dit_config.arch_config
+        chunk_size = int(getattr(arch_config, "num_frames_per_block", 3))
+        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
+        chunk_inputs = self._sample_chunk_inputs(session, chunk, action_chunk_size)
+        sampling_params = self._build_sampling_params(
+            session,
+            chunk,
+            chunk_inputs,
+            chunk_size,
+        )
+        batch = prepare_request(
+            server_args=server_args,
+            sampling_params=sampling_params,
+        )
+        batch.session = session.realtime_session
+        batch.realtime_session_id = session.id
+        batch.return_raw_frames = True
+        batch.block_idx = chunk.index
+        batch.realtime_event_id = self._state(session).latest_sampled_event_id
+        if session.request is not None:
+            # Forward the full transport config like the LingBot adapter does —
+            # the shared RawRGB output adapter / realtime_video_api consume
+            # preview width + pacing too; dropping them silently disabled both
+            # features for SANA-WM sessions.
+            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)
+        return batch
+
+    async def send_output(
+        self,
+        ws: WebSocket,
+        session: GenerateSession,
+        result: OutputBatch,
+        batch: Req,
+    ) -> RealtimeFrameSendStats:
+        return await self.output_adapter.send(ws, session, result, batch)
+
+    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 dispose(self, session: GenerateSession) -> None:
+        state = session.adapter_state
+        if isinstance(state, SanaWMRealtimeAdapterState):
+            state.clear()
+        self.output_adapter.reset()

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

@@ -124,6 +124,7 @@ def _frame_shape_from_metadata(
 
 
 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
@@ -139,12 +140,15 @@ class _TransportPayload:
     metadata: dict[str, int | str | bool | list[int]]
 
 
-def _split_frame_batch(frames: list[bytes]) -> list[list[bytes]]:
+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 + RAW_RGB_FRAMES_PER_WS_MESSAGE]
-        for i in range(0, len(frames), RAW_RGB_FRAMES_PER_WS_MESSAGE)
+        frames[i : i + frames_per_message]
+        for i in range(0, len(frames), frames_per_message)
     ]
 
 
@@ -495,7 +499,7 @@ class RawRGBRealtimeOutputAdapter:
         stats = empty_frame_send_stats(content_type)
         for frames in frame_batches:
             split_batches = (
-                [frames]
+                _split_frame_batch(frames, ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE)
                 if _is_encoded_preview_transport(
                     content_type=content_type,
                     output_format=output_format,

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

@@ -31,14 +31,24 @@ def _register_builtin_realtime_model_adapters() -> None:
     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
 
 

python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py

@@ -0,0 +1,906 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Callable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.configs.models.dits.sana_wm import SanaWMConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+
+# Re-exported for back-compat: callers import these names from this module path.
+from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import (  # noqa: F401
+    _CACHE_TYPE_CONCAT,
+    _CACHE_TYPE_STATE,
+    _INT32_SAFE_CONV_ELEMENTS,
+    _NUM_STREAM_CACHE_SLOTS,
+    _SLOT_CAM_K,
+    _SLOT_CAM_V,
+    _SLOT_FFN_TCONV,
+    _SLOT_K,
+    _SLOT_SHORTCONV,
+    _SLOT_TYPE_FLAG,
+    _SLOT_V,
+    BidirectionalGDNUCPESinglePathLiteLA,
+    CaptionEmbedder,
+    GLUMBConvTemp,
+    MultiHeadCrossAttention,
+    PatchEmbedMS3D,
+    T2IFinalLayer,
+    TimestepEmbedder,
+    WanRotaryPosEmbed,
+    _apply_block_diagonal,
+    _apply_complex_rope,
+    _apply_ray_projmat,
+    _apply_rotary_emb_bhnd,
+    _apply_rotary_emb_dn,
+    _bidirectional_short_conv,
+    _build_ucpe_apply_fns,
+    _compute_fov_from_focal,
+    _compute_frame_gates,
+    _ConvLayer,
+    _downscale_to_reference_rms,
+    _flip_and_shift,
+    _gdn_chunk_scan_forward,
+    _gdn_scan_bidirectional,
+    _gdn_scan_cached,
+    _gdn_scan_forward,
+    _gdn_scan_forward_stateful,
+    _invert_SE3,
+    _log_sana_wm_triton_cam_gdn_fallback,
+    _log_sana_wm_triton_gdn_fallback,
+    _RMSNorm,
+    _sana_wm_chunk_boundaries_for_attention,
+    _sana_wm_chunk_index_from_chunk_size,
+    _sana_wm_chunked_attention,
+    _sana_wm_normalize_chunk_index,
+    _sana_wm_padded_scale,
+    _sana_wm_sdpa,
+    _ShortConvolution,
+    _single_path_delta_chunk_scan_forward,
+    _single_path_delta_scan_bidirectional,
+    _single_path_delta_scan_cached,
+    _single_path_delta_scan_forward,
+    _single_path_delta_scan_forward_stateful,
+    _sinusoidal_timestep_embedding,
+    _slice_rope_for_cam,
+    _slice_rope_to_current_chunk,
+    _temporal_short_conv_cached,
+    _tensor_cache_key,
+    _UpstreamMlp,
+    compute_chunk_plucker,
+    process_camera_conditions_ucpe,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import (
+    parity_probe,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SanaWMBlock(nn.Module):
+    """One transformer block of SANA-WM."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        head_dim: int,
+        mlp_ratio: float,
+        t_kernel_size: int,
+        qk_norm: bool,
+        cross_norm: bool,
+        conv_kernel_size: int,
+        k_conv_only: bool,
+        softmax_main: bool,
+        use_chunk_plucker_post_attn: bool,
+        chunk_size: Optional[int] = None,
+        chunk_split_strategy: str = "uniform",
+        update_rule: str = "torch_chunk",
+        cam_update_rule: str = "torch_chunk",
+        chunk_gdn_chunk_size: int = 21,
+        use_chunked_softmax_attention: bool = False,
+        gdn_backend: str = "auto",
+    ) -> None:
+        super().__init__()
+        self.softmax_main = softmax_main
+        self.chunk_size = chunk_size
+        self.chunk_split_strategy = chunk_split_strategy
+
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.attn = BidirectionalGDNUCPESinglePathLiteLA(
+            in_dim=hidden_size,
+            heads=num_heads,
+            head_dim=head_dim,
+            qk_norm=qk_norm,
+            conv_kernel_size=conv_kernel_size,
+            k_conv_only=k_conv_only,
+            softmax_main=softmax_main,
+            update_rule=update_rule,
+            cam_update_rule=cam_update_rule,
+            chunk_gdn_chunk_size=chunk_gdn_chunk_size,
+            use_chunked_softmax_attention=use_chunked_softmax_attention,
+            gdn_backend=gdn_backend,
+        )
+
+        self.cross_attn = MultiHeadCrossAttention(
+            d_model=hidden_size,
+            num_heads=num_heads,
+            qk_norm=cross_norm,
+        )
+
+        self.mlp = GLUMBConvTemp(
+            in_features=hidden_size,
+            hidden_features=int(hidden_size * mlp_ratio),
+            t_kernel_size=t_kernel_size,
+        )
+
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(6, hidden_size) / hidden_size**0.5
+        )
+
+        if use_chunk_plucker_post_attn:
+            self.plucker_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+            nn.init.zeros_(self.plucker_proj.weight)
+            nn.init.zeros_(self.plucker_proj.bias)
+        else:
+            self.plucker_proj = None
+
+    @staticmethod
+    def _modulate(
+        x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
+    ) -> torch.Tensor:
+        return x * (1 + scale) + shift
+
+    @staticmethod
+    def _reshape_framewise_modulation(
+        x: torch.Tensor,
+        num_frames: int,
+    ) -> tuple[torch.Tensor, int]:
+        B, N, C = x.shape
+        tokens_per_frame = N // num_frames
+        return x.reshape(B, num_frames, tokens_per_frame, C), tokens_per_frame
+
+    def forward(
+        self,
+        x: torch.Tensor,  # (B, N, D)
+        y: torch.Tensor,  # (B, L, D) text embeds
+        t: torch.Tensor,  # (B, 6*D) AdaLN-single
+        HW: Tuple[int, int, int],
+        rotary_emb: Optional[torch.Tensor],
+        prope_fns: Optional[Tuple[Callable, Callable, Callable]],
+        plucker_emb: Optional[torch.Tensor],
+        mask: Optional[torch.Tensor],
+        chunk_size: Optional[int] = None,
+        chunk_split_strategy: Optional[str] = None,
+        chunk_index: Optional[List[int]] = None,
+    ) -> torch.Tensor:
+        B = x.shape[0]
+        if t.dim() == 2:
+            num_frames = None
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + t.reshape(B, 6, -1)
+            ).chunk(6, dim=1)
+        else:
+            num_frames = t.reshape(B, -1, 6, t.shape[-1] // 6).shape[1]
+            t = t.reshape(B, num_frames, 6, -1)
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None, None, :, :] + t
+            ).chunk(6, dim=2)
+
+        # Self-attention with UCPE camera branch
+        if num_frames is None:
+            x_in = self._modulate(self.norm1(x), shift_msa, scale_msa)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm1(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_msa, scale_msa).reshape_as(x)
+        attn_out = self.attn(
+            x_in,
+            HW=HW,
+            rotary_emb=rotary_emb,
+            prope_fns=prope_fns,
+            chunk_size=self.chunk_size if chunk_size is None else chunk_size,
+            chunk_split_strategy=(
+                self.chunk_split_strategy
+                if chunk_split_strategy is None
+                else chunk_split_strategy
+            ),
+            chunk_index=chunk_index,
+        )
+        if num_frames is None:
+            x = x + gate_msa * attn_out
+        else:
+            attn_out = attn_out.reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_msa * attn_out).reshape_as(x)
+
+        # Plücker post-attn injection (zero-init linear)
+        if self.plucker_proj is not None and plucker_emb is not None:
+            x = x + self.plucker_proj(plucker_emb)
+
+        # Cross-attention
+        x = x + self.cross_attn(x, y, mask=mask)
+
+        # FFN
+        if num_frames is None:
+            x_in = self._modulate(self.norm2(x), shift_mlp, scale_mlp)
+            x = x + gate_mlp * self.mlp(x_in, HW=HW)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm2(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_mlp, scale_mlp).reshape_as(x)
+            mlp_out = self.mlp(x_in, HW=HW).reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_mlp * mlp_out).reshape_as(x)
+        return x
+
+    def forward_long(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        t: torch.Tensor,
+        HW: Tuple[int, int, int],
+        rotary_emb: Optional[torch.Tensor],
+        prope_fns: Optional[Tuple[Callable, Callable, Callable]],
+        plucker_emb: Optional[torch.Tensor],
+        mask: Optional[torch.Tensor],
+        *,
+        kv_cache: list,
+        save_kv_cache: bool,
+    ) -> Tuple[torch.Tensor, list]:
+        """Streaming counterpart of ``forward``: threads the per-block 10-slot ``kv_cache`` through cached attention + FFN."""
+        B = x.shape[0]
+        if t.dim() == 2:
+            num_frames = None
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + t.reshape(B, 6, -1)
+            ).chunk(6, dim=1)
+        else:
+            num_frames = t.reshape(B, -1, 6, t.shape[-1] // 6).shape[1]
+            t = t.reshape(B, num_frames, 6, -1)
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None, None, :, :] + t
+            ).chunk(6, dim=2)
+
+        if num_frames is None:
+            x_in = self._modulate(self.norm1(x), shift_msa, scale_msa)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm1(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_msa, scale_msa).reshape_as(x)
+
+        attn_out, kv_cache = self.attn.forward_long(
+            x_in,
+            HW=HW,
+            rotary_emb=rotary_emb,
+            prope_fns=prope_fns,
+            kv_cache=kv_cache,
+            save_kv_cache=save_kv_cache,
+        )
+        if num_frames is None:
+            x = x + gate_msa * attn_out
+        else:
+            attn_out = attn_out.reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_msa * attn_out).reshape_as(x)
+
+        if self.plucker_proj is not None and plucker_emb is not None:
+            x = x + self.plucker_proj(plucker_emb)
+
+        x = x + self.cross_attn(x, y, mask=mask)
+
+        if num_frames is None:
+            x_in = self._modulate(self.norm2(x), shift_mlp, scale_mlp)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm2(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_mlp, scale_mlp).reshape_as(x)
+
+        # GLUMBConvTemp returns a tuple whenever the streaming path is active
+        # (ffn_tail set OR save requested); branch on tuple-ness, never on
+        # save_kv_cache alone (a read-only pass with a populated slot 9 still
+        # returns a tuple).
+        mlp_out = self.mlp(
+            x_in,
+            HW=HW,
+            ffn_tail=kv_cache[_SLOT_FFN_TCONV],
+            save_ffn_tail=save_kv_cache,
+        )
+        if isinstance(mlp_out, tuple):
+            mlp_out, ffn_tail = mlp_out
+            if save_kv_cache:
+                kv_cache[_SLOT_FFN_TCONV] = ffn_tail
+        if num_frames is None:
+            x = x + gate_mlp * mlp_out
+        else:
+            mlp_out = mlp_out.reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_mlp * mlp_out).reshape_as(x)
+        return x, kv_cache
+
+
+class SanaWMTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """SANA-WM 2.6B TI2V world model.
+
+    Forward inputs:
+        hidden_states:           (B, C, T, H, W)        128-ch LTX-2 latent
+        encoder_hidden_states:   (B, L, 2304)           Gemma-2 embeddings
+        timestep:                (B,)
+        encoder_attention_mask:  (B, L) optional bool
+        camera_conditions:       (B, T, 20)             latent-frame raymap:
+                                                        16 c2w + (fx,fy,cx,cy)
+        chunk_plucker:           (B, 48, T, H, W)       optional, computed
+                                                        from camera_conditions
+                                                        if absent.
+
+    Returns: ``(B, C, T, H, W)`` predicted velocity / noise.
+    """
+
+    _fsdp_shard_conditions = SanaWMConfig()._fsdp_shard_conditions
+    _compile_conditions = SanaWMConfig()._compile_conditions
+    _supported_attention_backends = SanaWMConfig()._supported_attention_backends
+    param_names_mapping = SanaWMConfig().param_names_mapping
+    reverse_param_names_mapping = SanaWMConfig().reverse_param_names_mapping
+    lora_param_names_mapping: dict = {}
+
+    def __init__(self, config: SanaWMConfig, hf_config=None, **kwargs) -> None:
+        super().__init__(config, hf_config=hf_config or {}, **kwargs)
+        arch = config.arch_config
+
+        self.patch_size = (arch.patch_size_t, arch.patch_size, arch.patch_size)
+        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
+        self.hidden_size = self.inner_dim
+        self.num_attention_heads = arch.num_attention_heads
+        self.attention_head_dim = arch.attention_head_dim
+        self.out_channels = arch.out_channels
+        self.num_channels_latents = arch.num_channels_latents
+        self.vae_temporal_stride = arch.vae_temporal_stride
+        self.timestep_norm_scale_factor = getattr(
+            arch, "timestep_norm_scale_factor", 1.0
+        )
+
+        # --- Embedders ---
+        self.x_embedder = PatchEmbedMS3D(
+            self.patch_size,
+            arch.in_channels,
+            self.inner_dim,
+            bias=True,
+        )
+
+        self.t_embedder = TimestepEmbedder(self.inner_dim, frequency_embedding_size=256)
+        self.t_block = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True),
+        )
+
+        self.y_embedder = CaptionEmbedder(
+            in_channels=arch.caption_channels,
+            hidden_size=self.inner_dim,
+            token_num=arch.model_max_length,
+        )
+        self.y_norm = bool(getattr(arch, "y_norm", True))
+        self.attention_y_norm = _RMSNorm(
+            self.inner_dim,
+            scale_factor=getattr(arch, "y_norm_scale_factor", 1.0),
+            eps=getattr(arch, "y_norm_eps", 1e-5),
+        )
+
+        # 3-channel raymap embedder -- kept for state_dict compatibility but
+        # only invoked when ``use_chunk_plucker_post_attn`` is False.
+        # When ``True`` (the case for the released checkpoint) the absmap
+        # path is skipped entirely.
+        self.raymap_embedder = PatchEmbedMS3D(
+            self.patch_size,
+            3,
+            self.inner_dim,
+            bias=True,
+        )
+        # 48-channel plucker embedder (chunk-packed)
+        if arch.use_chunk_plucker_post_attn or arch.use_chunk_plucker_input:
+            self.plucker_embedder = PatchEmbedMS3D(
+                self.patch_size,
+                arch.chunk_plucker_channels,
+                self.inner_dim,
+                bias=True,
+            )
+            nn.init.zeros_(self.plucker_embedder.proj.weight)
+            nn.init.zeros_(self.plucker_embedder.proj.bias)
+        else:
+            self.plucker_embedder = None
+        self.use_chunk_plucker_post_attn = arch.use_chunk_plucker_post_attn
+        self.use_chunk_plucker_input = arch.use_chunk_plucker_input
+        self.chunk_size = getattr(arch, "chunk_size", None)
+        self.chunk_split_strategy = getattr(arch, "chunk_split_strategy", "uniform")
+
+        # --- RoPE ---
+        self.rope = WanRotaryPosEmbed(
+            attention_head_dim=arch.linear_head_dim,
+            patch_size=self.patch_size,
+            max_seq_len=1024,
+        )
+
+        # --- Transformer blocks ---
+        depth = arch.num_layers
+        self.softmax_every_n = arch.softmax_every_n
+        softmax_idx = set(
+            i
+            for i in range(depth)
+            if arch.softmax_every_n > 0 and (i + 1) % arch.softmax_every_n == 0
+        )
+        self.softmax_block_indices = tuple(sorted(softmax_idx))
+
+        self.blocks = nn.ModuleList(
+            [
+                SanaWMBlock(
+                    hidden_size=self.inner_dim,
+                    num_heads=arch.num_attention_heads,
+                    head_dim=arch.linear_head_dim,
+                    mlp_ratio=arch.mlp_ratio,
+                    t_kernel_size=arch.t_kernel_size,
+                    qk_norm=arch.qk_norm,
+                    cross_norm=arch.cross_norm,
+                    conv_kernel_size=arch.conv_kernel_size,
+                    k_conv_only=arch.k_conv_only,
+                    softmax_main=(i in softmax_idx),
+                    use_chunk_plucker_post_attn=(
+                        arch.use_chunk_plucker_post_attn
+                        and (
+                            arch.chunk_plucker_post_attn_blocks < 0
+                            or i < arch.chunk_plucker_post_attn_blocks
+                        )
+                    ),
+                    chunk_size=self.chunk_size,
+                    chunk_split_strategy=self.chunk_split_strategy,
+                    update_rule=getattr(arch, "update_rule", "torch_chunk"),
+                    cam_update_rule=getattr(arch, "cam_update_rule", "torch_chunk"),
+                    chunk_gdn_chunk_size=getattr(arch, "chunk_gdn_chunk_size", 21),
+                    use_chunked_softmax_attention=getattr(
+                        arch, "use_chunked_softmax_attention", False
+                    ),
+                    gdn_backend=getattr(arch, "gdn_backend", "auto"),
+                )
+                for i in range(depth)
+            ]
+        )
+
+        self.final_layer = T2IFinalLayer(
+            self.inner_dim, self.patch_size, self.out_channels
+        )
+
+        # Cache RoPE freqs per shape -- avoids recomputation across denoising
+        # steps with constant latent shapes.
+        self._freqs_cache: dict = {}
+        self._ucpe_apply_fns_cache: Optional[
+            Tuple[Tuple, torch.Tensor, Tuple[Callable, Callable, Callable]]
+        ] = None
+        self._plucker_emb_cache: Optional[Tuple[Tuple, torch.Tensor, torch.Tensor]] = (
+            None
+        )
+
+        # FSDP shard targets
+        self.layer_names = ["blocks"]
+
+    def post_load_weights(self) -> None:
+        # FSDP loader initializes the model on meta and only materializes
+        # tensors that appear in the checkpoint. WanRotaryPosEmbed._freqs is a
+        # derived, non-persistent constant, so recompute it deterministically.
+        for module in self.modules():
+            if isinstance(module, WanRotaryPosEmbed):
+                if module._freqs.is_meta:
+                    module._init_freqs_buffer()
+
+    # ------------------------------------------------------------------ #
+    # Forward
+    # ------------------------------------------------------------------ #
+
+    def _get_freqs(self, T: int, H: int, W: int, device: torch.device) -> torch.Tensor:
+        key = (T, H, W, str(device))
+        if key not in self._freqs_cache:
+            self._freqs_cache[key] = self.rope((T, H, W), device)
+        return self._freqs_cache[key]
+
+    def _get_freqs_window(
+        self,
+        start: int,
+        end: int,
+        H: int,
+        W: int,
+        device: torch.device,
+        frame_index: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """RoPE freqs for a streaming chunk at GLOBAL frame positions.
+
+        ``frame_index`` (per-token global positions) overrides ``(start, end)``;
+        the count branch is cached, the tensor branch is computed fresh.
+        """
+        if frame_index is not None:
+            return self.rope((end - start, H, W), device, frame_index=frame_index)
+        key = ("win", int(start), int(end), H, W, str(device))
+        if key not in self._freqs_cache:
+            self._freqs_cache[key] = self.rope(((int(start), int(end)), H, W), device)
+        return self._freqs_cache[key]
+
+    def _get_ucpe_apply_fns(
+        self,
+        camera_conditions: torch.Tensor,
+        *,
+        HW: Tuple[int, int, int],
+        freqs: torch.Tensor,
+    ) -> Tuple[Callable, Callable, Callable]:
+        head_dim = self.attention_head_dim
+        if torch.is_grad_enabled():
+            raymats = process_camera_conditions_ucpe(
+                camera_conditions,
+                HW=HW,
+                patch_size=self.patch_size,
+            )
+            raymats_flat = raymats.reshape(camera_conditions.shape[0], -1, 4, 4)
+            return _build_ucpe_apply_fns(head_dim, raymats_flat, freqs)
+
+        key = (
+            "ucpe",
+            HW,
+            self.patch_size,
+            head_dim,
+            _tensor_cache_key(camera_conditions),
+            _tensor_cache_key(freqs),
+        )
+        cached = self._ucpe_apply_fns_cache
+        if cached is not None and cached[0] == key:
+            return cached[2]
+
+        raymats = process_camera_conditions_ucpe(
+            camera_conditions,
+            HW=HW,
+            patch_size=self.patch_size,
+        )
+        raymats_flat = raymats.reshape(camera_conditions.shape[0], -1, 4, 4)
+        prope_fns = _build_ucpe_apply_fns(head_dim, raymats_flat, freqs)
+        self._ucpe_apply_fns_cache = (key, camera_conditions, prope_fns)
+        return prope_fns
+
+    def _get_plucker_emb(
+        self,
+        chunk_plucker: torch.Tensor,
+        *,
+        latent_token_count: int,
+    ) -> torch.Tensor:
+        if self.plucker_embedder is None:
+            raise ValueError("SANA-WM plucker_embedder is not initialized.")
+
+        weight = self.plucker_embedder.proj.weight
+        bias = self.plucker_embedder.proj.bias
+        key = (
+            "plucker_emb",
+            latent_token_count,
+            self.patch_size,
+            _tensor_cache_key(chunk_plucker),
+            _tensor_cache_key(weight),
+            None if bias is None else _tensor_cache_key(bias),
+        )
+        if not torch.is_grad_enabled():
+            cached = self._plucker_emb_cache
+            if cached is not None and cached[0] == key:
+                return cached[2]
+
+        plucker_emb = self.plucker_embedder(chunk_plucker.to(weight.dtype))
+        if plucker_emb.shape[1] != latent_token_count:
+            raise ValueError(
+                f"plucker_emb token count {plucker_emb.shape[1]} != "
+                f"latent token count {latent_token_count}; "
+                "expected chunk_plucker shape (B, 48, T, H, W)."
+            )
+
+        if not torch.is_grad_enabled():
+            self._plucker_emb_cache = (key, chunk_plucker, plucker_emb)
+        return plucker_emb
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        camera_conditions: Optional[torch.Tensor] = None,
+        chunk_plucker: Optional[torch.Tensor] = None,
+        guidance: Optional[torch.Tensor] = None,  # kept for compat
+        **kwargs,
+    ) -> torch.Tensor:
+        if encoder_hidden_states is None:
+            raise ValueError("SANA-WM forward requires encoder_hidden_states.")
+        if timestep is None:
+            raise ValueError("SANA-WM forward requires timestep.")
+
+        B, C, T_raw, H_raw, W_raw = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        T = T_raw // p_t
+        H = H_raw // p_h
+        W = W_raw // p_w
+        chunk_size = kwargs.get("chunk_size", self.chunk_size)
+        chunk_split_strategy = kwargs.get(
+            "chunk_split_strategy", self.chunk_split_strategy
+        )
+        chunk_index = kwargs.get("chunk_index", None)
+
+        # Patch embed: (B, C, T, H, W) -> (B, T*H*W, D)
+        x = self.x_embedder(hidden_states.to(dtype=self.x_embedder.proj.weight.dtype))
+
+        # Timestep AdaLN-single. SANA-WM's LTX sampler passes per-frame
+        # timesteps shaped (B, 1, T) so the clean first-frame condition can stay
+        # at timestep 0 while remaining latent frames denoise. Keep the scalar
+        # path for generic scheduler compatibility.
+        if self.timestep_norm_scale_factor != 1.0:
+            timestep_for_embed = (
+                timestep.float() / self.timestep_norm_scale_factor
+            ).to(torch.float32)
+        else:
+            timestep_for_embed = timestep.long().to(torch.float32)
+
+        if timestep_for_embed.dim() == 1:
+            t_emb = self.t_embedder(timestep_for_embed)  # (B, D)
+            t6 = self.t_block(t_emb)  # (B, 6D)
+        else:
+            timestep_shape = tuple(timestep_for_embed.shape)
+            t_flat = self.t_embedder(timestep_for_embed.flatten())
+            t6_flat = self.t_block(t_flat)
+            t_emb = t_flat.unflatten(0, timestep_shape)
+            t6 = t6_flat.unflatten(0, timestep_shape)
+
+        if isinstance(encoder_attention_mask, (list, tuple)):
+            encoder_attention_mask = encoder_attention_mask[0]
+        y = encoder_hidden_states
+        if y.dim() == 3:
+            y = y.unsqueeze(1)
+        y = self.y_embedder(y).squeeze(1)  # (B, L, D)
+        if y.shape[0] != B:
+            y = y.expand(B, -1, -1).contiguous()
+        if self.y_norm:
+            y = self.attention_y_norm(y)
+        if encoder_attention_mask is not None and encoder_attention_mask.shape[0] != B:
+            encoder_attention_mask = encoder_attention_mask.expand(B, -1).contiguous()
+
+        freqs = self._get_freqs(T, H, W, x.device)
+
+        # Camera conditioning: UCPE prope_fns + Plücker
+        prope_fns = None
+        if camera_conditions is not None:
+            if camera_conditions.shape[1] != T:
+                raise ValueError(
+                    "SANA-WM camera_conditions must be sampled at latent "
+                    f"frames: got {camera_conditions.shape[1]} frames, "
+                    f"expected T={T}."
+                )
+            prope_fns = self._get_ucpe_apply_fns(
+                camera_conditions,
+                HW=(T, H, W),
+                freqs=freqs,
+            )
+
+        # Plücker post-attn embedding (shared across all blocks)
+        plucker_emb = None
+        needs_plucker_emb = (
+            chunk_plucker is not None
+            and self.plucker_embedder is not None
+            and (self.use_chunk_plucker_post_attn or self.use_chunk_plucker_input)
+        )
+        if needs_plucker_emb:
+            plucker_emb = self._get_plucker_emb(
+                chunk_plucker,
+                latent_token_count=x.shape[1],
+            )  # (B, T*H*W, D)
+
+        if self.use_chunk_plucker_input and plucker_emb is not None:
+            x = x + plucker_emb
+
+        if not self.use_chunk_plucker_post_attn:
+            plucker_emb = None
+
+        # --- 6. Transformer blocks ---
+        HW = (T, H, W)
+        for block in self.blocks:
+            x = block(
+                x,
+                y=y,
+                t=t6,
+                HW=HW,
+                rotary_emb=freqs,
+                prope_fns=prope_fns,
+                plucker_emb=plucker_emb,
+                mask=encoder_attention_mask,
+                chunk_size=chunk_size,
+                chunk_split_strategy=chunk_split_strategy,
+                chunk_index=chunk_index,
+            )
+
+        x = self.final_layer(x, t_emb)  # (B, N, p_t*p_h*p_w*C_out)
+
+        # Un-patch
+        x = x.reshape(B, T, H, W, p_t, p_h, p_w, self.out_channels)
+        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).contiguous()
+        x = x.reshape(B, self.out_channels, T * p_t, H * p_h, W * p_w)
+        return x
+
+    def forward_long(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        camera_conditions: Optional[torch.Tensor] = None,
+        chunk_plucker: Optional[torch.Tensor] = None,
+        *,
+        kv_cache: Optional[list] = None,
+        save_kv_cache: bool = True,
+        start_f: Optional[int] = None,
+        end_f: Optional[int] = None,
+        frame_index: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, list]:
+        """Streaming autoregressive forward over a chunk of latent frames.
+
+        RoPE / camera / plücker are windowed to the chunk's GLOBAL frame range
+        ``[start_f, end_f)``; a per-block 10-slot ``kv_cache`` carries recurrent
+        state / concat-windows across chunks. Returns ``(out, new_cache)``.
+        """
+        if encoder_hidden_states is None:
+            raise ValueError("SANA-WM forward_long requires encoder_hidden_states.")
+        if timestep is None:
+            raise ValueError("SANA-WM forward_long requires timestep.")
+
+        if kv_cache is None:
+            kv_cache = [[None] * _NUM_STREAM_CACHE_SLOTS for _ in self.blocks]
+
+        B, C, T_raw, H_raw, W_raw = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        T = T_raw // p_t
+        H = H_raw // p_h
+        W = W_raw // p_w
+        start = 0 if start_f is None else int(start_f)
+        end = start + T if end_f is None else int(end_f)
+
+        x = self.x_embedder(hidden_states.to(dtype=self.x_embedder.proj.weight.dtype))
+
+        # Timestep AdaLN-single: force the framewise (B, 1, T) path so blocks
+        # always apply per-frame modulation.
+        if timestep.dim() == 1:
+            timestep = timestep[:, None, None].expand(-1, 1, T)
+        elif timestep.dim() == 2:
+            timestep = timestep[:, None, :]
+        if self.timestep_norm_scale_factor != 1.0:
+            timestep_for_embed = (
+                timestep.float() / self.timestep_norm_scale_factor
+            ).to(torch.float32)
+        else:
+            timestep_for_embed = timestep.long().to(torch.float32)
+        timestep_shape = tuple(timestep_for_embed.shape)
+        t_flat = self.t_embedder(timestep_for_embed.flatten())
+        t6_flat = self.t_block(t_flat)
+        t_emb = t_flat.unflatten(0, timestep_shape)
+        t6 = t6_flat.unflatten(0, timestep_shape)
+
+        if isinstance(encoder_attention_mask, (list, tuple)):
+            encoder_attention_mask = encoder_attention_mask[0]
+        y = encoder_hidden_states
+        if y.dim() == 3:
+            y = y.unsqueeze(1)
+        y = self.y_embedder(y).squeeze(1)
+        if y.shape[0] != B:
+            y = y.expand(B, -1, -1).contiguous()
+        if self.y_norm:
+            y = self.attention_y_norm(y)
+        if encoder_attention_mask is not None and encoder_attention_mask.shape[0] != B:
+            encoder_attention_mask = encoder_attention_mask.expand(B, -1).contiguous()
+
+        # RoPE windowed to global frame positions [start, end)
+        freqs = self._get_freqs_window(
+            start, end, H, W, x.device, frame_index=frame_index
+        )
+
+        # Camera conditioning: slice to the chunk, co-windowed w/ freqs
+        prope_fns = None
+        if camera_conditions is not None:
+            if camera_conditions.shape[1] != T:
+                # .contiguous(): canonical layout regardless of how the caller
+                # built the full-length tensor, so batch and realtime windows are
+                # kernel-level identical (slice offset/stride changes the reduction
+                # order otherwise — measured 1e-7 seeds amplifying to %-level drift
+                # through the bf16 block stack).
+                camera_conditions = camera_conditions[:, start:end].contiguous()
+            if camera_conditions.shape[0] != B:
+                camera_conditions = camera_conditions.repeat(
+                    B // camera_conditions.shape[0], 1, 1
+                )
+            prope_fns = self._get_ucpe_apply_fns(
+                camera_conditions, HW=(T, H, W), freqs=freqs
+            )
+
+        # Plücker post-attn / input embedding, sliced to the chunk.
+        if chunk_plucker is not None and chunk_plucker.shape[2] != T:
+            chunk_plucker = chunk_plucker[
+                :, :, start:end
+            ].contiguous()  # see camera note
+        if chunk_plucker is not None and chunk_plucker.shape[0] != B:
+            chunk_plucker = chunk_plucker.repeat(
+                B // chunk_plucker.shape[0], 1, 1, 1, 1
+            )
+        plucker_emb = None
+        needs_plucker_emb = (
+            chunk_plucker is not None
+            and self.plucker_embedder is not None
+            and (self.use_chunk_plucker_post_attn or self.use_chunk_plucker_input)
+        )
+        if needs_plucker_emb:
+            plucker_emb = self._get_plucker_emb(
+                chunk_plucker, latent_token_count=x.shape[1]
+            )
+        if self.use_chunk_plucker_input and plucker_emb is not None:
+            x = x + plucker_emb
+        if not self.use_chunk_plucker_post_attn:
+            plucker_emb = None
+
+        # parity harness (env-gated, no-op in prod): on the FIRST sink-path call
+        # (frame_index not None), checksum the pre-block tensors and x after every
+        # block to localize where the two execution paths first diverge.
+        _probe_path = os.environ.get(parity_probe.ENV_BLOCK_PROBE)
+        _probe = None
+        if (
+            _probe_path
+            and frame_index is not None
+            and not getattr(self, "_block_probe_done", False)
+        ):
+            _ck = parity_probe.checksum
+            _probe = {
+                "x_embed": _ck(x),
+                "t6": _ck(t6),
+                "y": _ck(y),
+                "freqs": (
+                    (
+                        tuple(freqs.shape),
+                        float(freqs.real.detach().double().sum().item()),
+                        float(freqs.imag.detach().double().sum().item()),
+                    )
+                    if freqs is not None
+                    else None
+                ),
+                "plucker_emb": _ck(plucker_emb),
+                "frame_index": frame_index.tolist(),
+            }
+
+        HW = (T, H, W)
+        new_cache = []
+        for i, block in enumerate(self.blocks):
+            x, block_cache = block.forward_long(
+                x,
+                y,
+                t6,
+                HW,
+                freqs,
+                prope_fns,
+                plucker_emb,
+                encoder_attention_mask,
+                kv_cache=kv_cache[i],
+                save_kv_cache=save_kv_cache,
+            )
+            new_cache.append(block_cache)
+            if _probe is not None:
+                _probe[f"x_after_block_{i:02d}"] = parity_probe.checksum(x)
+        if _probe is not None:
+            torch.save(_probe, _probe_path)
+            self._block_probe_done = True
+
+        x = self.final_layer(x, t_emb)
+        x = x.reshape(B, T, H, W, p_t, p_h, p_w, self.out_channels)
+        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).contiguous()
+        x = x.reshape(B, self.out_channels, T * p_t, H * p_h, W * p_w)
+        return x, new_cache
+
+
+EntryClass = SanaWMTransformer3DModel

python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py

@@ -0,0 +1,426 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.configs.models.dits.sana_wm_refiner import (
+    SanaWMRefinerArchConfig,
+    SanaWMRefinerConfig,
+)
+from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.models.dits.ltx_2 import (
+    LTX2AdaLayerNormSingle,
+    LTX2Attention,
+    LTX2AudioVideoRotaryPosEmbed,
+    LTX2FeedForward,
+    LTX2TextProjection,
+)
+
+
+def pack_latents(
+    latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1
+) -> torch.Tensor:
+    """Pack a 5D latent (B, C, T, H, W) into a 3D token sequence (B, L, in_dim)."""
+    B, _, T, H, W = latents.shape
+    pT = T // patch_size_t
+    pH = H // patch_size
+    pW = W // patch_size
+    latents = latents.reshape(B, -1, pT, patch_size_t, pH, patch_size, pW, patch_size)
+    return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+
+
+def unpack_latents(
+    tokens: torch.Tensor,
+    num_frames: int,
+    height: int,
+    width: int,
+    patch_size: int = 1,
+    patch_size_t: int = 1,
+) -> torch.Tensor:
+    """Inverse of `pack_latents`: (B, L, out_dim) -> (B, C, T, H, W)."""
+    B = tokens.size(0)
+    tokens = tokens.reshape(
+        B,
+        num_frames // patch_size_t,
+        height // patch_size,
+        width // patch_size,
+        -1,
+        patch_size_t,
+        patch_size,
+        patch_size,
+    )
+    return (
+        tokens.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
+    )
+
+
+def _slice_rope(
+    rope: tuple[torch.Tensor, torch.Tensor], start: int, end: Optional[int] = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Slice along token axis for either interleaved (rank-3) or split (rank-4)."""
+    cos, sin = rope
+    end_ = end if end is not None else cos.shape[-2 if cos.ndim == 4 else 1]
+    if cos.ndim == 3:
+        return cos[:, start:end_], sin[:, start:end_]
+    if cos.ndim == 4:
+        return cos[:, :, start:end_, :], sin[:, :, start:end_, :]
+    raise ValueError(f"Unexpected RoPE rank: {cos.ndim}")
+
+
+def _streaming_self_attention(
+    attn: LTX2Attention,
+    hidden_states: torch.Tensor,
+    video_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    n_context_tokens: int,
+) -> torch.Tensor:
+    """Streaming SLA: context attends to context only, current attends to context+current.
+
+    Mirrors NVlabs `inference_sana_wm.py::_streaming_self_attention`.
+    """
+    seq_len = hidden_states.shape[1]
+    if n_context_tokens <= 0 or n_context_tokens >= seq_len:
+        return attn(hidden_states, context=None, pe=video_rotary_emb)
+
+    ctx_rope = _slice_rope(video_rotary_emb, 0, n_context_tokens)
+    out_ctx = attn(
+        hidden_states[:, :n_context_tokens],
+        context=None,
+        pe=ctx_rope,
+    )
+
+    cur_rope = _slice_rope(video_rotary_emb, n_context_tokens, seq_len)
+    out_cur = attn(
+        hidden_states[:, n_context_tokens:],
+        context=hidden_states,
+        pe=cur_rope,
+        k_pe=video_rotary_emb,
+    )
+    return torch.cat([out_ctx, out_cur], dim=1)
+
+
+class SanaWMRefinerBlock(nn.Module):
+    """Video-only LTX-2 transformer block.
+
+    Diffusers-compatible layout: `norm1 -> attn1 (self) -> norm2 -> attn2 (cross) -> norm3 -> ff`,
+    each modulated via per-block `scale_shift_table` + token-wise `temb`.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        qk_norm: bool = True,
+        norm_eps: float = 1e-6,
+        apply_gated_attention: bool = False,
+        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__()
+        self.dim = int(dim)
+
+        self.norm1 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False)
+        self.attn1 = LTX2Attention(
+            query_dim=self.dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            norm_eps=norm_eps,
+            qk_norm=qk_norm,
+            apply_gated_attention=apply_gated_attention,
+            prefix=f"{prefix}.attn1",
+            quant_config=quant_config,
+        )
+
+        self.norm2 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False)
+        self.attn2 = LTX2Attention(
+            query_dim=self.dim,
+            context_dim=cross_attention_dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            norm_eps=norm_eps,
+            qk_norm=qk_norm,
+            use_local_attention=True,
+            apply_gated_attention=apply_gated_attention,
+            prefix=f"{prefix}.attn2",
+            quant_config=quant_config,
+        )
+
+        self.norm3 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False)
+        self.ff = LTX2FeedForward(self.dim, dim_out=self.dim, quant_config=quant_config)
+
+        self.scale_shift_table = nn.Parameter(torch.randn(6, self.dim) / self.dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        video_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        n_context_tokens: int = 0,
+    ) -> torch.Tensor:
+        B = hidden_states.size(0)
+        T = temb.size(1)
+        D = self.dim
+        ada = self.scale_shift_table[None, None].to(
+            device=temb.device, dtype=temb.dtype
+        ) + temb.reshape(B, T, 6, D)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada.unbind(
+            dim=2
+        )
+
+        normed = self.norm1(hidden_states) * (1 + scale_msa) + shift_msa
+        attn_out = _streaming_self_attention(
+            self.attn1,
+            normed,
+            video_rotary_emb,
+            n_context_tokens=n_context_tokens,
+        )
+        hidden_states = hidden_states + attn_out * gate_msa
+
+        normed = self.norm2(hidden_states)
+        ca_out = self.attn2(
+            normed,
+            context=encoder_hidden_states,
+            mask=encoder_attention_mask,
+            pe=None,
+        )
+        hidden_states = hidden_states + ca_out
+
+        normed = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp
+        hidden_states = hidden_states + self.ff(normed) * gate_mlp
+        return hidden_states
+
+
+class SanaWMLTX2VideoRefiner(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """SANA-WM stage-2 LTX-2 video-only refiner.
+
+    Loads Diffusers-format refiner weights from `<model_path>/refiner/transformer/`.
+    Audio params present in the checkpoint are silently dropped by the loader's
+    `strict=False` state_dict load.
+    """
+
+    _fsdp_shard_conditions = SanaWMRefinerArchConfig()._fsdp_shard_conditions
+    _compile_conditions = SanaWMRefinerArchConfig()._compile_conditions
+    _supported_attention_backends = (
+        SanaWMRefinerArchConfig()._supported_attention_backends
+    )
+    param_names_mapping = SanaWMRefinerArchConfig().param_names_mapping
+    reverse_param_names_mapping: dict = {}
+    lora_param_names_mapping: dict = {}
+
+    def __init__(
+        self,
+        config: SanaWMRefinerConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__(config, hf_config=hf_config)
+        arch = config.arch_config
+
+        self.in_channels = int(arch.in_channels)
+        self.out_channels = int(arch.out_channels)
+        self.patch_size = int(arch.patch_size)
+        self.patch_size_t = int(arch.patch_size_t)
+        self.hidden_size = int(arch.hidden_size)
+        self.num_attention_heads = int(arch.num_attention_heads)
+        self.num_channels_latents = int(arch.num_channels_latents)
+        self.attention_head_dim = int(arch.attention_head_dim)
+        self.timestep_scale_multiplier = float(arch.timestep_scale_multiplier)
+        self.rope_type = str(arch.rope_type)
+
+        in_dim = (
+            self.in_channels * self.patch_size_t * self.patch_size * self.patch_size
+        )
+        out_dim = (
+            self.out_channels * self.patch_size_t * self.patch_size * self.patch_size
+        )
+
+        self.proj_in = ColumnParallelLinear(
+            in_dim,
+            self.hidden_size,
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
+        )
+
+        self.time_embed = LTX2AdaLayerNormSingle(
+            self.hidden_size, embedding_coefficient=6
+        )
+        self.caption_projection = LTX2TextProjection(
+            in_features=int(arch.caption_channels),
+            hidden_size=self.hidden_size,
+            out_features=self.hidden_size,
+            act_fn="gelu_tanh",
+        )
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                SanaWMRefinerBlock(
+                    dim=self.hidden_size,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                    cross_attention_dim=int(arch.cross_attention_dim),
+                    qk_norm=bool(arch.qk_norm),
+                    norm_eps=float(arch.norm_eps),
+                    apply_gated_attention=bool(arch.apply_gated_attention),
+                    prefix=f"transformer_blocks.{i}",
+                    quant_config=quant_config,
+                )
+                for i in range(int(arch.num_layers))
+            ]
+        )
+
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, self.hidden_size) / self.hidden_size**0.5
+        )
+        self.norm_out = nn.LayerNorm(
+            self.hidden_size, eps=float(arch.norm_eps), elementwise_affine=False
+        )
+        self.proj_out = ColumnParallelLinear(
+            self.hidden_size,
+            out_dim,
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
+        )
+
+        # LTX2AudioVideoRotaryPosEmbed expects `dim` to be the *total* hidden
+        # size (num_heads * head_dim), not the per-head dim. It internally
+        # reshapes cos/sin to (B, T, num_heads, head_dim/2). Passing
+        # `attention_head_dim` here would size the RoPE to head_dim/num_heads
+        # and produce a (1, num_heads, L, 2) cos/sin that won't match
+        # LTX2Attention's q/k. See LTX2Transformer3DAVModel.__init__ in
+        # ltx_2.py for the canonical convention (`dim=self.hidden_size`).
+        self.rope = LTX2AudioVideoRotaryPosEmbed(
+            dim=self.hidden_size,
+            patch_size=self.patch_size,
+            patch_size_t=self.patch_size_t,
+            base_num_frames=int(arch.base_num_frames),
+            base_height=int(arch.base_height),
+            base_width=int(arch.base_width),
+            sampling_rate=int(arch.sampling_rate),
+            hop_length=int(arch.hop_length),
+            scale_factors=tuple(arch.scale_factors),
+            causal_offset=int(arch.causal_offset),
+            modality="video",
+            rope_type=self.rope_type,
+            num_attention_heads=self.num_attention_heads,
+        )
+
+        self.layer_names = ["transformer_blocks"]
+
+    def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor:
+        return timestep * self.timestep_scale_multiplier
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states_image=None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        fps: float = 24.0,
+        n_context_tokens: int = 0,
+        guidance=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        # Accept either packed (B, L, in_dim) or raw 5D (B, C, T, H, W).
+        if hidden_states.dim() == 5:
+            B_, _, T_, H_, W_ = hidden_states.shape
+            if num_frames is None:
+                num_frames = T_
+            if height is None:
+                height = H_
+            if width is None:
+                width = W_
+            hidden_states = pack_latents(
+                hidden_states,
+                patch_size=self.patch_size,
+                patch_size_t=self.patch_size_t,
+            )
+            packed_input = True
+        else:
+            if num_frames is None or height is None or width is None:
+                raise ValueError(
+                    "num_frames/height/width are required when hidden_states is pre-packed."
+                )
+            packed_input = False
+
+        B = hidden_states.size(0)
+
+        video_coords = self.rope.prepare_video_coords(
+            batch_size=B,
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            device=hidden_states.device,
+            fps=fps,
+        )
+        video_rotary_emb = self.rope(
+            video_coords,
+            device=hidden_states.device,
+            out_dtype=hidden_states.dtype,
+        )
+
+        hidden_states, _ = self.proj_in(hidden_states)
+
+        scaled_t = self._scale_timestep_for_adaln(timestep)
+        temb, embedded_timestep = self.time_embed(
+            scaled_t.flatten(), hidden_dtype=hidden_states.dtype
+        )
+        if timestep.dim() >= 2:
+            temb = temb.view(B, -1, temb.size(-1))
+            embedded_timestep = embedded_timestep.view(
+                B, -1, embedded_timestep.size(-1)
+            )
+        else:
+            temb = temb.view(B, 1, temb.size(-1))
+            embedded_timestep = embedded_timestep.view(B, 1, embedded_timestep.size(-1))
+
+        encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states.view(B, -1, self.hidden_size)
+
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                video_rotary_emb=video_rotary_emb,
+                encoder_attention_mask=encoder_attention_mask,
+                n_context_tokens=n_context_tokens,
+            )
+
+        scale_shift_values = self.scale_shift_table[None, None].to(
+            device=hidden_states.device, dtype=hidden_states.dtype
+        ) + embedded_timestep[:, :, None].to(hidden_states.dtype)
+        shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+        hidden_states = self.norm_out(hidden_states) * (1 + scale) + shift
+        hidden_states, _ = self.proj_out(hidden_states)
+
+        if packed_input:
+            return hidden_states
+        return unpack_latents(
+            hidden_states,
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            patch_size=self.patch_size,
+            patch_size_t=self.patch_size_t,
+        )
+
+
+EntryClass = SanaWMLTX2VideoRefiner

python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py

@@ -87,14 +87,47 @@ class LTX2VideoCausalConv3d(nn.Module):
             padding_mode=spatial_padding_mode,
         )
 
-    def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor:
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        causal: bool = True,
+        conv_cache: Optional[dict] = None,
+        cache_key: Optional[str] = None,
+    ) -> torch.Tensor:
         time_kernel_size = self.kernel_size[0]
 
         if causal:
-            pad_left = hidden_states[:, :, :1, :, :].repeat(
-                (1, 1, time_kernel_size - 1, 1, 1)
-            )
-            hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
+            if (
+                conv_cache is not None
+                and cache_key is not None
+                and time_kernel_size > 1
+            ):
+                # Streaming: prepend the previous chunk's last (k-1) frames of the
+                # PADDED conv input (first chunk: replicate frame 0, the "sink"),
+                # then store the last (k-1) frames of THIS padded input for the next
+                # chunk. Storing the padded tail (not the raw input) keeps exactly
+                # k-1 frames even for short chunks. Left-pad k-1 with temporal pad 0 /
+                # stride 1 keeps the output T invariant.
+                prev = conv_cache.get(cache_key)
+                if prev is None:
+                    pad_left = hidden_states[:, :, :1, :, :].repeat(
+                        (1, 1, time_kernel_size - 1, 1, 1)
+                    )
+                else:
+                    pad_left = prev.to(
+                        device=hidden_states.device, dtype=hidden_states.dtype
+                    )
+                hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
+                conv_cache[cache_key] = (
+                    hidden_states[:, :, -(time_kernel_size - 1) :, :, :]
+                    .detach()
+                    .clone()
+                )
+            else:
+                pad_left = hidden_states[:, :, :1, :, :].repeat(
+                    (1, 1, time_kernel_size - 1, 1, 1)
+                )
+                hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
         else:
             pad_left = hidden_states[:, :, :1, :, :].repeat(
                 (1, 1, (time_kernel_size - 1) // 2, 1, 1)
@@ -200,6 +233,8 @@ class LTX2VideoResnetBlock3d(nn.Module):
         temb: Optional[torch.Tensor] = None,
         generator: Optional[torch.Generator] = None,
         causal: bool = True,
+        conv_cache: Optional[dict] = None,
+        cache_key: Optional[str] = None,
     ) -> torch.Tensor:
         hidden_states = inputs
 
@@ -214,7 +249,12 @@ class LTX2VideoResnetBlock3d(nn.Module):
             hidden_states = hidden_states * (1 + scale_1) + shift_1
 
         hidden_states = self.nonlinearity(hidden_states)
-        hidden_states = self.conv1(hidden_states, causal=causal)
+        hidden_states = self.conv1(
+            hidden_states,
+            causal=causal,
+            conv_cache=conv_cache,
+            cache_key=None if cache_key is None else f"{cache_key}.conv1",
+        )
 
         if self.per_channel_scale1 is not None:
             spatial_shape = hidden_states.shape[-2:]
@@ -236,7 +276,12 @@ class LTX2VideoResnetBlock3d(nn.Module):
 
         hidden_states = self.nonlinearity(hidden_states)
         hidden_states = self.dropout(hidden_states)
-        hidden_states = self.conv2(hidden_states, causal=causal)
+        hidden_states = self.conv2(
+            hidden_states,
+            causal=causal,
+            conv_cache=conv_cache,
+            cache_key=None if cache_key is None else f"{cache_key}.conv2",
+        )
 
         if self.per_channel_scale2 is not None:
             spatial_shape = hidden_states.shape[-2:]
@@ -343,9 +388,27 @@ class LTXVideoUpsampler3d(nn.Module):
             spatial_padding_mode=spatial_padding_mode,
         )
 
-    def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor:
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        causal: bool = True,
+        conv_cache: Optional[dict] = None,
+        cache_key: Optional[str] = None,
+    ) -> torch.Tensor:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
 
+        # The temporal pixel-shuffle expands 1 latent frame -> stride[0] sample
+        # frames, then drops the first (stride[0]-1) frames as the causal "anchor".
+        # Under chunking that drop must happen ONLY at the true clip start, else a
+        # frame is lost at every chunk boundary. trim_start is read once and applied
+        # to BOTH the residual and shuffle branches (mirrors the reference).
+        trim_start = self.stride[0] - 1
+        if conv_cache is not None and cache_key is not None:
+            if conv_cache.get(f"{cache_key}.trim_applied", False):
+                trim_start = 0
+            else:
+                conv_cache[f"{cache_key}.trim_applied"] = True
+
         if self.residual:
             residual = hidden_states.reshape(
                 batch_size,
@@ -367,9 +430,11 @@ class LTXVideoUpsampler3d(nn.Module):
                 self.stride[0] * self.stride[1] * self.stride[2]
             ) // self.upscale_factor
             residual = residual.repeat(1, repeats, 1, 1, 1)
-            residual = residual[:, :, self.stride[0] - 1 :]
+            residual = residual[:, :, trim_start:]
 
-        hidden_states = self.conv(hidden_states, causal=causal)
+        hidden_states = self.conv(
+            hidden_states, causal=causal, conv_cache=conv_cache, cache_key=cache_key
+        )
         hidden_states = hidden_states.reshape(
             batch_size,
             -1,
@@ -386,7 +451,7 @@ class LTXVideoUpsampler3d(nn.Module):
             .flatten(4, 5)
             .flatten(2, 3)
         )
-        hidden_states = hidden_states[:, :, self.stride[0] - 1 :]
+        hidden_states = hidden_states[:, :, trim_start:]
 
         if self.residual:
             hidden_states = hidden_states + residual
@@ -597,6 +662,8 @@ class LTX2VideoMidBlock3d(nn.Module):
         temb: Optional[torch.Tensor] = None,
         generator: Optional[torch.Generator] = None,
         causal: bool = True,
+        conv_cache: Optional[dict] = None,
+        cache_key: Optional[str] = None,
     ) -> torch.Tensor:
         r"""Forward method of the `LTXMidBlock3D` class."""
 
@@ -616,7 +683,14 @@ class LTX2VideoMidBlock3d(nn.Module):
                     resnet, hidden_states, temb, generator, causal
                 )
             else:
-                hidden_states = resnet(hidden_states, temb, generator, causal=causal)
+                hidden_states = resnet(
+                    hidden_states,
+                    temb,
+                    generator,
+                    causal=causal,
+                    conv_cache=conv_cache,
+                    cache_key=None if cache_key is None else f"{cache_key}.resnets.{i}",
+                )
 
         return hidden_states
 
@@ -782,9 +856,18 @@ class LTX2VideoUpBlock3d(nn.Module):
         temb: Optional[torch.Tensor] = None,
         generator: Optional[torch.Generator] = None,
         causal: bool = True,
+        conv_cache: Optional[dict] = None,
+        cache_key: Optional[str] = None,
     ) -> torch.Tensor:
         if self.conv_in is not None:
-            hidden_states = self.conv_in(hidden_states, temb, generator, causal=causal)
+            hidden_states = self.conv_in(
+                hidden_states,
+                temb,
+                generator,
+                causal=causal,
+                conv_cache=conv_cache,
+                cache_key=None if cache_key is None else f"{cache_key}.conv_in",
+            )
 
         if self.time_embedder is not None:
             temb = self.time_embedder(
@@ -797,8 +880,15 @@ class LTX2VideoUpBlock3d(nn.Module):
             temb = temb.view(hidden_states.size(0), -1, 1, 1, 1)
 
         if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states, causal=causal)
+            for j, upsampler in enumerate(self.upsamplers):
+                hidden_states = upsampler(
+                    hidden_states,
+                    causal=causal,
+                    conv_cache=conv_cache,
+                    cache_key=(
+                        None if cache_key is None else f"{cache_key}.upsamplers.{j}"
+                    ),
+                )
 
         for i, resnet in enumerate(self.resnets):
             if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -806,7 +896,14 @@ class LTX2VideoUpBlock3d(nn.Module):
                     resnet, hidden_states, temb, generator, causal
                 )
             else:
-                hidden_states = resnet(hidden_states, temb, generator, causal=causal)
+                hidden_states = resnet(
+                    hidden_states,
+                    temb,
+                    generator,
+                    causal=causal,
+                    conv_cache=conv_cache,
+                    cache_key=None if cache_key is None else f"{cache_key}.resnets.{i}",
+                )
 
         return hidden_states
 
@@ -1116,10 +1213,17 @@ class LTX2VideoDecoder3d(nn.Module):
         hidden_states: torch.Tensor,
         temb: Optional[torch.Tensor] = None,
         causal: Optional[bool] = None,
+        conv_cache: Optional[dict] = None,
     ) -> torch.Tensor:
         causal = causal or self.is_causal
+        _ck = (lambda k: k) if conv_cache is not None else (lambda k: None)
 
-        hidden_states = self.conv_in(hidden_states, causal=causal)
+        hidden_states = self.conv_in(
+            hidden_states,
+            causal=causal,
+            conv_cache=conv_cache,
+            cache_key=_ck("conv_in"),
+        )
 
         if self.timestep_scale_multiplier is not None:
             temb = temb * self.timestep_scale_multiplier
@@ -1134,10 +1238,22 @@ class LTX2VideoDecoder3d(nn.Module):
                     up_block, hidden_states, temb, None, causal
                 )
         else:
-            hidden_states = self.mid_block(hidden_states, temb, causal=causal)
+            hidden_states = self.mid_block(
+                hidden_states,
+                temb,
+                causal=causal,
+                conv_cache=conv_cache,
+                cache_key=_ck("mid_block"),
+            )
 
-            for up_block in self.up_blocks:
-                hidden_states = up_block(hidden_states, temb, causal=causal)
+            for i, up_block in enumerate(self.up_blocks):
+                hidden_states = up_block(
+                    hidden_states,
+                    temb,
+                    causal=causal,
+                    conv_cache=conv_cache,
+                    cache_key=_ck(f"up_blocks.{i}"),
+                )
 
         hidden_states = self.norm_out(hidden_states)
 
@@ -1155,7 +1271,12 @@ class LTX2VideoDecoder3d(nn.Module):
             hidden_states = hidden_states * (1 + scale) + shift
 
         hidden_states = self.conv_act(hidden_states)
-        hidden_states = self.conv_out(hidden_states, causal=causal)
+        hidden_states = self.conv_out(
+            hidden_states,
+            causal=causal,
+            conv_cache=conv_cache,
+            cache_key=_ck("conv_out"),
+        )
 
         p = self.patch_size
         p_t = self.patch_size_t
@@ -1990,4 +2111,38 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE):
         return dec
 
 
-EntryClass = AutoencoderKLLTX2Video
+class AutoencoderKLCausalLTX2Video(AutoencoderKLLTX2Video):
+    """Streaming causal LTX-2 VAE.
+
+    Same weights / architecture as ``AutoencoderKLLTX2Video`` but loaded with
+    ``decoder_causal=True`` (the checkpoint config.json flips it automatically),
+    plus a chunk-by-chunk ``decode_chunk`` that threads a per-conv ``conv_cache``
+    across chunks so frames decode causally as the streaming denoise produces
+    them -- the ``decode_per_frame_with_cache`` equivalent at chunk granularity.
+    The streaming VAE config.json's ``_class_name`` is
+    ``AutoencoderKLCausalLTX2Video``, so this name must be registered.
+    """
+
+    @staticmethod
+    def reset_decoder_cache() -> dict:
+        """Fresh decoder cache (per-conv feat tails + upsampler trim flags)."""
+        return {}
+
+    def decode_chunk(
+        self,
+        z_chunk: torch.Tensor,
+        conv_cache: dict,
+        *,
+        temb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Causally decode one chunk of latents, carrying ``conv_cache`` across calls.
+
+        ``z_chunk`` is ``(B, C_latent, T_chunk, h, w)`` and must already be
+        de-normalized (the stage applies the latents scale/shift). Returns the
+        pixel chunk ``(B, 3, T_out, h*sf, w*sf)`` and mutates ``conv_cache`` in
+        place. Bypasses the tiled/framewise ``decode`` dispatch on purpose.
+        """
+        return self.decoder(z_chunk, temb, causal=True, conv_cache=conv_cache)
+
+
+EntryClass = [AutoencoderKLLTX2Video, AutoencoderKLCausalLTX2Video]

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -5,6 +5,9 @@ import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
 
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
+    STAGE_2_DISTILLED_SIGMA_VALUES as _SHARED_STAGE_2_DISTILLED_SIGMA_VALUES,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
     LTX2PipelineConfig,
     is_ltx23_native_variant,
@@ -779,7 +782,7 @@ class LTX2TwoStageResidencyController:
 
 class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
     pipeline_name = "LTX2TwoStagePipeline"
-    STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
+    STAGE_2_DISTILLED_SIGMA_VALUES = list(_SHARED_STAGE_2_DISTILLED_SIGMA_VALUES)
     STAGE_1_DISTILLED_LORA_STRENGTH = 0.0
     STAGE_2_DISTILLED_LORA_STRENGTH = 1.0
     STAGE_1_DENOISING_SAMPLER_NAME = "euler"

python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py

@@ -0,0 +1,322 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig
+from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams
+from sglang.multimodal_gen.runtime.loader.utils import get_memory_usage_of_component
+from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    InputValidationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import (
+    SanaWMBeforeDenoisingStage,
+    SanaWMDecodingStage,
+    SanaWMDenoisingStage,
+    SanaWMTextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.refiner import (
+    OfficialDiffusersLTX2RefinerModule,
+    OfficialGemma3TextEncoderModule,
+    SanaWMLTX2RefinerStage,
+    SanaWMRefinerDecodingStage,
+    default_sana_wm_refiner_dtype,
+    sana_wm_skip_refiner_enabled,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import (
+    SanaWMStreamingDecodingStage,
+    SanaWMStreamingDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming_refiner import (
+    SanaWMStreamingRefinerStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+# Stage-2 refiner sub-modules live under `<model_path>/refiner/...`, not at the
+# model root. They're loaded manually in `initialize_pipeline` rather than via
+# `_required_config_modules`, because the framework verifier resolves every
+# required module key as a literal top-level subdir of the materialized model.
+
+
+logger = init_logger(__name__)
+
+
+class SanaWMPipeline(LoRAPipeline, ComposedPipelineBase):
+    """SANA-WM TI2V pipeline (single-stage)."""
+
+    pipeline_name = "SanaWMPipeline"
+    pipeline_config_cls = SanaWMPipelineConfig
+    sampling_params_cls = SanaWMSamplingParams
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    @staticmethod
+    def _validate_parallelism_args(server_args: ServerArgs) -> None:
+        tp_size = getattr(server_args, "tp_size", 1) or 1
+        if tp_size != 1:
+            raise ValueError(
+                "SANA-WM does not support tensor parallelism yet. "
+                "Use --num-gpus with FSDP/CFG parallelism instead of "
+                f"--tp-size {tp_size}."
+            )
+
+        sp_degree = getattr(server_args, "sp_degree", 1) or 1
+        if sp_degree != 1:
+            raise ValueError(
+                "SANA-WM does not support temporal sequence parallelism yet. "
+                "Stage-1 GDN/GLUMBConvTemp span frames and require halo/state "
+                "exchange before latents can be sharded. Use --num-gpus with "
+                "FSDP/CFG parallelism instead of "
+                f"--sp-degree {sp_degree}."
+            )
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        self._validate_parallelism_args(server_args)
+        self.add_stage(InputValidationStage())
+
+        self.add_stage(
+            SanaWMTextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+            "prompt_encoding_stage",
+        )
+
+        self.add_stage(
+            SanaWMBeforeDenoisingStage(
+                vae=self.get_module("vae"),
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                pipeline_config=server_args.pipeline_config,
+            ),
+            "sana_wm_before_denoising",
+        )
+
+        if getattr(server_args.pipeline_config, "streaming", False):
+            DenoiseStage = SanaWMStreamingDenoisingStage
+        else:
+            DenoiseStage = SanaWMDenoisingStage
+        self.add_stage(
+            DenoiseStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        # Subclasses (e.g. SanaWMTwoStagePipeline) insert latent-domain stages
+        # between denoising and VAE decoding.
+        self._maybe_add_refiner_stage(server_args)
+
+        self._add_decoding_stage(server_args)
+
+    def _add_decoding_stage(self, server_args: ServerArgs = None) -> None:
+        if server_args is not None and getattr(
+            server_args.pipeline_config, "streaming", False
+        ):
+            DecodeStage = SanaWMStreamingDecodingStage
+        else:
+            DecodeStage = SanaWMDecodingStage
+        self.add_stage(
+            DecodeStage(
+                vae=self.get_module("vae"),
+                pipeline=self,
+                component_name="vae",
+            ),
+            "decoding_stage",
+        )
+
+    def _maybe_add_refiner_stage(self, server_args: ServerArgs) -> None:
+        """Hook for subclasses; single-stage pipeline is a no-op."""
+        return None
+
+
+class SanaWMTwoStagePipeline(SanaWMPipeline):
+    """SANA-WM two-stage pipeline: SANA-WM DiT + LTX-2 latent refiner.
+
+    Stage-1 produces a coarse 720p latent; the LTX-2 refiner runs 3 Euler steps
+    on it before VAE decode, matching the NVlabs ``inference_sana_wm.py`` default.
+    """
+
+    pipeline_name = "SanaWMTwoStagePipeline"
+
+    # Stage-2 refiner sub-modules and their on-disk layout. Loaded through the
+    # official Diffusers/Transformers classes because NVlabs' reference refiner
+    # is a narrow video-only wrapper around those modules.
+    _REFINER_SUB_MODULES: tuple[tuple[str, str], ...] = (
+        ("transformer_2", "refiner/transformer"),
+        ("connectors", "refiner/connectors"),
+        ("text_encoder_2", "refiner/text_encoder"),
+        # The refiner Gemma-3 ships its tokenizer files alongside the encoder.
+        ("tokenizer_2", "refiner/text_encoder"),
+    )
+
+    def initialize_pipeline(self, server_args: ServerArgs) -> None:
+        super().initialize_pipeline(server_args)
+        if sana_wm_skip_refiner_enabled():
+            logger.info(
+                "SANA-WM refiner component loading skipped by "
+                "SGLANG_SANA_WM_SKIP_REFINER."
+            )
+            return
+        self._load_refiner_modules(server_args)
+
+    def _resolve_refiner_paths(self, server_args: ServerArgs) -> tuple[str, str]:
+        component_paths = getattr(server_args, "component_paths", {}) or {}
+        refiner_root = component_paths.get(
+            "refiner", os.path.join(self.model_path, "refiner")
+        )
+        refiner_gemma_root = component_paths.get(
+            "refiner_text_encoder",
+            component_paths.get(
+                "text_encoder_2", os.path.join(refiner_root, "text_encoder")
+            ),
+        )
+        return refiner_root, refiner_gemma_root
+
+    def _resolve_refiner_component_path(
+        self, server_args: ServerArgs, module_name: str, subpath: str
+    ) -> str:
+        component_paths = getattr(server_args, "component_paths", {}) or {}
+        if module_name in component_paths:
+            return self._resolve_component_path(server_args, module_name, subpath)
+
+        if (
+            "refiner" not in component_paths
+            and "refiner_text_encoder" not in component_paths
+        ):
+            return self._resolve_component_path(server_args, module_name, subpath)
+
+        refiner_root, refiner_gemma_root = self._resolve_refiner_paths(server_args)
+        if module_name in ("text_encoder_2", "tokenizer_2"):
+            return refiner_gemma_root
+
+        rel_subpath = subpath.removeprefix("refiner/")
+        return os.path.join(refiner_root, rel_subpath)
+
+    def _load_refiner_modules(self, server_args: ServerArgs) -> None:
+        for module_name, subpath in self._REFINER_SUB_MODULES:
+            component_path = self._resolve_refiner_component_path(
+                server_args, module_name, subpath
+            )
+            logger.info(
+                "SANA-WM loading refiner component %s from %s",
+                module_name,
+                component_path,
+            )
+            module, memory_usage = self._load_official_refiner_component(
+                module_name,
+                component_path,
+                server_args,
+            )
+            self.modules[module_name] = module
+            self.memory_usages[module_name] = memory_usage
+
+    @staticmethod
+    def _load_official_refiner_component(
+        module_name: str,
+        component_path: str,
+        server_args: ServerArgs,
+    ):
+        """Load SANA-WM refiner modules through the same libraries as NVlabs.
+
+        The upstream wrapper (``diffusion/refiner/diffusers_ltx2_refiner.py``)
+        keeps the LTX-2 transformer/connectors as Diffusers modules and only
+        customizes the video-only forward surface; use that path for the
+        quality-critical stage-2 refiner instead of the experimental native port.
+        """
+
+        dtype = default_sana_wm_refiner_dtype(server_args)
+        if module_name == "transformer_2":
+            from diffusers.models.transformers.transformer_ltx2 import (
+                LTX2VideoTransformer3DModel,
+            )
+
+            module = LTX2VideoTransformer3DModel.from_pretrained(
+                component_path,
+                torch_dtype=dtype,
+            ).eval()
+            module = OfficialDiffusersLTX2RefinerModule(module)
+        elif module_name == "connectors":
+            from diffusers.pipelines.ltx2 import LTX2TextConnectors
+
+            module = LTX2TextConnectors.from_pretrained(
+                component_path,
+                torch_dtype=dtype,
+            ).eval()
+        elif module_name == "text_encoder_2":
+            from transformers import Gemma3ForConditionalGeneration
+
+            module = Gemma3ForConditionalGeneration.from_pretrained(
+                component_path,
+                torch_dtype=dtype,
+                low_cpu_mem_usage=True,
+            ).eval()
+            module = OfficialGemma3TextEncoderModule(module)
+        elif module_name == "tokenizer_2":
+            from transformers import AutoTokenizer
+
+            module = AutoTokenizer.from_pretrained(component_path)
+        else:
+            raise ValueError(f"Unsupported SANA-WM refiner component: {module_name}")
+
+        memory_usage = get_memory_usage_of_component(module)
+        logger.info(
+            "Loaded %s: %s (official native version). model size: %s GB",
+            module_name,
+            module.__class__.__name__,
+            memory_usage if memory_usage is not None else "NA",
+        )
+        return module, memory_usage or 0.0
+
+    def _maybe_add_refiner_stage(self, server_args: ServerArgs) -> None:
+        if sana_wm_skip_refiner_enabled():
+            return
+        pc = server_args.pipeline_config
+        common = dict(
+            transformer=self.get_module("transformer_2"),
+            connectors=self.get_module("connectors"),
+            text_encoder=self.get_module("text_encoder_2"),
+            tokenizer=self.get_module("tokenizer_2"),
+            dtype=default_sana_wm_refiner_dtype(server_args),
+        )
+        if getattr(pc, "streaming", False) and getattr(pc, "refiner_chunked", True):
+            stage = SanaWMStreamingRefinerStage(
+                **common,
+                block_size=int(getattr(pc, "refiner_block_size", 3)),
+                kv_max_frames=int(getattr(pc, "refiner_kv_max_frames", 11)),
+                sink_size=int(getattr(pc, "sink_size", 1)),
+                seed=int(getattr(pc, "refiner_seed", 42)),
+            )
+        else:
+            stage = SanaWMLTX2RefinerStage(**common)
+        self.add_stage(stage, "sana_wm_refiner")
+
+    def _add_decoding_stage(self, server_args: ServerArgs = None) -> None:
+        # Streaming and skip-refiner both route to the base decode
+        # (SanaWMStreamingDecodingStage / dense decode); otherwise dense refiner-decode.
+        streaming = server_args is not None and getattr(
+            server_args.pipeline_config, "streaming", False
+        )
+        if streaming or sana_wm_skip_refiner_enabled():
+            return super()._add_decoding_stage(server_args)
+        self.add_stage(
+            SanaWMRefinerDecodingStage(
+                vae=self.get_module("vae"),
+                pipeline=self,
+                component_name="vae",
+            ),
+            "decoding_stage",
+        )
+
+
+EntryClass = [SanaWMPipeline, SanaWMTwoStagePipeline]

python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import (
+    SanaWMRealtimeConfig,
+)
+from sglang.multimodal_gen.runtime.pipelines.sana_wm_pipeline import (
+    SanaWMTwoStagePipeline,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    RealtimeInputValidationStage,
+    RealtimeTextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import (
+    SanaWMTextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_chain import (
+    SanaWMCameraCondStage,
+    SanaWMCausalDecodeChainStage,
+    SanaWMChunkedRefinerChainStage,
+    SanaWMCondFrameEncodeStage,
+    SanaWMRealtimeLatentPrepStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.refiner import (
+    default_sana_wm_refiner_dtype,
+    sana_wm_skip_refiner_enabled,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import (
+    SanaWMStreamingDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming_refiner import (
+    SanaWMStreamingRefinerStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
+
+DEFAULT_SANA_WM_TEXT_ENCODER = "Efficient-Large-Model/gemma-2-2b-it"
+
+
+class SanaWMRealtimeTextEncodingStage(
+    RealtimeTextEncodingStage, SanaWMTextEncodingStage
+):
+    """Realtime text encoding using SANA-WM prompt processing.
+
+    MRO contract: ``RealtimeTextEncodingStage.forward`` (per-session cache) calls
+    ``super().forward``, which must resolve to ``SanaWMTextEncodingStage.forward``.
+    This preserves the chi prompt prefix and official prompt window used by the
+    batch path.
+    """
+
+
+class SanaWMRealtimePipeline(SanaWMTwoStagePipeline):
+    """SANA-WM realtime interactive pipeline.
+
+    Extends the two-stage pipeline to inherit refiner sub-module loading (``transformer_2`` /
+    ``connectors`` / ``text_encoder_2`` / ``tokenizer_2``). The streaming refiner stage built
+    here is purely a carrier of those modules handed to ``SanaWMRealtimeStage``, not added to
+    the stage list (the realtime stage drives the incremental stage-1 session + chunked refiner
+    runner per user action).
+    """
+
+    pipeline_name = "SanaWMRealtimePipeline"
+    is_video_pipeline = True
+    # Must be the realtime config so get_realtime_model_adapter() resolves the
+    # SANA-WM adapter (the realtime registry keys on SanaWMRealtimeConfig).
+    pipeline_config_cls = SanaWMRealtimeConfig
+
+    def _resolve_component_path(
+        self, server_args: ServerArgs, module_name: str, load_module_name: str
+    ) -> str:
+        if (
+            module_name in {"text_encoder", "tokenizer"}
+            and module_name not in server_args.component_paths
+        ):
+            return maybe_download_model(DEFAULT_SANA_WM_TEXT_ENCODER)
+        return super()._resolve_component_path(
+            server_args,
+            module_name,
+            load_module_name,
+        )
+
+    def _build_realtime_refiner_stage(self, server_args: ServerArgs):
+        """Build the chunked streaming refiner carrier when refiner modules exist."""
+        if sana_wm_skip_refiner_enabled():
+            return None
+        if self.get_module("transformer_2") is None:
+            return None
+
+        pc = server_args.pipeline_config
+        return SanaWMStreamingRefinerStage(
+            transformer=self.get_module("transformer_2"),
+            connectors=self.get_module("connectors"),
+            text_encoder=self.get_module("text_encoder_2"),
+            tokenizer=self.get_module("tokenizer_2"),
+            dtype=default_sana_wm_refiner_dtype(server_args),
+            block_size=int(getattr(pc, "refiner_block_size", 3)),
+            kv_max_frames=int(getattr(pc, "refiner_kv_max_frames", 11)),
+            sink_size=int(getattr(pc, "sink_size", 1)),
+            seed=int(getattr(pc, "refiner_seed", 42)),
+        )
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        refiner_stage = self._build_realtime_refiner_stage(server_args)
+        common = dict(
+            transformer=self.get_module("transformer"),
+            vae=self.get_module("vae"),
+            model_path=self.model_path,
+        )
+        self.add_stage(RealtimeInputValidationStage())
+        self.add_stage(
+            SanaWMRealtimeTextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            )
+        )
+        self.add_stage(SanaWMCondFrameEncodeStage(**common))
+        self.add_stage(
+            SanaWMRealtimeLatentPrepStage(
+                use_refiner=refiner_stage is not None, **common
+            )
+        )
+        self.add_stage(SanaWMCameraCondStage(**common))
+        self.add_stage(
+            SanaWMStreamingDenoisingStage(
+                transformer=self.get_module("transformer"), keep_resident=True
+            )
+        )
+        if refiner_stage is not None:
+            self.add_stage(
+                SanaWMChunkedRefinerChainStage(refiner_stage=refiner_stage, **common)
+            )
+        self.add_stage(SanaWMCausalDecodeChainStage(**common))
+
+
+EntryClass = SanaWMRealtimePipeline

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -436,4 +436,5 @@ class OutputBatch:
         self.rollout_trajectory_data = None
         self.trajectory_decoded = None
         self.output_file_paths = None
+        self.raw_frame_batches = None
         self.noise_pred = None

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

@@ -65,6 +65,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av i
 from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import (
     LTX2DenoisingStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import (
+    RealtimeDiffusionStage,
+    RealtimeStageComponent,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_input_validation import (
     RealtimeInputValidationStage,
 )
@@ -95,6 +99,8 @@ __all__ = [
     "PipelineStage",
     "InputValidationStage",
     "RealtimeInputValidationStage",
+    "RealtimeDiffusionStage",
+    "RealtimeStageComponent",
     "TimestepPreparationStage",
     "DMDTimestepPreparationStage",
     "LatentPreparationStage",

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

@@ -31,6 +31,33 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 logger = init_logger(__name__)
 
 
+def scale_and_shift_latents(latents: torch.Tensor, server_args, vae) -> torch.Tensor:
+    """De-normalize latents before VAE decode (single shared implementation).
+
+    Used by DecodingStage.scale_and_shift and by realtime stages that decode
+    outside a DecodingStage instance.
+    """
+    scaling_factor, shift_factor = (
+        server_args.pipeline_config.get_decode_scale_and_shift(
+            latents.device, latents.dtype, vae
+        )
+    )
+
+    # 1. scale
+    if isinstance(scaling_factor, torch.Tensor):
+        latents = latents / scaling_factor.to(latents.device, latents.dtype)
+    else:
+        latents = latents / scaling_factor
+
+    # 2. apply shifting if needed
+    if shift_factor is not None:
+        if isinstance(shift_factor, torch.Tensor):
+            latents = latents + shift_factor.to(latents.device, latents.dtype)
+        else:
+            latents = latents + shift_factor
+    return latents
+
+
 def _ensure_tensor_decode_output(decode_output):
     """
     Ensure VAE decode output is a tensor.
@@ -106,25 +133,7 @@ class DecodingStage(PipelineStage):
         return result
 
     def scale_and_shift(self, latents: torch.Tensor, server_args):
-        scaling_factor, shift_factor = (
-            server_args.pipeline_config.get_decode_scale_and_shift(
-                latents.device, latents.dtype, self.vae
-            )
-        )
-
-        # 1. scale
-        if isinstance(scaling_factor, torch.Tensor):
-            latents = latents / scaling_factor.to(latents.device, latents.dtype)
-        else:
-            latents = latents / scaling_factor
-
-        # 2. apply shifting if needed
-        if shift_factor is not None:
-            if isinstance(shift_factor, torch.Tensor):
-                latents += shift_factor.to(latents.device, latents.dtype)
-            else:
-                latents += shift_factor
-        return latents
+        return scale_and_shift_latents(latents, server_args, self.vae)
 
     @torch.no_grad()
     def decode(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/__init__.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SANA-WM pipeline stages (package).
+
+The realtime serving framework drives ``SanaWMRealtimeStage`` over the
+``/v1/realtime_video`` WebSocket. Base stages + helpers are re-exported here
+for back-compat.
+"""
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base import (
+    SanaWMBeforeDenoisingStage,
+    SanaWMDecodingStage,
+    SanaWMDenoisingStage,
+    SanaWMTextEncodingStage,
+    _align_sana_wm_cfg_text_conditions,
+    configure_sana_wm_ltx2_vae_for_long_video,
+    parse_sana_wm_action_string,
+    sana_wm_action_to_camera_to_world,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_stage import (
+    SanaWMRealtimeStage,
+)
+
+__all__ = [
+    "SanaWMBeforeDenoisingStage",
+    "SanaWMDecodingStage",
+    "SanaWMDenoisingStage",
+    "SanaWMTextEncodingStage",
+    "_align_sana_wm_cfg_text_conditions",
+    "configure_sana_wm_ltx2_vae_for_long_video",
+    "parse_sana_wm_action_string",
+    "sana_wm_action_to_camera_to_world",
+    "SanaWMRealtimeStage",
+]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/parity_probe.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Env-gated parity dump/probe helpers (no-op in production).
+
+Single home for the debug harness that localized the realtime<->batch parity
+bugs (previously copy-pasted across streaming/realtime/refiner/sana_wm).
+
+Env vars:
+- ``SANAWM_FORK_DUMP_DIR``  — batch-path dumps (init_noise, conditioning,
+  per-chunk stage-1, kv/refiner probes, weights fingerprint)
+- ``SANAWM_RT_DUMP_DIR``    — realtime-path dumps (same set)
+- ``SANAWM_BLOCK_PROBE``    — file path: per-block forward_long checksums on
+  the first sink-path call
+- ``SANAWM_INJECT_DIR``     — batch-path injection inputs (read, not written)
+
+All checksums reduce in float64 (bf16/fp32 upcasts are lossless, so the
+formula is bitwise-stable across both paths).
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+ENV_FORK_DUMP = "SANAWM_FORK_DUMP_DIR"
+ENV_RT_DUMP = "SANAWM_RT_DUMP_DIR"
+ENV_BLOCK_PROBE = "SANAWM_BLOCK_PROBE"
+ENV_INJECT = "SANAWM_INJECT_DIR"
+
+
+def probe_dir(*env_names: str) -> Path | None:
+    """First configured dump dir among ``env_names`` (created on demand)."""
+    for name in env_names:
+        value = os.environ.get(name)
+        if value:
+            path = Path(value)
+            path.mkdir(parents=True, exist_ok=True)
+            return path
+    return None
+
+
+def dump_tensor(dirpath: Path | str | None, name: str, tensor) -> None:
+    """Save ``tensor`` as float32 CPU under ``dirpath/name.pt`` (None-safe)."""
+    if dirpath is None or tensor is None:
+        return
+    torch.save(tensor.detach().float().cpu(), Path(dirpath) / f"{name}.pt")
+
+
+def dump_obj(dirpath: Path | str | None, name: str, obj) -> None:
+    """Save a picklable object (e.g. a checksum dict) under ``dirpath/name.pt``."""
+    if dirpath is None or obj is None:
+        return
+    torch.save(obj, Path(dirpath) / f"{name}.pt")
+
+
+def checksum(tensor) -> tuple[tuple[int, ...], float] | None:
+    """(shape, float64 sum) of a tensor; None passes through."""
+    if tensor is None:
+        return None
+    return (
+        tuple(tensor.shape),
+        float(tensor.detach().float().double().sum().item()),
+    )
+
+
+def kv_cache_checksums(chunk_kv, sink_num: int) -> dict:
+    """Per-block, per-slot checksums of an accumulated KV cache."""
+    probe: dict = {"sink_num": sink_num}
+    for block_id, slots in enumerate(chunk_kv):
+        for slot_id, tensor in enumerate(slots):
+            if tensor is not None:
+                probe[f"b{block_id:02d}s{slot_id}"] = checksum(tensor)
+    return probe
+
+
+def weights_fingerprint(module: torch.nn.Module) -> dict[str, float]:
+    """Per-parameter |.|-sum fingerprint (catches wrong/mutated weights)."""
+    return {
+        name: float(param.detach().float().abs().sum().item())
+        for name, param in module.named_parameters()
+    }

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py

@@ -0,0 +1,556 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SANA-WM realtime stage chain.
+
+Per-tick pipeline:
+
+    RealtimeInputValidationStage          (framework)
+    SanaWMRealtimeTextEncodingStage       (framework cache via MRO)
+    SanaWMCondFrameEncodeStage            -> batch.image_latent (+ inputs snapshot)
+    SanaWMRealtimeLatentPrepStage         -> batch.latents = this tick's pre-noised
+                                             chunk(s); batch.extra["sana_wm_chunk_plan"]
+    SanaWMCameraCondStage                 -> batch.extra camera_conditions/chunk_plucker
+    SanaWMStreamingDenoisingStage         (session path; SanaWMStreamCacheState)
+    SanaWMChunkedRefinerChainStage        -> batch.latents = refined buffer
+    SanaWMCausalDecodeChainStage          -> OutputBatch (decodes past its frontier)
+
+Chain stages share ``SanaWMRealtimeStage`` for first-frame image handling and
+causal VAE decode helpers; model-specific chunk planning stays in this module.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import (
+    compute_chunk_plucker,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
+from sglang.multimodal_gen.runtime.realtime.causal_state import (
+    RealtimeCausalDecodeState,
+)
+from sglang.multimodal_gen.runtime.realtime.session import BaseRealtimeState
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+from . import parity_probe
+from .base import (
+    _SANA_WM_DEFAULT_ROTATION_SPEED_DEG,
+    _SANA_WM_DEFAULT_TRANSLATION_SPEED,
+    SanaWMBeforeDenoisingStage,
+    configure_sana_wm_ltx2_vae_for_long_video,
+    snap_sana_wm_num_frames,
+)
+from .realtime_stage import (
+    DEFAULT_REFINER_BLOCK_SIZE,
+    DEFAULT_REFINER_KV_MAX_FRAMES,
+    SANA_WM_HEIGHT,
+    SANA_WM_WIDTH,
+    SanaWMRealtimeStage,
+    _motion_param,
+)
+from .streaming import SanaWMStreamCacheState, SanaWMStreamingDenoisingStage
+from .streaming_refiner import RefinerChunkRunner
+
+
+# --------------------------------------------------------------------- #
+# Per-session state blobs (one per owning stage)
+# --------------------------------------------------------------------- #
+class SanaWMSessionInputsState(BaseRealtimeState):
+    """Session input snapshot — written by the cond-encode stage, read by the
+    camera stage (crop geometry) and the refiner stage (prompt)."""
+
+    def __init__(self):
+        super().__init__()
+        self.image = None
+        self.intrinsics_image = None
+        self.src_size = None
+        self.resized_size = None
+        self.crop_offset = None
+        self.target_height: int | None = None
+        self.target_width: int | None = None
+        self.prompt: str = ""
+        self.open_ended = False
+        self.latent_t: int | None = None  # None => open-ended
+        self.num_frame_per_block = DEFAULT_REFINER_BLOCK_SIZE
+        self.sink_size = 1
+        self.camera_actions: list[list[str]] = []
+        self.max_camera_actions = 0  # unlimited (sessions stream past any horizon)
+        self.static_c2w: np.ndarray | None = None
+        self.intrinsics_raw: np.ndarray | None = None
+        self.translation_speed = _SANA_WM_DEFAULT_TRANSLATION_SPEED
+        self.rotation_speed_deg = _SANA_WM_DEFAULT_ROTATION_SPEED_DEG
+        self.first_latent: torch.Tensor | None = None
+
+    def dispose(self):
+        super().dispose()
+        self.__init__()
+
+
+class SanaWMNoiseState(BaseRealtimeState):
+    """Seeded noise discipline — full-horizon buffer for fixed N (sliced per
+    chunk, matching the offline draw bitwise), seeded fallback otherwise."""
+
+    def __init__(self):
+        super().__init__()
+        self.noise_buffer: torch.Tensor | None = None
+        self.generator: torch.Generator | None = None
+        self.segments: list[int] | None = None  # front-loaded grid (fixed N)
+
+    def dispose(self):
+        super().dispose()
+        self.noise_buffer = None
+        self.generator = None
+        self.segments = None
+
+
+class SanaWMRefinerChainState(BaseRealtimeState):
+    def __init__(self):
+        super().__init__()
+        self.runner: RefinerChunkRunner | None = None
+        self.refined_full: torch.Tensor | None = None
+        self.next_ref_idx = 0
+        self.block_size = DEFAULT_REFINER_BLOCK_SIZE
+        self.sink_size = 1
+
+    def dispose(self):
+        super().dispose()
+        self.runner = None
+        self.refined_full = None
+        self.next_ref_idx = 0
+        self.block_size = DEFAULT_REFINER_BLOCK_SIZE
+        self.sink_size = 1
+
+
+# --------------------------------------------------------------------- #
+# Chain stages
+# --------------------------------------------------------------------- #
+class SanaWMCondFrameEncodeStage(SanaWMRealtimeStage):
+    """Encode the first frame and write the session input snapshot."""
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        session = self.require_session(batch, context="SANA-WM realtime chain")
+        self._pipeline_config = server_args.pipeline_config
+        device = get_local_torch_device()
+        weight_dtype = PRECISION_TO_TYPE.get(
+            getattr(server_args.pipeline_config, "dit_precision", "bf16"),
+            torch.bfloat16,
+        )
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+
+        self.vae = self.vae.to(device=device, dtype=vae_dtype).eval()
+        configure_sana_wm_ltx2_vae_for_long_video(self.vae, server_args.pipeline_config)
+
+        st = session.get_or_create_state(SanaWMSessionInputsState)
+        if batch.block_idx == 0 or st.image is None:
+            st.dispose()
+            (
+                st.image,
+                st.intrinsics_image,
+                st.src_size,
+                st.resized_size,
+                st.crop_offset,
+            ) = self._prepare_image(batch)
+            st.target_width, st.target_height = st.image.size
+            st.prompt = str(batch.prompt)
+            st.open_ended = self._is_open_ended(batch)
+            if st.open_ended:
+                st.latent_t = None
+            else:
+                num_frames = snap_sana_wm_num_frames(int(batch.num_frames), stride=8)
+                batch.num_frames = num_frames
+                st.latent_t = (num_frames - 1) // 8 + 1
+            st.num_frame_per_block = int(
+                batch.extra.get("sana_wm_num_frame_per_block", 3)
+            )
+            st.sink_size = int(batch.extra.get("sana_wm_sink_size", 1))
+            st.translation_speed = _motion_param(
+                batch, "translation_speed", _SANA_WM_DEFAULT_TRANSLATION_SPEED
+            )
+            st.rotation_speed_deg = _motion_param(
+                batch, "rotation_speed_deg", _SANA_WM_DEFAULT_ROTATION_SPEED_DEG
+            )
+            st.static_c2w = self._prepare_static_camera(
+                batch,
+                num_frames=(
+                    ((st.latent_t - 1) * 8 + 1)
+                    if st.latent_t is not None
+                    else st.num_frame_per_block * 8 + 1
+                ),
+                translation_speed=st.translation_speed,
+                rotation_speed_deg=st.rotation_speed_deg,
+            )
+            st.first_latent = self._get_first_frame_latent(
+                batch,
+                st.image,
+                device=device,
+                vae_dtype=vae_dtype,
+                latent_dtype=weight_dtype,
+            )
+        batch.image_latent = st.first_latent
+        return batch
+
+
+class SanaWMRealtimeLatentPrepStage(SanaWMRealtimeStage):
+    """This tick's chunk PLAN + pre-noised latents.
+
+    Fixed N: front-loaded segments + a full-horizon seeded buffer sliced per
+    chunk (bitwise-matching the offline draw); past the horizon / open-ended:
+    uniform chunks from the seeded fallback generator. The plan holds more than
+    one chunk only when the refiner block grid lags the stage-1 chunk grid (the
+    option-b boundary tick) so every tick emits frames."""
+
+    def __init__(self, *, use_refiner: bool, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.use_refiner = bool(use_refiner)
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        device = get_local_torch_device()
+        weight_dtype = PRECISION_TO_TYPE.get(
+            getattr(server_args.pipeline_config, "dit_precision", "bf16"),
+            torch.bfloat16,
+        )
+        session = self.require_session(batch, context="SANA-WM realtime chain")
+        inputs = session.get_or_create_state(SanaWMSessionInputsState)
+        noise = session.get_or_create_state(SanaWMNoiseState)
+        cache = session.get_or_create_state(SanaWMStreamCacheState)
+        first_latent = batch.image_latent
+        if first_latent is None:
+            raise ValueError("cond-frame latent missing (run the encode stage first)")
+
+        nfpb = int(inputs.num_frame_per_block)
+        if batch.block_idx == 0 or (
+            noise.noise_buffer is None and noise.generator is None
+        ):
+            noise.dispose()
+            seed = batch.seed[0] if isinstance(batch.seed, list) else batch.seed
+            if seed is not None:
+                noise.generator = torch.Generator(device=device).manual_seed(int(seed))
+            if inputs.latent_t is not None:
+                noise.segments = SanaWMStreamingDenoisingStage._autoregressive_segments(
+                    int(inputs.latent_t), nfpb
+                )
+                gen = (
+                    batch.generator[0]
+                    if isinstance(batch.generator, list)
+                    else batch.generator
+                ) or noise.generator
+                buf = torch.randn(
+                    1,
+                    first_latent.shape[1],
+                    int(inputs.latent_t),
+                    first_latent.shape[-2],
+                    first_latent.shape[-1],
+                    device=device,
+                    dtype=weight_dtype,
+                    generator=gen,
+                )
+                buf[:, :, :1] = first_latent
+                noise.noise_buffer = buf
+                parity_probe.dump_tensor(  # parity harness (no-op in prod)
+                    parity_probe.probe_dir(parity_probe.ENV_RT_DUMP),
+                    "noise_buffer",
+                    buf,
+                )
+
+        # Chunk plan: enough stage-1 chunks for >=1 refiner block.
+        stage1_cur = cache.chunk_indices[-1]
+        chunk_idx = cache.chunk_idx
+
+        def _next_len(idx: int) -> int:
+            if noise.segments is not None and idx + 1 < len(noise.segments):
+                # Front-loaded grid (chunk 0's segment already includes the
+                # conditioning frame).
+                return noise.segments[idx + 1] - noise.segments[idx]
+            n = nfpb
+            if idx == 0:
+                # Uniform (open-ended) grid: chunk 0 = cond frame + nfpb new
+                # frames, matching the retired engine's step(n_frames=nfpb).
+                n += first_latent.shape[2]
+            return n
+
+        plan: list[int] = []
+        sim = stage1_cur
+        idx = chunk_idx
+        if self.use_refiner and chunk_idx > 0:
+            sink = int(inputs.sink_size)
+            block = int(inputs.num_frame_per_block)
+            refined_cur = sink + max(0, (sim - sink)) // block * block
+            target = refined_cur + block
+            while sim < target:
+                n = _next_len(idx)
+                plan.append(n)
+                sim += n
+                idx += 1
+        else:
+            plan.append(_next_len(idx))
+
+        pieces = []
+        pos = stage1_cur
+        for i, n in enumerate(plan):
+            is_chunk0 = (chunk_idx + i) == 0
+            new_frames = n - (first_latent.shape[2] if is_chunk0 else 0)
+            lo = pos + (first_latent.shape[2] if is_chunk0 else 0)
+            hi = lo + new_frames
+            if noise.noise_buffer is not None and hi <= noise.noise_buffer.shape[2]:
+                draw = noise.noise_buffer[:, :, lo:hi].to(device, weight_dtype)
+            else:
+                draw = torch.randn(
+                    first_latent.shape[0],
+                    first_latent.shape[1],
+                    new_frames,
+                    first_latent.shape[-2],
+                    first_latent.shape[-1],
+                    device=device,
+                    dtype=weight_dtype,
+                    generator=noise.generator,
+                )
+            pieces.append(
+                torch.cat([first_latent.to(device, weight_dtype), draw], dim=2)
+                if is_chunk0
+                else draw
+            )
+            pos = hi
+        batch.latents = torch.cat(pieces, dim=2)
+        batch.extra["sana_wm_chunk_plan"] = plan
+        return batch
+
+
+class SanaWMCameraCondStage(SanaWMRealtimeStage):
+    """Camera conditioning windows — appends this tick's actions and rebuilds
+    the full-length raymap/chunk_plucker through the END of the planned chunks
+    (single-sourced through the batch helpers; parity RC#3)."""
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        device = get_local_torch_device()
+        weight_dtype = PRECISION_TO_TYPE.get(
+            getattr(server_args.pipeline_config, "dit_precision", "bf16"),
+            torch.bfloat16,
+        )
+        session = self.require_session(batch, context="SANA-WM realtime chain")
+        inputs = session.get_or_create_state(SanaWMSessionInputsState)
+        cache = session.get_or_create_state(SanaWMStreamCacheState)
+        plan = list(batch.extra.get("sana_wm_chunk_plan") or [])
+        target_latent = cache.chunk_indices[-1] + sum(plan)
+        if cache.chunk_idx == 0 and plan:
+            target_latent = max(target_latent, plan[0])
+
+        self._append_realtime_camera_actions(batch, inputs)
+        camera, plucker = self._build_camera_windows(
+            batch,
+            inputs,
+            target_latent=target_latent,
+            device=device,
+            dtype=weight_dtype,
+        )
+        batch.extra["camera_conditions"] = camera
+        batch.extra["chunk_plucker"] = plucker
+        return batch
+
+    def _build_camera_windows(
+        self,
+        batch: Req,
+        inputs: SanaWMSessionInputsState,
+        *,
+        target_latent: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Port of the mega-stage's ``_update_camera_tensors`` (value-identical
+        ops; returns instead of writing legacy state). Coverage runs through the
+        END of the planned chunks; within a fixed horizon ``max()`` keeps it
+        EXACTLY at num_frames — bitwise-identical camera tensors."""
+        if (
+            inputs.src_size is None
+            or inputs.resized_size is None
+            or inputs.crop_offset is None
+        ):
+            raise ValueError("SANA-WM crop metadata is not initialized")
+        needed = (int(target_latent) - 1) * 8 + 1
+        if inputs.open_ended:
+            num_frames = needed
+        else:
+            num_frames = max(int(batch.num_frames), needed)
+        condition_inputs = batch.condition_inputs or {}
+        if "translation_speed" in condition_inputs:
+            inputs.translation_speed = float(condition_inputs["translation_speed"])
+        if "rotation_speed_deg" in condition_inputs:
+            inputs.rotation_speed_deg = float(condition_inputs["rotation_speed_deg"])
+        c2w = self._camera_from_state(
+            inputs,
+            num_frames=num_frames,
+            translation_speed=inputs.translation_speed,
+            rotation_speed_deg=inputs.rotation_speed_deg,
+        )
+        intrinsics_raw = self._prepare_intrinsics(
+            batch, inputs, num_frames=num_frames, device=device
+        )
+        inputs.intrinsics_raw = intrinsics_raw
+        vae_time_stride = 8
+        pixel_h = int(inputs.target_height or batch.height or SANA_WM_HEIGHT)
+        pixel_w = int(inputs.target_width or batch.width or SANA_WM_WIDTH)
+        latent_h = pixel_h // 32
+        latent_w = pixel_w // 32
+        latent_t = (num_frames - 1) // vae_time_stride + 1
+        camera_to_world = (
+            torch.from_numpy(np.asarray(c2w, dtype=np.float32))
+            .unsqueeze(0)
+            .to(device=device, dtype=torch.float32)
+        )
+        intrinsics_vec4 = (
+            torch.from_numpy(np.asarray(intrinsics_raw, dtype=np.float32))
+            .unsqueeze(0)
+            .to(device=device, dtype=torch.float32)
+        )
+        intrinsics_vec4 = (
+            SanaWMBeforeDenoisingStage._transform_intrinsics_for_condition_image(
+                intrinsics_vec4,
+                {
+                    "source_size": inputs.src_size,
+                    "resized_size": inputs.resized_size,
+                    "crop_offset": inputs.crop_offset,
+                },
+            )
+        )
+        rel_poses = SanaWMBeforeDenoisingStage._relative_camera_poses(camera_to_world)
+        intrinsics_latent = SanaWMBeforeDenoisingStage._scale_intrinsics_to_latent(
+            intrinsics_vec4,
+            pixel_h=pixel_h,
+            pixel_w=pixel_w,
+            latent_h=latent_h,
+            latent_w=latent_w,
+        )
+        original_camera = SanaWMBeforeDenoisingStage._flatten_camera_conditions(
+            rel_poses, intrinsics_latent
+        )
+        raymap = SanaWMBeforeDenoisingStage._latent_frame_camera_conditions(
+            original_camera,
+            num_frames=num_frames,
+            latent_frames=latent_t,
+            vae_temporal_stride=vae_time_stride,
+        ).to(device=device, dtype=dtype)
+        chunk_plucker = compute_chunk_plucker(
+            original_camera,
+            HW=(latent_t, latent_h, latent_w),
+            vae_temporal_stride=vae_time_stride,
+            patch_size=(1, 1, 1),
+        ).to(device=device, dtype=dtype)
+        return raymap, chunk_plucker
+
+
+class SanaWMChunkedRefinerChainStage(SanaWMRealtimeStage):
+    """Chunked LTX-2 refiner on the uniform ``sink + i*block`` grid — refines
+    COMPLETE blocks only (option b) and hands the refined buffer downstream."""
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        device = get_local_torch_device()
+        session = self.require_session(batch, context="SANA-WM realtime chain")
+        inputs = session.get_or_create_state(SanaWMSessionInputsState)
+        st = session.get_or_create_state(SanaWMRefinerChainState)
+        stage1 = batch.latents  # growing stage-1 buffer from the denoise stage
+        if stage1 is None:
+            raise ValueError("refiner stage expects the stage-1 latent buffer")
+
+        with set_forward_context(
+            current_timestep=batch.block_idx, attn_metadata=None, forward_batch=batch
+        ):
+            if st.runner is None:
+                st.block_size = int(inputs.num_frame_per_block)
+                st.sink_size = int(inputs.sink_size)
+                first_latent = batch.image_latent
+                st.refined_full = first_latent.detach().clone()
+                st.runner = self._build_chain_refiner_runner(
+                    inputs,
+                    st,
+                    batch,
+                    device=device,
+                    spatial_shape=(
+                        int(first_latent.shape[3]),
+                        int(first_latent.shape[4]),
+                    ),
+                )
+
+            frontier = st.refined_full.shape[2]
+            end_f = stage1.shape[2]
+            while True:
+                block_start = frontier
+                block_end = block_start + st.block_size  # complete blocks only
+                if block_end > end_f:
+                    break
+                sink_seed = (
+                    stage1[:, :, : st.sink_size]
+                    if block_start == st.sink_size
+                    else None
+                )
+                refined = st.runner.refine_block(
+                    block_idx=st.next_ref_idx,
+                    clean_block=stage1[:, :, block_start:block_end].contiguous(),
+                    block_start=block_start,
+                    block_end=block_end,
+                    sink_seed_frames=sink_seed,
+                )
+                st.refined_full = torch.cat(
+                    [st.refined_full, refined.to(st.refined_full.dtype)], dim=2
+                )
+                st.next_ref_idx += 1
+                frontier = block_end
+
+        batch.latents = st.refined_full
+        return batch
+
+    def _build_chain_refiner_runner(
+        self, inputs, st, batch, *, device, spatial_shape
+    ) -> RefinerChunkRunner:
+        # Reuse the parity-validated builder via a state adapter that exposes
+        # prompt/sink/block/kv_max in the legacy state shape it expects.
+        legacy = SimpleNamespace(
+            prompt=inputs.prompt,
+            sink_size=st.sink_size,
+            refiner_block_size=st.block_size,
+            refiner_kv_max_frames=int(
+                batch.extra.get(
+                    "sana_wm_refiner_kv_max_frames", DEFAULT_REFINER_KV_MAX_FRAMES
+                )
+            ),
+        )
+        return self._build_refiner_runner(
+            legacy,
+            batch,
+            device=device,
+            spatial_shape=spatial_shape,
+            seed=int(batch.extra.get("sana_wm_refiner_seed", batch.seed)),
+            fps=float(batch.fps),
+        )
+
+
+class SanaWMCausalDecodeChainStage(SanaWMRealtimeStage):
+    """Causal-VAE chunk decode past this session's decode frontier."""
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
+        device = get_local_torch_device()
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+
+        self.vae = self.vae.to(device=device, dtype=vae_dtype).eval()
+        configure_sana_wm_ltx2_vae_for_long_video(self.vae, server_args.pipeline_config)
+        session = self.require_session(batch, context="SANA-WM realtime chain")
+        st = session.get_or_create_state(RealtimeCausalDecodeState)
+        src = batch.latents
+        if src is None or src.shape[2] <= st.next_dec_idx:
+            return self._empty_output(batch)
+        with set_forward_context(
+            current_timestep=batch.block_idx, attn_metadata=None, forward_batch=batch
+        ):
+            frames = self._decode_chunk(
+                src[:, :, st.next_dec_idx :], st, server_args, vae_dtype=vae_dtype
+            )
+        st.next_dec_idx = src.shape[2]
+        return OutputBatch(output=frames.to(torch.float32), metrics=batch.metrics)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py

@@ -0,0 +1,448 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import (
+    RealtimeDiffusionStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+from .base import (
+    SANA_WM_TARGET_HEIGHT,
+    SANA_WM_TARGET_WIDTH,
+    SanaWMBeforeDenoisingStage,
+    normalize_sana_wm_camera_actions,
+    sana_wm_action_to_camera_to_world_array,
+    sana_wm_load_camera,
+    sana_wm_load_intrinsics,
+    sana_wm_normalize_vae_latents,
+    sana_wm_pil_to_model_tensor,
+    sana_wm_resize_and_center_crop,
+)
+from .refiner import (
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+    SanaWMLTX2RefinerStage,
+    _unwrap_diffusers_ltx2_refiner,
+)
+from .streaming_refiner import (
+    RefinerChunkRunner,
+    _RefinerCore,
+)
+
+logger = init_logger(__name__)
+
+SANA_WM_HEIGHT = SANA_WM_TARGET_HEIGHT
+SANA_WM_WIDTH = SANA_WM_TARGET_WIDTH
+DEFAULT_REFINER_BLOCK_SIZE = 3
+DEFAULT_REFINER_KV_MAX_FRAMES = 11
+
+
+@contextmanager
+def _deterministic_vae_encode_context():
+    prev_benchmark = torch.backends.cudnn.benchmark
+    prev_deterministic = torch.backends.cudnn.deterministic
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    try:
+        yield
+    finally:
+        torch.backends.cudnn.benchmark = prev_benchmark
+        torch.backends.cudnn.deterministic = prev_deterministic
+
+
+def _normalize_camera_actions(payload: Any) -> list[list[str]]:
+    return normalize_sana_wm_camera_actions(payload, allow_none=True)
+
+
+def _actions_to_action_string(actions: list[list[str]]) -> str:
+    if not actions:
+        return "none-1"
+
+    segments: list[str] = []
+    current = tuple(sorted(set(actions[0])))
+    count = 0
+    for frame_actions in actions:
+        normalized = tuple(sorted(set(frame_actions)))
+        if normalized == current:
+            count += 1
+            continue
+        key = "".join(current) if current else "none"
+        segments.append(f"{key}-{count}")
+        current = normalized
+        count = 1
+    key = "".join(current) if current else "none"
+    segments.append(f"{key}-{count}")
+    return ",".join(segments)
+
+
+def _normalize_intrinsics_array(arr: Any, num_frames: int) -> np.ndarray:
+    intrinsics = np.asarray(arr, dtype=np.float32)
+    if intrinsics.shape == (4,):
+        return np.broadcast_to(intrinsics, (num_frames, 4)).copy()
+    if intrinsics.shape == (3, 3):
+        vec = np.array(
+            [
+                intrinsics[0, 0],
+                intrinsics[1, 1],
+                intrinsics[0, 2],
+                intrinsics[1, 2],
+            ],
+            dtype=np.float32,
+        )
+        return np.broadcast_to(vec, (num_frames, 4)).copy()
+    if intrinsics.ndim == 2 and intrinsics.shape[1] == 4:
+        if intrinsics.shape[0] < num_frames:
+            pad = np.broadcast_to(
+                intrinsics[-1:], (num_frames - intrinsics.shape[0], 4)
+            )
+            intrinsics = np.concatenate([intrinsics, pad], axis=0)
+        return intrinsics[:num_frames].copy()
+    if intrinsics.ndim == 3 and intrinsics.shape[1:] == (3, 3):
+        if intrinsics.shape[0] < num_frames:
+            pad = np.broadcast_to(
+                intrinsics[-1:], (num_frames - intrinsics.shape[0], 3, 3)
+            )
+            intrinsics = np.concatenate([intrinsics, pad], axis=0)
+        intrinsics = intrinsics[:num_frames]
+        return np.stack(
+            [
+                intrinsics[:, 0, 0],
+                intrinsics[:, 1, 1],
+                intrinsics[:, 0, 2],
+                intrinsics[:, 1, 2],
+            ],
+            axis=1,
+        ).astype(np.float32)
+    raise ValueError(
+        "intrinsics must have shape (4,), (3,3), (F,4), or (F,3,3), "
+        f"got {intrinsics.shape}"
+    )
+
+
+def _motion_param(batch: Req, name: str, default: float) -> float:
+    value = (batch.condition_inputs or {}).get(name)
+    if value is None:
+        value = batch.extra.get(f"sana_wm_{name}", default)
+    return float(value)
+
+
+def _default_source_intrinsics(image: Image.Image, num_frames: int) -> np.ndarray:
+    width, height = image.size
+    focal = 0.8 * float(max(width, height))
+    intrinsics = np.array(
+        [focal, focal, width / 2.0, height / 2.0],
+        dtype=np.float32,
+    )
+    return np.broadcast_to(intrinsics, (num_frames, 4)).copy()
+
+
+class SanaWMRealtimeStage(RealtimeDiffusionStage):
+    def __init__(
+        self,
+        *,
+        transformer: torch.nn.Module,
+        vae: torch.nn.Module,
+        model_path: str,
+        refiner_stage: SanaWMLTX2RefinerStage | None = None,
+    ) -> None:
+        super().__init__(
+            transformer=transformer,
+            vae=vae,
+            model_path=model_path,
+            default_height=SANA_WM_HEIGHT,
+            default_width=SANA_WM_WIDTH,
+        )
+        # Injected by the pipeline; we reuse its transformer / text encoder to
+        # drive the RefinerChunkRunner. None -> stage-1-only output.
+        self.refiner_stage = refiner_stage
+        self.first_frame_latent_cache = None
+
+    def _prepare_image(
+        self, batch: Req
+    ) -> tuple[
+        Image.Image, Image.Image, tuple[int, int], tuple[int, int], tuple[int, int]
+    ]:
+        if isinstance(batch.condition_image, Image.Image):
+            original = batch.condition_image.convert("RGB")
+        elif batch.image_path is not None and isinstance(batch.image_path, str):
+            original = Image.open(batch.image_path).convert("RGB")
+        else:
+            raise ValueError("SANA-WM realtime requires a first-frame image")
+        target_h, target_w = self.target_pixel_size(batch)
+        cropped, src_size, resized_size, crop_offset = sana_wm_resize_and_center_crop(
+            original, target_h, target_w
+        )
+        return cropped, original, src_size, resized_size, crop_offset
+
+    def _prepare_static_camera(
+        self,
+        batch: Req,
+        *,
+        num_frames: int,
+        translation_speed: float,
+        rotation_speed_deg: float,
+    ) -> np.ndarray | None:
+        condition_inputs = batch.condition_inputs or {}
+        camera_path = condition_inputs.get("camera_path")
+        if camera_path is not None:
+            c2w = sana_wm_load_camera(Path(str(camera_path)))
+        elif condition_inputs.get("camera") is not None:
+            c2w = np.asarray(condition_inputs["camera"], dtype=np.float32)
+        elif condition_inputs.get("action") is not None:
+            c2w = sana_wm_action_to_camera_to_world_array(
+                str(condition_inputs["action"]),
+                translation_speed=translation_speed,
+                rotation_speed_deg=rotation_speed_deg,
+            )
+        else:
+            return None
+
+        if c2w.ndim != 3 or c2w.shape[1:] != (4, 4):
+            raise ValueError(
+                f"camera trajectory must have shape (F,4,4), got {c2w.shape}"
+            )
+        if c2w.shape[0] < num_frames:
+            pad = np.broadcast_to(c2w[-1:], (num_frames - c2w.shape[0], 4, 4))
+            c2w = np.concatenate([c2w, pad], axis=0)
+        return c2w[:num_frames].astype(np.float32)
+
+    def _prepare_intrinsics(
+        self,
+        batch: Req,
+        state,
+        *,
+        num_frames: int,
+        device: torch.device,
+    ) -> np.ndarray:
+        condition_inputs = batch.condition_inputs or {}
+        if condition_inputs.get("intrinsics_path") is not None:
+            return sana_wm_load_intrinsics(
+                Path(str(condition_inputs["intrinsics_path"])), num_frames
+            )
+        if condition_inputs.get("intrinsics") is not None:
+            return _normalize_intrinsics_array(
+                condition_inputs["intrinsics"], num_frames
+            )
+        if state.intrinsics_raw is not None:
+            cached = state.intrinsics_raw
+            if cached.shape[0] < num_frames:
+                # Open-ended growth: hold the last row (fixed-horizon sessions
+                # cache the full-length array, so this never triggers).
+                pad = np.broadcast_to(
+                    cached[-1:], (num_frames - cached.shape[0],) + cached.shape[1:]
+                )
+                cached = np.concatenate([cached, pad], axis=0)
+            return cached
+        if state.intrinsics_image is None:
+            raise ValueError("SANA-WM image is not initialized")
+        logger.info("No intrinsics provided; using heuristic centered intrinsics.")
+        return _default_source_intrinsics(state.intrinsics_image, num_frames)
+
+    def _append_realtime_camera_actions(self, batch: Req, state) -> None:
+        actions = _normalize_camera_actions(
+            (batch.condition_inputs or {}).get("camera_actions")
+        )
+        if actions:
+            state.camera_actions.extend(actions)
+            if (
+                state.max_camera_actions > 0
+                and len(state.camera_actions) > state.max_camera_actions
+            ):
+                del state.camera_actions[state.max_camera_actions :]
+
+    def _camera_from_state(
+        self,
+        state,
+        *,
+        num_frames: int,
+        translation_speed: float,
+        rotation_speed_deg: float,
+    ) -> np.ndarray:
+        if state.static_c2w is not None:
+            c2w = state.static_c2w
+            if c2w.shape[0] < num_frames:
+                # Open-ended growth: hold the last pose (fixed-horizon sessions
+                # always precompute the full trajectory, so this never triggers).
+                pad = np.broadcast_to(c2w[-1:], (num_frames - c2w.shape[0], 4, 4))
+                c2w = np.concatenate([c2w, pad], axis=0)
+            return c2w[:num_frames]
+
+        num_actions = max(0, num_frames - 1)
+        actions = list(state.camera_actions[:num_actions])
+        if len(actions) < num_actions:
+            actions.extend([[] for _ in range(num_actions - len(actions))])
+        return sana_wm_action_to_camera_to_world_array(
+            _actions_to_action_string(actions),
+            translation_speed=translation_speed,
+            rotation_speed_deg=rotation_speed_deg,
+        )[:num_frames]
+
+    @torch.inference_mode()
+    def _encode_first_frame(
+        self,
+        image: Image.Image,
+        *,
+        device: torch.device,
+        vae_dtype: torch.dtype,
+        latent_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if hasattr(self.vae, "enable_tiling"):
+            self.vae.enable_tiling()
+        image_tensor = sana_wm_pil_to_model_tensor(
+            image, device=device, dtype=vae_dtype
+        )
+        with _deterministic_vae_encode_context():
+            # Must use the SAME shared core as the batch path's _vae_encode_image:
+            # drifting this was a parity root cause.
+            z = SanaWMBeforeDenoisingStage._extract_vae_latents(
+                self.vae.encode(image_tensor.to(device=device, dtype=vae_dtype))
+            ).float()
+        z = sana_wm_normalize_vae_latents(
+            self.vae, z, getattr(self, "_pipeline_config", None)
+        )
+        return z.to(device=device, dtype=latent_dtype)
+
+    def _first_frame_cache_key(
+        self,
+        batch: Req,
+        *,
+        device: torch.device,
+        vae_dtype: torch.dtype,
+        latent_dtype: torch.dtype,
+    ) -> tuple | None:
+        if batch.image_path is None or not isinstance(batch.image_path, str):
+            return None
+        stat = os.stat(batch.image_path)
+        target_h, target_w = self.target_pixel_size(batch)
+        return (
+            batch.image_path,
+            stat.st_mtime_ns,
+            stat.st_size,
+            target_h,
+            target_w,
+            device.type,
+            device.index,
+            vae_dtype,
+            latent_dtype,
+        )
+
+    @torch.inference_mode()
+    def _get_first_frame_latent(
+        self,
+        batch: Req,
+        image: Image.Image,
+        *,
+        device: torch.device,
+        vae_dtype: torch.dtype,
+        latent_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        cache_key = self._first_frame_cache_key(
+            batch,
+            device=device,
+            vae_dtype=vae_dtype,
+            latent_dtype=latent_dtype,
+        )
+        if (
+            cache_key is not None
+            and self.first_frame_latent_cache is not None
+            and self.first_frame_latent_cache[0] == cache_key
+        ):
+            return self.first_frame_latent_cache[1]
+
+        first_latent = self._encode_first_frame(
+            image,
+            device=device,
+            vae_dtype=vae_dtype,
+            latent_dtype=latent_dtype,
+        )
+        if cache_key is not None:
+            self.first_frame_latent_cache = (cache_key, first_latent.detach())
+        return first_latent
+
+    @torch.inference_mode()
+    def _decode_chunk(
+        self,
+        latents: torch.Tensor,
+        state,
+        server_args: ServerArgs,
+        *,
+        vae_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if latents.shape[2] == 0:
+            vae_stride = getattr(server_args.pipeline_config, "vae_stride", (8, 32, 32))
+            return torch.empty(
+                (
+                    latents.shape[0],
+                    3,
+                    0,
+                    int(latents.shape[3]) * int(vae_stride[1]),
+                    int(latents.shape[4]) * int(vae_stride[2]),
+                ),
+                dtype=torch.float32,
+                device=latents.device,
+            )
+        conv_cache = self.ensure_causal_vae_conv_cache(state)
+        z = self.scale_and_shift_latents(latents.to(vae_dtype), server_args)
+        decoded = self.vae.decode_chunk(z, conv_cache)
+        return (decoded / 2 + 0.5).clamp(0, 1)
+
+    def _build_refiner_runner(
+        self,
+        state,
+        batch: Req,
+        *,
+        device: torch.device,
+        spatial_shape: tuple[int, int],
+        seed: int,
+        fps: float,
+    ) -> RefinerChunkRunner:
+        """Build the chunked LTX-2 refiner runner that carries sink/history KV across blocks."""
+        rs = self.refiner_stage
+        if rs is None:
+            raise RuntimeError("SANA-WM realtime refiner stage is not initialized")
+        # Keep refiner sub-modules resident: the realtime stage runs the refiner
+        # directly, outside the offline stage's use_declared_component context.
+        for _mod in (rs.text_encoder, rs.connectors, rs.transformer):
+            if _mod is not None:
+                _mod.to(device)
+        prompt_embeds, prompt_mask = rs._encode_prompt(state.prompt, device)
+        unwrapped = _unwrap_diffusers_ltx2_refiner(rs.transformer)
+        core = _RefinerCore(unwrapped, device, rs.dtype)
+        sigmas = torch.tensor(
+            STAGE_2_DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device
+        )
+        return RefinerChunkRunner(
+            core,
+            prompt_embeds=prompt_embeds,
+            prompt_attention_mask=prompt_mask,
+            fps=fps,
+            sigmas=sigmas,
+            source_sink_frames=state.sink_size,
+            block_size=state.refiner_block_size,
+            kv_max_frames=state.refiner_kv_max_frames,
+            seed=seed,
+            spatial_shape=spatial_shape,
+        )
+
+    @staticmethod
+    def _is_open_ended(batch: Req) -> bool:
+        """Open-ended = the init request carried no num_frames.
+
+        The adapter flags this via condition_inputs (build_sampling_params strips
+        None fields, so batch.num_frames would otherwise carry the SamplingParams
+        default and be indistinguishable from an explicit request)."""
+        condition_inputs = batch.condition_inputs or {}
+        if bool(condition_inputs.get("sana_wm_open_ended")):
+            return True
+        num_frames = getattr(batch, "num_frames", None)
+        return num_frames is None or int(num_frames) <= 0

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py

@@ -0,0 +1,799 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
+    STAGE_2_DISTILLED_SIGMA_VALUES as _STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
+    pack_text_embeds,
+)
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_classifier_free_guidance_rank,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+    ComponentUse,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.sana_wm_refiner_transformer import (
+    pack_latents,
+    unpack_latents,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+    PipelineStage,
+    StageParallelismType,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+from .base import (
+    SanaWMDecodingStage,
+    log_sana_wm_tensor_stats,
+    sana_wm_diagnostics_enabled,
+)
+
+logger = init_logger(__name__)
+
+# Distilled 3-step sigma schedule, matches NVlabs `inference_sana_wm.py`.
+# Canonical value lives in the LTX-2 pipeline config (shared with LTX2TwoStagePipeline).
+STAGE_2_DISTILLED_SIGMA_VALUES: tuple[float, ...] = _STAGE_2_DISTILLED_SIGMA_VALUES
+
+# Default Gemma-3 token budget for the refiner prompt encoder.
+_REFINER_TEXT_MAX_LENGTH = 1024
+
+
+class _OfficialLayerwiseModule(nn.Module, LayerwiseOffloadableModuleMixin):
+    def __init__(self, module: nn.Module) -> None:
+        super().__init__()
+        self.module = module
+        self.layerwise_offload_managers = []
+
+    def forward(self, *args, **kwargs):
+        return self.module(*args, **kwargs)
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError as exc:
+            wrapped = self.__dict__.get("module")
+            if wrapped is not None:
+                return getattr(wrapped, name)
+            modules = self.__dict__.get("_modules", {})
+            wrapped = modules.get("module")
+            if wrapped is not None:
+                return getattr(wrapped, name)
+            raise exc
+
+
+class OfficialDiffusersLTX2RefinerModule(_OfficialLayerwiseModule):
+    """Thin offload wrapper around Diffusers' official LTX-2 refiner module."""
+
+    layer_names = ["module.transformer_blocks"]
+
+
+class OfficialGemma3TextEncoderModule(_OfficialLayerwiseModule):
+    """Thin offload wrapper around HF Gemma-3 used by the official refiner."""
+
+    layer_names = [
+        "module.language_model.layers",
+        "module.model.language_model.layers",
+    ]
+
+
+def _truthy_flag(value: Any) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, (int, float)):
+        return bool(value)
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return False
+
+
+def sana_wm_skip_refiner_enabled(batch: Req | None = None) -> bool:
+    if os.getenv("SGLANG_SANA_WM_SKIP_REFINER", "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }:
+        return True
+    if batch is None:
+        return False
+    extra = getattr(batch, "extra", None) or {}
+    diffusers_kwargs = extra.get("diffusers_kwargs", {})
+    if not isinstance(diffusers_kwargs, dict):
+        diffusers_kwargs = {}
+    return any(
+        _truthy_flag(value)
+        for value in (
+            extra.get("skip_refiner"),
+            extra.get("sana_wm_skip_refiner"),
+            diffusers_kwargs.get("skip_refiner"),
+            diffusers_kwargs.get("sana_wm_skip_refiner"),
+        )
+    )
+
+
+def default_sana_wm_refiner_dtype(server_args: ServerArgs) -> torch.dtype:
+    precision = getattr(server_args.pipeline_config, "dit_precision", "bf16")
+    return PRECISION_TO_TYPE.get(precision, torch.bfloat16)
+
+
+def _is_current_cfg_main_rank() -> bool:
+    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+        return True
+    try:
+        return get_classifier_free_guidance_rank() == 0
+    except AssertionError:
+        return True
+
+
+def _pack_text_embeds(
+    text_hidden_states: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    *,
+    padding_side: str = "left",
+    scale_factor: int = 8,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Gemma-3 masked min-max text-embed pooling for the LTX-2 refiner.
+
+    Delegates to framework-wide ``pack_text_embeds`` (verified BITWISE-identical,
+    incl. bf16 and both padding sides) — a private fork would drift, like the
+    realtime<->batch parity bugs.
+    """
+    return pack_text_embeds(
+        text_hidden_states,
+        sequence_lengths,
+        padding_side=padding_side,
+        scale_factor=scale_factor,
+        eps=eps,
+    )
+
+
+def _refiner_config_value(transformer: nn.Module, name: str) -> Any:
+    transformer = _unwrap_diffusers_ltx2_refiner(transformer)
+    config = getattr(transformer, "config", None)
+    if config is not None:
+        if isinstance(config, dict) and name in config:
+            return config[name]
+        if hasattr(config, name):
+            return getattr(config, name)
+    return getattr(transformer, name)
+
+
+def _unwrap_diffusers_ltx2_refiner(transformer: nn.Module) -> nn.Module:
+    if isinstance(transformer, OfficialDiffusersLTX2RefinerModule):
+        return transformer.module
+    return transformer
+
+
+def _uses_diffusers_ltx2_refiner(transformer: nn.Module) -> bool:
+    transformer = _unwrap_diffusers_ltx2_refiner(transformer)
+    return transformer.__class__.__name__ == "LTX2VideoTransformer3DModel"
+
+
+def _as_additive_attention_mask(
+    attention_mask: torch.Tensor | None,
+    dtype: torch.dtype,
+) -> torch.Tensor | None:
+    if attention_mask is None:
+        return None
+    if attention_mask.ndim == 2:
+        return ((1 - attention_mask.to(dtype)) * -10000.0).unsqueeze(1)
+    return attention_mask.to(dtype)
+
+
+def _forward_diffusers_video_only(
+    transformer: nn.Module,
+    *,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    encoder_attention_mask: torch.Tensor | None,
+    num_frames: int,
+    height: int,
+    width: int,
+    fps: float,
+    n_context_tokens: int,
+) -> torch.Tensor:
+    """Official SANA-WM LTX-2 video-only forward adapted to injected modules."""
+
+    batch_size = hidden_states.size(0)
+    encoder_attention_mask = _as_additive_attention_mask(
+        encoder_attention_mask, hidden_states.dtype
+    )
+
+    video_coords = transformer.rope.prepare_video_coords(
+        batch_size, num_frames, height, width, hidden_states.device, fps=fps
+    )
+    video_rotary_emb = transformer.rope(video_coords, device=hidden_states.device)
+
+    hidden_states = transformer.proj_in(hidden_states)
+    temb, embedded_timestep = transformer.time_embed(
+        timestep.flatten(),
+        batch_size=batch_size,
+        hidden_dtype=hidden_states.dtype,
+    )
+    temb = temb.view(batch_size, -1, temb.size(-1))
+    embedded_timestep = embedded_timestep.view(
+        batch_size, -1, embedded_timestep.size(-1)
+    )
+
+    encoder_hidden_states = transformer.caption_projection(encoder_hidden_states)
+    encoder_hidden_states = encoder_hidden_states.view(
+        batch_size, -1, hidden_states.size(-1)
+    )
+
+    for block in transformer.transformer_blocks:
+        hidden_states = _forward_diffusers_video_block(
+            block=block,
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            temb=temb,
+            video_rotary_emb=video_rotary_emb,
+            encoder_attention_mask=encoder_attention_mask,
+            n_context_tokens=n_context_tokens,
+        )
+
+    scale_shift_values = (
+        transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None]
+    )
+    shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+    hidden_states = transformer.norm_out(hidden_states)
+    hidden_states = hidden_states * (1 + scale) + shift
+    return transformer.proj_out(hidden_states)
+
+
+def _forward_diffusers_video_block(
+    *,
+    block: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    temb: torch.Tensor,
+    video_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    encoder_attention_mask: torch.Tensor | None,
+    n_context_tokens: int,
+) -> torch.Tensor:
+    batch_size = hidden_states.size(0)
+
+    norm_hidden_states = block.norm1(hidden_states)
+    num_ada_params = block.scale_shift_table.shape[0]
+    ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape(
+        batch_size, temb.size(1), num_ada_params, -1
+    )
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(
+        dim=2
+    )
+    norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+
+    attn_hidden_states = _streaming_diffusers_self_attention(
+        attn=block.attn1,
+        hidden_states=norm_hidden_states,
+        query_rotary_emb=video_rotary_emb,
+        n_context_tokens=n_context_tokens,
+    )
+    hidden_states = hidden_states + attn_hidden_states * gate_msa
+
+    norm_hidden_states = block.norm2(hidden_states)
+    attn_hidden_states = block.attn2(
+        norm_hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        query_rotary_emb=None,
+        attention_mask=encoder_attention_mask,
+    )
+    hidden_states = hidden_states + attn_hidden_states
+
+    norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp
+    hidden_states = hidden_states + block.ff(norm_hidden_states) * gate_mlp
+    return hidden_states
+
+
+def _streaming_diffusers_self_attention(
+    *,
+    attn: nn.Module,
+    hidden_states: torch.Tensor,
+    query_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    n_context_tokens: int,
+) -> torch.Tensor:
+    """SANA-WM sink/current streaming mask using Diffusers LTX-2 attention."""
+
+    sequence_length = hidden_states.shape[1]
+    if n_context_tokens <= 0 or n_context_tokens >= sequence_length:
+        return attn(
+            hidden_states=hidden_states,
+            encoder_hidden_states=None,
+            query_rotary_emb=query_rotary_emb,
+        )
+
+    from diffusers.models.attention_dispatch import dispatch_attention_fn
+    from diffusers.models.transformers.transformer_ltx2 import (
+        apply_interleaved_rotary_emb,
+        apply_split_rotary_emb,
+    )
+
+    # Diffusers 0.38+ always defines `to_gate_logits`, while 0.37 only has it
+    # on gated variants. The public SANA-WM refiner config is ungated, so a
+    # missing attribute means the same thing as `None`.
+    to_gate_logits = getattr(attn, "to_gate_logits", None)
+    gate_logits = to_gate_logits(hidden_states) if to_gate_logits is not None else None
+
+    query = attn.to_q(hidden_states)
+    key = attn.to_k(hidden_states)
+    value = attn.to_v(hidden_states)
+
+    query = attn.norm_q(query)
+    key = attn.norm_k(key)
+
+    if attn.rope_type == "interleaved":
+        query = apply_interleaved_rotary_emb(query, query_rotary_emb)
+        key = apply_interleaved_rotary_emb(key, query_rotary_emb)
+    elif attn.rope_type == "split":
+        query = apply_split_rotary_emb(query, query_rotary_emb)
+        key = apply_split_rotary_emb(key, query_rotary_emb)
+    else:
+        raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}")
+
+    query = query.unflatten(2, (attn.heads, -1))
+    key = key.unflatten(2, (attn.heads, -1))
+    value = value.unflatten(2, (attn.heads, -1))
+
+    processor = attn.processor
+    backend = getattr(processor, "_attention_backend", None)
+    parallel_config = getattr(processor, "_parallel_config", None)
+    context_hidden_states = dispatch_attention_fn(
+        query[:, :n_context_tokens],
+        key[:, :n_context_tokens],
+        value[:, :n_context_tokens],
+        attn_mask=None,
+        dropout_p=0.0,
+        is_causal=False,
+        backend=backend,
+        parallel_config=parallel_config,
+    )
+    current_hidden_states = dispatch_attention_fn(
+        query[:, n_context_tokens:],
+        key,
+        value,
+        attn_mask=None,
+        dropout_p=0.0,
+        is_causal=False,
+        backend=backend,
+        parallel_config=parallel_config,
+    )
+
+    hidden_states = torch.cat([context_hidden_states, current_hidden_states], dim=1)
+    hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+    if gate_logits is not None:
+        hidden_states = hidden_states.unflatten(2, (attn.heads, -1))
+        gates = 2.0 * torch.sigmoid(gate_logits)
+        hidden_states = hidden_states * gates.unsqueeze(-1)
+        hidden_states = hidden_states.flatten(2, 3)
+
+    hidden_states = attn.to_out[0](hidden_states)
+    hidden_states = attn.to_out[1](hidden_states)
+    return hidden_states
+
+
+class SanaWMLTX2RefinerStage(PipelineStage):
+    """Run the SANA-WM stage-2 LTX-2 refiner before VAE decode.
+
+    Modules are injected by `SanaWMTwoStagePipeline`:
+      * `transformer` (`SanaWMLTX2VideoRefiner`) -- video-only LTX-2 forward
+      * `connectors` (`LTX2TextConnectors`)
+      * `text_encoder` (`Gemma3ForConditionalGeneration`)
+      * `tokenizer` (HF `AutoTokenizer` for the refiner Gemma-3)
+    """
+
+    def __init__(
+        self,
+        *,
+        transformer: nn.Module,
+        connectors: nn.Module,
+        text_encoder: nn.Module,
+        tokenizer: Any,
+        dtype: torch.dtype,
+        text_max_sequence_length: int = _REFINER_TEXT_MAX_LENGTH,
+    ) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.connectors = connectors
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.dtype = dtype
+        self.text_max_sequence_length = int(text_max_sequence_length)
+
+    @property
+    def role_affinity(self) -> RoleType:
+        return RoleType.DENOISER
+
+    @property
+    def parallelism_type(self) -> StageParallelismType:
+        if getattr(self.server_args, "enable_cfg_parallel", False):
+            return StageParallelismType.MAIN_RANK_ONLY
+        return StageParallelismType.REPLICATED
+
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        if sana_wm_skip_refiner_enabled():
+            return []
+        if (
+            getattr(server_args, "enable_cfg_parallel", False)
+            and not _is_current_cfg_main_rank()
+        ):
+            return []
+
+        # Declare every component this stage forwards through so
+        # ComponentResidencyManager moves them onto GPU before the stage runs.
+        # Without this, `dit_cpu_offload=True` keeps refiner sub-modules on CPU
+        # and the first matmul fails with "mat2 is on cpu" vs cuda inputs.
+        # The tokenizer stays on CPU (no nn.Module weights to ferry).
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name=stage_name,
+                component_name="text_encoder_2",
+                target_dtype=self.dtype,
+            ),
+            ComponentUse(
+                stage_name=stage_name,
+                component_name="connectors",
+                target_dtype=self.dtype,
+            ),
+            ComponentUse(
+                stage_name=stage_name,
+                component_name="transformer_2",
+                target_dtype=self.dtype,
+                memory_intensive=True,
+            ),
+        ]
+
+    @staticmethod
+    def _prompts_for_batch(batch: Req, batch_size: int) -> list[str]:
+        prompt = batch.extra.get("refiner_prompt") if batch.extra else None
+        if prompt is None:
+            prompt = batch.prompt
+        if isinstance(prompt, str):
+            return [prompt] * batch_size
+        if isinstance(prompt, list) and all(isinstance(p, str) for p in prompt):
+            if len(prompt) == batch_size:
+                return prompt
+            if len(prompt) == 1:
+                return prompt * batch_size
+        raise ValueError(
+            "SANA-WM refiner requires a string prompt or one prompt per batch item."
+        )
+
+    @torch.inference_mode()
+    def _encode_prompt(
+        self,
+        prompt: str,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        tokenizer = self.tokenizer
+        if getattr(tokenizer, "padding_side", "right") != "left":
+            tokenizer.padding_side = "left"
+        if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        text_inputs = tokenizer(
+            [prompt.strip()],
+            padding="max_length",
+            max_length=self.text_max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = text_inputs.input_ids.to(device)
+        attention_mask = text_inputs.attention_mask.to(device)
+
+        # Diffusers-backed official path loads HF Gemma3ForConditionalGeneration.
+        # NVlabs encodes through `.model`; the fallback SGLang-native encoder is
+        # still callable directly, so keep both surfaces.
+        with self.use_declared_component(
+            component_name="text_encoder_2", module=self.text_encoder
+        ):
+            text_backbone = getattr(self.text_encoder, "model", self.text_encoder)
+            outputs = text_backbone(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+        per_layer_hidden = getattr(outputs, "hidden_states", None)
+        if per_layer_hidden is None:
+            raise RuntimeError(
+                "SANA-WM refiner text encoder must return per-layer hidden_states."
+            )
+        stacked = torch.stack(per_layer_hidden, dim=-1)  # (B, L, D, n_layers)
+        seq_lengths = attention_mask.sum(dim=-1)
+        log_sana_wm_tensor_stats("refiner.text_hidden_states_stacked", stacked)
+        prompt_embeds = _pack_text_embeds(
+            stacked,
+            seq_lengths,
+            padding_side=tokenizer.padding_side,
+        ).to(dtype=self.dtype)
+        log_sana_wm_tensor_stats("refiner.prompt_embeds_packed", prompt_embeds)
+
+        with self.use_declared_component(
+            component_name="connectors", module=self.connectors
+        ):
+            video_text_embedding, _, video_attention_mask = self.connectors(
+                prompt_embeds, attention_mask
+            )
+        log_sana_wm_tensor_stats("refiner.video_text_embedding", video_text_embedding)
+        log_sana_wm_tensor_stats("refiner.video_attention_mask", video_attention_mask)
+        return (
+            video_text_embedding.to(device=device, dtype=self.dtype),
+            video_attention_mask.to(device=device),
+        )
+
+    def _predict_current_x0(
+        self,
+        *,
+        sink: torch.Tensor,
+        noisy_current: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        prompt_attention_mask: torch.Tensor,
+        sigma: torch.Tensor,
+        fps: float,
+        n_context_tokens: int,
+        step_idx: int,
+    ) -> torch.Tensor:
+        full_latent = torch.cat([sink, noisy_current], dim=2)
+        batch_size, _, num_frames, height, width = full_latent.shape
+        patch_size = int(_refiner_config_value(self.transformer, "patch_size"))
+        patch_size_t = int(_refiner_config_value(self.transformer, "patch_size_t"))
+        latent_tokens = pack_latents(full_latent, patch_size, patch_size_t)
+
+        raw_timestep = torch.zeros(
+            batch_size,
+            latent_tokens.shape[1],
+            1,
+            dtype=torch.float32,
+            device=latent_tokens.device,
+        )
+        raw_timestep[:, n_context_tokens:, 0] = sigma.float()
+
+        with self.use_declared_component(
+            component_name="transformer_2", module=self.transformer
+        ):
+            if _uses_diffusers_ltx2_refiner(self.transformer):
+                model_timestep = raw_timestep.squeeze(-1) * float(
+                    _refiner_config_value(self.transformer, "timestep_scale_multiplier")
+                )
+                velocity_tokens = _forward_diffusers_video_only(
+                    self.transformer,
+                    hidden_states=latent_tokens.to(self.dtype),
+                    encoder_hidden_states=prompt_embeds,
+                    timestep=model_timestep,
+                    encoder_attention_mask=prompt_attention_mask,
+                    num_frames=num_frames,
+                    height=height,
+                    width=width,
+                    fps=fps,
+                    n_context_tokens=n_context_tokens,
+                )
+            else:
+                additive_mask = _as_additive_attention_mask(
+                    prompt_attention_mask, self.dtype
+                )
+                with set_forward_context(
+                    current_timestep=step_idx,
+                    attn_metadata=None,
+                ):
+                    velocity_tokens = self.transformer(
+                        hidden_states=latent_tokens.to(self.dtype),
+                        encoder_hidden_states=prompt_embeds,
+                        timestep=raw_timestep.squeeze(-1),
+                        encoder_attention_mask=additive_mask,
+                        num_frames=num_frames,
+                        height=height,
+                        width=width,
+                        fps=fps,
+                        n_context_tokens=n_context_tokens,
+                    )
+
+        denoised = latent_tokens.float() - velocity_tokens.float() * raw_timestep
+        return denoised[:, n_context_tokens:, :].to(self.dtype)
+
+    @torch.inference_mode()
+    def _refine_one(
+        self,
+        latent: torch.Tensor,
+        prompt: str,
+        *,
+        fps: float,
+        seed: int,
+        sink_size: int = 1,
+    ) -> torch.Tensor:
+        device = get_local_torch_device()
+        z = latent.to(device=device, dtype=self.dtype)
+        if z.shape[2] <= sink_size:
+            raise ValueError(
+                f"Stage-1 latent has {z.shape[2]} frames but sink_size={sink_size}."
+            )
+        self.log_info(
+            "SANA-WM refiner start: latent=%s, fps=%.3f, seed=%d, "
+            "sink_size=%d, sigmas=%s, diagnostics=%s",
+            tuple(z.shape),
+            fps,
+            seed,
+            sink_size,
+            STAGE_2_DISTILLED_SIGMA_VALUES,
+            "on" if sana_wm_diagnostics_enabled() else "off",
+        )
+        log_sana_wm_tensor_stats("refiner.input_latent", z)
+
+        prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt, device)
+
+        sigmas = torch.tensor(
+            STAGE_2_DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device
+        )
+        start_sigma = float(sigmas[0])
+        sink = z[:, :, :sink_size].contiguous()
+        current = z[:, :, sink_size:].contiguous()
+        log_sana_wm_tensor_stats("refiner.sink_latent", sink)
+        log_sana_wm_tensor_stats("refiner.current_latent_clean", current)
+        gen = torch.Generator(device=device).manual_seed(int(seed))
+        eps = torch.randn(current.shape, generator=gen, device=device, dtype=self.dtype)
+        noisy = (1.0 - start_sigma) * current + start_sigma * eps
+        log_sana_wm_tensor_stats("refiner.current_latent_noisy_initial", noisy)
+
+        patch_size = int(_refiner_config_value(self.transformer, "patch_size"))
+        patch_size_t = int(_refiner_config_value(self.transformer, "patch_size_t"))
+
+        sink_tokens = pack_latents(sink, patch_size, patch_size_t)
+        n_context_tokens = sink_tokens.shape[1]
+
+        for step_idx in range(len(sigmas) - 1):
+            sigma = sigmas[step_idx]
+            denoised = self._predict_current_x0(
+                sink=sink,
+                noisy_current=noisy,
+                prompt_embeds=prompt_embeds,
+                prompt_attention_mask=prompt_attention_mask,
+                sigma=sigma,
+                fps=fps,
+                n_context_tokens=n_context_tokens,
+                step_idx=step_idx,
+            )
+            noisy_tokens = pack_latents(noisy, patch_size, patch_size_t)
+            velocity_tokens = (noisy_tokens.float() - denoised.float()) / sigma.float()
+            next_tokens = (
+                noisy_tokens.float()
+                + velocity_tokens * (sigmas[step_idx + 1] - sigma).float()
+            )
+            noisy = unpack_latents(
+                next_tokens.to(self.dtype),
+                num_frames=noisy.shape[2],
+                height=noisy.shape[3],
+                width=noisy.shape[4],
+                patch_size=patch_size,
+                patch_size_t=patch_size_t,
+            )
+            velocity_5d = unpack_latents(
+                velocity_tokens,
+                num_frames=noisy.shape[2],
+                height=noisy.shape[3],
+                width=noisy.shape[4],
+                patch_size=patch_size,
+                patch_size_t=patch_size_t,
+            )
+            log_sana_wm_tensor_stats(
+                f"refiner.step_{step_idx}.velocity_current",
+                velocity_5d.to(self.dtype),
+            )
+            log_sana_wm_tensor_stats(f"refiner.step_{step_idx}.current_latent", noisy)
+
+        refined = torch.cat([sink, noisy], dim=2)
+        log_sana_wm_tensor_stats("refiner.output_latent", refined)
+        return refined
+
+    @torch.inference_mode()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        if batch.latents is None:
+            raise ValueError("SANA-WM refiner requires batch.latents from stage 1.")
+        if batch.latents.ndim != 5:
+            raise ValueError(
+                "SANA-WM refiner expects 5D latents shaped (B, C, T, H, W), "
+                f"got {tuple(batch.latents.shape)}."
+            )
+
+        if sana_wm_skip_refiner_enabled(batch):
+            if batch.extra is None:
+                batch.extra = {}
+            batch.extra["sana_wm_refiner_applied"] = False
+            self.log_info(
+                "SANA-WM LTX-2 refiner skipped by SGLANG_SANA_WM_SKIP_REFINER."
+            )
+            return batch
+
+        batch_size = int(batch.latents.shape[0])
+        prompts = self._prompts_for_batch(batch, batch_size)
+        fps = float(getattr(batch, "fps", 16) or 16)
+
+        seeds: list[int]
+        if batch.seeds is not None and len(batch.seeds) == batch_size:
+            seeds = [int(s) for s in batch.seeds]
+        elif batch.seeds is not None and len(batch.seeds) == 1:
+            seeds = [int(batch.seeds[0])] * batch_size
+        else:
+            seeds = [int(getattr(batch, "seed", 0) or 0)] * batch_size
+
+        refined: list[torch.Tensor] = []
+        for idx, (prompt, seed) in enumerate(zip(prompts, seeds, strict=True)):
+            refined.append(
+                self._refine_one(
+                    batch.latents[idx : idx + 1],
+                    prompt,
+                    fps=fps,
+                    seed=seed,
+                )
+            )
+        batch.latents = torch.cat(refined, dim=0).to(
+            device=batch.latents.device, dtype=batch.latents.dtype
+        )
+        if batch.extra is None:
+            batch.extra = {}
+        batch.extra["sana_wm_refiner_applied"] = True
+        self.log_info("SANA-WM LTX-2 refiner applied to stage-1 latents.")
+        return batch
+
+
+class SanaWMRefinerDecodingStage(SanaWMDecodingStage):
+    """Decode refined latents and drop the clean sink anchor frame."""
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs):
+        self._drop_refiner_sink = bool(
+            (getattr(batch, "extra", None) or {}).get("sana_wm_refiner_applied", True)
+        )
+        try:
+            return super().forward(batch, server_args)
+        finally:
+            self._drop_refiner_sink = True
+
+    @torch.no_grad()
+    def decode(
+        self,
+        latents: torch.Tensor,
+        server_args: ServerArgs,
+        *,
+        vae_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        frames = super().decode(latents, server_args, vae_dtype=vae_dtype)
+        log_sana_wm_tensor_stats("refiner.decode.frames_with_sink", frames)
+        if frames.ndim != 5:
+            raise ValueError(
+                "SANA-WM refiner decoding expects decoded video shaped "
+                f"(B, C, T, H, W), got {tuple(frames.shape)}."
+            )
+        if frames.shape[2] <= 1:
+            raise ValueError(
+                "SANA-WM refiner decoding expected a sink frame plus refined "
+                f"frames, got temporal length {frames.shape[2]}."
+            )
+        if not getattr(self, "_drop_refiner_sink", True):
+            log_sana_wm_tensor_stats("refiner.decode.frames_output", frames)
+            return frames
+        # Match NVlabs `inference_sana_wm.py`: decode with the clean sink anchor,
+        # then drop the first frame from the returned video.
+        frames = frames[:, :, 1:].contiguous()
+        log_sana_wm_tensor_stats("refiner.decode.frames_output", frames)
+        return frames

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SANA-WM self-forcing chunk SAMPLER utilities (not a diffusers scheduler).
+
+Owns the autoregressive chunk grid, per-block KV-cache accumulation /
+eviction, and the explicit per-chunk sigma list for streaming (self-forcing)
+SANA-WM generation. The actual stepping scheduler is the shared
+``FlowMatchEulerDiscreteScheduler`` (``per_token_timesteps`` path).
+
+Why not ``SelfForcingFlowMatchScheduler`` (the LingBot causal-DMD scheduler):
+
+  * SANA-WM's distilled sigma grid is an explicit NON-uniform list
+    ((1000, 960, 889, 727, 0)/1000); its ``set_timesteps`` only expresses
+    ``linspace(sigma_max -> sigma_min)`` + shift.
+  * SANA-WM pins the condition frame at timestep 0 INSIDE chunk 0 (per-frame
+    timesteps within one step); its ``step`` applies a single per-sample
+    sigma. LingBot instead warms the cond frames into the KV cache, so a
+    scalar sigma suffices there.
+
+The model-calling denoise loop stays in the stage.
+
+  * segmentation FRONT-LOADS the remainder into chunk 0;
+  * the KV accumulator concatenates softmax cache slots on **dim=1** (our
+    ``(B, N, H, D)`` softmax cache layout);
+  * GDN/STATE blocks copy-forward the previous chunk's recurrent state.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import (
+    _NUM_STREAM_CACHE_SLOTS,
+    _SLOT_CAM_K,
+    _SLOT_CAM_V,
+    _SLOT_FFN_TCONV,
+    _SLOT_K,
+    _SLOT_SHORTCONV,
+    _SLOT_TYPE_FLAG,
+    _SLOT_V,
+)
+
+
+@dataclass
+class SanaWMSelfForcingSamplerConfig:
+    """Streaming self-forcing knobs (defaults mirror ``SanaWMPipelineConfig``)."""
+
+    num_frame_per_block: int = 3
+    num_cached_blocks: int = 2
+    sink_token: bool = True
+    denoising_step_list: tuple = (1000, 960, 889, 727, 0)
+    streaming_cfg_scale: float = 1.0
+
+    @classmethod
+    def from_pipeline_config(cls, pcfg) -> "SanaWMSelfForcingSamplerConfig":
+        """Read the streaming knobs off a pipeline config (note the ``or 1.0`` cfg-scale guard)."""
+        return cls(
+            num_frame_per_block=int(getattr(pcfg, "num_frame_per_block", 3)),
+            num_cached_blocks=int(getattr(pcfg, "num_cached_blocks", 2)),
+            sink_token=bool(getattr(pcfg, "sink_token", True)),
+            denoising_step_list=tuple(
+                getattr(pcfg, "denoising_step_list", (1000, 960, 889, 727, 0))
+            ),
+            streaming_cfg_scale=float(getattr(pcfg, "streaming_cfg_scale", 1.0) or 1.0),
+        )
+
+
+class SanaWMSelfForcingSampler:
+    """Self-forcing chunk scheduler: segmentation + per-block KV-cache carry.
+
+    Stateless; methods are static (every input is passed explicitly). A
+    ``config`` may be attached for callers that prefer the instance form.
+    """
+
+    def __init__(self, config: SanaWMSelfForcingSamplerConfig | None = None):
+        self.config = config or SanaWMSelfForcingSamplerConfig()
+
+    # ----------------------------------------------------------------- #
+    # Chunk schedule + KV cache accumulation
+    # ----------------------------------------------------------------- #
+    @staticmethod
+    def create_autoregressive_segments(
+        total_frames: int, num_frame_per_block: int
+    ) -> list[int]:
+        base = int(num_frame_per_block)
+        remained = total_frames % base
+        num_chunks = total_frames // base
+        chunk_indices = [0]
+        for idx in range(num_chunks):
+            cur = chunk_indices[-1] + base + (remained if idx == 0 else 0)
+            chunk_indices.append(cur)
+        return chunk_indices
+
+    @staticmethod
+    def accumulate_kv_cache(
+        kv_cache: list,
+        chunk_idx: int,
+        chunk_indices: list[int],
+        num_cached_blocks: int,
+        sink_token: bool,
+        num_blocks: int,
+    ) -> tuple[list, int]:
+        """Build chunk ``chunk_idx``'s read-only KV prefix from prior chunks.
+
+        GDN/STATE blocks (type flag > 0.5) copy-forward the PREVIOUS chunk's
+        recurrent state; softmax/CONCAT blocks concatenate the rolling-window +
+        sink K/V along **dim=1** (token axis of our (B,N,H,D) softmax cache)."""
+        if chunk_idx == 0:
+            return kv_cache[0], 0
+
+        cur = kv_cache[chunk_idx]
+        start_chunk = (
+            max(chunk_idx - num_cached_blocks, 0) if num_cached_blocks > 0 else 0
+        )
+        valid = list(range(start_chunk, chunk_idx))
+        sink_num = 0
+        if sink_token and num_cached_blocks > 0:
+            sink_start = max(chunk_idx - num_cached_blocks + 1, 0)
+            if sink_start > 0:
+                valid = [0] + list(range(sink_start, chunk_idx))
+                sink_num = chunk_indices[1] - chunk_indices[0]
+
+        for block_id in range(num_blocks):
+            prev_last = kv_cache[chunk_idx - 1][block_id]
+            type_flag = prev_last[_SLOT_TYPE_FLAG]
+            if type_flag is not None and float(type_flag.item()) > 0.5:
+                # STATE (GDN) block: carry the previous chunk's recurrent state.
+                cur[block_id] = [
+                    prev_last[_SLOT_K],
+                    prev_last[_SLOT_V],
+                    prev_last[_SLOT_CAM_K],
+                    prev_last[_SLOT_CAM_V],
+                    prev_last[_SLOT_SHORTCONV],
+                    None,
+                    prev_last[_SLOT_TYPE_FLAG],
+                    None,
+                    None,
+                    prev_last[_SLOT_FFN_TCONV],
+                ]
+                continue
+
+            # CONCAT (softmax) block: concat cached K/V over the valid window.
+            acc: list[torch.Tensor | None] = [None] * _NUM_STREAM_CACHE_SLOTS
+            for idx in valid:
+                prev = kv_cache[idx][block_id]
+                if prev[_SLOT_K] is None:
+                    continue
+                for slot in (_SLOT_K, _SLOT_V, _SLOT_CAM_K, _SLOT_CAM_V):
+                    if prev[slot] is None:
+                        continue
+                    acc[slot] = (
+                        prev[slot].clone()
+                        if acc[slot] is None
+                        else torch.cat(
+                            [acc[slot], prev[slot]], dim=1
+                        )  # (B,N,H,D) token axis
+                    )
+            cur[block_id] = [
+                acc[_SLOT_K],
+                acc[_SLOT_V],
+                acc[_SLOT_CAM_K],
+                acc[_SLOT_CAM_V],
+                prev_last[_SLOT_SHORTCONV],
+                None,
+                prev_last[_SLOT_TYPE_FLAG],
+                None,
+                None,
+                prev_last[_SLOT_FFN_TCONV],
+            ]
+
+        SanaWMSelfForcingSampler.evict_stale_kv_cache(
+            kv_cache, chunk_idx, valid, num_cached_blocks, num_blocks
+        )
+        return cur, sink_num
+
+    @staticmethod
+    def evict_stale_kv_cache(
+        kv_cache: list,
+        chunk_idx: int,
+        valid: list[int],
+        num_cached_blocks: int,
+        num_blocks: int,
+    ) -> None:
+        if num_cached_blocks <= 0:
+            return
+        keep = set(valid)
+        keep.add(chunk_idx)
+        for stale in range(chunk_idx):
+            if stale in keep:
+                continue
+            kv_cache[stale] = [
+                [None] * _NUM_STREAM_CACHE_SLOTS for _ in range(num_blocks)
+            ]
+
+    # ----------------------------------------------------------------- #
+    # Per-chunk flow-Euler sigma schedule
+    # ----------------------------------------------------------------- #
+    @staticmethod
+    def build_per_chunk_sigmas(denoising_step_list) -> list[float]:
+        """Explicit flow-Euler sigmas for one chunk's short self-forcing schedule.
+
+        The ``denoising_step_list`` (e.g. (1000, 960, 889, 727, 0)) must end with
+        0; sigmas are the non-terminal steps divided by 1000."""
+        schedule = list(denoising_step_list)
+        if len(schedule) < 2 or schedule[-1] != 0:
+            raise ValueError(f"denoising_step_list must end with 0, got {schedule}")
+        return [float(t) / 1000.0 for t in schedule[:-1]]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py

@@ -0,0 +1,786 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SANA-WM streaming (self-forcing) denoising stage — S1c.
+
+Generates video chunk-by-chunk via the DiT's chunk-causal ``forward_long``, a
+rolling per-block KV cache carrying recurrent GDN state + softmax K/V window.
+
+Gotchas vs the dense ``SanaWMDenoisingStage`` (one-shot bidirectional):
+  * fresh ``FlowMatchEulerDiscreteScheduler(shift=1.0)`` with explicit sigmas
+    (NOT the request scheduler, whose shift would warp the sigmas);
+  * KV accumulator concats softmax cache slots on **dim=1** (our
+    ``(B, N, H, D)`` softmax cache layout), not the reference's dim=2;
+  * the condition frame (frame 0) is re-pinned every step in chunk 0.
+"""
+
+from __future__ import annotations
+
+# --- debug parity harness (gated by env; no-op in production) ---
+import os as _os
+import time
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+    ComponentUse,
+)
+from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
+    _NUM_STREAM_CACHE_SLOTS,
+)
+from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import (
+    CausalDMDDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.self_forcing import (
+    SanaWMSelfForcingSampler,
+    SanaWMSelfForcingSamplerConfig,
+)
+from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+from . import parity_probe
+from .base import (
+    _align_sana_wm_cfg_text_conditions,
+    _cat_optional_tensors,
+    _first_tensor,
+    _to_device_dtype,
+    log_sana_wm_tensor_stats,
+)
+
+_SANAWM_INJECT_DIR = _os.environ.get(parity_probe.ENV_INJECT)
+
+
+def self_forcing_denoise_chunk(
+    *,
+    transformer,
+    scheduler,
+    sigmas,
+    get_chunk,
+    set_chunk,
+    cond_mask: torch.Tensor,
+    embeds: torch.Tensor,
+    mask,
+    camera_conditions,
+    chunk_plucker,
+    chunk_kv,
+    start_f: int,
+    end_f: int,
+    frame_index,
+    do_cfg: bool,
+    cfg_scale: float,
+    target_dtype: torch.dtype,
+    device,
+    forward_ctx=None,
+    on_noise_pred=None,
+) -> list:
+    """Self-forcing denoise of ONE chunk + the clean t=0 KV pass; returns the chunk's updated per-block KV cache.
+
+    SINGLE parity-locked implementation shared by the offline and realtime paths
+    (it used to be copy-pasted — the same drift failure mode as the four parity bugs).
+
+    ``get_chunk``/``set_chunk`` are caller closures so each path keeps its EXACT
+    tensor lifecycle (kernel-visible layouts, and therefore bitwise behavior,
+    unchanged); ``set_chunk`` owns the dtype cast and cond-frame re-pinning.
+    ``forward_ctx`` (optional) maps ``int timestep -> context manager`` (batch
+    wraps forward_long in set_forward_context; realtime does not).
+    """
+    ctx = forward_ctx if forward_ctx is not None else (lambda _ts: nullcontext())
+
+    scheduler.set_timesteps(sigmas=sigmas, device=device)
+    for t in scheduler.timesteps:
+        chunk_lat = get_chunk()
+        B, C = chunk_lat.shape[0], chunk_lat.shape[1]
+        lat_in = torch.cat([chunk_lat, chunk_lat], dim=0) if do_cfg else chunk_lat
+        ts_tensor = (1.0 - cond_mask) * t.to(device=device, dtype=torch.float32).view(
+            1, 1, 1, 1, 1
+        )
+        ts_in = torch.cat([ts_tensor, ts_tensor], dim=0) if do_cfg else ts_tensor
+        model_ts = ts_in[:, :1, :, 0, 0]  # (B|2B, chunk_frames)
+
+        with ctx(int(t.item()) if t.ndim == 0 else 0):
+            noise_pred, _ = transformer.forward_long(
+                hidden_states=lat_in.to(target_dtype),
+                encoder_hidden_states=embeds,
+                timestep=model_ts,
+                encoder_attention_mask=mask,
+                camera_conditions=camera_conditions,
+                chunk_plucker=chunk_plucker,
+                kv_cache=chunk_kv,
+                save_kv_cache=False,
+                start_f=start_f,
+                end_f=end_f,
+                frame_index=frame_index,
+            )
+        if do_cfg:
+            noise_uncond, noise_text = noise_pred.chunk(2)
+            noise_pred = noise_uncond + cfg_scale * (noise_text - noise_uncond)
+
+        if on_noise_pred is not None:
+            on_noise_pred(int(t.item()), noise_pred)
+
+        denoised = scheduler.step(
+            -noise_pred.reshape(B, C, -1).transpose(1, 2),
+            t,
+            chunk_lat.reshape(B, C, -1).transpose(1, 2),
+            per_token_timesteps=ts_tensor.reshape(B, C, -1)[:, 0],
+            return_dict=False,
+        )[0]
+        set_chunk(denoised.transpose(1, 2).reshape(chunk_lat.shape))
+
+    # Clean pass (t=0) to write this chunk's KV for the next chunk.
+    chunk_lat = get_chunk()
+    lat_in = torch.cat([chunk_lat, chunk_lat], dim=0) if do_cfg else chunk_lat
+    ts_zero = torch.zeros(
+        lat_in.shape[0], 1, chunk_lat.shape[2], device=device, dtype=torch.float32
+    )
+    with ctx(0):
+        _, updated_cache = transformer.forward_long(
+            hidden_states=lat_in.to(target_dtype),
+            encoder_hidden_states=embeds,
+            timestep=ts_zero,
+            encoder_attention_mask=mask,
+            camera_conditions=camera_conditions,
+            chunk_plucker=chunk_plucker,
+            kv_cache=chunk_kv,
+            save_kv_cache=True,
+            start_f=start_f,
+            end_f=end_f,
+            frame_index=frame_index,
+        )
+    return updated_cache
+
+
+class SanaWMStreamCacheState(RealtimeCausalDiTState):
+    """Per-session streaming DiT state, framework-pattern (cf. LingBot).
+
+    The Wan-shaped ``kv_cache`` field stays None — SANA-WM's cache is the
+    heterogeneous per-block 10-slot list (GDN recurrent matrix states + softmax
+    concat windows + conv tails), carried in the fields below."""
+
+    def __init__(self):
+        super().__init__()
+        # kv[chunk][block][slot] — grown per chunk, stale chunks evicted.
+        self.stream_kv_cache: list = []
+        self.chunk_indices: list[int] = [0]
+        # Growing stage-1 latent buffer (cond frame + denoised chunks).
+        self.latents: torch.Tensor | None = None
+        # Per-session flow-Euler scheduler (fresh shift=1.0; persists across ticks).
+        self.scheduler: FlowMatchEulerDiscreteScheduler | None = None
+
+    def dispose(self) -> None:
+        super().dispose()
+        self.stream_kv_cache = []
+        self.chunk_indices = [0]
+        self.latents = None
+        self.scheduler = None
+
+
+class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage):
+    """Autoregressive self-forcing streaming denoise — SANA-WM's causal-DMD variant.
+
+    Same family as ``LingBotWorldCausalDMDDenoisingStage`` (offline whole-clip vs
+    realtime per-chunk). Differences from the Wan-style base are the cache contract
+    (heterogeneous 10-slot list instead of preallocated ``CausalSelfAttentionKVCache``),
+    the front-loaded chunk grid (chunk 0 carries the remainder), and in-chunk
+    condition-frame pinning instead of KV warm-up — hence the cache-management overrides.
+    """
+
+    def __init__(
+        self, transformer, scheduler=None, *, keep_resident: bool = False
+    ) -> None:
+        # Skip CausalDMDDenoisingStage.__init__: it reads Wan-specific arch
+        # fields (num_frames_per_block / sliding_window_num_frames / sink_size)
+        # that SANA-WM defines per-run via the pipeline config instead.
+        DenoisingStage.__init__(self, transformer, scheduler)
+        self.num_transformer_blocks = len(transformer.blocks)
+        # Realtime pipelines keep the DiT device-resident for the session's
+        # lifetime (the per-tick offload round-trip would dominate latency);
+        # the offline pipeline keeps the default offload behavior.
+        self._keep_resident = bool(keep_resident)
+
+    def component_uses(self, server_args: ServerArgs, stage_name: str | None = None):
+        if not self._keep_resident:
+            return super().component_uses(server_args, stage_name)
+
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name,
+                "transformer",
+                target_dtype=PRECISION_TO_TYPE[
+                    server_args.pipeline_config.dit_precision
+                ],
+                memory_intensive=True,
+                keep_ready_after_warmup=True,
+            ),
+        ]
+
+    # Chunk schedule + KV cache: delegate to SanaWMSelfForcingSampler.
+    @staticmethod
+    def _autoregressive_segments(
+        total_frames: int, num_frame_per_block: int
+    ) -> list[int]:
+        return SanaWMSelfForcingSampler.create_autoregressive_segments(
+            total_frames, num_frame_per_block
+        )
+
+    @staticmethod
+    def _accumulate_kv_cache(
+        kv_cache: list,
+        chunk_idx: int,
+        chunk_indices: list[int],
+        num_cached_blocks: int,
+        sink_token: bool,
+        num_blocks: int,
+    ) -> tuple[list, int]:
+        return SanaWMSelfForcingSampler.accumulate_kv_cache(
+            kv_cache,
+            chunk_idx,
+            chunk_indices,
+            num_cached_blocks,
+            sink_token,
+            num_blocks,
+        )
+
+    @staticmethod
+    def _evict_stale_kv_cache(
+        kv_cache: list,
+        chunk_idx: int,
+        valid: list[int],
+        num_cached_blocks: int,
+        num_blocks: int,
+    ) -> None:
+        SanaWMSelfForcingSampler.evict_stale_kv_cache(
+            kv_cache, chunk_idx, valid, num_cached_blocks, num_blocks
+        )
+
+    # Realtime per-chunk path (sessions): per-session state in SanaWMStreamCacheState.
+    @torch.no_grad()
+    def _forward_realtime_chunk(self, batch: Req, server_args: ServerArgs) -> Req:
+        if batch.latents is None or batch.latents.ndim != 5:
+            raise ValueError(
+                "SANA-WM realtime denoising expects this tick's pre-noised chunk "
+                "latents (B, C, n, H, W) from the latent-preparation stage."
+            )
+        pcfg = server_args.pipeline_config
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE.get(
+            getattr(pcfg, "dit_precision", "bf16"), torch.bfloat16
+        )
+        state = batch.session.get_or_create_state(SanaWMStreamCacheState)
+        if batch.block_idx == 0 and state.latents is not None:
+            state.dispose()  # session restart on chunk 0 (mirrors the base stage)
+
+        sc = self._resolve_stream_conditioning(
+            batch, server_args, device=device, target_dtype=target_dtype
+        )
+        sampler_cfg = sc.sampler_cfg
+        incoming = batch.latents.to(device=device, dtype=target_dtype).clone()
+        plan = list(batch.extra.get("sana_wm_chunk_plan") or [incoming.shape[2]])
+        if sum(plan) != incoming.shape[2]:
+            raise ValueError(
+                f"chunk plan {plan} does not cover the incoming {incoming.shape[2]} frames"
+            )
+        if state.scheduler is None:
+            state.scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0)
+
+        # Device-only move, NO dtype cast: Module.to(dtype=...) would cast the
+        # DiT's complex RoPE buffers to real, discarding the imaginary part
+        # (parity root cause #1). No use_declared_component round-trip either —
+        # the DiT stays device-resident for the session's lifetime.
+        transformer = self.transformer.to(device=device).eval()
+        num_blocks = len(transformer.blocks)
+        _dump_dir = parity_probe.probe_dir(parity_probe.ENV_RT_DUMP)
+        if _dump_dir and state.chunk_idx == 0:  # parity harness
+            parity_probe.dump_tensor(_dump_dir, "cond_embeds", sc.embeds)
+            parity_probe.dump_tensor(_dump_dir, "cond_mask", sc.mask)
+            parity_probe.dump_obj(
+                _dump_dir,
+                "dit_fingerprint",
+                parity_probe.weights_fingerprint(transformer),
+            )
+
+        offset = 0
+        for n in plan:
+            chunk_lat = incoming[:, :, offset : offset + n]
+            offset += n
+            chunk_idx = state.chunk_idx
+            # chunk 0's incoming includes the conditioning frame at index 0.
+            start_f = state.chunk_indices[-1] if chunk_idx > 0 else 0
+            end_f = (start_f + n) if chunk_idx > 0 else n
+            state.chunk_indices.append(end_f)
+            state.stream_kv_cache.append(
+                [[None] * _NUM_STREAM_CACHE_SLOTS for _ in range(num_blocks)]
+            )
+            chunk_kv, sink_num = self._accumulate_kv_cache(
+                state.stream_kv_cache,
+                chunk_idx,
+                state.chunk_indices,
+                sampler_cfg.num_cached_blocks,
+                sampler_cfg.sink_token,
+                num_blocks,
+            )
+            # Evict entries outside the accumulate window (sink + last
+            # num_cached_blocks) — unbounded sessions leak concat K/V otherwise.
+            if chunk_idx > 0 and sampler_cfg.num_cached_blocks > 0:
+                start_chunk = max(chunk_idx - sampler_cfg.num_cached_blocks, 0)
+                valid = list(range(start_chunk, chunk_idx))
+                if sampler_cfg.sink_token:
+                    sink_start = max(chunk_idx - sampler_cfg.num_cached_blocks + 1, 0)
+                    if sink_start > 0:
+                        valid = [0] + list(range(sink_start, chunk_idx))
+                self._evict_stale_kv_cache(
+                    state.stream_kv_cache,
+                    chunk_idx,
+                    valid,
+                    sampler_cfg.num_cached_blocks,
+                    num_blocks,
+                )
+            if _dump_dir:  # parity harness
+                parity_probe.dump_obj(
+                    _dump_dir,
+                    f"kv_probe_{chunk_idx:03d}",
+                    parity_probe.kv_cache_checksums(chunk_kv, sink_num),
+                )
+                parity_probe.dump_tensor(
+                    _dump_dir, f"cond_camera_{chunk_idx:03d}", sc.camera
+                )
+                parity_probe.dump_tensor(
+                    _dump_dir, f"cond_plucker_{chunk_idx:03d}", sc.plucker
+                )
+            frame_index = (
+                torch.arange(start_f, end_f, device=device, dtype=torch.long)
+                if sink_num > 0
+                else None
+            )
+
+            B, C = chunk_lat.shape[0], chunk_lat.shape[1]
+            cond_mask = torch.zeros(
+                B,
+                C,
+                end_f - start_f,
+                *chunk_lat.shape[3:],
+                device=device,
+                dtype=torch.float32,
+            )
+            cond_local = []
+            if chunk_idx == 0:
+                cond_mask[:, :, 0] = 1.0
+                cond_local = [0]
+            init_chunk = chunk_lat.clone()
+            _local = {"lat": chunk_lat}
+
+            def _get_chunk(_local=_local):
+                return _local["lat"]
+
+            def _set_chunk(
+                denoised, _local=_local, init_chunk=init_chunk, cond_local=cond_local
+            ):
+                lat = denoised.to(target_dtype)
+                for loc in cond_local:
+                    lat[:, :, loc] = init_chunk[:, :, loc]
+                _local["lat"] = lat
+
+            def _on_noise_pred(ts_int, noise_pred, chunk_idx=chunk_idx):
+                if _dump_dir and chunk_idx == 0:  # parity harness
+                    parity_probe.dump_tensor(
+                        _dump_dir, f"noise_pred_c0_t{ts_int}", noise_pred
+                    )
+
+            # Same forward context the mega-stage held across the tick (the
+            # per-chunk denoise core relies on an ambient context here; the
+            # offline path supplies per-step contexts instead).
+            with set_forward_context(
+                current_timestep=batch.block_idx,
+                attn_metadata=None,
+                forward_batch=batch,
+            ):
+                state.stream_kv_cache[chunk_idx] = self_forcing_denoise_chunk(
+                    transformer=transformer,
+                    scheduler=state.scheduler,
+                    sigmas=sc.explicit_sigmas,
+                    get_chunk=_get_chunk,
+                    set_chunk=_set_chunk,
+                    cond_mask=cond_mask,
+                    embeds=sc.embeds,
+                    mask=sc.mask,
+                    camera_conditions=sc.camera,
+                    chunk_plucker=sc.plucker,
+                    chunk_kv=chunk_kv,
+                    start_f=start_f,
+                    end_f=end_f,
+                    frame_index=frame_index,
+                    do_cfg=sc.do_cfg,
+                    cfg_scale=sc.cfg_scale,
+                    target_dtype=target_dtype,
+                    device=device,
+                    on_noise_pred=_on_noise_pred,
+                )
+            chunk_lat = _local["lat"]
+            state.latents = (
+                chunk_lat
+                if state.latents is None
+                else torch.cat([state.latents, chunk_lat], dim=2)
+            )
+            parity_probe.dump_tensor(  # parity harness
+                _dump_dir, f"stage1_{chunk_idx:03d}_{start_f}_{end_f}", chunk_lat
+            )
+            state.chunk_idx += 1
+            state.current_chunk_start_frame = end_f
+
+        # Downstream chain stages (refiner/decode) consume the growing buffer.
+        batch.latents = state.latents
+        return batch
+
+    def _resolve_stream_conditioning(
+        self,
+        batch: Req,
+        server_args: ServerArgs,
+        *,
+        device: torch.device,
+        target_dtype: torch.dtype,
+        iload=None,
+    ):
+        """Resolve sampler config + text/camera conditioning (shared by the offline loop and realtime path)."""
+        pcfg = server_args.pipeline_config
+        sampler_cfg = SanaWMSelfForcingSamplerConfig.from_pipeline_config(pcfg)
+        explicit_sigmas = SanaWMSelfForcingSampler.build_per_chunk_sigmas(
+            sampler_cfg.denoising_step_list
+        )
+
+        # Streaming uses its OWN cfg scale (official StreamingGenerationConfig.cfg_scale=1.0
+        # => no CFG on the distilled 4-step model). The general guidance_scale (e.g. 4.5)
+        # is for the dense path; using it here ran CFG=4.5 vs the reference's none.
+        cfg_scale = sampler_cfg.streaming_cfg_scale
+        do_cfg = bool(batch.do_classifier_free_guidance) and cfg_scale > 1.0
+        if server_args.enable_cfg_parallel and do_cfg:
+            raise NotImplementedError(
+                "SANA-WM streaming does not support CFG parallel; run replicated."
+            )
+
+        # --- text conditioning ---
+        pos_embeds = _to_device_dtype(
+            _first_tensor(pcfg.get_pos_prompt_embeds(batch)),
+            device=device,
+            dtype=target_dtype,
+        )
+        pos_mask = _to_device_dtype(
+            _first_tensor(batch.prompt_attention_mask), device=device
+        )
+        if pos_embeds is None:
+            raise ValueError("SANA-WM streaming requires positive prompt embeds.")
+        neg_embeds = neg_mask = None
+        if do_cfg:
+            neg_embeds = _to_device_dtype(
+                _first_tensor(pcfg.get_neg_prompt_embeds(batch)),
+                device=device,
+                dtype=target_dtype,
+            )
+            neg_mask = _to_device_dtype(
+                _first_tensor(batch.negative_attention_mask), device=device
+            )
+            if neg_embeds is None:
+                raise ValueError(
+                    "SANA-WM streaming CFG requires negative prompt embeds."
+                )
+            pos_embeds, neg_embeds, pos_mask, neg_mask = (
+                _align_sana_wm_cfg_text_conditions(
+                    pos_embeds, neg_embeds, pos_mask, neg_mask
+                )
+            )
+        embeds_in = torch.cat([neg_embeds, pos_embeds], dim=0) if do_cfg else pos_embeds
+        mask_in = _cat_optional_tensors(neg_mask, pos_mask) if do_cfg else pos_mask
+        if _SANAWM_INJECT_DIR and iload is not None and not do_cfg:
+            _cond = iload("cond").to(device=device, dtype=target_dtype)
+            while _cond.dim() > embeds_in.dim():
+                _cond = _cond.squeeze(1)
+            embeds_in = _cond
+            mask_in = iload("cond_mask").to(device=device)
+
+        # --- camera / plücker (FULL length; forward_long windows internally) ---
+        extra = batch.extra or {}
+        camera_conditions = _to_device_dtype(
+            extra.get("camera_conditions"), device=device, dtype=target_dtype
+        )
+        chunk_plucker = _to_device_dtype(
+            extra.get("chunk_plucker"), device=device, dtype=target_dtype
+        )
+        cam_in = (
+            torch.cat([camera_conditions, camera_conditions], dim=0)
+            if do_cfg and camera_conditions is not None
+            else camera_conditions
+        )
+        plk_in = (
+            torch.cat([chunk_plucker, chunk_plucker], dim=0)
+            if do_cfg and chunk_plucker is not None
+            else chunk_plucker
+        )
+        if _SANAWM_INJECT_DIR and iload is not None:
+            cam_in = iload("raymap").to(device=device, dtype=target_dtype)
+            plk_in = iload("chunk_plucker").to(device=device, dtype=target_dtype)
+            if do_cfg:
+                cam_in = torch.cat([cam_in, cam_in], dim=0)
+                plk_in = torch.cat([plk_in, plk_in], dim=0)
+
+        return SimpleNamespace(
+            sampler_cfg=sampler_cfg,
+            explicit_sigmas=explicit_sigmas,
+            cfg_scale=cfg_scale,
+            do_cfg=do_cfg,
+            embeds=embeds_in,
+            mask=mask_in,
+            camera=cam_in,
+            plucker=plk_in,
+        )
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        # LingBot-style dispatch: realtime sessions denoise ONE chunk per call
+        # with per-session state; otherwise run the whole clip offline.
+        if batch.session is not None:
+            return self._forward_realtime_chunk(batch, server_args)
+        return self._forward_offline(batch, server_args)
+
+    @torch.no_grad()
+    def _forward_offline(self, batch: Req, server_args: ServerArgs) -> Req:
+        if batch.latents is None or batch.latents.ndim != 5:
+            raise ValueError(
+                "SANA-WM streaming denoising expects 5D latents (B, C, T, H, W)."
+            )
+
+        pcfg = server_args.pipeline_config
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE.get(
+            getattr(pcfg, "dit_precision", "bf16"), torch.bfloat16
+        )
+
+        # .clone() detaches from the loader's InferenceMode tensor so the
+        # per-chunk in-place latent updates below are allowed.
+        latents = batch.latents.to(device=device, dtype=target_dtype).clone()
+        init_latents = latents.clone()
+        B, C, total_frames, H, W = latents.shape
+
+        def _iload(_name):
+            return torch.load(f"{_SANAWM_INJECT_DIR}/{_name}.pt", map_location=device)
+
+        _dump_dir = parity_probe.probe_dir(parity_probe.ENV_FORK_DUMP)
+
+        def _fdump(_name, _t):
+            parity_probe.dump_tensor(_dump_dir, _name, _t)
+
+        if (
+            _SANAWM_INJECT_DIR
+        ):  # parity harness: run the OFFICIAL's exact stage-1 inputs
+            latents = (
+                _iload("z_full_initial").to(device=device, dtype=target_dtype).clone()
+            )
+            init_latents = latents.clone()
+            B, C, total_frames, H, W = latents.shape
+
+        _fdump(
+            "init_noise", init_latents
+        )  # parity harness: seeded pre-noise (cond @ frame 0)
+
+        sc = self._resolve_stream_conditioning(
+            batch, server_args, device=device, target_dtype=target_dtype, iload=_iload
+        )
+        sampler_cfg = sc.sampler_cfg
+        num_frame_per_block = sampler_cfg.num_frame_per_block
+        num_cached_blocks = sampler_cfg.num_cached_blocks
+        sink_token = sampler_cfg.sink_token
+        explicit_sigmas = sc.explicit_sigmas
+        cfg_scale = sc.cfg_scale
+        do_cfg = sc.do_cfg
+        embeds_in, mask_in, cam_in, plk_in = sc.embeds, sc.mask, sc.camera, sc.plucker
+
+        # parity harness: full-length conditioning fed to forward_long (windowed
+        # internally per chunk via [start_f:end_f]).
+        _fdump("cond_embeds", embeds_in)
+        _fdump("cond_mask", mask_in)
+        _fdump("cond_camera", cam_in)
+        _fdump("cond_plucker", plk_in)
+
+        scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0)
+
+        chunk_indices = self._autoregressive_segments(total_frames, num_frame_per_block)
+        num_chunks = len(chunk_indices) - 1
+        if num_chunks < 1:
+            raise ValueError(
+                f"streaming needs >= {num_frame_per_block} latent frames, got {total_frames}."
+            )
+
+        start_time = time.perf_counter()
+        with self.use_declared_component(
+            component_name="transformer", module=self.transformer
+        ) as transformer:
+            assert transformer is not None
+            self.transformer = transformer
+            num_blocks = len(transformer.blocks)
+            if _dump_dir:  # parity harness: weights fingerprint
+                parity_probe.dump_obj(
+                    _dump_dir,
+                    "dit_fingerprint",
+                    parity_probe.weights_fingerprint(transformer),
+                )
+            kv_cache = [
+                [[None] * _NUM_STREAM_CACHE_SLOTS for _ in range(num_blocks)]
+                for _ in range(num_chunks)
+            ]
+
+            self.log_info(
+                "SANA-WM streaming denoise: latent=%s, chunks=%d (block=%d frames), "
+                "steps/chunk=%d, cfg=%s",
+                tuple(latents.shape),
+                num_chunks,
+                num_frame_per_block,
+                len(explicit_sigmas),
+                do_cfg,
+            )
+
+            for chunk_idx in self.progress_bar(range(num_chunks)):
+                chunk_kv, sink_num = self._accumulate_kv_cache(
+                    kv_cache,
+                    chunk_idx,
+                    chunk_indices,
+                    num_cached_blocks,
+                    sink_token,
+                    num_blocks,
+                )
+                if _dump_dir:  # parity harness: accumulated-KV checksums
+                    parity_probe.dump_obj(
+                        _dump_dir,
+                        f"kv_probe_{chunk_idx:03d}",
+                        parity_probe.kv_cache_checksums(chunk_kv, sink_num),
+                    )
+                start_f = chunk_indices[chunk_idx]
+                end_f = chunk_indices[chunk_idx + 1]
+                chunk_frames = end_f - start_f
+                frame_index = (
+                    torch.arange(start_f, end_f, device=device, dtype=torch.long)
+                    if sink_num > 0
+                    else None
+                )
+
+                # Condition mask: frame 0 only (chunk 0). Re-pinned each step.
+                cond_mask = torch.zeros(
+                    B, C, chunk_frames, H, W, device=device, dtype=torch.float32
+                )
+                cond_local = []
+                if start_f == 0:
+                    cond_mask[:, :, 0] = 1.0
+                    cond_local = [0]
+
+                # Buffer-view closures: the shared loop reads/writes views of
+                # the full-latent buffer (kernel-visible layouts unchanged).
+                def _get_chunk(start_f=start_f, end_f=end_f):
+                    return latents[:, :, start_f:end_f]
+
+                def _set_chunk(
+                    denoised, start_f=start_f, end_f=end_f, cond_local=cond_local
+                ):
+                    latents[:, :, start_f:end_f] = denoised.to(latents.dtype)
+                    for loc in cond_local:
+                        latents[:, :, start_f + loc] = init_latents[:, :, start_f + loc]
+
+                def _forward_ctx(ts_int):
+                    return set_forward_context(
+                        current_timestep=ts_int,
+                        attn_metadata=None,
+                        forward_batch=batch,
+                    )
+
+                def _on_noise_pred(ts_int, noise_pred, chunk_idx=chunk_idx):
+                    if chunk_idx == 0:  # parity harness: per-step model output
+                        _fdump(f"noise_pred_c0_t{ts_int}", noise_pred)
+
+                kv_cache[chunk_idx] = self_forcing_denoise_chunk(
+                    transformer=transformer,
+                    scheduler=scheduler,
+                    sigmas=explicit_sigmas,
+                    get_chunk=_get_chunk,
+                    set_chunk=_set_chunk,
+                    cond_mask=cond_mask,
+                    embeds=embeds_in,
+                    mask=mask_in,
+                    camera_conditions=cam_in,
+                    chunk_plucker=plk_in,
+                    chunk_kv=chunk_kv,
+                    start_f=start_f,
+                    end_f=end_f,
+                    frame_index=frame_index,
+                    do_cfg=do_cfg,
+                    cfg_scale=cfg_scale,
+                    target_dtype=target_dtype,
+                    device=device,
+                    forward_ctx=_forward_ctx,
+                    on_noise_pred=_on_noise_pred,
+                )
+                _fdump(
+                    f"stage1_{chunk_idx:03d}_{start_f}_{end_f}",
+                    latents[:, :, start_f:end_f],
+                )
+
+        log_sana_wm_tensor_stats("stream.output_latents", latents)
+        self.log_info(
+            "SANA-WM streaming denoise finished in %.4f s; first_frame_max_delta=%.6g",
+            time.perf_counter() - start_time,
+            float((latents[:, :, :1] - init_latents[:, :, :1]).abs().max().item()),
+        )
+        batch.latents = pcfg.post_denoising_loop(latents, batch)
+        return batch
+
+
+class SanaWMStreamingDecodingStage(DecodingStage):
+    """Streaming causal-VAE decode over the SAME autoregressive grid the denoise stage used.
+
+    Carries a per-conv decoder cache across chunks so the causal LTX-2 VAE produces
+    seam-free frames (the `decode_per_frame_with_cache` equivalent at chunk granularity).
+    Subclasses DecodingStage directly (NOT SanaWMDecodingStage, whose long-video config
+    re-enables the stateless tiled decode).
+    """
+
+    @torch.no_grad()
+    def decode(
+        self,
+        latents: torch.Tensor,
+        server_args: ServerArgs,
+        *,
+        vae_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if not hasattr(self.vae, "decode_chunk"):
+            raise ValueError(
+                "SANA-WM streaming decode requires AutoencoderKLCausalLTX2Video "
+                "(decode_chunk). Point --component_paths.vae at the ltx2_causal_vae "
+                "weights when streaming."
+            )
+        device = get_local_torch_device()
+        latents = latents.to(device)
+        pcfg = server_args.pipeline_config
+        num_frame_per_block = int(getattr(pcfg, "num_frame_per_block", 3))
+        total_frames = latents.shape[2]
+        segments = SanaWMStreamingDenoisingStage._autoregressive_segments(
+            total_frames, num_frame_per_block
+        )
+        conv_cache = self.vae.reset_decoder_cache()
+        chunks = []
+        for i in range(len(segments) - 1):
+            s, e = segments[i], segments[i + 1]
+            z = self.scale_and_shift(latents[:, :, s:e].to(vae_dtype), server_args)
+            # No autocast: match the official decode (VAE already runs in vae_dtype,
+            # z is cast once above) for consistent rounding across chunk boundaries.
+            pixel = self.vae.decode_chunk(z, conv_cache)
+            pixel = (pixel / 2 + 0.5).clamp(0, 1)
+            chunks.append(pixel.float().cpu())
+
+        frames = torch.cat(chunks, dim=2)
+        log_sana_wm_tensor_stats("stream.decode.frames", frames)
+        return frames

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py

@@ -0,0 +1,726 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SANA-WM chunked streaming LTX-2 refiner — S2b.
+
+Runs the refiner block-by-block carrying a refiner KV cache: a ``sink`` prefix
+(captured pre-RoPE, re-RoPE'd per block at shifted positions) plus a sliding
+``history`` of refined-block K/V (post-RoPE), injected into each ``attn1`` as a
+KV prefix. Ports the reference NVlabs ``RefinerChunkRunner``
+(minimal-sanawm/refiner.py), operating on mg's already-loaded diffusers LTX-2
+refiner transformer.
+"""
+
+from __future__ import annotations
+
+import math
+import time
+from typing import Any
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.models.dits.sana_wm_refiner_transformer import (
+    pack_latents,
+    unpack_latents,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+from . import parity_probe
+from .refiner import (
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+    SanaWMLTX2RefinerStage,
+    _as_additive_attention_mask,
+    _unwrap_diffusers_ltx2_refiner,
+    log_sana_wm_tensor_stats,
+    sana_wm_skip_refiner_enabled,
+)
+
+
+# --------------------------------------------------------------------------- #
+# Per-attn1 KV-prefix / capture hooks (port of refiner.py:753-788)
+# --------------------------------------------------------------------------- #
+def set_kv_prefix_on_blocks(transformer: nn.Module, kv_prefix_per_layer) -> None:
+    if kv_prefix_per_layer is None:
+        clear_kv_prefix_on_blocks(transformer)
+        return
+    for block, prefix in zip(transformer.transformer_blocks, kv_prefix_per_layer):
+        block.attn1._tf_kv_prefix = prefix
+
+
+def clear_kv_prefix_on_blocks(transformer: nn.Module) -> None:
+    for block in transformer.transformer_blocks:
+        block.attn1._tf_kv_prefix = None
+
+
+def set_capture_flag_on_blocks(
+    transformer: nn.Module, mode: str, *, enable: bool
+) -> None:
+    if mode == "pre_rope":
+        attr, clear_attr = "_kv_cache_capture", "_cached_kv_pre"
+    elif mode == "post_rope":
+        attr, clear_attr = "_tf_capture_kv", "_cached_kv_post"
+    else:
+        raise ValueError(f"unsupported capture mode: {mode}")
+    for block in transformer.transformer_blocks:
+        setattr(block.attn1, attr, bool(enable))
+        if enable and hasattr(block.attn1, clear_attr):
+            setattr(block.attn1, clear_attr, None)
+
+
+def collect_captured_kv_from_blocks(transformer: nn.Module, mode: str):
+    attr = "_cached_kv_pre" if mode == "pre_rope" else "_cached_kv_post"
+    out = []
+    for block in transformer.transformer_blocks:
+        cached = getattr(block.attn1, attr, None)
+        if cached is None:
+            raise RuntimeError(f"missing captured KV on {attr}")
+        out.append(cached)
+        setattr(block.attn1, attr, None)
+    return out
+
+
+# --------------------------------------------------------------------------- #
+# Absolute-position RoPE (port of refiner.py:721-750)
+# --------------------------------------------------------------------------- #
+def build_rotary_emb_for_absolute_positions(
+    *, transformer, batch_size, frame_positions, height, width, device, fps
+):
+    rope = transformer.rope
+    patch_size_t = int(rope.patch_size_t)
+    patch_size = int(rope.patch_size)
+    f_positions = torch.tensor(frame_positions, dtype=torch.float32, device=device)
+    if patch_size_t > 1:
+        f_positions = f_positions[::patch_size_t]
+    grid_h = torch.arange(0, height, patch_size, dtype=torch.float32, device=device)
+    grid_w = torch.arange(0, width, patch_size, dtype=torch.float32, device=device)
+    grid = torch.meshgrid(f_positions, grid_h, grid_w, indexing="ij")
+    grid = torch.stack(grid, dim=0)
+    patch_delta = torch.tensor(
+        (patch_size_t, patch_size, patch_size), dtype=grid.dtype, device=device
+    )
+    patch_ends = grid + patch_delta.view(3, 1, 1, 1)
+    latent_coords = (
+        torch.stack([grid, patch_ends], dim=-1)
+        .flatten(1, 3)
+        .unsqueeze(0)
+        .repeat(batch_size, 1, 1, 1)
+    )
+    scale = torch.tensor(rope.scale_factors, device=device)
+    broadcast_shape = [1] * latent_coords.ndim
+    broadcast_shape[1] = -1
+    pixel_coords = latent_coords * scale.view(*broadcast_shape)
+    pixel_coords[:, 0, ...] = (
+        pixel_coords[:, 0, ...] + rope.causal_offset - rope.scale_factors[0]
+    ).clamp(min=0)
+    pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / float(fps)
+    return rope(pixel_coords, device=device)
+
+
+# --------------------------------------------------------------------------- #
+# Self-attention with KV-prefix injection + capture (port of refiner.py:464-576)
+# --------------------------------------------------------------------------- #
+def streaming_self_attention(
+    *, attn, hidden_states, query_rotary_emb, n_context_tokens
+):
+    sequence_length = hidden_states.shape[1]
+    has_streaming_hooks = (
+        getattr(attn, "_kv_cache_capture", False)
+        or getattr(attn, "_tf_capture_kv", False)
+        or getattr(attn, "_tf_kv_prefix", None) is not None
+    )
+    if n_context_tokens >= sequence_length and not has_streaming_hooks:
+        return attn(
+            hidden_states=hidden_states,
+            encoder_hidden_states=None,
+            query_rotary_emb=query_rotary_emb,
+        )
+
+    from diffusers.models.attention_dispatch import dispatch_attention_fn
+    from diffusers.models.transformers.transformer_ltx2 import (
+        apply_interleaved_rotary_emb,
+        apply_split_rotary_emb,
+    )
+
+    # diffusers 0.38+ always defines `to_gate_logits`; 0.37 only on gated variants
+    # (the SANA-WM refiner is ungated) -> getattr so 0.37 works too.
+    _to_gate_logits = getattr(attn, "to_gate_logits", None)
+    gate_logits = (
+        _to_gate_logits(hidden_states) if _to_gate_logits is not None else None
+    )
+    query = attn.to_q(hidden_states)
+    key = attn.to_k(hidden_states)
+    value = attn.to_v(hidden_states)
+    query = attn.norm_q(query)
+    key = attn.norm_k(key)
+    if getattr(attn, "_kv_cache_capture", False):
+        attn._cached_kv_pre = (key.detach().clone(), value.detach().clone())
+
+    if attn.rope_type == "interleaved":
+        query = apply_interleaved_rotary_emb(query, query_rotary_emb)
+        key = apply_interleaved_rotary_emb(key, query_rotary_emb)
+    elif attn.rope_type == "split":
+        query = apply_split_rotary_emb(query, query_rotary_emb)
+        key = apply_split_rotary_emb(key, query_rotary_emb)
+    else:
+        raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}")
+    if getattr(attn, "_tf_capture_kv", False):
+        attn._cached_kv_post = (key.detach().clone(), value.detach().clone())
+
+    tf_prefix = getattr(attn, "_tf_kv_prefix", None)
+    if isinstance(tf_prefix, dict) and tf_prefix.get("mode") == "rf_shifted_sink":
+        prefix_k_parts = []
+        prefix_v_parts = []
+        sink_k_pre = tf_prefix.get("sink_k_pre")
+        sink_v = tf_prefix.get("sink_v")
+        if sink_k_pre is not None and sink_v is not None and sink_k_pre.shape[1] > 0:
+            sink_pe = tf_prefix.get("sink_pe")
+            if sink_pe is None:
+                raise RuntimeError("rf_shifted_sink prefix requires sink_pe")
+            if attn.rope_type == "interleaved":
+                sink_k = apply_interleaved_rotary_emb(sink_k_pre.to(key.dtype), sink_pe)
+            else:
+                sink_k = apply_split_rotary_emb(sink_k_pre.to(key.dtype), sink_pe)
+            prefix_k_parts.append(sink_k)
+            prefix_v_parts.append(sink_v.to(value.dtype))
+        history_k = tf_prefix.get("history_k")
+        history_v = tf_prefix.get("history_v")
+        if history_k is not None and history_v is not None and history_k.shape[1] > 0:
+            prefix_k_parts.append(history_k.to(key.dtype))
+            prefix_v_parts.append(history_v.to(value.dtype))
+        if prefix_k_parts:
+            key = torch.cat([*prefix_k_parts, key], dim=1)
+            value = torch.cat([*prefix_v_parts, value], dim=1)
+
+    query = query.unflatten(2, (attn.heads, -1))
+    key = key.unflatten(2, (attn.heads, -1))
+    value = value.unflatten(2, (attn.heads, -1))
+    processor = attn.processor
+    backend = getattr(processor, "_attention_backend", None)
+    parallel_config = getattr(processor, "_parallel_config", None)
+    if n_context_tokens <= 0 or n_context_tokens >= query.shape[1]:
+        hidden_states = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=False,
+            backend=backend,
+            parallel_config=parallel_config,
+        )
+    else:
+        context = dispatch_attention_fn(
+            query[:, :n_context_tokens],
+            key[:, :n_context_tokens],
+            value[:, :n_context_tokens],
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=False,
+            backend=backend,
+            parallel_config=parallel_config,
+        )
+        current = dispatch_attention_fn(
+            query[:, n_context_tokens:],
+            key,
+            value,
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=False,
+            backend=backend,
+            parallel_config=parallel_config,
+        )
+        hidden_states = torch.cat([context, current], dim=1)
+    hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+    if gate_logits is not None:
+        hidden_states = hidden_states.unflatten(2, (attn.heads, -1))
+        hidden_states = hidden_states * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)
+        hidden_states = hidden_states.flatten(2, 3)
+    hidden_states = attn.to_out[0](hidden_states)
+    return attn.to_out[1](hidden_states)
+
+
+def forward_video_block(
+    *,
+    block,
+    hidden_states,
+    encoder_hidden_states,
+    temb,
+    video_rotary_emb,
+    encoder_attention_mask,
+    n_context_tokens,
+):
+    batch = hidden_states.size(0)
+    norm_hidden_states = block.norm1(hidden_states)
+    num_ada_params = block.scale_shift_table.shape[0]
+    ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape(
+        batch, temb.size(1), num_ada_params, -1
+    )
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(
+        dim=2
+    )
+    norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+    attn_hidden_states = streaming_self_attention(
+        attn=block.attn1,
+        hidden_states=norm_hidden_states,
+        query_rotary_emb=video_rotary_emb,
+        n_context_tokens=n_context_tokens,
+    )
+    hidden_states = hidden_states + attn_hidden_states * gate_msa
+    norm_hidden_states = block.norm2(hidden_states)
+    attn_hidden_states = block.attn2(
+        norm_hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        query_rotary_emb=None,
+        attention_mask=encoder_attention_mask,
+    )
+    hidden_states = hidden_states + attn_hidden_states
+    norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp
+    return hidden_states + block.ff(norm_hidden_states) * gate_mlp
+
+
+class _RefinerCore:
+    """Adapts the unwrapped diffusers refiner transformer to the
+    DiffusersLTX2Refiner interface RefinerChunkRunner expects."""
+
+    def __init__(
+        self, transformer: nn.Module, device: torch.device, dtype: torch.dtype
+    ):
+        self.transformer = transformer
+        self.device = device
+        self.dtype = dtype
+
+    def _forward_video_only_with_rope(
+        self,
+        *,
+        hidden_states,
+        encoder_hidden_states,
+        timestep,
+        encoder_attention_mask,
+        video_rotary_emb,
+        n_context_tokens,
+    ):
+        transformer = self.transformer
+        batch = hidden_states.size(0)
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = _as_additive_attention_mask(
+                encoder_attention_mask, hidden_states.dtype
+            )
+        hidden_states = transformer.proj_in(hidden_states)
+        temb, embedded_timestep = transformer.time_embed(
+            timestep.flatten(),
+            batch_size=batch,
+            hidden_dtype=hidden_states.dtype,
+        )
+        temb = temb.view(batch, -1, temb.size(-1))
+        embedded_timestep = embedded_timestep.view(
+            batch, -1, embedded_timestep.size(-1)
+        )
+        encoder_hidden_states = transformer.caption_projection(encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states.view(
+            batch, -1, hidden_states.size(-1)
+        )
+        for block in transformer.transformer_blocks:
+            hidden_states = forward_video_block(
+                block=block,
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                video_rotary_emb=video_rotary_emb,
+                encoder_attention_mask=encoder_attention_mask,
+                n_context_tokens=n_context_tokens,
+            )
+        scale_shift = (
+            transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None]
+        )
+        shift, scale = scale_shift[:, :, 0], scale_shift[:, :, 1]
+        hidden_states = transformer.norm_out(hidden_states)
+        hidden_states = hidden_states * (1 + scale) + shift
+        return transformer.proj_out(hidden_states)
+
+    def _predict_x0_active_block(
+        self,
+        *,
+        active,
+        active_positions,
+        sigma_cur,
+        prompt_embeds,
+        prompt_attention_mask,
+        fps,
+        kv_prefix_per_layer,
+    ):
+        ps = int(self.transformer.config.patch_size)
+        pst = int(self.transformer.config.patch_size_t)
+        latent_tokens = pack_latents(active, ps, pst)
+        batch, seq_len, _ = latent_tokens.shape
+        timestep = torch.full(
+            (batch, seq_len),
+            float(sigma_cur) * float(self.transformer.config.timestep_scale_multiplier),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        video_rotary_emb = build_rotary_emb_for_absolute_positions(
+            transformer=self.transformer,
+            batch_size=batch,
+            frame_positions=active_positions,
+            height=int(active.shape[3]),
+            width=int(active.shape[4]),
+            device=self.device,
+            fps=float(fps),
+        )
+        set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer)
+        try:
+            velocity = self._forward_video_only_with_rope(
+                hidden_states=latent_tokens,
+                encoder_hidden_states=prompt_embeds,
+                timestep=timestep,
+                encoder_attention_mask=prompt_attention_mask,
+                video_rotary_emb=video_rotary_emb,
+                n_context_tokens=0,
+            )
+        finally:
+            clear_kv_prefix_on_blocks(self.transformer)
+        raw_sigma = torch.full(
+            (batch, seq_len, 1),
+            float(sigma_cur),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        denoised = latent_tokens.float() - velocity.float() * raw_sigma
+        return unpack_latents(
+            denoised.to(self.dtype),
+            num_frames=int(active.shape[2]),
+            height=int(active.shape[3]),
+            width=int(active.shape[4]),
+            patch_size=ps,
+            patch_size_t=pst,
+        )
+
+    def _capture_block_kv(
+        self,
+        *,
+        clean_block,
+        frame_positions,
+        prompt_embeds,
+        prompt_attention_mask,
+        fps,
+        capture_mode,
+        kv_prefix_per_layer,
+    ):
+        ps = int(self.transformer.config.patch_size)
+        pst = int(self.transformer.config.patch_size_t)
+        latent_tokens = pack_latents(clean_block, ps, pst)
+        batch, seq_len, _ = latent_tokens.shape
+        timestep = torch.zeros(batch, seq_len, dtype=torch.float32, device=self.device)
+        video_rotary_emb = build_rotary_emb_for_absolute_positions(
+            transformer=self.transformer,
+            batch_size=batch,
+            frame_positions=frame_positions,
+            height=int(clean_block.shape[3]),
+            width=int(clean_block.shape[4]),
+            device=self.device,
+            fps=float(fps),
+        )
+        set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer)
+        set_capture_flag_on_blocks(self.transformer, capture_mode, enable=True)
+        try:
+            _ = self._forward_video_only_with_rope(
+                hidden_states=latent_tokens,
+                encoder_hidden_states=prompt_embeds,
+                timestep=timestep,
+                encoder_attention_mask=prompt_attention_mask,
+                video_rotary_emb=video_rotary_emb,
+                n_context_tokens=0,
+            )
+        finally:
+            set_capture_flag_on_blocks(self.transformer, capture_mode, enable=False)
+            clear_kv_prefix_on_blocks(self.transformer)
+        return collect_captured_kv_from_blocks(self.transformer, capture_mode)
+
+
+class RefinerChunkRunner:
+    """Port of the reference RefinerChunkRunner (refiner.py:579-718)."""
+
+    def __init__(
+        self,
+        refiner: _RefinerCore,
+        *,
+        prompt_embeds,
+        prompt_attention_mask,
+        fps,
+        sigmas,
+        source_sink_frames,
+        block_size,
+        kv_max_frames,
+        seed,
+        spatial_shape,
+    ):
+        self.refiner = refiner
+        self.prompt_embeds = prompt_embeds
+        self.prompt_attention_mask = prompt_attention_mask
+        self.fps = float(fps)
+        self.sigmas = sigmas
+        self.source_sink_frames = int(source_sink_frames)
+        self.block_size = int(block_size)
+        self.kv_max_frames = int(kv_max_frames)
+        self.max_history_frames = self.kv_max_frames - self.source_sink_frames
+        self.generator = torch.Generator(device=refiner.device).manual_seed(int(seed))
+        self.device = refiner.device
+        self.dtype = refiner.dtype
+        self.height, self.width = int(spatial_shape[0]), int(spatial_shape[1])
+        transformer = refiner.transformer
+        self.tokens_per_frame = (
+            int(self.height // transformer.config.patch_size)
+            * int(self.width // transformer.config.patch_size)
+            * int(transformer.config.patch_size_t)
+        )
+        self.sink_kv_pre = None
+        self.history_kv_post = [None] * len(transformer.transformer_blocks)
+        self.history_frames = 0
+
+    @torch.inference_mode()
+    def refine_block(
+        self, *, block_idx, clean_block, block_start, block_end, sink_seed_frames=None
+    ):
+        # parity harness (env-gated, no-op in prod): per-block input/config/output
+        # checksums; both the batch stage and the realtime stage run through here.
+        _probe_dir = parity_probe.probe_dir(
+            parity_probe.ENV_RT_DUMP, parity_probe.ENV_FORK_DUMP
+        )
+        _probe = None
+        if _probe_dir:
+            _ck = parity_probe.checksum
+            _probe = {
+                "block_idx": int(block_idx),
+                "block_start": int(block_start),
+                "block_end": int(block_end),
+                "clean_block": _ck(clean_block),
+                "sink_seed_frames": _ck(sink_seed_frames),
+                "prompt_embeds": _ck(self.prompt_embeds),
+                "prompt_attention_mask": _ck(self.prompt_attention_mask),
+                "fps": self.fps,
+                "sigmas": [float(s) for s in self.sigmas],
+                "source_sink_frames": self.source_sink_frames,
+                "block_size": self.block_size,
+                "kv_max_frames": self.kv_max_frames,
+                "history_frames": self.history_frames,
+                "generator_state": float(
+                    self.generator.get_state().double().sum().item()
+                ),
+            }
+        del block_idx
+        refiner = self.refiner
+        if block_start < self.source_sink_frames:
+            raise ValueError("refiner block overlaps source sink")
+        if self.sink_kv_pre is None:
+            if self.source_sink_frames == 0:
+                self.sink_kv_pre = [(None, None) for _ in self.history_kv_post]
+            elif sink_seed_frames is None:
+                raise ValueError("first refine_block call requires sink_seed_frames")
+            else:
+                self.sink_kv_pre = refiner._capture_block_kv(
+                    clean_block=sink_seed_frames.contiguous(),
+                    frame_positions=list(range(self.source_sink_frames)),
+                    prompt_embeds=self.prompt_embeds,
+                    prompt_attention_mask=self.prompt_attention_mask,
+                    fps=self.fps,
+                    capture_mode="pre_rope",
+                    kv_prefix_per_layer=None,
+                )
+        batch = int(clean_block.shape[0])
+        sink_rope_offset = block_start - self.history_frames - self.source_sink_frames
+        sink_pe = None
+        if self.source_sink_frames > 0:
+            sink_pe = build_rotary_emb_for_absolute_positions(
+                transformer=refiner.transformer,
+                batch_size=batch,
+                frame_positions=list(
+                    range(sink_rope_offset, sink_rope_offset + self.source_sink_frames)
+                ),
+                height=self.height,
+                width=self.width,
+                device=self.device,
+                fps=self.fps,
+            )
+        kv_prefix_per_layer = []
+        for layer_idx, sink_kv in enumerate(self.sink_kv_pre):
+            history = self.history_kv_post[layer_idx]
+            kv_prefix_per_layer.append(
+                {
+                    "mode": "rf_shifted_sink",
+                    "sink_k_pre": sink_kv[0],
+                    "sink_v": sink_kv[1],
+                    "sink_pe": sink_pe,
+                    "history_k": history[0] if history is not None else None,
+                    "history_v": history[1] if history is not None else None,
+                }
+            )
+        sigma0 = float(self.sigmas[0].item())
+        eps = torch.randn(
+            clean_block.shape,
+            generator=self.generator,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        x_t = ((1.0 - sigma0) * clean_block.float() + sigma0 * eps.float()).to(
+            self.dtype
+        )
+        active_positions = list(range(int(block_start), int(block_end)))
+        for level in range(int(self.sigmas.numel()) - 1):
+            sigma_cur = float(self.sigmas[level].item())
+            sigma_next = float(self.sigmas[level + 1].item())
+            pred_x0 = refiner._predict_x0_active_block(
+                active=x_t,
+                active_positions=active_positions,
+                sigma_cur=sigma_cur,
+                prompt_embeds=self.prompt_embeds,
+                prompt_attention_mask=self.prompt_attention_mask,
+                fps=self.fps,
+                kv_prefix_per_layer=kv_prefix_per_layer,
+            )
+            if sigma_cur <= 1.0e-6:
+                x_t = pred_x0.to(self.dtype)
+            else:
+                ratio = sigma_next / sigma_cur
+                x_t = (ratio * x_t.float() + (1.0 - ratio) * pred_x0.float()).to(
+                    self.dtype
+                )
+        block_kv_post = refiner._capture_block_kv(
+            clean_block=x_t,
+            frame_positions=active_positions,
+            prompt_embeds=self.prompt_embeds,
+            prompt_attention_mask=self.prompt_attention_mask,
+            fps=self.fps,
+            capture_mode="post_rope",
+            kv_prefix_per_layer=kv_prefix_per_layer,
+        )
+        for layer_idx, new_kv in enumerate(block_kv_post):
+            old = self.history_kv_post[layer_idx]
+            self.history_kv_post[layer_idx] = (
+                new_kv
+                if old is None
+                else (
+                    torch.cat([old[0], new_kv[0]], dim=1),
+                    torch.cat([old[1], new_kv[1]], dim=1),
+                )
+            )
+        self.history_frames += int(block_end - block_start)
+        if (
+            self.max_history_frames > 0
+            and self.history_frames > self.max_history_frames
+        ):
+            keep_tokens = self.max_history_frames * self.tokens_per_frame
+            for layer_idx, old in enumerate(self.history_kv_post):
+                if old is not None:
+                    self.history_kv_post[layer_idx] = (
+                        old[0][:, -keep_tokens:],
+                        old[1][:, -keep_tokens:],
+                    )
+            self.history_frames = self.max_history_frames
+        if _probe is not None:
+            _probe["refined"] = parity_probe.checksum(x_t)
+            parity_probe.dump_obj(
+                _probe_dir, f"refiner_probe_{_probe['block_idx']:03d}", _probe
+            )
+        return x_t
+
+
+class SanaWMStreamingRefinerStage(SanaWMLTX2RefinerStage):
+    """Chunked streaming LTX-2 refiner — refines stage-1 latents block-by-block
+    carrying a sink/history KV cache. Inherits loading/encode/residency from
+    SanaWMLTX2RefinerStage; only the refine loop changes."""
+
+    def __init__(
+        self,
+        *,
+        block_size: int = 3,
+        kv_max_frames: int = 11,
+        sink_size: int = 1,
+        seed: int = 42,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.block_size = int(block_size)
+        self.kv_max_frames = int(kv_max_frames)
+        self.sink_size = int(sink_size)
+        self.seed = int(seed)
+
+    @torch.inference_mode()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        if sana_wm_skip_refiner_enabled():
+            return batch
+        if batch.latents is None or batch.latents.ndim != 5:
+            raise ValueError("SANA-WM streaming refiner expects 5D stage-1 latents.")
+        device = get_local_torch_device()
+        latents = batch.latents.to(device=device, dtype=self.dtype).clone()
+        B, C, T, H, W = latents.shape
+        prompt = self._prompts_for_batch(batch, B)[0]
+        fps = float((batch.extra or {}).get("fps", getattr(batch, "fps", 16.0)) or 16.0)
+
+        prompt_embeds, prompt_mask = self._encode_prompt(prompt, device)
+        sigmas = torch.tensor(
+            STAGE_2_DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device
+        )
+
+        n_active = T - self.sink_size
+        if n_active <= 0:
+            self.log_info(
+                "SANA-WM streaming refiner: no active frames (T=%d <= sink=%d); skipping.",
+                T,
+                self.sink_size,
+            )
+            return batch
+        n_blocks = math.ceil(n_active / self.block_size)
+        self.log_info(
+            "SANA-WM streaming refiner: latent=%s, sink=%d, block=%d, blocks=%d, kv_max=%d, seed=%d",
+            tuple(latents.shape),
+            self.sink_size,
+            self.block_size,
+            n_blocks,
+            self.kv_max_frames,
+            self.seed,
+        )
+
+        t0 = time.perf_counter()
+        with self.use_declared_component(
+            component_name="transformer_2", module=self.transformer
+        ) as transformer_mod:
+            self.transformer = transformer_mod
+            unwrapped = _unwrap_diffusers_ltx2_refiner(self.transformer)
+            core = _RefinerCore(unwrapped, device, self.dtype)
+            runner = RefinerChunkRunner(
+                core,
+                prompt_embeds=prompt_embeds,
+                prompt_attention_mask=prompt_mask,
+                fps=fps,
+                sigmas=sigmas,
+                source_sink_frames=self.sink_size,
+                block_size=self.block_size,
+                kv_max_frames=self.kv_max_frames,
+                seed=self.seed,
+                spatial_shape=(H, W),
+            )
+            for i in range(n_blocks):
+                start_f = self.sink_size + i * self.block_size
+                end_f = min(start_f + self.block_size, T)
+                sink_seed = latents[:, :, : self.sink_size] if i == 0 else None
+                refined = runner.refine_block(
+                    block_idx=i,
+                    clean_block=latents[:, :, start_f:end_f].contiguous(),
+                    block_start=start_f,
+                    block_end=end_f,
+                    sink_seed_frames=sink_seed,
+                )
+                latents[:, :, start_f:end_f] = refined.to(latents.dtype)
+
+        log_sana_wm_tensor_stats("stream.refiner.latents", latents)
+        self.log_info(
+            "SANA-WM streaming refiner applied (%d blocks) in %.4f s.",
+            n_blocks,
+            time.perf_counter() - t0,
+        )
+        batch.latents = latents
+        if batch.extra is None:
+            batch.extra = {}
+        batch.extra["sana_wm_refiner_applied"] = True
+        return batch

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

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+    ComponentUse,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+    PipelineStage,
+    StageParallelismType,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import (
+    scale_and_shift_latents,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+
+@dataclass(frozen=True)
+class RealtimeStageComponent:
+    component_name: str
+    precision_attr: str | None = None
+    memory_intensive: bool = False
+    keep_ready_after_warmup: bool = True
+
+
+class RealtimeDiffusionStage(PipelineStage):
+    """common contract for stateful realtime diffusion stages
+
+    model-specific subclasses keep conditioning, latent scheduling, and model
+    forward semantics
+    """
+
+    component_specs: tuple[RealtimeStageComponent, ...] = (
+        RealtimeStageComponent("transformer", "dit_precision", memory_intensive=True),
+        RealtimeStageComponent("vae", "vae_precision"),
+    )
+
+    def __init__(
+        self,
+        *,
+        transformer: torch.nn.Module | None = None,
+        vae: torch.nn.Module | None = None,
+        model_path: str | None = None,
+        default_height: int | None = None,
+        default_width: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.vae = vae
+        self.model_path = model_path
+        self.default_height = default_height
+        self.default_width = default_width
+
+    @property
+    def role_affinity(self):
+        return RoleType.MONOLITHIC
+
+    @property
+    def parallelism_type(self) -> StageParallelismType:
+        return StageParallelismType.REPLICATED
+
+    def require_session(self, batch: Req, *, context: str | None = None):
+        if batch.session is None:
+            label = context or self.__class__.__name__
+            raise ValueError(f"{label} requires a realtime session")
+        return batch.session
+
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        uses: list[ComponentUse] = []
+        for spec in self.component_specs:
+            target_dtype = None
+            if spec.precision_attr is not None:
+                precision = getattr(server_args.pipeline_config, spec.precision_attr)
+                target_dtype = PRECISION_TO_TYPE[precision]
+            uses.append(
+                ComponentUse(
+                    stage_name,
+                    spec.component_name,
+                    target_dtype=target_dtype,
+                    memory_intensive=spec.memory_intensive,
+                    keep_ready_after_warmup=spec.keep_ready_after_warmup,
+                )
+            )
+        return uses
+
+    def target_pixel_size(self, batch: Req) -> tuple[int, int]:
+        height = batch.height if batch.height is not None else self.default_height
+        width = batch.width if batch.width is not None else self.default_width
+        if height is None or width is None:
+            raise ValueError(
+                f"{self.__class__.__name__} needs batch.height/batch.width or defaults"
+            )
+        return int(height), int(width)
+
+    def _empty_output(self, batch: Req) -> OutputBatch:
+        target_h, target_w = self.target_pixel_size(batch)
+        output = torch.empty(
+            (1, 3, 0, target_h, target_w),
+            dtype=torch.float32,
+            device=get_local_torch_device(),
+        )
+        return OutputBatch(output=output, metrics=batch.metrics)
+
+    def scale_and_shift_latents(
+        self,
+        latents: torch.Tensor,
+        server_args: ServerArgs,
+    ) -> torch.Tensor:
+        return scale_and_shift_latents(latents, server_args, self.vae)
+
+    def ensure_causal_vae_conv_cache(self, state) -> dict:
+        if state.conv_cache is None:
+            state.conv_cache = self.vae.reset_decoder_cache()
+        return state.conv_cache

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

@@ -2,6 +2,9 @@
 
 """session-scoped realtime state, control events, and runtime-only helpers"""
 
+from sglang.multimodal_gen.runtime.realtime.camera_controls import (
+    RealtimeCameraControlState,
+)
 from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState
 from sglang.multimodal_gen.runtime.realtime.condition_events import (
     ConditionEvent,
@@ -25,6 +28,7 @@ __all__ = [
     "ControlSignal",
     "ControlStateSamplingQueue",
     "ControlStateTransition",
+    "RealtimeCameraControlState",
     "RealtimeCausalDiTState",
     "RealtimeSession",
     "RealtimeSessionCache",

python/sglang/multimodal_gen/runtime/realtime/camera_controls.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Callable
+from typing import Any
+
+from sglang.multimodal_gen.runtime.realtime.condition_events import (
+    ControlSignal,
+    ControlStateSamplingQueue,
+    ControlStateTransition,
+)
+
+CameraActionNormalizer = Callable[[list[Any]], list[str]]
+CameraActionValidator = Callable[[Any], list[list[str]]]
+
+
+def _identity_actions(actions: list[Any]) -> list[str]:
+    return list(actions)
+
+
+class RealtimeCameraControlState:
+    """Session-local camera-control buffer shared by realtime model adapters.
+
+    Camera controls arrive in two shapes:
+
+    1. Script mode: ``list[list[str]]`` where each item is one output-frame's
+       held actions. The script is consumed once from a FIFO and padded with
+       neutral ``[]`` frames after it runs out.
+    2. State mode: timestamped transitions such as "W is currently held".
+       ``ControlStateSamplingQueue`` expands that continuous state into the next
+       chunk and can pulse a short key press for a minimum number of frames.
+
+    The two modes are intentionally exclusive. A new script clears state mode,
+    and new state transitions clear script mode, so adapters never merge two
+    camera timelines accidentally. ``sample_camera_actions`` returns ``None``
+    only when no control should be sent; otherwise it returns exactly
+    ``chunk_size`` frames, with ``[]`` meaning neutral/no-op for that frame.
+    """
+
+    def __init__(
+        self,
+        *,
+        min_pulse_items: int = 1,
+        script_maxlen: int = 512,
+        max_transitions: int = 512,
+        normalize_state_actions: CameraActionNormalizer = _identity_actions,
+    ) -> None:
+        self.camera_state = ControlStateSamplingQueue(
+            default_item=[],
+            min_pulse_items=min_pulse_items,
+            max_transitions=max_transitions,
+        )
+        self.camera_script_queue: deque[ControlSignal] = deque(maxlen=script_maxlen)
+        self.latest_sampled_event_id: int | None = None
+        self._normalize_state_actions = normalize_state_actions
+
+    def clear(self) -> None:
+        """Reset all camera controls owned by this realtime session."""
+        self.camera_state.clear()
+        self.camera_script_queue.clear()
+        self.latest_sampled_event_id = None
+
+    def receive_camera_script(
+        self,
+        camera_actions: list[list[str]],
+        *,
+        event_id: int | None = None,
+    ) -> None:
+        """Replace active controls with a finite per-frame script."""
+        self.camera_script_queue.clear()
+        self.camera_state.clear()
+        for actions in camera_actions:
+            self.camera_script_queue.append(
+                ControlSignal(
+                    kind="camera_actions",
+                    payload=list(actions),
+                    seq_id=event_id,
+                )
+            )
+
+    def receive_camera_state_transitions(
+        self,
+        transitions: list[ControlStateTransition],
+    ) -> None:
+        """Replace the script with continuous state transitions."""
+        self.camera_script_queue.clear()
+        self.camera_state.push_many(transitions)
+
+    def receive_camera_actions(
+        self,
+        camera_actions: list[list[str]],
+        *,
+        event_id: int | None = None,
+    ) -> None:
+        self.receive_camera_script(camera_actions, event_id=event_id)
+
+    def receive_camera_state(
+        self,
+        actions: list[str],
+        *,
+        event_id: int | None = None,
+        timestamp_ms: int | None = None,
+    ) -> None:
+        self.receive_camera_state_transitions(
+            [
+                self._camera_state_transition(
+                    actions,
+                    event_id=event_id,
+                    timestamp_ms=timestamp_ms,
+                )
+            ]
+        )
+
+    def receive_camera_event_payload(
+        self,
+        payload: Any,
+        *,
+        event_id: int | None,
+        validate_camera_actions: CameraActionValidator,
+    ) -> str:
+        """Parse an external camera event and install it as script or state."""
+        if isinstance(payload, dict) and payload.get("mode") == "state":
+            transitions = self._camera_transitions_from_event_payload(
+                payload,
+                event_id=event_id,
+            )
+            self.receive_camera_state_transitions(transitions)
+            return f"kind=camera_actions, mode=state, transitions={len(transitions)}"
+
+        camera_actions = validate_camera_actions(payload)
+        self.receive_camera_script(camera_actions, event_id=event_id)
+        return f"kind=camera_actions, mode=script, frames={len(camera_actions)}"
+
+    def sample_camera_actions(self, chunk_size: int) -> list[list[str]] | None:
+        """Return the next chunk-sized camera action window.
+
+        Script mode has priority because it represents an explicit finite
+        timeline. State mode is sampled only when no script is pending.
+        """
+        if self.camera_script_queue:
+            return self._sample_camera_script(chunk_size)
+        action_list = self.camera_state.sample_chunk(chunk_size)
+        if action_list is None:
+            return None
+        self.latest_sampled_event_id = self.camera_state.latest_sampled_seq_id()
+        return [list(actions) for actions in action_list]
+
+    def _sample_camera_script(self, chunk_size: int) -> list[list[str]]:
+        chunk: list[list[str]] = []
+        latest_event_id = self.latest_sampled_event_id
+        while self.camera_script_queue and len(chunk) < chunk_size:
+            signal = self.camera_script_queue.popleft()
+            chunk.append(list(signal.payload))
+            latest_event_id = signal.seq_id
+        while len(chunk) < chunk_size:
+            chunk.append([])
+        self.latest_sampled_event_id = latest_event_id
+        return chunk
+
+    def _camera_state_transition(
+        self,
+        actions: list[Any],
+        *,
+        event_id: int | None,
+        timestamp_ms: int | None,
+    ) -> ControlStateTransition:
+        return ControlStateTransition(
+            payload=self._normalize_state_actions(actions),
+            seq_id=event_id,
+            timestamp_ms=timestamp_ms,
+        )
+
+    def _camera_transitions_from_event_payload(
+        self,
+        payload: dict[str, Any],
+        *,
+        event_id: int | None,
+    ) -> list[ControlStateTransition]:
+        transitions = payload.get("transitions")
+        if not isinstance(transitions, list):
+            raise ValueError("camera_actions state payload requires transitions")
+        result = []
+        for transition in transitions:
+            if not isinstance(transition, dict):
+                raise ValueError("camera_actions transition must be a map")
+            actions = transition.get("actions")
+            if not isinstance(actions, list):
+                raise ValueError("camera_actions transition actions must be a list")
+            timestamp_ms = transition.get("client_ts_ms")
+            if timestamp_ms is not None:
+                timestamp_ms = int(timestamp_ms)
+            result.append(
+                self._camera_state_transition(
+                    actions,
+                    event_id=event_id,
+                    timestamp_ms=timestamp_ms,
+                )
+            )
+        return result

python/sglang/multimodal_gen/runtime/realtime/causal_state.py

@@ -20,3 +20,16 @@ class RealtimeCausalDiTState(BaseRealtimeState):
         self.runtime_cache.clear()
         self.current_chunk_start_frame = 0
         self.chunk_idx = 0
+
+
+class RealtimeCausalDecodeState(BaseRealtimeState):
+    """persist causal VAE decode cache and output frontier across chunks"""
+
+    def __init__(self):
+        super().__init__()
+        self.conv_cache: dict | None = None
+        self.next_dec_idx: int = 0
+
+    def dispose(self) -> None:
+        self.conv_cache = None
+        self.next_dec_idx = 0

python/sglang/multimodal_gen/runtime/server_args.py

@@ -778,9 +778,11 @@ class ServerArgs(DisaggServerArgsMixin):
         # SamplingParams use classifier-free guidance (negative_prompt is not None),
         # because non-CFG models (e.g. FLUX) crash when CFG parallel splits ranks.
         if cfg_unspecified:
+            deployment_config = self.pipeline_config.get_model_deployment_config()
             cfg_group_size = self.dp_size * self.tp_size * 2
             if (
                 self.performance_mode != "manual"
+                and deployment_config.auto_enable_cfg_parallel
                 and self.num_gpus >= 2
                 and self.num_gpus % cfg_group_size == 0
                 and sp_unspecified

python/sglang/multimodal_gen/runtime/server_args_auto_tune.py

@@ -490,8 +490,10 @@ class ServerArgsAutoTuner:
 
     def _enable_cfg_parallel_if_supported(self) -> None:
         args = self.server_args
+        deployment_config = self._deployment_config()
         if (
-            args.enable_cfg_parallel is None
+            deployment_config.auto_enable_cfg_parallel
+            and args.enable_cfg_parallel is None
             and not self._has_explicit_parallel_policy()
             and args._model_default_uses_cfg()
         ):

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

@@ -29,6 +29,14 @@ BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
         "overlay_repo_id": "MickJ/LTX-2.3-overlay",
         "overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e",
     },
+    "Efficient-Large-Model/SANA-WM_bidirectional": {
+        "overlay_repo_id": "sjmshsh/SANA-WM_bidirectional-overlay",
+        "overlay_revision": "e611beacbcc0cf33c676306ae0eb89f149e044ad",
+    },
+    "Efficient-Large-Model/SANA-WM_streaming": {
+        "overlay_repo_id": "AgainstEntropy/SANA-WM_streaming-overlay",
+        "overlay_revision": "62c6840871ecc3559189047513ba0670e1bf62e7",
+    },
 }
 
 
@@ -150,9 +158,17 @@ def resolve_model_overlay_target(
     if os.path.exists(model_name_or_path):
         # Local source dirs do not have a repo id, so match them by basename.
         base_name = os.path.basename(os.path.normpath(model_name_or_path))
+        normalized_path = (
+            os.path.normpath(model_name_or_path).lower().replace(os.sep, "/")
+        )
         for source_model_id, spec in registry.items():
             if base_name == source_model_id.rsplit("/", 1)[-1]:
                 return source_model_id, spec
+            cache_repo_fragment = (
+                f"models--{source_model_id.lower().replace('/', '--')}"
+            )
+            if cache_repo_fragment in normalized_path:
+                return source_model_id, spec
 
     return None
 
@@ -577,7 +593,6 @@ def maybe_load_overlay_model_index(
         # A local overlay repo already contains the model_index we need.
         if load_overlay_manifest_if_present(model_name_or_path) is not None:
             return load_model_index_from_dir(model_name_or_path)
-        return None
 
     overlay_target = resolve_model_overlay_target(model_name_or_path)
     if overlay_target is not None:

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -26,6 +26,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
     MULTI_FRAME_I2I_sampling_params,
     MULTI_IMAGE_TI2I_sampling_params,
     MULTI_IMAGE_TI2I_UPLOAD_sampling_params,
+    SANA_WM_TI2V_CI_sampling_params,
     T2I_sampling_params,
     T2V_sampling_params,
     _make_modelopt_ci_case,
@@ -44,6 +45,7 @@ from sglang.multimodal_gen.test.test_utils import (
     DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST,
     DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST,
     DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
+    DEFAULT_SANA_WM_STREAMING_MODEL_NAME_FOR_TEST,
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
     DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST,
     DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST,
@@ -376,6 +378,17 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
             model_path="FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
         ),
     ),
+    DiffusionTestCase(
+        "sana_wm_ti2v",
+        DiffusionServerArgs(
+            model_path=DEFAULT_SANA_WM_STREAMING_MODEL_NAME_FOR_TEST,
+        ),
+        SANA_WM_TI2V_CI_sampling_params,
+        run_perf_check=False,
+        run_component_accuracy_check=False,
+        run_models_api_check=False,
+        run_t2v_input_reference_check=False,
+    ),
     # flaky
     # === Helios T2V ===
     # DiffusionTestCase(

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -545,6 +545,15 @@ TI2V_sampling_params = DiffusionSamplingParams(
     direct_url_test=True,
 )
 
+SANA_WM_TI2V_CI_sampling_params = DiffusionSamplingParams(
+    prompt=TI2V_sampling_params.prompt,
+    image_path=TI2V_sampling_params.image_path,
+    direct_url_test=True,
+    output_size="384x640",
+    num_frames=17,
+    extras={"num_inference_steps": 12, "seed": 0, "guidance_scale": 4.5},
+)
+
 TURBOWAN_I2V_sampling_params = DiffusionSamplingParams(
     prompt="The man in the picture slowly turns his head, his expression enigmatic and otherworldly. The camera performs a slow, cinematic dolly out, focusing on his face. Moody lighting, neon signs glowing in the background, shallow depth of field.",
     image_path="https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/5f/fa/56/5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/825646124206.jpg/600x600bb.jpg",

python/sglang/multimodal_gen/test/test_utils.py

@@ -33,7 +33,7 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-SGL_TEST_FILES_CI_DATA_REVISION = "af6e712a2c49ab5fcd81dde58e2f54c78e77683b"
+SGL_TEST_FILES_CI_DATA_REVISION = "caa56302ccf2d289e4488ed06d952edf5d2314cf"
 SGL_TEST_FILES_CONSISTENCY_GT_ROOT = (
     "https://raw.githubusercontent.com/"
     f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/"
@@ -154,6 +154,12 @@ DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-I2V-A14B-Diffusers
 # MOVA video generation models
 DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST = "OpenMOSS-Team/MOVA-360p"
 
+# SANA-WM world model (TI2V with optional camera conditioning)
+DEFAULT_SANA_WM_MODEL_NAME_FOR_TEST = "Efficient-Large-Model/SANA-WM_bidirectional"
+DEFAULT_SANA_WM_STREAMING_MODEL_NAME_FOR_TEST = (
+    "Efficient-Large-Model/SANA-WM_streaming"
+)
+
 
 def print_value_formatted(description: str, value: int | float | str):
     """Helper function to print a metric value formatted."""

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

@@ -562,6 +562,7 @@ def test_raw_rgb_realtime_output_adapter_offloads_preview_encoding(monkeypatch):
     async def run():
         ws = _WebSocket()
         adapter = RawRGBRealtimeOutputAdapter()
+        frame_count = realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE + 1
         batch = SimpleNamespace(
             block_idx=0,
             request_id="req-webp-offload",
@@ -575,8 +576,8 @@ def test_raw_rgb_realtime_output_adapter_offloads_preview_encoding(monkeypatch):
         result = OutputBatch(
             raw_frame_batches=[
                 [
-                    bytes([255, 0, 0, 0, 255, 0]),
-                    bytes([0, 0, 255, 255, 255, 0]),
+                    bytes([idx % 256, 0, 0, 0, 255, idx % 256])
+                    for idx in range(frame_count)
                 ]
             ],
             raw_frame_content_type=RAW_RGB_CONTENT_TYPE,
@@ -595,18 +596,31 @@ def test_raw_rgb_realtime_output_adapter_offloads_preview_encoding(monkeypatch):
     payloads = asyncio.run(run())
 
     assert [call[0] for call in calls] == [
-        realtime_output_adapter._encode_rgb_frame_to_webp,
-        realtime_output_adapter._encode_rgb_frame_to_webp,
-    ]
-    [(header, payload)] = _unpack_frame_batch_messages(payloads)
-    assert header["content_type"] == WEBP_FRAME_CONTENT_TYPE
-    assert header["encoding"] == "webp"
-    assert header["num_frames"] == 2
-    assert header["frame_batch_index"] == 0
-    assert header["num_frame_batches"] == 1
-    assert header["is_final_frame_batch"] is True
-    assert len(header["payload_lengths"]) == 2
-    assert payload.startswith(b"RIFF")
+        realtime_output_adapter._encode_rgb_frame_to_webp
+    ] * (realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE + 1)
+    (first_header, first_payload), (second_header, second_payload) = (
+        _unpack_frame_batch_messages(payloads)
+    )
+    assert first_header["content_type"] == WEBP_FRAME_CONTENT_TYPE
+    assert first_header["encoding"] == "webp"
+    assert first_header["num_frames"] == (
+        realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE
+    )
+    assert first_header["frame_batch_index"] == 0
+    assert first_header["num_frame_batches"] == 2
+    assert first_header["is_final_frame_batch"] is False
+    assert len(first_header["payload_lengths"]) == (
+        realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE
+    )
+    assert second_header["content_type"] == WEBP_FRAME_CONTENT_TYPE
+    assert second_header["encoding"] == "webp"
+    assert second_header["num_frames"] == 1
+    assert second_header["frame_batch_index"] == 1
+    assert second_header["num_frame_batches"] == 2
+    assert second_header["is_final_frame_batch"] is True
+    assert len(second_header["payload_lengths"]) == 1
+    assert first_payload.startswith(b"RIFF")
+    assert second_payload.startswith(b"RIFF")
 
 
 def test_raw_rgb_realtime_output_adapter_can_send_jpeg_preview_frames():

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

@@ -22,6 +22,9 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.realtime import (
 from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters import (
     lingbot_world_realtime_adapter as lingbot_realtime,
 )
+from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters import (
+    sana_wm_realtime_adapter as sana_wm_realtime,
+)
 from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import (
     GenerateSession,
 )
@@ -35,10 +38,16 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import (
     LingBotWorldCausalDMDDenoisingStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import (
+    RealtimeDiffusionStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_input_validation import (
     RealtimeInputValidationStage,
     RealtimeInputValidationState,
 )
+from sglang.multimodal_gen.runtime.realtime.causal_state import (
+    RealtimeCausalDecodeState,
+)
 from sglang.multimodal_gen.runtime.realtime.condition_events import (
     ControlStateTransition,
 )
@@ -65,6 +74,51 @@ class _State(BaseRealtimeState):
         self.disposed = True
 
 
+def test_realtime_diffusion_stage_declares_long_lived_components():
+    stage = RealtimeDiffusionStage()
+    server_args = SimpleNamespace(
+        pipeline_config=SimpleNamespace(dit_precision="bf16", vae_precision="fp32")
+    )
+
+    uses = stage.component_uses(server_args, stage_name="realtime")
+
+    assert [use.component_name for use in uses] == ["transformer", "vae"]
+    assert [use.stage_name for use in uses] == ["realtime", "realtime"]
+    assert uses[0].target_dtype == torch.bfloat16
+    assert uses[0].memory_intensive
+    assert uses[0].keep_ready_after_warmup
+    assert uses[1].target_dtype == torch.float32
+    assert not uses[1].memory_intensive
+    assert uses[1].keep_ready_after_warmup
+
+
+def test_realtime_diffusion_stage_requires_session():
+    stage = RealtimeDiffusionStage(default_height=480, default_width=832)
+    req = _Req(session=None)
+
+    try:
+        stage.require_session(req, context="test realtime")
+    except ValueError as exc:
+        assert "test realtime requires a realtime session" in str(exc)
+    else:
+        raise AssertionError("expected missing realtime session to fail")
+
+    session = object()
+    req.session = session
+    assert stage.require_session(req) is session
+
+
+def test_realtime_causal_decode_state_dispose_resets_frontier():
+    state = RealtimeCausalDecodeState()
+    state.conv_cache = {"cache": object()}
+    state.next_dec_idx = 7
+
+    state.dispose()
+
+    assert state.conv_cache is None
+    assert state.next_dec_idx == 0
+
+
 def test_realtime_session_cache_reuses_and_releases_state():
     cache = RealtimeSessionCache(max_sessions=1)
     first = _Req(realtime_session_id="session-a", block_idx=0, session=None)
@@ -148,6 +202,53 @@ def test_lingbot_realtime_camera_state_compacts_multiple_pending_updates():
     assert state.latest_sampled_event_id == 9
 
 
+def test_sana_wm_realtime_camera_state_uses_sana_normalizer():
+    state = sana_wm_realtime.SanaWMRealtimeAdapterState()
+    result = state.receive_camera_event_payload(
+        {
+            "mode": "state",
+            "transitions": [
+                {"actions": ["W"], "client_ts_ms": 100},
+                {"actions": [], "client_ts_ms": 120},
+            ],
+        },
+        event_id=11,
+    )
+
+    assert result == "kind=camera_actions, mode=state, transitions=2"
+    assert state.sample_camera_actions(10) == [["w"]] * 8 + [[], []]
+    assert state.latest_sampled_event_id == 11
+
+
+def test_sana_wm_realtime_adapter_preserves_requested_size():
+    async def fake_save_image_to_path(image, target_path):
+        return target_path
+
+    old_save_image_to_path = sana_wm_realtime.save_image_to_path
+    old_get_global_server_args = sana_wm_realtime.get_global_server_args
+    sana_wm_realtime.save_image_to_path = fake_save_image_to_path
+    sana_wm_realtime.get_global_server_args = lambda: SimpleNamespace(
+        input_save_path=None
+    )
+    try:
+        adapter = sana_wm_realtime.SanaWMRealtimeAdapter()
+        session = GenerateSession()
+        session.set_adapter(adapter)
+        request = RealtimeVideoGenerationsRequest(
+            type="init",
+            prompt="walk forward",
+            first_frame=b"fake-image",
+            size="832x480",
+        )
+
+        asyncio.run(adapter.on_init(session, request))
+
+        assert request.size == "832x480"
+    finally:
+        sana_wm_realtime.save_image_to_path = old_save_image_to_path
+        sana_wm_realtime.get_global_server_args = old_get_global_server_args
+
+
 def test_lingbot_realtime_adapter_ingests_generic_events():
     adapter = lingbot_realtime.LingBotWorldRealtimeAdapter()
     session = GenerateSession()

python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Realtime chain stages — latent-prep plan arithmetic + noise discipline.
+
+The chunk PLAN must reproduce the engine-era tick behavior exactly:
+front-loaded segments within a fixed horizon (chunk 0 carries the remainder
+and the conditioning frame), uniform chunks past the horizon (seamless
+continuation), and the noise buffer sliced bitwise within the horizon.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+
+from sglang.multimodal_gen.runtime import server_args as _sa_mod
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_chain import (
+    SanaWMCameraCondStage,
+    SanaWMNoiseState,
+    SanaWMRealtimeLatentPrepStage,
+    SanaWMSessionInputsState,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_stage import (
+    SanaWMRealtimeStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import (
+    SanaWMStreamCacheState,
+)
+from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession
+from sglang.multimodal_gen.runtime.server_args import set_global_server_args
+
+MC = 8
+
+
+class _TestRealtimeStage(SanaWMRealtimeStage):
+    def forward(self, batch, server_args):
+        raise NotImplementedError
+
+
+@pytest.fixture
+def _global_args():
+    prev = _sa_mod._global_server_args
+    set_global_server_args(
+        SimpleNamespace(
+            comfyui_mode=False,
+            enable_cfg_parallel=False,
+            enable_torch_compile=False,
+            attention_backend=None,
+        )
+    )
+    try:
+        yield
+    finally:
+        set_global_server_args(prev)
+
+
+def _prep_stage():
+    return SanaWMRealtimeLatentPrepStage(
+        use_refiner=True, transformer=None, vae=None, model_path=""
+    )
+
+
+def _realtime_stage():
+    return object.__new__(_TestRealtimeStage)
+
+
+def _camera_stage():
+    return object.__new__(SanaWMCameraCondStage)
+
+
+def _batch(session, block_idx, image_latent):
+    return SimpleNamespace(
+        session=session,
+        block_idx=block_idx,
+        image_latent=image_latent,
+        seed=7,
+        generator=None,
+        extra={},
+        latents=None,
+    )
+
+
+def _server_args():
+    return SimpleNamespace(pipeline_config=SimpleNamespace(dit_precision="fp32"))
+
+
+def test_realtime_intrinsics_default_to_centered_heuristic():
+    stage = _realtime_stage()
+    state = SimpleNamespace(
+        intrinsics_raw=None,
+        intrinsics_image=Image.new("RGB", (832, 480)),
+    )
+
+    intrinsics = stage._prepare_intrinsics(
+        SimpleNamespace(condition_inputs={}),
+        state,
+        num_frames=3,
+        device=torch.device("cpu"),
+    )
+
+    assert intrinsics.shape == (3, 4)
+    assert intrinsics[0].tolist() == pytest.approx([665.6, 665.6, 416.0, 240.0])
+    assert intrinsics[2].tolist() == pytest.approx([665.6, 665.6, 416.0, 240.0])
+
+
+def test_realtime_first_frame_uses_requested_size():
+    stage = _realtime_stage()
+    batch = SimpleNamespace(
+        condition_image=Image.new("RGB", (640, 360)),
+        image_path=None,
+        height=480,
+        width=832,
+    )
+
+    cropped, original, src_size, resized_size, crop_offset = stage._prepare_image(batch)
+
+    assert cropped.size == (832, 480)
+    assert original.size == (640, 360)
+    assert src_size == (640, 360)
+    assert resized_size == (853, 480)
+    assert crop_offset == (10, 0)
+
+
+def test_realtime_camera_conditioning_uses_requested_size():
+    stage = _camera_stage()
+    inputs = SanaWMSessionInputsState()
+    inputs.src_size = (640, 360)
+    inputs.resized_size = (853, 480)
+    inputs.crop_offset = (10, 0)
+    inputs.target_height = 480
+    inputs.target_width = 832
+    inputs.intrinsics_image = Image.new("RGB", (640, 360))
+    inputs.open_ended = True
+    batch = SimpleNamespace(
+        condition_inputs={},
+        extra={},
+        height=480,
+        width=832,
+        num_frames=17,
+    )
+
+    camera, plucker = stage._build_camera_windows(
+        batch,
+        inputs,
+        target_latent=3,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    assert camera.shape == (1, 3, 20)
+    assert plucker.shape == (1, 48, 3, 15, 26)
+
+
+def test_latent_prep_plan_and_noise_discipline(_global_args):
+    stage = _prep_stage()
+    session = RealtimeSession()
+    fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
+
+    inputs = session.get_or_create_state(SanaWMSessionInputsState)
+    inputs.latent_t = 5  # fixed horizon: segments [0, 3, 5] for nfpb=2
+    inputs.num_frame_per_block = 2
+    inputs.sink_size = 1
+    cache = session.get_or_create_state(SanaWMStreamCacheState)
+
+    # Tick 0: front-loaded chunk 0 (cond + remainder) + full-horizon buffer.
+    batch = stage.forward(_batch(session, 0, fl), _server_args())
+    noise = session.get_or_create_state(SanaWMNoiseState)
+    assert noise.segments == [0, 3, 5]
+    assert noise.noise_buffer is not None and noise.noise_buffer.shape[2] == 5
+    assert batch.extra["sana_wm_chunk_plan"] == [3]
+    assert batch.latents.shape[2] == 3
+    assert torch.equal(batch.latents[:, :, :1].cpu(), fl)  # cond frame in front
+    assert torch.equal(  # noise sliced from the buffer, bitwise
+        batch.latents[:, :, 1:].cpu(), noise.noise_buffer[:, :, 1:3].cpu()
+    )
+
+    # Tick 1: stage-1 advanced to frame 3 (simulated) -> next segment [3, 5).
+    cache.chunk_indices = [0, 3]
+    cache.chunk_idx = 1
+    batch = stage.forward(_batch(session, 1, fl), _server_args())
+    assert batch.extra["sana_wm_chunk_plan"] == [2]
+    assert torch.equal(batch.latents.cpu(), noise.noise_buffer[:, :, 3:5].cpu())
+
+    # Tick 2: horizon exhausted -> seamless continuation, uniform chunk from
+    # the seeded fallback generator (no reset, no rollover).
+    cache.chunk_indices = [0, 3, 5]
+    cache.chunk_idx = 2
+    batch = stage.forward(_batch(session, 2, fl), _server_args())
+    assert batch.extra["sana_wm_chunk_plan"] == [2]
+    assert batch.latents.shape[2] == 2
+    assert torch.isfinite(batch.latents).all()
+
+
+def test_latent_prep_open_ended_uniform_chunk0(_global_args):
+    stage = _prep_stage()
+    session = RealtimeSession()
+    fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
+
+    inputs = session.get_or_create_state(SanaWMSessionInputsState)
+    inputs.latent_t = None  # open-ended
+    inputs.num_frame_per_block = 3
+    inputs.sink_size = 1
+    session.get_or_create_state(SanaWMStreamCacheState)
+
+    batch = stage.forward(_batch(session, 0, fl), _server_args())
+    noise = session.get_or_create_state(SanaWMNoiseState)
+    assert noise.segments is None and noise.noise_buffer is None
+    # Uniform grid chunk 0 = cond frame + nfpb new frames.
+    assert batch.extra["sana_wm_chunk_plan"] == [4]
+    assert batch.latents.shape[2] == 4
+    assert torch.equal(batch.latents[:, :, :1].cpu(), fl)

python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_cached.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+"""S1a-2a tests — chunk-causal cached scan combiners for streaming `forward_long`.
+
+`_gdn_scan_cached` / `_single_path_delta_scan_cached` are the per-chunk combiners
+the streaming path calls: a FORWARD (inclusive) pass that carries recurrent state
+across chunks plus a BACKWARD (exclusive) pass recomputed intra-chunk. Two
+invariants pin them down:
+
+1. **Reduce-to-bidirectional** (#26153-only): a single chunk with no carried
+   state must equal the validated dense `_gdn_scan_bidirectional` /
+   `_single_path_delta_scan_bidirectional`.
+2. **Port-fidelity** (vs the reference `minimal-sanawm`): output AND carried state
+   must match `recurrent_gdn_cached` / `recurrent_delta_cached` chunk-for-chunk,
+   including a 2-chunk state-carry run.
+
+FP64, CPU.
+"""
+
+from __future__ import annotations
+
+import pathlib
+import sys
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
+    BidirectionalGDNUCPESinglePathLiteLA,
+    GLUMBConvTemp,
+    _gdn_scan_bidirectional,
+    _gdn_scan_cached,
+    _ShortConvolution,
+    _single_path_delta_scan_bidirectional,
+    _single_path_delta_scan_cached,
+    _temporal_short_conv_cached,
+)
+
+B, H, D, T, S = 1, 2, 4, 6, 3
+N = T * S
+HW = (T, S, 1)  # H_sp * W_sp == S
+
+
+def _rand(*shape):
+    return torch.randn(*shape, dtype=torch.float64)
+
+
+def _inputs():
+    torch.manual_seed(0)
+    q, k, v = _rand(B, H, D, N), _rand(B, H, D, N), _rand(B, H, D, N)
+    q_rot, k_rot = _rand(B, H, D, N), _rand(B, H, D, N)
+    beta = torch.rand(B, H, T, S, dtype=torch.float64)
+    decay = torch.rand(B, H, T, dtype=torch.float64) * 0.5 + 0.4
+    return q, k, v, q_rot, k_rot, beta, decay
+
+
+def _frame_slice(x, f0, f1):
+    # x is (B, H, D, N=T*S), frame-major (T outer, S inner)
+    return x[..., f0 * S : f1 * S].contiguous()
+
+
+# --------------------------------------------------------------------------- #
+# 1. Reduce-to-bidirectional (uses only #26153 functions)
+# --------------------------------------------------------------------------- #
+
+
+def test_gdn_cached_single_chunk_reduces_to_bidirectional():
+    q, k, v, q_rot, k_rot, beta, decay = _inputs()
+    ref = _gdn_scan_bidirectional(q, k, v, q_rot, k_rot, beta, decay, HW)
+    out, (kv, z) = _gdn_scan_cached(q, k, v, q_rot, k_rot, beta, decay)
+    torch.testing.assert_close(out, ref, atol=1e-9, rtol=0)
+    assert kv.shape == (B, H, D, D) and z.shape == (B, H, D, 1)
+
+
+def test_cam_cached_single_chunk_reduces_to_bidirectional():
+    _, _, v, q_rot, k_rot, beta, decay = _inputs()
+    ref = _single_path_delta_scan_bidirectional(q_rot, k_rot, v, beta, decay, HW)
+    out, kv = _single_path_delta_scan_cached(q_rot, k_rot, v, beta, decay)
+    torch.testing.assert_close(out, ref, atol=1e-9, rtol=0)
+    assert kv.shape == (B, H, D, D)
+
+
+# --------------------------------------------------------------------------- #
+# 2. Port-fidelity vs reference (gold standard; skip if reference absent)
+# --------------------------------------------------------------------------- #
+
+_REF_DIR = pathlib.Path("/sgl-workspace/sglang/myUtils/sana/minimal-sanawm")
+
+
+def _load_reference():
+    if not (_REF_DIR / "components.py").exists():
+        pytest.skip(f"reference impl not found at {_REF_DIR}")
+    if str(_REF_DIR) not in sys.path:
+        sys.path.insert(0, str(_REF_DIR))
+    try:
+        import components as ref  # type: ignore
+    except Exception as exc:  # pragma: no cover - environment dependent
+        pytest.skip(f"reference components.py not importable: {exc}")
+    return ref
+
+
+def test_gdn_cached_matches_reference_single_and_chunked():
+    ref = _load_reference()
+    q, k, v, q_rot, k_rot, beta, decay = _inputs()
+
+    # Single chunk, no carried state.
+    mine, (kv_m, z_m) = _gdn_scan_cached(q, k, v, q_rot, k_rot, beta, decay)
+    r_out, kv_r, z_r, _ = ref.recurrent_gdn_cached(
+        q, k, v, q_rot, k_rot, beta, decay, eps=1e-6, kv_state=None, z_state=None
+    )
+    torch.testing.assert_close(mine, r_out, atol=1e-9, rtol=0)
+    torch.testing.assert_close(kv_m, kv_r, atol=1e-9, rtol=0)
+    torch.testing.assert_close(z_m, z_r, atol=1e-9, rtol=0)
+
+    # Two chunks carrying state — mine vs reference, chunk-for-chunk.
+    split = 2
+
+    def chunk(x, f0, f1):
+        return _frame_slice(x, f0, f1)
+
+    args0 = (
+        chunk(q, 0, split),
+        chunk(k, 0, split),
+        chunk(v, 0, split),
+        chunk(q_rot, 0, split),
+        chunk(k_rot, 0, split),
+        beta[:, :, :split],
+        decay[:, :, :split],
+    )
+    args1 = (
+        chunk(q, split, T),
+        chunk(k, split, T),
+        chunk(v, split, T),
+        chunk(q_rot, split, T),
+        chunk(k_rot, split, T),
+        beta[:, :, split:],
+        decay[:, :, split:],
+    )
+
+    m0, (mkv, mz) = _gdn_scan_cached(*args0)
+    m1, _ = _gdn_scan_cached(*args1, init_state_kv=mkv, init_state_z=mz)
+
+    r0, rkv, rz, _ = ref.recurrent_gdn_cached(
+        *args0, eps=1e-6, kv_state=None, z_state=None
+    )
+    r1, _, _, _ = ref.recurrent_gdn_cached(*args1, eps=1e-6, kv_state=rkv, z_state=rz)
+
+    torch.testing.assert_close(m0, r0, atol=1e-9, rtol=0)
+    torch.testing.assert_close(m1, r1, atol=1e-9, rtol=0)
+
+
+def test_cam_cached_matches_reference_single_and_chunked():
+    ref = _load_reference()
+    _, _, v, q_rot, k_rot, beta, decay = _inputs()
+
+    mine, kv_m = _single_path_delta_scan_cached(q_rot, k_rot, v, beta, decay)
+    r_out, kv_r = ref.recurrent_delta_cached(q_rot, k_rot, v, beta, decay, state=None)
+    torch.testing.assert_close(mine, r_out, atol=1e-9, rtol=0)
+    torch.testing.assert_close(kv_m, kv_r, atol=1e-9, rtol=0)
+
+    split = 3
+    a0 = (
+        _frame_slice(q_rot, 0, split),
+        _frame_slice(k_rot, 0, split),
+        _frame_slice(v, 0, split),
+        beta[:, :, :split],
+        decay[:, :, :split],
+    )
+    a1 = (
+        _frame_slice(q_rot, split, T),
+        _frame_slice(k_rot, split, T),
+        _frame_slice(v, split, T),
+        beta[:, :, split:],
+        decay[:, :, split:],
+    )
+
+    m0, mkv = _single_path_delta_scan_cached(*a0)
+    m1, _ = _single_path_delta_scan_cached(*a1, init_state_kv=mkv)
+    r0, rkv = ref.recurrent_delta_cached(*a0, state=None)
+    r1, _ = ref.recurrent_delta_cached(*a1, state=rkv)
+
+    torch.testing.assert_close(m0, r0, atol=1e-9, rtol=0)
+    torch.testing.assert_close(m1, r1, atol=1e-9, rtol=0)
+
+
+# --------------------------------------------------------------------------- #
+# 3. Cached short conv on K (cache slot 4)
+# --------------------------------------------------------------------------- #
+
+C_CONV, KERN = 5, 4
+
+
+def _conv():
+    conv = _ShortConvolution(C_CONV, KERN).double()
+    with torch.no_grad():  # randomize off the identity init for a non-trivial filter
+        conv.weight.copy_(torch.randn(C_CONV, 1, KERN, dtype=torch.float64))
+    return conv
+
+
+def _nslice(x, f0, f1):
+    # x is (B, N=T*S, C), frame-major
+    return x[:, f0 * S : f1 * S, :].contiguous()
+
+
+def test_short_conv_cached_single_chunk_reduces_to_bidirectional():
+    torch.manual_seed(1)
+    x = torch.randn(B, N, C_CONV, dtype=torch.float64)
+    conv = _conv()
+    ref = BidirectionalGDNUCPESinglePathLiteLA._temporal_short_conv(
+        x, conv, (T, S, 1), bidirectional=True
+    )
+    out, prefix = _temporal_short_conv_cached(x, conv, (T, S, 1))
+    torch.testing.assert_close(out, ref, atol=1e-9, rtol=0)
+    assert prefix.shape == (B * S, KERN - 1, C_CONV)
+
+
+def test_short_conv_cached_forward_continuity_across_chunks():
+    # The forward (causal) direction must be chunk-invariant via the prefix carry.
+    torch.manual_seed(2)
+    x = torch.randn(B, N, C_CONV, dtype=torch.float64)
+    conv = _conv()
+    whole, _ = _temporal_short_conv_cached(x, conv, (T, S, 1), bidirectional=False)
+
+    split = 2
+    o0, p0 = _temporal_short_conv_cached(
+        _nslice(x, 0, split), conv, (split, S, 1), bidirectional=False
+    )
+    o1, _ = _temporal_short_conv_cached(
+        _nslice(x, split, T), conv, (T - split, S, 1), prefix=p0, bidirectional=False
+    )
+    chunked = torch.cat([o0, o1], dim=1)  # reassemble frame-major
+    torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0)
+
+
+# --------------------------------------------------------------------------- #
+# 4. Cached FFN temporal tail (GLUMBConvTemp, cache slot 9)
+# --------------------------------------------------------------------------- #
+
+C_FFN, HID = 6, 8
+HFFN, WFFN = S, 1  # S spatial tokens per frame
+
+
+def _ffn():
+    m = GLUMBConvTemp(C_FFN, HID, t_kernel_size=3).double().eval()
+    with torch.no_grad():  # zero-init t_conv -> randomize for a non-trivial temporal filter
+        m.t_conv.weight.copy_(torch.randn_like(m.t_conv.weight))
+    return m
+
+
+def test_ffn_cached_no_prefix_reduces_to_dense():
+    torch.manual_seed(3)
+    x = torch.randn(B, N, C_FFN, dtype=torch.float64)
+    m = _ffn()
+    dense = m(x, (T, HFFN, WFFN))
+    cached, tail = m(x, (T, HFFN, WFFN), save_ffn_tail=True)
+    torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0)
+    assert tail.shape[2] == m.t_conv.kernel_size[0] // 2
+
+
+def test_ffn_cached_final_chunk_matches_whole():
+    # The prefix supplies real left context, so the LAST chunk (right edge = end
+    # of sequence, same as the whole pass) matches the whole-sequence output.
+    torch.manual_seed(4)
+    x = torch.randn(B, N, C_FFN, dtype=torch.float64)
+    m = _ffn()
+    whole, _ = m(x, (T, HFFN, WFFN), save_ffn_tail=True)
+
+    split = 2
+    _, tail0 = m(_nslice(x, 0, split), (split, HFFN, WFFN), save_ffn_tail=True)
+    o1, _ = m(
+        _nslice(x, split, T),
+        (T - split, HFFN, WFFN),
+        ffn_tail=tail0,
+        save_ffn_tail=True,
+    )
+    torch.testing.assert_close(whole[:, split * S :, :], o1, atol=1e-9, rtol=0)

python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py

@@ -0,0 +1,397 @@
+# SPDX-License-Identifier: Apache-2.0
+"""S1a-2c tests — streaming `forward_long` assembly on the #26153 SANA-WM DiT.
+
+Built up in stages, each pinned at fp64 / CPU / atol=1e-9:
+  Stage 0 — RoPE windowing: `WanRotaryPosEmbed` must produce freqs at GLOBAL
+    frame positions for a chunk `[start, end)`, equal to the corresponding slice
+    of the full table. This is the prerequisite that keeps a chunk's queries
+    aligned with the carried K and the softmax concat-window.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.sana_wm import (
+    SanaWMArchConfig,
+    SanaWMConfig,
+)
+from sglang.multimodal_gen.runtime import server_args as _sa_mod
+from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
+    _CACHE_TYPE_STATE,
+    _SLOT_CAM_K,
+    _SLOT_FFN_TCONV,
+    _SLOT_K,
+    _SLOT_TYPE_FLAG,
+    _SLOT_V,
+    BidirectionalGDNUCPESinglePathLiteLA,
+    SanaWMBlock,
+    SanaWMTransformer3DModel,
+    WanRotaryPosEmbed,
+    _build_ucpe_apply_fns,
+    _slice_rope_to_current_chunk,
+    process_camera_conditions_ucpe,
+)
+from sglang.multimodal_gen.runtime.server_args import set_global_server_args
+
+HEAD_DIM = 112
+H, W = 2, 3
+S = H * W
+T = 6
+
+
+def _rope():
+    return WanRotaryPosEmbed(attention_head_dim=HEAD_DIM, patch_size=(1, 1, 1))
+
+
+def test_rope_full_equals_range_from_zero():
+    rope = _rope()
+    dense = rope((T, H, W), torch.device("cpu"))
+    ranged = rope(((0, T), H, W), torch.device("cpu"))
+    torch.testing.assert_close(ranged, dense, atol=1e-9, rtol=0)
+
+
+def test_rope_window_equals_slice_of_full():
+    rope = _rope()
+    start, end = 2, 5
+    full = rope((end, H, W), torch.device("cpu"))  # frames 0..end-1
+    window = rope(((start, end), H, W), torch.device("cpu"))  # frames start..end-1
+    # full is frame-major (1, 1, end*S, D/2); slice frames [start:end].
+    torch.testing.assert_close(
+        window, full[:, :, start * S : end * S, :], atol=1e-9, rtol=0
+    )
+
+
+def test_rope_frame_index_overrides_to_global_positions():
+    rope = _rope()
+    start, end = 2, 5
+    fidx = torch.arange(start, end)
+    by_index = rope((end - start, H, W), torch.device("cpu"), frame_index=fidx)
+    by_range = rope(((start, end), H, W), torch.device("cpu"))
+    torch.testing.assert_close(by_index, by_range, atol=1e-9, rtol=0)
+
+
+def test_slice_rope_to_current_chunk_is_noop_when_sized():
+    rope = _rope()
+    freqs = rope(((1, 4), H, W), torch.device("cpu"))  # 3 frames -> 3*S tokens
+    assert _slice_rope_to_current_chunk(freqs, 3 * S) is freqs
+    assert _slice_rope_to_current_chunk(None, 3 * S) is None
+    # wider table trimmed to the trailing chunk
+    wide = rope(((0, 4), H, W), torch.device("cpu"))  # 4 frames
+    trimmed = _slice_rope_to_current_chunk(wide, 3 * S)
+    torch.testing.assert_close(trimmed, wide[:, :, -3 * S :, :], atol=1e-9, rtol=0)
+
+
+# --------------------------------------------------------------------------- #
+# Stage 1 — attention-level forward_long: cached branch methods
+# --------------------------------------------------------------------------- #
+
+AB, AHEADS, ADIM = (
+    1,
+    2,
+    16,
+)  # head_dim/2 divisible by 4 (UCPE homog. 4-vectors), like 112
+AC = AHEADS * ADIM  # in_dim == heads*head_dim
+AHW = (T, H, W)
+AN = T * S
+
+
+def _arope():
+    # RoPE table sized to the attention's head_dim (NOT the Stage-0 HEAD_DIM).
+    return WanRotaryPosEmbed(attention_head_dim=ADIM, patch_size=(1, 1, 1))
+
+
+def _attn(softmax_main=False):
+    m = (
+        BidirectionalGDNUCPESinglePathLiteLA(
+            in_dim=AC,
+            heads=AHEADS,
+            head_dim=ADIM,
+            update_rule="torch_recurrent",  # match the cached recurrent scan exactly
+            cam_update_rule="torch_recurrent",
+            softmax_main=softmax_main,
+            use_chunked_softmax_attention=softmax_main,
+        )
+        .double()
+        .eval()
+    )
+    return m
+
+
+def _softmax_attn():
+    # Build a GDN module (no LocalAttention -> no server-args/backend needed),
+    # then flip to softmax mode: softmax blocks have no short conv, and the
+    # cached softmax path only reads softmax_attn.softmax_scale (a scalar).
+    m = _attn(softmax_main=False)
+    m.softmax_main = True
+    m.conv_k = None
+    m.conv_k_cam = None
+    m.softmax_attn = SimpleNamespace(softmax_scale=ADIM**-0.5)
+    return m
+
+
+def _x():
+    torch.manual_seed(5)
+    return torch.randn(AB, AN, AC, dtype=torch.float64)
+
+
+def _prope(start=0, end=T):
+    # Build UCPE apply fns co-windowed with freqs for frames [start, end).
+    torch.manual_seed(7)
+    cam = torch.randn(AB, T, 20, dtype=torch.float64)[:, start:end]
+    raymats = process_camera_conditions_ucpe(
+        cam, HW=(end - start, H, W), patch_size=(1, 1, 1)
+    )
+    raymats_flat = raymats.reshape(AB, -1, 4, 4)
+    freqs = _arope()(((start, end), H, W), torch.device("cpu"))
+    return _build_ucpe_apply_fns(ADIM, raymats_flat, freqs), freqs
+
+
+def _empty_cache():
+    return [None] * 10
+
+
+def test_main_gdn_cached_single_chunk_reduces_to_dense():
+    attn = _attn(softmax_main=False)
+    x = _x()
+    rope_emb = _arope()((T, H, W), torch.device("cpu"))
+    dense, _, _ = attn._main_branch_gdn(x, AHW, rope_emb)
+    cache = _empty_cache()
+    cached, _, _ = attn._main_branch_gdn_cached(x, AHW, rope_emb, cache, True)
+    torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0)
+    assert cache[_SLOT_K].shape == (AB, AHEADS, ADIM, ADIM)
+    assert cache[_SLOT_V].shape == (AB, AHEADS, ADIM, 1)
+    assert float(cache[_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_STATE
+
+
+def test_cam_gdn_cached_single_chunk_reduces_to_dense():
+    attn = _attn(softmax_main=False)
+    x = _x()
+    (apply_q, apply_kv, apply_o), _ = _prope()
+    _, beta, decay = attn._main_branch_gdn(x, AHW, None)
+    dense = attn._cam_branch(x, AHW, apply_q, apply_kv, apply_o, beta, decay)
+    cache = _empty_cache()
+    cached = attn._cam_branch_cached(
+        x, AHW, apply_q, apply_kv, apply_o, beta, decay, cache, True
+    )
+    torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0)
+    assert cache[_SLOT_CAM_K].shape == (AB, AHEADS, ADIM, ADIM)
+
+
+def test_softmax_main_cached_last_chunk_matches_whole():
+    attn = _softmax_attn()
+    x = _x()
+    rope_full = _arope()(((0, T), H, W), torch.device("cpu"))
+    whole, _, _ = attn._main_branch_softmax_cached(
+        x, AHW, rope_full, _empty_cache(), True
+    )
+
+    split = 2
+    rope0 = _arope()(((0, split), H, W), torch.device("cpu"))
+    rope1 = _arope()(((split, T), H, W), torch.device("cpu"))
+    cache = _empty_cache()
+    _o0, _, _ = attn._main_branch_softmax_cached(
+        x[:, : split * S, :], (split, H, W), rope0, cache, True
+    )
+    o1, _, _ = attn._main_branch_softmax_cached(
+        x[:, split * S :, :], (T - split, H, W), rope1, cache, True
+    )
+    # Last chunk's queries attend to [chunk0_K || chunk1_K] == full K.
+    torch.testing.assert_close(o1, whole[:, split * S :, :], atol=1e-9, rtol=0)
+
+
+def test_forward_long_gdn_reduces_to_dense_no_camera():
+    attn = _attn(softmax_main=False)
+    x = _x()
+    rope_emb = _arope()((T, H, W), torch.device("cpu"))
+    dense = attn(x, AHW, rope_emb, None)
+    cache = _empty_cache()
+    cached, ret = attn.forward_long(
+        x, AHW, rope_emb, None, kv_cache=cache, save_kv_cache=True
+    )
+    torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0)
+    assert ret is cache and cache[_SLOT_K] is not None
+
+
+def test_forward_long_gdn_reduces_to_dense_with_camera():
+    attn = _attn(softmax_main=False)
+    x = _x()
+    (apply_q, apply_kv, apply_o), freqs = _prope()
+    prope_fns = (apply_q, apply_kv, apply_o)
+    dense = attn(x, AHW, freqs, prope_fns)
+    cache = _empty_cache()
+    cached, _ = attn.forward_long(
+        x, AHW, freqs, prope_fns, kv_cache=cache, save_kv_cache=True
+    )
+    torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0)
+    assert cache[_SLOT_CAM_K] is not None
+
+
+# --------------------------------------------------------------------------- #
+# Stage 2 — block-level forward_long (cross-attn stubbed; needs server args
+# only to construct the cross-attn's LocalAttention)
+# --------------------------------------------------------------------------- #
+
+
+@pytest.fixture
+def _global_args():
+    prev = _sa_mod._global_server_args
+    set_global_server_args(
+        SimpleNamespace(
+            comfyui_mode=False,
+            enable_cfg_parallel=False,
+            enable_torch_compile=False,
+            attention_backend=None,
+        )
+    )
+    try:
+        yield
+    finally:
+        set_global_server_args(prev)
+
+
+class _ZeroCross(torch.nn.Module):
+    def forward(self, x, y, mask=None):
+        return torch.zeros_like(x)
+
+
+def _block():
+    b = (
+        SanaWMBlock(
+            hidden_size=AC,
+            num_heads=AHEADS,
+            head_dim=ADIM,
+            mlp_ratio=2.0,
+            t_kernel_size=3,
+            qk_norm=True,
+            cross_norm=True,
+            conv_kernel_size=4,
+            k_conv_only=True,
+            softmax_main=False,
+            use_chunk_plucker_post_attn=False,
+            update_rule="torch_recurrent",
+            cam_update_rule="torch_recurrent",
+        )
+        .double()
+        .eval()
+    )
+    # Stub cross-attn (unchanged/uncached in forward_long) to avoid the CUDA
+    # attention backend; both forward and forward_long add the same zero.
+    b.cross_attn = _ZeroCross()
+    return b
+
+
+def test_block_forward_long_reduces_to_dense(_global_args):
+    block = _block()
+    x = _x()
+    y = torch.randn(AB, 4, AC, dtype=torch.float64)
+    t6 = torch.randn(AB, 1, T, 6 * AC, dtype=torch.float64)
+    (apply_q, apply_kv, apply_o), freqs = _prope()
+    prope_fns = (apply_q, apply_kv, apply_o)
+
+    dense = block(x, y, t6, AHW, freqs, prope_fns, None, None)
+    cache = _empty_cache()
+    out, ret = block.forward_long(
+        x, y, t6, AHW, freqs, prope_fns, None, None, kv_cache=cache, save_kv_cache=True
+    )
+    torch.testing.assert_close(out, dense, atol=1e-9, rtol=0)
+    assert ret is cache
+    assert cache[_SLOT_K] is not None  # main GDN state
+    assert cache[_SLOT_CAM_K] is not None  # cam state
+    assert cache[_SLOT_FFN_TCONV] is not None  # FFN temporal tail
+
+
+# --------------------------------------------------------------------------- #
+# Stage 3 — model-level forward_long (tiny CPU model; depth 2 => all-GDN blocks,
+# no softmax/LocalAttention in the main path; cross-attn stubbed)
+# --------------------------------------------------------------------------- #
+
+MB, MC, MT, MH, MW = 1, 8, 4, 2, 2
+
+
+def _tiny_model():
+    arch = SanaWMArchConfig(
+        in_channels=MC,
+        out_channels=MC,
+        num_layers=2,  # blocks 0,1 -> (i+1)%4 != 0 -> no softmax blocks
+        num_attention_heads=2,
+        attention_head_dim=16,
+        linear_head_dim=16,
+        num_cross_attention_heads=2,
+        cross_attention_head_dim=16,
+        cross_attention_dim=32,
+        caption_channels=32,
+        model_max_length=8,
+        softmax_every_n=4,
+        update_rule="torch_recurrent",
+        cam_update_rule="torch_recurrent",
+        chunk_size=None,
+    )
+    m = SanaWMTransformer3DModel(SanaWMConfig(arch_config=arch)).double().eval()
+    for b in m.blocks:
+        b.cross_attn = _ZeroCross()
+    return m
+
+
+def _model_inputs():
+    torch.manual_seed(0)
+    return dict(
+        hidden_states=torch.randn(MB, MC, MT, MH, MW, dtype=torch.float64),
+        encoder_hidden_states=torch.randn(MB, 4, 32, dtype=torch.float64),
+        timestep=torch.randint(
+            1, 1000, (MB, 1, MT)
+        ).double(),  # framewise for both paths
+        camera_conditions=torch.randn(MB, MT, 20, dtype=torch.float64),
+        chunk_plucker=torch.randn(MB, 48, MT, MH, MW, dtype=torch.float64),
+    )
+
+
+def test_model_forward_long_single_chunk_reduces_to_dense(_global_args):
+    m = _tiny_model()
+    inp = _model_inputs()
+    with torch.no_grad():
+        dense = m(**inp)
+        out, cache = m.forward_long(
+            **inp, kv_cache=None, save_kv_cache=True, start_f=0, end_f=MT
+        )
+    torch.testing.assert_close(out, dense, atol=1e-9, rtol=0)
+    assert len(cache) == len(m.blocks)
+    assert cache[0][_SLOT_K] is not None
+    assert cache[0][_SLOT_CAM_K] is not None
+    assert cache[0][_SLOT_FFN_TCONV] is not None
+
+
+def test_model_forward_long_two_chunks_runs_and_windows(_global_args):
+    m = _tiny_model()
+    inp = _model_inputs()
+    split = 2
+    with torch.no_grad():
+        o0, c0 = m.forward_long(
+            hidden_states=inp["hidden_states"][:, :, :split],
+            encoder_hidden_states=inp["encoder_hidden_states"],
+            timestep=inp["timestep"][:, :, :split],
+            camera_conditions=inp["camera_conditions"],  # full; sliced internally
+            chunk_plucker=inp["chunk_plucker"],
+            kv_cache=None,
+            save_kv_cache=True,
+            start_f=0,
+            end_f=split,
+        )
+        o1, c1 = m.forward_long(
+            hidden_states=inp["hidden_states"][:, :, split:],
+            encoder_hidden_states=inp["encoder_hidden_states"],
+            timestep=inp["timestep"][:, :, split:],
+            camera_conditions=inp["camera_conditions"],
+            chunk_plucker=inp["chunk_plucker"],
+            kv_cache=c0,
+            save_kv_cache=True,
+            start_f=split,
+            end_f=MT,
+        )
+    assert o0.shape == (MB, MC, split, MH, MW)
+    assert o1.shape == (MB, MC, MT - split, MH, MW)
+    assert torch.isfinite(o0).all() and torch.isfinite(o1).all()
+    assert c1 is not None and len(c1) == len(m.blocks)

python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Realtime (per-chunk) path of SanaWMStreamingDenoisingStage.
+
+Successor of the retired SanaWMChunkGenerator tests: a tiny CPU-safe model is
+ticked through stage.forward with a session attached (one chunk per call,
+chunk noise supplied directly — the latent-prep stage is exercised separately).
+The stage must carry the per-session 10-slot KV cache + the growing latent
+across ticks, evict stale chunk entries, stay deterministic for identical
+inputs, and reset on block_idx == 0.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.sana_wm import (
+    SanaWMArchConfig,
+    SanaWMConfig,
+)
+from sglang.multimodal_gen.runtime import server_args as _sa_mod
+from sglang.multimodal_gen.runtime.models.dits.sana_wm import SanaWMTransformer3DModel
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import (
+    SanaWMStreamCacheState,
+    SanaWMStreamingDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession
+from sglang.multimodal_gen.runtime.server_args import set_global_server_args
+
+MC = 8
+
+
+class _ZeroCross(torch.nn.Module):
+    def forward(self, x, y, mask=None):
+        return torch.zeros_like(x)
+
+
+@pytest.fixture
+def _global_args():
+    prev = _sa_mod._global_server_args
+    set_global_server_args(
+        SimpleNamespace(
+            comfyui_mode=False,
+            enable_cfg_parallel=False,
+            enable_torch_compile=False,
+            attention_backend=None,
+            # DenoisingStage.__init__ reads this for its CFG-parallel plumbing.
+            pipeline_config=SimpleNamespace(
+                dit_config=SimpleNamespace(hidden_size=32, num_attention_heads=2)
+            ),
+        )
+    )
+    try:
+        yield
+    finally:
+        set_global_server_args(prev)
+
+
+def _tiny_model():
+    arch = SanaWMArchConfig(
+        in_channels=MC,
+        out_channels=MC,
+        num_layers=2,  # GDN-only -> CPU-safe main path
+        num_attention_heads=2,
+        attention_head_dim=16,
+        linear_head_dim=16,
+        num_cross_attention_heads=2,
+        cross_attention_head_dim=16,
+        cross_attention_dim=32,
+        caption_channels=32,
+        model_max_length=8,
+        softmax_every_n=4,
+        update_rule="torch_recurrent",
+        cam_update_rule="torch_recurrent",
+        chunk_size=None,
+    )
+    m = SanaWMTransformer3DModel(SanaWMConfig(arch_config=arch)).float().eval()
+    for b in m.blocks:
+        b.cross_attn = _ZeroCross()
+    return m
+
+
+def _stage_and_args(nfpb: int = 2):
+    stage = SanaWMStreamingDenoisingStage(transformer=_tiny_model())
+    prompt = torch.zeros(1, 4, 32, dtype=torch.float32)
+    pcfg = SimpleNamespace(
+        dit_precision="fp32",
+        num_frame_per_block=nfpb,
+        num_cached_blocks=2,
+        sink_token=True,
+        denoising_step_list=(1000, 700, 0),
+        streaming_cfg_scale=1.0,
+        get_pos_prompt_embeds=lambda batch: [prompt],
+        get_neg_prompt_embeds=lambda batch: [],
+    )
+    server_args = SimpleNamespace(pipeline_config=pcfg, enable_cfg_parallel=False)
+    return stage, server_args
+
+
+def _tick(session, block_idx: int, chunk_lat: torch.Tensor, plan: list[int]):
+    return SimpleNamespace(
+        session=session,
+        block_idx=block_idx,
+        latents=chunk_lat,
+        extra={"sana_wm_chunk_plan": plan},
+        prompt_attention_mask=None,
+        negative_attention_mask=None,
+        do_classifier_free_guidance=False,
+    )
+
+
+def _noise(*shape, seed: int):
+    g = torch.Generator().manual_seed(seed)
+    return torch.randn(*shape, dtype=torch.float32, generator=g)
+
+
+def test_realtime_path_multi_tick_carries_state(_global_args):
+    stage, server_args = _stage_and_args(nfpb=2)
+    session = RealtimeSession()
+    fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
+
+    # Tick 0 (chunk 0): conditioning frame + 2 new frames.
+    chunk0 = torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], dim=2)
+    out = stage.forward(_tick(session, 0, chunk0, [3]), server_args)
+    state = session.get_or_create_state(SanaWMStreamCacheState)
+    assert out.latents.shape[2] == 3
+    assert state.chunk_idx == 1 and state.chunk_indices == [0, 3]
+    assert state.stream_kv_cache[0][0][0] is not None  # GDN state stored
+    # The condition frame is held fixed.
+    assert torch.allclose(state.latents[:, :, 0].cpu(), fl[:, :, 0])
+
+    # Tick 1: 2 more frames, KV carried.
+    out = stage.forward(
+        _tick(session, 1, _noise(1, MC, 2, 2, 2, seed=2), [2]), server_args
+    )
+    assert out.latents.shape[2] == 5
+    assert state.chunk_idx == 2 and state.chunk_indices == [0, 3, 5]
+    assert torch.isfinite(state.latents).all()
+
+    # Boundary-style tick: a TWO-chunk plan in one call.
+    out = stage.forward(
+        _tick(session, 2, _noise(1, MC, 4, 2, 2, seed=3), [2, 2]), server_args
+    )
+    assert out.latents.shape[2] == 9
+    assert state.chunk_idx == 4 and state.chunk_indices == [0, 3, 5, 7, 9]
+
+
+def test_realtime_path_evicts_stale_kv(_global_args):
+    stage, server_args = _stage_and_args(nfpb=2)
+    session = RealtimeSession()
+    fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
+    stage.forward(
+        _tick(session, 0, torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], 2), [3]),
+        server_args,
+    )
+    for i in range(1, 5):
+        stage.forward(
+            _tick(session, i, _noise(1, MC, 2, 2, 2, seed=10 + i), [2]), server_args
+        )
+    state = session.get_or_create_state(SanaWMStreamCacheState)
+    assert state.chunk_indices == [0, 3, 5, 7, 9, 11]
+
+    def _has_any(entry):
+        return any(slot is not None for block in entry for slot in block)
+
+    kept = [i for i, e in enumerate(state.stream_kv_cache) if _has_any(e)]
+    # Sink chunk + the last num_cached_blocks chunks (accumulate's read window).
+    assert kept == [0, 3, 4]
+
+
+def test_realtime_path_is_deterministic(_global_args):
+    def _run():
+        stage, server_args = _stage_and_args(nfpb=2)
+        torch.manual_seed(0)  # tiny-model init inside _stage_and_args uses global RNG
+        session = RealtimeSession()
+        fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
+        stage.forward(
+            _tick(session, 0, torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=5)], 2), [3]),
+            server_args,
+        )
+        stage.forward(
+            _tick(session, 1, _noise(1, MC, 2, 2, 2, seed=6), [2]), server_args
+        )
+        return session.get_or_create_state(SanaWMStreamCacheState).latents.cpu()
+
+    torch.manual_seed(1234)
+    a = _run()
+    torch.manual_seed(1234)
+    b = _run()
+    assert torch.equal(a, b)
+
+
+def test_realtime_path_resets_on_block_zero(_global_args):
+    stage, server_args = _stage_and_args(nfpb=2)
+    session = RealtimeSession()
+    fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
+    chunk0 = torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], 2)
+    stage.forward(_tick(session, 0, chunk0, [3]), server_args)
+    stage.forward(_tick(session, 1, _noise(1, MC, 2, 2, 2, seed=2), [2]), server_args)
+    state = session.get_or_create_state(SanaWMStreamCacheState)
+    assert state.chunk_idx == 2
+
+    # block_idx == 0 restarts the session in place.
+    stage.forward(_tick(session, 0, chunk0, [3]), server_args)
+    assert state.chunk_idx == 1 and state.chunk_indices == [0, 3]
+    assert state.latents.shape[2] == 3

python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py

@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+"""S1c tests — streaming denoising stage: KV accumulator + chunk-loop composition.
+
+The accumulator is the load-bearing new logic: GDN/STATE blocks copy-forward the
+previous chunk's recurrent state; softmax/CONCAT blocks concatenate the rolling
++ sink K/V along **dim=1** (our (B,N,H,D) softmax cache layout — the reference's
+dim=2 would concat the head axis and silently corrupt every softmax block).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.sana_wm import (
+    SanaWMArchConfig,
+    SanaWMConfig,
+)
+from sglang.multimodal_gen.runtime import server_args as _sa_mod
+from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
+    _CACHE_TYPE_CONCAT,
+    _CACHE_TYPE_STATE,
+    _NUM_STREAM_CACHE_SLOTS,
+    _SLOT_CAM_K,
+    _SLOT_CAM_V,
+    _SLOT_FFN_TCONV,
+    _SLOT_K,
+    _SLOT_SHORTCONV,
+    _SLOT_TYPE_FLAG,
+    _SLOT_V,
+    SanaWMTransformer3DModel,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import (
+    SanaWMStreamingDenoisingStage as Stage,
+)
+from sglang.multimodal_gen.runtime.server_args import set_global_server_args
+
+B, Hh, D = 1, 2, 4
+
+
+# --------------------------------------------------------------------------- #
+# 1. Chunk schedule
+# --------------------------------------------------------------------------- #
+
+
+def test_autoregressive_segments_first_chunk_absorbs_remainder():
+    assert Stage._autoregressive_segments(13, 3) == [0, 4, 7, 10, 13]
+    seg = Stage._autoregressive_segments(121, 3)
+    assert seg[0] == 0 and seg[1] == 4 and seg[-1] == 121
+    assert len(seg) - 1 == 40  # num chunks
+    assert all(seg[i + 1] - seg[i] == 3 for i in range(1, len(seg) - 1))
+
+
+# --------------------------------------------------------------------------- #
+# 2. KV accumulator
+# --------------------------------------------------------------------------- #
+
+
+def _state_block(seed):
+    g = torch.Generator().manual_seed(seed)
+    blk = [None] * _NUM_STREAM_CACHE_SLOTS
+    blk[_SLOT_K] = torch.randn(B, Hh, D, D, generator=g)  # recurrent state_kv
+    blk[_SLOT_V] = torch.randn(B, Hh, D, 1, generator=g)  # state_z
+    blk[_SLOT_CAM_K] = torch.randn(B, Hh, D, D, generator=g)  # cam state
+    blk[_SLOT_SHORTCONV] = torch.randn(B * 6, 3, Hh * D, generator=g)
+    blk[_SLOT_TYPE_FLAG] = torch.tensor([_CACHE_TYPE_STATE])
+    blk[_SLOT_FFN_TCONV] = torch.randn(B, Hh * D, 1, 6, generator=g)
+    return blk
+
+
+def _concat_block(n_tok, seed):
+    g = torch.Generator().manual_seed(seed)
+    blk = [None] * _NUM_STREAM_CACHE_SLOTS
+    for slot in (_SLOT_K, _SLOT_V, _SLOT_CAM_K, _SLOT_CAM_V):
+        blk[slot] = torch.randn(B, n_tok, Hh, D, generator=g)  # (B, N, H, D)
+    blk[_SLOT_TYPE_FLAG] = torch.tensor([_CACHE_TYPE_CONCAT])
+    blk[_SLOT_FFN_TCONV] = torch.randn(B, Hh * D, 1, 6, generator=g)
+    return blk
+
+
+def test_accumulate_state_block_copies_previous_chunk():
+    # 3 chunks, 1 STATE block. chunk2 must copy-forward chunk1's state.
+    kv = [[_state_block(c)] for c in range(3)]
+    cur, sink_num = Stage._accumulate_kv_cache(
+        kv,
+        chunk_idx=2,
+        chunk_indices=[0, 4, 7, 10],
+        num_cached_blocks=2,
+        sink_token=True,
+        num_blocks=1,
+    )
+    assert cur[0][_SLOT_K] is kv[1][0][_SLOT_K]  # carried from chunk 1
+    assert cur[0][_SLOT_V] is kv[1][0][_SLOT_V]
+    assert cur[0][_SLOT_CAM_K] is kv[1][0][_SLOT_CAM_K]
+    assert float(cur[0][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_STATE
+
+
+def test_accumulate_concat_block_concats_on_token_axis_dim1():
+    # chunk0 K has 4 tokens, chunk1 K has 3 -> chunk2 prefix = 4+3=7 on dim=1.
+    n0, n1 = 4, 3
+    kv = [
+        [_concat_block(n0, 0)],
+        [_concat_block(n1, 1)],
+        [[None] * _NUM_STREAM_CACHE_SLOTS],
+    ]
+    cur, _ = Stage._accumulate_kv_cache(
+        kv,
+        chunk_idx=2,
+        chunk_indices=[0, n0, n0 + n1, n0 + n1 + 3],
+        num_cached_blocks=5,
+        sink_token=False,
+        num_blocks=1,
+    )
+    acc_k = cur[0][_SLOT_K]
+    assert acc_k.shape == (B, n0 + n1, Hh, D)  # dim=1 grew; head axis (dim 2) intact
+    # exact content: chunk0 then chunk1 along dim=1
+    torch.testing.assert_close(acc_k[:, :n0], kv[0][0][_SLOT_K], atol=0, rtol=0)
+    torch.testing.assert_close(acc_k[:, n0:], kv[1][0][_SLOT_K], atol=0, rtol=0)
+    assert float(cur[0][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_CONCAT
+
+
+def test_accumulate_sink_includes_chunk_zero():
+    # num_cached_blocks=2 at chunk 3 -> sink_start=2>0 -> valid = [0, 2]
+    n = 3
+    kv = [[_concat_block(n, c)] for c in range(3)] + [
+        [[None] * _NUM_STREAM_CACHE_SLOTS]
+    ]
+    cur, sink_num = Stage._accumulate_kv_cache(
+        kv,
+        chunk_idx=3,
+        chunk_indices=[0, n, 2 * n, 3 * n, 4 * n],
+        num_cached_blocks=2,
+        sink_token=True,
+        num_blocks=1,
+    )
+    # valid = [0, 2] -> 2 chunks of n tokens each
+    assert cur[0][_SLOT_K].shape == (B, 2 * n, Hh, D)
+    assert sink_num == n  # chunk_indices[1] - chunk_indices[0]
+
+
+def test_accumulate_chunk_zero_is_empty():
+    kv = [[[None] * _NUM_STREAM_CACHE_SLOTS]]
+    cur, sink_num = Stage._accumulate_kv_cache(
+        kv, 0, [0, 3], num_cached_blocks=2, sink_token=True, num_blocks=1
+    )
+    assert sink_num == 0 and cur[0][_SLOT_K] is None
+
+
+# --------------------------------------------------------------------------- #
+# 3. Model-level chunk-loop composition (depth >= 4 -> a softmax/CONCAT block)
+# --------------------------------------------------------------------------- #
+
+MC, MT, MHt, MWt = 8, 9, 2, 2  # 9 latent frames, block=3 -> chunks [0,3),[3,6),[6,9)
+
+
+class _ZeroCross(torch.nn.Module):
+    def forward(self, x, y, mask=None):
+        return torch.zeros_like(x)
+
+
+@pytest.fixture
+def _global_args():
+    prev = _sa_mod._global_server_args
+    set_global_server_args(
+        SimpleNamespace(
+            comfyui_mode=False,
+            enable_cfg_parallel=False,
+            enable_torch_compile=False,
+            attention_backend=None,
+        )
+    )
+    try:
+        yield
+    finally:
+        set_global_server_args(prev)
+
+
+def _depth4_model():
+    arch = SanaWMArchConfig(
+        in_channels=MC,
+        out_channels=MC,
+        num_layers=4,  # block 3 -> softmax/CONCAT
+        num_attention_heads=2,
+        attention_head_dim=16,
+        linear_head_dim=16,
+        num_cross_attention_heads=2,
+        cross_attention_head_dim=16,
+        cross_attention_dim=32,
+        caption_channels=32,
+        model_max_length=8,
+        softmax_every_n=4,
+        update_rule="torch_recurrent",
+        cam_update_rule="torch_recurrent",
+        chunk_size=None,
+    )
+    m = SanaWMTransformer3DModel(SanaWMConfig(arch_config=arch)).double().eval()
+    for b in m.blocks:
+        b.cross_attn = _ZeroCross()
+    return m
+
+
+def test_streaming_loop_runs_and_accumulates_concat_block(_global_args):
+    """Drive forward_long chunk-by-chunk (the stage's core loop) on a depth-4
+    model: accumulate -> denoise(save=False) -> clean(save=True). Verify finite
+    output, threaded GDN state, and that the softmax block (idx 3) accumulates a
+    growing K window across chunks (the CONCAT path the depth-2 fixture can't)."""
+    m = _depth4_model()
+    assert [b.softmax_main for b in m.blocks] == [False, False, False, True]
+    torch.manual_seed(0)
+    latents = torch.randn(B, MC, MT, MHt, MWt, dtype=torch.float64)
+    y = torch.randn(B, 4, 32, dtype=torch.float64)
+    cam = torch.randn(B, MT, 20, dtype=torch.float64)
+    plk = torch.randn(B, 48, MT, MHt, MWt, dtype=torch.float64)
+
+    seg = Stage._autoregressive_segments(MT, 3)  # [0,3,6,9]
+    num_chunks = len(seg) - 1
+    kv = [
+        [[None] * _NUM_STREAM_CACHE_SLOTS for _ in range(4)] for _ in range(num_chunks)
+    ]
+    concat_k_lens = []
+
+    for ci in range(num_chunks):
+        chunk_kv, sink_num = Stage._accumulate_kv_cache(
+            kv, ci, seg, num_cached_blocks=2, sink_token=True, num_blocks=4
+        )
+        s, e = seg[ci], seg[ci + 1]
+        # softmax block 3 prefix K length entering this chunk
+        pk = chunk_kv[3][_SLOT_K]
+        concat_k_lens.append(0 if pk is None else pk.shape[1])
+        fidx = torch.arange(s, e) if sink_num > 0 else None
+        lat = latents[:, :, s:e]
+        ts = torch.full((B, e - s), 500.0, dtype=torch.float64)
+        # denoise step (save=False reads the accumulated prefix)
+        out, _ = m.forward_long(
+            hidden_states=lat,
+            encoder_hidden_states=y,
+            timestep=ts,
+            camera_conditions=cam,
+            chunk_plucker=plk,
+            kv_cache=chunk_kv,
+            save_kv_cache=False,
+            start_f=s,
+            end_f=e,
+            frame_index=fidx,
+        )
+        assert out.shape == (B, MC, e - s, MHt, MWt)
+        assert torch.isfinite(out).all()
+        # clean pass writes this chunk's KV
+        ts0 = torch.zeros(B, 1, e - s, dtype=torch.float64)
+        _, updated = m.forward_long(
+            hidden_states=lat,
+            encoder_hidden_states=y,
+            timestep=ts0,
+            camera_conditions=cam,
+            chunk_plucker=plk,
+            kv_cache=chunk_kv,
+            save_kv_cache=True,
+            start_f=s,
+            end_f=e,
+            frame_index=fidx,
+        )
+        kv[ci] = updated
+        # block 3 (softmax) stores current-chunk K as (B, N_tok, H, D) where
+        # N_tok = frames * H * W; blocks 0-2 (GDN) store recurrent state (B,heads,hd,hd).
+        n_tok = (e - s) * MHt * MWt
+        assert kv[ci][3][_SLOT_K].shape == (B, n_tok, 2, 16)  # heads=2, head_dim=16
+        assert kv[ci][0][_SLOT_K].shape == (B, 2, 16, 16)  # GDN state_kv
+        assert float(kv[ci][3][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_CONCAT
+        assert float(kv[ci][0][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_STATE
+
+    # softmax prefix grows (in TOKENS = frames*H*W): chunk0 sees 0, chunk1 sees
+    # chunk0's 3*4=12, chunk2 sees the rolling+sink window.
+    S_tok = MHt * MWt
+    assert concat_k_lens[0] == 0
+    assert concat_k_lens[1] == 3 * S_tok
+    assert concat_k_lens[2] >= 3 * S_tok

python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_vae.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+"""S2 tests — streaming causal-VAE decode: conv_cache + upsampler trim.
+
+The streaming decode threads a per-conv `conv_cache` dict so a chunked causal
+decode (carrying the cache across chunks) is bit-comparable to a monolithic
+causal decode of the whole clip. The two logic-bearing pieces are the conv leaf
+(`LTX2VideoCausalConv3d`, prepend prev tail / store last k-1 frames) and the
+upsampler temporal trim (`LTXVideoUpsampler3d`, drop the anchor frame ONLY at the
+true clip start). FP64, CPU.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.runtime.models.vaes.ltx_2_vae import (
+    LTX2VideoCausalConv3d,
+    LTX2VideoResnetBlock3d,
+    LTXVideoUpsampler3d,
+)
+
+torch.manual_seed(0)
+
+
+def _chunks(t, splits):
+    out, i = [], 0
+    for s in splits:
+        out.append(t[:, :, i : i + s])
+        i += s
+    return out
+
+
+@pytest.mark.parametrize("splits", [[6], [3, 3], [1, 2, 3], [2, 4]])
+def test_causal_conv_chunked_equals_monolithic(splits):
+    conv = (
+        LTX2VideoCausalConv3d(in_channels=4, out_channels=5, kernel_size=3)
+        .double()
+        .eval()
+    )
+    x = torch.randn(1, 4, 6, 3, 3, dtype=torch.float64)
+    with torch.no_grad():
+        whole = conv(x, causal=True)
+        cache = {}
+        parts = [
+            conv(c, causal=True, conv_cache=cache, cache_key="c")
+            for c in _chunks(x, splits)
+        ]
+    chunked = torch.cat(parts, dim=2)
+    assert chunked.shape == whole.shape
+    torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0)
+
+
+@pytest.mark.parametrize("splits", [[6], [3, 3], [2, 4], [1, 2, 3]])
+def test_upsampler_chunked_equals_monolithic(splits):
+    # residual + temporal upsample (stride[0]=2) -> exercises the trim gating.
+    up = (
+        LTXVideoUpsampler3d(
+            in_channels=8, stride=(2, 2, 2), residual=True, upscale_factor=1
+        )
+        .double()
+        .eval()
+    )
+    x = torch.randn(1, 8, 6, 4, 4, dtype=torch.float64)
+    with torch.no_grad():
+        whole = up(x, causal=True)
+        cache = {}
+        parts = [
+            up(c, causal=True, conv_cache=cache, cache_key="u")
+            for c in _chunks(x, splits)
+        ]
+    chunked = torch.cat(parts, dim=2)
+    # monolithic trims the anchor frame once; chunked must reproduce that exactly.
+    assert chunked.shape == whole.shape
+    torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0)
+
+
+@pytest.mark.parametrize("splits", [[6], [3, 3], [2, 4]])
+def test_resnet_chunked_equals_monolithic(splits):
+    blk = LTX2VideoResnetBlock3d(in_channels=4, out_channels=4).double().eval()
+    x = torch.randn(1, 4, 6, 3, 3, dtype=torch.float64)
+    with torch.no_grad():
+        whole = blk(x, causal=True)
+        cache = {}
+        parts = [
+            blk(c, causal=True, conv_cache=cache, cache_key="r")
+            for c in _chunks(x, splits)
+        ]
+    chunked = torch.cat(parts, dim=2)
+    torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0)
+
+
+def test_conv_cache_none_is_unchanged_dense_behavior():
+    # Backward-compat: with no conv_cache, output == the original causal forward.
+    conv = (
+        LTX2VideoCausalConv3d(in_channels=4, out_channels=4, kernel_size=3)
+        .double()
+        .eval()
+    )
+    x = torch.randn(1, 4, 5, 2, 2, dtype=torch.float64)
+    with torch.no_grad():
+        a = conv(x, causal=True)
+        b = conv(x, causal=True, conv_cache=None, cache_key=None)
+    torch.testing.assert_close(a, b, atol=0, rtol=0)

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

@@ -279,6 +279,32 @@ class TestSamplingParamsCliArgs(unittest.TestCase):
         self.assertIn("width", explicit_fields)
         self.assertIn("height", explicit_fields)
 
+    def test_cli_path_preserves_diffusers_kwargs_in_request_extra(self):
+        server_args = MagicMock()
+        server_args.backend = "sglang"
+        server_args.model_id = None
+        server_args.pipeline_config = MagicMock()
+        diffusers_kwargs = {"camera_to_world_path": "/tmp/camera.npy"}
+
+        with patch.object(
+            SamplingParams,
+            "from_pretrained",
+            side_effect=lambda *args, **kwargs: Flux2SamplingParams(),
+        ):
+            params = SamplingParams.from_user_sampling_params_args(
+                "dummy-model",
+                server_args=server_args,
+                prompt="p",
+                image_path="/tmp/in.png",
+                diffusers_kwargs=diffusers_kwargs,
+            )
+
+        self.assertEqual(params.diffusers_kwargs, diffusers_kwargs)
+        self.assertEqual(
+            params.build_request_extra()["diffusers_kwargs"],
+            diffusers_kwargs,
+        )
+
     def test_dataclasses_replace_preserves_explicit_fields(self):
         """`dataclasses.replace` drops `_explicit_fields`; DiffGenerator must restore it."""
         import dataclasses

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

@@ -20,6 +20,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConf
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImagePipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import (
+    SanaWMPipelineConfig,
+    SanaWMRealtimeConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     FastWan2_2_TI2V_5B_Config,
     TurboWanT2V480PConfig,
@@ -665,6 +669,7 @@ class TestOffloadDefaults(unittest.TestCase):
         mova_deployment = MOVAPipelineConfig().get_model_deployment_config()
         zimage_deployment = ZImagePipelineConfig().get_model_deployment_config()
         ltx_deployment = LTX2PipelineConfig().get_model_deployment_config()
+        sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config()
 
         self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb)
         self.assertFalse(qwen_deployment.auto_dit_layerwise_offload)
@@ -686,6 +691,35 @@ class TestOffloadDefaults(unittest.TestCase):
             ltx_deployment.auto_disable_component_offload_components, ("dit",)
         )
 
+        self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60)
+        self.assertTrue(sana_wm_deployment.auto_dit_layerwise_offload)
+
+    def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self):
+        args = self._from_dict_with_pipeline_config(
+            SanaWMPipelineConfig(),
+            kwargs={
+                "model_path": "Efficient-Large-Model/SANA-WM_bidirectional",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertTrue(args.use_fsdp_inference)
+        self.assertTrue(args.enable_cfg_parallel)
+
+    def test_auto_multi_gpu_sana_wm_realtime_disables_cfg_parallel(self):
+        args = self._from_dict_with_pipeline_config(
+            SanaWMRealtimeConfig(),
+            kwargs={
+                "model_path": "Efficient-Large-Model/SANA-WM_streaming",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertFalse(args.enable_cfg_parallel)
+
     def test_manual_mode_preserves_unset_performance_args(self):
         args = self._from_dict_with_pipeline_config(
             QwenImagePipelineConfig(),
@@ -1346,6 +1380,10 @@ class TestModelIdResolution(unittest.TestCase):
         info = _get_config_info(path)
         self.assertIsNotNone(info)
 
+    def test_sana_wm_model_path_resolves_registry(self):
+        info = _get_config_info("Efficient-Large-Model/SANA-WM_bidirectional")
+        self.assertIs(info.pipeline_config_cls, SanaWMPipelineConfig)
+
     def test_model_id_unknown_falls_back_without_crash(self):
         # unrecognized model_id: should warn and fall back to path-based detection
         # with an unresolvable path, expect RuntimeError from the detector step