From 9795b4cd5b1d723804f444081c6b7bfbc8914e69 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Thu, 5 Mar 2026 10:18:00 +0800
Subject: [PATCH] [Diffusion] Open t5 encoder parallel folding for wan2.2 and
 mova video (#18493)

Co-authored-by: Mick <mickjagger19@icloud.com>
---
 .../configs/models/encoders/t5.py             | 16 --------------
 .../multimodal_gen/runtime/server_args.py     | 22 +++++++++++++++++++
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/python/sglang/multimodal_gen/configs/models/encoders/t5.py b/python/sglang/multimodal_gen/configs/models/encoders/t5.py
index b2ff5cc53..658569088 100644
--- a/python/sglang/multimodal_gen/configs/models/encoders/t5.py
+++ b/python/sglang/multimodal_gen/configs/models/encoders/t5.py
@@ -8,7 +8,6 @@ from sglang.multimodal_gen.configs.models.encoders.base import (
     TextEncoderArchConfig,
     TextEncoderConfig,
 )
-from sglang.multimodal_gen.utils import StoreBoolean
 
 
 def _is_transformer_layer(n: str, m) -> bool:
@@ -95,19 +94,4 @@ class T5Config(TextEncoderConfig):
     def add_cli_args(
         parser: argparse.ArgumentParser, prefix: str = "t5-config"
     ) -> argparse.ArgumentParser:
-        parser.add_argument(
-            f"--{prefix}.parallel-folding",
-            action=StoreBoolean,
-            dest=f"{prefix.replace('-', '_')}.parallel_folding",
-            default=T5Config.parallel_folding,
-            help="Whether to use parallel folding for T5",
-        )
-        parser.add_argument(
-            f"--{prefix}.parallel-folding-mode",
-            type=str,
-            choices=["sp", "ulysses", "ring"],
-            dest=f"{prefix.replace('-', '_')}.parallel_folding_mode",
-            default=T5Config.parallel_folding_mode,
-            help="Parallel folding mode for T5",
-        )
         return parser
diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py
index 65f58c718..85ee65de5 100644
--- a/python/sglang/multimodal_gen/runtime/server_args.py
+++ b/python/sglang/multimodal_gen/runtime/server_args.py
@@ -21,6 +21,7 @@ import addict
 import yaml
 
 from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.configs.models.encoders import T5Config
 from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
 from sglang.multimodal_gen.configs.quantization import NunchakuSVDQuantArgs
 from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
@@ -353,6 +354,7 @@ class ServerArgs:
         self._adjust_attention_backend()
         self._adjust_platform_specific()
         self._adjust_autocast()
+        self.adjust_pipeline_config()
 
     def _validate_parameters(self):
         """check consistency and raise errors for invalid configs"""
@@ -391,6 +393,26 @@ class ServerArgs:
                 transformer_weights_path=self.nunchaku_config.transformer_weights_path,
             )
 
+    def adjust_pipeline_config(self):
+        # enable parallel folding when SP is enabled
+        if self.tp_size != 1 or self.sp_degree <= 1:
+            return
+
+        enabled = False
+        for text_encoder_config in self.pipeline_config.text_encoder_configs:
+            if isinstance(text_encoder_config, T5Config):
+                text_encoder_config.parallel_folding = True
+                enabled = True
+                text_encoder_config.parallel_folding_mode = "sp"
+
+        if enabled:
+            logger.info(
+                "Enabled T5 text encoder parallel folding (mode=sp) for %s (tp_size=%s, sp_degree=%s).",
+                self.__class__.__name__,
+                self.tp_size,
+                self.sp_degree,
+            )
+
     def _adjust_offload(self):
         # TODO: to be handled by each platform
         if current_platform.get_device_total_memory() / BYTES_PER_GB < 30: