chore: import upstream snapshot with attribution

python/sglang/srt/sampling/sampling_batch_info.py

@@ -0,0 +1,464 @@
+from __future__ import annotations
+
+import dataclasses
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+
+import sglang.srt.sampling.penaltylib as penaltylib
+from sglang.srt.runtime_context import get_server_args
+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] = None
+    rids_int: Optional[torch.Tensor] = None
+    bootstrap_room_ids_int: Optional[torch.Tensor] = None
+    vocab_mask: Optional[torch.Tensor] = None
+    apply_mask_func: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = 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
+
+    # Device
+    device: str = "cuda"
+
+    # Handle logit bias
+    logit_bias: Optional[torch.Tensor] = None
+
+    @classmethod
+    def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
+        global_server_args = get_server_args()
+        enable_deterministic = global_server_args.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 = (
+            global_server_args.enable_custom_logit_processor
+            and any(r.custom_logit_processor for r in reqs)  # check the flag first.
+        )  # then check the requests.
+
+        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,
+        )
+        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.vocab_mask = None
+            self.apply_mask_func = None
+            return
+
+        # Find a grammar from the list
+        first_grammar = next(grammar for grammar in self.grammars if grammar)
+
+        # TODO(lianmin): Maybe we can reuse the existing mask?
+        self.vocab_mask = first_grammar.allocate_vocab_mask(
+            vocab_size=self.vocab_size,
+            batch_size=len(self.temperatures),
+            device=self.device,
+        )
+        self.apply_mask_func = (
+            first_grammar.apply_vocab_mask
+        )  # force to use static method
+
+        # Apply the mask
+        for i, grammar in enumerate(self.grammars):
+            if grammar and not grammar.finished and not grammar.is_terminated():
+                grammar.fill_vocab_mask(self.vocab_mask, i)
+
+        # Move the mask to the device if needed
+        self.vocab_mask = first_grammar.move_vocab_mask(self.vocab_mask, self.device)
+
+    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.vocab_mask is not None:
+            self.apply_mask_func(logits=logits, vocab_mask=self.vocab_mask)
+
+        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]
+
+        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
+
+        # 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, len(self), len(other), self.device, 0.0
+        )
+
+        # 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])