Wang fuyinandHongsheng Liu
af3faa6c1e
Signed-off-by: Acerak01-fy <wfy2003324@163.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Scheduler-to-Worker allocation payloads for diffusion KV.
|
|
|
|
Like native vLLM's SchedulerOutput DTOs, these classes carry trusted internal
|
|
Scheduler results. Cache geometry and block-table validation stay with the
|
|
native cache builders and Worker installation path.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
class DiffusionKVContextMetadata:
|
|
"""Request-scoped Worker allocation for an independent K/V context.
|
|
|
|
``context_id`` identifies this allocation within the enclosing
|
|
:class:`DiffusionKVMetadata`.
|
|
"""
|
|
|
|
context_id: str
|
|
cache_role: str
|
|
num_tokens: int
|
|
block_ids: tuple[list[int], ...]
|
|
|
|
|
|
@dataclass
|
|
class DiffusionKVSequenceMetadata:
|
|
"""Worker allocation for one Scheduler-owned diffusion sequence.
|
|
|
|
``context_ids`` reference request-scoped context allocations from the
|
|
enclosing :class:`DiffusionKVMetadata`.
|
|
"""
|
|
|
|
sequence_id: int
|
|
prefix_len: int
|
|
target_len: int
|
|
seq_len: int
|
|
block_ids: tuple[list[int], ...]
|
|
context_ids: tuple[str, ...] = ()
|
|
|
|
|
|
@dataclass
|
|
class DiffusionKVMetadata:
|
|
"""Scheduler allocation sent with a newly scheduled public request.
|
|
|
|
Independent contexts live once at request scope so multiple sequences can
|
|
reference the same allocation.
|
|
"""
|
|
|
|
request_id: str
|
|
allocation_generation: int
|
|
sequences: tuple[DiffusionKVSequenceMetadata, ...]
|
|
contexts: tuple[DiffusionKVContextMetadata, ...] = ()
|