[Bugfix][HunyuanImage3.0] Fix KV reuse compatibility in SP scenarios (#3546)

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

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -316,6 +316,13 @@ class SequenceParallelSplitHook(ModelHook):
             logger.warning_once(f"Expected tensor with {sp_input.expected_dims} dims, got {x.dim()}. Skipping split.")
             return x
 
+        split_dim = (
+            sp_input.split_dim if isinstance(sp_input, (SequenceParallelInput, SequenceParallelPartialInput)) else None
+        )
+        if split_dim is not None and x.size(split_dim) == 0:
+            logger.warning_once("Skip sharding for zero-sized tensors.")
+            return x
+
         def _raise_strict_divisibility_error(*, dim: int, seq_len: int, sp_size: int) -> None:
             # Keep message actionable: strict mode must be evenly shardable at the split hook level.
             msg = (

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -1032,11 +1032,12 @@ class ImageKVCacheManager:
             self.image_kv_cache_map = None  # reset first
             key, value = self._cache_prompt_kv(key, value, seq_len, shard_image_size)
             if self.sp_size > 1:
-                local_prompt_len = q_len - shard_image_size
-                joint_text_query = query[:, :local_prompt_len, :, :]
+                local_prompt_len = seq_len - shard_image_size
+                join_query_len = query.shape[1] - shard_image_size
+                joint_text_query = query[:, :join_query_len, :, :]
                 joint_text_key = key[:, :local_prompt_len, :, :]
                 joint_text_value = value[:, :local_prompt_len, :, :]
-                query = query[:, local_prompt_len:, :, :]
+                query = query[:, join_query_len:, :, :]
                 key = key[:, local_prompt_len:, :, :]
                 value = value[:, local_prompt_len:, :, :]
         else:
@@ -2250,6 +2251,7 @@ class HunyuanImage3Model(nn.Module):
         num_image_tokens: int | None = None,
         gen_timestep_scatter_index: torch.Tensor | None = None,
         uncond_cfg_prefill: bool = False,
+        ar_kv_reuse_len: int = 0,
         full_attn_spans: list[list[tuple[int, int]]] | None = None,
     ) -> tuple | BaseModelOutputWithPast:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@@ -2305,7 +2307,7 @@ class HunyuanImage3Model(nn.Module):
             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]
+                seq_lens = [prompt_size + shard_image_size + ar_kv_reuse_len 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]
@@ -2734,6 +2736,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             query_lens=[prefill_query_len],
             seq_lens=[prefill_seq_len],
             num_image_tokens=0,
+            ar_kv_reuse_len=negative_reuse_len,
             full_attn_spans=model_kwargs["full_attn_spans"][batch_slice]
             if model_kwargs.get("full_attn_spans")
             else None,
@@ -2787,7 +2790,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
     ):
         ar_kv_data = model_kwargs.pop("ar_kv_data", None)
         if ar_kv_data is None:
-            return input_ids
+            return input_ids, 0
 
         # 1. positive prefix len
         positive_reuse_len, negative_reuse_len = self._get_kv_reuse_len(model_kwargs, batch_size)
@@ -2796,7 +2799,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             f"negative_reuse_len={negative_reuse_len}"
         )
         if positive_reuse_len <= 0:
-            return input_ids
+            return input_ids, 0
 
         # 2. inject positive kv
         self.model.inject_ar_kv_into_layers(ar_kv_data, positive_reuse_len)
@@ -2821,7 +2824,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             positive_reuse_len=positive_reuse_len,
         )
 
-        return input_ids
+        return input_ids, positive_reuse_len
 
     @torch.no_grad()
     def __call__(
@@ -2975,10 +2978,13 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
 
         # Attempt to reuse KV cache from the AR stage.
         # Note: the reusable KV length may differ between positive and negative prompts.
-        input_ids = self._maybe_handle_ar_kv_reuse(
+        input_ids, ar_kv_reuse_len = self._maybe_handle_ar_kv_reuse(
             input_ids, model_kwargs, batch_size, cfg_parallel_ready, cfg_rank, device
         )
 
+        # Store ar_kv_reuse_len in model_kwargs for use in forward method (SP mode)
+        model_kwargs["ar_kv_reuse_len"] = ar_kv_reuse_len
+
         # Sampling loop
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -1011,6 +1011,7 @@ class HunyuanImage3Pipeline(
                 "query_lens": kwargs.get("query_lens"),
                 "seq_lens": kwargs.get("seq_lens"),
                 "num_image_tokens": kwargs.get("num_image_tokens"),
+                "ar_kv_reuse_len": kwargs.get("ar_kv_reuse_len", 0),
                 "full_attn_spans": kwargs.get("full_attn_spans"),
             }
         )
@@ -1162,6 +1163,7 @@ class HunyuanImage3Pipeline(
         seq_lens: list[int] | None = None,
         num_image_tokens: int | None = None,
         uncond_cfg_prefill: bool = False,
+        ar_kv_reuse_len: int = 0,
         full_attn_spans: list[list[tuple[int, int]]] | None = None,
     ) -> tuple | CausalMMOutputWithPast:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -1258,6 +1260,7 @@ class HunyuanImage3Pipeline(
                 num_image_tokens=num_image_tokens,
                 gen_timestep_scatter_index=gen_timestep_scatter_index,
                 uncond_cfg_prefill=uncond_cfg_prefill,
+                ar_kv_reuse_len=ar_kv_reuse_len,
                 full_attn_spans=full_attn_spans,
             )
         hidden_states = outputs[0]