diff --git a/python/minisgl/core.py b/python/minisgl/core.py index be4d643e..32a3fc6a 100644 --- a/python/minisgl/core.py +++ b/python/minisgl/core.py @@ -1,4 +1,4 @@ -from __future__ import annotations +from __future__ import annotations from contextlib import contextmanager from dataclasses import dataclass, field @@ -19,6 +19,9 @@ class SamplingParams: top_p: float = 1.0 ignore_eos: bool = False max_tokens: int = 1024 + stop: List[str] = field(default_factory=list) + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 @property def is_greedy(self) -> bool: diff --git a/python/minisgl/engine/engine.py b/python/minisgl/engine/engine.py index ea29a96b..c31512b0 100644 --- a/python/minisgl/engine/engine.py +++ b/python/minisgl/engine/engine.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import timedelta -from typing import Any, Dict, NamedTuple, Tuple +from typing import Any, Dict, NamedTuple, Tuple import torch from minisgl.attention import create_attention_backend @@ -199,7 +199,7 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: for req in batch.reqs: req.complete_one() - next_tokens_gpu = self.sampler.sample(logits[: batch.size], args).to(torch.int32) + next_tokens_gpu = self.sampler.sample(logits[: batch.size], args, batch).to(torch.int32) next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) copy_done_event = torch.cuda.Event() copy_done_event.record(self.stream) diff --git a/python/minisgl/engine/sample.py b/python/minisgl/engine/sample.py index cb6c7ee7..4fa89446 100644 --- a/python/minisgl/engine/sample.py +++ b/python/minisgl/engine/sample.py @@ -7,7 +7,7 @@ from minisgl.utils import is_sm90_supported, nvtx_annotate if TYPE_CHECKING: - from minisgl.core import Batch + from minisgl.core import Batch @dataclass @@ -15,6 +15,8 @@ class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None + presence_penalties: torch.Tensor | None = None + frequency_penalties: torch.Tensor | None = None def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: @@ -60,16 +62,32 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_ks = [p.top_k if p.top_k >= 1 else self.vocab_size for p in params] top_ps = [min(max(p.top_p, MIN_P), 1.0) for p in params] temperatures = make_device_tensor(ts, torch.float32, self.device) - top_k, top_p = None, None + top_k, top_p, presence_penalties, frequency_penalties = None, None, None, None if any(k != self.vocab_size for k in top_ks): top_k = make_device_tensor(top_ks, torch.int32, self.device) if any(p < 1.0 for p in top_ps): top_p = make_device_tensor(top_ps, torch.float32, self.device) - return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p) + if any(p.presence_penalty != 0.0 for p in params): + presence_penalties = make_device_tensor( + [p.presence_penalty for p in params], torch.float32, self.device + ) + if any(p.frequency_penalty != 0.0 for p in params): + frequency_penalties = make_device_tensor( + [p.frequency_penalty for p in params], torch.float32, self.device + ) + return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p, presence_penalties=presence_penalties, frequency_penalties=frequency_penalties) @nvtx_annotate("Sampler") - def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: + def sample(self, logits: torch.Tensor, args: BatchSamplingArgs, batch: Batch) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): - if args.temperatures is None: # greedy sampling + if args.presence_penalties is not None: + for i, req in enumerate(batch.reqs): + unique_tokens = torch.unique(req.input_ids) + logits[i, unique_tokens] -= args.presence_penalties[i] + if args.frequency_penalties is not None: + for i, req in enumerate(batch.reqs): + token_counts = torch.bincount(req.input_ids, minlength=self.vocab_size).float() + logits[i] -= token_counts * args.frequency_penalties[i] + if args.temperatures is None: return torch.argmax(logits, dim=-1) return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) diff --git a/python/minisgl/scheduler/scheduler.py b/python/minisgl/scheduler/scheduler.py index d0c08d83..8235ba97 100644 --- a/python/minisgl/scheduler/scheduler.py +++ b/python/minisgl/scheduler/scheduler.py @@ -153,6 +153,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: finished = not req.can_decode if not req.sampling_params.ignore_eos: finished |= next_token == self.eos_token_id + if not finished and req.sampling_params.stop: + text = self.tokenizer.decode(req.input_ids.tolist()) + finished = any(text.endswith(s) for s in req.sampling_params.stop) reply.append(DetokenizeMsg(uid=req.uid, next_token=next_token, finished=finished)) # NOTE: overlap scheduling may make the request freed twice, skip second free diff --git a/python/minisgl/server/api_server.py b/python/minisgl/server/api_server.py index a229ae68..fe4a8171 100644 --- a/python/minisgl/server/api_server.py +++ b/python/minisgl/server/api_server.py @@ -158,7 +158,7 @@ async def stream_generate(self, uid: int): logger.debug("Finished streaming response for user %s", uid) async def stream_chat_completions(self, uid: int): - first_chunk = True + first_chunk = True async for ack in self.wait_for_ack(uid): delta = {} if first_chunk: @@ -261,7 +261,6 @@ async def v1_completions(req: OpenAICompletionRequest, request: Request): assert req.prompt is not None, "Either 'messages' or 'prompt' must be provided" prompt = req.prompt - # TODO: support more sampling parameters uid = state.new_user() await state.send_one( TokenizeMsg( @@ -273,6 +272,9 @@ async def v1_completions(req: OpenAICompletionRequest, request: Request): temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, + stop=req.stop, + presence_penalty=req.presence_penalty, + frequency_penalty=req.frequency_penalty, ), ) ) @@ -294,7 +296,6 @@ async def shell_completion(req: OpenAICompletionRequest): assert req.messages is not None, "Shell completion only supports chat-completions" prompt = [msg.model_dump() for msg in req.messages] - # TODO: support more sampling parameters uid = state.new_user() await state.send_one( TokenizeMsg( @@ -306,6 +307,9 @@ async def shell_completion(req: OpenAICompletionRequest): temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, + stop=req.stop, + presence_penalty=req.presence_penalty, + frequency_penalty=req.frequency_penalty, ), ) )