From 36154ba9aca3d9d2ef6d3fe47a7b4c028e92b384 Mon Sep 17 00:00:00 2001
From: Wangzheee <634486483@qq.com>
Date: Fri, 12 Nov 2021 11:11:18 +0800
Subject: [PATCH] [Paddle-Inference] fix_qkv_plugin: fix half scale (#37096)

* fix_qkv_plugin: half_scale

* [Paddle-Inference] fix_qkv_plugin: fix half scale
---
 .../fluid/inference/tensorrt/plugin/qkv_to_context_plugin.cu  | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/inference/tensorrt/plugin/qkv_to_context_plugin.cu b/paddle/fluid/inference/tensorrt/plugin/qkv_to_context_plugin.cu
index 6bae3606afe..8e59fc1355a 100644
--- a/paddle/fluid/inference/tensorrt/plugin/qkv_to_context_plugin.cu
+++ b/paddle/fluid/inference/tensorrt/plugin/qkv_to_context_plugin.cu
@@ -229,7 +229,9 @@ template <typename T>
 __global__ void apply_scale(T *data, T scale, int n) {
 #if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
   int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  data[tid] = data[tid] * scale;
+  if (tid < n) {
+    data[tid] = data[tid] * scale;
+  }
 #endif
 }
 
-- 
GitLab