From eb964dcec5678f6d54347d900f4bb31e9fbed158 Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Thu, 11 Jun 2026 19:06:22 +0800 Subject: [PATCH] [HunyuanImage][Perf] opt prepare_attention_mask for e2e latency 6% reduction (#4333) Signed-off-by: dengyunyang <584797741@qq.com> --- .../models/hunyuan_image3/pipeline_hunyuan_image3.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 988b91920..5d15162fb 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -970,11 +970,14 @@ class HunyuanImage3Pipeline( # This implementation can handle sequences with text and image modalities, where text tokens use causal # attention and image tokens use full attention. bsz, seq_len = inputs_tensor.shape + device = inputs_tensor.device tokenizer_output = model_kwargs["tokenizer_output"] batch_image_slices = [ tokenizer_output.joint_image_slices[i] + tokenizer_output.gen_image_slices[i] for i in range(bsz) ] - attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril(diagonal=0).repeat(bsz, 1, 1) + attention_mask = ( + torch.ones(seq_len, seq_len, dtype=torch.bool, device=device).tril(diagonal=0).repeat(bsz, 1, 1) + ) full_attn_spans: list[list[tuple[int, int]]] = [[] for _ in range(bsz)] for i in range(bsz): for j, image_slice in enumerate(batch_image_slices[i]):