8939eb835d
Signed-off-by: Rein Yang <ruiruyang2@gmail.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Rein Yang <ruiruyang2@gmail.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
from dataclasses import dataclass
|
|
|
|
from vllm.v1.core.sched.output import NewRequestData
|
|
from vllm.v1.request import Request
|
|
|
|
from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
|
|
|
|
|
|
@dataclass
|
|
class OmniNewRequestData(NewRequestData):
|
|
"""New request data for omni models with embeddings support.
|
|
|
|
Extends NewRequestData to include prompt embeddings and additional
|
|
information for direct transfer between pipeline stages.
|
|
|
|
Args:
|
|
prompt_embeds: Optional serialized prompt embeddings payload
|
|
additional_information: Optional serialized additional information
|
|
dictionary containing tensors or lists
|
|
"""
|
|
|
|
# Optional serialized prompt embeddings
|
|
prompt_embeds: PromptEmbedsPayload | None = None
|
|
# Optional serialized additional information
|
|
additional_information: AdditionalInformationPayload | None = None
|
|
|
|
@classmethod
|
|
def from_request(
|
|
cls,
|
|
request: Request,
|
|
block_ids: tuple[list[int], ...],
|
|
prefill_token_ids: list[int] | None = None,
|
|
) -> "OmniNewRequestData":
|
|
"""Create OmniNewRequestData from a Request object.
|
|
|
|
Args:
|
|
request: Request object to convert
|
|
block_ids: Tuple of block ID lists for KV cache allocation
|
|
|
|
Returns:
|
|
OmniNewRequestData instance with data from the request
|
|
"""
|
|
return cls(
|
|
req_id=request.request_id,
|
|
prompt_token_ids=request.prompt_token_ids,
|
|
mm_features=request.mm_features,
|
|
sampling_params=request.sampling_params,
|
|
pooling_params=request.pooling_params,
|
|
block_ids=block_ids,
|
|
num_computed_tokens=request.num_computed_tokens,
|
|
lora_request=request.lora_request,
|
|
prompt_embeds=request.prompt_embeds,
|
|
prefill_token_ids=prefill_token_ids,
|
|
additional_information=request.additional_information,
|
|
)
|