Skip to content

[WebGPU] Share MatMulNBits input preparation and tune narrow DP4A #33635

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Jiajia Qin (qjia7) wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[WebGPU] Share MatMulNBits input preparation and tune narrow DP4A #33635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0910366
Refactor WebGPU MatMulNBits input preparation for reuse
qjia7 Oct 27, 2026
abec758
Optimize narrow WebGPU MatMulNBits DP4A projections
qjia7 Oct 27, 2026
92fe6cd
Keep tiled DP4A unchanged in kernel extraction
qjia7 Oct 27, 2026
afa23fa
Ensure WebGPU DP4A regression tests execute on WebGPU
qjia7 Oct 27, 2026
573e54b
Keep narrow WebGPU DP4A regressions FP32 only
qjia7 Oct 27, 2026
a541ed1
Test FP16 overflow in narrow WebGPU DP4A accumulation
qjia7 Oct 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 13 additions & 26 deletions ...ization/dp4a_matmul_small_m.wgsl.template → ...tion/dp4a_matmul_k_parallel.wgsl.template
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,15 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

// K-parallel DP4A projection: each workgroup computes one input row and
// tile_size output columns. Its 128 threads process four columns concurrently,
// using 32 K lanes per column and 32 elements per lane in each K iteration.
// The group stages 1024 activations and their scales in workgroup memory,
// reusing them across output columns, including additional four-column groups.
// Each lane accumulates scaled integer dot products; the 32 lane partials are
// reduced through inter_results to produce each output value after processing K.
// This exposes parallelism along K when small M or N yields few output tiles.

#param acc_f32
#param tile_size
#param tile_size_k_vec
Expand All @@ -15,35 +24,13 @@

#use .getByOffset .setByOffset


#include "quantization/dp4a_matmul_common.wgsl.template"

// This algorithm works to compute dot product of k in parallel, by processing k at each step amongst tile_size_k_vec threads,
// and utilizing the remaining threads in the workgroup to process additional rows of b in parallel (such that the values in shared memory for A can be reused).
// For each load of k, the tile_size_k_vec threads also reload B tile_size/num_concurrent_b_rows times to compute partial dot products of other B rows
// in order to complete all tile_size b rows in this workgroup and also reusing the loaded in register values of a.

// 1. Each workgroup handles tile_size_k_vec * k_vectorization_in_b (32) columns and num_concurrent_b_rows of matrix B at a time,
// iterating over the columns to compute a partial dot product.
// 2. Uses vec4 vectorization where each K represents 32 elements of matrix B

// 1. Workgroup Responsibility:
// - Processes one row of matrix A
// - Handles tile_size rows of matrix B
//
// 2. Computation Process:
// - Reads [tile_size][tile_size_k_vec] block of B data at a time
// - Each thread within workgroup computes dot products of 32 A*B elements since each K represents 32 elements of matrix B
// - Stores intermediate results in shared memory (inter_results)
// - Iterates through columns accumulating results in inter_results
// - Performs final reduction sum in inter_results for output
// sub_tile_count is the number of concurrent b rows processed by the workgroup.

const double_tile_size_k_vec = 2 * tile_size_k_vec;

// Precision of the cross-tile accumulator, selected by the
// "enableMatmulFp32Accumulation" provider option. SDP8AI itself is an exact integer
// dot product; it is the accumulation over K that can overflow the f16 max (~65504).
// Narrow outputs use FP32 cross-tile accumulation to limit reduction error;
// other outputs follow "enableMatmulFp32Accumulation". SDP8AI itself is an exact
// integer dot product; accumulation over K can overflow the f16 max (~65504).
#if acc_f32
alias acc_element_t = f32;
#else
Expand All @@ -64,7 +51,7 @@ fn loadSHMA(batch: u32, a_global: u32, kidx_v: u32, col: u32)
}

tile_A[col] = a.getByOffset(batch*uniforms.M*uniforms.K16+a_global*uniforms.K16+k_offset);
if (col < scale_a_size_in_tile_a)
if (col < scale_a_size_in_tile_a && kidx_v/8 + col < uniforms.K16/8)
{
// kidx_v - covers 16 values of k in input_a
scale_A[col] = scales_a.getByOffset(batch*uniforms.M*(uniforms.K/128) + a_global*(uniforms.K/128) + kidx_v/8 + col);
Expand Down
90 changes: 56 additions & 34 deletions onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,12 @@
namespace webgpu {

Status DP4AMatMulQuantizeProgram::GenerateShaderCode(ShaderHelper& shader) const {
const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias |
ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform);
const auto& scales = shader.AddOutput("scales", ShaderUsage::UseUniform);
return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_quantize.wgsl.template",
WGSL_TEMPLATE_VARIABLE(a, a),
WGSL_TEMPLATE_VARIABLE(output, output),
WGSL_TEMPLATE_VARIABLE(scales, scales));
return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_quantize.wgsl.template", WGSL_TEMPLATE_VARIABLE(a, a),
WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales, scales));
}

Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const {
Expand Down Expand Up @@ -52,7 +51,7 @@
}

// scale_A components = 1, b components = 4, output components = 1
Status DP4AMatMulNBitsSmallMProgram::GenerateShaderCode(ShaderHelper& shader) const {
Status DP4AMatMulNBitsKParallelProgram::GenerateShaderCode(ShaderHelper& shader) const {
const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform);
const auto& scales_a = shader.AddInput("scales_a", ShaderUsage::UseUniform);
const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform);
Expand All @@ -72,7 +71,7 @@
const uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec_;
ORT_ENFORCE(tile_size_ % sub_tile_count == 0, "tile_size_ must be divisible by sub_tile_count");

return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_small_m.wgsl.template",
return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_k_parallel.wgsl.template",
WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_),
WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_),
WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_),
Expand Down Expand Up @@ -106,51 +105,74 @@
onnxruntime::webgpu::ComputeContext& context,
Tensor* y,
const uint32_t weight_index,
const Tensor* weight_index_indirect) {
const Tensor* weight_index_indirect, DP4AQuantizedInput* quantized_input) {
constexpr uint32_t kVec4Components = 4;
constexpr uint32_t kVec2Components = 2;
constexpr uint32_t kU32Components = 4;

constexpr uint32_t kBlockSizeA = 128;
DP4AMatMulQuantizeProgram quantize_program;
quantize_program.SetWorkgroupSize(64);
uint32_t tile_size = 64 * kVec4Components;
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({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({batch_count * M * K / kU32Components});
ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program));
DP4AQuantizedInput local_input;
auto& input = quantized_input ? *quantized_input : local_input;
if (!input.values) {
constexpr uint32_t kBlockSizeA = 128;
DP4AMatMulQuantizeProgram quantize_program;
quantize_program.SetWorkgroupSize(64);
uint32_t tile_size = 64 * kVec4Components;
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({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({batch_count * M * K / kU32Components});
ORT_RETURN_IF_ERROR(context.RunProgram(quantize_program));
input.values.emplace(std::move(a_quant));
input.scales.emplace(std::move(a_scale));

Check warning on line 131 in onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Add #include <utility> for move  [build/include_what_you_use] [4]

Raw Output:
onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc:131:  Add #include <utility> for move  [build/include_what_you_use] [4]
}
const auto& a_quant = *input.values;
const auto& a_scale = *input.scales;

const bool has_zero_points = zero_points != nullptr;
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;
const bool single_scale_weights = (block_size == K * N);
if (has_weight_idx_indirect || M < min_M_for_tile_optimization) {
uint32_t tile_size_k_vec = 32;
uint32_t tile_size_n = 4;
const bool small_n = N <= 32;
// Each K-parallel iteration covers 32 lanes * 32 elements = 1024 K values.
// For narrow N, conservatively require two iterations to amortize reduction
// overhead; shorter K benefits more from the tiled kernel's reuse across M.
const bool use_k_parallel = M < min_M_for_tile_optimization || (small_n && K >= 2048);
// Indirect expert weights can differ by row, so they also require the row-wise kernel.
if (use_k_parallel || has_weight_idx_indirect) {
const uint32_t tile_size_k_vec = 32;
// Keep more workgroups available when there are only a few input rows.
const uint32_t tile_size_n = small_n && M >= min_M_for_tile_optimization ? N : 4;

const uint32_t b_components = (nbits == 2 ? kVec2Components : kVec4Components);
const bool broadcast_a = dispatch_M > M;
const bool acc_f32 = context.EnableMatmulFp32Accumulation();
DP4AMatMulNBitsSmallMProgram mul_program{tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32};
// The narrow projection changes the K reduction order. Keep partial sums
// and the final reduction in FP32, then cast once when storing the output.
const bool acc_f32 = small_n || context.EnableMatmulFp32Accumulation();
Comment thread
qjia7 marked this conversation as resolved.
DP4AMatMulNBitsKParallelProgram mul_program{
tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx,
has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32};
uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n;
mul_program.SetWorkgroupSize(128);
mul_program.SetDispatchGroupSize(batch_count * dispatch_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({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col, weight_index, dispatch_M})
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({batch_count, M, N, K, K / 16, K / 32, block_size, num_N_tile, zero_blocks_per_col,
weight_index, dispatch_M})
.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, has_weight_idx_indirect, broadcast_a, acc_f32);
.CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx,
has_weight_idx_indirect, broadcast_a, acc_f32);
if (has_zero_points) {
mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4});
mul_program.AddInput(
{zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4});
}
if (has_bias) {
mul_program.AddInput({bias, ProgramTensorMetadataDependency::None});
Expand Down
53 changes: 28 additions & 25 deletions onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#pragma once

#include <limits>
#include <optional>

#include "core/providers/webgpu/program.h"
#include "core/providers/webgpu/webgpu_kernel.h"
Expand All @@ -14,6 +15,12 @@ namespace webgpu {

using namespace onnxruntime::webgpu;

// Invocation-local workspace, reusable only by projections of the same input tensor.
struct DP4AQuantizedInput {
std::optional<Tensor> values;
std::optional<Tensor> scales;
};

class DP4AMatMulQuantizeProgram final : public Program<DP4AMatMulQuantizeProgram> {
public:
DP4AMatMulQuantizeProgram() : Program{"DP4AMatMulQuantize"} {}
Expand Down Expand Up @@ -59,23 +66,23 @@ class DP4AMatMulNBitsProgram final : public Program<DP4AMatMulNBitsProgram> {
bool acc_f32_;
};

class DP4AMatMulNBitsSmallMProgram final : public Program<DP4AMatMulNBitsSmallMProgram> {
class DP4AMatMulNBitsKParallelProgram final : public Program<DP4AMatMulNBitsKParallelProgram> {
public:
DP4AMatMulNBitsSmallMProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits,
bool has_zero_points, bool has_bias,
bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights,
bool broadcast_a_row = false,
bool acc_f32 = false) : Program{"DP4AMatMulNBitsSmallMProgram"},
tile_size_k_vec_(tile_size_k_vec),
tile_size_(tile_size),
nbits_(nbits),
has_bias_(has_bias),
has_zero_points_(has_zero_points),
has_weight_idx_(has_weight_idx),
has_weight_idx_indirect_(has_weight_idx_indirect),
single_scale_weights_(single_scale_weights),
broadcast_a_row_(broadcast_a_row),
acc_f32_(acc_f32) {}
DP4AMatMulNBitsKParallelProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits,
bool has_zero_points, bool has_bias,
bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights,
bool broadcast_a_row = false,
bool acc_f32 = false) : Program{"DP4AMatMulNBitsKParallelProgram"},
tile_size_k_vec_(tile_size_k_vec),
tile_size_(tile_size),
nbits_(nbits),
has_bias_(has_bias),
has_zero_points_(has_zero_points),
has_weight_idx_(has_weight_idx),
has_weight_idx_indirect_(has_weight_idx_indirect),
single_scale_weights_(single_scale_weights),
broadcast_a_row_(broadcast_a_row),
acc_f32_(acc_f32) {}
Status GenerateShaderCode(ShaderHelper& sh) const override;
WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
{"batch_count", ProgramUniformVariableDataType::Uint32},
Expand Down Expand Up @@ -117,22 +124,18 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
onnxruntime::webgpu::ComputeContext& context,
Tensor* y,
const uint32_t weight_index,
const Tensor* weight_index_indirect = nullptr);
const Tensor* weight_index_indirect = nullptr,
DP4AQuantizedInput* quantized_input = nullptr);

// The optional M / has_weight_idx_indirect / y arguments fold the original
// dispatch-precondition (DP4A is preferred when M is large enough, or
// unconditionally on FP32 outputs and Qualcomm GPUs) into the feasibility check
// so callers don't need a separate wrapper. Defaults make the precondition
// trivially satisfied for callers that only want the feasibility check.
bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context,
uint64_t accuracy_level,
uint32_t block_size,
uint32_t N,
uint32_t K,
uint32_t components_k,
bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level,
uint32_t block_size, uint32_t N, uint32_t K, uint32_t components_k,
uint32_t M = std::numeric_limits<uint32_t>::max(),
bool has_weight_idx_indirect = false,
const Tensor* y = nullptr);
bool has_weight_idx_indirect = false, const Tensor* y = nullptr);

} // namespace webgpu
} // namespace contrib
Expand Down
Loading
Loading