anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
This commit is contained in:
Yukim1 authored and GitHub committed 2026-09-23 15:34:56 +00:00
1 parent 9a1a05facd
commit 205aa01c7b
49 files changed
+2889 -127
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/data.py
+23 -1
View File
@@ -846,7 +846,7 @@ class OmniDiffusionConfig:
parallel_config: DiffusionParallelConfig = field(default_factory=DiffusionParallelConfig)
# Cache backend configuration (NEW)
cache_backend: str = "none" # "tea_cache", "deep_cache", etc.
cache_backend: str | None = "none" # "tea_cache", "deep_cache", etc.
cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
video_output_transport: VideoOutputTransportConfig = field(default_factory=VideoOutputTransportConfig)
enable_cache_dit_summary: bool = False
@@ -883,6 +883,9 @@ class OmniDiffusionConfig:
# Local Diffusion KV ownership and cache-layout mode.
diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
# Reuse block-aligned, immutable diffusion prefixes across requests when
# Scheduler-owned paged KV is active.
enable_prefix_caching: bool = False
# Maximum number of native BlockTable rows one public request can own
# (sequences plus independent contexts). The model adapter defines it.
diffusion_kv_max_rows_per_request: int | None = None
@@ -1204,6 +1207,19 @@ class OmniDiffusionConfig:
if not isinstance(self.diffusion_compile_dynamic, bool):
raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
if not isinstance(self.enable_prefix_caching, bool):
raise TypeError("enable_prefix_caching must be a bool")
if self.enable_prefix_caching and self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
raise ValueError(
"enable_prefix_caching=True requires diffusion_kv_mode='paged_scheduler'; "
"set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching"
)
if self.enable_prefix_caching and self.enable_sleep_mode:
raise ValueError(
"Diffusion prefix caching cannot be combined with sleep mode: "
"sleep discards KV pages without invalidating cached prefixes; "
"disable enable_prefix_caching or enable_sleep_mode"
)
if self.diffusion_kv_max_rows_per_request is not None and (
type(self.diffusion_kv_max_rows_per_request) is not int or self.diffusion_kv_max_rows_per_request <= 0
):
@@ -1247,6 +1263,11 @@ class OmniDiffusionConfig:
if self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
raise ValueError("native kv_transfer_config requires diffusion_kv_mode='paged_scheduler'")
self.kv_transfer_config = parse_kv_transfer_config(self.kv_transfer_config)
if self.enable_prefix_caching and self.kv_transfer_config is not None:
raise ValueError(
"Diffusion prefix caching cannot be combined with native kv_transfer_config; "
"disable enable_prefix_caching for AR KV import or remove kv_transfer_config for local DiT reuse"
)
if self.enable_sleep_mode:
raise ValueError("Native KV transfer does not support sleep mode: registered pages must remain mapped")
@@ -1722,6 +1743,7 @@ class OmniDiffusionConfig:
or DiffusionModelRegistry._try_load_model_cls(architecture) is not None
):
self.model_class_name = architecture
self.update_multimodal_support()
else:
raise
vllm_omni/diffusion/diffusion_engine.py
+25 -3
View File
@@ -54,6 +54,7 @@ from vllm_omni.diffusion.registry import (
DiffusionModelRegistry,
get_diffusion_post_process_func,
get_diffusion_pre_process_func,
get_diffusion_prefix_cache_func,
)
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
from vllm_omni.diffusion.sched import BaseScheduler, RequestScheduler, StepScheduler
@@ -130,7 +131,7 @@ __all__ = [
def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool:
if func is None:
if not callable(func):
return False
parameters = inspect.signature(func).parameters
return parameter_name in parameters or any(
@@ -300,6 +301,13 @@ class DiffusionEngine:
def _init_process_hooks(self, od_config: OmniDiffusionConfig) -> None:
self.post_process_func = get_diffusion_post_process_func(od_config)
self.pre_process_func = get_diffusion_pre_process_func(od_config)
self.prefix_cache_func = get_diffusion_prefix_cache_func(od_config) if self._prefix_cache_enabled() else None
if self._prefix_cache_enabled() and self.prefix_cache_func is None:
raise ValueError(
"enable_prefix_caching=True requires a registered prefix-cache hook for "
f"{od_config.model_class_name!r}; "
"disable enable_prefix_caching or use a supported native pipeline such as HunyuanImage3ForCausalMM"
)
# Cache whether the model-specific postprocess accepts request-level
# sampling params so step() can support both legacy and extended hooks.
self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params")
@@ -585,6 +593,7 @@ class DiffusionEngine:
)
def _busy_loop(self):
assert self.stop_event is not None
while not self.stop_event.is_set():
self._process_aborts_queue()
self._process_rpc_queue()
@@ -696,6 +705,7 @@ class DiffusionEngine:
Caller must hold ``self._cv``.
"""
assert self.stop_event is not None
start = time.monotonic()
decision = self.scheduler.get_admission_wait_decision(
now=start,
@@ -1045,6 +1055,12 @@ class DiffusionEngine:
engine.run_startup_warmup()
return engine
def _prefix_cache_enabled(self) -> bool:
config = getattr(self, "od_config", None)
return getattr(config, "diffusion_kv_mode", None) is DiffusionKVCacheMode.PAGED_SCHEDULER and bool(
getattr(config, "enable_prefix_caching", False)
)
def _prepare_request_for_admission(self, request: OmniDiffusionRequest) -> OmniDiffusionRequest:
"""Run model-owned preprocessing once, before entering Engine locks."""
@@ -1052,6 +1068,12 @@ class DiffusionEngine:
if pre_process_func is not None:
request = pre_process_func(request)
self._validate_diffusion_kv_profile_limits(request)
# Gate cache-input preparation itself: disabled caching must not inspect
# tensors / RNG state or copy token IDs just to discard their hashes.
# Both dependency hashing and preprocessing stay outside Engine locks.
prefix_cache_func = getattr(self, "prefix_cache_func", None)
if self._prefix_cache_enabled() and prefix_cache_func is not None and request.diffusion_kv_requests:
prefix_cache_func(request)
return request
def _validate_diffusion_kv_profile_limits(self, request: OmniDiffusionRequest) -> None:
@@ -1192,7 +1214,7 @@ class DiffusionEngine:
# sync func should receive one result
if (
sched_output.scheduled_request_ids
and not isinstance(runner_output, RunnerOutput)
and isinstance(runner_output, BatchRunnerOutput)
and len(runner_output) != 1
):
raise ValueError("Sync func should receive one result at one time")
@@ -1249,7 +1271,7 @@ class DiffusionEngine:
num_inference_steps: int = 1,
num_frames: int | None = None,
) -> OmniDiffusionRequest | None:
"""Build a minimal model request for startup profiling or warmup."""
"""Build a startup request; explicit frame counts bypass the warmup policy."""
prompt = OmniTextPrompt(prompt="dummy run")
model_class_name = self.od_config.model_class_name
if model_class_name is None:
vllm_omni/diffusion/diffusion_kv/kv_cache_utils.py
+94
View File
@@ -0,0 +1,94 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Small diffusion identity helpers alongside vLLM's native KV cache utilities.
Content hashing runs before Scheduler admission, only with prefix caching on.
Requests carry native MultiModalFeatureSpec / PlaceholderRange values; native
generate_block_hash_extra_keys handles block intersections, with no per-token
extra-key expansion or intermediate dependency objects.
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING
import numpy as np
import torch
from vllm.multimodal.hasher import MultiModalHasher
if TYPE_CHECKING:
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
def _normalize_identity(value: object) -> object:
# MultiModalHasher handles tensor bytes (including BF16). Its generic
# recursive serialization alone does not canonicalize nested dict order
# or distinguish every container/scalar boundary. A typed JSON envelope
# supplies those guarantees and rejects the upstream pickle fallback.
if value is None:
return ["none"]
if isinstance(value, np.generic):
scalar = value.item()
if isinstance(scalar, np.generic):
# Extended-precision scalars have no lossless Python scalar;
# item() returns another NumPy scalar instead of making progress.
raise TypeError(f"Unsupported prefix-cache identity scalar dtype: {value.dtype}")
return _normalize_identity(scalar)
if isinstance(value, bool):
return ["bool", value]
if isinstance(value, int):
return ["int", str(value)]
if isinstance(value, float):
return ["float", value.hex()]
if isinstance(value, str):
return ["str", value]
if isinstance(value, bytes):
return ["bytes", value.hex()]
if isinstance(value, torch.Tensor):
if value.layout != torch.strided or value.is_quantized:
raise TypeError("Prefix-cache identity requires a dense, non-quantized tensor")
tensor = value.detach().cpu().contiguous()
# A byte view also covers dtypes NumPy cannot represent directly.
data = tensor.reshape(-1).view(torch.uint8).numpy()
return ["tensor", str(tensor.dtype), list(tensor.shape), MultiModalHasher.hash_kwargs("sha256", data=data)]
if isinstance(value, np.ndarray):
if value.dtype.hasobject or value.dtype.fields is not None:
raise TypeError("Prefix-cache identity does not support object or structured arrays")
return [
"array",
value.dtype.str,
list(value.shape),
MultiModalHasher.hash_kwargs("sha256", data=np.ascontiguousarray(value)),
]
if isinstance(value, dict):
if any(not isinstance(key, str) for key in value):
raise TypeError("Prefix-cache identity dictionaries require string keys")
return ["dict", [[key, _normalize_identity(value[key])] for key in sorted(value)]]
if isinstance(value, list | tuple):
return ["sequence", [_normalize_identity(item) for item in value]]
raise TypeError(f"Unsupported prefix-cache identity value: {type(value).__name__}")
def hash_prefix_cache_value(value: object) -> bytes:
"""Deterministic, type-framed identity using vLLM's multimodal hasher."""
encoded = json.dumps(_normalize_identity(value), ensure_ascii=True, separators=(",", ":"))
return bytes.fromhex(MultiModalHasher.hash_kwargs("sha256", identity=encoded))
def get_cache_namespace(model_namespace: str, sampling: OmniDiffusionSamplingParams) -> str:
"""Isolate model semantics and the adapter actually activated by Worker.
Native LoRA block keys use lora_name, but diffusion also supports a scale
and identifies loaded adapters by lora_int_id. Keep that difference here.
Seeds are NOT global identity inputs: models attach random state only to
the multimodal features whose computation depends on it.
"""
if not isinstance(model_namespace, str) or not model_namespace:
raise ValueError("Model cache namespace must be non-empty")
lora = sampling.lora_request
adapter = None if lora is None else (lora.lora_int_id, float(sampling.lora_scale))
return hash_prefix_cache_value(("diffusion-prefix-identity-v2", model_namespace, adapter)).hex()
vllm_omni/diffusion/diffusion_kv/manager.py
+138 -22
View File
@@ -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
vllm_omni/diffusion/diffusion_kv/metadata.py
+2 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Scheduler-to-Worker allocation payloads for diffusion KV.
@@ -40,6 +40,7 @@ class DiffusionKVSequenceMetadata:
target_len: int
seq_len: int
block_ids: tuple[list[int], ...]
cached_prefix_len: int = 0
context_ids: tuple[str, ...] = ()
num_computed_tokens: int = 0
vllm_omni/diffusion/diffusion_kv/model_runner_backend.py
+13 -2
View File
@@ -4,7 +4,7 @@
from __future__ import annotations
import copy
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any
@@ -154,6 +154,7 @@ class DiffusionKVModelRunnerBackend:
if attention_geometry is not None:
num_heads, num_kv_heads, head_size = attention_geometry
assert callable(set_attention_geometry)
try:
set_attention_geometry(
num_heads=num_heads,
@@ -396,6 +397,15 @@ class DiffusionKVModelRunnerBackend:
for sequence in metadata.sequences:
if type(sequence.sequence_id) is not int or sequence.sequence_id < 0:
raise ValueError(f"Diffusion KV sequence_id must be a non-negative integer: {sequence.sequence_id!r}")
if (
type(sequence.cached_prefix_len) is not int
or sequence.cached_prefix_len < 0
or sequence.cached_prefix_len > sequence.prefix_len
):
raise ValueError(
"Diffusion KV cached_prefix_len must lie within the stable prefix: "
f"cached={sequence.cached_prefix_len!r}, prefix={sequence.prefix_len!r}"
)
sequence_identity = (metadata.request_id, sequence.sequence_id, None)
sequence_install = self._validate_row(
identity=sequence_identity,
@@ -410,6 +420,7 @@ class DiffusionKVModelRunnerBackend:
sequence.prefix_len,
sequence.target_len,
sequence.seq_len,
sequence.cached_prefix_len,
sequence_install.block_ids,
sequence.context_ids,
)
@@ -593,7 +604,7 @@ class DiffusionKVModelRunnerBackend:
)
raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> int:
def remove_diffusion_kv_requests(self, request_ids: Sequence[str | tuple[str, int]]) -> int:
"""Retire Worker rows without logically freeing Scheduler-owned blocks."""
if (
not is_scheduler_paged_kv_mode(
vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py
+7 -3
View File
@@ -84,14 +84,18 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
attention_config = vllm_config.attention_config
previous_backend = attention_config.backend
previous_backend_per_kind = attention_config.backend_per_kind
previous_pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
parallel_config = vllm_config.parallel_config
previous_pcp_size = parallel_config.prefill_context_parallel_size
try:
# This is a portable vLLM backend request. The active platform
# resolves it to its native implementation, such as FlashAttention
# on CUDA or AscendAttentionBackend on NPU.
attention_config.backend = AttentionBackendEnum.FLASH_ATTN
attention_config.backend_per_kind = {}
vllm_config.parallel_config.prefill_context_parallel_size = 1
# Diffusion maps SP to PCP for MoE, not for attention. Omni's
# Ulysses hooks gather the sequence and shard heads before this
# native kernel, so backend selection must not request PCP.
parallel_config.prefill_context_parallel_size = 1
with set_current_vllm_config(vllm_config):
attn_backend = get_attn_backend(
head_size=spec.head_size,
@@ -103,7 +107,7 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
finally:
attention_config.backend = previous_backend
attention_config.backend_per_kind = previous_backend_per_kind
vllm_config.parallel_config.prefill_context_parallel_size = previous_pcp_size
parallel_config.prefill_context_parallel_size = previous_pcp_size
attn_backend = current_omni_platform.get_diffusion_paged_kv_attn_backend(
attn_backend,
ulysses_degree=ulysses_degree,
vllm_omni/diffusion/diffusion_kv/request.py
+93 -5
View File
@@ -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."""
vllm_omni/diffusion/executor/abstract.py
+7 -2
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import Callable
from concurrent.futures import Future
from dataclasses import replace
from typing import TYPE_CHECKING, Any
@@ -12,7 +13,7 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
from vllm.v1.outputs import KVConnectorOutput
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.sched.interface import CachedRequestData
if TYPE_CHECKING:
@@ -132,6 +133,10 @@ class DiffusionExecutor(ABC):
"""
return None
def wait_output_ready(self, async_output_id: str) -> Future[DiffusionOutput]:
"""Resolve deferred output; only asynchronous executors implement this."""
raise NotImplementedError(f"{type(self).__name__} does not support asynchronous output")
def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
"""Collect rank-local native specs after every Worker loads its model."""
@@ -202,7 +207,7 @@ class DiffusionExecutor(ABC):
and not scheduler_output.kv_required_request_ids
and not scheduler_output.kv_poll_only
):
return
return None
transfer_output = replace(
scheduler_output,
scheduled_new_reqs=[],
vllm_omni/diffusion/forward_context.py
+12 -3
View File
@@ -19,6 +19,8 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
if TYPE_CHECKING:
import torch
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import DiffusionPagedAttentionRuntime
@dataclass
class ForwardContext:
@@ -32,7 +34,10 @@ class ForwardContext:
# 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
paged_kv_runtime: DiffusionPagedAttentionRuntime | None = None
# Block-aligned prefix already resident in Scheduler-owned pages for the
# active request-level prefill. Zero keeps the cold/full-prefill path.
paged_kv_cached_prefix_len: int = 0
# 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
@@ -177,7 +182,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,
paged_kv_runtime: DiffusionPagedAttentionRuntime | None = None,
paged_kv_cached_prefix_len: int = 0,
in_diffusion_kv_memory_profile: bool = False,
split_text_embed_in_sp: bool = False,
denoise_step_idx: int | None = None,
@@ -187,6 +193,7 @@ def create_forward_context(
omni_diffusion_config=omni_diffusion_config,
attn_metadata=attn_metadata,
paged_kv_runtime=paged_kv_runtime,
paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
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,
@@ -213,7 +220,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,
paged_kv_runtime: DiffusionPagedAttentionRuntime | None = None,
paged_kv_cached_prefix_len: int = 0,
in_diffusion_kv_memory_profile: bool = False,
split_text_embed_in_sp: bool = False,
denoise_step_idx: int | None = None,
@@ -227,6 +235,7 @@ def set_forward_context(
omni_diffusion_config=omni_diffusion_config,
attn_metadata=attn_metadata,
paged_kv_runtime=paged_kv_runtime,
paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
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,
vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+10 -1
View File
@@ -76,7 +76,9 @@ from vllm_omni.diffusion.distributed.sp_plan import (
)
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import (
get_forward_context,
get_paged_kv_computed_tokens,
is_forward_context_available,
paged_kv_prefill,
set_forward_context_denoise_step_idx,
)
@@ -2481,6 +2483,7 @@ class HunyuanImage3Model(nn.Module):
gen_timestep_scatter_index: torch.Tensor | None = None,
uncond_cfg_prefill: bool = False,
ar_kv_reuse_len: int = 0,
paged_kv_cached_prefix_len: int = 0,
full_attn_spans: list[list[tuple[int, int]]] | None = None,
) -> tuple | BaseModelOutputWithPast:
current_omni_platform.reset_diffusion_fused_moe_forward_context()
@@ -2539,7 +2542,9 @@ class HunyuanImage3Model(nn.Module):
else:
shard_padding_size = shard_image_size * sp_world_size - num_image_tokens
if first_step:
seq_lens = [prompt_size + shard_image_size + ar_kv_reuse_len for _ in seq_lens]
seq_lens = [
prompt_size + shard_image_size + ar_kv_reuse_len + paged_kv_cached_prefix_len for _ in seq_lens
]
else:
seq_lens = [x - y for x, y in zip(seq_lens, query_lens)]
seq_lens = [seq_len + shard_image_size for seq_len in seq_lens]
@@ -3023,6 +3028,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
cfg_rank,
device,
):
if is_forward_context_available() and get_forward_context().paged_kv_cached_prefix_len:
# Local hits are sliced after conditional-image embeddings are
# prepared in forward_call; the AR path must not truncate them.
return input_ids, 0
ar_kv_data = model_kwargs.pop("ar_kv_data", None)
computed_tokens = get_paged_kv_computed_tokens()
if computed_tokens and computed_tokens[0] > 0:
vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+114 -1
View File
@@ -23,7 +23,11 @@ from vllm.transformers_utils.config import get_config
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.forward_context import (
get_forward_context,
is_forward_context_available,
set_forward_context_denoise_step_idx,
)
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
@@ -105,6 +109,10 @@ def to_device(data, device):
return data
def get_hunyuan_image_3_prefix_cache_func(od_config: OmniDiffusionConfig):
return request_layout_utils.prepare_hunyuan_prefix_cache
def _to_pil_image(image: Any) -> PILImage.Image:
if isinstance(image, PILImage.Image):
return image
@@ -1408,6 +1416,13 @@ class HunyuanImage3Pipeline(
)
full_attn_spans: list[list[tuple[int, int]]] = [[] for _ in range(bsz)]
for i in range(bsz):
for reference_slice in tokenizer_output.joint_image_slices[i]:
logger.debug(
"Hunyuan reference span: sequence_id=%d start=%d end=%d",
i,
int(reference_slice.start or 0),
int(reference_slice.stop or seq_len),
)
for j, image_slice in enumerate(batch_image_slices[i]):
attention_mask[i, image_slice, image_slice] = True
start = image_slice.start if image_slice.start is not None else 0
@@ -1420,6 +1435,69 @@ class HunyuanImage3Pipeline(
model_kwargs["full_attn_spans"] = full_attn_spans
return attention_mask
@staticmethod
def _slice_cached_prefix_inputs(
inputs_embeds: torch.Tensor,
attention_mask: torch.Tensor | None,
position_ids: torch.Tensor | None,
custom_pos_emb: tuple[torch.Tensor, ...] | None,
image_mask: torch.Tensor | None,
gen_timestep_scatter_index: torch.Tensor | None,
query_lens: list[int] | None,
seq_lens: list[int] | None,
cached_prefix_len: int,
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
tuple[torch.Tensor, ...],
torch.Tensor,
torch.Tensor,
list[int],
int,
]:
"""Remove a cached first-step prefix while preserving full KV lengths."""
seq_len = int(inputs_embeds.shape[1])
if (
query_lens is None
or seq_lens is None
or len(query_lens) != inputs_embeds.shape[0]
or len(seq_lens) != len(query_lens)
):
raise ValueError("Hunyuan paged prefix execution requires one query/sequence length per row")
if any(query_len != seq_len for query_len in query_lens) or any(
logical_len != seq_len for logical_len in seq_lens
):
raise ValueError(
"Hunyuan paged prefix slicing requires an untrimmed first-step tensor: "
f"query_lens={query_lens}, tensor_len={seq_len}"
)
if not 0 < cached_prefix_len < seq_len:
raise ValueError(
"Hunyuan cached prefix must leave a non-empty first-step suffix: "
f"cached={cached_prefix_len}, seq_len={seq_len}"
)
if (
attention_mask is None
or position_ids is None
or custom_pos_emb is None
or image_mask is None
or gen_timestep_scatter_index is None
):
raise ValueError("Hunyuan paged prefix slicing is missing first-step layout tensors")
return (
inputs_embeds[:, cached_prefix_len:],
attention_mask[:, :, cached_prefix_len:, :],
position_ids[:, cached_prefix_len:],
tuple(pos[:, cached_prefix_len:] for pos in custom_pos_emb),
image_mask[:, cached_prefix_len:],
gen_timestep_scatter_index - cached_prefix_len,
[query_len - cached_prefix_len for query_len in query_lens],
seq_len - cached_prefix_len,
)
def prepare_inputs_for_generation(
self,
input_ids,
@@ -1747,6 +1825,40 @@ class HunyuanImage3Pipeline(
assert inputs_embeds is not None
bsz, seq_len, n_embd = inputs_embeds.shape
paged_kv_cached_prefix_len = 0
if first_step and self._uses_scheduler_paged_kv():
if not is_forward_context_available():
raise RuntimeError("Hunyuan paged prefix execution requires an Omni ForwardContext")
paged_kv_cached_prefix_len = int(get_forward_context().paged_kv_cached_prefix_len)
if paged_kv_cached_prefix_len:
if ar_kv_reuse_len:
raise ValueError("Hunyuan cannot combine imported AR KV and Scheduler prefix-cache hits")
(
inputs_embeds,
attention_mask,
position_ids,
custom_pos_emb,
image_mask,
gen_timestep_scatter_index,
query_lens,
seq_len,
) = self._slice_cached_prefix_inputs(
inputs_embeds,
attention_mask,
position_ids,
custom_pos_emb,
image_mask,
gen_timestep_scatter_index,
query_lens,
seq_lens,
paged_kv_cached_prefix_len,
)
logger.debug(
"Hunyuan prefix slice: cached_prefix_len=%d query_len=%d",
paged_kv_cached_prefix_len,
seq_len,
)
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
from vllm.forward_context import set_forward_context
@@ -1770,6 +1882,7 @@ class HunyuanImage3Pipeline(
gen_timestep_scatter_index=gen_timestep_scatter_index,
uncond_cfg_prefill=uncond_cfg_prefill,
ar_kv_reuse_len=ar_kv_reuse_len,
paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
full_attn_spans=full_attn_spans,
)
hidden_states = outputs[0]
vllm_omni/diffusion/models/hunyuan_image3/request_layout.py
+99 -2
View File
@@ -8,7 +8,9 @@ from typing import TYPE_CHECKING, Any
import numpy as np
import torch
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -357,7 +359,7 @@ def build_hunyuan_diffusion_kv_requests(
request: OmniDiffusionRequest,
prepared_layout: HunyuanPreparedLayout,
) -> tuple[DiffusionKVRequest, ...]:
"""Build one persistent Scheduler KV request per Hunyuan execution row."""
"""Build allocation-only KV requests, even when prefix caching is disabled."""
tokenizer_output = prepared_layout.tokenizer_output
cfg_factor = hunyuan_cfg_factor(
@@ -408,7 +410,15 @@ def build_hunyuan_diffusion_kv_requests(
prefix_len=int(prefix_row[-1].item()),
target_len=target_len,
seq_len=int(valid_row[-1].item()),
prompt_token_ids=tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist(),
# Native AR -> DiT transfer needs token IDs to describe the
# transferred prefix. Local paged prefix caching does not: its
# identity is attached later by ``prepare_hunyuan_prefix_cache``.
# Keep this conversion out of the disabled/local-only path.
prompt_token_ids=(
tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist()
if request.kv_transfer_params is not None
else None
),
# Prompt and reference-image tokens are already embedded in this
# row's primary self-attention sequence. Hunyuan therefore has no
# independently projected cross/joint-attention KV context.
@@ -418,6 +428,93 @@ def build_hunyuan_diffusion_kv_requests(
)
def prepare_hunyuan_prefix_cache(request: OmniDiffusionRequest) -> None:
"""Attach native MM identities / positions when Engine enables caching.
This is model input adaptation, not a hashing framework. All reference
inputs are hashed once, then reused by VAE/ViT subspans and CFG rows. The
native block hasher consumes these ranges directly, without token extras.
"""
layout = get_hunyuan_prepared_layout(request)
if layout is None or not request.diffusion_kv_requests:
raise ValueError("Hunyuan prefix caching requires prepared layout and KV requests")
sampling = request.sampling_params
_, _, _, images, _ = extract_hunyuan_prompt_inputs(
[request.prompt], sampling.extra_args or {}, request_id=request.request_id, allow_cond_image=True
)
reference_digest = None
if images:
# Match Worker: an explicitly supplied generator overrides the seed.
# Conditional VAE samples latents, so image bytes alone are not enough.
if sampling.generator is not None:
generators = sampling.generator if isinstance(sampling.generator, list) else [sampling.generator]
random_state: object = (
"generator-state",
tuple((str(generator.device), generator.get_state()) for generator in generators),
)
elif sampling.seed is not None:
random_state = ("seed", int(sampling.seed), str(sampling.generator_device or "worker-default"))
else:
random_state = ("request-local-random-state", request.request_id)
# Preserve the conservative whole-reference-set policy. Splitting
# independent image identities also needs the VAE RNG sequence modeled.
reference_digest = hash_prefix_cache_value(
(
"hunyuan-reference-v1",
[[joint_image_info_to_payload(image) for image in row] for row in images],
random_state,
)
)
namespace = get_cache_namespace("hunyuan-image3-primary-v4", sampling)
output = layout.tokenizer_output
joint_rows = output.joint_image_slices or [[] for _ in range(layout.num_branches)]
prepared = []
for row in request.diffusion_kv_requests:
if row.block_hashes or row.num_computed_tokens:
raise ValueError("Hunyuan cache inputs must be prepared before KV execution")
joint_spans = [(int(span.start or 0), int(span.stop or 0)) for span in joint_rows[row.sequence_id]]
features = []
for image_slice, (height, width) in layout.rope_image_info[row.sequence_id]:
start = int(image_slice.start or 0)
end = int(image_slice.stop or output.tokens.shape[1])
if start >= row.prefix_len or end <= 0:
continue
# Joint spans enclose VAE + separator + ViT; RoPE lists subspans.
is_reference = any(joint_start <= start and end <= joint_end for joint_start, joint_end in joint_spans)
if is_reference and reference_digest is None:
raise ValueError("Hunyuan reference-image prefix is missing its canonical inputs")
identifier = hash_prefix_cache_value(
(
"hunyuan-image-span-v2",
start,
end,
int(height),
int(width),
is_reference,
reference_digest if is_reference else None,
)
).hex()
features.append(
MultiModalFeatureSpec(
data=None,
modality="image",
identifier=identifier,
mm_position=PlaceholderRange(offset=start, length=end - start),
)
)
# Raw model inputs remain in request.prompt/prepared_layout; data=None
# here does not imply an encoder hit. These are Scheduler-only KV keys.
tokens = tuple(output.tokens[row.sequence_id, : row.prefix_len].tolist())
prepared.append((row, tokens, sorted(features, key=lambda feature: feature.mm_position.offset)))
for row, tokens, features in prepared:
row.cache_token_ids = tokens
row.mm_features = features
row.cache_namespace = namespace
def get_hunyuan_prepared_layout(source: Any) -> HunyuanPreparedLayout | None:
prepared_layout = getattr(source, "prepared_layout", None)
if prepared_layout is None:
vllm_omni/diffusion/registry.py
+19
View File
@@ -749,6 +749,11 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
}
_DIFFUSION_PREFIX_CACHE_FUNCS = {
"HunyuanImage3ForCausalMM": "get_hunyuan_image_3_prefix_cache_func",
}
def register_diffusion_model(
model_arch: str,
module_name: str,
@@ -757,6 +762,7 @@ def register_diffusion_model(
post_process_func_name: str | None = None,
ir_op_priority_func_name: str | None = None,
action_post_process_func_name: str | None = None,
prefix_cache_func_name: str | None = None,
) -> None:
"""Register a diffusion model pipeline from an out-of-tree plugin.
@@ -782,6 +788,9 @@ def register_diffusion_model(
for out-of-tree plugins. Action postprocess hooks are no longer
registered separately; move action handling into
``post_process_func_name`` and return a payload/metadata envelope.
prefix_cache_func_name: Optional factory for a CPU cache-input hook,
called after preprocessing only when paged prefix caching is
enabled. It fills existing DiffusionKVRequest cache inputs in place.
"""
if action_post_process_func_name is not None:
logger.warning(
@@ -810,6 +819,8 @@ def register_diffusion_model(
_DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name
if ir_op_priority_func_name is not None:
_DIFFUSION_IR_OP_PRIORITY_FUNCS[model_arch] = ir_op_priority_func_name
if prefix_cache_func_name is not None:
_DIFFUSION_PREFIX_CACHE_FUNCS[model_arch] = prefix_cache_func_name
logger.info(
"Registered diffusion model %s -> %s.%s",
@@ -860,3 +871,11 @@ def get_diffusion_pre_process_func(od_config: OmniDiffusionConfig):
return None # Return None if no pre-processing function is registered (for backward compatibility)
func_name = _DIFFUSION_PRE_PROCESS_FUNCS[od_config.model_class_name]
return _load_process_func(od_config, func_name)
def get_diffusion_prefix_cache_func(od_config: OmniDiffusionConfig):
"""Load optional model preparation for native multimodal KV identities."""
if uses_diffusers_adapter(od_config) or od_config.model_class_name is None:
return None
func_name = _DIFFUSION_PREFIX_CACHE_FUNCS.get(od_config.model_class_name)
return None if func_name is None else _load_process_func(od_config, func_name)
vllm_omni/diffusion/sched/base_scheduler.py
+20 -3
View File
@@ -38,7 +38,7 @@ from vllm_omni.diffusion.sched.interface import (
StepBatchSamplingParamsKey,
_AdmissionWaitDecision,
)
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
logger = init_logger(__name__)
@@ -120,6 +120,12 @@ class BaseScheduler(ABC):
scheduler_block_size=scheduler_block_size,
hash_block_size=hash_block_size,
max_in_flight_tokens=kv_vllm_config.max_in_flight_tokens,
enable_prefix_caching=bool(getattr(kv_vllm_config.cache_config, "enable_prefix_caching", False)),
prefix_caching_hash_algo=getattr(
kv_vllm_config.cache_config,
"prefix_caching_hash_algo",
"sha256",
),
)
else:
if any(
@@ -182,7 +188,8 @@ class BaseScheduler(ABC):
diffusion_kv_metadata: DiffusionKVMetadata | None = None
if self._diffusion_kv_manager is not None:
if self._diffusion_kv_manager.has_request(request_id):
already_reserved = self._diffusion_kv_manager.has_request(request_id)
if already_reserved:
diffusion_kv_metadata = self._diffusion_kv_manager.get_metadata(request_id)
else:
matched_tokens: list[int] = []
@@ -210,6 +217,11 @@ class BaseScheduler(ABC):
continue
if allocation is None:
break
# Check the lookup result before registering destination
# pages with the connector: deferred pages can be freed.
if not self._can_schedule_waiting(state):
self._diffusion_kv_manager.free_request(request_id)
break
diffusion_kv_metadata = allocation
self._kv_request_generations[request_id] = allocation.allocation_generation
if self._kv_connector is not None:
@@ -397,7 +409,7 @@ class BaseScheduler(ABC):
)
@abstractmethod
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]:
pass
def has_requests(self) -> bool:
@@ -633,6 +645,11 @@ class BaseScheduler(ABC):
# Also surface admission failures recorded while schedule() built this
# output. Older finished ids retained only for Worker cleanup have
# already been popped by the Engine and are deliberately ignored.
if self._diffusion_kv_manager is not None:
for request_id, status in statuses.items():
if status == DiffusionRequestStatus.FINISHED_COMPLETED:
self._diffusion_kv_manager.publish_request(request_id)
finished_req_ids = {
request_id for request_id in sched_output.finished_req_ids if request_id in self._request_states
}
vllm_omni/diffusion/sched/request_scheduler.py
+16 -3
View File
@@ -12,11 +12,12 @@ from vllm_omni.diffusion.sched.interface import (
DiffusionRequestStatus,
DiffusionSchedulerOutput,
RequestBatchSamplingParamsKey,
SchedulerRequestState,
_AdmissionWaitDecision,
)
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
# LoRA identity is derived from `sampling.lora_request`, not a same-named field
# on sampling params, so it must be resolved separately from the bulk lookup.
@@ -32,7 +33,7 @@ def _normalize_explicit_sample_solver(value: object | None) -> str | None:
return str(value).strip().lower()
def _normalize_explicit_flow_shift(value: object | None) -> float | None:
def _normalize_explicit_flow_shift(value: float | str | None) -> float | None:
"""Normalize an explicitly provided flow shift without selecting a default."""
if value is None:
return None
@@ -59,6 +60,18 @@ def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> Re
class RequestScheduler(BaseScheduler):
"""Scheduler for static request waves, including admission coalescing."""
def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool:
if not super()._can_schedule_waiting(state):
return False
manager = self._diffusion_kv_manager
if not self._running or manager is None or not manager.has_request(state.request_id):
return True
# The request-level runner uses one first-step query boundary per wave.
# The Manager already aligns all CFG sequences within each request.
current = manager.get_metadata(self._running[0]).sequences[0].cached_prefix_len
candidate = manager.get_metadata(state.request_id).sequences[0].cached_prefix_len
return candidate == current
def get_admission_wait_decision(
self,
*,
@@ -98,7 +111,7 @@ class RequestScheduler(BaseScheduler):
def _build_sampling_params_key(self, request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
return build_request_batch_sampling_params_key(request)
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]:
scheduled_request_ids = sched_output.scheduled_request_ids
if not scheduled_request_ids and not sched_output.finished_req_ids:
return set()
vllm_omni/diffusion/sched/step_scheduler.py
+3 -2
View File
@@ -16,7 +16,7 @@ from vllm_omni.diffusion.sched.interface import (
)
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
logger = init_logger(__name__)
@@ -62,7 +62,7 @@ class StepScheduler(BaseScheduler):
)
return request_id
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]:
scheduled_request_ids = sched_output.scheduled_request_ids
if not scheduled_request_ids and not sched_output.finished_req_ids:
return set()
@@ -125,6 +125,7 @@ class StepScheduler(BaseScheduler):
return self._sequence_length(sampling.timesteps)
if sampling.sigmas is not None:
return len(sampling.sigmas)
assert sampling.num_inference_steps is not None
return int(sampling.num_inference_steps)
@staticmethod
vllm_omni/diffusion/vllm_config.py
+3 -2
View File
@@ -152,6 +152,7 @@ class _DiffusionVllmModelConfig:
return self._require_attention_geometry()[2]
def get_total_num_kv_heads(self) -> int:
assert self.hf_text_config is not None
return self.hf_text_config.num_key_value_heads
@property
@@ -206,6 +207,7 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
"""Populate native model, parallel, cache, and Scheduler configuration."""
parallel_config = od_config.parallel_config
assert parallel_config.data_parallel_size is not None
vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size
vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size
if parallel_config.enable_expert_parallel and od_config.is_moe:
@@ -225,8 +227,7 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
vllm_config.parallel_config.pipeline_parallel_size = parallel_config.pipeline_parallel_size
vllm_config.cache_config.gpu_memory_utilization = float(getattr(od_config, "gpu_memory_utilization", 0.9))
vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None)
# Prefix hashes/publication are intentionally deferred to the next PR.
vllm_config.cache_config.enable_prefix_caching = False
vllm_config.cache_config.enable_prefix_caching = bool(getattr(od_config, "enable_prefix_caching", False))
max_rows_per_request = getattr(od_config, "diffusion_kv_max_rows_per_request", None)
if type(max_rows_per_request) is not int or max_rows_per_request <= 0:
raise ValueError("paged_scheduler requires a positive diffusion_kv_max_rows_per_request")
vllm_omni/diffusion/worker/diffusion_model_runner.py
+51 -9
View File
@@ -13,7 +13,7 @@ from __future__ import annotations
import copy
import gc
import time
from collections.abc import Callable
from collections.abc import Callable, Sequence
from contextlib import AbstractContextManager, nullcontext
from typing import TYPE_CHECKING, Any, cast
@@ -97,9 +97,10 @@ def _dit_any_rank_failed(local_failed: bool) -> bool:
if not torch.distributed.is_initialized():
return local_failed
try:
from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
from vllm_omni.diffusion.distributed import parallel_state
group = get_dit_group()
get_dit_group = getattr(parallel_state, "get_dit_group", None)
group = get_dit_group() if get_dit_group is not None else None
except (AssertionError, ImportError):
group = None
if group is None:
@@ -489,6 +490,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput):
from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load
assert self.od_config.kv_transfer_config is not None
timeout = self.od_config.kv_transfer_config.kv_connector_extra_config.get("transfer_timeout", 60.0)
if self._kv_receive_progress is not None:
return self._kv_receive_progress.prepare(self._kv_connector, scheduler_output, timeout)
@@ -505,7 +507,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
) -> int:
return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id)
def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> int:
def remove_diffusion_kv_requests(self, request_ids: Sequence[str | tuple[str, int]]) -> int:
return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids)
def refresh_diffusion_kv_block_table_layout(self) -> None:
@@ -538,13 +540,19 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
f"request={request_metadata.request_id!r}, sequence={sequence.sequence_id}, "
f"active={active_seq_len}, allocated={sequence.seq_len}"
)
# Imported AR KV and local hits have separate owners.
kv_start_pos = (
sequence.num_computed_tokens
if getattr(self.od_config, "kv_transfer_config", None) is not None
else sequence.cached_prefix_len
)
prefill_rows.append(
DiffusionPagedAttentionRow(
request_id=request_metadata.request_id,
sequence_id=sequence.sequence_id,
query_len=sequence.seq_len - sequence.num_computed_tokens,
kv_start_pos=kv_start_pos,
query_len=sequence.seq_len - kv_start_pos,
seq_len=sequence.seq_len,
kv_start_pos=sequence.num_computed_tokens,
)
)
denoise_rows.append(
@@ -800,17 +808,48 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
"Diffusion KV metadata count must match the request batch: "
f"metadata={len(diffusion_kv_metadata)}, requests={len(reqs)}"
)
for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True):
req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences)
native_kv_transfer = getattr(self.od_config, "kv_transfer_config", None) is not None
if native_kv_transfer:
for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True):
# This field is consumed by Hunyuan only for native
# AR->DiT transfer. Local prefix hits must still run
# their VAE/ViT conditioning path.
req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences)
paged_metadata = self._build_paged_attention_metadata(diffusion_kv_metadata)
paged_kv_cached_prefix_len = 0
if not native_kv_transfer:
cached_prefix_lens = {row.kv_start_pos for row in paged_metadata.prefill_rows}
if len(cached_prefix_lens) != 1:
raise ValueError(
"One paged request-level forward requires a uniform cached prefix boundary; "
f"got {sorted(cached_prefix_lens)}"
)
paged_kv_cached_prefix_len = next(iter(cached_prefix_lens))
paged_kv_runtime, paged_kv_context = self.diffusion_kv_backend.activate_paged_attention_metadata(
paged_metadata
)
if is_primary:
# Trace the boundary actually passed to the model, not a
# speculative lookup. Useful for warm-cache regressions.
for request_metadata in diffusion_kv_metadata:
for sequence in request_metadata.sequences:
logger.debug(
"Diffusion prefix prefill: request_id=%s sequence_id=%d "
"cached_prefix_len=%d prefix_len=%d query_len=%d",
request_metadata.request_id,
sequence.sequence_id,
sequence.cached_prefix_len,
sequence.prefix_len,
sequence.seq_len - sequence.cached_prefix_len,
)
else:
paged_kv_cached_prefix_len = 0
with (
set_forward_context(
vllm_config=self.vllm_config,
omni_diffusion_config=od_config,
paged_kv_runtime=paged_kv_runtime,
paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
),
paged_kv_context,
@@ -1054,7 +1093,10 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
prepared_layout=getattr(sched_new_req.req, "prepared_layout", None),
external_req_id=getattr(sched_new_req.req, "external_req_id", None),
)
if sched_new_req.diffusion_kv_metadata is not None:
if (
sched_new_req.diffusion_kv_metadata is not None
and getattr(self.od_config, "kv_transfer_config", None) is not None
):
new_state.extra["kv_computed_tokens"] = tuple(
seq.num_computed_tokens for seq in sched_new_req.diffusion_kv_metadata.sequences
)