From f32a6ab985156cf4f8b29f00578ae1bfce3ec7bd Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:11:04 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/runtime/WebGPUUtils.h | 22 ++ .../ops/quantized_linear/QuantizedLinear.cpp | 155 +++++++++--- .../quantized_linear/q4gsw_steel_bk64.wgsl | 117 +++++++++ .../quantized_linear/q4gsw_steel_bk64_wgsl.h | 141 +++++++++++ .../webgpu/test/native/test_dispatch_2d.cpp | 44 ++++ .../webgpu/test/native/test_dynamic_shape.cpp | 224 +++++++++++++++++- .../test_dynamic_shape_export.py | 161 +++++++++++++ 7 files changed, 830 insertions(+), 34 deletions(-) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64.wgsl create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index da37c7c0438..a623e923762 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -82,6 +82,28 @@ constexpr bool should_record_sdpa_dual_route( return fd_eligible && has_dynamic_sequence; } +constexpr bool is_q4gsw_bk64_eligible( + uint32_t k, + uint32_t n, + uint32_t group_size, + bool has_bias, + bool shader_f16_supported, + uint32_t max_invocations, + uint32_t max_workgroup_storage_bytes) { + constexpr uint32_t kRequiredInvocations = 256u; + constexpr uint32_t kRequiredStorageBytes = 2u * 64u * 64u * sizeof(uint16_t); + const bool ordinary_llama_projection = (k == 2048u && n == 8192u) || + (k == 8192u && n == 2048u) || (k == 2048u && n == 2048u); + return ordinary_llama_projection && k % 64u == 0u && group_size == 64u && + !has_bias && shader_f16_supported && + max_invocations >= kRequiredInvocations && + max_workgroup_storage_bytes >= kRequiredStorageBytes; +} + +constexpr bool is_q4gsw_bk64_live_m(uint32_t m) { + return m == 128u || m == 508u || m == 512u; +} + class DispatchRouteRegistry { public: template diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index ffeeffae5d1..9ed995fb748 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include @@ -60,6 +61,7 @@ constexpr uint32_t kQ4gswShmemNMinDim = 2048u; // steel GEMM: 64x64 tile, 256 threads (16x16), fixed wg (no override). constexpr uint32_t kQ4gswSteelTile = 64u; constexpr uint32_t kQ4gswSteelBK = 16u; +constexpr uint32_t kQ4gswSteelBK64 = 64u; constexpr uint32_t kQ4gswSteelInvocations = 256u; // One workgroup per (tile_m x tile_n) tile, no grid-stride: throw when the tile @@ -90,6 +92,27 @@ bool steel_supported(WGPUDevice device) { limits.maxComputeInvocationsPerWorkgroup >= kQ4gswSteelInvocations; } +bool steel_bk64_eligible( + WGPUDevice device, + uint32_t K, + uint32_t N, + uint32_t group_size, + bool has_bias) { + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(device, &limits) != WGPUStatus_Success) { + return false; + } + const WebGPUContext* context = get_default_webgpu_context(); + return utils::is_q4gsw_bk64_eligible( + K, + N, + group_size, + has_bias, + context != nullptr && context->shader_f16_supported, + limits.maxComputeInvocationsPerWorkgroup, + limits.maxComputeWorkgroupStorageSize); +} + // Not grid-strided: 0 (fall back) when K%BK != 0 or over the 1D dispatch limit. uint32_t steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { @@ -103,6 +126,17 @@ steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { return (total == 0u || total > max_count) ? 0u : static_cast(total); } +uint32_t steel_bk64_workgroup_count( + WGPUDevice device, + uint32_t m, + uint32_t n, + uint32_t K) { + if (K % kQ4gswSteelBK64 != 0u) { + return 0u; + } + return steel_workgroup_count(device, m, n, K); +} + // Workgroup count for a linear_q4gsw dispatch (bicol GEMV / shmem GEMM / tiled // GEMM), with the range/limit guards shared by the build-time path and the // resize hook. use_gemv/use_shmem_gemm are the build-time routing decision (the @@ -110,10 +144,12 @@ steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { uint32_t compute_q4gsw_workgroup_count( WGPUDevice device, bool use_gemv, + bool use_bk64, bool use_steel, bool use_shmem_gemm, uint32_t m, uint32_t n, + uint32_t K, uint32_t wg_size, const char* op_name) { if (use_gemv) { @@ -132,6 +168,14 @@ uint32_t compute_q4gsw_workgroup_count( } return wgc; } + if (use_bk64) { + const uint32_t count = steel_bk64_workgroup_count(device, m, n, K); + if (count == 0u) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": invalid BK64 dispatch"); + } + return count; + } if (use_steel) { // steel: one workgroup per 64x64 tile. Over-limit THROWS here -- unlike the // build-time steel_workgroup_count, which returns 0 so the caller falls @@ -170,7 +214,8 @@ struct Q4gswExecutionState { }; constexpr size_t kQ4gswBicolRoute = 0; -constexpr size_t kQ4gswPrefillRoute = 1; +constexpr size_t kQ4gswBk64Route = 1; +constexpr size_t kQ4gswPrefillRoute = 2; Q4gswExecutionState make_q4gsw_execution_state( WGPUDevice device, @@ -185,6 +230,8 @@ Q4gswExecutionState make_q4gsw_execution_state( uint32_t wg_size, bool use_single_gemv, bool use_dual_route, + bool record_bk64_route, + bool bk64_eligible, bool prefill_use_steel, bool prefill_use_shmem_gemm) { if (input_dims.empty()) { @@ -205,13 +252,18 @@ Q4gswExecutionState make_q4gsw_execution_state( } const uint32_t m = static_cast(live_m); const bool use_gemv = use_single_gemv || (use_dual_route && m == 1u); + const bool use_bk64 = !use_gemv && bk64_eligible && + utils::is_q4gsw_bk64_live_m(m) && + steel_bk64_workgroup_count(device, m, N, K) > 0u; const uint32_t workgroup_count = compute_q4gsw_workgroup_count( device, use_gemv, - !use_gemv && prefill_use_steel, - !use_gemv && prefill_use_shmem_gemm, + use_bk64, + !use_gemv && !use_bk64 && prefill_use_steel, + !use_gemv && !use_bk64 && prefill_use_shmem_gemm, m, N, + K, wg_size, "linear_q4gsw(resize)"); @@ -225,8 +277,12 @@ Q4gswExecutionState make_q4gsw_execution_state( state.params.has_bias = has_bias; state.output_dims = input_dims; state.output_dims.back() = static_cast(N); - state.active_route = - use_dual_route ? (use_gemv ? kQ4gswBicolRoute : kQ4gswPrefillRoute) : 0u; + state.active_route = use_dual_route + ? (use_gemv ? kQ4gswBicolRoute + : (record_bk64_route + ? (use_bk64 ? kQ4gswBk64Route : kQ4gswPrefillRoute) + : 1u)) + : 0u; state.active_grid = {workgroup_count, 1u}; return state; } @@ -309,13 +365,38 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { "WebGPU linear_q4gsw: scales dims too small for K/N"); } - // M==1 -> bicol GEMV; M>1 -> steel GEMM (preferred) else shmem else tiled. + // Optional bias: real buffer if present, else a dummy for the fixed layout. + uint32_t has_bias = 0; + WGPUBuffer bias_buffer = nullptr; + uint64_t bias_size = 4; + if (graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor) { + const auto& bias = graph.get_tensor(bias_id); + if (bias.buffer == nullptr || bias.nbytes < N * sizeof(float)) { + throw std::runtime_error( + "WebGPU linear_q4gsw: bias present but null/undersized"); + } + has_bias = 1; + bias_buffer = bias.buffer; + bias_size = bias.nbytes; + } + if (bias_buffer == nullptr) { + bias_buffer = graph.create_scratch_buffer(4); + } + + // M==1 -> bicol GEMV; M>1 -> BK64 for exact Llama projections/M values, + // otherwise steel GEMM (preferred), shmem, or tiled. const uint32_t wg_size = utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX); const bool bicol_eligible = K % 8u == 0u && gs % 8u == 0u; const bool use_gemv = M == 1u && bicol_eligible; const bool use_dual_route = utils::should_record_q4gsw_dual_route( M, bicol_eligible, graph.has_dynamic_shapes()); + const bool bk64_eligible = + steel_bk64_eligible(device, K, N, gs, has_bias != 0u); + const bool record_bk64_route = use_dual_route && bk64_eligible && M >= 128u; + const bool use_bk64 = !use_gemv && bk64_eligible && + utils::is_q4gsw_bk64_live_m(M) && + steel_bk64_workgroup_count(device, M, N, K) > 0u; // GEMV (bicol) is a pow2 tree reduction; compute its size only when used. const uint32_t gemv_wg_size = (use_gemv || use_dual_route) ? utils::clamp_workgroup_size_pow2( @@ -334,6 +415,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL : kQ4gswLinearWGSL; const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL + : use_bk64 ? kQ4gswSteelBk64WGSL : use_steel ? kQ4gswLinearGemmSteelWGSL : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL : kQ4gswLinearWGSL; @@ -349,7 +431,9 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { prefill_shader_src = (gs % kQ4gswSteelBK == 0u) ? kQ4gswLinearGemmSteelHalfPwdqWGSL : kQ4gswLinearGemmSteelHalfWGSL; - shader_src = prefill_shader_src; + if (!use_bk64) { + shader_src = prefill_shader_src; + } } } // f16-accumulate: pwdq staging with an f16 register accumulator. @@ -361,28 +445,12 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const WebGPUContext* ctx = get_default_webgpu_context(); if (ctx != nullptr && ctx->shader_f16_supported) { prefill_shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; - shader_src = prefill_shader_src; + if (!use_bk64) { + shader_src = prefill_shader_src; + } } } - // Optional bias: real buffer if present, else a dummy for the fixed layout. - uint32_t has_bias = 0; - WGPUBuffer bias_buffer = nullptr; - uint64_t bias_size = 4; - if (graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor) { - const auto& bias = graph.get_tensor(bias_id); - if (bias.buffer == nullptr || bias.nbytes < N * sizeof(float)) { - throw std::runtime_error( - "WebGPU linear_q4gsw: bias present but null/undersized"); - } - has_bias = 1; - bias_buffer = bias.buffer; - bias_size = bias.nbytes; - } - if (bias_buffer == nullptr) { - bias_buffer = graph.create_scratch_buffer(4); - } - const Q4gswExecutionState initial_state = make_q4gsw_execution_state( device, in.dims, @@ -396,6 +464,8 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { wg_size, use_gemv, use_dual_route, + record_bk64_route, + bk64_eligible, use_steel, use_shmem_gemm); @@ -496,6 +566,9 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { if (use_dual_route) { // Each recorded dispatch owns one bind-group reference. wgpuBindGroupAddRef(bind_group); + if (record_bk64_route) { + wgpuBindGroupAddRef(bind_group); + } WGPUComputePipeline bicol_pipeline = make_pipeline(kQ4gswLinearCoop4BicolWGSL, false, gemv_wg_size); const size_t bicol_idx = graph.add_dispatch( @@ -503,6 +576,16 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bind_group, initial_state.active_grid.x, "linear_q4gsw_coop4_bicol"}); + size_t bk64_idx = 0; + if (record_bk64_route) { + WGPUComputePipeline bk64_pipeline = + make_pipeline(kQ4gswSteelBk64WGSL, true, 0u); + bk64_idx = graph.add_dispatch( + {bk64_pipeline, + bind_group, + initial_state.active_grid.x, + "linear_q4gsw_bk64"}); + } WGPUComputePipeline prefill_pipeline = make_pipeline(prefill_shader_src, fixed_prefill_wg, wg_size); const size_t prefill_idx = graph.add_dispatch( @@ -510,19 +593,27 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bind_group, initial_state.active_grid.x, prefill_label}); - route_group = graph.register_dispatch_route_group( - {{bicol_idx, bicol_idx + 1}, {prefill_idx, prefill_idx + 1}}); + if (record_bk64_route) { + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, + {bk64_idx, bk64_idx + 1}, + {prefill_idx, prefill_idx + 1}}); + } else { + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, {prefill_idx, prefill_idx + 1}}); + } graph.select_dispatch_route( route_group, initial_state.active_route, {initial_state.active_grid}); } else { - const bool fixed_wg = use_gemv ? false : fixed_prefill_wg; + const bool fixed_wg = use_gemv ? false : (use_bk64 || fixed_prefill_wg); WGPUComputePipeline pipeline = make_pipeline(shader_src, fixed_wg, use_gemv ? gemv_wg_size : wg_size); dispatch_idx = graph.add_dispatch( {pipeline, bind_group, initial_state.active_grid.x, - use_gemv ? "linear_q4gsw_coop4_bicol" : prefill_label}); + use_gemv ? "linear_q4gsw_coop4_bicol" + : (use_bk64 ? "linear_q4gsw_bk64" : prefill_label)}); } // Dynamic shapes: recompute one shared Params block and select exactly one @@ -541,6 +632,8 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { wg_size, use_gemv, use_dual_route, + record_bk64_route, + bk64_eligible, use_steel, use_shmem_gemm, dispatch_idx, @@ -560,6 +653,8 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { wg_size, use_gemv, use_dual_route, + record_bk64_route, + bk64_eligible, use_steel, use_shmem_gemm); wgpuQueueWriteBuffer( diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64.wgsl new file mode 100644 index 00000000000..5c8b6d0802b --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64.wgsl @@ -0,0 +1,117 @@ +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; + +// BK64 prefill variant: group_size=64 keeps one scale valid for all eight packed words. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 64u; +var As: array; +var Bs: array; +@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; + 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; } + } + let ar = tid / 4u; + let ac = (tid % 4u) * 4u; + + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + let av0 = t_input[(base + 0u) >> 2u]; + let av1 = t_input[(base + 16u) >> 2u]; + let av2 = t_input[(base + 32u) >> 2u]; + let av3 = t_input[(base + 48u) >> 2u]; + As[ar * BK + ac + 0u] = f16(av0.x); As[ar * BK + ac + 1u] = f16(av0.y); + As[ar * BK + ac + 2u] = f16(av0.z); As[ar * BK + ac + 3u] = f16(av0.w); + As[ar * BK + ac + 16u] = f16(av1.x); As[ar * BK + ac + 17u] = f16(av1.y); + As[ar * BK + ac + 18u] = f16(av1.z); As[ar * BK + ac + 19u] = f16(av1.w); + As[ar * BK + ac + 32u] = f16(av2.x); As[ar * BK + ac + 33u] = f16(av2.y); + As[ar * BK + ac + 34u] = f16(av2.z); As[ar * BK + ac + 35u] = f16(av2.w); + As[ar * BK + ac + 48u] = f16(av3.x); As[ar * BK + ac + 49u] = f16(av3.y); + As[ar * BK + ac + 50u] = f16(av3.z); As[ar * BK + ac + 51u] = f16(av3.w); + } else { + for (var segment: u32 = 0u; segment < 4u; segment = segment + 1u) { + for (var ai: u32 = 0u; ai < 4u; ai = ai + 1u) { + As[ar * BK + ac + segment * 16u + ai] = 0.0h; + } + } + } + if (tid < BN) { + let c = tid; + let n = col0 + c; + if (n < params.N) { + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word + 0u]; + let w1 = t_weight[base_word + 1u]; + let w2 = t_weight[base_word + 2u]; + let w3 = t_weight[base_word + 3u]; + let w4 = t_weight[base_word + 4u]; + let w5 = t_weight[base_word + 5u]; + let w6 = t_weight[base_word + 6u]; + let w7 = t_weight[base_word + 7u]; + let words = array(w0, w1, w2, w3, w4, w5, w6, w7); + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = words[br >> 3u]; + 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) { + let r = row0 + lid.y * 4u + m; + let c0 = col0 + lid.x * 4u; + if (r < params.M && c0 < params.N) { + var vv = vec4( + f32(acc[m][0]), f32(acc[m][1]), f32(acc[m][2]), f32(acc[m][3])); + if (params.has_bias != 0u) { + vv = vv + vec4( + t_bias[c0], t_bias[c0 + 1u], t_bias[c0 + 2u], t_bias[c0 + 3u]); + } + t_out[(r * params.N + c0) >> 2u] = vv; + } + } +} diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h new file mode 100644 index 00000000000..1d6fb25cc3b --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_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_steel_bk64.wgsl - DO NOT EDIT. +// wgsl-sha256: f7425201059aa7d7d7671f405bba65d446627474d89db87a12fac4a406f6f2b2 +inline constexpr const char* kQ4gswSteelBk64WGSL = 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; + +// BK64 prefill variant: group_size=64 keeps one scale valid for all eight packed words. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 64u; +var As: array; +var Bs: array; +@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; + 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; } + } + let ar = tid / 4u; + let ac = (tid % 4u) * 4u; + + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + let av0 = t_input[(base + 0u) >> 2u]; + let av1 = t_input[(base + 16u) >> 2u]; + let av2 = t_input[(base + 32u) >> 2u]; + let av3 = t_input[(base + 48u) >> 2u]; + As[ar * BK + ac + 0u] = f16(av0.x); As[ar * BK + ac + 1u] = f16(av0.y); + As[ar * BK + ac + 2u] = f16(av0.z); As[ar * BK + ac + 3u] = f16(av0.w); + As[ar * BK + ac + 16u] = f16(av1.x); As[ar * BK + ac + 17u] = f16(av1.y); + As[ar * BK + ac + 18u] = f16(av1.z); As[ar * BK + ac + 19u] = f16(av1.w); + As[ar * BK + ac + 32u] = f16(av2.x); As[ar * BK + ac + 33u] = f16(av2.y); + As[ar * BK + ac + 34u] = f16(av2.z); As[ar * BK + ac + 35u] = f16(av2.w); + As[ar * BK + ac + 48u] = f16(av3.x); As[ar * BK + ac + 49u] = f16(av3.y); + As[ar * BK + ac + 50u] = f16(av3.z); As[ar * BK + ac + 51u] = f16(av3.w); + } else { + for (var segment: u32 = 0u; segment < 4u; segment = segment + 1u) { + for (var ai: u32 = 0u; ai < 4u; ai = ai + 1u) { + As[ar * BK + ac + segment * 16u + ai] = 0.0h; + } + } + } + if (tid < BN) { + let c = tid; + let n = col0 + c; + if (n < params.N) { + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word + 0u]; + let w1 = t_weight[base_word + 1u]; + let w2 = t_weight[base_word + 2u]; + let w3 = t_weight[base_word + 3u]; + let w4 = t_weight[base_word + 4u]; + let w5 = t_weight[base_word + 5u]; + let w6 = t_weight[base_word + 6u]; + let w7 = t_weight[base_word + 7u]; + let words = array(w0, w1, w2, w3, w4, w5, w6, w7); + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = words[br >> 3u]; + 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) { + let r = row0 + lid.y * 4u + m; + let c0 = col0 + lid.x * 4u; + if (r < params.M && c0 < params.N) { + var vv = vec4( + f32(acc[m][0]), f32(acc[m][1]), f32(acc[m][2]), f32(acc[m][3])); + if (params.has_bias != 0u) { + vv = vv + vec4( + t_bias[c0], t_bias[c0 + 1u], t_bias[c0 + 2u], t_bias[c0 + 3u]); + } + t_out[(r * params.N + c0) >> 2u] = vv; + } + } +} +)"; + +inline constexpr uint32_t kQ4gswSteelBk64WorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswSteelBk64WorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswSteelBk64WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index 305a6c39ee0..b80fdb3b377 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -16,6 +16,8 @@ #include #include +#include +#include using executorch::backends::webgpu::WebGPUDispatch; using executorch::backends::webgpu::WebGPUGraph; @@ -215,4 +217,46 @@ TEST(DispatchRoute, RecordsAlternatesOnlyForDynamicEligibleGraphs) { EXPECT_TRUE(should_record_sdpa_dual_route(true, true)); } +TEST(DispatchRoute, Bk64RequiresExactLlamaShapeAndCapabilities) { + using executorch::backends::webgpu::utils::is_q4gsw_bk64_eligible; + + constexpr uint32_t kRequiredInvocations = 256u; + constexpr uint32_t kRequiredStorageBytes = 16384u; + for (const auto& shape : std::vector>{ + {2048u, 8192u}, {8192u, 2048u}, {2048u, 2048u}}) { + EXPECT_TRUE(is_q4gsw_bk64_eligible( + shape.first, + shape.second, + 64u, + false, + true, + kRequiredInvocations, + kRequiredStorageBytes)); + } + + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 512u, 64u, false, true, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 32u, false, true, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, true, true, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, false, false, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, false, true, 255u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, false, true, 256u, kRequiredStorageBytes - 1u)); +} + +TEST(DispatchRoute, Bk64SelectsOnlyAcceptedExactLiveRows) { + using executorch::backends::webgpu::utils::is_q4gsw_bk64_live_m; + + EXPECT_TRUE(is_q4gsw_bk64_live_m(128u)); + EXPECT_TRUE(is_q4gsw_bk64_live_m(508u)); + EXPECT_TRUE(is_q4gsw_bk64_live_m(512u)); + for (uint32_t m : {1u, 127u, 129u, 507u, 509u, 511u, 513u}) { + EXPECT_FALSE(is_q4gsw_bk64_live_m(m)) << "M=" << m; + } +} + } // namespace diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 57b4aa1d8dd..9a0389cdcda 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -137,7 +137,11 @@ void run_linear( int m_rows, const char* prefix, int n, - int k = kLinK) { + int k = kLinK, + float atol = 5e-3f, + float rtol = 0.0f, + float nrmse_limit = -1.0f, + float tail_nrmse_limit = -1.0f) { const std::string base = g_dir + "/" + prefix + ".S" + std::to_string(m_rows); auto input = read_bin(base + ".input.bin"); auto golden = read_bin(base + ".golden.bin"); @@ -152,9 +156,56 @@ void run_linear( << prefix << " M=" << m_rows << " output numel mismatch"; std::vector got( out.const_data_ptr(), out.const_data_ptr() + numel); - const float e = max_err(got, golden); - // 4-bit quant: looser tol (the kernel mirrors the dequant-matmul reference). - EXPECT_LT(e, 5e-3f) << prefix << " M=" << m_rows << " max_err=" << e; + ASSERT_EQ(got.size(), golden.size()); + float max_abs = 0.0f; + float max_rel = 0.0f; + double error_sq_sum = 0.0; + double golden_sq_sum = 0.0; + bool within_tolerance = true; + for (size_t i = 0; i < got.size(); ++i) { + ASSERT_TRUE(std::isfinite(got[i])) + << prefix << " M=" << m_rows << " i=" << i; + ASSERT_TRUE(std::isfinite(golden[i])) + << prefix << " M=" << m_rows << " golden i=" << i; + const float abs_err = std::fabs(got[i] - golden[i]); + const float rel_err = abs_err / std::fmax(std::fabs(golden[i]), 1e-6f); + max_abs = std::fmax(max_abs, abs_err); + max_rel = std::fmax(max_rel, rel_err); + if (abs_err > atol && rel_err > rtol) { + within_tolerance = false; + } + error_sq_sum += static_cast(abs_err) * abs_err; + golden_sq_sum += static_cast(golden[i]) * golden[i]; + } + EXPECT_TRUE(within_tolerance) + << prefix << " M=" << m_rows << " max_abs=" << max_abs + << " max_rel=" << max_rel << " tolerances=" << atol << "/" << rtol; + if (nrmse_limit > 0.0f) { + const double nrmse = std::sqrt(error_sq_sum / golden_sq_sum); + EXPECT_LT(nrmse, nrmse_limit) + << prefix << " M=" << m_rows << " full-output NRMSE"; + std::printf( + "%s M=%d max_abs=%g max_rel=%g nrmse=%g\n", + prefix, + m_rows, + max_abs, + max_rel, + nrmse); + } + if (tail_nrmse_limit > 0.0f) { + double tail_error_sq_sum = 0.0; + double tail_golden_sq_sum = 0.0; + const size_t tail_begin = static_cast(m_rows - 1) * n; + for (size_t i = tail_begin; i < got.size(); ++i) { + const double error = static_cast(got[i]) - golden[i]; + tail_error_sq_sum += error * error; + tail_golden_sq_sum += static_cast(golden[i]) * golden[i]; + } + const double tail_nrmse = std::sqrt(tail_error_sq_sum / tail_golden_sq_sum); + EXPECT_LT(tail_nrmse, tail_nrmse_limit) + << prefix << " M=" << m_rows << " final-row NRMSE"; + std::printf("%s M=%d final_row_nrmse=%g\n", prefix, m_rows, tail_nrmse); + } } void check_linear(int m_rows) { @@ -175,6 +226,33 @@ void check_linear_tiled(int m_rows) { run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); } +constexpr int kBk64K = 2048; +constexpr int kBk64N = 2048; +constexpr int kBk64KvN = 512; +constexpr int kBk64GateN = 8192; +constexpr int kBk64DownK = 8192; +constexpr float kBk64Atol = 5e-2f; +constexpr float kBk64Rtol = 3e-2f; +constexpr float kBk64Nrmse = 7.5e-3f; +constexpr float kBk64TailNrmse = 1e-2f; +constexpr float kBk64DownAtol = 5e-2f; +constexpr float kBk64DownRtol = 8e-2f; +constexpr float kBk64DownNrmse = 1.5e-2f; +constexpr float kBk64DownTailNrmse = 2e-2f; + +void run_bk64_linear( + Module& module, + int m_rows, + const char* prefix, + int k = kBk64K, + int n = kBk64N, + float atol = kBk64Atol, + float rtol = kBk64Rtol, + float nrmse = kBk64Nrmse, + float tail_nrmse = kBk64TailNrmse) { + run_linear(module, m_rows, prefix, n, k, atol, rtol, nrmse, tail_nrmse); +} + constexpr int kSwiGluWidth = 8192; constexpr int kSwiGluSmallWidth = 64; constexpr int kSwiGluK = 64; @@ -600,7 +678,145 @@ TEST(DynamicShape, QuantizedLinearTiledReusedGraph) { } } +TEST(DynamicShape, QuantizedLinearBk64ReusedGraphAndFallbacks) { + Module candidate(g_dir + "/dyn_linear_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok) << "load dyn_linear_bk64.pte"; + for (int m_rows : {512, 511, 508, 128, 127, 1, 508, 512}) { + run_bk64_linear(candidate, m_rows, "dyn_linear_bk64"); + } + + Module gate(g_dir + "/dyn_linear_bk64_gate.pte"); + ASSERT_EQ(gate.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear(gate, m_rows, "dyn_linear_bk64_gate", kBk64K, kBk64GateN); + } + + Module down(g_dir + "/dyn_linear_bk64_down.pte"); + ASSERT_EQ(down.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear( + down, + m_rows, + "dyn_linear_bk64_down", + kBk64DownK, + kBk64N, + kBk64DownAtol, + kBk64DownRtol, + kBk64DownNrmse, + kBk64DownTailNrmse); + } + + Module group32(g_dir + "/dyn_linear_bk64_group32.pte"); + ASSERT_EQ(group32.load_forward(), Error::Ok); + run_bk64_linear(group32, 128, "dyn_linear_bk64_group32"); + + Module bias(g_dir + "/dyn_linear_bk64_bias.pte"); + ASSERT_EQ(bias.load_forward(), Error::Ok); + run_bk64_linear(bias, 128, "dyn_linear_bk64_bias"); + + Module kv_shape(g_dir + "/dyn_linear_bk64_kv_shape.pte"); + ASSERT_EQ(kv_shape.load_forward(), Error::Ok); + run_bk64_linear(kv_shape, 128, "dyn_linear_bk64_kv_shape", kBk64K, kBk64KvN); +} + #ifdef WGPU_BACKEND_ENABLE_PROFILING +void expect_single_q4_profile( + const char* expected_name, + uint32_t expected_x, + const char* fixture, + int m_rows) { + const auto* context = get_default_webgpu_context(); + ASSERT_NE(context, nullptr); + ASSERT_NE(context->querypool, nullptr); + const auto& profile = context->querypool->results(); + const auto is_q4 = [](const auto& duration) { + return duration.kernel_name.rfind("linear_q4gsw", 0) == 0; + }; + ASSERT_EQ(std::count_if(profile.begin(), profile.end(), is_q4), 1) + << fixture << " M=" << m_rows; + const auto active = std::find_if(profile.begin(), profile.end(), is_q4); + ASSERT_NE(active, profile.end()); + EXPECT_EQ(active->kernel_name, expected_name) << fixture << " M=" << m_rows; + EXPECT_EQ(active->global_wg[0], expected_x) << fixture << " M=" << m_rows; + EXPECT_EQ(active->global_wg[1], 1u) << fixture << " M=" << m_rows; +} + +TEST(DynamicShape, QuantizedLinearBk64ProfileSoleWriter) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !context->shader_f16_supported) { + GTEST_SKIP() << "BK64 timestamp or shader-f16 capability unavailable"; + } + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || + limits.maxComputeInvocationsPerWorkgroup < 256u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 1024u) { + GTEST_SKIP() << "BK64 workgroup limits unavailable"; + } + + Module candidate(g_dir + "/dyn_linear_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok); + for (int m_rows : {512, 511, 508, 128, 127, 1, 508, 512}) { + run_bk64_linear(candidate, m_rows, "dyn_linear_bk64"); + const char* expected = m_rows == 1 ? "linear_q4gsw_coop4_bicol" + : (m_rows == 128 || m_rows == 508 || m_rows == 512) + ? "linear_q4gsw_bk64" + : "linear_q4gsw_steel"; + const uint32_t expected_x = m_rows == 1 ? 1024u + : (m_rows == 128 || m_rows == 127) ? 64u + : 256u; + expect_single_q4_profile(expected, expected_x, "dyn_linear_bk64", m_rows); + } + + Module gate(g_dir + "/dyn_linear_bk64_gate.pte"); + ASSERT_EQ(gate.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear(gate, m_rows, "dyn_linear_bk64_gate", kBk64K, kBk64GateN); + expect_single_q4_profile( + "linear_q4gsw_bk64", + m_rows == 128 ? 256u : 1024u, + "dyn_linear_bk64_gate", + m_rows); + } + + Module down(g_dir + "/dyn_linear_bk64_down.pte"); + ASSERT_EQ(down.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear( + down, + m_rows, + "dyn_linear_bk64_down", + kBk64DownK, + kBk64N, + kBk64DownAtol, + kBk64DownRtol, + kBk64DownNrmse, + kBk64DownTailNrmse); + expect_single_q4_profile( + "linear_q4gsw_bk64", + m_rows == 128 ? 64u : 256u, + "dyn_linear_bk64_down", + m_rows); + } + + struct NegativeRoute { + const char* fixture; + int n; + uint32_t expected_x; + }; + for (const auto& negative : std::vector{ + {"dyn_linear_bk64_group32", kBk64N, 64u}, + {"dyn_linear_bk64_bias", kBk64N, 64u}, + {"dyn_linear_bk64_kv_shape", kBk64KvN, 16u}}) { + Module module(g_dir + "/" + negative.fixture + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << negative.fixture; + run_bk64_linear(module, 128, negative.fixture, kBk64K, negative.n); + expect_single_q4_profile( + "linear_q4gsw_steel", negative.expected_x, negative.fixture, 128); + } +} + TEST(DynamicShape, CombinedLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index 4d19b6928cf..f0ff22eb6f4 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -280,6 +280,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: ) _export_static_linear(out_dir, 1, "static_linear_m1") _export_static_linear(out_dir, 32, "static_linear_m32") + _export_dynamic_bk64_linear_cases(out_dir) # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0). _export_dynamic_sdpa(out_dir) @@ -332,6 +333,16 @@ def export_dynamic_shape_cases(out_dir: str) -> None: LIN_GROUP = 32 LIN_MAXM = 128 +BK64_K = 2048 +BK64_N = 2048 +BK64_KV_N = 512 +BK64_GATE_N = 8192 +BK64_DOWN_K = 8192 +BK64_GROUP = 64 +BK64_MAXM = 512 +BK64_LIVE_M = (BK64_MAXM, 511, 508, 128, 127, 1) +BK64_OPTIMIZED_M = (BK64_MAXM, 508, 128) + SWIGLU_MAXM = 512 SWIGLU_WIDTH = 8192 @@ -472,6 +483,129 @@ def _export_dynamic_linear( print(f" golden {prefix} M={m}") +def _make_bk64_model( + *, + k: int = BK64_K, + n: int = BK64_N, + group: int = BK64_GROUP, + bias: bool = False, +) -> torch.nn.Module: + from torchao.quantization.granularity import PerGroup + from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ + + torch.manual_seed(11) + model = torch.nn.Linear(k, n, bias=bias).eval() + if model.bias is not None: + with torch.no_grad(): + model.bias.copy_(torch.linspace(-0.25, 0.25, n)) + quantize_( + model, + IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(group)), + ) + return model + + +def _bk64_golden(model: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + golden = x.double() @ model.weight.dequantize().double().t() + if model.bias is not None: + golden = golden + model.bias.double() + return golden.to(torch.float32) + + +def _bk64_input(m: int, k: int) -> torch.Tensor: + flat = torch.arange(m * k, dtype=torch.int64) + hashed = (flat * 37 + torch.div(flat, 16, rounding_mode="floor") * 53) % 257 + return ((hashed.to(torch.float32) - 128.0) / 128.0).reshape(m, k) + + +class Bk64ShapeAwareLinear(torch.nn.Module): + def __init__(self, projection: torch.nn.Module, output_width: int) -> None: + super().__init__() + self.projection = projection + self.output_width = output_width + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.projection(x).reshape(1, x.shape[0], self.output_width) + + +def _export_bk64_program( + model: torch.nn.Module, + x: torch.Tensor, + n: int, + prefix: str, +): + export_model = Bk64ShapeAwareLinear(model, n).eval() + m_dim = torch.export.Dim("m", min=1, max=BK64_MAXM) + ep = torch.export.export(export_model, (x,), dynamic_shapes=({0: m_dim},)) + if not any(node.target == torch.ops.aten.sym_size.int for node in ep.graph.nodes): + raise RuntimeError(f"{prefix}: dynamic q4 fixture lost aten.sym_size.int") + return ep + + +def _export_dynamic_bk64_linear_case( + out_dir: str, + prefix: str, + *, + k: int = BK64_K, + n: int = BK64_N, + group: int = BK64_GROUP, + bias: bool = False, + live_m=BK64_LIVE_M, +) -> None: + model = _make_bk64_model(k=k, n=n, group=group, bias=bias) + x = _bk64_input(BK64_MAXM, k) + ep = _export_bk64_program(model, x, n, prefix) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + print(f"Exported {prefix}.pte") + for m in live_m: + xm = _bk64_input(m, k) + golden = _bk64_golden(model, xm) + xm.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) + _export_dynamic_bk64_linear_case(out_dir, "dyn_linear_bk64") + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_gate", + n=BK64_GATE_N, + live_m=BK64_OPTIMIZED_M, + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_down", + k=BK64_DOWN_K, + live_m=BK64_OPTIMIZED_M, + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_group32", + group=32, + live_m=[128], + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_bias", + bias=True, + live_m=[128], + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_kv_shape", + n=BK64_KV_N, + live_m=[128], + ) + + def _export_static_linear(out_dir: str, m: int, prefix: str) -> None: from executorch.backends.webgpu.test.ops.test_quantized_linear import ( _fp64_golden, @@ -798,6 +932,33 @@ def test_export_dynamic_rms(self) -> None: export_dynamic_shape_cases(d) self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.pte"))) self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.S1.golden.bin"))) + expected = [ + "dyn_linear_bk64.pte", + "dyn_linear_bk64.S512.input.bin", + "dyn_linear_bk64.S512.golden.bin", + "dyn_linear_bk64.S511.input.bin", + "dyn_linear_bk64.S511.golden.bin", + "dyn_linear_bk64.S508.golden.bin", + "dyn_linear_bk64.S128.golden.bin", + "dyn_linear_bk64.S127.golden.bin", + "dyn_linear_bk64.S1.golden.bin", + "dyn_linear_bk64_gate.pte", + "dyn_linear_bk64_gate.S512.input.bin", + "dyn_linear_bk64_gate.S512.golden.bin", + "dyn_linear_bk64_gate.S508.golden.bin", + "dyn_linear_bk64_gate.S128.golden.bin", + "dyn_linear_bk64_down.pte", + "dyn_linear_bk64_down.S512.input.bin", + "dyn_linear_bk64_down.S512.golden.bin", + "dyn_linear_bk64_down.S508.golden.bin", + "dyn_linear_bk64_down.S128.golden.bin", + "dyn_linear_bk64_group32.pte", + "dyn_linear_bk64_bias.pte", + "dyn_linear_bk64_kv_shape.pte", + ] + for name in expected: + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) if __name__ == "__main__":