tzhouamandClaude Opus 5
9a399adba8
Scoped pre-commit over the 27 files this rebase touches: ruff format reflow, the vLLM-Omni SPDX copyright line on the new modules, and wording the typos hook rejected. The Python 3.10 mypy hook still reports 24 errors across five of these files. All 24 sit on pre-existing lines; none fall on a line this rebase changed, which the repository documents as the CI-skipped condition. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0146Xu5Sgj4vXxLxvLjV5YPj Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
643 lines
30 KiB
Python
643 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.distributed import get_dcp_group
|
|
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec
|
|
from vllm.v1.worker import block_table as native_block_table
|
|
from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
|
|
from vllm.v1.worker.gpu.block_table import BlockTables
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
|
|
from vllm_omni.diffusion.diffusion_kv.layout import (
|
|
adopt_kv_cache_layout,
|
|
assert_backend_layout_supported,
|
|
)
|
|
from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
|
|
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
|
|
DiffusionPagedAttentionAdapter,
|
|
DiffusionPagedAttentionLayerAdapter,
|
|
DiffusionPagedAttentionMetadata,
|
|
DiffusionPagedAttentionRowBinding,
|
|
DiffusionPagedAttentionRuntime,
|
|
)
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
DiffusionKVIdentity = tuple[str, int | None, str | None]
|
|
DiffusionKVSnapshot = tuple[object, ...]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _DiffusionKVRowInstall:
|
|
identity: DiffusionKVIdentity
|
|
token_len: int
|
|
block_ids: tuple[tuple[int, ...], ...]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _DiffusionKVRequestState:
|
|
generation: int
|
|
snapshot: DiffusionKVSnapshot
|
|
row_installs: tuple[_DiffusionKVRowInstall, ...]
|
|
|
|
|
|
class DiffusionKVModelRunnerBackend:
|
|
"""Native paged-KV state and BlockTable operations for a model runner."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
vllm_config: VllmConfig,
|
|
od_config: OmniDiffusionConfig,
|
|
device: torch.device,
|
|
) -> None:
|
|
self.vllm_config = vllm_config
|
|
self.od_config = od_config
|
|
self.device = device
|
|
self.kv_cache_config: KVCacheConfig | None = None
|
|
self.kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
|
|
self.block_tables: BlockTables | None = None
|
|
self.paged_attention_adapter: DiffusionPagedAttentionAdapter | None = None
|
|
self._kv_cache_layer_adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {}
|
|
self._diffusion_kv_identity_to_row: dict[DiffusionKVIdentity, int] = {}
|
|
self._diffusion_kv_free_rows: list[int] = []
|
|
self._diffusion_kv_request_states: dict[str, _DiffusionKVRequestState] = {}
|
|
|
|
def register_kv_cache_layers(
|
|
self,
|
|
layers: Mapping[str, tuple[Any, KVCacheSpec]],
|
|
) -> dict[str, KVCacheSpec]:
|
|
"""Register native adapters and return their canonical cache specs."""
|
|
|
|
if self.kv_cache_config is not None:
|
|
raise RuntimeError("Diffusion KV cache layers cannot be changed after physical initialization")
|
|
|
|
if layers:
|
|
current_omni_platform.get_diffusion_kv_block_tables_cls()
|
|
if not callable(current_omni_platform.build_diffusion_kv_attn_metadata):
|
|
raise NotImplementedError(
|
|
"Diffusion paged KV requires platform-native BlockTables and attention metadata hooks"
|
|
)
|
|
|
|
forward_context = self.vllm_config.compilation_config.static_forward_context
|
|
previous_adapters = self._kv_cache_layer_adapters
|
|
for layer_name, (layer, spec) in layers.items():
|
|
if not isinstance(spec, AttentionSpec):
|
|
raise TypeError(
|
|
f"Diffusion KV layer {layer_name!r} produced unsupported spec {type(spec).__name__}; "
|
|
"only native attention specs are supported"
|
|
)
|
|
if not layer.attn_backend.supports_paged_kv:
|
|
raise NotImplementedError(
|
|
f"Diffusion paged KV layer {layer_name!r} requires an Omni backend with paged support; "
|
|
f"selected {layer.attn_backend.get_name()}"
|
|
)
|
|
existing = forward_context.get(layer_name)
|
|
if existing is not None and existing is not layer and existing is not previous_adapters.get(layer_name):
|
|
raise RuntimeError(f"Diffusion KV layer name {layer_name!r} collides with the native forward context")
|
|
|
|
previous_use_non_causal = self.vllm_config.attention_config.use_non_causal
|
|
self.vllm_config.attention_config.use_non_causal = any(
|
|
bool(getattr(spec, "non_causal", False)) for _, spec in layers.values()
|
|
)
|
|
adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {}
|
|
try:
|
|
with set_current_vllm_config(self.vllm_config):
|
|
for layer_name, (layer, spec) in layers.items():
|
|
adapters[layer_name] = DiffusionPagedAttentionLayerAdapter(
|
|
layer_name=layer_name,
|
|
layer=layer,
|
|
spec=spec,
|
|
vllm_config=self.vllm_config,
|
|
device=self.device,
|
|
ulysses_degree=self._get_paged_attention_ulysses_degree(layer),
|
|
)
|
|
except Exception:
|
|
self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
|
|
raise
|
|
|
|
geometries = {(adapter.num_heads, adapter.num_kv_heads, adapter.head_size) for adapter in adapters.values()}
|
|
if len(geometries) > 1:
|
|
self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
|
|
raise ValueError(
|
|
"Native paged attention metadata builders require one rank-local attention geometry; "
|
|
f"got {sorted(geometries)!r}"
|
|
)
|
|
attention_geometry = next(iter(geometries)) if geometries else None
|
|
set_attention_geometry = (
|
|
getattr(self.vllm_config.model_config, "set_attention_geometry", None)
|
|
if attention_geometry is not None
|
|
else None
|
|
)
|
|
if attention_geometry is not None and not callable(set_attention_geometry):
|
|
self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
|
|
raise RuntimeError(
|
|
"Diffusion native paged attention requires a model config that exposes set_attention_geometry"
|
|
)
|
|
|
|
updated_forward_context = dict(forward_context)
|
|
for layer_name in set(layers) | set(previous_adapters):
|
|
entry = updated_forward_context.get(layer_name)
|
|
layer_entry = layers.get(layer_name)
|
|
if entry is previous_adapters.get(layer_name) or (layer_entry is not None and entry is layer_entry[0]):
|
|
updated_forward_context.pop(layer_name, None)
|
|
updated_forward_context.update(adapters)
|
|
|
|
if attention_geometry is not None:
|
|
num_heads, num_kv_heads, head_size = attention_geometry
|
|
try:
|
|
set_attention_geometry(
|
|
num_heads=num_heads,
|
|
num_kv_heads=num_kv_heads,
|
|
head_size=head_size,
|
|
)
|
|
except Exception:
|
|
self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
|
|
raise
|
|
|
|
forward_context.clear()
|
|
forward_context.update(updated_forward_context)
|
|
self._kv_cache_layer_adapters = adapters
|
|
return {layer_name: adapter.spec for layer_name, adapter in adapters.items()}
|
|
|
|
def _get_paged_attention_ulysses_degree(self, layer: Any) -> int:
|
|
if bool(getattr(layer, "skip_sequence_parallel", False)):
|
|
return 1
|
|
|
|
parallel_config = getattr(self.od_config, "parallel_config", None)
|
|
ulysses_degree = int(getattr(parallel_config, "ulysses_degree", 1))
|
|
ring_degree = int(getattr(parallel_config, "ring_degree", 1))
|
|
allgather_degree = int(getattr(parallel_config, "allgather_degree", 1))
|
|
ulysses_mode = str(getattr(parallel_config, "ulysses_mode", "strict"))
|
|
if ring_degree > 1:
|
|
raise NotImplementedError(
|
|
"Diffusion paged attention does not support Ring or hybrid Ulysses+Ring attention"
|
|
)
|
|
if allgather_degree > 1:
|
|
raise NotImplementedError("Diffusion paged attention does not support AllGather-KV sequence parallelism")
|
|
if ulysses_degree > 1 and ulysses_mode != "strict":
|
|
raise NotImplementedError(
|
|
f"Diffusion paged attention currently supports only strict Ulysses; got ulysses_mode={ulysses_mode!r}"
|
|
)
|
|
return ulysses_degree
|
|
|
|
def _get_max_rows_per_request(self) -> int:
|
|
max_rows = getattr(self.od_config, "diffusion_kv_max_rows_per_request", None)
|
|
if type(max_rows) is not int or max_rows <= 0:
|
|
raise ValueError(
|
|
"paged_scheduler requires a positive integer diffusion_kv_max_rows_per_request adapter limit"
|
|
)
|
|
return max_rows
|
|
|
|
def _get_max_model_len(self) -> int:
|
|
max_len = getattr(self.vllm_config.model_config, "max_model_len", None)
|
|
if type(max_len) is not int or max_len <= 0:
|
|
raise ValueError("paged_scheduler requires a positive model_config.max_model_len")
|
|
return max_len
|
|
|
|
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
|
|
"""Allocate and bind native attention KV tensors for this rank."""
|
|
if self.kv_cache_config is not None:
|
|
raise RuntimeError("Diffusion native KV cache is already initialized")
|
|
if not self._kv_cache_layer_adapters:
|
|
raise RuntimeError("Diffusion KV cache layers must be registered before physical initialization")
|
|
|
|
configured_layers = {
|
|
layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names
|
|
}
|
|
registered_layers = set(self._kv_cache_layer_adapters)
|
|
if configured_layers != registered_layers:
|
|
raise ValueError(
|
|
"Rank-local Diffusion KVCacheConfig layer mismatch: "
|
|
f"expected={sorted(registered_layers)}, configured={sorted(configured_layers)}"
|
|
)
|
|
|
|
kv_cache_config = copy.deepcopy(kv_cache_config)
|
|
max_len = self._get_max_model_len()
|
|
max_rows_per_request = self._get_max_rows_per_request()
|
|
scheduler_config = self.vllm_config.scheduler_config
|
|
max_num_seqs = getattr(scheduler_config, "max_num_seqs", None)
|
|
if type(max_num_seqs) is not int or max_num_seqs <= 0:
|
|
raise ValueError("paged_scheduler requires a positive scheduler_config.max_num_seqs")
|
|
|
|
kv_cache_groups = kv_cache_config.kv_cache_groups
|
|
if not kv_cache_groups:
|
|
raise ValueError("paged_scheduler requires at least one native KV cache group")
|
|
block_sizes = [group.kv_cache_spec.block_size for group in kv_cache_groups]
|
|
if any(type(block_size) is not int or block_size <= 0 for block_size in block_sizes):
|
|
raise ValueError("native diffusion KV cache groups must have positive integer block sizes")
|
|
unaligned_max_num_blocks_per_group = [
|
|
group.kv_cache_spec.max_num_blocks_per_req(self.vllm_config, max_len) for group in kv_cache_groups
|
|
]
|
|
if any(type(capacity) is not int or capacity <= 0 for capacity in unaligned_max_num_blocks_per_group):
|
|
raise ValueError("native diffusion KV cache groups must have positive row capacities")
|
|
max_num_blocks_per_group = [
|
|
native_block_table.get_block_table_width(capacity, block_size)
|
|
for capacity, block_size in zip(unaligned_max_num_blocks_per_group, block_sizes, strict=True)
|
|
]
|
|
|
|
max_num_reqs = max_num_seqs * max_rows_per_request
|
|
max_num_batched_tokens = getattr(scheduler_config, "max_num_batched_tokens", None)
|
|
if type(max_num_batched_tokens) is not int or max_num_batched_tokens <= 0:
|
|
raise ValueError("scheduler_config.max_num_batched_tokens must be a positive integer")
|
|
parallel_config = self.vllm_config.parallel_config
|
|
cp_size = parallel_config.decode_context_parallel_size
|
|
cp_rank = get_dcp_group().rank_in_group if cp_size > 1 else 0
|
|
cp_interleave = parallel_config.cp_kv_cache_interleave_size
|
|
|
|
# Native metadata builders size per-row buffers from max_num_seqs.
|
|
# A diffusion public request can occupy several sequence/context rows.
|
|
scheduler_config.max_num_seqs = max_num_reqs
|
|
kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
|
|
previous_adapter_caches = {
|
|
layer_name: adapter.kv_cache for layer_name, adapter in self._kv_cache_layer_adapters.items()
|
|
}
|
|
try:
|
|
with set_current_vllm_config(self.vllm_config):
|
|
attn_groups, _, kernel_block_sizes = init_attn_backend(
|
|
kv_cache_config,
|
|
self.vllm_config,
|
|
self.device,
|
|
)
|
|
block_tables_cls = current_omni_platform.get_diffusion_kv_block_tables_cls()
|
|
block_tables = block_tables_cls(
|
|
block_sizes=block_sizes,
|
|
max_num_reqs=max_num_reqs,
|
|
max_num_batched_tokens=max_num_batched_tokens,
|
|
max_num_blocks_per_group=max_num_blocks_per_group,
|
|
device=self.device,
|
|
kernel_block_sizes=kernel_block_sizes,
|
|
cp_size=cp_size,
|
|
cp_rank=cp_rank,
|
|
cp_interleave=cp_interleave,
|
|
)
|
|
paged_attention_adapter = DiffusionPagedAttentionAdapter(
|
|
vllm_config=self.vllm_config,
|
|
device=self.device,
|
|
kv_cache_config=kv_cache_config,
|
|
block_tables=block_tables,
|
|
attn_groups=attn_groups,
|
|
layers=self._kv_cache_layer_adapters,
|
|
resolve_row=self._resolve_paged_attention_row,
|
|
)
|
|
# vLLM 0.29 reads the resolved physical layout when allocating,
|
|
# and dropped attn_groups/cache_dtype from the signature.
|
|
adopt_kv_cache_layout(self.vllm_config, kv_cache_config)
|
|
for layer_adapter in self._kv_cache_layer_adapters.values():
|
|
assert_backend_layout_supported(self.vllm_config, getattr(layer_adapter, "attn_backend", None))
|
|
init_kv_cache(
|
|
kv_caches,
|
|
self.vllm_config.compilation_config.static_forward_context,
|
|
kv_cache_config,
|
|
self.device,
|
|
kernel_block_sizes,
|
|
self.vllm_config,
|
|
)
|
|
except Exception:
|
|
scheduler_config.max_num_seqs = max_num_seqs
|
|
for layer_name, previous_cache in previous_adapter_caches.items():
|
|
self._kv_cache_layer_adapters[layer_name].kv_cache = previous_cache
|
|
raise
|
|
|
|
self.kv_caches = kv_caches
|
|
self.block_tables = block_tables
|
|
self.paged_attention_adapter = paged_attention_adapter
|
|
self._diffusion_kv_identity_to_row.clear()
|
|
self._diffusion_kv_free_rows = list(range(max_num_reqs - 1, -1, -1))
|
|
self._diffusion_kv_request_states.clear()
|
|
self.kv_cache_config = kv_cache_config
|
|
|
|
def _validate_row(
|
|
self,
|
|
*,
|
|
identity: DiffusionKVIdentity,
|
|
token_len: int,
|
|
block_ids: tuple[list[int], ...],
|
|
seen_identities: set[DiffusionKVIdentity],
|
|
) -> _DiffusionKVRowInstall:
|
|
if identity in seen_identities:
|
|
raise ValueError(f"Duplicate diffusion KV row identity: {identity!r}")
|
|
seen_identities.add(identity)
|
|
if type(token_len) is not int or token_len < 0:
|
|
raise ValueError(f"Diffusion KV row {identity!r} has invalid token length {token_len!r}")
|
|
|
|
assert self.kv_cache_config is not None
|
|
groups = self.kv_cache_config.kv_cache_groups
|
|
if not isinstance(block_ids, tuple) or len(block_ids) != len(groups):
|
|
actual_groups = len(block_ids) if isinstance(block_ids, (tuple, list)) else "invalid"
|
|
raise ValueError(f"Diffusion KV row {identity!r} has {actual_groups} block groups; expected {len(groups)}")
|
|
|
|
max_len = self._get_max_model_len()
|
|
normalized_groups: list[tuple[int, ...]] = []
|
|
for group_index, (group_block_ids, kv_cache_group) in enumerate(zip(block_ids, groups, strict=True)):
|
|
if not isinstance(group_block_ids, list):
|
|
raise ValueError(f"Diffusion KV row {identity!r} group {group_index} block IDs must be a list")
|
|
spec = kv_cache_group.kv_cache_spec
|
|
expected_count = spec.max_num_blocks_per_req(self.vllm_config, token_len)
|
|
if len(group_block_ids) != expected_count:
|
|
raise ValueError(
|
|
f"Diffusion KV row {identity!r} group {group_index} has {len(group_block_ids)} blocks; "
|
|
f"expected {expected_count} for {token_len} tokens"
|
|
)
|
|
row_capacity = spec.max_num_blocks_per_req(self.vllm_config, max_len)
|
|
if len(group_block_ids) > row_capacity:
|
|
raise ValueError(
|
|
f"Diffusion KV row {identity!r} group {group_index} exceeds row capacity {row_capacity}"
|
|
)
|
|
normalized_ids: list[int] = []
|
|
for block_id in group_block_ids:
|
|
if type(block_id) is not int:
|
|
raise ValueError(
|
|
f"Diffusion KV row {identity!r} group {group_index} has non-integer block ID {block_id!r}"
|
|
)
|
|
if not 0 <= block_id < self.kv_cache_config.num_blocks:
|
|
raise ValueError(
|
|
f"Diffusion KV row {identity!r} group {group_index} block ID {block_id} is outside "
|
|
f"[0, {self.kv_cache_config.num_blocks})"
|
|
)
|
|
normalized_ids.append(block_id)
|
|
normalized_groups.append(tuple(normalized_ids))
|
|
|
|
return _DiffusionKVRowInstall(
|
|
identity=identity,
|
|
token_len=token_len,
|
|
block_ids=tuple(normalized_groups),
|
|
)
|
|
|
|
def _prepare_install(
|
|
self,
|
|
metadata: DiffusionKVMetadata,
|
|
) -> tuple[list[_DiffusionKVRowInstall], DiffusionKVSnapshot]:
|
|
if self.block_tables is None or self.kv_cache_config is None:
|
|
raise RuntimeError("paged_scheduler native KV cache is not initialized")
|
|
if type(metadata.request_id) is not str or not metadata.request_id:
|
|
raise ValueError("Diffusion KV metadata request_id must be a non-empty string")
|
|
if type(metadata.allocation_generation) is not int or metadata.allocation_generation < 0:
|
|
raise ValueError("Diffusion KV allocation_generation must be a non-negative integer")
|
|
if type(self.kv_cache_config.num_blocks) is not int or self.kv_cache_config.num_blocks <= 0:
|
|
raise ValueError("native diffusion KV cache config must provide a positive num_blocks")
|
|
if not metadata.sequences:
|
|
raise ValueError(f"Diffusion KV request {metadata.request_id!r} must contain at least one sequence")
|
|
|
|
seen_identities: set[DiffusionKVIdentity] = set()
|
|
installs: list[_DiffusionKVRowInstall] = []
|
|
sequence_snapshots: list[object] = []
|
|
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}")
|
|
sequence_identity = (metadata.request_id, sequence.sequence_id, None)
|
|
sequence_install = self._validate_row(
|
|
identity=sequence_identity,
|
|
token_len=sequence.seq_len,
|
|
block_ids=sequence.block_ids,
|
|
seen_identities=seen_identities,
|
|
)
|
|
installs.append(sequence_install)
|
|
sequence_snapshots.append(
|
|
(
|
|
sequence.sequence_id,
|
|
sequence.prefix_len,
|
|
sequence.target_len,
|
|
sequence.seq_len,
|
|
sequence_install.block_ids,
|
|
sequence.context_ids,
|
|
)
|
|
)
|
|
|
|
context_snapshots: list[object] = []
|
|
for context in metadata.contexts:
|
|
if type(context.context_id) is not str or not context.context_id:
|
|
raise ValueError("Diffusion KV context_id must be a non-empty string")
|
|
context_identity = (metadata.request_id, None, context.context_id)
|
|
context_install = self._validate_row(
|
|
identity=context_identity,
|
|
token_len=context.num_tokens,
|
|
block_ids=context.block_ids,
|
|
seen_identities=seen_identities,
|
|
)
|
|
installs.append(context_install)
|
|
context_snapshots.append(
|
|
(
|
|
context.context_id,
|
|
context.cache_role,
|
|
context.num_tokens,
|
|
context_install.block_ids,
|
|
)
|
|
)
|
|
|
|
max_rows_per_request = self._get_max_rows_per_request()
|
|
if len(installs) > max_rows_per_request:
|
|
raise ValueError(
|
|
f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows; "
|
|
f"adapter limit is {max_rows_per_request}"
|
|
)
|
|
snapshot: DiffusionKVSnapshot = (
|
|
metadata.request_id,
|
|
tuple(sequence_snapshots),
|
|
tuple(context_snapshots),
|
|
)
|
|
return installs, snapshot
|
|
|
|
def _rollback_staged_writes(self, rows: list[int], previous_num_blocks: object | None) -> None:
|
|
block_tables = self.block_tables
|
|
if block_tables is None:
|
|
return
|
|
for group_table in getattr(block_tables, "block_tables", ()):
|
|
clear_staged_writes = getattr(group_table, "clear_staged_writes", None)
|
|
if clear_staged_writes is not None:
|
|
clear_staged_writes()
|
|
if previous_num_blocks is None:
|
|
return
|
|
native_num_blocks = getattr(block_tables, "num_blocks", None)
|
|
if native_num_blocks is None or not hasattr(native_num_blocks, "np"):
|
|
return
|
|
native_num_blocks.np[:, rows] = previous_num_blocks
|
|
copy_to_uva = getattr(native_num_blocks, "copy_to_uva", None)
|
|
if copy_to_uva is not None:
|
|
copy_to_uva()
|
|
|
|
def _apply_rows(self, rows: list[int], installs: list[_DiffusionKVRowInstall]) -> None:
|
|
assert self.block_tables is not None
|
|
if self.paged_attention_adapter is not None:
|
|
self.paged_attention_adapter.invalidate_prepared_batches()
|
|
native_num_blocks = getattr(self.block_tables, "num_blocks", None)
|
|
previous_num_blocks = None
|
|
if native_num_blocks is not None and hasattr(native_num_blocks, "np"):
|
|
previous_num_blocks = native_num_blocks.np[:, rows].copy()
|
|
try:
|
|
for row, install in zip(rows, installs, strict=True):
|
|
self.block_tables.append_block_ids(
|
|
row,
|
|
tuple(list(group_ids) for group_ids in install.block_ids),
|
|
overwrite=True,
|
|
)
|
|
self.block_tables.apply_staged_writes()
|
|
except Exception:
|
|
self._rollback_staged_writes(rows, previous_num_blocks)
|
|
raise
|
|
|
|
def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool:
|
|
"""Install one Scheduler allocation snapshot into native Worker rows."""
|
|
if not is_scheduler_paged_kv_mode(
|
|
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
|
):
|
|
raise ValueError("Dense diffusion execution must not install Diffusion KV metadata")
|
|
|
|
installs, snapshot = self._prepare_install(metadata)
|
|
current_state = self._diffusion_kv_request_states.get(metadata.request_id)
|
|
if current_state is not None:
|
|
if metadata.allocation_generation < current_state.generation:
|
|
raise ValueError(
|
|
f"Stale Diffusion KV allocation generation {metadata.allocation_generation} for request "
|
|
f"{metadata.request_id!r}; installed generation is {current_state.generation}"
|
|
)
|
|
if snapshot != current_state.snapshot:
|
|
raise ValueError(
|
|
f"Conflicting Diffusion KV allocation snapshot for active request {metadata.request_id!r}; "
|
|
"remove the request before replacing its rows"
|
|
)
|
|
if metadata.allocation_generation == current_state.generation:
|
|
return False
|
|
self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
|
|
generation=metadata.allocation_generation,
|
|
snapshot=snapshot,
|
|
row_installs=current_state.row_installs,
|
|
)
|
|
return False
|
|
|
|
existing_identities = [
|
|
install.identity for install in installs if install.identity in self._diffusion_kv_identity_to_row
|
|
]
|
|
if existing_identities:
|
|
raise RuntimeError(f"Diffusion KV row registry contains orphaned identities: {existing_identities!r}")
|
|
if len(installs) > len(self._diffusion_kv_free_rows):
|
|
raise ValueError(
|
|
f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows, but only "
|
|
f"{len(self._diffusion_kv_free_rows)} rows are free"
|
|
)
|
|
|
|
rows = list(reversed(self._diffusion_kv_free_rows[-len(installs) :])) if installs else []
|
|
if rows:
|
|
self._apply_rows(rows, installs)
|
|
|
|
if installs:
|
|
del self._diffusion_kv_free_rows[-len(installs) :]
|
|
self._diffusion_kv_identity_to_row.update(
|
|
(install.identity, row) for install, row in zip(installs, rows, strict=True)
|
|
)
|
|
self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
|
|
generation=metadata.allocation_generation,
|
|
snapshot=snapshot,
|
|
row_installs=tuple(installs),
|
|
)
|
|
return True
|
|
|
|
def get_diffusion_kv_row(
|
|
self,
|
|
request_id: str,
|
|
sequence_id: int | None,
|
|
context_id: str | None = None,
|
|
) -> int:
|
|
"""Resolve a Scheduler allocation identity to its native table row."""
|
|
if (
|
|
not is_scheduler_paged_kv_mode(
|
|
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
|
)
|
|
or self.block_tables is None
|
|
):
|
|
raise RuntimeError("paged_scheduler native BlockTables are not initialized")
|
|
identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
|
|
try:
|
|
return self._diffusion_kv_identity_to_row[identity]
|
|
except KeyError as exc:
|
|
raise KeyError(f"No native diffusion KV row is installed for {identity!r}") from exc
|
|
|
|
def get_paged_attention_adapter(self) -> DiffusionPagedAttentionAdapter:
|
|
if self.paged_attention_adapter is None:
|
|
raise RuntimeError("paged_scheduler native attention adapter is not initialized")
|
|
return self.paged_attention_adapter
|
|
|
|
def activate_paged_attention_metadata(self, metadata: DiffusionPagedAttentionMetadata):
|
|
"""Activate Runner-owned request metadata for an internal denoise loop."""
|
|
|
|
runtime = DiffusionPagedAttentionRuntime(self.get_paged_attention_adapter(), metadata)
|
|
return runtime, runtime.activate()
|
|
|
|
def _resolve_paged_attention_row(
|
|
self,
|
|
request_id: str,
|
|
sequence_id: int | None,
|
|
context_id: str | None,
|
|
) -> DiffusionPagedAttentionRowBinding:
|
|
row_index = self.get_diffusion_kv_row(request_id, sequence_id, context_id)
|
|
identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
|
|
request_state = self._diffusion_kv_request_states.get(request_id)
|
|
if request_state is not None:
|
|
for install in request_state.row_installs:
|
|
if install.identity == identity:
|
|
return DiffusionPagedAttentionRowBinding(
|
|
row_index=row_index,
|
|
max_seq_len=install.token_len,
|
|
block_ids=install.block_ids,
|
|
)
|
|
raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
|
|
|
|
def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
|
|
"""Retire Worker rows without logically freeing Scheduler-owned blocks."""
|
|
if (
|
|
not is_scheduler_paged_kv_mode(
|
|
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
|
)
|
|
or self.block_tables is None
|
|
):
|
|
return 0
|
|
request_id_set = set(request_ids)
|
|
identities_and_rows = [
|
|
(identity, row)
|
|
for identity, row in self._diffusion_kv_identity_to_row.items()
|
|
if identity[0] in request_id_set
|
|
]
|
|
rows = [row for _, row in identities_and_rows]
|
|
if rows:
|
|
assert self.kv_cache_config is not None
|
|
num_groups = len(self.kv_cache_config.kv_cache_groups)
|
|
installs = [
|
|
_DiffusionKVRowInstall(
|
|
identity=identity,
|
|
token_len=0,
|
|
block_ids=tuple(tuple() for _ in range(num_groups)),
|
|
)
|
|
for identity, _ in identities_and_rows
|
|
]
|
|
self._apply_rows(rows, installs)
|
|
|
|
for identity, _ in identities_and_rows:
|
|
del self._diffusion_kv_identity_to_row[identity]
|
|
for request_id in request_id_set:
|
|
self._diffusion_kv_request_states.pop(request_id, None)
|
|
self._diffusion_kv_free_rows.extend(rows)
|
|
self._diffusion_kv_free_rows.sort(reverse=True)
|
|
return len(rows)
|
|
|
|
def refresh_block_table_layout(self) -> None:
|
|
"""Refresh native pointer tensors after a CuMem KV-cache wake-up."""
|
|
if (
|
|
not is_scheduler_paged_kv_mode(
|
|
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
|
)
|
|
or self.block_tables is None
|
|
):
|
|
return
|
|
if self.paged_attention_adapter is not None:
|
|
self.paged_attention_adapter.invalidate_prepared_batches()
|
|
self.block_tables.init_block_table_layout_tensors()
|