[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/forward_context.py

@@ -29,11 +29,19 @@ class ForwardContext:
     vllm_config: VllmConfig | None = None
     omni_diffusion_config: OmniDiffusionConfig | None = None
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
+    # Runner-owned paged execution metadata/runtime. Attention resolves the
+    # active Worker adapter from it; model code must not construct BlockTable
+    # rows or activate the runtime directly.
+    paged_kv_runtime: object | None = None
     # Active Worker-side paged KV adapter.  The adapter is installed only for
     # the duration of a paged forward; dense forwards leave this as ``None``.
     # Keep the field opaque here to avoid coupling the common context module to
     # the diffusion_kv implementation.
     paged_kv_adapter: Any | None = None
+    # Startup-only memory profiling runs before Scheduler-owned pages exist.
+    # Attention layers use this explicit scope to distinguish that probe from
+    # a malformed paged request whose Worker adapter was not activated.
+    in_diffusion_kv_memory_profile: bool = False
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
     denoise_timestep: float | None = None
@@ -154,6 +162,8 @@ def create_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    paged_kv_runtime: object | None = None,
+    in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
 ):
@@ -161,6 +171,8 @@ def create_forward_context(
         vllm_config=vllm_config,
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
+        paged_kv_runtime=paged_kv_runtime,
+        in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,
     )
@@ -186,6 +198,8 @@ def set_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    paged_kv_runtime: object | None = None,
+    in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
 ):
@@ -197,6 +211,8 @@ def set_forward_context(
         vllm_config=vllm_config,
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
+        paged_kv_runtime=paged_kv_runtime,
+        in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,
     )
@@ -246,6 +262,11 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
     """Set the current diffusion denoise step on the active ForwardContext."""
     if _forward_context is not None:
         _forward_context.denoise_step_idx = step_idx
+        if step_idx is not None:
+            paged_kv_runtime = getattr(_forward_context, "paged_kv_runtime", None)
+            ensure_active = getattr(paged_kv_runtime, "ensure_active", None)
+            if callable(ensure_active):
+                ensure_active(step_idx)
 
 
 def set_forward_context_denoise_timestep(timestep: float | None) -> None: