72 lines
1.9 KiB
Python
72 lines
1.9 KiB
Python
|
||
|
||
from enum import Enum
|
||
import os
|
||
import torch
|
||
import vllm.envs as envs
|
||
|
||
zero_no_thread = os.environ.get('VLLM_ZERO_NO_THREAD') == '1'
|
||
|
||
def is_zero_no_thread():
|
||
return zero_no_thread and envs.VLLM_ZERO_OVERHEAD
|
||
|
||
class SpecStepKind(Enum):
|
||
KIND_DEFAULT = 0
|
||
PREFILL = 1
|
||
FIRST_PROPOSAL = 2
|
||
OTHER_PROPOSAL = 3
|
||
SCORE_DECODE = 4
|
||
|
||
class ZeroOverheadSpecContext():
|
||
def __init__(self):
|
||
self.step_kind = SpecStepKind.KIND_DEFAULT
|
||
self.last_step = SpecStepKind.KIND_DEFAULT
|
||
self.proposal_lens_list = None
|
||
self.proposal_token_ids = None
|
||
self.accepted_token_ids = None
|
||
self.accepted_seq_ids = None
|
||
|
||
spec_context = ZeroOverheadSpecContext()
|
||
|
||
def set_spec_step(_step):
|
||
global spec_context
|
||
spec_context.last_step = spec_context.step_kind
|
||
spec_context.step_kind = _step
|
||
|
||
def get_spec_step():
|
||
return spec_context.step_kind
|
||
|
||
def get_spec_last_step():
|
||
return spec_context.last_step
|
||
|
||
def record_proposal_lens_list(list):
|
||
global spec_context
|
||
spec_context.proposal_lens_list = list
|
||
|
||
def get_proposal_lens_list():
|
||
return spec_context.proposal_lens_list
|
||
|
||
def record_proposal_token_ids(tensor):
|
||
global spec_context
|
||
spec_context.proposal_token_ids = tensor
|
||
|
||
def get_proposal_token_ids():
|
||
return spec_context.proposal_token_ids
|
||
|
||
def record_accepted_token_ids(tensor, seq_ids):
|
||
global spec_context
|
||
spec_context.accepted_token_ids = tensor
|
||
spec_context.accepted_seq_ids = seq_ids
|
||
|
||
def get_accepted_token_ids():
|
||
return spec_context.accepted_token_ids, spec_context.accepted_seq_ids
|
||
|
||
# 零消耗调度不在默认流上推理，用以规避runtime引入的内存申请流同步问题。
|
||
alloc_stream = {}
|
||
|
||
def zero_overhead_stream(target_device):
|
||
"""Asynchronously create a tensor and copy it from host to device."""
|
||
if target_device not in alloc_stream.keys():
|
||
alloc_stream[target_device] = torch.cuda.Stream(device=target_device)
|
||
return alloc_stream[target_device]
|