[CPU/WebGPU] Support empty MatMulNBits weights

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
 <dt><tt>A</tt> : T1</dt>
 <dd>The input tensor, not quantized.</dd>
 <dt><tt>B</tt> : T2</dt>
-<dd>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.</dd>
+<dd>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.</dd>
 <dt><tt>scales</tt> : T1</dt>
 <dd>Per-block scaling factors for dequantization with shape (N, k_blocks) and same data type as input A.</dd>
 <dt><tt>zero_points</tt> (optional) : T3</dt>

onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

@@ -221,6 +221,12 @@
 
     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 @@
   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 @@
                                 /*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 @@
 Status MatMulNBits<MLFloat16>::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 @@
   const Tensor* reorder_idx = ctx->Input<Tensor>(InputIndex::g_idx);
   const Tensor* bias = ctx->Input<Tensor>(InputIndex::bias);
 
-  ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs<Tensor>(
-      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<int64_t>("weight_prepacked", 0)));
+  } else {
+    ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs<Tensor>(
+        a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_));
+  }
 
   TensorShape b_shape({static_cast<int64_t>(N_), static_cast<int64_t>(K_)});
   MatMulComputeHelper helper;
@@ -1540,6 +1558,19 @@
     return Status::OK();
   }
 
+  if (empty_weights) {
+    auto output = y->MutableDataAsSpan<T1>();
+    if (bias == nullptr) {
+      std::fill(output.begin(), output.end(), T1(0.0f));
+    } else {
+      const auto bias_values = bias->DataAsSpan<T1>();
+      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));
 

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 <typename T = Tensor>
 Status CheckInputs(const T* /*activation*/,
                    const T* quantized_weight,

onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc

@@ -92,6 +92,15 @@ Status MatMulNBits<T>::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<T>::ComputeInternal(OpKernelContext* ctx) const {
   const auto* scales_data = scales->Data<T>();
   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<T>();
 
   ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs<Tensor>(

onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc

@@ -4,12 +4,14 @@
 #include <string>
 #include <string_view>
 
+#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<uint8_t>(), "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<int64_t>("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<uint8_t>(),
+                  "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;

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;

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());
 

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")