[Feature] support sp for hunyuan (#2163)

Signed-off-by: dengyunyang <584797741@qq.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

examples/offline_inference/text_to_image/text_to_image.py

@@ -237,6 +237,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable logging of diffusion pipeline stats.",
+    )
     return parser.parse_args()
 
 
@@ -326,6 +331,7 @@ def main():
         "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
         "mode": "text-to-image",
+        "log_stats": args.log_stats,
         "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         **lora_args,
         **quant_kwargs,

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -373,8 +373,14 @@ class UlyssesParallelAttention:
             orig_head_cnt=int(orig_head_cnt) if mode == "advanced_uaa" else 0,
             joint_orig_head_cnt=int(joint_orig_head_cnt) if mode == "advanced_uaa" else 0,
         )
-
+        use_2d_mask = False
         if attn_metadata is not None:
+            if attn_metadata.attn_mask is not None and attn_metadata.attn_mask.ndim == 2:
+                use_2d_mask = True
+            if attn_metadata.joint_attn_mask is not None and attn_metadata.joint_attn_mask.ndim == 2:
+                use_2d_mask = True
+
+        if attn_metadata is not None and use_2d_mask:
             if is_joint:
                 if attn_metadata.joint_attn_mask is None and attn_metadata.attn_mask is None:
                     attn_metadata.attn_mask = None

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -242,6 +242,7 @@ class SequenceParallelSplitHook(ModelHook):
 
         if not is_tensor and not is_tensor_list:
             # No tensor outputs to shard
+            logger.warning_once("No tensor outputs to shard.")
             return output
 
         output_list = [output] if is_tensor else list(output)

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py

@@ -11,7 +11,6 @@ from typing import Any, cast
 import numpy as np
 import regex as re
 import torch
-import torch.nn.functional as F
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -57,8 +56,19 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.v1.attention.backend import AttentionType
 
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+)
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_pp_group,
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+)
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import HunyuanFusedMoE
@@ -411,14 +421,18 @@ def apply_rotary_pos_emb(
 
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    Equivalent to torch.repeat_interleave(x, dim=2, repeats=n_rep).
+    Input:  (batch, seqlen, num_key_value_heads, head_dim)
+    Output: (batch, seqlen, num_attention_heads, head_dim)
     """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    batch, slen, num_key_value_heads, head_dim = hidden_states.shape
+
     if n_rep == 1:
         return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+    hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim)
+
+    return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim)
 
 
 def default(value, default_value):
@@ -830,15 +844,31 @@ class ImageKVCacheManager:
     Manages specialized caching and updating of KV-Cache for image tokens in multimodal models.
     """
 
-    def __init__(self, image_token_len: int = 4097):
+    def __init__(self, num_heads: int, num_kv_heads: int, head_dim: int, scaling: float, image_token_len: int = 4097):
         """
         Args:
             image_token_len: Number of tokens per image (including special placeholders),
             default 4097 (timestamp + 4096 image tokens).
         """
+        # attention related
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.scaling = scaling
+        # cache related
         self.image_token_len: int = image_token_len
         self.image_kv_cache: tuple[torch.Tensor, torch.Tensor] = None
 
+        self.sp_size = get_sequence_parallel_world_size()
+        self.sp_rank = get_sequence_parallel_rank()
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            causal=False,
+            softmax_scale=self.scaling,
+            num_kv_heads=self.num_kv_heads,
+        )
+
     def _save_image_kv_caches(
         self,
         key: torch.Tensor,
@@ -910,6 +940,63 @@ class ImageKVCacheManager:
 
         return new_key.contiguous(), new_value.contiguous()
 
+    def _sp_save_prompt_kv_caches(
+        self,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        seq_len: int,
+        shard_image_size: int,
+    ) -> None:
+        """
+        We don't need to cached last token, since it all false and nevel attn to in non-first step.
+        With sp the key len is [text, image_shard], image_shard in rank include [image token, padding token]
+        The cache is [prompt0, prompt1], see _prepare_attention_mask_for_generation and
+        _update_model_kwargs_for_generation
+        """
+        bs, q_len, num_kv_heads, head_dim = key.shape
+        assert q_len == seq_len, f"for first-step, {q_len} != {seq_len}"
+        key = key.reshape(-1, num_kv_heads, head_dim)
+        value = value.reshape(-1, num_kv_heads, head_dim)
+        cached_prompt_len = seq_len - shard_image_size
+        if bs > 1:
+            assert bs == 2, "for cfg case, bs must be 2"
+
+        cached_key = []
+        cached_value = []
+        for b in range(bs):
+            base = b * seq_len
+            # cache text prompt
+            cached_key.append(key[base : base + cached_prompt_len])
+            cached_value.append(value[base : base + cached_prompt_len])
+
+        cached_key = torch.cat(cached_key, dim=0)
+        cached_value = torch.cat(cached_value, dim=0)
+        self.image_kv_cache_map = (cached_key, cached_value)
+
+    def _sp_get_prompt_kv_caches(
+        self,
+        key: torch.Tensor,
+        seq_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cached_key, cached_value = self.image_kv_cache_map
+        bs, q_len, kv_head_num, head_dim = key.shape
+
+        cached_prompt_len = cached_key.shape[0] // bs
+
+        assert cached_prompt_len == seq_len - q_len
+
+        joint_text_key = []
+        joint_text_value = []
+        for b in range(bs):
+            cache_base = b * cached_prompt_len
+            joint_text_key.append(cached_key[cache_base : cache_base + cached_prompt_len])
+            joint_text_value.append(cached_value[cache_base : cache_base + cached_prompt_len])
+
+        joint_text_key = torch.cat(joint_text_key, dim=0).reshape(bs, cached_prompt_len, kv_head_num, head_dim)
+        joint_text_value = torch.cat(joint_text_value, dim=0).reshape(bs, cached_prompt_len, kv_head_num, head_dim)
+
+        return joint_text_key.contiguous(), joint_text_value.contiguous()
+
     def __call__(
         self,
         query: torch.Tensor,
@@ -923,6 +1010,7 @@ class ImageKVCacheManager:
 
         query_lens = kwargs.get("query_lens")
         seq_lens = kwargs.get("seq_lens")
+        shard_image_size = kwargs.get("shard_image_size")
         bs = len(query_lens)
         q_len = query_lens[0]
         seq_len = seq_lens[0]
@@ -936,25 +1024,54 @@ class ImageKVCacheManager:
         query = query.reshape(bs, q_len, head_num_per_rank, head_dim)
         key = key.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
         value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
-
         if first_step:
             self.image_kv_cache_map = None
-            self._save_image_kv_caches(key, value, seq_len)
+            if self.sp_size <= 1:
+                self._save_image_kv_caches(key, value, seq_len)
+            else:
+                self._sp_save_prompt_kv_caches(key, value, seq_len, shard_image_size)
+                cached_prompt_len = self.image_kv_cache_map[0].shape[0] // bs
+                # joint text part
+                joint_text_query = query[:, :cached_prompt_len, :, :]
+                joint_text_key = key[:, :cached_prompt_len, :, :]
+                joint_text_value = value[:, :cached_prompt_len, :, :]
+                # image part
+                query = query[:, cached_prompt_len:, :, :]
+                key = key[:, cached_prompt_len:, :, :]
+                value = value[:, cached_prompt_len:, :, :]
         else:
-            key, value = self._update_image_kv_caches(key, value, seq_len)
-
-        query = query.transpose(1, 2).contiguous()
-        key = key.transpose(1, 2).contiguous()
-        value = value.transpose(1, 2).contiguous()
+            if self.sp_size <= 1:
+                key, value = self._update_image_kv_caches(key, value, seq_len)
+            else:
+                joint_text_query = query[:, :0, :, :]
+                joint_text_key, joint_text_value = self._sp_get_prompt_kv_caches(key, seq_len)
 
         key = repeat_kv(key, repeat_num)
         value = repeat_kv(value, repeat_num)
+        if self.sp_size > 1:
+            joint_text_key = repeat_kv(joint_text_key, repeat_num)
+            joint_text_value = repeat_kv(joint_text_value, repeat_num)
 
         attention_mask = attention_mask.contiguous()
 
-        attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0)
+        if self.sp_size <= 1:
+            attn_metadata = AttentionMetadata(
+                attn_mask=attention_mask,
+            )
+        else:
+            attn_metadata = AttentionMetadata(
+                joint_query=joint_text_query,
+                joint_key=joint_text_key,
+                joint_value=joint_text_value,
+                joint_strategy="front",
+                attn_mask=attention_mask,
+            )
+        # Compute attention using unified attention layer
+        attn_output = self.attn(query, key, value, attn_metadata)
 
-        attn_output = attn_output.transpose(1, 2).contiguous()  # [bs, q_len, heads, head_dim]
+        # attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0)
+
+        # attn_output = attn_output.transpose(1, 2).contiguous()  # [bs, q_len, heads, head_dim]
         attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim)
         return attn_output
 
@@ -1512,7 +1629,13 @@ class HunYuanAttention(nn.Module):
         )
 
         # default image_token_len = timestamp + 4096*image_tokes
-        self.image_attn = ImageKVCacheManager(image_token_len=4097)
+        self.image_attn = ImageKVCacheManager(
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+            num_kv_heads=self.num_kv_heads,
+            scaling=self.scaling,
+            image_token_len=4097,
+        )
         self.image_rope2d_emb = HunYuanRotary2DEmbedder(
             num_heads=self.num_heads,
             num_kv_heads=self.num_kv_heads,
@@ -1709,8 +1832,73 @@ class HunyuanImage3PreTrainedModel(PreTrainedModel):
                 module.weight.data[module.padding_idx].zero_()
 
 
+class HunyuanImagePreprocessor(nn.Module):
+    """
+    Prepares hidden_states for SP by separating to  text, image parts
+
+    Split hidden_states [B, seq_len, D] to
+        - text part: [B, text_seq_len, D]
+        - image part: [B, image_seq_len + last_token, D]
+    """
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        custom_pos_emb: tuple[torch.FloatTensor],
+        position_ids: torch.LongTensor,
+        image_token_len: int,
+        is_first_step: bool = False,
+    ):
+        # ---------- compute prompt length ----------
+        prompt_len = hidden_states.shape[1] - image_token_len - 1 if is_first_step else 0
+
+        # ---------- hidden_states ----------
+        text_hidden_states = hidden_states[:, :prompt_len, :]
+        image_hidden_states = hidden_states[:, prompt_len:, :]
+
+        # ---------- position_ids ----------
+        text_position_ids = position_ids[:, :prompt_len]
+        image_position_ids = position_ids[:, prompt_len:]
+
+        # ---------- custom_pos_emb ----------
+        text_custom_pos_emb = (custom_pos_emb[0][:, :prompt_len, :], custom_pos_emb[1][:, :prompt_len, :])
+        image_custom_pos_emb = (custom_pos_emb[0][:, prompt_len:, :], custom_pos_emb[1][:, prompt_len:, :])
+
+        return (
+            text_hidden_states,
+            image_hidden_states,
+            text_position_ids,
+            image_position_ids,
+            text_custom_pos_emb[0],
+            image_custom_pos_emb[0],
+            text_custom_pos_emb[1],
+            image_custom_pos_emb[1],
+        )
+
+
+class UnifiledCat(nn.Module):
+    def forward(self, x, y, dim: int = 1):
+        return torch.cat([x, y], dim=dim)
+
+
+class HunyuanImagePostprocessor(nn.Module):
+    def forward(self, x):
+        return x
+
+
 class HunyuanImage3Model(nn.Module):
-    def __init__(self, config: HunyuanImage3Config, quant_config: QuantizationConfig | None, prefix: str = ""):
+    _sp_plan = {
+        # Split custom_pos_emb tuple elements (cos, sin) at model forward input
+        "pre_processor": {
+            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            3: SequenceParallelInput(split_dim=1, expected_dims=2, split_output=True, auto_pad=True),
+            5: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),  # sin
+            7: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),  # cos
+        },
+        "post_processor": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
+    def __init__(self, config: HunyuanImage3Config, quant_config=None, prefix: str = ""):
         super().__init__()
         lora_config = None
         self.num_redundant_experts = 0
@@ -1748,6 +1936,10 @@ class HunyuanImage3Model(nn.Module):
         else:
             self.norm = PPMissingLayer()
 
+        self.pre_processor = HunyuanImagePreprocessor()
+        self.unifiled_cat = UnifiledCat()
+        self.post_processor = HunyuanImagePostprocessor()
+
     def _split_qkv_weight(self, qkv: torch.Tensor):
         num_attention_heads = self.config.num_attention_heads
         num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads)
@@ -1998,6 +2190,11 @@ class HunyuanImage3Model(nn.Module):
             loaded_params.add(name)
         return loaded_params
 
+    def _split_result(self, x, text_prompt_len):
+        text_output = x[:, 0:text_prompt_len, :].contiguous()
+        image_output = x[:, text_prompt_len:, :].contiguous()
+        return text_output, image_output
+
     def forward(
         self,
         input_ids: torch.LongTensor | None = None,
@@ -2035,6 +2232,71 @@ class HunyuanImage3Model(nn.Module):
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
         next_decoder_cache = None
+        sp_world_size = get_sequence_parallel_world_size()
+        prompt_size = 0
+        shard_image_size = 0
+        shard_padding_size = 0
+        origin_query_len = query_lens[0]
+        if sp_world_size > 1:
+            assert query_lens[0] == hidden_states.shape[1]
+            (
+                text_hidden_states,
+                image_hidden_states,
+                text_position_ids,
+                image_position_ids,
+                text_custom_pos_emb_sin,
+                image_custom_pos_emb_sin,
+                text_custom_pos_emb_cos,
+                image_custom_pos_emb_cos,
+            ) = self.pre_processor(
+                hidden_states,
+                custom_pos_emb,
+                position_ids,
+                num_image_tokens,
+                first_step,
+            )
+            assert len(set(query_lens)) == 1 and len(set(seq_lens)) == 1, (
+                "query_lens and seq_lens must be the same for sequence parallel"
+            )
+            prompt_size = text_hidden_states.shape[1]
+            shard_image_size = image_hidden_states.shape[1]
+
+            if first_step:  # [image tokens, last token]
+                shard_padding_size = shard_image_size * sp_world_size - (num_image_tokens + 1)
+            else:  # [image tokens, last token]
+                shard_padding_size = shard_image_size * sp_world_size - num_image_tokens
+            if first_step:
+                seq_lens = [prompt_size + shard_image_size for _ in seq_lens]
+            else:
+                seq_lens = [x - y for x, y in zip(seq_lens, query_lens)]
+                seq_lens = [seq_len + shard_image_size - 1 for seq_len in seq_lens]
+            query_lens = [prompt_size + shard_image_size for _ in query_lens]
+
+            if prompt_size > 0:
+                hidden_states = self.unifiled_cat(text_hidden_states, image_hidden_states, dim=1)
+                position_ids = self.unifiled_cat(text_position_ids, image_position_ids, dim=1)
+                custom_pos_emb = (
+                    self.unifiled_cat(text_custom_pos_emb_sin, image_custom_pos_emb_sin, dim=1),
+                    self.unifiled_cat(text_custom_pos_emb_cos, image_custom_pos_emb_cos, dim=1),
+                )
+            else:
+                hidden_states = image_hidden_states
+                position_ids = image_position_ids
+                custom_pos_emb = (image_custom_pos_emb_sin, image_custom_pos_emb_cos)
+
+            if not first_step:
+                assert torch.all(~attention_mask[..., -1]), "The last token should not be attended to"
+                attention_mask = attention_mask[..., :-1]  # we attention without last token
+
+            if shard_padding_size > 0:
+                B, H, Q, K = attention_mask.shape
+                pad = shard_padding_size
+
+                q_pad = attention_mask.new_zeros(B, H, pad, K)
+                attention_mask = torch.cat((attention_mask, q_pad), dim=2)
+
+                k_pad = attention_mask.new_zeros(B, H, Q + pad, pad)
+                attention_mask = torch.cat((attention_mask, k_pad), dim=3)
 
         for layer_idx, decoder_layer in enumerate(self.layers):
             if output_hidden_states:
@@ -2055,6 +2317,8 @@ class HunyuanImage3Model(nn.Module):
                 seq_lens=seq_lens,
                 num_image_tokens=num_image_tokens,
                 gen_timestep_scatter_index=gen_timestep_scatter_index,
+                shard_image_size=shard_image_size,
+                shard_padding_size=shard_padding_size,
             )
 
             hidden_states = layer_outputs[0]
@@ -2068,6 +2332,13 @@ class HunyuanImage3Model(nn.Module):
         # add hidden states from the last decoder layer
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
+        if sp_world_size > 1:
+            text_output, image_output = self._split_result(hidden_states, prompt_size)
+            hidden_states = self.post_processor(image_output)
+            hidden_states = self.unifiled_cat(text_output, hidden_states, dim=1)
+            # Since the SP system only records the primary padding information from the first stage,
+            # but different stages may introduce different padding sizes, we truncate it to ensure correct output.
+            hidden_states = hidden_states[:, :origin_query_len, :].contiguous()
 
         next_cache = None
         if use_cache:

vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import logging
+import os
 from collections.abc import Iterable
 from typing import Any
 
@@ -14,12 +15,13 @@ from transformers.generation.utils import ALL_CACHE_NAMES, GenerationMixin
 from transformers.models.siglip2 import Siglip2VisionConfig, Siglip2VisionModel
 from transformers.utils.generic import ModelOutput
 from vllm.config.vllm import get_current_vllm_config
-from vllm.distributed import get_tensor_model_parallel_rank
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.transformers_utils.config import get_config
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 from .autoencoder import AutoencoderKLConv3D
@@ -61,7 +63,16 @@ def to_device(data, device):
         return data
 
 
-class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
+class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin, DiffusionPipelineProfilerMixin):
+    _PROFILER_TARGETS = [
+        "model.forward",
+        "model.layers[0].forward",
+        "model.layers[0].self_attn.forward",
+        "model.layers[0].mlp.forward",
+        "vae.encode",
+        "vae.decode",
+    ]
+
     def __init__(self, od_config: OmniDiffusionConfig) -> None:
         self.hf_config = get_config(od_config.model, trust_remote_code=True)
         super().__init__(self.hf_config)
@@ -78,7 +89,13 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
             )
         ]
         quant_config = od_config.quantization_config
+        os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
+        logger.info(
+            "Setting attention backend to TORCH_SDPA. "
+            "HunyuanImage3Pipeline only supports TORCH_SDPA to handle mixed causal and full attention."
+        )
         self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config)
+        self.transformer = self.model
         self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae)
         self._pipeline = None
         self._tkwrapper = TokenizerWrapper(od_config.model)
@@ -112,6 +129,9 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
         self.lm_head = nn.Linear(self.hf_config.hidden_size, self.hf_config.vocab_size, bias=False)
         self.vllm_config = get_current_vllm_config()
         self.post_init()
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
+        )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else []
@@ -129,8 +149,8 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
             "time_embed_2",
             "final_layer.model",
         ]
-        tp_rank = get_tensor_model_parallel_rank()
-        device_str = f"{self.model.device.type}:{tp_rank}"
+
+        device_str = f"{get_local_device()}"
         named_modules = dict(self.named_modules())
         for prefix in non_model_layer_prefixes:
             mod = named_modules.get(prefix)
@@ -712,6 +732,7 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
         updated_model_kwargs = {
             "mode": mode,
             "custom_pos_emb": model_kwargs["custom_pos_emb"],
+            "num_image_tokens": model_kwargs["num_image_tokens"],
         }
 
         # update past_key_values keeping its naming used in model code
@@ -997,4 +1018,6 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
             guidance_scale=guidance_scale,
         )
         outputs = self._generate(**model_inputs, **kwargs)
-        return DiffusionOutput(output=outputs[0])
+        return DiffusionOutput(
+            output=outputs[0], stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )

vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py

@@ -1,4 +1,5 @@
 import functools
+import re
 import time
 from collections.abc import Callable
 from threading import Lock
@@ -30,20 +31,31 @@ def profiler(name: str, func: Callable, instance: Any) -> Callable:
             logger.info(f"[DiffusionPipelineProfiler] {name} took {duration:.6f}s")
             # record the profiling data: duration of stages
             with instance._profiler_lock:
-                instance._stage_durations[name] = duration
+                instance._stage_durations[name] = instance._stage_durations.get(name, 0.0) + duration
 
     return wrapper
 
 
+def _parse_part(part: str) -> tuple[str, int | None]:
+    """Parse 'att[num]' into ('att', num)."""
+    if m := re.compile(r"(\w+)\[(\d+)\]").fullmatch(part):
+        return m.group(1), int(m.group(2))
+    return part, None
+
+
 def _get_attribute_by_path(obj: Any, path: str) -> tuple[Any, str]:
     """Traverse an object by dotted path and return (parent_obj, attribute_name)."""
     parts = path.split(".")
     current = obj
 
     for part in parts[:-1]:
-        current = getattr(current, part, None)
+        attr, idx = _parse_part(part)
+
+        current = getattr(current, attr, None)
         if current is None:
             return None, None
+        if idx is not None:
+            current = current[idx]
 
     return current, parts[-1]