[Feature] support top_k_top_p sampling (#2753)

* support top_k_top_p sampling

* fix

* add api param

* add api para

* fix

* fix

* fix

* fix

* fix

* fix

* fix

fastdeploy/worker/gcu_model_runner.py

@@ -154,12 +154,29 @@ class GCUModelRunner(ModelRunnerBase):
                 -1].disaggregate_info["role"] == "prefill":
             os.environ['PREFILL_NODE_ONE_STEP_STOP'] = "1"
 
+        top_k_reqs = []
+        top_p_reqs = []
+        max_num_seqs = self.parallel_config.max_num_seqs
+        top_p_buffer = paddle.full([max_num_seqs, 1],
+                                    self.model_config.top_p,
+                                    dtype='float32')
+        top_k_buffer = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
+
         req_len = len(req_dicts)
         for i in range(req_len):
             request = req_dicts[i]
             idx = request.idx
             length = len(request.prompt_token_ids)
 
+            if sampling_params := request.sampling_params:
+                if sampling_params.top_p < 1:
+                    top_p_reqs.append(idx)
+                top_k = sampling_params.top_k
+                if top_k > 0:
+                    top_k_reqs.append(idx)
+
             prefill_tokens = []
             if (request.guided_json is not None
                     or request.guided_regex is not None
@@ -234,8 +251,8 @@ class GCUModelRunner(ModelRunnerBase):
                 request.eos_token_ids.append(request.eos_token_ids[0])
             self.share_inputs["eos_token_id"][:] = np.array(
                 request.eos_token_ids, dtype="int64").reshape(-1, 1)
-
-            self.share_inputs["top_p"][idx:idx + 1] = request.get("top_p", 0.7)
+            top_p_buffer[idx:idx + 1] = request.get("top_p", 1.0)
+            top_k_buffer[idx:idx + 1] = request.get("top_k", 0)
             self.share_inputs["temperature"][idx:idx + 1] = request.get(
                 "temperature", 0.95)
             self.share_inputs["penalty_score"][idx:idx + 1] = request.get(
@@ -286,6 +303,16 @@ class GCUModelRunner(ModelRunnerBase):
         if self.speculative_method in ["mtp"]:
             self.proposer.insert_prefill_inputs(req_dicts)
 
+        if len(top_k_reqs) == 0:
+            self.share_inputs["top_k"] = None
+        else:
+            self.share_inputs["top_k"] = top_k_buffer
+
+        if len(top_p_reqs) == 0:
+            self.share_inputs["top_p"] = None
+        else:
+            self.share_inputs["top_p"] = top_p_buffer
+
     def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int,
                               expected_decode_len: int):
         """ Set dummy prefill inputs to share_inputs """
@@ -340,8 +367,11 @@ class GCUModelRunner(ModelRunnerBase):
         self.share_inputs["eos_token_id"] = paddle.full(
             [self.parallel_config.eos_tokens_lens, 1], 0, dtype='int64')
         self.share_inputs["top_p"] = paddle.full([max_num_seqs, 1],
-                                                 self.model_config.top_p,
-                                                 dtype='float32')
+                                                self.model_config.top_p,
+                                                dtype='float32')
+        self.share_inputs["top_k"] = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
         self.share_inputs["temperature"] = paddle.full(
             [max_num_seqs, 1], self.model_config.temperature, dtype='float32')
         self.share_inputs["penalty_score"] = paddle.full(
@@ -563,6 +593,7 @@ class GCUModelRunner(ModelRunnerBase):
         self.sampling_metadata = SamplingMetadata(
             temperature=self.share_inputs["temperature"],
             top_p=self.share_inputs["top_p"],
+            top_k=self.share_inputs["top_k"],
             step_idx=self.share_inputs["step_idx"],
             pre_token_ids=self.share_inputs["pre_ids"],
             frequency_penalties=self.share_inputs["frequency_score"],

fastdeploy/worker/gpu_model_runner.py

@@ -161,6 +161,15 @@ class GPUModelRunner(ModelRunnerBase):
                 -1].disaggregate_info["role"] == "prefill":
             os.environ['PREFILL_NODE_ONE_STEP_STOP'] = "1"
 
+        top_k_reqs = []
+        top_p_reqs = []
+        max_num_seqs = self.parallel_config.max_num_seqs
+        top_p_buffer = paddle.full([max_num_seqs, 1],
+                                    self.model_config.top_p,
+                                    dtype='float32')
+        top_k_buffer = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
         req_len = len(req_dicts)
         for i in range(req_len):
             request = req_dicts[i]
@@ -168,6 +177,13 @@ class GPUModelRunner(ModelRunnerBase):
             length = len(request.prompt_token_ids)
             assert length > 0, "The prompt requested must not be empty."
 
+            if sampling_params := request.sampling_params:
+                if sampling_params.top_p < 1:
+                    top_p_reqs.append(idx)
+                top_k = sampling_params.top_k
+                if top_k > 0:
+                    top_k_reqs.append(idx)
+
             prefill_tokens = []
             if (request.guided_json is not None
                     or request.guided_regex is not None
@@ -242,8 +258,8 @@ class GPUModelRunner(ModelRunnerBase):
                 request.eos_token_ids.append(request.eos_token_ids[0])
             self.share_inputs["eos_token_id"][:] = np.array(
                 request.eos_token_ids, dtype="int64").reshape(-1, 1)
-
-            self.share_inputs["top_p"][idx:idx + 1] = request.get("top_p", 0.7)
+            top_p_buffer[idx:idx + 1] = request.get("top_p", 1.0)
+            top_k_buffer[idx:idx + 1] = request.get("top_k", 0)
             self.share_inputs["temperature"][idx:idx + 1] = request.get(
                 "temperature", 0.95)
             self.share_inputs["penalty_score"][idx:idx + 1] = request.get(
@@ -294,6 +310,16 @@ class GPUModelRunner(ModelRunnerBase):
         if self.speculative_method in ["mtp"]:
             self.proposer.insert_prefill_inputs(req_dicts)
 
+        if len(top_k_reqs) == 0:
+            self.share_inputs["top_k"] = None
+        else:
+            self.share_inputs["top_k"] = top_k_buffer
+
+        if len(top_p_reqs) == 0:
+            self.share_inputs["top_p"] = None
+        else:
+            self.share_inputs["top_p"] = top_p_buffer
+
     def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int,
                               expected_decode_len: int):
         """ Set dummy prefill inputs to share_inputs """
@@ -349,8 +375,11 @@ class GPUModelRunner(ModelRunnerBase):
         self.share_inputs["eos_token_id"] = paddle.full(
             [self.parallel_config.eos_tokens_lens, 1], 0, dtype='int64')
         self.share_inputs["top_p"] = paddle.full([max_num_seqs, 1],
-                                                 self.model_config.top_p,
-                                                 dtype='float32')
+                                                self.model_config.top_p,
+                                                dtype='float32')
+        self.share_inputs["top_k"] = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
         self.share_inputs["temperature"] = paddle.full(
             [max_num_seqs, 1], self.model_config.temperature, dtype='float32')
         self.share_inputs["penalty_score"] = paddle.full(
@@ -574,6 +603,7 @@ class GPUModelRunner(ModelRunnerBase):
         self.sampling_metadata = SamplingMetadata(
             temperature=self.share_inputs["temperature"],
             top_p=self.share_inputs["top_p"],
+            top_k=self.share_inputs["top_k"],
             step_idx=self.share_inputs["step_idx"],
             pre_token_ids=self.share_inputs["pre_ids"],
             frequency_penalties=self.share_inputs["frequency_score"],

fastdeploy/worker/iluvatar_model_runner.py

@@ -29,9 +29,8 @@ from fastdeploy.model_executor.layers.attention.base_attention_backend import \
     AttentionBackend
 from fastdeploy.model_executor.layers.rotary_embedding import get_rope
 from fastdeploy.model_executor.layers.sample.meta_data import SamplingMetadata
-from fastdeploy.model_executor.layers.sample.sampler import (Sampler,
-                                                             SpeculativeSampler
-                                                             )
+from fastdeploy.model_executor.layers.sample.sampler import (
+    Sampler, SpeculativeSampler)
 from fastdeploy.model_executor.model_loader import get_model_from_loader
 from fastdeploy.model_executor.ops.iluvatar import set_value_by_flags_and_idx
 from fastdeploy.model_executor.pre_and_post_process import (post_process,
@@ -145,12 +144,29 @@ class IluvatarModelRunner(ModelRunnerBase):
                 -1].disaggregate_info["role"] == "prefill":
             os.environ['PREFILL_NODE_ONE_STEP_STOP'] = "1"
 
+        top_k_reqs = []
+        top_p_reqs = []
+        max_num_seqs = self.parallel_config.max_num_seqs
+        top_p_buffer = paddle.full([max_num_seqs, 1],
+                                    self.model_config.top_p,
+                                    dtype='float32')
+        top_k_buffer = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
+
         req_len = len(req_dicts)
         for i in range(req_len):
             request = req_dicts[i]
             idx = request.idx
             length = len(request.prompt_token_ids)
 
+            if sampling_params := request.sampling_params:
+                if sampling_params.top_p < 1:
+                    top_p_reqs.append(idx)
+                top_k = sampling_params.top_k
+                if top_k > 0:
+                    top_k_reqs.append(idx)
+
             prefill_tokens = []
             if (request.guided_json is not None
                     or request.guided_regex is not None
@@ -225,8 +241,8 @@ class IluvatarModelRunner(ModelRunnerBase):
                 request.eos_token_ids.append(request.eos_token_ids[0])
             self.share_inputs["eos_token_id"][:] = np.array(
                 request.eos_token_ids, dtype="int64").reshape(-1, 1)
-
-            self.share_inputs["top_p"][idx:idx + 1] = request.get("top_p", 0.7)
+            top_p_buffer[idx:idx + 1] = request.get("top_p", 1.0)
+            top_k_buffer[idx:idx + 1] = request.get("top_k", 0)
             self.share_inputs["temperature"][idx:idx + 1] = request.get(
                 "temperature", 0.95)
             self.share_inputs["penalty_score"][idx:idx + 1] = request.get(
@@ -273,6 +289,15 @@ class IluvatarModelRunner(ModelRunnerBase):
                 idx, request.get("logits_processor"), prefill_tokens)
 
         self.share_inputs["not_need_stop"][0] = True
+        if len(top_k_reqs) == 0:
+            self.share_inputs["top_k"] = None
+        else:
+            self.share_inputs["top_k"] = top_k_buffer
+
+        if len(top_p_reqs) == 0:
+            self.share_inputs["top_p"] = None
+        else:
+            self.share_inputs["top_p"] = top_p_buffer
 
     def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int,
                               expected_decode_len: int):
@@ -329,8 +354,11 @@ class IluvatarModelRunner(ModelRunnerBase):
         self.share_inputs["eos_token_id"] = paddle.full(
             [self.parallel_config.eos_tokens_lens, 1], 0, dtype='int64')
         self.share_inputs["top_p"] = paddle.full([max_num_seqs, 1],
-                                                 self.model_config.top_p,
-                                                 dtype='float32')
+                                                self.model_config.top_p,
+                                                dtype='float32')
+        self.share_inputs["top_k"] = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
         self.share_inputs["temperature"] = paddle.full(
             [max_num_seqs, 1], self.model_config.temperature, dtype='float32')
         self.share_inputs["penalty_score"] = paddle.full(
@@ -558,6 +586,7 @@ class IluvatarModelRunner(ModelRunnerBase):
         self.sampling_metadata = SamplingMetadata(
             temperature=self.share_inputs["temperature"],
             top_p=self.share_inputs["top_p"],
+            top_k=self.share_inputs["top_k"],
             step_idx=self.share_inputs["step_idx"],
             pre_token_ids=self.share_inputs["pre_ids"],
             frequency_penalties=self.share_inputs["frequency_score"],

fastdeploy/worker/vl_model_runner_base.py

@@ -14,14 +14,14 @@
 # limitations under the License.
 """
 
-from abc import ABC, abstractmethod
 import argparse
+from abc import ABC, abstractmethod
 
 import paddle
 import paddle.distributed as dist
 import paddle.distributed.fleet as fleet
-from fastdeploy.config import ModelConfig
 
+from fastdeploy.config import ModelConfig
 from fastdeploy.utils import get_logger
 
 logger = get_logger("worker", "worker.log")

fastdeploy/worker/xpu_model_runner.py

@@ -282,11 +282,26 @@ class XPUModelRunner(ModelRunnerBase):
 
     def process_prefill_inputs(self, req_dicts: List[Request]):
         """ Process inputs for prefill tasks and update share_inputs buffer """
+        top_k_reqs = []
+        top_p_reqs = []
+        max_num_seqs = self.parallel_config.max_num_seqs
+        top_p_buffer = paddle.full([max_num_seqs, 1],
+                                    self.model_config.top_p,
+                                    dtype='float32')
+        top_k_buffer = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
         req_len = len(req_dicts)
         for i in range(req_len):
             request = req_dicts[i]
             idx = request.idx
             length = request.prompt_token_ids_len
+            if sampling_params := request.sampling_params:
+                if sampling_params.top_p < 1:
+                    top_p_reqs.append(idx)
+                top_k = sampling_params.top_k
+                if top_k > 0:
+                    top_k_reqs.append(idx)
             self.share_inputs["input_ids"][idx:idx + 1, :length] = np.array(
                 request.prompt_token_ids)
             if len(request.eos_token_ids
@@ -295,7 +310,8 @@ class XPUModelRunner(ModelRunnerBase):
             self.share_inputs["eos_token_id"][:] = np.array(
                 request.eos_token_ids, dtype="int64").reshape(-1, 1)
             self.share_inputs["pre_ids"][idx:idx + 1] = -1
-            self.share_inputs["top_p"][idx:idx + 1] = request.get("top_p", 0.7)
+            top_p_buffer[idx:idx + 1] = request.get("top_p", 1.0)
+            top_k_buffer[idx:idx + 1] = request.get("top_k", 0)
             self.share_inputs["temperature"][idx:idx + 1] = request.get(
                 "temperature", 0.95)
             self.share_inputs["penalty_score"][idx:idx + 1] = request.get(
@@ -344,6 +360,15 @@ class XPUModelRunner(ModelRunnerBase):
                         request.get("stop_token_ids"), dtype="int64")
 
         self.share_inputs["not_need_stop"][0] = True
+        if len(top_k_reqs) == 0:
+            self.share_inputs["top_k"] = None
+        else:
+            self.share_inputs["top_k"] = top_k_buffer
+
+        if len(top_p_reqs) == 0:
+            self.share_inputs["top_p"] = None
+        else:
+            self.share_inputs["top_p"] = top_p_buffer
 
     def _init_share_inputs(self, max_num_seqs: int):
         """Initialize all share buffers for model inputs.
@@ -363,8 +388,11 @@ class XPUModelRunner(ModelRunnerBase):
         self.share_inputs["eos_token_id"] = paddle.full(
             [self.parallel_config.eos_tokens_lens, 1], 0, dtype='int64')
         self.share_inputs["top_p"] = paddle.full([max_num_seqs, 1],
-                                                 self.model_config.top_p,
-                                                 dtype='float32')
+                                                self.model_config.top_p,
+                                                dtype='float32')
+        self.share_inputs["top_k"] = paddle.full([max_num_seqs, 1],
+                                                0,
+                                                dtype='int64')
         self.share_inputs["temperature"] = paddle.full(
             [max_num_seqs, 1], self.model_config.temperature, dtype='float32')
         self.share_inputs["penalty_score"] = paddle.full(
@@ -514,6 +542,7 @@ class XPUModelRunner(ModelRunnerBase):
         self.sampling_metadata = SamplingMetadata(
             temperature=self.share_inputs["temperature"],
             top_p=self.share_inputs["top_p"],
+            top_k=self.share_inputs["top_k"],
             step_idx=self.share_inputs["step_idx"],
             pre_token_ids=self.share_inputs["pre_ids"],
             frequency_penalties=self.share_inputs["frequency_score"],