diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index 90ead722ace..b41b02bf6e4 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -48,6 +48,16 @@ inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) { return desired; } +// Clamp to device limit, then floor to pow2 (reduction kernels halve stride). +inline uint32_t clamp_workgroup_size_pow2(WGPUDevice device, uint32_t desired) { + uint32_t v = clamp_workgroup_size(device, desired); + uint32_t p = 1u; + while (p <= (v >> 1u)) { + p <<= 1u; + } + return p; +} + struct WgCount { uint32_t x; uint32_t y; diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 3dd83b7ab3e..a931e22b08d 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -244,6 +244,11 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t wg_size = utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX); const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u); + // GEMV (bicol) is a pow2 tree reduction; compute its size only when used. + const uint32_t gemv_wg_size = use_gemv + ? utils::clamp_workgroup_size_pow2( + device, kQ4gswLinearCoop4BicolWorkgroupSizeX) + : 0u; // steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible. const bool use_steel = !use_gemv && steel_supported(device) && steel_workgroup_count(device, M, N, K) > 0u; @@ -362,16 +367,17 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { WGPUPipelineLayout pipeline_layout = wgpuDeviceCreatePipelineLayout(device, &pl_desc); + // GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed. + const bool fixed_wg = use_steel || use_shmem_gemm; WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); + wg_size_constant.value = + static_cast(use_gemv ? gemv_wg_size : wg_size); WGPUComputePipelineDescriptor pipeline_desc = {}; pipeline_desc.layout = pipeline_layout; pipeline_desc.compute.module = shader; pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - // Only tiled GEMM overrides wg_size; GEMV/shmem (64) + steel (256) are fixed. - const bool fixed_wg = use_gemv || use_steel || use_shmem_gemm; pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u; pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant; WGPUComputePipeline pipeline = diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl deleted file mode 100644 index af6f661279d..00000000000 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl +++ /dev/null @@ -1,87 +0,0 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_input: array; -@group(0) @binding(2) var t_weight: array; -@group(0) @binding(3) var t_scales: array; -@group(0) @binding(4) var t_bias: array; - -struct Params { - M: u32, - N: u32, - K: u32, - K_packed: u32, - group_size: u32, - padded_N: u32, - has_bias: u32, - _pad: u32, -} -@group(0) @binding(5) var params: Params; - -// Cooperative-over-K GEMV with u32-batched coalesced weight loads (64 lanes). -const WG: u32 = 64u; -var partial: array; - -@compute @workgroup_size(WG, 1, 1) -fn main( - @builtin(workgroup_id) wid: vec3, - @builtin(num_workgroups) ngrp: vec3, - @builtin(local_invocation_id) lid: vec3) { - let total = params.M * params.N; - let stride = ngrp.x; - let num_words = params.K >> 3u; // K / 8 words per row - let row_words = params.K_packed >> 2u; // u32s per weight row (= K/8) - var idx = wid.x; - loop { - if (idx >= total) { - break; - } - let m = idx / params.N; - let n = idx % params.N; - let in_base = m * params.K; - let wbase = n * row_words; - - var acc: f32 = 0.0; - var w: u32 = lid.x; - loop { - if (w >= num_words) { - break; - } - let word = t_weight[wbase + w]; - let k0 = w << 3u; // first K of this word - let scale = t_scales[(k0 / params.group_size) * params.padded_N + n]; - let ib = in_base + k0; - // 4 bytes, low+high nibble each -> 8 consecutive K. - for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) { - let byte = (word >> (bi * 8u)) & 0xFFu; - let lo = f32(i32(byte & 0x0Fu) - 8); - let hi = f32(i32((byte >> 4u) & 0x0Fu) - 8); - let kk = bi << 1u; - acc = acc + t_input[ib + kk] * lo * scale; - acc = acc + t_input[ib + kk + 1u] * hi * scale; - } - w = w + WG; - } - - partial[lid.x] = acc; - workgroupBarrier(); - var s: u32 = WG >> 1u; - loop { - if (s == 0u) { - break; - } - if (lid.x < s) { - partial[lid.x] = partial[lid.x] + partial[lid.x + s]; - } - workgroupBarrier(); - s = s >> 1u; - } - if (lid.x == 0u) { - var o = partial[0]; - if (params.has_bias != 0u) { - o = o + t_bias[n]; - } - t_out[idx] = o; - } - workgroupBarrier(); - idx = idx + stride; - } -} diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl index 69174269f38..0d33c5a3a65 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl @@ -17,10 +17,11 @@ struct Params { @group(0) @binding(5) var params: Params; // Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused. -const WG: u32 = 64u; -var partial: array, WG>; +// wg_size must be a power of two (the tree reduction halves the stride). +override wg_size: u32 = 64u; +var partial: array, wg_size>; -@compute @workgroup_size(WG, 1, 1) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @builtin(num_workgroups) ngrp: vec3, @@ -69,12 +70,12 @@ fn main( acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1; acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1; } - w = w + WG; + w = w + wg_size; } partial[lid.x] = vec2(acc0, acc1); workgroupBarrier(); - var s: u32 = WG >> 1u; + var s: u32 = wg_size >> 1u; loop { if (s == 0u) { break; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h index d6c51bcb893..864165da136 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_coop4_bicol.wgsl - DO NOT EDIT. -// wgsl-sha256: 44c53ca7879f7a3382a45703a67cbc6fc221ecbcada58f8c413df50abbc6e544 +// wgsl-sha256: ef677c2771c3d4a73704f4725bb490174149b971bc62e696595b8d13941c979e inline constexpr const char* kQ4gswLinearCoop4BicolWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_input: array; @@ -34,10 +34,11 @@ struct Params { @group(0) @binding(5) var params: Params; // Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused. -const WG: u32 = 64u; -var partial: array, WG>; +// wg_size must be a power of two (the tree reduction halves the stride). +override wg_size: u32 = 64u; +var partial: array, wg_size>; -@compute @workgroup_size(WG, 1, 1) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @builtin(num_workgroups) ngrp: vec3, @@ -86,12 +87,12 @@ fn main( acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1; acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1; } - w = w + WG; + w = w + wg_size; } partial[lid.x] = vec2(acc0, acc1); workgroupBarrier(); - var s: u32 = WG >> 1u; + var s: u32 = wg_size >> 1u; loop { if (s == 0u) { break; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h deleted file mode 100644 index 7bacedfcb17..00000000000 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h +++ /dev/null @@ -1,111 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include - -namespace executorch::backends::webgpu { - -// @generated from q4gsw_linear_coop4.wgsl - DO NOT EDIT. -// wgsl-sha256: 3031886e68c375e617dfb263da39c492c6de4d8c1fb4073d70b18823a3e6a4fe -inline constexpr const char* kQ4gswLinearCoop4WGSL = R"( -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_input: array; -@group(0) @binding(2) var t_weight: array; -@group(0) @binding(3) var t_scales: array; -@group(0) @binding(4) var t_bias: array; - -struct Params { - M: u32, - N: u32, - K: u32, - K_packed: u32, - group_size: u32, - padded_N: u32, - has_bias: u32, - _pad: u32, -} -@group(0) @binding(5) var params: Params; - -// Cooperative-over-K GEMV with u32-batched coalesced weight loads (64 lanes). -const WG: u32 = 64u; -var partial: array; - -@compute @workgroup_size(WG, 1, 1) -fn main( - @builtin(workgroup_id) wid: vec3, - @builtin(num_workgroups) ngrp: vec3, - @builtin(local_invocation_id) lid: vec3) { - let total = params.M * params.N; - let stride = ngrp.x; - let num_words = params.K >> 3u; // K / 8 words per row - let row_words = params.K_packed >> 2u; // u32s per weight row (= K/8) - var idx = wid.x; - loop { - if (idx >= total) { - break; - } - let m = idx / params.N; - let n = idx % params.N; - let in_base = m * params.K; - let wbase = n * row_words; - - var acc: f32 = 0.0; - var w: u32 = lid.x; - loop { - if (w >= num_words) { - break; - } - let word = t_weight[wbase + w]; - let k0 = w << 3u; // first K of this word - let scale = t_scales[(k0 / params.group_size) * params.padded_N + n]; - let ib = in_base + k0; - // 4 bytes, low+high nibble each -> 8 consecutive K. - for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) { - let byte = (word >> (bi * 8u)) & 0xFFu; - let lo = f32(i32(byte & 0x0Fu) - 8); - let hi = f32(i32((byte >> 4u) & 0x0Fu) - 8); - let kk = bi << 1u; - acc = acc + t_input[ib + kk] * lo * scale; - acc = acc + t_input[ib + kk + 1u] * hi * scale; - } - w = w + WG; - } - - partial[lid.x] = acc; - workgroupBarrier(); - var s: u32 = WG >> 1u; - loop { - if (s == 0u) { - break; - } - if (lid.x < s) { - partial[lid.x] = partial[lid.x] + partial[lid.x + s]; - } - workgroupBarrier(); - s = s >> 1u; - } - if (lid.x == 0u) { - var o = partial[0]; - if (params.has_bias != 0u) { - o = o + t_bias[n]; - } - t_out[idx] = o; - } - workgroupBarrier(); - idx = idx + stride; - } -} -)"; - -inline constexpr uint32_t kQ4gswLinearCoop4WorkgroupSizeX = 64; -inline constexpr uint32_t kQ4gswLinearCoop4WorkgroupSizeY = 1; -inline constexpr uint32_t kQ4gswLinearCoop4WorkgroupSizeZ = 1; - -} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl index a7384f36115..e68af88ad4c 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl @@ -16,7 +16,8 @@ struct Params { } @group(0) @binding(5) var params: Params; -override wg_size: u32 = 64u; +// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob). +const wg_size: u32 = 64u; // Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile. const WG_M: u32 = 32u; // output rows per workgroup diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h index 64bbd2cd15f..3255a05220d 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_shmem.wgsl - DO NOT EDIT. -// wgsl-sha256: 0a1219d3b6781315a21066089fca3f92235587e8af8eb734185f35ea4bfc8a52 +// wgsl-sha256: f0180cad41701807b34912b57d9715c02588c9447783bbf691f93836168fe981 inline constexpr const char* kQ4gswLinearGemmShmemWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_input: array; @@ -33,7 +33,8 @@ struct Params { } @group(0) @binding(5) var params: Params; -override wg_size: u32 = 64u; +// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob). +const wg_size: u32 = 64u; // Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile. const WG_M: u32 = 32u; // output rows per workgroup diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl index 4d7ab0b1d1e..f08cb39bd5c 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl @@ -34,6 +34,7 @@ struct Params { const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array<${buffer_scalar_type(DTYPE)}, 1024>; // BM*BK var Bs: array<${buffer_scalar_type(DTYPE)}, 1024>; // BK*BN +// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h index efefd7edce1..6d23ea7ac76 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: 36b3d3f9dd08a529909c13ec7d66cd0cf392c347ca047a4d38453b3c295f72ce +// wgsl-sha256: fe8b71afc5634e7db5607deb58f30c6517f93e1d66370ea017bcd9071201c6ca inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -50,6 +50,7 @@ struct Params { const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN +// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h index 46057de2340..e684b00a228 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: 1f916bcd30dbbbcc7eca37e795ecc26e3c72e645ccd2c361fa0ac4e66f1a174a +// wgsl-sha256: 1b3f5e384e09bde50c66d43db4797e01e5610d70af235355dbc4966643bf9103 inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -50,6 +50,7 @@ struct Params { const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN +// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h index ae03f63fb5c..d4c78406582 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: 00cdd2f2fb98a5c7343d16fdf7e59f1b840e180cec3f82bf9b569513c0a45396 +// wgsl-sha256: d77542e5975fab44c58644c771390e89598ce82ab25e38d57150b581b4904acd inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -50,6 +50,7 @@ struct Params { const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN +// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h index 71b0f45bdf7..db7a7689a6f 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: dd771b9ab096410f3ad0d9259bef7816e41330a434325ea28baa5abbfb2841d2 +// wgsl-sha256: dcb63b04d002899acfed88d273fe658618650d5dcbc108540cef086bb97f5bce inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_input: array; @@ -49,6 +49,7 @@ struct Params { const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN +// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp index 39aaae3a6bf..d9f6b89e9f2 100644 --- a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp +++ b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp @@ -178,11 +178,21 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { WGPUPipelineLayout pipeline_layout = wgpuDeviceCreatePipelineLayout(device, &pl_desc); + // Runtime-overridable workgroup size (mirrors add op); clamp only reduces. + // Pow2 required: the kernel halves the reduction stride (wg_size / 2u). + const uint32_t wg_size = + utils::clamp_workgroup_size_pow2(device, kRmsNormWorkgroupSizeX); + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + // Create compute pipeline WGPUComputePipelineDescriptor pipeline_desc = {}; pipeline_desc.layout = pipeline_layout; pipeline_desc.compute.module = shader; pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; WGPUComputePipeline pipeline = wgpuDeviceCreateComputePipeline(device, &pipeline_desc); @@ -217,10 +227,10 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { // One workgroup per row (kRmsNormWorkgroupSizeX threads cooperate per row) static_assert( kRmsNormWorkgroupSizeX == 64, - "must match @workgroup_size and WG_SIZE in rms_norm.wgsl"); + "kRmsNormWorkgroupSizeX must match override wg_size default (64)"); static_assert( kRmsNormVec4WorkgroupSizeX == 64, - "must match @workgroup_size and WG_SIZE in rms_norm_vec4.wgsl"); + "kRmsNormVec4WorkgroupSizeX must match override wg_size default (64)"); const size_t dispatch_idx = graph.add_dispatch({pipeline, bind_group, num_rows}); diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl index 1234174981e..11302434ec2 100644 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl @@ -12,13 +12,14 @@ struct Params { } @group(0) @binding(3) var params: Params; -const WG_SIZE: u32 = 64u; +// wg_size must be a power of two (the tree reduction halves the stride). +override wg_size: u32 = 64u; -var shared_sum: array<${accum_scalar_type(DTYPE)}, WG_SIZE>; +var shared_sum: array<${accum_scalar_type(DTYPE)}, wg_size>; fn reduce_shared(worker_id: u32) { workgroupBarrier(); - var stride: u32 = WG_SIZE / 2u; + var stride: u32 = wg_size / 2u; loop { if (stride == 0u) { break; @@ -32,10 +33,10 @@ fn reduce_shared(worker_id: u32) { } $if VEC == 4: - // vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4 + // vec4 variant of rms_norm: each lane strides by wg_size over rw4 = row_width/4 // texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq // divides by it (not rw4). The host selects this only when row_width % 4 == 0. -@compute @workgroup_size(64, 1, 1) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -64,7 +65,7 @@ fn main( local_sq_sum = local_sq_sum + dot(vec4(v), vec4(v)); $else: local_sq_sum = local_sq_sum + dot(v, v); - x4 = x4 + WG_SIZE; + x4 = x4 + wg_size; } $else: var x: u32 = worker_id; @@ -77,7 +78,7 @@ fn main( local_sq_sum = local_sq_sum + f32(v) * f32(v); $else: local_sq_sum = local_sq_sum + v * v; - x = x + WG_SIZE; + x = x + wg_size; } shared_sum[worker_id] = local_sq_sum; @@ -96,7 +97,7 @@ fn main( t_out[base4 + x4] = vec4(vec4(t_in[base4 + x4]) * rstd * vec4(t_weight[x4])); $else: t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4]; - x4 = x4 + WG_SIZE; + x4 = x4 + wg_size; } $else: x = worker_id; @@ -110,6 +111,6 @@ fn main( t_out[base + x] = f16(f32(v) * rstd * f32(w)); $else: t_out[base + x] = v * rstd * w; - x = x + WG_SIZE; + x = x + wg_size; } } diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h index 02213e3d7b0..2e17ac02aa0 100644 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from rms_norm.wgsl - DO NOT EDIT. -// wgsl-sha256: 4c0ba56708bf125a7ec6ea3c51d1288e05ac00a8e2cfa10e38e9a208e230b8df +// wgsl-sha256: 62fdfe03fc67eb44fa17ca5e91e433b3a45a96c76151777856d6f557fd829919 inline constexpr const char* kRmsNormVec4WGSL = R"( @group(0) @binding(0) var t_out: array>; @group(0) @binding(1) var t_in: array>; @@ -27,13 +27,14 @@ struct Params { } @group(0) @binding(3) var params: Params; -const WG_SIZE: u32 = 64u; +// wg_size must be a power of two (the tree reduction halves the stride). +override wg_size: u32 = 64u; -var shared_sum: array; +var shared_sum: array; fn reduce_shared(worker_id: u32) { workgroupBarrier(); - var stride: u32 = WG_SIZE / 2u; + var stride: u32 = wg_size / 2u; loop { if (stride == 0u) { break; @@ -46,10 +47,10 @@ fn reduce_shared(worker_id: u32) { } } -// vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4 +// vec4 variant of rms_norm: each lane strides by wg_size over rw4 = row_width/4 // texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq // divides by it (not rw4). The host selects this only when row_width % 4 == 0. -@compute @workgroup_size(64, 1, 1) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -71,7 +72,7 @@ fn main( } let v = t_in[base4 + x4]; local_sq_sum = local_sq_sum + dot(v, v); - x4 = x4 + WG_SIZE; + x4 = x4 + wg_size; } shared_sum[worker_id] = local_sq_sum; @@ -86,7 +87,7 @@ fn main( break; } t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4]; - x4 = x4 + WG_SIZE; + x4 = x4 + wg_size; } } )"; diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h index 5d9fc236e91..841b146b2ef 100644 --- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h +++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from rms_norm.wgsl - DO NOT EDIT. -// wgsl-sha256: 340dcbf3c06dc311e70bef953c1e9cbbdf4121fe177eedd3253549e614b55069 +// wgsl-sha256: 55bc862d64451f5fe8f530114f0318dfbc724017314091418a1fe41b180be7c6 inline constexpr const char* kRmsNormWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_in: array; @@ -27,13 +27,14 @@ struct Params { } @group(0) @binding(3) var params: Params; -const WG_SIZE: u32 = 64u; +// wg_size must be a power of two (the tree reduction halves the stride). +override wg_size: u32 = 64u; -var shared_sum: array; +var shared_sum: array; fn reduce_shared(worker_id: u32) { workgroupBarrier(); - var stride: u32 = WG_SIZE / 2u; + var stride: u32 = wg_size / 2u; loop { if (stride == 0u) { break; @@ -46,7 +47,7 @@ fn reduce_shared(worker_id: u32) { } } -@compute @workgroup_size(64, 1, 1) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -67,7 +68,7 @@ fn main( } let v = t_in[base + x]; local_sq_sum = local_sq_sum + v * v; - x = x + WG_SIZE; + x = x + wg_size; } shared_sum[worker_id] = local_sq_sum; @@ -84,7 +85,7 @@ fn main( let v = t_in[base + x]; let w = t_weight[x]; t_out[base + x] = v * rstd * w; - x = x + WG_SIZE; + x = x + wg_size; } } )"; diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 50321ba4bdf..f3345c89c77 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -188,7 +188,7 @@ void build_dispatch( WGPUPipelineLayout pipeline_layout = wgpuDeviceCreatePipelineLayout(device, &pl_desc); - // QK/AV/update_cache have an `override wg_size`; softmax (0) keeps a const. + // All callers pass an override wg_size; a 0 would keep the shader default. WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); @@ -529,6 +529,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { // One workgroup per (h,s) row; wg_size 1 keeps the device dispatch check. const utils::WgCount wgc = utils::compute_2d_workgroup_count( device, static_cast(Hq * S), 1, "softmax"); + const uint32_t sm_wg = + utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX); SoftmaxParams p = make_softmax_params(Hq, S, context_len); WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); BufferBinding bindings[2] = { @@ -542,7 +544,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { sizeof(p), wgc.x, wgc.y, - 0, + sm_wg, true, "sdpa_softmax"); softmax_buf = ubuf; diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl index f0e349ceaff..724e427073a 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl @@ -9,15 +9,16 @@ struct Params { } @group(0) @binding(2) var params: Params; -const WG_SIZE: u32 = 64u; +// wg_size must be a power of two (the tree reduction halves the stride). +override wg_size: u32 = 64u; // WGSL forbids literal -inf; a large finite negative inits the running max. const NEG_INF: f32 = -1.0e30; -var shared_max: array; -var shared_sum: array; +var shared_max: array; +var shared_sum: array; -@compute @workgroup_size(WG_SIZE, 1, 1) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3, @@ -41,14 +42,14 @@ fn main( break; } local_max = max(local_max, t_in[base + x]); - x = x + WG_SIZE; + x = x + wg_size; } } shared_max[worker_id] = local_max; // Reduce max. workgroupBarrier() calls are in uniform control flow. workgroupBarrier(); - var stride: u32 = WG_SIZE / 2u; + var stride: u32 = wg_size / 2u; loop { if (stride == 0u) { break; @@ -70,13 +71,13 @@ fn main( break; } local_sum = local_sum + exp(t_in[base + x] - row_max); - x = x + WG_SIZE; + x = x + wg_size; } } shared_sum[worker_id] = local_sum; workgroupBarrier(); - stride = WG_SIZE / 2u; + stride = wg_size / 2u; loop { if (stride == 0u) { break; @@ -98,7 +99,7 @@ fn main( break; } t_out[base + x] = exp(t_in[base + x] - row_max) * inv; - x = x + WG_SIZE; + x = x + wg_size; } } } diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h index f1e7396fafe..077f173a546 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sdpa_softmax.wgsl - DO NOT EDIT. -// wgsl-sha256: bdd45cf344663533b243153200c507f41a90295751924e70452abcd5da4cdd5a +// wgsl-sha256: 774a372c0eae85b6656618408e40fff4b8af5647ea5353b51d3a34b9b6a7ac8a inline constexpr const char* kSdpaSoftmaxWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_in: array; @@ -26,15 +26,16 @@ struct Params { } @group(0) @binding(2) var params: Params; -const WG_SIZE: u32 = 64u; +// wg_size must be a power of two (the tree reduction halves the stride). +override wg_size: u32 = 64u; // WGSL forbids literal -inf; a large finite negative inits the running max. const NEG_INF: f32 = -1.0e30; -var shared_max: array; -var shared_sum: array; +var shared_max: array; +var shared_sum: array; -@compute @workgroup_size(WG_SIZE, 1, 1) +@compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3, @@ -58,14 +59,14 @@ fn main( break; } local_max = max(local_max, t_in[base + x]); - x = x + WG_SIZE; + x = x + wg_size; } } shared_max[worker_id] = local_max; // Reduce max. workgroupBarrier() calls are in uniform control flow. workgroupBarrier(); - var stride: u32 = WG_SIZE / 2u; + var stride: u32 = wg_size / 2u; loop { if (stride == 0u) { break; @@ -87,13 +88,13 @@ fn main( break; } local_sum = local_sum + exp(t_in[base + x] - row_max); - x = x + WG_SIZE; + x = x + wg_size; } } shared_sum[worker_id] = local_sum; workgroupBarrier(); - stride = WG_SIZE / 2u; + stride = wg_size / 2u; loop { if (stride == 0u) { break; @@ -115,7 +116,7 @@ fn main( break; } t_out[base + x] = exp(t_in[base + x] - row_max) * inv; - x = x + WG_SIZE; + x = x + wg_size; } } } diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl index 60f86d217ff..6ba784a5d5a 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl @@ -19,6 +19,7 @@ const NEG_INF: f32 = -1.0e30; var sh_w: array; // FlashDecoding pass 2: online-softmax merge of the per-split partials, then normalize. +// 64 bound to MAX_D_PER_LANE output coverage (D<=WG_SIZE*2=128); not a knob. @compute @workgroup_size(64, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h index 6e5a515c065..e6eafe901ef 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sdpa_fd_reduce.wgsl - DO NOT EDIT. -// wgsl-sha256: 3debe9b52adece82b31a067a876105121afb4e09cb31d95596ea699d1a179d18 +// wgsl-sha256: 776e9605b309f7b5d796f28b70512ed030c7178075e6b3f843176b00494a03a6 inline constexpr const char* kSdpaFdReduceWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_part_o: array; @@ -36,6 +36,7 @@ const NEG_INF: f32 = -1.0e30; var sh_w: array; // FlashDecoding pass 2: online-softmax merge of the per-split partials, then normalize. +// 64 bound to MAX_D_PER_LANE output coverage (D<=WG_SIZE*2=128); not a knob. @compute @workgroup_size(64, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl index da67489cc4d..f5a003b4562 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl @@ -28,6 +28,7 @@ var sh_s: array; var sh_red: array; // FlashDecoding pass 1: per-(head,split) unnormalized softmax partial. +// 64 sizes sh_s/sh_red + MAX_D_PER_LANE output coverage (D<=128); not a knob. @compute @workgroup_size(64, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h index bc69a444edb..a0898512905 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sdpa_fd_split.wgsl - DO NOT EDIT. -// wgsl-sha256: 147fc5775b76f3626dc934b2df5e34148bef12c345789332f098d13143bb646e +// wgsl-sha256: 596f7376a2f15f788827b921b293de91541f53758d5ee7fe858bf4a72ab102df inline constexpr const char* kSdpaFdSplitHalfWGSL = R"( enable f16; @group(0) @binding(0) var t_part_o: array; @@ -44,6 +44,7 @@ var sh_s: array; var sh_red: array; // FlashDecoding pass 1: per-(head,split) unnormalized softmax partial. +// 64 sizes sh_s/sh_red + MAX_D_PER_LANE output coverage (D<=128); not a knob. @compute @workgroup_size(64, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h index 6b23e6fd36e..cb362c9cc82 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from sdpa_fd_split.wgsl - DO NOT EDIT. -// wgsl-sha256: 44bee2691f85839b638a51971a8966bcf8efbb72b8acc44a4c3e991f31cc8fcb +// wgsl-sha256: b5a926169f085c870ebb551510722dad5ee10abaf97f355eab2635268923c34c inline constexpr const char* kSdpaFdSplitWGSL = R"( @group(0) @binding(0) var t_part_o: array; @group(0) @binding(1) var t_part_ml: array; @@ -43,6 +43,7 @@ var sh_s: array; var sh_red: array; // FlashDecoding pass 1: per-(head,split) unnormalized softmax partial. +// 64 sizes sh_s/sh_red + MAX_D_PER_LANE output coverage (D<=128); not a knob. @compute @workgroup_size(64, 1, 1) fn main( @builtin(workgroup_id) wid: vec3, diff --git a/backends/webgpu/test/ops/test_quantized_linear.py b/backends/webgpu/test/ops/test_quantized_linear.py index 72945c37d6d..54040a54bd7 100644 --- a/backends/webgpu/test/ops/test_quantized_linear.py +++ b/backends/webgpu/test/ops/test_quantized_linear.py @@ -57,8 +57,8 @@ class Q4gswConfig: # decode GEMV: the handler routes M==1 -> bicol, so each reads its own per- # column scale (col0/col1) across many K-groups (down_proj: 256 groups). q4gsw # requires N % 8 == 0 (torchao pads N for the scale layout), so odd-N / N=1 are - # not exportable -- bicol's has1 odd-N guard is defensive (mirrors coop4's - # general-N robustness) and unreachable through this op. + # not exportable -- bicol's has1 odd-N guard is defensive and unreachable + # through this op. # M>1 prefill: prefer the steel GEMM (K%16==0) on a >=256-invocation device # (e.g. lvp); else shmem (K>=4096 or N>=2048) or register-tiled (SwiftShader # caps at 128). Same fp64 golden regardless of which kernel runs. diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index fbdfbd09076..5ba4a95f976 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -6,9 +6,11 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include +#include #include #include #include @@ -21,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -256,6 +259,185 @@ bool sdpa_within_tol( return ok; } +// Matches the WGSL Params struct in rms_norm.wgsl (16-byte aligned). +struct RmsNormProbeParams { + uint32_t num_rows; + uint32_t row_width; + float epsilon; + uint32_t _pad; +}; + +struct WgMapData { + WGPUMapAsyncStatus status = WGPUMapAsyncStatus_Error; +}; +void wg_map_cb( + WGPUMapAsyncStatus status, + WGPUStringView /*message*/, + void* userdata1, + void* /*userdata2*/) { + static_cast(userdata1)->status = status; +} + +// Run the rms_norm scalar kernel at an explicit override wg_size and map the +// output back. Module::forward can't set a second workgroup size (the handler +// always clamps to 64), so the pipeline is built directly here; the +// map/readback mirrors WebGPUGraph::copy_outputs. +std::vector run_rms_norm_at_wg( + const WebGPUContext& ctx, + uint32_t wg_size, + const std::vector& input, + const std::vector& weight, + uint32_t num_rows, + uint32_t row_width, + float epsilon) { + WGPUDevice device = ctx.device; + const uint64_t out_bytes = + static_cast(num_rows) * row_width * sizeof(float); + const uint64_t in_bytes = static_cast(input.size()) * sizeof(float); + const uint64_t w_bytes = static_cast(weight.size()) * sizeof(float); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kRmsNormWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry bgl_entries[4] = {}; + bgl_entries[0].binding = 0; + bgl_entries[0].visibility = WGPUShaderStage_Compute; + bgl_entries[0].buffer.type = WGPUBufferBindingType_Storage; + bgl_entries[1].binding = 1; + bgl_entries[1].visibility = WGPUShaderStage_Compute; + bgl_entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + bgl_entries[2].binding = 2; + bgl_entries[2].visibility = WGPUShaderStage_Compute; + bgl_entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + bgl_entries[3].binding = 3; + bgl_entries[3].visibility = WGPUShaderStage_Compute; + bgl_entries[3].buffer.type = WGPUBufferBindingType_Uniform; + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = bgl_entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUConstantEntry wg_const = {}; + wg_const.key = {"wg_size", WGPU_STRLEN}; + wg_const.value = static_cast(wg_size); + + WGPUComputePipelineDescriptor pipe_desc = {}; + pipe_desc.layout = pl; + pipe_desc.compute.module = shader; + pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipe_desc.compute.constantCount = 1; + pipe_desc.compute.constants = &wg_const; + WGPUComputePipeline pipe = + wgpuDeviceCreateComputePipeline(device, &pipe_desc); + + WGPUBufferDescriptor out_bd = {}; + out_bd.size = out_bytes; + out_bd.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopySrc; + WGPUBuffer out_buf = wgpuDeviceCreateBuffer(device, &out_bd); + + WGPUBufferDescriptor in_bd = {}; + in_bd.size = in_bytes; + in_bd.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + WGPUBuffer in_buf = wgpuDeviceCreateBuffer(device, &in_bd); + wgpuQueueWriteBuffer(ctx.queue, in_buf, 0, input.data(), in_bytes); + + WGPUBufferDescriptor w_bd = {}; + w_bd.size = w_bytes; + w_bd.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + WGPUBuffer w_buf = wgpuDeviceCreateBuffer(device, &w_bd); + wgpuQueueWriteBuffer(ctx.queue, w_buf, 0, weight.data(), w_bytes); + + RmsNormProbeParams params = {num_rows, row_width, epsilon, 0u}; + WGPUBufferDescriptor p_bd = {}; + p_bd.size = sizeof(RmsNormProbeParams); + p_bd.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + WGPUBuffer p_buf = wgpuDeviceCreateBuffer(device, &p_bd); + wgpuQueueWriteBuffer(ctx.queue, p_buf, 0, ¶ms, sizeof(params)); + + WGPUBufferDescriptor stg_bd = {}; + stg_bd.size = out_bytes; + stg_bd.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; + WGPUBuffer staging = wgpuDeviceCreateBuffer(device, &stg_bd); + + WGPUBindGroupEntry bg_entries[4] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = out_buf; + bg_entries[0].size = out_bytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = in_buf; + bg_entries[1].size = in_bytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = w_buf; + bg_entries[2].size = w_bytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = p_buf; + bg_entries[3].size = sizeof(RmsNormProbeParams); + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg_entries; + WGPUBindGroup bg = wgpuDeviceCreateBindGroup(device, &bg_desc); + + WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(device, nullptr); + WGPUComputePassDescriptor pass_desc = {}; + WGPUComputePassEncoder pass = + wgpuCommandEncoderBeginComputePass(enc, &pass_desc); + wgpuComputePassEncoderSetPipeline(pass, pipe); + wgpuComputePassEncoderSetBindGroup(pass, 0, bg, 0, nullptr); + wgpuComputePassEncoderDispatchWorkgroups(pass, num_rows, 1, 1); + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); + wgpuCommandEncoderCopyBufferToBuffer(enc, out_buf, 0, staging, 0, out_bytes); + WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(enc, nullptr); + wgpuQueueSubmit(ctx.queue, 1, &cmd); + wgpuCommandBufferRelease(cmd); + wgpuCommandEncoderRelease(enc); + + WgMapData cb = {}; + WGPUBufferMapCallbackInfo cb_info = {}; + cb_info.mode = WGPUCallbackMode_WaitAnyOnly; + cb_info.callback = wg_map_cb; + cb_info.userdata1 = &cb; + WGPUFuture fut = + wgpuBufferMapAsync(staging, WGPUMapMode_Read, 0, out_bytes, cb_info); + const WGPUWaitStatus wait = webgpu_wait(ctx.instance, fut); + + std::vector result(static_cast(num_rows) * row_width); + bool ok = false; + if (wait == WGPUWaitStatus_Success && + cb.status == WGPUMapAsyncStatus_Success) { + const void* mapped = wgpuBufferGetConstMappedRange(staging, 0, out_bytes); + std::memcpy(result.data(), mapped, out_bytes); + wgpuBufferUnmap(staging); + ok = true; + } + + wgpuBufferRelease(staging); + wgpuBufferRelease(p_buf); + wgpuBufferRelease(w_buf); + wgpuBufferRelease(in_buf); + wgpuBufferRelease(out_buf); + wgpuBindGroupRelease(bg); + wgpuComputePipelineRelease(pipe); + wgpuPipelineLayoutRelease(pl); + wgpuBindGroupLayoutRelease(bgl); + wgpuShaderModuleRelease(shader); + + if (!ok) { + throw std::runtime_error("rms_norm wg-size probe: output map failed"); + } + return result; +} + // linear_q4gsw sweep config; mirrors CONFIGS in test_quantized_linear.py. struct Q4gswConfig { const char* name; @@ -284,8 +466,7 @@ const Q4gswConfig kQ4gswConfigs[] = { // The M==1 configs above (q/kv/gate/down_proj) exercise the bicol 2-col // decode GEMV (handler routes M==1 -> bicol; each reads its own per-column // scale over 64-256 K-groups). q4gsw requires N % 8 == 0, so odd-N is not - // exportable; bicol's has1 odd-N guard is defensive (mirrors coop4 - // general-N robustness). + // exportable; bicol's has1 odd-N guard is defensive. // M>1: steel GEMM on a >=256-invocation device (K%16==0), else shmem/tiled. {"steel", 96, 2048, 256, 1e-4f, 1e-3f, true, false}, // steel-isolating // Same shape as "steel" run under the f16-multiply steel kernel; the f16 @@ -1415,6 +1596,56 @@ TEST(WebGPUNative, QueryPoolRoundtrip) { } #endif // WGPU_BACKEND_ENABLE_PROFILING +// The override wg_size must not change results: run the rms_norm scalar kernel +// at wg_size 64 and 128 (the handler always clamps to 64, so 128 is only +// reachable via a direct pipeline) and require element-wise agreement. Absolute +// rms_norm correctness is covered by the model-driven golden tests; this locks +// the runtime-configurability guarantee (same WGSL, different size -> same +// output). +TEST(WebGPUNative, RmsNormWorkgroupSizeConfigurable) { + const WebGPUContext* ctx = get_default_webgpu_context(); + if (ctx == nullptr || ctx->device == nullptr) { + GTEST_SKIP() << "no WebGPU device"; + } + constexpr uint32_t num_rows = 3, row_width = 256; + constexpr float epsilon = 1e-5f; + std::vector input(static_cast(num_rows) * row_width); + std::vector weight(row_width); + for (size_t i = 0; i < input.size(); i++) { + input[i] = std::sin(0.1f * static_cast(i)) * 2.0f + 0.5f; + } + for (uint32_t j = 0; j < row_width; j++) { + weight[j] = 0.5f + 0.001f * static_cast(j); + } + + const std::vector out64 = + run_rms_norm_at_wg(*ctx, 64, input, weight, num_rows, row_width, epsilon); + const std::vector out128 = run_rms_norm_at_wg( + *ctx, 128, input, weight, num_rows, row_width, epsilon); + ASSERT_EQ(out64.size(), static_cast(num_rows) * row_width); + ASSERT_EQ(out128.size(), out64.size()); + + double sumsq = 0.0; + for (float v : out64) { + sumsq += static_cast(v) * static_cast(v); + } + float max_abs = 0.0f, max_rel = 0.0f; + const bool consistent = sdpa_within_tol( + out64.data(), + out128.data(), + static_cast(out64.size()), + &max_abs, + &max_rel); + printf( + " rms_norm wg64-vs-wg128: max_abs=%e max_rel=%e sumsq=%f\n", + max_abs, + max_rel, + sumsq); + EXPECT_GT(sumsq, 0.0) << "wg64 output is all zero (kernel did not run)"; + EXPECT_TRUE(consistent) << "wg64 vs wg128 differ beyond tol (abs " << max_abs + << " rel " << max_rel << ")"; +} + TEST(WebGPUNative, UpdateCache) { if (g_update_cache_model_path.empty()) { GTEST_SKIP() << "WEBGPU_TEST_UPDATE_CACHE_MODEL not set";