211 lines
9.7 KiB
Python
211 lines
9.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Sequence
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
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
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DiffusionKVContext:
|
|
"""An independently managed K/V context outside the primary sequence.
|
|
|
|
``context_id`` identifies the logical context within one execution
|
|
sequence. ``cache_role`` binds it to a logical attention cache role exposed
|
|
by the Worker. Physical cache geometry remains native ``KVCacheSpec`` /
|
|
``KVCacheConfig`` state and is deliberately absent here.
|
|
"""
|
|
|
|
context_id: str
|
|
cache_role: str
|
|
num_tokens: int
|
|
block_hashes: tuple[BlockHash, ...] = ()
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.context_id:
|
|
raise ValueError("context_id must be non-empty")
|
|
if not self.cache_role:
|
|
raise ValueError("cache_role must be non-empty")
|
|
if self.num_tokens <= 0:
|
|
raise ValueError(f"num_tokens must be positive, got {self.num_tokens}")
|
|
object.__setattr__(self, "block_hashes", tuple(self.block_hashes))
|
|
|
|
|
|
class DiffusionKVRequest:
|
|
"""Scheduler-owned KV state for one diffusion execution sequence.
|
|
|
|
The primary sequence follows an ordered ``[prefix | target]`` policy, such
|
|
as one Hunyuan CFG row. ``prefix_len`` is the contiguous reusable prefix;
|
|
``target_len`` is overwritten by every denoise step; ``num_tokens`` is the
|
|
complete first-step allocation boundary.
|
|
|
|
Independent cross/joint-attention K/V does not belong to that token axis
|
|
and is described by ``kv_contexts``. This object also exposes the minimal
|
|
mutable Request surface consumed by native ``KVCacheManager``.
|
|
|
|
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``.
|
|
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__(
|
|
self,
|
|
request_id: str,
|
|
*,
|
|
sequence_id: int,
|
|
prefix_len: int,
|
|
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,
|
|
) -> None:
|
|
if not request_id:
|
|
raise ValueError("request_id must be non-empty")
|
|
if sequence_id < 0:
|
|
raise ValueError(f"sequence_id must be non-negative, got {sequence_id}")
|
|
if prefix_len < 0:
|
|
raise ValueError(f"prefix_len must be non-negative, got {prefix_len}")
|
|
if target_len <= 0:
|
|
raise ValueError(f"target_len must be positive, got {target_len}")
|
|
if seq_len <= 0:
|
|
raise ValueError(f"seq_len must be positive, got {seq_len}")
|
|
if prefix_len + target_len > seq_len:
|
|
raise ValueError(
|
|
"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):
|
|
raise TypeError("kv_contexts must contain only DiffusionKVContext values")
|
|
context_ids = [context.context_id for context in contexts]
|
|
if len(context_ids) != len(set(context_ids)):
|
|
raise ValueError(f"kv_contexts must have unique context_id values, got {context_ids}")
|
|
|
|
# Diffusion semantics consumed by the Scheduler-side facade.
|
|
self.sequence_id = sequence_id
|
|
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.
|
|
self.request_id = request_id
|
|
self.num_tokens = seq_len
|
|
self.num_prompt_tokens = prefix_len
|
|
self.num_computed_tokens = 0
|
|
self.block_hashes = list(block_hashes)
|
|
self.skip_reading_prefix_cache = not self.block_hashes
|
|
self.status = RequestStatus.WAITING
|
|
self.num_preemptions = 0
|
|
self.num_in_flight_tokens = 0
|
|
# vLLM 0.26+ uses this optional boundary when publishing cached blocks;
|
|
# zero means that no sparse-retention boundary is active.
|
|
self.shared_prefix_boundary = 0
|
|
# Opaque request-scoped parameters consumed by an upstream KVConnector.
|
|
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."""
|
|
|
|
return self.num_tokens
|