diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index a0036e6c4ffdb..754d0fa38d448 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -4181,13 +4181,20 @@ 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.
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
@@ -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/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/contrib_ops/webgpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc
index 404f06f858869..a1cfedc0785c6 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;
@@ -298,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/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..d6c09266d74a7
--- /dev/null
+++ b/onnxruntime/test/contrib_ops/matmul_nbits_empty_weights_test.cc
@@ -0,0 +1,498 @@
+// 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/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.";
+ }
+};
+
+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) {
+ if (GetParam() && !WebGpuSupportsFp16()) GTEST_SKIP() << "WebGPU shader-f16 support is unavailable.";
+ 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}}) {
+ 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) {
+ 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);
+ }
+}
+#endif
+
+} // namespace
+} // namespace onnxruntime::test
+#endif