Extend sampler tests

serve/mlc_serve/engine/sampling_params.py

@@ -88,6 +88,8 @@ def __post_init__(self):
             self._verify_greedy_sampling()
         if not self.logprobs:
             self.top_logprobs = 0
+        if self.top_k == -1:
+            self.top_k = self.vocab_size
 
     def verify(self) -> None:
         if not -2.0 <= self.presence_penalty <= 2.0:
@@ -99,15 +101,15 @@ def verify(self) -> None:
                 "frequency_penalty must be in [-2, 2], got "
                 f"{self.frequency_penalty}."
             )
-        if self.temperature < 0.0:
+        if not 0.0 <= self.temperature <= 2.0:
             raise ValueError(
-                f"temperature must be non-negative, got {self.temperature}."
+                f"temperature must be in [0, 2], got {self.temperature}."
             )
         if not 0.0 < self.top_p <= 1.0:
             raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
 
         if not isinstance(self.top_k, int):
-            raise ValueError(f"top_k must be integer.")
+            raise TypeError(f"top_k must be integer.")
 
         if self.top_k < -1 or self.top_k == 0:
             raise ValueError(
@@ -132,6 +134,10 @@ def verify(self) -> None:
                 raise ValueError(
                     f"top_logprobs must be between 0 and {LOGPROB_TOP_K_MAX}, got {self.top_logprobs}."
                 )
+            if not isinstance(self.top_logprobs, int):
+                raise TypeError(
+                    "top_logprobs must be integer"
+                )
 
     def _verify_greedy_sampling(self) -> None:
         if self.top_p < 1.0 - _SAMPLING_EPS:

serve/mlc_serve/model/sampler.py

@@ -540,8 +540,8 @@ def _is_safe_to_sample(prob_like):
             assert sampling_state.sampling_params[batch_idx].logprobs
             top_k = sampling_state.sampling_params[batch_idx].top_logprobs
             logprob_infos[batch_idx] = RawLogprobsInfo(
-                current_token_id=next_token,
-                current_logprob=logprobs[batch_idx][next_token],
+                current_token_id=int(next_token),
+                current_logprob=float(logprobs[batch_idx][next_token]),
                 top_token_ids=top_tokens[idx][:top_k],
                 top_logprobs=top_logprobs[idx][:top_k],
             )

serve/tests/unittest/test_engine_with_samplers.py

@@ -11,7 +11,7 @@
 from mlc_serve.utils import get_default_mlc_serve_argparser, postproc_mlc_serve_args, create_mlc_engine
 import random
 from pydantic import BaseModel
-from typing import List
+from typing import List, Callable
 
 
 def create_request(
@@ -22,6 +22,7 @@ def create_request(
     pre_pen,
     max_tokens,
     stop,
+    num_sequences=1,
     ignore_eos=False,
     top_logprobs=0,
     logprobs=False,
@@ -41,6 +42,7 @@ def create_request(
             json_schema=json_schema,
         ),
         stopping_criteria=StoppingCriteria(max_tokens=max_tokens, stop_sequences=stop),
+        num_sequences=num_sequences,
         debug_options=DebugOptions(ignore_eos=ignore_eos),
     )
 
@@ -257,6 +259,46 @@ def _test_logprobs(
                 )
                 generated[int(res.request_id)] += seq.delta
 
+    # If temperature is increasing then difference between
+    # boundaries of range of top logprobs in response must decrease
+    temperatures = [0.2, 1.1, 2.0]
+    mean_bounds_diff = [0 for _ in range(num_requests * len(temperatures))]
+    for idx, temp in enumerate(temperatures):
+        requests = [
+            create_request(
+                idx=str(n),
+                prompt=random.choice(prompts),
+                temp=temp,
+                freq_pen=0,
+                pre_pen=0,
+                max_tokens=300,
+                stop=None,
+                ignore_eos=True,
+                logprobs=True,
+                top_logprobs=5
+            )
+            for n in range(num_requests)
+        ]
+        engine.add(requests)
+
+        while engine.has_pending_requests():
+            results = engine.step()
+            for res in results.outputs:
+                seq = res.sequences[0]
+                req = requests[int(res.request_id)]
+
+                if not seq.is_finished:
+                    mean_bounds_diff[idx * num_requests + int(res.request_id)] += \
+                        seq.logprob_info[0].top_logprobs[0].logprob \
+                        - seq.logprob_info[0].top_logprobs[4].logprob
+                else:
+                    mean_bounds_diff[idx * num_requests + int(res.request_id)] /= seq.num_generated_tokens
+
+    for num_req_batch in range(num_requests):
+        for idx in range(1, len(temperatures)):
+            assert mean_bounds_diff[idx * num_requests + num_req_batch] < \
+                   mean_bounds_diff[(idx - 1) * num_requests + num_req_batch]
+
 
 def _test_logprobs_mixed_requests(
     engine,
@@ -301,6 +343,48 @@ def _test_logprobs_mixed_requests(
                     assert len(seq.logprob_info) == 0
                 generated[int(res.request_id)] += seq.delta
 
+def _test_num_sequences(
+    engine,
+    num_requests=5,
+):
+    prompt = "Write a merge sort program in Python."
+    requests = []
+    num_sequences = [2 * i for i in range(1, num_requests + 1)]
+    for n, num_seq in enumerate(num_sequences):
+        requests.append(
+            create_request(
+                idx=str(n),
+                prompt=prompt,
+                temp=0.6,
+                freq_pen=0,
+                pre_pen=0,
+                stop=None,
+                max_tokens=300,
+                ignore_eos=False,
+                num_sequences=num_seq
+            )
+        )
+    engine.add(requests)
+
+    generated = [[""] * num_seq for _, num_seq in zip(range(num_requests), num_sequences)]
+    unique_sequences = [set() for _ in range(num_requests)]
+    while engine.has_pending_requests():
+        results = engine.step()
+        for idx, res in enumerate(results.outputs):
+            assert len(res.sequences) == num_sequences[idx]
+            for seq_id, seq in enumerate(res.sequences):
+                req_id = int(res.request_id)
+
+                if seq.delta:
+                    generated[int(req_id)][seq_id] += seq.delta
+
+                if seq.is_finished:
+                    unique_sequences[req_id].add(generated[req_id][seq_id])
+
+    for idx, response in enumerate(unique_sequences):
+        assert num_sequences[idx] == len(response)
+
+
 
 # These three models are used in _test_json_mode
 class France(BaseModel):
@@ -407,6 +491,7 @@ def _test_json_mode(
     # _test_stop(staging_engine)
     _test_logprobs(staging_engine)
     _test_logprobs_mixed_requests(staging_engine)
+    _test_num_sequences(staging_engine)
     _test_json_mode(staging_engine)
     # These tests are broken since we are now imposing no length limit
     # if max_tokens = None. The tests do not finish in a reasonable time.