453 lines
17 KiB
Python
453 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
import logging
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
import triton
|
|
|
|
from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
|
|
from sglang.srt.server_args import get_global_server_args
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch.nn.functional as F
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
|
|
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
|
from sglang.srt.layers.sampler import apply_custom_logit_processor
|
|
from sglang.srt.managers.schedule_batch import ScheduleBatch
|
|
from sglang.srt.mem_cache.common import (
|
|
alloc_paged_token_slots_extend,
|
|
alloc_token_slots,
|
|
get_last_loc,
|
|
)
|
|
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
|
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
|
|
from sglang.srt.speculative.spec_utils import (
|
|
TREE_SPEC_KERNEL_AVAILABLE,
|
|
assign_req_to_token_pool,
|
|
get_src_tgt_cache_loc,
|
|
get_target_cache_loc,
|
|
)
|
|
from sglang.srt.utils import is_cuda, is_hip, next_power_of_2
|
|
|
|
if is_cuda():
|
|
from sgl_kernel import (
|
|
top_k_renorm_prob,
|
|
top_p_renorm_prob,
|
|
tree_speculative_sampling_target_only,
|
|
verify_tree_greedy,
|
|
)
|
|
elif is_hip():
|
|
from sgl_kernel import verify_tree_greedy
|
|
|
|
|
|
@dataclass
|
|
class NgramVerifyInput(SpecInput):
|
|
def __init__(
|
|
self,
|
|
draft_token: torch.Tensor,
|
|
tree_mask: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
retrive_index: torch.Tensor,
|
|
retrive_next_token: torch.Tensor,
|
|
retrive_next_sibling: torch.Tensor,
|
|
draft_token_num: int,
|
|
grammar: BaseGrammarObject = None,
|
|
):
|
|
super().__init__(SpecInputType.NGRAM_VERIFY)
|
|
self.draft_token = draft_token
|
|
self.custom_mask = tree_mask
|
|
self.positions = positions
|
|
self.retrive_index = retrive_index
|
|
self.retrive_next_token = retrive_next_token
|
|
self.retrive_next_sibling = retrive_next_sibling
|
|
self.draft_token_num = draft_token_num
|
|
self.device = self.custom_mask.device
|
|
self.grammar = grammar
|
|
|
|
def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
|
|
return self.draft_token_num, self.draft_token_num
|
|
|
|
def prepare_for_verify(self, batch: ScheduleBatch, page_size: int):
|
|
if batch.forward_mode.is_idle():
|
|
return
|
|
|
|
batch.input_ids = self.draft_token
|
|
|
|
if page_size == 1:
|
|
batch.out_cache_loc = alloc_token_slots(
|
|
batch.tree_cache,
|
|
len(batch.input_ids),
|
|
)
|
|
end_offset = batch.seq_lens + self.draft_token_num
|
|
else:
|
|
# TODO(lsyin): add prefix lens cpu here to support page size > 1
|
|
prefix_lens = batch.seq_lens
|
|
prefix_lens_cpu = batch.seq_lens_cpu
|
|
end_offset = prefix_lens + self.draft_token_num
|
|
end_offset_cpu = prefix_lens_cpu + self.draft_token_num
|
|
last_loc = get_last_loc(
|
|
batch.req_to_token_pool.req_to_token,
|
|
batch.req_pool_indices,
|
|
prefix_lens,
|
|
)
|
|
batch.out_cache_loc = alloc_paged_token_slots_extend(
|
|
batch.tree_cache,
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
end_offset,
|
|
end_offset_cpu,
|
|
last_loc,
|
|
len(batch.input_ids),
|
|
)
|
|
self.last_loc = last_loc
|
|
|
|
bs = batch.batch_size()
|
|
assign_req_to_token_pool[(bs,)](
|
|
batch.req_pool_indices,
|
|
batch.req_to_token_pool.req_to_token,
|
|
batch.seq_lens,
|
|
end_offset,
|
|
batch.out_cache_loc,
|
|
batch.req_to_token_pool.req_to_token.shape[1],
|
|
triton.next_power_of_2(bs),
|
|
)
|
|
|
|
def generate_attn_arg_prefill(
|
|
self,
|
|
req_pool_indices: torch.Tensor,
|
|
paged_kernel_lens: torch.Tensor,
|
|
paged_kernel_lens_sum: int,
|
|
req_to_token: torch.Tensor,
|
|
):
|
|
bs = len(req_pool_indices)
|
|
|
|
cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
|
|
|
|
paged_kernel_lens = paged_kernel_lens + self.draft_token_num
|
|
cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
|
|
|
|
self.qo_indptr = (
|
|
torch.arange(0, bs + 1, dtype=torch.int32, device=self.device)
|
|
* self.draft_token_num
|
|
)
|
|
|
|
kv_indices = torch.empty(
|
|
cum_kv_seq_len[-1], dtype=torch.int32, device=self.device
|
|
)
|
|
|
|
create_flashinfer_kv_indices_triton[(bs,)](
|
|
req_to_token,
|
|
req_pool_indices,
|
|
paged_kernel_lens,
|
|
cum_kv_seq_len,
|
|
None,
|
|
kv_indices,
|
|
req_to_token.size(1),
|
|
)
|
|
return kv_indices, cum_kv_seq_len, self.qo_indptr, self.custom_mask
|
|
|
|
def _fill_requests(
|
|
self,
|
|
batch: ScheduleBatch,
|
|
logits_output: torch.Tensor,
|
|
):
|
|
accept_index_cpu = self.accepted_indices.tolist()
|
|
predict_cpu = self.predict.tolist()
|
|
has_finished = False
|
|
|
|
# Iterate every accepted token and check if req has finished after append the token
|
|
# should be checked BEFORE free kv cache slots
|
|
for i, (req, accept_index_row) in enumerate(zip(batch.reqs, accept_index_cpu)):
|
|
for j, idx in enumerate(accept_index_row):
|
|
if idx == -1:
|
|
break
|
|
id = predict_cpu[idx]
|
|
req.output_ids.append(id)
|
|
req.check_finished()
|
|
if req.finished():
|
|
has_finished = True
|
|
# set all tokens after finished token to -1 and break
|
|
self.accepted_indices[i, j + 1 :] = -1
|
|
break
|
|
else:
|
|
if req.grammar is not None:
|
|
try:
|
|
req.grammar.accept_token(id)
|
|
except ValueError as e:
|
|
logger.info(
|
|
f"{i=}, {req=}\n"
|
|
f"{self.accepted_indices=}\n"
|
|
f"{self.predict=}\n"
|
|
)
|
|
raise e
|
|
req.spec_verify_ct += 1
|
|
accepted_draft_tokens = sum(1 for idx in accept_index_row if idx != -1) - 1
|
|
req.spec_accepted_tokens += accepted_draft_tokens
|
|
req.update_spec_acceptance_histogram(accepted_draft_tokens)
|
|
|
|
if has_finished:
|
|
self.accept_length = (self.accepted_indices != -1).sum(dim=1) - 1
|
|
self.accepted_indices = self.accepted_indices[self.accepted_indices != -1]
|
|
|
|
logits_output.next_token_logits = logits_output.next_token_logits[
|
|
self.accepted_indices
|
|
]
|
|
if logits_output.hidden_states:
|
|
logits_output.hidden_states = logits_output.hidden_states[
|
|
self.accepted_indices
|
|
]
|
|
self.verified_id = self.predict[self.accepted_indices]
|
|
|
|
def _free_cache(
|
|
self, batch: ScheduleBatch, page_size: int, accept_length_cpu: torch.Tensor
|
|
):
|
|
bs = batch.batch_size()
|
|
# Free the KV cache for unaccepted tokens
|
|
if page_size == 1:
|
|
# TODO: boolean array index leads to a device sync. Remove it.
|
|
evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool)
|
|
evict_mask[self.accepted_indices] = False
|
|
batch.token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
|
|
batch.out_cache_loc = batch.out_cache_loc[self.accepted_indices]
|
|
else:
|
|
# Shift the accepted tokens to the beginning.
|
|
# Only evict the last part
|
|
src_cache_loc, tgt_cache_loc, to_free_num_slots = get_src_tgt_cache_loc(
|
|
batch.seq_lens,
|
|
batch.out_cache_loc,
|
|
self.accepted_indices,
|
|
self.accept_length,
|
|
self.draft_token_num,
|
|
page_size,
|
|
)
|
|
to_free_slots = torch.empty(
|
|
(to_free_num_slots.sum().item(),),
|
|
dtype=torch.int64,
|
|
device=to_free_num_slots.device,
|
|
)
|
|
|
|
# out_cache_loc: [0 1 2, 3 4 5, 6 7 8]
|
|
# accept_index: [0 -1 2, 3 4 -1, 6 -1 -1]
|
|
# tgt_cache_loc: [0 1 , 3 4 , 6 ]
|
|
# to_free_slots: [ 2, 5, 7 8]
|
|
# to_free_slots also needs to be page-aligned without the first partial page
|
|
#
|
|
# split each row of out_cache_loc into two parts.
|
|
# 1. the first part goes to tgt_cache_loc. length = accept_length[i] + 1
|
|
# 2. the second part goes to to_free_slots.
|
|
get_target_cache_loc[(bs,)](
|
|
tgt_cache_loc,
|
|
to_free_slots,
|
|
self.accept_length,
|
|
to_free_num_slots,
|
|
batch.out_cache_loc,
|
|
self.draft_token_num,
|
|
next_power_of_2(self.draft_token_num),
|
|
next_power_of_2(bs),
|
|
)
|
|
|
|
# Free the kv cache
|
|
batch.token_to_kv_pool_allocator.free(to_free_slots)
|
|
|
|
# Copy the kv cache
|
|
batch.token_to_kv_pool_allocator.get_kvcache().move_kv_cache(
|
|
tgt_cache_loc, src_cache_loc
|
|
)
|
|
batch.out_cache_loc = tgt_cache_loc
|
|
|
|
accept_length_list = accept_length_cpu.tolist()
|
|
for i, req in enumerate(batch.reqs):
|
|
req.kv_committed_len += accept_length_list[i] + 1
|
|
req.kv_allocated_len = req.kv_committed_len
|
|
|
|
assign_req_to_token_pool[(bs,)](
|
|
batch.req_pool_indices,
|
|
batch.req_to_token_pool.req_to_token,
|
|
batch.seq_lens,
|
|
batch.seq_lens + self.accept_length + 1,
|
|
batch.out_cache_loc,
|
|
batch.req_to_token_pool.req_to_token.shape[1],
|
|
triton.next_power_of_2(bs),
|
|
)
|
|
|
|
def _greedy_verify(
|
|
self,
|
|
batch: ScheduleBatch,
|
|
logits_output: LogitsProcessorOutput,
|
|
):
|
|
bs = batch.batch_size()
|
|
target_predict = torch.argmax(logits_output.next_token_logits, dim=-1)
|
|
target_predict = target_predict.reshape(bs, self.draft_token_num)
|
|
|
|
candidates = self.draft_token.reshape(bs, self.draft_token_num)
|
|
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
|
|
predict_shape[-1] += 1
|
|
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
|
|
self.accepted_indices = torch.full(
|
|
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
|
|
)
|
|
self.accept_length = torch.empty((bs,), dtype=torch.int32, device=self.device)
|
|
|
|
verify_tree_greedy(
|
|
predicts=self.predict, # mutable
|
|
accept_index=self.accepted_indices, # mutable
|
|
accept_token_num=self.accept_length, # mutable
|
|
candidates=candidates,
|
|
retrive_index=self.retrive_index,
|
|
retrive_next_token=self.retrive_next_token,
|
|
retrive_next_sibling=self.retrive_next_sibling,
|
|
target_predict=target_predict,
|
|
)
|
|
|
|
def _sampling_verify(
|
|
self,
|
|
batch: ScheduleBatch,
|
|
logits_output: LogitsProcessorOutput,
|
|
sampling_info: SamplingBatchInfo,
|
|
):
|
|
bs = batch.batch_size()
|
|
candidates = self.draft_token.reshape(bs, self.draft_token_num)
|
|
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
|
|
predict_shape[-1] += 1
|
|
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
|
|
self.accepted_indices = torch.full(
|
|
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
|
|
)
|
|
self.accept_length = torch.empty((bs,), dtype=torch.int32, device=self.device)
|
|
# apply temperature and get target probs
|
|
expanded_temperature = torch.repeat_interleave(
|
|
sampling_info.temperatures, self.draft_token_num, dim=0
|
|
) # (bs * draft_token_num, 1)
|
|
|
|
target_probs = F.softmax(
|
|
logits_output.next_token_logits / expanded_temperature, dim=-1
|
|
) # (bs * draft_token_num, vocab_size)
|
|
|
|
# NOTE: The test shows that top_p_renorm_prob and top_k_renorm_prob are the key factors
|
|
# contributing to the poor performance of _sampling_verify.
|
|
target_probs = top_k_renorm_prob(
|
|
target_probs,
|
|
torch.repeat_interleave(sampling_info.top_ks, self.draft_token_num, dim=0),
|
|
) # (bs * draft_token_num, vocab_size)
|
|
|
|
if sampling_info.need_top_p_sampling:
|
|
# logger.info("Using top-p sampling in speculative decoding verification.")
|
|
target_probs = top_p_renorm_prob(
|
|
target_probs,
|
|
torch.repeat_interleave(
|
|
sampling_info.top_ps, self.draft_token_num, dim=0
|
|
),
|
|
)
|
|
|
|
target_probs = target_probs.reshape(bs, self.draft_token_num, -1)
|
|
draft_probs = torch.zeros(
|
|
target_probs.shape, dtype=torch.float32, device=self.device
|
|
)
|
|
|
|
# coins for rejection sampling
|
|
coins = torch.rand_like(candidates, dtype=torch.float32, device=self.device)
|
|
# coins for final sampling
|
|
coins_for_final_sampling = torch.rand(
|
|
(bs,), dtype=torch.float32, device=self.device
|
|
)
|
|
tree_speculative_sampling_target_only(
|
|
predicts=self.predict, # mutable
|
|
accept_index=self.accepted_indices, # mutable
|
|
accept_token_num=self.accept_length, # mutable
|
|
candidates=candidates.to(torch.int64),
|
|
retrive_index=self.retrive_index.to(torch.int64),
|
|
retrive_next_token=self.retrive_next_token.to(torch.int64),
|
|
retrive_next_sibling=self.retrive_next_sibling.to(torch.int64),
|
|
uniform_samples=coins,
|
|
uniform_samples_for_final_sampling=coins_for_final_sampling,
|
|
target_probs=target_probs,
|
|
draft_probs=draft_probs,
|
|
threshold_single=get_global_server_args().speculative_accept_threshold_single,
|
|
threshold_acc=get_global_server_args().speculative_accept_threshold_acc,
|
|
deterministic=True,
|
|
)
|
|
|
|
def verify(
|
|
self,
|
|
batch: ScheduleBatch,
|
|
logits_output: LogitsProcessorOutput,
|
|
page_size: int,
|
|
vocab_mask: Optional[torch.Tensor] = None, # For grammar
|
|
) -> torch.Tensor:
|
|
bs = self.retrive_index.shape[0]
|
|
sampling_info = batch.sampling_info
|
|
|
|
if bs != len(sampling_info):
|
|
sampling_info = copy.deepcopy(sampling_info)
|
|
# NOTE: retrive_index are the indices of the requests that are kept.
|
|
sampling_info.filter_batch(self.retrive_index.tolist(), self.retrive_index)
|
|
|
|
# Apply the custom logit processors if registered in the sampling info.
|
|
if sampling_info.has_custom_logit_processor:
|
|
apply_custom_logit_processor(
|
|
logits_output.next_token_logits,
|
|
sampling_info,
|
|
num_tokens_in_batch=self.draft_token_num,
|
|
)
|
|
|
|
# Apply penalty
|
|
if sampling_info.penalizer_orchestrator.is_required:
|
|
# This is a relaxed version of penalties for speculative decoding.
|
|
linear_penalty = torch.zeros(
|
|
(bs, logits_output.next_token_logits.shape[1]),
|
|
dtype=torch.float32,
|
|
device=self.device,
|
|
)
|
|
sampling_info.apply_logits_bias(linear_penalty)
|
|
logits_output.next_token_logits.add_(
|
|
torch.repeat_interleave(linear_penalty, self.draft_token_num, dim=0)
|
|
)
|
|
|
|
# Apply grammar mask
|
|
if vocab_mask is not None:
|
|
assert self.grammar is not None
|
|
self.grammar.apply_vocab_mask(
|
|
logits=logits_output.next_token_logits, vocab_mask=vocab_mask
|
|
)
|
|
|
|
# Sample tokens. Force greedy sampling on AMD
|
|
is_all_greedy = (
|
|
sampling_info.is_all_greedy or envs.SGLANG_NGRAM_FORCE_GREEDY_VERIFY.get()
|
|
)
|
|
if (not is_all_greedy) and (not TREE_SPEC_KERNEL_AVAILABLE):
|
|
logger.warning(
|
|
"Tree speculative sampling kernel unavailable (likely AMD/HIP build). "
|
|
"Falling back to greedy verification."
|
|
)
|
|
|
|
if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE:
|
|
self._greedy_verify(batch, logits_output)
|
|
else:
|
|
# NOTE: Compared with greedy_verify, the performance of _sampling_verify is relatively poor.
|
|
self._sampling_verify(batch, logits_output, sampling_info)
|
|
|
|
self._fill_requests(batch, logits_output)
|
|
|
|
accept_length_cpu = self.accept_length.cpu()
|
|
num_accepted_tokens = accept_length_cpu.sum().item()
|
|
|
|
self._free_cache(batch, page_size, accept_length_cpu)
|
|
|
|
batch.seq_lens.add_(self.accept_length + 1)
|
|
batch.seq_lens_cpu.add_(accept_length_cpu + 1)
|
|
|
|
return logits_output, self.verified_id, num_accepted_tokens
|
|
|
|
def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True):
|
|
pass
|
|
|
|
def merge_batch(self, spec_info: NgramVerifyInput):
|
|
pass
|