diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index ce4784a4798..f6ac8f3f3d1 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -30,6 +30,7 @@ set(WEBGPU_SRCS runtime/WebGPUExecutionOptions.cpp runtime/WebGPUGraph.cpp runtime/passes/SwiGLU.cpp + runtime/passes/QkvBk64.cpp runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp runtime/WebGPUQueryPool.cpp diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index abf6f027aa0..e729b3926a0 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -9,7 +9,6 @@ #include #include #include -#include #include #include @@ -17,10 +16,10 @@ #include #include #include +#include #include #include -#include #include #include #include @@ -166,6 +165,7 @@ std::vector canonical_constants( constexpr const char* kPrepackOpName = "et_vk.prepack.default"; constexpr const char* kQ4gswLinearOpName = "et_vk.linear_q4gsw.default"; constexpr size_t kQ4gswOutputArg = 5; + size_t vk_datatype_size(vkgraph::VkDataType dtype) { switch (dtype) { case vkgraph::VkDataType::BOOL: @@ -210,20 +210,6 @@ int normalize_dim(int dim, int rank, const char* op) { return dim; } -// Uniform layout matching the fused-QKV WGSL Params struct (16B-aligned, 32B); -// identical to QuantizedLinear.cpp's Q4gswParams (kept local to this TU). -struct QkvFusedParams { - uint32_t M; - uint32_t N; - uint32_t K; - uint32_t K_packed; - uint32_t group_size; - uint32_t padded_N; - uint32_t has_bias; - uint32_t _pad; -}; -static_assert(sizeof(QkvFusedParams) == 32, "QkvFusedParams must be 32 bytes"); - } // namespace std::string make_compute_pipeline_key( @@ -1071,7 +1057,20 @@ void WebGPUGraph::build( std::unordered_set swiglu_skipped_ops; std::unordered_set claimed_fusion_ops; + std::vector qkv_fusions; + std::unordered_map qkv_first_ops; + std::unordered_map qkv_last_ops; + std::unordered_map qkv_member_ops; + const auto* chain = graph->chain(); + passes::detect_qkv_bk64_fusions( + *this, + graph, + num_vals, + qkv_fusions, + qkv_first_ops, + qkv_last_ops, + qkv_member_ops); passes::detect_swiglu_fusions( *this, graph, @@ -1082,182 +1081,19 @@ void WebGPUGraph::build( swiglu_skipped_ops, claimed_fusion_ops); - // Phase 3: Build operator dispatch chain - - // QKV-concat fusion detection (auto-applied graph pass, no flag): the maps - // stay empty when no q/k/v triple matches -> the Phase-3 loop below runs - // verbatim. Find each attention q/k/v triple: EXACTLY 3 et_vk.linear_q4gsw - // ops sharing args[0] (the same input activation), in chain order q,k,v with - // the N-pattern {2048,512,512}, on the steel route (K%16==0), - // group_size%16==0, no bias. The fused kernel needs shader-f16 + a 256-thread - // WG, so gate on those (else leave the triple to the normal per-linear - // handlers). qkv_fused_skip holds all 3 op indices; qkv_anchor maps the FIRST - // op index -> its group, so the fused dispatch is emitted IN-PLACE at the - // anchor (correct execution order). - std::vector qkv_groups; - std::unordered_map - qkv_first; // first triple op -> group (repoint buffers) - std::unordered_map - qkv_last; // last triple op -> group (emit fused) - std::unordered_map - qkv_member; // any triple op -> group (record dispatch) - if (chain) { - bool device_ok = false; - { - WGPULimits limits = {}; - const bool have = - wgpuDeviceGetLimits(device_, &limits) == WGPUStatus_Success; - bool f16 = false; - if (auto* ctx = get_default_webgpu_context()) { - f16 = ctx->shader_f16_supported; - } - device_ok = - have && f16 && limits.maxComputeInvocationsPerWorkgroup >= 256u; - } - if (device_ok) { - // Group linear_q4gsw op indices by input id, preserving chain order. - std::unordered_map> by_input; - std::vector input_order; - for (unsigned i = 0; i < chain->size(); i++) { - const auto* oc = chain->Get(i); - if (oc->name()->str() != "et_vk.linear_q4gsw.default") { - continue; - } - const auto* a = oc->args(); - if (!a || a->size() < 6) { - continue; - } - const int inp = static_cast(a->Get(0)); - if (by_input.find(inp) == by_input.end()) { - input_order.push_back(inp); - } - by_input[inp].push_back(i); - } - auto op_arg = [&](unsigned oi, unsigned j) { - return static_cast(chain->Get(oi)->args()->Get(j)); - }; - for (int inp : input_order) { - const auto& ops = by_input[inp]; - if (ops.size() != 3) { - continue; // gate+up is a 2-group; o/down/lm_head are 1 each. - } - if (std::any_of(ops.begin(), ops.end(), [&](unsigned op) { - return claimed_fusion_ops.count(op) != 0; - })) { - continue; - } - // args: [in, weight, scales, group_size, bias, out]. - const int wq = op_arg(ops[0], 1), sqid = op_arg(ops[0], 2), - bq = op_arg(ops[0], 4), oq = op_arg(ops[0], 5); - const int wk = op_arg(ops[1], 1), skid = op_arg(ops[1], 2), - bk = op_arg(ops[1], 4), ok = op_arg(ops[1], 5); - const int wv = op_arg(ops[2], 1), svid = op_arg(ops[2], 2), - bv = op_arg(ops[2], 4), ov = op_arg(ops[2], 5); - const int gsid = op_arg(ops[0], 3); - if (op_arg(ops[1], 3) != gsid || op_arg(ops[2], 3) != gsid) { - continue; // all 3 must share the group_size scalar. - } - if (get_value_type(bq) == ValueType::Tensor || - get_value_type(bk) == ValueType::Tensor || - get_value_type(bv) == ValueType::Tensor) { - continue; // fused kernel path assumes has_bias == 0. - } - const auto& twq = tensors_[wq]; - const auto& twk = tensors_[wk]; - const auto& twv = tensors_[wv]; - if (twq.dims.size() != 2 || twk.dims.size() != 2 || - twv.dims.size() != 2) { - continue; - } - const uint32_t Nq = static_cast(twq.dims[0]); - const uint32_t Nk = static_cast(twk.dims[0]); - const uint32_t Nv = static_cast(twv.dims[0]); - if (Nq != 2048u || Nk != 512u || Nv != 512u) { - continue; // kernel hardcodes N_Q=2048, N_KV=512 (Llama-3.2 GQA). - } - const uint32_t K_packed = static_cast(twq.dims[1]); - if (static_cast(twk.dims[1]) != K_packed || - static_cast(twv.dims[1]) != K_packed) { - continue; - } - const auto& tin = tensors_[inp]; - if (tin.dims.empty()) { - continue; - } - const uint32_t K = static_cast(tin.dims.back()); - if (K == 0 || K % 16u != 0u || K_packed != (K + 1u) / 2u) { - continue; // steel route stages a full BK=16 K-tile with no K-mask. - } - if (get_value_type(gsid) != ValueType::Int) { - continue; - } - const int64_t gsv = get_int(gsid); - if (gsv <= 0 || static_cast(gsv) % 16u != 0u) { - continue; // hoisted scale must be constant across the BK tile. - } - const uint32_t gs = static_cast(gsv); - const auto& tsq = tensors_[sqid]; - const auto& tsk = tensors_[skid]; - const auto& tsv = tensors_[svid]; - if (tsq.dims.size() != 2 || tsk.dims.size() != 2 || - tsv.dims.size() != 2) { - continue; - } - const uint32_t num_groups = static_cast(tsq.dims[0]); - if (static_cast(tsk.dims[0]) != num_groups || - static_cast(tsv.dims[0]) != num_groups) { - continue; - } - const uint32_t pNq = static_cast(tsq.dims[1]); - const uint32_t pNk = static_cast(tsk.dims[1]); - const uint32_t pNv = static_cast(tsv.dims[1]); - if (pNq < Nq || pNk < Nk || pNv < Nv || - num_groups < (K + gs - 1u) / gs) { - continue; - } - // All source + destination buffers must be live (Phase 1/2 allocated). - if (!twq.buffer || !twk.buffer || !twv.buffer || !tsq.buffer || - !tsk.buffer || !tsv.buffer || !tin.buffer || !tensors_[oq].buffer || - !tensors_[ok].buffer || !tensors_[ov].buffer) { - continue; - } - - QkvFusionGroup grp; - grp.input_id = inp; - grp.out_q = oq; - grp.out_k = ok; - grp.out_v = ov; - grp.weight_q = wq; - grp.weight_k = wk; - grp.weight_v = wv; - grp.scales_q = sqid; - grp.scales_k = skid; - grp.scales_v = svid; - grp.Nq = Nq; - grp.Nk = Nk; - grp.Nv = Nv; - grp.K = K; - grp.K_packed = K_packed; - grp.group_size = gs; - grp.num_groups = num_groups; - grp.padded_N_q = pNq; - grp.padded_N_k = pNk; - grp.padded_N_v = pNv; - grp.op_idx[0] = ops[0]; - grp.op_idx[1] = ops[1]; - grp.op_idx[2] = ops[2]; - const size_t gidx = qkv_groups.size(); - qkv_groups.push_back(grp); - qkv_first[ops[0]] = gidx; - qkv_last[ops[2]] = gidx; - qkv_member[ops[0]] = gidx; - qkv_member[ops[1]] = gidx; - qkv_member[ops[2]] = gidx; - claimed_fusion_ops.insert(ops.begin(), ops.end()); - } - } - } + // SwiGLU keeps precedence when the exact QKV geometry is also formed by a + // q projection plus gate/up projections. QKV detection runs first because it + // validates constant geometry, but it has no side effects until Phase 3; now + // discard candidates claimed by the completed SwiGLU pass and rebuild the + // index maps for the retained groups. + passes::retain_unclaimed_qkv_fusions( + qkv_fusions, + qkv_first_ops, + qkv_last_ops, + qkv_member_ops, + claimed_fusion_ops); + // Phase 3: Build operator dispatch chain if (chain) { for (unsigned i = 0; i < chain->size(); i++) { const auto* op_call = chain->Get(i); @@ -1296,63 +1132,40 @@ void WebGPUGraph::build( continue; } - const size_t dispatch_begin = dispatches_.size(); - // QKV fusion (M-gated): keep the 3 separate q/k/v linears AND add a fused - // multi-output GEMM; the fused resize hook selects by LIVE M (prefill M>1 - // -> fused runs, the 3 zeroed; decode M==1 -> the 3 coop4 GEMVs run, - // fused zeroed -- the fused 64x64 tile is ~4x slower than coop4 at M=1). - // At the FIRST triple op, repoint the 3 outputs to FRESH distinct - // buffers: the planner reuse-aliases q/k/v (each dies right after RoPE), - // which is fatal for a simultaneous fused write, so BOTH paths use - // non-aliased storage. All maps empty when no triple matches (verbatim - // path). - { - auto fit = qkv_first.find(i); - if (fit != qkv_first.end()) { - const auto& g = qkv_groups[fit->second]; - tensors_[g.out_q].buffer = - create_scratch_buffer(tensors_[g.out_q].nbytes); - tensors_[g.out_k].buffer = - create_scratch_buffer(tensors_[g.out_k].nbytes); - tensors_[g.out_v].buffer = - create_scratch_buffer(tensors_[g.out_v].nbytes); + const auto qkv_first = qkv_first_ops.find(i); + if (qkv_first != qkv_first_ops.end()) { + passes::QkvBk64Fusion& fusion = qkv_fusions[qkv_first->second]; + for (int output_id : fusion.output_ids) { + tensors_[output_id].buffer = + create_scratch_buffer(tensors_[output_id].nbytes); } } + const size_t dispatch_begin = dispatches_.size(); webgpu_operator_registry().get_op_fn(op_name)(*this, args); + const size_t dispatch_end = dispatches_.size(); - { - auto mit = qkv_member.find(i); - if (mit != qkv_member.end()) { - QkvFusionGroup& g = qkv_groups[mit->second]; - 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] = 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 - // triple op would let a consumer (rope-q) read still-unwritten - // fresh_q -> garbage. - add_qkv_fused_dispatch(g); - } else if (i == g.op_idx[1]) { - g.sep_dispatch[1] = dispatch_range; - } else { - g.sep_dispatch[2] = dispatch_range; - } + const auto qkv_member = qkv_member_ops.find(i); + if (qkv_member != qkv_member_ops.end()) { + passes::QkvBk64Fusion& fusion = qkv_fusions[qkv_member->second]; + size_t member = 0; + while (member < 3 && fusion.op_indices[member] != i) { + member++; + } + if (member == 3 || dispatch_end <= dispatch_begin) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv: malformed member dispatch range"); } - auto lit = qkv_last.find(i); - if (lit != qkv_last.end()) { - // All 3 separate route ranges + the fused index are now known. - add_qkv_fused_hook(qkv_groups[lit->second]); + fusion.separate_begin[member] = dispatch_begin; + fusion.separate_end[member] = dispatch_end; + if (member == 0) { + passes::add_qkv_bk64_dispatch(*this, fusion); } } - - const size_t dispatch_end = dispatches_.size(); + const auto qkv_last = qkv_last_ops.find(i); + if (qkv_last != qkv_last_ops.end()) { + passes::add_qkv_bk64_resize_hook(*this, qkv_fusions[qkv_last->second]); + } if (i + 1 == chain->size() && op_name == kQ4gswLinearOpName && args.size() > kQ4gswOutputArg && dispatch_end > dispatch_begin) { @@ -1480,277 +1293,6 @@ WGPUBindGroupLayout WebGPUGraph::get_or_create_bgl( return bgl; } -void WebGPUGraph::add_qkv_fused_dispatch(QkvFusionGroup& g) { - const uint32_t N = g.Nq + g.Nk + g.Nv; // fused output width (3072) - - const auto& in = tensors_[g.input_id]; - const auto& out_q = tensors_[g.out_q]; - const auto& out_k = tensors_[g.out_k]; - const auto& out_v = tensors_[g.out_v]; - const auto& wq = tensors_[g.weight_q]; - const auto& wk = tensors_[g.weight_k]; - const auto& wv = tensors_[g.weight_v]; - const auto& sq = tensors_[g.scales_q]; - const auto& sk = tensors_[g.scales_k]; - const auto& sv = tensors_[g.scales_v]; - - // Buffers were repointed to FRESH distinct slots at the first triple op (see - // the build() op-walk), so out_q/k/v no longer alias. Live M from the shared - // input. - uint64_t in_numel = 1; - for (int64_t d : in.dims) { - in_numel *= static_cast(d); - } - const uint32_t M = static_cast(in_numel / g.K); - - // Fused weight [N, K_packed]: a byte-contiguous row-stack of Wq;Wk;Wv (q4gsw - // packs each output row independently along a shared K_packed, so stacking - // along N is a flat append -- bit-exact). Fused scales [num_groups, N]: a - // strided PER-GROUP-ROW gather (dest row stride N != the per-linear source - // strides padded_N_{q,k,v}), NOT a flat append. Both dtor-freed via scratch. - const uint64_t kp = static_cast(g.K_packed); // packed bytes / row - const uint64_t fs = sizeof(float); - WGPUBuffer fused_weight = create_scratch_buffer(static_cast(N) * kp); - WGPUBuffer fused_scales = create_scratch_buffer( - static_cast(g.num_groups) * N * sizeof(float)); - - // Sources are direct constants materialized in Phase 1 (or prepack outputs - // materialized earlier in Phase 3); all writes are already enqueued on - // queue_, so this build-time copy sees the materialized bytes. - WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(device_, nullptr); - wgpuCommandEncoderCopyBufferToBuffer( - enc, wq.buffer, 0, fused_weight, 0, static_cast(g.Nq) * kp); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - wk.buffer, - 0, - fused_weight, - static_cast(g.Nq) * kp, - static_cast(g.Nk) * kp); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - wv.buffer, - 0, - fused_weight, - static_cast(g.Nq + g.Nk) * kp, - static_cast(g.Nv) * kp); - for (uint32_t grp = 0; grp < g.num_groups; grp++) { - const uint64_t dst_row = static_cast(grp) * N * fs; - wgpuCommandEncoderCopyBufferToBuffer( - enc, - sq.buffer, - static_cast(grp) * g.padded_N_q * fs, - fused_scales, - dst_row, - static_cast(g.Nq) * fs); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - sk.buffer, - static_cast(grp) * g.padded_N_k * fs, - fused_scales, - dst_row + static_cast(g.Nq) * fs, - static_cast(g.Nk) * fs); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - sv.buffer, - static_cast(grp) * g.padded_N_v * fs, - fused_scales, - dst_row + static_cast(g.Nq + g.Nk) * fs, - static_cast(g.Nv) * fs); - } - WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(enc, nullptr); - wgpuQueueSubmit(queue_, 1, &cmd); - wgpuCommandBufferRelease(cmd); - wgpuCommandEncoderRelease(enc); - - // Params UBO (owned; rewritten by the resize hook). padded_N == N (fused - // scales row stride); has_bias == 0 (attention q/k/v are bias-less). - QkvFusedParams params = {}; - params.M = M; - params.N = N; - params.K = g.K; - params.K_packed = g.K_packed; - params.group_size = g.group_size; - params.padded_N = N; - params.has_bias = 0; - WGPUBufferDescriptor u_desc = {}; - u_desc.size = sizeof(QkvFusedParams); - u_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - u_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device_, &u_desc); - std::memcpy( - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(QkvFusedParams)), - ¶ms, - sizeof(QkvFusedParams)); - wgpuBufferUnmap(uniform_buffer); - add_uniform_buffer_bytes(sizeof(QkvFusedParams)); - - // 4-byte dummy for the fixed bias binding (has_bias == 0). - WGPUBuffer bias_dummy = create_scratch_buffer(4); - - // Bespoke 8-binding layout: 3 rw-storage outputs + 4 ro-storage + 1 uniform. - // One-off shader/bgl/pipeline owned by the dispatch (matches - // q4gsw_linear_impl). - WGPUBindGroupLayoutEntry entries[8] = {}; - for (uint32_t i = 0; i < 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_Storage; - } - for (uint32_t i = 3; i < 7; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[7].binding = 7; - entries[7].visibility = WGPUShaderStage_Compute; - entries[7].buffer.type = WGPUBufferBindingType_Uniform; - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 8; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc); - - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ4gswLinearGemmQkvFusedWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shader_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device_, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device_, &pipeline_desc); - - WGPUBindGroupEntry bg[8] = {}; - bg[0].binding = 0; - bg[0].buffer = out_q.buffer; - bg[0].size = out_q.nbytes; - bg[1].binding = 1; - bg[1].buffer = out_k.buffer; - bg[1].size = out_k.nbytes; - bg[2].binding = 2; - bg[2].buffer = out_v.buffer; - bg[2].size = out_v.nbytes; - bg[3].binding = 3; - bg[3].buffer = in.buffer; - bg[3].size = in.nbytes; - bg[4].binding = 4; - bg[4].buffer = fused_weight; - bg[4].size = static_cast(N) * kp; - bg[5].binding = 5; - bg[5].buffer = fused_scales; - bg[5].size = static_cast(g.num_groups) * N * fs; - bg[6].binding = 6; - bg[6].buffer = bias_dummy; - bg[6].size = 4; - bg[7].binding = 7; - bg[7].buffer = uniform_buffer; - bg[7].size = sizeof(QkvFusedParams); - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 8; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device_, &bg_desc); - - // 1D dispatch over ceil(M/BM) * ceil(N/BN) tiles (BM=BN=64), matching the - // kernel's nbN = ceil(N/64) tile decode (NOT grid-strided). - const uint32_t nbN = (N + 63u) / 64u; - const uint32_t nbM = (M + 63u) / 64u; - const size_t fused_idx = - add_dispatch({pipeline, bind_group, nbN * nbM, "linear_q4gsw_qkv_fused"}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); - own_uniform_buffer(uniform_buffer); - g.fused_dispatch = fused_idx; // consumed by add_qkv_fused_hook at the last op - g.fused_params = uniform_buffer; -} - -// 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 -// workgroup_count of 0 = no-op. -void WebGPUGraph::add_qkv_fused_hook(const QkvFusionGroup& g) { - const int input_id = g.input_id, out_q_id = g.out_q, out_k_id = g.out_k, - 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; - 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, - out_k_id, - out_v_id, - K, - Kp, - gs, - Nq, - Nk, - Nv, - Nf, - fused_idx, - separate_ranges, - params_buf](WebGPUGraph& gr) { - const auto& d = gr.cur_dims(input_id); - uint64_t numel = 1; - for (int64_t v : d) { - numel *= static_cast(v); - } - const uint32_t m = static_cast(numel / K); - std::vector oq = d; - oq.back() = static_cast(Nq); - std::vector ok = d; - ok.back() = static_cast(Nk); - std::vector ov = d; - ov.back() = static_cast(Nv); - gr.set_cur_dims(out_q_id, oq); - gr.set_cur_dims(out_k_id, ok); - gr.set_cur_dims(out_v_id, ov); - QkvFusedParams p = {}; - p.M = m; - p.N = Nf; - p.K = K; - p.K_packed = Kp; - p.group_size = gs; - p.padded_N = Nf; - p.has_bias = 0; - wgpuQueueWriteBuffer(gr.queue(), params_buf, 0, &p, sizeof(p)); - if (m > 1u) { - const uint32_t nbN2 = (Nf + 63u) / 64u; - 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; - 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; - } - }; - // Apply the max-shape route immediately. Resize hooks do not run before the - // first execution when cur_dims already equal the serialized max shape. - update_route(*this); - add_tensor_resize_hook(input_id, std::move(update_route)); -} - void WebGPUGraph::copy_inputs(const std::vector& inputs) { for (size_t i = 0; i < inputs.size() && i < input_ids_.size(); i++) { const InputData& in = inputs[i]; diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 705687c20a2..a70f099dc8e 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -523,6 +523,12 @@ class WebGPUGraph { return tensor_mem_obj_ids_[id]; } + // True if id is a prepack-routed constant with a recorded source (inline + // offset or named-data-map key); fusion passes require direct constants. + bool has_constant_source(int id) const { + return constant_sources_.count(id) != 0; + } + public: // True when the sdpa K/V cache is stored f16-packed (runtime opt-in). bool kv_f16() const { @@ -660,34 +666,6 @@ class WebGPUGraph { std::unordered_map bgl_cache_; size_t uniform_buffer_bytes_ = 0; - - // QKV-concat fusion: one detected attention q/k/v linear - // triple sharing an input activation (value ids + shapes), fused in build() - // into a single multi-output q4gsw GEMM that scatter-writes q/k/v. Only used - // during build(); inert (never populated) when no q/k/v triple matches. - struct QkvFusionGroup { - int input_id = -1; - int out_q = -1, out_k = -1, out_v = -1; - int weight_q = -1, weight_k = -1, weight_v = -1; - int scales_q = -1, scales_k = -1, scales_v = -1; - uint32_t Nq = 0, Nk = 0, Nv = 0; // 2048, 512, 512 - 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 - 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) - }; - // Concat the 3 packed weights (row-stack) + scales (strided gather) into - // fused buffers, then record ONE fused-GEMM dispatch (bespoke 8-binding - // layout) that writes the 3 original q/k/v output buffers, plus a 3-output - // resize hook. - void add_qkv_fused_dispatch(QkvFusionGroup& g); - void add_qkv_fused_hook(const QkvFusionGroup& g); }; } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index ff4064efa8f..d7ba7dbe2cd 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -91,13 +91,13 @@ #include #include #include -#include #include #include #include #include #include #include +#include #include #include #include @@ -709,13 +709,6 @@ constexpr std::array kShaderRegistry = {{ kQ4gswLinearCoop4BicolWorkgroupSizeY, kQ4gswLinearCoop4BicolWorkgroupSizeZ, }, - { - "q4gsw_linear_gemm_qkv_fused", - kQ4gswLinearGemmQkvFusedWGSL, - kQ4gswLinearGemmQkvFusedWorkgroupSizeX, - kQ4gswLinearGemmQkvFusedWorkgroupSizeY, - kQ4gswLinearGemmQkvFusedWorkgroupSizeZ, - }, { "q4gsw_linear_gemm_shmem", kQ4gswLinearGemmShmemWGSL, @@ -751,6 +744,13 @@ constexpr std::array kShaderRegistry = {{ kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY, kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ, }, + { + "q4gsw_qkv_bk64", + kQ4gswQkvBk64WGSL, + kQ4gswQkvBk64WorkgroupSizeX, + kQ4gswQkvBk64WorkgroupSizeY, + kQ4gswQkvBk64WorkgroupSizeZ, + }, { "q4gsw_requant", kQ4gswRequantWGSL, diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64.wgsl similarity index 63% rename from backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl rename to backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64.wgsl index c89924a67c1..aa977b65284 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64.wgsl @@ -1,12 +1,5 @@ enable f16; -// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM -// (vec4 activation load) that scatter-writes each output column range to a SEPARATE buffer -- -// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches; -// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned -// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE: -// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the -// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead -// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV). + @group(0) @binding(0) var t_out_q: array; @group(0) @binding(1) var t_out_k: array; @group(0) @binding(2) var t_out_v: array; @@ -25,10 +18,12 @@ struct Params { _pad: u32, } @group(0) @binding(7) var params: Params; -const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; + +// BK64 QKV 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; const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u; -var As: array; -var Bs: array; +var As: array; +var Bs: array; @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -44,18 +39,31 @@ fn main(@builtin(workgroup_id) wid: vec3, } 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 av = t_input[base >> 2u]; - As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y); - As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w); + 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 { - As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; - As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + 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; @@ -64,10 +72,17 @@ fn main(@builtin(workgroup_id) wid: vec3, 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]; + 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 = select(w1, w0, br < 8u); + let word = words[br >> 3u]; let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; Bs[br * BN + c] = f16(i32(nib) - 8) * scale; } @@ -91,7 +106,7 @@ fn main(@builtin(workgroup_id) wid: vec3, 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; // global fused column [0, 3072) + let c = col0 + lid.x * 4u + n; if (r < params.M && c < params.N) { var val = f32(acc[m][n]); if (params.has_bias != 0u) { val = val + t_bias[c]; } diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h similarity index 61% rename from backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h rename to backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h index 93243698a3b..371f79785ab 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h @@ -12,18 +12,11 @@ namespace executorch::backends::webgpu { -// @generated from q4gsw_linear_gemm_qkv_fused.wgsl - DO NOT EDIT. -// wgsl-sha256: 93e127e8ee4609d846015c8b75a600a29502e19a92bdf3a08e3429635f834085 -inline constexpr const char* kQ4gswLinearGemmQkvFusedWGSL = R"( +// @generated from q4gsw_qkv_bk64.wgsl - DO NOT EDIT. +// wgsl-sha256: d738762f00f79ca16cf1549d47e6d1f51155f50805eec5e7e6df3bc07ee309ee +inline constexpr const char* kQ4gswQkvBk64WGSL = R"( enable f16; -// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM -// (vec4 activation load) that scatter-writes each output column range to a SEPARATE buffer -- -// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches; -// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned -// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE: -// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the -// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead -// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV). + @group(0) @binding(0) var t_out_q: array; @group(0) @binding(1) var t_out_k: array; @group(0) @binding(2) var t_out_v: array; @@ -42,10 +35,12 @@ struct Params { _pad: u32, } @group(0) @binding(7) var params: Params; -const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; + +// BK64 QKV 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; const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u; -var As: array; -var Bs: array; +var As: array; +var Bs: array; @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -61,18 +56,31 @@ fn main(@builtin(workgroup_id) wid: vec3, } 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 av = t_input[base >> 2u]; - As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y); - As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w); + 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 { - As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; - As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + 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; @@ -81,10 +89,17 @@ fn main(@builtin(workgroup_id) wid: vec3, 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]; + 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 = select(w1, w0, br < 8u); + let word = words[br >> 3u]; let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; Bs[br * BN + c] = f16(i32(nib) - 8) * scale; } @@ -108,7 +123,7 @@ fn main(@builtin(workgroup_id) wid: vec3, 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; // global fused column [0, 3072) + let c = col0 + lid.x * 4u + n; if (r < params.M && c < params.N) { var val = f32(acc[m][n]); if (params.has_bias != 0u) { val = val + t_bias[c]; } @@ -121,8 +136,8 @@ fn main(@builtin(workgroup_id) wid: vec3, } )"; -inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeX = 16; -inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeY = 16; -inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeZ = 1; +inline constexpr uint32_t kQ4gswQkvBk64WorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswQkvBk64WorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswQkvBk64WorkgroupSizeZ = 1; } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/passes/QkvBk64.cpp b/backends/webgpu/runtime/passes/QkvBk64.cpp new file mode 100644 index 00000000000..73c129e45b9 --- /dev/null +++ b/backends/webgpu/runtime/passes/QkvBk64.cpp @@ -0,0 +1,444 @@ +/* + * 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. + */ + +#include + +#include +#include + +#include +#include +#include + +namespace executorch::backends::webgpu::passes { + +namespace { + +constexpr const char* kQ4gswLinearOpName = "et_vk.linear_q4gsw.default"; + +constexpr uint32_t kQkvQWidth = 2048u; +constexpr uint32_t kQkvKvWidth = 512u; +constexpr uint32_t kQkvFusedWidth = 3072u; +constexpr uint32_t kQkvK = 2048u; +constexpr uint32_t kQkvKPacked = 1024u; +constexpr uint32_t kQkvGroupSize = 64u; +constexpr uint32_t kQkvNumGroups = 32u; +constexpr uint32_t kQkvTile = 64u; + +// Uniform layout matching q4gsw_qkv_bk64.wgsl's Params struct. +struct QkvBk64Params { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t has_bias; + uint32_t _pad; +}; +static_assert(sizeof(QkvBk64Params) == 32); + +bool is_qkv_bk64_live_m(uint32_t m) { + return m == 128u || m == 508u || m == 512u; +} + +struct QkvBk64ResizeContext { + int input_id; + std::array output_ids; + std::array separate_begin; + std::array separate_end; + size_t fused_dispatch; + uint32_t max_m; + WGPUBuffer params_buffer; +}; + +void resize_qkv_bk64(WebGPUGraph& graph, const QkvBk64ResizeContext& context) { + const auto& input_dims = graph.cur_dims(context.input_id); + const uint64_t input_numel = utils::numel_of(input_dims); + if (input_dims.empty() || input_numel % kQkvK != 0u) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv(resize): malformed input shape"); + } + const uint64_t live_m = input_numel / kQkvK; + if (live_m == 0u || live_m > context.max_m) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv(resize): live M out of range"); + } + const uint32_t m = static_cast(live_m); + const uint32_t widths[3] = {kQkvQWidth, kQkvKvWidth, kQkvKvWidth}; + for (size_t i = 0; i < context.output_ids.size(); i++) { + std::vector output_dims = input_dims; + output_dims.back() = widths[i]; + graph.set_cur_dims(context.output_ids[i], output_dims); + } + + const QkvBk64Params params = { + m, + kQkvFusedWidth, + kQkvK, + kQkvKPacked, + kQkvGroupSize, + kQkvFusedWidth, + 0u, + 0u}; + wgpuQueueWriteBuffer( + graph.queue(), context.params_buffer, 0, ¶ms, sizeof(params)); + + const bool use_fused = is_qkv_bk64_live_m(m); + auto& fused = graph.dispatch_at(context.fused_dispatch); + fused.workgroup_count_x = use_fused + ? ((m + kQkvTile - 1u) / kQkvTile) * (kQkvFusedWidth / kQkvTile) + : 0u; + fused.workgroup_count_y = use_fused ? 1u : 0u; + if (use_fused) { + // The separate projections are inactive while the fused route is live. + for (size_t member = 0; member < context.separate_begin.size(); member++) { + for (size_t i = context.separate_begin[member]; + i < context.separate_end[member]; + i++) { + auto& dispatch = graph.dispatch_at(i); + dispatch.workgroup_count_x = 0u; + dispatch.workgroup_count_y = 0u; + } + } + } else { + // The separate projections' own resize hooks are registered before this + // one (Phase 3 processes each Q/K/V member before this fusion's combined + // hook) and unconditionally restore their live grids, so this hook + // normally has nothing to do here. That ordering isn't enforced by the + // type system, so fail loud rather than silently drop Q/K/V outputs if a + // future change ever violates it. + for (size_t member = 0; member < context.separate_begin.size(); member++) { + for (size_t i = context.separate_begin[member]; + i < context.separate_end[member]; + i++) { + const auto& dispatch = graph.dispatch_at(i); + if (dispatch.workgroup_count_x == 0u || + dispatch.workgroup_count_y == 0u) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv(resize): separate projection dispatch " + "was not restored before the QKV resize hook ran"); + } + } + } + } +} + +} // namespace + +bool qkv_bk64_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 16u && + limits.maxComputeWorkgroupStorageSize >= 16384u && + limits.maxComputeWorkgroupsPerDimension >= 384u; +} + +void detect_qkv_bk64_fusions( + const WebGPUGraph& graph, + const vkgraph::VkGraph* fb_graph, + int num_vals, + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops) { + const auto* chain = fb_graph->chain(); + if (!chain || !qkv_bk64_device_supported(graph.device())) { + return; + } + std::unordered_map> q4_ops_by_input; + std::vector input_order; + for (unsigned i = 0; i < chain->size(); i++) { + const auto* op = chain->Get(i); + const auto* args = op->args(); + if (op->name()->str() != kQ4gswLinearOpName || !args || args->size() != 6) { + continue; + } + const int input_id = static_cast(args->Get(0)); + if (q4_ops_by_input.count(input_id) == 0) { + input_order.push_back(input_id); + } + q4_ops_by_input[input_id].push_back(i); + } + + auto op_arg = [&](unsigned op_index, unsigned arg_index) { + return static_cast(chain->Get(op_index)->args()->Get(arg_index)); + }; + const auto& output_ids = graph.output_ids(); + auto is_graph_output = [&](int id) { + return std::find(output_ids.begin(), output_ids.end(), id) != + output_ids.end(); + }; + for (int input_id : input_order) { + const auto& ops = q4_ops_by_input.at(input_id); + if (ops.size() != 3 || input_id < 0 || input_id >= num_vals || + graph.get_value_type(input_id) != WebGPUGraph::ValueType::Tensor) { + continue; + } + + QkvBk64Fusion fusion; + fusion.input_id = input_id; + bool exact_args = true; + for (size_t member = 0; member < 3; member++) { + fusion.op_indices[member] = ops[member]; + fusion.weight_ids[member] = op_arg(ops[member], 1); + fusion.scale_ids[member] = op_arg(ops[member], 2); + fusion.output_ids[member] = op_arg(ops[member], 5); + const int group_size_id = op_arg(ops[member], 3); + const int bias_id = op_arg(ops[member], 4); + exact_args = exact_args && group_size_id >= 0 && + group_size_id < num_vals && + graph.get_value_type(group_size_id) == WebGPUGraph::ValueType::Int && + graph.get_int(group_size_id) == kQkvGroupSize && bias_id >= 0 && + bias_id < num_vals && + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Null; + } + if (!exact_args) { + continue; + } + + const std::array constant_ids = { + fusion.weight_ids[0], + fusion.weight_ids[1], + fusion.weight_ids[2], + fusion.scale_ids[0], + fusion.scale_ids[1], + fusion.scale_ids[2]}; + const std::unordered_set distinct_constants( + constant_ids.begin(), constant_ids.end()); + bool direct_constants = distinct_constants.size() == constant_ids.size(); + for (int id : constant_ids) { + direct_constants = direct_constants && id >= 0 && id < num_vals && + graph.get_value_type(id) == WebGPUGraph::ValueType::Tensor && + graph.has_constant_source(id) && + graph.get_tensor(id).buffer != nullptr; + } + if (!direct_constants) { + continue; + } + + const std::unordered_set distinct_outputs = { + fusion.output_ids[0], fusion.output_ids[1], fusion.output_ids[2]}; + bool outputs_ok = distinct_outputs.size() == 3; + for (int id : fusion.output_ids) { + outputs_ok = outputs_ok && id >= 0 && id < num_vals && + graph.get_value_type(id) == WebGPUGraph::ValueType::Tensor && + graph.mem_obj_id(id) >= 0 && !is_graph_output(id) && + utils::is_fp32_tensor(graph.get_tensor(id)); + } + if (!outputs_ok) { + continue; + } + + const auto& input = graph.get_tensor(input_id); + if (!utils::is_fp32_tensor(input) || input.dims.empty() || + input.dims.back() != kQkvK) { + continue; + } + const uint64_t input_numel = utils::numel_of(input.dims); + if (input_numel % kQkvK != 0u || input_numel / kQkvK < 128u || + input_numel / kQkvK > UINT32_MAX) { + continue; + } + fusion.max_m = static_cast(input_numel / kQkvK); + + const uint32_t widths[3] = {kQkvQWidth, kQkvKvWidth, kQkvKvWidth}; + bool exact_geometry = true; + for (size_t member = 0; member < 3; member++) { + const auto& weight = graph.get_tensor(fusion.weight_ids[member]); + const auto& scale = graph.get_tensor(fusion.scale_ids[member]); + const auto& output = graph.get_tensor(fusion.output_ids[member]); + exact_geometry = + exact_geometry && weight.dims.size() == 2 && + weight.dims[0] == widths[member] && weight.dims[1] == kQkvKPacked && + weight.nbytes == static_cast(widths[member]) * kQkvKPacked && + scale.dims.size() == 2 && scale.dims[0] == kQkvNumGroups && + scale.dims[1] == widths[member] && utils::is_fp32_tensor(scale) && + output.dims.size() == input.dims.size() && + std::equal( + input.dims.begin(), input.dims.end() - 1, output.dims.begin()) && + output.dims.back() == widths[member] && + utils::numel_of(output.dims) == + static_cast(fusion.max_m) * widths[member]; + } + if (!exact_geometry) { + continue; + } + + const size_t fusion_index = fusions.size(); + fusions.push_back(fusion); + first_ops[ops[0]] = fusion_index; + last_ops[ops[2]] = fusion_index; + for (unsigned op : ops) { + member_ops[op] = fusion_index; + } + } +} + +void retain_unclaimed_qkv_fusions( + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops, + std::unordered_set& claimed_ops) { + std::vector retained_fusions; + first_ops.clear(); + last_ops.clear(); + member_ops.clear(); + for (QkvBk64Fusion& fusion : fusions) { + bool overlaps = false; + for (unsigned op : fusion.op_indices) { + overlaps = overlaps || claimed_ops.count(op) != 0; + } + if (overlaps) { + continue; + } + const size_t fusion_index = retained_fusions.size(); + retained_fusions.push_back(std::move(fusion)); + const QkvBk64Fusion& retained = retained_fusions.back(); + first_ops[retained.op_indices[0]] = fusion_index; + last_ops[retained.op_indices[2]] = fusion_index; + for (unsigned op : retained.op_indices) { + member_ops[op] = fusion_index; + claimed_ops.insert(op); + } + } + fusions = std::move(retained_fusions); +} + +void add_qkv_bk64_dispatch(WebGPUGraph& graph, QkvBk64Fusion& fusion) { + const auto& input = graph.get_tensor(fusion.input_id); + const auto& output_q = graph.get_tensor(fusion.output_ids[0]); + const auto& output_k = graph.get_tensor(fusion.output_ids[1]); + const auto& output_v = graph.get_tensor(fusion.output_ids[2]); + const auto& weight_q = graph.get_tensor(fusion.weight_ids[0]); + const auto& weight_k = graph.get_tensor(fusion.weight_ids[1]); + const auto& weight_v = graph.get_tensor(fusion.weight_ids[2]); + const auto& scale_q = graph.get_tensor(fusion.scale_ids[0]); + const auto& scale_k = graph.get_tensor(fusion.scale_ids[1]); + const auto& scale_v = graph.get_tensor(fusion.scale_ids[2]); + + const size_t weight_row_bytes = kQkvKPacked; + WGPUBuffer fused_weight = graph.create_scratch_buffer( + static_cast(kQkvFusedWidth) * weight_row_bytes); + WGPUBuffer fused_scales = graph.create_scratch_buffer( + static_cast(kQkvNumGroups) * kQkvFusedWidth * sizeof(float)); + + WGPUCommandEncoder encoder = + wgpuDeviceCreateCommandEncoder(graph.device(), nullptr); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + weight_q.buffer, + 0, + fused_weight, + 0, + static_cast(kQkvQWidth) * weight_row_bytes); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + weight_k.buffer, + 0, + fused_weight, + static_cast(kQkvQWidth) * weight_row_bytes, + static_cast(kQkvKvWidth) * weight_row_bytes); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + weight_v.buffer, + 0, + fused_weight, + static_cast(kQkvQWidth + kQkvKvWidth) * weight_row_bytes, + static_cast(kQkvKvWidth) * weight_row_bytes); + for (uint32_t group = 0; group < kQkvNumGroups; group++) { + const uint64_t destination = + static_cast(group) * kQkvFusedWidth * sizeof(float); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + scale_q.buffer, + static_cast(group) * kQkvQWidth * sizeof(float), + fused_scales, + destination, + static_cast(kQkvQWidth) * sizeof(float)); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + scale_k.buffer, + static_cast(group) * kQkvKvWidth * sizeof(float), + fused_scales, + destination + static_cast(kQkvQWidth) * sizeof(float), + static_cast(kQkvKvWidth) * sizeof(float)); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + scale_v.buffer, + static_cast(group) * kQkvKvWidth * sizeof(float), + fused_scales, + destination + + static_cast(kQkvQWidth + kQkvKvWidth) * sizeof(float), + static_cast(kQkvKvWidth) * sizeof(float)); + } + WGPUCommandBuffer command = wgpuCommandEncoderFinish(encoder, nullptr); + wgpuQueueSubmit(graph.queue(), 1, &command); + wgpuCommandBufferRelease(command); + wgpuCommandEncoderRelease(encoder); + + const QkvBk64Params params = { + fusion.max_m, + kQkvFusedWidth, + kQkvK, + kQkvKPacked, + kQkvGroupSize, + kQkvFusedWidth, + 0u, + 0u}; + WGPUBuffer params_buffer = graph.create_params_buffer(params); + WGPUBuffer bias_dummy = graph.create_scratch_buffer(4); + + const bool initially_active = is_qkv_bk64_live_m(fusion.max_m); + const uint32_t workgroups = + ((fusion.max_m + kQkvTile - 1u) / kQkvTile) * (kQkvFusedWidth / kQkvTile); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "q4gsw_qkv_bk64"; + descriptor.kernel_name = "linear_q4gsw_bk64_qkv"; + descriptor.bindings = { + {output_q.buffer, 0u, output_q.nbytes}, + {output_k.buffer, 0u, output_k.nbytes}, + {output_v.buffer, 0u, output_v.nbytes}, + {input.buffer, 0u, input.nbytes}, + {fused_weight, + 0u, + static_cast(kQkvFusedWidth) * weight_row_bytes}, + {fused_scales, + 0u, + static_cast(kQkvNumGroups) * kQkvFusedWidth * sizeof(float)}, + {bias_dummy, 0u, 4u}, + {params_buffer, 0u, sizeof(QkvBk64Params)}}; + descriptor.grid = { + initially_active ? workgroups : 0u, initially_active ? 1u : 0u}; + fusion.fused_dispatch = graph.add_compute_dispatch(descriptor); + fusion.params_buffer = params_buffer; +} + +void add_qkv_bk64_resize_hook(WebGPUGraph& graph, const QkvBk64Fusion& fusion) { + const QkvBk64ResizeContext context = { + fusion.input_id, + {fusion.output_ids[0], fusion.output_ids[1], fusion.output_ids[2]}, + {fusion.separate_begin[0], + fusion.separate_begin[1], + fusion.separate_begin[2]}, + {fusion.separate_end[0], fusion.separate_end[1], fusion.separate_end[2]}, + fusion.fused_dispatch, + fusion.max_m, + fusion.params_buffer}; + resize_qkv_bk64(graph, context); + graph.add_tensor_resize_hook(fusion.input_id, [context](WebGPUGraph& g) { + resize_qkv_bk64(g, context); + }); +} + +} // namespace executorch::backends::webgpu::passes diff --git a/backends/webgpu/runtime/passes/QkvBk64.h b/backends/webgpu/runtime/passes/QkvBk64.h new file mode 100644 index 00000000000..5effc2ec04f --- /dev/null +++ b/backends/webgpu/runtime/passes/QkvBk64.h @@ -0,0 +1,71 @@ +/* + * 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 +#include + +#include +#include +#include + +namespace executorch::backends::webgpu::passes { + +// One matched QKV-BK64 (three q4gsw linears sharing one input, exact Llama +// Q/K/V geometry) pattern. +struct QkvBk64Fusion { + int input_id = -1; + int output_ids[3] = {-1, -1, -1}; + int weight_ids[3] = {-1, -1, -1}; + int scale_ids[3] = {-1, -1, -1}; + unsigned op_indices[3] = {0, 0, 0}; + size_t separate_begin[3] = {0, 0, 0}; + size_t separate_end[3] = {0, 0, 0}; + size_t fused_dispatch = SIZE_MAX; + WGPUBuffer params_buffer = nullptr; + uint32_t max_m = 0; +}; + +// True if the device meets the BK64 kernel's shader-f16/workgroup limits. +bool qkv_bk64_device_supported(WGPUDevice device); + +// Phase 2: scan fb_graph's op chain for three q4gsw-linear ops sharing one +// input in exact Q/K/V geometry. Populates `fusions` and the per-op index +// maps Phase 3 uses; does NOT filter against already-claimed op indices -- +// SwiGLU keeps precedence over an overlapping QKV candidate, so call +// retain_unclaimed_qkv_fusions after SwiGLU detection completes. +void detect_qkv_bk64_fusions( + const WebGPUGraph& graph, + const vkgraph::VkGraph* fb_graph, + int num_vals, + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops); + +// Drops any QKV candidate overlapping an op index already in `claimed_ops` +// (claimed by a higher-precedence pass), rebuilds the index maps for the +// retained set, and adds the retained candidates' op indices to +// `claimed_ops`. +void retain_unclaimed_qkv_fusions( + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops, + std::unordered_set& claimed_ops); + +// Emits the single fused q4gsw_qkv_bk64 dispatch for a matched pattern. +void add_qkv_bk64_dispatch(WebGPUGraph& graph, QkvBk64Fusion& fusion); + +// Registers the dynamic-resize hook that switches the fusion between its +// fused and separate-projection dispatches as the live M crosses the BK64 +// kernel's supported shapes. +void add_qkv_bk64_resize_hook(WebGPUGraph& graph, const QkvBk64Fusion& fusion); + +} // namespace executorch::backends::webgpu::passes diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 8c05fd379ea..59e816798db 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -299,6 +299,90 @@ constexpr float kBk64DownRtol = 8e-2f; constexpr float kBk64DownNrmse = 1.5e-2f; constexpr float kBk64DownTailNrmse = 2e-2f; +void expect_bk64_tensor( + const executorch::aten::Tensor& output, + const std::vector& golden, + int m_rows, + int width, + const std::string& label) { + const size_t numel = static_cast(m_rows) * width; + ASSERT_EQ(static_cast(output.numel()), numel) << label; + ASSERT_EQ(golden.size(), numel) << label; + const float* data = output.const_data_ptr(); + double error_sq_sum = 0.0; + double golden_sq_sum = 0.0; + double tail_error_sq_sum = 0.0; + double tail_golden_sq_sum = 0.0; + bool within_tolerance = true; + const size_t tail_begin = static_cast(m_rows - 1) * width; + for (size_t i = 0; i < numel; ++i) { + ASSERT_TRUE(std::isfinite(data[i])) << label << " i=" << i; + ASSERT_TRUE(std::isfinite(golden[i])) << label << " golden i=" << i; + const double error = static_cast(data[i]) - golden[i]; + const float abs_error = std::fabs(static_cast(error)); + const float rel_error = abs_error / std::fmax(std::fabs(golden[i]), 1e-6f); + within_tolerance &= abs_error <= kBk64Atol || rel_error <= kBk64Rtol; + error_sq_sum += error * error; + golden_sq_sum += static_cast(golden[i]) * golden[i]; + if (i >= tail_begin) { + tail_error_sq_sum += error * error; + tail_golden_sq_sum += static_cast(golden[i]) * golden[i]; + } + } + EXPECT_TRUE(within_tolerance) << label << " hybrid tolerance"; + ASSERT_GT(golden_sq_sum, 0.0) << label << " zero golden norm"; + EXPECT_LT(std::sqrt(error_sq_sum / golden_sq_sum), kBk64Nrmse) + << label << " full-output NRMSE"; + ASSERT_GT(tail_golden_sq_sum, 0.0) << label << " zero final-row golden norm"; + EXPECT_LT(std::sqrt(tail_error_sq_sum / tail_golden_sq_sum), kBk64TailNrmse) + << label << " final-row NRMSE"; + + for (size_t i : + {size_t{0}, static_cast(width - 1), tail_begin, numel - 1}) { + const float abs_error = std::fabs(data[i] - golden[i]); + const float rel_error = abs_error / std::fmax(std::fabs(golden[i]), 1e-6f); + EXPECT_TRUE(abs_error <= kBk64Atol || rel_error <= kBk64Rtol) + << label << " boundary i=" << i; + } +} + +void run_bk64_qkv( + Module& module, + int m_rows, + const char* prefix, + int q_width = kBk64N, + int k_width = kBk64KvN, + int v_width = kBk64KvN, + bool separate_v_input = false) { + const std::string base = + g_dir + "/" + prefix + ".S" + std::to_string(m_rows) + "."; + auto input = read_bin(base + "input.bin"); + ASSERT_EQ(input.size(), static_cast(m_rows) * kBk64K); + auto input_tensor = make_tensor_ptr({m_rows, kBk64K}, std::move(input)); + std::vector inputs{EValue(input_tensor)}; + decltype(input_tensor) v_input_tensor; + if (separate_v_input) { + auto v_input = read_bin(base + "v_input.bin"); + ASSERT_EQ(v_input.size(), static_cast(m_rows) * kBk64K); + v_input_tensor = make_tensor_ptr({m_rows, kBk64K}, std::move(v_input)); + inputs.emplace_back(v_input_tensor); + } + auto result = module.forward(inputs); + ASSERT_TRUE(result.ok()) << prefix << " M=" << m_rows; + ASSERT_EQ(result.get().size(), 3) << prefix << " M=" << m_rows; + const int widths[] = {q_width, k_width, v_width}; + const char* names[] = {"q", "k", "v"}; + for (size_t i = 0; i < 3; ++i) { + ASSERT_TRUE(result.get()[i].isTensor()) << prefix << " output " << names[i]; + expect_bk64_tensor( + result.get()[i].toTensor(), + read_bin(base + names[i] + ".bin"), + m_rows, + widths[i], + std::string(prefix) + " M=" + std::to_string(m_rows) + " " + names[i]); + } +} + void run_bk64_linear( Module& module, int m_rows, @@ -385,7 +469,7 @@ void run_swiglu_qkv_overlap(Module& module, int m_rows) { g_dir + "/dyn_swiglu_qkv_overlap.S" + std::to_string(m_rows) + "."; auto input = read_bin(base + "input.bin"); ASSERT_FALSE(input.empty()); - auto input_tensor = make_tensor_ptr({m_rows, kSwiGluK}, std::move(input)); + auto input_tensor = make_tensor_ptr({m_rows, kBk64K}, std::move(input)); auto result = module.forward({EValue(input_tensor)}); ASSERT_TRUE(result.ok()); ASSERT_EQ(result.get().size(), 2u); @@ -812,6 +896,38 @@ TEST(DynamicShape, QuantizedLinearBk64ReusedGraphAndFallbacks) { run_bk64_linear(kv_shape, 128, "dyn_linear_bk64_kv_shape", kBk64K, kBk64KvN); } +TEST(DynamicShape, QuantizedLinearBk64QkvReusedGraphAndFallbacks) { + Module candidate(g_dir + "/dyn_qkv_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok) << "load dyn_qkv_bk64.pte"; + for (int m_rows : {512, 511, 508, 128, 127, 16, 2, 1, 508, 512}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + } + + Module group32(g_dir + "/dyn_qkv_bk64_group32.pte"); + ASSERT_EQ(group32.load_forward(), Error::Ok); + run_bk64_qkv(group32, 128, "dyn_qkv_bk64_group32"); + + Module bias(g_dir + "/dyn_qkv_bk64_bias.pte"); + ASSERT_EQ(bias.load_forward(), Error::Ok); + run_bk64_qkv(bias, 128, "dyn_qkv_bk64_bias"); + + Module wrong_width(g_dir + "/dyn_qkv_bk64_wrong_width.pte"); + ASSERT_EQ(wrong_width.load_forward(), Error::Ok); + run_bk64_qkv( + wrong_width, 128, "dyn_qkv_bk64_wrong_width", kBk64N, kBk64N, kBk64KvN); + + Module different_input(g_dir + "/dyn_qkv_bk64_different_input.pte"); + ASSERT_EQ(different_input.load_forward(), Error::Ok); + run_bk64_qkv( + different_input, + 128, + "dyn_qkv_bk64_different_input", + kBk64N, + kBk64KvN, + kBk64KvN, + true); +} + #ifdef WGPU_BACKEND_ENABLE_PROFILING TEST(DynamicShape, QkvLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); @@ -825,14 +941,18 @@ TEST(DynamicShape, QkvLiveRoutesProfile) { run_qkv_routes(m, 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); + std::count(names.begin(), names.end(), "linear_q4gsw_bk64_qkv"), 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_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); + EXPECT_EQ( + std::count_if( + names.begin(), + names.end(), + [](const std::string& name) { + return name.rfind("linear_q4gsw", 0) == 0; + }), + 3); } } @@ -845,7 +965,10 @@ TEST(DynamicShape, QkvBk64LiveRoutesProfile) { WGPULimits limits = {}; if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || limits.maxComputeInvocationsPerWorkgroup < 256u || - limits.maxComputeWorkgroupStorageSize < 16384u) { + limits.maxComputeWorkgroupSizeX < 16u || + limits.maxComputeWorkgroupSizeY < 16u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 384u) { GTEST_SKIP() << "BK64 workgroup limits unavailable"; } Module module(g_dir + "/qkv_bk64_routes.pte"); @@ -854,7 +977,7 @@ TEST(DynamicShape, QkvBk64LiveRoutesProfile) { 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"), + std::count(names.begin(), names.end(), "linear_q4gsw_bk64_qkv"), m_rows > 1 ? 1 : 0); EXPECT_EQ( std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), @@ -963,6 +1086,101 @@ TEST(DynamicShape, QuantizedLinearBk64ProfileSoleWriter) { } } +TEST(DynamicShape, QuantizedLinearBk64QkvProfileSoleWriter) { + 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() << "QKV timestamp or shader-f16 capability unavailable"; + } + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || + limits.maxComputeInvocationsPerWorkgroup < 256u || + limits.maxComputeWorkgroupSizeX < 16u || + limits.maxComputeWorkgroupSizeY < 16u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 384u) { + GTEST_SKIP() << "QKV workgroup limits unavailable"; + } + const auto q4_profiles = [&]() { + std::vector names; + for (const auto& duration : context->querypool->results()) { + if (duration.kernel_name.rfind("linear_q4gsw", 0) == 0) { + names.push_back(duration.kernel_name); + } + } + return names; + }; + + Module candidate(g_dir + "/dyn_qkv_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 1) << "M=" << m_rows; + EXPECT_EQ(names[0], "linear_q4gsw_bk64_qkv") << "M=" << m_rows; + const auto& profile = context->querypool->results(); + const auto active = + std::find_if(profile.begin(), profile.end(), [](const auto& d) { + return d.kernel_name == "linear_q4gsw_bk64_qkv"; + }); + ASSERT_NE(active, profile.end()); + EXPECT_EQ(active->global_wg[0], m_rows == 128 ? 96u : 384u); + EXPECT_EQ(active->global_wg[1], 1u); + } + + for (int m_rows : {511, 127, 16, 2}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 3) << "M=" << m_rows; + EXPECT_FALSE(contains_name(names, "linear_q4gsw_bk64_qkv")); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 3) + << "M=" << m_rows; + } + + run_bk64_qkv(candidate, 1, "dyn_qkv_bk64"); + const auto decode_names = q4_profiles(); + ASSERT_EQ(decode_names.size(), 3); + EXPECT_EQ( + std::count( + decode_names.begin(), decode_names.end(), "linear_q4gsw_coop4_bicol"), + 3); + EXPECT_FALSE(contains_name(decode_names, "linear_q4gsw_bk64_qkv")); + + for (int m_rows : {508, 512}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 1) << "re-entry M=" << m_rows; + EXPECT_EQ(names[0], "linear_q4gsw_bk64_qkv") << "re-entry M=" << m_rows; + } + + struct NegativeRoute { + const char* fixture; + int q_width; + int k_width; + bool separate_v_input; + }; + for (const auto& negative : std::vector{ + {"dyn_qkv_bk64_group32", kBk64N, kBk64KvN, false}, + {"dyn_qkv_bk64_bias", kBk64N, kBk64KvN, false}, + {"dyn_qkv_bk64_wrong_width", kBk64N, kBk64N, false}, + {"dyn_qkv_bk64_different_input", kBk64N, kBk64KvN, true}}) { + Module module(g_dir + "/" + negative.fixture + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << negative.fixture; + run_bk64_qkv( + module, + 128, + negative.fixture, + negative.q_width, + negative.k_width, + kBk64KvN, + negative.separate_v_input); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 3) << negative.fixture; + EXPECT_FALSE(contains_name(names, "linear_q4gsw_bk64_qkv")) + << negative.fixture; + } +} + TEST(DynamicShape, CombinedLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || @@ -1370,12 +1588,20 @@ TEST(DynamicShape, SwiGluQkvOverlapProfile) { !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.maxComputeWorkgroupSizeX < 16u || + limits.maxComputeWorkgroupSizeY < 16u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 384u) { + GTEST_SKIP() << "QKV workgroup limits unavailable"; + } Module overlap(g_dir + "/dyn_swiglu_qkv_overlap.pte"); ASSERT_EQ(overlap.load_forward(), Error::Ok); run_swiglu_qkv_overlap(overlap, 128); const auto names = current_profile_names(); - EXPECT_EQ( - std::count(names.begin(), names.end(), "linear_q4gsw_qkv_fused"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_bk64_qkv"), 0); EXPECT_EQ(std::count(names.begin(), names.end(), "silu_mul_fused"), 1); EXPECT_EQ(std::count(names.begin(), names.end(), "mul"), 0); } 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 aa3a195713a..2cba504afac 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 @@ -110,6 +110,8 @@ def __init__( interleaved_projection: bool = False, qkv_overlap: bool = False, width: int = 8192, + input_width: int = 64, + group_size: int = 32, ) -> None: super().__init__() from torchao.quantization.granularity import PerGroup @@ -117,10 +119,12 @@ def __init__( def make_q4(seed: int, output_width: int = width): torch.manual_seed(seed) - linear = torch.nn.Linear(64, output_width, bias=False).eval() + linear = torch.nn.Linear(input_width, output_width, bias=False).eval() quantize_( linear, - IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)), + IntxWeightOnlyConfig( + weight_dtype=torch.int4, granularity=PerGroup(group_size) + ), ) return linear @@ -138,6 +142,7 @@ def make_q4(seed: int, output_width: int = width): self.separate_inputs = separate_inputs self.interleaved_projection = interleaved_projection self.qkv_overlap = qkv_overlap + self.input_width = input_width def forward(self, x: torch.Tensor, up_input: torch.Tensor | None = None): overlap_q = torch.sigmoid(self.overlap_q_proj(x)) if self.qkv_overlap else None @@ -372,6 +377,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: BK64_MAXM = 512 BK64_LIVE_M = (BK64_MAXM, 511, 508, 128, 127, 1) BK64_OPTIMIZED_M = (BK64_MAXM, 508, 128) +BK64_QKV_LIVE_M = (BK64_MAXM, 511, 508, 128, 127, 16, 2, 1) SWIGLU_MAXM = 512 @@ -382,7 +388,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: def _swiglu_inputs(model: SwiGluModule, m: int): - x = _ramp((m, SWIGLU_K)) + x = _ramp((m, model.input_width)) if model.separate_inputs: return x, torch.flip(x, dims=[-1]).contiguous() return (x,) @@ -467,7 +473,12 @@ def _export_dynamic_swiglu(out_dir: str) -> None: _export_swiglu_case( out_dir, "dyn_swiglu_qkv_overlap", - SwiGluModule(qkv_overlap=True, width=SWIGLU_QKV_OVERLAP_WIDTH), + SwiGluModule( + qkv_overlap=True, + width=SWIGLU_QKV_OVERLAP_WIDTH, + input_width=BK64_K, + group_size=BK64_GROUP, + ), [128], ) _export_swiglu_case( @@ -526,11 +537,12 @@ def _make_bk64_model( n: int = BK64_N, group: int = BK64_GROUP, bias: bool = False, + seed: int = 11, ) -> torch.nn.Module: from torchao.quantization.granularity import PerGroup from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ - torch.manual_seed(11) + torch.manual_seed(seed) model = torch.nn.Linear(k, n, bias=bias).eval() if model.bias is not None: with torch.no_grad(): @@ -565,6 +577,34 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.projection(x).reshape(1, x.shape[0], self.output_width) +class Bk64Qkv(torch.nn.Module): + def __init__( + self, + *, + widths=(BK64_N, BK64_KV_N, BK64_KV_N), + group: int = BK64_GROUP, + bias: bool = False, + separate_v_input: bool = False, + ) -> None: + super().__init__() + self.q = _make_bk64_model(n=widths[0], group=group, bias=bias, seed=21) + self.k = _make_bk64_model(n=widths[1], group=group, bias=bias, seed=22) + self.v = _make_bk64_model(n=widths[2], group=group, bias=bias, seed=23) + self.widths = widths + self.separate_v_input = separate_v_input + + def forward( + self, + x: torch.Tensor, + v_input: torch.Tensor | None = None, + ): + q = self.q(x).reshape(1, x.shape[0], self.widths[0]) + k = self.k(x).reshape(1, x.shape[0], self.widths[1]) + v_source = v_input if self.separate_v_input else x + v = self.v(v_source).reshape(1, v_source.shape[0], self.widths[2]) + return q, k, v + + def _export_bk64_program( model: torch.nn.Module, x: torch.Tensor, @@ -641,6 +681,66 @@ def _export_dynamic_bk64_linear_cases(out_dir: str) -> None: n=BK64_KV_N, live_m=[128], ) + _export_dynamic_bk64_qkv_cases(out_dir) + + +def _export_dynamic_bk64_qkv_case( + out_dir: str, + prefix: str, + model: Bk64Qkv, + live_m, +) -> None: + inputs = (_bk64_input(BK64_MAXM, BK64_K),) + if model.separate_v_input: + inputs += (torch.flip(inputs[0], dims=[-1]).contiguous(),) + m_dim = torch.export.Dim("m", min=1, max=BK64_MAXM) + dynamic_shapes = tuple({0: m_dim} for _ in inputs) + ep = torch.export.export(model.eval(), inputs, dynamic_shapes=dynamic_shapes) + if not any(node.target == torch.ops.aten.sym_size.int for node in ep.graph.nodes): + raise RuntimeError(f"{prefix}: dynamic QKV fixture lost aten.sym_size.int") + 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: + x = _bk64_input(m, BK64_K) + case_inputs = (x,) + if model.separate_v_input: + case_inputs += (torch.flip(x, dims=[-1]).contiguous(),) + outputs = ( + _bk64_golden(model.q, case_inputs[0]), + _bk64_golden(model.k, case_inputs[0]), + _bk64_golden(model.v, case_inputs[-1]), + ) + base = os.path.join(out_dir, f"{prefix}.S{m}.") + case_inputs[0].detach().numpy().astype(" None: + _export_dynamic_bk64_qkv_case( + out_dir, + "dyn_qkv_bk64", + Bk64Qkv(), + BK64_QKV_LIVE_M, + ) + for prefix, model in ( + ("dyn_qkv_bk64_group32", Bk64Qkv(group=32)), + ("dyn_qkv_bk64_bias", Bk64Qkv(bias=True)), + ( + "dyn_qkv_bk64_wrong_width", + Bk64Qkv(widths=(BK64_N, BK64_N, BK64_KV_N)), + ), + ( + "dyn_qkv_bk64_different_input", + Bk64Qkv(separate_v_input=True), + ), + ): + _export_dynamic_bk64_qkv_case(out_dir, prefix, model, [128]) def _export_static_linear(out_dir: str, m: int, prefix: str) -> None: @@ -1128,6 +1228,27 @@ def test_export_dynamic_rms(self) -> None: "dyn_linear_bk64_group32.pte", "dyn_linear_bk64_bias.pte", "dyn_linear_bk64_kv_shape.pte", + "dyn_qkv_bk64.pte", + "dyn_qkv_bk64.S512.input.bin", + "dyn_qkv_bk64.S512.q.bin", + "dyn_qkv_bk64.S512.k.bin", + "dyn_qkv_bk64.S512.v.bin", + "dyn_qkv_bk64.S511.input.bin", + "dyn_qkv_bk64.S511.q.bin", + "dyn_qkv_bk64.S511.k.bin", + "dyn_qkv_bk64.S511.v.bin", + "dyn_qkv_bk64.S508.q.bin", + "dyn_qkv_bk64.S508.k.bin", + "dyn_qkv_bk64.S508.v.bin", + "dyn_qkv_bk64.S128.q.bin", + "dyn_qkv_bk64.S127.q.bin", + "dyn_qkv_bk64.S16.q.bin", + "dyn_qkv_bk64.S2.q.bin", + "dyn_qkv_bk64.S1.q.bin", + "dyn_qkv_bk64_group32.pte", + "dyn_qkv_bk64_bias.pte", + "dyn_qkv_bk64_wrong_width.pte", + "dyn_qkv_bk64_different_input.pte", ] for name in expected: with self.subTest(artifact=name):