496 lines
18 KiB
Python
496 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import logging
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
|
|
import sglang.srt.sampling.penaltylib as penaltylib
|
|
from sglang.srt.constrained.base_grammar_backend import (
|
|
BaseGrammarObject,
|
|
GrammarMask,
|
|
GrammarRow,
|
|
)
|
|
from sglang.srt.runtime_context import get_exec
|
|
from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
|
|
from sglang.srt.sampling.penaltylib.repetition_penalty import apply_scaling_penalties
|
|
from sglang.srt.sampling.sampling_params import TOP_K_ALL
|
|
from sglang.srt.utils.common import is_pin_memory_available
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.schedule_batch import ScheduleBatch
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SamplingBatchInfo:
|
|
# Basic batched sampling params
|
|
temperatures: torch.Tensor
|
|
top_ps: torch.Tensor
|
|
top_ks: torch.Tensor
|
|
min_ps: torch.Tensor
|
|
|
|
# Whether all requests use greedy sampling
|
|
is_all_greedy: bool
|
|
|
|
is_any_greedy: bool
|
|
|
|
# Whether any requests use top_p sampling
|
|
need_top_p_sampling: bool
|
|
|
|
# Whether any requests use top_k sampling
|
|
need_top_k_sampling: bool
|
|
|
|
# Whether any request needs min_p sampling
|
|
need_min_p_sampling: bool
|
|
|
|
# Masking tensors for grammar-guided structured outputs
|
|
vocab_size: int
|
|
grammars: Optional[List[Optional[BaseGrammarObject]]] = None
|
|
rids_int: Optional[torch.Tensor] = None
|
|
bootstrap_room_ids_int: Optional[torch.Tensor] = None
|
|
grammar_mask: Optional[GrammarMask] = None
|
|
|
|
# Penalizer
|
|
penalizer_orchestrator: Optional[penaltylib.BatchedPenalizerOrchestrator] = None
|
|
acc_additive_penalties: Optional[torch.Tensor] = None # Used in the overlap mode
|
|
acc_scaling_penalties: Optional[torch.Tensor] = (
|
|
None # Used in the overlap mode for repetition penalty
|
|
)
|
|
|
|
# Whether any request has custom logit processor
|
|
has_custom_logit_processor: bool = False
|
|
# Custom parameters
|
|
custom_params: Optional[List[Optional[Dict[str, Any]]]] = None
|
|
# Custom logit processor
|
|
custom_logit_processor: Optional[
|
|
Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]]
|
|
] = None
|
|
|
|
# Used for deterministic sampling
|
|
sampling_seed: Optional[torch.Tensor] = None
|
|
|
|
# Per-request flag for returning sparse sampling support metadata.
|
|
return_sampling_masks: Optional[List[bool]] = None
|
|
sampling_mask_max_top_k: int = 0
|
|
|
|
# Device
|
|
device: str = "cuda"
|
|
|
|
# Handle logit bias
|
|
logit_bias: Optional[torch.Tensor] = None
|
|
|
|
@classmethod
|
|
def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
|
|
enable_deterministic = get_exec().deterministic.enable_deterministic_inference
|
|
|
|
reqs = batch.reqs
|
|
device = batch.device
|
|
_pin = is_pin_memory_available(device)
|
|
temperatures = (
|
|
torch.tensor(
|
|
[r.sampling_params.temperature for r in reqs],
|
|
dtype=torch.float,
|
|
pin_memory=_pin,
|
|
)
|
|
.to(device, non_blocking=True)
|
|
.view(-1, 1)
|
|
)
|
|
top_ps = torch.tensor(
|
|
[r.sampling_params.top_p for r in reqs],
|
|
dtype=torch.float,
|
|
pin_memory=_pin,
|
|
).to(device, non_blocking=True)
|
|
top_ks = torch.tensor(
|
|
[r.sampling_params.top_k for r in reqs],
|
|
dtype=torch.int32,
|
|
pin_memory=_pin,
|
|
).to(device, non_blocking=True)
|
|
min_ps = torch.tensor(
|
|
[r.sampling_params.min_p for r in reqs],
|
|
dtype=torch.float,
|
|
pin_memory=_pin,
|
|
).to(device, non_blocking=True)
|
|
sampling_seed = (
|
|
torch.tensor(
|
|
[
|
|
(
|
|
r.sampling_params.sampling_seed
|
|
if r.sampling_params.sampling_seed is not None
|
|
else 42
|
|
)
|
|
for r in reqs
|
|
],
|
|
dtype=torch.int64,
|
|
pin_memory=_pin,
|
|
).to(device, non_blocking=True)
|
|
if enable_deterministic
|
|
else None
|
|
)
|
|
|
|
logit_bias = None
|
|
if any(r.sampling_params.logit_bias is not None for r in reqs):
|
|
logit_bias = torch.zeros(len(reqs), vocab_size, device=device)
|
|
for i, r in enumerate(reqs):
|
|
if r.sampling_params.logit_bias is not None:
|
|
for key, value in r.sampling_params.logit_bias.items():
|
|
logit_bias[i, int(key)] = value
|
|
|
|
# Check if any request has custom logit processor
|
|
has_custom_logit_processor = (
|
|
get_exec().features.enable_custom_logit_processor
|
|
and any(r.custom_logit_processor for r in reqs) # check the flag first.
|
|
) # then check the requests.
|
|
return_sampling_masks = [r.return_sampling_mask for r in reqs]
|
|
sampling_mask_max_top_k = max(
|
|
(r.sampling_params.top_k for r in reqs if r.return_sampling_mask),
|
|
default=0,
|
|
)
|
|
|
|
if has_custom_logit_processor:
|
|
# Merge the same type of custom logit processors together
|
|
processor_dict = {}
|
|
for i, r in enumerate(reqs):
|
|
if r.custom_logit_processor is None:
|
|
continue
|
|
processor_str = r.custom_logit_processor
|
|
if processor_str not in processor_dict:
|
|
processor_dict[processor_str] = []
|
|
processor_dict[processor_str].append(i)
|
|
|
|
merged_custom_logit_processor = {
|
|
hash(processor_str): (
|
|
# The deserialized custom logit processor object
|
|
CustomLogitProcessor.from_str(processor_str),
|
|
# The mask tensor for the requests that use this custom logit processor
|
|
torch.zeros(len(reqs), dtype=torch.bool)
|
|
.scatter_(0, torch.tensor(true_indices), True)
|
|
.to(device, non_blocking=True),
|
|
)
|
|
for processor_str, true_indices in processor_dict.items()
|
|
}
|
|
custom_params = [r.sampling_params.custom_params for r in reqs]
|
|
else:
|
|
merged_custom_logit_processor = None
|
|
custom_params = None
|
|
|
|
# Each penalizers will do nothing if they evaluate themselves as not required by looking at
|
|
# the sampling_params of the requests (See {_is_required()} of each penalizers). So this
|
|
# should not add hefty computation overhead other than simple checks.
|
|
#
|
|
# While we can choose not to even create the class instances if they are not required, this
|
|
# could add additional complexity to the {ScheduleBatch} class, especially we need to
|
|
# handle {filter_batch()} and {merge_batch()} cases as well.
|
|
penalizer_orchestrator = penaltylib.BatchedPenalizerOrchestrator(
|
|
vocab_size=vocab_size,
|
|
batch=batch,
|
|
penalizers={
|
|
penaltylib.BatchedFrequencyPenalizer,
|
|
penaltylib.BatchedMinNewTokensPenalizer,
|
|
penaltylib.BatchedPresencePenalizer,
|
|
penaltylib.BatchedRepetitionPenalizer,
|
|
},
|
|
)
|
|
|
|
ret = cls(
|
|
temperatures=temperatures,
|
|
top_ps=top_ps,
|
|
top_ks=top_ks,
|
|
min_ps=min_ps,
|
|
sampling_seed=sampling_seed,
|
|
is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs),
|
|
is_any_greedy=any(r.sampling_params.top_k <= 1 for r in reqs),
|
|
need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs),
|
|
need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs),
|
|
need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs),
|
|
vocab_size=vocab_size,
|
|
penalizer_orchestrator=penalizer_orchestrator,
|
|
has_custom_logit_processor=has_custom_logit_processor,
|
|
custom_params=custom_params,
|
|
custom_logit_processor=merged_custom_logit_processor,
|
|
device=device,
|
|
logit_bias=logit_bias,
|
|
return_sampling_masks=return_sampling_masks,
|
|
sampling_mask_max_top_k=sampling_mask_max_top_k,
|
|
)
|
|
ret.adjusted_from_schedule_batch(batch, vocab_size)
|
|
return ret
|
|
|
|
# placeholder for override
|
|
def adjusted_from_schedule_batch(self, batch: ScheduleBatch, vocab_size: int):
|
|
pass
|
|
|
|
# placeholder for override
|
|
def adjusted_merge_batch(self, other: SamplingBatchInfo):
|
|
pass
|
|
|
|
# placeholder for override
|
|
def adjusted_filter_batch(
|
|
self, keep_indices: List[int], keep_indices_device: torch.Tensor
|
|
):
|
|
pass
|
|
|
|
def __len__(self):
|
|
return len(self.temperatures)
|
|
|
|
def update_regex_vocab_mask(self):
|
|
if not self.grammars:
|
|
self.grammar_mask = None
|
|
return
|
|
|
|
# Find a grammar from the list
|
|
first_grammar = next(grammar for grammar in self.grammars if grammar)
|
|
|
|
vocab_mask = first_grammar.allocate_vocab_mask(
|
|
vocab_size=self.vocab_size,
|
|
batch_size=len(self.temperatures),
|
|
device=self.device,
|
|
)
|
|
|
|
# Rows omitted here (finished / terminated / non-grammar requests) retain
|
|
# the freshly allocated buffer's unconstrained value.
|
|
entries = [
|
|
GrammarRow(row=row, grammar=grammar)
|
|
for row, grammar in enumerate(self.grammars)
|
|
if grammar and not grammar.finished and not grammar.is_terminated()
|
|
]
|
|
first_grammar.fill_vocab_mask_batched(entries, vocab_mask)
|
|
|
|
# Move the mask to the device if needed
|
|
vocab_mask = first_grammar.move_vocab_mask(vocab_mask, self.device)
|
|
self.grammar_mask = GrammarMask(first_grammar, vocab_mask)
|
|
|
|
def update_penalties(self):
|
|
if self.penalizer_orchestrator.is_required:
|
|
self.acc_additive_penalties = torch.zeros(
|
|
(len(self.temperatures), self.vocab_size),
|
|
dtype=torch.float32,
|
|
device=self.temperatures.device,
|
|
)
|
|
self.penalizer_orchestrator.accumulate_additive_penalties(
|
|
self.acc_additive_penalties
|
|
)
|
|
self.acc_scaling_penalties = (
|
|
self.penalizer_orchestrator.accumulate_scaling_penalties()
|
|
)
|
|
else:
|
|
self.acc_additive_penalties = None
|
|
self.acc_scaling_penalties = None
|
|
|
|
def apply_logits_bias(self, logits: torch.Tensor):
|
|
if self.acc_additive_penalties is not None:
|
|
# Used in the overlap mode
|
|
logits.add_(self.acc_additive_penalties)
|
|
|
|
if self.acc_scaling_penalties is not None:
|
|
# Used in the overlap mode
|
|
apply_scaling_penalties(logits, self.acc_scaling_penalties)
|
|
|
|
if self.penalizer_orchestrator and self.penalizer_orchestrator.is_required:
|
|
# Used in the non-overlap mode
|
|
self.penalizer_orchestrator.apply(logits)
|
|
|
|
if self.grammar_mask is not None:
|
|
self.grammar_mask.apply(logits)
|
|
|
|
if self.logit_bias is not None:
|
|
logits.add_(self.logit_bias)
|
|
|
|
def filter_batch(self, keep_indices: List[int], keep_indices_device: torch.Tensor):
|
|
self.penalizer_orchestrator.filter(keep_indices_device)
|
|
|
|
if self.has_custom_logit_processor:
|
|
self._filter_batch_custom_logit_processor(keep_indices, keep_indices_device)
|
|
|
|
for item in [
|
|
"temperatures",
|
|
"top_ps",
|
|
"top_ks",
|
|
"min_ps",
|
|
"sampling_seed",
|
|
]:
|
|
value = getattr(self, item, None)
|
|
if value is not None:
|
|
setattr(self, item, value[keep_indices_device])
|
|
|
|
if self.logit_bias is not None:
|
|
self.logit_bias = self.logit_bias[keep_indices_device]
|
|
if self.return_sampling_masks is not None:
|
|
self.return_sampling_masks = [
|
|
self.return_sampling_masks[i] for i in keep_indices
|
|
]
|
|
|
|
self.adjusted_filter_batch(keep_indices, keep_indices_device)
|
|
|
|
def _filter_batch_custom_logit_processor(
|
|
self, keep_indices: List[int], keep_indices_device: torch.Tensor
|
|
):
|
|
"""Filter the custom logit processor and custom params"""
|
|
self.custom_logit_processor = {
|
|
k: (p, mask[keep_indices_device])
|
|
for k, (p, mask) in self.custom_logit_processor.items()
|
|
if torch.any(
|
|
mask[keep_indices_device]
|
|
) # ignore the custom logit processor whose mask is all False
|
|
}
|
|
self.custom_params = [self.custom_params[i] for i in keep_indices]
|
|
|
|
# If the custom logit processor is an empty dict, set the flag to False,
|
|
# and set the custom logit processor and custom params to None.
|
|
if len(self.custom_logit_processor) == 0:
|
|
self.custom_logit_processor = None
|
|
self.custom_params = None
|
|
self.has_custom_logit_processor = False
|
|
|
|
@staticmethod
|
|
def merge_custom_logit_processor(
|
|
lhs: Optional[Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]]],
|
|
rhs: Optional[Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]]],
|
|
bs1: int,
|
|
bs2: int,
|
|
device: str,
|
|
):
|
|
if lhs is None and rhs is None:
|
|
return None
|
|
lhs, rhs = lhs or {}, rhs or {}
|
|
|
|
keys = set(lhs.keys()).union(set(rhs.keys()))
|
|
merged_dict = {}
|
|
|
|
for k in keys:
|
|
# Get the logit processor object
|
|
processor = lhs[k][0] if k in lhs else rhs[k][0]
|
|
# Get and merge the mask tensors from the two dicts
|
|
left_mask = (
|
|
lhs[k][1]
|
|
if k in lhs
|
|
else torch.zeros(bs1, dtype=torch.bool, device=device)
|
|
)
|
|
right_mask = (
|
|
rhs[k][1]
|
|
if k in rhs
|
|
else torch.zeros(bs2, dtype=torch.bool, device=device)
|
|
)
|
|
merged_dict[k] = (processor, torch.cat([left_mask, right_mask]))
|
|
|
|
assert merged_dict[k][1].shape[0] == bs1 + bs2, (
|
|
f"The batch size of merged mask ({merged_dict[k][1].shape[0]}) does not match "
|
|
f"the sum of the batch sizes of the two masks ({bs1 + bs2})"
|
|
f"\n{left_mask=}\n{right_mask=}\n{bs1=}\n{bs2=}"
|
|
f"\n{lhs=}\n{rhs=}"
|
|
)
|
|
|
|
return merged_dict
|
|
|
|
def merge_batch(self, other: SamplingBatchInfo):
|
|
self.penalizer_orchestrator.merge(other.penalizer_orchestrator)
|
|
|
|
# Merge the custom logit processors and custom params lists
|
|
if self.has_custom_logit_processor or other.has_custom_logit_processor:
|
|
# Merge the custom logit processors
|
|
self.custom_logit_processor = (
|
|
SamplingBatchInfo.merge_custom_logit_processor(
|
|
self.custom_logit_processor,
|
|
other.custom_logit_processor,
|
|
len(self),
|
|
len(other),
|
|
self.device,
|
|
)
|
|
)
|
|
# Merge the custom params lists
|
|
self.custom_params = self.custom_params or [None] * len(self)
|
|
other.custom_params = other.custom_params or [None] * len(other)
|
|
self.custom_params.extend(other.custom_params)
|
|
|
|
# Set the flag to True if any of the two has custom logit processor
|
|
self.has_custom_logit_processor = True
|
|
|
|
self_len = len(self)
|
|
other_len = len(other)
|
|
|
|
# Merge logit bias - note this has to come before the temperatures tensor update! Otherwise will cause crashes.
|
|
# See note below on len(self) and len(other).
|
|
self.logit_bias = merge_bias_tensor(
|
|
self.logit_bias, other.logit_bias, self_len, other_len, self.device, 0.0
|
|
)
|
|
if (
|
|
self.return_sampling_masks is not None
|
|
or other.return_sampling_masks is not None
|
|
):
|
|
self.return_sampling_masks = (
|
|
self.return_sampling_masks or [False] * self_len
|
|
) + (other.return_sampling_masks or [False] * other_len)
|
|
self.sampling_mask_max_top_k = max(
|
|
self.sampling_mask_max_top_k, other.sampling_mask_max_top_k
|
|
)
|
|
|
|
# Note: because the __len()__ operator is defined on the temperatures tensor,
|
|
# please make sure any merge operation with len(self) or len(other) is done before
|
|
# the merge operation of the temperatures tensor below.
|
|
for item in [
|
|
"temperatures",
|
|
"top_ps",
|
|
"top_ks",
|
|
"min_ps",
|
|
"sampling_seed",
|
|
]:
|
|
self_val = getattr(self, item, None)
|
|
other_val = getattr(other, item, None)
|
|
if self_val is not None and other_val is not None:
|
|
setattr(self, item, torch.cat([self_val, other_val]))
|
|
|
|
self.is_all_greedy &= other.is_all_greedy
|
|
self.is_any_greedy |= other.is_any_greedy
|
|
self.need_top_p_sampling |= other.need_top_p_sampling
|
|
self.need_top_k_sampling |= other.need_top_k_sampling
|
|
self.need_min_p_sampling |= other.need_min_p_sampling
|
|
|
|
self.adjusted_merge_batch(other)
|
|
|
|
def copy_for_forward(self):
|
|
# Accumulate the penalty into a pre-allocated buffer to get rid of the dependency of `penalizer_orchestrator` later
|
|
self.update_penalties()
|
|
return dataclasses.replace(self, penalizer_orchestrator=None)
|
|
|
|
|
|
def merge_bias_tensor(
|
|
lhs: Optional[torch.Tensor],
|
|
rhs: Optional[torch.Tensor],
|
|
bs1: int,
|
|
bs2: int,
|
|
device: str,
|
|
default: float,
|
|
):
|
|
"""Merge two bias tensors for batch merging.
|
|
|
|
Args:
|
|
lhs: Left-hand side tensor
|
|
rhs: Right-hand side tensor
|
|
bs1: Batch size of left-hand side tensor
|
|
bs2: Batch size of right-hand side tensor
|
|
device: Device to place the merged tensor on
|
|
default: Default value for missing tensor elements
|
|
|
|
Returns:
|
|
Merged tensor or None if both inputs are None
|
|
"""
|
|
if lhs is None and rhs is None:
|
|
return None
|
|
|
|
if lhs is not None and rhs is not None:
|
|
return torch.cat([lhs, rhs])
|
|
else:
|
|
if lhs is not None:
|
|
shape, dtype = lhs.shape[1:], lhs.dtype
|
|
else:
|
|
shape, dtype = rhs.shape[1:], rhs.dtype
|
|
|
|
if lhs is None:
|
|
lhs = torch.empty((bs1, *shape), device=device, dtype=dtype).fill_(default)
|
|
if rhs is None:
|
|
rhs = torch.empty((bs2, *shape), device=device, dtype=dtype).fill_(default)
|
|
return torch.cat([lhs, rhs])
|