chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support min-p sampling (#1167)

Browse Source
This commit is contained in:
intervitens
2024-08-22 01:49:32 +03:00
committed by GitHub
parent d6aeb9fa15
commit 068e9eae55
7 changed files with 58 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
python/sglang/srt/managers/schedule_batch.py
View File

@@ -21,7 +21,12 @@ from typing import List, Optional, Union
import torch
import torch.distributed as dist
from flashinfer.sampling import top_k_top_p_sampling_from_probs
from flashinfer.sampling import (
min_p_sampling_from_probs,
top_k_renorm_prob,
top_k_top_p_sampling_from_probs,
top_p_renorm_prob,
)
from vllm.distributed import get_tensor_model_parallel_group
import sglang.srt.sampling.penaltylib as penaltylib
@@ -339,6 +344,7 @@ class ScheduleBatch:
temperatures: torch.Tensor = None
top_ps: torch.Tensor = None
top_ks: torch.Tensor = None
min_ps: torch.Tensor = None
penalizer_orchestrator: penaltylib.BatchedPenalizerOrchestrator = None
logit_bias: torch.Tensor = None
@@ -403,6 +409,9 @@ class ScheduleBatch:
self.top_ks = torch.tensor(
[r.sampling_params.top_k for r in reqs], dtype=torch.int, device=device
)
self.min_ps = torch.tensor(
[r.sampling_params.min_p for r in reqs], dtype=torch.float, device=device
)
# 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
@@ -701,6 +710,7 @@ class ScheduleBatch:
"temperatures",
"top_ps",
"top_ks",
"min_ps",
"logit_bias",
]:
self_val = getattr(self, item, None)
@@ -730,6 +740,7 @@ class ScheduleBatch:
"temperatures",
"top_ps",
"top_ks",
"min_ps",
]:
self_val = getattr(self, item, None)
other_val = getattr(other, item, None)
@@ -780,13 +791,20 @@ class ScheduleBatch:
uniform_samples = torch.rand(
(max_top_k_round, batch_size), device=probs.device
)
batch_next_token_ids, success = top_k_top_p_sampling_from_probs(
probs, uniform_samples, self.top_ks, self.top_ps
)
if self.min_ps.any():
probs = top_k_renorm_prob(probs, self.top_ks)
probs = top_p_renorm_prob(probs, self.top_ps)
batch_next_token_ids, success = min_p_sampling_from_probs(
probs, uniform_samples, self.min_ps
)
else:
batch_next_token_ids, success = top_k_top_p_sampling_from_probs(
probs, uniform_samples, self.top_ks, self.top_ps
)
else:
# Here we provide a slower fallback implementation.
batch_next_token_ids, success = top_k_top_p_sampling_from_probs_torch(
probs, self.top_ks, self.top_ps
batch_next_token_ids, success = top_k_top_p_min_p_sampling_from_probs_torch(
probs, self.top_ks, self.top_ps, self.min_ps
)
if not torch.all(success):
@@ -810,17 +828,22 @@ class ScheduleBatch:
return batch_next_token_ids
def top_k_top_p_sampling_from_probs_torch(
probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor
def top_k_top_p_min_p_sampling_from_probs_torch(
probs: torch.Tensor,
top_ks: torch.Tensor,
top_ps: torch.Tensor,
min_ps: torch.Tensor,
):
"""A top-k and top-k sampling implementation with native pytorch operations."""
"""A top-k, top-p and min-p sampling implementation with native pytorch operations."""
probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
probs_sum = torch.cumsum(probs_sort, dim=-1)
min_p_thresholds = probs_sort[:, 0] * min_ps
probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0
probs_sort[
torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1)
>= top_ks.view(-1, 1)
] = 0.0
probs_sort[probs_sort < min_p_thresholds.view(-1, 1)] = 0.0
probs_sort.div_(probs_sort.max(dim=-1, keepdim=True)[0])
try:
sampled_index = torch.multinomial(probs_sort, num_samples=1)