diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan_image3.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan_image3.py index 1b402fe7d6..f7807fb13b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan_image3.py +++ b/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. diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_image3.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_image3.py index b26e6a34cb..2e7231667b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_image3.py +++ b/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( diff --git a/python/sglang/multimodal_gen/test/unit/test_hunyuan_image3_tp_broadcast.py b/python/sglang/multimodal_gen/test/unit/test_hunyuan_image3_tp_broadcast.py index 91d636341e..bebbcb58c5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_hunyuan_image3_tp_broadcast.py +++ b/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()