[Feature] support mtp logprob (#4464)

* support mtp logprob

* fix unitest

custom_ops/gpu_ops/cpp_extensions.cc

@@ -348,7 +348,9 @@ paddle::Tensor RebuildPaddingFunc(
     const paddle::Tensor &seq_lens_decoder,
     const paddle::Tensor &seq_lens_encoder,
     const paddle::optional<paddle::Tensor> &output_padding_offset,
-    int max_input_length);
+    const paddle::optional<paddle::Tensor> &first_token_out,
+    int max_input_length,
+    bool enable_logprob);
 
 void GetStopFlagsMulti(const paddle::Tensor &topk_ids,
                        const paddle::Tensor &stop_flags,
@@ -910,6 +912,32 @@ void SaveOutMmsgStatic(const paddle::Tensor& x,
                        int64_t rank_id,
                        bool save_each_rank);
 
+void SpeculateGetLogits(const paddle::Tensor &draft_logits,
+                        const paddle::Tensor &next_token_num,
+                        const paddle::Tensor &batch_token_num,
+                        const paddle::Tensor &cu_next_token_offset,
+                        const paddle::Tensor &cu_batch_token_offset,
+                        const paddle::Tensor &logits,
+                        const paddle::Tensor &first_token_logits,
+                        const paddle::Tensor &seq_lens_this_time,
+                        const paddle::Tensor &seq_lens_encoder);
+
+void SpeculateInsertFirstToken(const paddle::Tensor &token_ids,
+                               const paddle::Tensor &accept_tokens,
+                               const paddle::Tensor &next_tokens,
+                               const paddle::Tensor &cu_next_token_offset,
+                               const paddle::Tensor &cu_batch_token_offset,
+                               const paddle::Tensor &seq_lens_this_time,
+                               const paddle::Tensor &seq_lens_encoder);
+
+void SpeculateGetTargetLogits(const paddle::Tensor &target_logits,
+                              const paddle::Tensor &logits,
+                              const paddle::Tensor &cu_batch_token_offset,
+                              const paddle::Tensor &ori_cu_batch_token_offset,
+                              const paddle::Tensor &seq_lens_this_time,
+                              const paddle::Tensor &seq_lens_encoder,
+                              const paddle::Tensor &accept_num);
+
 PYBIND11_MODULE(fastdeploy_ops, m) {
 
   m.def("get_expert_token_num", &GetExpertTokenNum, py::arg("topk_ids"),
@@ -1291,4 +1319,10 @@ PYBIND11_MODULE(fastdeploy_ops, m) {
   m.def("min_p_sampling", &MinPSamplingFromProbs, "min_p_sampling function");
 
   m.def("save_output", &SaveOutMmsgStatic, "save_output function");
+
+  m.def("speculate_get_logits", &SpeculateGetLogits, "speculate_get_logits function");
+
+  m.def("speculate_insert_first_token", &SpeculateInsertFirstToken, "speculate_insert_first_token function");
+
+  m.def("speculate_get_target_logits", &SpeculateGetTargetLogits, "speculate_get_target_logits function");
 }

custom_ops/gpu_ops/rebuild_padding.cu

@@ -46,6 +46,7 @@ __global__ void RebuildPaddingKernel(T *output_data,
 
 template <typename T, int VecSize>
 __global__ void RebuildAppendPaddingKernel(T *output_data,
+                                           T *first_token_out,
                                            const T *input_data,
                                            const int *cu_seqlens_q,
                                            const int *seq_len_this_time,
@@ -55,7 +56,8 @@ __global__ void RebuildAppendPaddingKernel(T *output_data,
                                            const int max_input_length,
                                            const int dim_embed,
                                            const int64_t output_elem_nums,
-                                           const int bsz) {
+                                           const int bsz,
+                                           const bool enable_logprob) {
     AlignedVector<T, VecSize> src_vec;
     const int64_t global_idx = blockDim.x * blockIdx.x + threadIdx.x;
     for (int64_t i = global_idx * VecSize; i < output_elem_nums;
@@ -70,13 +72,20 @@ __global__ void RebuildAppendPaddingKernel(T *output_data,
         if (seq_len_decoder[bi] == 0 && seq_len_encoder[bi] == 0) continue;
 
         if (seq_len_encoder[bi] > 0) seq_id = seq_len_encoder[bi] - 1;
-         const int cum_offset_bi = bi * max_input_length - cu_seqlens_q[bi];
+        const int cum_offset_bi = bi * max_input_length - cu_seqlens_q[bi];
         const int input_token_id = ori_token_id - cum_offset_bi + seq_id;
         const int bias_idx = i % dim_embed;
 
         Load<T, VecSize>(&input_data[input_token_id * dim_embed + bias_idx],
                          &src_vec);
         Store<T, VecSize>(src_vec, &output_data[i]);
+
+        if (enable_logprob && seq_len_encoder[bi] > 0) {
+            const int first_input_token_id = input_token_id - 1;
+            Load<T, VecSize>(&input_data[first_input_token_id * dim_embed + bias_idx],
+                             &src_vec);
+            Store<T, VecSize>(src_vec, &first_token_out[bi * dim_embed + bias_idx]);
+        }
     }
 }
 
@@ -89,7 +98,9 @@ std::vector<paddle::Tensor> rebuild_padding(
     const paddle::Tensor &seq_lens_decoder,
     const paddle::Tensor &seq_lens_encoder,
     const paddle::optional<paddle::Tensor> &output_padding_offset,
-    int max_input_length) {
+    const paddle::optional<paddle::Tensor> &first_token_out,
+    int max_input_length,
+    bool enable_logprob) {
     typedef PDTraits<D> traits_;
     typedef typename traits_::DataType DataType_;
     typedef typename traits_::data_t data_t;
@@ -134,6 +145,10 @@ std::vector<paddle::Tensor> rebuild_padding(
         RebuildAppendPaddingKernel<DataType_, PackSize>
             <<<grid_size, blocksize, 0, cu_stream>>>(
                 reinterpret_cast<DataType_ *>(out.data<data_t>()),
+                first_token_out.is_initialized()
+                    ? reinterpret_cast<DataType_ *>(const_cast<data_t *>(
+                          first_token_out.get_ptr()->data<data_t>()))
+                    : nullptr,
                 reinterpret_cast<const DataType_ *>(tmp_out.data<data_t>()),
                 cu_seqlens_q.data<int>(),
                 seq_len_this_time.data<int>(),
@@ -143,7 +158,8 @@ std::vector<paddle::Tensor> rebuild_padding(
                 max_input_length,
                 dim_embed,
                 elem_nums,
-                bsz);
+                bsz,
+                enable_logprob);
     } else {
         RebuildPaddingKernel<DataType_, PackSize>
             <<<grid_size, blocksize, 0, cu_stream>>>(
@@ -168,7 +184,9 @@ paddle::Tensor RebuildPaddingFunc(
     const paddle::Tensor &seq_lens_decoder,
     const paddle::Tensor &seq_lens_encoder,
     const paddle::optional<paddle::Tensor> &output_padding_offset,
-    int max_input_length) {
+    const paddle::optional<paddle::Tensor> &first_token_out,
+    int max_input_length,
+    bool enable_logprob) {
     switch (tmp_out.type()) {
         case paddle::DataType::BFLOAT16: {
             return rebuild_padding<paddle::DataType::BFLOAT16>(
@@ -178,7 +196,9 @@ paddle::Tensor RebuildPaddingFunc(
                 seq_lens_decoder,
                 seq_lens_encoder,
                 output_padding_offset,
-                max_input_length)[0];
+                first_token_out,
+                max_input_length,
+                enable_logprob)[0];
         }
         case paddle::DataType::FLOAT16: {
             return rebuild_padding<paddle::DataType::FLOAT16>(
@@ -188,7 +208,9 @@ paddle::Tensor RebuildPaddingFunc(
                 seq_lens_decoder,
                 seq_lens_encoder,
                 output_padding_offset,
-                max_input_length)[0];
+                first_token_out,
+                max_input_length,
+                enable_logprob)[0];
         }
         case paddle::DataType::FLOAT32: {
             return rebuild_padding<paddle::DataType::FLOAT32>(
@@ -198,7 +220,9 @@ paddle::Tensor RebuildPaddingFunc(
                 seq_lens_decoder,
                 seq_lens_encoder,
                 output_padding_offset,
-                max_input_length)[0];
+                first_token_out,
+                max_input_length,
+                enable_logprob)[0];
         }
         default: {
             PD_THROW(
@@ -216,14 +240,18 @@ std::vector<paddle::Tensor> RebuildPadding(
     const paddle::Tensor &seq_lens_decoder,
     const paddle::Tensor &seq_lens_encoder,
     const paddle::optional<paddle::Tensor> &output_padding_offset,
-    int max_input_length) {
+    const paddle::optional<paddle::Tensor> &first_token_out,
+    int max_input_length,
+    bool enable_logprob) {
     return {RebuildPaddingFunc(tmp_out,
                                cu_seqlens_q,
                                seq_len_this_time,
                                seq_lens_decoder,
                                seq_lens_encoder,
                                output_padding_offset,
-                               max_input_length)};
+                               first_token_out,
+                               max_input_length,
+                               enable_logprob)};
 }
 
 std::vector<std::vector<int64_t>> RebuildPaddingInferShape(
@@ -259,9 +287,10 @@ PD_BUILD_STATIC_OP(rebuild_padding)
              "seq_len_this_time",
              "seq_lens_decoder",
              "seq_lens_encoder",
-             paddle::Optional("output_padding_offset")})
+             paddle::Optional("output_padding_offset"),
+             paddle::Optional("first_token_out")})
     .Outputs({"out"})
-    .Attrs({"max_input_length: int"})
+    .Attrs({"max_input_length: int", "enable_logprob: bool"})
     .SetKernelFn(PD_KERNEL(RebuildPadding))
     .SetInferShapeFn(PD_INFER_SHAPE(RebuildPaddingInferShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(RebuildPaddingInferDtype));

custom_ops/gpu_ops/speculate_decoding/speculate_get_output_with_topk.cc

@@ -0,0 +1,161 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <stdio.h>
+#include <string.h>
+#include <sys/ipc.h>
+#include <sys/msg.h>
+#include <sys/types.h>
+#include "paddle/extension.h"
+
+#ifndef PD_BUILD_STATIC_OP
+#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
+#endif
+
+#define MAX_BSZ 512
+#define K 20
+#define MAX_DRAFT_TOKEN_NUM 6
+
+struct batch_msgdata {
+    int tokens[MAX_DRAFT_TOKEN_NUM * (K + 1)];
+    float scores[MAX_DRAFT_TOKEN_NUM * (K + 1)];
+    int ranks[MAX_DRAFT_TOKEN_NUM];
+};
+
+struct msgdata {
+    long mtype;
+    int meta[3 + MAX_BSZ];  // stop_flag, message_flag, bsz, batch_token_nums
+    batch_msgdata mtext[MAX_BSZ];
+};
+
+void SpeculateGetOutMmsgTopK(const paddle::Tensor& output_tokens,
+                             const paddle::Tensor& output_scores,
+                             const paddle::Tensor& output_ranks,
+                             int real_k,
+                             int64_t rank_id,
+                             bool wait_flag) {
+    struct msgdata msg_rcv;
+    int msg_queue_id = 1;
+
+    if (const char* inference_msg_queue_id_env_p =
+            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+        std::string inference_msg_queue_id_env_str(
+            inference_msg_queue_id_env_p);
+        int inference_msg_queue_id_from_env =
+            std::stoi(inference_msg_queue_id_env_str);
+#ifdef SPECULATE_GET_WITH_OUTPUT_DEBUG
+        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                  << inference_msg_queue_id_from_env << std::endl;
+#endif
+        msg_queue_id = inference_msg_queue_id_from_env;
+    }
+    static key_t key = ftok("/dev/shm", msg_queue_id);
+
+    static int msgid = msgget(key, IPC_CREAT | 0666);
+#ifdef SPECULATE_GET_WITH_OUTPUT_DEBUG
+    std::cout << "get_output_key: " << key << std::endl;
+    std::cout << "get_output msgid: " << msgid << std::endl;
+#endif
+
+    int64_t* output_tokens_data =
+        const_cast<int64_t*>(output_tokens.data<int64_t>());
+    float* output_scores_data = const_cast<float*>(output_scores.data<float>());
+    int64_t* output_ranks_data =
+        const_cast<int64_t*>(output_ranks.data<int64_t>());
+    int ret = -1;
+    if (!wait_flag) {
+        ret = msgrcv(
+            msgid, &msg_rcv, sizeof(msg_rcv) - sizeof(long), 0, IPC_NOWAIT);
+    } else {
+        ret = msgrcv(msgid, &msg_rcv, sizeof(msg_rcv) - sizeof(long), 0, 0);
+    }
+    if (ret == -1) {
+        // read none
+        output_tokens_data[0] = -2;  // stop_flag
+        output_tokens_data[1] = 0;   // message_flag, Target: 3, Draft: 4
+        output_tokens_data[2] = 0;   // bsz
+        return;
+    }
+
+    int bsz = msg_rcv.meta[2];
+    output_tokens_data[0] = (int64_t)msg_rcv.meta[0];
+    output_tokens_data[1] = (int64_t)msg_rcv.meta[1];
+    output_tokens_data[2] = (int64_t)msg_rcv.meta[2];
+
+    int output_tokens_offset = 3 + MAX_BSZ;
+    for (int i = 0; i < bsz; i++) {
+        int cur_token_num = msg_rcv.meta[3 + i];
+        output_tokens_data[3 + i] = (int64_t)cur_token_num;  // batch_token_nums
+
+        auto* cur_output_token = output_tokens_data + output_tokens_offset +
+                                 i * (MAX_DRAFT_TOKEN_NUM * (K + 1));
+        auto* cur_output_score =
+            output_scores_data + i * (MAX_DRAFT_TOKEN_NUM * (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++) {
+                cur_output_token[j * (K + 1) + k] =
+                    (int64_t)cur_batch_msg_rcv->tokens[j * (K + 1) + k];
+                cur_output_score[j * (K + 1) + k] =
+                    cur_batch_msg_rcv->scores[j * (K + 1) + k];
+            }
+            output_ranks_data[i * MAX_DRAFT_TOKEN_NUM + j] =
+                (int64_t)cur_batch_msg_rcv->ranks[j];
+        }
+    }
+#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]
+              << ", 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];
+        std::cout << "batch " << i << " token_num: " << cur_token_num
+                  << std::endl;
+        for (int j = 0; j < cur_token_num; j++) {
+            std::cout << "tokens: ";
+            for (int k = 0; k < K + 1; k++) {
+                std::cout
+                    << output_tokens_data[output_tokens_offset +
+                                          i * MAX_DRAFT_TOKEN_NUM * (K + 1) +
+                                          j * (K + 1) + k]
+                    << " ";
+            }
+            std::cout << std::endl;
+            std::cout << "scores: ";
+            for (int k = 0; k < K + 1; k++) {
+                std::cout
+                    << output_scores_data[i * MAX_DRAFT_TOKEN_NUM * (K + 1) +
+                                          j * (K + 1) + k]
+                    << " ";
+            }
+            std::cout << std::endl;
+            std::cout << "ranks: "
+                      << output_ranks_data[i * MAX_DRAFT_TOKEN_NUM + j]
+                      << std::endl;
+        }
+    }
+    std::cout << std::endl;
+#endif
+    return;
+}
+
+PD_BUILD_STATIC_OP(speculate_get_output_topk)
+    .Inputs({"output_tokens", "output_scores", "output_ranks"})
+    .Attrs({"real_k: int", "rank_id: int64_t", "wait_flag: bool"})
+    .Outputs({"output_tokens_out", "output_scores_out", "output_ranks_out"})
+    .SetInplaceMap({{"output_tokens", "output_tokens_out"},
+                    {"output_scores", "output_scores_out"},
+                    {"output_ranks", "output_ranks_out"}})
+    .SetKernelFn(PD_KERNEL(SpeculateGetOutMmsgTopK));

custom_ops/gpu_ops/speculate_decoding/speculate_logprob_utils.cu

@@ -0,0 +1,290 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "helper.h"
+
+__global__ void get_token_num_per_batch_kernel(int* next_token_num,
+                                               int* batch_token_num,
+                                               const int* seq_lens_this_time,
+                                               const int* seq_lens_encoder,
+                                               const int real_bsz) {
+    int bid = threadIdx.x;
+    if (bid < real_bsz) {
+        next_token_num[bid] =
+            seq_lens_encoder[bid] > 0 ? 1 : seq_lens_this_time[bid];
+        batch_token_num[bid] =
+            seq_lens_encoder[bid] > 0 ? 2 : seq_lens_this_time[bid];
+    }
+}
+
+template <int VecSize>
+__global__ void speculate_get_logits_kernel(float* draft_logits,
+                                            const float* logits,
+                                            const float* first_token_logits,
+                                            const int* cu_next_token_offset,
+                                            const int* cu_batch_token_offset,
+                                            const int* seq_lens_this_time,
+                                            const int* seq_lens_encoder,
+                                            const int vocab_size,
+                                            const int real_bsz) {
+    AlignedVector<float, VecSize> src_vec;
+    const int bid = blockIdx.x;
+    const int tid = threadIdx.x;
+    if (bid < real_bsz) {
+        auto* draft_logits_now =
+            draft_logits + cu_batch_token_offset[bid] * vocab_size;
+        auto* logits_now = logits + cu_next_token_offset[bid] * vocab_size;
+        for (int i = tid * VecSize; i < vocab_size; i += blockDim.x * VecSize) {
+            if (seq_lens_encoder[bid] > 0) {
+                Load<float, VecSize>(&first_token_logits[bid * vocab_size + i],
+                                     &src_vec);
+                Store<float, VecSize>(src_vec, &draft_logits_now[i]);
+
+                Load<float, VecSize>(&logits_now[i], &src_vec);
+                Store<float, VecSize>(src_vec,
+                                      &draft_logits_now[vocab_size + i]);
+            } else {
+                for (int j = 0; j < seq_lens_this_time[bid]; j++) {
+                    Load<float, VecSize>(&logits_now[j * vocab_size + i],
+                                         &src_vec);
+                    Store<float, VecSize>(
+                        src_vec, &draft_logits_now[j * vocab_size + i]);
+                }
+            }
+        }
+    }
+}
+
+void SpeculateGetLogits(const paddle::Tensor& draft_logits,
+                        const paddle::Tensor& next_token_num,
+                        const paddle::Tensor& batch_token_num,
+                        const paddle::Tensor& cu_next_token_offset,
+                        const paddle::Tensor& cu_batch_token_offset,
+                        const paddle::Tensor& logits,
+                        const paddle::Tensor& first_token_logits,
+                        const paddle::Tensor& seq_lens_this_time,
+                        const paddle::Tensor& seq_lens_encoder) {
+    auto cu_stream = seq_lens_this_time.stream();
+    const int vocab_size = logits.shape()[1];
+    const int real_bsz = seq_lens_this_time.shape()[0];
+
+    get_token_num_per_batch_kernel<<<1, 512, 0, cu_stream>>>(
+        const_cast<int*>(next_token_num.data<int>()),
+        const_cast<int*>(batch_token_num.data<int>()),
+        seq_lens_this_time.data<int>(),
+        seq_lens_encoder.data<int>(),
+        real_bsz);
+
+    void* temp_storage1 = nullptr;
+    size_t temp_storage_bytes1 = 0;
+    cub::DeviceScan::InclusiveSum(
+        temp_storage1,
+        temp_storage_bytes1,
+        batch_token_num.data<int>(),
+        const_cast<int*>(&cu_batch_token_offset.data<int>()[1]),
+        real_bsz,
+        cu_stream);
+    cudaMalloc(&temp_storage1, temp_storage_bytes1);
+    cub::DeviceScan::InclusiveSum(
+        temp_storage1,
+        temp_storage_bytes1,
+        batch_token_num.data<int>(),
+        const_cast<int*>(&cu_batch_token_offset.data<int>()[1]),
+        real_bsz,
+        cu_stream);
+
+    void* temp_storage2 = nullptr;
+    size_t temp_storage_bytes2 = 0;
+    cub::DeviceScan::InclusiveSum(
+        temp_storage2,
+        temp_storage_bytes2,
+        next_token_num.data<int>(),
+        const_cast<int*>(&cu_next_token_offset.data<int>()[1]),
+        real_bsz,
+        cu_stream);
+    cudaMalloc(&temp_storage2, temp_storage_bytes2);
+    cub::DeviceScan::InclusiveSum(
+        temp_storage2,
+        temp_storage_bytes2,
+        next_token_num.data<int>(),
+        const_cast<int*>(&cu_next_token_offset.data<int>()[1]),
+        real_bsz,
+        cu_stream);
+
+    constexpr int PackSize = VEC_16B / sizeof(float);
+    dim3 grid_dim(real_bsz);
+    dim3 block_dim(128);
+    speculate_get_logits_kernel<PackSize>
+        <<<grid_dim, block_dim, 0, cu_stream>>>(
+            const_cast<float*>(draft_logits.data<float>()),
+            logits.data<float>(),
+            first_token_logits.data<float>(),
+            cu_next_token_offset.data<int>(),
+            cu_batch_token_offset.data<int>(),
+            seq_lens_this_time.data<int>(),
+            seq_lens_encoder.data<int>(),
+            vocab_size,
+            real_bsz);
+}
+
+__global__ void speculate_insert_first_token_kernel(
+    int64_t* token_ids,
+    const int64_t* accept_tokens,
+    const int64_t* next_tokens,
+    const int* cu_next_token_offset,
+    const int* cu_batch_token_offset,
+    const int* seq_lens_this_time,
+    const int* seq_lens_encoder,
+    const int max_draft_tokens,
+    const int real_bsz) {
+    const int bid = threadIdx.x;
+
+    auto* token_ids_now = token_ids + cu_batch_token_offset[bid];
+    auto* accept_tokens_now = accept_tokens + bid * max_draft_tokens;
+    auto* next_tokens_now = next_tokens + cu_next_token_offset[bid];
+    if (seq_lens_encoder[bid] != 0) {
+        token_ids_now[0] = accept_tokens_now[0];
+        token_ids_now[1] = next_tokens_now[0];
+    } else {
+        for (int i = 0; i < seq_lens_this_time[bid]; i++) {
+            token_ids_now[i] = next_tokens_now[i];
+        }
+    }
+}
+
+void SpeculateInsertFirstToken(const paddle::Tensor& token_ids,
+                               const paddle::Tensor& accept_tokens,
+                               const paddle::Tensor& next_tokens,
+                               const paddle::Tensor& cu_next_token_offset,
+                               const paddle::Tensor& cu_batch_token_offset,
+                               const paddle::Tensor& seq_lens_this_time,
+                               const paddle::Tensor& seq_lens_encoder) {
+    auto cu_stream = seq_lens_this_time.stream();
+    const int max_draft_tokens = accept_tokens.shape()[1];
+    const int real_bsz = seq_lens_this_time.shape()[0];
+
+    speculate_insert_first_token_kernel<<<1, real_bsz, 0, cu_stream>>>(
+        const_cast<int64_t*>(token_ids.data<int64_t>()),
+        accept_tokens.data<int64_t>(),
+        next_tokens.data<int64_t>(),
+        cu_next_token_offset.data<int>(),
+        cu_batch_token_offset.data<int>(),
+        seq_lens_this_time.data<int>(),
+        seq_lens_encoder.data<int>(),
+        max_draft_tokens,
+        real_bsz);
+}
+
+template <int VecSize>
+__global__ void speculate_get_target_logits_kernel(
+    float* target_logtis,
+    const float* logits,
+    const int* cu_batch_token_offset,
+    const int* ori_cu_batch_token_offset,
+    const int* seq_lens_this_time,
+    const int* seq_lens_encoder,
+    const int* accept_num,
+    const int vocab_size,
+    const int real_bsz) {
+    AlignedVector<float, VecSize> src_vec;
+    const int bid = blockIdx.x;
+    const int tid = threadIdx.x;
+    if (bid < real_bsz) {
+        auto* target_logtis_now =
+            target_logtis + cu_batch_token_offset[bid] * vocab_size;
+        auto* logits_now = logits + ori_cu_batch_token_offset[bid] * vocab_size;
+        for (int i = tid * VecSize; i < vocab_size; i += blockDim.x * VecSize) {
+            if (seq_lens_encoder[bid] > 0) {
+                Load<float, VecSize>(&logits_now[i], &src_vec);
+                Store<float, VecSize>(src_vec, &target_logtis_now[i]);
+            } else {
+                for (int j = 0; j < accept_num[bid]; j++) {
+                    Load<float, VecSize>(&logits_now[j * vocab_size + i],
+                                         &src_vec);
+                    Store<float, VecSize>(
+                        src_vec, &target_logtis_now[j * vocab_size + i]);
+                }
+            }
+        }
+    }
+}
+
+void SpeculateGetTargetLogits(const paddle::Tensor& target_logits,
+                              const paddle::Tensor& logits,
+                              const paddle::Tensor& cu_batch_token_offset,
+                              const paddle::Tensor& ori_cu_batch_token_offset,
+                              const paddle::Tensor& seq_lens_this_time,
+                              const paddle::Tensor& seq_lens_encoder,
+                              const paddle::Tensor& accept_num) {
+    auto cu_stream = seq_lens_this_time.stream();
+    const int vocab_size = logits.shape()[1];
+    const int real_bsz = seq_lens_this_time.shape()[0];
+
+    constexpr int PackSize = VEC_16B / sizeof(float);
+    dim3 grid_dim(real_bsz);
+    dim3 block_dim(128);
+    speculate_get_target_logits_kernel<PackSize>
+        <<<grid_dim, block_dim, 0, cu_stream>>>(
+            const_cast<float*>(target_logits.data<float>()),
+            logits.data<float>(),
+            cu_batch_token_offset.data<int>(),
+            ori_cu_batch_token_offset.data<int>(),
+            seq_lens_this_time.data<int>(),
+            seq_lens_encoder.data<int>(),
+            accept_num.data<int>(),
+            vocab_size,
+            real_bsz);
+}
+
+PD_BUILD_STATIC_OP(speculate_get_logits)
+    .Inputs({"draft_logits",
+             "next_token_num",
+             "batch_token_num",
+             "cu_next_token_offset",
+             "cu_batch_token_offset",
+             "logits",
+             "first_token_logits",
+             "seq_lens_this_time",
+             "seq_lens_encoder"})
+    .Outputs({"draft_logits_out",
+              "batch_token_num_out",
+              "cu_batch_token_offset_out"})
+    .SetInplaceMap({{"draft_logits", "draft_logits_out"},
+                    {"batch_token_num", "batch_token_num_out"},
+                    {"cu_batch_token_offset", "cu_batch_token_offset_out"}})
+    .SetKernelFn(PD_KERNEL(SpeculateGetLogits));
+
+PD_BUILD_STATIC_OP(speculate_insert_first_token)
+    .Inputs({"token_ids",
+             "accept_tokens",
+             "next_tokens",
+             "cu_next_token_offset",
+             "cu_batch_token_offset",
+             "seq_lens_this_time",
+             "seq_lens_encoder"})
+    .Outputs({"token_ids_out"})
+    .SetInplaceMap({{"token_ids", "token_ids_out"}})
+    .SetKernelFn(PD_KERNEL(SpeculateInsertFirstToken));
+
+PD_BUILD_STATIC_OP(speculate_get_target_logits)
+    .Inputs({"target_logits",
+             "logits",
+             "cu_batch_token_offset",
+             "ori_cu_batch_token_offset",
+             "seq_lens_this_time",
+             "seq_lens_encoder",
+             "accept_num"})
+    .Outputs({"target_logits_out"})
+    .SetInplaceMap({{"target_logits", "target_logits_out"}})
+    .SetKernelFn(PD_KERNEL(SpeculateGetTargetLogits));

custom_ops/gpu_ops/speculate_decoding/speculate_save_output_with_topk.cc

@@ -0,0 +1,202 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <stdio.h>
+#include <string.h>
+#include <sys/ipc.h>
+#include <sys/msg.h>
+#include <sys/types.h>
+#include "paddle/extension.h"
+
+#ifndef PD_BUILD_STATIC_OP
+#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
+#endif
+
+#define MAX_BSZ 512
+#define K 20
+#define MAX_DRAFT_TOKEN_NUM 6
+
+struct batch_msgdata {
+    int tokens[MAX_DRAFT_TOKEN_NUM * (K + 1)];
+    float scores[MAX_DRAFT_TOKEN_NUM * (K + 1)];
+    int ranks[MAX_DRAFT_TOKEN_NUM];
+};
+
+struct msgdata {
+    long mtype;
+    int meta[3 + MAX_BSZ];  // stop_flag, message_flag, bsz, batch_token_nums
+    batch_msgdata mtext[MAX_BSZ];
+};
+
+void SpeculateSaveOutMmsgTopK(const paddle::Tensor& sampled_token_ids,
+                              const paddle::Tensor& logprob_token_ids,
+                              const paddle::Tensor& logprob_scores,
+                              const paddle::Tensor& logprob_ranks,
+                              const paddle::Tensor& token_num_per_batch,
+                              const paddle::Tensor& cu_batch_token_offset,
+                              const paddle::Tensor& not_need_stop,
+                              int message_flag,  // Target: 3, Draft: 4
+                              int64_t rank_id) {
+    if (rank_id > 0) {
+        return;
+    }
+    auto sampled_token_ids_cpu =
+        sampled_token_ids.copy_to(paddle::CPUPlace(), false);
+    auto logprob_token_ids_cpu =
+        logprob_token_ids.copy_to(paddle::CPUPlace(), false);
+    auto logprob_scores_cpu = logprob_scores.copy_to(paddle::CPUPlace(), false);
+    auto logprob_ranks_cpu = logprob_ranks.copy_to(paddle::CPUPlace(), false);
+    auto token_num_per_batch_cpu =
+        token_num_per_batch.copy_to(paddle::CPUPlace(), false);
+    auto cu_batch_token_offset_cpu =
+        cu_batch_token_offset.copy_to(paddle::CPUPlace(), false);
+    int64_t* sampled_token_ids_data = sampled_token_ids_cpu.data<int64_t>();
+    int64_t* logprob_token_ids_data = logprob_token_ids_cpu.data<int64_t>();
+    float* logprob_scores_data = logprob_scores_cpu.data<float>();
+    int64_t* logprob_ranks_data = logprob_ranks_cpu.data<int64_t>();
+    int* token_num_per_batch_data = token_num_per_batch_cpu.data<int>();
+    int* cu_batch_token_offset_data = cu_batch_token_offset_cpu.data<int>();
+
+    static struct msgdata msg_sed;
+    int msg_queue_id = 1;
+    if (const char* inference_msg_queue_id_env_p =
+            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+        std::string inference_msg_queue_id_env_str(
+            inference_msg_queue_id_env_p);
+        int inference_msg_queue_id_from_env =
+            std::stoi(inference_msg_queue_id_env_str);
+        msg_queue_id = inference_msg_queue_id_from_env;
+#ifdef SPECULATE_SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                  << inference_msg_queue_id_from_env << std::endl;
+#endif
+    } else {
+#ifdef SPECULATE_SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
+                  << std::endl;
+#endif
+    }
+    int inference_msg_id_from_env = 1;
+    if (const char* inference_msg_id_env_p = std::getenv("INFERENCE_MSG_ID")) {
+        std::string inference_msg_id_env_str(inference_msg_id_env_p);
+        inference_msg_id_from_env = std::stoi(inference_msg_id_env_str);
+        if (inference_msg_id_from_env == 2) {
+            // 2 and -2 is perserve for no-output indication.
+            throw std::runtime_error(
+                " INFERENCE_MSG_ID cannot be 2, please use other number.");
+        }
+        if (inference_msg_id_from_env < 0) {
+            throw std::runtime_error(
+                " INFERENCE_MSG_ID cannot be negative, please use other "
+                "number.");
+        }
+#ifdef SPECULATE_SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Your INFERENCE_MSG_ID is: " << inference_msg_id_from_env
+                  << std::endl;
+#endif
+    } else {
+#ifdef SPECULATE_SAVE_WITH_OUTPUT_DEBUG
+        std::cout
+            << "Failed to got INFERENCE_MSG_ID at env, use (int)1 as default."
+            << std::endl;
+#endif
+    }
+    static key_t key = ftok("/dev/shm", msg_queue_id);
+    static int msgid = msgget(key, IPC_CREAT | 0666);
+#ifdef SPECULATE_SAVE_WITH_OUTPUT_DEBUG
+    std::cout << "save_output_key: " << key << std::endl;
+    std::cout << "save msgid: " << msgid << std::endl;
+#endif
+    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;
+    int bsz = token_num_per_batch.shape()[0];
+    msg_sed.meta[2] = bsz;
+    int max_num_logprobs = logprob_token_ids.shape()[1];
+    for (int i = 0; i < bsz; i++) {
+        int cur_token_num = token_num_per_batch_data[i];
+        msg_sed.meta[3 + i] = cur_token_num;
+        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++) {
+            auto* cur_tokens = &cur_batch_msg_sed->tokens[j * (K + 1)];
+            auto* cur_scores = &cur_batch_msg_sed->scores[j * (K + 1)];
+            for (int k = 0; k < K + 1; k++) {
+                if (k == 0) {
+                    cur_tokens[k] =
+                        (int)sampled_token_ids_data[token_offset + j];
+                    cur_scores[k] =
+                        logprob_scores_data[(token_offset + j) * (K + 1) + k];
+                } else if (k < max_num_logprobs) {
+                    cur_tokens[k] = (int)
+                        logprob_token_ids_data[(token_offset + j) * (K + 1) +
+                                               k];
+                    cur_scores[k] =
+                        logprob_scores_data[(token_offset + j) * (K + 1) + k];
+                } else {
+                    cur_tokens[k] = -1;
+                    cur_scores[k] = 0.0;
+                }
+            }
+            cur_batch_msg_sed->ranks[j] =
+                (int)logprob_ranks_data[token_offset + j];
+        }
+    }
+#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]
+              << ", bsz: " << msg_sed.meta[2] << std::endl;
+    for (int i = 0; i < bsz; i++) {
+        int cur_token_num = msg_sed.meta[3 + i];
+        auto* cur_batch_msg_sed = &msg_sed.mtext[i];
+        std::cout << "batch " << i << " token_num: " << cur_token_num
+                  << std::endl;
+        for (int j = 0; j < cur_token_num; j++) {
+            auto* cur_tokens = &cur_batch_msg_sed->tokens[j * (K + 1)];
+            auto* cur_scores = &cur_batch_msg_sed->scores[j * (K + 1)];
+            std::cout << "tokens: ";
+            for (int k = 0; k < K + 1; k++) {
+                std::cout << cur_tokens[k] << " ";
+            }
+            std::cout << std::endl;
+            std::cout << "scores: ";
+            for (int k = 0; k < K + 1; k++) {
+                std::cout << cur_scores[k] << " ";
+            }
+            std::cout << std::endl;
+            std::cout << "ranks: " << cur_batch_msg_sed->ranks[j] << std::endl;
+        }
+    }
+    std::cout << std::endl;
+#endif
+    if (msgsnd(msgid, &msg_sed, sizeof(msg_sed) - sizeof(long), 0) == -1) {
+        printf("full msg buffer\n");
+    }
+}
+
+PD_BUILD_STATIC_OP(speculate_save_output_topk)
+    .Inputs({
+        "sampled_token_ids",
+        "logprob_token_ids",
+        "logprob_scores",
+        "logprob_ranks",
+        "token_num_per_batch",
+        "cu_batch_token_offset",
+        "not_need_stop",
+    })
+    .Attrs({"message_flag: int", "rank_id: int64_t"})
+    .SetKernelFn(PD_KERNEL(SpeculateSaveOutMmsgTopK));