[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 <hgaoaf@connect.ust.hk>

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."""
 

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