NATUREandHongsheng Liu
0d4e975e1b
Signed-off-by: natureofnature <wzliu@connect.hku.hk> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
65 lines
2.6 KiB
Python
65 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Helpers for explicit thinker decode span metadata."""
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
THINKER_DECODE_EMBEDDINGS_KEY = "thinker_decode_embeddings"
|
|
THINKER_OUTPUT_TOKEN_IDS_KEY = "thinker_output_token_ids"
|
|
THINKER_DECODE_TOKEN_START_KEY = "thinker_decode_embeddings_token_start"
|
|
THINKER_DECODE_TOKEN_END_KEY = "thinker_decode_embeddings_token_end"
|
|
|
|
CACHED_THINKER_DECODE_EMBEDDINGS_KEY = "cached_thinker_decode_embeddings"
|
|
CACHED_THINKER_DECODE_TOKEN_START_KEY = "cached_thinker_decode_embeddings_token_start"
|
|
CACHED_THINKER_DECODE_TOKEN_END_KEY = "cached_thinker_decode_embeddings_token_end"
|
|
|
|
TensorSpan = tuple[torch.Tensor, int, int]
|
|
|
|
|
|
def get_tensor_span(payload: Mapping[str, Any], *, tensor_key: str, start_key: str, end_key: str) -> TensorSpan | None:
|
|
tensor = payload.get(tensor_key)
|
|
start = payload.get(start_key)
|
|
end = payload.get(end_key)
|
|
if not isinstance(tensor, torch.Tensor):
|
|
return None
|
|
if not isinstance(start, int) or not isinstance(end, int):
|
|
return None
|
|
if start < 0 or end < start or (end - start) != int(tensor.shape[0]):
|
|
return None
|
|
return tensor, start, end
|
|
|
|
|
|
def merge_tensor_spans(existing_span: TensorSpan | None, incoming_span: TensorSpan | None) -> TensorSpan | None:
|
|
if existing_span is None or incoming_span is None:
|
|
return None
|
|
|
|
existing_tensor, existing_start, existing_end = existing_span
|
|
incoming_tensor, incoming_start, incoming_end = incoming_span
|
|
if incoming_tensor.device != existing_tensor.device or incoming_tensor.dtype != existing_tensor.dtype:
|
|
incoming_tensor = incoming_tensor.to(device=existing_tensor.device, dtype=existing_tensor.dtype)
|
|
if incoming_start == existing_end:
|
|
return torch.cat([existing_tensor, incoming_tensor], dim=0), existing_start, incoming_end
|
|
if incoming_start < existing_end:
|
|
overlap = existing_end - incoming_start
|
|
if overlap >= int(incoming_tensor.shape[0]):
|
|
return existing_tensor, existing_start, existing_end
|
|
trimmed_tensor = incoming_tensor[overlap:]
|
|
return (
|
|
torch.cat([existing_tensor, trimmed_tensor], dim=0),
|
|
existing_start,
|
|
existing_end + int(trimmed_tensor.shape[0]),
|
|
)
|
|
return None
|
|
|
|
|
|
def get_tensor_span_row(span: TensorSpan | None, index: int) -> torch.Tensor | None:
|
|
if span is None:
|
|
return None
|
|
tensor, start, end = span
|
|
if index < start or index >= end:
|
|
return None
|
|
return tensor[index - start]
|