[Diffusion] Add paged KV cache support for HunyuanImage3 DiT on NPU and GPU (#6563)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <2831474076@qq.com>

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -1,10 +1,10 @@
 # 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
 
 import copy
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any
 
@@ -22,9 +22,9 @@ from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
     DiffusionPagedAttentionLayerAdapter,
-    DiffusionPagedAttentionRow,
+    DiffusionPagedAttentionMetadata,
     DiffusionPagedAttentionRowBinding,
-    PreparedDiffusionPagedAttentionBatch,
+    DiffusionPagedAttentionRuntime,
 )
 from vllm_omni.platforms import current_omni_platform
 
@@ -32,13 +32,6 @@ DiffusionKVIdentity = tuple[str, int | None, str | None]
 DiffusionKVSnapshot = tuple[object, ...]
 
 
-@dataclass(frozen=True)
-class _DiffusionKVRequestState:
-    generation: int
-    snapshot: DiffusionKVSnapshot
-    row_token_lens: tuple[tuple[DiffusionKVIdentity, int], ...]
-
-
 @dataclass(frozen=True)
 class _DiffusionKVRowInstall:
     identity: DiffusionKVIdentity
@@ -46,6 +39,13 @@ class _DiffusionKVRowInstall:
     block_ids: tuple[tuple[int, ...], ...]
 
 
+@dataclass(frozen=True)
+class _DiffusionKVRequestState:
+    generation: int
+    snapshot: DiffusionKVSnapshot
+    row_installs: tuple[_DiffusionKVRowInstall, ...]
+
+
 class DiffusionKVModelRunnerBackend:
     """Native paged-KV state and BlockTable operations for a model runner."""
 
@@ -503,7 +503,7 @@ class DiffusionKVModelRunnerBackend:
             self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
                 generation=metadata.allocation_generation,
                 snapshot=snapshot,
-                row_token_lens=current_state.row_token_lens,
+                row_installs=current_state.row_installs,
             )
             return False
 
@@ -530,7 +530,7 @@ class DiffusionKVModelRunnerBackend:
         self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
             generation=metadata.allocation_generation,
             snapshot=snapshot,
-            row_token_lens=tuple((install.identity, install.token_len) for install in installs),
+            row_installs=tuple(installs),
         )
         return True
 
@@ -559,24 +559,11 @@ class DiffusionKVModelRunnerBackend:
             raise RuntimeError("paged_scheduler native attention adapter is not initialized")
         return self.paged_attention_adapter
 
-    def prepare_paged_attention_batch(
-        self,
-        rows: Sequence[DiffusionPagedAttentionRow],
-    ) -> PreparedDiffusionPagedAttentionBatch:
-        """Build native page-table metadata for one Diffusion forward.
+    def activate_paged_attention_metadata(self, metadata: DiffusionPagedAttentionMetadata):
+        """Activate Runner-owned request metadata for an internal denoise loop."""
 
-        The scheduler allocation payload describes capacity, while the model
-        integration supplies the current ``query_len``/``kv_start_pos`` span
-        for each row.  Keeping this boundary explicit prevents the Worker from
-        guessing model-specific text/image layout.
-        """
-
-        return self.get_paged_attention_adapter().prepare_batch(rows)
-
-    def activate_paged_attention(self, batch: PreparedDiffusionPagedAttentionBatch):
-        """Return the context manager that exposes a prepared batch to Omni Attention."""
-
-        return self.get_paged_attention_adapter().activate(batch)
+        runtime = DiffusionPagedAttentionRuntime(self.get_paged_attention_adapter(), metadata)
+        return runtime, runtime.activate()
 
     def _resolve_paged_attention_row(
         self,
@@ -588,11 +575,12 @@ class DiffusionKVModelRunnerBackend:
         identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
         request_state = self._diffusion_kv_request_states.get(request_id)
         if request_state is not None:
-            for installed_identity, token_len in request_state.row_token_lens:
-                if installed_identity == identity:
+            for install in request_state.row_installs:
+                if install.identity == identity:
                     return DiffusionPagedAttentionRowBinding(
                         row_index=row_index,
-                        max_seq_len=token_len,
+                        max_seq_len=install.token_len,
+                        block_ids=install.block_ids,
                     )
         raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")