fix top_p_candidates and support separate setting of sampling params for mtp

custom_ops/gpu_ops/speculate_decoding/speculate_verify.cu

@@ -38,14 +38,20 @@ __device__ int64_t topp_sampling_kernel(const int64_t *candidate_ids,
   const int tid = threadIdx.x;
 
   float sum_scores = 0.0f;
-  float rand_top_p = curand_uniform(dev_curand_states + tid) * topp;
+  for (int i = 0; i < candidate_len; i++) {
+    sum_scores += candidate_scores[i];
+  }
+  float tgt_topp = sum_scores < topp ? sum_scores : topp;
+
+  sum_scores = 0.0f;
+  float rand_top_p = curand_uniform(dev_curand_states + tid) * tgt_topp;
   for (int i = 0; i < candidate_len; i++) {
     sum_scores += candidate_scores[i];
     if (rand_top_p <= sum_scores) {
-      return candidate_ids[i];
+            return candidate_ids[i];
     }
   }
-  return candidate_ids[0];
+    return candidate_ids[0];
 }
 
 __global__ void setup_kernel(curandState_t *state, const uint64_t seed,
@@ -136,6 +142,9 @@ __global__ void speculate_verify(
             // accept_num_now++;
             step_idx[bid]++;
             auto accept_token = draft_tokens_now[i + 1];
+            if (is_in_end(verify_tokens_now[i * max_candidate_len], end_tokens, end_length)) {
+              accept_token = end_tokens[0];
+            }
             accept_tokens[bid * max_draft_tokens + i] = accept_token;
 
             if (is_in_end(accept_token, end_tokens, end_length) ||

custom_ops/gpu_ops/speculate_decoding/top_p_candidates.cu

@@ -467,6 +467,9 @@ __global__ void KeMatrixTopPBeamTopKFt(
                 break;
             }
         }
+        if (top_p_value == 1.0 && actual_candidates_lens[token_id] == 0){
+            actual_candidates_lens[token_id] = max_cadidate_len;
+        }
     }
 }
 

fastdeploy/envs.py

@@ -95,6 +95,13 @@
     "FD_FOR_TORCH_MODEL_FORMAT": lambda: bool(int(os.getenv("FD_FOR_TORCH_MODEL_FORMAT", "0"))),
     # force disable default chunked prefill
     "FD_DISABLE_CHUNKED_PREFILL": lambda: bool(int(os.getenv("FD_DISABLE_CHUNKED_PREFILL", "0"))),
+    # For separate setting of sampling parameters for speculative decoding
+    "FD_SPECULATE_SAMPLING_TOP_P": lambda: (
+        None if "FD_SPECULATE_SAMPLING_TOP_P" not in os.environ else float(os.environ["FD_SPECULATE_SAMPLING_TOP_P"])
+    ),
+    "FD_SPECULATE_SAMPLING_TOP_K": lambda: (
+        None if "FD_SPECULATE_SAMPLING_TOP_K" not in os.environ else float(os.environ["FD_SPECULATE_SAMPLING_TOP_K"])
+    ),
 }
 
 

fastdeploy/spec_decode/mtp.py

@@ -303,8 +303,16 @@ def _init_model_inputs(self):
         )
         # self.model_inputs["caches"] = self.cache_kvs
         # Inherit generation hyperparameters from the main model for consistency
-        self.model_inputs["top_p"] = self.target_model_inputs["top_p"]
-        self.model_inputs["top_k"] = self.target_model_inputs["top_k"]
+        self.model_inputs["top_p"] = (
+            self.target_model_inputs["top_p"]
+            if envs.FD_SPECULATE_SAMPLING_TOP_P is None
+            else paddle.full_like(self.target_model_inputs["top_p"], envs.FD_SPECULATE_SAMPLING_TOP_P)
+        )
+        self.model_inputs["top_k"] = (
+            self.target_model_inputs["top_k"]
+            if envs.FD_SPECULATE_SAMPLING_TOP_K is None
+            else paddle.full_like(self.target_model_inputs["top_k"], envs.FD_SPECULATE_SAMPLING_TOP_K)
+        )
         self.model_inputs["temperature"] = self.target_model_inputs["temperature"]
         self.model_inputs["eos_token_id"] = self.target_model_inputs["eos_token_id"]
         self.model_inputs["penalty_score"] = self.target_model_inputs["penalty_score"]
@@ -645,6 +653,8 @@ def _propose(self, target_hidden_states):
                 self._initialize_forward_meta()
 
                 # Get sampling metadata
+                print(f"[MTPProposer] top_p: {self.model_inputs['top_p']}")
+                print(f"[MTPProposer] top_k: {self.model_inputs['top_k']}")
                 self.sampling_metadata = SamplingMetadata(
                     temperature=self.model_inputs["temperature"],
                     top_p=self.model_inputs["top_p"],