Skip to content

[WebGPU] Optimize DP4A for large Qualcomm subgroups #33560

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
Yang Gu (gyagp) wants to merge 2 commits into
base: main
Choose a base branch
from
+202 −28
Open

[WebGPU] Optimize DP4A for large Qualcomm subgroups #33560

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e75fd4c
Optimize WebGPU DP4A for large Qualcomm subgroups
gyagp Oct 23, 2026
4a43703
Reuse DP4A activation scales across K tiles
gyagp 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
49 changes: 30 additions & 19 deletions onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
#param has_bias
#param has_zero_points
#param is_qualcomm
#param qc_wave128
#param has_weight_idx
#param has_weight_idx_indirect

Expand Down Expand Up @@ -36,24 +37,30 @@
// 512B of registry from each lane.
//
// The medium size matmul is implemented using dot4I8Packed, so the inputs for
// this shader require A to be int8 quantized with block size 64. B is regular
// this shader require A to be int8 quantized with block size 128. B is regular
// matmulnbits input with block size 32.

const tile_size = 64;
#if qc_wave128
const tile_m = 128;
const tile_n = 32;
#else
const tile_m = 64;
const tile_n = 64;
#endif
const subtile_size = 16;
const tile_size_k = 32;
const vec_factor = 4;
const u32_factor = 4;
const tile_size_k_vec = 2;

// Shared memory
var<workgroup> tile_A : array<array<vec4<u32>, tile_size>, tile_size_k_vec>; // 64 x 32
var<workgroup> scale_A : array<output_element_t, tile_size>; // 64 x 1
var<workgroup> tile_B : array<array<vec4<u32>, tile_size>, tile_size_k_vec>; // 64 x 32
var<workgroup> scale_B : array<output_element_t, tile_size>; // 64 x 1
var<workgroup> tile_A : array<array<vec4<u32>, tile_m>, tile_size_k_vec>; // tile_m x 32
var<workgroup> scale_A : array<output_element_t, tile_m>; // tile_m x 1
var<workgroup> tile_B : array<array<vec4<u32>, tile_n>, tile_size_k_vec>; // tile_n x 32
var<workgroup> scale_B : array<output_element_t, tile_n>; // tile_n x 1

#if has_zero_points
var<workgroup> zeroes : array<i32, tile_size>;
var<workgroup> zeroes : array<i32, tile_n>;
#endif

fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32)
Expand All @@ -64,9 +71,9 @@ fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32)
return;
}
tile_A[col][row] = a.getByOffset(batch*uniforms.M*uniforms.K16+a_global*uniforms.K16+kidx_v+col);
if (col == 0)
// A 128-value quantization block shares its scale across four 32-value K tiles.
if (col == 0 && kidx_v % 8u == 0u)
{
// kidx_v - covers 16 values of k
scale_A[row] = scales_a.getByOffset(batch*uniforms.M*(uniforms.K/128) + a_global*(uniforms.K/128) + kidx_v/8);
}
}
Expand Down Expand Up @@ -197,8 +204,8 @@ $MAIN {
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 a_global_base = u32((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * tile_m;
let b_global_base = (workgroup_idx % uniforms.num_N_tile) * tile_n;
let load_AorB = u32(local_idx/128);
let load_row = u32((local_idx%128)/2);
let load_col = u32(local_idx%2);
Expand All @@ -214,12 +221,9 @@ $MAIN {
var a_idx = u32(local_idx % subtile_size);

#if is_qualcomm
// subtile_idx is always 0
// subtile_idy is one of {0,1,2,3}
// The subtile is now rectangular 64x16 for qualcomm case and we have 4 subtiles, this way we don't need to
// increase the number of lane_output each thread needs to track. That is if we want to use a subtile that is 64x64
// we would need var lane_outputs: array<output_element_t, 64>;
if (sg_size == 64) {
// Qualcomm uses four 64x16 subtiles or two 128x16 subtiles. Each lane
// owns one activation row and accumulates 16 output columns.
if (sg_size == tile_m) {
subtile_id = u32(local_idx / sg_size);
subtile_idx = u32(subtile_id / 4);
subtile_idy = u32(subtile_id % 4);
Expand All @@ -240,6 +244,12 @@ $MAIN {
for (var kidx_v:u32 = 0; kidx_v < uniforms.K16; kidx_v+=tile_size_k_vec)
{
// Load Phase: Populate shared memory for the workgroup.
#if qc_wave128
loadSHMA(batch, a_global_base, kidx_v, local_idx / 2u, local_idx % 2u);
Comment thread
gyagp marked this conversation as resolved.
if (local_idx < 64u) {
loadSHMB(b_global_base, kidx_v, local_idx / 2u, local_idx % 2u);
}
#else
if (load_AorB == 0)
{
loadSHMA(batch, a_global_base, kidx_v, load_row, load_col);
Expand All @@ -248,6 +258,7 @@ $MAIN {
{
loadSHMB(b_global_base, kidx_v, load_row, load_col);
}
#endif
workgroupBarrier();

// Compute phase: Perform matmul for this subtile 16 x 32 x 16.
Expand All @@ -258,7 +269,7 @@ $MAIN {

#if has_zero_points && n_bits == 8
#if is_qualcomm
if (sg_size == 64)
if (sg_size == tile_m)
{
var own_b0: vec4<u32>;
var own_b1: vec4<u32>;
Expand Down Expand Up @@ -340,7 +351,7 @@ $MAIN {
#endif
#else
#if is_qualcomm
if (sg_size == 64)
if (sg_size == tile_m)
{
var own_b0: vec4<u32>;
var own_b1: vec4<u32>;
Expand Down
17 changes: 14 additions & 3 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 @@ -44,6 +44,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const {
WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_),
WGSL_TEMPLATE_PARAMETER(n_bits, nbits_),
WGSL_TEMPLATE_PARAMETER(output_type_i32, true),
WGSL_TEMPLATE_PARAMETER(qc_wave128, SubgroupSize() == 128),
WGSL_TEMPLATE_VARIABLE(a, a),
WGSL_TEMPLATE_VARIABLE(b, b),
WGSL_TEMPLATE_VARIABLE(output, output),
Expand Down Expand Up @@ -161,14 +162,24 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
return context.RunProgram(mul_program);
}

constexpr uint32_t kTileSize = 64;
// Two 128-lane subgroups cover a 128x32 tile, with one activation row per lane.
const bool qc_wave128 = context.AdapterInfo().vendor == std::string_view{"qualcomm"} &&
context.HasFeature(wgpu::FeatureName::SubgroupSizeControl) &&
context.AdapterInfo().subgroupMinSize <= 128 &&
context.AdapterInfo().subgroupMaxSize >= 128 &&
nbits == 4 && block_size == 32 && M >= 128;
const uint32_t tile_m = qc_wave128 ? 128 : 64;
const uint32_t tile_n = qc_wave128 ? 32 : 64;
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;
uint32_t num_M_tile = (M + tile_m - 1) / tile_m;
uint32_t num_N_tile = (N + tile_n - 1) / tile_n;
bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"};
const bool acc_f32 = context.EnableMatmulFp32Accumulation();
DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32};
mul_program.SetWorkgroupSize(256);
if (qc_wave128) {
mul_program.SetSubgroupSize(128);
}
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},
Expand Down
19 changes: 19 additions & 0 deletions onnxruntime/contrib_ops/webgpu/quantization/dp4a_quantize.wgsl.template
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,25 @@ $MAIN {
// 127 is the max value of signed int8 [-127,127] used by pack4x8snorm for 1.0f.
scales.setByOffset(workgroup_idx * 2 + local_idx / 32, scale/127);
}
} else if (sg_size > 32) {
let local_a = readInput(global_idx);
// Reduce components before shuffling; each 32-lane partition is one 128-value block.
let abs_a = abs(local_a);
var max_val = max(max(abs_a.x, abs_a.y), max(abs_a.z, abs_a.w));
max_val = max(max_val, subgroupShuffleXor(max_val, 16u));
max_val = max(max_val, subgroupShuffleXor(max_val, 8u));
max_val = max(max_val, subgroupShuffleXor(max_val, 4u));
max_val = max(max_val, subgroupShuffleXor(max_val, 2u));
max_val = max(max_val, subgroupShuffleXor(max_val, 1u));
if (global_idx >= uniforms.output_size) {
return;
}
let scale = max_val;
let norm_a = local_a / scale;
output.setByOffset(global_idx, pack4x8snorm(vec4<f32>(norm_a)));
if (local_idx % 32 == 0) {
scales.setByOffset(workgroup_idx * 2 + local_idx / 32, scale / 127);
}
} else if (sg_size == 16) {
let local_a = readInput(global_idx);
let sub_max_value = subgroupMax(abs(local_a));
Expand Down
145 changes: 139 additions & 6 deletions onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS)

#include <cstdint>
#include <string_view>
#include <utility>

#include "gtest/gtest.h"
Expand All @@ -21,7 +22,7 @@ namespace test {
namespace {

void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K,
bool has_zero_points, bool has_bias) {
bool has_zero_points, bool has_bias, int64_t batch_count = 1) {
SCOPED_TRACE(testing::Message() << "bits=" << bits << ", M=" << M << ", N=" << N << ", K=" << K
<< ", zero_points=" << has_zero_points << ", bias=" << has_bias);

Expand All @@ -40,7 +41,7 @@ void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K,
const int32_t default_zero_point = 1 << (bits - 1);

RandomValueGenerator random{1234};
const auto a = FloatsToMLFloat16s(random.Gaussian<float>(AsSpan({M, K}), 0.0f, 0.25f));
const auto a = FloatsToMLFloat16s(random.Gaussian<float>(AsSpan({batch_count * M, K}), 0.0f, 0.25f));
const auto weights = random.Uniform<int32_t>(AsSpan({N, K}), 0, (1 << bits) - 1);
const auto bias = FloatsToMLFloat16s(random.Uniform<float>(AsSpan({N}), -0.5f, 0.5f));
InlinedVector<uint8_t> packed_weights(N * blocks * blob_size, 0);
Expand Down Expand Up @@ -69,8 +70,8 @@ void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K,
}
}

InlinedVector<MLFloat16> expected(M * N);
for (int64_t m = 0; m < M; ++m) {
InlinedVector<MLFloat16> expected(batch_count * M * N);
for (int64_t m = 0; m < batch_count * M; ++m) {
for (int64_t n = 0; n < N; ++n) {
float sum = 0.0f;
for (int64_t k = 0; k < K; ++k) {
Expand All @@ -86,7 +87,11 @@ void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K,
test.AddAttribute<int64_t>("bits", bits);
test.AddAttribute<int64_t>("block_size", block_size);
test.AddAttribute<int64_t>("accuracy_level", 4);
test.AddInput<MLFloat16>("A", {M, K}, a);
const TensorShapeVector a_shape = batch_count == 1 ? TensorShapeVector{M, K}
: TensorShapeVector{batch_count, M, K};
const TensorShapeVector y_shape = batch_count == 1 ? TensorShapeVector{M, N}
: TensorShapeVector{batch_count, M, N};
test.AddInput<MLFloat16>("A", a_shape, a);
test.AddInput<uint8_t>("B", {N, blocks, blob_size}, packed_weights.data(), packed_weights.size(), true);
test.AddInput<MLFloat16>("scales", {N, blocks}, scales.data(), scales.size(), true);
if (has_zero_points) {
Expand All @@ -98,7 +103,7 @@ void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K,
if (has_bias) {
test.AddInput<MLFloat16>("bias", {N}, bias, true);
}
test.AddOutput<MLFloat16>("Y", {M, N}, expected.data(), expected.size());
test.AddOutput<MLFloat16>("Y", y_shape, expected.data(), expected.size());
test.SetOutputAbsErr("Y", 0.1f);
test.SetOutputRelErr("Y", 0.02f);

Expand Down Expand Up @@ -146,6 +151,134 @@ TEST(WebGpuMatMulNBitsTest, Float16_8Bits_SubgroupMatrixTiles) {
RunMatMulNBitsTileTests(8);
}

TEST(WebGpuMatMulNBitsTest, DP4AQualcommWave128Tiles) {
auto provider = DefaultWebGpuExecutionProvider();
ASSERT_NE(provider, nullptr);
const auto& context = webgpu::WebGpuContextFactory::GetContext(provider->GetDeviceId());
if (context.AdapterInfo().vendor != std::string_view{"qualcomm"} ||
!context.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl) ||
context.AdapterInfo().subgroupMinSize > 128 || context.AdapterInfo().subgroupMaxSize < 128) {
GTEST_SKIP() << "Requires Qualcomm 128-lane subgroup size control";
}
for (bool has_zero_points : {false, true}) {
for (bool has_bias : {false, true}) {
// Exercise both column subgroups and partially filled M and N tiles.
RunMatMulNBitsTileTest(4, 128, 32, 128, has_zero_points, has_bias);
RunMatMulNBitsTileTest(4, 129, 48, 256, has_zero_points, has_bias, 2);
RunMatMulNBitsTileTest(4, 257, 16, 128, has_zero_points, has_bias);
}
}
}

TEST(WebGpuMatMulNBitsTest, DP4AActivationScaleReuse) {
auto provider = DefaultWebGpuExecutionProvider();
ASSERT_NE(provider, nullptr);
const auto& context = webgpu::WebGpuContextFactory::GetContext(provider->GetDeviceId());
if (!context.DeviceHasFeature(wgpu::FeatureName::Subgroups) ||
context.AdapterInfo().vendor == std::string_view{"apple"}) {
GTEST_SKIP() << "Requires the DP4A path";
}

constexpr int64_t batch_count = 2;
constexpr int64_t N = 48;
constexpr int64_t K = 384;
constexpr int64_t block_size = 32;
for (int64_t M : {65, 129}) {
for (int64_t bits : {4, 8}) {
SCOPED_TRACE(testing::Message() << "M=" << M << ", bits=" << bits);
InlinedVector<float> a(batch_count * M * K);
InlinedVector<float> expected(batch_count * M * N);
for (int64_t batch = 0; batch < batch_count; ++batch) {
for (int64_t row = 0; row < M; ++row) {
for (int64_t k = 0; k < K; ++k) {
// Powers of two give distinct scales across blocks, rows, and batches.
a[(batch * M + row) * K + k] = static_cast<float>(1 << ((batch * 3 + row % 3 + (k / 128) * 2) % 7));
}
for (int64_t n = 0; n < N; ++n) {
const int64_t k = (n % (K / block_size)) * block_size + block_size - 1;
expected[(batch * M + row) * N + n] = a[(batch * M + row) * K + k];
}
}
}

// Each column selects the last activation of one K tile, including all
// four tiles sharing a scale and the transitions to the next scale.
const int64_t values_per_byte = 8 / bits;
InlinedVector<uint8_t> weights(N * K / values_per_byte, bits == 4 ? 0x88 : 0x80);
for (int64_t n = 0; n < N; ++n) {
const int64_t k = (n % (K / block_size)) * block_size + block_size - 1;
weights[(n * K + k) / values_per_byte] |= static_cast<uint8_t>(1 << ((k % values_per_byte) * bits));
}
InlinedVector<float> scales(N * K / block_size, 1.0f);
OpTester test("MatMulNBits", 1, kMSDomain);
test.AddAttribute<int64_t>("K", K);
test.AddAttribute<int64_t>("N", N);
test.AddAttribute<int64_t>("bits", bits);
test.AddAttribute<int64_t>("block_size", block_size);
test.AddAttribute<int64_t>("accuracy_level", 4);
test.AddInput<float>("A", {batch_count, M, K}, a.data(), a.size());
test.AddInput<uint8_t>("B", {N, K / block_size, block_size / values_per_byte}, weights.data(), weights.size(), true);
test.AddInput<float>("scales", {N, K / block_size}, scales.data(), scales.size(), true);
test.AddOutput<float>("Y", {batch_count, M, N}, expected.data(), expected.size());
test.SetOutputAbsErr("Y", 0.0001f);
test.SetOutputRelErr("Y", 0.00001f);
SessionOptions options;
ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"));
test.Config(options).ConfigEp(DefaultWebGpuExecutionProvider()).RunWithConfig();
}
}
}

TEST(WebGpuMatMulNBitsTest, DP4AQuantizationLargeSubgroups) {
auto provider = DefaultWebGpuExecutionProvider();
ASSERT_NE(provider, nullptr);
const auto& context = webgpu::WebGpuContextFactory::GetContext(provider->GetDeviceId());
if (!context.DeviceHasFeature(wgpu::FeatureName::ShaderF16) ||
!context.DeviceHasFeature(wgpu::FeatureName::Subgroups) ||
context.AdapterInfo().subgroupMinSize <= 32) {
GTEST_SKIP() << "Requires subgroups wider than a 128-value quantization block";
}

constexpr int64_t M = 5;
constexpr int64_t N = 16;
for (int64_t K : {128, 256}) {
SCOPED_TRACE(K);
InlinedVector<MLFloat16> a(M * K, MLFloat16(0.0f));
InlinedVector<MLFloat16> expected(M * N);
for (int64_t m = 0; m < M; ++m) {
for (int64_t block = 0; block < K / 128; ++block) {
const bool large = (m * (K / 128) + block) % 2 != 0;
a[m * K + block * 128] = MLFloat16(large ? 64.0f : 0.5f);
a[m * K + block * 128 + 127] = MLFloat16(large ? 127.0f : 1.0f);
}
for (int64_t n = 0; n < N; ++n) {
expected[m * N + n] = a[m * K];
}
}
// Select the first activation. A reduction across the whole subgroup would
// mix the small and large blocks and quantize 0.5 to 1 instead of about 0.504.
InlinedVector<uint8_t> weights(N * K / 2, 0x88);
for (int64_t n = 0; n < N; ++n) {
weights[n * K / 2] = 0x89;
}
InlinedVector<MLFloat16> scales(N * K / 32, MLFloat16(1.0f));
OpTester test("MatMulNBits", 1, kMSDomain);
test.AddAttribute<int64_t>("K", K);
test.AddAttribute<int64_t>("N", N);
test.AddAttribute<int64_t>("bits", 4);
test.AddAttribute<int64_t>("block_size", 32);
test.AddAttribute<int64_t>("accuracy_level", 4);
test.AddInput<MLFloat16>("A", {M, K}, a.data(), a.size());
test.AddInput<uint8_t>("B", {N, K / 32, 16}, weights.data(), weights.size(), true);
test.AddInput<MLFloat16>("scales", {N, K / 32}, scales.data(), scales.size(), true);
test.AddOutput<MLFloat16>("Y", {M, N}, expected.data(), expected.size());
test.SetOutputAbsErr("Y", 0.01f);
SessionOptions options;
ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"));
test.Config(options).ConfigEp(DefaultWebGpuExecutionProvider()).RunWithConfig();
}
}

} // namespace test
} // namespace onnxruntime

Expand Down
Loading