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_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/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";