[Diffusion] Add native paged KV backend for diffusion workers (#6102)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/forward_context.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 import vllm.ir
@@ -26,6 +26,11 @@ class ForwardContext:
     vllm_config: VllmConfig | None = None
     omni_diffusion_config: OmniDiffusionConfig | None = None
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | 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
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
     denoise_timestep: float | None = None
@@ -209,6 +214,30 @@ def set_forward_context(
                 yield
 
 
+@contextmanager
+def override_paged_kv_adapter(adapter: Any | None):
+    """Temporarily expose a Worker paged-KV adapter to Omni Attention.
+
+    This is deliberately a small context override instead of a second global
+    forward context.  The model runner owns the outer context and the adapter
+    only replaces one opaque field while its prepared native metadata is live.
+    """
+
+    if _forward_context is None:
+        # Unit-level adapter users can prepare/activate metadata without an
+        # Omni model forward context.  In that case there is nothing to
+        # override and the adapter's explicit ``forward`` API remains usable.
+        yield
+        return
+
+    previous = _forward_context.paged_kv_adapter
+    _forward_context.paged_kv_adapter = adapter
+    try:
+        yield
+    finally:
+        _forward_context.paged_kv_adapter = previous
+
+
 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: