172 lines
8.7 KiB
Python
172 lines
8.7 KiB
Python
|
|
|
|
import torch
|
|
import itertools
|
|
from typing import List, Optional, Set
|
|
from vllm.lora.layers import LoRAMapping
|
|
from vllm.multimodal.inputs import MultiModalKwargs
|
|
from vllm.prompt_adapter.layers import PromptAdapterMapping
|
|
from vllm.prompt_adapter.request import PromptAdapterRequest
|
|
from vllm.sequence import SequenceGroupMetadata
|
|
from vllm.utils import async_tensor_h2d, flatten_2d_lists
|
|
from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder
|
|
from vllm.zero_overhead.sampler import get_last_sampler
|
|
from vllm.zero_overhead.utils import SpecStepKind, get_accepted_token_ids, get_proposal_token_ids, get_spec_last_step, get_spec_step
|
|
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
@triton.jit
|
|
def _update_input_tokens(
|
|
accepted_req_ids,
|
|
accepted_req_ids_len,
|
|
accepted_token_ids,
|
|
accepted_token_len,
|
|
chidren_req_ids,
|
|
chidren_req_ids_len,
|
|
input_tokens,
|
|
input_tokens_len,
|
|
input_positions,
|
|
seq_lens,
|
|
seq_lens_meta,
|
|
seq_lens_tensor,
|
|
slot_mapping,
|
|
seq_start_loc,
|
|
context_lens_tensor,
|
|
):
|
|
chidren_req_ids_ = tl.load(chidren_req_ids + tl.arange(0, chidren_req_ids_len))
|
|
accepted_req_ids_ = tl.load(accepted_req_ids + tl.arange(0, chidren_req_ids_len))
|
|
|
|
for seq_id_idx in range(chidren_req_ids_len / 2):
|
|
seq_id = chidren_req_ids_[2 * seq_id_idx]
|
|
for i in range(accepted_req_ids_len):
|
|
if seq_id == accepted_req_ids_[i]:
|
|
accepted_token_ids_ = tl.load(accepted_token_ids + tl.arange(i * accepted_token_len, tl.arange(0, accepted_token_len)))
|
|
accepted_token_counter = 0
|
|
for j in range(accepted_token_len):
|
|
if accepted_token_ids_[j] == -1:
|
|
break
|
|
accepted_token_counter += 1
|
|
if accepted_token_counter == accepted_token_len:
|
|
tl.store(input_tokens + seq_id_idx * 2 + tl.arange(0, 2), accepted_token_ids_[-2:])
|
|
else:
|
|
tl.store(input_tokens + seq_id_idx * 2, 0)
|
|
tl.store(input_tokens + seq_id_idx * 2 + 1, accepted_token_ids_[accepted_token_counter - 1])
|
|
input_pos = tl.load(input_positions + seq_id_idx * 2 + tl.arange(0, 2))
|
|
input_pos[0] = 0
|
|
input_pos[1] = input_pos[1] - (accepted_req_ids_len - accepted_token_counter)
|
|
tl.store(input_positions + seq_id_idx * 2 + tl.arange(0, 2), input_pos)
|
|
tl.store(context_lens_tensor + seq_id_idx * 2 + tl.arange(0, 2), input_pos)
|
|
input_pos[0] = -1
|
|
tl.store(slot_mapping + seq_id_idx * 2 + tl.arange(0, 2), input_pos)
|
|
input_pos[0] = 1
|
|
input_pos[1] = input_pos[1] + 1
|
|
tl.store(seq_lens + seq_id_idx * 2 + tl.arange(0, 2), input_pos)
|
|
tl.store(seq_lens_meta + seq_id_idx * 2 + tl.arange(0, 2), input_pos)
|
|
tl.store(seq_lens_tensor + seq_id_idx * 2 + tl.arange(0, 2), input_pos)
|
|
seq_lens_ = tl.load(seq_lens + tl.arange(0, input_tokens_len))
|
|
seq_start_loc_ = tl.zero_like(seq_start_loc)
|
|
for i in range(input_tokens_len):
|
|
seq_start_loc_[i + 1] = seq_start_loc_[i] + seq_lens_[i]
|
|
tl.store(seq_start_loc + tl.arange(0, input_tokens_len + 1), seq_start_loc_)
|
|
|
|
|
|
|
|
class ZeroOverheadModelInputForGpuBuilder(ModelInputForGPUBuilder):
|
|
def __init__(self, runner, finished_requests_ids = None):
|
|
super().__init__(runner, finished_requests_ids)
|
|
self.req_ids = []
|
|
|
|
def prepare(self,
|
|
finished_requests_ids: Optional[List[str]] = None) -> None:
|
|
self.req_ids.clear()
|
|
return super().prepare(finished_requests_ids)
|
|
|
|
def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
|
|
seq_ids = seq_group_metadata.seq_data.keys()
|
|
n_seqs = len(seq_ids)
|
|
seq_ids = list(seq_ids)
|
|
for seq_idx in range(n_seqs):
|
|
self.req_ids.append(seq_ids[seq_idx])
|
|
return super().add_seq_group(seq_group_metadata)
|
|
|
|
def build(self) -> ModelInputForGPU:
|
|
model_input = super().build()
|
|
last_sampler = get_last_sampler()
|
|
spec_step = get_spec_step()
|
|
last_step = get_spec_last_step()
|
|
if last_sampler is not None:
|
|
if spec_step == SpecStepKind.KIND_DEFAULT:
|
|
update_indices = []
|
|
select_indices = []
|
|
query_idx = 0
|
|
for i, seq_id in enumerate(self.req_ids):
|
|
for j, seq_id_ in enumerate(last_sampler.seq_ids):
|
|
if seq_id == seq_id_:
|
|
select_indices.append(j)
|
|
update_indices.append(query_idx)
|
|
break
|
|
query_idx += model_input.query_lens[i]
|
|
if len(select_indices) > 0 and last_sampler.sampled_token_ids_tensor is not None:
|
|
select_indices = async_tensor_h2d(select_indices, torch.long,
|
|
self.runner.device,
|
|
self.runner.pin_memory)
|
|
update_indices = async_tensor_h2d(update_indices, torch.long,
|
|
self.runner.device,
|
|
self.runner.pin_memory)
|
|
model_input.input_tokens[update_indices] = last_sampler.sampled_token_ids_tensor[select_indices, 0]
|
|
if spec_step == SpecStepKind.OTHER_PROPOSAL:
|
|
if last_step == SpecStepKind.OTHER_PROPOSAL: # copy last sampled token ids to input tokens directly.
|
|
update_indices = [i for i in range(len(self.req_ids))]
|
|
update_indices = async_tensor_h2d(update_indices, torch.long,
|
|
self.runner.device,
|
|
self.runner.pin_memory)
|
|
model_input.input_tokens[update_indices] = last_sampler.sampled_token_ids_tensor[update_indices, 0]
|
|
if last_step == SpecStepKind.FIRST_PROPOSAL: # TODO: ajust input tokens number to 1 per request.
|
|
update_indices = [i for i in range(len(self.req_ids))]
|
|
update_indices = async_tensor_h2d(update_indices, torch.long,
|
|
self.runner.device,
|
|
self.runner.pin_memory)
|
|
model_input.input_tokens[update_indices] = last_sampler.sampled_token_ids_tensor[update_indices, 0]
|
|
|
|
if spec_step == SpecStepKind.SCORE_DECODE:
|
|
proposal_token_ids = get_proposal_token_ids()
|
|
shape = proposal_token_ids.shape
|
|
batch_size = shape[0]
|
|
proposal_len = shape[1]
|
|
update_indices = []
|
|
for i in range(batch_size):
|
|
for j in range(proposal_len):
|
|
update_indices.append(i * (proposal_len + 1) + j + 1)
|
|
|
|
update_indices = async_tensor_h2d(update_indices, torch.long,
|
|
self.runner.device,
|
|
self.runner.pin_memory)
|
|
model_input.input_tokens[update_indices] = proposal_token_ids.view(-1)
|
|
if spec_step == SpecStepKind.FIRST_PROPOSAL:
|
|
if last_step == SpecStepKind.PREFILL:# TODO: when last step is prefill, just update the input ids for last seqence_id onely.
|
|
pass
|
|
if last_step == SpecStepKind.SCORE_DECODE:# TODO: when last step is score decode, fix input ids、seq_lens、input_positions use accepte token ids
|
|
accept_token_ids, accept_seq_ids = get_accepted_token_ids()
|
|
|
|
chidren_req_ids = async_tensor_h2d(self.req_ids, torch.long,
|
|
self.runner.device,
|
|
self.runner.pin_memory)
|
|
grid = [1, 1, 1]
|
|
_update_input_tokens[grid](
|
|
accept_seq_ids, accept_seq_ids.shape[0],
|
|
accept_token_ids, accept_token_ids.shape[1],
|
|
chidren_req_ids, chidren_req_ids.shape[0],
|
|
model_input.input_tokens, model_input.input_tokens.shape[0],
|
|
model_input.input_positions,
|
|
model_input.seq_lens,
|
|
model_input.attn_metadata.seq_lens_tensor,
|
|
model_input.attn_metadata.seq_lens,
|
|
model_input.attn_metadata.slot_mapping,
|
|
model_input.attn_metadata.seq_start_loc,
|
|
model_input.attn_metadata.context_lens_tensor,
|
|
)
|
|
|
|
|
|
return model_input
|