diff --git a/cmake/onnxruntime_rocm_hipify.cmake b/cmake/onnxruntime_rocm_hipify.cmake index 238f09f0ae..525d112b24 100644 --- a/cmake/onnxruntime_rocm_hipify.cmake +++ b/cmake/onnxruntime_rocm_hipify.cmake @@ -9,8 +9,8 @@ set(contrib_ops_excluded_files "bert/attention.h" "bert/attention_impl.cu" "bert/attention_softmax.h" - "bert/cross_attention.cc" - "bert/cross_attention.h" + "bert/multihead_attention.cc" + "bert/multihead_attention.h" "bert/embed_layer_norm.cc" "bert/embed_layer_norm.h" "bert/embed_layer_norm_impl.cu" diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 3fe9fb1d3a..814a076fae 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -17,7 +17,6 @@ Do not modify directly.* * com.microsoft.ComplexMulConj * com.microsoft.ConvTransposeWithDynamicPads * com.microsoft.CropAndResize - * com.microsoft.CrossAttention * com.microsoft.DecoderAttention * com.microsoft.DequantizeBFP * com.microsoft.DequantizeLinear @@ -42,6 +41,7 @@ Do not modify directly.* * com.microsoft.MatMulIntegerToFloat * com.microsoft.MaxpoolWithMask * com.microsoft.MulInteger + * com.microsoft.MultiHeadAttention * com.microsoft.MurmurHash3 * com.microsoft.NGramRepeatBlock * com.microsoft.NhwcConv @@ -955,57 +955,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.CrossAttention** - - Multi-Head Cross Attention. Bias from input projection is included. - - The key padding mask is optional. When its shape is (batch_size, kv_sequence_length), value 0 - means padding or 1 otherwise. When key has right-side padding, its shape could be (batch_size): it is actual length of - each key sequence excluding paddings. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
num_heads : int (required)
-
Number of attention heads
-
- -#### Inputs (4 - 5) - -
-
query : T
-
Query with shape (batch_size, sequence_length, hidden_size) when weights is not available.
-
key : T
-
Key with shape (batch_size, kv_sequence_length, hidden_size)
-
value : T
-
Value with shape (batch_size, kv_sequence_length, v_hidden_size)
-
bias : T
-
Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection
-
key_padding_mask (optional) : M
-
Key padding mask with shape (batch_size) or (batch_size, kv_sequence_length)
-
- -#### Outputs - -
-
output : T
-
3D output tensor with shape (batch_size, sequence_length, v_hidden_size)
-
- -#### Type Constraints - -
-
T : tensor(float), tensor(float16)
-
Constrain input and output to float tensors.
-
M : tensor(int32)
-
Constrain mask to integer types
-
- - ### **com.microsoft.DecoderAttention** This DecoderAttention supports self attention and cross attention, key and value cache, and key_padding_mask. The attention mask is not support at the moment. @@ -2156,6 +2105,57 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.MultiHeadAttention** + + Multi-Head Self/Cross Attention. Bias from input projection is included. + + The key padding mask is optional. When its shape is (batch_size, kv_sequence_length), value 0 + means padding or 1 otherwise. When key has right-side padding, its shape could be (batch_size): it is actual length of + each key sequence excluding paddings. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
num_heads : int (required)
+
Number of attention heads
+
+ +#### Inputs (4 - 5) + +
+
query : T
+
Query with shape (batch_size, sequence_length, hidden_size) when weights is not available.
+
key : T
+
Key with shape (batch_size, kv_sequence_length, hidden_size)
+
value : T
+
Value with shape (batch_size, kv_sequence_length, v_hidden_size)
+
bias : T
+
Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection
+
key_padding_mask (optional) : M
+
Key padding mask with shape (batch_size) or (batch_size, kv_sequence_length)
+
+ +#### Outputs + +
+
output : T
+
3D output tensor with shape (batch_size, sequence_length, v_hidden_size)
+
+ +#### Type Constraints + +
+
T : tensor(float), tensor(float16)
+
Constrain input and output to float tensors.
+
M : tensor(int32)
+
Constrain mask to integer types
+
+ + ### **com.microsoft.MurmurHash3** The underlying implementation is MurmurHash3_x86_32 generating low latency 32bits hash suitable for implementing lookup tables, Bloom filters, count min sketch or feature hashing. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index f575c1a01b..a1cdb7e30d 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -784,7 +784,6 @@ Do not modify directly.* |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| -|CrossAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |DecoderAttention|*in* query:**T**
*in* key:**T**
*in* q_weight:**T**
*in* kv_weight:**T**
*in* bias:**T**
*in* key_padding_mask:**B**
*in* key_cache:**T**
*in* value_cache:**T**
*in* static_kv:**B**
*in* use_past:**B**
*in* has_layer_state:**B**
*in* has_key_padding_mask:**B**
*out* output:**T**
*out* new_key_cache:**T**
*out* new_value_cache:**T**|1+|**T** = tensor(float), tensor(float16)| |DequantizeLinear|*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**|1+|**T1** = tensor(int8), tensor(uint8)
**T2** = tensor(float16)| |DequantizeWithOrder|*in* input:**Q**
*in* scale_input:**S**
*out* output:**F**|1+|**F** = tensor(float), tensor(float16)
**Q** = tensor(int8)
**S** = tensor(float)| @@ -798,6 +797,7 @@ Do not modify directly.* |Inverse|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |Irfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |LongformerAttention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask:**T**
*in* global_weight:**T**
*in* global_bias:**T**
*in* global:**G**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| +|MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(int8)
**T2** = tensor(int8)
**T3** = tensor(float), tensor(float16)
**T4** = tensor(int32)| |QOrderedAttention|*in* input:**Q**
*in* scale_input:**S**
*in* scale_Q_gemm:**S**
*in* scale_K_gemm:**S**
*in* scale_V_gemm:**S**
*in* Q_weight:**Q**
*in* K_weight:**Q**
*in* V_weight:**Q**
*in* scale_Q_weight:**S**
*in* scale_K_weight:**S**
*in* scale_V_weight:**S**
*in* Q_bias:**S**
*in* K_bias:**S**
*in* V_bias:**S**
*in* scale_QKT_gemm:**S**
*in* scale_QKT_softmax:**S**
*in* scale_values_gemm:**S**
*in* mask_index:**G**
*in* past:**Q**
*in* extra_add:**S**
*out* output:**Q**|1+|**G** = tensor(int32)
**Q** = tensor(int8)
**S** = tensor(float)| diff --git a/onnxruntime/contrib_ops/cpu/bert/cross_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h similarity index 98% rename from onnxruntime/contrib_ops/cpu/bert/cross_attention_helper.h rename to onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h index 78ebd23d98..a1b4156f72 100644 --- a/onnxruntime/contrib_ops/cpu/bert/cross_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h @@ -9,7 +9,7 @@ namespace onnxruntime { namespace contrib { -namespace cross_attention_helper { +namespace multihead_attention_helper { template Status CheckInputs(const T* query, @@ -114,6 +114,6 @@ Status CheckInputs(const T* query, return Status::OK(); } -} // namespace cross_attention_helper +} // namespace multihead_attention_helper } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/cross_attention.cc b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc similarity index 92% rename from onnxruntime/contrib_ops/cuda/bert/cross_attention.cc rename to onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc index 8fcb041f9f..3d168ca453 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cross_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc @@ -4,8 +4,8 @@ #include "core/providers/cuda/cuda_common.h" #include "core/platform/env_var_utils.h" #include "contrib_ops/cuda/bert/attention_impl.h" -#include "contrib_ops/cuda/bert/cross_attention.h" -#include "contrib_ops/cpu/bert/cross_attention_helper.h" +#include "contrib_ops/cuda/bert/multihead_attention.h" +#include "contrib_ops/cpu/bert/multihead_attention_helper.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; @@ -17,20 +17,20 @@ namespace cuda { #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ - CrossAttention, \ + MultiHeadAttention, \ kMSDomain, \ 1, \ T, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ - CrossAttention); + MultiHeadAttention); REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(MLFloat16) template -CrossAttention::CrossAttention(const OpKernelInfo& info) : CudaKernel(info) { +MultiHeadAttention::MultiHeadAttention(const OpKernelInfo& info) : CudaKernel(info) { int64_t num_heads = 0; ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); num_heads_ = static_cast(num_heads); @@ -43,7 +43,7 @@ CrossAttention::CrossAttention(const OpKernelInfo& info) : CudaKernel(info) { } template -Status CrossAttention::ComputeInternal(OpKernelContext* context) const { +Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) const { const Tensor* query = context->Input(0); const Tensor* key = context->Input(1); const Tensor* value = context->Input(2); @@ -52,7 +52,7 @@ Status CrossAttention::ComputeInternal(OpKernelContext* context) const { auto& device_prop = GetDeviceProp(); AttentionParameters parameters; - ORT_RETURN_IF_ERROR(cross_attention_helper::CheckInputs(query, + ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckInputs(query, key, value, bias, @@ -84,7 +84,7 @@ Status CrossAttention::ComputeInternal(OpKernelContext* context) const { enable_flash_attention_, false); if (use_fused_runner) { - // Here we assume that num_heads and head_size does not change for an CrossAttention node. + // Here we assume that num_heads and head_size does not change for an MultiHeadAttention node. if (nullptr == fused_fp16_runner_.get()) { constexpr bool is_unidirectional = false; fused_fp16_runner_.reset(new FusedMHARunnerFP16v2( diff --git a/onnxruntime/contrib_ops/cuda/bert/cross_attention.h b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h similarity index 87% rename from onnxruntime/contrib_ops/cuda/bert/cross_attention.h rename to onnxruntime/contrib_ops/cuda/bert/multihead_attention.h index 53019f03e6..6eeecc697c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cross_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h @@ -14,9 +14,9 @@ namespace cuda { using namespace onnxruntime::cuda; template -class CrossAttention final : public CudaKernel { +class MultiHeadAttention final : public CudaKernel { public: - CrossAttention(const OpKernelInfo& info); + MultiHeadAttention(const OpKernelInfo& info); Status ComputeInternal(OpKernelContext* context) const override; protected: diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index c1460a6a56..38bcbc298b 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -62,8 +62,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, Crop); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, Crop); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, Crop); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, CrossAttention); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, CrossAttention); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, MultiHeadAttention); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, MultiHeadAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, DecoderAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, DecoderAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int32_t, DynamicSlice); @@ -183,8 +183,8 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 34b1317f05..ecd4b4f5cf 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -125,7 +125,7 @@ void RestorePaddingTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) } } -void CrossAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { +void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { // Input 0 (query) has shape (batch_size, sequence_length, hidden_size) // Input 1 (key) has shape (batch_size, kv_sequence_length, hidden_size) // Input 2 (value) has shape (batch_size, kv_sequence_length, v_hidden_size) @@ -258,8 +258,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( AttentionTypeAndShapeInference(ctx, past_input_index); })); -constexpr const char* CrossAttention_ver1_doc = R"DOC( -Multi-Head Cross Attention. Bias from input projection is included. +constexpr const char* MultiHeadAttention_ver1_doc = R"DOC( +Multi-Head Self/Cross Attention. Bias from input projection is included. The key padding mask is optional. When its shape is (batch_size, kv_sequence_length), value 0 means padding or 1 otherwise. When key has right-side padding, its shape could be (batch_size): it is actual length of @@ -267,9 +267,9 @@ each key sequence excluding paddings. )DOC"; ONNX_MS_OPERATOR_SET_SCHEMA( - CrossAttention, 1, + MultiHeadAttention, 1, OpSchema() - .SetDoc(CrossAttention_ver1_doc) + .SetDoc(MultiHeadAttention_ver1_doc) .Attr("num_heads", "Number of attention heads", AttributeProto::INT) .Input(0, "query", @@ -299,7 +299,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output to float tensors.") .TypeConstraint("M", {"tensor(int32)"}, "Constrain mask to integer types") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - CrossAttentionTypeAndShapeInference(ctx); + MultiHeadAttentionTypeAndShapeInference(ctx); })); constexpr const char* Longformer_Attention_doc = R"DOC( @@ -436,7 +436,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( schema.BuildFunction(functionProto); return true; })); - + ONNX_MS_OPERATOR_SET_SCHEMA( RelativePositionBias, 1, OpSchema() diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 673cbbaf8e..1f0af31a4b 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -56,7 +56,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ComplexMul); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ComplexMulConj); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ConvTransposeWithDynamicPads); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, CropAndResize); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, CrossAttention); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MultiHeadAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, DecoderAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, EmbedLayerNormalization); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ExpandDims); @@ -142,7 +142,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index dd8143913c..53937bc7f6 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -190,7 +190,7 @@ class SymbolicShapeInference: # contrib ops: "Attention": self._infer_Attention, "BiasGelu": self._infer_BiasGelu, - "CrossAttention": self._infer_CrossAttention, + "MultiHeadAttention": self._infer_MultiHeadAttention, "EmbedLayerNormalization": self._infer_EmbedLayerNormalization, "FastGelu": self._infer_FastGelu, "Gelu": self._infer_Gelu, @@ -1994,7 +1994,7 @@ class SymbolicShapeInference: def _infer_BiasGelu(self, node): self._propagate_shape_and_type(node) - def _infer_CrossAttention(self, node): + def _infer_MultiHeadAttention(self, node): # Input 0 (query) has shape (batch_size, sequence_length, hidden_size) # Input 1 (key) has shape (batch_size, kv_sequence_length, hidden_size) # Input 2 (value) has shape (batch_size, kv_sequence_length, v_hidden_size) diff --git a/onnxruntime/test/contrib_ops/cross_attention_op_test.cc b/onnxruntime/test/contrib_ops/cross_attention_op_test.cc index ae1e395a80..62a150ec92 100644 --- a/onnxruntime/test/contrib_ops/cross_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/cross_attention_op_test.cc @@ -13,7 +13,7 @@ namespace onnxruntime { using contrib::AttentionMaskType; namespace test { -static void RunCrossAttentionTest( +static void RunMultiHeadAttentionTest( const std::vector& query_data, // query: [batch_size, sequence_length, hidden_size] const std::vector& key_data, // key: [batch_size, kv_sequence_length, hidden_size] const std::vector& value_data, // value: [batch_size, kv_sequence_length, v_hidden_size] @@ -39,7 +39,7 @@ static void RunCrossAttentionTest( bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get()) && !use_float16 && !disable_cpu; if (enable_cpu || enable_cuda || enable_rocm) { - OpTester tester("CrossAttention", 1, onnxruntime::kMSDomain); + OpTester tester("MultiHeadAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(number_of_heads)); std::vector query_dims = {batch_size, sequence_length, hidden_size}; @@ -112,7 +112,7 @@ static void RunCrossAttentionTest( } } -TEST(CrossAttentionTest, CrossAttentionBatch1) { +TEST(MultiHeadAttentionTest, MultiHeadAttentionBatch1) { int batch_size = 1; int sequence_length = 2; int hidden_size = 4; @@ -149,7 +149,7 @@ TEST(CrossAttentionTest, CrossAttentionBatch1) { constexpr bool disable_cuda = false; constexpr bool disable_rocm = true; // not supported in rocm right now. - RunCrossAttentionTest( + RunMultiHeadAttentionTest( query_data, key_data, value_data, bias_data, key_padding_mask_data, mask_type, output_data, number_of_heads, batch_size, sequence_length, kv_sequence_length, hidden_size, v_hidden_size, use_float16, disable_cpu, disable_cuda, disable_rocm);