GithubStar/ncnn
1
0
Fork 0
mirror of https://github.com/Tencent/ncnn.git synced 2026-10-26 10:49:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

pnnx fuse whisper style attention (#6331)

Browse Source
This commit is contained in:
nihui
2025-09-29 16:39:38 +08:00
committed by GitHub
parent 0ecf814d52
commit 00f816ffd5
1 changed files with 166 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp
+166 -1
View File
@@ -1100,6 +1100,54 @@ pnnx.Output output 1 0 out
}
};
class fuse_multiheadattention_pass_12_2 : public fuse_multiheadattention_pass_sameqkv
{
public:
const char* match_pattern_graph() const
{
return R"PNNXIR(7767517
22 21
pnnx.Input input 0 1 input
nn.Linear op_0 1 1 input 9 bias=%qbias in_features=%embed_dim out_features=%embed_dim @bias @weight
nn.Linear op_1 1 1 input 14 bias=%kbias in_features=%embed_dim out_features=%embed_dim @bias @weight
nn.Linear op_2 1 1 input 16 bias=%vbias in_features=%embed_dim out_features=%embed_dim @bias @weight
pnnx.Expression op_3 1 1 9 10 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
Tensor.view op_4 1 1 10 11 shape=(%batch,%size,%num_heads,%feat_per_head)
Tensor.view op_5 1 1 14 15 shape=(%batch,%size,%num_heads,%feat_per_head)
Tensor.view op_6 1 1 16 17 shape=(%batch,%size,%num_heads,%feat_per_head)
torch.transpose op_7 1 1 11 12 dim0=1 dim1=2
torch.transpose op_8 1 1 15 18 dim0=1 dim1=2
torch.transpose op_9 1 1 17 20 dim0=1 dim1=2
Tensor.contiguous op_10 1 1 12 13 memory_format=*
Tensor.contiguous op_11 1 1 18 19 memory_format=*
Tensor.contiguous op_12 1 1 20 21 memory_format=*
F.scaled_dot_product_attention op_13 3 1 13 19 21 22 attn_mask=None dropout_p=0.000000e+00 is_causal=False scale=1.000000e+00
torch.transpose op_14 1 1 22 23 dim0=1 dim1=2
Tensor.contiguous op_15 1 1 23 24 memory_format=*
Tensor.reshape op_16 1 1 24 25 shape=(%batch,%size,%embed_dim)
Tensor.contiguous op_17 1 1 25 26 memory_format=*
nn.Linear out_proj 1 1 26 out bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight
pnnx.Output output 1 0 out
)PNNXIR";
}
bool match(const std::map<std::string, const Operator*>& /*matched_operators*/, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
{
const int embed_dim = captured_params.at("embed_dim").i;
const int num_heads = captured_params.at("num_heads").i;
const int feat_per_head = captured_params.at("feat_per_head").i;
const float inv_sqrt_embed_dim_per_head = captured_params.at("inv_sqrt_embed_dim_per_head").f;
if (embed_dim != num_heads * feat_per_head)
return false;
if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001))
return false;
return true;
}
};
class fuse_multiheadattention_pass_13 : public fuse_multiheadattention_pass_qkv
{
public:
@@ -1323,6 +1371,117 @@ pnnx.Output output 1 0 out
}
};
class fuse_multiheadattention_pass_16_2 : public fuse_multiheadattention_pass_sameqkv
{
public:
const char* match_pattern_graph() const
{
return R"PNNXIR(7767517
22 21
pnnx.Input input_0 0 1 input
pnnx.Input input_1 0 1 attn_mask $attn_mask=(1,1,%size,%size)f32
nn.Linear op_0 1 1 input 7 bias=%qbias in_features=%embed_dim out_features=%embed_dim @bias @weight
nn.Linear op_1 1 1 input 12 bias=%kbias in_features=%embed_dim out_features=%embed_dim @bias @weight
nn.Linear op_2 1 1 input 14 bias=%vbias in_features=%embed_dim out_features=%embed_dim @bias @weight
pnnx.Expression op_3 1 1 7 8 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
Tensor.view op_4 1 1 8 9 shape=(%batch,%size,%num_heads,%feat_per_head)
Tensor.view op_5 1 1 12 13 shape=(%batch,%size,%num_heads,%feat_per_head)
Tensor.view op_6 1 1 14 15 shape=(%batch,%size,%num_heads,%feat_per_head)
torch.transpose op_7 1 1 9 10 dim0=1 dim1=2
torch.transpose op_8 1 1 13 16 dim0=1 dim1=2
torch.transpose op_9 1 1 15 18 dim0=1 dim1=2
Tensor.contiguous op_10 1 1 10 11 memory_format=*
Tensor.contiguous op_11 1 1 16 17 memory_format=*
Tensor.contiguous op_12 1 1 18 19 memory_format=*
F.scaled_dot_product_attention op_13 4 1 11 17 19 attn_mask 20 dropout_p=0.000000e+00 is_causal=False scale=1.000000e+00
torch.transpose op_14 1 1 20 21 dim0=1 dim1=2
Tensor.contiguous op_15 1 1 21 22 memory_format=*
Tensor.reshape op_16 1 1 22 23 shape=(%batch,%size,%embed_dim)
Tensor.contiguous op_17 1 1 23 24 memory_format=*
nn.Linear out_proj 1 1 24 out bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight
pnnx.Output output 1 0 out
)PNNXIR";
}
const char* replace_pattern_graph() const
{
return R"PNNXIR(7767517
5 4
pnnx.Input input_0 0 1 input
pnnx.Input input_1 0 1 attn_mask
Tensor.reshape reshape 1 1 attn_mask attn_mask2 shape=(%size,%size)
nn.MultiheadAttention attention 2 1 input attn_mask2 out embed_dim=%embed_dim kdim=%embed_dim vdim=%embed_dim num_heads=%num_heads batch_first=True add_zero_attn=False add_bias_kv=False $attn_mask=attn_mask2
pnnx.Output output 1 0 out
)PNNXIR";
}
bool match(const std::map<std::string, const Operator*>& /*matched_operators*/, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
{
const int embed_dim = captured_params.at("embed_dim").i;
const int num_heads = captured_params.at("num_heads").i;
const int feat_per_head = captured_params.at("feat_per_head").i;
const float inv_sqrt_embed_dim_per_head = captured_params.at("inv_sqrt_embed_dim_per_head").f;
if (embed_dim != num_heads * feat_per_head)
return false;
if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001))
return false;
return true;
}
};
class fuse_multiheadattention_pass_16_3 : public fuse_multiheadattention_pass_qkv
{
public:
const char* match_pattern_graph() const
{
return R"PNNXIR(7767517
25 24
pnnx.Input input_q 0 1 query
pnnx.Input input_k 0 1 key
pnnx.Input input_v 0 1 value
nn.Linear op_0 1 1 query 7 bias=%qbias in_features=%embed_dim out_features=%embed_dim @bias @weight
nn.Linear op_1 1 1 key 12 bias=%kbias in_features=%kdim out_features=%embed_dim @bias @weight
nn.Linear op_2 1 1 value 14 bias=%vbias in_features=%vdim out_features=%embed_dim @bias @weight
pnnx.Expression op_3 1 1 7 8 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
Tensor.view op_4 1 1 8 9 shape=(%batch,%qsize,%num_heads,%feat_per_head)
Tensor.view op_5 1 1 12 13 shape=(%batch,%kvsize,%num_heads,%feat_per_head)
Tensor.view op_6 1 1 14 15 shape=(%batch,%kvsize,%num_heads,%feat_per_head)
torch.transpose op_7 1 1 9 10 dim0=1 dim1=2
torch.transpose op_8 1 1 13 16 dim0=1 dim1=2
torch.transpose op_9 1 1 15 18 dim0=1 dim1=2
Tensor.contiguous op_10 1 1 10 11 memory_format=*
Tensor.contiguous op_11 1 1 16 17 memory_format=*
Tensor.contiguous op_12 1 1 18 19 memory_format=*
F.scaled_dot_product_attention op_13 3 1 11 17 19 20 attn_mask=None dropout_p=0.000000e+00 is_causal=False scale=1.000000e+00
torch.transpose op_14 1 1 20 21 dim0=1 dim1=2
Tensor.contiguous op_15 1 1 21 22 memory_format=*
Tensor.reshape op_16 1 1 22 23 shape=(%batch,%qsize,%embed_dim)
Tensor.contiguous op_17 1 1 23 24 memory_format=*
nn.Linear out_proj 1 1 24 out bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight
pnnx.Output output 1 0 out
)PNNXIR";
}
bool match(const std::map<std::string, const Operator*>& /*matched_operators*/, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
{
const int embed_dim = captured_params.at("embed_dim").i;
const int num_heads = captured_params.at("num_heads").i;
const int feat_per_head = captured_params.at("feat_per_head").i;
const float inv_sqrt_embed_dim_per_head = captured_params.at("inv_sqrt_embed_dim_per_head").f;
if (embed_dim != num_heads * feat_per_head)
return false;
if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001))
return false;
return true;
}
};
class fuse_multiheadattention_pass_17 : public fuse_multiheadattention_pass
{
public:
@@ -1730,7 +1889,7 @@ pnnx.Output output 1 0 out
)PNNXIR";
}
bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
bool match(const std::map<std::string, const Operator*>& /*matched_operators*/, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
{
const int embed_dim = captured_params.at("embed_dim").i;
const int qkv_out_features = captured_params.at("qkv_out_features").i;
@@ -2048,10 +2207,13 @@ void fuse_multiheadattention(Graph& graph)
fuse_multiheadattention_pass_10 j;
fuse_multiheadattention_pass_12 k;
fuse_multiheadattention_pass_12_1 k1;
fuse_multiheadattention_pass_12_2 k2;
fuse_multiheadattention_pass_13 l;
fuse_multiheadattention_pass_14 m;
fuse_multiheadattention_pass_16 o;
fuse_multiheadattention_pass_16_1 o1;
fuse_multiheadattention_pass_16_2 o2;
fuse_multiheadattention_pass_16_3 o3;
fuse_multiheadattention_pass_17 p;
fuse_multiheadattention_pass_17_1 p1;
fuse_multiheadattention_pass_18 q;
@@ -2087,10 +2249,13 @@ void fuse_multiheadattention(Graph& graph)
pnnx_graph_rewrite(graph, &j, opindex);
pnnx_graph_rewrite(graph, &k, opindex);
pnnx_graph_rewrite(graph, &k1, opindex);
pnnx_graph_rewrite(graph, &k2, opindex);
pnnx_graph_rewrite(graph, &l, opindex);
pnnx_graph_rewrite(graph, &m, opindex);
pnnx_graph_rewrite(graph, &o, opindex);
pnnx_graph_rewrite(graph, &o1, opindex);
pnnx_graph_rewrite(graph, &o2, opindex);
pnnx_graph_rewrite(graph, &o3, opindex);
pnnx_graph_rewrite(graph, &p, opindex);
pnnx_graph_rewrite(graph, &p1, opindex);
pnnx_graph_rewrite(graph, &q, opindex);