mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-11 17:48:34 +00:00
6856619b18
* add kernel interface * register kernel * add self/cross qkv projection without cache * add LaunchTransQkv2 for (S,B,X,N,H) -> (X,B,N,S,H) * refactor ConcatPastToPresent * DecoderQkvToContext interface * q,k,v buffer and cache as output * qk, pv and transctx * fix compiler error on linux machine * key_padding_mask * add test_parity file. However not runnable * add partial unittest * made partial attributes to inputs * --gen_doc * change kernel interface, add more tests * morre parity tests * fix test * fix typo * transpose optimizer has bug. remove it temporarily * add input shape checks * add type/shape inference * fix cache shape check * fix rocm build failure * fix rocm build error * review comments * review comments
423 lines
No EOL
19 KiB
Python
423 lines
No EOL
19 KiB
Python
# --------------------------------------------------------------------------
|
|
# Copyright 2020 The HuggingFace Inc. team
|
|
#
|
|
# 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
|
|
# --------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License. See License.txt in the project root for
|
|
# license information.
|
|
# -------------------------------------------------------------------------
|
|
|
|
import math
|
|
import numpy
|
|
import torch
|
|
from torch import Tensor, nn
|
|
from torch.nn import functional as F
|
|
from typing import Dict, List, Optional, Tuple
|
|
import os
|
|
|
|
torch.manual_seed(0)
|
|
|
|
"""
|
|
This is an example of export bart decoder attention with huggingface v3.5.1
|
|
def my_bart_attention_forward(
|
|
self,
|
|
query,
|
|
key: Tensor,
|
|
key_padding_mask: Optional[Tensor],
|
|
layer_state: Optional[List[Tensor]],
|
|
attn_mask: Optional[Tensor] = None,
|
|
output_attentions: bool=False,
|
|
use_past=torch.tensor(False),
|
|
):
|
|
static_kv: bool = self.encoder_decoder_attention
|
|
q_weight = self.q_proj.weight.transpose(0,1)
|
|
q_weight = q_weight.reshape(self.embed_dim, self.embed_dim)
|
|
|
|
kv_weight = torch.stack((self.k_v_proj.k_proj.weight.transpose(0,1), self.k_v_proj.v_proj.weight.transpose(0,1)), dim=1)
|
|
kv_weight = kv_weight.reshape(self.embed_dim, 2 * self.embed_dim)
|
|
|
|
bias = torch.stack((self.q_proj.bias, self.k_v_proj.k_proj.bias, self.k_v_proj.v_proj.bias), dim=0)
|
|
bias = bias.reshape(3 * self.embed_dim)
|
|
|
|
self_p_k, self_p_v, enc_dec_p_k, enc_dec_p_v = layer_state
|
|
if static_kv:
|
|
key_cache, value_cache = enc_dec_p_k, enc_dec_p_v
|
|
else:
|
|
key_cache, value_cache = self_p_k, self_p_v
|
|
|
|
if not static_kv:
|
|
key_padding_mask = torch.tensor(False)
|
|
|
|
attn_output, new_key_cache, new_value_cache = torch.ops.onnxruntime.DecoderAttention(
|
|
query,
|
|
key,
|
|
q_weight,
|
|
kv_weight,
|
|
bias,
|
|
key_padding_mask,
|
|
key_cache,
|
|
value_cache,
|
|
torch.tensor(static_kv), #static_kv
|
|
use_past, #use_past
|
|
torch.tensor(True), #has_layer_state
|
|
torch.tensor(static_kv), #has_key_padding_mask
|
|
self.num_heads)
|
|
|
|
if not use_past:
|
|
if self.encoder_decoder_attention:
|
|
layer_state[2] = new_key_cache
|
|
layer_state[3] = new_value_cache
|
|
else:
|
|
layer_state[0] = new_key_cache
|
|
layer_state[1] = new_value_cache
|
|
else:
|
|
if not self.encoder_decoder_attention:
|
|
layer_state[0] = new_key_cache
|
|
layer_state[1] = new_value_cache
|
|
|
|
attn_output = self.out_proj(attn_output)
|
|
|
|
return attn_output, None, layer_state
|
|
"""
|
|
|
|
class Config:
|
|
batch_size = 0
|
|
sequence_length = 0
|
|
kv_sequence_length = 0
|
|
num_heads = 0
|
|
head_size = 0
|
|
embed_dim = 0
|
|
|
|
def __init__(self, b, s, s2, n, h):
|
|
self.batch_size = b
|
|
self.sequence_length = s
|
|
self.kv_sequence_length = s2
|
|
self.num_heads = n
|
|
self.head_size = h
|
|
self.embed_dim = self.num_heads * self.head_size
|
|
|
|
class AttentionProjection(nn.Module):
|
|
def __init__(self, num_heads, head_dim, embed_dim, bias=True):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = head_dim
|
|
self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
|
|
def shape_state(self, state, batch_size):
|
|
return state.view(batch_size * self.num_heads, -1, self.head_dim)
|
|
|
|
def shape_proj(self, proj, batch_size):
|
|
return proj.view(-1, batch_size * self.num_heads, self.head_dim).transpose(0, 1)
|
|
|
|
def forward(
|
|
self,
|
|
query,
|
|
key,
|
|
layer_state: Optional[List[Tensor]],
|
|
encoder_decoder_attention: bool,
|
|
use_past=torch.tensor(False),
|
|
):
|
|
bsz = torch._shape_as_tensor(query)[1]
|
|
if layer_state is None or not use_past:
|
|
if not encoder_decoder_attention:
|
|
k = self.k_proj(query)
|
|
v = self.v_proj(query)
|
|
k = self.shape_proj(k, bsz)
|
|
v = self.shape_proj(v, bsz)
|
|
else:
|
|
k = self.k_proj(key)
|
|
v = self.v_proj(key)
|
|
k = self.shape_proj(k, bsz)
|
|
v = self.shape_proj(v, bsz)
|
|
else:
|
|
self_p_k, self_p_v, enc_dec_p_k, enc_dec_p_v = layer_state
|
|
if not encoder_decoder_attention:
|
|
k = self.k_proj(query)
|
|
v = self.v_proj(query)
|
|
k = self.shape_proj(k, bsz)
|
|
v = self.shape_proj(v, bsz)
|
|
k = torch.cat([self.shape_state(self_p_k, bsz), k], dim=1)
|
|
v = torch.cat([self.shape_state(self_p_v, bsz), v], dim=1)
|
|
else:
|
|
k = self.shape_state(enc_dec_p_k, bsz)
|
|
v = self.shape_state(enc_dec_p_v, bsz)
|
|
|
|
return k, v
|
|
|
|
class AttentionForONNX(nn.Module):
|
|
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
|
|
|
def __init__(
|
|
self,
|
|
embed_dim,
|
|
num_heads,
|
|
dropout=0.0,
|
|
bias=True,
|
|
encoder_decoder_attention=False, # otherwise self_attention
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.num_heads = num_heads
|
|
self.dropout = dropout
|
|
self.head_dim = embed_dim // num_heads
|
|
assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
|
|
self.scaling = self.head_dim ** -0.5
|
|
|
|
self.encoder_decoder_attention = encoder_decoder_attention
|
|
self.k_v_proj = torch.jit.script(AttentionProjection(num_heads, self.head_dim, embed_dim, bias))
|
|
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
self.cache_key = "encoder_decoder" if self.encoder_decoder_attention else "self"
|
|
|
|
def _shape(self, tensor, seq_len, bsz):
|
|
return tensor.contiguous().view(seq_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
|
|
|
|
def forward(
|
|
self,
|
|
query,
|
|
key: Tensor,
|
|
key_padding_mask: Optional[Tensor] = None,
|
|
layer_state: Optional[List[Tensor]] = None,
|
|
attn_mask: Optional[Tensor] = None,
|
|
output_attentions: bool=False,
|
|
use_past=torch.tensor(False),
|
|
has_key_padding_mask: bool=False
|
|
) -> Tuple[Tensor, Optional[Tensor]]:
|
|
"""Input shape: Time(SeqLen) x Batch x Channel"""
|
|
static_kv: bool = self.encoder_decoder_attention
|
|
tgt_len, bsz, embed_dim = query.size()
|
|
# get here for encoder decoder cause of static_kv
|
|
k, v = self.k_v_proj(query, key, layer_state, self.encoder_decoder_attention, use_past)
|
|
|
|
q = self.q_proj(query) * self.scaling
|
|
q = self._shape(q, tgt_len, bsz)
|
|
|
|
# Update cache
|
|
if layer_state is not None:
|
|
cached_shape = (bsz, self.num_heads, -1, self.head_dim) # bsz must be first for reorder_cache
|
|
if static_kv:
|
|
# cross-attn
|
|
new_key_cache = k.view(*cached_shape)
|
|
new_value_cache = v.view(*cached_shape)
|
|
else:
|
|
# self-attn
|
|
new_key_cache = k.view(*cached_shape)
|
|
new_value_cache = v.view(*cached_shape)
|
|
|
|
src_len = k.size(1)
|
|
assert key_padding_mask is None or key_padding_mask.shape == (bsz, src_len)
|
|
attn_weights = torch.bmm(q, k.transpose(1, 2))
|
|
assert attn_weights.size() == (bsz * self.num_heads, tgt_len, src_len)
|
|
|
|
if has_key_padding_mask: # don't attend to padding symbols
|
|
attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
|
|
reshaped = key_padding_mask.unsqueeze(1).unsqueeze(2)
|
|
attn_weights = attn_weights.masked_fill(reshaped, float("-inf"))
|
|
attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
|
|
|
|
attn_weights = F.softmax(attn_weights, dim=-1)
|
|
attn_probs = attn_weights
|
|
|
|
assert v is not None
|
|
attn_output = torch.bmm(attn_probs, v)
|
|
assert attn_output.size() == (bsz * self.num_heads, tgt_len, self.head_dim)
|
|
attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
|
|
attn_output = self.out_proj(attn_output)
|
|
|
|
return attn_output, new_key_cache, new_value_cache
|
|
|
|
def ORT_forward(
|
|
self,
|
|
query,
|
|
key: Tensor,
|
|
key_padding_mask: Optional[Tensor] = None,
|
|
layer_state: Optional[List[Tensor]] = None,
|
|
attn_mask: Optional[Tensor] = None,
|
|
output_attentions: bool=False,
|
|
use_past=torch.tensor(False),
|
|
has_key_padding_mask: bool=False
|
|
) -> Tuple[Tensor, Optional[Tensor]]:
|
|
"""Input shape: Time(SeqLen) x Batch x Channel"""
|
|
# For readability
|
|
static_kv = True if self.encoder_decoder_attention else False
|
|
has_layer_state = True if layer_state is not None else False
|
|
use_past_cache = True if use_past else False
|
|
|
|
q_weight = self.q_proj.weight.transpose(0,1)
|
|
q_weight = q_weight.reshape(self.embed_dim, self.embed_dim)
|
|
|
|
kv_weight = torch.stack((self.k_v_proj.k_proj.weight.transpose(0,1), self.k_v_proj.v_proj.weight.transpose(0,1)), dim=1)
|
|
kv_weight = kv_weight.reshape(self.embed_dim, 2 * self.embed_dim)
|
|
|
|
bias = torch.stack((self.q_proj.bias, self.k_v_proj.k_proj.bias, self.k_v_proj.v_proj.bias), dim=0)
|
|
bias = bias.reshape(3 * self.embed_dim)
|
|
|
|
onnx_model_str = create_decoder_attention_graph(query, key, q_weight, kv_weight, bias, self.num_heads, static_kv, use_past_cache, has_layer_state, has_key_padding_mask)
|
|
|
|
self_p_k, self_p_v, enc_dec_p_k, enc_dec_p_v = layer_state
|
|
if self.encoder_decoder_attention:
|
|
key_cache, value_cache = enc_dec_p_k, enc_dec_p_v
|
|
else:
|
|
key_cache, value_cache = self_p_k, self_p_v
|
|
|
|
ort_inputs = {
|
|
'query': numpy.ascontiguousarray(query.cpu().numpy()),
|
|
'key': numpy.ascontiguousarray(key.cpu().numpy()),
|
|
'key_padding_mask': numpy.ascontiguousarray(key_padding_mask.cpu().numpy()),
|
|
'key_cache': numpy.ascontiguousarray(key_cache.detach().cpu().numpy()),
|
|
'value_cache': numpy.ascontiguousarray(value_cache.detach().cpu().numpy())
|
|
}
|
|
|
|
from onnxruntime import SessionOptions, InferenceSession
|
|
sess_options = SessionOptions()
|
|
ort_session = InferenceSession(onnx_model_str, sess_options, providers=['CUDAExecutionProvider'])
|
|
ort_output = ort_session.run(None, ort_inputs)
|
|
output, new_key_cache, new_value_cache = ort_output
|
|
|
|
output = torch.tensor(output)
|
|
attn_output = self.out_proj(output)
|
|
|
|
return attn_output, torch.tensor(new_key_cache), torch.tensor(new_value_cache)
|
|
|
|
|
|
def create_decoder_attention_graph(query, key, q_weight, kv_weight, bias, num_heads_, static_kv, use_past, has_layer_state, has_key_padding_mask):
|
|
from onnx import helper, TensorProto
|
|
|
|
S, B, NH = query.size()
|
|
S2 = key.size()[0]
|
|
N = num_heads_
|
|
H = int(NH / N)
|
|
|
|
nodes = [
|
|
helper.make_node("DecoderAttention",
|
|
["query", "key", "q_weight", "kv_weight", "bias", "key_padding_mask", "key_cache", "value_cache", "static_kv", "use_past", "has_layer_state", "has_key_padding_mask"],
|
|
["output", "new_key_cache", "new_value_cache"],
|
|
"DecoderAttention_0",
|
|
num_heads=num_heads_,
|
|
domain="com.microsoft"),
|
|
]
|
|
|
|
initializers = [
|
|
helper.make_tensor('q_weight', TensorProto.FLOAT, [NH, NH],
|
|
q_weight.flatten().tolist()),
|
|
helper.make_tensor('kv_weight', TensorProto.FLOAT, [NH, 2 * NH],
|
|
kv_weight.flatten().tolist()),
|
|
helper.make_tensor('bias', TensorProto.FLOAT, [3 * NH],
|
|
bias.flatten().tolist()),
|
|
helper.make_tensor('static_kv', TensorProto.BOOL, [1],
|
|
[static_kv]),
|
|
helper.make_tensor('use_past', TensorProto.BOOL, [1],
|
|
[use_past]),
|
|
helper.make_tensor('has_layer_state', TensorProto.BOOL, [1],
|
|
[has_layer_state]),
|
|
helper.make_tensor('has_key_padding_mask', TensorProto.BOOL, [1],
|
|
[has_key_padding_mask]),
|
|
]
|
|
|
|
graph = helper.make_graph(nodes, "DecoderAttention_Graph", [
|
|
helper.make_tensor_value_info('query', TensorProto.FLOAT, [S, B, NH]),
|
|
helper.make_tensor_value_info('key', TensorProto.FLOAT, [S2, B, NH]),
|
|
helper.make_tensor_value_info('key_padding_mask', TensorProto.BOOL, [B, "mask_len"]),
|
|
helper.make_tensor_value_info('key_cache', TensorProto.FLOAT, [B, N, "cache_len", H]),
|
|
helper.make_tensor_value_info('value_cache', TensorProto.FLOAT, [B, N, "cache_len", H]),
|
|
], [
|
|
helper.make_tensor_value_info('output', TensorProto.FLOAT, [S, B, NH]),
|
|
helper.make_tensor_value_info('new_key_cache', TensorProto.FLOAT, [B, N, "new_cache_len", H]),
|
|
helper.make_tensor_value_info('new_value_cache', TensorProto.FLOAT, [B, N, "new_cache_len", H]),
|
|
], initializers)
|
|
|
|
model = helper.make_model(graph)
|
|
return model.SerializeToString()
|
|
|
|
|
|
def create_inputs(config: Config, has_layer_state: bool, use_past: bool, encoder_decoder_attention:bool):
|
|
query = torch.normal(mean=0.0,
|
|
std=0.1,
|
|
size=(config.sequence_length,
|
|
config.batch_size,
|
|
config.embed_dim)
|
|
).to(torch.float32)
|
|
key = torch.normal(mean=0.0,
|
|
std=0.1,
|
|
size=(config.kv_sequence_length,
|
|
config.batch_size,
|
|
config.embed_dim)
|
|
).to(torch.float32)
|
|
|
|
key_length = None
|
|
if not has_layer_state or not use_past:
|
|
if not encoder_decoder_attention:
|
|
key_length = config.sequence_length
|
|
else:
|
|
key_length = config.kv_sequence_length
|
|
else:
|
|
if not encoder_decoder_attention:
|
|
key_length = config.sequence_length + config.kv_sequence_length
|
|
else:
|
|
key_length = config.kv_sequence_length
|
|
|
|
key_padding_mask = torch.normal(mean=0.0,
|
|
std=0.1,
|
|
size=(config.batch_size,
|
|
key_length)
|
|
) > 0
|
|
# The following line ensure not all the mask are true
|
|
key_padding_mask[0][0] = False
|
|
|
|
cache = torch.normal(mean=0.0,
|
|
std=0.1,
|
|
size=(config.batch_size,
|
|
config.num_heads,
|
|
config.kv_sequence_length,
|
|
config.head_size)
|
|
).to(torch.float32)
|
|
layer_state = [cache, cache, cache, cache]
|
|
|
|
return query, key, key_padding_mask, layer_state, torch.tensor(use_past)
|
|
|
|
|
|
def parity_check(config, has_layer_state, use_past, static_kv, has_key_padding_mask, rtol = 1e-4, atol = 1e-4):
|
|
query, key, key_padding_mask, layer_state, use_past = create_inputs(config,
|
|
has_layer_state,
|
|
use_past,
|
|
static_kv)
|
|
attn = AttentionForONNX(config.embed_dim,
|
|
config.num_heads,
|
|
encoder_decoder_attention = static_kv)
|
|
attn_output, new_key_cache, new_value_cache = attn.forward(query, key, key_padding_mask, layer_state, None, False, use_past, has_key_padding_mask)
|
|
attn_output_ort, new_key_cache_ort, new_value_cache_ort = attn.ORT_forward(query, key, key_padding_mask, layer_state, None, False, use_past, has_key_padding_mask)
|
|
attn_output_ort_1, _, _ = attn.ORT_forward(query, key, key_padding_mask, layer_state, None, False, use_past, has_key_padding_mask)
|
|
print(" B:", config.batch_size,
|
|
" S:", config.sequence_length,
|
|
" S*:", config.kv_sequence_length,
|
|
" h:", config.embed_dim,
|
|
" has_layer_state:", has_layer_state,
|
|
" use_past:", use_past,
|
|
" static_kv:", static_kv,
|
|
" has_key_padding_mask:", has_key_padding_mask,
|
|
"[attn_output, randomness, key, value] parity:",
|
|
numpy.allclose(attn_output.detach().numpy(), attn_output_ort.detach().numpy(), rtol = rtol, atol = atol, equal_nan = True),
|
|
numpy.allclose(attn_output_ort_1.detach().numpy(), attn_output_ort.detach().numpy(), rtol = rtol, atol = atol, equal_nan = True),
|
|
numpy.allclose(new_key_cache.detach().numpy(), new_key_cache_ort.detach().numpy(), rtol = rtol, atol = atol, equal_nan = True),
|
|
numpy.allclose(new_value_cache.detach().numpy(), new_value_cache_ort.detach().numpy(), rtol = rtol, atol = atol, equal_nan = True))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
for b in [1, 32, 128]:
|
|
for s in [1, 2, 128]:
|
|
for s2 in [1, 64, 256]:
|
|
for n in [8]:
|
|
for h in [64]:
|
|
config = Config(b, s, s2, n, h)
|
|
parity_check(config, has_layer_state = True, use_past = True, static_kv = True, has_key_padding_mask = False)
|
|
parity_check(config, has_layer_state = True, use_past = True, static_kv = False, has_key_padding_mask = False)
|
|
parity_check(config, has_layer_state = True, use_past = False, static_kv = True, has_key_padding_mask = False)
|
|
parity_check(config, has_layer_state = True, use_past = False, static_kv = False, has_key_padding_mask = False)
|
|
parity_check(config, has_layer_state = True, use_past = True, static_kv = True, has_key_padding_mask = True)
|
|
parity_check(config, has_layer_state = True, use_past = True, static_kv = False, has_key_padding_mask = True)
|
|
parity_check(config, has_layer_state = True, use_past = False, static_kv = True, has_key_padding_mask = True)
|
|
parity_check(config, has_layer_state = True, use_past = False, static_kv = False, has_key_padding_mask = True) |