[diffusion] refactor: refactor realtime and model-specific stage modules (#27697)

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

@@ -28,7 +28,7 @@ from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d import (
     Hunyuan3DPaintPostprocessStage,
     Hunyuan3DPaintPreprocessStage,
     Hunyuan3DPaintTexGenStage,

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

@@ -15,17 +15,19 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
 from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     AuxiliaryConditionEncodingStage,
-    CausalVaeDecodingStage,
     DMDTimestepPreparationStage,
     ImageEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import (
+    LingBotWorldCausalDMDDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import (
+    CausalVaeDecodingStage,
     RealtimeChunkLatentPreparationStage,
     RealtimeImageVAEEncodingStage,
     RealtimeInputValidationStage,
     RealtimeTextEncodingStage,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import (
-    LingBotWorldCausalDMDDenoisingStage,
-)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 
 

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

@@ -35,18 +35,22 @@ from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipel
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     InputValidationStage,
+    TextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import (
+    LTX2ImageEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2 import (
     LTX2AVDecodingStage,
     LTX2AVDenoisingStage,
     LTX2AVLatentPreparationStage,
     LTX2HalveResolutionStage,
-    LTX2ImageEncodingStage,
     LTX2LoRASwitchStage,
     LTX2RefinementStage,
     LTX2TextConnectorStage,
     LTX2UpsampleStage,
-    TextEncodingStage,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import (
     LTX2_RESIDENT_AUTO_ENABLE_MEM_GB,

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

@@ -7,10 +7,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import (
 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,
 )
@@ -31,6 +27,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.s
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming_refiner import (
     SanaWMStreamingRefinerStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import (
+    RealtimeInputValidationStage,
+    RealtimeTextEncodingStage,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
 

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

@@ -20,67 +20,20 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.condition_encoding impo
     ConditionEncodingStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
-from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import (
-    LTX2AVDecodingStage,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
-from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import (
-    LTX2AVDenoisingStage,
-    LTX2RefinementStage,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import (
     DmdDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.encoding import EncodingStage
-
-# Hunyuan3D paint stages
-from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_paint import (
-    Hunyuan3DPaintPostprocessStage,
-    Hunyuan3DPaintPreprocessStage,
-    Hunyuan3DPaintTexGenStage,
-)
-
-# Hunyuan3D shape stages
-from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape import (
-    Hunyuan3DShapeBeforeDenoisingStage,
-    Hunyuan3DShapeDenoisingStage,
-    Hunyuan3DShapeExportStage,
-    Hunyuan3DShapeSaveStage,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import (
     ImageEncodingStage,
     ImageVAEEncodingStage,
-    LTX2ImageEncodingStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
     InputValidationStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
     LatentPreparationStage,
-    RealtimeChunkLatentPreparationStage,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import (
-    LTX2AVLatentPreparationStage,
-)
-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,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_text_encoding import (
-    RealtimeTextEncodingStage,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_vae import (
-    CausalVaeDecodingStage,
-    RealtimeImageVAEEncodingStage,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import (
-    LTX2TextConnectorStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
     TextEncodingStage,
@@ -89,54 +42,22 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation im
     DMDTimestepPreparationStage,
     TimestepPreparationStage,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages.upsampling import (
-    LTX2HalveResolutionStage,
-    LTX2LoRASwitchStage,
-    LTX2UpsampleStage,
-)
 
 __all__ = [
     "PipelineStage",
     "InputValidationStage",
-    "RealtimeInputValidationStage",
-    "RealtimeDiffusionStage",
-    "RealtimeStageComponent",
     "TimestepPreparationStage",
     "DMDTimestepPreparationStage",
     "LatentPreparationStage",
-    "RealtimeChunkLatentPreparationStage",
     "ComfyUILatentPreparationStage",
-    "LTX2AVLatentPreparationStage",
     "DenoisingStage",
     "DmdDenoisingStage",
-    "LTX2DenoisingStage",
-    "LTX2AVDenoisingStage",
     "CausalDMDDenoisingStage",
     "EncodingStage",
     "ConditionEncodingStage",
     "AuxiliaryConditionEncodingStage",
     "DecodingStage",
-    "CausalVaeDecodingStage",
-    "LTX2AVDecodingStage",
     "ImageEncodingStage",
     "ImageVAEEncodingStage",
-    "RealtimeImageVAEEncodingStage",
-    "LTX2ImageEncodingStage",
     "TextEncodingStage",
-    "RealtimeTextEncodingStage",
-    "LTX2TextConnectorStage",
-    # Hunyuan3D shape stages
-    "Hunyuan3DShapeBeforeDenoisingStage",
-    "Hunyuan3DShapeDenoisingStage",
-    "Hunyuan3DShapeExportStage",
-    "Hunyuan3DShapeSaveStage",
-    # Hunyuan3D paint stages
-    "Hunyuan3DPaintPreprocessStage",
-    "Hunyuan3DPaintTexGenStage",
-    "Hunyuan3DPaintPostprocessStage",
-    # LTX-2 two-stage
-    "LTX2RefinementStage",
-    "LTX2HalveResolutionStage",
-    "LTX2LoRASwitchStage",
-    "LTX2UpsampleStage",
 ]

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

@@ -377,59 +377,3 @@ class LatentPreparationStage(PipelineStage):
         # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
         result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple)
         return result
-
-
-class RealtimeChunkLatentPreparationStage(LatentPreparationStage):
-    """Prepare one realtime causal DiT chunk from the encoded condition shape."""
-
-    def get_forward_latent_num_frames(
-        self,
-        batch: Req,
-        server_args: ServerArgs,
-    ) -> int:
-        return int(
-            batch.realtime_chunk_size
-            or self.transformer.config.arch_config.num_frames_per_block
-        )
-
-    def get_latent_preparation_spec(
-        self,
-        batch: Req,
-        server_args: ServerArgs,
-        batch_size: int,
-        num_frames: int,
-        device: torch.device | str,
-    ) -> LatentPreparationSpec:
-        condition_latent = batch.image_latent
-        assert condition_latent is not None, (
-            "Realtime chunk latent preparation requires image_latent. "
-            "Ensure the condition VAE encoding stage runs before this stage."
-        )
-        return LatentPreparationSpec(
-            shape=(
-                condition_latent.shape[0],
-                self.transformer.config.arch_config.out_channels,
-                num_frames,
-                condition_latent.shape[3],
-                condition_latent.shape[4],
-            ),
-            dtype=condition_latent.dtype,
-            device=device,
-            prepare_latent_ids=False,
-            pack_latents=False,
-        )
-
-    def should_scale_initial_noise(self, batch: Req, server_args: ServerArgs) -> bool:
-        return False
-
-    def requires_batch_height_width(self, batch: Req, server_args: ServerArgs) -> bool:
-        return False
-
-    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check(
-            "image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)]
-        )
-        result.add_check("generator", batch.generator, V.generator_or_list_generators)
-        result.add_check("latents", batch.latents, V.none_or_tensor)
-        return result

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

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Hunyuan3D-specific pipeline stages"""
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.paint import (
+    Hunyuan3DPaintPostprocessStage,
+    Hunyuan3DPaintPreprocessStage,
+    Hunyuan3DPaintTexGenStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.shape import (
+    Hunyuan3DShapeBeforeDenoisingStage,
+    Hunyuan3DShapeDenoisingStage,
+    Hunyuan3DShapeExportStage,
+    Hunyuan3DShapeSaveStage,
+)
+
+__all__ = [
+    "Hunyuan3DPaintPostprocessStage",
+    "Hunyuan3DPaintPreprocessStage",
+    "Hunyuan3DPaintTexGenStage",
+    "Hunyuan3DShapeBeforeDenoisingStage",
+    "Hunyuan3DShapeDenoisingStage",
+    "Hunyuan3DShapeExportStage",
+    "Hunyuan3DShapeSaveStage",
+]

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

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""LTX-2-specific pipeline stages"""
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.decoding_av import (
+    LTX2AVDecodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising import (
+    LTX2DenoisingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising_av import (
+    LTX2AVDenoisingStage,
+    LTX2RefinementStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.latent_preparation_av import (
+    LTX2AVLatentPreparationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.text_connector import (
+    LTX2TextConnectorStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.upsampling import (
+    LTX2HalveResolutionStage,
+    LTX2LoRASwitchStage,
+    LTX2UpsampleStage,
+)
+
+__all__ = [
+    "LTX2AVDecodingStage",
+    "LTX2AVDenoisingStage",
+    "LTX2AVLatentPreparationStage",
+    "LTX2DenoisingStage",
+    "LTX2HalveResolutionStage",
+    "LTX2LoRASwitchStage",
+    "LTX2RefinementStage",
+    "LTX2TextConnectorStage",
+    "LTX2UpsampleStage",
+]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py

@@ -12,7 +12,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     StageParallelismType,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import (
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising import (
     LTX2DenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs

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

@@ -11,7 +11,7 @@ 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 (
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.base import (
     RealtimeDiffusionStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs

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

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""realtime pipeline stages shared by interactive diffusion models"""
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.base import (
+    RealtimeDiffusionStage,
+    RealtimeStageComponent,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.input_validation import (
+    RealtimeInputValidationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.latent_preparation import (
+    RealtimeChunkLatentPreparationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.text_encoding import (
+    RealtimeTextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.vae import (
+    CausalVaeDecodingStage,
+    RealtimeImageVAEEncodingStage,
+)
+
+__all__ = [
+    "CausalVaeDecodingStage",
+    "RealtimeChunkLatentPreparationStage",
+    "RealtimeDiffusionStage",
+    "RealtimeImageVAEEncodingStage",
+    "RealtimeInputValidationStage",
+    "RealtimeStageComponent",
+    "RealtimeTextEncodingStage",
+]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
+    LatentPreparationSpec,
+    LatentPreparationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    VerificationResult,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+class RealtimeChunkLatentPreparationStage(LatentPreparationStage):
+    """Prepare one realtime causal DiT chunk from the encoded condition shape."""
+
+    def get_forward_latent_num_frames(
+        self,
+        batch: Req,
+        server_args: ServerArgs,
+    ) -> int:
+        return int(
+            batch.realtime_chunk_size
+            or self.transformer.config.arch_config.num_frames_per_block
+        )
+
+    def get_latent_preparation_spec(
+        self,
+        batch: Req,
+        server_args: ServerArgs,
+        batch_size: int,
+        num_frames: int,
+        device: torch.device | str,
+    ) -> LatentPreparationSpec:
+        condition_latent = batch.image_latent
+        assert condition_latent is not None, (
+            "Realtime chunk latent preparation requires image_latent. "
+            "Ensure the condition VAE encoding stage runs before this stage."
+        )
+        return LatentPreparationSpec(
+            shape=(
+                condition_latent.shape[0],
+                self.transformer.config.arch_config.out_channels,
+                num_frames,
+                condition_latent.shape[3],
+                condition_latent.shape[4],
+            ),
+            dtype=condition_latent.dtype,
+            device=device,
+            prepare_latent_ids=False,
+            pack_latents=False,
+        )
+
+    def should_scale_initial_noise(self, batch: Req, server_args: ServerArgs) -> bool:
+        return False
+
+    def requires_batch_height_width(self, batch: Req, server_args: ServerArgs) -> bool:
+        return False
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check(
+            "image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)]
+        )
+        result.add_check("generator", batch.generator, V.generator_or_list_generators)
+        result.add_check("latents", batch.latents, V.none_or_tensor)
+        return result

python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/text_encoding.py

@@ -49,6 +49,34 @@ def _copy_seq_lens(
     return [list(seq_lens) for seq_lens in value]
 
 
+_TEXT_CACHE_TENSOR_LIST_FIELDS = (
+    "prompt_embeds",
+    "pooled_embeds",
+    "prompt_attention_mask",
+    "prompt_embeds_mask",
+    "negative_prompt_embeds",
+    "neg_pooled_embeds",
+    "negative_attention_mask",
+    "negative_prompt_embeds_mask",
+)
+_TEXT_CACHE_SEQ_LENS_FIELDS = (
+    "prompt_seq_lens",
+    "negative_prompt_seq_lens",
+)
+_TEXT_CACHE_FIELDS = _TEXT_CACHE_TENSOR_LIST_FIELDS + _TEXT_CACHE_SEQ_LENS_FIELDS
+_TEXT_CACHE_DEFAULT_EMPTY_LIST_FIELDS = {
+    "prompt_embeds",
+    "pooled_embeds",
+    "neg_pooled_embeds",
+}
+
+
+def _copy_text_cache_field(name: str, value):
+    if name in _TEXT_CACHE_SEQ_LENS_FIELDS:
+        return _copy_seq_lens(value)
+    return _copy_tensor_list(value)
+
+
 class RealtimeTextState(BaseRealtimeState):
     def __init__(self):
         super().__init__()
@@ -66,16 +94,8 @@ class RealtimeTextState(BaseRealtimeState):
 
     def clear_text_cache(self):
         self.cache_key = None
-        self.prompt_embeds = None
+        for field in _TEXT_CACHE_FIELDS:
-        self.pooled_embeds = None
+            setattr(self, field, None)
-        self.prompt_attention_mask = None
-        self.prompt_embeds_mask = None
-        self.prompt_seq_lens = None
-        self.negative_prompt_embeds = None
-        self.neg_pooled_embeds = None
-        self.negative_attention_mask = None
-        self.negative_prompt_embeds_mask = None
-        self.negative_prompt_seq_lens = None
 
     def dispose(self):
         super().dispose()
@@ -97,33 +117,16 @@ class RealtimeTextEncodingStage(TextEncodingStage):
         )
 
     def _restore_cached_outputs(self, batch: Req, state: RealtimeTextState) -> Req:
-        batch.prompt_embeds = _copy_tensor_list(state.prompt_embeds) or []
+        for field in _TEXT_CACHE_FIELDS:
-        batch.pooled_embeds = _copy_tensor_list(state.pooled_embeds) or []
+            value = _copy_text_cache_field(field, getattr(state, field))
-        batch.prompt_attention_mask = _copy_tensor_list(state.prompt_attention_mask)
+            if value is None and field in _TEXT_CACHE_DEFAULT_EMPTY_LIST_FIELDS:
-        batch.prompt_embeds_mask = _copy_tensor_list(state.prompt_embeds_mask)
+                value = []
-        batch.prompt_seq_lens = _copy_seq_lens(state.prompt_seq_lens)
+            setattr(batch, field, value)
-        batch.negative_prompt_embeds = _copy_tensor_list(state.negative_prompt_embeds)
-        batch.neg_pooled_embeds = _copy_tensor_list(state.neg_pooled_embeds) or []
-        batch.negative_attention_mask = _copy_tensor_list(state.negative_attention_mask)
-        batch.negative_prompt_embeds_mask = _copy_tensor_list(
-            state.negative_prompt_embeds_mask
-        )
-        batch.negative_prompt_seq_lens = _copy_seq_lens(state.negative_prompt_seq_lens)
         return batch
 
     def _store_outputs(self, batch: Req, state: RealtimeTextState) -> None:
-        state.prompt_embeds = _copy_tensor_list(batch.prompt_embeds)
+        for field in _TEXT_CACHE_FIELDS:
-        state.pooled_embeds = _copy_tensor_list(batch.pooled_embeds)
+            setattr(state, field, _copy_text_cache_field(field, getattr(batch, field)))
-        state.prompt_attention_mask = _copy_tensor_list(batch.prompt_attention_mask)
-        state.prompt_embeds_mask = _copy_tensor_list(batch.prompt_embeds_mask)
-        state.prompt_seq_lens = _copy_seq_lens(batch.prompt_seq_lens)
-        state.negative_prompt_embeds = _copy_tensor_list(batch.negative_prompt_embeds)
-        state.neg_pooled_embeds = _copy_tensor_list(batch.neg_pooled_embeds)
-        state.negative_attention_mask = _copy_tensor_list(batch.negative_attention_mask)
-        state.negative_prompt_embeds_mask = _copy_tensor_list(
-            batch.negative_prompt_embeds_mask
-        )
-        state.negative_prompt_seq_lens = _copy_seq_lens(batch.negative_prompt_seq_lens)
 
     @torch.no_grad()
     def forward(

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

@@ -38,10 +38,10 @@ 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 (
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.base import (
     RealtimeDiffusionStage,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_input_validation import (
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.input_validation import (
     RealtimeInputValidationStage,
     RealtimeInputValidationState,
 )
@@ -74,8 +74,14 @@ class _State(BaseRealtimeState):
         self.disposed = True
 
 
+class _TestRealtimeDiffusionStage(RealtimeDiffusionStage):
+    def forward(self, batch, component_manager=None):
+        del batch, component_manager
+        raise NotImplementedError
+
+
 def test_realtime_diffusion_stage_declares_long_lived_components():
-    stage = RealtimeDiffusionStage()
+    stage = _TestRealtimeDiffusionStage()
     server_args = SimpleNamespace(
         pipeline_config=SimpleNamespace(dit_precision="bf16", vae_precision="fp32")
     )
@@ -93,7 +99,7 @@ def test_realtime_diffusion_stage_declares_long_lived_components():
 
 
 def test_realtime_diffusion_stage_requires_session():
-    stage = RealtimeDiffusionStage(default_height=480, default_width=832)
+    stage = _TestRealtimeDiffusionStage(default_height=480, default_width=832)
     req = _Req(session=None)
 
     try:
@@ -884,7 +890,7 @@ def test_realtime_chunk_latent_preparation_uses_chunk_spec():
 
     import torch
 
-    from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import (
         RealtimeChunkLatentPreparationStage,
     )
 

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

@@ -4,7 +4,7 @@ from types import SimpleNamespace
 
 import torch
 
-from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_vae import (
+from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.vae import (
     CausalVaeDecodingStage,
     RealtimeVAEDecodeState,
 )
@@ -22,7 +22,9 @@ def test_realtime_vae_decode_state_clears_model_cache_on_dispose():
 
 
 def test_causal_vae_decoding_stage_keeps_wan_decoder_cache(monkeypatch):
-    from sglang.multimodal_gen.runtime.pipelines_core.stages import realtime_vae
+    from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import (
+        vae as realtime_vae,
+    )
 
     class _WanVAE:
         def __init__(self):
@@ -104,7 +106,9 @@ def test_causal_vae_decoding_stage_keeps_wan_decoder_cache(monkeypatch):
 
 
 def test_causal_vae_decoding_stage_prefers_native_causal_decode(monkeypatch):
-    from sglang.multimodal_gen.runtime.pipelines_core.stages import realtime_vae
+    from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import (
+        vae as realtime_vae,
+    )
 
     class _NativeCausalVAE:
         def __init__(self):

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

@@ -39,20 +39,20 @@ from sglang.multimodal_gen.runtime.pipelines.wan_i2v_pipeline import (
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import (
-    LTX2RefinementStage,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape import (
-    Hunyuan3DShapeBeforeDenoisingStage,
-    Hunyuan3DShapeExportStage,
-    Hunyuan3DShapeSaveStage,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import (
     ImageVAEEncodingStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_denoising import (
     HeliosChunkedDenoisingStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.shape import (
+    Hunyuan3DShapeBeforeDenoisingStage,
+    Hunyuan3DShapeExportStage,
+    Hunyuan3DShapeSaveStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising_av import (
+    LTX2RefinementStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.mova import (
     MOVADecodingStage,
     MOVADenoisingStage,
@@ -628,7 +628,7 @@ class TestHunyuan3DShapeStageRuntimeDtype(_GlobalStageArgsMixin, unittest.TestCa
         )
 
         with patch(
-            "sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape.logger.warning"
+            "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.shape.logger.warning"
         ) as mock_warning:
             dtype = stage._resolve_runtime_dtype(torch.zeros(1, dtype=torch.float16))