[Bugfix][HunyuanImage3.0] Fix default guidance_scale from 1.0 to 4.0 and port GPU MoE ForwardContext fix from NPU (#2142)

Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>

examples/offline_inference/text_to_image/text_to_image.py

@@ -64,8 +64,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--guidance-scale",
         type=float,
-        default=1.0,
-        help="Classifier-free guidance scale.",
+        default=4.0,
+        help="Classifier-free guidance scale. HunyuanImage3 recommends 4.0-5.0.",
     )
     parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
     parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")

tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py

@@ -7,6 +7,49 @@ import pytest
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+class TestSetForwardContextNumTokens:
+    """Test _set_forward_context_num_tokens defensive fix for vLLM 0.18.0."""
+
+    def test_sets_num_tokens_when_context_available(self, mocker):
+        """num_tokens should be set on ForwardContext when available."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        del mock_ctx.in_profile_run  # simulate missing attr
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        hunyuan_moe._set_forward_context_num_tokens(1024)
+
+        assert mock_ctx.num_tokens == 1024
+        assert mock_ctx.in_profile_run is False
+
+    def test_sets_in_profile_run_only_if_missing(self, mocker):
+        """in_profile_run should not be overwritten if already set."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.in_profile_run = True  # already set
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        hunyuan_moe._set_forward_context_num_tokens(512)
+
+        assert mock_ctx.num_tokens == 512
+        assert mock_ctx.in_profile_run is True  # not overwritten
+
+    def test_noop_when_context_unavailable(self, mocker):
+        """Should do nothing when ForwardContext is not available."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False)
+        mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context")
+
+        hunyuan_moe._set_forward_context_num_tokens(256)
+
+        mock_get.assert_not_called()
+
+
 class TestHunyuanFusedMoEPlatformDispatch:
     """Test platform dispatch via platform qualname hooks."""
 

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py

@@ -3,12 +3,28 @@
 
 from typing import Any
 
+import vllm.forward_context as _vllm_fc
 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
 from vllm_omni.platforms import current_omni_platform
 
 
+def _set_forward_context_num_tokens(num_tokens: int) -> None:
+    """Set num_tokens on the vLLM ForwardContext for MoE routing.
+
+    After the rebase to vLLM 0.18.0, SharedFusedMoE expects
+    ForwardContext.num_tokens to be set. Without it, MoE expert
+    routing may produce incorrect results (silent correctness bug).
+    """
+    if not _vllm_fc.is_forward_context_available():
+        return
+    forward_context = _vllm_fc.get_forward_context()
+    forward_context.num_tokens = num_tokens
+    if not hasattr(forward_context, "in_profile_run"):
+        forward_context.in_profile_run = False
+
+
 class HunyuanFusedMoEDefault(SharedFusedMoE):
     def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
         super().__init__(prefix=prefix, **kwargs)
@@ -21,6 +37,7 @@ class HunyuanFusedMoEDefault(SharedFusedMoE):
         self._init_hook_handle.remove()
 
     def forward(self, hidden_states: Any, router_logits: Any) -> Any:
+        _set_forward_context_num_tokens(hidden_states.shape[0])
         return super().forward(hidden_states, router_logits)