mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
[DML EP] Add MatMulNBits (#20308)
This commit is contained in:
Patrice Vignola
GitHub
parent
55e0aaeeef
commit
8fbb8a149f
8 changed files with 288 additions and 6 deletions
|
|
@ -1296,6 +1296,7 @@ Do not modify directly.*
|
|||
|GroupNorm|*in* X:**T**<br> *in* gamma:**M**<br> *in* beta:**M**<br> *out* Y:**T**|1+|**M** = tensor(float), tensor(float16)<br/> **T** = tensor(float), tensor(float16)|
|
||||
|GroupQueryAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(float), tensor(float16)|
|
||||
|MatMulIntegerToFloat|*in* A:**T1**<br> *in* B:**T2**<br> *in* a_scale:**T3**<br> *in* b_scale:**T3**<br> *in* a_zero_point:**T1**<br> *in* b_zero_point:**T2**<br> *in* bias:**T3**<br> *out* Y:**T3**|1+|**T1** = tensor(int8), tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(float), tensor(float16)|
|
||||
|MatMulNBits|*in* A:**T1**<br> *in* B:**T2**<br> *in* scales:**T1**<br> *in* zero_points:**T3**<br> *in* g_idx:**T4**<br> *out* Y:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)|
|
||||
|MultiHeadAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* bias:**T**<br> *in* key_padding_mask:**M**<br> *in* relative_position_bias:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(float), tensor(float16)|
|
||||
|NhwcConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|QAttention|*in* input:**T1**<br> *in* weight:**T2**<br> *in* bias:**T3**<br> *in* input_scale:**T3**<br> *in* weight_scale:**T3**<br> *in* mask_index:**T4**<br> *in* input_zero_point:**T1**<br> *in* weight_zero_point:**T2**<br> *in* past:**T3**<br> *out* output:**T3**<br> *out* present:**T3**|1+|**T1** = tensor(int8), tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(float), tensor(float16)<br/> **T4** = tensor(int32)|
|
||||
|
|
|
|||
|
|
@ -0,0 +1,200 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "precomp.h"
|
||||
|
||||
namespace Dml
|
||||
{
|
||||
class DmlOperatorMatMulNBits : public DmlOperator
|
||||
{
|
||||
public:
|
||||
DmlOperatorMatMulNBits(const MLOperatorKernelCreationContext& kernelInfo)
|
||||
: DmlOperator(kernelInfo)
|
||||
{
|
||||
ML_CHECK_VALID_ARGUMENT(kernelInfo.GetInputCount() >= 3 && kernelInfo.GetInputCount() <= 4);
|
||||
ML_CHECK_VALID_ARGUMENT(kernelInfo.GetOutputCount() == 1);
|
||||
DmlOperator::Initialize(kernelInfo);
|
||||
|
||||
const bool hasZeroPoint = kernelInfo.IsInputValid(3);
|
||||
const uint32_t bRowCount = gsl::narrow_cast<uint32_t>(kernelInfo.GetAttribute<int64_t>(AttrName::UppercaseN));
|
||||
const uint32_t bColCount = gsl::narrow_cast<uint32_t>(kernelInfo.GetAttribute<int64_t>(AttrName::UppercaseK));
|
||||
const auto bitCount = kernelInfo.GetAttribute<int64_t>(AttrName::Bits);
|
||||
|
||||
MLOperatorTensorDataType mlDataType = kernelInfo.GetInputEdgeDescription(0).tensorDataType;
|
||||
|
||||
ML_CHECK_VALID_ARGUMENT(bitCount == 4 || bitCount == 8);
|
||||
const DML_TENSOR_DATA_TYPE quantizedDataType = bitCount == 4 ? DML_TENSOR_DATA_TYPE_UINT4 : DML_TENSOR_DATA_TYPE_UINT8;
|
||||
|
||||
std::vector<DimensionType> aShape = kernelInfo.GetTensorShapeDescription().GetInputTensorShape(0);
|
||||
ML_CHECK_VALID_ARGUMENT(aShape.size() >= 2);
|
||||
|
||||
// The quantized input to MatMulNBits always comes as uint8, but the real shape is provided through the N and K attributes
|
||||
std::vector<DimensionType> bBroadcastedShape = aShape;
|
||||
bBroadcastedShape[bBroadcastedShape.size() - 2] = bRowCount;
|
||||
bBroadcastedShape[bBroadcastedShape.size() - 1] = bColCount;
|
||||
|
||||
// The B tensor always has a batch size and channel of 1 since it's a weight, but DML requires it to have the same channels
|
||||
// and channel count as the A tensor so we need to broadcast it
|
||||
std::vector<DimensionType> bShape(aShape.size(), 1);
|
||||
bShape[bShape.size() - 2] = bRowCount;
|
||||
bShape[bShape.size() - 1] = bColCount;
|
||||
|
||||
std::vector<DimensionType> outputShape = kernelInfo.GetTensorShapeDescription().GetOutputTensorShape(0);
|
||||
|
||||
std::vector<DimensionType> scaleShape = bShape;
|
||||
|
||||
uint32_t scaleElementCount = ComputeElementCountFromDimensions(m_inputTensorDescs[2].GetSizes());
|
||||
scaleShape[scaleShape.size() - 1] = scaleElementCount / bRowCount;
|
||||
|
||||
std::vector<DimensionType> scaleBroadcastedShape = bBroadcastedShape;
|
||||
scaleBroadcastedShape.back() = scaleShape.back();
|
||||
|
||||
// The quantized input and zero point to MatMulNBits always comes as uint8, but DML will expect the real data type (int4 or int8)
|
||||
m_inputTensorDescs[0] = TensorDesc::ConstructDefaultTensorDesc(mlDataType, aShape);
|
||||
m_inputTensorDescs[1] = TensorDesc::ConstructBroadcastedTensorDesc(GetMlDataTypeFromDmlDataType(quantizedDataType), bBroadcastedShape, bShape);
|
||||
m_inputTensorDescs[2] = TensorDesc::ConstructBroadcastedTensorDesc(mlDataType, scaleBroadcastedShape, scaleShape);
|
||||
|
||||
if (hasZeroPoint)
|
||||
{
|
||||
m_inputTensorDescs[3] = TensorDesc::ConstructBroadcastedTensorDesc(GetMlDataTypeFromDmlDataType(quantizedDataType), scaleBroadcastedShape, scaleShape);
|
||||
|
||||
// Zero points are stored in uint8 data in ORT, which means that it will not be tightly packed if the last dimension is an odd number. To account for that,
|
||||
// we round the strides of the next multiple of 2.
|
||||
auto lastDimension = m_inputTensorDescs[3].GetSizes().back();
|
||||
|
||||
if (lastDimension % 2 != 0)
|
||||
{
|
||||
std::vector<uint32_t> strides(m_inputTensorDescs[3].GetDimensionCount());
|
||||
strides[strides.size() - 1] = 1;
|
||||
strides[strides.size() - 2] = lastDimension + 1;
|
||||
m_inputTensorDescs[3].SetStrides(strides);
|
||||
}
|
||||
}
|
||||
|
||||
auto dequantizedInputTensorDesc = TensorDesc::ConstructDefaultTensorDesc(mlDataType, bBroadcastedShape);
|
||||
auto dequantizedInputDmlTensorDesc = dequantizedInputTensorDesc.GetDmlDesc();
|
||||
|
||||
// Initialize the output description while overriding the shape
|
||||
m_outputTensorDescs[0] = TensorDesc::ConstructDefaultTensorDesc(mlDataType, outputShape);
|
||||
|
||||
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
|
||||
std::vector<DML_TENSOR_DESC> outputDescs = GetDmlOutputDescs();
|
||||
|
||||
std::vector<DML_TENSOR_DESC> quantizationTensors;
|
||||
quantizationTensors.push_back(inputDescs[2]);
|
||||
|
||||
if (hasZeroPoint)
|
||||
{
|
||||
quantizationTensors.push_back(inputDescs[3]);
|
||||
}
|
||||
|
||||
DML_DEQUANTIZE_OPERATOR_DESC dequantizeDesc = {};
|
||||
dequantizeDesc.InputTensor = &inputDescs[1];
|
||||
dequantizeDesc.QuantizationType = hasZeroPoint ? DML_QUANTIZATION_TYPE_SCALE_ZERO_POINT : DML_QUANTIZATION_TYPE_SCALE;
|
||||
dequantizeDesc.QuantizationTensorCount = gsl::narrow_cast<uint32_t>(quantizationTensors.size());
|
||||
dequantizeDesc.QuantizationTensors = quantizationTensors.data();
|
||||
dequantizeDesc.OutputTensor = &dequantizedInputDmlTensorDesc;
|
||||
DML_OPERATOR_DESC dequantizeOpDesc = { DML_OPERATOR_DEQUANTIZE, &dequantizeDesc };
|
||||
|
||||
DML_GEMM_OPERATOR_DESC gemmDesc = {};
|
||||
gemmDesc.ATensor = &inputDescs[0];
|
||||
gemmDesc.BTensor = &dequantizedInputDmlTensorDesc;
|
||||
gemmDesc.CTensor = nullptr;
|
||||
gemmDesc.OutputTensor = &outputDescs[0];
|
||||
gemmDesc.TransA = DML_MATRIX_TRANSFORM_NONE;
|
||||
gemmDesc.TransB = DML_MATRIX_TRANSFORM_TRANSPOSE;
|
||||
gemmDesc.Alpha = 1.0f;
|
||||
gemmDesc.Beta = 0.0f;
|
||||
DML_OPERATOR_DESC gemmOpDesc = { DML_OPERATOR_GEMM, &gemmDesc };
|
||||
|
||||
// Construct the graph
|
||||
std::vector<DML_INPUT_GRAPH_EDGE_DESC> inputEdges;
|
||||
std::vector<DML_INTERMEDIATE_GRAPH_EDGE_DESC> intermediateEdges;
|
||||
std::vector<DML_OUTPUT_GRAPH_EDGE_DESC> outputEdges;
|
||||
|
||||
std::vector<const DML_OPERATOR_DESC*> opDescs = {
|
||||
&dequantizeOpDesc,
|
||||
&gemmOpDesc,
|
||||
};
|
||||
|
||||
DML_INPUT_GRAPH_EDGE_DESC secondInputToDequantizeEdge = {};
|
||||
secondInputToDequantizeEdge.GraphInputIndex = 1;
|
||||
secondInputToDequantizeEdge.ToNodeIndex = 0;
|
||||
secondInputToDequantizeEdge.ToNodeInputIndex = 0;
|
||||
inputEdges.push_back(secondInputToDequantizeEdge);
|
||||
|
||||
DML_INPUT_GRAPH_EDGE_DESC scaleToDequantizeEdge = {};
|
||||
scaleToDequantizeEdge.GraphInputIndex = 2;
|
||||
scaleToDequantizeEdge.ToNodeIndex = 0;
|
||||
scaleToDequantizeEdge.ToNodeInputIndex = 1;
|
||||
inputEdges.push_back(scaleToDequantizeEdge);
|
||||
|
||||
if (hasZeroPoint)
|
||||
{
|
||||
DML_INPUT_GRAPH_EDGE_DESC zeroPointToDequantizeEdge = {};
|
||||
zeroPointToDequantizeEdge.GraphInputIndex = 3;
|
||||
zeroPointToDequantizeEdge.ToNodeIndex = 0;
|
||||
zeroPointToDequantizeEdge.ToNodeInputIndex = 2;
|
||||
inputEdges.push_back(zeroPointToDequantizeEdge);
|
||||
}
|
||||
|
||||
DML_INPUT_GRAPH_EDGE_DESC firstInputToGemmEdge = {};
|
||||
firstInputToGemmEdge.GraphInputIndex = 0;
|
||||
firstInputToGemmEdge.ToNodeIndex = 1;
|
||||
firstInputToGemmEdge.ToNodeInputIndex = 0;
|
||||
inputEdges.push_back(firstInputToGemmEdge);
|
||||
|
||||
DML_INTERMEDIATE_GRAPH_EDGE_DESC dequantizeToGemmEdge = {};
|
||||
dequantizeToGemmEdge.FromNodeIndex = 0;
|
||||
dequantizeToGemmEdge.FromNodeOutputIndex = 0;
|
||||
dequantizeToGemmEdge.ToNodeIndex = 1;
|
||||
dequantizeToGemmEdge.ToNodeInputIndex = 1;
|
||||
intermediateEdges.push_back(dequantizeToGemmEdge);
|
||||
|
||||
DML_OUTPUT_GRAPH_EDGE_DESC gemmToOutputEdge = {};
|
||||
gemmToOutputEdge.FromNodeIndex = 1;
|
||||
gemmToOutputEdge.FromNodeOutputIndex = 0;
|
||||
gemmToOutputEdge.GraphOutputIndex = 0;
|
||||
outputEdges.push_back(gemmToOutputEdge);
|
||||
|
||||
MLOperatorGraphDesc operatorGraphDesc = {};
|
||||
operatorGraphDesc.inputEdgeCount = gsl::narrow_cast<uint32_t>(inputEdges.size());
|
||||
operatorGraphDesc.inputEdges = inputEdges.data();
|
||||
operatorGraphDesc.intermediateEdgeCount = gsl::narrow_cast<uint32_t>(intermediateEdges.size());
|
||||
operatorGraphDesc.intermediateEdges = intermediateEdges.data();
|
||||
operatorGraphDesc.outputEdgeCount = gsl::narrow_cast<uint32_t>(outputEdges.size());
|
||||
operatorGraphDesc.outputEdges = outputEdges.data();
|
||||
operatorGraphDesc.nodeCount = gsl::narrow_cast<uint32_t>(opDescs.size());
|
||||
operatorGraphDesc.nodes = opDescs.data();
|
||||
|
||||
SetDmlOperatorGraphDesc(std::move(operatorGraphDesc), kernelInfo);
|
||||
}
|
||||
};
|
||||
|
||||
void CALLBACK QueryMatMulNBits(IMLOperatorSupportQueryContextPrivate* context, /*out*/ bool* isSupported)
|
||||
{
|
||||
*isSupported = false;
|
||||
|
||||
MLOperatorAttributes attributes(context);
|
||||
const auto bitCount = attributes.GetAttribute<int64_t>(AttrName::Bits);
|
||||
if (bitCount != 4 && bitCount != 8)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
// DML only supports perfect blocks at the moment, which is what most quantization tools produce anyways since they pad
|
||||
// as part of the quantization process
|
||||
const uint32_t bColCount = gsl::narrow_cast<uint32_t>(attributes.GetAttribute<int64_t>(AttrName::UppercaseK));
|
||||
const auto blockSize = attributes.GetAttribute<int64_t>(AttrName::MatMulNBitsBlockSize);
|
||||
ML_CHECK_VALID_ARGUMENT(blockSize > 0);
|
||||
|
||||
if (bColCount % blockSize != 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
*isSupported = true;
|
||||
}
|
||||
|
||||
DML_OP_DEFINE_CREATION_FUNCTION(MatMulNBits, DmlOperatorMatMulNBits);
|
||||
} // namespace Dml
|
||||
|
|
@ -530,6 +530,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(BitwiseOr);
|
|||
DML_OP_EXTERN_CREATION_FUNCTION(BitwiseXor);
|
||||
DML_OP_EXTERN_CREATION_FUNCTION(BitwiseNot);
|
||||
DML_OP_EXTERN_CREATION_FUNCTION(RotaryEmbedding);
|
||||
DML_OP_EXTERN_CREATION_FUNCTION(MatMulNBits);
|
||||
|
||||
DML_OP_EXTERN_QUERY_FUNCTION(MaxPool);
|
||||
DML_OP_EXTERN_QUERY_FUNCTION(Slice);
|
||||
|
|
@ -544,6 +545,7 @@ DML_OP_EXTERN_QUERY_FUNCTION(SkipLayerNormalization);
|
|||
DML_OP_EXTERN_QUERY_FUNCTION(QLinearSigmoid);
|
||||
DML_OP_EXTERN_QUERY_FUNCTION(QAttention);
|
||||
DML_OP_EXTERN_QUERY_FUNCTION(Attention);
|
||||
DML_OP_EXTERN_QUERY_FUNCTION(MatMulNBits);
|
||||
|
||||
constexpr static std::array<const char*, 1> typeNameListDefault = {"T"};
|
||||
constexpr static std::array<const char*, 1> typeNameListDefaultV = {"V"};
|
||||
|
|
@ -632,6 +634,7 @@ constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListAttent
|
|||
constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListRotaryEmbedding = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Int64};
|
||||
constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListGroupNorm = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Float16to32};
|
||||
constexpr static std::array<SupportedTensorDataTypes, 1> supportedTypeListNonZero = {SupportedTensorDataTypes::Float16to32 | SupportedTensorDataTypes::Ints8Bit | SupportedTensorDataTypes::Ints16Bit | SupportedTensorDataTypes::Ints32Bit | SupportedTensorDataTypes::Bool};
|
||||
constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListMatMulNBits = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::UInt8};
|
||||
|
||||
constexpr static std::array<SupportedTensorDataTypes, 3> supportedTypeListQLinearMatMul = {
|
||||
SupportedTensorDataTypes::Ints8Bit,
|
||||
|
|
@ -1140,6 +1143,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation
|
|||
{REG_INFO_MS( 1, BiasAdd, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
{REG_INFO_MS( 1, QuickGelu, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},
|
||||
{REG_INFO_MS( 1, GroupNorm, typeNameListGroupNorm, supportedTypeListGroupNorm, DmlGraphSupport::Supported)},
|
||||
{REG_INFO_MS( 1, MatMulNBits, typeNameListTwo, supportedTypeListMatMulNBits, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryMatMulNBits)},
|
||||
|
||||
// Operators that need to alias an input with an output
|
||||
{REG_INFO_MS_ALIAS(1, GroupQueryAttention, Aliases(std::make_pair(3, 1), std::make_pair(4, 2)), typeNameListAttention, supportedTypeListAttention, DmlGraphSupport::Supported, requiredConstantCpuInputs(6))},
|
||||
|
|
|
|||
|
|
@ -126,6 +126,10 @@ namespace AttrName
|
|||
static constexpr const char* GraphFusedActivation = "activation";
|
||||
static constexpr const char* GraphFusedAxis = "activation_axis";
|
||||
static constexpr const char* Interleaved = "interleaved";
|
||||
static constexpr const char* Bits = "bits";
|
||||
static constexpr const char* UppercaseN = "N";
|
||||
static constexpr const char* UppercaseK = "K";
|
||||
static constexpr const char* MatMulNBitsBlockSize = "block_size";
|
||||
|
||||
} // namespace AttrName
|
||||
|
||||
|
|
|
|||
|
|
@ -2919,4 +2919,29 @@ namespace OperatorHelper
|
|||
return { EdgeShapes(std::move(outputShape)) };
|
||||
}
|
||||
|
||||
std::vector<EdgeShapes> MatMulNBitsHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const
|
||||
{
|
||||
auto inputShape = shapeInfo.GetInputTensorShape(0);
|
||||
onnxruntime::TensorShape aShape(std::vector<int64_t>(inputShape.begin(), inputShape.end()));
|
||||
onnxruntime::TensorShape bShape({m_bRowCount, m_bColCount});
|
||||
|
||||
onnxruntime::MatMulComputeHelper helper;
|
||||
|
||||
// The B tensor is always transposed
|
||||
ML_CHECK_VALID_ARGUMENT(helper.Compute(aShape, bShape, false, true).IsOK());
|
||||
const auto outputShape = helper.OutputShape().GetDims();
|
||||
|
||||
std::vector<uint32_t> uint32OutputShape;
|
||||
uint32OutputShape.reserve(outputShape.size());
|
||||
std::transform(outputShape.begin(), outputShape.end(), std::back_inserter(uint32OutputShape), [](int64_t dimSize){ return static_cast<uint32_t>(dimSize); });
|
||||
|
||||
return { EdgeShapes(uint32OutputShape) };
|
||||
}
|
||||
|
||||
void MatMulNBitsHelper::Initialize(const IKernelInformationAdapter& kernelInformation)
|
||||
{
|
||||
m_bRowCount = kernelInformation.GetAttributes().GetAttribute<int64_t>(AttrName::UppercaseN);
|
||||
m_bColCount = kernelInformation.GetAttributes().GetAttribute<int64_t>(AttrName::UppercaseK);
|
||||
}
|
||||
|
||||
} // namespace OperatorHelper
|
||||
|
|
|
|||
|
|
@ -1606,6 +1606,24 @@ public:
|
|||
std::vector<EdgeShapes> GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const;
|
||||
};
|
||||
|
||||
class MatMulNBitsHelper
|
||||
{
|
||||
public:
|
||||
template <typename Info_t, typename Shape_t>
|
||||
MatMulNBitsHelper(const Info_t& info, const Shape_t& shapeInfo)
|
||||
{
|
||||
Initialize(KernelInformationAdapter(info));
|
||||
}
|
||||
|
||||
std::vector<EdgeShapes> GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const;
|
||||
|
||||
private:
|
||||
void Initialize(const IKernelInformationAdapter& kernelInformation);
|
||||
|
||||
int64_t m_bRowCount;
|
||||
int64_t m_bColCount;
|
||||
};
|
||||
|
||||
using ShapeInferenceHelper_Conv = ConvHelper;
|
||||
using ShapeInferenceHelper_NhwcConv = NhwcConvHelper;
|
||||
using ShapeInferenceHelper_ConvTranspose = ConvTransposeHelper;
|
||||
|
|
@ -1859,5 +1877,6 @@ using ShapeInferenceHelper_DmlFusedSum = GetBroadcastedOutputShapeHelper;
|
|||
using ShapeInferenceHelper_Shape = ShapeHelper;
|
||||
using ShapeInferenceHelper_Size = SizeHelper;
|
||||
using ShapeInferenceHelper_BiasSplitGelu = BiasSplitGeluHelper;
|
||||
using ShapeInferenceHelper_MatMulNBits = MatMulNBitsHelper;
|
||||
|
||||
} // namespace OperatorHelper
|
||||
|
|
|
|||
|
|
@ -467,6 +467,7 @@ namespace OperatorHelper
|
|||
static const int sc_sinceVer_RotaryEmbedding = 1;
|
||||
static const int sc_sinceVer_QLinearAveragePool = 1;
|
||||
static const int sc_sinceVer_QLinearGlobalAveragePool = 1;
|
||||
static const int sc_sinceVer_MatMulNBits = 1;
|
||||
static const int sc_sinceVer_DynamicQuantizeMatMul = 1;
|
||||
} // namespace MsftOperatorSet1
|
||||
|
||||
|
|
|
|||
|
|
@ -156,7 +156,15 @@ void RunTest(int64_t M, int64_t N, int64_t K, int64_t block_size, int64_t accura
|
|||
test.SetOutputAbsErr("Y", fp16_abs_error);
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
|
||||
#ifdef USE_CUDA
|
||||
execution_providers.push_back(DefaultCudaExecutionProvider());
|
||||
#endif
|
||||
|
||||
#ifdef USE_DML
|
||||
execution_providers.push_back(DefaultDmlExecutionProvider());
|
||||
#endif
|
||||
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
|
||||
} else {
|
||||
test.AddInput<float>("A", {M, K}, input0_vals, false);
|
||||
|
|
@ -205,7 +213,7 @@ TEST(MatMulNBits, Float32) {
|
|||
for (auto N : {1, 2, 32, 288}) {
|
||||
for (auto K : {16, 32, 64, 128, 256, 1024, 93, 1234}) {
|
||||
for (auto block_size : {16, 32, 64, 128}) {
|
||||
#ifdef ORT_NEURAL_SPEED
|
||||
#if defined(ORT_NEURAL_SPEED) || defined(USE_DML)
|
||||
for (auto accuracy_level : {0, 1, 4}) {
|
||||
RunTest(M, N, K, block_size, accuracy_level, false, false);
|
||||
RunTest(M, N, K, block_size, accuracy_level, true, false);
|
||||
|
|
@ -224,15 +232,25 @@ TEST(MatMulNBits, Float32) {
|
|||
}
|
||||
}
|
||||
|
||||
#if defined(USE_CUDA)
|
||||
#if defined(USE_CUDA) || defined(USE_DML)
|
||||
TEST(MatMulNBits, Float16) {
|
||||
#ifdef USE_CUDA
|
||||
auto has_gidx_options = {true, false};
|
||||
#else
|
||||
auto has_gidx_options = {false};
|
||||
#endif
|
||||
|
||||
for (auto M : {1, 2, 100}) {
|
||||
for (auto N : {1, 2, 32, 288}) {
|
||||
for (auto K : {16, 32, 64, 128, 256, 1024, 93, 1234}) {
|
||||
for (auto block_size : {16, 32, 64, 128}) {
|
||||
for (auto has_gidx : {true, false}) {
|
||||
for (auto has_gidx : has_gidx_options) {
|
||||
#ifdef USE_DML
|
||||
RunTest(M, N, K, block_size, 0, false, true, has_gidx, true, 0.04f);
|
||||
#else
|
||||
RunTest(M, N, K, block_size, 0, false, true, has_gidx);
|
||||
RunTest(M, N, K, block_size, 0, true, true, has_gidx, false);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -241,11 +259,21 @@ TEST(MatMulNBits, Float16) {
|
|||
}
|
||||
|
||||
TEST(MatMulNBits, Float16Large) {
|
||||
#ifdef USE_DML
|
||||
// For some reason, the A10 machine that runs these tests during CI has a much bigger error than all retail
|
||||
// machines we tested on. All consumer-grade machines from Nvidia/AMD/Intel seem to pass these tests with an
|
||||
// absolute error of 0.08, but the A10 has errors going as high as 0.22. Ultimately, given the large number
|
||||
// of elements in this test, ULPs should probably be used instead of absolute/relative tolerances.
|
||||
float abs_error = 0.3f;
|
||||
#else
|
||||
float abs_error = 0.05f;
|
||||
#endif
|
||||
|
||||
for (auto block_size : {16, 32, 64, 128}) {
|
||||
for (auto symmetric : {false, true}) {
|
||||
RunTest(1, 4096, 4096, block_size, 0, symmetric, true, false, true, 0.05f);
|
||||
RunTest(1, 4096, 11008, block_size, 0, symmetric, true, false, true, 0.05f);
|
||||
RunTest(1, 11008, 4096, block_size, 0, symmetric, true, false, true, 0.05f);
|
||||
RunTest(1, 4096, 4096, block_size, 0, symmetric, true, false, true, abs_error);
|
||||
RunTest(1, 4096, 11008, block_size, 0, symmetric, true, false, true, abs_error);
|
||||
RunTest(1, 11008, 4096, block_size, 0, symmetric, true, false, true, abs_error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue