Merge remote-tracking branch 'upstream/develop' into fix/guard-sm90-pybind-bindings

Author: ShigureNyako

.github/workflows/_build_xpu.yml

@@ -159,7 +159,8 @@ jobs:
               python -m pip install paddlepaddle-xpu==${PADDLEVERSION} -i https://www.paddlepaddle.org.cn/packages/stable/xpu-p800/
             else
               python -m pip uninstall paddlepaddle-xpu fastdeploy-xpu -y
-              python -m pip install --pre paddlepaddle-xpu -i https://www.paddlepaddle.org.cn/packages/nightly/xpu-p800/
+              # python -m pip install --pre paddlepaddle-xpu -i https://www.paddlepaddle.org.cn/packages/nightly/xpu-p800/
+              python -m pip install https://paddle-whl.bj.bcebos.com/nightly/xpu-p800/paddlepaddle-xpu/paddlepaddle_xpu-3.5.0.dev20260524-cp310-cp310-linux_x86_64.whl
             fi
 
 

.github/workflows/_xpu_4cards_case_test.yml

@@ -185,6 +185,7 @@ jobs:
           else
             python -m pip uninstall paddlepaddle-xpu fastdeploy-xpu -y
             python -m pip install --pre paddlepaddle-xpu -i https://www.paddlepaddle.org.cn/packages/nightly/xpu-p800/
+            python -m pip install https://paddle-whl.bj.bcebos.com/nightly/xpu-p800/paddlepaddle-xpu/paddlepaddle_xpu-3.5.0.dev20260524-cp310-cp310-linux_x86_64.whl
           fi
           echo "安装上游任务编译的fastdeploy-xpu..."
           python -m pip install ${FASTDEPLOY_WHEEL_URL}

.github/workflows/_xpu_8cards_case_test.yml

@@ -173,7 +173,7 @@ jobs:
             python -m pip install paddlepaddle-xpu==${PADDLEVERSION} -i https://www.paddlepaddle.org.cn/packages/stable/xpu-p800/
           else
             python -m pip uninstall paddlepaddle-xpu fastdeploy-xpu -y
-            python -m pip install --pre paddlepaddle-xpu -i https://www.paddlepaddle.org.cn/packages/nightly/xpu-p800/
+            python -m pip install https://paddle-whl.bj.bcebos.com/nightly/xpu-p800/paddlepaddle-xpu/paddlepaddle_xpu-3.5.0.dev20260524-cp310-cp310-linux_x86_64.whl
           fi
           echo "安装上游任务编译的fastdeploy-xpu..."
           python -m pip install ${FASTDEPLOY_WHEEL_URL}

.github/workflows/_xpu_unit_test.yml

@@ -172,6 +172,7 @@ jobs:
           else
             python -m pip uninstall paddlepaddle-xpu fastdeploy-xpu -y
             python -m pip install --pre paddlepaddle-xpu -i https://www.paddlepaddle.org.cn/packages/nightly/xpu-p800/
+            python -m pip install https://paddle-whl.bj.bcebos.com/nightly/xpu-p800/paddlepaddle-xpu/paddlepaddle_xpu-3.5.0.dev20260524-cp310-cp310-linux_x86_64.whl
           fi
           echo "安装上游任务编译的fastdeploy-xpu..."
           python -m pip install ${FASTDEPLOY_WHEEL_URL}

.github/workflows/ci_metax.yml

@@ -22,12 +22,9 @@ jobs:
     environment: Metax_ci
 
     steps:
-      - name: Cleanup
-        run: |
-          rm -rf * .[^.]*
       - name: Trigger Jenkins job
         timeout-minutes: 60
-        uses: MetaX-MACA/simple-jenkins-githubaction@v1.1
+        uses: MetaX-MACA/simple-jenkins-githubaction@v1.2
         with:
           job_name: paddle_fastdeploy_metax_smoketest
           username: fastdeploy_builder

custom_ops/gpu_ops/speculate_decoding/draft_model/mtp_save_first_token_with_topk.cc

@@ -119,9 +119,11 @@ void MTPSaveFirstTokenWithTopK(const paddle::Tensor& sampled_token_ids,
   msg_sed.mtype = 1;
   msg_sed.meta[0] = not_need_stop.data<bool>()[0] ? inference_msg_id_from_env
                                                   : -inference_msg_id_from_env;
-  msg_sed.meta[1] = message_flag;
-  msg_sed.meta[2] = bsz;
+  // Pack message_flag (low 8 bits) and max_num_logprobs (high 24 bits) into
+  // meta[1]. Receiver unpacks both to avoid reading unused topk slots.
   int max_num_logprobs = logprob_token_ids.shape()[1];
+  msg_sed.meta[1] = message_flag | (max_num_logprobs << 8);
+  msg_sed.meta[2] = bsz;
   for (int i = 0; i < bsz; i++) {
     int cur_token_num;
     if (seq_lens_decoder_data[i] < prompt_lens_data[i] ||
@@ -139,29 +141,24 @@ void MTPSaveFirstTokenWithTopK(const paddle::Tensor& sampled_token_ids,
     auto* cur_batch_msg_sed = &msg_sed.mtext[i];
     int token_offset = cu_batch_token_offset_data[i];
     for (int j = 0; j < cur_token_num; j++) {
+      // Use SPEC_LOGPROB_K+1 as stride (fixed struct layout), but only write
+      // max_num_logprobs columns to avoid filling unused topk slots.
       auto* cur_tokens = &cur_batch_msg_sed->tokens[j * (SPEC_LOGPROB_K + 1)];
       auto* cur_scores = &cur_batch_msg_sed->scores[j * (SPEC_LOGPROB_K + 1)];
       if (j == 0) {
         // first token has full logprobs
-        for (int k = 0; k < SPEC_LOGPROB_K + 1; k++) {
+        for (int k = 0; k < max_num_logprobs; k++) {
           if (k == 0) {
             cur_tokens[k] =
                 (int)sampled_token_ids_data[i * max_draft_tokens + j];
             cur_scores[k] =
-                logprob_scores_data[(token_offset + j) * (SPEC_LOGPROB_K + 1) +
-                                    k];
-          } else if (k < max_num_logprobs) {
-            // only for first token
-            cur_tokens[k] =
-                (int)logprob_token_ids_data[(token_offset + j) *
-                                                (SPEC_LOGPROB_K + 1) +
-                                            k];
-            cur_scores[k] =
-                logprob_scores_data[(token_offset + j) * (SPEC_LOGPROB_K + 1) +
-                                    k];
+                logprob_scores_data[(token_offset + j) * max_num_logprobs + k];
           } else {
-            cur_tokens[k] = -1;
-            cur_scores[k] = 0.0;
+            cur_tokens[k] = (int)
+                logprob_token_ids_data[(token_offset + j) * max_num_logprobs +
+                                       k];
+            cur_scores[k] =
+                logprob_scores_data[(token_offset + j) * max_num_logprobs + k];
           }
         }
         cur_batch_msg_sed->ranks[j] = (int)logprob_ranks_data[token_offset + j];
@@ -174,7 +171,8 @@ void MTPSaveFirstTokenWithTopK(const paddle::Tensor& sampled_token_ids,
 #ifdef SPECULATE_SAVE_WITH_OUTPUT_DEBUG
   std::cout << "msg data: " << std::endl;
   std::cout << "stop_flag: " << msg_sed.meta[0]
-            << ", message_flag: " << msg_sed.meta[1]
+            << ", message_flag: " << (msg_sed.meta[1] & 0xFF)
+            << ", max_num_logprobs: " << (msg_sed.meta[1] >> 8)
             << ", bsz: " << msg_sed.meta[2] << std::endl;
   for (int i = 0; i < bsz; i++) {
     int cur_token_num = msg_sed.meta[3 + i];

custom_ops/gpu_ops/speculate_decoding/speculate_get_output_with_topk.cc

@@ -75,8 +75,11 @@ void SpeculateGetOutMmsgTopK(const paddle::Tensor& output_tokens,
 
   int bsz = msg_rcv.meta[2];
   output_tokens_data[0] = (int64_t)msg_rcv.meta[0];
+  // Unpack message_flag (low 8 bits) and actual_topk (high 24 bits) from
+  // meta[1]. Keep packed value; Python unpacks message_flag and actual_topk.
   output_tokens_data[1] = (int64_t)msg_rcv.meta[1];
   output_tokens_data[2] = (int64_t)msg_rcv.meta[2];
+  int actual_topk = msg_rcv.meta[1] >> 8;
 
   int output_tokens_offset = 3 + SPEC_LOGPROB_MAX_BSZ;
   for (int i = 0; i < bsz; i++) {
@@ -89,7 +92,7 @@ void SpeculateGetOutMmsgTopK(const paddle::Tensor& output_tokens,
         output_scores_data + i * (MAX_DRAFT_TOKEN_NUM * (SPEC_LOGPROB_K + 1));
     auto* cur_batch_msg_rcv = &msg_rcv.mtext[i];
     for (int j = 0; j < cur_token_num; j++) {
-      for (int k = 0; k < real_k + 1; k++) {
+      for (int k = 0; k < actual_topk; k++) {
         cur_output_token[j * (SPEC_LOGPROB_K + 1) + k] =
             (int64_t)cur_batch_msg_rcv->tokens[j * (SPEC_LOGPROB_K + 1) + k];
         cur_output_score[j * (SPEC_LOGPROB_K + 1) + k] =
@@ -102,7 +105,8 @@ void SpeculateGetOutMmsgTopK(const paddle::Tensor& output_tokens,
 #ifdef SPECULATE_GET_WITH_OUTPUT_DEBUG
   std::cout << "msg data: " << std::endl;
   std::cout << "stop_flag: " << output_tokens_data[0]
-            << ", message_flag: " << output_tokens_data[1]
+            << ", message_flag: " << (output_tokens_data[1] & 0xFF)
+            << ", max_num_logprobs: " << (output_tokens_data[1] >> 8)
             << ", bsz: " << output_tokens_data[2] << std::endl;
   for (int i = 0; i < output_tokens_data[2]; i++) {
     int cur_token_num = output_tokens_data[3 + i];

custom_ops/gpu_ops/speculate_decoding/speculate_save_output_with_topk.cc

@@ -121,9 +121,11 @@ void SpeculateSaveOutMmsgTopK(const paddle::Tensor& sampled_token_ids,
   msg_sed.mtype = 1;
   msg_sed.meta[0] = not_need_stop.data<bool>()[0] ? inference_msg_id_from_env
                                                   : -inference_msg_id_from_env;
-  msg_sed.meta[1] = message_flag;
-  msg_sed.meta[2] = bsz;
+  // Pack message_flag (low 8 bits) and max_num_logprobs (high 24 bits) into
+  // meta[1]. Receiver unpacks both to avoid reading unused topk slots.
   int max_num_logprobs = logprob_token_ids.shape()[1];
+  msg_sed.meta[1] = message_flag | (max_num_logprobs << 8);
+  msg_sed.meta[2] = bsz;
   for (int i = 0; i < bsz; i++) {
     int cur_token_num;
     if (seq_lens_decoder_data[i] < prompt_lens_data[i]) {
@@ -139,24 +141,20 @@ void SpeculateSaveOutMmsgTopK(const paddle::Tensor& sampled_token_ids,
     auto* cur_batch_msg_sed = &msg_sed.mtext[i];
     int token_offset = cu_batch_token_offset_data[i];
     for (int j = 0; j < cur_token_num; j++) {
+      // Use SPEC_LOGPROB_K+1 as stride (fixed struct layout), but only write
+      // max_num_logprobs columns to avoid filling unused topk slots.
       auto* cur_tokens = &cur_batch_msg_sed->tokens[j * (SPEC_LOGPROB_K + 1)];
       auto* cur_scores = &cur_batch_msg_sed->scores[j * (SPEC_LOGPROB_K + 1)];
-      for (int k = 0; k < SPEC_LOGPROB_K + 1; k++) {
+      for (int k = 0; k < max_num_logprobs; k++) {
         if (k == 0) {
           cur_tokens[k] = (int)sampled_token_ids_data[i * max_draft_tokens + j];
           cur_scores[k] =
-              logprob_scores_data[(token_offset + j) * (SPEC_LOGPROB_K + 1) +
-                                  k];
-        } else if (k < max_num_logprobs) {
+              logprob_scores_data[(token_offset + j) * max_num_logprobs + k];
+        } else {
           cur_tokens[k] = (int)
-              logprob_token_ids_data[(token_offset + j) * (SPEC_LOGPROB_K + 1) +
-                                     k];
+              logprob_token_ids_data[(token_offset + j) * max_num_logprobs + k];
           cur_scores[k] =
-              logprob_scores_data[(token_offset + j) * (SPEC_LOGPROB_K + 1) +
-                                  k];
-        } else {
-          cur_tokens[k] = -1;
-          cur_scores[k] = 0.0;
+              logprob_scores_data[(token_offset + j) * max_num_logprobs + k];
         }
       }
       cur_batch_msg_sed->ranks[j] = (int)logprob_ranks_data[token_offset + j];
@@ -165,7 +163,8 @@ void SpeculateSaveOutMmsgTopK(const paddle::Tensor& sampled_token_ids,
 #ifdef SPECULATE_SAVE_WITH_OUTPUT_DEBUG
   std::cout << "msg data: " << std::endl;
   std::cout << "stop_flag: " << msg_sed.meta[0]
-            << ", message_flag: " << msg_sed.meta[1]
+            << ", message_flag: " << (msg_sed.meta[1] & 0xFF)
+            << ", max_num_logprobs: " << (msg_sed.meta[1] >> 8)
             << ", bsz: " << msg_sed.meta[2] << std::endl;
   for (int i = 0; i < bsz; i++) {
     int cur_token_num = msg_sed.meta[3 + i];

fastdeploy/model_executor/layers/normalization.py

@@ -243,14 +243,21 @@ def forward(
 
         if residual_input is None:
             residual_out = x
+        use_allreduce_fused = (
+            self.enable_all_reduce_fusion
+            and self.tp_size > 1
+            and x.shape[0] <= 2048
+            and residual_input is not None
+            and current_platform.is_cuda()
+        )
         if proxy_rmsnorm is None:
             if current_platform.is_gcu():
                 if residual_input is None:
                     norm_out = rms_norm(x, self.weight, self.eps)
                     return norm_out.astype(x_dtype), residual_out
                 norm_out = self.norm_func(x, residual_input, self.weight, self.eps)
             # enable trtllm all reduce fusion
-            elif self.enable_all_reduce_fusion and x.shape[0] <= 2048:
+            elif use_allreduce_fused:
                 norm_out = flashinfer_allreduce_residual_rmsnorm(
                     fd_config=self.fd_config, input_tensor=x, residual=residual_input, weight=self.weight, eps=self.eps
                 )
@@ -276,9 +283,19 @@ def forward(
                         quant_min_bound=self.quant_min_bound,
                     )
         else:
-            if residual_input is not None:
-                x = x + residual_input
-            norm_out = proxy_rmsnorm(x, self.weight, self.eps), x
+            if use_allreduce_fused:
+                norm_out = flashinfer_allreduce_residual_rmsnorm(
+                    fd_config=self.fd_config,
+                    input_tensor=x,
+                    residual=residual_input,
+                    weight=self.weight,
+                    eps=self.eps,
+                )
+                assert norm_out[0] is not None, "Trtllm-all-reduce fusion failed!"
+            else:
+                if residual_input is not None:
+                    x = x + residual_input
+                norm_out = proxy_rmsnorm(x, self.weight, self.eps), x
 
         out = norm_out[0].astype(x_dtype)
         if residual_input is not None:

fastdeploy/output/token_processor.py

@@ -739,12 +739,15 @@ def _process_batch_draft_tokens(self, mtype, batch, accept_num, tokens, scores,
                 metrics=None,
             )
 
-            token_ids = tokens[i][:, 0].tolist()[: accept_num[i]]
+            tokens_i = tokens[i].tolist()
+            scores_i = scores[i].tolist()
+            ranks_i = ranks[i].tolist()
+            token_ids = [row[0] for row in tokens_i[: accept_num[i]]]
             for batch_token_index in range(len(token_ids)):
-                result.outputs.logprob = float(scores[i, batch_token_index, 0])
-                topk_token_ids = tokens[i, batch_token_index, :].tolist()
-                topk_logprobs = scores[i, batch_token_index, :].tolist()
-                sampled_rank = ranks[i, batch_token_index].item()
+                result.outputs.logprob = scores_i[batch_token_index][0]
+                topk_token_ids = tokens_i[batch_token_index]
+                topk_logprobs = scores_i[batch_token_index]
+                sampled_rank = ranks_i[batch_token_index]
 
                 if result.outputs.draft_top_logprobs is None:
                     result.outputs.draft_top_logprobs = LogprobsLists(
@@ -771,16 +774,19 @@ def _process_batch_output(self):
         mtype = 3
         if self.cfg.speculative_config.method:
             if self.use_logprobs:
-                mtype = int(self.output_tokens[1, 0].item())
+                # meta[1] packs message_flag (low 8 bits) and actual_topk (high 24 bits).
+                packed_meta1 = int(self.output_tokens[1, 0].item())
+                mtype = packed_meta1 & 0xFF
+                actual_topk = packed_meta1 >> 8
                 batch = self.output_tokens[2, 0]
                 accept_num = [int(num[0]) for num in self.output_tokens[3 : batch + 3]]
                 tokens = tokens[3 + MAX_BSZ : 3 + MAX_BSZ + batch * MAX_DRAFT_TOKENS * (K + 1)].reshape(
                     [batch, MAX_DRAFT_TOKENS, K + 1]
-                )
+                )[:, :, :actual_topk]
                 scores = (
                     self.output_scores[: batch * MAX_DRAFT_TOKENS * (K + 1)]
                     .numpy()
-                    .reshape([batch, MAX_DRAFT_TOKENS, K + 1])
+                    .reshape([batch, MAX_DRAFT_TOKENS, K + 1])[:, :, :actual_topk]
                 )
                 ranks = self.output_ranks[: batch * MAX_DRAFT_TOKENS].numpy().reshape([batch, MAX_DRAFT_TOKENS])
 
@@ -789,6 +795,10 @@ def _process_batch_output(self):
                     batch_result = self._process_batch_draft_tokens(mtype, batch, accept_num, tokens, scores, ranks)
                     self.postprocess(batch_result, mtype)
                     return
+                # Pre-convert full arrays to Python lists once for MTP target token path.
+                tokens_lists = tokens.tolist()
+                scores_lists = scores.tolist()
+                ranks_list = ranks.tolist()
             else:
                 batch = self.output_tokens[1]
                 accept_num = tokens[2 : batch + 2]
@@ -856,7 +866,7 @@ def _process_batch_output(self):
                         )
                     token_ids = [RECOVERY_STOP_SIGNAL]
                 elif self.use_logprobs:
-                    token_ids = tokens[i][:, 0].tolist()[: accept_num[i]]
+                    token_ids = [row[0] for row in tokens_lists[i][: accept_num[i]]]
                 else:
                     token_ids = tokens[
                         2
@@ -988,10 +998,10 @@ def _process_batch_output(self):
                     task.output_token_ids.append(token_id)
                     if self.use_logprobs:
                         if self.cfg.speculative_config.method:
-                            result.outputs.logprob = float(scores[i, batch_token_index, 0])
-                            topk_token_ids = tokens[i, batch_token_index, :].tolist()
-                            topk_logprobs = scores[i, batch_token_index, :].tolist()
-                            sampled_rank = ranks[i, batch_token_index].item()
+                            result.outputs.logprob = scores_lists[i][batch_token_index][0]
+                            topk_token_ids = tokens_lists[i][batch_token_index]
+                            topk_logprobs = scores_lists[i][batch_token_index]
+                            sampled_rank = ranks_list[i][batch_token_index]
                         else:
                             # Use pre-converted lists (batch .tolist() done before the loop).
                             result.outputs.logprob = scores_lists[i][0]