From d71cacc74ca93b9de50adeb952619ca94273b485 Mon Sep 17 00:00:00 2001 From: TaffyOfficial <2587297563@qq.com> Date: Wed, 24 Jun 2026 11:05:35 +0800 Subject: [PATCH] [perf] avoid HunyuanImage3 timestep scalar sync (#4363) Signed-off-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: Gao Han --- .../test_hunyuan_image3_sampler.py | 62 +++++++++++++++++++ .../models/hunyuan_image3/hunyuan_image3.py | 7 +-- 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py index 51f6a85f5..34f86347a 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py @@ -20,6 +20,7 @@ RATIO_START = 200 RATIO_END = 210 RATIO_OTHER_START = 220 RATIO_OTHER_END = 223 +TIMESTEP = 224 class FakeSamplerModel: @@ -60,6 +61,67 @@ class FakeSamplerModel: _apply_ratio_restriction = _Real._apply_ratio_restriction +class FakeTokenEmbedding: + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return torch.stack( + [ + input_ids.to(dtype=torch.float32), + input_ids.to(dtype=torch.float32) + 0.5, + ], + dim=-1, + ) + + +class FakeEmbeddingModel: + def __init__(self): + self.model = FakeTokenEmbedding() + self._timestep_token_id = TIMESTEP + + from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import ( + HunyuanImage3ForConditionalGeneration as _Real, + ) + + embed_input_ids = _Real.embed_input_ids + + def _timestep_encode(self, timestep: torch.Tensor) -> torch.Tensor: + return torch.full( + (1, 2), + -7.0, + device=timestep.device, + dtype=timestep.dtype, + ) + + +class TestTimestepEmbedding: + def test_timestep_embedding_avoids_scalar_item_sync(self, monkeypatch): + model = FakeEmbeddingModel() + input_ids = torch.tensor([1, TIMESTEP, 2, TIMESTEP], dtype=torch.long) + + def fail_item(_self): + raise AssertionError("embed_input_ids should not synchronize timestep count") + + monkeypatch.setattr(torch.Tensor, "item", fail_item) + + actual = model.embed_input_ids(input_ids) + expected = model.model.embed_input_ids(input_ids) + expected[input_ids == TIMESTEP] = -7.0 + + torch.testing.assert_close(actual, expected) + + def test_timestep_embedding_preserves_no_timestep_inputs(self, monkeypatch): + model = FakeEmbeddingModel() + input_ids = torch.tensor([1, 2, 3], dtype=torch.long) + + def fail_item(_self): + raise AssertionError("embed_input_ids should not synchronize timestep count") + + monkeypatch.setattr(torch.Tensor, "item", fail_item) + + actual = model.embed_input_ids(input_ids) + + torch.testing.assert_close(actual, model.model.embed_input_ids(input_ids)) + + class TestGetForcedToken: """Tests for the stateless _get_forced_token method.""" diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index f2321ba22..d9efccd9a 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -2008,10 +2008,9 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo # the bidirectional MM region and biased AR ratio prediction toward # the first image's bucket. timestep_mask = input_ids == self._timestep_token_id - n_timestep = int(timestep_mask.sum().item()) - if n_timestep > 0: - timestep_input = torch.zeros((n_timestep,), device=inputs_embeds.device, dtype=inputs_embeds.dtype) - inputs_embeds[timestep_mask] = self._timestep_encode(timestep_input) + timestep_input = torch.zeros((), device=inputs_embeds.device, dtype=inputs_embeds.dtype) + timestep_embed = self._timestep_encode(timestep_input).to(inputs_embeds.dtype) + inputs_embeds = torch.where(timestep_mask.unsqueeze(-1), timestep_embed, inputs_embeds) if multimodal_embeddings is None or len(multimodal_embeddings) == 0: return inputs_embeds