Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
PaddlePaddle
Paddle
提交
a48b8e2c
P
Paddle
项目概览
PaddlePaddle
/
Paddle
大约 1 年 前同步成功
通知
2299
Star
20931
Fork
5422
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
1423
列表
看板
标记
里程碑
合并请求
543
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
P
Paddle
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
1,423
Issue
1,423
列表
看板
标记
里程碑
合并请求
543
合并请求
543
Pages
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
提交
Issue看板
未验证
提交
a48b8e2c
编写于
1月 13, 2023
作者:
W
Wang Bojun
提交者:
GitHub
1月 13, 2023
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
add oss flash fmha and fmhca support (#49438)
* add fmha_flashattention oss plugin
上级
650a0836
变更
14
展开全部
隐藏空白更改
内联
并排
Showing
14 changed file
with
2483 addition
and
0 deletion
+2483
-0
paddle/fluid/framework/ir/CMakeLists.txt
paddle/fluid/framework/ir/CMakeLists.txt
+2
-0
paddle/fluid/framework/ir/trt_cross_multihead_matmul_fuse_pass.cc
...luid/framework/ir/trt_cross_multihead_matmul_fuse_pass.cc
+587
-0
paddle/fluid/framework/ir/trt_cross_multihead_matmul_fuse_pass.h
...fluid/framework/ir/trt_cross_multihead_matmul_fuse_pass.h
+92
-0
paddle/fluid/framework/ir/trt_flash_multihead_matmul_fuse_pass.cc
...luid/framework/ir/trt_flash_multihead_matmul_fuse_pass.cc
+579
-0
paddle/fluid/framework/ir/trt_flash_multihead_matmul_fuse_pass.h
...fluid/framework/ir/trt_flash_multihead_matmul_fuse_pass.h
+91
-0
paddle/fluid/framework/ir/trt_multihead_matmul_fuse_pass.cc
paddle/fluid/framework/ir/trt_multihead_matmul_fuse_pass.cc
+3
-0
paddle/fluid/inference/api/analysis_predictor.cc
paddle/fluid/inference/api/analysis_predictor.cc
+4
-0
paddle/fluid/inference/api/paddle_pass_builder.cc
paddle/fluid/inference/api/paddle_pass_builder.cc
+3
-0
paddle/fluid/inference/tensorrt/convert/CMakeLists.txt
paddle/fluid/inference/tensorrt/convert/CMakeLists.txt
+2
-0
paddle/fluid/inference/tensorrt/convert/cross_multihead_matmul_op.cc
...d/inference/tensorrt/convert/cross_multihead_matmul_op.cc
+277
-0
paddle/fluid/inference/tensorrt/convert/flash_multihead_matmul_op.cc
...d/inference/tensorrt/convert/flash_multihead_matmul_op.cc
+190
-0
paddle/fluid/inference/tensorrt/op_teller.cc
paddle/fluid/inference/tensorrt/op_teller.cc
+6
-0
python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_cross_multihead_matmul.py
...s/ir/inference/test_trt_convert_cross_multihead_matmul.py
+326
-0
python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_flash_multihead_matmul.py
...s/ir/inference/test_trt_convert_flash_multihead_matmul.py
+321
-0
未找到文件。
paddle/fluid/framework/ir/CMakeLists.txt
浏览文件 @
a48b8e2c
...
@@ -130,6 +130,8 @@ target_link_libraries(generate_pass pass_desc_proto)
...
@@ -130,6 +130,8 @@ target_link_libraries(generate_pass pass_desc_proto)
if
(
WITH_TENSORRT
)
if
(
WITH_TENSORRT
)
pass_library
(
trt_map_matmul_to_mul_pass inference
)
pass_library
(
trt_map_matmul_to_mul_pass inference
)
pass_library
(
trt_multihead_matmul_fuse_pass inference
)
pass_library
(
trt_multihead_matmul_fuse_pass inference
)
pass_library
(
trt_flash_multihead_matmul_fuse_pass inference
)
pass_library
(
trt_cross_multihead_matmul_fuse_pass inference
)
pass_library
(
trt_skip_layernorm_fuse_pass inference
)
pass_library
(
trt_skip_layernorm_fuse_pass inference
)
pass_library
(
merge_layernorm_fuse_pass inference
)
pass_library
(
merge_layernorm_fuse_pass inference
)
pass_library
(
preln_skip_layernorm_fuse_pass inference
)
pass_library
(
preln_skip_layernorm_fuse_pass inference
)
...
...
paddle/fluid/framework/ir/trt_cross_multihead_matmul_fuse_pass.cc
0 → 100644
浏览文件 @
a48b8e2c
此差异已折叠。
点击以展开。
paddle/fluid/framework/ir/trt_cross_multihead_matmul_fuse_pass.h
0 → 100644
浏览文件 @
a48b8e2c
// Copyright (c) 2022 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.
#pragma once
#include <memory>
#include <string>
#include "paddle/fluid/framework/ir/fuse_pass_base.h"
namespace
paddle
{
namespace
framework
{
namespace
ir
{
namespace
patterns
{
struct
TrtCrossMultiHeadMatmulPattern
:
public
PatternBase
{
TrtCrossMultiHeadMatmulPattern
(
PDPattern
*
pattern
,
const
std
::
string
&
name_scope
)
:
PatternBase
(
pattern
,
name_scope
,
"cross_multihead_matmul"
)
{}
PDNode
*
operator
()();
// declare operator node's name
PATTERN_DECL_NODE
(
input0
);
PATTERN_DECL_NODE
(
input1
);
PATTERN_DECL_NODE
(
mul0
);
PATTERN_DECL_NODE
(
mul1
);
PATTERN_DECL_NODE
(
mul2
);
PATTERN_DECL_NODE
(
mul0_w
);
PATTERN_DECL_NODE
(
mul1_w
);
PATTERN_DECL_NODE
(
mul2_w
);
PATTERN_DECL_NODE
(
mul0_out
);
PATTERN_DECL_NODE
(
mul1_out
);
PATTERN_DECL_NODE
(
mul2_out
);
PATTERN_DECL_NODE
(
scale
);
PATTERN_DECL_NODE
(
scale_out
);
PATTERN_DECL_NODE
(
reshape2_0
);
PATTERN_DECL_NODE
(
reshape2_1
);
PATTERN_DECL_NODE
(
reshape2_2
);
PATTERN_DECL_NODE
(
reshape2_qkv
);
PATTERN_DECL_NODE
(
reshape2_0_out
);
PATTERN_DECL_NODE
(
reshape2_1_out
);
PATTERN_DECL_NODE
(
reshape2_2_out
);
PATTERN_DECL_NODE
(
reshape2_qkv_out
);
PATTERN_DECL_NODE
(
transpose2_0
);
PATTERN_DECL_NODE
(
transpose2_1
);
PATTERN_DECL_NODE
(
transpose2_2
);
PATTERN_DECL_NODE
(
transpose2_qkv
);
PATTERN_DECL_NODE
(
transpose2_0_out
);
PATTERN_DECL_NODE
(
transpose2_1_out
);
PATTERN_DECL_NODE
(
transpose2_2_out
);
PATTERN_DECL_NODE
(
transpose2_qkv_out
);
PATTERN_DECL_NODE
(
matmul_qk
);
PATTERN_DECL_NODE
(
matmul_qk_out
);
PATTERN_DECL_NODE
(
softmax_qk
);
PATTERN_DECL_NODE
(
softmax_qk_out
);
PATTERN_DECL_NODE
(
matmul_qkv
);
PATTERN_DECL_NODE
(
matmul_qkv_out
);
};
}
// namespace patterns
class
TrtCrossMultiHeadMatmulFusePass
:
public
FusePassBase
{
public:
TrtCrossMultiHeadMatmulFusePass
();
protected:
void
ApplyImpl
(
Graph
*
graph
)
const
;
const
std
::
string
name_scope_
{
"trt_cross_multihead_matmul_fuse"
};
private:
int
BuildCrossFusion
(
Graph
*
graph
,
const
std
::
string
&
name_scope
,
Scope
*
scope
)
const
;
};
}
// namespace ir
}
// namespace framework
}
// namespace paddle
paddle/fluid/framework/ir/trt_flash_multihead_matmul_fuse_pass.cc
0 → 100644
浏览文件 @
a48b8e2c
此差异已折叠。
点击以展开。
paddle/fluid/framework/ir/trt_flash_multihead_matmul_fuse_pass.h
0 → 100644
浏览文件 @
a48b8e2c
// Copyright (c) 2022 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.
#pragma once
#include <memory>
#include <string>
#include "paddle/fluid/framework/ir/fuse_pass_base.h"
namespace
paddle
{
namespace
framework
{
namespace
ir
{
namespace
patterns
{
struct
TrtFlashMultiHeadMatmulPattern
:
public
PatternBase
{
TrtFlashMultiHeadMatmulPattern
(
PDPattern
*
pattern
,
const
std
::
string
&
name_scope
)
:
PatternBase
(
pattern
,
name_scope
,
"flash_multihead_matmul"
)
{}
PDNode
*
operator
()();
// declare operator node's name
PATTERN_DECL_NODE
(
input0
);
PATTERN_DECL_NODE
(
mul0
);
PATTERN_DECL_NODE
(
mul1
);
PATTERN_DECL_NODE
(
mul2
);
PATTERN_DECL_NODE
(
mul0_w
);
PATTERN_DECL_NODE
(
mul1_w
);
PATTERN_DECL_NODE
(
mul2_w
);
PATTERN_DECL_NODE
(
mul0_out
);
PATTERN_DECL_NODE
(
mul1_out
);
PATTERN_DECL_NODE
(
mul2_out
);
PATTERN_DECL_NODE
(
scale
);
PATTERN_DECL_NODE
(
scale_out
);
PATTERN_DECL_NODE
(
reshape2_0
);
PATTERN_DECL_NODE
(
reshape2_1
);
PATTERN_DECL_NODE
(
reshape2_2
);
PATTERN_DECL_NODE
(
reshape2_qkv
);
PATTERN_DECL_NODE
(
reshape2_0_out
);
PATTERN_DECL_NODE
(
reshape2_1_out
);
PATTERN_DECL_NODE
(
reshape2_2_out
);
PATTERN_DECL_NODE
(
reshape2_qkv_out
);
PATTERN_DECL_NODE
(
transpose2_0
);
PATTERN_DECL_NODE
(
transpose2_1
);
PATTERN_DECL_NODE
(
transpose2_2
);
PATTERN_DECL_NODE
(
transpose2_qkv
);
PATTERN_DECL_NODE
(
transpose2_0_out
);
PATTERN_DECL_NODE
(
transpose2_1_out
);
PATTERN_DECL_NODE
(
transpose2_2_out
);
PATTERN_DECL_NODE
(
transpose2_qkv_out
);
PATTERN_DECL_NODE
(
matmul_qk
);
PATTERN_DECL_NODE
(
matmul_qk_out
);
PATTERN_DECL_NODE
(
softmax_qk
);
PATTERN_DECL_NODE
(
softmax_qk_out
);
PATTERN_DECL_NODE
(
matmul_qkv
);
PATTERN_DECL_NODE
(
matmul_qkv_out
);
};
}
// namespace patterns
class
TrtFlashMultiHeadMatmulFusePass
:
public
FusePassBase
{
public:
TrtFlashMultiHeadMatmulFusePass
();
protected:
void
ApplyImpl
(
Graph
*
graph
)
const
;
const
std
::
string
name_scope_
{
"trt_flash_multihead_matmul_fuse"
};
private:
int
BuildFlashFusion
(
Graph
*
graph
,
const
std
::
string
&
name_scope
,
Scope
*
scope
)
const
;
};
}
// namespace ir
}
// namespace framework
}
// namespace paddle
paddle/fluid/framework/ir/trt_multihead_matmul_fuse_pass.cc
浏览文件 @
a48b8e2c
...
@@ -631,6 +631,7 @@ PDNode* TrtMultiHeadMatmulV3Pattern::operator()() {
...
@@ -631,6 +631,7 @@ PDNode* TrtMultiHeadMatmulV3Pattern::operator()() {
return
transpose2_2_out_var
;
return
transpose2_2_out_var
;
}
}
}
// namespace patterns
}
// namespace patterns
void
TrtMultiHeadMatmulFusePass
::
ApplyImpl
(
Graph
*
graph
)
const
{
void
TrtMultiHeadMatmulFusePass
::
ApplyImpl
(
Graph
*
graph
)
const
{
...
@@ -1667,6 +1668,7 @@ REGISTER_PASS(trt_multihead_matmul_fuse_pass_v2,
...
@@ -1667,6 +1668,7 @@ REGISTER_PASS(trt_multihead_matmul_fuse_pass_v2,
paddle
::
framework
::
ir
::
TrtMultiHeadMatmulV2FusePass
);
paddle
::
framework
::
ir
::
TrtMultiHeadMatmulV2FusePass
);
REGISTER_PASS
(
trt_multihead_matmul_fuse_pass_v3
,
REGISTER_PASS
(
trt_multihead_matmul_fuse_pass_v3
,
paddle
::
framework
::
ir
::
TrtMultiHeadMatmulV3FusePass
);
paddle
::
framework
::
ir
::
TrtMultiHeadMatmulV3FusePass
);
REGISTER_PASS_CAPABILITY
(
trt_multihead_matmul_fuse_pass_v2
)
REGISTER_PASS_CAPABILITY
(
trt_multihead_matmul_fuse_pass_v2
)
.
AddCombination
(
.
AddCombination
(
paddle
::
framework
::
compatible
::
OpVersionComparatorCombination
()
paddle
::
framework
::
compatible
::
OpVersionComparatorCombination
()
...
@@ -1677,6 +1679,7 @@ REGISTER_PASS_CAPABILITY(trt_multihead_matmul_fuse_pass_v2)
...
@@ -1677,6 +1679,7 @@ REGISTER_PASS_CAPABILITY(trt_multihead_matmul_fuse_pass_v2)
.
EQ
(
"scale"
,
0
)
.
EQ
(
"scale"
,
0
)
.
LE
(
"matmul"
,
1
)
.
LE
(
"matmul"
,
1
)
.
EQ
(
"softmax"
,
0
));
.
EQ
(
"softmax"
,
0
));
REGISTER_PASS_CAPABILITY
(
trt_multihead_matmul_fuse_pass_v3
)
REGISTER_PASS_CAPABILITY
(
trt_multihead_matmul_fuse_pass_v3
)
.
AddCombination
(
.
AddCombination
(
paddle
::
framework
::
compatible
::
OpVersionComparatorCombination
()
paddle
::
framework
::
compatible
::
OpVersionComparatorCombination
()
...
...
paddle/fluid/inference/api/analysis_predictor.cc
浏览文件 @
a48b8e2c
...
@@ -2427,6 +2427,10 @@ USE_TRT_CONVERTER(expand_v2)
...
@@ -2427,6 +2427,10 @@ USE_TRT_CONVERTER(expand_v2)
USE_TRT_CONVERTER
(
take_along_axis
)
USE_TRT_CONVERTER
(
take_along_axis
)
USE_TRT_CONVERTER
(
skip_groupnorm_act
)
USE_TRT_CONVERTER
(
skip_groupnorm_act
)
USE_TRT_CONVERTER
(
preln_groupnorm_act
)
USE_TRT_CONVERTER
(
preln_groupnorm_act
)
#if IS_TRT_VERSION_GE(8522)
USE_TRT_CONVERTER
(
flash_multihead_matmul
)
USE_TRT_CONVERTER
(
cross_multihead_matmul
)
#endif
#if PADDLE_WITH_CUSPARSELT && IS_TRT_VERSION_GE(8000)
#if PADDLE_WITH_CUSPARSELT && IS_TRT_VERSION_GE(8000)
USE_TRT_CONVERTER
(
sparse_fc
)
USE_TRT_CONVERTER
(
sparse_fc
)
USE_TRT_CONVERTER
(
sparse_multihead_matmul
)
USE_TRT_CONVERTER
(
sparse_multihead_matmul
)
...
...
paddle/fluid/inference/api/paddle_pass_builder.cc
浏览文件 @
a48b8e2c
...
@@ -19,6 +19,7 @@
...
@@ -19,6 +19,7 @@
#ifdef PADDLE_WITH_HIP
#ifdef PADDLE_WITH_HIP
#include <miopen/miopen.h>
#include <miopen/miopen.h>
#endif
#endif
#include <glog/logging.h>
#include <glog/logging.h>
#include <algorithm>
#include <algorithm>
...
@@ -103,6 +104,8 @@ const std::vector<std::string> kTRTSubgraphPasses({
...
@@ -103,6 +104,8 @@ const std::vector<std::string> kTRTSubgraphPasses({
"trt_multihead_matmul_fuse_pass_v3"
,
//
"trt_multihead_matmul_fuse_pass_v3"
,
//
"multihead_matmul_roformer_fuse_pass"
,
//
"multihead_matmul_roformer_fuse_pass"
,
//
"constant_folding_pass"
,
//
"constant_folding_pass"
,
//
"trt_flash_multihead_matmul_fuse_pass"
,
//
"trt_cross_multihead_matmul_fuse_pass"
,
//
"vit_attention_fuse_pass"
,
//
"vit_attention_fuse_pass"
,
//
#if defined _WIN32 // Windows CI is TensorRT7.0. Remove this after upgrading.
#if defined _WIN32 // Windows CI is TensorRT7.0. Remove this after upgrading.
#else
#else
...
...
paddle/fluid/inference/tensorrt/convert/CMakeLists.txt
浏览文件 @
a48b8e2c
...
@@ -25,6 +25,8 @@ list(
...
@@ -25,6 +25,8 @@ list(
layer_norm_op.cc
layer_norm_op.cc
multihead_matmul_op.cc
multihead_matmul_op.cc
multihead_matmul_roformer_op.cc
multihead_matmul_roformer_op.cc
flash_multihead_matmul_op.cc
cross_multihead_matmul_op.cc
shuffle_channel_op.cc
shuffle_channel_op.cc
fill_any_like_op.cc
fill_any_like_op.cc
where_op.cc
where_op.cc
...
...
paddle/fluid/inference/tensorrt/convert/cross_multihead_matmul_op.cc
0 → 100644
浏览文件 @
a48b8e2c
/* Copyright (c) 2022 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 "paddle/fluid/inference/tensorrt/convert/op_converter.h"
namespace
paddle
{
namespace
inference
{
namespace
tensorrt
{
class
CrossMultiheadMatMulOpConverter
:
public
OpConverter
{
public:
void
operator
()(
const
framework
::
proto
::
OpDesc
&
op
,
const
framework
::
Scope
&
scope
,
bool
test_mode
)
override
{
VLOG
(
3
)
<<
"convert a cross_multihead_mamul op to a corresponding tensorrt "
"network structure"
;
bool
with_fp16
=
engine_
->
WithFp16
()
&&
!
engine_
->
disable_trt_plugin_fp16
();
if
(
engine_
->
precision
()
==
AnalysisConfig
::
Precision
::
kInt8
)
{
with_fp16
=
true
;
}
PADDLE_ENFORCE_EQ
(
with_fp16
,
true
,
platform
::
errors
::
Unimplemented
(
"Trt cross attention oss plugin only support fp16 mode yet."
));
framework
::
OpDesc
op_desc
(
op
,
nullptr
);
auto
*
input_q
=
engine_
->
GetITensor
(
op_desc
.
Input
(
"Input_q"
).
front
());
auto
*
input_kv
=
engine_
->
GetITensor
(
op_desc
.
Input
(
"Input_kv"
).
front
());
// auto input_dims = input->getDimensions();
auto
output_name
=
op_desc
.
Output
(
"Out"
)[
0
];
auto
weight_q_name
=
op_desc
.
Input
(
"W_q"
).
front
();
auto
*
weight_q_v
=
scope
.
FindVar
(
weight_q_name
);
auto
*
weight_q_t
=
weight_q_v
->
GetMutable
<
phi
::
DenseTensor
>
();
float
*
weight_q_data
=
nullptr
;
weight_q_data
=
const_cast
<
float
*>
(
static_cast
<
const
float
*>
(
engine_
->
GetFp32TrtWeight
(
weight_q_name
,
*
weight_q_t
).
get
().
values
));
const
auto
&
weight_q_dims
=
weight_q_t
->
dims
();
int
hidden_in_q
=
weight_q_dims
[
0
];
int
hidden_out_q
=
weight_q_dims
[
1
];
int
head_number_q
=
PADDLE_GET_CONST
(
int
,
op_desc
.
GetAttr
(
"head_number"
));
int
head_size_q
=
hidden_out_q
/
head_number_q
;
int
n_q
=
hidden_out_q
;
auto
transpose_weight_q
=
[](
const
float
*
src
,
float
*
dst
,
int
head_number
,
int
head_size
,
int
hidden_in
)
{
for
(
int
hn
=
0
;
hn
<
head_number
;
hn
++
)
{
for
(
int
hs
=
0
;
hs
<
head_size
;
hs
++
)
{
for
(
int
hi
=
0
;
hi
<
hidden_in
;
hi
++
)
{
int
out_index
=
hn
*
head_size
*
hidden_in
+
hs
*
hidden_in
+
hi
;
int
in_index
=
hi
*
head_number
*
head_size
+
hn
*
head_size
+
hs
;
dst
[
out_index
]
=
src
[
in_index
];
}
}
}
};
std
::
vector
<
float
>
weight_q_data_tmp
;
weight_q_data_tmp
.
reserve
(
weight_q_t
->
numel
());
memcpy
(
weight_q_data_tmp
.
data
(),
weight_q_data
,
weight_q_t
->
numel
()
*
sizeof
(
float
));
transpose_weight_q
(
weight_q_data_tmp
.
data
(),
weight_q_data
,
head_number_q
,
head_size_q
,
hidden_in_q
);
nvinfer1
::
Weights
weight_q
{
nvinfer1
::
DataType
::
kFLOAT
,
static_cast
<
void
*>
(
weight_q_data
),
static_cast
<
int32_t
>
(
weight_q_t
->
numel
())};
nvinfer1
::
Weights
bias_q
{};
// add shuffle for FullyConnected layer
std
::
vector
<
nvinfer1
::
ITensor
*>
reshape_before_fc_q_shape_tensor
;
nvinfer1
::
ITensor
*
input_q_shape_tensor
=
Shape
(
input_q
);
for
(
int
i
=
0
;
i
<
5
;
i
++
)
{
reshape_before_fc_q_shape_tensor
.
push_back
(
Add1DConstantLayer
(
1
));
}
for
(
int
i
=
0
;
i
<
3
;
i
++
)
{
reshape_before_fc_q_shape_tensor
[
i
]
=
GetEleTensorOfShape
(
input_q_shape_tensor
,
i
);
}
auto
*
reshape_before_fc_q_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
input_q
);
reshape_before_fc_q_layer
->
setInput
(
1
,
*
Concat
(
reshape_before_fc_q_shape_tensor
));
reshape_before_fc_q_layer
->
setName
(
(
"shuffle_before_fc_q_multihead_matmul(Output: "
+
output_name
+
")"
)
.
c_str
());
nvinfer1
::
ILayer
*
fc_q_layer
=
nullptr
;
fc_q_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
FullyConnected
,
*
reshape_before_fc_q_layer
->
getOutput
(
0
),
n_q
,
weight_q
,
bias_q
);
fc_q_layer
->
setName
(
(
"multihead_mamul_fc_q(Output: "
+
output_name
+
")"
).
c_str
());
// add shuffle for fc layer
auto
*
reshape_after_fc_q_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
fc_q_layer
->
getOutput
(
0
));
std
::
vector
<
nvinfer1
::
ITensor
*>
mha_input_q_tensor_shape
;
for
(
int
i
=
0
;
i
<
4
;
i
++
)
{
mha_input_q_tensor_shape
.
push_back
(
Add1DConstantLayer
(
1
));
}
mha_input_q_tensor_shape
[
0
]
=
GetEleTensorOfShape
(
input_q_shape_tensor
,
0
);
mha_input_q_tensor_shape
[
1
]
=
GetEleTensorOfShape
(
input_q_shape_tensor
,
1
);
mha_input_q_tensor_shape
[
2
]
=
Add1DConstantLayer
(
head_number_q
);
mha_input_q_tensor_shape
[
3
]
=
Add1DConstantLayer
(
head_size_q
);
reshape_after_fc_q_layer
->
setInput
(
1
,
*
Concat
(
mha_input_q_tensor_shape
));
reshape_after_fc_q_layer
->
setName
(
(
"shuffle_after_fc_q_multihead_matmul(Output: "
+
output_name
+
")"
)
.
c_str
());
auto
weight_kv_name
=
op_desc
.
Input
(
"W_kv"
).
front
();
auto
*
weight_kv_v
=
scope
.
FindVar
(
weight_kv_name
);
auto
*
weight_kv_t
=
weight_kv_v
->
GetMutable
<
phi
::
DenseTensor
>
();
float
*
weight_kv_data
=
nullptr
;
weight_kv_data
=
const_cast
<
float
*>
(
static_cast
<
const
float
*>
(
engine_
->
GetFp32TrtWeight
(
weight_kv_name
,
*
weight_kv_t
).
get
().
values
));
// (hidden_in, 2, hidden_out)
const
auto
&
weight_kv_dims
=
weight_kv_t
->
dims
();
int
hidden_in
=
weight_kv_dims
[
0
];
// channels_in
int
two
=
weight_kv_dims
[
1
];
// three
int
hidden_out
=
weight_kv_dims
[
2
];
// channels_out
int
head_number
=
PADDLE_GET_CONST
(
int
,
op_desc
.
GetAttr
(
"head_number"
));
int
head_size
=
hidden_out
/
head_number
;
int
n
=
two
*
hidden_out
;
nvinfer1
::
ILayer
*
layer
=
nullptr
;
// [hidden_in, 3, head_number, head_size]
// -> [head_number, 3, head_size, hidden_in]
auto
transpose_weight
=
[](
const
float
*
src
,
float
*
dst
,
int
two
,
int
head_number
,
int
head_size
,
int
hidden_in
)
{
for
(
int
hn
=
0
;
hn
<
head_number
;
hn
++
)
{
for
(
int
t
=
0
;
t
<
two
;
t
++
)
{
for
(
int
hs
=
0
;
hs
<
head_size
;
hs
++
)
{
for
(
int
hi
=
0
;
hi
<
hidden_in
;
hi
++
)
{
int
out_index
=
hn
*
two
*
head_size
*
hidden_in
+
t
*
head_size
*
hidden_in
+
hs
*
hidden_in
+
hi
;
int
in_index
=
hi
*
two
*
head_number
*
head_size
+
t
*
head_number
*
head_size
+
hn
*
head_size
+
hs
;
dst
[
out_index
]
=
src
[
in_index
];
}
}
}
}
};
std
::
vector
<
float
>
weight_kv_data_tmp
;
weight_kv_data_tmp
.
reserve
(
weight_kv_t
->
numel
());
memcpy
(
weight_kv_data_tmp
.
data
(),
weight_kv_data
,
weight_kv_t
->
numel
()
*
sizeof
(
float
));
transpose_weight
(
weight_kv_data_tmp
.
data
(),
weight_kv_data
,
two
,
head_number
,
head_size
,
hidden_in
);
nvinfer1
::
Weights
weight_kv
{
nvinfer1
::
DataType
::
kFLOAT
,
static_cast
<
void
*>
(
weight_kv_data
),
static_cast
<
int32_t
>
(
weight_kv_t
->
numel
())};
nvinfer1
::
Weights
bias_kv
{};
// add shuffle for FullyConnected layer
std
::
vector
<
nvinfer1
::
ITensor
*>
reshape_before_fc_shape_tensor
;
nvinfer1
::
ITensor
*
input_shape_tensor
=
Shape
(
input_kv
);
for
(
int
i
=
0
;
i
<
5
;
i
++
)
{
reshape_before_fc_shape_tensor
.
push_back
(
Add1DConstantLayer
(
1
));
}
for
(
int
i
=
0
;
i
<
3
;
i
++
)
{
reshape_before_fc_shape_tensor
[
i
]
=
GetEleTensorOfShape
(
input_shape_tensor
,
i
);
}
auto
*
reshape_before_fc_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
input_kv
);
reshape_before_fc_layer
->
setInput
(
1
,
*
Concat
(
reshape_before_fc_shape_tensor
));
reshape_before_fc_layer
->
setName
(
(
"shuffle_before_fc_multihead_matmul(Output: "
+
output_name
+
")"
)
.
c_str
());
nvinfer1
::
ILayer
*
fc_layer
=
nullptr
;
fc_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
FullyConnected
,
*
reshape_before_fc_layer
->
getOutput
(
0
),
n
,
weight_kv
,
bias_kv
);
fc_layer
->
setName
(
(
"multihead_mamul_fc(Output: "
+
output_name
+
")"
).
c_str
());
// add shuffle for fc layer
auto
*
reshape_after_fc_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
fc_layer
->
getOutput
(
0
));
std
::
vector
<
nvinfer1
::
ITensor
*>
mha_input_tensor_shape
;
for
(
int
i
=
0
;
i
<
5
;
i
++
)
{
mha_input_tensor_shape
.
push_back
(
Add1DConstantLayer
(
1
));
}
mha_input_tensor_shape
[
0
]
=
GetEleTensorOfShape
(
input_shape_tensor
,
0
);
mha_input_tensor_shape
[
1
]
=
GetEleTensorOfShape
(
input_shape_tensor
,
1
);
mha_input_tensor_shape
[
2
]
=
Add1DConstantLayer
(
head_number
);
mha_input_tensor_shape
[
3
]
=
Add1DConstantLayer
(
2
);
mha_input_tensor_shape
[
4
]
=
Add1DConstantLayer
(
head_size
);
reshape_after_fc_layer
->
setInput
(
1
,
*
Concat
(
mha_input_tensor_shape
));
reshape_after_fc_layer
->
setName
(
(
"shuffle_after_fc_multihead_matmul(Output: "
+
output_name
+
")"
)
.
c_str
());
auto
creator
=
GetPluginRegistry
()
->
getPluginCreator
(
"fMHCA"
,
"1"
);
assert
(
creator
!=
nullptr
);
std
::
vector
<
nvinfer1
::
PluginField
>
fields
{};
nvinfer1
::
PluginFieldCollection
*
plugin_collection
=
static_cast
<
nvinfer1
::
PluginFieldCollection
*>
(
malloc
(
sizeof
(
*
plugin_collection
)
+
fields
.
size
()
*
sizeof
(
nvinfer1
::
PluginField
)));
// remember to free
plugin_collection
->
nbFields
=
static_cast
<
int
>
(
fields
.
size
());
plugin_collection
->
fields
=
fields
.
data
();
auto
plugin
=
creator
->
createPlugin
(
"fMHA_V2"
,
plugin_collection
);
free
(
plugin_collection
);
std
::
vector
<
nvinfer1
::
ITensor
*>
plugin_inputs
;
plugin_inputs
.
emplace_back
(
reshape_after_fc_q_layer
->
getOutput
(
0
));
plugin_inputs
.
emplace_back
(
reshape_after_fc_layer
->
getOutput
(
0
));
auto
plugin_layer
=
engine_
->
network
()
->
addPluginV2
(
plugin_inputs
.
data
(),
plugin_inputs
.
size
(),
*
plugin
);
// add shuffle
nvinfer1
::
ITensor
*
batch_tensor
=
GetEleTensorOfShape
(
input_q_shape_tensor
,
0
);
nvinfer1
::
ITensor
*
length_tensor
=
GetEleTensorOfShape
(
input_q_shape_tensor
,
1
);
auto
*
reshape_after_mha_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
plugin_layer
->
getOutput
(
0
));
std
::
vector
<
nvinfer1
::
ITensor
*>
reshape_tensor
;
reshape_tensor
.
push_back
(
batch_tensor
);
reshape_tensor
.
push_back
(
length_tensor
);
reshape_tensor
.
push_back
(
Add1DConstantLayer
(
-
1
));
reshape_after_mha_layer
->
setInput
(
1
,
*
Concat
(
reshape_tensor
));
reshape_after_mha_layer
->
setName
(
(
"shuffle_last_multihead_matmul(Output: "
+
output_name
+
")"
).
c_str
());
// return
layer
=
reshape_after_mha_layer
;
RreplenishLayerAndOutput
(
layer
,
"cross_multihead_matmul"
,
{
output_name
},
test_mode
);
}
};
}
// namespace tensorrt
}
// namespace inference
}
// namespace paddle
REGISTER_TRT_OP_CONVERTER
(
cross_multihead_matmul
,
CrossMultiheadMatMulOpConverter
);
paddle/fluid/inference/tensorrt/convert/flash_multihead_matmul_op.cc
0 → 100644
浏览文件 @
a48b8e2c
/* Copyright (c) 2022 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 "paddle/fluid/inference/tensorrt/convert/op_converter.h"
namespace
paddle
{
namespace
inference
{
namespace
tensorrt
{
class
FlashMultiheadMatMulOpConverter
:
public
OpConverter
{
public:
void
operator
()(
const
framework
::
proto
::
OpDesc
&
op
,
const
framework
::
Scope
&
scope
,
bool
test_mode
)
override
{
VLOG
(
3
)
<<
"convert a flash_multihead_mamul op to a corresponding tensorrt "
"network structure"
;
bool
with_fp16
=
engine_
->
WithFp16
()
&&
!
engine_
->
disable_trt_plugin_fp16
();
if
(
engine_
->
precision
()
==
AnalysisConfig
::
Precision
::
kInt8
)
{
with_fp16
=
true
;
}
PADDLE_ENFORCE_EQ
(
with_fp16
,
true
,
platform
::
errors
::
Unimplemented
(
"Trt flash attention oss plugin only support fp16 mode yet."
));
framework
::
OpDesc
op_desc
(
op
,
nullptr
);
auto
*
input
=
engine_
->
GetITensor
(
op_desc
.
Input
(
"Input"
).
front
());
auto
weight_name
=
op_desc
.
Input
(
"W"
).
front
();
auto
*
weight_v
=
scope
.
FindVar
(
weight_name
);
auto
*
weight_t
=
weight_v
->
GetMutable
<
phi
::
DenseTensor
>
();
float
*
weight_data
=
nullptr
;
weight_data
=
const_cast
<
float
*>
(
static_cast
<
const
float
*>
(
engine_
->
GetFp32TrtWeight
(
weight_name
,
*
weight_t
).
get
().
values
));
// (hidden_in, 3, hidden_out)
const
auto
&
weight_dims
=
weight_t
->
dims
();
int
hidden_in
=
weight_dims
[
0
];
// channels_in
int
three
=
weight_dims
[
1
];
// three
int
hidden_out
=
weight_dims
[
2
];
// channels_out
int
head_number
=
PADDLE_GET_CONST
(
int
,
op_desc
.
GetAttr
(
"head_number"
));
int
head_size
=
hidden_out
/
head_number
;
int
n
=
three
*
hidden_out
;
nvinfer1
::
ILayer
*
layer
=
nullptr
;
auto
output_name
=
op_desc
.
Output
(
"Out"
)[
0
];
// [hidden_in, 3, head_number, head_size]
// -> [head_number, 3, head_size, hidden_in]
auto
transpose_weight
=
[](
const
float
*
src
,
float
*
dst
,
int
three
,
int
head_number
,
int
head_size
,
int
hidden_in
)
{
for
(
int
hn
=
0
;
hn
<
head_number
;
hn
++
)
{
for
(
int
t
=
0
;
t
<
three
;
t
++
)
{
for
(
int
hs
=
0
;
hs
<
head_size
;
hs
++
)
{
for
(
int
hi
=
0
;
hi
<
hidden_in
;
hi
++
)
{
int
out_index
=
hn
*
three
*
head_size
*
hidden_in
+
t
*
head_size
*
hidden_in
+
hs
*
hidden_in
+
hi
;
int
in_index
=
hi
*
three
*
head_number
*
head_size
+
t
*
head_number
*
head_size
+
hn
*
head_size
+
hs
;
dst
[
out_index
]
=
src
[
in_index
];
}
}
}
}
};
std
::
vector
<
float
>
weight_data_tmp
;
weight_data_tmp
.
reserve
(
weight_t
->
numel
());
memcpy
(
weight_data_tmp
.
data
(),
weight_data
,
weight_t
->
numel
()
*
sizeof
(
float
));
transpose_weight
(
weight_data_tmp
.
data
(),
weight_data
,
three
,
head_number
,
head_size
,
hidden_in
);
nvinfer1
::
Weights
weight
{
nvinfer1
::
DataType
::
kFLOAT
,
static_cast
<
void
*>
(
weight_data
),
static_cast
<
int32_t
>
(
weight_t
->
numel
())};
nvinfer1
::
Weights
bias
{};
// add shuffle for FullyConnected layer
std
::
vector
<
nvinfer1
::
ITensor
*>
reshape_before_fc_shape_tensor
;
nvinfer1
::
ITensor
*
input_shape_tensor
=
Shape
(
input
);
for
(
int
i
=
0
;
i
<
5
;
i
++
)
{
reshape_before_fc_shape_tensor
.
push_back
(
Add1DConstantLayer
(
1
));
}
for
(
int
i
=
0
;
i
<
3
;
i
++
)
{
reshape_before_fc_shape_tensor
[
i
]
=
GetEleTensorOfShape
(
input_shape_tensor
,
i
);
}
auto
*
reshape_before_fc_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
input
);
reshape_before_fc_layer
->
setInput
(
1
,
*
Concat
(
reshape_before_fc_shape_tensor
));
reshape_before_fc_layer
->
setName
(
(
"shuffle_before_fc_multihead_matmul(Output: "
+
output_name
+
")"
)
.
c_str
());
nvinfer1
::
ILayer
*
fc_layer
=
nullptr
;
fc_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
FullyConnected
,
*
reshape_before_fc_layer
->
getOutput
(
0
),
n
,
weight
,
bias
);
fc_layer
->
setName
(
(
"multihead_mamul_fc(Output: "
+
output_name
+
")"
).
c_str
());
// add shuffle for fc layer
auto
*
reshape_after_fc_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
fc_layer
->
getOutput
(
0
));
std
::
vector
<
nvinfer1
::
ITensor
*>
mha_input_tensor_shape
;
for
(
int
i
=
0
;
i
<
5
;
i
++
)
{
mha_input_tensor_shape
.
push_back
(
Add1DConstantLayer
(
1
));
}
mha_input_tensor_shape
[
0
]
=
GetEleTensorOfShape
(
input_shape_tensor
,
0
);
mha_input_tensor_shape
[
1
]
=
GetEleTensorOfShape
(
input_shape_tensor
,
1
);
mha_input_tensor_shape
[
2
]
=
Add1DConstantLayer
(
head_number
);
mha_input_tensor_shape
[
3
]
=
Add1DConstantLayer
(
3
);
mha_input_tensor_shape
[
4
]
=
Add1DConstantLayer
(
head_size
);
reshape_after_fc_layer
->
setInput
(
1
,
*
Concat
(
mha_input_tensor_shape
));
reshape_after_fc_layer
->
setName
(
(
"shuffle_after_fc_multihead_matmul(Output: "
+
output_name
+
")"
)
.
c_str
());
auto
creator
=
GetPluginRegistry
()
->
getPluginCreator
(
"fMHA_V2"
,
"1"
);
assert
(
creator
!=
nullptr
);
std
::
vector
<
nvinfer1
::
PluginField
>
fields
{};
nvinfer1
::
PluginFieldCollection
*
plugin_collection
=
static_cast
<
nvinfer1
::
PluginFieldCollection
*>
(
malloc
(
sizeof
(
*
plugin_collection
)
+
fields
.
size
()
*
sizeof
(
nvinfer1
::
PluginField
)));
// remember to free
plugin_collection
->
nbFields
=
static_cast
<
int
>
(
fields
.
size
());
plugin_collection
->
fields
=
fields
.
data
();
auto
plugin
=
creator
->
createPlugin
(
"fMHA_V2"
,
plugin_collection
);
free
(
plugin_collection
);
std
::
vector
<
nvinfer1
::
ITensor
*>
plugin_inputs
;
plugin_inputs
.
emplace_back
(
reshape_after_fc_layer
->
getOutput
(
0
));
auto
plugin_layer
=
engine_
->
network
()
->
addPluginV2
(
plugin_inputs
.
data
(),
plugin_inputs
.
size
(),
*
plugin
);
// add shuffle
nvinfer1
::
ITensor
*
batch_tensor
=
GetEleTensorOfShape
(
input_shape_tensor
,
0
);
nvinfer1
::
ITensor
*
length_tensor
=
GetEleTensorOfShape
(
input_shape_tensor
,
1
);
auto
*
reshape_after_mha_layer
=
TRT_ENGINE_ADD_LAYER
(
engine_
,
Shuffle
,
*
plugin_layer
->
getOutput
(
0
));
std
::
vector
<
nvinfer1
::
ITensor
*>
reshape_tensor
;
reshape_tensor
.
push_back
(
batch_tensor
);
reshape_tensor
.
push_back
(
length_tensor
);
reshape_tensor
.
push_back
(
Add1DConstantLayer
(
-
1
));
reshape_after_mha_layer
->
setInput
(
1
,
*
Concat
(
reshape_tensor
));
reshape_after_mha_layer
->
setName
(
(
"shuffle_last_multihead_matmul(Output: "
+
output_name
+
")"
).
c_str
());
// return
layer
=
reshape_after_mha_layer
;
RreplenishLayerAndOutput
(
layer
,
"flash_multihead_matmul"
,
{
output_name
},
test_mode
);
}
};
}
// namespace tensorrt
}
// namespace inference
}
// namespace paddle
REGISTER_TRT_OP_CONVERTER
(
flash_multihead_matmul
,
FlashMultiheadMatMulOpConverter
);
paddle/fluid/inference/tensorrt/op_teller.cc
浏览文件 @
a48b8e2c
...
@@ -66,6 +66,12 @@ struct SimpleOpTypeSetTeller : public Teller {
...
@@ -66,6 +66,12 @@ struct SimpleOpTypeSetTeller : public Teller {
teller_set
.
insert
(
"sparse_multihead_matmul"
);
teller_set
.
insert
(
"sparse_multihead_matmul"
);
int8_teller_set
.
insert
(
"sparse_multihead_matmul"
);
int8_teller_set
.
insert
(
"sparse_multihead_matmul"
);
#endif
#endif
#if IS_TRT_VERSION_GE(8522)
teller_set
.
insert
(
"flash_multihead_matmul"
);
int8_teller_set
.
insert
(
"flash_multihead_matmul"
);
teller_set
.
insert
(
"cross_multihead_matmul"
);
int8_teller_set
.
insert
(
"cross_multihead_matmul"
);
#endif
#if IS_TRT_VERSION_GE(8200)
#if IS_TRT_VERSION_GE(8200)
teller_set
.
insert
(
"round"
);
teller_set
.
insert
(
"round"
);
int8_teller_set
.
insert
(
"round"
);
int8_teller_set
.
insert
(
"round"
);
...
...
python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_cross_multihead_matmul.py
0 → 100644
浏览文件 @
a48b8e2c
# Copyright (c) 2021 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.
import
unittest
from
functools
import
partial
from
typing
import
List
import
numpy
as
np
from
program_config
import
ProgramConfig
,
TensorConfig
from
trt_layer_auto_scan_test
import
SkipReasons
,
TrtLayerAutoScanTest
import
paddle.inference
as
paddle_infer
class
TrtConvertCrossMultiHeadMatmulTest
(
TrtLayerAutoScanTest
):
def
is_program_valid
(
self
,
program_config
:
ProgramConfig
)
->
bool
:
ver
=
paddle_infer
.
get_trt_compile_version
()
if
ver
[
0
]
*
1000
+
ver
[
1
]
*
100
+
ver
[
2
]
*
10
<
8520
:
return
False
return
True
def
sample_program_configs
(
self
):
def
generate_input1
(
batch
,
dim1
):
return
np
.
random
.
random
((
batch
,
dim1
,
320
)).
astype
(
np
.
float32
)
/
10
def
generate_input2
(
batch
,
dim2
):
return
np
.
random
.
random
((
batch
,
dim2
,
768
)).
astype
(
np
.
float32
)
/
10
def
generate_weight1
():
return
np
.
random
.
random
((
320
,
320
)).
astype
(
np
.
float32
)
/
10
def
generate_weight2
():
return
np
.
random
.
random
((
768
,
320
)).
astype
(
np
.
float32
)
/
10
for
batch
in
[
1
,
2
]:
self
.
batch
=
batch
for
reshape_shape
in
[[
0
,
0
,
8
,
40
]]:
for
dim1
in
[
4096
]:
for
dim2
in
[
768
]:
dics
=
[
{
"trans_x"
:
False
,
"trans_y"
:
False
},
{
"shape"
:
reshape_shape
},
{
"axis"
:
[
0
,
2
,
1
,
3
]},
{
"trans_x"
:
False
,
"trans_y"
:
False
},
{
"shape"
:
reshape_shape
},
{
"axis"
:
[
0
,
2
,
1
,
3
]},
{
"trans_x"
:
False
,
"trans_y"
:
False
},
{
"shape"
:
reshape_shape
},
{
"axis"
:
[
0
,
2
,
1
,
3
]},
{
"trans_x"
:
False
,
"trans_y"
:
True
,
},
{
"scale"
:
0.15811388194561005
,
"bias"
:
0.0
,
"bias_after_scale"
:
True
,
},
{
"axis"
:
-
1
,
"is_test"
:
True
},
{
"trans_x"
:
False
,
"trans_y"
:
False
},
{
"axis"
:
[
0
,
2
,
1
,
3
]},
{
"shape"
:
[
0
,
0
,
320
]},
]
ops_config
=
[
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"input_data1"
],
"Y"
:
[
"mul1_weight"
],
},
"op_outputs"
:
{
"Out"
:
[
"mul1_output"
]},
"op_attrs"
:
dics
[
0
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"mul1_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"reshape21_output"
],
"XShape"
:
[
"reshape21_output_xshape"
],
},
"op_attrs"
:
dics
[
1
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"reshape21_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose21_output"
],
"XShape"
:
[
"transpose21_output_xshape"
],
},
"op_attrs"
:
dics
[
2
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"input_data2"
],
"Y"
:
[
"mul2_weight"
],
},
"op_outputs"
:
{
"Out"
:
[
"mul2_output"
]},
"op_attrs"
:
dics
[
3
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"mul2_output"
]},
"op_outputs"
:
{
"Out"
:
[
"reshape22_output"
],
"XShape"
:
[
"reshape22_output_xshape"
],
},
"op_attrs"
:
dics
[
4
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"reshape22_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose22_output"
],
"XShape"
:
[
"transpose22_output_xshape"
],
},
"op_attrs"
:
dics
[
5
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"input_data2"
],
"Y"
:
[
"mul3_weight"
],
},
"op_outputs"
:
{
"Out"
:
[
"mul3_output"
]},
"op_attrs"
:
dics
[
6
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"mul3_output"
]},
"op_outputs"
:
{
"Out"
:
[
"reshape23_output"
],
"XShape"
:
[
"reshape23_output_xshape"
],
},
"op_attrs"
:
dics
[
7
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"reshape23_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose23_output"
],
"XShape"
:
[
"transpose23_output_xshape"
],
},
"op_attrs"
:
dics
[
8
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"transpose21_output"
],
"Y"
:
[
"transpose22_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"matmul1_output"
]},
"op_attrs"
:
dics
[
9
],
},
{
"op_type"
:
"scale"
,
"op_inputs"
:
{
"X"
:
[
"matmul1_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"scale_output"
]},
"op_attrs"
:
dics
[
10
],
},
{
"op_type"
:
"softmax"
,
"op_inputs"
:
{
"X"
:
[
"scale_output"
]},
"op_outputs"
:
{
"Out"
:
[
"softmax_output"
]},
"op_attrs"
:
dics
[
11
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"softmax_output"
],
"Y"
:
[
"transpose23_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"matmul2_output"
]},
"op_attrs"
:
dics
[
12
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"matmul2_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose24_output"
],
"XShape"
:
[
"transpose24_output_xshape"
],
},
"op_attrs"
:
dics
[
13
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"transpose24_output"
]},
"op_outputs"
:
{
"Out"
:
[
"reshape24_output"
],
"XShape"
:
[
"reshape24_output_xshape"
],
},
"op_attrs"
:
dics
[
14
],
},
]
ops
=
self
.
generate_op_config
(
ops_config
)
program_config
=
ProgramConfig
(
ops
=
ops
,
weights
=
{
"mul1_weight"
:
TensorConfig
(
data_gen
=
partial
(
generate_weight1
)
),
"mul2_weight"
:
TensorConfig
(
data_gen
=
partial
(
generate_weight2
)
),
"mul3_weight"
:
TensorConfig
(
data_gen
=
partial
(
generate_weight2
)
),
},
inputs
=
{
"input_data1"
:
TensorConfig
(
data_gen
=
partial
(
generate_input1
,
batch
,
dim1
)
),
"input_data2"
:
TensorConfig
(
data_gen
=
partial
(
generate_input2
,
batch
,
dim2
)
),
},
outputs
=
[
"reshape24_output"
],
)
yield
program_config
def
sample_predictor_configs
(
self
,
program_config
)
->
(
paddle_infer
.
Config
,
List
[
int
],
float
):
def
generate_dynamic_shape
(
attrs
):
# The last dim of input1 and input2 should be static.
self
.
dynamic_shape
.
min_input_shape
=
{
"input_data1"
:
[
1
,
4096
,
320
],
"input_data2"
:
[
1
,
77
,
768
],
}
self
.
dynamic_shape
.
max_input_shape
=
{
"input_data1"
:
[
8
,
4096
,
320
],
"input_data2"
:
[
8
,
77
,
768
],
}
self
.
dynamic_shape
.
opt_input_shape
=
{
"input_data1"
:
[
2
,
4096
,
320
],
"input_data2"
:
[
2
,
77
,
768
],
}
def
clear_dynamic_shape
():
self
.
dynamic_shape
.
max_input_shape
=
{}
self
.
dynamic_shape
.
min_input_shape
=
{}
self
.
dynamic_shape
.
opt_input_shape
=
{}
attrs
=
[
program_config
.
ops
[
i
].
attrs
for
i
in
range
(
len
(
program_config
.
ops
))
]
# for static_shape
clear_dynamic_shape
()
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Float32
self
.
trt_param
.
workspace_size
=
2013265920
yield
self
.
create_inference_config
(),
(
1
,
4
),
(
1e-5
,
1e-5
)
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Half
yield
self
.
create_inference_config
(),
(
1
,
4
),
(
1e-2
,
1e-3
)
# for dynamic_shape
generate_dynamic_shape
(
attrs
)
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Float32
self
.
trt_param
.
workspace_size
=
2013265920
yield
self
.
create_inference_config
(),
(
1
,
3
),
(
1e-5
,
1e-4
)
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Half
yield
self
.
create_inference_config
(),
(
1
,
3
),
(
1e-2
,
1e-3
)
def
add_skip_trt_case
(
self
):
def
teller1
(
program_config
,
predictor_config
):
if
self
.
dynamic_shape
.
min_input_shape
==
{}:
return
True
return
False
self
.
add_skip_case
(
teller1
,
SkipReasons
.
TRT_NOT_IMPLEMENTED
,
"TThe cross attention trt oss plugin do not support static shape yet"
,
)
def
teller2
(
program_config
,
predictor_config
):
if
self
.
trt_param
.
precision
==
paddle_infer
.
PrecisionType
.
Float32
:
return
True
return
False
self
.
add_skip_case
(
teller2
,
SkipReasons
.
TRT_NOT_IMPLEMENTED
,
"The cross attention trt oss plugin do not support fp32 yet"
,
)
def
teller3
(
program_config
,
predictor_config
):
if
self
.
trt_param
.
precision
==
paddle_infer
.
PrecisionType
.
Int8
:
return
True
return
False
self
.
add_skip_case
(
teller3
,
SkipReasons
.
TRT_NOT_IMPLEMENTED
,
"The cross attention trt oss plugin do not support int8 yet."
,
)
def
test
(
self
):
self
.
add_skip_trt_case
()
self
.
run_test
()
if
__name__
==
"__main__"
:
unittest
.
main
()
python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_flash_multihead_matmul.py
0 → 100644
浏览文件 @
a48b8e2c
# Copyright (c) 2021 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.
import
unittest
from
functools
import
partial
from
typing
import
List
import
numpy
as
np
from
program_config
import
ProgramConfig
,
TensorConfig
from
trt_layer_auto_scan_test
import
SkipReasons
,
TrtLayerAutoScanTest
import
paddle.inference
as
paddle_infer
class
TrtConvertFlashMultiHeadMatmulTest
(
TrtLayerAutoScanTest
):
def
is_program_valid
(
self
,
program_config
:
ProgramConfig
)
->
bool
:
ver
=
paddle_infer
.
get_trt_compile_version
()
if
ver
[
0
]
*
1000
+
ver
[
1
]
*
100
+
ver
[
2
]
*
10
<
8520
:
return
False
return
True
def
sample_program_configs
(
self
):
def
generate_input1
(
batch
,
dim1
):
return
np
.
random
.
rand
(
batch
,
dim1
,
320
).
astype
(
np
.
float32
)
/
10
def
generate_weight1
():
return
np
.
random
.
rand
(
320
,
320
).
astype
(
np
.
float32
)
/
10
for
batch
in
[
1
,
2
]:
self
.
batch
=
batch
for
reshape_shape
in
[[
0
,
0
,
8
,
40
]]:
for
dim1
in
[
4096
]:
dics
=
[
{
"trans_x"
:
False
,
"trans_y"
:
False
},
# 0,matmul_v2_q
{
"shape"
:
reshape_shape
},
# 1,reshape_q
{
"axis"
:
[
0
,
2
,
1
,
3
],
"data_format"
:
"AnyLayout"
,
},
# 2,trans_q
{
"trans_x"
:
False
,
"trans_y"
:
False
},
# 3,matmul_v2_k
{
"shape"
:
reshape_shape
},
# 4,reshape_k
{
"axis"
:
[
0
,
2
,
1
,
3
],
"data_format"
:
"AnyLayout"
,
},
# 5,trans_k
{
"trans_x"
:
False
,
"trans_y"
:
False
},
# 6,matmul_v2_q
{
"shape"
:
reshape_shape
},
# 7,reshape_q
{
"axis"
:
[
0
,
2
,
1
,
3
],
"data_format"
:
"AnyLayout"
,
},
# 8,trans_q
{
# 9,matmul_qk
"trans_x"
:
False
,
"trans_y"
:
True
,
},
{
# 10,scale
"scale"
:
0.15811388194561005
,
"bias"
:
0.0
,
"bias_after_scale"
:
True
,
},
{
"axis"
:
-
1
,
"is_test"
:
True
},
# 11,softmax
{
"trans_x"
:
False
,
"trans_y"
:
False
},
# 12,matmul_qkv
{
"axis"
:
[
0
,
2
,
1
,
3
],
"data_format"
:
"AnyLayout"
,
},
# 13,trans_qkv
{
"shape"
:
[
0
,
0
,
320
]},
# 14,reshape_qkv
]
ops_config
=
[
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"input_data1"
],
"Y"
:
[
"mul1_weight"
],
},
"op_outputs"
:
{
"Out"
:
[
"mul1_output"
]},
"op_attrs"
:
dics
[
0
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"mul1_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"reshape21_output"
],
"XShape"
:
[
"reshape21_output_xshape"
],
},
"op_attrs"
:
dics
[
1
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"reshape21_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose21_output"
],
"XShape"
:
[
"transpose21_output_xshape"
],
},
"op_attrs"
:
dics
[
2
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"input_data1"
],
"Y"
:
[
"mul2_weight"
],
},
"op_outputs"
:
{
"Out"
:
[
"mul2_output"
]},
"op_attrs"
:
dics
[
3
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"mul2_output"
]},
"op_outputs"
:
{
"Out"
:
[
"reshape22_output"
],
"XShape"
:
[
"reshape22_output_xshape"
],
},
"op_attrs"
:
dics
[
4
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"reshape22_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose22_output"
],
"XShape"
:
[
"transpose22_output_xshape"
],
},
"op_attrs"
:
dics
[
5
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"input_data1"
],
"Y"
:
[
"mul3_weight"
],
},
"op_outputs"
:
{
"Out"
:
[
"mul3_output"
]},
"op_attrs"
:
dics
[
6
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"mul3_output"
]},
"op_outputs"
:
{
"Out"
:
[
"reshape23_output"
],
"XShape"
:
[
"reshape23_output_xshape"
],
},
"op_attrs"
:
dics
[
7
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"reshape23_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose23_output"
],
"XShape"
:
[
"transpose23_output_xshape"
],
},
"op_attrs"
:
dics
[
8
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"transpose21_output"
],
"Y"
:
[
"transpose22_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"matmul1_output"
]},
"op_attrs"
:
dics
[
9
],
},
{
"op_type"
:
"scale"
,
"op_inputs"
:
{
"X"
:
[
"matmul1_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"scale_output"
]},
"op_attrs"
:
dics
[
10
],
},
{
"op_type"
:
"softmax"
,
"op_inputs"
:
{
"X"
:
[
"scale_output"
]},
"op_outputs"
:
{
"Out"
:
[
"softmax_output"
]},
"op_attrs"
:
dics
[
11
],
},
{
"op_type"
:
"matmul_v2"
,
"op_inputs"
:
{
"X"
:
[
"softmax_output"
],
"Y"
:
[
"transpose23_output"
],
},
"op_outputs"
:
{
"Out"
:
[
"matmul2_output"
]},
"op_attrs"
:
dics
[
12
],
},
{
"op_type"
:
"transpose2"
,
"op_inputs"
:
{
"X"
:
[
"matmul2_output"
]},
"op_outputs"
:
{
"Out"
:
[
"transpose24_output"
],
"XShape"
:
[
"transpose24_output_xshape"
],
},
"op_attrs"
:
dics
[
13
],
},
{
"op_type"
:
"reshape2"
,
"op_inputs"
:
{
"X"
:
[
"transpose24_output"
]},
"op_outputs"
:
{
"Out"
:
[
"reshape24_output"
],
"XShape"
:
[
"reshape24_output_xshape"
],
},
"op_attrs"
:
dics
[
14
],
},
]
ops
=
self
.
generate_op_config
(
ops_config
)
program_config
=
ProgramConfig
(
ops
=
ops
,
weights
=
{
"mul1_weight"
:
TensorConfig
(
data_gen
=
partial
(
generate_weight1
)
),
"mul2_weight"
:
TensorConfig
(
data_gen
=
partial
(
generate_weight1
)
),
"mul3_weight"
:
TensorConfig
(
data_gen
=
partial
(
generate_weight1
)
),
},
inputs
=
{
"input_data1"
:
TensorConfig
(
data_gen
=
partial
(
generate_input1
,
batch
,
dim1
)
)
},
outputs
=
[
"reshape24_output"
],
)
yield
program_config
def
sample_predictor_configs
(
self
,
program_config
)
->
(
paddle_infer
.
Config
,
List
[
int
],
float
):
def
generate_dynamic_shape
(
attrs
):
# The last dim of input1 and input2 should be static.
self
.
dynamic_shape
.
min_input_shape
=
{
"input_data1"
:
[
1
,
4096
,
320
],
}
self
.
dynamic_shape
.
max_input_shape
=
{
"input_data1"
:
[
16
,
4096
,
320
],
}
self
.
dynamic_shape
.
opt_input_shape
=
{
"input_data1"
:
[
2
,
4096
,
320
],
}
def
clear_dynamic_shape
():
self
.
dynamic_shape
.
max_input_shape
=
{}
self
.
dynamic_shape
.
min_input_shape
=
{}
self
.
dynamic_shape
.
opt_input_shape
=
{}
attrs
=
[
program_config
.
ops
[
i
].
attrs
for
i
in
range
(
len
(
program_config
.
ops
))
]
# for static_shape
clear_dynamic_shape
()
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Float32
self
.
trt_param
.
workspace_size
=
2013265920
yield
self
.
create_inference_config
(),
(
1
,
2
),
(
1e-5
,
1e-5
)
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Half
yield
self
.
create_inference_config
(),
(
1
,
2
),
(
1e-3
,
1e-3
)
# for dynamic_shape
generate_dynamic_shape
(
attrs
)
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Float32
self
.
trt_param
.
workspace_size
=
2013265920
yield
self
.
create_inference_config
(),
(
1
,
2
),
(
1e-5
,
1e-4
)
self
.
trt_param
.
precision
=
paddle_infer
.
PrecisionType
.
Half
yield
self
.
create_inference_config
(),
(
1
,
2
),
(
1e-2
,
1e-3
)
def
add_skip_trt_case
(
self
):
def
teller1
(
program_config
,
predictor_config
):
if
self
.
dynamic_shape
.
min_input_shape
==
{}:
return
True
return
False
self
.
add_skip_case
(
teller1
,
SkipReasons
.
TRT_NOT_IMPLEMENTED
,
"TThe flash attention trt oss plugin do not support static shape yet"
,
)
def
teller2
(
program_config
,
predictor_config
):
if
self
.
trt_param
.
precision
==
paddle_infer
.
PrecisionType
.
Float32
:
return
True
return
False
self
.
add_skip_case
(
teller2
,
SkipReasons
.
TRT_NOT_IMPLEMENTED
,
"The flash attention trt oss plugin do not support fp32 yet"
,
)
def
teller3
(
program_config
,
predictor_config
):
if
self
.
trt_param
.
precision
==
paddle_infer
.
PrecisionType
.
Int8
:
return
True
return
False
self
.
add_skip_case
(
teller3
,
SkipReasons
.
TRT_NOT_IMPLEMENTED
,
"The flash attention trt oss plugin do not support int8 yet."
,
)
def
test
(
self
):
self
.
add_skip_trt_case
()
self
.
run_test
()
if
__name__
==
"__main__"
:
unittest
.
main
()
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录