[LLM] First commit the llm deployment code

custom_ops/gpu_ops/speculate_decoding/ngram_match.cc

@@ -0,0 +1,212 @@
+// 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.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include <vector>
+#include <string>
+#include <algorithm>
+#include <chrono>
+#include <cstdlib>
+#include "paddle/extension.h"
+
+#ifndef PD_BUILD_STATIC_OP
+#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
+#endif
+
+int sum(const int *value, int num) {
+    int sum_value = 0;
+    for (int i = 0; i <= num; i++) {
+        sum_value += value[i];
+    }
+    return sum_value;
+}
+
+void find_candidate_pred_tokens(const int64_t *input_ids,
+        const int64_t *input_ids_len,
+        const int64_t *pre_ids,
+        const int64_t *step_idx,
+        const int *draft_token_num,
+        int64_t *draft_tokens,
+        int32_t *seq_lens_this_time,
+        int32_t *seq_lens_encoder,
+        int32_t *seq_lens_decoder,
+        int64_t *max_dec_len,
+        int64_t input_ids_stride,
+        int64_t pre_ids_stride,
+        int64_t draft_tokens_stride,
+        const int real_batch_size,
+        int max_ngram_size = 3,
+        int max_draft_tokens = 10) {
+    int threshold = 128;
+    char *env_var = getenv("INFER_WITH_REFERENCE_TOKENUM_THRESHOLD");
+    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;
+        }
+    }
+    for (int batch_idx = 0; batch_idx < real_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) {
+            continue;
+        } else if (seq_lens_decoder[batch_idx] == 0) {
+            seq_lens_this_time[batch_idx] = 0;
+            continue;
+        }
+        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;
+
+            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
+            if (cur_step_idx < ngram_size) {
+                continue;
+            }
+            const int64_t *ngram = cur_pre_ids + (cur_step_idx + 1 - ngram_size);
+
+            // Iterate through sliding windows of size ngram_size
+            bool match_input = false;
+            for (int64_t i = 0; i <= cur_input_ids_len - ngram_size; ++i) {
+                // Check if the current window matches the ngram
+                bool match = true;
+                for (int j = 0; j < ngram_size; j++) {
+                    if (ngram[j] != cur_input_ids[i + j]) {
+                        match = false;
+                        break;
+                    }
+                }
+                if (match) {
+                    int64_t start_idx = i + ngram_size;
+                    int64_t end_idx = std::min(start_idx + max_draft_tokens, cur_input_ids_len);
+                    if (start_idx >= end_idx)
+                        continue;
+
+                    int64_t cur_draft_token_num = end_idx - start_idx;
+
+                    seq_lens_this_time[batch_idx] = cur_draft_token_num + 1;
+                    memcpy(cur_draft_tokens + 1, cur_input_ids + start_idx, sizeof(int64_t) * cur_draft_token_num);
+                    // To break the current batch_idx for-loop
+                    ngram_size = 0;
+                    match_input = true;
+                    break;
+                    // }
+                }
+            }
+            if (!match_input) {
+                for (int64_t i = 0; i <= cur_step_idx - ngram_size; ++i) {
+                    // Check if the current window matches the ngram
+                    bool match = true;
+
+                    for (int j = 0; j < ngram_size; j++) {
+                        if (ngram[j] != cur_pre_ids[i + j]) {
+                            match = false;
+                            break;
+                        }
+                    }
+
+                    if (match) {
+                        int64_t start_idx = i + ngram_size;
+                        int64_t end_idx = std::min(start_idx + max_draft_tokens, cur_step_idx);
+                        int64_t cur_draft_token_num = end_idx - start_idx;
+                        if (start_idx >= end_idx)
+                            continue;
+
+                        seq_lens_this_time[batch_idx] = cur_draft_token_num + 1;
+                        memcpy(cur_draft_tokens + 1, cur_pre_ids + start_idx, sizeof(int64_t) * cur_draft_token_num);
+                        ngram_size = 0;
+                        break;
+                    }
+                }
+            }
+        }
+    }
+}
+
+void NgramMatch(const paddle::Tensor &input_ids,
+        const paddle::Tensor &input_ids_len,
+        const paddle::Tensor &pre_ids,
+        const paddle::Tensor &step_idx,
+        const paddle::Tensor &draft_token_num,
+        const paddle::Tensor &draft_tokens,
+        const paddle::Tensor &seq_lens_this_time,
+        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) {
+
+    auto input_ids_shape = input_ids.shape();
+    const int64_t input_ids_stride = input_ids_shape[1];
+
+    auto pre_ids_shape = pre_ids.shape();
+    const int64_t pre_ids_stride = pre_ids_shape[1];
+
+    auto draft_tokens_shape = draft_tokens.shape();
+    const int64_t draft_tokens_stride = draft_tokens_shape[1];
+
+    find_candidate_pred_tokens(input_ids.data<int64_t>(),
+            input_ids_len.data<int64_t>(),
+            pre_ids.data<int64_t>(),
+            step_idx.data<int64_t>(),
+            draft_token_num.data<int>(),
+            const_cast<int64_t *>(draft_tokens.data<int64_t>()),
+            const_cast<int32_t *>(seq_lens_this_time.data<int32_t>()),
+            const_cast<int32_t *>(seq_lens_encoder.data<int32_t>()),
+            const_cast<int32_t *>(seq_lens_decoder.data<int32_t>()),
+            const_cast<int64_t *>(max_dec_len.data<int64_t>()),
+            input_ids_stride,
+            pre_ids_stride,
+            draft_tokens_stride,
+            real_batch_size,
+            max_ngram_size,
+            max_draft_tokens);
+}
+
+PD_BUILD_STATIC_OP(ngram_match)
+        .Inputs({"input_ids",
+                "input_ids_len",
+                "pre_ids",
+                "step_idx",
+                "draft_token_num",
+                "draft_tokens",
+                "seq_lens_this_time",
+                "seq_lens_encoder",
+                "seq_lens_decoder",
+                "max_dec_len"})
+        .Attrs({"real_batch_size: int", "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"}});