saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

pack qkv in t5 decoder (#15801)

Browse source 
### Description
<!-- Describe your changes. -->

V100, b_4_s_128, max_output_len=64, beam=4

before:
t5_small: 101.28ms
t5_base:  200.07ms

after:
t5_small: 87.65ms
t5_base: 174.44ms



### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang 2023-05-15 13:45:39 -07:00 committed by GitHub
parent a8d9b29cd2
commit 3418ca28a8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 129 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
docs/ContribOperators.md
View file

@ -1130,14 +1130,14 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Custom scale will be used if specified. Default value is 1/sqrt(head_size)</dd>
</dl>
#### Inputs (3 - 10)
#### Inputs (1 - 10)
<dl>
<dt><tt>query</tt> : T</dt>
<dd>Query with shape (batch_size, 1, hidden_size)</dd>
<dt><tt>key</tt> : T</dt>
<dd>Query with shape (batch_size, 1, hidden_size) or packed QKV with shape (batch_size, 1, 2 * hidden_size + v_hidden_size)</dd>
<dt><tt>key</tt> (optional) : T</dt>
<dd>Key with shape (batch_size, 1, hidden_size) for self attention or past_key with shape (batch_size, num_heads, kv_sequence_length, head_size) for cross attention</dd>
<dt><tt>value</tt> : T</dt>
<dt><tt>value</tt> (optional) : T</dt>
<dd>Value with shape (batch_size, 1, v_hidden_size) for self attention or past_value with shape (batch_size, num_heads, kv_sequence_length, head_size) for cross attention</dd>
<dt><tt>mask_index</tt> (optional) : M</dt>
<dd>Mask values of shape (batch_size, total_sequence_length) or (batch_size, kv_sequence_length)</dd>

3
onnxruntime/contrib_ops/cpu/bert/multihead_attention.cc
View file

@ -259,7 +259,8 @@ Status MultiHeadAttention<T>::Compute(OpKernelContext* context) const {
num_heads_,
scale,
mask_filter_value_,
past_present_share_buffer));
past_present_share_buffer,
false));
const int batch_size = parameters.batch_size;
const int q_sequence_length = parameters.sequence_length;

13
onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h
View file

@ -25,7 +25,8 @@ Status CheckInputs(const T* query,
int num_heads,
float mask_filter_value,
float scale,
bool past_present_share_buffer) {
bool past_present_share_buffer,
bool dmmha_packing) {
// key_padding_mask (K/V) : (B) or (2*B + 1) or (B, L) or None
// relative_position_bias : (B, 1, S, L)
// past_key : (B, N, S*, H)
@ -41,7 +42,7 @@ Status CheckInputs(const T* query,
// value (V) : None
// bias (Q/K/V) : None
// When packed qkv is used:
// query (Q) : (B, L, N, 3, H)
// query (Q) : (B, L, N, 3, H) or (B, S, 3*D)
// key (K) : None
// value (V) : None
// bias (Q/K/V) : None or (D + D + D_v)
@ -56,7 +57,9 @@ Status CheckInputs(const T* query,
int batch_size = static_cast<int>(query_dims[0]);
int sequence_length = static_cast<int>(query_dims[1]);
int hidden_size = query_dims.size() == 3 ? static_cast<int>(query_dims[2]) : (num_heads * static_cast<int>(query_dims[4]));
int hidden_size = (query_dims.size() == 3)
? (dmmha_packing ? (static_cast<int>(query_dims[2]) / 3) : static_cast<int>(query_dims[2]))
: (num_heads * static_cast<int>(query_dims[4]));
int head_size = static_cast<int>(hidden_size) / num_heads;
int kv_sequence_length = sequence_length;
@ -331,13 +334,15 @@ Status CheckInputs(const T* query,
float mask_filter_value,
float scale,
bool past_present_share_buffer,
bool dmmha_packing,
int max_threads_per_block) {
if (max_threads_per_block > 0 && num_heads > max_threads_per_block) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "num_heads should be no larger than ", max_threads_per_block);
}
return CheckInputs(query, key, value, bias, key_padding_mask, relative_position_bias, past_key, past_value,
past_seq_len, parameters, num_heads, mask_filter_value, scale, past_present_share_buffer);
past_seq_len, parameters, num_heads, mask_filter_value, scale, past_present_share_buffer,
dmmha_packing);
}
} // namespace multihead_attention_helper

11
onnxruntime/contrib_ops/cuda/bert/decoder_masked_multihead_attention.cc
View file

@ -66,6 +66,7 @@ Status DecoderMaskedMultiHeadAttention<T1, T2>::ComputeInternal(OpKernelContext*
auto& device_prop = GetDeviceProp();
DecoderMaskedMultiHeadAttentionParams parameters;
bool is_dmmha_packing = (key == nullptr && value == nullptr);
ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckInputs<Tensor>(query,
key,
value,
@ -80,6 +81,7 @@ Status DecoderMaskedMultiHeadAttention<T1, T2>::ComputeInternal(OpKernelContext*
mask_filter_value_,
scale_,
past_present_share_buffer_,
is_dmmha_packing, // dmmha_packing
device_prop.maxThreadsPerBlock));
int batch_size = parameters.batch_size;
@ -165,9 +167,14 @@ Status DecoderMaskedMultiHeadAttention<T1, T2>::ComputeInternal(OpKernelContext*
}
parameters.is_cross_attention = false;
parameters.is_packed_qkv = is_dmmha_packing;
parameters.k = const_cast<T1*>(key->Data<T1>());
parameters.v = const_cast<T1*>(value->Data<T1>());
parameters.k = is_dmmha_packing
? const_cast<T1*>(query->Data<T1>() + parameters.hidden_size)
: const_cast<T1*>(key->Data<T1>());
parameters.v = is_dmmha_packing
? const_cast<T1*>(query->Data<T1>() + 2 * static_cast<size_t>(parameters.hidden_size))
: const_cast<T1*>(value->Data<T1>());
parameters.k_cache = present_key_data;
parameters.v_cache = present_value_data;
}

16
onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu
View file

@ -137,9 +137,9 @@ __global__ void masked_multihead_attention_kernel(DecoderMaskedMultiHeadAttentio
float qk = 0.0F;
int qkv_base_offset = params.is_mha
? bi * params.hidden_size + hi * head_size
: bi * (3 * params.hidden_size) + hi * head_size;
int qkv_base_offset = params.is_mha && !params.is_packed_qkv
? bi * params.hidden_size + hi * head_size
: bi * (3 * params.hidden_size) + hi * head_size;
const size_t bi_total_seq_length = bi * params.total_sequence_length;
@ -230,7 +230,7 @@ __global__ void masked_multihead_attention_kernel(DecoderMaskedMultiHeadAttentio
qk = dot<Qk_vec_acum, Qk_vec_k>(q, k);
if (QK_VECS_PER_WARP <= WARP_SIZE) {
#pragma unroll
#pragma unroll
for (int mask = QK_VECS_PER_WARP / 2; mask >= 1; mask /= 2) {
qk += __shfl_xor_sync(shfl_mask(QK_VECS_PER_WARP), qk, mask);
}
@ -248,9 +248,7 @@ __global__ void masked_multihead_attention_kernel(DecoderMaskedMultiHeadAttentio
qk *= inv_sqrt_dh;
if (params.relative_attention_bias != nullptr) {
qk = add_vec(qk,
reinterpret_cast<T*>(params.relative_attention_bias)[hi * params.sequence_length
* params.total_sequence_length
+ tlength]);
reinterpret_cast<T*>(params.relative_attention_bias)[hi * params.sequence_length * params.total_sequence_length + tlength]);
}
qk_max = qk;
qk_smem[tlength] = qk;
@ -351,9 +349,7 @@ __global__ void masked_multihead_attention_kernel(DecoderMaskedMultiHeadAttentio
if (ti < tlength && tidx % THREADS_PER_KEY == 0) {
if (params.relative_attention_bias != nullptr) {
qk = add_vec(qk,
reinterpret_cast<T*>(params.relative_attention_bias)[hi * params.sequence_length
* params.total_sequence_length
+ ti]);
reinterpret_cast<T*>(params.relative_attention_bias)[hi * params.sequence_length * params.total_sequence_length + ti]);
}
qk_max = fmaxf(qk_max, qk);
qk_smem[ti] = qk;

1
onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h
View file

@ -16,6 +16,7 @@ struct DecoderMaskedMultiHeadAttentionParams : AttentionParameters {
// Weather to use multihead attention(excludes matmul and bias)
bool is_mha = false;
bool is_cross_attention = false;
bool is_packed_qkv = false;
void* q = nullptr;
void* q_bias = nullptr;

1
onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
View file

@ -94,6 +94,7 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
mask_filter_value_,
scale_,
false, // past_present_share_buffer
false, // dmmha_packing
device_prop.maxThreadsPerBlock));
int sequence_length = parameters.sequence_length;

2
onnxruntime/contrib_ops/rocm/bert/multihead_attention.cu
View file

@ -71,7 +71,7 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
past_key, past_value, /*past_seq_len=*/nullptr,
&attn,
num_heads_, mask_filter_value_, scale_,
false, device_prop.maxThreadsPerBlock));
false, false, device_prop.maxThreadsPerBlock));
// TODO: support more qkv formats
ORT_ENFORCE(attn.qkv_format == Q_KV_BSNH_BSN2H || attn.qkv_format == QKV_BSN3H, "Got ", attn.qkv_format);

24
onnxruntime/core/graph/contrib_ops/bert_defs.cc
View file

@ -125,7 +125,9 @@ void RestorePaddingTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx)
}
}
void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int past_input_index) {
void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx,
int past_input_index,
bool dmmha_packing = false) {
// Output 0 has shape (batch_size, sequence_length, v_hidden_size)
// Q, K and V without packing:
@ -144,7 +146,9 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c
// Input 2 nullptr
// Packed QKV:
// Input 0 (batch_size, sequence_length, num_heads, 3, head_size)
// Input 0 (batch_size, sequence_length, num_heads, 3, head_size) or
// (batch_size, sequence_length, 3 * hidden_size))
// for DecoderMaskedMultiHeadAttention.
// Input 1 nullptr
// Input 2 nullptr
@ -184,7 +188,9 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c
ONNX_NAMESPACE::TensorShapeProto output_shape;
*output_shape.add_dim() = query_dims[0];
*output_shape.add_dim() = query_dims[1];
*output_shape.add_dim() = value_dims.size() == 3 ? value_dims[2] : value_dims[1] * value_dims[3];
*output_shape.add_dim() = value_dims.size() == 3
? (dmmha_packing ? value_dims[2] / 3 : value_dims[2])
: value_dims[1] * value_dims[3];
updateOutputShape(ctx, 0, output_shape);
return;
}
@ -609,18 +615,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
OPTIONAL_VALUE)
.Input(0,
"query",
"Query with shape (batch_size, 1, hidden_size)",
"Query with shape (batch_size, 1, hidden_size) or packed QKV with shape "
"(batch_size, 1, 2 * hidden_size + v_hidden_size)",
"T")
.Input(1,
"key",
"Key with shape (batch_size, 1, hidden_size) for self attention "
"or past_key with shape (batch_size, num_heads, kv_sequence_length, head_size) for cross attention",
"T")
"T",
OpSchema::Optional)
.Input(2,
"value",
"Value with shape (batch_size, 1, v_hidden_size) for self attention "
"or past_value with shape (batch_size, num_heads, kv_sequence_length, head_size) for cross attention",
"T")
"T",
OpSchema::Optional)
.Input(3,
"mask_index",
"Mask values of shape (batch_size, total_sequence_length) or (batch_size, kv_sequence_length)",
@ -694,7 +703,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
{"tensor(int32)"},
"Constrain mask index to integer types")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
MultiHeadAttentionTypeAndShapeInference(ctx, 5);
bool is_dmmha_packing = !hasInputShape(ctx, 1) && !hasInputShape(ctx, 2);
MultiHeadAttentionTypeAndShapeInference(ctx, 5, is_dmmha_packing);
}));
constexpr const char* MultiHeadAttention_ver1_doc = R"DOC(

94
onnxruntime/python/tools/transformers/convert_generation.py
View file

@ -1330,6 +1330,65 @@ def update_decoder_subgraph_share_buffer_and_use_decoder_masked_mha(subg: GraphP
return True
def pack_qkv_for_decoder_masked_mha(model_proto: ModelProto):
onnx_model = OnnxModel(model_proto)
output_name_to_node = onnx_model.output_name_to_node()
nodes_to_add = []
nodes_to_remove = []
for node in onnx_model.nodes():
if node.op_type == "DecoderMaskedMultiHeadAttention":
if "past_key_cross" in node.input[1] and "past_value_cross" in node.input[2]:
continue
q_matmul = output_name_to_node[node.input[0]]
k_matmul = output_name_to_node[node.input[1]]
v_matmul = output_name_to_node[node.input[2]]
q_weight = onnx_model.get_initializer(q_matmul.input[1])
k_weight = onnx_model.get_initializer(k_matmul.input[1])
v_weight = onnx_model.get_initializer(v_matmul.input[1])
if not (q_weight and k_weight and v_weight):
return False
qw = NumpyHelper.to_array(q_weight)
kw = NumpyHelper.to_array(k_weight)
vw = NumpyHelper.to_array(v_weight)
qkv_weight = np.concatenate([qw, kw, vw], axis=1)
matmul_node_name = onnx_model.create_node_name("MatMul", name_prefix="MatMul_QKV")
weight = onnx.helper.make_tensor(
name=matmul_node_name + "_weight",
data_type=TensorProto.FLOAT if q_weight.data_type == 1 else TensorProto.FLOAT16,
dims=[qkv_weight.shape[0], qkv_weight.shape[1]],
vals=qkv_weight.flatten().tolist(),
)
model_proto.graph.initializer.extend([weight])
matmul_node = onnx.helper.make_node(
"MatMul",
inputs=[q_matmul.input[0], matmul_node_name + "_weight"],
outputs=[matmul_node_name + "_out"],
name=matmul_node_name,
)
node.input[0] = matmul_node.output[0]
node.input[1] = ""
node.input[2] = ""
nodes_to_add.extend([matmul_node])
nodes_to_remove.extend([q_matmul, k_matmul, v_matmul])
onnx_model.add_nodes(nodes_to_add)
onnx_model.remove_nodes(nodes_to_remove)
onnx_model.update_graph()
onnx_model.topological_sort()
return True
def update_input_shapes_for_gpt2_decoder_model(decoder_onnx_path: str, use_external_data_format: bool = True):
"""Update the input shapes for the inputs "input_ids" and "position_ids" and make the sequence length dim value 1 for each of them.
The decoder model will be over-written.
@ -1955,21 +2014,6 @@ def convert_generation_model(args: argparse.Namespace, generation_type: Generati
encoder_model.graph.name = f"{args.model_type} encoder and decoder init"
verify_t5_encoder_decoder_init_subgraph(encoder_model.graph, args.precision)
if not args.disable_shared_initializers:
# Unique shared initializers from the decoder and decoder_init could reduce memory usage in inference.
initializers = get_shared_initializers(encoder_model, decoder_model)
logger.info(
f"{len(initializers)} shared initializers ({[i.name for i in initializers]}) in encoder and decoder subgraphs are moved to the main graph"
)
# TODO(tianleiwu): investigate the following which causes error in inference
# Move initializer from subgraph to main graph could reduce memory usage in inference.
# moved_initializers = move_initializers(encoder_model.graph)
# logger.info(
# f"{len(moved_initializers)} initializers ({[i.name for i in moved_initializers]}) from the encoder are moved to the main graph"
# )
# initializers.extend(moved_initializers)
make_dim_proto_numeric_t5(encoder_model, config)
make_dim_proto_numeric_t5(decoder_model, config)
@ -1986,6 +2030,26 @@ def convert_generation_model(args: argparse.Namespace, generation_type: Generati
else:
logger.info("*****DecoderMaskedMultiHeadAttention is not applied to T5 decoder*****")
if pack_qkv_for_decoder_masked_mha(decoder_model):
logger.info("*****pack qkv for decoder masked mha successfully!!!*****")
else:
logger.info("*****pack qkv for decoder masked mha failed!!!*****")
if not args.disable_shared_initializers:
# Unique shared initializers from the decoder and decoder_init could reduce memory usage in inference.
initializers = get_shared_initializers(encoder_model, decoder_model)
logger.info(
f"{len(initializers)} shared initializers ({[i.name for i in initializers]}) in encoder and decoder subgraphs are moved to the main graph"
)
# TODO(tianleiwu): investigate the following which causes error in inference
# Move initializer from subgraph to main graph could reduce memory usage in inference.
# moved_initializers = move_initializers(encoder_model.graph)
# logger.info(
# f"{len(moved_initializers)} initializers ({[i.name for i in moved_initializers]}) from the encoder are moved to the main graph"
# )
# initializers.extend(moved_initializers)
node.attribute.extend(
[
onnx.helper.make_attribute("encoder", encoder_model.graph),