diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 7f0e7f1927a..9b87990aa8f 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -100,6 +100,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/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 81e524bf47c..abf6f027aa0 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1325,9 +1325,14 @@ void WebGPUGraph::build( auto mit = qkv_member.find(i); if (mit != qkv_member.end()) { QkvFusionGroup& g = qkv_groups[mit->second]; - const size_t di = num_dispatches() - 1; // this linear's dispatch + const utils::DispatchRange dispatch_range = { + dispatch_begin, num_dispatches()}; + if (dispatch_range.begin == dispatch_range.end) { + throw std::runtime_error( + "WebGPU QKV fusion member emitted no dispatch"); + } if (i == g.op_idx[0]) { - g.sep_dispatch[0] = di; + g.sep_dispatch[0] = dispatch_range; // Emit the fused dispatch RIGHT AFTER the q-linear (the anchor) so // at M>1 it writes q/k/v BEFORE any consumer. q/k/v may be // interleaved with rope in the chain, so emitting it at the LAST @@ -1335,14 +1340,14 @@ void WebGPUGraph::build( // fresh_q -> garbage. add_qkv_fused_dispatch(g); } else if (i == g.op_idx[1]) { - g.sep_dispatch[1] = di; + g.sep_dispatch[1] = dispatch_range; } else { - g.sep_dispatch[2] = di; + g.sep_dispatch[2] = dispatch_range; } } auto lit = qkv_last.find(i); if (lit != qkv_last.end()) { - // All 3 sep dispatch indices + the fused index are now known. + // All 3 separate route ranges + the fused index are now known. add_qkv_fused_hook(qkv_groups[lit->second]); } } @@ -1670,8 +1675,9 @@ void WebGPUGraph::add_qkv_fused_dispatch(QkvFusionGroup& g) { g.fused_params = uniform_buffer; } -// M-gate coordinator: registered at the LAST triple op (all dispatch indices -// known). Prefill (M>1): run the fused GEMM, zero the 3 separate linears. +// M-gate coordinator: registered at the LAST triple op (all dispatch ranges +// known). Prefill (M>1): run the fused GEMM, zero every route of the 3 separate +// linears. // Decode (M==1): zero the fused, leave the 3 coop4 GEMVs (their own hooks set // the decode wg) -- the fused 64x64 tile wastes 63/64 rows at M=1. Recomputes // live M + the 3 output cur_dims + fused params. Inert on a static graph; a @@ -1681,8 +1687,9 @@ void WebGPUGraph::add_qkv_fused_hook(const QkvFusionGroup& g) { out_v_id = g.out_v; const uint32_t K = g.K, Kp = g.K_packed, gs = g.group_size, Nq = g.Nq, Nk = g.Nk, Nv = g.Nv, Nf = g.Nq + g.Nk + g.Nv; - const size_t fused_idx = g.fused_dispatch, sep0 = g.sep_dispatch[0], - sep1 = g.sep_dispatch[1], sep2 = g.sep_dispatch[2]; + const size_t fused_idx = g.fused_dispatch; + const std::array separate_ranges = { + g.sep_dispatch[0], g.sep_dispatch[1], g.sep_dispatch[2]}; WGPUBuffer params_buf = g.fused_params; auto update_route = [input_id, out_q_id, @@ -1696,9 +1703,7 @@ void WebGPUGraph::add_qkv_fused_hook(const QkvFusionGroup& g) { Nv, Nf, fused_idx, - sep0, - sep1, - sep2, + separate_ranges, params_buf](WebGPUGraph& gr) { const auto& d = gr.cur_dims(input_id); uint64_t numel = 1; @@ -1729,12 +1734,12 @@ void WebGPUGraph::add_qkv_fused_hook(const QkvFusionGroup& g) { const uint32_t nbM2 = (m + 63u) / 64u; gr.dispatch_at(fused_idx).workgroup_count_x = nbN2 * nbM2; gr.dispatch_at(fused_idx).workgroup_count_y = 1u; - gr.dispatch_at(sep0).workgroup_count_x = 0u; - gr.dispatch_at(sep0).workgroup_count_y = 0u; - gr.dispatch_at(sep1).workgroup_count_x = 0u; - gr.dispatch_at(sep1).workgroup_count_y = 0u; - gr.dispatch_at(sep2).workgroup_count_x = 0u; - gr.dispatch_at(sep2).workgroup_count_y = 0u; + for (const auto& range : separate_ranges) { + for (size_t i = range.begin; i < range.end; i++) { + gr.dispatch_at(i).workgroup_count_x = 0u; + gr.dispatch_at(i).workgroup_count_y = 0u; + } + } } else { gr.dispatch_at(fused_idx).workgroup_count_x = 0u; gr.dispatch_at(fused_idx).workgroup_count_y = 0u; diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 261c29a282d..705687c20a2 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -674,10 +674,10 @@ class WebGPUGraph { uint32_t K = 0, K_packed = 0, group_size = 0, num_groups = 0; uint32_t padded_N_q = 0, padded_N_k = 0, padded_N_v = 0; unsigned op_idx[3] = {0, 0, 0}; // the 3 q/k/v linear op-chain indices - size_t sep_dispatch[3] = { - 0, - 0, - 0}; // their dispatch indices (filled in build()) + utils::DispatchRange sep_dispatch[3] = { + {0, 0}, + {0, 0}, + {0, 0}}; // each linear's complete route range (filled in build()) size_t fused_dispatch = 0; // the fused GEMM dispatch index WGPUBuffer fused_params = nullptr; // the fused params UBO (rewritten by the hook) diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 36722aa2c39..ff4064efa8f 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -99,6 +99,7 @@ #include #include #include +#include #include #include #include @@ -147,7 +148,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -757,6 +758,13 @@ constexpr std::array kShaderRegistry = {{ kQ4gswRequantWorkgroupSizeY, kQ4gswRequantWorkgroupSizeZ, }, + { + "q4gsw_steel_bk64", + kQ4gswSteelBk64WGSL, + kQ4gswSteelBk64WorkgroupSizeX, + kQ4gswSteelBk64WorkgroupSizeY, + kQ4gswSteelBk64WorkgroupSizeZ, + }, { "q8ta_add", kQ8taAddWGSL, diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 16d78a6f923..fe60e8377a1 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -8,15 +8,9 @@ #include #include +#include #include #include -#include -#include -#include -#include -#include -#include -#include #include @@ -59,8 +53,20 @@ 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; +constexpr const char* kQ4gswLinearShader = "q4gsw_linear"; +constexpr const char* kQ4gswBicolShader = "q4gsw_linear_coop4_bicol"; +constexpr const char* kQ4gswShmemShader = "q4gsw_linear_gemm_shmem"; +constexpr const char* kQ4gswSteelShader = "q4gsw_linear_gemm_steel"; +constexpr const char* kQ4gswSteelHalfShader = "q4gsw_linear_gemm_steel_half"; +constexpr const char* kQ4gswSteelHalfPwdqShader = + "q4gsw_linear_gemm_steel_half_pwdq"; +constexpr const char* kQ4gswSteelHalfPwdqF16accShader = + "q4gsw_linear_gemm_steel_half_pwdq_f16acc"; +constexpr const char* kQ4gswSteelBk64Shader = "q4gsw_steel_bk64"; + // One workgroup per (tile_m x tile_n) tile, no grid-stride: throw when the tile // count would exceed the 1D dispatch limit. Shared by the steel + shmem GEMM // routes; `kind` names the route in the error message. @@ -89,6 +95,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) { @@ -102,6 +129,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 @@ -109,10 +147,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) { @@ -131,6 +171,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 @@ -169,7 +217,11 @@ struct Q4gswExecutionState { }; constexpr size_t kQ4gswBicolRoute = 0; -constexpr size_t kQ4gswPrefillRoute = 1; +constexpr size_t kQ4gswBk64Route = 1; +constexpr size_t kQ4gswPrefillRoute = 2; +// 2-route (bicol + prefill) layout, used when the BK64 route is not recorded: +// the prefill dispatch sits at index 1, not kQ4gswPrefillRoute (the 3-route 2). +constexpr size_t kQ4gswPrefillRoute2Way = 1; Q4gswExecutionState make_q4gsw_execution_state( WGPUDevice device, @@ -184,6 +236,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()) { @@ -204,13 +258,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)"); @@ -224,12 +283,73 @@ 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) + : kQ4gswPrefillRoute2Way)) + : 0u; state.active_grid = {workgroup_count, 1u}; return state; } +struct Q4gswResizeContext { + int in_id; + int out_id; + uint32_t max_m; + uint32_t K; + uint32_t N; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t has_bias; + 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; + size_t dispatch_idx; + size_t route_group; + WGPUBuffer params_buffer; +}; + +void resize_q4gsw(WebGPUGraph& graph, const Q4gswResizeContext& context) { + const Q4gswExecutionState state = make_q4gsw_execution_state( + graph.device(), + graph.cur_dims(context.in_id), + context.max_m, + context.K, + context.N, + context.K_packed, + context.group_size, + context.padded_N, + context.has_bias, + context.wg_size, + context.use_single_gemv, + context.use_dual_route, + context.record_bk64_route, + context.bk64_eligible, + context.prefill_use_steel, + context.prefill_use_shmem_gemm); + wgpuQueueWriteBuffer( + graph.queue(), + context.params_buffer, + 0, + &state.params, + sizeof(state.params)); + if (context.use_dual_route) { + graph.select_dispatch_route( + context.route_group, state.active_route, {state.active_grid}); + } else { + auto& dispatch = graph.dispatch_at(context.dispatch_idx); + dispatch.workgroup_count_x = state.active_grid.x; + dispatch.workgroup_count_y = state.active_grid.y; + } + graph.set_cur_dims(context.out_id, state.output_dims); +} + // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out]. void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const int in_id = args.at(0); @@ -308,9 +428,28 @@ 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. - const uint32_t wg_size = - utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX); + // 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, get_webgpu_shader_info(kQ4gswLinearShader).workgroup_size_x); 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( @@ -318,10 +457,16 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bicol_eligible, graph.has_dynamic_shapes(), graph.config().record_q4gsw_decode_route); + 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( - device, kQ4gswLinearCoop4BicolWorkgroupSizeX) + device, get_webgpu_shader_info(kQ4gswBicolShader).workgroup_size_x) : 0u; // steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible. const bool use_steel = !use_gemv && steel_supported(device) && @@ -332,13 +477,14 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // large K/N thresholds; otherwise the register-tiled path handles it. const bool use_shmem_gemm = !use_gemv && !use_steel && (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim); - const char* prefill_shader_src = use_steel ? kQ4gswLinearGemmSteelWGSL - : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL - : kQ4gswLinearWGSL; - const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL - : use_steel ? kQ4gswLinearGemmSteelWGSL - : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL - : kQ4gswLinearWGSL; + const char* prefill_shader_name = use_steel ? kQ4gswSteelShader + : use_shmem_gemm ? kQ4gswShmemShader + : kQ4gswLinearShader; + const char* shader_name = use_gemv ? kQ4gswBicolShader + : use_bk64 ? kQ4gswSteelBk64Shader + : use_steel ? kQ4gswSteelShader + : use_shmem_gemm ? kQ4gswShmemShader + : kQ4gswLinearShader; // f16-multiply steel: only when the device negotiated shader-f16; else the // f32 steel kernel runs (fail-closed). Same bindings and tile. if (use_steel) { @@ -348,10 +494,12 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // 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. - prefill_shader_src = (gs % kQ4gswSteelBK == 0u) - ? kQ4gswLinearGemmSteelHalfPwdqWGSL - : kQ4gswLinearGemmSteelHalfWGSL; - shader_src = prefill_shader_src; + prefill_shader_name = (gs % kQ4gswSteelBK == 0u) + ? kQ4gswSteelHalfPwdqShader + : kQ4gswSteelHalfShader; + if (!use_bk64) { + shader_name = prefill_shader_name; + } } } // f16-accumulate: pwdq staging with an f16 register accumulator. @@ -362,29 +510,13 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { if (use_steel && graph.f16_accumulate_gemm() && (gs % kQ4gswSteelBK == 0u)) { const WebGPUContext* ctx = get_default_webgpu_context(); if (ctx != nullptr && ctx->shader_f16_supported) { - prefill_shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; - shader_src = prefill_shader_src; + prefill_shader_name = kQ4gswSteelHalfPwdqF16accShader; + if (!use_bk64) { + shader_name = prefill_shader_name; + } } } - // 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, @@ -398,11 +530,12 @@ 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); WGPUBuffer params_buffer = graph.create_params_buffer(initial_state.params); - const std::vector bindings = { {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, @@ -444,84 +577,90 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { size_t dispatch_idx = 0; size_t route_group = 0; if (use_dual_route) { - utils::ComputePipelineBundle bicol_bundle = - make_bundle(kQ4gswLinearCoop4BicolWGSL, false, gemv_wg_size); + utils::ComputePipelineBundle bicol_bundle = make_bundle( + get_webgpu_shader_info(kQ4gswBicolShader).source, false, gemv_wg_size); const size_t bicol_idx = graph.add_dispatch( {bicol_bundle.pipeline, bicol_bundle.bind_group, initial_state.active_grid.x, - "linear_q4gsw_coop4_bicol"}); + "linear_q4gsw_coop4_bicol", + initial_state.active_grid.y}); + size_t bk64_idx = 0; + if (record_bk64_route) { + utils::ComputePipelineBundle bk64_bundle = make_shared_bundle( + get_webgpu_shader_info(kQ4gswSteelBk64Shader).source, + bicol_bundle, + true, + 0u); + bk64_idx = graph.add_dispatch( + {bk64_bundle.pipeline, + bk64_bundle.bind_group, + initial_state.active_grid.x, + "linear_q4gsw_bk64", + initial_state.active_grid.y}); + } utils::ComputePipelineBundle prefill_bundle = make_shared_bundle( - prefill_shader_src, bicol_bundle, fixed_prefill_wg, wg_size); + get_webgpu_shader_info(prefill_shader_name).source, + bicol_bundle, + fixed_prefill_wg, + wg_size); const size_t prefill_idx = graph.add_dispatch( {prefill_bundle.pipeline, prefill_bundle.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}}); + prefill_label, + initial_state.active_grid.y}); + 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; - utils::ComputePipelineBundle bundle = - make_bundle(shader_src, fixed_wg, use_gemv ? gemv_wg_size : wg_size); + const bool fixed_wg = use_gemv ? false : (use_bk64 || fixed_prefill_wg); + utils::ComputePipelineBundle bundle = make_bundle( + get_webgpu_shader_info(shader_name).source, + fixed_wg, + use_gemv ? gemv_wg_size : wg_size); dispatch_idx = graph.add_dispatch( {bundle.pipeline, bundle.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), + initial_state.active_grid.y}); } // Dynamic shapes: recompute one shared Params block and select exactly one // writer. The prefill pipeline remains the route chosen from max M. - graph.add_tensor_resize_hook( + const Q4gswResizeContext resize_context = { in_id, - [in_id, - out_id, - M, - K, - N, - K_packed, - gs, - padded_N, - has_bias, - wg_size, - use_gemv, - use_dual_route, - use_steel, - use_shmem_gemm, - dispatch_idx, - route_group, - params_buffer](WebGPUGraph& g) { - const auto& d = g.cur_dims(in_id); - const Q4gswExecutionState state = make_q4gsw_execution_state( - g.device(), - d, - M, - K, - N, - K_packed, - gs, - padded_N, - has_bias, - wg_size, - use_gemv, - use_dual_route, - use_steel, - use_shmem_gemm); - wgpuQueueWriteBuffer( - g.queue(), params_buffer, 0, &state.params, sizeof(state.params)); - if (use_dual_route) { - g.select_dispatch_route( - route_group, state.active_route, {state.active_grid}); - } else { - auto& dispatch = g.dispatch_at(dispatch_idx); - dispatch.workgroup_count_x = state.active_grid.x; - dispatch.workgroup_count_y = state.active_grid.y; - } - g.set_cur_dims(out_id, state.output_dims); - }); + out_id, + M, + K, + N, + K_packed, + gs, + padded_N, + has_bias, + wg_size, + use_gemv, + use_dual_route, + record_bk64_route, + bk64_eligible, + use_steel, + use_shmem_gemm, + dispatch_idx, + route_group, + params_buffer}; + graph.add_tensor_resize_hook(in_id, [resize_context](WebGPUGraph& g) { + resize_q4gsw(g, resize_context); + }); } } // namespace 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 e2dfc610976..19d2b9aff77 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 @@ -1,6 +1,6 @@ $if DTYPE == "half": enable f16; -@group(0) @binding(0) var t_out: array; +@group(0) @binding(0) var t_out: array<${"vec4" if BK == 64 else "f32"}>; @group(0) @binding(1) var t_input: array>; @group(0) @binding(2) var t_weight: array; @group(0) @binding(3) var t_scales: array; @@ -31,9 +31,10 @@ struct Params { // 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 +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = ${BK}u; +var As: array<${buffer_scalar_type(DTYPE)}, ${64 * BK}>; // BM*BK +var Bs: array<${buffer_scalar_type(DTYPE)}, ${BK * 64}>; // BK*BN // 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @@ -64,14 +65,35 @@ fn main(@builtin(workgroup_id) wid: vec3, if (arow < params.M) { let base = arow * params.K + k0 + ac; // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). - let av = t_input[base >> 2u]; - As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(av.x); - As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(av.y); - As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(av.z); - As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(av.w); + $if BK == 64: + 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: + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(av.x); + As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(av.y); + As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(av.z); + As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(av.w); } else { - 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"}; + $if BK == 64: + 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; + } + } + $else: + 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"}; } $if PWDQ: // Packed-word dequant: threads [0,BN) each stage one full BK-column of Bs. @@ -83,16 +105,28 @@ fn main(@builtin(workgroup_id) wid: vec3, // 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. + // Column n's BK-nibble K-slice starts at this packed word. // 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; - } + $if BK == 64: + let words = array( + t_weight[base_word + 0u], t_weight[base_word + 1u], + t_weight[base_word + 2u], t_weight[base_word + 3u], + t_weight[base_word + 4u], t_weight[base_word + 5u], + t_weight[base_word + 6u], t_weight[base_word + 7u]); + 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: + 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; } } @@ -135,15 +169,30 @@ fn main(@builtin(workgroup_id) wid: vec3, workgroupBarrier(); k0 = k0 + BK; } - for (var m: u32 = 0u; m < 4u; m = m + 1u) { - for (var n: u32 = 0u; n < 4u; n = n + 1u) { + $if BK == 64: + for (var m: u32 = 0u; m < 4u; m = m + 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 ACC == "half" else "acc[m][n]"}; - if (params.has_bias != 0u) { v = v + t_bias[c]; } - t_out[r * params.N + c] = v; + 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; + } + } + $else: + 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 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 5a2cae5e499..9c16e924a0c 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 @@ -3,6 +3,7 @@ q4gsw_linear_gemm_steel: DTYPE: float PWDQ: false ACC: float + BK: 16 shader_variants: - NAME: q4gsw_linear_gemm_steel DTYPE: float @@ -20,3 +21,8 @@ q4gsw_linear_gemm_steel: DTYPE: half PWDQ: true ACC: half + - NAME: q4gsw_steel_bk64 + DTYPE: half + PWDQ: true + ACC: half + BK: 64 diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h index ab4d3c06915..a57292ac0f1 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: a4346ddef028036f29aa73f0620c586467627626a745159fccef816a32f9475a +// wgsl-sha256: 99b460d282c04a4f624e03f8216d3f36ba5f5562b6f60ff752ed9a7f4becffdc inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -47,6 +47,7 @@ struct Params { // 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. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN @@ -94,7 +95,7 @@ fn main(@builtin(workgroup_id) wid: vec3, // 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. + // Column n's BK-nibble K-slice starts at this packed word. // 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]; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h index 251a78b9fc8..2cb01ef7b53 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: d772a385edf91547a51fa3f5c4bdb4101da933f4c47828578c2818ff9eadf5fc +// wgsl-sha256: ea3edb89ec54946b392d5b4f09dd027ff19b3baad8352f1da567796e4eabc047 inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -47,6 +47,7 @@ struct Params { // 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. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN @@ -94,7 +95,7 @@ fn main(@builtin(workgroup_id) wid: vec3, // 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. + // Column n's BK-nibble K-slice starts at this packed word. // 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]; 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 4741e13b74a..37b0b8b471f 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: 363916e1aacae9635b4573d2ce35d0ee22916ce225537a3717df98a0ed3353da +// wgsl-sha256: 95cd59521a07d686a0eb396f79fe4e260e16e81d61f3c185ed7ed9be72963d59 inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -47,6 +47,7 @@ struct Params { // 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. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. 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 197b4f0ae92..04db8ea7a5b 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: f8106cb3f18424a9db8e04464e52a7fd522ca78f49e19ba6692f1c8fa83474e6 +// wgsl-sha256: 11de59057b19656a2b0ab6bf8c78c37d56997f2a3664634b4a1adf07bddc45ee inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_input: array>; @@ -46,6 +46,7 @@ struct Params { // 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. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. 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_steel_bk64_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h new file mode 100644 index 00000000000..e502bf5a5ff --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h @@ -0,0 +1,158 @@ +/* + * 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: dc25d61a4fe98a3fcfe9c006e7488ad84e2e873e6f66f26ac6cf41fd400da148 +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; + +// "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. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 64u; +var As: array; // BM*BK +var Bs: array; // BK*BN +// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + 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; + // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). + 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; + } + } + } + // 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 BK-nibble K-slice starts at this packed word. + // 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 words = array( + t_weight[base_word + 0u], t_weight[base_word + 1u], + t_weight[base_word + 2u], t_weight[base_word + 3u], + t_weight[base_word + 4u], t_weight[base_word + 5u], + t_weight[base_word + 6u], t_weight[base_word + 7u]); + 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 4478c08dbdc..67ce74659eb 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -17,6 +17,8 @@ #include #include +#include +#include using executorch::backends::webgpu::WebGPUDispatch; using executorch::backends::webgpu::WebGPUGraph; @@ -272,4 +274,46 @@ TEST(WebGPUGraphConfig, RejectsMalformedRouteCompileOption) { EXPECT_EQ(parsed_null.error(), Error::DelegateInvalidCompatibility); } +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 a7aa99ceccf..8c05fd379ea 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,60 @@ 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) { + ASSERT_GT(golden_sq_sum, 0.0) + << prefix << " M=" << m_rows << " zero golden norm (NRMSE undefined)"; + 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]; + } + ASSERT_GT(tail_golden_sq_sum, 0.0) + << prefix << " M=" << m_rows << " zero final-row golden norm"; + 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) { @@ -208,6 +263,55 @@ void run_qkv_routes(Module& m, int m_rows) { } } +constexpr int kQkvBk64K = 2048; +constexpr int kQkvBk64MaxM = 128; + +void run_qkv_bk64_routes(Module& module, int m_rows) { + const std::string path = + g_dir + "/qkv_bk64_routes.S" + std::to_string(m_rows) + ".input.bin"; + auto input = read_bin(path); + ASSERT_EQ(input.size(), static_cast(m_rows) * kQkvBk64K); + auto tensor = make_tensor_ptr({m_rows, kQkvBk64K}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE(result.ok()) << "qkv_bk64_routes M=" << m_rows; + ASSERT_EQ(result.get().size(), 3u); + const int widths[] = {kQkvNq, kQkvNk, kQkvNv}; + for (size_t i = 0; i < 3; i++) { + ASSERT_TRUE(result.get()[i].isTensor()); + const auto& output = result.get()[i].toTensor(); + EXPECT_EQ( + static_cast(output.numel()), + static_cast(m_rows) * widths[i]); + } +} + +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 kSwiGluQkvOverlapWidth = 512; @@ -667,6 +771,47 @@ TEST(DynamicShape, QkvRoutesReusedGraph) { } } +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 TEST(DynamicShape, QkvLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); @@ -691,6 +836,133 @@ TEST(DynamicShape, QkvLiveRoutesProfile) { } } +TEST(DynamicShape, QkvBk64LiveRoutesProfile) { + 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() << "timestamp queries or shader-f16 unavailable"; + } + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || + limits.maxComputeInvocationsPerWorkgroup < 256u || + limits.maxComputeWorkgroupStorageSize < 16384u) { + GTEST_SKIP() << "BK64 workgroup limits unavailable"; + } + Module module(g_dir + "/qkv_bk64_routes.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + for (int m_rows : {kQkvBk64MaxM, 1, kQkvBk64MaxM}) { + run_qkv_bk64_routes(module, m_rows); + const auto names = current_profile_names(); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_qkv_fused"), + m_rows > 1 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), + m_rows == 1 ? 3 : 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_bk64"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_shmem"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_tiled"), 0); + } +} + +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 1e4fe67b14e..aa3a195713a 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 @@ -309,6 +309,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) @@ -362,6 +363,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 @@ -509,6 +520,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, @@ -655,6 +789,9 @@ def forward(self, x, q, k, v, k_cache, v_cache): QKV_NK = 512 QKV_NV = 512 QKV_MAXM = 16 +QKV_BK64_K = 2048 +QKV_BK64_GROUP = 64 +QKV_BK64_MAXM = 128 def _export_dynamic_qkv_routes(out_dir: str) -> None: @@ -663,11 +800,11 @@ def _export_dynamic_qkv_routes(out_dir: str) -> None: ) class QkvRoutes(torch.nn.Module): - def __init__(self): + def __init__(self, k: int = LIN_K, group: int = LIN_GROUP): super().__init__() - self.q = _make_quantized_model(LIN_K, QKV_NQ, LIN_GROUP, seed=0) - self.k = _make_quantized_model(LIN_K, QKV_NK, LIN_GROUP, seed=1) - self.v = _make_quantized_model(LIN_K, QKV_NV, LIN_GROUP, seed=2) + self.q = _make_quantized_model(k, QKV_NQ, group, seed=0) + self.k = _make_quantized_model(k, QKV_NK, group, seed=1) + self.v = _make_quantized_model(k, QKV_NV, group, seed=2) def forward(self, x): # Keep the linears internal so graph-output copies do not capture @@ -701,6 +838,29 @@ def forward(self, x): tensor.detach().numpy().astype(" None: from executorch.backends.webgpu.test.ops.test_sdpa import ( @@ -936,6 +1096,7 @@ def test_q4_route_compile_specs(self) -> None: with tempfile.TemporaryDirectory() as d: _export_dynamic_qkv_routes(d) self.assertTrue(os.path.exists(os.path.join(d, "qkv_routes.pte"))) + self.assertTrue(os.path.exists(os.path.join(d, "qkv_bk64_routes.pte"))) def test_export_dynamic_rms(self) -> None: import tempfile @@ -944,6 +1105,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__":