diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index 0b0b5254908..e1497bc4ed8 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -89,13 +89,22 @@ Result WebGPUBackend::init( enable_f16_kv_cache = spec.get(); } } + bool enable_f16_accumulate_gemm = false; + { + Result spec = + context.get_runtime_spec("enable_f16_accumulate_gemm"); + if (spec.ok()) { + enable_f16_accumulate_gemm = spec.get(); + } + } try { graph->build( flatbuffer_data, constant_data, context.get_named_data_map(), - enable_f16_kv_cache); + enable_f16_kv_cache, + enable_f16_accumulate_gemm); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU graph build failed: %s", e.what()); graph->~WebGPUGraph(); diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index f7bc4c58660..d9ab6467882 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -90,6 +90,49 @@ WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { return buffer; } +WGPUBuffer WebGPUGraph::acquire_scratch(size_t nbytes) { + nbytes = nbytes > 0 ? nbytes : 4; + // Best-fit reuse: smallest free slot with size in [nbytes, 2*nbytes] -- the + // 2x cap stops a large Cmax-sized buffer from backing a tiny request. Never + // reuse an in_use slot (co-live safety). + ScratchSlot* best = nullptr; + for (auto& s : scratch_pool_) { + // s.size - nbytes (safe: s.size >= nbytes) avoids overflowing 2 * nbytes. + if (!s.in_use && s.size >= nbytes && s.size - nbytes <= nbytes) { + if (best == nullptr || s.size < best->size) { + best = &s; + } + } + } + if (best != nullptr) { + best->in_use = true; + return best->buffer; + } + // None reusable -> create a new slot (freed in the dtor, like + // scratch_buffers_). + WGPUBufferDescriptor buf_desc = {}; + buf_desc.size = nbytes; + buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | + WGPUBufferUsage_CopySrc; + buf_desc.mappedAtCreation = false; + WGPUBuffer buffer = wgpuDeviceCreateBuffer(device_, &buf_desc); + scratch_pool_.push_back({buffer, nbytes, true}); + return buffer; +} + +void WebGPUGraph::release_scratch(WGPUBuffer buffer) { + if (!buffer) { + return; + } + for (auto& s : scratch_pool_) { + if (s.buffer == buffer) { + s.in_use = false; + return; + } + } + // Not a pooled buffer -> no-op; the dtor frees it via scratch_buffers_. +} + WGPUBuffer WebGPUGraph::make_uniform_buffer(const void* data, size_t size) { WGPUBufferDescriptor desc = {}; desc.size = size; @@ -267,6 +310,11 @@ WebGPUGraph::~WebGPUGraph() { wgpuBufferRelease(buf); } } + for (auto& s : scratch_pool_) { + if (s.buffer) { + wgpuBufferRelease(s.buffer); + } + } for (auto& buf : owned_uniform_buffers_) { if (buf) { wgpuBufferRelease(buf); @@ -311,7 +359,8 @@ void WebGPUGraph::build( const void* flatbuffer_data, const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map, - bool f16_kv_cache) { + bool f16_kv_cache, + bool f16_accumulate_gemm) { if (!device_) { auto* ctx = get_default_webgpu_context(); if (ctx) { @@ -337,6 +386,10 @@ void WebGPUGraph::build( const WebGPUContext* kv_ctx = get_default_webgpu_context(); kv_f16_ = f16_kv_cache && (kv_ctx != nullptr && kv_ctx->shader_f16_supported); + // f16-accumulate q4gsw steel prefill GEMM (runtime opt-in). QuantizedLinear + // additionally gates the kernel on the negotiated shader-f16 feature. + f16_accumulate_gemm_ = f16_accumulate_gemm; + // Phase 1: Create all values const auto* values = graph->values(); const int num_vals = values ? values->size() : 0; diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 7e7d13fde88..66f0e401de5 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -105,7 +105,8 @@ class WebGPUGraph { const void* flatbuffer_data, const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map = nullptr, - bool f16_kv_cache = false); + bool f16_kv_cache = false, + bool f16_accumulate_gemm = false); // Copy input tensor data from host pointers into GPU buffers. void copy_inputs(const std::vector& inputs); @@ -268,6 +269,35 @@ class WebGPUGraph { // Graph-owned scratch storage buffer for fused-op intermediates (e.g. SDPA). WGPUBuffer create_scratch_buffer(size_t nbytes); + // Reusable scratch pool for SINGLE-OP-LIFETIME fused-op scratch (SDPA + // attn_weights/softmax, FlashDecoding partials). acquire_scratch() reuses a + // free slot (best-fit, size in [n,2n]) or creates one; the caller RELEASES it + // at op-lowering scope exit (use ScopedScratch), so N layers' scratch reuses + // a small constant of buffers instead of N x held to graph teardown. + // Correctness: WebGPU/Dawn auto-inserts RAW hazard barriers between + // dispatches on a shared storage buffer regardless of pass structure -- the + // SAME guarantee mem_obj_id aliasing already relies on -- so reuse is + // bit-identical. Never hand a still-in_use slot to a co-live requester. + WGPUBuffer acquire_scratch(size_t nbytes); + void release_scratch(WGPUBuffer buffer); + // RAII: releases an acquired scratch slot when the op-lowering scope exits + // (leak-safe vs early returns). + struct ScopedScratch { + WebGPUGraph* g = nullptr; + WGPUBuffer buf = nullptr; + ScopedScratch(WebGPUGraph* graph, WGPUBuffer b) : g(graph), buf(b) {} + ~ScopedScratch() { + if (g && buf) { + g->release_scratch(buf); + } + } + ScopedScratch(const ScopedScratch&) = delete; + ScopedScratch& operator=(const ScopedScratch&) = delete; + operator WGPUBuffer() const { + return buf; + } + }; + // Create a mapped-at-creation uniform buffer from `size` bytes and track it // in the memory stats. Shared helper for ops needing a uniform Params buffer. WGPUBuffer make_uniform_buffer(const void* data, size_t size); @@ -321,9 +351,16 @@ class WebGPUGraph { return kv_f16_; } + // True when the q4gsw steel prefill GEMM uses the lossy f16-accumulate kernel + // (runtime opt-in; perplexity-gated, not bit-exact). + bool f16_accumulate_gemm() const { + return f16_accumulate_gemm_; + } + private: bool kv_f16_ = false; std::unordered_set kv_cache_ids_; + bool f16_accumulate_gemm_ = false; private: WGPUInstance instance_ = nullptr; @@ -377,6 +414,16 @@ class WebGPUGraph { // Long-lived scratch storage buffers for fused ops (e.g. SDPA temporaries). std::vector scratch_buffers_; + // Reusable scratch pool: single-op-lifetime buffers recycled across ops + // (acquire_scratch/release_scratch). Each slot is freed in the dtor. See + // acquire_scratch() for the reuse policy. + struct ScratchSlot { + WGPUBuffer buffer = nullptr; + size_t size = 0; + bool in_use = false; + }; + std::vector scratch_pool_; + // Uniform buffers owned for the graph's lifetime; released in the dtor. std::vector owned_uniform_buffers_; diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 20f7024ca82..b0728764310 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include #include #include #include @@ -270,7 +272,24 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { if (use_steel) { const WebGPUContext* ctx = get_default_webgpu_context(); if (ctx != nullptr && ctx->shader_f16_supported) { - shader_src = kQ4gswLinearGemmSteelHalfWGSL; + // Packed-word dequant: bit-exact to the steel `half` kernel but loads + // each u32 weight word once + hoists the per-column scale (half re-reads + // them ~8x/~16x). Needs group_size % BK == 0 so the hoisted scale is + // constant across the BK tile; else the per-nibble `half` kernel. + shader_src = (gs % kQ4gswSteelBK == 0u) + ? kQ4gswLinearGemmSteelHalfPwdqWGSL + : kQ4gswLinearGemmSteelHalfWGSL; + } + } + // f16-accumulate: pwdq staging with an f16 register accumulator. + // Lossy (f16 accumulate over K) -> opt-in via the enable_f16_accumulate_gemm + // runtime spec (default off), gated on the negotiated shader-f16 feature and + // group_size % BK == 0 (same hoisted-scale requirement as pwdq). Overrides + // the f32-accumulate steel kernels. + if (use_steel && graph.f16_accumulate_gemm() && (gs % kQ4gswSteelBK == 0u)) { + const WebGPUContext* ctx = get_default_webgpu_context(); + if (ctx != nullptr && ctx->shader_f16_supported) { + shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; } } const uint32_t workgroup_count = compute_q4gsw_workgroup_count( 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 788c9ffd941..4d7ab0b1d1e 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 @@ -21,7 +21,16 @@ struct Params { // "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. // The "steel" name + register-tiled dequant-to-shared GEMM structure are // inspired by MLX's steel GEMM kernels (github.com/ml-explore/mlx, -// mlx/backend/metal/kernels/steel). +// mlx/backend/metal/kernels/steel). One template, four variants: +// DTYPE=float f32 storage/multiply, per-nibble weight staging. +// DTYPE=half f16 storage/multiply, per-nibble weight staging. +// PWDQ (half only) packed-word dequant: load each u32 weight word ONCE, +// unpack all 16 nibbles of a column + hoist the per-column scale to one read +// (the per-nibble path re-reads each word ~8x). Requires K%BK==0 (steel +// route guarantees it) and group_size%BK==0 (hoisted scale across the tile). +// ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue +// -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 +// accumulate -- BIT-EXACT to the per-nibble half kernel. 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 @@ -34,16 +43,17 @@ fn main(@builtin(workgroup_id) wid: vec3, let row0 = by * BM; let col0 = bx * BN; let tid = lid.y * 16u + lid.x; - var acc: array, 4>; + var acc: array, 4>; for (var m: u32 = 0u; m < 4u; m = m + 1u) { - for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = ${"0.0h" if ACC == "half" else "0.0"}; } } // A staging coords: 256 threads load 64x16 = 1024 f32 -> 4 rows each (4 contiguous K). let ar = tid / 4u; // 0..63 (row in tile) let ac = (tid % 4u) * 4u; // 0,4,8,12 (K offset, 4 contiguous) - // B staging coords: 256 threads load 16x64 = 1024 dequant weights -> 4 cols each. - let br = tid / 16u; // 0..15 (K within BK) - let bc = (tid % 16u) * 4u; // 0,4,..60 (N offset, 4 contiguous) + $if not PWDQ: + // B staging coords: 256 threads load 16x64 = 1024 dequant weights -> 4 cols each. + let br = tid / 16u; // 0..15 (K within BK) + let bc = (tid % 16u) * 4u; // 0,4,..60 (N offset, 4 contiguous) var k0: u32 = 0u; loop { @@ -57,28 +67,53 @@ fn main(@builtin(workgroup_id) wid: vec3, As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 2u]); As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 3u]); } else { - As[ar * BK + ac + 0u] = 0.0; As[ar * BK + ac + 1u] = 0.0; - As[ar * BK + ac + 2u] = 0.0; As[ar * BK + ac + 3u] = 0.0; + As[ar * BK + ac + 0u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 1u] = ${"0.0h" if PWDQ else "0.0"}; + As[ar * BK + ac + 2u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 3u] = ${"0.0h" if PWDQ else "0.0"}; } - // stage DEQUANTIZED weights into Bs[k][n]: 4 contiguous N per thread. - let kk = k0 + br; // K index for this shmem row - let scale_row = (kk / params.group_size) * params.padded_N; - for (var j: u32 = 0u; j < 4u; j = j + 1u) { - let n = col0 + bc + j; - var dqv: ${buffer_scalar_type(DTYPE)} = 0.0; - if (n < params.N) { - let byte_idx = n * params.K_packed + (kk >> 1u); - let word = t_weight[byte_idx >> 2u]; - let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; - var nib: u32; - if ((kk & 1u) == 0u) { nib = b & 0x0Fu; } else { nib = (b >> 4u) & 0x0Fu; } - $if DTYPE == "half": - dqv = f16(i32(nib) - 8) * f16(t_scales[scale_row + n]); - $else: - dqv = f32(i32(nib) - 8) * t_scales[scale_row + n]; + $if PWDQ: + // Packed-word dequant: threads [0,BN) each stage one full BK-column of Bs. + if (tid < BN) { + let c = tid; // Bs column within this tile + let n = col0 + c; // global output column + if (n < params.N) { + // Scale is constant across the BK tile (group_size % BK == 0 for all real + // group sizes; K%BK==0 on the steel route), so hoist it to one read. + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + // Column n's 16-nibble K-slice for this tile = two consecutive words. + // K_packed multiple of 8 => base_word stays inside column n's own region. + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word]; + let w1 = t_weight[base_word + 1u]; + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = select(w1, w0, br < 8u); // word0 holds K-slice [0,8) + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + } else { + for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } + } + } + $else: + // stage DEQUANTIZED weights into Bs[k][n]: 4 contiguous N per thread. + let kk = k0 + br; // K index for this shmem row + let scale_row = (kk / params.group_size) * params.padded_N; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let n = col0 + bc + j; + var dqv: ${buffer_scalar_type(DTYPE)} = 0.0; + if (n < params.N) { + let byte_idx = n * params.K_packed + (kk >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + if ((kk & 1u) == 0u) { nib = b & 0x0Fu; } else { nib = (b >> 4u) & 0x0Fu; } + $if DTYPE == "half": + dqv = f16(i32(nib) - 8) * f16(t_scales[scale_row + n]); + $else: + dqv = f32(i32(nib) - 8) * t_scales[scale_row + n]; + } + Bs[br * BN + bc + j] = dqv; } - Bs[br * BN + bc + j] = dqv; - } workgroupBarrier(); for (var k: u32 = 0u; k < BK; k = k + 1u) { var a: array<${buffer_scalar_type(DTYPE)}, 4>; @@ -86,7 +121,9 @@ fn main(@builtin(workgroup_id) wid: vec3, for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; } for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; } for (var m: u32 = 0u; m < 4u; m = m + 1u) { - $if DTYPE == "half": + $if ACC == "half": + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + $elif DTYPE == "half": for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + f32(a[m] * bvec[n]); } $else: for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + a[m] * bvec[n]; } @@ -100,7 +137,7 @@ fn main(@builtin(workgroup_id) wid: vec3, let r = row0 + lid.y * 4u + m; let c = col0 + lid.x * 4u + n; if (r < params.M && c < params.N) { - var v = acc[m][n]; + var v = ${"f32(acc[m][n])" if ACC == "half" else "acc[m][n]"}; if (params.has_bias != 0u) { v = v + t_bias[c]; } t_out[r * params.N + c] = v; } diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml index 1505d8b9924..5a2cae5e499 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml @@ -1,11 +1,22 @@ q4gsw_linear_gemm_steel: parameter_names_with_default_values: DTYPE: float - generate_variant_forall: - DTYPE: - - VALUE: float - SUFFIX: "" - - VALUE: half - SUFFIX: half + PWDQ: false + ACC: float shader_variants: - NAME: q4gsw_linear_gemm_steel + DTYPE: float + PWDQ: false + ACC: float + - NAME: q4gsw_linear_gemm_steel_half + DTYPE: half + PWDQ: false + ACC: float + - NAME: q4gsw_linear_gemm_steel_half_pwdq + DTYPE: half + PWDQ: true + ACC: float + - NAME: q4gsw_linear_gemm_steel_half_pwdq_f16acc + DTYPE: half + PWDQ: true + ACC: half 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 new file mode 100644 index 00000000000..efefd7edce1 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h @@ -0,0 +1,141 @@ +/* + * 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_gemm_steel.wgsl - DO NOT EDIT. +// wgsl-sha256: 36b3d3f9dd08a529909c13ec7d66cd0cf392c347ca047a4d38453b3c295f72ce +inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"( +enable f16; +@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; + +// "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. +// The "steel" name + register-tiled dequant-to-shared GEMM structure are +// inspired by MLX's steel GEMM kernels (github.com/ml-explore/mlx, +// mlx/backend/metal/kernels/steel). One template, four variants: +// DTYPE=float f32 storage/multiply, per-nibble weight staging. +// DTYPE=half f16 storage/multiply, per-nibble weight staging. +// PWDQ (half only) packed-word dequant: load each u32 weight word ONCE, +// unpack all 16 nibbles of a column + hoist the per-column scale to one read +// (the per-nibble path re-reads each word ~8x). Requires K%BK==0 (steel +// route guarantees it) and group_size%BK==0 (hoisted scale across the tile). +// ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue +// -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 +// accumulate -- BIT-EXACT to the per-nibble half kernel. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; +var As: array; // BM*BK +var Bs: array; // BK*BN +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let nbN = (params.N + BN - 1u) / BN; + let bx = wid.x % nbN; // decode 2D tile id from 1D dispatch + let by = wid.x / nbN; + let row0 = by * BM; + let col0 = bx * BN; + let tid = lid.y * 16u + lid.x; + var acc: array, 4>; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0h; } + } + // A staging coords: 256 threads load 64x16 = 1024 f32 -> 4 rows each (4 contiguous K). + let ar = tid / 4u; // 0..63 (row in tile) + let ac = (tid % 4u) * 4u; // 0,4,8,12 (K offset, 4 contiguous) + + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + // stage activations (edge-masked on M; K is a multiple of BK for our shapes) + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + As[ar * BK + ac + 0u] = f16(t_input[base]); + As[ar * BK + ac + 1u] = f16(t_input[base + 1u]); + As[ar * BK + ac + 2u] = f16(t_input[base + 2u]); + As[ar * BK + ac + 3u] = f16(t_input[base + 3u]); + } else { + As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; + As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + } + // Packed-word dequant: threads [0,BN) each stage one full BK-column of Bs. + if (tid < BN) { + let c = tid; // Bs column within this tile + let n = col0 + c; // global output column + if (n < params.N) { + // Scale is constant across the BK tile (group_size % BK == 0 for all real + // group sizes; K%BK==0 on the steel route), so hoist it to one read. + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + // Column n's 16-nibble K-slice for this tile = two consecutive words. + // K_packed multiple of 8 => base_word stays inside column n's own region. + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word]; + let w1 = t_weight[base_word + 1u]; + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = select(w1, w0, br < 8u); // word0 holds K-slice [0,8) + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + } else { + for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < BK; k = k + 1u) { + var a: array; + var bvec: array; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + } + } + workgroupBarrier(); + k0 = k0 + BK; + } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { + let r = row0 + lid.y * 4u + m; + let c = col0 + lid.x * 4u + n; + if (r < params.M && c < params.N) { + var v = f32(acc[m][n]); + if (params.has_bias != 0u) { v = v + t_bias[c]; } + t_out[r * params.N + c] = v; + } + } + } +} +)"; + +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeX = + 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY = + 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu 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 new file mode 100644 index 00000000000..46057de2340 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h @@ -0,0 +1,139 @@ +/* + * 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_gemm_steel.wgsl - DO NOT EDIT. +// wgsl-sha256: 1f916bcd30dbbbcc7eca37e795ecc26e3c72e645ccd2c361fa0ac4e66f1a174a +inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = R"( +enable f16; +@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; + +// "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. +// The "steel" name + register-tiled dequant-to-shared GEMM structure are +// inspired by MLX's steel GEMM kernels (github.com/ml-explore/mlx, +// mlx/backend/metal/kernels/steel). One template, four variants: +// DTYPE=float f32 storage/multiply, per-nibble weight staging. +// DTYPE=half f16 storage/multiply, per-nibble weight staging. +// PWDQ (half only) packed-word dequant: load each u32 weight word ONCE, +// unpack all 16 nibbles of a column + hoist the per-column scale to one read +// (the per-nibble path re-reads each word ~8x). Requires K%BK==0 (steel +// route guarantees it) and group_size%BK==0 (hoisted scale across the tile). +// ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue +// -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 +// accumulate -- BIT-EXACT to the per-nibble half kernel. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; +var As: array; // BM*BK +var Bs: array; // BK*BN +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let nbN = (params.N + BN - 1u) / BN; + let bx = wid.x % nbN; // decode 2D tile id from 1D dispatch + let by = wid.x / nbN; + let row0 = by * BM; + let col0 = bx * BN; + let tid = lid.y * 16u + lid.x; + var acc: array, 4>; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0; } + } + // A staging coords: 256 threads load 64x16 = 1024 f32 -> 4 rows each (4 contiguous K). + let ar = tid / 4u; // 0..63 (row in tile) + let ac = (tid % 4u) * 4u; // 0,4,8,12 (K offset, 4 contiguous) + + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + // stage activations (edge-masked on M; K is a multiple of BK for our shapes) + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + As[ar * BK + ac + 0u] = f16(t_input[base]); + As[ar * BK + ac + 1u] = f16(t_input[base + 1u]); + As[ar * BK + ac + 2u] = f16(t_input[base + 2u]); + As[ar * BK + ac + 3u] = f16(t_input[base + 3u]); + } else { + As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; + As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + } + // Packed-word dequant: threads [0,BN) each stage one full BK-column of Bs. + if (tid < BN) { + let c = tid; // Bs column within this tile + let n = col0 + c; // global output column + if (n < params.N) { + // Scale is constant across the BK tile (group_size % BK == 0 for all real + // group sizes; K%BK==0 on the steel route), so hoist it to one read. + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + // Column n's 16-nibble K-slice for this tile = two consecutive words. + // K_packed multiple of 8 => base_word stays inside column n's own region. + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word]; + let w1 = t_weight[base_word + 1u]; + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = select(w1, w0, br < 8u); // word0 holds K-slice [0,8) + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + } else { + for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < BK; k = k + 1u) { + var a: array; + var bvec: array; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + f32(a[m] * bvec[n]); } + } + } + workgroupBarrier(); + k0 = k0 + BK; + } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { + let r = row0 + lid.y * 4u + m; + let c = col0 + lid.x * 4u + n; + if (r < params.M && c < params.N) { + var v = acc[m][n]; + if (params.has_bias != 0u) { v = v + t_bias[c]; } + t_out[r * params.N + c] = v; + } + } + } +} +)"; + +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu 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 7e9363e3b36..ae03f63fb5c 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: e3c21e7db7c18f6e085de71e283988f0bd3b2543807ddc17774a1c607e69c766 +// wgsl-sha256: 00cdd2f2fb98a5c7343d16fdf7e59f1b840e180cec3f82bf9b569513c0a45396 inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -37,7 +37,16 @@ struct Params { // "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. // The "steel" name + register-tiled dequant-to-shared GEMM structure are // inspired by MLX's steel GEMM kernels (github.com/ml-explore/mlx, -// mlx/backend/metal/kernels/steel). +// mlx/backend/metal/kernels/steel). One template, four variants: +// DTYPE=float f32 storage/multiply, per-nibble weight staging. +// DTYPE=half f16 storage/multiply, per-nibble weight staging. +// PWDQ (half only) packed-word dequant: load each u32 weight word ONCE, +// unpack all 16 nibbles of a column + hoist the per-column scale to one read +// (the per-nibble path re-reads each word ~8x). Requires K%BK==0 (steel +// route guarantees it) and group_size%BK==0 (hoisted scale across the tile). +// ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue +// -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 +// accumulate -- BIT-EXACT to the per-nibble half kernel. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN 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 9e73a0c9a66..71b0f45bdf7 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: 43536b16026d3b62d77087b86289885606c04834d9783c3c871512d6789ee6f6 +// wgsl-sha256: dd771b9ab096410f3ad0d9259bef7816e41330a434325ea28baa5abbfb2841d2 inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_input: array; @@ -36,7 +36,16 @@ struct Params { // "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. // The "steel" name + register-tiled dequant-to-shared GEMM structure are // inspired by MLX's steel GEMM kernels (github.com/ml-explore/mlx, -// mlx/backend/metal/kernels/steel). +// mlx/backend/metal/kernels/steel). One template, four variants: +// DTYPE=float f32 storage/multiply, per-nibble weight staging. +// DTYPE=half f16 storage/multiply, per-nibble weight staging. +// PWDQ (half only) packed-word dequant: load each u32 weight word ONCE, +// unpack all 16 nibbles of a column + hoist the per-column scale to one read +// (the per-nibble path re-reads each word ~8x). Requires K%BK==0 (steel +// route guarantees it) and group_size%BK==0 (hoisted scale across the tile). +// ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue +// -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 +// accumulate -- BIT-EXACT to the per-nibble half kernel. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 0b951717775..50321ba4bdf 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -481,8 +481,11 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { } // QK/softmax scratch — allocated only on the non-FD path (Hq*S*Cmax prefill). - WGPUBuffer attn_weights = graph.create_scratch_buffer(aw_bytes); - WGPUBuffer attn_weights_softmax = graph.create_scratch_buffer(aw_bytes); + WGPUBuffer attn_weights = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_guard(&graph, attn_weights); + WGPUBuffer attn_weights_softmax = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_softmax_guard( + &graph, attn_weights_softmax); // --- Dispatch 3: QK -> attn_weights. One thread per TM x TN tile. { diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index a443beb2149..ffd3b24dce3 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -186,8 +186,10 @@ void sdpa_fd_decode_dispatch( static_cast(kSdpaFdMaxSplits) * static_cast(D); const uint64_t pml_floats = static_cast(Hq) * static_cast(kSdpaFdMaxSplits) * 2ull; - WGPUBuffer part_o = graph.create_scratch_buffer(po_floats * sizeof(float)); - WGPUBuffer part_ml = graph.create_scratch_buffer(pml_floats * sizeof(float)); + WGPUBuffer part_o = graph.acquire_scratch(po_floats * sizeof(float)); + WebGPUGraph::ScopedScratch part_o_guard(&graph, part_o); + WGPUBuffer part_ml = graph.acquire_scratch(pml_floats * sizeof(float)); + WebGPUGraph::ScopedScratch part_ml_guard(&graph, part_ml); // Pass 1: split (Hq*num_splits WGs) -> writes part_o, part_ml. FdSplitParams sp = {}; diff --git a/backends/webgpu/test/native/test_scratch_buffer.cpp b/backends/webgpu/test/native/test_scratch_buffer.cpp index 98cf3648c6b..1a8f4fcd96f 100644 --- a/backends/webgpu/test/native/test_scratch_buffer.cpp +++ b/backends/webgpu/test/native/test_scratch_buffer.cpp @@ -6,7 +6,8 @@ * LICENSE file in the root directory of this source tree. */ -// White-box unit tests for WebGPUGraph::create_scratch_buffer. +// White-box unit tests for WebGPUGraph scratch buffers: +// create_scratch_buffer and the acquire_scratch/release_scratch reuse pool. #include #include @@ -19,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -222,6 +224,59 @@ TEST(ScratchBuffer, Tier3Lifecycle) { } // each graph's dtor releases its 256 buffers here } +// Tier 4: reuse-pool semantics (acquire_scratch / release_scratch / +// ScopedScratch). The pool recycles single-op-lifetime scratch across ops so N +// layers reuse a small constant of buffers instead of N x. + +// A released slot is handed back on the next same-size acquire (the reuse win). +TEST(ScratchPool, ReuseAfterRelease) { + WebGPUGraph g; + g.set_device(g_device); + WGPUBuffer a = g.acquire_scratch(64 * sizeof(float)); + g.release_scratch(a); + WGPUBuffer b = g.acquire_scratch(64 * sizeof(float)); + EXPECT_EQ(a, b) << "released slot should be reused for a same-size request"; +} + +// A still-in_use slot is never handed to a co-live requester (RAW-safety). +TEST(ScratchPool, NoReuseWhileInUse) { + WebGPUGraph g; + g.set_device(g_device); + WGPUBuffer a = g.acquire_scratch(64 * sizeof(float)); + WGPUBuffer b = g.acquire_scratch(64 * sizeof(float)); // a not released + EXPECT_TRUE(a && b && a != b) << "co-live acquires must be distinct buffers"; +} + +// Best-fit 2x cap: a large free slot must not back a much smaller request, but +// a request it does fit (size in [n, 2n]) reuses it. +TEST(ScratchPool, BestFitSizeCap) { + WebGPUGraph g; + g.set_device(g_device); + WGPUBuffer big = g.acquire_scratch(1024 * sizeof(float)); + g.release_scratch(big); + // 1024*4 bytes is outside [4, 8], so the big slot is ineligible for 4 bytes. + WGPUBuffer tiny = g.acquire_scratch(4); + EXPECT_NE(big, tiny) + << "oversized slot must not back a tiny request (2x cap)"; + g.release_scratch(tiny); + WGPUBuffer same = g.acquire_scratch(1024 * sizeof(float)); + EXPECT_EQ(big, same) << "an in-range request should reuse the big slot"; +} + +// ScopedScratch releases its slot at scope exit, so the next acquire reuses it. +TEST(ScratchPool, ScopedScratchReleasesOnScopeExit) { + WebGPUGraph g; + g.set_device(g_device); + WGPUBuffer first = nullptr; + { + WebGPUGraph::ScopedScratch s(&g, g.acquire_scratch(64 * sizeof(float))); + first = s; // operator WGPUBuffer + EXPECT_NE(first, nullptr); + } // s releases the slot here + WGPUBuffer second = g.acquire_scratch(64 * sizeof(float)); + EXPECT_EQ(first, second) << "slot freed by ScopedScratch should be reused"; +} + int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); diff --git a/backends/webgpu/test/ops/test_quantized_linear.py b/backends/webgpu/test/ops/test_quantized_linear.py index e9a785dcd6d..72945c37d6d 100644 --- a/backends/webgpu/test/ops/test_quantized_linear.py +++ b/backends/webgpu/test/ops/test_quantized_linear.py @@ -70,6 +70,19 @@ class Q4gswConfig: # Partial M and N steel tiles under the f16 kernel; exercises f16 boundary # masking (the exact-N "steel_f16" shape does not). N%8==0, steel-isolating. Q4gswConfig("steel_f16_edge", 70, 1024, 136), # f16 partial-tile + # pwdq (packed-word dequant) backs the f16 steel path at group_size % BK(16) + # == 0 (bit-exact to steel_half; steel_f16 above runs it at gs=32). These lock + # the gs gate at group sizes those omit: gs=64 stays on pwdq; gs=8 (< BK) falls + # back to the per-nibble steel_half kernel (its hoisted-per-BK scale is invalid + # there). Same fp64 golden regardless of which kernel runs. + Q4gswConfig("pwdq_gs64", 96, 2048, 256, group_size=64), # pwdq, non-32 group + Q4gswConfig("pwdq_gs8", 96, 2048, 256, group_size=8), # steel_half fallback + # pwdqf16acc (f16-accumulate) runs when the enable_f16_accumulate_gemm runtime + # spec is set and gs % BK == 0 (perplexity-gated; see the kernel diff). Same + # .pte as the f32 configs -- only the accumulator dtype differs -- goldened at a + # looser f16-accumulate tol in the native test; deep-K stresses the worst case. + Q4gswConfig("pwdqf16acc", 96, 2048, 256), # f16-accumulate steel (runtime) + Q4gswConfig("pwdqf16acc_down", 128, 8192, 2048), # deep-K f16-accum worst case Q4gswConfig("gate_proj_pf", 128, 2048, 8192), # gate/up prefill (shmem via N) Q4gswConfig("down_proj_pf", 128, 8192, 2048), # down prefill (shmem via K) Q4gswConfig("shmem_edge", 130, 4096, 2056), # partial 32-tile bounds diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index 5f3d6d788ec..fbdfbd09076 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -11,6 +11,8 @@ #include #include #include +#include +#include #include @@ -229,6 +231,15 @@ bool sdpa_within_tol( int n, float* ma, float* mr) { + float atol = 1e-4f, rtol = 1e-3f; + // f16 KV (runtime opt-in) reads K/V at reduced precision; loosen the tol on a + // shader-f16 device to cover that rounding. Harmless for f32 KV (looser + // gate). + const WebGPUContext* kv_ctx = get_default_webgpu_context(); + if (kv_ctx != nullptr && kv_ctx->shader_f16_supported) { + atol = 2e-3f; + rtol = 1e-2f; + } float max_abs = 0.0f, max_rel = 0.0f; bool ok = true; for (int i = 0; i < n; i++) { @@ -236,7 +247,7 @@ bool sdpa_within_tol( const float re = ae / std::max(std::abs(golden[i]), 1e-6f); max_abs = std::max(max_abs, ae); max_rel = std::max(max_rel, re); - if (ae > 1e-4f && re > 1e-3f) { + if (ae > atol && re > rtol) { ok = false; } } @@ -284,6 +295,21 @@ const Q4gswConfig kQ4gswConfigs[] = { {"steel_f16", 96, 2048, 256, 2.3e-4f, 1e-3f, true, false}, // Partial M and N steel tiles under the f16 kernel (f16 boundary masking). {"steel_f16_edge", 70, 1024, 136, 2.3e-4f, 1e-3f, true, false}, + // pwdq (packed-word dequant) backs the f16 steel path at group_size % BK == + // 0 + // (bit-exact to steel_half; the steel_f16 configs above run it at gs=32). + // These lock the gs gate at group sizes those omit: gs=64 stays on pwdq; + // gs=8 (< BK=16) falls back to the per-nibble steel_half kernel. + {"pwdq_gs64", 96, 2048, 256, 2.3e-4f, 1e-3f, true, false}, + {"pwdq_gs8", 96, 2048, 256, 2.3e-4f, 1e-3f, true, false}, + // f16-ACCUMULATE steel (pwdqf16acc): lossy, so a wider gate than the + // f16-multiply steel_f16 (2.3e-4). f16 accumulation error grows with K, so + // the deep-K down shape (K=8192) gets the loosest tol. Perplexity is the + // primary quality gate (see the kernel diff); this catches gross bit/index + // bugs. gs=32 (% BK == 0) selects pwdqf16acc; the sweep loads these rows + // with the enable_f16_accumulate_gemm runtime spec set. + {"pwdqf16acc", 96, 2048, 256, 2e-2f, 3e-2f, true, false}, + {"pwdqf16acc_down", 128, 8192, 2048, 5e-2f, 8e-2f, true, false}, {"gate_proj_pf", 128, 2048, 8192, 1e-4f, 1e-3f, true, false}, // shmem via N {"down_proj_pf", 128, 8192, 2048, 1e-3f, 1e-2f, true, false}, // shmem via K {"shmem_edge", 130, 4096, 2056, 1e-4f, 1e-3f, true, false}, // partial tiles @@ -547,7 +573,18 @@ void test_q4gsw_config( cfg.n); Module module(pte); - ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load " << pte; + // pwdqf16acc rows exercise the lossy f16-accumulate kernel, a runtime opt-in + // (default off); enable it via the backend option keyed by the registered id. + if (std::string(cfg.name).rfind("pwdqf16acc", 0) == 0) { + BackendOptions<1> opts; + opts.set_option("enable_f16_accumulate_gemm", true); + LoadBackendOptionsMap map; + ASSERT_EQ(map.set_options("VulkanBackend", opts.view()), Error::Ok); + ASSERT_EQ(module.load_forward(nullptr, nullptr, &map), Error::Ok) + << "could not load " << pte; + } else { + ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load " << pte; + } const int in_numel = cfg.m * cfg.k; const int out_numel = cfg.m * cfg.n;