diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 58dc2a159..927b0f0b0 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/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.") diff --git a/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py b/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py index 2aa1adf44..2cda9116c 100644 --- a/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py +++ b/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.""" diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py index 5ada5ceb8..cb44717c5 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py +++ b/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)