[Perf][HunyuanImage] remove sync logic (#4401)

Signed-off-by: dengyunyang <584797741@qq.com>

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -71,6 +71,7 @@ def _assert_cot_structural_markers(cot_text: str, label: str) -> None:
 
 
 os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
+os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG"
 
 pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -1065,9 +1065,10 @@ class ImageKVCacheManager:
         assert self.image_kv_cache_lens is not None
         if position_ids is not None:
             assert position_ids.shape == (bs, q_len)
-            assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens), (
-                "The first current position must immediately follow each sample's cached prompt KV."
-            )
+            if logger.isEnabledFor(logging.DEBUG):
+                assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens.to(position_ids.device)), (
+                    "The first current position must immediately follow each sample's cached prompt KV."
+                )
         new_key = torch.cat([cached_key, key], dim=1)
         new_value = torch.cat([cached_value, value], dim=1)
         return new_key.contiguous(), new_value.contiguous()

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -987,7 +987,7 @@ class HunyuanImage3Pipeline(
         timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd)
         x.scatter_(
             dim=1,
-            index=timestep_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
+            index=timestep_scatter_index.to(x.device).unsqueeze(-1).repeat(1, 1, n_embd),
             src=timestep_scatter_src,
         )
 
@@ -1370,7 +1370,7 @@ class HunyuanImage3Pipeline(
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             image_mask=to_device(output.gen_image_mask, device),
-            gen_timestep_scatter_index=to_device(output.gen_timestep_scatter_index, device),
+            gen_timestep_scatter_index=output.gen_timestep_scatter_index,
             cond_vae_images=to_device(cond_vae_images, device),
             cond_timestep=to_device(cond_timestep, device),
             cond_vae_image_mask=to_device(output.cond_vae_image_mask, device),