diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index 19108b848c420..501a8d1d9dc09 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 @@ -36,10 +37,16 @@ // 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; @@ -47,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>; // 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_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) @@ -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); } } @@ -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); @@ -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_m) { 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_m) { 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_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 d9c5f5839a3bf..3d326a8374344 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 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(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..a7f32168df624 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,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 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); + 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