[webgpu] Support any batch size for dp4a matmul path (#26884)

This pull request adds support for batched matrix multiplication in the
DP4A quantized matmul WebGPU kernels and their associated C++ code and
tests. The changes update the kernel code, tensor shapes, dispatch
logic, and test infrastructure to properly handle a `batch_count`
greater than 1, enabling efficient batched execution.

Author: qjia7

onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template

@@ -54,18 +54,18 @@ var<workgroup> scale_B : array<output_element_t, tile_size>;
     var<workgroup> zeroes : array<i32, tile_size>;
 #endif
 
-fn loadSHMA(a_global_base:u32, kidx_v:u32, row: u32, col: u32)
+fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32)
 {
     let a_global = a_global_base + row;
     if (a_global >= uniforms.M)
     {
         return;
     }
-    tile_A[col][row] = a.getByOffset(a_global*uniforms.K16+kidx_v+col);
+    tile_A[col][row] = a.getByOffset(batch*uniforms.M*uniforms.K16+a_global*uniforms.K16+kidx_v+col);
     if (col == 0)
     {
         // kidx_v - covers 16 values of k
-        scale_A[row] = scales_a.getByOffset(a_global*(uniforms.K/128) + kidx_v/8);
+        scale_A[row] = scales_a.getByOffset(batch*uniforms.M*(uniforms.K/128) + a_global*(uniforms.K/128) + kidx_v/8);
     }
 }
 
@@ -154,7 +154,11 @@ $MAIN {
 #endif
     // During the load phase we use all 256 threads to load 64 rows of A/B.
     // For each row we load tile_size_k_vec (2) vectorized elements, which are 32 elements of K.
-    let a_global_base = u32(workgroup_idx / uniforms.num_N_tile) * tile_size;
+    let batch = workgroup_idx / (uniforms.num_M_tile * uniforms.num_N_tile);
+    if (batch >= uniforms.batch_count) {
+        return;
+    }
+    let a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_size;
     let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size;
     let load_AorB = u32(local_idx/128);
     let load_row = u32((local_idx%128)/2);
@@ -199,7 +203,7 @@ $MAIN {
         // Load Phase: Populate shared memory for the workgroup.
         if (load_AorB == 0)
         {
-            loadSHMA(a_global_base, kidx_v, load_row, load_col);
+            loadSHMA(batch, a_global_base, kidx_v, load_row, load_col);
         }
         else
         {
@@ -380,7 +384,7 @@ $MAIN {
 
     let a_global = a_global_base + base_A + a_idx;
     let b_global = b_global_base + base_B;
-    let output_idx = ((a_global) * uniforms.N + b_global)/4;
+    let output_idx = (batch * uniforms.M * uniforms.N + a_global * uniforms.N + b_global)/4;
 #if has_bias
 #if has_weight_idx
     let b_bias_offset = uniforms.weight_idx * uniforms.N;

onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc

@@ -83,6 +83,7 @@ Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) co
 
 Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales,
                                   const Tensor* zero_points, const Tensor* bias,
+                                  uint32_t batch_count,
                                   uint32_t M,
                                   uint32_t N,
                                   uint32_t K,
@@ -101,15 +102,15 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
   DP4AMatMulQuantizeProgram quantize_program;
   quantize_program.SetWorkgroupSize(64);
   uint32_t tile_size = 64 * kVec4Components;
-  quantize_program.SetDispatchGroupSize((M * K + tile_size - 1) / tile_size, 1, 1);
-  TensorShape a_quant_shape{1, M, K / kU32Components};
+  quantize_program.SetDispatchGroupSize((batch_count * M * K + tile_size - 1) / tile_size, 1, 1);
+  TensorShape a_quant_shape{batch_count, M, K / kU32Components};
   Tensor a_quant = context.CreateGPUTensor(DataTypeImpl::GetType<uint32_t>(), a_quant_shape);
-  TensorShapeVector a_scales_dims({1, 1, M, K / kBlockSizeA});
+  TensorShapeVector a_scales_dims({batch_count, 1, M, K / kBlockSizeA});
   Tensor a_scale = context.CreateGPUTensor(a->DataType(), a_scales_dims);
   quantize_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>(kVec4Components)}})
       .AddOutputs({{&a_quant, ProgramTensorMetadataDependency::Rank, a_quant.Shape(), 1},
                    {&a_scale, ProgramTensorMetadataDependency::Rank, 1}})
-      .AddUniformVariable({M * K / kU32Components});
+      .AddUniformVariable({batch_count * M * K / kU32Components});
   ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program));
 
   const bool has_zero_points = zero_points != nullptr;
@@ -128,12 +129,12 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
     DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, single_scale_weights};
     uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n;
     mul_program.SetWorkgroupSize(128);
-    mul_program.SetDispatchGroupSize(M * num_N_tile);
+    mul_program.SetDispatchGroupSize(batch_count * M * num_N_tile);
     mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>(kVec4Components)},
                            {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1},
                            {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>(b_components * kU32Components)},
                            {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}})
-        .AddUniformVariables({M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index})
+        .AddUniformVariables({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index})
         .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1})
         .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx);
     if (has_zero_points) {
@@ -146,22 +147,24 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
   }
 
   constexpr uint32_t kTileSize = 64;
-  TensorShape reshaped_y_shape{1, M, N / kVec4Components};
+  TensorShape reshaped_y_shape{batch_count, M, N / kVec4Components};
   uint32_t num_M_tile = (M + kTileSize - 1) / kTileSize;
   uint32_t num_N_tile = (N + kTileSize - 1) / kTileSize;
   bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"};
   DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, is_qualcomm};
   mul_program.SetWorkgroupSize(256);
-  mul_program.SetDispatchGroupSize(num_M_tile * num_N_tile);
+  mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile);
   mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>(kVec4Components)},
                          {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1},
                          {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>((nbits / 2) * kU32Components)},
                          {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}})
-      .AddUniformVariables({{static_cast<uint32_t>(M)},
+      .AddUniformVariables({{static_cast<uint32_t>(batch_count)},
+                            {static_cast<uint32_t>(M)},
                             {static_cast<uint32_t>(N)},
                             {static_cast<uint32_t>(K)},
                             {static_cast<uint32_t>(K / 8)},
                             {static_cast<uint32_t>(K / 16)},
+                            {num_M_tile},
                             {num_N_tile},
                             {zero_blocks_per_col},
                             {weight_index}})
@@ -179,7 +182,6 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
 bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context,
                                    uint64_t accuracy_level,
                                    uint32_t block_size,
-                                   uint32_t batch_count,
                                    uint32_t N,
                                    uint32_t K,
                                    uint32_t components_k) {
@@ -189,7 +191,7 @@ bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context,
   bool use_dp4a = context.HasFeature(wgpu::FeatureName::Subgroups) &&
                   context.AdapterInfo().vendor != std::string_view{"apple"};
   return (accuracy_level == 4 && block_size % 32 == 0 &&
-          batch_count == 1 && components_k == 4 && K % 128 == 0 && N % 16 == 0 &&
+          components_k == 4 && K % 128 == 0 && N % 16 == 0 &&
           use_dp4a);
 }
 

onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h

@@ -32,11 +32,13 @@ class DP4AMatMulNBitsProgram final : public Program<DP4AMatMulNBitsProgram> {
                                                                   is_qualcomm_(is_qualcomm) {}
   Status GenerateShaderCode(ShaderHelper& sh) const override;
   WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
+      {"batch_count", ProgramUniformVariableDataType::Uint32},
       {"M", ProgramUniformVariableDataType::Uint32},
       {"N", ProgramUniformVariableDataType::Uint32},
       {"K", ProgramUniformVariableDataType::Uint32},
       {"K8", ProgramUniformVariableDataType::Uint32},
       {"K16", ProgramUniformVariableDataType::Uint32},
+      {"num_M_tile", ProgramUniformVariableDataType::Uint32},
       {"num_N_tile", ProgramUniformVariableDataType::Uint32},
       {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32},
       {"weight_idx", ProgramUniformVariableDataType::Uint32});
@@ -64,6 +66,7 @@ class DP4AMatMulNBitsSmallMProgram final : public Program<DP4AMatMulNBitsSmallMP
                                                                                  single_scale_weights_(single_scale_weights) {}
   Status GenerateShaderCode(ShaderHelper& sh) const override;
   WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
+      {"batch_count", ProgramUniformVariableDataType::Uint32},
       {"M", ProgramUniformVariableDataType::Uint32},
       {"N", ProgramUniformVariableDataType::Uint32},
       {"K", ProgramUniformVariableDataType::Uint32},
@@ -86,6 +89,7 @@ class DP4AMatMulNBitsSmallMProgram final : public Program<DP4AMatMulNBitsSmallMP
 
 Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales,
                                   const Tensor* zero_points, const Tensor* bias,
+                                  uint32_t batch_count,
                                   uint32_t M,
                                   uint32_t N,
                                   uint32_t K,
@@ -100,7 +104,6 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
 bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context,
                                    uint64_t accuracy_level,
                                    uint32_t block_size,
-                                   uint32_t batch_count,
                                    uint32_t N,
                                    uint32_t K,
                                    uint32_t components_k);

onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template

@@ -45,23 +45,27 @@ var<workgroup> tile_A : array<vec4<u32>, double_tile_size_k_vec>;
 const scale_a_size_in_tile_a = double_tile_size_k_vec / 8;
 var<workgroup> scale_A : array<output_element_t, scale_a_size_in_tile_a>;
 
-fn loadSHMA(a_global: u32, kidx_v: u32, col: u32)
+fn loadSHMA(batch: u32, a_global: u32, kidx_v: u32, col: u32)
 {
     let k_offset = kidx_v + col;
     if (k_offset >= uniforms.K16) {
     return;
     }
 
-    tile_A[col] = a.getByOffset(a_global*uniforms.K16+k_offset);
+    tile_A[col] = a.getByOffset(batch*uniforms.M*uniforms.K16+a_global*uniforms.K16+k_offset);
     if (col < scale_a_size_in_tile_a)
     {
     // kidx_v - covers 16 values of k in input_a
-    scale_A[col] = scales_a.getByOffset(a_global*(uniforms.K/128) + kidx_v/8 + col);
+    scale_A[col] = scales_a.getByOffset(batch*uniforms.M*(uniforms.K/128) + a_global*(uniforms.K/128) + kidx_v/8 + col);
     }
 }
 
 $MAIN {
-    let a_global = u32(workgroup_idx / uniforms.num_N_tile);
+    let batch = workgroup_idx / (uniforms.M * uniforms.num_N_tile);
+    if (batch >= uniforms.batch_count) {
+        return;
+    }
+    let a_global = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.M);
     let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_size;
     // Handle each workgroup threads as a block of [sub_tile_count][tile_size_k_vec]
     let local_col = local_idx % tile_size_k_vec;
@@ -95,7 +99,7 @@ $MAIN {
         // Load Phase: Populate shared memory for the workgroup.
         if (local_idx < double_tile_size_k_vec)
         {
-        loadSHMA(a_global, kidx_v * 2, local_idx);
+        loadSHMA(batch, a_global, kidx_v * 2, local_idx);
         }
         workgroupBarrier();
         var own_a: vec4<u32> = tile_A[local_col * 2];
@@ -153,7 +157,7 @@ $MAIN {
         output_value += inter_results[local_idx][b];
       }
       let b_global =  b_global_base + local_idx;
-      let output_idx = a_global * uniforms.N + b_global;
+      let output_idx = batch * uniforms.M * uniforms.N + a_global * uniforms.N + b_global;
       if (b_global < uniforms.N) {
 #if has_bias
         let bias_value = bias[b_global + b_bias_offset];

onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc

@@ -220,8 +220,8 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales,
 
   // On FP32 only GPUs, integer math is faster than FP32 therefore always use DP4A independent of length of M.
   if ((M >= kMinMForTileOptimization || y->DataType() == DataTypeImpl::GetType<float>() || context.AdapterInfo().vendor == std::string_view{"qualcomm"}) &&
-      CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, batch_count, N, K, components_a)) {
-    return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast<uint32_t>(nbits), context, y, weight_index);
+      CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a)) {
+    return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast<uint32_t>(nbits), context, y, weight_index);
   }
 
   // WideTileProgram