309 lines
14 KiB
Python
309 lines
14 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 Sequence
|
|
|
|
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
|
|
from vllm_omni.diffusion.diffusion_kv.metadata import (
|
|
DiffusionKVMetadata,
|
|
DiffusionKVSequenceMetadata,
|
|
)
|
|
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
|
|
|
|
|
|
class DiffusionKVAdmissionError(ValueError):
|
|
"""A request-specific KV admission failure safe to return to the caller."""
|
|
|
|
|
|
class DiffusionKVCacheManager:
|
|
"""Atomic public-request facade over native vLLM KVCacheManager."""
|
|
|
|
def __init__(
|
|
self,
|
|
kv_cache_config: KVCacheConfig,
|
|
*,
|
|
max_model_len: int,
|
|
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}")
|
|
if kv_cache_config.num_blocks <= 0 or not kv_cache_config.kv_cache_groups:
|
|
raise ValueError("Diffusion KVCacheConfig must contain a positive block pool and at least one group")
|
|
if scheduler_block_size <= 0 or hash_block_size <= 0:
|
|
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=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()
|
|
self._requests: dict[str, tuple[DiffusionKVRequest, ...]] = {}
|
|
self._metadata: dict[str, DiffusionKVMetadata] = {}
|
|
self._internal_request_ids: set[str] = set()
|
|
self._next_allocation_generation = 1
|
|
|
|
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)
|
|
|
|
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
|
|
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=hit.blocks,
|
|
num_encoder_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,
|
|
)
|
|
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
|
|
|
|
def reserve_request(
|
|
self,
|
|
public_request_id: str,
|
|
kv_requests: Sequence[DiffusionKVRequest],
|
|
) -> DiffusionKVMetadata | None:
|
|
"""Reserve all CFG sequences or roll back the complete public request."""
|
|
|
|
if not public_request_id:
|
|
raise ValueError("public_request_id must be non-empty")
|
|
if public_request_id in self._requests:
|
|
raise ValueError(f"Diffusion KV request {public_request_id!r} is already allocated")
|
|
requests = tuple(kv_requests)
|
|
if not requests:
|
|
raise ValueError("Diffusion KV allocation requires at least one sequence")
|
|
|
|
sequence_ids = [request.sequence_id for request in requests]
|
|
internal_ids = [request.request_id for request in requests]
|
|
expected_sequence_ids = list(range(len(requests)))
|
|
if sequence_ids != expected_sequence_ids:
|
|
raise ValueError(
|
|
"Diffusion KV requests must be ordered by contiguous sequence_id values: "
|
|
f"expected={expected_sequence_ids}, got={sequence_ids}"
|
|
)
|
|
if len(internal_ids) != len(set(internal_ids)):
|
|
raise ValueError(f"Diffusion KV internal request IDs must be unique, got {internal_ids}")
|
|
conflicts = self._internal_request_ids.intersection(internal_ids)
|
|
if conflicts:
|
|
raise ValueError(f"Diffusion KV internal request IDs are already allocated: {sorted(conflicts)}")
|
|
for request in requests:
|
|
if request.kv_contexts:
|
|
raise DiffusionKVAdmissionError("independent DiffusionKVContext allocation is not implemented yet")
|
|
if request.seq_len > self.max_model_len:
|
|
raise DiffusionKVAdmissionError(
|
|
f"Diffusion KV sequence {request.request_id!r} exceeds max_model_len: "
|
|
f"seq_len={request.seq_len}, max_model_len={self.max_model_len}"
|
|
)
|
|
|
|
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: "
|
|
f"required_blocks={required_blocks}, available_blocks={self._empty_pool_num_free_blocks}; "
|
|
"increase KV cache capacity or reduce the request sequence count/length"
|
|
)
|
|
|
|
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, request_computed_blocks in zip(requests, computed_blocks, strict=True):
|
|
blocks = self.native_manager.allocate_slots(
|
|
request,
|
|
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,
|
|
)
|
|
if blocks is None:
|
|
self._rollback(allocated)
|
|
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=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,
|
|
allocation_generation=self._next_allocation_generation,
|
|
sequences=tuple(sequence_metadata),
|
|
)
|
|
self._next_allocation_generation += 1
|
|
self._requests[public_request_id] = requests
|
|
self._metadata[public_request_id] = metadata
|
|
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]
|
|
except KeyError as exc:
|
|
raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc
|
|
|
|
def free_request(self, public_request_id: str) -> None:
|
|
requests = self._requests.pop(public_request_id, ())
|
|
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:
|
|
for public_request_id in tuple(self._requests):
|
|
self.free_request(public_request_id)
|
|
|
|
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
|