Make FG adapt to latest paddle (#6727)

guoshengCS · web-flow · commit 8c01eebfba35 · 2023-08-15T11:57:11.000+08:00
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_bart_decoding_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_bart_decoding_op.cu
@@ -22,7 +22,7 @@ limitations under the License. */
 #include <vector>
 
 #include "cublas_handle.h"
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fusion_bart_decoding_op.h"
 #include "pd_traits.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_decoder_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_decoder_op.cu
@@ -21,7 +21,7 @@ limitations under the License. */
 #include <sstream>
 #include <vector>
 
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fusion_decoder_op.h"
 #include "pd_traits.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_decoding_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_decoding_op.cu
@@ -22,7 +22,7 @@ limitations under the License. */
 #include <sstream>
 #include <vector>
 
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fusion_decoding_op.h"
 #include "pd_traits.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_encoder_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_encoder_op.cu
@@ -24,7 +24,7 @@ limitations under the License. */
 
 #include "cublas_handle.h"
 #include "fastertransformer/bert_encoder_transformer.h"
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fastertransformer/cuda/cuda_kernels.h"
 #include "fastertransformer/standard_encoder.h"
 #include "fusion_encoder_op.h"
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_force_decoding_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_force_decoding_op.cu
@@ -22,7 +22,7 @@ limitations under the License. */
 #include <sstream>
 #include <vector>
 
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fusion_force_decoding_op.h"
 #include "pd_traits.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_gpt_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_gpt_op.cu
@@ -28,7 +28,7 @@
 #ifdef HOST
 #undef HOST
 #endif
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fastertransformer/gpt.h"
 #include "fastertransformer/utils/common.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_gptj_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_gptj_op.cu
@@ -31,7 +31,7 @@ limitations under the License. */
 #ifdef HOST
 #undef HOST
 #endif
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fastertransformer/utils/common.h"
 
 #ifdef BUILD_GPT  // consistent with FasterTransformer
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_mbart_decoding_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_mbart_decoding_op.cu
@@ -22,7 +22,7 @@ limitations under the License. */
 #include <vector>
 
 #include "cublas_handle.h"
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fusion_mbart_decoding_op.h"
 #include "pd_traits.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_miro_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_miro_op.cu
@@ -27,7 +27,7 @@ limitations under the License. */
 #ifdef HOST
 #undef HOST
 #endif
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fastertransformer/decoding_beamsearch.h"
 #include "fastertransformer/decoding_sampling.h"
 #include "fastertransformer/utils/common.h"
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_opt_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_opt_op.cu
@@ -28,7 +28,7 @@
 #ifdef HOST
 #undef HOST
 #endif
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fastertransformer/opt.h"
 #include "fastertransformer/utils/common.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_pegasus_decoding_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_pegasus_decoding_op.cu
@@ -22,7 +22,7 @@ limitations under the License. */
 #include <vector>
 
 #include "cublas_handle.h"
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fusion_pegasus_decoding_op.h"
 #include "pd_traits.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_t5_decoding_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_t5_decoding_op.cu
@@ -22,7 +22,7 @@ limitations under the License. */
 #include <sstream>
 #include <vector>
 
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fusion_t5_decoding_op.h"
 #include "pd_traits.h"
 
diff --git a/paddlenlp/ops/fast_transformer/src/fusion_unified_decoding_op.cu b/paddlenlp/ops/fast_transformer/src/fusion_unified_decoding_op.cu
@@ -27,7 +27,7 @@ limitations under the License. */
 #ifdef HOST
 #undef HOST
 #endif
-#include "fastertransformer/cuda/cub/cub.cuh"
+
 #include "fastertransformer/decoding_beamsearch.h"
 #include "fastertransformer/decoding_sampling.h"
 #include "fastertransformer/utils/common.h"
diff --git a/paddlenlp/ops/fast_transformer/transformer/decoding.py b/paddlenlp/ops/fast_transformer/transformer/decoding.py
@@ -34,19 +34,14 @@ def run_custom(op_name, inputs_names, inputs_var, attrs_names, attrs_val, output
     ret = []
 
     if paddle.in_dynamic_mode():
-        # TODO(guosheng): wrong _run_custom_op usage in latest paddle, fix it
-        ctx = core.CustomOpKernelContext()
-
-        for ins in inputs_var:
-            ctx.add_inputs(ins)
-        for ats in attrs_val:
-            ctx.add_attr(ats)
-
-        for name in outputs_names:
-            ret.append(core.eager.Tensor())
-            ctx.add_outputs(ret[-1])
-
-        core.eager._run_custom_op(ctx, op_name, True)
+        new_inputs_var = []
+        for k, v in zip(inputs_names, inputs_var):
+            if not k.endswith("@VECTOR") and isinstance(v, (list, tuple)) and len(v) == 1:
+                new_inputs_var.append(v[0])
+            else:
+                new_inputs_var.append(v)
+        outs = core.eager._run_custom_op(op_name, *new_inputs_var, *attrs_val)
+        return outs[0] if len(outs) == 1 else outs
     else:
         inputs = dict(zip(inputs_names, inputs_var))
         attrs = dict(zip(attrs_names, attrs_val))
