From feb39f776805ac59415b45d225e69ebcd5dddcfb Mon Sep 17 00:00:00 2001
From: HuangJi <32611516+IPostYellow@users.noreply.github.com>
Date: Mon, 12 Jan 2026 13:55:38 +0800
Subject: [PATCH] [diffusion] model: Support TurboWan2.2-I2V SLA && add CI test
for TurboWan (#16536)
---
.../configs/pipeline_configs/wan.py | 12 +++
.../multimodal_gen/configs/sample/wan.py | 8 ++
.../multimodal_gen/docs/support_matrix.md | 30 +++---
python/sglang/multimodal_gen/registry.py | 10 ++
.../runtime/pipelines/wan_i2v_dmd_pipeline.py | 20 ++--
.../pipelines_core/stages/denoising_dmd.py | 95 ++++++++++++++++++-
.../test/server/perf_baselines.json | 43 +++++++++
.../test/server/test_server_utils.py | 10 +-
.../test/server/testcase_configs.py | 44 +++++++++
9 files changed, 247 insertions(+), 25 deletions(-)
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py
index b9f5ea6a2..6a824e678 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py
@@ -148,6 +148,18 @@ class WanI2V720PConfig(WanI2V480PConfig):
flow_shift: float | None = 5.0
+@dataclass
+class TurboWanI2V720Config(WanI2V720PConfig):
+ flow_shift: float | None = 8.0
+ dmd_denoising_steps: list[int] | None = field(
+ default_factory=lambda: [996, 932, 852, 608]
+ )
+ boundary_ratio: float | None = 0.9
+
+ def __post_init__(self) -> None:
+ self.dit_config.boundary_ratio = self.boundary_ratio
+
+
@dataclass
class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
"""Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD"""
diff --git a/python/sglang/multimodal_gen/configs/sample/wan.py b/python/sglang/multimodal_gen/configs/sample/wan.py
index 839dc4049..efe86c4b1 100644
--- a/python/sglang/multimodal_gen/configs/sample/wan.py
+++ b/python/sglang/multimodal_gen/configs/sample/wan.py
@@ -273,6 +273,14 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
)
+@dataclass
+class Turbo_Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
+ guidance_scale: float = 3.5 # high_noise
+ guidance_scale_2: float = 3.5 # low_noise
+ num_inference_steps: int = 4
+ fps: int = 16
+
+
# =============================================
# ============= Causal Self-Forcing =============
# =============================================
diff --git a/python/sglang/multimodal_gen/docs/support_matrix.md b/python/sglang/multimodal_gen/docs/support_matrix.md
index 7d89f0b3b..d7af4756f 100644
--- a/python/sglang/multimodal_gen/docs/support_matrix.md
+++ b/python/sglang/multimodal_gen/docs/support_matrix.md
@@ -16,19 +16,23 @@ default parameters when initializing and generating videos.
### Video Generation Models
-| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) |
-|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|
-| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ |
-| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ |
-| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ |
-| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ |
-| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ |
-| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ |
-| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ |
-| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ |
-| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ |
-| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ |
-| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ |
+| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) | Sparse Linear Attention（SLA）
+|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:----------------------------:|
+| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
+| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
+| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
+| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ |
+| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ |
+| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ |
+| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ |
+| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ |
+| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ | ❌ |
+| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ |
+| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ | ❌ |
+| TurboWan2.1 T2V 1.3B | `IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ |
+| TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ |
+| TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ |
+| TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ |
**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py
index 1f779dd5b..113ad0e5f 100644
--- a/python/sglang/multimodal_gen/registry.py
+++ b/python/sglang/multimodal_gen/registry.py
@@ -49,6 +49,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
from sglang.multimodal_gen.configs.pipeline_configs.wan import (
FastWan2_1_T2V_480P_Config,
FastWan2_2_TI2V_5B_Config,
+ TurboWanI2V720Config,
TurboWanT2V480PConfig,
Wan2_2_I2V_A14B_Config,
Wan2_2_T2V_A14B_Config,
@@ -67,6 +68,7 @@ from sglang.multimodal_gen.configs.sample.qwenimage import (
)
from sglang.multimodal_gen.configs.sample.wan import (
FastWanT2V480PConfig,
+ Turbo_Wan2_2_I2V_A14B_SamplingParam,
Wan2_1_Fun_1_3B_InP_SamplingParams,
Wan2_2_I2V_A14B_SamplingParam,
Wan2_2_T2V_A14B_SamplingParam,
@@ -426,6 +428,7 @@ def _register_configs():
pipeline_config_cls=TurboWanT2V480PConfig,
hf_model_paths=[
"IPostYellow/TurboWan2.1-T2V-14B-Diffusers",
+ "IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers",
],
)
register_configs(
@@ -443,6 +446,13 @@ def _register_configs():
"Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
],
)
+ register_configs(
+ sampling_param_cls=Turbo_Wan2_2_I2V_A14B_SamplingParam,
+ pipeline_config_cls=TurboWanI2V720Config,
+ hf_model_paths=[
+ "IPostYellow/TurboWan2.2-I2V-A14B-Diffusers",
+ ],
+ )
register_configs(
sampling_param_cls=Wan2_1_Fun_1_3B_InP_SamplingParams,
pipeline_config_cls=WanI2V480PConfig,
diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py
index 7a49eef9b..23d560711 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py
@@ -68,14 +68,17 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
tokenizers=[self.get_module("tokenizer")],
),
)
-
- self.add_stage(
- stage_name="image_encoding_stage",
- stage=ImageEncodingStage(
- image_encoder=self.get_module("image_encoder"),
- image_processor=self.get_module("image_processor"),
- ),
- )
+ if (
+ self.get_module("image_encoder") is not None
+ and self.get_module("image_processor") is not None
+ ):
+ self.add_stage(
+ stage_name="image_encoding_stage",
+ stage=ImageEncodingStage(
+ image_encoder=self.get_module("image_encoder"),
+ image_processor=self.get_module("image_processor"),
+ ),
+ )
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
@@ -102,6 +105,7 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
stage=DmdDenoisingStage(
transformer=self.get_module("transformer"),
scheduler=self.get_module("scheduler"),
+ transformer_2=self.get_module("transformer_2"),
),
)
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py
index c67e9d017..057ba2d03 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py
@@ -26,8 +26,10 @@ class DmdDenoisingStage(DenoisingStage):
Denoising stage for DMD.
"""
- def __init__(self, transformer, scheduler) -> None:
- super().__init__(transformer, scheduler)
+ def __init__(self, transformer, scheduler, transformer_2=None) -> None:
+ super().__init__(
+ transformer=transformer, scheduler=scheduler, transformer_2=transformer_2
+ )
self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs):
@@ -103,6 +105,20 @@ class DmdDenoisingStage(DenoisingStage):
timings=batch.timings,
perf_dump_path_provided=batch.perf_dump_path is not None,
):
+ t_int = int(t.item())
+ if self.transformer_2 is not None:
+ current_model, current_guidance_scale = (
+ self._select_and_manage_model(
+ t_int=t_int,
+ boundary_timestep=self._handle_boundary_ratio(
+ server_args, batch
+ ),
+ server_args=server_args,
+ batch=batch,
+ )
+ )
+ else:
+ current_model = self.transformer
# Expand latents for I2V
noise_latents = latents.clone()
latent_model_input = latents.to(target_dtype)
@@ -143,7 +159,7 @@ class DmdDenoisingStage(DenoisingStage):
forward_batch=batch,
):
# Run transformer
- pred_noise = self.transformer(
+ pred_noise = current_model(
hidden_states=latent_model_input.permute(0, 2, 1, 3, 4),
timestep=t_expand,
guidance=guidance_expand,
@@ -202,3 +218,76 @@ class DmdDenoisingStage(DenoisingStage):
)
return batch
+
+ def _select_and_manage_model(
+ self,
+ t_int: int,
+ boundary_timestep: float | None,
+ server_args: ServerArgs,
+ batch: Req,
+ ):
+ if boundary_timestep is None or t_int >= boundary_timestep:
+ # High-noise stage
+ current_model = self.transformer
+ model_to_offload = self.transformer_2
+ current_guidance_scale = batch.guidance_scale
+ else:
+ # Low-noise stage
+ current_model = self.transformer_2
+ model_to_offload = self.transformer
+ current_guidance_scale = batch.guidance_scale_2
+
+ self._manage_device_placement(current_model, model_to_offload, server_args)
+
+ assert current_model is not None, "The model for the current step is not set."
+ return current_model, current_guidance_scale
+
+ def _manage_device_placement(
+ self,
+ model_to_use: torch.nn.Module,
+ model_to_offload: torch.nn.Module | None,
+ server_args: ServerArgs,
+ ):
+ """
+ Manages the offload / load behavior of dit
+ """
+ if not server_args.dit_cpu_offload:
+ return
+
+ # Offload the unused model if it's on CUDA
+ if (
+ model_to_offload is not None
+ and next(model_to_offload.parameters()).device.type == "cuda"
+ ):
+ model_to_offload.to("cpu")
+
+ # Load the model to use if it's on CPU
+ if (
+ model_to_use is not None
+ and next(model_to_use.parameters()).device.type == "cpu"
+ ):
+ model_to_use.to(get_local_torch_device())
+
+ def _handle_boundary_ratio(
+ self,
+ server_args,
+ batch,
+ ):
+ """
+ (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
+ """
+ boundary_ratio = server_args.pipeline_config.dit_config.boundary_ratio
+ if batch.boundary_ratio is not None:
+ logger.info(
+ "Overriding boundary ratio from %s to %s",
+ boundary_ratio,
+ batch.boundary_ratio,
+ )
+ boundary_ratio = batch.boundary_ratio
+
+ if boundary_ratio is not None:
+ boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
+ else:
+ boundary_timestep = None
+
+ return boundary_timestep
diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json
index 769d5adf3..2a7034217 100644
--- a/python/sglang/multimodal_gen/test/server/perf_baselines.json
+++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json
@@ -804,6 +804,27 @@
"expected_avg_denoise_ms": 260.76,
"expected_median_denoise_ms": 247.84
},
+ "turbo_wan2_1_t2v_1.3b": {
+ "stages_ms": {
+ "InputValidationStage": 0.06,
+ "TextEncodingStage": 2508.95,
+ "ConditioningStage": 0.04,
+ "TimestepPreparationStage": 73.51,
+ "LatentPreparationStage": 1.34,
+ "DmdDenoisingStage": 1285.25,
+ "DecodingStage": 805.04,
+ "per_frame_generation": null
+ },
+ "denoise_step_ms": {
+ "0": 897.62,
+ "1": 126.04,
+ "2": 126.52,
+ "3": 128.26
+ },
+ "expected_e2e_ms": 4686.66,
+ "expected_avg_denoise_ms": 319.61,
+ "expected_median_denoise_ms": 127.39
+ },
"wan2_2_ti2v_5b": {
"stages_ms": {
"InputValidationStage": 96.27,
@@ -1092,6 +1113,28 @@
"expected_avg_denoise_ms": 2831.00,
"expected_median_denoise_ms": 1600.09
},
+ "turbo_wan2_2_i2v_a14b_2gpu": {
+ "stages_ms": {
+ "InputValidationStage": 25.01,
+ "TextEncodingStage": 5198.6,
+ "ConditioningStage": 0.04,
+ "TimestepPreparationStage": 56.26,
+ "LatentPreparationStage": 1.4,
+ "ImageVAEEncodingStage": 1001.89,
+ "DmdDenoisingStage": 4487.79,
+ "DecodingStage": 821.01,
+ "per_frame_generation": null
+ },
+ "denoise_step_ms": {
+ "0": 3042.56,
+ "1": 485.88,
+ "2": 477.59,
+ "3": 475.58
+ },
+ "expected_e2e_ms": 11605.97,
+ "expected_avg_denoise_ms": 1120.4,
+ "expected_median_denoise_ms": 481.74
+ },
"wan2_1_i2v_14b_480P_2gpu": {
"stages_ms": {
"InputValidationStage": 38.23,
diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py
index b4c170b2e..4bdcb7fc9 100644
--- a/python/sglang/multimodal_gen/test/server/test_server_utils.py
+++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py
@@ -1078,7 +1078,11 @@ def get_generate_fn(
prompt=sampling_params.prompt,
size=sampling_params.output_size,
seconds=video_seconds,
- extra_body={"reference_url": sampling_params.image_path},
+ extra_body={
+ "reference_url": sampling_params.image_path,
+ "fps": sampling_params.fps,
+ "num_frames": sampling_params.num_frames,
+ },
)
def generate_text_image_to_video(case_id, client) -> str:
@@ -1102,6 +1106,10 @@ def get_generate_fn(
size=output_size,
seconds=video_seconds,
input_reference=fh,
+ extra_body={
+ "fps": sampling_params.fps,
+ "num_frames": sampling_params.num_frames,
+ },
)
if modality == "video":
diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py
index 562435511..bc097eb61 100644
--- a/python/sglang/multimodal_gen/test/server/testcase_configs.py
+++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py
@@ -25,6 +25,7 @@ from dataclasses import dataclass
from pathlib import Path
from typing import Sequence
+from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
@@ -300,6 +301,15 @@ TI2V_sampling_params = DiffusionSamplingParams(
direct_url_test=True,
)
+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",
+ direct_url_test=True,
+ output_size="960x960",
+ num_frames=4,
+ fps=4,
+)
+
# All test cases with clean default values
# To test different models, simply add more DiffusionCase entries
ONE_GPU_CASES_A: list[DiffusionTestCase] = [
@@ -496,6 +506,23 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
),
]
+# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.
+if not current_platform.is_hip():
+ ONE_GPU_CASES_B.append(
+ DiffusionTestCase(
+ "turbo_wan2_1_t2v_1.3b",
+ DiffusionServerArgs(
+ model_path="IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
+ modality="video",
+ warmup=0,
+ custom_validator="video",
+ ),
+ DiffusionSamplingParams(
+ prompt=T2V_PROMPT,
+ ),
+ )
+ )
+
TWO_GPU_CASES_A = [
DiffusionTestCase(
"wan2_2_i2v_a14b_2gpu",
@@ -551,6 +578,23 @@ TWO_GPU_CASES_A = [
),
]
+# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.
+if not current_platform.is_hip():
+ TWO_GPU_CASES_A.append(
+ DiffusionTestCase(
+ "turbo_wan2_2_i2v_a14b_2gpu",
+ DiffusionServerArgs(
+ model_path="IPostYellow/TurboWan2.2-I2V-A14B-Diffusers",
+ modality="video",
+ warmup=0,
+ custom_validator="video",
+ num_gpus=2,
+ tp_size=2,
+ ),
+ TURBOWAN_I2V_sampling_params,
+ )
+ )
+
TWO_GPU_CASES_B = [
DiffusionTestCase(
"wan2_1_i2v_14b_480P_2gpu",