Sync v2.0 version of code to github repo

This commit is contained in:
Jiang-Jia-Jun
2025-06-29 23:29:37 +00:00
parent d151496038
commit 92c2cfa2e7
597 changed files with 78776 additions and 22905 deletions

View File

@@ -1,11 +1,11 @@
// Copyright (c) 2024 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.
@@ -45,7 +45,7 @@ void find_candidate_pred_tokens(const int64_t *input_ids,
int64_t input_ids_stride,
int64_t pre_ids_stride,
int64_t draft_tokens_stride,
const int real_batch_size,
int64_t max_batch_size,
int max_ngram_size = 3,
int max_draft_tokens = 10) {
int threshold = 128;
@@ -53,13 +53,13 @@ void find_candidate_pred_tokens(const int64_t *input_ids,
if (env_var) {
threshold = std::stoi(env_var);
}
bool is_insert = false;
for (int batch_idx = 0; batch_idx < real_batch_size; batch_idx++) {
if (seq_lens_encoder[batch_idx] > 0) {
is_insert = true;
int unprocessed_batch_size = 0;
for (int batch_idx = 0; batch_idx < max_batch_size; batch_idx++) {
if (seq_lens_encoder[batch_idx] > 0 || seq_lens_decoder[batch_idx] > 0) {
unprocessed_batch_size++;
}
}
for (int batch_idx = 0; batch_idx < real_batch_size; batch_idx++) {
for (int batch_idx = 0; batch_idx < max_batch_size; batch_idx++) {
max_draft_tokens = std::min(static_cast<int64_t>(
draft_token_num[batch_idx]), max_dec_len[batch_idx] - step_idx[batch_idx] - 1);
if (seq_lens_encoder[batch_idx] > 0) {
@@ -68,26 +68,27 @@ void find_candidate_pred_tokens(const int64_t *input_ids,
seq_lens_this_time[batch_idx] = 0;
continue;
}
// printf("bid: %d. enc: %d. dec. %d\n", batch_idx, seq_lens_encoder[batch_idx], seq_lens_decoder[batch_idx]);
const int64_t *cur_input_ids = input_ids + batch_idx * input_ids_stride;
int64_t *cur_draft_tokens = draft_tokens + batch_idx * draft_tokens_stride;
const int64_t *cur_pre_ids = pre_ids + batch_idx * pre_ids_stride;
const int64_t cur_step_idx = step_idx[batch_idx];
const int64_t cur_input_ids_len = input_ids_len[batch_idx];
seq_lens_this_time[batch_idx] = 1;
if (!is_insert) {
auto sum_token_num = sum(seq_lens_this_time, batch_idx);
int left_min_token_num = real_batch_size - batch_idx;
unprocessed_batch_size--;
if (sum_token_num + max_draft_tokens + left_min_token_num > threshold) {
int tmp_max_draft_tokens = threshold - sum_token_num - left_min_token_num;
max_draft_tokens = tmp_max_draft_tokens < max_draft_tokens ? tmp_max_draft_tokens : max_draft_tokens;
}
auto sum_token_num = sum(seq_lens_this_time, batch_idx);
int left_min_token_num = unprocessed_batch_size;
if (sum_token_num + left_min_token_num >= threshold - 1) {
continue;
}
if (sum_token_num + max_draft_tokens + left_min_token_num > threshold) {
int tmp_max_draft_tokens = threshold - sum_token_num - left_min_token_num;
max_draft_tokens = tmp_max_draft_tokens < max_draft_tokens ? tmp_max_draft_tokens : max_draft_tokens;
}
if (sum_token_num + left_min_token_num >= threshold - 1) {
continue;
}
for (int ngram_size = max_ngram_size; ngram_size > 0; --ngram_size) {
// Extract the last n tokens as our search ngram
@@ -164,7 +165,6 @@ void NgramMatch(const paddle::Tensor &input_ids,
const paddle::Tensor &seq_lens_encoder,
const paddle::Tensor &seq_lens_decoder,
const paddle::Tensor &max_dec_len,
const int real_batch_size,
const int max_ngram_size,
const int max_draft_tokens) {
@@ -177,6 +177,8 @@ void NgramMatch(const paddle::Tensor &input_ids,
auto draft_tokens_shape = draft_tokens.shape();
const int64_t draft_tokens_stride = draft_tokens_shape[1];
const int64_t max_batch_size = seq_lens_this_time.shape()[0];
find_candidate_pred_tokens(input_ids.data<int64_t>(),
input_ids_len.data<int64_t>(),
pre_ids.data<int64_t>(),
@@ -190,7 +192,7 @@ void NgramMatch(const paddle::Tensor &input_ids,
input_ids_stride,
pre_ids_stride,
draft_tokens_stride,
real_batch_size,
max_batch_size,
max_ngram_size,
max_draft_tokens);
}
@@ -206,7 +208,7 @@ PD_BUILD_STATIC_OP(ngram_match)
"seq_lens_encoder",
"seq_lens_decoder",
"max_dec_len"})
.Attrs({"real_batch_size: int", "max_ngram_size: int", "max_draft_tokens: int"})
.Attrs({"max_ngram_size: int", "max_draft_tokens: int"})
.Outputs({"draft_tokens_out", "seq_lens_this_time_out"})
.SetKernelFn(PD_KERNEL(NgramMatch))
.SetInplaceMap({{"draft_tokens", "draft_tokens_out"}, {"seq_lens_this_time", "seq_lens_this_time_out"}});