[Diffusion][HunyuanImage3] Add request paged prefix caching (#7719)

Signed-off-by: zwhzzz0821 <2831474076@qq.com>

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -76,7 +76,9 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
     get_paged_kv_computed_tokens,
+    is_forward_context_available,
     paged_kv_prefill,
     set_forward_context_denoise_step_idx,
 )
@@ -2481,6 +2483,7 @@ class HunyuanImage3Model(nn.Module):
         gen_timestep_scatter_index: torch.Tensor | None = None,
         uncond_cfg_prefill: bool = False,
         ar_kv_reuse_len: int = 0,
+        paged_kv_cached_prefix_len: int = 0,
         full_attn_spans: list[list[tuple[int, int]]] | None = None,
     ) -> tuple | BaseModelOutputWithPast:
         current_omni_platform.reset_diffusion_fused_moe_forward_context()
@@ -2539,7 +2542,9 @@ class HunyuanImage3Model(nn.Module):
             else:
                 shard_padding_size = shard_image_size * sp_world_size - num_image_tokens
             if first_step:
-                seq_lens = [prompt_size + shard_image_size + ar_kv_reuse_len for _ in seq_lens]
+                seq_lens = [
+                    prompt_size + shard_image_size + ar_kv_reuse_len + paged_kv_cached_prefix_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 for seq_len in seq_lens]
@@ -3023,6 +3028,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         cfg_rank,
         device,
     ):
+        if is_forward_context_available() and get_forward_context().paged_kv_cached_prefix_len:
+            # Local hits are sliced after conditional-image embeddings are
+            # prepared in forward_call; the AR path must not truncate them.
+            return input_ids, 0
         ar_kv_data = model_kwargs.pop("ar_kv_data", None)
         computed_tokens = get_paged_kv_computed_tokens()
         if computed_tokens and computed_tokens[0] > 0:

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -23,7 +23,11 @@ from vllm.transformers_utils.config import get_config
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
+    is_forward_context_available,
+    set_forward_context_denoise_step_idx,
+)
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
@@ -105,6 +109,10 @@ def to_device(data, device):
         return data
 
 
+def get_hunyuan_image_3_prefix_cache_func(od_config: OmniDiffusionConfig):
+    return request_layout_utils.prepare_hunyuan_prefix_cache
+
+
 def _to_pil_image(image: Any) -> PILImage.Image:
     if isinstance(image, PILImage.Image):
         return image
@@ -1408,6 +1416,13 @@ class HunyuanImage3Pipeline(
         )
         full_attn_spans: list[list[tuple[int, int]]] = [[] for _ in range(bsz)]
         for i in range(bsz):
+            for reference_slice in tokenizer_output.joint_image_slices[i]:
+                logger.debug(
+                    "Hunyuan reference span: sequence_id=%d start=%d end=%d",
+                    i,
+                    int(reference_slice.start or 0),
+                    int(reference_slice.stop or seq_len),
+                )
             for j, image_slice in enumerate(batch_image_slices[i]):
                 attention_mask[i, image_slice, image_slice] = True
                 start = image_slice.start if image_slice.start is not None else 0
@@ -1420,6 +1435,69 @@ class HunyuanImage3Pipeline(
         model_kwargs["full_attn_spans"] = full_attn_spans
         return attention_mask
 
+    @staticmethod
+    def _slice_cached_prefix_inputs(
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor | None,
+        custom_pos_emb: tuple[torch.Tensor, ...] | None,
+        image_mask: torch.Tensor | None,
+        gen_timestep_scatter_index: torch.Tensor | None,
+        query_lens: list[int] | None,
+        seq_lens: list[int] | None,
+        cached_prefix_len: int,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        tuple[torch.Tensor, ...],
+        torch.Tensor,
+        torch.Tensor,
+        list[int],
+        int,
+    ]:
+        """Remove a cached first-step prefix while preserving full KV lengths."""
+
+        seq_len = int(inputs_embeds.shape[1])
+        if (
+            query_lens is None
+            or seq_lens is None
+            or len(query_lens) != inputs_embeds.shape[0]
+            or len(seq_lens) != len(query_lens)
+        ):
+            raise ValueError("Hunyuan paged prefix execution requires one query/sequence length per row")
+        if any(query_len != seq_len for query_len in query_lens) or any(
+            logical_len != seq_len for logical_len in seq_lens
+        ):
+            raise ValueError(
+                "Hunyuan paged prefix slicing requires an untrimmed first-step tensor: "
+                f"query_lens={query_lens}, tensor_len={seq_len}"
+            )
+        if not 0 < cached_prefix_len < seq_len:
+            raise ValueError(
+                "Hunyuan cached prefix must leave a non-empty first-step suffix: "
+                f"cached={cached_prefix_len}, seq_len={seq_len}"
+            )
+        if (
+            attention_mask is None
+            or position_ids is None
+            or custom_pos_emb is None
+            or image_mask is None
+            or gen_timestep_scatter_index is None
+        ):
+            raise ValueError("Hunyuan paged prefix slicing is missing first-step layout tensors")
+
+        return (
+            inputs_embeds[:, cached_prefix_len:],
+            attention_mask[:, :, cached_prefix_len:, :],
+            position_ids[:, cached_prefix_len:],
+            tuple(pos[:, cached_prefix_len:] for pos in custom_pos_emb),
+            image_mask[:, cached_prefix_len:],
+            gen_timestep_scatter_index - cached_prefix_len,
+            [query_len - cached_prefix_len for query_len in query_lens],
+            seq_len - cached_prefix_len,
+        )
+
     def prepare_inputs_for_generation(
         self,
         input_ids,
@@ -1747,6 +1825,40 @@ class HunyuanImage3Pipeline(
         assert inputs_embeds is not None
         bsz, seq_len, n_embd = inputs_embeds.shape
 
+        paged_kv_cached_prefix_len = 0
+        if first_step and self._uses_scheduler_paged_kv():
+            if not is_forward_context_available():
+                raise RuntimeError("Hunyuan paged prefix execution requires an Omni ForwardContext")
+            paged_kv_cached_prefix_len = int(get_forward_context().paged_kv_cached_prefix_len)
+        if paged_kv_cached_prefix_len:
+            if ar_kv_reuse_len:
+                raise ValueError("Hunyuan cannot combine imported AR KV and Scheduler prefix-cache hits")
+            (
+                inputs_embeds,
+                attention_mask,
+                position_ids,
+                custom_pos_emb,
+                image_mask,
+                gen_timestep_scatter_index,
+                query_lens,
+                seq_len,
+            ) = self._slice_cached_prefix_inputs(
+                inputs_embeds,
+                attention_mask,
+                position_ids,
+                custom_pos_emb,
+                image_mask,
+                gen_timestep_scatter_index,
+                query_lens,
+                seq_lens,
+                paged_kv_cached_prefix_len,
+            )
+            logger.debug(
+                "Hunyuan prefix slice: cached_prefix_len=%d query_len=%d",
+                paged_kv_cached_prefix_len,
+                seq_len,
+            )
+
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         from vllm.forward_context import set_forward_context
 
@@ -1770,6 +1882,7 @@ class HunyuanImage3Pipeline(
                 gen_timestep_scatter_index=gen_timestep_scatter_index,
                 uncond_cfg_prefill=uncond_cfg_prefill,
                 ar_kv_reuse_len=ar_kv_reuse_len,
+                paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
                 full_attn_spans=full_attn_spans,
             )
         hidden_states = outputs[0]

vllm_omni/diffusion/models/hunyuan_image3/request_layout.py

@@ -8,7 +8,9 @@ from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
 
+from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
@@ -357,7 +359,7 @@ def build_hunyuan_diffusion_kv_requests(
     request: OmniDiffusionRequest,
     prepared_layout: HunyuanPreparedLayout,
 ) -> tuple[DiffusionKVRequest, ...]:
-    """Build one persistent Scheduler KV request per Hunyuan execution row."""
+    """Build allocation-only KV requests, even when prefix caching is disabled."""
 
     tokenizer_output = prepared_layout.tokenizer_output
     cfg_factor = hunyuan_cfg_factor(
@@ -408,7 +410,15 @@ def build_hunyuan_diffusion_kv_requests(
             prefix_len=int(prefix_row[-1].item()),
             target_len=target_len,
             seq_len=int(valid_row[-1].item()),
-            prompt_token_ids=tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist(),
+            # Native AR -> DiT transfer needs token IDs to describe the
+            # transferred prefix.  Local paged prefix caching does not: its
+            # identity is attached later by ``prepare_hunyuan_prefix_cache``.
+            # Keep this conversion out of the disabled/local-only path.
+            prompt_token_ids=(
+                tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist()
+                if request.kv_transfer_params is not None
+                else None
+            ),
             # Prompt and reference-image tokens are already embedded in this
             # row's primary self-attention sequence. Hunyuan therefore has no
             # independently projected cross/joint-attention KV context.
@@ -418,6 +428,93 @@ def build_hunyuan_diffusion_kv_requests(
     )
 
 
+def prepare_hunyuan_prefix_cache(request: OmniDiffusionRequest) -> None:
+    """Attach native MM identities / positions when Engine enables caching.
+
+    This is model input adaptation, not a hashing framework. All reference
+    inputs are hashed once, then reused by VAE/ViT subspans and CFG rows. The
+    native block hasher consumes these ranges directly, without token extras.
+    """
+
+    layout = get_hunyuan_prepared_layout(request)
+    if layout is None or not request.diffusion_kv_requests:
+        raise ValueError("Hunyuan prefix caching requires prepared layout and KV requests")
+    sampling = request.sampling_params
+    _, _, _, images, _ = extract_hunyuan_prompt_inputs(
+        [request.prompt], sampling.extra_args or {}, request_id=request.request_id, allow_cond_image=True
+    )
+    reference_digest = None
+    if images:
+        # Match Worker: an explicitly supplied generator overrides the seed.
+        # Conditional VAE samples latents, so image bytes alone are not enough.
+        if sampling.generator is not None:
+            generators = sampling.generator if isinstance(sampling.generator, list) else [sampling.generator]
+            random_state: object = (
+                "generator-state",
+                tuple((str(generator.device), generator.get_state()) for generator in generators),
+            )
+        elif sampling.seed is not None:
+            random_state = ("seed", int(sampling.seed), str(sampling.generator_device or "worker-default"))
+        else:
+            random_state = ("request-local-random-state", request.request_id)
+        # Preserve the conservative whole-reference-set policy. Splitting
+        # independent image identities also needs the VAE RNG sequence modeled.
+        reference_digest = hash_prefix_cache_value(
+            (
+                "hunyuan-reference-v1",
+                [[joint_image_info_to_payload(image) for image in row] for row in images],
+                random_state,
+            )
+        )
+
+    namespace = get_cache_namespace("hunyuan-image3-primary-v4", sampling)
+    output = layout.tokenizer_output
+    joint_rows = output.joint_image_slices or [[] for _ in range(layout.num_branches)]
+    prepared = []
+    for row in request.diffusion_kv_requests:
+        if row.block_hashes or row.num_computed_tokens:
+            raise ValueError("Hunyuan cache inputs must be prepared before KV execution")
+        joint_spans = [(int(span.start or 0), int(span.stop or 0)) for span in joint_rows[row.sequence_id]]
+        features = []
+        for image_slice, (height, width) in layout.rope_image_info[row.sequence_id]:
+            start = int(image_slice.start or 0)
+            end = int(image_slice.stop or output.tokens.shape[1])
+            if start >= row.prefix_len or end <= 0:
+                continue
+            # Joint spans enclose VAE + separator + ViT; RoPE lists subspans.
+            is_reference = any(joint_start <= start and end <= joint_end for joint_start, joint_end in joint_spans)
+            if is_reference and reference_digest is None:
+                raise ValueError("Hunyuan reference-image prefix is missing its canonical inputs")
+            identifier = hash_prefix_cache_value(
+                (
+                    "hunyuan-image-span-v2",
+                    start,
+                    end,
+                    int(height),
+                    int(width),
+                    is_reference,
+                    reference_digest if is_reference else None,
+                )
+            ).hex()
+            features.append(
+                MultiModalFeatureSpec(
+                    data=None,
+                    modality="image",
+                    identifier=identifier,
+                    mm_position=PlaceholderRange(offset=start, length=end - start),
+                )
+            )
+        # Raw model inputs remain in request.prompt/prepared_layout; data=None
+        # here does not imply an encoder hit. These are Scheduler-only KV keys.
+        tokens = tuple(output.tokens[row.sequence_id, : row.prefix_len].tolist())
+        prepared.append((row, tokens, sorted(features, key=lambda feature: feature.mm_position.offset)))
+
+    for row, tokens, features in prepared:
+        row.cache_token_ids = tokens
+        row.mm_features = features
+        row.cache_namespace = namespace
+
+
 def get_hunyuan_prepared_layout(source: Any) -> HunyuanPreparedLayout | None:
     prepared_layout = getattr(source, "prepared_layout", None)
     if prepared_layout is None: