merge custom op and fix code style

Author: Wanglongzhi2001

csrc/cpu/src/stop_generation_multi_ends.cc

@@ -15,20 +15,9 @@
 #include <stdlib.h>
 #include <string.h>
 
-#include "paddle/extension.h"
+#include "helper.h"
 #include <stdio.h>
 
-
-bool is_in_end(const int64_t id, const int64_t* end_ids, int length) {
- bool flag = false;
- for (int i = 0; i < length; i++) {
- if (id == end_ids[i]) {
- return true;
- }
- }
- return flag;
-}
-
 void set_value_by_flags(const bool* stop_flags,
  const int64_t* end_ids,
  int64_t* topk_ids,

csrc/gpu/get_output.cc

@@ -28,23 +28,38 @@ struct msgdata {
  int mtext[MAX_BSZ + 2]; // stop_flag, bsz, tokens
 };
 
-void GetOutputFunc(const paddle::Tensor& x,
+struct SpeculateMsgData {
+ long mtype;
+ int mtext[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2]; // stop_flag, bsz, tokens
+};
+
+static struct msgdata msg_rcv;
+static struct SpeculateMsgData specu_msg_rcv;
+
+void GetOutput(const paddle::Tensor& x,
  int64_t rank_id,
- bool wait_flag) {
+ bool wait_flag,
+ bool speculative_decoding) {
  if (rank_id > 0) return;
 
- static struct msgdata msg_rcv;
-
  static key_t key = ftok("./", 1);
 
  static int msgid = msgget(key, IPC_CREAT | 0666);
 
  int64_t *out_data = const_cast<int64_t*>(x.data<int64_t>());
  int ret = -1;
  if (!wait_flag) {
- ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, IPC_NOWAIT);
+ if (!speculative_decoding) {
+ ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, IPC_NOWAIT);
+ } else {
+ ret = msgrcv(msgid, &specu_msg_rcv, (SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2) * 4, 0, IPC_NOWAIT);
+ }
  } else {
- ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, 0);
+ if (!speculative_decoding) {
+ ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, 0);
+ } else{
+ ret = msgrcv(msgid, &specu_msg_rcv, (SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2) * 4, 0, 0);
+ }
  }
  if(ret == -1)
 {
@@ -54,61 +69,20 @@ void GetOutputFunc(const paddle::Tensor& x,
 return;
 }
 
- int bsz = msg_rcv.mtext[1];
-
- for (int64_t i = 0; i < bsz + 2; i++) {
- out_data[i] = (int64_t)msg_rcv.mtext[i];
- }
- return;
-}
-
-struct SpeculateMsgData {
- long mtype;
- int mtext[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2]; // stop_flag, bsz, tokens
-};
-
-
-void SpeculateGetOutputFunc(const paddle::Tensor& x,
- int64_t rank_id,
- bool wait_flag) {
- if (rank_id > 0) {
- return;
- }
- static struct SpeculateMsgData msg_rcv;
-
- static key_t key = ftok("./", 1);
 
- static int msgid = msgget(key, IPC_CREAT | 0666);
-
- int64_t *out_data = const_cast<int64_t*>(x.data<int64_t>());
- int ret = -1;
- if (!wait_flag) {
- ret = msgrcv(msgid, &msg_rcv, (SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2) * 4, 0, IPC_NOWAIT);
- } else {
- ret = msgrcv(msgid, &msg_rcv, (SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2) * 4, 0, 0);
- }
- if(ret == -1) {
- out_data[0] = -2;
- out_data[1] = 0;
- return;
- }
+ if (!speculative_decoding) {
  int bsz = msg_rcv.mtext[1];
-
- for (int64_t i = 0; i < SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2; i++) {
- out_data[i] = (int64_t)msg_rcv.mtext[i];
+ for (int64_t i = 0; i < bsz + 2; i++) {
+ out_data[i] = (int64_t)msg_rcv.mtext[i];
  }
- return;
-}
-
-void GetOutput(const paddle::Tensor& x,
- int64_t rank_id,
- bool wait_flag,
- bool speculative_decoding){
- if (speculative_decoding) {
- SpeculateGetOutputFunc(x, rank_id, wait_flag);
  } else {
- GetOutputFunc(x, rank_id, wait_flag);
+ int bsz = specu_msg_rcv.mtext[1];
+ for (int64_t i = 0; i < SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2; i++) {
+ out_data[i] = (int64_t)specu_msg_rcv.mtext[i];
+ }
  }
+
+ return;
 }
 
 PD_BUILD_OP(get_output)

csrc/gpu/get_padding_offset_v2.cu

@@ -23,6 +23,9 @@ __global__ void GetPaddingOffsetV2Kernel(int *padding_offset,
  const int64_t *input_data,
  const int *cum_offsets,
  const int *seq_lens,
+ const int64_t *draft_tokens,
+ const int *seq_lens_encoder,
+ const int max_draft_tokens,
  const int max_seq_len) {
  // get padding offset of each batch
  const int bi = blockIdx.x;
@@ -31,8 +34,18 @@ __global__ void GetPaddingOffsetV2Kernel(int *padding_offset,
  for (int i = ti; i < seq_lens[bi]; i += blockDim.x) {
  padding_offset[bi * max_seq_len - cum_offset + i] = cum_offset;
  const int tgt_seq_id = bi * max_seq_len - cum_offset + i;
- const int src_seq_id = bi * max_seq_len + i;
- output_data[tgt_seq_id] = input_data[src_seq_id];
+ if (draft_tokens == nullptr) {
+ const int src_seq_id = bi * max_seq_len + i;
+ output_data[tgt_seq_id] = input_data[src_seq_id];
+ } else { // speculative decoding
+ if (seq_lens_encoder[bi] > 0) {
+ const int src_seq_id = bi * max_seq_len + i;
+ output_data[tgt_seq_id] = input_data[src_seq_id];
+ } else {
+ const int src_seq_id = bi * max_draft_tokens + i;
+ output_data[tgt_seq_id] = draft_tokens[src_seq_id];
+ }
+ }
  }
  if (ti == 0) {
  if (bi == 0) {
@@ -50,7 +63,9 @@ __global__ void GetPaddingOffsetV2Kernel(int *padding_offset,
 std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
  const paddle::Tensor& cum_offsets,
  const paddle::Tensor& token_num,
- const paddle::Tensor& seq_len) {
+ const paddle::Tensor& seq_len,
+ const paddle::optional<paddle::Tensor>& draft_tokens,
+ const paddle::optional<paddle::Tensor>& seq_lens_encoder) {
  auto cu_stream = input_ids.stream();
  std::vector<int64_t> input_ids_shape = input_ids.shape();
  const int bsz = seq_len.shape()[0];
@@ -65,23 +80,46 @@ std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
  auto cu_seqlens_q = GetEmptyTensor({bsz + 1}, paddle::DataType::INT32, input_ids.place());
  auto cu_seqlens_k = GetEmptyTensor({bsz + 1}, paddle::DataType::INT32, input_ids.place());
 
- GetPaddingOffsetV2Kernel<<<bsz, 128, 0, cu_stream>>>(
- padding_offset.data<int>(), 
- cum_offsets_out.data<int>(),
- cu_seqlens_q.data<int>(),
- cu_seqlens_k.data<int>(),
- x_remove_padding.data<int64_t>(), 
- input_ids.data<int64_t>(), 
- cum_offsets.data<int>(),
- seq_len.data<int>(),
- seq_length);
+ int max_draft_tokens = 0;
+ if (draft_tokens) { // speculative decoding
+ max_draft_tokens = draft_tokens.get().shape()[1];
+ GetPaddingOffsetV2Kernel<<<bsz, 128, 0, cu_stream>>>(
+ padding_offset.data<int>(), 
+ cum_offsets_out.data<int>(),
+ cu_seqlens_q.data<int>(),
+ cu_seqlens_k.data<int>(),
+ x_remove_padding.data<int64_t>(), 
+ input_ids.data<int64_t>(), 
+ cum_offsets.data<int>(),
+ seq_len.data<int>(),
+ draft_tokens.get_ptr()->data<int64_t>(),
+ seq_lens_encoder.get_ptr()->data<int>(),
+ max_draft_tokens,
+ seq_length);
+ } else {
+ GetPaddingOffsetV2Kernel<<<bsz, 128, 0, cu_stream>>>(
+ padding_offset.data<int>(), 
+ cum_offsets_out.data<int>(),
+ cu_seqlens_q.data<int>(),
+ cu_seqlens_k.data<int>(),
+ x_remove_padding.data<int64_t>(), 
+ input_ids.data<int64_t>(), 
+ cum_offsets.data<int>(),
+ seq_len.data<int>(),
+ nullptr,
+ nullptr,
+ max_draft_tokens,
+ seq_length);
+ }
  return {x_remove_padding, cum_offsets_out, padding_offset, cu_seqlens_q, cu_seqlens_k}; // , enc_token_num, dec_token_num};
 }
 
 std::vector<std::vector<int64_t>> GetPaddingOffsetV2InferShape(const std::vector<int64_t>& input_ids_shape,
  const std::vector<int64_t>& cum_offsets_shape,
  const std::vector<int64_t>& token_num_shape,
- const std::vector<int64_t>& seq_len_shape) {
+ const std::vector<int64_t>& seq_len_shape,
+ const std::vector<int64_t>& draft_tokens_shape,
+ const std::vector<int64_t>& seq_lens_encoder_shape) {
  int64_t bsz = seq_len_shape[0];
  int64_t seq_len = input_ids_shape[1];
  return {{-1}, {bsz}, {-1}, {bsz + 1}, {bsz + 1}};
@@ -90,12 +128,14 @@ std::vector<std::vector<int64_t>> GetPaddingOffsetV2InferShape(const std::vector
 std::vector<paddle::DataType> GetPaddingOffsetV2InferDtype(const paddle::DataType& input_ids_dtype,
  const paddle::DataType& cum_offsets_dtype,
  const paddle::DataType& token_num_dtype,
- const paddle::DataType& seq_len_dtype) {
+ const paddle::DataType& seq_len_dtype,
+ const paddle::DataType& draft_tokens_dtype,
+ const paddle::DataType& seq_lens_encoder_dtype) {
  return {input_ids_dtype, seq_len_dtype, seq_len_dtype, seq_len_dtype, seq_len_dtype};
 }
 
 PD_BUILD_OP(get_padding_offset_v2)
- .Inputs({"input_ids", "cum_offsets", "token_num", "seq_len"})
+ .Inputs({"input_ids", "cum_offsets", "token_num", "seq_len", paddle::Optional("draft_tokens"), paddle::Optional("seq_lens_encoder"),})
  .Outputs({"x_remove_padding", "cum_offsets_out", "padding_offset", "cu_seqlens_q", "cu_seqlens_k"})
  .SetKernelFn(PD_KERNEL(GetPaddingOffsetV2))
  .SetInferShapeFn(PD_INFER_SHAPE(GetPaddingOffsetV2InferShape))