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

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

vllm_omni/diffusion/diffusion_kv/manager.py

@@ -1,11 +1,13 @@
 # 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 vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+from vllm.v1.core.kv_cache_utils import KVCacheBlock
 from vllm.v1.kv_cache_interface import KVCacheConfig
 
 from vllm_omni.core.sched.utils import free_kv_blocks_in_physical_order
@@ -31,6 +33,8 @@ class DiffusionKVCacheManager:
         scheduler_block_size: int,
         hash_block_size: int,
         max_in_flight_tokens: int | None = None,
+        enable_prefix_caching: bool = False,
+        prefix_caching_hash_algo: str = "sha256",
     ) -> None:
         if max_model_len <= 0:
             raise ValueError(f"max_model_len must be positive, got {max_model_len}")
@@ -40,15 +44,18 @@ class DiffusionKVCacheManager:
             raise ValueError("scheduler_block_size and hash_block_size must be positive")
         if max_in_flight_tokens is not None and max_in_flight_tokens <= 0:
             raise ValueError(f"max_in_flight_tokens must be positive when provided, got {max_in_flight_tokens}")
+        self._hash_function = get_hash_fn_by_name(prefix_caching_hash_algo)
         self.native_manager = KVCacheManager(
             kv_cache_config=kv_cache_config,
             max_model_len=max_model_len,
             scheduler_block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
             max_in_flight_tokens=max_in_flight_tokens,
-            enable_caching=False,
+            enable_caching=enable_prefix_caching,
         )
         self.max_model_len = max_model_len
+        self.hash_block_size = hash_block_size
+        self.enable_prefix_caching = enable_prefix_caching
         # Native vLLM may reserve a null block, so an idle BlockPool does not
         # necessarily report ``kv_cache_config.num_blocks`` free blocks.
         self._empty_pool_num_free_blocks = self.native_manager.block_pool.get_num_free_blocks()
@@ -57,31 +64,87 @@ class DiffusionKVCacheManager:
         self._internal_request_ids: set[str] = set()
         self._next_allocation_generation = 1
 
-    def _get_empty_pool_required_blocks(self, requests: Sequence[DiffusionKVRequest]) -> int:
-        """Return the native full-sequence reservation for one public request.
+    def _prepare_block_hashes(self, requests: Sequence[DiffusionKVRequest]) -> None:
+        if not self.enable_prefix_caching:
+            return
+        for request in requests:
+            request.build_block_hashes(self.hash_block_size, self._hash_function)
 
-        vLLM performs this calculation inside ``allocate_slots`` when
-        ``full_sequence_must_fit`` is enabled. Diffusion CFG adds one semantic
-        requirement that native requests do not have: every execution sequence
-        must fit atomically. Sum the native per-sequence result so a request
-        that can never fit is rejected independently of current pool usage.
+    def _clip_computed_blocks(
+        self,
+        blocks: KVCacheBlocks,
+        *,
+        hit_len: int,
+        cached_prefix_len: int,
+    ) -> KVCacheBlocks:
+        if cached_prefix_len == hit_len:
+            return blocks
+        clipped_groups: list[list[KVCacheBlock]] = []
+        for group_blocks, cache_group in zip(
+            blocks.blocks,
+            self.native_manager.kv_cache_config.kv_cache_groups,
+            strict=True,
+        ):
+            block_size = int(cache_group.kv_cache_spec.block_size)
+            if cached_prefix_len % block_size:
+                raise RuntimeError(
+                    "Diffusion prefix hit is not aligned to a native cache group: "
+                    f"cached_prefix_len={cached_prefix_len}, block_size={block_size}"
+                )
+            keep_blocks = cached_prefix_len // block_size
+            if keep_blocks > len(group_blocks):
+                raise RuntimeError(
+                    "Native prefix lookup returned too few blocks for the common CFG boundary: "
+                    f"required={keep_blocks}, returned={len(group_blocks)}"
+                )
+            clipped_groups.append(list(group_blocks[:keep_blocks]))
+        return self.native_manager.create_kv_cache_blocks(tuple(clipped_groups))
+
+    @staticmethod
+    def _unique_blocks(block_sets: Sequence[KVCacheBlocks]) -> list[KVCacheBlock]:
+        unique: dict[int, KVCacheBlock] = {}
+        for blocks in block_sets:
+            for group in blocks.blocks:
+                for block in group:
+                    if not block.is_null:
+                        unique[id(block)] = block
+        return list(unique.values())
+
+    def _get_empty_pool_required_blocks(
+        self,
+        requests: Sequence[DiffusionKVRequest],
+        computed_blocks: Sequence[KVCacheBlocks],
+        cached_prefix_len: int,
+    ) -> int:
+        """Return the minimum physical capacity for one public request.
+
+        Cached CFG rows may point to the same physical blocks. Count those hit
+        blocks once, then ask the native coordinator how many additional blocks
+        each row needs for its suffix. This gives an admission bound that still
+        works for warm requests and does not over-count shared prefix pages.
         """
 
         coordinator = self.native_manager.coordinator
-        empty_blocks = self.native_manager.empty_kv_cache_blocks.blocks
-        return sum(
-            coordinator.get_num_blocks_to_allocate(
+        unique_hit_blocks = len(self._unique_blocks(computed_blocks))
+        suffix_blocks = 0
+        for request, hit in zip(requests, computed_blocks, strict=True):
+            native_required = coordinator.get_num_blocks_to_allocate(
                 request_id=request.request_id,
                 num_tokens=request.num_tokens,
-                new_computed_blocks=empty_blocks,
+                new_computed_blocks=hit.blocks,
                 num_encoder_tokens=0,
-                total_computed_tokens=0,
-                num_local_computed_tokens=0,
+                total_computed_tokens=cached_prefix_len,
+                num_local_computed_tokens=cached_prefix_len,
                 num_tokens_main_model=request.num_tokens,
                 apply_admission_cap=True,
             )
-            for request in requests
-        )
+            evictable_hits = sum(block.ref_cnt == 0 and not block.is_null for group in hit.blocks for block in group)
+            suffix_blocks += max(native_required - evictable_hits, 0)
+        # ``get_num_blocks_to_allocate`` includes evictable hit pages because
+        # the live free-capacity check must first remove them from the free
+        # queue. They are already part of the pool's physical capacity, so do
+        # not count them a second time in this empty-pool bound.
+        return unique_hit_blocks + suffix_blocks
 
     def has_request(self, public_request_id: str) -> bool:
         return public_request_id in self._requests
@@ -123,7 +186,25 @@ class DiffusionKVCacheManager:
                     f"seq_len={request.seq_len}, max_model_len={self.max_model_len}"
                 )
 
-        required_blocks = self._get_empty_pool_required_blocks(requests)
+        self._prepare_block_hashes(requests)
+        hit_results = [self.native_manager.get_computed_blocks(request) for request in requests]
+        cached_prefix_len = min(result[1] for result in hit_results)
+        computed_blocks = [
+            self._clip_computed_blocks(
+                result[0],
+                hit_len=result[1],
+                cached_prefix_len=cached_prefix_len,
+            )
+            for result in hit_results
+        ]
+        for request, result in zip(requests, hit_results, strict=True):
+            request.shared_prefix_boundary = result[2]
+
+        required_blocks = self._get_empty_pool_required_blocks(
+            requests,
+            computed_blocks,
+            cached_prefix_len,
+        )
         if required_blocks > self._empty_pool_num_free_blocks:
             raise DiffusionKVAdmissionError(
                 f"Diffusion KV request {public_request_id!r} cannot fit even when the block pool is empty: "
@@ -133,11 +214,20 @@ class DiffusionKVCacheManager:
 
         allocated: list[DiffusionKVRequest] = []
         sequence_metadata: list[DiffusionKVSequenceMetadata] = []
+        pinned_blocks = self._unique_blocks(computed_blocks)
+        if pinned_blocks:
+            # All CFG lookups happen before allocation so they can share one
+            # execution boundary. Pin every distinct hit until each native
+            # allocation has taken its own reference; otherwise an earlier
+            # branch's miss allocation could evict a later branch's hit.
+            self.native_manager.block_pool.touch(pinned_blocks)
         try:
-            for request in requests:
+            for request, request_computed_blocks in zip(requests, computed_blocks, strict=True):
                 blocks = self.native_manager.allocate_slots(
                     request,
-                    num_new_tokens=request.seq_len,
+                    num_new_tokens=request.seq_len - cached_prefix_len,
+                    num_new_computed_tokens=cached_prefix_len,
+                    new_computed_blocks=request_computed_blocks,
                     delay_cache_blocks=True,
                     full_sequence_must_fit=True,
                 )
@@ -146,18 +236,23 @@ class DiffusionKVCacheManager:
                     allocated.clear()
                     return None
                 allocated.append(request)
+                request.num_computed_tokens = cached_prefix_len
                 sequence_metadata.append(
                     DiffusionKVSequenceMetadata(
                         sequence_id=request.sequence_id,
                         prefix_len=request.prefix_len,
                         target_len=request.target_len,
                         seq_len=request.seq_len,
-                        block_ids=blocks.get_block_ids(),
+                        block_ids=self.native_manager.get_block_ids(request.request_id),
+                        cached_prefix_len=cached_prefix_len,
                     )
                 )
         except Exception:
             self._rollback(allocated)
             raise
+        finally:
+            if pinned_blocks:
+                self.native_manager.block_pool.free_blocks(reversed(pinned_blocks))
 
         metadata = DiffusionKVMetadata(
             request_id=public_request_id,
@@ -170,6 +265,23 @@ class DiffusionKVCacheManager:
         self._internal_request_ids.update(internal_ids)
         return metadata
 
+    def publish_request(self, public_request_id: str) -> None:
+        """Publish successfully materialized, block-aligned stable prefixes."""
+
+        if not self.enable_prefix_caching:
+            return
+        try:
+            requests = self._requests[public_request_id]
+        except KeyError as exc:
+            raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc
+        for request in requests:
+            publishable_prefix_len = min(
+                request.prefix_len // self.hash_block_size * self.hash_block_size,
+                len(request.block_hashes) * self.hash_block_size,
+            )
+            if publishable_prefix_len > 0:
+                self.native_manager.cache_blocks(request, publishable_prefix_len)
+
     def get_metadata(self, public_request_id: str) -> DiffusionKVMetadata:
         try:
             return self._metadata[public_request_id]
@@ -181,6 +293,8 @@ class DiffusionKVCacheManager:
         self._metadata.pop(public_request_id, None)
         for request in reversed(requests):
             free_kv_blocks_in_physical_order(self.native_manager, request)
+            request.num_computed_tokens = 0
+            request.shared_prefix_boundary = 0
             self._internal_request_ids.discard(request.request_id)
 
     def close(self) -> None:
@@ -190,3 +304,5 @@ class DiffusionKVCacheManager:
     def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None:
         for request in reversed(requests):
             free_kv_blocks_in_physical_order(self.native_manager, request)
+            request.num_computed_tokens = 0
+            request.shared_prefix_boundary = 0