Sy03andYueqian Lin
9eaca72382
Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
34 lines
963 B
Python
34 lines
963 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from typing import NamedTuple, Protocol
|
|
|
|
|
|
class RunnerAssistedFullAttentionMetadataRequest(NamedTuple):
|
|
num_reqs_padded: int
|
|
for_cudagraph_capture: bool
|
|
|
|
|
|
class RunnerAssistedAttentionMetadataProvider(Protocol):
|
|
"""Optional model hook for runner-built full attention metadata."""
|
|
|
|
def get_runner_assisted_full_attention_metadata_request(
|
|
self,
|
|
*,
|
|
req_ids: Sequence[str],
|
|
num_reqs: int,
|
|
num_scheduled_tokens: Sequence[int],
|
|
num_computed_tokens: Sequence[int],
|
|
max_num_scheduled_tokens: int,
|
|
) -> RunnerAssistedFullAttentionMetadataRequest | None: ...
|
|
|
|
def set_runner_assisted_full_attention_metadata_context(
|
|
self,
|
|
*,
|
|
enabled: bool,
|
|
num_reqs: int = 0,
|
|
) -> None: ...
|