mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2026-04-23 00:17:25 +08:00
@@ -30,30 +30,30 @@ __global__ void GetPaddingOffsetSystemKernel(int64_t* output_data,
|
||||
const int* seq_lens,
|
||||
const int* seq_lens_merged,
|
||||
const int max_seq_len) {
|
||||
// get padding offset of each batch
|
||||
const int bi = blockIdx.x;
|
||||
const int real_bi = seq_mapping[bi];
|
||||
const int ti = threadIdx.x;
|
||||
const int cum_offset = bi == 0 ? 0 : cum_offsets[bi - 1];
|
||||
// get padding offset of each batch
|
||||
const int bi = blockIdx.x;
|
||||
const int real_bi = seq_mapping[bi];
|
||||
const int ti = threadIdx.x;
|
||||
const int cum_offset = bi == 0 ? 0 : cum_offsets[bi - 1];
|
||||
|
||||
const int base_offset = bi * max_seq_len - cum_offset;
|
||||
const int src_base_offset = real_bi * max_seq_len;
|
||||
for (int i = ti; i < seq_lens[bi]; i += blockDim.x) {
|
||||
const int offset_now = base_offset + i;
|
||||
padding_offset[offset_now] = cum_offset;
|
||||
output_data[offset_now] = input_data[src_base_offset + i];
|
||||
}
|
||||
const int base_offset = bi * max_seq_len - cum_offset;
|
||||
const int src_base_offset = real_bi * max_seq_len;
|
||||
for (int i = ti; i < seq_lens[bi]; i += blockDim.x) {
|
||||
const int offset_now = base_offset + i;
|
||||
padding_offset[offset_now] = cum_offset;
|
||||
output_data[offset_now] = input_data[src_base_offset + i];
|
||||
}
|
||||
|
||||
const int seq_len_merged = seq_lens_merged[bi];
|
||||
const int cum_offset_merged = bi == 0 ? 0 : cum_offsets_merged[bi - 1];
|
||||
const int merged_base_offset = bi * max_seq_len - cum_offset_merged;
|
||||
for (int i = ti; i < seq_len_merged; i += blockDim.x) {
|
||||
padding_offset_merged[merged_base_offset + i] = cum_offset_merged;
|
||||
}
|
||||
if (ti == 0) {
|
||||
cum_offsets_out[bi] = cum_offset;
|
||||
cum_offsets_out_merged[bi] = cum_offset_merged;
|
||||
}
|
||||
const int seq_len_merged = seq_lens_merged[bi];
|
||||
const int cum_offset_merged = bi == 0 ? 0 : cum_offsets_merged[bi - 1];
|
||||
const int merged_base_offset = bi * max_seq_len - cum_offset_merged;
|
||||
for (int i = ti; i < seq_len_merged; i += blockDim.x) {
|
||||
padding_offset_merged[merged_base_offset + i] = cum_offset_merged;
|
||||
}
|
||||
if (ti == 0) {
|
||||
cum_offsets_out[bi] = cum_offset;
|
||||
cum_offsets_out_merged[bi] = cum_offset_merged;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<paddle::Tensor> GetPaddingOffsetSystem(
|
||||
@@ -64,41 +64,41 @@ std::vector<paddle::Tensor> GetPaddingOffsetSystem(
|
||||
const paddle::Tensor& seq_len,
|
||||
const paddle::Tensor& seq_len_merged,
|
||||
const paddle::Tensor& seq_mapping) {
|
||||
auto cu_stream = input_ids.stream();
|
||||
std::vector<int64_t> input_ids_shape = input_ids.shape();
|
||||
const int bsz = seq_len.shape()[0];
|
||||
const int seq_length = input_ids_shape[1];
|
||||
auto cum_offsets_out = cum_offsets.copy_to(cum_offsets.place(), false);
|
||||
auto cum_offsets_merged_out =
|
||||
cum_offsets_merged.copy_to(cum_offsets_merged.place(), false);
|
||||
auto cpu_token_num = token_num.copy_to(paddle::CPUPlace(), false);
|
||||
auto cu_stream = input_ids.stream();
|
||||
std::vector<int64_t> input_ids_shape = input_ids.shape();
|
||||
const int bsz = seq_len.shape()[0];
|
||||
const int seq_length = input_ids_shape[1];
|
||||
auto cum_offsets_out = cum_offsets.copy_to(cum_offsets.place(), false);
|
||||
auto cum_offsets_merged_out =
|
||||
cum_offsets_merged.copy_to(cum_offsets_merged.place(), false);
|
||||
auto cpu_token_num = token_num.copy_to(paddle::CPUPlace(), false);
|
||||
|
||||
const int token_num_data = cpu_token_num.data<int64_t>()[0];
|
||||
auto x_remove_padding = paddle::full(
|
||||
{token_num_data}, 0, paddle::DataType::INT64, input_ids.place());
|
||||
auto padding_offset = paddle::full(
|
||||
{token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
|
||||
auto padding_offset_merged = paddle::full(
|
||||
{token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
|
||||
int blockSize = min((token_num_data + 32 - 1) / 32 * 32, 128);
|
||||
GetPaddingOffsetSystemKernel<<<bsz, 1024, 0, cu_stream>>>(
|
||||
x_remove_padding.data<int64_t>(),
|
||||
padding_offset.data<int>(),
|
||||
padding_offset_merged.data<int>(),
|
||||
cum_offsets_out.data<int>(),
|
||||
cum_offsets_merged_out.data<int>(),
|
||||
input_ids.data<int64_t>(),
|
||||
seq_mapping.data<int>(),
|
||||
cum_offsets.data<int>(),
|
||||
cum_offsets_merged.data<int>(),
|
||||
seq_len.data<int>(),
|
||||
seq_len_merged.data<int>(),
|
||||
seq_length);
|
||||
return {x_remove_padding,
|
||||
cum_offsets_out,
|
||||
cum_offsets_merged_out,
|
||||
padding_offset,
|
||||
padding_offset_merged}; // , enc_token_num, dec_token_num};
|
||||
const int token_num_data = cpu_token_num.data<int64_t>()[0];
|
||||
auto x_remove_padding = paddle::full(
|
||||
{token_num_data}, 0, paddle::DataType::INT64, input_ids.place());
|
||||
auto padding_offset = paddle::full(
|
||||
{token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
|
||||
auto padding_offset_merged = paddle::full(
|
||||
{token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
|
||||
int blockSize = min((token_num_data + 32 - 1) / 32 * 32, 128);
|
||||
GetPaddingOffsetSystemKernel<<<bsz, 1024, 0, cu_stream>>>(
|
||||
x_remove_padding.data<int64_t>(),
|
||||
padding_offset.data<int>(),
|
||||
padding_offset_merged.data<int>(),
|
||||
cum_offsets_out.data<int>(),
|
||||
cum_offsets_merged_out.data<int>(),
|
||||
input_ids.data<int64_t>(),
|
||||
seq_mapping.data<int>(),
|
||||
cum_offsets.data<int>(),
|
||||
cum_offsets_merged.data<int>(),
|
||||
seq_len.data<int>(),
|
||||
seq_len_merged.data<int>(),
|
||||
seq_length);
|
||||
return {x_remove_padding,
|
||||
cum_offsets_out,
|
||||
cum_offsets_merged_out,
|
||||
padding_offset,
|
||||
padding_offset_merged}; // , enc_token_num, dec_token_num};
|
||||
}
|
||||
|
||||
std::vector<std::vector<int64_t>> GetPaddingOffsetSystemInferShape(
|
||||
@@ -109,8 +109,8 @@ std::vector<std::vector<int64_t>> GetPaddingOffsetSystemInferShape(
|
||||
const std::vector<int64_t>& seq_len_shape,
|
||||
const std::vector<int64_t>& seq_len_merged_shape,
|
||||
const std::vector<int64_t>& seq_mapping_shape) {
|
||||
int64_t bsz = seq_len_shape[0];
|
||||
return {{-1}, {bsz}, {bsz}, {-1}, {-1}};
|
||||
int64_t bsz = seq_len_shape[0];
|
||||
return {{-1}, {bsz}, {bsz}, {-1}, {-1}};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> GetPaddingOffsetSystemInferDtype(
|
||||
@@ -121,11 +121,11 @@ std::vector<paddle::DataType> GetPaddingOffsetSystemInferDtype(
|
||||
const paddle::DataType& seq_len_dtype,
|
||||
const paddle::DataType& seq_len_merged_dtype,
|
||||
const paddle::DataType& seq_mapping_dtype) {
|
||||
return {input_ids_dtype,
|
||||
seq_len_dtype,
|
||||
seq_len_dtype,
|
||||
seq_len_dtype,
|
||||
seq_len_dtype};
|
||||
return {input_ids_dtype,
|
||||
seq_len_dtype,
|
||||
seq_len_dtype,
|
||||
seq_len_dtype,
|
||||
seq_len_dtype};
|
||||
}
|
||||
|
||||
PD_BUILD_STATIC_OP(get_padding_offset_system)
|
||||
|
||||
Reference in New Issue
Block a user