From e75fd4c33213c94bd99c6699c064200ede39e18a Mon Sep 17 00:00:00 2001 From: Yang Gu Date: Fri, 23 Oct 2026 13:37:33 +0800 Subject: [PATCH 1/2] Optimize WebGPU DP4A for large Qualcomm subgroups Reduce each lane's four activation components before XOR shuffles on wide subgroups, keeping each 128-value quantization block independent. Use a 128x32 matrix tile with two explicitly requested 128-lane subgroups on compatible Qualcomm GPUs for 4-bit weights, block size 32, and M >= 128. A fresh source-built ORT GenAI native benchmark on Adreno X1-85 measures Aion batch-1, 1024-token prefill at 3.76782 s before and 3.63298 s after: 3.58% lower latency and 3.71% higher throughput (271.78 -> 281.86 TPS). This uses KV capacity 8192, identical build settings, D3D12, generator reuse, graph capture disabled, and one generated token. Reordered runs collect 20 measured prefills per variant with two warmups per process; loaded ORT DLL paths and hashes are verified. At capacity 1025, ten samples per variant measure 4.08% lower latency and 4.25% more TPS. Earlier GPU profiling measured quantization at 244.4 -> 82.9 ms (2.95x faster) and matrix execution at 5.586 -> 5.342 s (4.4% less time). Those profiles are from a separate session, not a breakdown of the GenAI run. Performance results are limited to this device and workload. Add regression coverage for independent quantization blocks, partial workgroups, M/N tile tails, batching, bias, and zero points. Validation: Windows ARM64 Release builds, four focused provider tests, C++ lint, and whitespace checks passed. Fresh GenAI prefill logits and sampled KV values match the baseline bit-for-bit. Earlier Aion prefill and one-token decode checks also matched at 128 and 1024 tokens. --- .../quantization/dp4a_matmul.wgsl.template | 37 +++++--- .../webgpu/quantization/dp4a_matmul_nbits.cc | 15 +++- .../quantization/dp4a_quantize.wgsl.template | 19 ++++ .../webgpu/quantization/matmul_nbits_test.cc | 86 +++++++++++++++++-- 4 files changed, 136 insertions(+), 21 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index 19108b848c420..fbb72efd8ed84 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -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 @@ -39,7 +40,13 @@ // this shader require A to be int8 quantized with block size 64. B is regular // matmulnbits input with block size 32. +#if qc_wave128 +const tile_size = 128; +const tile_n = 32; +#else const tile_size = 64; +const tile_n = 64; +#endif const subtile_size = 16; const tile_size_k = 32; const vec_factor = 4; @@ -47,10 +54,10 @@ const u32_factor = 4; const tile_size_k_vec = 2; // Shared memory -var tile_A : array, tile_size>, tile_size_k_vec>; // 64 x 32 -var scale_A : array; // 64 x 1 -var tile_B : array, tile_size>, tile_size_k_vec>; // 64 x 32 -var scale_B : array; // 64 x 1 +var tile_A : array, tile_size>, tile_size_k_vec>; // tile_size x 32 +var scale_A : array; // tile_size x 1 +var tile_B : array, tile_n>, tile_size_k_vec>; // tile_n x 32 +var scale_B : array; // tile_n x 1 #if has_zero_points var zeroes : array; @@ -198,7 +205,7 @@ $MAIN { 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 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); @@ -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; - 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_size) { subtile_id = u32(local_idx / sg_size); subtile_idx = u32(subtile_id / 4); subtile_idy = u32(subtile_id % 4); @@ -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); + 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); @@ -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. @@ -258,7 +269,7 @@ $MAIN { #if has_zero_points && n_bits == 8 #if is_qualcomm - if (sg_size == 64) + if (sg_size == tile_size) { var own_b0: vec4; var own_b1: vec4; @@ -340,7 +351,7 @@ $MAIN { #endif #else #if is_qualcomm - if (sg_size == 64) + if (sg_size == tile_size) { var own_b0: vec4; var own_b1: vec4; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index d9c5f5839a3bf..954a715aa6139 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -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), @@ -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 kTileSize = 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_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(kVec4Components)}, {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_quantize.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_quantize.wgsl.template index 09cbd78fd6ccd..2398cfb5c8056 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_quantize.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_quantize.wgsl.template @@ -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(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)); diff --git a/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc b/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc index 99050394fd8a5..4abc5a58ab7d4 100644 --- a/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc +++ b/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc @@ -4,6 +4,7 @@ #if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS) #include +#include #include #include "gtest/gtest.h" @@ -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); @@ -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(AsSpan({M, K}), 0.0f, 0.25f)); + const auto a = FloatsToMLFloat16s(random.Gaussian(AsSpan({batch_count * M, K}), 0.0f, 0.25f)); const auto weights = random.Uniform(AsSpan({N, K}), 0, (1 << bits) - 1); const auto bias = FloatsToMLFloat16s(random.Uniform(AsSpan({N}), -0.5f, 0.5f)); InlinedVector packed_weights(N * blocks * blob_size, 0); @@ -69,8 +70,8 @@ void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K, } } - InlinedVector expected(M * N); - for (int64_t m = 0; m < M; ++m) { + InlinedVector 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) { @@ -86,7 +87,11 @@ void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K, test.AddAttribute("bits", bits); test.AddAttribute("block_size", block_size); test.AddAttribute("accuracy_level", 4); - test.AddInput("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("A", a_shape, a); test.AddInput("B", {N, blocks, blob_size}, packed_weights.data(), packed_weights.size(), true); test.AddInput("scales", {N, blocks}, scales.data(), scales.size(), true); if (has_zero_points) { @@ -98,7 +103,7 @@ void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K, if (has_bias) { test.AddInput("bias", {N}, bias, true); } - test.AddOutput("Y", {M, N}, expected.data(), expected.size()); + test.AddOutput("Y", y_shape, expected.data(), expected.size()); test.SetOutputAbsErr("Y", 0.1f); test.SetOutputRelErr("Y", 0.02f); @@ -146,6 +151,75 @@ 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, 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 a(M * K, MLFloat16(0.0f)); + InlinedVector 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 weights(N * K / 2, 0x88); + for (int64_t n = 0; n < N; ++n) { + weights[n * K / 2] = 0x89; + } + InlinedVector scales(N * K / 32, MLFloat16(1.0f)); + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("bits", 4); + test.AddAttribute("block_size", 32); + test.AddAttribute("accuracy_level", 4); + test.AddInput("A", {M, K}, a.data(), a.size()); + test.AddInput("B", {N, K / 32, 16}, weights.data(), weights.size(), true); + test.AddInput("scales", {N, K / 32}, scales.data(), scales.size(), true); + test.AddOutput("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 From 4a437032623f3e2757fba465c2aecd4eac0a71af Mon Sep 17 00:00:00 2001 From: Yang Gu Date: Tue, 27 Oct 2026 23:02:02 +0800 Subject: [PATCH 2/2] Reuse DP4A activation scales across K tiles Rename the row tile dimension to tile_m in host dispatch and WGSL, matching tile_n and clarifying rectangular tiles. Load each activation scale only at the start of its 128-value block. The following three 32-value K tiles reuse the workgroup scale; batch and row mapping stay fixed and both workgroup barriers are unchanged. This reduces scale fetches from 512 to 128 per full 128-row tile and 128 K values. No additional end-to-end speedup is claimed. Add a deterministic regression selecting every K tile across three activation blocks with different scales per block, row, and batch. Cover regular and Qualcomm wave128 tile shapes, M/N tails, and 4/8-bit weights. Use tile_n for the B zero-point array and correct the comment on activation quantization block size. Validation: Windows ARM64 Release build, six focused provider tests, C++ lintrunner, and git diff --check passed on Adreno X1-85. --- .../quantization/dp4a_matmul.wgsl.template | 24 ++++---- .../webgpu/quantization/dp4a_matmul_nbits.cc | 4 +- .../webgpu/quantization/matmul_nbits_test.cc | 59 +++++++++++++++++++ 3 files changed, 73 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index fbb72efd8ed84..501a8d1d9dc09 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -37,14 +37,14 @@ // 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. #if qc_wave128 -const tile_size = 128; +const tile_m = 128; const tile_n = 32; #else -const tile_size = 64; +const tile_m = 64; const tile_n = 64; #endif const subtile_size = 16; @@ -54,13 +54,13 @@ const u32_factor = 4; const tile_size_k_vec = 2; // Shared memory -var tile_A : array, tile_size>, tile_size_k_vec>; // tile_size x 32 -var scale_A : array; // tile_size x 1 +var tile_A : array, tile_m>, tile_size_k_vec>; // tile_m x 32 +var scale_A : array; // tile_m x 1 var tile_B : array, tile_n>, tile_size_k_vec>; // tile_n x 32 var scale_B : array; // tile_n x 1 #if has_zero_points - var zeroes : array; + var zeroes : array; #endif fn loadSHMA(batch:u32, a_global_base:u32, kidx_v:u32, row: u32, col: u32) @@ -71,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); } } @@ -204,7 +204,7 @@ $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 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); @@ -223,7 +223,7 @@ $MAIN { #if is_qualcomm // Qualcomm uses four 64x16 subtiles or two 128x16 subtiles. Each lane // owns one activation row and accumulates 16 output columns. - if (sg_size == tile_size) { + if (sg_size == tile_m) { subtile_id = u32(local_idx / sg_size); subtile_idx = u32(subtile_id / 4); subtile_idy = u32(subtile_id % 4); @@ -269,7 +269,7 @@ $MAIN { #if has_zero_points && n_bits == 8 #if is_qualcomm - if (sg_size == tile_size) + if (sg_size == tile_m) { var own_b0: vec4; var own_b1: vec4; @@ -351,7 +351,7 @@ $MAIN { #endif #else #if is_qualcomm - if (sg_size == tile_size) + if (sg_size == tile_m) { var own_b0: vec4; var own_b1: vec4; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 954a715aa6139..3d326a8374344 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -168,10 +168,10 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor context.AdapterInfo().subgroupMinSize <= 128 && context.AdapterInfo().subgroupMaxSize >= 128 && nbits == 4 && block_size == 32 && M >= 128; - const uint32_t kTileSize = qc_wave128 ? 128 : 64; + 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_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(); diff --git a/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc b/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc index 4abc5a58ab7d4..a7f32168df624 100644 --- a/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc +++ b/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc @@ -170,6 +170,65 @@ TEST(WebGpuMatMulNBitsTest, DP4AQualcommWave128Tiles) { } } +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 a(batch_count * M * K); + InlinedVector 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(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 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(1 << ((k % values_per_byte) * bits)); + } + InlinedVector scales(N * K / block_size, 1.0f); + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("bits", bits); + test.AddAttribute("block_size", block_size); + test.AddAttribute("accuracy_level", 4); + test.AddInput("A", {batch_count, M, K}, a.data(), a.size()); + test.AddInput("B", {N, K / block_size, block_size / values_per_byte}, weights.data(), weights.size(), true); + test.AddInput("scales", {N, K / block_size}, scales.data(), scales.size(), true); + test.AddOutput("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);