[WebGPU] QKV and MLP layer fusions for Qwen3-style models

cmake/onnxruntime_unittests.cmake

@@ -1373,7 +1373,8 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
       ${BENCHMARK_DIR}/activation.cc
       ${BENCHMARK_DIR}/quantize.cc
       ${BENCHMARK_DIR}/reduceminmax.cc
-      ${BENCHMARK_DIR}/layer_normalization.cc)
+      ${BENCHMARK_DIR}/layer_normalization.cc
+      ${BENCHMARK_DIR}/webgpu_matmul_nbits_decode.cc)
     target_include_directories(onnxruntime_benchmark PRIVATE ${ONNXRUNTIME_ROOT} ${onnxruntime_graph_header} ${ONNXRUNTIME_ROOT}/core/mlas/inc)
     target_compile_definitions(onnxruntime_benchmark PRIVATE BENCHMARK_STATIC_DEFINE)
     target_compile_definitions(onnxruntime_benchmark PRIVATE ${mlas_private_compile_definitions})

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

@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#param tile_size
+#param tile_size_k_vec
+#param single_scale_weights
+#param has_gate_bias
+#param has_up_bias
+
+#use .getByOffset .setByOffset
+
+#include "quantization/dp4a_matmul_common.wgsl.template"
+
+const double_tile_size_k_vec = 2 * tile_size_k_vec;
+const scale_a_size_in_tile_a = double_tile_size_k_vec / 8;
+
+var<workgroup> gate_inter_results: array<array<output_element_t, tile_size_k_vec>, tile_size>;
+var<workgroup> up_inter_results: array<array<output_element_t, tile_size_k_vec>, tile_size>;
+var<workgroup> tile_A: array<vec4<u32>, double_tile_size_k_vec>;
+var<workgroup> scale_A: array<output_element_t, scale_a_size_in_tile_a>;
+
+fn loadSHMA(batch: u32, kidx_v: u32, col: u32) {
+  let k_offset = kidx_v + col;
+  if (k_offset >= uniforms.K16) {
+    return;
+  }
+
+  tile_A[col] = a.getByOffset(batch * uniforms.K16 + k_offset);
+  if (col < scale_a_size_in_tile_a) {
+    scale_A[col] = scales_a.getByOffset(batch * (uniforms.K / 128) + kidx_v / 8 + col);
+  }
+}
+
+$MAIN {
+  let batch = workgroup_id.z;
+  if (batch >= uniforms.batch_count) {
+    return;
+  }
+
+  let b_global_base = workgroup_id.x * tile_size;
+  let local_col = local_idx % tile_size_k_vec;
+  let local_row = local_idx / tile_size_k_vec;
+
+  if (local_idx < tile_size) {
+    for (var lane = 0u; lane < tile_size_k_vec; lane++) {
+      gate_inter_results[local_idx][lane] = output_element_t(0);
+      up_inter_results[local_idx][lane] = output_element_t(0);
+    }
+  }
+  workgroupBarrier();
+
+#if single_scale_weights
+  let gate_scale_b = gate_scales_b.getByOffset(0);
+  let up_scale_b = up_scales_b.getByOffset(0);
+#endif
+
+  for (var kidx_v: u32 = 0u; kidx_v < uniforms.K32; kidx_v += tile_size_k_vec) {
+    if (local_idx < double_tile_size_k_vec) {
+      loadSHMA(batch, kidx_v * 2u, local_idx);
+    }
+    workgroupBarrier();
+
+    let own_a0 = tile_A[local_col * 2u];
+    let own_a1 = tile_A[local_col * 2u + 1u];
+    let own_scale_a = scale_A[local_col / 4u];
+    let k_offset = kidx_v + local_col;
+    let block_idx = k_offset * 32u / uniforms.block_size;
+
+    let b_global = b_global_base + local_row;
+    if (b_global < uniforms.N && k_offset < uniforms.K32) {
+#if !single_scale_weights
+      let gate_scale_b = gate_scales_b.getByOffset(b_global * uniforms.blocks_per_col + block_idx);
+      let up_scale_b = up_scales_b.getByOffset(b_global * uniforms.blocks_per_col + block_idx);
+#endif
+      let gate_b_value = gate_b.getByOffset(b_global * uniforms.K32 + k_offset);
+      let up_b_value = up_b.getByOffset(b_global * uniforms.K32 + k_offset);
+      let gate_b0 = DequantizedFrom4BitsTo8Bits(gate_b_value.xy, default_zero_point);
+      let gate_b1 = DequantizedFrom4BitsTo8Bits(gate_b_value.zw, default_zero_point);
+      let up_b0 = DequantizedFrom4BitsTo8Bits(up_b_value.xy, default_zero_point);
+      let up_b1 = DequantizedFrom4BitsTo8Bits(up_b_value.zw, default_zero_point);
+      let gate_scale = own_scale_a * gate_scale_b;
+      let up_scale = own_scale_a * up_scale_b;
+      gate_inter_results[local_row][local_col] += SDP8AI(own_a0, gate_b0, own_a1, gate_b1, gate_scale);
+      up_inter_results[local_row][local_col] += SDP8AI(own_a0, up_b0, own_a1, up_b1, up_scale);
+    }
+    workgroupBarrier();
+  }
+
+  if (local_idx < tile_size) {
+    var gate_output_value = output_element_t(0);
+    var up_output_value = output_element_t(0);
+    for (var lane = 0u; lane < tile_size_k_vec; lane++) {
+      gate_output_value += gate_inter_results[local_idx][lane];
+      up_output_value += up_inter_results[local_idx][lane];
+    }
+
+    let b_global = b_global_base + local_idx;
+    if (b_global < uniforms.N) {
+#if has_gate_bias
+      gate_output_value += gate_bias[b_global];
+#endif
+#if has_up_bias
+      up_output_value += up_bias[b_global];
+#endif
+      let one = output_element_t(1.0);
+      let silu_value = gate_output_value * (one / (one + exp(-gate_output_value)));
+      output.setByOffset(batch * uniforms.N + b_global, silu_value * up_output_value);
+    }
+  }
+}

onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc

@@ -18,10 +18,6 @@ namespace onnxruntime {
 namespace contrib {
 namespace webgpu {
 
-namespace {
-constexpr unsigned int kMinMForTileOptimization = 4;
-}  // namespace
-
 ONNX_OPERATOR_KERNEL_EX(
     MatMulNBits,
     kMSDomain,
@@ -226,29 +222,44 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales,
   uint32_t zero_blocks_per_col = (n_blocks_per_col + zp_elements_per_byte - 1) / zp_elements_per_byte * zp_elements_per_byte;
 
 #if !defined(__wasm__)
+  // apple|intel - Experimental dawn support for subgroup matrix matmul.
   int32_t subgroup_matrix_config_index = -1;
-  // Experimental dawn support for subgroup matrix matmul (vendor-agnostic).
-  if ((M >= kMinMForTileOptimization && !has_weight_idx_indirect) &&
-      CanApplySubgroupMatrixMatMulNBits(context, accuracy_level, block_size, batch_count, N, K, static_cast<uint32_t>(nbits), y->DataType() == DataTypeImpl::GetType<MLFloat16>(), subgroup_matrix_config_index)) {
+  if (WouldApplySubgroupMatrixMatMulNBitsInCurrentDispatch(a,
+                                                           K_op,
+                                                           N_op,
+                                                           block_size_op,
+                                                           accuracy_level,
+                                                           nbits,
+                                                           context,
+                                                           y,
+                                                           has_weight_idx_indirect,
+                                                           &subgroup_matrix_config_index,
+                                                           override_M)) {
     return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, static_cast<uint32_t>(nbits), zero_blocks_per_col, subgroup_matrix_config_index, context, y, weight_index, weight_index_indirect);
   }
 #endif
 
   // On FP32 only GPUs and Qualcomm GPUs, integer math is faster than FP32 therefore always use DP4A independent of length of M.
   // DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte values).
-  if (((M >= kMinMForTileOptimization && !has_weight_idx_indirect) || y->DataType() == DataTypeImpl::GetType<float>() || context.AdapterInfo().vendor == std::string_view{"qualcomm"}) &&
-      CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a)) {
+  if (WouldApplyDP4AMatMulNBitsInCurrentDispatch(a,
+                                                 K_op,
+                                                 N_op,
+                                                 block_size_op,
+                                                 accuracy_level,
+                                                 context,
+                                                 y,
+                                                 has_weight_idx_indirect)) {
     return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast<uint32_t>(nbits), context, y, weight_index, weight_index_indirect);
   }
 
   // WideTileProgram
   // This program is optimized for Block32 prefill using Tile16x128.
-  const bool use_wide_tile_program = !has_weight_idx_indirect &&
-                                     block_size == 32 &&
-                                     components_a == 4 &&
-                                     components_b == 4 &&
-                                     nbits != 2 &&
-                                     M >= kMinMForTileOptimization;
+  const bool use_wide_tile_program = WouldApplyWideTileMatMulNBitsInCurrentDispatch(a,
+                                                                                     K_op,
+                                                                                     N_op,
+                                                                                     block_size_op,
+                                                                                     nbits,
+                                                                                     has_weight_idx_indirect);
 
   if (use_wide_tile_program) {
     // Enforce output components to 1.
@@ -308,7 +319,8 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales,
 
   // Use tile_size_k_vec=32 by default for better K-dimension parallelism.
   // Intel devices use 16 as they have different subgroup/cache characteristics.
-  const uint32_t tile_size_k_vec = (context.AdapterInfo().vendor == std::string_view{"intel"}) ? 16u : 32u;
+  const uint32_t tile_size_k_vec =
+      (context.AdapterInfo().vendor == std::string_view{"intel"}) ? 16u : 32u;
 
   constexpr uint32_t workgroup_size = 128;
   constexpr uint32_t tile_size = 8;

onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_common.cc

@@ -2,9 +2,16 @@
 // Licensed under the MIT License.
 
 #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h"
+
 #include <sstream>
+
 #include "core/common/common.h"
+#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h"
+#include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h"
+#include "core/providers/cpu/math/matmul_helper.h"
 #include "core/providers/webgpu/webgpu_context.h"
+#include "core/providers/webgpu/webgpu_utils.h"
+#include "core/framework/tensor_shape.h"
 
 namespace onnxruntime {
 namespace contrib {
@@ -61,6 +68,109 @@ bool HasDP4ADeviceSupport(int context_id) {
          ctx.AdapterInfo().vendor != std::string_view{"apple"};
 }
 
+bool WouldApplySubgroupMatrixMatMulNBitsInCurrentDispatch(const Tensor* a,
+                                                          int64_t K_op,
+                                                          int64_t N_op,
+                                                          int64_t block_size_op,
+                                                          int64_t accuracy_level,
+                                                          int64_t nbits,
+                                                          onnxruntime::webgpu::ComputeContext& context,
+                                                          Tensor* y,
+                                                          bool has_weight_idx_indirect,
+                                                          int32_t* subgroup_matrix_config_index,
+                                                          uint32_t override_M) {
+  TensorShape b_shape({N_op, K_op});
+  MatMulComputeHelper helper;
+  if (!helper.Compute(a->Shape(), b_shape, false, true).IsOK()) {
+    return false;
+  }
+
+  const uint32_t batch_count = onnxruntime::narrow<uint32_t>(helper.OutputOffsets().size());
+  const uint32_t M = onnxruntime::narrow<uint32_t>(helper.M());
+  const uint32_t dispatch_M = override_M > 0 ? override_M : M;
+  const uint32_t N = onnxruntime::narrow<uint32_t>(helper.N());
+  const uint32_t K = onnxruntime::narrow<uint32_t>(helper.K());
+  const uint32_t block_size = onnxruntime::narrow<uint32_t>(block_size_op);
+
+#if !defined(__wasm__)
+  int32_t local_subgroup_matrix_config_index = -1;
+  if (dispatch_M != M) {
+    return false;
+  }
+
+  return (M >= kMinMForTileOptimization && !has_weight_idx_indirect) &&
+         CanApplySubgroupMatrixMatMulNBits(context,
+                                           accuracy_level,
+                                           block_size,
+                                           batch_count,
+                                           N,
+                                           K,
+                                           static_cast<uint32_t>(nbits),
+                                           y->DataType() == DataTypeImpl::GetType<MLFloat16>(),
+                                           subgroup_matrix_config_index != nullptr ? *subgroup_matrix_config_index : local_subgroup_matrix_config_index);
+#endif
+
+  return false;
+}
+
+bool WouldApplyDP4AMatMulNBitsInCurrentDispatch(const Tensor* a,
+                                                int64_t K_op,
+                                                int64_t N_op,
+                                                int64_t block_size_op,
+                                                int64_t accuracy_level,
+                                                onnxruntime::webgpu::ComputeContext& context,
+                                                Tensor* y,
+                                                bool has_weight_idx_indirect) {
+  TensorShape b_shape({N_op, K_op});
+  MatMulComputeHelper helper;
+  if (!helper.Compute(a->Shape(), b_shape, false, true).IsOK()) {
+    return false;
+  }
+
+  const uint32_t M = onnxruntime::narrow<uint32_t>(helper.M());
+  const uint32_t N = onnxruntime::narrow<uint32_t>(helper.N());
+  const uint32_t K = onnxruntime::narrow<uint32_t>(helper.K());
+  const uint32_t block_size = onnxruntime::narrow<uint32_t>(block_size_op);
+  const uint32_t components_a = GetMaxComponents(K);
+
+  return ((M >= kMinMForTileOptimization && !has_weight_idx_indirect) ||
+          y->DataType() == DataTypeImpl::GetType<float>() ||
+          context.AdapterInfo().vendor == std::string_view{"qualcomm"}) &&
+         CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a);
+}
+
+bool WouldApplyWideTileMatMulNBitsInCurrentDispatch(const Tensor* a,
+                                                    int64_t K_op,
+                                                    int64_t N_op,
+                                                    int64_t block_size_op,
+                                                    int64_t nbits,
+                                                    bool has_weight_idx_indirect) {
+  if (has_weight_idx_indirect) {
+    return false;
+  }
+
+  TensorShape b_shape({N_op, K_op});
+  MatMulComputeHelper helper;
+  if (!helper.Compute(a->Shape(), b_shape, false, true).IsOK()) {
+    return false;
+  }
+
+  const uint32_t M = onnxruntime::narrow<uint32_t>(helper.M());
+  const uint32_t K = onnxruntime::narrow<uint32_t>(helper.K());
+  const uint32_t block_size = onnxruntime::narrow<uint32_t>(block_size_op);
+  const uint32_t components_a = GetMaxComponents(K);
+  const uint32_t block_size_per_col = block_size;
+  const uint32_t blob_size = (block_size_per_col / 8) * static_cast<uint32_t>(nbits);
+  const uint32_t blob_size_in_words = blob_size / 4;
+  const uint32_t components_b = GetMaxComponents(blob_size_in_words);
+
+  return block_size == 32 &&
+         components_a == 4 &&
+         components_b == 4 &&
+         nbits != 2 &&
+         M >= kMinMForTileOptimization;
+}
+
 }  // namespace webgpu
 }  // namespace contrib
 }  // namespace onnxruntime

onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_common.h

@@ -6,10 +6,20 @@
 #include <string>
 #include <cstdint>
 
+namespace onnxruntime {
+class Tensor;
+
+namespace webgpu {
+class ComputeContext;
+}  // namespace webgpu
+}
+
 namespace onnxruntime {
 namespace contrib {
 namespace webgpu {
 
+inline constexpr uint32_t kMinMForTileOptimization = 4u;
+
 /**
  * Generates WebGPU shader code for reading zero points in quantized matrix multiplication
  *
@@ -26,6 +36,34 @@ std::string GenerateZeroPointReadingCode(uint32_t nbits, bool has_zero_points,
 /// \p context_id is the WebGpuContext slot (0 for the default context).
 bool HasDP4ADeviceSupport(int context_id = 0);
 
+bool WouldApplySubgroupMatrixMatMulNBitsInCurrentDispatch(const Tensor* a,
+                                                          int64_t K_op,
+                                                          int64_t N_op,
+                                                          int64_t block_size_op,
+                                                          int64_t accuracy_level,
+                                                          int64_t nbits,
+                                                          onnxruntime::webgpu::ComputeContext& context,
+                                                          Tensor* y,
+                                                          bool has_weight_idx_indirect = false,
+                                                          int32_t* subgroup_matrix_config_index = nullptr,
+                                                          uint32_t override_M = 0);
+
+bool WouldApplyDP4AMatMulNBitsInCurrentDispatch(const Tensor* a,
+                                                int64_t K_op,
+                                                int64_t N_op,
+                                                int64_t block_size_op,
+                                                int64_t accuracy_level,
+                                                onnxruntime::webgpu::ComputeContext& context,
+                                                Tensor* y,
+                                                bool has_weight_idx_indirect = false);
+
+bool WouldApplyWideTileMatMulNBitsInCurrentDispatch(const Tensor* a,
+                                                    int64_t K_op,
+                                                    int64_t N_op,
+                                                    int64_t block_size_op,
+                                                    int64_t nbits,
+                                                    bool has_weight_idx_indirect = false);
+
 }  // namespace webgpu
 }  // namespace contrib
 }  // namespace onnxruntime