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..2a3e1a9ce77 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -284,8 +284,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