From 7035470541384f7ecefca0eae402dc7c7332955b Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 23 Oct 2026 16:05:10 +0800 Subject: [PATCH 1/4] Support empty MatMulNBits weights on CPU and WebGPU --- docs/ContribOperators.md | 11 +- .../cpu/quantization/matmul_nbits.cc | 35 ++- .../cpu/quantization/matmul_nbits_helper.h | 23 ++ .../webgpu/quantization/matmul_nbits.cc | 25 +- .../core/graph/contrib_ops/contrib_defs.cc | 10 +- .../matmul_nbits_empty_weights_test.cc | 227 ++++++++++++++++++ 6 files changed, 321 insertions(+), 10 deletions(-) create mode 100644 onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index a0036e6c4ffdb..c6809ce5f9d73 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4181,7 +4181,7 @@ This version of the operator has been available since version 1 of the 'com.micr dequantized_weight = (quantized_weight - zero_point) * scale 2. Matrix multiplication between the input matrix A and the dequantized weight matrix. - The weight matrix is a 2D constant matrix with the input feature count and output feature count specified by attributes 'K' and 'N'. + The logical weight matrix is 2D, with the input feature count and output feature count specified by attributes 'K' and 'N'. It is quantized block-wise along the K dimension with a block size specified by the 'block_size' attribute. The block size must be a power of 2 and not smaller than 16 (e.g., 16, 32, 64, 128). Each block has its own scale and zero-point. The quantization is performed using a bit-width specified by the 'bits' attribute, which can take values from 2 to 8. @@ -4189,6 +4189,13 @@ This version of the operator has been available since version 1 of the 'com.micr The quantized weights are stored in a bit-packed format along the K dimension, with each block being represented by a blob of uint8. For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a byte, and the second 4 bits are stored in the higher 4 bits of a byte. + When weight_prepacked is 0, B may represent an omitted projection by replacing one or more dimensions of its + packed shape (N, k_blocks, blob_size) with 0. All remaining dimensions must retain their normal values. + In this case, Y is filled with zeros (or the broadcast bias, if provided). Its shape is still determined by A, K and N. + The contents of A, scales, zero_points and g_idx are not used; the quantization parameters may also be empty. + This permits zero-sized defaults for optional weights, with dynamic input dimensions allowing nonempty overrides. + Empty packed weights are currently supported by the CPU and WebGPU execution providers. + #### Version This version of the operator has been available since version 1 of the 'com.microsoft' operator set. @@ -4216,7 +4223,7 @@ This version of the operator has been available since version 1 of the 'com.micr
A : T1
The input tensor, not quantized.
B : T2
-
Packed uint8 tensor of shape (N, k_blocks, blob_size), where k_blocks = ceil(K / block_size) and blob_size = (block_size * bits / 8). The quantized weights are stored in a bit-packed format along the K dimension, packed within each block_size.
+
Packed uint8 tensor of shape (N, k_blocks, blob_size), where k_blocks = ceil(K / block_size) and blob_size = (block_size * bits / 8). With weight_prepacked=0, one or more dimensions may be 0 to omit the projection as described above. The quantized weights are stored in a bit-packed format along the K dimension, packed within each block_size.
scales : T1
Per-block scaling factors for dequantization with shape (N, k_blocks) and same data type as input A.
zero_points (optional) : T3
diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc index d42dbae635047..e781655f287e0 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc @@ -221,6 +221,12 @@ class MatMulNBits final : public OpKernel { const Tensor* tensor_zero_point = nullptr; has_zp_input_ = info.TryGetConstantInput(InputIndex::zero_points, &tensor_zero_point); + + const Tensor* weight = nullptr; + // Runtime B may be empty. Defer parameter validation/conversion until Compute + // in that case, and never pack unused parameters for a constant empty B. + skip_prepack_ = !info.TryGetConstantInput(InputIndex::B, &weight) || + weight == nullptr || weight->Shape().Size() == 0; } Status Compute(OpKernelContext* context) const override; @@ -242,6 +248,7 @@ class MatMulNBits final : public OpKernel { const bool has_g_idx_; const bool has_bias_; bool scales_are_packed_{false}; + bool skip_prepack_{false}; const bool force_fp32_; const bool prefer_lut_gemm_{false}; const MLAS_QNBIT_GEMM_COMPUTE_TYPE compute_type_; @@ -344,6 +351,8 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; + if (skip_prepack_) return Status::OK(); + ORT_RETURN_IF_ERROR(ValidatePrePackInitializerShape(OpKernel::Info(), tensor, input_idx, N_, K_, block_size_, nbits_, has_zp_arg_, has_zp_input_)); @@ -728,6 +737,9 @@ template <> Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc, /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights) { + is_packed = false; + if (skip_prepack_) return Status::OK(); + ORT_RETURN_IF_ERROR(ValidatePrePackInitializerShape(OpKernel::Info(), tensor, input_idx, N_, K_, block_size_, nbits_, has_zp_arg_, has_zp_input_)); @@ -1526,8 +1538,14 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { const Tensor* reorder_idx = ctx->Input(InputIndex::g_idx); const Tensor* bias = ctx->Input(InputIndex::bias); - ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( - a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_)); + const bool empty_weights = b != nullptr && b->Shape().Size() == 0; + if (empty_weights) { + ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckEmptyInputs(*b, bias, N_, K_, block_size_, nbits_, + Info().template GetAttrOrDefault("weight_prepacked", 0))); + } else { + ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( + a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_)); + } TensorShape b_shape({static_cast(N_), static_cast(K_)}); MatMulComputeHelper helper; @@ -1540,6 +1558,19 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const { return Status::OK(); } + if (empty_weights) { + auto output = y->MutableDataAsSpan(); + if (bias == nullptr) { + std::fill(output.begin(), output.end(), T1(0.0f)); + } else { + const auto bias_values = bias->DataAsSpan(); + for (size_t offset = 0; offset < output.size(); offset += N_) { + std::copy(bias_values.begin(), bias_values.end(), output.begin() + offset); + } + } + return Status::OK(); + } + AllocatorPtr allocator; ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator)); diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h index 0114e20182308..eeb3117d21d45 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h @@ -34,6 +34,29 @@ constexpr int64_t GetUint8ZeroPointBlobSize(int64_t k_blocks, int64_t bits) { return (k_blocks * bits + 7) / 8; } +// Empty packed weights represent an omitted projection, not a change to logical K/N. +// Keep this separate from CheckInputs: EPs must implement the zero/bias fast path +// before accepting empty weights, and must not dequantize unused parameters. +inline Status CheckEmptyInputs(const Tensor& quantized_weight, const Tensor* bias, + int64_t n, int64_t k, int64_t block_size, int64_t bits, + int64_t weight_prepacked) { + ORT_RETURN_IF_NOT(weight_prepacked == 0, "MatMulNBits: empty B requires weight_prepacked=0"); + ORT_RETURN_IF_NOT(bits == 2 || bits == 4 || bits == 8, "MatMulNBits: unsupported bits: ", bits); + ORT_RETURN_IF_NOT(block_size >= 16 && (block_size & (block_size - 1)) == 0 && + (bits != 2 || block_size <= 256), + "MatMulNBits: unsupported block_size: ", block_size); + const auto& shape = quantized_weight.Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() == 3 && shape.Size() == 0, + "MatMulNBits: empty B must have rank 3"); + const int64_t expected[] = {n, GetKBlocks(k, block_size), GetBlobSize(block_size, bits)}; + for (size_t i = 0; i < 3; ++i) { + ORT_RETURN_IF_NOT(shape[i] == 0 || shape[i] == expected[i], + "MatMulNBits: empty B dimension ", i, " must be 0 or ", expected[i], ", got ", shape[i]); + } + ASSERT_TENSOR_SHAPE(bias, make_shape(n)); + return Status::OK(); +} + template Status CheckInputs(const T* /*activation*/, const T* quantized_weight, diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 404f06f858869..54d515489ddab 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -4,12 +4,14 @@ #include #include +#include "contrib_ops/cpu/quantization/matmul_nbits_helper.h" #include "contrib_ops/webgpu/quantization/matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "core/providers/cpu/math/matmul_helper.h" +#include "core/providers/webgpu/math/matmul.h" #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" #include "core/providers/webgpu/webgpu_utils.h" @@ -229,11 +231,16 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context const Tensor* g_idx = context.Input(4); const Tensor* bias = context.Input(5); - ORT_ENFORCE(g_idx == nullptr, "group_idx as input is not supported yet."); - - const bool has_zero_points = zero_points != nullptr; - if (has_zero_points) { - ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); + const bool empty_weights = b->Shape().Size() == 0; + if (empty_weights) { + ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckEmptyInputs(*b, bias, N_, K_, block_size_, bits_, + Info().GetAttrOrDefault("weight_prepacked", 0))); + } else { + ORT_ENFORCE(g_idx == nullptr, "group_idx as input is not supported yet."); + if (zero_points != nullptr) { + ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), + "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); + } } MatMulComputeHelper helper; @@ -246,6 +253,14 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context return Status::OK(); } + if (empty_weights) { + // Record a fill dispatch so graph replay also overwrites the output. The + // MatMul shader uses [batch, row, column] coordinates, even for a vector A. + auto program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram( + bias, y, false, 1.0f, 1, TensorShape({1, y->Shape().Size() / N_, N_})); + return context.RunProgram(program); + } + return ApplyMatMulNBits(a, b, scales, zero_points, bias, K_, N_, block_size_, accuracy_level_, bits_, context, y, 0); } diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index e05119bd05dfa..fe19adb003419 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -4069,13 +4069,20 @@ It is a fusion of two operations: dequantized_weight = (quantized_weight - zero_point) * scale 2. Matrix multiplication between the input matrix A and the dequantized weight matrix. -The weight matrix is a 2D constant matrix with the input feature count and output feature count specified by attributes 'K' and 'N'. +The logical weight matrix is 2D, with the input feature count and output feature count specified by attributes 'K' and 'N'. It is quantized block-wise along the K dimension with a block size specified by the 'block_size' attribute. The block size must be a power of 2 and not smaller than 16 (e.g., 16, 32, 64, 128). Each block has its own scale and zero-point. The quantization is performed using a bit-width specified by the 'bits' attribute, which can take values from 2 to 8. The quantized weights are stored in a bit-packed format along the K dimension, with each block being represented by a blob of uint8. For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a byte, and the second 4 bits are stored in the higher 4 bits of a byte. + +When weight_prepacked is 0, B may represent an omitted projection by replacing one or more dimensions of its +packed shape (N, k_blocks, blob_size) with 0. All remaining dimensions must retain their normal values. +In this case, Y is filled with zeros (or the broadcast bias, if provided). Its shape is still determined by A, K and N. +The contents of A, scales, zero_points and g_idx are not used; the quantization parameters may also be empty. +This permits zero-sized defaults for optional weights, with dynamic input dimensions allowing nonempty overrides. +Empty packed weights are currently supported by the CPU and WebGPU execution providers. )DOC"; ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBits) @@ -4106,6 +4113,7 @@ For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a by .Input(1, "B", "Packed uint8 tensor of shape (N, k_blocks, blob_size), " "where k_blocks = ceil(K / block_size) and blob_size = (block_size * bits / 8). " + "With weight_prepacked=0, one or more dimensions may be 0 to omit the projection as described above. " "The quantized weights are stored in a bit-packed format along the K dimension, packed within each block_size.", "T2") .Input(2, "scales", "Per-block scaling factors for dequantization with shape (N, k_blocks) and same data type as input A.", "T1") diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc new file mode 100644 index 0000000000000..4bc982fa64248 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc @@ -0,0 +1,227 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS) + +#include + +#include "core/session/IOBinding.h" +#include "core/session/onnxruntime_run_options_config_keys.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/providers/provider_test_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +namespace { + +class MatMulNBitsEmptyWeights : public ::testing::TestWithParam { + void SetUp() override { + if (GetParam() && !DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + } +}; + +template +void RunEmptyCases(bool webgpu) { + struct Case { + int64_t bits, block_size; + std::vector a_shape; + int empty_axis; + bool constant, bias; + }; + const Case cases[] = { + {8, 32, {2, 3, 64}, 0, false, false}, // down projection, batched + {8, 32, {2, 64}, 1, true, false}, // up projection, constant empty B + {4, 16, {64}, 1, true, true}, // vector A and bias + {2, 64, {2, 64}, 2, false, true}, // empty packed-blob axis + {4, 32, {0, 64}, 0, false, false}, // empty output + }; + for (const auto& c : cases) { + SCOPED_TRACE(MakeString("bits=", c.bits, ", block_size=", c.block_size, ", axis=", c.empty_axis)); + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 3); + test.AddAttribute("bits", c.bits); + test.AddAttribute("block_size", c.block_size); + test.AddInput("A", c.a_shape, + std::vector(TensorShape(c.a_shape).Size(), T(std::numeric_limits::quiet_NaN()))); + std::vector b_shape{3, 64 / c.block_size, c.block_size * c.bits / 8}; + b_shape[c.empty_axis] = 0; + test.AddInput("B", b_shape, {}, c.constant); + test.AddInput("scales", {0}, {}, true); + // Unused quantization inputs need neither their normal shapes nor data. + if (!c.constant) { + test.AddInput("zero_points", {0}, {}); + test.AddInput("g_idx", {0}, {}); + } else { + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + } + if (c.bias) test.AddInput("bias", {3}, {T(1.0f), T(-2.0f), T(3.0f)}, c.constant); + auto y_shape = c.a_shape; + y_shape.back() = 3; + std::vector expected(TensorShape(y_shape).Size(), T(0.0f)); + const T bias[] = {T(1.0f), T(-2.0f), T(3.0f)}; + if (c.bias) { + for (size_t i = 0; i < expected.size(); ++i) expected[i] = bias[i % 3]; + } + test.AddOutput("Y", y_shape, expected); + test.SetOutputTolerance(0.0f, 0.0f); + SessionOptions options; + options.use_per_session_threads = false; + if (webgpu) ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + auto ep = webgpu ? DefaultWebGpuExecutionProvider() : DefaultCpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + size_t packed_count = 0; + test.Config(options).ConfigEp(std::move(ep)).RunWithConfig(&packed_count); + EXPECT_EQ(packed_count, 0u); + } +} + +TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp32) { RunEmptyCases(GetParam()); } +TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp16) { RunEmptyCases(GetParam()); } + +TEST_P(MatMulNBitsEmptyWeights, RejectMalformedPackedShape) { + for (bool constant : {false, true}) { + for (const auto& shape : {std::vector{0, 2, 33}, std::vector{0, 64}}) { + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", 3); + test.AddAttribute("bits", 8); + test.AddAttribute("block_size", 32); + test.AddInput("A", {1, 64}, std::vector(64, 1.0f)); + test.AddInput("B", shape, {}, constant); + test.AddInput("scales", {0}, {}, constant); + test.AddOutput("Y", {1, 3}, std::vector(3)); + SessionOptions options; + options.use_per_session_threads = false; + if (GetParam()) + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(options) + .Config(OpTester::ExpectResult::kExpectFailure, "MatMulNBits: empty B") + .ConfigEp(GetParam() ? DefaultWebGpuExecutionProvider() : DefaultCpuExecutionProvider()) + .RunWithConfig(); + } + } +} + +template +void RunSwitching(bool webgpu, bool capture, bool with_bias = true) { + Model model("empty MatMulNBits defaults", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + auto* a = builder.MakeInput({2, 64}, T(0.5f), T(0.5f)); + auto* b = builder.MakeInput({3, 2, 32}, uint8_t{129}, uint8_t{130}); + auto* s = builder.MakeInput({3, 2}, T(0.25f), T(0.25f)); + for (auto* arg : {b, s}) { + ONNX_NAMESPACE::TensorProto init; + init.set_name(arg->Name()); + init.set_data_type(arg == b ? ONNX_NAMESPACE::TensorProto::UINT8 : utils::ToTensorProtoElementType()); + init.add_dims(3); + init.add_dims(0); + if (arg == b) init.add_dims(32); + graph.AddInitializedTensor(init); + auto shape = *arg->Shape(); + shape.mutable_dim(1)->Clear(); + arg->SetShape(shape); + } + graph.SetInputs({a, b, s}); + auto* y = builder.MakeOutput(std::vector{2, 3}); + std::vector inputs{a, b, s}; + if (with_bias) { + auto* bias = builder.MakeInitializer({3}, T(1.0f), T(1.0f)); + inputs.insert(inputs.end(), {builder.MakeEmptyInput(), builder.MakeEmptyInput(), bias}); + } + auto& node = builder.AddNode("MatMulNBits", inputs, {y}, kMSDomain); + node.AddAttribute("K", int64_t{64}); + node.AddAttribute("N", int64_t{3}); + node.AddAttribute("bits", int64_t{8}); + node.AddAttribute("block_size", int64_t{32}); + node.AddAttribute("accuracy_level", int64_t{0}); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + + ConfigOptions config; + if (capture) ASSERT_STATUS_OK(config.AddConfigEntry("ep.webgpuexecutionprovider.enableGraphCapture", "1")); + auto ep = webgpu ? WebGpuExecutionProviderWithOptions(config) : DefaultCpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + auto* provider = ep.get(); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + if (webgpu) ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(ep))); + const auto data = model.ToProto().SerializeAsString(); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + auto off_feeds = builder.feeds_; + off_feeds.erase(b->Name()); + off_feeds.erase(s->Name()); + auto check = [with_bias](const OrtValue& output, int state) { + const T expected((with_bias ? 1.0f : 0.0f) + (state ? 8.0f : 0.0f)); + for (T value : output.Get().DataAsSpan()) EXPECT_EQ(value, expected); + }; + if (!capture) { + for (int state : {0, 1, 0, 1}) { + std::vector output; + ASSERT_STATUS_OK(session.Run(RunOptions{}, state ? builder.feeds_ : off_feeds, builder.output_names_, &output)); + check(output[0], state); + } + return; + } + + const auto allocator = provider->CreatePreferredAllocators().back(); + ASSERT_NE(allocator->Info().device.Type(), OrtDevice::CPU); + std::unique_ptr bindings[2]; + OrtValue outputs[2]; + for (int state : {0, 1}) { + ASSERT_STATUS_OK(session.NewIOBinding(&bindings[state])); + for (const auto& [name, value] : state ? builder.feeds_ : off_feeds) { + const auto& source = value.template Get(); + OrtValue input; + Tensor::InitOrtValue(source.DataType(), source.Shape(), allocator, input); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(source, *input.GetMutable())); + ASSERT_STATUS_OK(bindings[state]->BindInput(name, input)); + } + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({2, 3}), allocator, outputs[state]); + ASSERT_STATUS_OK(bindings[state]->BindOutput(y->Name(), outputs[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeInputs()); + } + OrtValue actual; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({2, 3}), CPUAllocator::DefaultInstance(), actual); + for (int state : {0, 0, 0, 1, 1, 1, 0, 1}) { + auto* cpu = actual.GetMutable(); + auto values = cpu->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), T(-99.0f)); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(*cpu, *outputs[state].GetMutable())); + RunOptions run; + ASSERT_STATUS_OK(run.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, state ? "2" : "1")); + ASSERT_STATUS_OK(session.Run(run, *bindings[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeOutputs()); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(outputs[state].Get(), *cpu)); + check(actual, state); + } + EXPECT_TRUE(provider->IsGraphCaptured(1)); + EXPECT_TRUE(provider->IsGraphCaptured(2)); +} + +TEST_P(MatMulNBitsEmptyWeights, OverridesInSameSessionFp32) { RunSwitching(GetParam(), false); } +TEST_P(MatMulNBitsEmptyWeights, OverridesInSameSessionFp16) { RunSwitching(GetParam(), false); } + +INSTANTIATE_TEST_SUITE_P(Cpu, MatMulNBitsEmptyWeights, ::testing::Values(false)); +#if defined(USE_WEBGPU) +INSTANTIATE_TEST_SUITE_P(WebGpu, MatMulNBitsEmptyWeights, ::testing::Values(true)); +TEST(MatMulNBitsEmptyWeightsCapture, ZeroAndBiasReplay) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (bool bias : {false, true}) { + RunSwitching(true, true, bias); + RunSwitching(true, true, bias); + } +} +#endif + +} // namespace +} // namespace onnxruntime::test +#endif From 42e288346de862d13943871c88d9345f4319eea7 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 23 Oct 2026 17:17:29 +0800 Subject: [PATCH 2/4] Handle empty weights in fused WebGPU MatMulNBits paths --- .../webgpu/quantization/matmul_nbits.cc | 19 +- .../webgpu/quantization/matmul_nbits_mlp.cc | 3 + .../webgpu/quantization/matmul_nbits_qkv.cc | 5 +- .../matmul_nbits_empty_weights_test.cc | 240 +++++++++++++++++- 4 files changed, 256 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc index 54d515489ddab..a1cfedc0785c6 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc @@ -253,14 +253,6 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context return Status::OK(); } - if (empty_weights) { - // Record a fill dispatch so graph replay also overwrites the output. The - // MatMul shader uses [batch, row, column] coordinates, even for a vector A. - auto program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram( - bias, y, false, 1.0f, 1, TensorShape({1, y->Shape().Size() / N_, N_})); - return context.RunProgram(program); - } - return ApplyMatMulNBits(a, b, scales, zero_points, bias, K_, N_, block_size_, accuracy_level_, bits_, context, y, 0); } @@ -313,6 +305,17 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, MatMulComputeHelper helper; ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true)); + if (b->Shape().Size() == 0) { + ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckEmptyInputs(*b, bias, N_op, K_op, block_size_op, nbits, 0)); + if (y->Shape().Size() == 0) return Status::OK(); + // Shared by standalone MatMulNBits and fused QKV/MLP fallbacks. Record a + // fill dispatch so capture replay overwrites the output without reading A + // or unused quantization parameters. The fill uses [batch, row, column]. + auto program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram( + bias, y, false, 1.0f, 1, TensorShape({1, y->Shape().Size() / N_op, N_op})); + return context.RunProgram(program); + } + const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc index 3b94812c481f9..ab500c7742bb4 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc @@ -293,8 +293,11 @@ Status MatMulNBitsMlp::ComputeInternal(onnxruntime::webgpu::ComputeContext& cont const bool has_skip_input = skip != nullptr; const bool has_skip_output = input_skip_bias_sum != nullptr; + // The fallback fills empty projections with zero/bias before applying the + // activation and product; an empty gate with bias need not yield zero. const bool is_decode_fast_path_candidate = M == 1 && + gate_b->Shape().Size() != 0 && up_b->Shape().Size() != 0 && bits_ == kFusedDecodeFastPathBits && block_size == kFusedDecodeFastPathBlockSize; const bool has_norm_input = norm_scale != nullptr; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc index 7137b6a3f73ab..8c62a0d2c5fc3 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc @@ -233,7 +233,10 @@ Status MatMulNBitsQkv::ComputeInternal(onnxruntime::webgpu::ComputeContext& cont const bool exceeds_storage_buffer_limit = required_storage_buffers > context.DeviceLimits().maxStorageBuffersPerShaderStage; - if (would_use_subgroup_unfused || would_use_dp4a_unfused || would_use_wide_tile_unfused || + // The decode shader reads every projection's packed weights. The fallback + // handles empty projections individually and keeps nonempty projections active. + const bool has_empty_weights = q_b->Shape().Size() == 0 || k_b->Shape().Size() == 0 || v_b->Shape().Size() == 0; + if (has_empty_weights || would_use_subgroup_unfused || would_use_dp4a_unfused || would_use_wide_tile_unfused || M != 1) { Tensor normalized_a = context.CreateGPUTensor(a->DataType(), a->Shape()); diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc index 4bc982fa64248..3e95a4509a6ef 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc @@ -9,12 +9,29 @@ #include "core/session/onnxruntime_run_options_config_keys.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "test/providers/provider_test_utils.h" +#include "test/compare_ortvalue.h" #include "test/unittest_util/graph_transform_test_builder.h" #include "test/util/include/default_providers.h" +#include "test/util/include/inference_session_wrapper.h" + +#if defined(USE_WEBGPU) && !defined(ORT_USE_EP_API_ADAPTERS) +#include "core/providers/webgpu/webgpu_context.h" +#endif namespace onnxruntime::test { namespace { +bool WebGpuSupportsFp16() { +#if defined(USE_WEBGPU) && !defined(ORT_USE_EP_API_ADAPTERS) + auto ep = DefaultWebGpuExecutionProvider(); + return ep && webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()).DeviceHasFeature(wgpu::FeatureName::ShaderF16); +#else + // The plugin EP interface does not expose Dawn device features. Keep CPU FP16 + // and WebGPU FP32 coverage without assuming that the plugin device supports FP16. + return false; +#endif +} + class MatMulNBitsEmptyWeights : public ::testing::TestWithParam { void SetUp() override { if (GetParam() && !DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; @@ -79,7 +96,10 @@ void RunEmptyCases(bool webgpu) { } TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp32) { RunEmptyCases(GetParam()); } -TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp16) { RunEmptyCases(GetParam()); } +TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp16) { + if (GetParam() && !WebGpuSupportsFp16()) GTEST_SKIP() << "WebGPU shader-f16 support is unavailable."; + RunEmptyCases(GetParam()); +} TEST_P(MatMulNBitsEmptyWeights, RejectMalformedPackedShape) { for (bool constant : {false, true}) { @@ -208,15 +228,231 @@ void RunSwitching(bool webgpu, bool capture, bool with_bias = true) { } TEST_P(MatMulNBitsEmptyWeights, OverridesInSameSessionFp32) { RunSwitching(GetParam(), false); } -TEST_P(MatMulNBitsEmptyWeights, OverridesInSameSessionFp16) { RunSwitching(GetParam(), false); } +TEST_P(MatMulNBitsEmptyWeights, OverridesInSameSessionFp16) { + if (GetParam() && !WebGpuSupportsFp16()) GTEST_SKIP() << "WebGPU shader-f16 support is unavailable."; + RunSwitching(GetParam(), false); +} INSTANTIATE_TEST_SUITE_P(Cpu, MatMulNBitsEmptyWeights, ::testing::Values(false)); #if defined(USE_WEBGPU) INSTANTIATE_TEST_SUITE_P(WebGpu, MatMulNBitsEmptyWeights, ::testing::Values(true)); + +enum class WeightSource { Constant, + Input, + Overridable }; + +void RunFusedEmptyWeights(bool qkv, int64_t m, int empty_mask, WeightSource source, bool with_skip, + bool with_bias, bool capture = false) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + SCOPED_TRACE(MakeString("qkv=", qkv, ", M=", m, ", mask=", empty_mask, + ", source=", static_cast(source), ", skip=", with_skip, ", bias=", with_bias)); + Model model("fused empty MatMulNBits", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + constexpr int64_t k = 64; + auto* x = builder.MakeInput({1, m, k}, 0.5f, 0.5f); + std::vector graph_inputs{x}; + auto* norm_scale = builder.MakeInitializer({k}, 1.0f, 1.0f); + auto* normalized = builder.MakeIntermediate(std::vector{1, m, k}); + if (with_skip) { + auto* skip = builder.MakeInput({1, m, k}, 0.25f, 0.25f); + graph_inputs.push_back(skip); + auto* residual = builder.MakeOutput(std::vector{1, m, k}); + auto& norm = builder.AddNode("SkipSimplifiedLayerNormalization", {x, skip, norm_scale}, + {normalized, builder.MakeEmptyInput(), builder.MakeEmptyInput(), residual}, kMSDomain); + norm.AddAttribute("epsilon", 1e-6f); + } else { + builder.AddNode("SimplifiedLayerNormalization", {x, norm_scale}, {normalized}).AddAttribute("epsilon", 1e-6f); + } + + std::vector projections; + InlinedHashMap empty_inputs; + for (int i = 0; i < (qkv ? 3 : 2); ++i) { + const int64_t n = qkv && i == 0 ? 128 : 64; + const bool empty = (empty_mask & (1 << i)) != 0; + const int axis = i; + std::vector shape{n, 2, 16}; + auto empty_shape = shape; + empty_shape[axis] = 0; + NodeArg* weights; + if (source == WeightSource::Constant) { + weights = builder.MakeInitializer(empty ? empty_shape : shape, uint8_t{0x99}, uint8_t{0x9a}); + } else { + weights = builder.MakeInput(shape, uint8_t{0x99}, uint8_t{0x9a}); + graph_inputs.push_back(weights); + auto dynamic_shape = *weights->Shape(); + dynamic_shape.mutable_dim(axis)->Clear(); + weights->SetShape(dynamic_shape); + if (source == WeightSource::Overridable) { + ONNX_NAMESPACE::TensorProto init; + init.set_name(weights->Name()); + init.set_data_type(ONNX_NAMESPACE::TensorProto::UINT8); + for (auto dim : empty ? empty_shape : shape) init.add_dims(dim); + if (!empty) init.set_raw_data(std::string(static_cast(n * 2 * 16), static_cast(0x99))); + graph.AddInitializedTensor(init); + } else if (empty) { + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(empty_shape), + CPUAllocator::DefaultInstance(), empty_inputs[weights->Name()]); + } + } + auto* scales = builder.MakeInitializer({n, 2}, 0.125f, 0.125f); + auto* output = qkv ? builder.MakeOutput(std::vector{1, m, n}) + : builder.MakeIntermediate(std::vector{1, m, n}); + std::vector inputs{normalized, weights, scales}; + if (with_bias) { + auto* bias = builder.MakeInitializer({n}, 1.0f, 1.0f); + inputs.insert(inputs.end(), {builder.MakeEmptyInput(), builder.MakeEmptyInput(), bias}); + } + auto& node = builder.AddNode("MatMulNBits", inputs, {output}, kMSDomain); + node.AddAttribute("K", k); + node.AddAttribute("N", n); + node.AddAttribute("bits", int64_t{4}); + node.AddAttribute("block_size", int64_t{32}); + node.AddAttribute("accuracy_level", int64_t{0}); + projections.push_back(output); + } + if (!qkv) { + auto* activated = builder.MakeIntermediate(std::vector{1, m, 64}); + builder.AddNode("QuickGelu", {projections[0]}, {activated}, kMSDomain).AddAttribute("alpha", 1.0f); + auto* y = builder.MakeOutput(std::vector{1, m, 64}); + builder.AddNode("Mul", {activated, projections[1]}, {y}); + } + graph.SetInputs(graph_inputs); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + const auto data = model.ToProto().SerializeAsString(); + auto off_feeds = builder.feeds_; + if (source == WeightSource::Overridable) { + for (const auto* input : graph.GetOverridableInitializers()) off_feeds.erase(input->Name()); + } else { + for (const auto& [name, value] : empty_inputs) off_feeds[name] = value; + } + + std::vector expected[2]; + for (bool optimized : {false, true}) { + ConfigOptions config; + if (capture && optimized) ASSERT_STATUS_OK(config.AddConfigEntry("ep.webgpuexecutionprovider.enableGraphCapture", "1")); + auto ep = WebGpuExecutionProviderWithOptions(config); + ASSERT_NE(ep, nullptr); + auto* provider = ep.get(); + SessionOptions options; + options.graph_optimization_level = optimized ? TransformerLevel::Level2 : TransformerLevel::Level1; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(ep))); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + size_t fused_count = 0; + for (const auto& node : session.GetGraph().Nodes()) { + fused_count += node.OpType() == (qkv ? "MatMulNBitsQkv" : "MatMulNBitsMlp"); + } + ASSERT_EQ(fused_count, optimized ? 1u : 0u); + + if (!capture || !optimized) { + for (int state : {0, 1, 0, 1}) { + if (state && source == WeightSource::Constant) continue; + std::vector outputs; + ASSERT_STATUS_OK(session.Run(RunOptions{}, state ? builder.feeds_ : off_feeds, builder.output_names_, &outputs)); + if (!optimized) { + expected[state] = std::move(outputs); + } else { + ASSERT_EQ(outputs.size(), expected[state].size()); + for (size_t i = 0; i < outputs.size(); ++i) { + auto comparison = CompareOrtValue(outputs[i], expected[state][i], 1e-4, 1e-4, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + } + } + continue; + } + + const auto allocator = provider->CreatePreferredAllocators().back(); + ASSERT_NE(allocator->Info().device.Type(), OrtDevice::CPU); + std::unique_ptr bindings[2]; + std::vector gpu_outputs[2]; + for (int state : {0, 1}) { + ASSERT_STATUS_OK(session.NewIOBinding(&bindings[state])); + for (const auto& [name, value] : state ? builder.feeds_ : off_feeds) { + const auto& tensor = value.Get(); + OrtValue input; + Tensor::InitOrtValue(tensor.DataType(), tensor.Shape(), allocator, input); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(tensor, *input.GetMutable())); + ASSERT_STATUS_OK(bindings[state]->BindInput(name, input)); + } + gpu_outputs[state].resize(expected[state].size()); + for (size_t i = 0; i < expected[state].size(); ++i) { + const auto& tensor = expected[state][i].Get(); + Tensor::InitOrtValue(tensor.DataType(), tensor.Shape(), allocator, gpu_outputs[state][i]); + ASSERT_STATUS_OK(bindings[state]->BindOutput(builder.output_names_[i], gpu_outputs[state][i])); + } + ASSERT_STATUS_OK(bindings[state]->SynchronizeInputs()); + } + for (int state : {0, 0, 0, 1, 1, 1, 0, 1}) { + std::vector actual(expected[state].size()); + for (size_t i = 0; i < actual.size(); ++i) { + const auto& tensor = expected[state][i].Get(); + Tensor::InitOrtValue(tensor.DataType(), tensor.Shape(), CPUAllocator::DefaultInstance(), actual[i]); + auto values = actual[i].GetMutable()->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), -99.0f); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(actual[i].Get(), + *gpu_outputs[state][i].GetMutable())); + } + RunOptions run; + ASSERT_STATUS_OK(run.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, state ? "2" : "1")); + ASSERT_STATUS_OK(session.Run(run, *bindings[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeOutputs()); + for (size_t i = 0; i < actual.size(); ++i) { + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(gpu_outputs[state][i].Get(), + *actual[i].GetMutable())); + auto comparison = CompareOrtValue(actual[i], expected[state][i], 1e-4, 1e-4, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + } + EXPECT_TRUE(provider->IsGraphCaptured(1)); + EXPECT_TRUE(provider->IsGraphCaptured(2)); + } +} + +TEST(MatMulNBitsEmptyWeightsFusion, Qkv) { + for (auto source : {WeightSource::Constant, WeightSource::Input, WeightSource::Overridable}) { + for (int64_t m : {1, 3}) { + for (int mask : {1, 2, 4, 7}) { + ASSERT_NO_FATAL_FAILURE(RunFusedEmptyWeights(true, m, mask, source, m == 1, false)); + } + } + } +} + +TEST(MatMulNBitsEmptyWeightsFusion, Mlp) { + for (auto source : {WeightSource::Constant, WeightSource::Input, WeightSource::Overridable}) { + for (int64_t m : {1, 3}) { + for (int mask : {1, 2, 3}) { + for (bool bias : {false, true}) { + ASSERT_NO_FATAL_FAILURE(RunFusedEmptyWeights(false, m, mask, source, m == 1, bias)); + } + } + } + } +} + +TEST(MatMulNBitsEmptyWeightsFusion, QkvReplay) { + ASSERT_NO_FATAL_FAILURE(RunFusedEmptyWeights(true, 1, 1, WeightSource::Overridable, true, false, true)); +} + +TEST(MatMulNBitsEmptyWeightsFusion, MlpReplay) { + ASSERT_NO_FATAL_FAILURE(RunFusedEmptyWeights(false, 1, 1, WeightSource::Overridable, true, true, true)); +} + TEST(MatMulNBitsEmptyWeightsCapture, ZeroAndBiasReplay) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; for (bool bias : {false, true}) { RunSwitching(true, true, bias); + } +} +TEST(MatMulNBitsEmptyWeightsCapture, ZeroAndBiasReplayFp16) { + if (!WebGpuSupportsFp16()) GTEST_SKIP() << "WebGPU shader-f16 support is unavailable."; + for (bool bias : {false, true}) { RunSwitching(true, true, bias); } } From a9cfff87627d53aa4d96e64e0eea4d0314a26aeb Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 26 Oct 2026 09:41:52 +0800 Subject: [PATCH 3/4] Fix generated MatMulNBits documentation formatting --- docs/ContribOperators.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index c6809ce5f9d73..754d0fa38d448 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4188,7 +4188,7 @@ This version of the operator has been available since version 1 of the 'com.micr The quantized weights are stored in a bit-packed format along the K dimension, with each block being represented by a blob of uint8. For example, for 4 bits, the first 4 bits are stored in the lower 4 bits of a byte, and the second 4 bits are stored in the higher 4 bits of a byte. - + When weight_prepacked is 0, B may represent an omitted projection by replacing one or more dimensions of its packed shape (N, k_blocks, blob_size) with 0. All remaining dimensions must retain their normal values. In this case, Y is filled with zeros (or the broadcast bias, if provided). Its shape is still determined by A, K and N. From 93fcef5ef556a84be4181bccceba6a12c8bf2631 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Oct 2026 15:30:56 +0800 Subject: [PATCH 4/4] Reject empty MatMulNBits weights on CUDA --- .../cuda/quantization/matmul_nbits.cc | 13 +++++++ .../matmul_nbits_empty_weights_test.cc | 35 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 51a0f89625fa7..5bee3dc87918e 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -92,6 +92,15 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr /*out*/ bool& is_packed, /*out*/ PrePackedWeights* /*prepacked_weights*/) { is_packed = false; + if (tensor.Shape().Size() == 0) { + // Reject empty weights before fpA_intB prepacking can read their buffer. + if (input_idx == MatMulNBits_Input_B) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "CUDA MatMulNBits does not support empty B weights yet."); + } + // Leave empty quantization parameters for runtime validation: B may be a runtime input. + return Status::OK(); + } constexpr int kInputIndexGroupIndex = 4; if (input_idx == kInputIndexGroupIndex) { const int64_t k_blocks = K_ / block_size_ + (K_ % block_size_ != 0); @@ -1227,6 +1236,10 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { const auto* scales_data = scales->Data(); const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw(); #endif + if (b != nullptr && b->Shape().Size() == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "CUDA MatMulNBits does not support empty B weights yet."); + } const auto* bias_data = bias == nullptr ? nullptr : bias->Data(); ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc index 3e95a4509a6ef..d6c09266d74a7 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc @@ -101,6 +101,41 @@ TEST_P(MatMulNBitsEmptyWeights, ZeroAndBiasFp16) { RunEmptyCases(GetParam()); } +#if defined(USE_CUDA) +TEST(MatMulNBitsEmptyWeightsCuda, RejectsEmptyWeights) { + for (bool constant : {false, true}) { + for (int empty_axis : {0, 1, 2}) { + SCOPED_TRACE(MakeString("constant=", constant, ", empty_axis=", empty_axis)); + auto ep = DefaultCudaExecutionProvider(); + if (!ep) GTEST_SKIP() << "CUDA unavailable."; + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 128); + test.AddAttribute("bits", 4); + test.AddAttribute("block_size", 32); + test.AddInput("A", {1, 128}, std::vector(128, MLFloat16(1.0f))); + std::vector b_shape{128, 4, 16}; + b_shape[empty_axis] = 0; + test.AddInput("B", b_shape, {}, constant); + test.AddInput("scales", {0}, {}, true); + test.AddOutput("Y", {1, 128}, std::vector(128)); + + SessionOptions options; + options.use_per_session_threads = false; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + // Exercise the prepacking rejection in fpA_intB builds as well as the runtime input check. + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsCudaFpAIntBGemm, "1")); + test.Config(options) + .Config(OpTester::ExpectResult::kExpectFailure, + "CUDA MatMulNBits does not support empty B weights yet.") + .ConfigEp(std::move(ep)) + .RunWithConfig(); + } + } +} +#endif + TEST_P(MatMulNBitsEmptyWeights, RejectMalformedPackedShape) { for (bool constant : {false, true}) { for (const auto& shape : {std::vector{0, 2, 33}, std::vector{0, 64}}) {