[diffusion] fix: fix image encoder parallel folding proposal (#36863)

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -756,9 +756,10 @@ class ServerArgs(DisaggServerArgsMixin):
         # propose the fold group from the parallelism alone; the loader keeps it
         # only for encoders worth folding at their real post-load size
         # (finalize_encoder_folding)
-        encoder_configs = list(self.pipeline_config.text_encoder_configs) + list(
-            getattr(self.pipeline_config, "image_encoder_configs", ()) or ()
-        )
+        encoder_configs = [
+            *self.pipeline_config.text_encoder_configs,
+            self.pipeline_config.image_encoder_config,
+        ]
         for encoder_config in encoder_configs:
             encoder_config.parallel_folding_mode = mode
 

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

@@ -14,6 +14,7 @@ from sglang.multimodal_gen.configs.models.encoders import (
     TextEncoderConfig,
 )
 from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
+from sglang.multimodal_gen.configs.pipeline_configs import PipelineConfig
 from sglang.multimodal_gen.runtime.models.encoders import base as _base_mod
 from sglang.multimodal_gen.runtime.models.encoders.base import (
     FOLD_MIN_HIDDEN_SIZE,
@@ -33,7 +34,7 @@ def _run(
     dp=1,
     disagg=False,
     num_gpus=None,
-    image=(),
+    image=None,
     policy="auto",
     batching_max_size=1,
     explicit=(),
@@ -48,9 +49,9 @@ def _run(
         batching_max_size=batching_max_size,
         is_arg_explicitly_set=lambda name: name in explicit,
         num_gpus=num_gpus if num_gpus is not None else tp * sp * cfg * dp,
-        pipeline_config=SimpleNamespace(
+        pipeline_config=PipelineConfig(
             text_encoder_configs=tuple(encoders),
-            image_encoder_configs=tuple(image),
+            image_encoder_config=(image if image is not None else ImageEncoderConfig()),
         ),
     )
     ServerArgs.adjust_pipeline_config(self)
@@ -130,12 +131,20 @@ def test_all_encoders_get_the_same_proposed_mode():
     img = ImageEncoderConfig()
     for e in (t5, clip, img):
         e.parallel_folding_mode = None
-    _run([t5, clip], tp=1, sp=2, cfg=1, image=[img])
+    _run([t5, clip], tp=1, sp=2, cfg=1, image=img)
     assert t5.parallel_folding_mode == "world"
     assert clip.parallel_folding_mode == "world"
     assert img.parallel_folding_mode == "world"
 
 
+def test_image_encoder_gets_each_policy_proposal():
+    expected_modes = {"auto": "world", "fold": "replica", "replicate": "world"}
+    for policy, expected_mode in expected_modes.items():
+        image = ImageEncoderConfig()
+        _run([], tp=1, sp=2, cfg=1, image=image, policy=policy)
+        assert image.parallel_folding_mode == expected_mode
+
+
 def test_adjust_proposal_policy_dependence():
     # adjust reads the parallelism only for auto/dp/replicate; finalize owns
     # those policy decisions. An explicit fold is the one exception: it widens
@@ -193,6 +202,22 @@ def test_indivisible_dims_not_folded():
     assert encoder_folding_worthwhile(_enc(4096, 64, 10250), group_size=4) is False
 
 
+def test_image_encoder_fold_requires_divisible_dims(monkeypatch):
+    monkeypatch.setattr(
+        _base_mod,
+        "get_folding_tp_group",
+        lambda config: SimpleNamespace(world_size=4),
+    )
+    for heads, expected_mode in ((64, "world"), (6, None)):
+        image = ImageEncoderConfig()
+        image.hidden_size = 4096
+        image.num_attention_heads = heads
+        image.intermediate_size = 10240
+        image.parallel_folding_mode = "world"
+        finalize_encoder_folding(image, "fold")
+        assert image.parallel_folding_mode == expected_mode
+
+
 def test_group_size_one_not_folded():
     assert encoder_folding_worthwhile(_enc(4096, 64, 10240), group_size=1) is False