[Fix] diffusion: HunyuanImage-3 MoE shared MLP and timestep schedule on CUDA (#46600)

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

@@ -520,13 +520,18 @@ class HunYuanSparseMoeBlock(nn.Module):
         hidden_dim = hidden_states.shape[-1]
         hidden_states = hidden_states.view(-1, hidden_dim)
 
+        # Must run before the experts: the CUDA FusedMoE kernel runs in place
+        # and overwrites hidden_states with the expert output.
+        shared_out = None
+        if self.shared_mlp is not None:
+            shared_out = self.shared_mlp(hidden_states)
+
         router_logits, _ = self.gate(hidden_states)
         topk_output = self.topk(hidden_states, router_logits)
         final_hidden_states = self.experts(hidden_states, topk_output)
 
-        if self.shared_mlp is not None:
-            _shared_out = self.shared_mlp(hidden_states)
-            final_hidden_states = final_hidden_states + _shared_out
+        if shared_out is not None:
+            final_hidden_states = final_hidden_states + shared_out
 
         # NPU: AscendTPDispatcher all-gathers the FusedMoE output internally;
         # the shared MLP's down_proj all-reduces its own output.

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

@@ -121,7 +121,7 @@ class HunyuanImage3Pipeline(LoRAPipeline, ComposedPipelineBase):
                 server_args, pipeline_config, model_path, vae_config_dict
             ),
             "tokenizer": self._load_tokenizer(server_args, model_path),
-            "scheduler": FlowMatchEulerDiscreteScheduler(shift=flow_shift),
+            "scheduler": self._build_scheduler(flow_shift),
             "processor": self._load_processor(server_args, model_path, hf_config),
             "vision_model": self._load_component(
                 HunyuanImage3VisionModel(config_dict["vit"]),
@@ -326,6 +326,14 @@ class HunyuanImage3Pipeline(LoRAPipeline, ComposedPipelineBase):
                 state_dict[name[len(prefix) :]] = tensor
         return state_dict
 
+    @staticmethod
+    def _build_scheduler(flow_shift: float) -> FlowMatchEulerDiscreteScheduler:
+        # Official schedule: shift(linspace(1, 0, N + 1))[:-1], i.e. uniform
+        # sigmas shifted once, not the double-shifted training-sigma range.
+        return FlowMatchEulerDiscreteScheduler(
+            shift=flow_shift, use_uniform_sigmas=True
+        )
+
     @staticmethod
     def _read_flow_shift(model_path: str) -> float:
         return float(

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

@@ -1,5 +1,5 @@
-"""HunyuanImage-3 denoising steps: tensor-parallel input broadcasts and the
-host work each step skips."""
+"""HunyuanImage-3 denoising: tensor-parallel input broadcasts, per-step host
+work, the MoE block and the timestep schedule."""
 
 import math
 import unittest
@@ -8,7 +8,13 @@ from unittest.mock import MagicMock, patch
 
 import torch
 
-from sglang.multimodal_gen.runtime.models.dits import hunyuan_image3_utils
+from sglang.multimodal_gen.runtime.models.dits import (
+    hunyuan_image3,
+    hunyuan_image3_utils,
+)
+from sglang.multimodal_gen.runtime.pipelines.hunyuan_image3 import (
+    HunyuanImage3Pipeline,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan_image3 import (
     ar_stage,
 )
@@ -146,5 +152,43 @@ class TestHunyuanImage3DenoiseStep(CustomTestCase):
         torch.testing.assert_close(actual, (scattered, extracted), atol=0, rtol=0)
 
 
+class TestHunyuanImage3MoEAndSchedule(CustomTestCase):
+    def test_shared_mlp_sees_layer_input_when_experts_run_in_place(self):
+        """The CUDA FusedMoE kernel overwrites its input with the expert
+        output; the shared MLP must still see the layer input."""
+        block = hunyuan_image3.HunYuanSparseMoeBlock.__new__(
+            hunyuan_image3.HunYuanSparseMoeBlock
+        )
+        torch.nn.Module.__init__(block)
+
+        def experts_in_place(x, topk_output):
+            out = x * 10
+            x.copy_(out)
+            return out
+
+        block.gate = lambda x: (torch.zeros(x.shape[0], 2), None)
+        block.topk = lambda x, logits: None
+        block.experts = experts_in_place
+        block.shared_mlp = lambda x: x.clone()
+        x = torch.randn(4, 8)
+
+        out = block(x.clone())
+
+        torch.testing.assert_close(out, 11 * x)
+
+    def test_timesteps_follow_official_schedule(self):
+        """Official FlowMatchDiscreteScheduler: t_i = 1000 * shift(1 - i / N)
+        for i < N, with shift(s) = k * s / (1 + (k - 1) * s)."""
+        shift, steps = 3.0, 10
+        scheduler = HunyuanImage3Pipeline._build_scheduler(shift)
+        scheduler.set_timesteps(steps)
+
+        sigmas = torch.linspace(1, 0, steps + 1)[:-1]
+        expected = 1000 * shift * sigmas / (1 + (shift - 1) * sigmas)
+        torch.testing.assert_close(
+            scheduler.timesteps.float().cpu(), expected, rtol=0, atol=1e-3
+        )
+
+
 if __name__ == "__main__":
     unittest.main()