Feature/support more sampling parameters

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:

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)

python/minisgl/engine/sample.py

@@ -7,14 +7,16 @@
 from minisgl.utils import is_sm90_supported, nvtx_annotate
 
 if TYPE_CHECKING:
-    from minisgl.core import Batch
+    from minisgl.core import Batch 
 
 
 @dataclass
 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)

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

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,
             ),
         )
     )