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

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

vllm_omni/diffusion/diffusion_kv/request.py

@@ -1,13 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from typing import Any
 
-from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys, hash_block_tokens
 from vllm.v1.request import RequestStatus
 
 
@@ -51,8 +52,10 @@ class DiffusionKVRequest:
     An empty ``block_hashes`` sequence means the prefix has no canonical cache
     identity yet. Such a request may use native request-local page allocation,
     but consumers must not publish it through ``KVCacheManager.cache_blocks``.
-    Prefix publication becomes valid only after model preprocessing supplies
-    one canonical hash for every cacheable full block.
+    With prefix caching enabled, Engine invokes model input preparation to
+    attach token IDs and native multimodal identities / positions. The Manager
+    builds one canonical hash for every cacheable full block before lookup;
+    publication requires those hashes and successfully materialized KV.
     """
 
     def __init__(
@@ -64,6 +67,9 @@ class DiffusionKVRequest:
         target_len: int,
         seq_len: int,
         block_hashes: Sequence[BlockHash] = (),
+        cache_token_ids: Sequence[int] = (),
+        mm_features: Sequence[MultiModalFeatureSpec] = (),
+        cache_namespace: str = "diffusion-primary-v1",
         kv_contexts: Sequence[DiffusionKVContext] = (),
         kv_transfer_params: dict[str, Any] | None = None,
         prompt_token_ids: list[int] | None = None,
@@ -83,6 +89,15 @@ class DiffusionKVRequest:
                 "prefix_len + target_len must not exceed seq_len: "
                 f"prefix_len={prefix_len}, target_len={target_len}, seq_len={seq_len}"
             )
+        token_ids = tuple(cache_token_ids)
+        if token_ids and len(token_ids) < prefix_len:
+            raise ValueError(
+                f"cache_token_ids must cover the reusable prefix: tokens={len(token_ids)}, prefix_len={prefix_len}"
+            )
+        if any(type(token_id) is not int for token_id in token_ids):
+            raise TypeError("cache_token_ids must contain only integers")
+        if not isinstance(cache_namespace, str) or not cache_namespace:
+            raise ValueError("cache_namespace must be a non-empty string")
 
         contexts = tuple(kv_contexts)
         if any(not isinstance(context, DiffusionKVContext) for context in contexts):
@@ -96,6 +111,15 @@ class DiffusionKVRequest:
         self.prefix_len = prefix_len
         self.target_len = target_len
         self.kv_contexts = contexts
+        self.cache_token_ids = token_ids
+        self.mm_features = list(mm_features)
+        self.cache_namespace = cache_namespace
+        # Native generate_block_hash_extra_keys consumes this Request surface.
+        # Diffusion adapter ID + scale are already isolated in cache_namespace;
+        # there are no separate LoRA-name, salt, or prompt-embedding inputs.
+        self.lora_request = None
+        self.cache_salt = None
+        self.prompt_embeds = None
 
         # Native vLLM Request surface. Keep this list intentionally small and
         # cover it with real-KVCacheManager conformance tests.
@@ -115,6 +139,70 @@ class DiffusionKVRequest:
         self.kv_transfer_params = kv_transfer_params
         self.prompt_token_ids = prompt_token_ids
 
+    def _validate_mm_features(self) -> None:
+        # Native block traversal expects sorted, non-overlapping placeholders.
+        # Do not silently accept nested/overlapping ranges: the traversal can
+        # skip a feature and incorrectly reuse KV. Combine such dependencies
+        # in the feature identifier instead.
+        previous_end = 0
+        for feature in self.mm_features:
+            if not isinstance(feature, MultiModalFeatureSpec):
+                raise TypeError("mm_features must contain native MultiModalFeatureSpec values")
+            position = feature.mm_position
+            end = position.offset + position.length
+            if position.offset < previous_end or position.length <= 0 or end > self.seq_len:
+                raise ValueError("mm_features must have valid sorted, non-overlapping positions within seq_len")
+            if not isinstance(feature.identifier, str) or not feature.identifier:
+                raise ValueError("mm_features must have non-empty content identifiers")
+            previous_end = end
+
+    def build_block_hashes(
+        self,
+        hash_block_size: int,
+        hash_function: Callable[[object], bytes],
+    ) -> None:
+        """Build native chained hashes for complete reusable-prefix blocks."""
+
+        if self.block_hashes:
+            self.skip_reading_prefix_cache = False
+            return
+        if hash_block_size <= 0:
+            raise ValueError(f"hash_block_size must be positive, got {hash_block_size}")
+        num_cache_tokens = self.prefix_len // hash_block_size * hash_block_size
+        if num_cache_tokens == 0 or not self.cache_token_ids:
+            return
+        if len(self.cache_token_ids) < num_cache_tokens:
+            raise ValueError(
+                "cache_token_ids do not cover the block-aligned reusable prefix: "
+                f"tokens={len(self.cache_token_ids)}, required={num_cache_tokens}"
+            )
+        self._validate_mm_features()
+
+        # Keep the first hash deterministic without mutating vLLM's process-wide
+        # ``NONE_HASH``.  A diffusion manager can be created more than once in
+        # one process (for example in tests or worker reinitialization); using
+        # the global random root would make requests built by different manager
+        # instances impossible to match.
+        parent_hash: BlockHash | None = BlockHash(hash_function(("diffusion-prefix-root-v1", self.cache_namespace)))
+        hashes: list[BlockHash] = []
+        mm_index = 0
+        for start in range(0, num_cache_tokens, hash_block_size):
+            end = start + hash_block_size
+            # Reuse upstream MM intersection/offset logic directly. Keep our
+            # stable-prefix bound: native get_request_block_hasher hashes the
+            # entire token axis, including diffusion's mutable target region.
+            extra_keys, mm_index = generate_block_hash_extra_keys(self, start, end, mm_index)
+            parent_hash = hash_block_tokens(
+                hash_function=hash_function,
+                parent_block_hash=parent_hash,
+                curr_block_token_ids=self.cache_token_ids[start:end],
+                extra_keys=extra_keys,
+            )
+            hashes.append(parent_hash)
+
+        self.block_hashes = hashes
+        self.skip_reading_prefix_cache = not hashes
+
     @property
     def seq_len(self) -> int:
         """Complete first-step sequence length and allocation boundary."""