From cb607c5d9ff6ffd716803f708139d26d16bb1af9 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:28 -0700 Subject: [PATCH 01/12] [ExecuTorch][WebGPU] Replace broad QKV fusion with BK64 kernel Pull Request resolved: https://github.com/pytorch/executorch/pull/21131 The previously landed broad QKV fusion applied too widely and did not match the BK64 schedule now used for the ordinary projections. This corrective diff replaces it with a capability- and geometry-qualified BK64 kernel that fuses the exact Llama Q/K/V projection triple at live M128, M508, and M512, packing the constant weights and scales once and scattering the result into three distinct planner-safe outputs. Outside the accepted shapes it switches atomically back to the ordinary Steel and bicol routes, and it deletes the obsolete broad shader and header so only one QKV path remains. Mirrors Vulkan xplat/executorch/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.glsl for the per-projection GEMM; the three-output fusion itself is WebGPU-specific. Key changes: - Adds runtime/ops/quantized_linear/q4gsw_qkv_bk64.wgsl (+ generated header); removes the obsolete q4gsw_linear_gemm_qkv_fused shader and header. - WebGPUGraph.{cpp,h}: geometry and capability predicate, one-time constant packing, distinct Q/K/V outputs, and atomic fallback to Steel/bicol. ghstack-source-id: 411961443 @exported-using-ghexport Differential Revision: [D113171749](https://our.internmc.facebook.com/intern/diff/D113171749/) --- backends/webgpu/CMakeLists.txt | 1 + backends/webgpu/runtime/WebGPUGraph.cpp | 564 ++---------------- backends/webgpu/runtime/WebGPUGraph.h | 34 +- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 16 +- ...emm_qkv_fused.wgsl => q4gsw_qkv_bk64.wgsl} | 53 +- ...qkv_fused_wgsl.h => q4gsw_qkv_bk64_wgsl.h} | 65 +- backends/webgpu/runtime/passes/QkvBk64.cpp | 444 ++++++++++++++ backends/webgpu/runtime/passes/QkvBk64.h | 71 +++ .../webgpu/test/native/test_dynamic_shape.cpp | 246 +++++++- .../test_dynamic_shape_export.py | 131 +++- 10 files changed, 1019 insertions(+), 606 deletions(-) rename backends/webgpu/runtime/ops/quantized_linear/{q4gsw_linear_gemm_qkv_fused.wgsl => q4gsw_qkv_bk64.wgsl} (63%) rename backends/webgpu/runtime/ops/quantized_linear/{q4gsw_linear_gemm_qkv_fused_wgsl.h => q4gsw_qkv_bk64_wgsl.h} (61%) create mode 100644 backends/webgpu/runtime/passes/QkvBk64.cpp create mode 100644 backends/webgpu/runtime/passes/QkvBk64.h 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): From 2bd9d69aff95c319e97a812f3240fcdfa626e354 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:29 -0700 Subject: [PATCH 02/12] [ExecuTorch][WebGPU] Add Llama K16 online causal attention Pull Request resolved: https://github.com/pytorch/executorch/pull/21132 Materializing the full attention-weight matrix for Llama prefill is memory- and bandwidth-heavy and does not scale to longer sequences. This adds a single-pass online-softmax K16 causal-attention kernel for the exact Llama Hq32/Hkv8/G4/D64 geometry with fp16 KV, guarded by explicit adapter limits, so accepted prefill shapes through S512 never materialize attention weights. S1 keeps the existing FlashDecoding path, and unsupported geometry, storage, or capabilities fall back to the materialized implementation. No Vulkan analogue (WebGPU-specific): the Vulkan backend has only a materialized attention (compute the weights, then a separate multi-pass softmax), with no online-softmax kernel. Key changes: - runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl (+ generated header): the online-softmax single-pass causal kernel. - Sdpa.cpp: exact Llama geometry and capability guards, S512 transition handling, and fallback to the materialized route. ghstack-source-id: 411961446 @exported-using-ghexport Differential Revision: [D113171738](https://our.internmc.facebook.com/intern/diff/D113171738/) --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 10 +- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 418 +++++++++++------- .../streaming_attention_k16_causal_bound.wgsl | 268 +++++++++++ ...treaming_attention_k16_causal_bound_wgsl.h | 292 ++++++++++++ .../webgpu/test/native/test_dynamic_shape.cpp | 214 ++++++++- .../test_dynamic_shape_export.py | 219 +++++++++ 6 files changed, 1249 insertions(+), 172 deletions(-) create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index d7ba7dbe2cd..4a282f8266e 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -113,6 +113,7 @@ #include #include #include +#include #include #include #include @@ -148,7 +149,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -1003,6 +1004,13 @@ constexpr std::array kShaderRegistry = {{ kSqrtWorkgroupSizeY, kSqrtWorkgroupSizeZ, }, + { + "streaming_attention_k16_causal_bound", + kStreamingAttentionK16CausalBoundWGSL, + kStreamingAttentionK16CausalBoundWorkgroupSizeX, + kStreamingAttentionK16CausalBoundWorkgroupSizeY, + kStreamingAttentionK16CausalBoundWorkgroupSizeZ, + }, { "tanh", kTanhWGSL, diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 3c06d60747f..f36a84cc54f 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -6,17 +6,12 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include #include #include -#include -#include -#include -#include -#include #include -#include -#include #include @@ -34,6 +29,16 @@ namespace { constexpr int64_t kSdpaTileM = 4; constexpr int64_t kSdpaTileN = 4; +constexpr const char* kUpdateCacheShader = "update_cache"; +constexpr const char* kUpdateCacheHalfShader = "update_cache_half"; +constexpr const char* kAttnWeightsShader = "sdpa_compute_attn_weights"; +constexpr const char* kAttnWeightsHalfShader = "sdpa_compute_attn_weights_half"; +constexpr const char* kSoftmaxShader = "sdpa_softmax"; +constexpr const char* kComputeOutShader = "sdpa_compute_out"; +constexpr const char* kComputeOutHalfShader = "sdpa_compute_out_half"; +constexpr const char* kStreamingK16Shader = + "streaming_attention_k16_causal_bound"; + // Uniform param structs (all 16-byte aligned, matching the WGSL Params). struct UpdateCacheParams { uint32_t numel; @@ -75,6 +80,24 @@ struct ComputeOutParams { }; static_assert(sizeof(ComputeOutParams) == 32, "ComputeOutParams must be 32B"); +struct StreamingAttentionK16Params { + uint32_t S; + uint32_t context_len; + uint32_t input_pos; + uint32_t q_token_stride4; + uint32_t q_head_stride4; + uint32_t kv_token_stride4; + uint32_t kv_head_stride4; + uint32_t o_token_stride4; + uint32_t o_head_stride4; + uint32_t _pad0; + uint32_t _pad1; + uint32_t _pad2; +}; +static_assert( + sizeof(StreamingAttentionK16Params) == 48, + "StreamingAttentionK16Params must be 48B"); + struct SdpaLiveState { int64_t s; int64_t pos; @@ -83,10 +106,12 @@ struct SdpaLiveState { AttnWeightsParams attn_weights; SoftmaxParams softmax; ComputeOutParams compute_out; + StreamingAttentionK16Params streaming_k16; utils::WgCount update_cache_grid; utils::WgCount qk_grid; utils::WgCount softmax_grid; utils::WgCount av_grid; + utils::WgCount streaming_k16_grid; bool use_fd; SdpaFdDecodeState fd; }; @@ -148,111 +173,102 @@ static ComputeOutParams make_compute_out_params( return p; } -// A buffer + its byte size, for binding. -struct BufferBinding { - WGPUBuffer buffer; - uint64_t size; -}; +static StreamingAttentionK16Params make_streaming_attention_k16_params( + int64_t S, + int64_t context_len, + int64_t input_pos, + int64_t Hq, + int64_t Hkv, + int64_t D) { + StreamingAttentionK16Params p = {}; + p.S = static_cast(S); + p.context_len = static_cast(context_len); + p.input_pos = static_cast(input_pos); + p.q_token_stride4 = static_cast(Hq * D / 4); + p.q_head_stride4 = static_cast(D / 4); + p.kv_token_stride4 = static_cast(Hkv * D / 4); + p.kv_head_stride4 = static_cast(D / 4); + p.o_token_stride4 = static_cast(Hq * D / 4); + p.o_head_stride4 = static_cast(D / 4); + return p; +} + +static bool streaming_attention_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 4u && + limits.maxComputeWorkgroupStorageSize >= 14720u; +} -// Build one dispatch (pipeline + bind group) and record it on the graph. -void build_dispatch( +static utils::WgCount streaming_attention_k16_grid( + WGPUDevice device, + int64_t S, + int64_t Hkv, + int64_t g) { + const uint64_t groups_per_kv = + (static_cast(S) * static_cast(g) + 31u) / 32u; + const uint64_t workgroups = static_cast(Hkv) * groups_per_kv; + if (workgroups == 0u || workgroups > UINT32_MAX) { + throw std::runtime_error("WebGPU sdpa: K16 workgroup count exceeds uint32"); + } + if (workgroups > utils::queried_max_workgroups(device)) { + throw std::runtime_error( + "WebGPU sdpa: K16 workgroup count exceeds the 1D dispatch limit"); + } + return {static_cast(workgroups), 1u}; +} + +size_t add_sdpa_compute_dispatch( WebGPUGraph& graph, - const char* wgsl_source, - const BufferBinding* storage_bindings, - uint32_t n_storage, // includes the rw output at index 0 + const char* shader_name, + std::vector bindings, WGPUBuffer uniform_buffer, uint64_t uniform_size, - uint32_t workgroup_count_x, - uint32_t workgroup_count_y, + utils::WgCount grid, uint32_t wg_size, - bool retain_uniform = false, const char* kernel_name = "") { - WGPUDevice device = graph.device(); - - // Bind group layout: storage entries then the uniform. - constexpr uint32_t kMaxEntries = 8; - if (n_storage + 1 > kMaxEntries) { - throw std::runtime_error("WebGPU sdpa: n_storage exceeds kMaxEntries"); - } - const uint32_t uniform_binding = n_storage; - std::vector bindings; - bindings.reserve(n_storage + 1u); - for (uint32_t i = 0; i < n_storage; i++) { - bindings.push_back( - {i, - (i == 0) ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage, - storage_bindings[i].buffer, - storage_bindings[i].size}); - } - bindings.push_back( - {uniform_binding, - WGPUBufferBindingType_Uniform, - uniform_buffer, - uniform_size}); - - // All callers pass an override wg_size; a 0 would keep the shader default. - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - device, - wgsl_source, - bindings, - wg_size != 0 ? &wg_size_constant : nullptr, - wg_size != 0 ? 1u : 0u); - - graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count_x, - kernel_name, - workgroup_count_y}); - - if (retain_uniform) { - // Graph owns it so a resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); - } else { - // Drop our ref; the bind group keeps the uniform alive. - wgpuBufferRelease(uniform_buffer); + bindings.push_back({uniform_buffer, 0u, uniform_size}); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = shader_name; + descriptor.kernel_name = kernel_name; + descriptor.bindings = std::move(bindings); + if (wg_size != 0) { + descriptor.constants = {{"wg_size", static_cast(wg_size)}}; } + descriptor.grid = {grid.x, grid.y}; + return graph.add_compute_dispatch(descriptor); } // Dispatch one update_cache (K or V); returns the retained uniform buffer. static WGPUBuffer record_update_cache_dispatch( WebGPUGraph& graph, - WGPUDevice device, const WebGPUTensor& cache, const WebGPUTensor& src, uint64_t kv_numel, uint32_t kv_dst_offset, uint64_t cache_numel, uint32_t uc_wg, - bool retain_uniform, const char* label) { const uint32_t wgc = utils::compute_1d_workgroup_count( - device, static_cast(kv_numel), uc_wg, label); - UpdateCacheParams uc = + graph.device(), static_cast(kv_numel), uc_wg, label); + const UpdateCacheParams uc = make_update_cache_params(kv_numel, kv_dst_offset, cache_numel); - WGPUBuffer ubuf = graph.make_uniform_buffer(&uc, sizeof(uc)); - BufferBinding bindings[2] = { - {cache.buffer, cache.nbytes}, {src.buffer, src.nbytes}}; - const char* uc_src = kUpdateCacheWGSL; - if (graph.kv_f16()) { - uc_src = kUpdateCacheHalfWGSL; - } - build_dispatch( + WGPUBuffer ubuf = graph.create_params_buffer(uc); + const std::vector bindings = { + {cache.buffer, 0u, cache.nbytes}, {src.buffer, 0u, src.nbytes}}; + add_sdpa_compute_dispatch( graph, - uc_src, + graph.kv_f16() ? kUpdateCacheHalfShader : kUpdateCacheShader, bindings, - 2, ubuf, sizeof(uc), - wgc, - 1, + {wgc, 1u}, uc_wg, - retain_uniform, "update_cache"); return ubuf; } @@ -296,10 +312,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const size_t cn = k_cache.dims.size(); const int64_t Cmax = k_cache.dims[cn - 3]; - // Validate B == 1 (leading dims must all be 1). - for (size_t i = 0; i + 3 < qn; i++) { - if (q.dims[i] != 1) { - throw std::runtime_error("WebGPU sdpa: only batch size 1 is supported"); + // Validate B == 1 for every tensor (leading dims must all be 1). Rank-3 + // tensors are the equivalent squeezed-batch representation. + for (const WebGPUTensor* tensor : {&q, &k, &v, &k_cache, &v_cache, &out}) { + for (size_t i = 0; i + 3 < tensor->dims.size(); i++) { + if (tensor->dims[i] != 1) { + throw std::runtime_error("WebGPU sdpa: only batch size 1 is supported"); + } } } if (S <= 0 || Hq <= 0 || D <= 0 || Hkv <= 0 || Cmax <= 0) { @@ -335,6 +354,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (k_cache.dims != v_cache.dims) { throw std::runtime_error("WebGPU sdpa: k_cache and v_cache shape mismatch"); } + if (k_cache.dims[cn - 2] != Hkv) { + throw std::runtime_error( + "WebGPU sdpa: cache num_heads must match projected k/v"); + } + if (out.dims != q.dims) { + throw std::runtime_error("WebGPU sdpa: output shape must match q"); + } // fp32-only: validate byte counts against fp32 element counts. auto numel = [](const WebGPUTensor& t) { @@ -400,14 +426,26 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { } const WGPUDevice device = graph.device(); - const uint32_t uc_wg = - utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX); + const WGPUBuffer k16_buffers[] = { + q.buffer, k.buffer, v.buffer, k_cache.buffer, v_cache.buffer, out.buffer}; + bool k16_buffers_distinct = true; + for (size_t i = 0; i < 6; i++) { + for (size_t j = i + 1; j < 6; j++) { + k16_buffers_distinct = + k16_buffers_distinct && k16_buffers[i] != k16_buffers[j]; + } + } + const bool k16_eligible = graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && + D == 64 && scale == 0.125f && out.dims == q.dims && + k16_buffers_distinct && streaming_attention_k16_device_supported(device); + const uint32_t uc_wg = utils::clamp_workgroup_size( + device, get_webgpu_shader_info(kUpdateCacheShader).workgroup_size_x); const uint32_t qk_wg = utils::clamp_workgroup_size( - device, kSdpaComputeAttnWeightsWorkgroupSizeX); - const uint32_t av_wg = - utils::clamp_workgroup_size(device, kSdpaComputeOutWorkgroupSizeX); - const uint32_t sm_wg = - utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX); + device, get_webgpu_shader_info(kAttnWeightsShader).workgroup_size_x); + const uint32_t av_wg = utils::clamp_workgroup_size( + device, get_webgpu_shader_info(kComputeOutShader).workgroup_size_x); + const uint32_t sm_wg = utils::clamp_workgroup_size_pow2( + device, get_webgpu_shader_info(kSoftmaxShader).workgroup_size_x); const bool fd_eligible = D <= kSdpaFdMaxHeadDim; const int64_t pos_const = input_pos; @@ -430,7 +468,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { uc_wg, qk_wg, av_wg, - fd_eligible](WebGPUGraph& gr) { + fd_eligible, + k16_eligible](WebGPUGraph& gr) { SdpaLiveState state = {}; const auto& q_live_dims = gr.cur_dims(q_id); state.s = q_live_dims[qn - 3]; @@ -467,21 +506,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { static_cast(Hkv) * static_cast(D); const uint64_t cache_numel = static_cast(Cmax) * static_cast(Hkv) * static_cast(D); - const uint64_t aw_floats = static_cast(Hq) * - static_cast(state.s) * - static_cast(state.context_len); - const uint64_t qk_tiles = static_cast(Hq) * - static_cast(utils::div_up(state.s, kSdpaTileM)) * - static_cast(utils::div_up(state.context_len, kSdpaTileN)); - const uint64_t softmax_rows = - static_cast(Hq) * static_cast(state.s); - const uint64_t av_tiles = static_cast(Hq) * - static_cast(utils::div_up(state.s, kSdpaTileM)) * - static_cast(utils::div_up(D, kSdpaTileN)); if (kv_numel > UINT32_MAX || kv_offset > UINT32_MAX || - cache_numel > UINT32_MAX || aw_floats > UINT32_MAX || - qk_tiles > UINT32_MAX || softmax_rows > UINT32_MAX || - av_tiles > UINT32_MAX) { + cache_numel > UINT32_MAX) { throw std::runtime_error("WebGPU sdpa: live workload exceeds uint32"); } @@ -492,16 +518,43 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { state.softmax = make_softmax_params(Hq, state.s, state.context_len); state.compute_out = make_compute_out_params(state.s, Hq, Hkv, D, state.context_len, g); + if (k16_eligible) { + state.streaming_k16 = make_streaming_attention_k16_params( + state.s, state.context_len, state.pos, Hq, Hkv, D); + state.streaming_k16_grid = + streaming_attention_k16_grid(gr.device(), state.s, Hkv, g); + } state.update_cache_grid = { utils::compute_1d_workgroup_count( gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"), 1u}; - state.qk_grid = utils::compute_2d_workgroup_count( - gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); - state.softmax_grid = utils::compute_2d_workgroup_count( - gr.device(), static_cast(softmax_rows), 1, "softmax(resize)"); - state.av_grid = utils::compute_2d_workgroup_count( - gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + if (!k16_eligible) { + const uint64_t aw_floats = static_cast(Hq) * + static_cast(state.s) * + static_cast(state.context_len); + const uint64_t qk_tiles = static_cast(Hq) * + static_cast(utils::div_up(state.s, kSdpaTileM)) * + static_cast(utils::div_up(state.context_len, kSdpaTileN)); + const uint64_t softmax_rows = + static_cast(Hq) * static_cast(state.s); + const uint64_t av_tiles = static_cast(Hq) * + static_cast(utils::div_up(state.s, kSdpaTileM)) * + static_cast(utils::div_up(D, kSdpaTileN)); + if (aw_floats > UINT32_MAX || qk_tiles > UINT32_MAX || + softmax_rows > UINT32_MAX || av_tiles > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa: materialized workload exceeds uint32"); + } + state.qk_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); + state.softmax_grid = utils::compute_2d_workgroup_count( + gr.device(), + static_cast(softmax_rows), + 1, + "softmax(resize)"); + state.av_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + } state.use_fd = fd_eligible && state.s == 1; // make_sdpa_fd_decode_state requires D % 4 == 0; the op-level guard above // ("head_dim (D) must be a multiple of 4") rejects any other D before this @@ -514,32 +567,29 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { }; const SdpaLiveState initial_state = compute_live_state(graph); - const uint64_t aw_cap_floats = static_cast(Hq) * - static_cast(S) * - static_cast(dynamic_pos ? Cmax : context_len); + const uint64_t aw_cap_floats = k16_eligible + ? 0u + : static_cast(Hq) * static_cast(S) * + static_cast(dynamic_pos ? Cmax : context_len); const uint64_t aw_bytes = aw_cap_floats * sizeof(float); WGPUBuffer uc_k_buf = record_update_cache_dispatch( graph, - device, k_cache, k, initial_state.update_cache.numel, initial_state.update_cache.dst_offset, initial_state.update_cache.cache_numel, uc_wg, - true, "update_cache(K)"); WGPUBuffer uc_v_buf = record_update_cache_dispatch( graph, - device, v_cache, v, initial_state.update_cache.numel, initial_state.update_cache.dst_offset, initial_state.update_cache.cache_numel, uc_wg, - true, "update_cache(V)"); const size_t uc_k_idx = graph.num_dispatches() - 2; const size_t uc_v_idx = graph.num_dispatches() - 1; @@ -548,7 +598,9 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { graph.tensor_has_dynamic_dims(v_id); const bool dual_route = utils::should_record_sdpa_dual_route(fd_eligible, dynamic_sequence); - const bool record_materialized = dual_route || !initial_state.use_fd; + const bool record_k16 = k16_eligible && (dual_route || !initial_state.use_fd); + const bool record_materialized = + !k16_eligible && (dual_route || !initial_state.use_fd); const bool record_fd = dual_route || initial_state.use_fd; WGPUBuffer qk_buf = nullptr; @@ -566,70 +618,81 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { &graph, attn_weights_softmax); materialized_range.begin = graph.num_dispatches(); - qk_buf = graph.make_uniform_buffer( - &initial_state.attn_weights, sizeof(AttnWeightsParams)); - BufferBinding qk_bindings[3] = { - {attn_weights, aw_bytes}, - {q.buffer, q.nbytes}, - {k_cache.buffer, k_cache.nbytes}}; - const char* qk_src = graph.kv_f16() ? kSdpaComputeAttnWeightsHalfWGSL - : kSdpaComputeAttnWeightsWGSL; - build_dispatch( + qk_buf = graph.create_params_buffer(initial_state.attn_weights); + const std::vector qk_bindings = { + {attn_weights, 0u, aw_bytes}, + {q.buffer, 0u, q.nbytes}, + {k_cache.buffer, 0u, k_cache.nbytes}}; + add_sdpa_compute_dispatch( graph, - qk_src, + graph.kv_f16() ? kAttnWeightsHalfShader : kAttnWeightsShader, qk_bindings, - 3, qk_buf, sizeof(AttnWeightsParams), - initial_state.qk_grid.x, - initial_state.qk_grid.y, + initial_state.qk_grid, qk_wg, - true, "sdpa_compute_attn_weights"); qk_idx = graph.num_dispatches() - 1; - softmax_buf = graph.make_uniform_buffer( - &initial_state.softmax, sizeof(SoftmaxParams)); - BufferBinding softmax_bindings[2] = { - {attn_weights_softmax, aw_bytes}, {attn_weights, aw_bytes}}; - build_dispatch( + softmax_buf = graph.create_params_buffer(initial_state.softmax); + const std::vector softmax_bindings = { + {attn_weights_softmax, 0u, aw_bytes}, {attn_weights, 0u, aw_bytes}}; + add_sdpa_compute_dispatch( graph, - kSdpaSoftmaxWGSL, + kSoftmaxShader, softmax_bindings, - 2, softmax_buf, sizeof(SoftmaxParams), - initial_state.softmax_grid.x, - initial_state.softmax_grid.y, + initial_state.softmax_grid, sm_wg, - true, "sdpa_softmax"); softmax_idx = graph.num_dispatches() - 1; - av_buf = graph.make_uniform_buffer( - &initial_state.compute_out, sizeof(ComputeOutParams)); - BufferBinding av_bindings[3] = { - {out.buffer, out.nbytes}, - {attn_weights_softmax, aw_bytes}, - {v_cache.buffer, v_cache.nbytes}}; - const char* av_src = - graph.kv_f16() ? kSdpaComputeOutHalfWGSL : kSdpaComputeOutWGSL; - build_dispatch( + av_buf = graph.create_params_buffer(initial_state.compute_out); + const std::vector av_bindings = { + {out.buffer, 0u, out.nbytes}, + {attn_weights_softmax, 0u, aw_bytes}, + {v_cache.buffer, 0u, v_cache.nbytes}}; + add_sdpa_compute_dispatch( graph, - av_src, + graph.kv_f16() ? kComputeOutHalfShader : kComputeOutShader, av_bindings, - 3, av_buf, sizeof(ComputeOutParams), - initial_state.av_grid.x, - initial_state.av_grid.y, + initial_state.av_grid, av_wg, - true, "sdpa_compute_out"); av_idx = graph.num_dispatches() - 1; materialized_range.end = graph.num_dispatches(); } + WGPUBuffer k16_buf = nullptr; + size_t k16_idx = 0; + utils::DispatchRange k16_range = {}; + if (record_k16) { + k16_range.begin = graph.num_dispatches(); + k16_buf = graph.create_params_buffer(initial_state.streaming_k16); + const std::vector k16_bindings = { + {out.buffer, 0u, out.nbytes}, + {q.buffer, 0u, q.nbytes}, + {k_cache.buffer, 0u, k_cache.nbytes}, + {v_cache.buffer, 0u, v_cache.nbytes}}; + const utils::WgCount initial_grid = initial_state.use_fd + ? utils::WgCount{0u, 0u} + : initial_state.streaming_k16_grid; + add_sdpa_compute_dispatch( + graph, + kStreamingK16Shader, + k16_bindings, + k16_buf, + sizeof(StreamingAttentionK16Params), + initial_grid, + 0, + "sdpa_streaming_attention_k16_causal_bound"); + k16_idx = graph.num_dispatches() - 1; + k16_range.end = graph.num_dispatches(); + } + SdpaFdDecodeResources fd_resources = {}; size_t route_group = 0; if (record_fd) { @@ -637,14 +700,17 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { graph, q, k_cache, v_cache, out, initial_state.fd); } if (dual_route) { + const utils::DispatchRange prefill_range = + record_k16 ? k16_range : materialized_range; route_group = graph.register_dispatch_route_group( - {materialized_range, fd_resources.dispatch_range}); + {prefill_range, fd_resources.dispatch_range}); } auto refresh_state = [compute_live_state, q_id, out_id, dual_route, + record_k16, record_materialized, record_fd, fixed_use_fd = initial_state.use_fd, @@ -654,11 +720,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { qk_idx, softmax_idx, av_idx, + k16_idx, uc_k_buf, uc_v_buf, qk_buf, softmax_buf, av_buf, + k16_buf, fd_resources](WebGPUGraph& gr) { const SdpaLiveState state = compute_live_state(gr); @@ -686,6 +754,14 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { wgpuQueueWriteBuffer( gr.queue(), av_buf, 0, &state.compute_out, sizeof(state.compute_out)); } + if (record_k16) { + wgpuQueueWriteBuffer( + gr.queue(), + k16_buf, + 0, + &state.streaming_k16, + sizeof(state.streaming_k16)); + } if (record_fd) { write_sdpa_fd_decode_uniforms(gr.queue(), fd_resources, state.fd); } @@ -699,8 +775,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const std::vector active_grids = state.use_fd ? std::vector< utils::WgCount>{state.fd.split_grid, state.fd.reduce_grid} - : std::vector{ - state.qk_grid, state.softmax_grid, state.av_grid}; + : (record_k16 + ? std::vector{state.streaming_k16_grid} + : std::vector{ + state.qk_grid, state.softmax_grid, state.av_grid}); gr.select_dispatch_route(route_group, active_route, active_grids); } else if (state.use_fd) { if (!fixed_use_fd) { @@ -714,6 +792,12 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { state.fd.reduce_grid.x; gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_y = state.fd.reduce_grid.y; + } else if (record_k16) { + if (fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(k16_idx).workgroup_count_x = state.streaming_k16_grid.x; + gr.dispatch_at(k16_idx).workgroup_count_y = state.streaming_k16_grid.y; } else { if (fixed_use_fd) { throw std::runtime_error("WebGPU sdpa: static route changed"); diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl new file mode 100644 index 00000000000..5b8cf1de03c --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl @@ -0,0 +1,268 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 32u; +const HKV: u32 = 8u; +const G: u32 = 4u; +const D: u32 = 64u; +const D4: u32 = 16u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.125; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_k_tile: array, 256>; +var t_v_tile: array, 256>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_k_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max4(v: vec4) -> f32 { + return max(max(v.x, v.y), max(v.z, v.w)); +} + +fn exp_sum(v: vec4, maximum: f32) -> f32 { + let p = exp(v - vec4(maximum)); + return p.x + p.y + p.z + p.w; +} + +@compute @workgroup_size(32, 4, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc: vec4; + var output_acc: array, 4>; + score_acc = vec4(0.0); + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_k_tile[tile_index] = t_k_cache[cache_index]; + t_v_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_k_tile[tile_index] = vec4(0.0h); + t_v_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 4u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + let key2 = key0 + 2u; + let key3 = key0 + 3u; + score_acc = vec4( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + score_for(row, score_key_base + 2u, key2, row_valid, key2 < params.context_len, token), + score_for(row, score_key_base + 3u, key3, row_valid, key3 < params.context_len, token), + ); + let score_store = row * 4u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 4u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let tile_max = max(max(max4(s0), max4(s1)), max(max4(s2), max4(s3))); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 4u; + var score_block = 0u; + loop { + if (score_block >= 4u) { + break; + } + let probabilities = exp(t_scores[row_score_base + score_block] - vec4(new_m)); + let value_key_base = score_block * 4u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim0]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim0]) * probabilities.w; + output_acc[1] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim1]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim1]) * probabilities.w; + output_acc[2] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim2]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim2]) * probabilities.w; + output_acc[3] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim3]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim3]) * probabilities.w; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..255c6afc6e0 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h @@ -0,0 +1,292 @@ +/* + * 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 streaming_attention_k16_causal_bound.wgsl - DO NOT EDIT. +// wgsl-sha256: b1435c4f72834cb896eb4248a899aedeabbdcc772818afc419feca03e8957ffd +inline constexpr const char* kStreamingAttentionK16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 32u; +const HKV: u32 = 8u; +const G: u32 = 4u; +const D: u32 = 64u; +const D4: u32 = 16u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.125; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_k_tile: array, 256>; +var t_v_tile: array, 256>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_k_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max4(v: vec4) -> f32 { + return max(max(v.x, v.y), max(v.z, v.w)); +} + +fn exp_sum(v: vec4, maximum: f32) -> f32 { + let p = exp(v - vec4(maximum)); + return p.x + p.y + p.z + p.w; +} + +@compute @workgroup_size(32, 4, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc: vec4; + var output_acc: array, 4>; + score_acc = vec4(0.0); + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_k_tile[tile_index] = t_k_cache[cache_index]; + t_v_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_k_tile[tile_index] = vec4(0.0h); + t_v_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 4u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + let key2 = key0 + 2u; + let key3 = key0 + 3u; + score_acc = vec4( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + score_for(row, score_key_base + 2u, key2, row_valid, key2 < params.context_len, token), + score_for(row, score_key_base + 3u, key3, row_valid, key3 < params.context_len, token), + ); + let score_store = row * 4u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 4u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let tile_max = max(max(max4(s0), max4(s1)), max(max4(s2), max4(s3))); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 4u; + var score_block = 0u; + loop { + if (score_block >= 4u) { + break; + } + let probabilities = exp(t_scores[row_score_base + score_block] - vec4(new_m)); + let value_key_base = score_block * 4u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim0]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim0]) * probabilities.w; + output_acc[1] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim1]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim1]) * probabilities.w; + output_acc[2] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim2]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim2]) * probabilities.w; + output_acc[3] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim3]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim3]) * probabilities.w; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeX = 32; +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeY = 4; +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 59e816798db..9b0cf65a723 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -27,6 +27,8 @@ #include #include #include +#include +#include #include @@ -37,6 +39,7 @@ #include #include #include +#include #include #include @@ -497,7 +500,8 @@ void run_sdpa_case( int hq, int hkv, int d, - int cmax) { + int cmax, + float max_error_limit = 2e-3f) { const std::string b = g_dir + "/" + prefix + ".S" + std::to_string(s) + "."; auto q = read_bin(b + "q.bin"); auto k = read_bin(b + "k.bin"); @@ -535,7 +539,8 @@ void run_sdpa_case( << s << "," << hq << "," << d << "]"; std::vector got(attn, attn + numel); const float e = max_err(got, golden); - EXPECT_LT(e, 2e-3f) << "sdpa_dyn S=" << s << " max_err=" << e; + EXPECT_LT(e, max_error_limit) + << prefix << " S=" << s << " full-output max_err=" << e; } void run_sdpa(Module& m, int s) { @@ -548,6 +553,125 @@ void check_sdpa(int s) { run_sdpa(m, s); } +constexpr int kK16Hq = 32; +constexpr int kK16Hkv = 8; +constexpr int kK16D = 64; + +bool k16_device_supported() { + const auto* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 4u && + limits.maxComputeWorkgroupStorageSize >= 14720u; +} + +void load_sdpa_module(Module& module, bool f16_kv) { + if (!f16_kv) { + ASSERT_EQ(module.load_forward(), Error::Ok); + return; + } + executorch::runtime::BackendOptions<1> options; + ASSERT_EQ(options.set_option("enable_f16_kv_cache", true), Error::Ok); + executorch::runtime::LoadBackendOptionsMap option_map; + ASSERT_EQ(option_map.set_options("VulkanBackend", options.view()), Error::Ok); + ASSERT_EQ(module.load_forward(nullptr, nullptr, &option_map), Error::Ok); +} + +void run_k16_sdpa( + Module& module, + int s, + const char* prefix, + int hq = kK16Hq, + int hkv = kK16Hkv, + int d = kK16D, + bool prime = false) { + const std::string base = g_dir + "/" + prefix + + (prime ? ".prime." : ".S" + std::to_string(s) + "."); + auto q = read_bin(base + "q.bin"); + auto k = read_bin(base + "k.bin"); + auto v = read_bin(base + "v.bin"); + auto control = read_bin(base + "control.bin"); + auto golden = read_bin(base + "golden.bin"); + ASSERT_FALSE( + q.empty() || k.empty() || v.empty() || control.empty() || golden.empty()); + auto tq = make_tensor_ptr({1, s, hq, d}, std::move(q)); + auto tk = make_tensor_ptr({1, s, hkv, d}, std::move(k)); + auto tv = make_tensor_ptr({1, s, hkv, d}, std::move(v)); + const int control_size = static_cast(control.size()); + auto tcontrol = make_tensor_ptr({1, control_size}, std::move(control)); + auto result = + module.forward({EValue(tq), EValue(tk), EValue(tv), EValue(tcontrol)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << prefix << " S=" << s << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const size_t token_width = static_cast(hq) * d; + const size_t numel = static_cast(s) * token_width; + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + ASSERT_EQ(got.size(), golden.size()); + constexpr float kMaxError = 3e-3f; + EXPECT_LT(max_err(got, golden), kMaxError) + << prefix << " S=" << s << " full output"; + const std::set tokens = { + 0, std::min(15, s - 1), std::min(16, s - 1), s - 1}; + for (int token : tokens) { + float token_error = 0.0f; + const size_t begin = static_cast(token) * token_width; + for (size_t i = begin; i < begin + token_width; ++i) { + token_error = std::fmax(token_error, std::fabs(got[i] - golden[i])); + } + EXPECT_LT(token_error, kMaxError) + << prefix << " S=" << s << " causal token=" << token; + } +} + +void prime_k16_sdpa( + Module& module, + const char* prefix, + int hq = kK16Hq, + int hkv = kK16Hkv, + int d = kK16D) { + run_k16_sdpa(module, 12, prefix, hq, hkv, d, true); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void expect_sdpa_route( + const std::vector& names, + int s, + bool expect_k16) { + const bool expect_fd = s == 1; + const bool expect_materialized = !expect_fd && !expect_k16; + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ( + std::count( + names.begin(), + names.end(), + "sdpa_streaming_attention_k16_causal_bound"), + expect_k16 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "fd_split"), expect_fd ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "fd_reduce"), expect_fd ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_softmax"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + names.size(), + static_cast(2 + (expect_k16 ? 1 : (expect_fd ? 2 : 3)))); +} +#endif + void run_combined_routes(Module& m, int s) { const std::string b = g_dir + "/combined_routes.S" + std::to_string(s) + "."; auto x = read_bin(b + "x.bin"); @@ -1184,8 +1308,8 @@ TEST(DynamicShape, QuantizedLinearBk64QkvProfileSoleWriter) { TEST(DynamicShape, CombinedLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || - !context->timestamp_supported) { - GTEST_SKIP() << "timestamp queries unavailable"; + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; } Module m(g_dir + "/combined_routes.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; @@ -1301,6 +1425,27 @@ TEST(DynamicShape, SdpaWideMaterializedOnly) { } } +TEST(DynamicShape, K16CausalNumericsReusedGraph) { + if (!k16_device_supported()) { + GTEST_SKIP() << "K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa(module, "sdpa_k16_llama"); + for (int s : {512, 1, 508, 128, 127, 16, 1, 512}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + } +} + +TEST(DynamicShape, K16CacheHeadMismatchRejectedAtLoad) { + executorch::runtime::BackendOptions<1> options; + ASSERT_EQ(options.set_option("enable_f16_kv_cache", true), Error::Ok); + executorch::runtime::LoadBackendOptionsMap option_map; + ASSERT_EQ(option_map.set_options("VulkanBackend", options.view()), Error::Ok); + Module module(g_dir + "/sdpa_k16_bad_cache_heads.pte"); + EXPECT_NE(module.load_forward(nullptr, nullptr, &option_map), Error::Ok); +} + #ifdef WGPU_BACKEND_ENABLE_PROFILING TEST(DynamicShape, SdpaLiveRoutesProfile) { const auto* context = get_default_webgpu_context(); @@ -1325,6 +1470,67 @@ TEST(DynamicShape, SdpaLiveRoutesProfile) { std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); } } + +TEST(DynamicShape, K16CausalLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa(module, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), 12, true); + for (int s : {512, 128, 1, 508, 127, 1, 512}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), s, s > 1); + } +} + +TEST(DynamicShape, K16F32KvFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, false); + prime_k16_sdpa(module, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), s, false); + } +} + +TEST(DynamicShape, K16MetadataFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + struct NegativeCase { + const char* prefix; + int hq; + int hkv; + int d; + }; + for (const auto& negative : std::vector{ + {"sdpa_k16_wrong_geometry", 14, 2, 64}, + {"sdpa_k16_wrong_d", 32, 8, 128}, + {"sdpa_k16_wrong_scale", 32, 8, 64}}) { + Module module(g_dir + "/" + negative.prefix + ".pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, negative.prefix, negative.hq, negative.hkv, negative.d); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa( + module, s, negative.prefix, negative.hq, negative.hkv, negative.d); + expect_sdpa_route(current_profile_names(), s, false); + } + } +} #endif // K: dynamic embedding (int64 token ids) at several token counts. 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 2cba504afac..89a19f1fabc 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 @@ -318,6 +318,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0). _export_dynamic_sdpa(out_dir) + _export_dynamic_k16_sdpa_cases(out_dir) _export_combined_routes(out_dir) _export_dynamic_qkv_routes(out_dir) _export_dynamic_sdpa_wide(out_dir) @@ -771,6 +772,17 @@ def _export_static_linear(out_dir: str, m: int, prefix: str) -> None: SD_CMAX = 64 SD_MAXS = 64 +K16_HQ = 32 +K16_HKV = 8 +K16_D = 64 +K16_CMAX = 525 +K16_MAXS = 512 +K16_INPUT_POS = 13 +K16_PRIME_POS = 1 +K16_PRIME_S = K16_INPUT_POS - K16_PRIME_POS +K16_POS_CONTROL = K16_INPUT_POS +K16_LIVE_S = (K16_MAXS, 508, 128, 127, 16, 1) + def _export_dynamic_sdpa(out_dir: str) -> None: from executorch.backends.webgpu.test.ops.test_sdpa import ( @@ -810,6 +822,191 @@ def cfg(s: int) -> "SdpaConfig": print(f" golden sdpa_dyn S={s} (golden shape {tuple(g.shape)})") +def _write_f32_tensors(base, tensors) -> None: + for name, tensor in tensors: + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + ) + + def cfg(s: int) -> "SdpaConfig": + return SdpaConfig( + prefix, + hq, + hkv, + d, + s, + K16_CMAX, + K16_INPUT_POS, + ) + + q, k, v, kc, vc = _det_inputs(cfg(K16_MAXS)) + if cache_hkv is not None: + kc = torch.zeros(1, K16_CMAX, cache_hkv, d) + vc = torch.zeros_like(kc) + kc.zero_() + vc.zero_() + + class K16SdpaModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("k_cache", kc) + self.register_buffer("v_cache", vc) + + def forward(self, q, k, v, pos_control): + input_pos = pos_control.shape[1] + return torch.ops.llama.sdpa_with_kv_cache( + q, + k, + v, + self.k_cache, + self.v_cache, + input_pos, + q.shape[1], + None, + 0.0, + True, + scale, + ) + + model = K16SdpaModule().eval() + inputs = (q, k, v, torch.zeros(1, K16_POS_CONTROL)) + s_dim = torch.export.Dim(f"{prefix}_s", min=1, max=K16_MAXS) + pos_dim = torch.export.Dim(f"{prefix}_pos_control", min=1, max=K16_POS_CONTROL) + ep = torch.export.export( + model, + inputs, + dynamic_shapes=({1: s_dim}, {1: s_dim}, {1: s_dim}, {1: pos_dim}), + ) + sym_sizes = [ + node for node in ep.graph.nodes if node.target == torch.ops.aten.sym_size.int + ] + if len(sym_sizes) < 2: + raise RuntimeError(f"{prefix}: dynamic K16 fixture lost S/position symbols") + 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") + if expect_runtime_reject: + return + + def reference(live_cfg, q, k, v, kc, vc): + if scale is None: + return _golden(live_cfg, q, k, v, kc, vc) + context_len = live_cfg.s + live_cfg.input_pos + g = hq // hkv + k_full = torch.cat((kc[0, : live_cfg.input_pos].double(), k[0].double()), dim=0) + v_full = torch.cat((vc[0, : live_cfg.input_pos].double(), v[0].double()), dim=0) + q_heads = q[0].double().transpose(0, 1) + k_heads = k_full.repeat_interleave(g, dim=1).transpose(0, 1) + v_heads = v_full.repeat_interleave(g, dim=1).transpose(0, 1) + mask = torch.full((live_cfg.s, context_len), float("-inf"), dtype=torch.float64) + for token in range(live_cfg.s): + mask[token, : live_cfg.input_pos + token + 1] = 0.0 + golden = torch.nn.functional.scaled_dot_product_attention( + q_heads, k_heads, v_heads, attn_mask=mask, scale=scale + ) + return golden.transpose(0, 1).reshape(1, live_cfg.s, hq, d).float().contiguous() + + prime_cfg = SdpaConfig(prefix, hq, hkv, d, K16_PRIME_S, K16_CMAX, K16_PRIME_POS) + prime_q, prime_k, prime_v, prime_kc, prime_vc = _det_inputs(prime_cfg) + prime_kc.zero_() + prime_vc.zero_() + prime_golden = reference(prime_cfg, prime_q, prime_k, prime_v, prime_kc, prime_vc) + prime_base = os.path.join(out_dir, f"{prefix}.prime.") + _write_f32_tensors( + prime_base, + ( + ("q", prime_q), + ("k", prime_k), + ("v", prime_v), + ("control", torch.zeros(1, 1)), + ("golden", prime_golden), + ), + ) + + for s in live_s: + live_cfg = cfg(s) + q, k, v, kc, vc = _det_inputs(live_cfg) + kc.zero_() + vc.zero_() + kc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_k[0] + vc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_v[0] + golden = reference(live_cfg, q, k, v, kc, vc) + base = os.path.join(out_dir, f"{prefix}.S{s}.") + _write_f32_tensors( + base, + ( + ("q", q), + ("k", k), + ("v", v), + ("control", torch.zeros(1, K16_POS_CONTROL)), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ), + ) + print(f" golden {prefix} S={s}") + + +def _export_dynamic_k16_sdpa_cases(out_dir: str) -> None: + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_llama", + K16_HQ, + K16_HKV, + K16_LIVE_S, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_geometry", + 14, + 2, + (128, 1), + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_d", + K16_HQ, + K16_HKV, + (128, 1), + d=128, + scale=0.125, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_scale", + K16_HQ, + K16_HKV, + (128, 1), + scale=0.25, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_bad_cache_heads", + K16_HQ, + K16_HKV, + (), + cache_hkv=K16_HKV - 1, + expect_runtime_reject=True, + ) + + def _export_combined_routes(out_dir: str) -> None: from executorch.backends.webgpu.test.ops.test_quantized_linear import ( _make_quantized_model, @@ -1249,10 +1446,32 @@ def test_export_dynamic_rms(self) -> None: "dyn_qkv_bk64_bias.pte", "dyn_qkv_bk64_wrong_width.pte", "dyn_qkv_bk64_different_input.pte", + "sdpa_k16_bad_cache_heads.pte", ] for name in expected: with self.subTest(artifact=name): self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) + for prefix, live_s in ( + ("sdpa_k16_llama", K16_LIVE_S), + ("sdpa_k16_wrong_geometry", (128, 1)), + ("sdpa_k16_wrong_d", (128, 1)), + ("sdpa_k16_wrong_scale", (128, 1)), + ): + with self.subTest(artifact=f"{prefix}.pte"): + self.assertGreater( + os.path.getsize(os.path.join(d, f"{prefix}.pte")), 0 + ) + for kind in ("q", "k", "v", "control", "golden"): + name = f"{prefix}.prime.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) + for s in live_s: + for kind in ("q", "k", "v", "control", "golden"): + name = f"{prefix}.S{s}.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater( + os.path.getsize(os.path.join(d, name)), 0 + ) if __name__ == "__main__": From 4c723f541aaa1a5b180875deea0d654904f3df14 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:30 -0700 Subject: [PATCH 03/12] [ExecuTorch][WebGPU] Select emdawnwebgpu for Emscripten builds Pull Request resolved: https://github.com/pytorch/executorch/pull/21133 Building the WebGPU backend for the browser previously required a manual source-tree patch to pick the right WebGPU implementation. This makes the CMake configuration select the emdawnwebgpu port for Emscripten builds while native builds keep Dawn and their platform libraries, so the production kernel stack builds for the browser unpatched. Build-system integration only; no model or kernel behavior changes. Key changes: - CMakeLists.txt: choose emdawnwebgpu under Emscripten, Dawn otherwise. - test/test_cmake_configuration.py: contract test asserting the selection. ghstack-source-id: 411961450 @exported-using-ghexport Differential Revision: [D113171750](https://our.internmc.facebook.com/intern/diff/D113171750/) --- backends/webgpu/CMakeLists.txt | 35 ++++---- .../webgpu/test/test_cmake_configuration.py | 87 +++++++++++++++++++ 2 files changed, 107 insertions(+), 15 deletions(-) create mode 100644 backends/webgpu/test/test_cmake_configuration.py diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index f6ac8f3f3d1..536348ca69f 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -64,22 +64,27 @@ target_include_directories( target_link_libraries(webgpu_backend PRIVATE vulkan_schema executorch_core) -# Native WebGPU backend: Dawn (Tint) + SwiftShader; deps script sets Dawn_DIR. -# Native-only: browser/Emscripten builds use the system webgpu.h and never reach -# this find_package (root CMake gates it via EXECUTORCH_BUILD_WEBGPU). -# dawn::webgpu_dawn's link interface references Threads::Threads. -find_package(Threads REQUIRED) -find_package(Dawn REQUIRED) -set(WEBGPU_GPU_LIB dawn::webgpu_dawn) -target_link_libraries(webgpu_backend PUBLIC ${WEBGPU_GPU_LIB}) - -if(APPLE) - target_link_libraries( - webgpu_backend PRIVATE "-framework Metal" "-framework QuartzCore" - "-framework CoreGraphics" "-framework Foundation" - ) +# WASM gets its WebGPU implementation from emdawnwebgpu at executable link time. +# Native builds link Dawn (Tint) and the platform GPU libraries. +if(EMSCRIPTEN) + target_compile_options(webgpu_backend PUBLIC "--use-port=emdawnwebgpu") + # --use-port is also required at link time: the link step pulls in the port's + # headers, JS glue, and libraries for consumers that link webgpu_backend. + target_link_options(webgpu_backend PUBLIC "--use-port=emdawnwebgpu") else() - target_link_libraries(webgpu_backend PRIVATE dl m pthread) + find_package(Threads REQUIRED) + find_package(Dawn REQUIRED) + set(WEBGPU_GPU_LIB dawn::webgpu_dawn) + target_link_libraries(webgpu_backend PUBLIC ${WEBGPU_GPU_LIB}) + + if(APPLE) + target_link_libraries( + webgpu_backend PRIVATE "-framework Metal" "-framework QuartzCore" + "-framework CoreGraphics" "-framework Foundation" + ) + else() + target_link_libraries(webgpu_backend PRIVATE dl m pthread) + endif() endif() target_compile_options(webgpu_backend PRIVATE -fexceptions) diff --git a/backends/webgpu/test/test_cmake_configuration.py b/backends/webgpu/test/test_cmake_configuration.py new file mode 100644 index 00000000000..b42df74b5be --- /dev/null +++ b/backends/webgpu/test/test_cmake_configuration.py @@ -0,0 +1,87 @@ +# 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. + +import pathlib +import re +import unittest + + +def _conditional_branches(source: str, condition: str) -> tuple[str, str]: + lines = source.splitlines() + condition_pattern = re.compile( + rf"^\s*if\s*\(\s*{re.escape(condition)}\s*\)\s*$", + re.IGNORECASE, + ) + command_pattern = re.compile(r"^\s*([A-Za-z_][A-Za-z0-9_]*)\s*\(") + start = next( + (index for index, line in enumerate(lines) if condition_pattern.match(line)), + None, + ) + if start is None: + raise AssertionError(f"if({condition}) branch not found") + + depth = 1 + else_index = None + for index in range(start + 1, len(lines)): + match = command_pattern.match(lines[index]) + if match is None: + continue + command = match.group(1).lower() + if command == "if": + depth += 1 + elif command == "endif": + depth -= 1 + if depth == 0: + if else_index is None: + raise AssertionError(f"if({condition}) has no else() branch") + return ( + "\n".join(lines[start + 1 : else_index]), + "\n".join(lines[else_index + 1 : index]), + ) + elif command == "else" and depth == 1: + if else_index is not None: + raise AssertionError(f"if({condition}) has multiple else() branches") + else_index = index + raise AssertionError(f"if({condition}) has no matching endif()") + + +class TestCMakeConfiguration(unittest.TestCase): + def test_branch_parser_keeps_nested_conditionals_in_native_branch(self) -> None: + source = """ +if ( EMSCRIPTEN ) + wasm_command() +else() + if(APPLE) + apple_command() + else() + linux_command() + endif() +endif() +""" + wasm_branch, native_branch = _conditional_branches(source, "EMSCRIPTEN") + + self.assertIn("wasm_command()", wasm_branch) + self.assertIn("apple_command()", native_branch) + self.assertIn("linux_command()", native_branch) + + def test_emscripten_uses_port_instead_of_native_dawn(self) -> None: + cmake = pathlib.Path(__file__).parents[1] / "CMakeLists.txt" + wasm_branch, native_branch = _conditional_branches( + cmake.read_text(), "EMSCRIPTEN" + ) + port_flag = r'"--use-port=emdawnwebgpu"' + + self.assertRegex( + wasm_branch, + rf"target_compile_options\s*\(\s*webgpu_backend\s+PUBLIC\s+{port_flag}\s*\)", + ) + self.assertRegex( + wasm_branch, + rf"target_link_options\s*\(\s*webgpu_backend\s+PUBLIC\s+{port_flag}\s*\)", + ) + self.assertNotIn("find_package(Dawn", wasm_branch) + self.assertNotIn("--use-port=emdawnwebgpu", native_branch) + self.assertRegex(native_branch, r"find_package\s*\(\s*Dawn\s+REQUIRED\s*\)") From e535bd388ac76326610bf6560e4f8edeeafbd822 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:31 -0700 Subject: [PATCH 04/12] [ExecuTorch][WebGPU] Repair dynamic SDPA routing, add attestation Pull Request resolved: https://github.com/pytorch/executorch/pull/21134 With dynamic sequence positions, the backend could record only one SDPA route, so a decode following a dynamic-position prefill could dispatch the wrong kernel, and there was no way to confirm which kernel actually ran. This treats a dynamic SymInt position (not just a dynamic sequence dimension) as sufficient to record both SDPA routes, exposes active-kernel route attestation for correctness and performance harnesses, and reads the timestamp-query gate per execution so diagnostics can be enabled after module initialization. The attestation state and query are compiled only under the WGPU_BACKEND_ENABLE_PROFILING build flag, so production builds carry no additional state or cost, mirroring how the Vulkan backend gates its QueryPool behind ET_EVENT_TRACER_ENABLED. No Vulkan analogue for the routing fix (WebGPU runtime routing and observability). Key changes: - WebGPUGraph.{cpp,h}: expose the active-route attestation query (compile-gated behind WGPU_BACKEND_ENABLE_PROFILING) and read the timestamp-query gate per execution. - WebGPUUtils.h, Sdpa.cpp: treat a dynamic SymInt position (not only a dynamic sequence dimension) as sufficient to record both SDPA routes. ghstack-source-id: 411961454 @exported-using-ghexport Differential Revision: [D113171745](https://our.internmc.facebook.com/intern/diff/D113171745/) --- backends/webgpu/runtime/WebGPUDispatchMath.h | 5 +- backends/webgpu/runtime/WebGPUGraph.cpp | 81 ++++++++++++++++++- backends/webgpu/runtime/WebGPUGraph.h | 9 +++ backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 4 +- .../webgpu/test/native/test_dispatch_2d.cpp | 7 +- 5 files changed, 97 insertions(+), 9 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 9b87990aa8f..561bebc6b87 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -96,8 +96,9 @@ constexpr bool should_record_q4gsw_dual_route( constexpr bool should_record_sdpa_dual_route( bool fd_eligible, - bool has_dynamic_sequence) { - return fd_eligible && has_dynamic_sequence; + bool has_dynamic_sequence, + bool has_dynamic_position) { + return fd_eligible && (has_dynamic_sequence || has_dynamic_position); } constexpr bool is_q4gsw_bk64_eligible( diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index e729b3926a0..859e01e7a40 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1340,19 +1340,86 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { } } +#ifdef WGPU_BACKEND_ENABLE_PROFILING +// Profiling/attestation only; never compiled into a production build. Written +// during WebGPUGraph::execute without synchronization: the attestation +// harnesses that read them run one graph on one thread, so no atomics or +// locking are needed. To support concurrent profiled execution, make these +// per-instance behind a whole-record mutex (per-field atomics would not cover +// the conflict check's read-modify-write across both globals). +uint32_t g_last_route_mask = 0; +uint32_t g_last_route_conflict_count = 0; +#endif // WGPU_BACKEND_ENABLE_PROFILING + namespace { +#ifdef WGPU_BACKEND_ENABLE_PROFILING +constexpr uint32_t kRoutePrefill = 1u << 0; +constexpr uint32_t kRouteK16 = 1u << 1; +constexpr uint32_t kRouteMaterializedAttention = 1u << 2; +constexpr uint32_t kRouteT0Steel = 1u << 3; +constexpr uint32_t kRouteT1Bk64 = 1u << 4; +constexpr uint32_t kRouteT1Bk64Qkv = 1u << 5; +constexpr uint32_t kRouteT2PairedGateUp = 1u << 6; +constexpr uint32_t kRouteFusedSwiGlu = 1u << 7; +constexpr uint32_t kRouteGenericFallback = 1u << 8; +// bit 1u << 9 is intentionally reserved (a retired route) and left unused. +constexpr uint32_t kRouteFlashDecoding = 1u << 10; +constexpr uint32_t kRouteK16CausalBound = 1u << 11; +constexpr uint32_t kRouteBicolSubgroup = 1u << 12; +#endif // WGPU_BACKEND_ENABLE_PROFILING + // Bench gate: compiled out unless WGPU_BACKEND_ENABLE_PROFILING; then the // WEBGPU_TIMESTAMP_QUERY env var enables per-pass GPU timestamp queries. bool should_timestamp_query() { #ifdef WGPU_BACKEND_ENABLE_PROFILING - static const bool enabled = std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr; - return enabled; + return std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr; #else return false; #endif } } // namespace +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void WebGPUGraph::record_active_route(const std::string& kernel_name) { + uint32_t bits = 0; + if (kernel_name == "sdpa_streaming_attention_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound; + } else if (kernel_name == "sdpa_streaming_attention_k16") { + bits = kRoutePrefill | kRouteK16; + } else if ( + kernel_name.rfind("sdpa_compute_", 0) == 0 || + kernel_name == "sdpa_softmax") { + bits = kRoutePrefill | kRouteMaterializedAttention; + } else if (kernel_name == "fd_split" || kernel_name == "fd_reduce") { + bits = kRouteFlashDecoding; + } else if (kernel_name == "linear_q4gsw_coop4_bicol_subgroup") { + bits = kRouteBicolSubgroup; + } else if (kernel_name.rfind("linear_q4gsw_bk64_qkv", 0) == 0) { + bits = kRouteT1Bk64Qkv; + } else if (kernel_name.rfind("linear_q4gsw_bk64", 0) == 0) { + bits = kRouteT1Bk64; + } else if (kernel_name.rfind("linear_q4gsw_paired_gate_up", 0) == 0) { + bits = kRouteT2PairedGateUp; + } else if (kernel_name == "silu_mul_fused") { + bits = kRouteFusedSwiGlu; + } else if (kernel_name.rfind("linear_q4gsw_steel", 0) == 0) { + bits = kRouteT0Steel; + } else if (kernel_name.rfind("linear_q4gsw", 0) == 0) { + bits = kRouteGenericFallback; + } + + constexpr uint32_t kAttentionRoutes = kRouteK16 | kRouteK16CausalBound | + kRouteMaterializedAttention | kRouteFlashDecoding; + const uint32_t new_attention = bits & kAttentionRoutes; + const uint32_t prior_attention = g_last_route_mask & kAttentionRoutes; + if (new_attention != 0 && prior_attention != 0 && + (new_attention & prior_attention) == 0) { + ++g_last_route_conflict_count; + } + g_last_route_mask |= bits; +} +#endif // WGPU_BACKEND_ENABLE_PROFILING + WebGPUExecutionPlan WebGPUGraph::make_execution_plan( const WebGPUGraphExecutionOptions& options) const { const size_t n = dispatches_.size(); @@ -1378,6 +1445,10 @@ WebGPUExecutionPlan WebGPUGraph::make_execution_plan( } size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + g_last_route_mask = 0; + g_last_route_conflict_count = 0; +#endif // WGPU_BACKEND_ENABLE_PROFILING const size_t n = dispatches_.size(); const size_t chunk = execute_config_.chunk_size; if (plan.copy_outputs.size() != output_copies_.size()) { @@ -1443,6 +1514,9 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { dispatch.copy_nbytes); continue; } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + record_active_route(dispatch.kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING WGPUComputePassDescriptor pass_desc = {}; #ifdef WGPU_BACKEND_ENABLE_PROFILING // tw must outlive BeginComputePass (the descriptor points at it). @@ -1530,6 +1604,9 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { dispatches_[i].copy_nbytes); continue; } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + record_active_route(dispatches_[i].kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING WGPUComputePassDescriptor pass_desc = {}; WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index a70f099dc8e..292f2ee0cf4 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -546,6 +546,10 @@ class WebGPUGraph { } private: +#ifdef WGPU_BACKEND_ENABLE_PROFILING + void record_active_route(const std::string& kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING + bool kv_f16_ = false; std::unordered_set kv_cache_ids_; WebGPUGraphConfig config_; @@ -668,4 +672,9 @@ class WebGPUGraph { size_t uniform_buffer_bytes_ = 0; }; +#ifdef WGPU_BACKEND_ENABLE_PROFILING +extern uint32_t g_last_route_mask; +extern uint32_t g_last_route_conflict_count; +#endif // WGPU_BACKEND_ENABLE_PROFILING + } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index f36a84cc54f..9a51f853caf 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -596,8 +596,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const bool dynamic_sequence = graph.tensor_has_dynamic_dims(q_id) || graph.tensor_has_dynamic_dims(k_id) || graph.tensor_has_dynamic_dims(v_id); - const bool dual_route = - utils::should_record_sdpa_dual_route(fd_eligible, dynamic_sequence); + const bool dual_route = utils::should_record_sdpa_dual_route( + fd_eligible, dynamic_sequence, dynamic_pos); const bool record_k16 = k16_eligible && (dual_route || !initial_state.use_fd); const bool record_materialized = !k16_eligible && (dual_route || !initial_state.use_fd); diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index 67ce74659eb..c9ce9648b43 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -217,9 +217,10 @@ TEST(DispatchRoute, RecordsEligibleQ4AndDynamicSdpaAlternates) { EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, true, false)); EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, false, true)); - EXPECT_FALSE(should_record_sdpa_dual_route(true, false)); - EXPECT_FALSE(should_record_sdpa_dual_route(false, true)); - EXPECT_TRUE(should_record_sdpa_dual_route(true, true)); + EXPECT_FALSE(should_record_sdpa_dual_route(true, false, false)); + EXPECT_FALSE(should_record_sdpa_dual_route(false, true, true)); + EXPECT_TRUE(should_record_sdpa_dual_route(true, true, false)); + EXPECT_TRUE(should_record_sdpa_dual_route(true, false, true)); } TEST(WebGPUGraphConfig, ParsesExactBooleanCompileOption) { From 26d4fb6169d8d3d465820c74fff232759a18452c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:32 -0700 Subject: [PATCH 05/12] [ExecuTorch][WebGPU] Add HuggingFace rotate-half RoPE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21135 **Add HuggingFace rotate-half RoPE and migrate RoPE to shared dispatch construction** Qwen models pair the first and second halves of each head vector rather than adjacent elements. This adds the rotate-half operator with dynamic start-position updates, full-dimension Q/K handling, generated WGSL, and strict frequency-table bounds. The shared RoPE handler also closes the dispatch-boilerplate review: it uses `graph.device()`, typed graph-owned parameter buffers, descriptor-driven bindings and pipelines, named validation and resize callbacks, generated shader-registry lookup, scoped ownership, and graph-owned workgroup recomputation. Interleaved shader math and output are unchanged. Key changes: - `rotary_embedding_hf.wgsl` and generated registry entry — one thread per pair for HuggingFace rotate-half. - `RotaryEmbedding.cpp` — named validation, typed resize contexts, shared dispatch descriptors, and one initial/resize grid picker per route. - Native tests — malformed input rejection, dynamic bounds, shrink/regrow reuse, and HF/interleaved lifecycle coverage. Co-authored-with: Claude Code. ghstack-source-id: 411961455 @exported-using-ghexport Differential Revision: [D113171746](https://our.internmc.facebook.com/intern/diff/D113171746/) --- .../runtime/ops/rope/RotaryEmbedding.cpp | 953 ++++++++---------- .../runtime/ops/rope/rotary_embedding_hf.wgsl | 24 +- .../ops/rope/rotary_embedding_hf_wgsl.h | 26 +- .../webgpu/scripts/test_webgpu_native_ci.sh | 8 + .../test/native/test_compute_dispatch.cpp | 137 +++ backends/webgpu/test/ops/test_rope_hf.py | 282 +++++- backends/webgpu/test/test_build_webgpu.sh | 12 + .../webgpu/test/test_native_ci_contract.py | 9 + backends/webgpu/test/test_webgpu_native.cpp | 479 +++++++++ backends/webgpu/test/test_wgsl_codegen.py | 22 + 10 files changed, 1388 insertions(+), 564 deletions(-) diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index 012875b21e7..29f00c5823f 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -7,23 +7,22 @@ */ #include +#include #include #include -#include -#include #include #include -#include -#include #include -#include namespace executorch::backends::webgpu { namespace { +constexpr const char* kRotaryShader = "rotary_embedding"; +constexpr const char* kRotaryHfShader = "rotary_embedding_hf"; + // Uniform layout matching the WGSL Params struct (16-byte aligned, 32 bytes). struct RotaryParams { uint32_t n_heads; @@ -37,184 +36,109 @@ struct RotaryParams { }; static_assert(sizeof(RotaryParams) == 32, "RotaryParams must be 32 bytes"); -// A rope dispatch: its param-uniform (rewritten on resize) and its index in the -// graph's dispatch list (so a resize hook can update the workgroup count). -struct RopeDispatch { - WGPUBuffer uniform; - size_t dispatch_index; +enum class RopeGridPolicy { + OneDimensional, + FoldedTwoDimensional, +}; + +struct RopeGridContext { + int tensor_id; + uint32_t workgroup_size; + RopeGridPolicy policy; + const char* op_name; }; -// Rotate one (x->out) with the shared shader; freqs shared between xq and xk. -RopeDispatch add_rope_dispatch( +WebGPUDispatchGrid pick_rope_grid( + const WebGPUGraph& graph, + const RopeGridContext& context) { + const uint32_t num_pairs = static_cast( + utils::numel_of(graph.cur_dims(context.tensor_id)) / 2u); + if (context.policy == RopeGridPolicy::OneDimensional) { + return { + utils::compute_1d_workgroup_count( + graph.device(), num_pairs, context.workgroup_size, context.op_name), + 1u}; + } + const utils::WgCount grid = utils::compute_2d_workgroup_count( + graph.device(), num_pairs, context.workgroup_size, context.op_name); + return {grid.x, grid.y}; +} + +void preflight_rope_grids( + const WebGPUGraph& graph, + const RopeGridContext& q_context, + const RopeGridContext& k_context) { + (void)pick_rope_grid(graph, q_context); + (void)pick_rope_grid(graph, k_context); +} + +template +WGPUBuffer add_rope_dispatch( WebGPUGraph& graph, - WGPUDevice device, - std::optional& shared_resources, - uint32_t wg_size, + const char* shader_name, + const char* kernel_name, const WebGPUTensor& x, const WebGPUTensor& out, const WebGPUTensor& freqs_cos, const WebGPUTensor& freqs_sin, - uint32_t n_heads, - uint32_t seq, - uint32_t head_dim, - uint32_t workgroup_count) { - const uint32_t half_dim = head_dim / 2u; - // out.dims == in.dims (asserted in impl), so this matches the caller's wgc. - const uint32_t num_pairs = - static_cast(utils::numel_of(out.dims) / 2u); - - RotaryParams params = {}; - params.n_heads = n_heads; - params.seq = seq; - params.head_dim = head_dim; - params.half_dim = half_dim; - params.num_pairs = num_pairs; - - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(RotaryParams); - uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - uniform_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); - void* mapped = - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(RotaryParams)); - std::memcpy(mapped, ¶ms, sizeof(RotaryParams)); - wgpuBufferUnmap(uniform_buffer); - graph.add_uniform_buffer_bytes(sizeof(RotaryParams)); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - const std::vector bindings = { - {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, - {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, - {2, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_cos.buffer, - freqs_cos.nbytes}, - {3, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_sin.buffer, - freqs_sin.nbytes}, - {4, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(RotaryParams)}, - }; - utils::ComputePipelineBundle bundle = shared_resources.has_value() - ? utils::make_compute_pipeline( - device, *shared_resources, bindings, &wg_size_constant, 1) - : utils::make_compute_pipeline( - device, kRotaryEmbeddingWGSL, bindings, &wg_size_constant, 1); - - const size_t dispatch_index = graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count, - "apply_rotary_emb"}); - if (!shared_resources.has_value()) { - shared_resources.emplace(std::move(bundle)); - } - - // Graph owns it so a resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); - return {uniform_buffer, dispatch_index}; + const Params& params, + int trigger_tensor_id, + const RopeGridContext& grid_context, + uint32_t wg_size) { + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = shader_name; + descriptor.kernel_name = kernel_name; + descriptor.bindings = { + {out.buffer, 0u, out.nbytes}, + {x.buffer, 0u, x.nbytes}, + {freqs_cos.buffer, 0u, freqs_cos.nbytes}, + {freqs_sin.buffer, 0u, freqs_sin.nbytes}, + {uniform_buffer, 0u, sizeof(Params)}}; + descriptor.constants = {{"wg_size", static_cast(wg_size)}}; + graph.add_dynamic_compute_dispatch( + descriptor, trigger_tensor_id, pick_rope_grid, grid_context); + return uniform_buffer; } -// Resize hook body: recompute S/num_pairs + both dispatches; out follows xq/xk. -void resize_rope( - WebGPUGraph& g, - int xq_id, - int xk_id, - int xq_out_id, - int xk_out_id, - uint32_t n_heads_q, - uint32_t n_heads_k, - uint32_t head_dim, - uint32_t half_dim, - uint32_t wg_size, - size_t q_idx, - size_t k_idx, - WGPUBuffer q_ubuf, - WGPUBuffer k_ubuf) { - const auto& qd = g.cur_dims(xq_id); - const auto& kd = g.cur_dims(xk_id); - if (qd.size() < 3 || kd.size() < 3) { - throw std::runtime_error("apply_rotary_emb(resize): q/k rank must be >= 3"); - } - const uint32_t s = static_cast(qd[qd.size() - 3]); - const uint64_t qn = utils::numel_of(qd); - const uint64_t kn = utils::numel_of(kd); - // pk = pq (seq=s); require k's seq == s, not silently q's. - if (static_cast(kd[kd.size() - 3]) != s) { - throw std::runtime_error( - "apply_rotary_emb(resize): q and k seq lengths differ"); - } - // freqs stay max-allocated; shader indexes by position (S = prefix). - RotaryParams pq = {}; - pq.n_heads = n_heads_q; - pq.seq = s; - pq.head_dim = head_dim; - pq.half_dim = half_dim; - pq.num_pairs = static_cast(qn / 2u); - RotaryParams pk = pq; - pk.n_heads = n_heads_k; - pk.num_pairs = static_cast(kn / 2u); - wgpuQueueWriteBuffer(g.queue(), q_ubuf, 0, &pq, sizeof(pq)); - wgpuQueueWriteBuffer(g.queue(), k_ubuf, 0, &pk, sizeof(pk)); - g.dispatch_at(q_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(qn / 2u), - wg_size, - "apply_rotary_emb(resize)"); - g.dispatch_at(k_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(kn / 2u), - wg_size, - "apply_rotary_emb(resize)"); - g.set_cur_dims(xq_out_id, qd); - g.set_cur_dims(xk_out_id, kd); -} - -// args: [xq, xk, freqs_cos, freqs_sin, out_list(ValueList[xq_out, xk_out])]. -void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { - const int xq_id = args.at(0); - const int xk_id = args.at(1); - const int freqs_cos_id = args.at(2); - const int freqs_sin_id = args.at(3); - - const std::vector& out_list = graph.get_value_list(args.at(4)); - if (out_list.size() != 2) { - throw std::runtime_error( - "WebGPU apply_rotary_emb: expected an output ValueList of size 2"); - } - - WGPUDevice device = graph.device(); - - const auto& xq = graph.get_tensor(xq_id); - const auto& xk = graph.get_tensor(xk_id); - const auto& freqs_cos = graph.get_tensor(freqs_cos_id); - const auto& freqs_sin = graph.get_tensor(freqs_sin_id); - const auto& xq_out = graph.get_tensor(out_list[0]); - const auto& xk_out = graph.get_tensor(out_list[1]); +struct RotaryGeometry { + uint32_t head_dim; + uint32_t seq; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t half_dim; + uint64_t xq_numel; + uint64_t xk_numel; +}; - // Vulkan shape contract: xq/xk (B,S,n_heads,head_dim), freqs (S,head_dim/2). +RotaryGeometry validate_rope_inputs( + const WebGPUTensor& xq, + const WebGPUTensor& xk, + const WebGPUTensor& freqs_cos, + const WebGPUTensor& freqs_sin, + const WebGPUTensor& xq_out, + const WebGPUTensor& xk_out) { if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { throw std::runtime_error("WebGPU apply_rotary_emb: malformed dims"); } - const uint32_t head_dim = static_cast(xq.dims.back()); - const uint32_t seq = static_cast(xq.dims[xq.dims.size() - 3]); - const uint32_t n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); - const uint32_t n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); + RotaryGeometry geometry = {}; + geometry.head_dim = static_cast(xq.dims.back()); + geometry.seq = static_cast(xq.dims[xq.dims.size() - 3]); + geometry.n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); + geometry.n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); const uint32_t seq_k = static_cast(xk.dims[xk.dims.size() - 3]); - const uint32_t half_dim = static_cast(freqs_cos.dims.back()); + geometry.half_dim = static_cast(freqs_cos.dims.back()); - if (head_dim == 0 || head_dim % 2 != 0) { + if (geometry.head_dim == 0 || geometry.head_dim % 2 != 0) { throw std::runtime_error( "WebGPU apply_rotary_emb: head_dim must be a nonzero multiple of 2"); } - if (static_cast(xk.dims.back()) != head_dim || seq_k != seq) { + if (static_cast(xk.dims.back()) != geometry.head_dim || + seq_k != geometry.seq) { throw std::runtime_error( "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"); } - if (half_dim * 2u != head_dim) { + if (geometry.half_dim * 2u != geometry.head_dim) { throw std::runtime_error( "WebGPU apply_rotary_emb: head_dim != 2 * freqs_cos last dim"); } @@ -222,127 +146,163 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "WebGPU apply_rotary_emb: freqs_cos and freqs_sin shapes differ"); } - if (xq.buffer == nullptr || xk.buffer == nullptr || freqs_cos.buffer == nullptr || freqs_sin.buffer == nullptr || xq_out.buffer == nullptr || xk_out.buffer == nullptr) { throw std::runtime_error("WebGPU apply_rotary_emb: null buffer binding"); } - // All tensors are fp32; output shapes equal their inputs. - const uint64_t xq_numel = utils::numel_of(xq.dims); - const uint64_t xk_numel = utils::numel_of(xk.dims); + geometry.xq_numel = utils::numel_of(xq.dims); + geometry.xk_numel = utils::numel_of(xk.dims); const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims); - if (freqs_numel != static_cast(seq) * half_dim || - xq.nbytes != xq_numel * sizeof(float) || - xk.nbytes != xk_numel * sizeof(float) || + if (freqs_numel != static_cast(geometry.seq) * geometry.half_dim || + xq.nbytes != geometry.xq_numel * sizeof(float) || + xk.nbytes != geometry.xk_numel * sizeof(float) || freqs_cos.nbytes != freqs_numel * sizeof(float) || freqs_sin.nbytes != freqs_numel * sizeof(float) || - xq_out.nbytes != xq_numel * sizeof(float) || - xk_out.nbytes != xk_numel * sizeof(float)) { + xq_out.nbytes != geometry.xq_numel * sizeof(float) || + xk_out.nbytes != geometry.xk_numel * sizeof(float)) { throw std::runtime_error( "WebGPU apply_rotary_emb: dtype/byte-size mismatch (all fp32) or " "freqs shape != [seq, head_dim/2]"); } + if (geometry.xq_numel > UINT32_MAX || geometry.xk_numel > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb: element index exceeds uint32 range"); + } + return geometry; +} - if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { +struct RotaryResizeContext { + int xq_id; + int xk_id; + int xq_out_id; + int xk_out_id; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t head_dim; + uint32_t half_dim; + WGPUBuffer q_uniform; + WGPUBuffer k_uniform; +}; + +// Resize hook body: update parameters and outputs; the graph owns grid refresh. +void resize_rope(WebGPUGraph& graph, const RotaryResizeContext& context) { + const auto& q_dims = graph.cur_dims(context.xq_id); + const auto& k_dims = graph.cur_dims(context.xk_id); + if (q_dims.size() < 3 || k_dims.size() < 3) { + throw std::runtime_error("apply_rotary_emb(resize): q/k rank must be >= 3"); + } + const uint32_t seq = static_cast(q_dims[q_dims.size() - 3]); + const uint64_t q_numel = utils::numel_of(q_dims); + const uint64_t k_numel = utils::numel_of(k_dims); + // pk = pq (seq=s); require k's seq == s, not silently q's. + if (static_cast(k_dims[k_dims.size() - 3]) != seq) { throw std::runtime_error( - "WebGPU apply_rotary_emb: pair count exceeds uint32 dispatch range"); + "apply_rotary_emb(resize): q and k seq lengths differ"); } + // freqs stay max-allocated; shader indexes by position (S = prefix). + RotaryParams q_params = {}; + q_params.n_heads = context.n_heads_q; + q_params.seq = seq; + q_params.head_dim = context.head_dim; + q_params.half_dim = context.half_dim; + q_params.num_pairs = static_cast(q_numel / 2u); + RotaryParams k_params = q_params; + k_params.n_heads = context.n_heads_k; + k_params.num_pairs = static_cast(k_numel / 2u); + wgpuQueueWriteBuffer( + graph.queue(), context.q_uniform, 0, &q_params, sizeof(q_params)); + wgpuQueueWriteBuffer( + graph.queue(), context.k_uniform, 0, &k_params, sizeof(k_params)); + graph.set_cur_dims(context.xq_out_id, q_dims); + graph.set_cur_dims(context.xk_out_id, k_dims); +} - const uint32_t wg_size = - utils::clamp_workgroup_size(device, kRotaryEmbeddingWorkgroupSizeX); - // Validate both dispatches before any GPU-object alloc (no leak on throw). - const uint32_t xq_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xq_numel / 2u), - wg_size, - "apply_rotary_emb"); - const uint32_t xk_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xk_numel / 2u), - wg_size, - "apply_rotary_emb"); +// args: [xq, xk, freqs_cos, freqs_sin, out_list(ValueList[xq_out, xk_out])]. +void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { + const int xq_id = args.at(0); + const int xk_id = args.at(1); + const int freqs_cos_id = args.at(2); + const int freqs_sin_id = args.at(3); + + const std::vector& out_list = graph.get_value_list(args.at(4)); + if (out_list.size() != 2) { + throw std::runtime_error( + "WebGPU apply_rotary_emb: expected an output ValueList of size 2"); + } - std::optional shared_resources; - RopeDispatch q_disp = add_rope_dispatch( + const auto& xq = graph.get_tensor(xq_id); + const auto& xk = graph.get_tensor(xk_id); + const auto& freqs_cos = graph.get_tensor(freqs_cos_id); + const auto& freqs_sin = graph.get_tensor(freqs_sin_id); + const auto& xq_out = graph.get_tensor(out_list[0]); + const auto& xk_out = graph.get_tensor(out_list[1]); + + const RotaryGeometry geometry = + validate_rope_inputs(xq, xk, freqs_cos, freqs_sin, xq_out, xk_out); + + const uint32_t wg_size = utils::clamp_workgroup_size( + graph.device(), get_webgpu_shader_info(kRotaryShader).workgroup_size_x); + const RopeGridContext q_grid = { + xq_id, wg_size, RopeGridPolicy::OneDimensional, "apply_rotary_emb"}; + const RopeGridContext k_grid = { + xk_id, wg_size, RopeGridPolicy::OneDimensional, "apply_rotary_emb"}; + preflight_rope_grids(graph, q_grid, k_grid); + + RotaryParams q_params = {}; + q_params.n_heads = geometry.n_heads_q; + q_params.seq = geometry.seq; + q_params.head_dim = geometry.head_dim; + q_params.half_dim = geometry.half_dim; + q_params.num_pairs = static_cast(geometry.xq_numel / 2u); + RotaryParams k_params = q_params; + k_params.n_heads = geometry.n_heads_k; + k_params.num_pairs = static_cast(geometry.xk_numel / 2u); + const WGPUBuffer q_uniform = add_rope_dispatch( graph, - device, - shared_resources, - wg_size, + kRotaryShader, + "apply_rotary_emb", xq, xq_out, freqs_cos, freqs_sin, - n_heads_q, - seq, - head_dim, - xq_wgc); - RopeDispatch k_disp = add_rope_dispatch( + q_params, + xq_id, + q_grid, + wg_size); + const WGPUBuffer k_uniform = add_rope_dispatch( graph, - device, - shared_resources, - wg_size, + kRotaryShader, + "apply_rotary_emb", xk, xk_out, freqs_cos, freqs_sin, - n_heads_k, - seq, - head_dim, - xk_wgc); - WGPUBuffer q_ubuf = q_disp.uniform; - WGPUBuffer k_ubuf = k_disp.uniform; - const size_t q_idx = q_disp.dispatch_index; - const size_t k_idx = k_disp.dispatch_index; - - // Dynamic shapes: recompute S/num_pairs + both dispatches; out follows xq/xk. - const int xq_out_id = out_list[0]; - const int xk_out_id = out_list[1]; + k_params, + xk_id, + k_grid, + wg_size); + // Register on both xq and xk so the recompute fires whichever is marked dirty // (q and k co-resize on S; resize_rope is idempotent, so a double-fire when // both are dirty is harmless). - auto rope_hook = [xq_id, - xk_id, - xq_out_id, - xk_out_id, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf](WebGPUGraph& g) { - resize_rope( - g, - xq_id, - xk_id, - xq_out_id, - xk_out_id, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf); - }; - graph.add_tensor_resize_hook(xq_id, rope_hook); - graph.add_tensor_resize_hook(xk_id, rope_hook); + const RotaryResizeContext resize_context = { + xq_id, + xk_id, + out_list[0], + out_list[1], + geometry.n_heads_q, + geometry.n_heads_k, + geometry.head_dim, + geometry.half_dim, + q_uniform, + k_uniform}; + graph.add_tensor_resize_hook(xq_id, resize_rope, resize_context); + graph.add_tensor_resize_hook(xk_id, resize_rope, resize_context); } -// HuggingFace rotate-half RoPE (Qwen3 etc.). Structural sibling of the -// interleaved path above (same one-thread-per-pair scalar dispatch, wg_size, -// and resize hook); differs only in element pairing (i with i+half_dim vs -// even/odd), a full [max_seq, rotary_dim] freqs table, and a start_pos offset. -// Mirrors Vulkan's et_vk.apply_rotary_emb_hf -// (backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp:211). - -// Uniform layout matching the HF WGSL Params struct (32 bytes). +// Mirrors Vulkan's full-dimension HuggingFace rotate-half RoPE. struct RotaryHfParams { uint32_t n_heads; uint32_t seq; @@ -355,187 +315,65 @@ struct RotaryHfParams { }; static_assert(sizeof(RotaryHfParams) == 32, "RotaryHfParams must be 32 bytes"); -RopeDispatch add_rope_hf_dispatch( - WebGPUGraph& graph, - uint32_t wg_size, - const WebGPUTensor& x, - const WebGPUTensor& out, - const WebGPUTensor& freqs_cos, - const WebGPUTensor& freqs_sin, - uint32_t n_heads, - uint32_t seq, - uint32_t head_dim, - uint32_t half_dim, - uint32_t rotary_dim, - uint32_t start_pos, - uint32_t workgroup_count) { - const uint32_t num_pairs = - static_cast(utils::numel_of(out.dims) / 2u); - - RotaryHfParams params = {}; - params.n_heads = n_heads; - params.seq = seq; - params.head_dim = head_dim; - params.half_dim = half_dim; - params.num_pairs = num_pairs; - params.rotary_dim = rotary_dim; - params.start_pos = start_pos; - - WGPUBuffer uniform_buffer = - utils::make_uniform(graph.device(), ¶ms, sizeof(RotaryHfParams)); - graph.add_uniform_buffer_bytes(sizeof(RotaryHfParams)); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - graph.device(), - kRotaryEmbeddingHfWGSL, - { - {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, - {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, - {2, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_cos.buffer, - freqs_cos.nbytes}, - {3, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_sin.buffer, - freqs_sin.nbytes}, - {4, - WGPUBufferBindingType_Uniform, - uniform_buffer, - sizeof(RotaryHfParams)}, - }, - &wg_size_constant, - 1); - - const size_t dispatch_index = graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count, - "apply_rotary_emb_hf"}); - - graph.own_uniform_buffer(uniform_buffer); - return {uniform_buffer, dispatch_index}; -} - -// Resize hook body: recompute S/num_pairs + (dynamic) start_pos for both -// dispatches; out follows xq/xk. Fires on xq/xk seq resize and, when start_pos -// is a runtime SymInt (KV-cache decode), on each start_pos change; idempotent. -void resize_rope_hf( - WebGPUGraph& g, - int xq_id, - int xk_id, - int xq_out_id, - int xk_out_id, - int start_pos_id, - bool dynamic_pos, - uint32_t baked_start_pos, - uint32_t n_heads_q, - uint32_t n_heads_k, - uint32_t head_dim, - uint32_t half_dim, - uint32_t rotary_dim, - uint32_t wg_size, - size_t q_idx, - size_t k_idx, - WGPUBuffer q_ubuf, - WGPUBuffer k_ubuf) { - const auto& qd = g.cur_dims(xq_id); - const auto& kd = g.cur_dims(xk_id); - if (qd.size() < 3 || kd.size() < 3) { - throw std::runtime_error( - "apply_rotary_emb_hf(resize): q/k rank must be >= 3"); - } - const uint32_t s = static_cast(qd[qd.size() - 3]); - const uint64_t qn = utils::numel_of(qd); - const uint64_t kn = utils::numel_of(kd); - if (static_cast(kd[kd.size() - 3]) != s) { - throw std::runtime_error( - "apply_rotary_emb_hf(resize): q and k seq lengths differ"); - } - uint32_t start_pos = baked_start_pos; - if (dynamic_pos) { - const int32_t pos = g.read_symint(start_pos_id); - if (pos < 0) { - throw std::runtime_error( - "apply_rotary_emb_hf(resize): start_pos must be non-negative"); - } - start_pos = static_cast(pos); - } - RotaryHfParams pq = {}; - pq.n_heads = n_heads_q; - pq.seq = s; - pq.head_dim = head_dim; - pq.half_dim = half_dim; - pq.num_pairs = static_cast(qn / 2u); - pq.rotary_dim = rotary_dim; - pq.start_pos = start_pos; - RotaryHfParams pk = pq; - pk.n_heads = n_heads_k; - pk.num_pairs = static_cast(kn / 2u); - wgpuQueueWriteBuffer(g.queue(), q_ubuf, 0, &pq, sizeof(pq)); - wgpuQueueWriteBuffer(g.queue(), k_ubuf, 0, &pk, sizeof(pk)); - g.dispatch_at(q_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(qn / 2u), - wg_size, - "apply_rotary_emb_hf(resize)"); - g.dispatch_at(k_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(kn / 2u), - wg_size, - "apply_rotary_emb_hf(resize)"); - g.set_cur_dims(xq_out_id, qd); - g.set_cur_dims(xk_out_id, kd); -} - -// Validated HF-rope shape, derived from the input tensors. -struct RopeHfShape { +struct RotaryHfGeometry { uint32_t head_dim; uint32_t seq; uint32_t n_heads_q; uint32_t n_heads_k; + uint32_t max_seq; uint32_t rotary_dim; uint32_t half_dim; uint64_t xq_numel; uint64_t xk_numel; }; -// Derive and validate the HF-rope input shapes; throws on any malformed input. -RopeHfShape validate_rope_hf_inputs( - const WebGPUTensor& xq, +RotaryHfGeometry validate_rope_hf_inputs( + const WebGPUTensor& x, const WebGPUTensor& xk, const WebGPUTensor& freqs_cos, const WebGPUTensor& freqs_sin, - const WebGPUTensor& xq_out, + const WebGPUTensor& x_out, const WebGPUTensor& xk_out) { - // Shape contract: xq/xk (B,S,n_heads,head_dim), freqs (max_seq, rotary_dim). - if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { + if (x.dims.size() < 3 || xk.dims.size() != x.dims.size() || + freqs_cos.dims.size() != 2) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: malformed dims"); } - const uint32_t head_dim = static_cast(xq.dims.back()); - const uint32_t seq = static_cast(xq.dims[xq.dims.size() - 3]); - const uint32_t n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); - const uint32_t n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); - const uint32_t seq_k = static_cast(xk.dims[xk.dims.size() - 3]); - const uint32_t max_seq = - static_cast(freqs_cos.dims[freqs_cos.dims.size() - 2]); - const uint32_t rotary_dim = static_cast(freqs_cos.dims.back()); - - if (head_dim == 0 || head_dim % 2 != 0) { + if (x_out.dims != x.dims || xk_out.dims != xk.dims) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: output shapes must match q/k inputs"); + } + for (size_t i = 0; i + 3 < x.dims.size(); i++) { + if (x.dims[i] != xk.dims[i]) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: q/k batch dimensions differ"); + } + } + const auto positive_u32 = [](int64_t value, const char* label) { + if (value <= 0 || static_cast(value) > UINT32_MAX) { + throw std::runtime_error( + std::string("WebGPU apply_rotary_emb_hf: invalid ") + label); + } + return static_cast(value); + }; + RotaryHfGeometry geometry = {}; + geometry.head_dim = positive_u32(x.dims.back(), "head_dim"); + geometry.seq = positive_u32(x.dims[x.dims.size() - 3], "sequence length"); + geometry.n_heads_q = + positive_u32(x.dims[x.dims.size() - 2], "query head count"); + geometry.n_heads_k = + positive_u32(xk.dims[xk.dims.size() - 2], "key head count"); + geometry.max_seq = positive_u32(freqs_cos.dims[0], "frequency row count"); + geometry.rotary_dim = positive_u32(freqs_cos.dims[1], "rotary_dim"); + if (geometry.head_dim % 2 != 0) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: head_dim must be a nonzero multiple of 2"); } - if (static_cast(xk.dims.back()) != head_dim || seq_k != seq) { + if (xk.dims.back() != static_cast(geometry.head_dim) || + xk.dims[xk.dims.size() - 3] != static_cast(geometry.seq)) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: xq/xk head_dim and seq must match"); } - // Full rotary only (rotary_dim == head_dim); partial-rotary passthrough is a - // documented follow-up (Qwen3 uses full RoPE). Throw rather than mis-rotate. - if (rotary_dim != head_dim) { + if (geometry.rotary_dim != geometry.head_dim) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: partial rotary (rotary_dim != head_dim) " "not supported"); @@ -544,51 +382,134 @@ RopeHfShape validate_rope_hf_inputs( throw std::runtime_error( "WebGPU apply_rotary_emb_hf: freqs_cos and freqs_sin shapes differ"); } - if (max_seq < seq) { + if (geometry.max_seq < geometry.seq) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: freqs max_seq < seq"); } - - if (xq.buffer == nullptr || xk.buffer == nullptr || + if (x.buffer == nullptr || xk.buffer == nullptr || freqs_cos.buffer == nullptr || freqs_sin.buffer == nullptr || - xq_out.buffer == nullptr || xk_out.buffer == nullptr) { + x_out.buffer == nullptr || xk_out.buffer == nullptr) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: null buffer binding"); } + const WebGPUTensor* tensors[] = { + &x, &xk, &freqs_cos, &freqs_sin, &x_out, &xk_out}; + for (const WebGPUTensor* tensor : tensors) { + if (tensor->is_int || tensor->elem_size != sizeof(float)) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: all tensors must be fp32"); + } + } - // All tensors are fp32; output shapes equal their inputs. - const uint64_t xq_numel = utils::numel_of(xq.dims); - const uint64_t xk_numel = utils::numel_of(xk.dims); + geometry.half_dim = geometry.rotary_dim / 2u; + geometry.xq_numel = utils::numel_of(x.dims); + geometry.xk_numel = utils::numel_of(xk.dims); const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims); - if (freqs_numel != static_cast(max_seq) * rotary_dim || - xq.nbytes != xq_numel * sizeof(float) || - xk.nbytes != xk_numel * sizeof(float) || + if (freqs_numel != + static_cast(geometry.max_seq) * geometry.rotary_dim || + x.nbytes != geometry.xq_numel * sizeof(float) || + xk.nbytes != geometry.xk_numel * sizeof(float) || freqs_cos.nbytes != freqs_numel * sizeof(float) || freqs_sin.nbytes != freqs_numel * sizeof(float) || - xq_out.nbytes != xq_numel * sizeof(float) || - xk_out.nbytes != xk_numel * sizeof(float)) { + x_out.nbytes != geometry.xq_numel * sizeof(float) || + xk_out.nbytes != geometry.xk_numel * sizeof(float)) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: dtype/byte-size mismatch (all fp32) or " "freqs shape != [max_seq, rotary_dim]"); } + if (geometry.xq_numel == 0 || geometry.xk_numel == 0 || + geometry.xq_numel > UINT32_MAX || geometry.xk_numel > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: element index exceeds uint32 range"); + } + return geometry; +} + +struct RotaryHfResizeContext { + int xq_id; + int xk_id; + int xq_out_id; + int xk_out_id; + int start_pos_id; + bool dynamic_pos; + uint32_t baked_start_pos; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t head_dim; + uint32_t half_dim; + uint32_t rotary_dim; + uint32_t max_seq; + WGPUBuffer q_uniform; + WGPUBuffer k_uniform; +}; - if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { +void resize_rope_hf(WebGPUGraph& graph, const RotaryHfResizeContext& context) { + const auto& q_dims = graph.cur_dims(context.xq_id); + const auto& k_dims = graph.cur_dims(context.xk_id); + if (q_dims.size() < 3 || k_dims.size() != q_dims.size()) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k rank must be >= 3"); + } + const int64_t seq_value = q_dims[q_dims.size() - 3]; + if (seq_value <= 0 || static_cast(seq_value) > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): invalid sequence length"); + } + const uint32_t seq = static_cast(seq_value); + if (k_dims[k_dims.size() - 3] != seq_value) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q and k seq lengths differ"); + } + if (q_dims.back() != static_cast(context.head_dim) || + k_dims.back() != static_cast(context.head_dim) || + q_dims[q_dims.size() - 2] != static_cast(context.n_heads_q) || + k_dims[k_dims.size() - 2] != static_cast(context.n_heads_k)) { throw std::runtime_error( - "WebGPU apply_rotary_emb_hf: pair count exceeds uint32 dispatch range"); - } - - return { - head_dim, - seq, - n_heads_q, - n_heads_k, - rotary_dim, - rotary_dim / 2u, - xq_numel, - xk_numel}; + "apply_rotary_emb_hf(resize): q/k head geometry changed"); + } + for (size_t i = 0; i + 3 < q_dims.size(); i++) { + if (q_dims[i] != k_dims[i]) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k batch dimensions differ"); + } + } + const uint64_t q_numel = utils::numel_of(q_dims); + const uint64_t k_numel = utils::numel_of(k_dims); + if (q_numel == 0 || k_numel == 0 || q_numel > UINT32_MAX || + k_numel > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): element index exceeds uint32 range"); + } + uint32_t start_pos = context.baked_start_pos; + if (context.dynamic_pos) { + const int64_t pos = graph.read_symint(context.start_pos_id); + if (pos < 0 || static_cast(pos) > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos must be non-negative"); + } + start_pos = static_cast(pos); + } + if (static_cast(start_pos) + seq > context.max_seq) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos + seq exceeds freqs max_seq"); + } + RotaryHfParams q_params = {}; + q_params.n_heads = context.n_heads_q; + q_params.seq = seq; + q_params.head_dim = context.head_dim; + q_params.half_dim = context.half_dim; + q_params.num_pairs = static_cast(q_numel / 2u); + q_params.rotary_dim = context.rotary_dim; + q_params.start_pos = start_pos; + RotaryHfParams k_params = q_params; + k_params.n_heads = context.n_heads_k; + k_params.num_pairs = static_cast(k_numel / 2u); + wgpuQueueWriteBuffer( + graph.queue(), context.q_uniform, 0, &q_params, sizeof(q_params)); + wgpuQueueWriteBuffer( + graph.queue(), context.k_uniform, 0, &k_params, sizeof(k_params)); + graph.set_cur_dims(context.xq_out_id, q_dims); + graph.set_cur_dims(context.xk_out_id, k_dims); } -// args: [xq, xk, freqs_cos, freqs_sin, start_pos, out_list(ValueList[xq_out, -// xk_out])]. freqs is the FULL [max_seq, rotary_dim] table (start_pos offsets -// into it), unlike the pre-sliced interleaved freqs. void apply_rotary_emb_hf_impl( WebGPUGraph& graph, const std::vector& args) { @@ -604,8 +525,6 @@ void apply_rotary_emb_hf_impl( "WebGPU apply_rotary_emb_hf: expected an output ValueList of size 2"); } - WGPUDevice device = graph.device(); - const auto& xq = graph.get_tensor(xq_id); const auto& xk = graph.get_tensor(xk_id); const auto& freqs_cos = graph.get_tensor(freqs_cos_id); @@ -613,24 +532,15 @@ void apply_rotary_emb_hf_impl( const auto& xq_out = graph.get_tensor(out_list[0]); const auto& xk_out = graph.get_tensor(out_list[1]); - const RopeHfShape shp = + const RotaryHfGeometry geometry = validate_rope_hf_inputs(xq, xk, freqs_cos, freqs_sin, xq_out, xk_out); - const uint32_t head_dim = shp.head_dim; - const uint32_t seq = shp.seq; - const uint32_t n_heads_q = shp.n_heads_q; - const uint32_t n_heads_k = shp.n_heads_k; - const uint32_t rotary_dim = shp.rotary_dim; - const uint32_t half_dim = shp.half_dim; - const uint64_t xq_numel = shp.xq_numel; - const uint64_t xk_numel = shp.xk_numel; - - // start_pos: build-time Int (baked) OR runtime SymInt (dynamic decode); - // mirrors sdpa's input_pos handling. + + // Decode uses a SymInt position; static graphs use an Int. int64_t start_pos = 0; const auto start_pos_type = graph.get_value_type(start_pos_id); const bool dynamic_pos = start_pos_type == WebGPUGraph::ValueType::SymInt; if (dynamic_pos) { - start_pos = graph.read_symint(start_pos_id); // build placeholder (e.g. 0) + start_pos = graph.read_symint(start_pos_id); } else if (start_pos_type == WebGPUGraph::ValueType::Int) { start_pos = graph.get_int(start_pos_id); } else { @@ -641,99 +551,82 @@ void apply_rotary_emb_hf_impl( throw std::runtime_error( "WebGPU apply_rotary_emb_hf: start_pos must be non-negative"); } + if (static_cast(start_pos) + geometry.seq > geometry.max_seq) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos + seq exceeds freqs max_seq"); + } - const uint32_t wg_size = - utils::clamp_workgroup_size(device, kRotaryEmbeddingHfWorkgroupSizeX); - // Validate both dispatches before any GPU-object alloc (no leak on throw). - const uint32_t xq_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xq_numel / 2u), + const uint32_t wg_size = utils::clamp_workgroup_size( + graph.device(), get_webgpu_shader_info(kRotaryHfShader).workgroup_size_x); + const RopeGridContext q_grid = { + xq_id, wg_size, - "apply_rotary_emb_hf"); - const uint32_t xk_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xk_numel / 2u), + RopeGridPolicy::FoldedTwoDimensional, + "apply_rotary_emb_hf"}; + const RopeGridContext k_grid = { + xk_id, wg_size, - "apply_rotary_emb_hf"); - - RopeDispatch q_disp = add_rope_hf_dispatch( + RopeGridPolicy::FoldedTwoDimensional, + "apply_rotary_emb_hf"}; + preflight_rope_grids(graph, q_grid, k_grid); + + RotaryHfParams q_params = {}; + q_params.n_heads = geometry.n_heads_q; + q_params.seq = geometry.seq; + q_params.head_dim = geometry.head_dim; + q_params.half_dim = geometry.half_dim; + q_params.num_pairs = static_cast(geometry.xq_numel / 2u); + q_params.rotary_dim = geometry.rotary_dim; + q_params.start_pos = static_cast(start_pos); + RotaryHfParams k_params = q_params; + k_params.n_heads = geometry.n_heads_k; + k_params.num_pairs = static_cast(geometry.xk_numel / 2u); + + const WGPUBuffer q_uniform = add_rope_dispatch( graph, - wg_size, + kRotaryHfShader, + "apply_rotary_emb_hf", xq, xq_out, freqs_cos, freqs_sin, - n_heads_q, - seq, - head_dim, - half_dim, - rotary_dim, - static_cast(start_pos), - xq_wgc); - RopeDispatch k_disp = add_rope_hf_dispatch( + q_params, + xq_id, + q_grid, + wg_size); + const WGPUBuffer k_uniform = add_rope_dispatch( graph, - wg_size, + kRotaryHfShader, + "apply_rotary_emb_hf", xk, xk_out, freqs_cos, freqs_sin, - n_heads_k, - seq, - head_dim, - half_dim, - rotary_dim, + k_params, + xk_id, + k_grid, + wg_size); + + const RotaryHfResizeContext resize_context = { + xq_id, + xk_id, + out_list[0], + out_list[1], + start_pos_id, + dynamic_pos, static_cast(start_pos), - xk_wgc); - WGPUBuffer q_ubuf = q_disp.uniform; - WGPUBuffer k_ubuf = k_disp.uniform; - const size_t q_idx = q_disp.dispatch_index; - const size_t k_idx = k_disp.dispatch_index; - - const int xq_out_id = out_list[0]; - const int xk_out_id = out_list[1]; - const uint32_t baked_start_pos = static_cast(start_pos); - auto rope_hook = [xq_id, - xk_id, - xq_out_id, - xk_out_id, - start_pos_id, - dynamic_pos, - baked_start_pos, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - rotary_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf](WebGPUGraph& g) { - resize_rope_hf( - g, - xq_id, - xk_id, - xq_out_id, - xk_out_id, - start_pos_id, - dynamic_pos, - baked_start_pos, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - rotary_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf); - }; - graph.add_tensor_resize_hook(xq_id, rope_hook); - graph.add_tensor_resize_hook(xk_id, rope_hook); - // Dynamic decode: re-fire when the runtime start_pos SymInt changes. + geometry.n_heads_q, + geometry.n_heads_k, + geometry.head_dim, + geometry.half_dim, + geometry.rotary_dim, + geometry.max_seq, + q_uniform, + k_uniform}; + graph.add_tensor_resize_hook(xq_id, resize_rope_hf, resize_context); + graph.add_tensor_resize_hook(xk_id, resize_rope_hf, resize_context); if (dynamic_pos) { - graph.add_resize_hook(start_pos_id, rope_hook); + graph.add_resize_hook(start_pos_id, resize_rope_hf, resize_context); } } diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl index 14a6853afa3..859b24ca6fc 100644 --- a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl @@ -18,11 +18,13 @@ struct Params { override wg_size: u32 = 64u; // One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared -// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated -// halves) indexed at row (start_pos + s); only the first-half column is read. +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table indexed at row +// (start_pos + s); each output half uses its corresponding frequency column. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let pair = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let pair = gid.x + gid.y * (num_workgroups.x * wg_size); if (pair >= params.num_pairs) { return; } @@ -38,12 +40,16 @@ fn main(@builtin(global_invocation_id) gid: vec3) { ((b * params.seq + s) * params.n_heads + head) * params.head_dim; let a_idx = head_base + pair_i; let b_idx = head_base + pair_i + half_dim; - let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + let freqs_base = (s + params.start_pos) * params.rotary_dim; + let freqs_a_idx = freqs_base + pair_i; + let freqs_b_idx = freqs_a_idx + half_dim; - let c = t_freqs_cos[freqs_idx]; - let si = t_freqs_sin[freqs_idx]; + let c_a = t_freqs_cos[freqs_a_idx]; + let si_a = t_freqs_sin[freqs_a_idx]; + let c_b = t_freqs_cos[freqs_b_idx]; + let si_b = t_freqs_sin[freqs_b_idx]; let x_a = t_in[a_idx]; let x_b = t_in[b_idx]; - t_out[a_idx] = x_a * c - x_b * si; - t_out[b_idx] = x_b * c + x_a * si; + t_out[a_idx] = x_a * c_a - x_b * si_a; + t_out[b_idx] = x_b * c_b + x_a * si_b; } diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h index 191ec710e66..21242fb480a 100644 --- a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from rotary_embedding_hf.wgsl - DO NOT EDIT. -// wgsl-sha256: 5ba8d45925f00f12af17bf3092a1af9513a9e501c5c35e6b0d48cfb3dac7b5d6 +// wgsl-sha256: 4f081ed4c8165f021cbb722d379e437f30b8dfb08bf03bfcbaa406ed7799c7b6 inline constexpr const char* kRotaryEmbeddingHfWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_in: array; @@ -35,11 +35,13 @@ struct Params { override wg_size: u32 = 64u; // One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared -// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated -// halves) indexed at row (start_pos + s); only the first-half column is read. +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table indexed at row +// (start_pos + s); each output half uses its corresponding frequency column. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let pair = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let pair = gid.x + gid.y * (num_workgroups.x * wg_size); if (pair >= params.num_pairs) { return; } @@ -55,14 +57,18 @@ fn main(@builtin(global_invocation_id) gid: vec3) { ((b * params.seq + s) * params.n_heads + head) * params.head_dim; let a_idx = head_base + pair_i; let b_idx = head_base + pair_i + half_dim; - let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + let freqs_base = (s + params.start_pos) * params.rotary_dim; + let freqs_a_idx = freqs_base + pair_i; + let freqs_b_idx = freqs_a_idx + half_dim; - let c = t_freqs_cos[freqs_idx]; - let si = t_freqs_sin[freqs_idx]; + let c_a = t_freqs_cos[freqs_a_idx]; + let si_a = t_freqs_sin[freqs_a_idx]; + let c_b = t_freqs_cos[freqs_b_idx]; + let si_b = t_freqs_sin[freqs_b_idx]; let x_a = t_in[a_idx]; let x_b = t_in[b_idx]; - t_out[a_idx] = x_a * c - x_b * si; - t_out[b_idx] = x_b * c + x_a * si; + t_out[a_idx] = x_a * c_a - x_b * si_a; + t_out[b_idx] = x_b * c_b + x_a * si_b; } )"; diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 75af3ff84e6..0d5d7a9e799 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -65,6 +65,7 @@ DISPATCH_ORDER_DIR="/tmp/dispatch_order" UPDATE_CACHE_DIR="/tmp/update_cache" INDEX_DIR="/tmp/index" DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape" +ROPE_HF_DIR="/tmp/webgpu_rope_hf" SYMINT_BLOB="/tmp/sdpa_dyn_small.pte" OUTPUT_SUPPRESSION_DIR="/tmp/output_suppression" EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte" @@ -104,6 +105,11 @@ export_rope_model('${ROPE_MODEL}', '${ROPE_XQ_GOLDEN}', '${ROPE_XK_GOLDEN}') export_rope_model('${ROPE_DECODE_MODEL}', '${ROPE_DECODE_XQ_GOLDEN}', '${ROPE_DECODE_XK_GOLDEN}', 'decode') " +$PYTHON_EXECUTABLE -c " +from executorch.backends.webgpu.test.ops.test_rope_hf import export_rope_hf_dynamic +export_rope_hf_dynamic('${ROPE_HF_DIR}') +" + $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_prepack import export_prepack_model, export_prepack_two_const_model, export_prepack_tied_const_model export_prepack_model('${PREPACK_MODEL}', '${PREPACK_GOLDEN}') @@ -150,6 +156,7 @@ export_dynamic_decode('/tmp') export_incache_decode('/tmp') " +require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte" require_file "${SYMINT_BLOB}" require_file "${OUTPUT_SUPPRESSION_DIR}/input.bin" @@ -201,6 +208,7 @@ run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \ WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \ WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \ WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \ + WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}" \ WEBGPU_TEST_SYMINT_BLOB="${SYMINT_BLOB}" \ WEBGPU_TEST_PREPACK_MODEL="${PREPACK_MODEL}" \ WEBGPU_TEST_PREPACK_GOLDEN="${PREPACK_GOLDEN}" \ diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 0963168ccd5..feb6854c59d 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -20,6 +21,8 @@ #include #include #include +#include +#include namespace executorch::backends::webgpu { namespace { @@ -653,6 +656,140 @@ TEST(WebGPUDynamicDispatch, RejectsRouteOverlapWithoutPoisoningRegistry) { expect_dispatch_grid(graph, 2u, 13u, 17u); } +struct InvalidRopeGraphCase { + const char* name; + std::vector xq_dims; + std::vector xk_dims; + std::vector cos_dims; + std::vector sin_dims; + std::vector xq_out_dims; + std::vector xk_out_dims; + vkgraph::VkDataType xq_dtype; + const char* expected_error; +}; + +void expect_invalid_rope_graph(const InvalidRopeGraphCase& test_case) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, + const std::vector& dims, + int mem_obj_id) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + }; + add_tensor(test_case.xq_dtype, test_case.xq_dims, 0); + add_tensor(vk::VkDataType::FLOAT32, test_case.xk_dims, 1); + add_tensor(vk::VkDataType::FLOAT32, test_case.cos_dims, 2); + add_tensor(vk::VkDataType::FLOAT32, test_case.sin_dims, 3); + add_tensor(vk::VkDataType::FLOAT32, test_case.xq_out_dims, 4); + add_tensor(vk::VkDataType::FLOAT32, test_case.xk_out_dims, 5); + std::vector output_value_ids = {4, 5}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &output_value_ids).Union())); + + std::vector args = {0, 1, 2, 3, 6}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.apply_rotary_emb.default", &args)); + std::vector input_ids = {0, 1, 2, 3}; + std::vector output_ids = {4, 5}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + std::string error; + try { + graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + } catch (const std::exception& exception) { + error = exception.what(); + } + EXPECT_FALSE(error.empty()) << test_case.name << " unexpectedly built"; + EXPECT_EQ(error, test_case.expected_error) + << test_case.name << " rejected for the wrong reason"; + const WebGPUMemoryStats stats = graph.memory_stats(); + EXPECT_EQ(stats.num_dispatches, 0) << test_case.name; + EXPECT_EQ(stats.uniform_buffer_bytes, 0u) << test_case.name; + EXPECT_EQ(stats.num_cached_shaders, 0) << test_case.name; + EXPECT_EQ(stats.num_cached_pipelines, 0) << test_case.name; +} + +TEST(WebGPURopeValidation, RejectsMalformedGraphsBeforeDispatchAllocation) { + ASSERT_TRUE( + webgpu_operator_registry().has_op("et_vk.apply_rotary_emb.default")); + const std::vector xq = {1, 2, 2, 4}; + const std::vector xk = {1, 2, 1, 4}; + const std::vector freqs = {2, 2}; + const InvalidRopeGraphCase cases[] = { + {"query rank", + {2, 4}, + xk, + freqs, + freqs, + {2, 4}, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: malformed dims"}, + {"sequence mismatch", + xq, + {1, 3, 1, 4}, + freqs, + freqs, + xq, + {1, 3, 1, 4}, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"}, + {"head dimension mismatch", + xq, + {1, 2, 1, 6}, + freqs, + freqs, + xq, + {1, 2, 1, 6}, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"}, + {"frequency width mismatch", + xq, + xk, + {2, 3}, + {2, 3}, + xq, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: head_dim != 2 * freqs_cos last dim"}, + {"cosine/sine shape mismatch", + xq, + xk, + freqs, + {2, 1}, + xq, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: freqs_cos and freqs_sin shapes differ"}, + {"query byte size mismatch", + xq, + xk, + freqs, + freqs, + xq, + xk, + vkgraph::VkDataType::INT64, + "WebGPU apply_rotary_emb: dtype/byte-size mismatch (all fp32) or " + "freqs shape != [seq, head_dim/2]"}, + }; + for (const InvalidRopeGraphCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + expect_invalid_rope_graph(test_case); + } +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/ops/test_rope_hf.py b/backends/webgpu/test/ops/test_rope_hf.py index 6e26fe53b17..f1929ac8d3b 100644 --- a/backends/webgpu/test/ops/test_rope_hf.py +++ b/backends/webgpu/test/ops/test_rope_hf.py @@ -19,6 +19,7 @@ compare (it has no ATen). """ +import os import unittest from collections import namedtuple @@ -26,8 +27,12 @@ import torch from executorch.backends.vulkan import VulkanPartitioner -from executorch.examples.models.llama.rope import hf_apply_rotary_emb +from executorch.examples.models.llama.rope import ( + hf_apply_rotary_emb, + hf_precompute_freqs_cis, +) from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes # B batch, S tokens, NH query heads, NKV kv heads (NH != NKV so the two outputs # are distinguishable by numel), HD head dim (even; full rotary, rotary_dim==HD). @@ -39,6 +44,20 @@ Shape("decode", 1, 1, 16, 8, 128), ] +DYNAMIC_BATCH = 1 +DYNAMIC_SEQ = 1 +DYNAMIC_N_HEADS_Q = 16 +DYNAMIC_N_HEADS_K = 8 +DYNAMIC_HEAD_DIM = 128 +DYNAMIC_MAX_SEQ = 16 +DYNAMIC_POSITIONS = (0, 7, 15) +DYNAMIC_SEQUENCE_CASES = ( + (DYNAMIC_MAX_SEQ, 0), + (5, 7), + (1, DYNAMIC_MAX_SEQ - 1), + (DYNAMIC_MAX_SEQ, 0), +) + class HfRope(torch.nn.Module): # unsqueeze_dim=1: freqs [S, HD] -> [S, 1, HD] broadcasts over (B, NH) of the @@ -47,6 +66,26 @@ def forward(self, xq, xk, freqs_cos, freqs_sin): return hf_apply_rotary_emb(xq, xk, freqs_cos, freqs_sin, unsqueeze_dim=1) +class DynamicHfRope(torch.nn.Module): + def forward( + self, + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + input_pos: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + start_pos = input_pos[0].item() + torch._check_is_size(start_pos) + torch._check(start_pos + xq.shape[1] <= freqs_cos.shape[0]) + return hf_apply_rotary_emb( + xq, + xk, + freqs_cos.narrow(0, start_pos, xq.shape[1]), + freqs_sin.narrow(0, start_pos, xq.shape[1]), + ) + + def _ramp(numel: int, mod: int, off: int) -> torch.Tensor: # ((i % mod) - off) / 16: exact in fp32, matches test_webgpu_native.cpp. idx = torch.arange(numel, dtype=torch.int64) @@ -68,6 +107,32 @@ def _inputs( return xq, xk, freqs_cos, freqs_sin +def _dynamic_inputs(seq: int = DYNAMIC_SEQ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + xq = _ramp( + DYNAMIC_BATCH * seq * DYNAMIC_N_HEADS_Q * DYNAMIC_HEAD_DIM, + 17, + 8, + ).reshape(DYNAMIC_BATCH, seq, DYNAMIC_N_HEADS_Q, DYNAMIC_HEAD_DIM) + xk = _ramp( + DYNAMIC_BATCH * seq * DYNAMIC_N_HEADS_K * DYNAMIC_HEAD_DIM, + 13, + 6, + ).reshape(DYNAMIC_BATCH, seq, DYNAMIC_N_HEADS_K, DYNAMIC_HEAD_DIM) + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + DYNAMIC_HEAD_DIM, + DYNAMIC_MAX_SEQ, + theta=10000.0, + ) + input_pos = torch.tensor([0], dtype=torch.long) + return xq, xk, freqs_cos, freqs_sin, input_pos + + def _golden( xq: torch.Tensor, xk: torch.Tensor, @@ -78,27 +143,123 @@ def _golden( return torch.ops.et_vk.apply_rotary_emb_hf.default(xq, xk, freqs_cos, freqs_sin, 0) -def _export(inputs): +def _dynamic_golden( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + position: int, +) -> tuple[torch.Tensor, torch.Tensor]: + return hf_apply_rotary_emb( + xq, + xk, + freqs_cos[position : position + xq.shape[1]], + freqs_sin[position : position + xq.shape[1]], + ) + + +def _assert_fully_delegated(edge) -> None: + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + portable = get_non_lowered_nodes(graph) + if len(delegates) != 1: + raise AssertionError(f"expected one delegate, got {len(delegates)}") + if portable: + raise AssertionError(f"unexpected non-lowered nodes: {portable}") + + +def _lower(inputs): ep = torch.export.export(HfRope().eval(), inputs) - return to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _export(inputs): + return _lower(inputs).to_executorch() + + +def _lower_dynamic_program(): + inputs = _dynamic_inputs() + with torch._dynamo.config.patch(capture_scalar_outputs=True): + ep = torch.export.export(DynamicHfRope().eval(), inputs) + + symints = [ + node + for node in ep.graph_module.graph.nodes + if isinstance(node.meta.get("val"), torch.SymInt) + ] + if not symints: + raise AssertionError("input_pos did not lower to a SymInt") + + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _lower_dynamic_sequence_program(): + inputs = _dynamic_inputs(DYNAMIC_MAX_SEQ) + s_dim = torch.export.Dim("rope_hf_s", min=1, max=DYNAMIC_MAX_SEQ) + dynamic_shapes = ({1: s_dim}, {1: s_dim}, None, None, None) + with torch._dynamo.config.patch(capture_scalar_outputs=True): + ep = torch.export.export( + DynamicHfRope().eval(), + inputs, + dynamic_shapes=dynamic_shapes, + ) + + scalar_symints = [ + node + for node in ep.graph_module.graph.nodes + if isinstance(node.meta.get("val"), torch.SymInt) + ] + if not scalar_symints: + raise AssertionError("input_pos did not lower to a SymInt") + xq_placeholder = next( + node + for node in ep.graph_module.graph.nodes + if node.op == "placeholder" and node.target == "xq" + ) + if not isinstance(xq_placeholder.meta["val"].shape[1], torch.SymInt): + raise AssertionError("query sequence dimension did not remain symbolic") + + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _export_dynamic_program(): + edge = _lower_dynamic_program() + + et = edge.to_executorch() + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + if delegate_ids != ["VulkanBackend"]: + raise AssertionError(f"unexpected delegates: {delegate_ids}") + return et + + +def _export_dynamic_sequence_program(): + edge = _lower_dynamic_sequence_program() + et = edge.to_executorch() + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + if delegate_ids != ["VulkanBackend"]: + raise AssertionError(f"unexpected delegates: {delegate_ids}") + return et class TestRopeHf(unittest.TestCase): def test_export_delegates(self) -> None: for shape in SHAPES: with self.subTest(shape=shape.name): - et = _export(_inputs(shape)) - found = any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ) - self.assertTrue( - found, - "Expected a VulkanBackend delegate (apply_rotary_emb_hf " "fusion)", - ) + self.assertIsNotNone(_lower(_inputs(shape))) def test_golden_matches_eager(self) -> None: # The et_vk golden must equal the real HF rotate-half apply_rotary_emb, @@ -112,6 +273,44 @@ def test_golden_matches_eager(self) -> None: torch.testing.assert_close(gq, eq, atol=1e-5, rtol=1e-5) torch.testing.assert_close(gk, ek, atol=1e-5, rtol=1e-5) + def test_dynamic_export_is_fully_delegated(self) -> None: + self.assertIsNotNone(_lower_dynamic_program()) + + def test_dynamic_position_goldens_match_custom_op(self) -> None: + xq, xk, freqs_cos, freqs_sin, _ = _dynamic_inputs() + self.assertNotEqual(xq.shape[2], xk.shape[2]) + position_outputs = [] + for position in DYNAMIC_POSITIONS: + with self.subTest(position=position): + expected_q, expected_k = _dynamic_golden( + xq, xk, freqs_cos, freqs_sin, position + ) + position_outputs.append(expected_q) + actual_q, actual_k = torch.ops.et_vk.apply_rotary_emb_hf.default( + xq, xk, freqs_cos, freqs_sin, position + ) + torch.testing.assert_close(actual_q, expected_q) + torch.testing.assert_close(actual_k, expected_k) + self.assertFalse(torch.allclose(position_outputs[0], position_outputs[1])) + self.assertFalse(torch.allclose(position_outputs[1], position_outputs[2])) + + def test_dynamic_sequence_export_is_fully_delegated(self) -> None: + self.assertIsNotNone(_lower_dynamic_sequence_program()) + + def test_dynamic_sequence_goldens_match_custom_op(self) -> None: + _, _, freqs_cos, freqs_sin, _ = _dynamic_inputs(DYNAMIC_MAX_SEQ) + for seq, position in dict.fromkeys(DYNAMIC_SEQUENCE_CASES): + with self.subTest(seq=seq, position=position): + xq, xk, _, _, _ = _dynamic_inputs(seq) + expected_q, expected_k = _dynamic_golden( + xq, xk, freqs_cos, freqs_sin, position + ) + actual_q, actual_k = torch.ops.et_vk.apply_rotary_emb_hf.default( + xq, xk, freqs_cos, freqs_sin, position + ) + torch.testing.assert_close(actual_q, expected_q) + torch.testing.assert_close(actual_k, expected_k) + def export_rope_hf_model( pte_path: str, xq_golden_path: str, xk_golden_path: str, shape_name: str = "multi" @@ -132,5 +331,58 @@ def export_rope_hf_model( ) +def export_rope_hf_dynamic(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) + xq, xk, freqs_cos, freqs_sin, _ = _dynamic_inputs() + et = _export_dynamic_program() + with open(os.path.join(out_dir, "rope_hf_dynamic.pte"), "wb") as output: + output.write(et.buffer) + + for name, tensor in ( + ("xq", xq), + ("xk", xk), + ("freqs_cos", freqs_cos), + ("freqs_sin", freqs_sin), + ): + tensor.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) + _, _, freqs_cos, freqs_sin, _ = _dynamic_inputs(DYNAMIC_MAX_SEQ) + et = _export_dynamic_sequence_program() + prefix = "rope_hf_dynamic_sequence" + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as output: + output.write(et.buffer) + + for name, tensor in (("freqs_cos", freqs_cos), ("freqs_sin", freqs_sin)): + tensor.detach().numpy().astype(".pte + .golden.bin to /tmp) ===" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import export_all_sdpa_models @@ -96,6 +105,7 @@ cmake \ "${EXECUTORCH_ROOT}" cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_native_test -j${NPROC} +cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_compute_dispatch_test -j${NPROC} cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_dispatch_order_test -j${NPROC} cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_scratch_buffer_test -j${NPROC} @@ -108,9 +118,11 @@ else fi env \ ${UPDATE_CACHE_ENV_VAR} \ + WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}" \ WEBGPU_TEST_SDPA_DIR=/tmp/ \ "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_native_test" +"${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_compute_dispatch_test" "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_scratch_buffer_test" diff --git a/backends/webgpu/test/test_native_ci_contract.py b/backends/webgpu/test/test_native_ci_contract.py index 011354a0e0f..43b8ec6f56e 100644 --- a/backends/webgpu/test/test_native_ci_contract.py +++ b/backends/webgpu/test/test_native_ci_contract.py @@ -64,3 +64,12 @@ def test_requires_symint_and_suppression_fixtures(self) -> None: "${OUTPUT_SUPPRESSION_DIR}/input.bin", ): self.assertIn(f'require_file "{fixture}"', script) + + def test_requires_dynamic_rope_fixture(self) -> None: + script = ( + pathlib.Path(__file__).parents[1] / "scripts/test_webgpu_native_ci.sh" + ).read_text() + + self.assertIn("export_rope_hf_dynamic('${ROPE_HF_DIR}')", script) + self.assertIn('WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}"', script) + self.assertIn('require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte"', script) diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index f64cccec550..a57c41e53a8 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -513,6 +514,152 @@ std::vector run_rms_norm_at_wg( return result; } +struct RotaryHfProbeParams { + uint32_t n_heads; + uint32_t seq; + uint32_t head_dim; + uint32_t half_dim; + uint32_t num_pairs; + uint32_t rotary_dim; + uint32_t start_pos; + uint32_t _pad; +}; + +std::vector run_rope_hf_2d_probe(const WebGPUContext& ctx) { + constexpr uint32_t kWorkgroupSize = 2; + constexpr uint32_t kWorkgroupsX = 2; + constexpr uint32_t kWorkgroupsY = 2; + constexpr uint32_t kNumPairs = kWorkgroupSize * kWorkgroupsX * kWorkgroupsY; + constexpr uint32_t kHeadDim = kNumPairs * 2; + + std::vector input(kHeadDim); + std::vector output(kHeadDim, 0.0f); + std::vector freqs_cos(kHeadDim, 1.0f); + std::vector freqs_sin(kHeadDim, 0.0f); + for (uint32_t i = 0; i < kHeadDim; i++) { + input[i] = static_cast(i + 1u); + if (i >= kNumPairs) { + freqs_cos[i] = 2.0f; + } + } + + WGPUDevice device = ctx.device; + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kRotaryEmbeddingHfWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUConstantEntry wg_const = {}; + wg_const.key = {"wg_size", WGPU_STRLEN}; + wg_const.value = static_cast(kWorkgroupSize); + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_const; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + WGPUBindGroupLayout layout = + wgpuComputePipelineGetBindGroupLayout(pipeline, 0); + + auto make_buffer = + [device](const void* data, uint64_t size, WGPUBufferUsage usage) { + WGPUBufferDescriptor desc = {}; + desc.size = size; + desc.usage = usage; + desc.mappedAtCreation = true; + WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &desc); + std::memcpy(wgpuBufferGetMappedRange(buffer, 0, size), data, size); + wgpuBufferUnmap(buffer); + return buffer; + }; + + const uint64_t data_bytes = kHeadDim * sizeof(float); + WGPUBuffer out_buffer = make_buffer( + output.data(), + data_bytes, + WGPUBufferUsage_Storage | WGPUBufferUsage_CopySrc); + WGPUBuffer in_buffer = + make_buffer(input.data(), data_bytes, WGPUBufferUsage_Storage); + WGPUBuffer cos_buffer = + make_buffer(freqs_cos.data(), data_bytes, WGPUBufferUsage_Storage); + WGPUBuffer sin_buffer = + make_buffer(freqs_sin.data(), data_bytes, WGPUBufferUsage_Storage); + const RotaryHfProbeParams params = { + 1u, 1u, kHeadDim, kNumPairs, kNumPairs, kHeadDim, 0u, 0u}; + WGPUBuffer params_buffer = + make_buffer(¶ms, sizeof(params), WGPUBufferUsage_Uniform); + + WGPUBindGroupEntry entries[5] = {}; + const WGPUBuffer buffers[] = { + out_buffer, in_buffer, cos_buffer, sin_buffer, params_buffer}; + const uint64_t sizes[] = { + data_bytes, data_bytes, data_bytes, data_bytes, sizeof(params)}; + for (uint32_t i = 0; i < 5; i++) { + entries[i].binding = i; + entries[i].buffer = buffers[i]; + entries[i].size = sizes[i]; + } + WGPUBindGroupDescriptor bind_group_desc = {}; + bind_group_desc.layout = layout; + bind_group_desc.entryCount = 5; + bind_group_desc.entries = entries; + WGPUBindGroup bind_group = + wgpuDeviceCreateBindGroup(device, &bind_group_desc); + + WGPUBufferDescriptor staging_desc = {}; + staging_desc.size = data_bytes; + staging_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; + WGPUBuffer staging = wgpuDeviceCreateBuffer(device, &staging_desc); + + WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device, nullptr); + WGPUComputePassDescriptor pass_desc = {}; + WGPUComputePassEncoder pass = + wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + wgpuComputePassEncoderSetPipeline(pass, pipeline); + wgpuComputePassEncoderSetBindGroup(pass, 0, bind_group, 0, nullptr); + wgpuComputePassEncoderDispatchWorkgroups(pass, kWorkgroupsX, kWorkgroupsY, 1); + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, out_buffer, 0, staging, 0, data_bytes); + WGPUCommandBuffer command = wgpuCommandEncoderFinish(encoder, nullptr); + wgpuQueueSubmit(ctx.queue, 1, &command); + wgpuCommandBufferRelease(command); + wgpuCommandEncoderRelease(encoder); + + WgMapData callback = {}; + WGPUBufferMapCallbackInfo callback_info = {}; + callback_info.mode = WGPUCallbackMode_WaitAnyOnly; + callback_info.callback = wg_map_cb; + callback_info.userdata1 = &callback; + WGPUFuture future = wgpuBufferMapAsync( + staging, WGPUMapMode_Read, 0, data_bytes, callback_info); + const WGPUWaitStatus wait = webgpu_wait(ctx.instance, future); + if (wait == WGPUWaitStatus_Success && + callback.status == WGPUMapAsyncStatus_Success) { + const void* mapped = wgpuBufferGetConstMappedRange(staging, 0, data_bytes); + std::memcpy(output.data(), mapped, data_bytes); + wgpuBufferUnmap(staging); + } else { + output.clear(); + } + + wgpuBufferRelease(staging); + wgpuBindGroupRelease(bind_group); + wgpuBufferRelease(params_buffer); + wgpuBufferRelease(sin_buffer); + wgpuBufferRelease(cos_buffer); + wgpuBufferRelease(in_buffer); + wgpuBufferRelease(out_buffer); + wgpuBindGroupLayoutRelease(layout); + wgpuComputePipelineRelease(pipeline); + wgpuShaderModuleRelease(shader); + return output; +} + // linear_q4gsw sweep config; mirrors CONFIGS in test_quantized_linear.py. struct Q4gswConfig { const char* name; @@ -767,6 +914,220 @@ void test_rope( << "apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)"; } +bool has_shape( + const executorch::aten::Tensor& tensor, + const std::vector& expected) { + if (tensor.dim() != static_cast(expected.size())) { + return false; + } + for (size_t i = 0; i < expected.size(); i++) { + if (tensor.size(static_cast(i)) != expected[i]) { + return false; + } + } + return true; +} + +void test_rope_hf_dynamic(const std::string& dir) { + constexpr int S = 1; + constexpr int NH = 16; + constexpr int NKV = 8; + constexpr int HD = 128; + constexpr int MAXS = 16; + constexpr int positions[] = {0, 7, 15}; + constexpr int xq_numel = S * NH * HD; + constexpr int xk_numel = S * NKV * HD; + constexpr int freqs_numel = MAXS * HD; + + Module module(dir + "rope_hf_dynamic.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load HF RoPE dynamic model"; + + std::vector xq = load_golden(dir + "rope_hf_dynamic.xq.bin", xq_numel); + std::vector xk = load_golden(dir + "rope_hf_dynamic.xk.bin", xk_numel); + std::vector freqs_cos = + load_golden(dir + "rope_hf_dynamic.freqs_cos.bin", freqs_numel); + std::vector freqs_sin = + load_golden(dir + "rope_hf_dynamic.freqs_sin.bin", freqs_numel); + ASSERT_FALSE( + xq.empty() || xk.empty() || freqs_cos.empty() || freqs_sin.empty()) + << "could not load HF RoPE input binaries from " << dir; + + for (const int position : positions) { + auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::vector(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::vector(freqs_sin)); + auto post = make_tensor_ptr( + {1}, std::vector{static_cast(position)}); + auto result = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + ASSERT_TRUE(result.ok()) + << "HF RoPE forward failed at position " << position << " (error " + << static_cast(result.error()) << ")"; + const auto& outputs = result.get(); + ASSERT_TRUE( + outputs.size() == 2 && outputs[0].isTensor() && outputs[1].isTensor()) + << "expected exactly two HF RoPE tensor outputs"; + const auto& xq_out = outputs[0].toTensor(); + const auto& xk_out = outputs[1].toTensor(); + ASSERT_TRUE(has_shape(xq_out, {1, S, NH, HD})) + << "HF RoPE query output has the wrong shape at position " << position; + ASSERT_TRUE(has_shape(xk_out, {1, S, NKV, HD})) + << "HF RoPE key output has the wrong shape at position " << position; + + const std::string prefix = + dir + "rope_hf_dynamic.pos" + std::to_string(position); + const std::vector golden_q = + load_golden(prefix + ".xq.golden.bin", xq_numel); + const std::vector golden_k = + load_golden(prefix + ".xk.golden.bin", xk_numel); + ASSERT_FALSE(golden_q.empty() || golden_k.empty()) + << "could not load HF RoPE goldens for position " << position; + + float q_abs = 0.0f, q_rel = 0.0f, k_abs = 0.0f, k_rel = 0.0f; + const bool q_ok = quant_within_tol( + xq_out.const_data_ptr(), + golden_q.data(), + xq_numel, + 1e-4f, + 1e-3f, + &q_abs, + &q_rel); + const bool k_ok = quant_within_tol( + xk_out.const_data_ptr(), + golden_k.data(), + xk_numel, + 1e-4f, + 1e-3f, + &k_abs, + &k_rel); + EXPECT_TRUE(q_ok && k_ok) + << "HF RoPE mismatch at position " << position << ": q abs=" << q_abs + << " rel=" << q_rel << ", k abs=" << k_abs << " rel=" << k_rel; + } + + auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::move(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::move(freqs_sin)); + + auto overflow_post = + make_tensor_ptr({1}, std::vector{INT64_C(1) << 32}); + auto overflow = module.forward({ + EValue(xqt), + EValue(xkt), + EValue(fct), + EValue(fst), + EValue(overflow_post), + }); + EXPECT_FALSE(overflow.ok()) + << "HF RoPE accepted a start_pos that aliases to zero when narrowed"; + + auto post = make_tensor_ptr({1}, std::vector{MAXS}); + auto out_of_range = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + EXPECT_FALSE(out_of_range.ok()) + << "HF RoPE accepted start_pos + seq beyond the frequency table"; + + auto negative_post = make_tensor_ptr({1}, std::vector{-1}); + auto negative = module.forward({ + EValue(xqt), + EValue(xkt), + EValue(fct), + EValue(fst), + EValue(negative_post), + }); + EXPECT_FALSE(negative.ok()) << "HF RoPE accepted a negative start_pos"; +} + +void test_rope_hf_dynamic_sequence_reused_graph(const std::string& dir) { + constexpr int NH = 16; + constexpr int NKV = 8; + constexpr int HD = 128; + constexpr int MAXS = 16; + struct Case { + int seq; + int position; + }; + constexpr Case cases[] = {{16, 0}, {5, 7}, {1, 15}, {16, 0}}; + + const std::string prefix = dir + "rope_hf_dynamic_sequence"; + Module module(prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load HF RoPE dynamic-sequence model"; + + const int freqs_numel = MAXS * HD; + const std::vector freqs_cos = + load_golden(prefix + ".freqs_cos.bin", freqs_numel); + const std::vector freqs_sin = + load_golden(prefix + ".freqs_sin.bin", freqs_numel); + ASSERT_FALSE(freqs_cos.empty() || freqs_sin.empty()) + << "could not load HF RoPE dynamic-sequence frequencies"; + + for (const Case& c : cases) { + const int xq_numel = c.seq * NH * HD; + const int xk_numel = c.seq * NKV * HD; + const std::string case_prefix = prefix + ".S" + std::to_string(c.seq) + + ".pos" + std::to_string(c.position); + const std::vector xq = + load_golden(case_prefix + ".xq.bin", xq_numel); + const std::vector xk = + load_golden(case_prefix + ".xk.bin", xk_numel); + const std::vector golden_q = + load_golden(case_prefix + ".xq.golden.bin", xq_numel); + const std::vector golden_k = + load_golden(case_prefix + ".xk.golden.bin", xk_numel); + ASSERT_FALSE( + xq.empty() || xk.empty() || golden_q.empty() || golden_k.empty()) + << "could not load HF RoPE dynamic-sequence case " << case_prefix; + + auto xqt = make_tensor_ptr({1, c.seq, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, c.seq, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::vector(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::vector(freqs_sin)); + auto post = make_tensor_ptr( + {1}, std::vector{static_cast(c.position)}); + auto result = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + ASSERT_TRUE(result.ok()) + << "HF RoPE dynamic-sequence forward failed for " << case_prefix + << " (error " << static_cast(result.error()) << ")"; + const auto& outputs = result.get(); + ASSERT_TRUE( + outputs.size() == 2 && outputs[0].isTensor() && outputs[1].isTensor()) + << "expected exactly two HF RoPE dynamic-sequence tensor outputs"; + const auto& xq_out = outputs[0].toTensor(); + const auto& xk_out = outputs[1].toTensor(); + ASSERT_TRUE(has_shape(xq_out, {1, c.seq, NH, HD})) + << "HF RoPE query output has the wrong shape for " << case_prefix; + ASSERT_TRUE(has_shape(xk_out, {1, c.seq, NKV, HD})) + << "HF RoPE key output has the wrong shape for " << case_prefix; + + float q_abs = 0.0f, q_rel = 0.0f, k_abs = 0.0f, k_rel = 0.0f; + const bool q_ok = quant_within_tol( + xq_out.const_data_ptr(), + golden_q.data(), + xq_numel, + 1e-4f, + 1e-3f, + &q_abs, + &q_rel); + const bool k_ok = quant_within_tol( + xk_out.const_data_ptr(), + golden_k.data(), + xk_numel, + 1e-4f, + 1e-3f, + &k_abs, + &k_rel); + EXPECT_TRUE(q_ok && k_ok) + << "HF RoPE dynamic-sequence mismatch for " << case_prefix + << ": q abs=" << q_abs << " rel=" << q_rel << ", k abs=" << k_abs + << " rel=" << k_rel; + } +} + void test_prepack( const std::string& model_path, const std::string& golden_path, @@ -2125,6 +2486,82 @@ static bool test_select_double_scalars() { return ok; } +void expect_rope_hf_resize_numel_overflow(uint32_t q_heads, uint32_t k_heads) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + + const std::vector q_dims = {1u, 1u, q_heads, 2u}; + const std::vector k_dims = {1u, 1u, k_heads, 2u}; + const std::vector freqs_dims = {2u, 2u}; + std::vector<::flatbuffers::Offset> values; + const auto add_tensor = [&](const std::vector& dims, int mem_id) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + /*mem_obj_id=*/mem_id) + .Union())); + }; + add_tensor(q_dims, 0); + add_tensor(k_dims, 1); + add_tensor(freqs_dims, 2); + add_tensor(freqs_dims, 3); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + add_tensor(q_dims, 4); + add_tensor(k_dims, 5); + const std::vector output_items = {5, 6}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &output_items).Union())); + + const std::vector args = {0, 1, 2, 3, 4, 7}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.apply_rotary_emb_hf.default", &args)); + const std::vector input_ids = {0, 1, 2, 3}; + const std::vector output_ids = {5, 6}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, nullptr)); + ASSERT_EQ(graph.num_dispatches(), 2u); + const uint32_t q_x = graph.dispatch_at(0).workgroup_count_x; + const uint32_t q_y = graph.dispatch_at(0).workgroup_count_y; + const uint32_t k_x = graph.dispatch_at(1).workgroup_count_x; + const uint32_t k_y = graph.dispatch_at(1).workgroup_count_y; + + constexpr int64_t kLargeBatch = INT64_C(1) << 30; + const std::vector q_live = { + kLargeBatch, 1, static_cast(q_heads), 2}; + const std::vector k_live = { + kLargeBatch, 1, static_cast(k_heads), 2}; + graph.get_tensor(0).dims = q_live; + graph.get_tensor(1).dims = k_live; + ASSERT_NO_THROW(graph.resize_input(0, q_live)); + ASSERT_NO_THROW(graph.resize_input(1, k_live)); + + try { + graph.propagate_resize(); + FAIL() << "accepted q/k element count outside uint32 range"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "apply_rotary_emb_hf(resize): element index exceeds uint32 range"); + } + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_x, q_x); + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_y, q_y); + EXPECT_EQ(graph.dispatch_at(1).workgroup_count_x, k_x); + EXPECT_EQ(graph.dispatch_at(1).workgroup_count_y, k_y); +} + // apply_rotary_emb on-GPU configs: multi + decode (env-gated, run-if-present). struct RopeConfig { const char* name; @@ -2298,6 +2735,48 @@ TEST(WebGPUNative, Rope) { } } +TEST(WebGPUNative, RopeHfDynamic) { + const char* env = std::getenv("WEBGPU_TEST_ROPE_HF_DIR"); + if (env == nullptr || *env == '\0') { + GTEST_SKIP() << "WEBGPU_TEST_ROPE_HF_DIR not set"; + } + std::string dir = env; + if (dir.back() != '/') { + dir += '/'; + } + test_rope_hf_dynamic(dir); +} + +TEST(WebGPUNative, RopeHfDynamicSequenceReusedGraph) { + const char* env = std::getenv("WEBGPU_TEST_ROPE_HF_DIR"); + if (env == nullptr || *env == '\0') { + GTEST_SKIP() << "WEBGPU_TEST_ROPE_HF_DIR not set"; + } + std::string dir = env; + if (dir.back() != '/') { + dir += '/'; + } + test_rope_hf_dynamic_sequence_reused_graph(dir); +} + +TEST(WebGPUNative, RopeHfUsesFull2DGridStride) { + const WebGPUContext* ctx = get_default_webgpu_context(); + ASSERT_NE(ctx, nullptr); + const std::vector output = run_rope_hf_2d_probe(*ctx); + ASSERT_EQ(output.size(), 16u) << "HF RoPE probe output map failed"; + for (size_t i = 0; i < output.size(); i++) { + const float expected = + static_cast(i + 1u) * (i < output.size() / 2u ? 1.0f : 2.0f); + EXPECT_EQ(output[i], expected) + << "HF RoPE 2D grid or second-half frequency mismatch at element " << i; + } +} + +TEST(WebGPUNative, RopeHfResizeRejectsQOrKNumelOverflow) { + expect_rope_hf_resize_numel_overflow(/*q_heads=*/2, /*k_heads=*/1); + expect_rope_hf_resize_numel_overflow(/*q_heads=*/1, /*k_heads=*/2); +} + TEST(WebGPUNative, Prepack) { if (g_prepack_model_path.empty() || g_prepack_golden_path.empty()) { GTEST_SKIP() << "WEBGPU_TEST_PREPACK_MODEL/GOLDEN not set"; diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index fe8340d6528..0cd110d0ff7 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -178,6 +178,28 @@ def test_committed_headers_match_generator(self) -> None: got, want, f"{header.name} stale; run scripts/gen_wgsl_headers.py" ) + def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: + shader = ( + g.BACKEND_ROOT / "runtime" / "ops" / "rope" / "rotary_embedding_hf.wgsl" + ).read_text() + self.assertIn("@builtin(num_workgroups) num_workgroups", shader) + self.assertIn( + "gid.x + gid.y * (num_workgroups.x * wg_size)", + shader, + ) + self.assertIn("let freqs_b_idx = freqs_a_idx + half_dim;", shader) + self.assertIn("t_out[b_idx] = x_b * c_b + x_a * si_b;", shader) + + wg_size = 2 + workgroups_x = 2 + indices = [ + group_x * wg_size + lane + group_y * (workgroups_x * wg_size) + for group_y in range(2) + for group_x in range(workgroups_x) + for lane in range(wg_size) + ] + self.assertEqual(indices, list(range(8))) + def test_parse_workgroup_allows_space(self) -> None: # @workgroup_size (64) — the spec-legal spaced form must still parse. self.assertEqual( From 2c9ca643df3610257a84a4e7cdeaec424192b2f2 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:33 -0700 Subject: [PATCH 06/12] [ExecuTorch][WebGPU] Add Qwen3 K16 attention and fp16 KV cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21136 Qwen3's attention geometry differs from Llama's, and its KV cache is produced in fp32 on the host but must be consumed in fp16 on the device. This adds guarded Qwen3 K16 streaming online-softmax attention schedules — a Q16 schedule that is the automatic default whenever the exact geometry and capability guards pass, plus a Q32 candidate that is opt-in through the `sdpa_query_tile` runtime spec (BackendOption) for future autotuning — together with the exact fp32-host to fp16-device KV-cache boundary conversion. Selection requires the exact Qwen3 geometry, fp16 KV storage, adapter limits, a valid workgroup count, and an exact 2:1 byte ratio; the established Llama, materialized, and FlashDecoding routes remain fallbacks. This builds on the HuggingFace rotate-half RoPE operator. No Vulkan analogue (WebGPU-specific online-softmax attention; Vulkan has only a materialized attention). It also makes the long generated WGSL provenance and constant declarations format-stable and covers them with a generator regression test. Key changes: - runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl and streaming_attention_qwen3_q32_k16_causal_bound.wgsl (+ generated headers): the Q16 and Q32 online-softmax Qwen3 kernels. - Sdpa.cpp, WebGPUGraph.{cpp,h}: exact Qwen3 geometry and limit guards, Q16 default route selection, and the fp32-host to fp16-device KV-cache conversion. - WebGPUBackend.cpp: read the optional `sdpa_query_tile` runtime spec and thread it into graph build so the Q32 tile can be requested without a rebuild. - scripts/gen_wgsl_headers.py (+ test_wgsl_codegen.py): format-stable generated headers with a regression test. ghstack-source-id: 411961459 @exported-using-ghexport Differential Revision: [D113171744](https://our.internmc.facebook.com/intern/diff/D113171744/) --- backends/webgpu/runtime/WebGPUBackend.cpp | 39 +- .../webgpu/runtime/WebGPUDelegateHeader.cpp | 17 +- .../webgpu/runtime/WebGPUDelegateHeader.h | 5 +- backends/webgpu/runtime/WebGPUGraph.cpp | 222 +++++++- backends/webgpu/runtime/WebGPUGraph.h | 20 +- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 18 +- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 165 +++++- ...ming_attention_qwen3_k16_causal_bound.wgsl | 283 ++++++++++ ...ng_attention_qwen3_k16_causal_bound_wgsl.h | 311 +++++++++++ ..._attention_qwen3_q32_k16_causal_bound.wgsl | 283 ++++++++++ ...ttention_qwen3_q32_k16_causal_bound_wgsl.h | 311 +++++++++++ backends/webgpu/scripts/gen_wgsl_headers.py | 27 +- .../test/native/test_compute_dispatch.cpp | 8 +- .../webgpu/test/native/test_dispatch_2d.cpp | 1 + .../webgpu/test/native/test_dynamic_shape.cpp | 145 +++++- .../test_dynamic_shape_export.py | 97 +++- backends/webgpu/test/ops/test_sdpa.py | 130 ++++- backends/webgpu/test/test_webgpu_native.cpp | 485 +++++++++++++++++- backends/webgpu/test/test_wgsl_codegen.py | 43 ++ 19 files changed, 2488 insertions(+), 122 deletions(-) create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index bedde5abdb7..35d4225bc7c 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -90,7 +90,7 @@ Result WebGPUBackend::init( // Parse header to locate flatbuffer and constant data Result header = - WebGPUDelegateHeader::parse(processed->data()); + WebGPUDelegateHeader::parse(processed->data(), processed->size()); if (!header.ok()) { ET_LOG(Error, "WebGPUDelegateHeader may be corrupt"); return header.error(); @@ -101,6 +101,17 @@ Result WebGPUBackend::init( const uint8_t* flatbuffer_data = buffer_start + header->flatbuffer_offset; const uint8_t* constant_data = buffer_start + header->bytes_offset; + size_t constant_data_size = header->bytes_size; + if (constant_data_size == 0 && processed->size() > header->bytes_offset) { + constant_data_size = processed->size() - header->bytes_offset; + } + + flatbuffers::Verifier verifier(flatbuffer_data, header->flatbuffer_size); + if (!vkgraph::VerifyVkGraphBuffer(verifier)) { + ET_LOG(Error, "WebGPU delegate FlatBuffer verification failed"); + return Error::DelegateInvalidCompatibility; + } + // Verify FlatBuffer identifier if (!vkgraph::VkGraphBufferHasIdentifier(flatbuffer_data)) { ET_LOG( @@ -125,10 +136,20 @@ Result WebGPUBackend::init( config.f16_accumulate_gemm = spec.get(); } } + { + Result spec = context.get_runtime_spec("sdpa_query_tile"); + if (spec.ok()) { + config.sdpa_query_tile = spec.get(); + } + } try { graph->build( - flatbuffer_data, constant_data, context.get_named_data_map(), config); + flatbuffer_data, + constant_data, + constant_data_size, + context.get_named_data_map(), + config); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU graph build failed: %s", e.what()); graph->~WebGPUGraph(); @@ -163,8 +184,13 @@ Error WebGPUBackend::execute( const auto& tensor = args[i]->toTensor(); const bool host_is_int64 = tensor.scalar_type() == executorch::aten::ScalarType::Long; + const bool host_is_fp32 = + tensor.scalar_type() == executorch::aten::ScalarType::Float; inputs.push_back( - {tensor.const_data_ptr(), tensor.nbytes(), host_is_int64}); + {tensor.const_data_ptr(), + tensor.nbytes(), + host_is_int64, + host_is_fp32}); const auto sizes = tensor.sizes(); std::vector new_dims(sizes.begin(), sizes.end()); graph->resize_input(graph->input_ids()[i], new_dims); @@ -205,12 +231,15 @@ Error WebGPUBackend::execute( graph->execute(plan); // Copy outputs from GPU staging buffers to EValue tensor data pointers - std::vector> outputs; + std::vector outputs; outputs.reserve(num_outputs); for (size_t i = 0; i < num_outputs; i++) { const size_t arg_idx = num_inputs + i; auto& tensor = args[arg_idx]->toTensor(); - outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); + const bool host_is_fp32 = + tensor.scalar_type() == executorch::aten::ScalarType::Float; + outputs.push_back( + {tensor.mutable_data_ptr(), tensor.nbytes(), host_is_fp32}); } graph->copy_outputs(outputs, plan); } catch (const std::exception& e) { diff --git a/backends/webgpu/runtime/WebGPUDelegateHeader.cpp b/backends/webgpu/runtime/WebGPUDelegateHeader.cpp index d1e8b2110a7..69e6dd70536 100644 --- a/backends/webgpu/runtime/WebGPUDelegateHeader.cpp +++ b/backends/webgpu/runtime/WebGPUDelegateHeader.cpp @@ -65,13 +65,19 @@ bool WebGPUDelegateHeader::is_valid() const { if (flatbuffer_size == 0) { return false; } - if (bytes_offset < flatbuffer_offset + flatbuffer_size) { + if (bytes_offset < flatbuffer_offset || + flatbuffer_size > bytes_offset - flatbuffer_offset) { return false; } return true; } -Result WebGPUDelegateHeader::parse(const void* data) { +Result WebGPUDelegateHeader::parse( + const void* data, + size_t buffer_size) { + if (data == nullptr || buffer_size < kExpectedSize) { + return Error::InvalidArgument; + } const uint8_t* header_data = (const uint8_t*)data; const uint8_t* magic_start = header_data + kMagic.offset; @@ -91,6 +97,13 @@ Result WebGPUDelegateHeader::parse(const void* data) { return Error::InvalidArgument; } + if (header.flatbuffer_offset > buffer_size || + header.flatbuffer_size > buffer_size - header.flatbuffer_offset || + header.bytes_offset > buffer_size || + header.bytes_size > buffer_size - header.bytes_offset) { + return Error::InvalidArgument; + } + return header; } diff --git a/backends/webgpu/runtime/WebGPUDelegateHeader.h b/backends/webgpu/runtime/WebGPUDelegateHeader.h index 6f2f65130c7..14dcfd14a9a 100644 --- a/backends/webgpu/runtime/WebGPUDelegateHeader.h +++ b/backends/webgpu/runtime/WebGPUDelegateHeader.h @@ -8,6 +8,8 @@ #pragma once +#include + #include namespace executorch { @@ -18,7 +20,8 @@ struct WebGPUDelegateHeader { bool is_valid() const; static executorch::runtime::Result parse( - const void* data); + const void* data, + size_t buffer_size); uint32_t header_size; uint32_t flatbuffer_offset; diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 859e01e7a40..58d7178ddbf 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -34,6 +35,19 @@ namespace executorch::backends::webgpu { namespace { +const uint8_t* checked_inline_constant( + const uint8_t* data, + size_t data_size, + uint64_t offset, + size_t required_size, + const char* error_message) { + if (data == nullptr || offset > data_size || + required_size > data_size - static_cast(offset)) { + throw std::runtime_error(error_message); + } + return data + static_cast(offset); +} + class ScopedBindGroupLayout final { public: explicit ScopedBindGroupLayout(WGPUBindGroupLayout handle) @@ -713,6 +727,7 @@ WebGPUGraph::~WebGPUGraph() { void WebGPUGraph::build( const void* flatbuffer_data, const uint8_t* constant_data, + size_t constant_data_size, const executorch::runtime::NamedDataMap* named_data_map, WebGPUGraphConfig config) { if (!device_) { @@ -733,6 +748,7 @@ void WebGPUGraph::build( // .pte byte sources for prepack-time constant materialization (build-only). constant_data_ = constant_data; + constant_data_size_ = constant_data_size; named_data_map_ = named_data_map; // f16 KV cache (runtime opt-in): store K/V caches as f16 iff the opt-in is @@ -817,11 +833,27 @@ void WebGPUGraph::build( continue; } for (unsigned j = 0; j < a->size(); j++) { - if (kv_cache_ids_.count(static_cast(a->Get(j))) != 0) { + const int id = static_cast(a->Get(j)); + if (kv_cache_ids_.count(id) != 0) { throw std::runtime_error( "WebGPU f16 KV: cache tensor consumed by non-sdpa op '" + nm + "' would misread the f16 buffer"); } + const auto* value = values ? values->Get(id) : nullptr; + if (value && value->value_type() == vkgraph::GraphTypes::ValueList) { + const auto* items = value->value_as_ValueList()->items(); + if (!items) { + continue; + } + for (unsigned k = 0; k < items->size(); k++) { + if (kv_cache_ids_.count(static_cast(items->Get(k))) != 0) { + throw std::runtime_error( + "WebGPU f16 KV: cache tensor consumed through a ValueList " + "by non-sdpa op '" + + nm + "' would misread the f16 buffer"); + } + } + } } } } @@ -843,11 +875,20 @@ void WebGPUGraph::build( size_t numel = 1; if (dims) { for (unsigned j = 0; j < dims->size(); j++) { - tensor.dims.push_back(static_cast(dims->Get(j))); - numel *= dims->Get(j); + const uint32_t dim = dims->Get(j); + tensor.dims.push_back(static_cast(dim)); + if (dim != 0 && numel > std::numeric_limits::max() / dim) { + throw std::runtime_error( + "WebGPU: tensor element count overflows"); + } + numel *= dim; } } tensor.elem_size = vk_datatype_size(vk_tensor->datatype()); + if (tensor.elem_size != 0 && + numel > std::numeric_limits::max() / tensor.elem_size) { + throw std::runtime_error("WebGPU: tensor byte size overflows"); + } tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8; tensor.nbytes = numel * tensor.elem_size; @@ -860,6 +901,11 @@ void WebGPUGraph::build( // zero-initializes freshly-created buffers, so no explicit clear is // needed. Inert unless kv_f16_ (runtime opt-in) is set. if (kv_f16_ && kv_cache_ids_.count(i) != 0) { + if (tensor.is_int || tensor.elem_size != sizeof(float) || + tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error( + "WebGPU f16 KV: serialized cache tensor must be fp32"); + } tensor.elem_size = 2; tensor.nbytes = numel * 2; tensor.cur_nbytes = tensor.nbytes; @@ -870,6 +916,63 @@ void WebGPUGraph::build( WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; tensor.buffer = wgpuDeviceCreateBuffer(device_, &buf_desc); + + // Mutable caches normally start empty. If the serialized graph owns + // an initialized cache constant, preserve it while changing storage + // representation instead of silently replacing it with zeros. + const int cache_constant_id = vk_tensor->constant_id(); + if (cache_constant_id >= 0) { + const auto* constants = graph->constants(); + if (!constants || + cache_constant_id >= static_cast(constants->size())) { + throw std::runtime_error( + "WebGPU f16 KV: cache constant id is out of range"); + } + const auto* bytes = constants->Get(cache_constant_id); + auto write_fp16_cache = [&](const uint8_t* src) { + std::vector converted(numel); + for (size_t e = 0; e < numel; e++) { + float value = 0.0f; + std::memcpy(&value, src + e * sizeof(float), sizeof(float)); + converted[e] = executorch::runtime::etensor::Half(value); + } + wgpuQueueWriteBuffer( + queue_, + tensor.buffer, + 0, + converted.data(), + converted.size() * sizeof(converted[0])); + }; + if (bytes->offset() != UINT64_MAX) { + write_fp16_cache(checked_inline_constant( + constant_data_, + constant_data_size_, + bytes->offset(), + numel * sizeof(float), + "WebGPU f16 KV: inline cache constant exceeds constant " + "data")); + } else if ( + bytes->named_key() != nullptr && named_data_map_ != nullptr) { + const std::string key = bytes->named_key()->str(); + auto data = named_data_map_->get_data(key.c_str()); + if (!data.ok()) { + throw std::runtime_error( + "WebGPU f16 KV: named cache constant '" + key + + "' not found"); + } + if (data->size() < numel * sizeof(float)) { + data->Free(); + throw std::runtime_error( + "WebGPU f16 KV: named cache constant '" + key + + "' is undersized"); + } + write_fp16_cache(static_cast(data->data())); + data->Free(); + } else { + throw std::runtime_error( + "WebGPU f16 KV: cache constant has no readable source"); + } + } break; } @@ -1190,6 +1293,7 @@ void WebGPUGraph::build( // The .pte bytes are freed right after build() returns (WebGPUBackend // processed->Free()), so clear the build-only source pointers. constant_data_ = nullptr; + constant_data_size_ = 0; named_data_map_ = nullptr; } @@ -1201,15 +1305,18 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { std::to_string(const_value_id)); } const ConstantSource& cs = it->second; - if (cs.nbytes == 0) { - return; - } if (cs.inline_offset != UINT64_MAX) { - if (constant_data_ == nullptr) { - throw std::runtime_error("WebGPU: inline constant data is null"); - } - wgpuQueueWriteBuffer( - queue_, dst, 0, constant_data_ + cs.inline_offset, cs.nbytes); + const uint8_t* data = checked_inline_constant( + constant_data_, + constant_data_size_, + cs.inline_offset, + cs.nbytes, + "WebGPU: inline constant exceeds constant data"); + if (cs.nbytes != 0) { + wgpuQueueWriteBuffer(queue_, dst, 0, data, cs.nbytes); + } + } else if (cs.nbytes == 0) { + return; } else if (!cs.named_key.empty() && named_data_map_ != nullptr) { auto buf = named_data_map_->get_data(cs.named_key.c_str()); if (!buf.ok()) { @@ -1304,6 +1411,11 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { // Upload only the live (cur) bytes, not the max allocation; cur_nbytes == // nbytes on a static graph, so this is byte-identical there. const size_t live_nbytes = tensor.cur_nbytes; + const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2; + if (buffer_is_fp16 && !in.host_is_fp32) { + throw std::runtime_error( + "WebGPU: fp16 device input requires an fp32 host tensor"); + } // Fast path: host and GPU element types match byte-for-byte. if (in.nbytes == live_nbytes) { @@ -1332,6 +1444,21 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { continue; } + // Require an explicit fp32 host dtype, not merely "not int64": inferring + // the narrow from the 2:1 byte ratio alone would silently reinterpret a + // same-sized non-fp32 host buffer (e.g. a stale int32) as fp32. + if (in.host_is_fp32 && buffer_is_fp16 && in.nbytes == live_nbytes * 2) { + const size_t numel = live_nbytes / sizeof(uint16_t); + const float* src = static_cast(in.data); + std::vector narrowed(numel); + for (size_t e = 0; e < numel; e++) { + narrowed[e] = executorch::runtime::etensor::Half(src[e]); + } + wgpuQueueWriteBuffer( + queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + continue; + } + throw std::runtime_error( "WebGPU: unsupported input copy for input " + std::to_string(i) + " (host " + std::to_string(in.nbytes) + " bytes" + @@ -1366,6 +1493,8 @@ constexpr uint32_t kRouteGenericFallback = 1u << 8; constexpr uint32_t kRouteFlashDecoding = 1u << 10; constexpr uint32_t kRouteK16CausalBound = 1u << 11; constexpr uint32_t kRouteBicolSubgroup = 1u << 12; +constexpr uint32_t kRouteQwen3Q16K16 = 1u << 13; +constexpr uint32_t kRouteQwen3Q32K16 = 1u << 14; #endif // WGPU_BACKEND_ENABLE_PROFILING // Bench gate: compiled out unless WGPU_BACKEND_ENABLE_PROFILING; then the @@ -1382,7 +1511,13 @@ bool should_timestamp_query() { #ifdef WGPU_BACKEND_ENABLE_PROFILING void WebGPUGraph::record_active_route(const std::string& kernel_name) { uint32_t bits = 0; - if (kernel_name == "sdpa_streaming_attention_k16_causal_bound") { + if (kernel_name == "sdpa_streaming_attention_qwen3_q32_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound | kRouteQwen3Q32K16; + } else if (kernel_name == "sdpa_streaming_attention_qwen3_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound | kRouteQwen3Q16K16; + } else if ( + kernel_name.rfind("sdpa_streaming_attention_", 0) == 0 && + kernel_name.find("k16_causal_bound") != std::string::npos) { bits = kRoutePrefill | kRouteK16CausalBound; } else if (kernel_name == "sdpa_streaming_attention_k16") { bits = kRoutePrefill | kRouteK16; @@ -1665,28 +1800,59 @@ void buffer_map_callback( } // namespace void WebGPUGraph::copy_outputs( - std::vector>& outputs, + std::vector& outputs, const WebGPUExecutionPlan& plan) { if (plan.copy_outputs.size() != output_copies_.size()) { throw std::runtime_error("WebGPU: execution plan output count mismatch"); } const size_t count = std::min(outputs.size(), output_staging_buffers_.size()); + // Reject all dtype/size mismatches before issuing an asynchronous map. for (size_t i = 0; i < count; i++) { - if (!plan.copy_outputs[i] || outputs[i].second == 0) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { continue; } - const size_t map_nbytes = tensors_[output_ids_[i]].cur_nbytes; + const auto& tensor = tensors_[output_ids_[i]]; + const size_t map_nbytes = tensor.cur_nbytes; if (map_nbytes == 0) { continue; } - const size_t dst_nbytes = outputs[i].second; - const bool widen_int32 = dst_nbytes == 2 * map_nbytes && - tensors_[output_ids_[i]].is_int && - tensors_[output_ids_[i]].elem_size == 4; - if (dst_nbytes != map_nbytes && !widen_int32) { + const size_t dst_nbytes = outputs[i].nbytes; + const bool is_double_width = + dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + const bool widen_fp16 = + is_double_width && !tensor.is_int && tensor.elem_size == 2; + const bool widen_int32 = + is_double_width && tensor.is_int && tensor.elem_size == 4; + const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2; + if (buffer_is_fp16 && !outputs[i].host_is_fp32) { + throw std::runtime_error( + "WebGPU: fp16 device output requires an fp32 host tensor"); + } + if (outputs[i].host_is_fp32 && buffer_is_fp16 && !widen_fp16) { + throw std::runtime_error("WebGPU: fp16 output buffer size mismatch"); + } + if (dst_nbytes != map_nbytes && !widen_fp16 && !widen_int32) { throw std::runtime_error("WebGPU: output buffer size mismatch"); } + } + + for (size_t i = 0; i < count; i++) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { + continue; + } + const auto& tensor = tensors_[output_ids_[i]]; + const size_t map_nbytes = tensor.cur_nbytes; + if (map_nbytes == 0) { + continue; + } + const size_t dst_nbytes = outputs[i].nbytes; + const bool is_double_width = + dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + const bool widen_fp16 = + is_double_width && !tensor.is_int && tensor.elem_size == 2; + const bool widen_int32 = + is_double_width && tensor.is_int && tensor.elem_size == 4; const auto cb_data = std::make_shared(); WGPUBufferMapCallbackInfo cb_info = {}; @@ -1717,16 +1883,24 @@ void WebGPUGraph::copy_outputs( wgpuBufferUnmap(output_staging_buffers_[i]); throw std::runtime_error("WebGPU mapped output range is null"); } - if (!widen_int32) { - std::memcpy(outputs[i].first, mapped, map_nbytes); - } else { + if (widen_fp16) { + const auto* src = + static_cast(mapped); + auto* dst = static_cast(outputs[i].data); + const size_t n = map_nbytes / sizeof(*src); + for (size_t k = 0; k < n; k++) { + dst[k] = static_cast(src[k]); + } + } else if (widen_int32) { // int64 host output backed by an int32 GPU buffer: widen (sign-extend). const int32_t* src = static_cast(mapped); - int64_t* dst = static_cast(outputs[i].first); + int64_t* dst = static_cast(outputs[i].data); const size_t n = map_nbytes / sizeof(int32_t); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } + } else { + std::memcpy(outputs[i].data, mapped, map_nbytes); } wgpuBufferUnmap(output_staging_buffers_[i]); } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 292f2ee0cf4..a634ebb4172 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -48,6 +48,15 @@ struct InputData { const void* data = nullptr; size_t nbytes = 0; bool host_is_int64 = false; + bool host_is_fp32 = false; +}; + +// Host destination for a graph output. host_is_fp32 gates the fp16->fp32 widen +// on readback (mirrors InputData's guard on the copy_inputs narrow path). +struct OutputData { + void* data = nullptr; + size_t nbytes = 0; + bool host_is_fp32 = false; }; struct WebGPUDispatch { @@ -130,6 +139,7 @@ struct WebGPUMemoryStats { struct WebGPUGraphConfig { bool f16_kv_cache = false; bool f16_accumulate_gemm = false; + int sdpa_query_tile = 0; bool record_q4gsw_decode_route = false; }; @@ -143,6 +153,7 @@ class WebGPUGraph { void build( const void* flatbuffer_data, const uint8_t* constant_data, + size_t constant_data_size, const executorch::runtime::NamedDataMap* named_data_map = nullptr, WebGPUGraphConfig config = {}); @@ -159,7 +170,7 @@ class WebGPUGraph { // Copy output tensor data from GPU buffers back to host pointers. // Uses mapAsync + ASYNCIFY in Wasm. void copy_outputs( - std::vector>& outputs, + std::vector& outputs, const WebGPUExecutionPlan& plan); const std::vector& input_ids() const { @@ -541,6 +552,12 @@ class WebGPUGraph { return config_.f16_accumulate_gemm; } + // Runtime-selected SDPA query-tile candidate; 0 = geometry default (Q16), + // 32 = Q32 candidate. + int sdpa_query_tile() const { + return config_.sdpa_query_tile; + } + const WebGPUGraphConfig& config() const { return config_; } @@ -659,6 +676,7 @@ class WebGPUGraph { // materializes these once. constant_data_/named_data_map_ point at the .pte // bytes and are valid only during build(). const uint8_t* constant_data_ = nullptr; + size_t constant_data_size_ = 0; const executorch::runtime::NamedDataMap* named_data_map_ = nullptr; std::unordered_map constant_sources_; diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 4a282f8266e..1cf4b952a5f 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -114,6 +114,8 @@ #include #include #include +#include +#include #include #include #include @@ -149,7 +151,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -1011,6 +1013,20 @@ constexpr std::array kShaderRegistry = {{ kStreamingAttentionK16CausalBoundWorkgroupSizeY, kStreamingAttentionK16CausalBoundWorkgroupSizeZ, }, + { + "streaming_attention_qwen3_k16_causal_bound", + kStreamingAttentionQwen3K16CausalBoundWGSL, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeX, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeY, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeZ, + }, + { + "streaming_attention_qwen3_q32_k16_causal_bound", + kStreamingAttentionQwen3Q32K16CausalBoundWGSL, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeY, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeZ, + }, { "tanh", kTanhWGSL, diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 9a51f853caf..bcb5deaa770 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -38,6 +38,10 @@ constexpr const char* kComputeOutShader = "sdpa_compute_out"; constexpr const char* kComputeOutHalfShader = "sdpa_compute_out_half"; constexpr const char* kStreamingK16Shader = "streaming_attention_k16_causal_bound"; +constexpr const char* kStreamingQwen3K16Shader = + "streaming_attention_qwen3_k16_causal_bound"; +constexpr const char* kStreamingQwen3Q32K16Shader = + "streaming_attention_qwen3_q32_k16_causal_bound"; // Uniform param structs (all 16-byte aligned, matching the WGSL Params). struct UpdateCacheParams { @@ -205,13 +209,87 @@ static bool streaming_attention_k16_device_supported(WGPUDevice device) { limits.maxComputeWorkgroupStorageSize >= 14720u; } +constexpr uint32_t kLlamaK16QueryTile = 32u; +constexpr uint32_t kQwen3K16QueryTile = 16u; +constexpr uint32_t kQwen3Q32K16QueryTile = 32u; +constexpr uint32_t kQwen3Q16K16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); +constexpr uint32_t kQwen3K16StorageBytes = kQwen3Q16K16StorageBytes; +// Mirrors the Q32 shader's workgroup arrays (t_q_tile vec4x1024, t_kv_tile +// vec4x512, t_scores vec2x256, t_m/t_d/t_alpha f32x32) so the +// device-support gate stays tied to the declared storage, not a literal. +constexpr uint32_t kQwen3Q32K16StorageBytes = 1024u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 256u * 2u * sizeof(float) + + 3u * 32u * sizeof(float); + +constexpr bool streaming_attention_k16_workgroup_count_fits( + int64_t S, + int64_t Hkv, + int64_t g, + uint32_t query_tile, + uint32_t max_workgroups) { + if (S <= 0 || Hkv <= 0 || g <= 0 || query_tile == 0u || + max_workgroups == 0u) { + return false; + } + if (static_cast(S) > UINT64_MAX / static_cast(g)) { + return false; + } + const uint64_t logical_rows = + static_cast(S) * static_cast(g); + if (logical_rows > UINT64_MAX - (query_tile - 1u)) { + return false; + } + const uint64_t groups_per_kv = (logical_rows + query_tile - 1u) / query_tile; + if (groups_per_kv > UINT64_MAX / static_cast(Hkv)) { + return false; + } + const uint64_t workgroups = groups_per_kv * static_cast(Hkv); + return workgroups > 0u && workgroups <= UINT32_MAX && + workgroups <= max_workgroups; +} + +static_assert( + streaming_attention_k16_workgroup_count_fits(65528, 8, 2, 16, 65535)); +static_assert( + !streaming_attention_k16_workgroup_count_fits(65529, 8, 2, 16, 65535)); + +bool qwen3_q16_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQwen3K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +bool qwen3_q32_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupStorageSize >= kQwen3Q32K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + static utils::WgCount streaming_attention_k16_grid( WGPUDevice device, int64_t S, int64_t Hkv, - int64_t g) { + int64_t g, + uint32_t query_tile) { const uint64_t groups_per_kv = - (static_cast(S) * static_cast(g) + 31u) / 32u; + (static_cast(S) * static_cast(g) + query_tile - 1u) / + query_tile; const uint64_t workgroups = static_cast(Hkv) * groups_per_kv; if (workgroups == 0u || workgroups > UINT32_MAX) { throw std::runtime_error("WebGPU sdpa: K16 workgroup count exceeds uint32"); @@ -362,7 +440,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("WebGPU sdpa: output shape must match q"); } - // fp32-only: validate byte counts against fp32 element counts. + // q/k/v/out are serialized fp32. KV caches are fp32 by default and use + // dedicated fp16 storage only when the graph-level option is active. auto numel = [](const WebGPUTensor& t) { uint64_t n = 1; for (int64_t d : t.dims) { @@ -370,11 +449,23 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { } return n; }; - if (q.nbytes != numel(q) * sizeof(float) || - k.nbytes != numel(k) * sizeof(float) || - v.nbytes != numel(v) * sizeof(float) || - out.nbytes != numel(out) * sizeof(float)) { - throw std::runtime_error("WebGPU sdpa: fp32-only (byte-size mismatch)"); + auto is_fp32 = [&numel](const WebGPUTensor& t) { + return !t.is_int && t.elem_size == sizeof(float) && + t.nbytes == numel(t) * sizeof(float); + }; + if (!is_fp32(q) || !is_fp32(k) || !is_fp32(v) || !is_fp32(out)) { + throw std::runtime_error("WebGPU sdpa: q/k/v/output must be fp32"); + } + const size_t cache_elem_size = + graph.kv_f16() ? sizeof(uint16_t) : sizeof(float); + auto cache_storage_is_valid = [&numel, + cache_elem_size](const WebGPUTensor& t) { + return !t.is_int && t.elem_size == cache_elem_size && + t.nbytes == numel(t) * cache_elem_size; + }; + if (!cache_storage_is_valid(k_cache) || !cache_storage_is_valid(v_cache)) { + throw std::runtime_error( + "WebGPU sdpa: cache dtype does not match the selected storage mode"); } // input_pos: build-time Int (baked) OR runtime SymInt (dynamic decode). @@ -435,9 +526,50 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { k16_buffers_distinct && k16_buffers[i] != k16_buffers[j]; } } - const bool k16_eligible = graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && - D == 64 && scale == 0.125f && out.dims == q.dims && - k16_buffers_distinct && streaming_attention_k16_device_supported(device); + // The specialized shaders bake the standard Qwen3 scale, so eligibility must + // be exact. A nearby explicit scale has different operator semantics and must + // use the general path. + const float qwen3_expected_scale = 1.0f / std::sqrt(128.0f); + const bool qwen3_k16_geometry = Hq == 16 && Hkv == 8 && g == 2 && D == 128 && + scale == qwen3_expected_scale && out.dims == q.dims; + // Q16 is the default route for exact Qwen3 geometry; Q32 is an explicit + // autotuning candidate requested via the sdpa_query_tile RuntimeSpec. Support + // is evaluated per-tile so an unsupported Q32 request falls back to the Q16 + // streaming route instead of dropping to the materialized path. + const uint32_t device_max_workgroups = utils::queried_max_workgroups(device); + const bool qwen3_q16_supported = + qwen3_k16_geometry && graph.kv_f16() && + qwen3_q16_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kQwen3K16QueryTile, device_max_workgroups); + const bool qwen3_q32_requested = qwen3_k16_geometry && + graph.sdpa_query_tile() == static_cast(kQwen3Q32K16QueryTile); + const bool qwen3_q32_supported = + qwen3_q32_requested && graph.kv_f16() && + qwen3_q32_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kQwen3Q32K16QueryTile, device_max_workgroups); + const bool qwen3_q32_selected = qwen3_q32_supported; + const bool qwen3_k16_selected = qwen3_q32_selected || qwen3_q16_supported; + const uint32_t qwen3_query_tile = + qwen3_q32_selected ? kQwen3Q32K16QueryTile : kQwen3K16QueryTile; + const bool llama_k16_eligible = + graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && D == 64 && + scale == 0.125f && out.dims == q.dims && + streaming_attention_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kLlamaK16QueryTile, device_max_workgroups); + const bool k16_eligible = + k16_buffers_distinct && (llama_k16_eligible || qwen3_k16_selected); + const uint32_t k16_query_tile = + qwen3_k16_selected ? qwen3_query_tile : kLlamaK16QueryTile; + const char* k16_shader = qwen3_q32_selected ? kStreamingQwen3Q32K16Shader + : qwen3_k16_selected ? kStreamingQwen3K16Shader + : kStreamingK16Shader; + const char* k16_label = qwen3_q32_selected + ? "sdpa_streaming_attention_qwen3_q32_k16_causal_bound" + : qwen3_k16_selected ? "sdpa_streaming_attention_qwen3_k16_causal_bound" + : "sdpa_streaming_attention_k16_causal_bound"; const uint32_t uc_wg = utils::clamp_workgroup_size( device, get_webgpu_shader_info(kUpdateCacheShader).workgroup_size_x); const uint32_t qk_wg = utils::clamp_workgroup_size( @@ -469,7 +601,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { qk_wg, av_wg, fd_eligible, - k16_eligible](WebGPUGraph& gr) { + k16_eligible, + k16_query_tile](WebGPUGraph& gr) { SdpaLiveState state = {}; const auto& q_live_dims = gr.cur_dims(q_id); state.s = q_live_dims[qn - 3]; @@ -521,8 +654,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (k16_eligible) { state.streaming_k16 = make_streaming_attention_k16_params( state.s, state.context_len, state.pos, Hq, Hkv, D); - state.streaming_k16_grid = - streaming_attention_k16_grid(gr.device(), state.s, Hkv, g); + state.streaming_k16_grid = streaming_attention_k16_grid( + gr.device(), state.s, Hkv, g, k16_query_tile); } state.update_cache_grid = { utils::compute_1d_workgroup_count( @@ -682,13 +815,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { : initial_state.streaming_k16_grid; add_sdpa_compute_dispatch( graph, - kStreamingK16Shader, + k16_shader, k16_bindings, k16_buf, sizeof(StreamingAttentionK16Params), initial_grid, 0, - "sdpa_streaming_attention_k16_causal_bound"); + k16_label); k16_idx = graph.num_dispatches() - 1; k16_range.end = graph.num_dispatches(); } diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl new file mode 100644 index 00000000000..5e27346af6f --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..f69a4d72dae --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * 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 streaming_attention_qwen3_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: b26e81d92ba58c833f7dfbf2da6bce1b9cb468d70d893cef9df3a7fae16ee1fc +inline constexpr const char* kStreamingAttentionQwen3K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeX = + 16; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeY = + 8; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeZ = + 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl new file mode 100644 index 00000000000..6135b15b75d --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..e58a9d94bac --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * 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 streaming_attention_qwen3_q32_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: 6d4396945b86cc3445698b1182fc4535ca786ad509ba128a6b6b08977e6a55e0 +inline constexpr const char* kStreamingAttentionQwen3Q32K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeY = 8; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index cf044232dde..dd8cc4cc635 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -437,14 +437,17 @@ def embedded_sha256(header_text: str) -> str: def _wg_size_const(base: str, axis: str, val: int) -> str: """One WorkgroupSize constant; wrap to <=80 cols so CLANGFORMAT accepts it. - Long shader names push the single-line form past the 80-col limit (clang-format - then breaks after '=' with a 4-space continuation indent); emit that wrapped - form up front so the generated header matches lintrunner's CLANGFORMAT. + Long shader names push the single-line form past the 80-col limit. Emit the + wrapped form that clang-format selects so generated headers stay byte-stable. """ - decl = f"inline constexpr uint32_t k{base}WorkgroupSize{axis} =" - if len(decl) + len(f" {val};") > 80: - return f"{decl}\n {val};\n" - return f"{decl} {val};\n" + name = f"k{base}WorkgroupSize{axis}" + prefix = f"inline constexpr uint32_t {name} =" + decl = f"{prefix} {val};" + if len(decl) > 85: + return f"inline constexpr uint32_t\n {name} = {val};\n" + if len(decl) > 80: + return f"{prefix}\n {val};\n" + return f"{decl}\n" def render_header( @@ -469,6 +472,14 @@ def render_header( raise ValueError('shader contains )" which would close the R"( literal') base = symbol_base(name) x, y, z = parse_workgroup_size(wgsl_text) + provenance = f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT." + if len(provenance) > 80: + provenance_lines = [ + f"// @generated from {provenance_stem}.wgsl", + "// DO NOT EDIT.", + ] + else: + provenance_lines = [provenance] head = [ _BSD_HEADER, @@ -479,7 +490,7 @@ def render_header( "", "namespace executorch::backends::webgpu {", "", - f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT.", + *provenance_lines, f"// wgsl-sha256: {wgsl_sha256(wgsl_text)}", f'inline constexpr const char* k{base}WGSL = R"(', ] diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index feb6854c59d..b2b0d028567 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -120,7 +120,7 @@ void build_q4_route_graph(WebGPUGraph& graph, Q4RouteSignal signal) { if (signal == Q4RouteSignal::ExplicitOption) { config.record_q4gsw_decode_route = true; } - graph.build(fbb.GetBufferPointer(), nullptr, nullptr, config); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr, config); } std::vector q4_dispatches(WebGPUGraph& graph) { @@ -407,7 +407,7 @@ void build_resize_test_graph(WebGPUGraph& graph) { vk::FinishVkGraphBuffer(fbb, root); graph.set_device(g_device); - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } WebGPUComputeDispatchDescriptor make_dynamic_test_descriptor( @@ -707,7 +707,7 @@ void expect_invalid_rope_graph(const InvalidRopeGraphCase& test_case) { WebGPUGraph graph; std::string error; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& exception) { error = exception.what(); } @@ -801,7 +801,7 @@ TEST(WebGPUExecution, RejectsPlanOutputCountMismatch) { WebGPUGraph graph; WebGPUExecutionPlan plan; plan.copy_outputs = {true}; - std::vector> outputs; + std::vector outputs; EXPECT_THROW(graph.execute(plan), std::runtime_error); EXPECT_THROW(graph.copy_outputs(outputs, plan), std::runtime_error); diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index c9ce9648b43..7c5a5e141f3 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -232,6 +232,7 @@ TEST(WebGPUGraphConfig, ParsesExactBooleanCompileOption) { EXPECT_FALSE(absent->record_q4gsw_decode_route); EXPECT_FALSE(absent->f16_kv_cache); EXPECT_FALSE(absent->f16_accumulate_gemm); + EXPECT_EQ(absent->sdpa_query_tile, 0); uint8_t false_value = 0; CompileSpec false_spec = { diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 9b0cf65a723..0e72ce6234d 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -556,6 +556,9 @@ void check_sdpa(int s) { constexpr int kK16Hq = 32; constexpr int kK16Hkv = 8; constexpr int kK16D = 64; +constexpr int kQwen3Hq = 16; +constexpr int kQwen3Hkv = 8; +constexpr int kQwen3D = 128; bool k16_device_supported() { const auto* context = get_default_webgpu_context(); @@ -568,6 +571,21 @@ bool k16_device_supported() { limits.maxComputeWorkgroupStorageSize >= 14720u; } +bool qwen3_q16_device_supported() { + constexpr uint32_t kQ16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); + const auto* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQ16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + void load_sdpa_module(Module& module, bool f16_kv) { if (!f16_kv) { ASSERT_EQ(module.load_forward(), Error::Ok); @@ -587,9 +605,13 @@ void run_k16_sdpa( int hq = kK16Hq, int hkv = kK16Hkv, int d = kK16D, - bool prime = false) { - const std::string base = g_dir + "/" + prefix + - (prime ? ".prime." : ".S" + std::to_string(s) + "."); + bool prime = false, + float max_error_limit = 3e-3f, + bool initial = false) { + const std::string suffix = initial ? ".initial." + : prime ? ".prime." + : ".S" + std::to_string(s) + "."; + const std::string base = g_dir + "/" + prefix + suffix; auto q = read_bin(base + "q.bin"); auto k = read_bin(base + "k.bin"); auto v = read_bin(base + "v.bin"); @@ -614,8 +636,7 @@ void run_k16_sdpa( std::vector got( output.const_data_ptr(), output.const_data_ptr() + numel); ASSERT_EQ(got.size(), golden.size()); - constexpr float kMaxError = 3e-3f; - EXPECT_LT(max_err(got, golden), kMaxError) + EXPECT_LT(max_err(got, golden), max_error_limit) << prefix << " S=" << s << " full output"; const std::set tokens = { 0, std::min(15, s - 1), std::min(16, s - 1), s - 1}; @@ -625,7 +646,7 @@ void run_k16_sdpa( for (size_t i = begin; i < begin + token_width; ++i) { token_error = std::fmax(token_error, std::fabs(got[i] - golden[i])); } - EXPECT_LT(token_error, kMaxError) + EXPECT_LT(token_error, max_error_limit) << prefix << " S=" << s << " causal token=" << token; } } @@ -635,23 +656,22 @@ void prime_k16_sdpa( const char* prefix, int hq = kK16Hq, int hkv = kK16Hkv, - int d = kK16D) { - run_k16_sdpa(module, 12, prefix, hq, hkv, d, true); + int d = kK16D, + float max_error_limit = 3e-3f) { + run_k16_sdpa(module, 12, prefix, hq, hkv, d, true, max_error_limit); } #ifdef WGPU_BACKEND_ENABLE_PROFILING void expect_sdpa_route( const std::vector& names, int s, - bool expect_k16) { + bool expect_k16, + const char* k16_kernel_name = "sdpa_streaming_attention_k16_causal_bound") { const bool expect_fd = s == 1; const bool expect_materialized = !expect_fd && !expect_k16; EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); EXPECT_EQ( - std::count( - names.begin(), - names.end(), - "sdpa_streaming_attention_k16_causal_bound"), + std::count(names.begin(), names.end(), k16_kernel_name), expect_k16 ? 1 : 0); EXPECT_EQ( std::count(names.begin(), names.end(), "fd_split"), expect_fd ? 1 : 0); @@ -1437,6 +1457,46 @@ TEST(DynamicShape, K16CausalNumericsReusedGraph) { } } +TEST(DynamicShape, Qwen3K16CausalNumericsReusedGraph) { + if (!qwen3_q16_device_supported()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, "sdpa_k16_qwen3", kQwen3Hq, kQwen3Hkv, kQwen3D, kQwen3MaxError); + for (int s : {128, 1, 17, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + } +} + +TEST(DynamicShape, Qwen3InitializedConstantCachePreserved) { + if (!qwen3_q16_device_supported()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + run_k16_sdpa( + module, + 17, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + 1e-2f, + true); +} + TEST(DynamicShape, K16CacheHeadMismatchRejectedAtLoad) { executorch::runtime::BackendOptions<1> options; ASSERT_EQ(options.set_option("enable_f16_kv_cache", true), Error::Ok); @@ -1487,6 +1547,65 @@ TEST(DynamicShape, K16CausalLiveRoutesProfile) { } } +TEST(DynamicShape, Qwen3K16CausalLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !qwen3_q16_device_supported()) { + GTEST_SKIP() << "timestamp queries or Qwen3 Q16 K16 limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + constexpr const char* kQwen3Kernel = + "sdpa_streaming_attention_qwen3_k16_causal_bound"; + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, "sdpa_k16_qwen3", kQwen3Hq, kQwen3Hkv, kQwen3D, kQwen3MaxError); + expect_sdpa_route(current_profile_names(), 12, true, kQwen3Kernel); + for (int s : {128, 1, 17, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), s, s > 1, kQwen3Kernel); + } +} + +TEST(DynamicShape, Qwen3NearScaleFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !qwen3_q16_device_supported()) { + GTEST_SKIP() << "timestamp queries or Qwen3 Q16 K16 limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + Module module(g_dir + "/sdpa_k16_qwen3_near_scale.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, + "sdpa_k16_qwen3_near_scale", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3_near_scale", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), s, false); + } +} + TEST(DynamicShape, K16F32KvFallsBackToExistingRoutes) { 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 89a19f1fabc..a99c9baf46e 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 @@ -13,6 +13,7 @@ goldens. """ +import math import os import unittest @@ -782,6 +783,10 @@ def _export_static_linear(out_dir: str, m: int, prefix: str) -> None: K16_PRIME_S = K16_INPUT_POS - K16_PRIME_POS K16_POS_CONTROL = K16_INPUT_POS K16_LIVE_S = (K16_MAXS, 508, 128, 127, 16, 1) +QWEN3_HQ = 16 +QWEN3_HKV = 8 +QWEN3_D = 128 +QWEN3_LIVE_S = (128, 17, 1) def _export_dynamic_sdpa(out_dir: str) -> None: @@ -827,6 +832,12 @@ def _write_f32_tensors(base, tensors) -> None: tensor.detach().numpy().astype(" None: + if not initialized_cache: + k_cache.zero_() + v_cache.zero_() + + def _export_dynamic_k16_sdpa_case( out_dir: str, prefix: str, @@ -837,10 +848,14 @@ def _export_dynamic_k16_sdpa_case( scale: float | None = None, cache_hkv: int | None = None, expect_runtime_reject: bool = False, + denom: float = 16.0, + kv_f16_golden: bool = False, + initialized_cache: bool = False, ) -> None: from executorch.backends.webgpu.test.ops.test_sdpa import ( _det_inputs, _golden, + _round_kv_for_storage, SdpaConfig, ) @@ -853,14 +868,15 @@ def cfg(s: int) -> "SdpaConfig": s, K16_CMAX, K16_INPUT_POS, + denom, + kv_f16=kv_f16_golden, ) q, k, v, kc, vc = _det_inputs(cfg(K16_MAXS)) if cache_hkv is not None: kc = torch.zeros(1, K16_CMAX, cache_hkv, d) vc = torch.zeros_like(kc) - kc.zero_() - vc.zero_() + _zero_uninitialized_cache(initialized_cache, kc, vc) class K16SdpaModule(torch.nn.Module): def __init__(self) -> None: @@ -908,6 +924,8 @@ def forward(self, q, k, v, pos_control): def reference(live_cfg, q, k, v, kc, vc): if scale is None: return _golden(live_cfg, q, k, v, kc, vc) + k, v, kc, vc = _round_kv_for_storage(live_cfg, k, v, kc, vc) + runtime_scale = float(torch.tensor(scale, dtype=torch.float32).item()) context_len = live_cfg.s + live_cfg.input_pos g = hq // hkv k_full = torch.cat((kc[0, : live_cfg.input_pos].double(), k[0].double()), dim=0) @@ -919,14 +937,48 @@ def reference(live_cfg, q, k, v, kc, vc): for token in range(live_cfg.s): mask[token, : live_cfg.input_pos + token + 1] = 0.0 golden = torch.nn.functional.scaled_dot_product_attention( - q_heads, k_heads, v_heads, attn_mask=mask, scale=scale + q_heads, k_heads, v_heads, attn_mask=mask, scale=runtime_scale ) return golden.transpose(0, 1).reshape(1, live_cfg.s, hq, d).float().contiguous() - prime_cfg = SdpaConfig(prefix, hq, hkv, d, K16_PRIME_S, K16_CMAX, K16_PRIME_POS) + if initialized_cache: + initial_cfg = cfg(17) + initial_q, initial_k, initial_v, initial_kc, initial_vc = _det_inputs( + initial_cfg + ) + initial_golden = reference( + initial_cfg, + initial_q, + initial_k, + initial_v, + initial_kc, + initial_vc, + ) + initial_base = os.path.join(out_dir, f"{prefix}.initial.") + _write_f32_tensors( + initial_base, + ( + ("q", initial_q), + ("k", initial_k), + ("v", initial_v), + ("control", torch.zeros(1, K16_POS_CONTROL)), + ("golden", initial_golden), + ), + ) + + prime_cfg = SdpaConfig( + prefix, + hq, + hkv, + d, + K16_PRIME_S, + K16_CMAX, + K16_PRIME_POS, + denom, + kv_f16=kv_f16_golden, + ) prime_q, prime_k, prime_v, prime_kc, prime_vc = _det_inputs(prime_cfg) - prime_kc.zero_() - prime_vc.zero_() + _zero_uninitialized_cache(initialized_cache, prime_kc, prime_vc) prime_golden = reference(prime_cfg, prime_q, prime_k, prime_v, prime_kc, prime_vc) prime_base = os.path.join(out_dir, f"{prefix}.prime.") _write_f32_tensors( @@ -943,8 +995,7 @@ def reference(live_cfg, q, k, v, kc, vc): for s in live_s: live_cfg = cfg(s) q, k, v, kc, vc = _det_inputs(live_cfg) - kc.zero_() - vc.zero_() + _zero_uninitialized_cache(initialized_cache, kc, vc) kc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_k[0] vc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_v[0] golden = reference(live_cfg, q, k, v, kc, vc) @@ -972,6 +1023,28 @@ def _export_dynamic_k16_sdpa_cases(out_dir: str) -> None: K16_HKV, K16_LIVE_S, ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_qwen3", + QWEN3_HQ, + QWEN3_HKV, + QWEN3_LIVE_S, + d=QWEN3_D, + denom=10.0, + kv_f16_golden=True, + initialized_cache=True, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_qwen3_near_scale", + QWEN3_HQ, + QWEN3_HKV, + (128, 1), + d=QWEN3_D, + scale=1.0 / math.sqrt(float(QWEN3_D)) + 5e-7, + denom=10.0, + kv_f16_golden=True, + ) _export_dynamic_k16_sdpa_case( out_dir, "sdpa_k16_wrong_geometry", @@ -1447,12 +1520,16 @@ def test_export_dynamic_rms(self) -> None: "dyn_qkv_bk64_wrong_width.pte", "dyn_qkv_bk64_different_input.pte", "sdpa_k16_bad_cache_heads.pte", + "sdpa_k16_qwen3.pte", + "sdpa_k16_qwen3_near_scale.pte", ] for name in expected: with self.subTest(artifact=name): self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) for prefix, live_s in ( ("sdpa_k16_llama", K16_LIVE_S), + ("sdpa_k16_qwen3", QWEN3_LIVE_S), + ("sdpa_k16_qwen3_near_scale", (128, 1)), ("sdpa_k16_wrong_geometry", (128, 1)), ("sdpa_k16_wrong_d", (128, 1)), ("sdpa_k16_wrong_scale", (128, 1)), @@ -1472,6 +1549,10 @@ def test_export_dynamic_rms(self) -> None: self.assertGreater( os.path.getsize(os.path.join(d, name)), 0 ) + for kind in ("q", "k", "v", "control", "golden"): + name = f"sdpa_k16_qwen3.initial.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) if __name__ == "__main__": diff --git a/backends/webgpu/test/ops/test_sdpa.py b/backends/webgpu/test/ops/test_sdpa.py index 960c62e5203..b2606256e28 100644 --- a/backends/webgpu/test/ops/test_sdpa.py +++ b/backends/webgpu/test/ops/test_sdpa.py @@ -23,7 +23,7 @@ import torch import torch.nn.functional as F -from executorch.backends.vulkan import VulkanPartitioner +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower from executorch.extension.llm.custom_ops import ( # noqa: F401 registers llama ops custom_ops, @@ -40,6 +40,7 @@ class SdpaConfig: cmax: int # kv-cache capacity input_pos: int # number of prior tokens already in the cache (decode) denom: float = 16.0 # ramp divisor; small denom -> large logits (softmax stress) + kv_f16: bool = False # Single source of truth, mirrored by the C++ CONFIGS table in the native test. @@ -64,6 +65,10 @@ class SdpaConfig: # 2D-dispatch cap (>65535 wg): S=512 folds QK; S=2048 folds QK+softmax+AV (cap+1). SdpaConfig("llama1b_prefill_512", 32, 8, 64, 512, 512, 0), SdpaConfig("llama1b_prefill_2048", 32, 8, 64, 2048, 2048, 0), + # denom=10 intentionally makes K/V values lossy in fp16, so native + # execution exercises the real fp32->fp16->fp32 cache conversion path. + SdpaConfig("qwen3_prefill", 16, 8, 128, 128, 256, 0, 10.0, kv_f16=True), + SdpaConfig("qwen3_odd_boundary", 16, 8, 128, 17, 64, 31, 10.0, kv_f16=True), ] @@ -83,6 +88,7 @@ class ReplaySeq: d: int # head dim cmax: int # kv-cache capacity (>= sum(seq_lens)) seq_lens: tuple[int, ...] + kv_f16: bool = False # Mirror Vulkan sdpa_test.cpp:856/867/875 (3 param sets); cmax = sum rounded up. @@ -90,11 +96,18 @@ class ReplaySeq: ReplaySeq("small", 8, 4, 4, 16, (3, 1, 1, 5, 1, 1, 2)), ReplaySeq("small_d", 6, 2, 8, 16, (3, 1, 1, 5, 1, 1)), ReplaySeq("llama3", 24, 8, 128, 256, (111, 1, 1, 1, 57, 1, 1)), + ReplaySeq("qwen3_fd", 16, 8, 128, 64, (17, 1), kv_f16=True), ] +DYNAMIC_REPLAY_SEQS = [seq for seq in REPLAY_SEQS if not seq.kv_f16] -# (head_dim, num_heads, num_kv_heads) from sdpa_test.cpp:856/867/875 -- guards a -# transposition of the (hq, hkv, d) field order against the Vulkan source. -VULKAN_PARAMS = {"small": (4, 8, 4), "small_d": (8, 6, 2), "llama3": (128, 24, 8)} +# Guards transposition of the (hq, hkv, d) field order. The first three values +# mirror Vulkan sdpa_test.cpp:856/867/875; Qwen3 extends the same contract. +VULKAN_PARAMS = { + "small": (4, 8, 4), + "small_d": (8, 6, 2), + "llama3": (128, 24, 8), + "qwen3_fd": (128, 16, 8), +} class SdpaModule(torch.nn.Module): @@ -178,6 +191,12 @@ def _det_inputs(cfg: SdpaConfig): return q, k, v, k_cache, v_cache +def _round_kv_for_storage(cfg: SdpaConfig, *tensors: torch.Tensor): + if not cfg.kv_f16: + return tensors + return tuple(tensor.to(torch.float16).to(torch.float32) for tensor in tensors) + + def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """Reference attention output [1,S,Hq,D], computed in fp64 then cast to fp32. @@ -189,6 +208,7 @@ def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """ context_len = cfg.s + cfg.input_pos g = cfg.hq // cfg.hkv + k, v, k_cache, v_cache = _round_kv_for_storage(cfg, k, v, k_cache, v_cache) qd, kd, vd = q.double(), k.double(), v.double() kcd, vcd = k_cache.double(), v_cache.double() @@ -228,6 +248,50 @@ def _export_pte(cfg: SdpaConfig, q, k, v, kc, vc): class TestSdpa(unittest.TestCase): + def test_qwen3_fixture_contract(self) -> None: + configs = {cfg.name: cfg for cfg in CONFIGS} + expected_geometries = { + "qwen3_prefill": (16, 8, 128, 128, 256, 0), + "qwen3_odd_boundary": (16, 8, 128, 17, 64, 31), + } + for name, geometry in expected_geometries.items(): + with self.subTest(config=name): + self.assertIn(name, configs) + cfg = configs[name] + self.assertEqual( + (cfg.hq, cfg.hkv, cfg.d, cfg.s, cfg.cmax, cfg.input_pos), + geometry, + ) + self.assertTrue(cfg.kv_f16) + + replays = {seq.name: seq for seq in REPLAY_SEQS} + self.assertIn("qwen3_fd", replays) + qwen3_fd = replays["qwen3_fd"] + self.assertEqual((qwen3_fd.hq, qwen3_fd.hkv, qwen3_fd.d), (16, 8, 128)) + self.assertEqual(qwen3_fd.seq_lens, (17, 1)) + self.assertTrue(qwen3_fd.kv_f16) + + probe = torch.tensor([0.1], dtype=torch.float32) + (rounded,) = _round_kv_for_storage(configs["qwen3_prefill"], probe) + expected = probe.to(torch.float16).to(torch.float32) + torch.testing.assert_close(rounded, expected, atol=0.0, rtol=0.0) + self.assertFalse(torch.equal(rounded, probe)) + + boundary = configs["qwen3_odd_boundary"] + q, k, v, k_cache, v_cache = _det_inputs(boundary) + self.assertGreater(torch.count_nonzero(k_cache).item(), 0) + self.assertGreater(torch.count_nonzero(v_cache).item(), 0) + initialized = _golden(boundary, q, k, v, k_cache, v_cache) + cleared = _golden( + boundary, + q, + k, + v, + torch.zeros_like(k_cache), + torch.zeros_like(v_cache), + ) + self.assertFalse(torch.equal(initialized, cleared)) + def test_sdpa_export_delegates(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): @@ -248,7 +312,12 @@ def test_golden_matches_eager_op(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): q, k, v, kc, vc = _det_inputs(cfg) - eager = SdpaModule(cfg.input_pos)(q, k, v, kc.clone(), vc.clone()) + eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( + cfg, k, v, kc, vc + ) + eager = SdpaModule(cfg.input_pos)( + q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() + ) golden = _golden(cfg, q, k, v, kc, vc) torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) @@ -277,10 +346,16 @@ def test_replay_golden_matches_eager(self) -> None: s, seq.cmax, input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) golden = _golden(cfg, q, k, v, kc, vc) - eager = SdpaModule(input_pos)(q, k, v, kc.clone(), vc.clone()) + eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( + cfg, k, v, kc, vc + ) + eager = SdpaModule(input_pos)( + q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() + ) torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) kc[0, input_pos : input_pos + s] = k[0] vc[0, input_pos : input_pos + s] = v[0] @@ -302,6 +377,7 @@ def test_replay_export_delegates(self) -> None: s, seq.cmax, input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) et = _export_pte(cfg, q, k, v, kc, vc) @@ -347,7 +423,14 @@ def export_replay_sequences(out_dir: str) -> None: input_pos = 0 for t, s in enumerate(seq.seq_lens): cfg = SdpaConfig( - f"{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, s, seq.cmax, input_pos + f"{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + s, + seq.cmax, + input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) et = _export_pte(cfg, q, k, v, ref_kc, ref_vc) @@ -421,7 +504,7 @@ def export_dynamic_decode(out_dir: str) -> None: Mirrors the host accumulation the native test threads: at step t the golden attends over input_pos=t prior tokens plus the new token. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" et = _export_dyn_pte(seq, 1) pte_path = os.path.join(out_dir, f"sdpa_dyn_{seq.name}.pte") @@ -431,7 +514,14 @@ def export_dynamic_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" None: def test_dynamic_decode_golden_matches_eager(self) -> None: # The threaded-cache decode golden must equal the eager op step-by-step. - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: ref_kc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc) @@ -499,7 +596,7 @@ def export_incache_decode(out_dir: str) -> None: """One sdpa_incache_.pte (mutable-buffer KV cache) + per-step decode goldens. forward() feeds only q/k/v + input_pos; the cache persists in-graph. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" m = DecodeCacheModule(seq.hkv, seq.d, seq.cmax) q, k, v = _step_inputs(seq, 0, 1) @@ -517,7 +614,14 @@ def export_incache_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"incache_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"incache_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -309,13 +311,12 @@ bool sdpa_within_tol( const float* golden, int n, float* ma, - float* mr) { + float* mr, + bool kv_f16 = false) { float atol = 1e-4f, rtol = 1e-3f; - // f16 KV (runtime opt-in) reads K/V at reduced precision; loosen the tol on a - // shader-f16 device to cover that rounding. Harmless for f32 KV (looser - // gate). - const WebGPUContext* kv_ctx = get_default_webgpu_context(); - if (kv_ctx != nullptr && kv_ctx->shader_f16_supported) { + // Only fp16-KV cases receive the tolerance needed for storage rounding; + // device capability alone must not weaken unrelated fp32 tests. + if (kv_f16) { atol = 2e-3f; rtol = 1e-2f; } @@ -1250,6 +1251,7 @@ struct SdpaConfig { float denom; // ramp divisor (mirrors Python); small -> large logits bool required = false; // CI (SDPA dir set): absent .pte = FAIL, not skip bool expect_reject = false; // load MUST fail (e.g. D%4 guard), no golden + bool kv_f16 = false; }; const SdpaConfig kSdpaConfigs[] = { @@ -1292,6 +1294,28 @@ const SdpaConfig kSdpaConfigs[] = { 0, 16.0f, /*required=*/true}, + {"qwen3_prefill", + 16, + 8, + 128, + 128, + 256, + 0, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_odd_boundary", + 16, + 8, + 128, + 17, + 64, + 31, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, }; // Ramp denominator; mirror of test_sdpa.py::_RAMP_DENOM (keep in sync). @@ -1314,9 +1338,8 @@ float sdpa_ramp_t( return static_cast(((i + 31 * t) % mod) - off) / denom; } -// Multi-step replay sequences. Mirror the Python REPLAY_SEQS / Vulkan param -// sets (sdpa_test.cpp:856/867/875). Each seq_lens entry is one step replayed on -// a host-threaded KV cache (big=prefill, mid=multi-token, 1=decode). +// Multi-step replay sequences. The first three mirror Vulkan param sets; Qwen3 +// extends the same Python REPLAY_SEQS contract. struct SdpaSequence { const char* name; int hq; @@ -1324,18 +1347,90 @@ struct SdpaSequence { int d; int cmax; std::vector seq_lens; + bool kv_f16 = false; }; const SdpaSequence kSdpaSequences[] = { {"small", 8, 4, 4, 16, {3, 1, 1, 5, 1, 1, 2}}, {"small_d", 6, 2, 8, 16, {3, 1, 1, 5, 1, 1}}, {"llama3", 24, 8, 128, 256, {111, 1, 1, 1, 57, 1, 1}}, + {"qwen3_fd", 16, 8, 128, 64, {17, 1}, /*kv_f16=*/true}, }; +Error load_sdpa_forward(Module& module, bool kv_f16, int sdpa_query_tile = 0) { + if (!kv_f16 && sdpa_query_tile == 0) { + return module.load_forward(); + } + BackendOptions<2> options; + Error error = Error::Ok; + if (kv_f16) { + error = options.set_option("enable_f16_kv_cache", true); + if (error != Error::Ok) { + return error; + } + } + if (sdpa_query_tile != 0) { + error = options.set_option("sdpa_query_tile", sdpa_query_tile); + if (error != Error::Ok) { + return error; + } + } + LoadBackendOptionsMap option_map; + error = option_map.set_options("VulkanBackend", options.view()); + if (error != Error::Ok) { + return error; + } + return module.load_forward(nullptr, nullptr, &option_map); +} + +bool shader_f16_supported_on_test_device() { + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->shader_f16_supported; +} + +bool qwen3_q16_supported_on_test_device() { + constexpr uint32_t kQ16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); + const WebGPUContext* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQ16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +bool qwen3_q32_supported_on_test_device() { + constexpr uint32_t kQ32StorageBytes = 1024u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 256u * 2u * sizeof(float) + + 3u * 32u * sizeof(float); + const WebGPUContext* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupStorageSize >= kQ32StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +constexpr uint32_t kTestRouteMaterializedAttention = 1u << 2; +constexpr uint32_t kTestRouteFlashDecoding = 1u << 10; +constexpr uint32_t kTestRouteK16CausalBound = 1u << 11; +constexpr uint32_t kTestRouteQwen3Q16K16 = 1u << 13; +constexpr uint32_t kTestRouteQwen3Q32K16 = 1u << 14; +#endif // WGPU_BACKEND_ENABLE_PROFILING + void test_sdpa_config( const SdpaConfig& cfg, const std::string& model_path, - const std::string& golden_path) { + const std::string& golden_path, + int sdpa_query_tile = 0) { // Inputs reconstruct test_sdpa.py::_det_inputs bit-for-bit (/16 exact fp32). printf( "\n--- Test: sdpa_with_kv_cache (%s: Hq=%d,Hkv=%d,D=%d,S=%d,Cmax=%d,pos=%d) ---\n", @@ -1347,8 +1442,13 @@ void test_sdpa_config( cfg.cmax, cfg.input_pos); + if (cfg.kv_f16 && !shader_f16_supported_on_test_device()) { + printf("SKIP: %s requires shader-f16\n", cfg.name); + return; + } + Module module(model_path); - auto err = module.load_forward(); + auto err = load_sdpa_forward(module, cfg.kv_f16, sdpa_query_tile); if (cfg.expect_reject) { // D not a multiple of 4 must be rejected at load by the head_dim guard. ASSERT_NE(err, Error::Ok) @@ -1393,6 +1493,38 @@ void test_sdpa_config( {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() << ")"; + if (cfg.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // Exact Qwen3 geometry + fp16 KV selects the K16 streaming (causal-bound) + // route by default. The sdpa_query_tile RuntimeSpec only swaps the Q16/Q32 + // kernel variant; both map to the K16CausalBound bit. A non-Qwen3 fp16-KV + // shape falls back to the materialized path (or flash-decoding at S==1). + const bool qwen3_geometry = cfg.hq == 16 && cfg.hkv == 8 && cfg.d == 128; + const bool qwen3_streaming = + qwen3_geometry && cfg.s > 1 && qwen3_q16_supported_on_test_device(); + const uint32_t expected_route = qwen3_streaming + ? kTestRouteK16CausalBound + : (cfg.s == 1 ? kTestRouteFlashDecoding + : kTestRouteMaterializedAttention); + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route); + EXPECT_EQ(g_last_route_conflict_count, 0u); + const uint32_t qwen3_tile_routes = + g_last_route_mask & (kTestRouteQwen3Q16K16 | kTestRouteQwen3Q32K16); + if (qwen3_streaming) { + const uint32_t expected_tile_route = + sdpa_query_tile == 32 && qwen3_q32_supported_on_test_device() + ? kTestRouteQwen3Q32K16 + : kTestRouteQwen3Q16K16; + EXPECT_EQ(qwen3_tile_routes, expected_tile_route); + } else { + EXPECT_EQ(qwen3_tile_routes, 0u); + } +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outputs = result.get(); // Select the attention output [1,S,Hq,D] by shape; the op returns @@ -1424,8 +1556,8 @@ void test_sdpa_config( ASSERT_FALSE(golden.empty()) << "could not load golden " << golden_path; float max_abs_err = 0.0f, max_rel_err = 0.0f; - const bool pass = - sdpa_within_tol(out_data, golden.data(), on, &max_abs_err, &max_rel_err); + const bool pass = sdpa_within_tol( + out_data, golden.data(), on, &max_abs_err, &max_rel_err, cfg.kv_f16); printf( "Max abs error: %e Max rel error: %e (checked %d elements)\n", max_abs_err, @@ -1447,6 +1579,10 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { seq.d, seq.cmax, seq.seq_lens.size()); + if (seq.kv_f16 && !shader_f16_supported_on_test_device()) { + printf("SKIP: %s requires shader-f16\n", seq.name); + return; + } const int cn = seq.cmax * seq.hkv * seq.d; std::vector kc(cn, 0.0f), vc(cn, 0.0f); @@ -1460,7 +1596,7 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { std::to_string(t) + "_S" + std::to_string(s) + "_pos" + std::to_string(input_pos); Module module(base + ".pte"); - ASSERT_EQ(module.load_forward(), Error::Ok) + ASSERT_EQ(load_sdpa_forward(module, seq.kv_f16), Error::Ok) << "could not load " << base << ".pte"; const int qn = s * seq.hq * seq.d; @@ -1486,6 +1622,26 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward " << base << ".pte (error " << (int)result.error() << ")"; + if (seq.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // S==1 decode -> flash-decoding; a multi-token exact-Qwen3-geometry + // prefill -> the K16 streaming (causal-bound) route by default (no env); + // any other multi-token fp16-KV shape -> materialized. + const bool qwen3_geometry = seq.hq == 16 && seq.hkv == 8 && seq.d == 128; + const bool qwen3_streaming = + qwen3_geometry && qwen3_q16_supported_on_test_device(); + const uint32_t expected_route = s == 1 ? kTestRouteFlashDecoding + : qwen3_streaming ? kTestRouteK16CausalBound + : kTestRouteMaterializedAttention; + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route) + << seq.name << " step" << t; + EXPECT_EQ(g_last_route_conflict_count, 0u) << seq.name << " step" << t; +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outs = result.get(); // The op returns [k_cache, v_cache, attn_output]: attn has a unique numel; @@ -1533,7 +1689,8 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { ASSERT_FALSE(golden.empty()) << "could not load " << base << ".golden.bin"; const float* ad = outs[attn_idx].toTensor().const_data_ptr(); float ma = 0.0f, mr = 0.0f; - const bool step_ok = sdpa_within_tol(ad, golden.data(), qn, &ma, &mr); + const bool step_ok = + sdpa_within_tol(ad, golden.data(), qn, &ma, &mr, seq.kv_f16); printf( " step%zu (S=%d pos=%d ctx=%d): max abs %e rel %e\n", t, @@ -1906,12 +2063,160 @@ void test_symint_input_narrowing() { vk::FinishVkGraphBuffer(fbb, root); WebGPUGraph graph; - ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, nullptr)); + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr)); ASSERT_EQ(graph.symint_sources().size(), 1u); const auto& source = graph.symint_sources().front(); exercise_symint_host_inputs(graph, source.symint_id, source.input_tensor_id); } +void write_u16_le(std::vector& data, size_t offset, uint16_t value) { + data.at(offset) = static_cast(value); + data.at(offset + 1) = static_cast(value >> 8); +} + +void write_u32_le(std::vector& data, size_t offset, uint32_t value) { + for (size_t i = 0; i < sizeof(value); i++) { + data.at(offset + i) = static_cast(value >> (8 * i)); + } +} + +void write_u64_le(std::vector& data, size_t offset, uint64_t value) { + for (size_t i = 0; i < sizeof(value); i++) { + data.at(offset + i) = static_cast(value >> (8 * i)); + } +} + +std::vector make_delegate_header_test_blob() { + std::vector blob(44, 0); + std::memcpy(blob.data() + 4, "VH00", 4); + write_u16_le(blob, 8, 30); + write_u32_le(blob, 10, 32); + write_u32_le(blob, 14, 8); + write_u32_le(blob, 18, 40); + write_u64_le(blob, 22, 4); + return blob; +} + +void finish_inline_constant_graph( + ::flatbuffers::FlatBufferBuilder& fbb, + bool mark_as_kv_cache, + const std::vector& dims, + uint64_t inline_offset = 0) { + namespace vk = vkgraph; + std::vector<::flatbuffers::Offset> values; + const int tensor_count = mark_as_kv_cache ? 5 : 1; + for (int i = 0; i < tensor_count; i++) { + const bool is_cache = mark_as_kv_cache && i >= 3; + const bool is_constant = !mark_as_kv_cache || is_cache; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + is_constant ? (is_cache ? i - 3 : 0) : -1, + is_constant ? -1 : i) + .Union())); + } + + std::vector<::flatbuffers::Offset> chain; + if (mark_as_kv_cache) { + const std::vector args = {0, 1, 2, 3, 4}; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "sdpa_with_kv_cache.default", &args)); + } + std::vector<::flatbuffers::Offset> constants; + constants.push_back( + vk::CreateVkBytesDirect(fbb, inline_offset, sizeof(float))); + if (mark_as_kv_cache) { + constants.push_back(vk::CreateVkBytesDirect(fbb, 0, sizeof(float))); + } + const std::vector output_ids = {0}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, nullptr, &output_ids, &constants); + vk::FinishVkGraphBuffer(fbb, root); +} + +TEST(WebGPUNative, DelegateHeaderRejectsTruncatedRanges) { + const auto blob = make_delegate_header_test_blob(); + EXPECT_TRUE(WebGPUDelegateHeader::parse(blob.data(), blob.size()).ok()); + EXPECT_FALSE(WebGPUDelegateHeader::parse(blob.data(), 29).ok()); + EXPECT_FALSE(WebGPUDelegateHeader::parse(blob.data(), blob.size() - 1).ok()); +} + +TEST(WebGPUNative, InlineConstantExtentIsBounded) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, {1u}); + const std::array data = {0, 0, 0, 0}; + + WebGPUGraph exact_graph; + EXPECT_NO_THROW(exact_graph.build( + fbb.GetBufferPointer(), data.data(), data.size(), nullptr)); + + WebGPUGraph short_graph; + EXPECT_THROW( + short_graph.build( + fbb.GetBufferPointer(), data.data(), data.size() - 1, nullptr), + std::runtime_error); +} + +TEST(WebGPUNative, ZeroByteInlineConstantOffsetIsBounded) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, {0u}, 1); + const std::array data = {0}; + + WebGPUGraph graph; + EXPECT_THROW( + graph.build(fbb.GetBufferPointer(), data.data(), 0, nullptr), + std::runtime_error); +} + +TEST(WebGPUNative, F16KvInlineConstantExtentIsBounded) { + const auto* context = get_default_webgpu_context(); + if (context == nullptr || !context->shader_f16_supported) { + GTEST_SKIP() << "shader-f16 unavailable"; + } + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, true, {1u}); + const std::array data = {0, 0, 0, 0}; + + WebGPUGraph graph; + WebGPUGraphConfig config; + config.f16_kv_cache = true; + try { + graph.build( + fbb.GetBufferPointer(), data.data(), data.size() - 1, nullptr, config); + FAIL() << "undersized inline fp16 KV constant was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "WebGPU f16 KV: inline cache constant exceeds constant data"); + } +} + +void expect_tensor_extent_error( + const std::vector& dims, + const char* expected_error) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, dims); + const std::array data = {0, 0, 0, 0}; + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), data.data(), data.size(), nullptr); + ADD_FAILURE() << "overflowing tensor extent was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), expected_error); + } +} + +TEST(WebGPUNative, TensorExtentOverflowIsRejected) { + expect_tensor_extent_error( + {UINT32_MAX, UINT32_MAX, 2u}, "WebGPU: tensor element count overflows"); + expect_tensor_extent_error( + {UINT32_MAX, UINT32_MAX}, "WebGPU: tensor byte size overflows"); +} + struct DelegateBlobView { size_t base_offset; WebGPUDelegateHeader header; @@ -1932,7 +2237,7 @@ std::optional find_delegate_blob( if (std::memcmp(base + kMagicOffset, kMagic, sizeof(kMagic)) != 0) { continue; } - auto header = WebGPUDelegateHeader::parse(base); + auto header = WebGPUDelegateHeader::parse(base, blob.size() - base_offset); if (!header.ok()) { continue; } @@ -1950,6 +2255,43 @@ std::optional find_delegate_blob( return std::nullopt; } +TEST(WebGPUNative, StructurallyInvalidVkGraphIsRejectedAtLoad) { + if (g_symint_blob.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SYMINT_BLOB not set"; + } + FILE* input = std::fopen(g_symint_blob.c_str(), "rb"); + ASSERT_NE(input, nullptr); + std::fseek(input, 0, SEEK_END); + const long file_size = std::ftell(input); + std::fseek(input, 0, SEEK_SET); + ASSERT_GT(file_size, 0); + std::vector blob(static_cast(file_size)); + ASSERT_EQ(std::fread(blob.data(), 1, blob.size(), input), blob.size()); + std::fclose(input); + + const auto delegate = find_delegate_blob(blob); + ASSERT_TRUE(delegate.has_value()); + ASSERT_GE(delegate->header.flatbuffer_size, sizeof(uint32_t)); + const size_t root_offset = + delegate->base_offset + delegate->header.flatbuffer_offset; + std::fill_n(blob.begin() + root_offset, sizeof(uint32_t), UINT8_MAX); + + const std::string malformed_path = "/tmp/webgpu_invalid_vkgraph_" + + std::to_string(reinterpret_cast(blob.data())) + ".pte"; + FILE* output = std::fopen(malformed_path.c_str(), "wb"); + ASSERT_NE(output, nullptr); + ASSERT_EQ(std::fwrite(blob.data(), 1, blob.size(), output), blob.size()); + std::fclose(output); + + Error load_result = Error::Ok; + { + Module module(malformed_path); + load_result = module.load_forward(); + } + EXPECT_NE(load_result, Error::Ok); + EXPECT_EQ(std::remove(malformed_path.c_str()), 0); +} + // S1 SymInt round-trip: confirm a dynamic input_pos stays live. void test_symint_roundtrip(const std::string& blob_path) { printf("\n--- Test: symint round-trip (%s) ---\n", blob_path.c_str()); @@ -1972,6 +2314,7 @@ void test_symint_roundtrip(const std::string& blob_path) { graph.build( base + delegate->header.flatbuffer_offset, base + delegate->header.bytes_offset, + delegate->header.bytes_size, nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); @@ -2029,6 +2372,7 @@ void test_resize_hook(const std::string& blob_path) { graph.build( base + delegate->header.flatbuffer_offset, base + delegate->header.bytes_offset, + delegate->header.bytes_size, nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); @@ -2178,7 +2522,7 @@ static bool test_slice_double_start_case(double start_d, int out_len) { WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("FAIL: graph build threw: %s\n", e.what()); return false; @@ -2191,8 +2535,8 @@ static bool test_slice_double_start_case(double start_d, int out_len) { std::vector inputs(1); inputs[0] = {in.data(), in.size() * sizeof(float), false}; std::vector out(out_len, -1.0f); - std::vector> outputs(1); - outputs[0] = {out.data(), out.size() * sizeof(float)}; + std::vector outputs(1); + outputs[0] = {out.data(), out.size() * sizeof(float), true}; try { graph.copy_inputs(inputs); const WebGPUExecutionPlan plan = graph.make_execution_plan({}); @@ -2276,7 +2620,7 @@ static bool test_slice_double_start_rejects(double bad_start) { WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("PASS: rejected as expected: %s\n", e.what()); return true; @@ -2374,7 +2718,7 @@ static bool test_select_double_scalar_case( WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("FAIL: graph build threw: %s\n", e.what()); return false; @@ -2384,8 +2728,8 @@ static bool test_select_double_scalar_case( std::vector inputs = { {in.data(), in.size() * sizeof(float), false}}; std::vector out(expected.size(), -1.0f); - std::vector> outputs = { - {out.data(), out.size() * sizeof(float)}}; + std::vector outputs = { + {out.data(), out.size() * sizeof(float), true}}; try { graph.copy_inputs(inputs); const WebGPUExecutionPlan plan = graph.make_execution_plan({}); @@ -2414,7 +2758,7 @@ static bool test_select_scalar_build_error( WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { const std::string error = e.what(); if (error.find(expected_error) != std::string::npos) { @@ -2531,7 +2875,7 @@ void expect_rope_hf_resize_numel_overflow(uint32_t q_heads, uint32_t k_heads) { vk::FinishVkGraphBuffer(fbb, root); WebGPUGraph graph; - ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, nullptr)); + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr)); ASSERT_EQ(graph.num_dispatches(), 2u); const uint32_t q_x = graph.dispatch_at(0).workgroup_count_x; const uint32_t q_y = graph.dispatch_at(0).workgroup_count_y; @@ -2812,6 +3156,95 @@ TEST(WebGPUNative, PrepackTied) { } // SDPA sweep: configs self-discover sdpa_.pte; required=FAIL else skip. +TEST(WebGPUNative, Qwen3SdpaFixtureContract) { + const auto find_config = [](const char* name) { + return std::find_if( + std::begin(kSdpaConfigs), + std::end(kSdpaConfigs), + [name](const SdpaConfig& cfg) { + return std::strcmp(cfg.name, name) == 0; + }); + }; + const auto prefill = find_config("qwen3_prefill"); + const auto boundary = find_config("qwen3_odd_boundary"); + ASSERT_NE(prefill, std::end(kSdpaConfigs)); + ASSERT_NE(boundary, std::end(kSdpaConfigs)); + EXPECT_EQ( + std::vector( + {prefill->hq, + prefill->hkv, + prefill->d, + prefill->s, + prefill->cmax, + prefill->input_pos}), + std::vector({16, 8, 128, 128, 256, 0})); + EXPECT_EQ( + std::vector( + {boundary->hq, + boundary->hkv, + boundary->d, + boundary->s, + boundary->cmax, + boundary->input_pos}), + std::vector({16, 8, 128, 17, 64, 31})); + EXPECT_TRUE(prefill->kv_f16 && boundary->kv_f16); + + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + EXPECT_EQ( + std::vector({replay->hq, replay->hkv, replay->d, replay->cmax}), + std::vector({16, 8, 128, 64})); + EXPECT_EQ(replay->seq_lens, std::vector({17, 1})); + EXPECT_TRUE(replay->kv_f16); +} + +TEST(WebGPUNative, Qwen3SdpaRoutes) { + if (g_sdpa_dir.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SDPA_DIR not set"; + } + if (!qwen3_q16_supported_on_test_device()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + // Default route: exact-Qwen3-geometry fp16-KV configs select the Q16 K16 + // streaming (causal-bound) route by geometry (no runtime config needed) -- + // the per-config assertions live in test_sdpa_config / test_sdpa_replay. + for (const auto& cfg : kSdpaConfigs) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + continue; + } + const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; + test_sdpa_config(cfg, base + ".pte", base + ".golden.bin"); + } + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + test_sdpa_replay(*replay, g_sdpa_dir); + + // Run Q32 over both an aligned prefill and the S=17/nonzero-position case so + // the partial final workgroup's row mask is covered. Unsupported Q32 devices + // intentionally fall back to the already-qualified Q16 route. + for (const auto& cfg : kSdpaConfigs) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + continue; + } + const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; + test_sdpa_config( + cfg, + base + ".pte", + base + ".golden.bin", + /*sdpa_query_tile=*/32); + } +} + TEST(WebGPUNative, SdpaSweep) { const std::string& dir = g_sdpa_dir; bool ran = false; diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 0cd110d0ff7..32d59501a3a 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -158,6 +158,22 @@ def test_render_header_embeds_sha256(self) -> None: self.assertEqual(g.embedded_sha256(h), want) self.assertEqual(g.wgsl_sha256(wgsl), want) + def test_render_header_long_name_is_clang_format_stable(self) -> None: + stem = "streaming_attention_qwen3_q32_k16_causal_bound" + wgsl = "@compute @workgroup_size(32, 8, 1)\nfn main(){}\n" + h = g.render_header(Path(f"runtime/ops/sdpa/{stem}.wgsl"), wgsl) + + self.assertIn( + f"// @generated from {stem}.wgsl\n// DO NOT EDIT.", + h, + ) + self.assertIn( + "inline constexpr uint32_t\n" + " kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32;", + h, + ) + self.assertEqual(g.embedded_sha256(h), g.wgsl_sha256(wgsl)) + def test_embedded_sha256_missing_returns_empty(self) -> None: self.assertEqual(g.embedded_sha256("no sha line here\n"), "") @@ -200,6 +216,33 @@ def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: ] self.assertEqual(indices, list(range(8))) + def test_qwen3_runtime_eligibility_is_exact(self) -> None: + sdpa = (g.BACKEND_ROOT / "runtime/ops/sdpa/Sdpa.cpp").read_text() + self.assertIn("q/k/v/output must be fp32", sdpa) + self.assertIn("cache dtype does not match the selected storage mode", sdpa) + self.assertIn("scale == qwen3_expected_scale", sdpa) + self.assertNotIn("std::fabs(scale - qwen3_expected_scale)", sdpa) + + def test_fp16_kv_graph_guards_transfer_and_topology(self) -> None: + graph = (g.BACKEND_ROOT / "runtime/WebGPUGraph.cpp").read_text() + self.assertIn("serialized cache tensor must be fp32", graph) + self.assertIn("consumed through a ValueList", graph) + self.assertIn("preserve it while changing storage", graph) + + copy_inputs = graph.index("void WebGPUGraph::copy_inputs") + input_guard = graph.index( + "fp16 device input requires an fp32 host tensor", copy_inputs + ) + fast_path = graph.index("// Fast path", copy_inputs) + self.assertLess(input_guard, fast_path) + + copy_outputs = graph.index("void WebGPUGraph::copy_outputs") + output_guard = graph.index( + "fp16 device output requires an fp32 host tensor", copy_outputs + ) + map_request = graph.index("wgpuBufferMapAsync", copy_outputs) + self.assertLess(output_guard, map_request) + def test_parse_workgroup_allows_space(self) -> None: # @workgroup_size (64) — the spec-legal spaced form must still parse. self.assertEqual( From f98d92e94e319144337bebe08b96019b4bf24026 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:35 -0700 Subject: [PATCH 07/12] [ExecuTorch][WebGPU] Make WGSL generation rollback-safe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21448 **Make WGSL publication rollback-safe without changing generated runtime output** Generated headers and the registry were previously written incrementally, so later validation, interruption, or write failure could leave a partial generated tree. This change stages and validates the complete output set before sequential publication, rejects path/name/C++ symbol collisions, and restores replaced outputs on failure or interruption. The family model follows Vulkan's YAML/template generation. The WebGPU publisher adds staging, backup, rollback, and temporary-file cleanup; it does not claim filesystem-atomic multi-file replacement. Key changes: - `gen_wgsl_headers.py` — stage, validate, publish, restore on failure, and clean temporary outputs. - `test_wgsl_codegen.py` — cover collisions, orphan repair, interruption, rollback, and idempotence. All generated shader bytes and runtime registry entries remain unchanged. Co-authored-with: Claude Code. ghstack-source-id: 411961465 @exported-using-ghexport Differential Revision: [D113979658](https://our.internmc.facebook.com/intern/diff/D113979658/) --- backends/webgpu/scripts/gen_wgsl_headers.py | 270 +++++++++++--- backends/webgpu/test/test_wgsl_codegen.py | 368 ++++++++++++++++++++ 2 files changed, 582 insertions(+), 56 deletions(-) diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index dd8cc4cc635..1e35b205888 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -26,11 +26,14 @@ import copy import hashlib import io +import os import re +import stat import sys +import tempfile from itertools import product from pathlib import Path -from typing import Any, Dict, List, NamedTuple, Optional, Set +from typing import Any, Dict, List, NamedTuple, Optional, Set, Tuple import yaml from yaml.constructor import ConstructorError @@ -523,23 +526,51 @@ def registry_path() -> Path: return BACKEND_ROOT / "runtime/WebGPUShaderRegistry.cpp" -def registry_entries() -> List[RegistryEntry]: - """Return one registry entry for every concrete generated shader.""" - entries = [] +def _registry_entry(header: Path) -> RegistryEntry: + suffix = "_wgsl.h" + if not header.name.endswith(suffix): + raise ValueError(f"unexpected generated header name: {header.name}") + name = header.name[: -len(suffix)] + return RegistryEntry( + name=name, + include=header.relative_to(BACKEND_ROOT).as_posix(), + symbol=symbol_base(name), + ) + + +def _collect_header_outputs() -> Tuple[Dict[Path, str], List[RegistryEntry]]: + """Render every concrete header once and reject global collisions.""" + outputs: Dict[Path, str] = {} + entries: List[RegistryEntry] = [] + registry_names: Set[str] = set() + registry_symbols: Set[str] = set() for wgsl in discover(): - for header, _ in headers_for_shader(wgsl): - suffix = "_wgsl.h" - if not header.name.endswith(suffix): - raise ValueError(f"unexpected generated header name: {header.name}") - name = header.name[: -len(suffix)] - entries.append( - RegistryEntry( - name=name, - include=header.relative_to(BACKEND_ROOT).as_posix(), - symbol=symbol_base(name), + try: + rendered_headers = list(headers_for_shader(wgsl)) + except Exception as error: + raise ValueError(f"{wgsl.relative_to(BACKEND_ROOT)}: {error}") from error + for header, rendered in rendered_headers: + if header in outputs: + raise ValueError( + "duplicate generated header path: " + f"{header.relative_to(BACKEND_ROOT)}" ) - ) - return sorted(entries) + entry = _registry_entry(header) + if entry.name in registry_names: + raise ValueError(f"duplicate shader registry name: {entry.name}") + if entry.symbol in registry_symbols: + raise ValueError(f"duplicate shader registry symbol: {entry.symbol}") + outputs[header] = rendered + entries.append(entry) + registry_names.add(entry.name) + registry_symbols.add(entry.symbol) + return outputs, sorted(entries) + + +def registry_entries() -> List[RegistryEntry]: + """Return one registry entry for every concrete generated shader.""" + _, entries = _collect_header_outputs() + return entries def render_registry(entries: List[RegistryEntry]) -> str: @@ -630,7 +661,143 @@ def headers_for_shader(wgsl): yield header, render_header(stem, text, stem) -def _report_drift(missing, stale) -> None: +def collect_outputs() -> Tuple[Dict[Path, bytes], List[Path]]: + """Render the complete output tree and report unexpected old headers.""" + header_outputs, entries = _collect_header_outputs() + outputs = { + path: rendered.encode("utf-8") for path, rendered in header_outputs.items() + } + registry = registry_path() + if registry in outputs: + raise ValueError(f"duplicate generated output path: {registry}") + outputs[registry] = render_registry(entries).encode("utf-8") + + expected_headers = set(header_outputs) + actual_headers = set((BACKEND_ROOT / "runtime/ops").glob("**/*_wgsl.h")) + return outputs, sorted(actual_headers - expected_headers) + + +class _OriginalOutput(NamedTuple): + existed: bool + contents: bytes + mode: int + + +def _stage_bytes(destination: Path, contents: bytes, mode: int) -> Path: + """Write one same-directory candidate without changing its destination.""" + fd, name = tempfile.mkstemp( + prefix=f".{destination.name}.wgsl-gen-", + suffix=".tmp", + dir=destination.parent, + ) + temporary = Path(name) + try: + with os.fdopen(fd, "wb") as output: + output.write(contents) + temporary.chmod(mode) + except BaseException: + try: + temporary.unlink(missing_ok=True) + except OSError: + pass + raise + return temporary + + +def _cleanup_temporaries(temporaries) -> List[str]: + errors = [] + for temporary in temporaries: + try: + temporary.unlink(missing_ok=True) + except OSError as error: + errors.append(f"cannot remove temporary {temporary}: {error}") + return errors + + +def _stage_outputs( + outputs: Dict[Path, bytes], changed: List[Path] +) -> Tuple[Dict[Path, _OriginalOutput], Dict[Path, Path], List[str]]: + originals: Dict[Path, _OriginalOutput] = {} + staged: Dict[Path, Path] = {} + try: + for destination in sorted(changed): + if destination.exists(): + original = _OriginalOutput( + existed=True, + contents=destination.read_bytes(), + mode=stat.S_IMODE(destination.stat().st_mode), + ) + else: + original = _OriginalOutput(False, b"", 0o644) + originals[destination] = original + staged[destination] = _stage_bytes( + destination, outputs[destination], original.mode + ) + except BaseException as error: + cleanup_errors = _cleanup_temporaries(staged.values()) + if isinstance(error, OSError): + errors = [f"cannot stage generated output: {error}"] + cleanup_errors + return originals, staged, errors + raise + return originals, staged, [] + + +def _rollback_outputs( + originals: Dict[Path, _OriginalOutput], + replaced: List[Path], + staged: Dict[Path, Path], +) -> List[str]: + errors = [] + for destination in reversed(replaced): + original = originals[destination] + restore_temporary: Optional[Path] = None + try: + if original.existed: + restore_temporary = _stage_bytes( + destination, original.contents, original.mode + ) + os.replace(restore_temporary, destination) + else: + destination.unlink(missing_ok=True) + except OSError as error: + errors.append(f"cannot roll back {destination}: {error}") + finally: + if restore_temporary is not None: + errors.extend(_cleanup_temporaries([restore_temporary])) + errors.extend(_cleanup_temporaries(staged.values())) + return errors + + +def _publish_outputs(outputs: Dict[Path, bytes], changed: List[Path]) -> List[str]: + """Stage and publish changed outputs, rolling back reported failures.""" + originals, staged, stage_errors = _stage_outputs(outputs, changed) + if stage_errors: + return stage_errors + + replaced: List[Path] = [] + try: + for destination in sorted(changed): + try: + os.replace(staged[destination], destination) + except OSError: + raise + except BaseException: + replaced.append(destination) + raise + else: + replaced.append(destination) + except OSError as commit_error: + return [f"cannot publish generated output: {commit_error}"] + _rollback_outputs( + originals, replaced, staged + ) + except BaseException: + _rollback_outputs(originals, replaced, staged) + raise + + return _cleanup_temporaries(staged.values()) + + +def _report_drift(missing, stale, orphans) -> None: """Print the --check report for missing/stale committed headers.""" if missing: print("Missing embedded WGSL headers (run scripts/gen_wgsl_headers.py):") @@ -640,16 +807,10 @@ def _report_drift(missing, stale) -> None: print("Stale embedded WGSL headers (run scripts/gen_wgsl_headers.py):") for h in stale: print(f" {h.relative_to(BACKEND_ROOT)}") - - -def _sync_generated_output(output, want, check, missing, stale) -> None: - """Write one generated file, or record its --check drift.""" - if output.exists() and output.read_text() == want: - return - if check: - (missing if not output.exists() else stale).append(output) - else: - output.write_text(want) + if orphans: + print("Orphan embedded WGSL headers (remove or restore their sources):") + for h in orphans: + print(f" {h.relative_to(BACKEND_ROOT)}") def main(argv=None) -> int: @@ -661,38 +822,35 @@ def main(argv=None) -> int: ) args = parser.parse_args(argv) - stale = [] - missing = [] - errors = [] - for wgsl in discover(): - try: - rendered = list(headers_for_shader(wgsl)) - # A malformed spec raises yaml.YAMLError (incl. UniqueKeyLoader's - # ConstructorError) / ValueError / KeyError from parse_template_spec, and - # a malformed template raises AssertionError from preprocess; catch them - # all so a bad shader is a clean --check report, not a traceback. - except (ValueError, KeyError, AssertionError, yaml.YAMLError) as e: - errors.append(f"{wgsl.relative_to(BACKEND_ROOT)}: {e}") - continue - for header, want in rendered: - # Full-content compare (not just the sha) catches generator-logic drift too. - _sync_generated_output(header, want, args.check, missing, stale) + try: + outputs, orphans = collect_outputs() + missing = [] + stale = [] + for output, want in sorted(outputs.items()): + if not output.exists(): + missing.append(output) + elif output.read_bytes() != want: + stale.append(output) + except Exception as error: + print("Cannot generate WGSL outputs:") + print(f" {error}") + return 1 - if not errors: - try: - registry = render_registry(registry_entries()) - output = registry_path() - _sync_generated_output(output, registry, args.check, missing, stale) - except ValueError as e: - errors.append(f"shader registry: {e}") + if orphans: + _report_drift([], [], orphans) + return 1 + + if args.check: + if stale or missing: + _report_drift(missing, stale, []) + return 1 + return 0 + errors = _publish_outputs(outputs, missing + stale) if errors: - print("Cannot generate header (malformed shader):") - for e in errors: - print(f" {e}") - return 1 - if args.check and (stale or missing): - _report_drift(missing, stale) + print("Cannot publish WGSL outputs:") + for error in errors: + print(f" {error}") return 1 return 0 diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 32d59501a3a..831d2ce0747 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -13,10 +13,13 @@ import hashlib import importlib.util import io +import os import re +import stat import tempfile import unittest from pathlib import Path +from unittest import mock import yaml @@ -194,6 +197,25 @@ def test_committed_headers_match_generator(self) -> None: got, want, f"{header.name} stale; run scripts/gen_wgsl_headers.py" ) + def test_generated_output_manifest_digest(self) -> None: + outputs = sorted( + [ + *(g.BACKEND_ROOT / "runtime/ops").glob("**/*_wgsl.h"), + g.registry_path(), + ] + ) + digest = hashlib.sha256() + for output in outputs: + digest.update(output.relative_to(g.BACKEND_ROOT).as_posix().encode()) + digest.update(b"\0") + digest.update(output.read_bytes()) + digest.update(b"\0") + self.assertEqual(len(outputs), 134) + self.assertEqual( + digest.hexdigest(), + "19a0baf9345bec02fe2091a0e6320b81d966724bedc33feb0779c6a824925972", + ) + def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: shader = ( g.BACKEND_ROOT / "runtime" / "ops" / "rope" / "rotary_embedding_hf.wgsl" @@ -334,6 +356,352 @@ def test_render_header_3d_emits_xyz(self) -> None: self.assertIn("inline constexpr uint32_t kFooWorkgroupSizeZ = 2;", h) +class WgslGenerationTransactionTest(unittest.TestCase): + _VALID_SHADER = "@compute @workgroup_size(1)\nfn main() {}\n" + + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + (self.root / "runtime/ops").mkdir(parents=True) + self._original_root = g.BACKEND_ROOT + g.BACKEND_ROOT = self.root + + def tearDown(self) -> None: + g.BACKEND_ROOT = self._original_root + self._tmp.cleanup() + + def _write_shader( + self, directory: str, stem: str, text: str = _VALID_SHADER + ) -> Path: + op_dir = self.root / "runtime/ops" / directory + op_dir.mkdir(parents=True, exist_ok=True) + shader = op_dir / f"{stem}.wgsl" + shader.write_text(text) + return shader + + def _write_template( + self, directory: str, stem: str, text: str, names: list[str] + ) -> Path: + shader = self._write_shader(directory, stem, text) + spec = { + stem: { + "parameter_names_with_default_values": {}, + "shader_variants": [{"NAME": name} for name in names], + } + } + shader.with_suffix(".yaml").write_text(yaml.safe_dump(spec)) + return shader + + def _snapshot(self): + return { + path.relative_to(self.root).as_posix(): ( + path.read_bytes(), + stat.S_IMODE(path.stat().st_mode), + ) + for path in sorted(self.root.rglob("*")) + if path.is_file() + } + + def _run(self, *args: str): + output = io.StringIO() + with contextlib.redirect_stdout(output): + result = g.main(list(args)) + return result, output.getvalue() + + def _assert_no_temps(self) -> None: + self.assertEqual(list(self.root.rglob("*.tmp")), []) + + @staticmethod + def _fail_nth(real_fn, n: int): + calls = 0 + + def wrapped(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == n: + raise OSError(f"injected failure on call {n}") + return real_fn(*args, **kwargs) + + return wrapped + + def test_late_malformed_shader_leaves_tree_unchanged(self) -> None: + good = self._write_shader("a", "good") + good.with_name("good_wgsl.h").write_text("stale\n") + self._write_shader("z", "bad", "${MISSING\n") + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_duplicate_registry_name_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "shared") + self._write_shader("b", "shared") + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_duplicate_registry_symbol_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "foo_bar") + self._write_shader("b", "foo__bar") + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_duplicate_output_path_leaves_tree_unchanged(self) -> None: + self._write_template("op", "op", self._VALID_SHADER, ["duplicate", "duplicate"]) + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_second_stage_failure_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "first") + self._write_shader("b", "second") + before = self._snapshot() + real_mkstemp = tempfile.mkstemp + + with mock.patch( + "tempfile.mkstemp", side_effect=self._fail_nth(real_mkstemp, 2) + ): + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_staging_interrupt_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "first") + self._write_shader("b", "second") + before = self._snapshot() + real_chmod = Path.chmod + + def interrupt_second(path, mode, **kwargs): + interrupt_second.calls += 1 + if interrupt_second.calls == 2: + raise KeyboardInterrupt("injected staging interruption") + return real_chmod(path, mode, **kwargs) + + interrupt_second.calls = 0 + with mock.patch.object( + Path, "chmod", autospec=True, side_effect=interrupt_second + ): + with self.assertRaises(KeyboardInterrupt): + self._run() + + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_replace_failure_restores_existing_destination(self) -> None: + self._write_shader("op", "op") + registry = g.registry_path() + registry.write_text("old registry\n") + registry.chmod(0o600) + before = self._snapshot() + real_replace = os.replace + + with mock.patch("os.replace", side_effect=self._fail_nth(real_replace, 2)): + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_replace_failure_removes_new_destination(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + real_replace = os.replace + + with mock.patch("os.replace", side_effect=self._fail_nth(real_replace, 2)): + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_multiple_rollback_errors_do_not_stop_later_restores(self) -> None: + headers = [] + for directory in ("a", "b", "c"): + shader = self._write_shader(directory, directory) + header = shader.with_name(f"{directory}_wgsl.h") + header.write_text(f"old {directory}\n") + headers.append(header) + registry = g.registry_path() + registry.write_text("old registry\n") + real_replace = os.replace + calls = [] + + def fail_commit_and_two_rollbacks(source, destination): + calls.append(Path(destination)) + if len(calls) in (4, 5, 6): + raise OSError(f"injected failure on replace {len(calls)}") + return real_replace(source, destination) + + with mock.patch("os.replace", side_effect=fail_commit_and_two_rollbacks): + result, output = self._run() + + self.assertEqual(result, 1) + self.assertEqual(len(calls), 7) + self.assertEqual(calls[-3:], [headers[1], headers[0], registry]) + self.assertIn(f"cannot roll back {headers[1]}", output) + self.assertIn(f"cannot roll back {headers[0]}", output) + self.assertEqual(registry.read_text(), "old registry\n") + self.assertNotEqual(headers[0].read_text(), "old a\n") + self.assertNotEqual(headers[1].read_text(), "old b\n") + self.assertEqual(headers[2].read_text(), "old c\n") + self._assert_no_temps() + + def test_success_preserves_existing_mode_and_creates_0644(self) -> None: + shader = self._write_shader("op", "op") + registry = g.registry_path() + registry.write_text("old registry\n") + registry.chmod(0o600) + + result, _ = self._run() + + self.assertEqual(result, 0) + self.assertEqual(stat.S_IMODE(registry.stat().st_mode), 0o600) + self.assertEqual( + stat.S_IMODE(shader.with_name("op_wgsl.h").stat().st_mode), 0o644 + ) + + def test_orphans_are_sorted_reported_and_never_deleted(self) -> None: + self._write_shader("new", "new") + orphan_z = self.root / "runtime/ops/z/old_z_wgsl.h" + orphan_a = self.root / "runtime/ops/a/old_a_wgsl.h" + orphan_z.parent.mkdir(parents=True) + orphan_a.parent.mkdir(parents=True) + orphan_z.write_text("// @generated\n") + orphan_a.write_text("// @generated\n") + before = self._snapshot() + + check_result, check_output = self._run("--check") + normal_result, normal_output = self._run() + + self.assertEqual(check_result, 1) + self.assertEqual(normal_result, 1) + for output in (check_output, normal_output): + self.assertIn("Orphan", output) + self.assertLess(output.index("old_a_wgsl.h"), output.index("old_z_wgsl.h")) + self.assertEqual(self._snapshot(), before) + + def test_check_fails_read_only_when_outputs_are_only_missing(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + + result, output = self._run("--check") + + self.assertEqual(result, 1) + self.assertIn("Missing embedded WGSL headers", output) + self.assertEqual(self._snapshot(), before) + + def test_check_catches_template_syntax_error_without_writing(self) -> None: + self._write_template( + "a", "syntax", "$if :\n " + self._VALID_SHADER, ["syntax"] + ) + before = self._snapshot() + + result, output = self._run("--check") + + self.assertEqual(result, 1) + self.assertIn("runtime/ops/a/syntax.wgsl", output) + self.assertEqual(self._snapshot(), before) + + def test_check_catches_template_name_error_without_writing(self) -> None: + self._write_template( + "op", "name", "$if MISSING:\n " + self._VALID_SHADER, ["name"] + ) + before = self._snapshot() + + result, output = self._run("--check") + + self.assertEqual(result, 1) + self.assertIn("runtime/ops/op/name.wgsl", output) + self.assertEqual(self._snapshot(), before) + + def test_interrupted_commit_is_detected_and_repaired(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + real_replace = os.replace + + def interrupt_second(source, destination): + interrupt_second.calls += 1 + if interrupt_second.calls == 2: + raise KeyboardInterrupt("injected interruption") + return real_replace(source, destination) + + interrupt_second.calls = 0 + with mock.patch("os.replace", side_effect=interrupt_second): + with self.assertRaises(KeyboardInterrupt): + self._run() + + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + check_result, check_output = self._run("--check") + self.assertEqual(check_result, 1) + self.assertNotIn("Orphan", check_output) + + normal_result, _ = self._run() + self.assertEqual(normal_result, 0) + final_check_result, _ = self._run("--check") + self.assertEqual(final_check_result, 0) + self._assert_no_temps() + + def test_interrupt_after_replace_restores_tree(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + real_replace = os.replace + + def interrupt_after_second(source, destination): + interrupt_after_second.calls += 1 + result = real_replace(source, destination) + if interrupt_after_second.calls == 2: + raise KeyboardInterrupt("injected post-replace interruption") + return result + + interrupt_after_second.calls = 0 + with mock.patch("os.replace", side_effect=interrupt_after_second): + with self.assertRaises(KeyboardInterrupt): + self._run() + + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_generation_renders_once_and_second_run_does_no_io(self) -> None: + shaders = [ + self._write_shader("a", "first"), + self._write_shader("b", "second"), + ] + render_counts = {shader: 0 for shader in shaders} + real_headers_for_shader = g.headers_for_shader + + def counted(shader): + render_counts[shader] += 1 + return real_headers_for_shader(shader) + + with mock.patch.object(g, "headers_for_shader", side_effect=counted): + first_result, _ = self._run() + self.assertEqual(first_result, 0) + self.assertEqual(render_counts, {shader: 1 for shader in shaders}) + + with mock.patch( + "tempfile.mkstemp", wraps=tempfile.mkstemp + ) as mkstemp, mock.patch("os.replace", wraps=os.replace) as replace: + second_result, _ = self._run() + self.assertEqual(second_result, 0) + mkstemp.assert_not_called() + replace.assert_not_called() + + class WgslTemplateEngineTest(unittest.TestCase): """Coverage for the $-block template engine + DTYPE/VEC variant matrix.""" From 4c55df40ddee1accfe131fd21c35d050705c525f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:36 -0700 Subject: [PATCH 08/12] [ExecuTorch][WebGPU] Generate typed to_copy shader variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21449 **Generate typed `to_copy` variants from one byte-identical template** The two conversion directions duplicated the same WGSL structure and could drift independently. This moves fp32-to-int32 and int32-to-fp32 into one typed family while preserving both expanded payloads. Mirrors Vulkan `backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.{glsl,yaml}`. Key changes: - `to_copy_convert.wgsl` and YAML — declare the two typed variants. - Generated headers and codegen locks — preserve symbols, workgroups, registry entries, and payload hashes. - Structural/native tests — lock both directions and the round trip. Host routing and dispatch remain unchanged; tests expand fixture coverage. Co-authored-with: Claude Code. ghstack-source-id: 411961470 @exported-using-ghexport Differential Revision: [D113979712](https://our.internmc.facebook.com/intern/diff/D113979712/) --- ...float_to_int.wgsl => to_copy_convert.wgsl} | 6 +- .../runtime/ops/to_copy/to_copy_convert.yaml | 15 ++++ .../ops/to_copy/to_copy_float_to_int_wgsl.h | 2 +- .../ops/to_copy/to_copy_int_to_float.wgsl | 18 ---- .../ops/to_copy/to_copy_int_to_float_wgsl.h | 2 +- backends/webgpu/test/op_tests/cases.py | 39 +++++++++ backends/webgpu/test/ops/test_to_copy.py | 87 ++++++++++++++++++- backends/webgpu/test/test_wgsl_codegen.py | 50 ++++++++++- 8 files changed, 191 insertions(+), 28 deletions(-) rename backends/webgpu/runtime/ops/to_copy/{to_copy_float_to_int.wgsl => to_copy_convert.wgsl} (62%) create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml delete mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl similarity index 62% rename from backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl rename to backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl index 3eb6cb44595..f1113f0e14c 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl @@ -1,5 +1,5 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; +@group(0) @binding(0) var input: array<${IN_TYPE}>; +@group(0) @binding(1) var output: array<${OUT_TYPE}>; struct Params { num_elements: u32, @@ -14,5 +14,5 @@ fn main(@builtin(global_invocation_id) gid: vec3) { if (idx >= params.num_elements) { return; } - output[idx] = i32(input[idx]); + output[idx] = ${OUT_TYPE}(input[idx]); } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml new file mode 100644 index 00000000000..219b2d8edd3 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml @@ -0,0 +1,15 @@ +# 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. + +to_copy_convert: + parameter_names_with_default_values: + IN_TYPE: f32 + OUT_TYPE: i32 + shader_variants: + - NAME: to_copy_float_to_int + - NAME: to_copy_int_to_float + IN_TYPE: i32 + OUT_TYPE: f32 diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h index e7e0391dd13..1a384c747d8 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from to_copy_float_to_int.wgsl - DO NOT EDIT. +// @generated from to_copy_convert.wgsl - DO NOT EDIT. // wgsl-sha256: c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f inline constexpr const char* kToCopyFloatToIntWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl deleted file mode 100644 index 87affe78290..00000000000 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl +++ /dev/null @@ -1,18 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; - if (idx >= params.num_elements) { - return; - } - output[idx] = f32(input[idx]); -} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h index fd18700f17c..6fdf37ec2b7 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from to_copy_int_to_float.wgsl - DO NOT EDIT. +// @generated from to_copy_convert.wgsl - DO NOT EDIT. // wgsl-sha256: e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195 inline constexpr const char* kToCopyIntToFloatWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 19438d6acbc..be46d35b520 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -145,6 +145,13 @@ SqueezeModule, ) +from executorch.backends.webgpu.test.ops.test_to_copy import ( + to_copy_float_input, + to_copy_int_input, + ToCopyFloatToIntToFloatModule, + ToCopyIntToFloatModule, +) + from executorch.backends.webgpu.test.ops.test_unary_activations import ( _lin as _unary_lin, CLAMP_CONFIGS, @@ -845,6 +852,38 @@ def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) +def _to_copy_factory(variant: str) -> torch.nn.Module: + return { + "int_to_float": ToCopyIntToFloatModule, + "float_roundtrip": ToCopyFloatToIntToFloatModule, + }[variant]() + + +@register_op_test("to_copy") +def _to_copy_suite() -> WebGPUTestSuite: + cases = [] + for n in (63, 64, 65, 257): + cases.extend( + [ + Case( + name=f"int_to_float_{n}", + construct={"variant": "int_to_float"}, + inputs=(InputSpec(shape=(n,), gen=to_copy_int_input),), + ), + Case( + name=f"float_roundtrip_{n}", + construct={"variant": "float_roundtrip"}, + inputs=(InputSpec(shape=(n,), gen=to_copy_float_input),), + ), + ] + ) + return WebGPUTestSuite( + module_factory=_to_copy_factory, + cases=cases, + golden_dtype="float32", + ) + + @register_op_test("select") def _select_suite() -> WebGPUTestSuite: return _fn_config_suite(SelectModule, _SELECT_CONFIGS) diff --git a/backends/webgpu/test/ops/test_to_copy.py b/backends/webgpu/test/ops/test_to_copy.py index baf068a3438..1fa2375f248 100644 --- a/backends/webgpu/test/ops/test_to_copy.py +++ b/backends/webgpu/test/ops/test_to_copy.py @@ -13,12 +13,14 @@ lvp golden (yolo11n / Depth-Anything-V2). """ +import math import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import get_lowered_submodules class ToCopyIntToFloatModule(torch.nn.Module): @@ -27,6 +29,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.float32) +class ToCopyFloatToIntModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.int32) + + +class ToCopyFloatToIntToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.int32).to(torch.float32) + + class ToCopyFloatModule(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: # Same-dtype copy (flat byte-copy path); copy=True keeps the op from @@ -34,11 +46,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.float32, copy=True) -def _export(model: torch.nn.Module, x: torch.Tensor): +def to_copy_int_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + return (torch.arange(n, dtype=torch.int32) - n // 2).reshape(shape) + + +def to_copy_float_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor( + [-8.75, -3.0, -1.5, -0.25, 0.0, 0.25, 1.5, 3.0, 8.75], + dtype=torch.float32, + ) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def _lower(model: torch.nn.Module, x: torch.Tensor): ep = torch.export.export(model.eval(), (x,)) - return to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + return ep, edge + + +def _export(model: torch.nn.Module, x: torch.Tensor): + _, edge = _lower(model, x) + return edge.to_executorch() def _delegated(et) -> bool: @@ -49,6 +80,29 @@ def _delegated(et) -> bool: ) +def _prepartition_cast_dtypes(ep) -> list[torch.dtype]: + return [ + node.args[1] + for node in ep.graph_module.graph.nodes + if node.op == "call_function" and node.target == torch.ops.aten.to.dtype + ] + + +def _delegated_cast_dtypes(edge) -> list[torch.dtype]: + graph_module = edge.exported_program().graph_module + if any( + "_to_dim_order_copy" in str(getattr(node, "target", "")) + for node in graph_module.graph.nodes + ): + return [] + return [ + node.kwargs["dtype"] + for _, lowered, _ in get_lowered_submodules(graph_module) + for node in lowered.original_module.graph_module.graph.nodes + if "_to_dim_order_copy" in str(getattr(node, "target", "")) + ] + + class ToCopyTest(unittest.TestCase): def test_int_to_float_delegates(self) -> None: x = torch.tensor([1, 2, 3], dtype=torch.int32) @@ -57,6 +111,31 @@ def test_int_to_float_delegates(self) -> None: _delegated(et), "Expected a VulkanBackend delegate (to_copy int->float)" ) + def test_float_to_int_delegates(self) -> None: + x = torch.tensor([-3.75, -1.0, 0.0, 1.9, 63.0], dtype=torch.float32) + et = _export(ToCopyFloatToIntModule(), x) + self.assertTrue( + _delegated(et), "Expected a VulkanBackend delegate (to_copy float->int)" + ) + + def test_roundtrip_keeps_both_casts_in_delegate(self) -> None: + x = torch.tensor([-3.75, -1.0, 0.0, 1.9, 63.0], dtype=torch.float32) + ep, edge = _lower(ToCopyFloatToIntToFloatModule(), x) + expected = [torch.int32, torch.float32] + self.assertEqual(_prepartition_cast_dtypes(ep), expected) + self.assertEqual(_delegated_cast_dtypes(edge), expected) + self.assertTrue(_delegated(edge.to_executorch())) + + for module, one_direction_input in ( + (ToCopyFloatToIntModule(), x), + ( + ToCopyIntToFloatModule(), + torch.tensor([-3, -1, 0, 1, 63], dtype=torch.int32), + ), + ): + _, one_direction_edge = _lower(module, one_direction_input) + self.assertNotEqual(_delegated_cast_dtypes(one_direction_edge), expected) + def test_float_passthrough_delegates(self) -> None: x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) et = _export(ToCopyFloatModule(), x) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 831d2ce0747..b8c239483f2 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -213,7 +213,7 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 134) self.assertEqual( digest.hexdigest(), - "19a0baf9345bec02fe2091a0e6320b81d966724bedc33feb0779c6a824925972", + "ab15be30e7cfa2cb2f6fa7743d3b9f03535f5cc0b88d64d9b6bad8f777efda25", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -907,6 +907,54 @@ def test_rms_norm_template_roundtrip_byte_identical(self) -> None: got, want, f"{header_name} not reproduced from rms_norm.wgsl template" ) + def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: + to_copy_dir = g.BACKEND_ROOT / "runtime/ops/to_copy" + template_path = to_copy_dir / "to_copy_convert.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "to_copy_float_to_int": ( + "f32", + "i32", + "c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f", + ), + "to_copy_int_to_float": ( + "i32", + "f32", + "e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (in_type, out_type, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual( + (params["IN_TYPE"], params["OUT_TYPE"]), (in_type, out_type) + ) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header = (to_copy_dir / f"{name}_wgsl.h").read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + self.assertEqual( + entries["to_copy_float_to_int"].include, + "runtime/ops/to_copy/to_copy_float_to_int_wgsl.h", + ) + self.assertEqual( + entries["to_copy_int_to_float"].include, + "runtime/ops/to_copy/to_copy_int_to_float_wgsl.h", + ) + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "ce1777820bffe77e7cdda312f86a3fd41a090f8f282a6add66666446d06b1608", + ) + def test_rms_norm_half_variant_is_type_correct(self) -> None: # A DTYPE=half expansion must emit compilable WGSL: `enable f16;`, an f32 # accumulator, loads widened to f32 for the reduction, and the store From 1b73aa5d07a3a3b55201fc60fbb7158385bad796 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:37 -0700 Subject: [PATCH 09/12] [ExecuTorch][WebGPU] Generate extrema and unary shader variants Pull Request resolved: https://github.com/pytorch/executorch/pull/21450 **Generate extrema and unary shader families** The extrema reductions and ten no-parameter unary kernels duplicated shader skeletons that could drift independently. This consolidates amax/amin behind one extrema template and abs/cos/exp/hardswish/neg/round/rsqrt/sin/sqrt/tanh behind one unary template while preserving the generated runtime payloads. Key changes: - Generate amax/amin from one extrema manifest. - Generate ten unary payloads from one operator-expression manifest. - Lock expanded bytes, registry entries, delegation, and boundary numerics. The attempted Unary lifecycle migration is intentionally not part of the stack: its performance campaign did not produce an authoritative passing result, so the Unary builder, interface, and activation/sigmoid call sites are restored to their pre-migration bytes. Co-authored-with: Claude Code. ghstack-source-id: 411961475 @exported-using-ghexport Differential Revision: [D113979760](https://our.internmc.facebook.com/intern/diff/D113979760/) --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 4 +- backends/webgpu/runtime/ops/amax/Reduce.cpp | 2 +- backends/webgpu/runtime/ops/amin/Reduce.cpp | 2 +- backends/webgpu/runtime/ops/amin/amin.wgsl | 49 -------- .../runtime/ops/{amax => extrema}/amax_wgsl.h | 2 +- .../runtime/ops/{amin => extrema}/amin_wgsl.h | 2 +- .../{amax/amax.wgsl => extrema/extrema.wgsl} | 4 +- .../webgpu/runtime/ops/extrema/extrema.yaml | 13 +++ backends/webgpu/runtime/ops/unary/abs.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/abs_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/cos.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/cos_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/exp_wgsl.h | 2 +- .../webgpu/runtime/ops/unary/hardswish.wgsl | 21 ---- .../webgpu/runtime/ops/unary/hardswish_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/neg.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/neg_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/round.wgsl | 21 ---- .../webgpu/runtime/ops/unary/round_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/rsqrt.wgsl | 21 ---- .../webgpu/runtime/ops/unary/rsqrt_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/sin.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/sin_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/sqrt.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/sqrt_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/tanh.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/tanh_wgsl.h | 2 +- .../ops/unary/{exp.wgsl => unary.wgsl} | 2 +- backends/webgpu/runtime/ops/unary/unary.yaml | 29 +++++ .../test/native/test_compute_dispatch.cpp | 109 ++++++++++++++++++ backends/webgpu/test/op_tests/cases.py | 55 +++++---- backends/webgpu/test/ops/test_reduce.py | 108 ++++++++++++++++- backends/webgpu/test/test_wgsl_codegen.py | 90 ++++++++++++++- 33 files changed, 397 insertions(+), 283 deletions(-) delete mode 100644 backends/webgpu/runtime/ops/amin/amin.wgsl rename backends/webgpu/runtime/ops/{amax => extrema}/amax_wgsl.h (97%) rename backends/webgpu/runtime/ops/{amin => extrema}/amin_wgsl.h (97%) rename backends/webgpu/runtime/ops/{amax/amax.wgsl => extrema/extrema.wgsl} (94%) create mode 100644 backends/webgpu/runtime/ops/extrema/extrema.yaml delete mode 100644 backends/webgpu/runtime/ops/unary/abs.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/cos.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/hardswish.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/neg.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/round.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/rsqrt.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/sin.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/sqrt.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/tanh.wgsl rename backends/webgpu/runtime/ops/unary/{exp.wgsl => unary.wgsl} (94%) create mode 100644 backends/webgpu/runtime/ops/unary/unary.yaml diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 1cf4b952a5f..b475a6cd999 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -13,8 +13,6 @@ #include #include #include -#include -#include #include #include #include @@ -48,6 +46,8 @@ #include #include #include +#include +#include #include #include #include diff --git a/backends/webgpu/runtime/ops/amax/Reduce.cpp b/backends/webgpu/runtime/ops/amax/Reduce.cpp index 14675d97e96..c1dff62e035 100644 --- a/backends/webgpu/runtime/ops/amax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amax/Reduce.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/amin/Reduce.cpp b/backends/webgpu/runtime/ops/amin/Reduce.cpp index fbe574fdf0b..24a5ebeb826 100644 --- a/backends/webgpu/runtime/ops/amin/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amin/Reduce.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/amin/amin.wgsl b/backends/webgpu/runtime/ops/amin/amin.wgsl deleted file mode 100644 index 4778800ab3d..00000000000 --- a/backends/webgpu/runtime/ops/amin/amin.wgsl +++ /dev/null @@ -1,49 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_rows: u32, - reduce_size: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -// Cooperative shared-memory reduction; mirrors Vulkan reduce.glsl (a group of -// threads co-operates per reduction row, partials aggregated in shared memory). -// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. -var partials: array; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(workgroup_id) wid: vec3, - @builtin(local_invocation_id) lid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. - let row = wid.x + wid.y * num_workgroups.x; - if (row >= params.num_rows) { - return; - } - let base = row * params.reduce_size; - - // Each thread reduces a strided slice of the row into a partial. Seed with - // the row's first element (always valid; reduce_size >= 1) so threads that - // own no element contribute a real value, not an out-of-range identity. - var acc = input[base]; - var i = lid.x; - while (i < params.reduce_size) { - acc = min(acc, input[base + i]); - i = i + wg_size; - } - partials[lid.x] = acc; - workgroupBarrier(); - - // Thread 0 aggregates the wg_size partials (mirrors Vulkan's group aggregate). - if (lid.x == 0u) { - var m = partials[0]; - for (var t = 1u; t < wg_size; t = t + 1u) { - m = min(m, partials[t]); - } - output[row] = m; - } -} diff --git a/backends/webgpu/runtime/ops/amax/amax_wgsl.h b/backends/webgpu/runtime/ops/extrema/amax_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/amax/amax_wgsl.h rename to backends/webgpu/runtime/ops/extrema/amax_wgsl.h index 48ec8f20e27..0b7a4d2238e 100644 --- a/backends/webgpu/runtime/ops/amax/amax_wgsl.h +++ b/backends/webgpu/runtime/ops/extrema/amax_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from amax.wgsl - DO NOT EDIT. +// @generated from extrema.wgsl - DO NOT EDIT. // wgsl-sha256: 35fc059d7c72caa17f9cb1128823ecfd8f75be4ce24b6cd4f9629a97b52f64c0 inline constexpr const char* kAmaxWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/amin/amin_wgsl.h b/backends/webgpu/runtime/ops/extrema/amin_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/amin/amin_wgsl.h rename to backends/webgpu/runtime/ops/extrema/amin_wgsl.h index 40a97c67a63..8b8bf5456b1 100644 --- a/backends/webgpu/runtime/ops/amin/amin_wgsl.h +++ b/backends/webgpu/runtime/ops/extrema/amin_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from amin.wgsl - DO NOT EDIT. +// @generated from extrema.wgsl - DO NOT EDIT. // wgsl-sha256: 8cb6035ae4d34eb2a6cc973d93d9847905722e967239c96033fccfe3a1943cb2 inline constexpr const char* kAminWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/amax/amax.wgsl b/backends/webgpu/runtime/ops/extrema/extrema.wgsl similarity index 94% rename from backends/webgpu/runtime/ops/amax/amax.wgsl rename to backends/webgpu/runtime/ops/extrema/extrema.wgsl index 2f23b9c38fa..6c55d60d9b7 100644 --- a/backends/webgpu/runtime/ops/amax/amax.wgsl +++ b/backends/webgpu/runtime/ops/extrema/extrema.wgsl @@ -32,7 +32,7 @@ fn main( var acc = input[base]; var i = lid.x; while (i < params.reduce_size) { - acc = max(acc, input[base + i]); + acc = ${REDUCE_FN}(acc, input[base + i]); i = i + wg_size; } partials[lid.x] = acc; @@ -42,7 +42,7 @@ fn main( if (lid.x == 0u) { var m = partials[0]; for (var t = 1u; t < wg_size; t = t + 1u) { - m = max(m, partials[t]); + m = ${REDUCE_FN}(m, partials[t]); } output[row] = m; } diff --git a/backends/webgpu/runtime/ops/extrema/extrema.yaml b/backends/webgpu/runtime/ops/extrema/extrema.yaml new file mode 100644 index 00000000000..a85da849c60 --- /dev/null +++ b/backends/webgpu/runtime/ops/extrema/extrema.yaml @@ -0,0 +1,13 @@ +# 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. + +extrema: + parameter_names_with_default_values: + REDUCE_FN: max + shader_variants: + - NAME: amax + - NAME: amin + REDUCE_FN: min diff --git a/backends/webgpu/runtime/ops/unary/abs.wgsl b/backends/webgpu/runtime/ops/unary/abs.wgsl deleted file mode 100644 index e3e10c75dd9..00000000000 --- a/backends/webgpu/runtime/ops/unary/abs.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = abs(x); -} diff --git a/backends/webgpu/runtime/ops/unary/abs_wgsl.h b/backends/webgpu/runtime/ops/unary/abs_wgsl.h index 3d2873e69c4..0efb72d5189 100644 --- a/backends/webgpu/runtime/ops/unary/abs_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/abs_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from abs.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 39d3c163fdf6a92286828f4b3217e00294e3ca5634a878ed5fd34e3b1cdf0a27 inline constexpr const char* kAbsWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/cos.wgsl b/backends/webgpu/runtime/ops/unary/cos.wgsl deleted file mode 100644 index c2bafe7b248..00000000000 --- a/backends/webgpu/runtime/ops/unary/cos.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = cos(x); -} diff --git a/backends/webgpu/runtime/ops/unary/cos_wgsl.h b/backends/webgpu/runtime/ops/unary/cos_wgsl.h index 422b5618146..4ca99df88f3 100644 --- a/backends/webgpu/runtime/ops/unary/cos_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/cos_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from cos.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 9df78873e5fae98d347c26db2a02b047ea3d5d2c93f0761cb9ac6995f9a71ab2 inline constexpr const char* kCosWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/exp_wgsl.h b/backends/webgpu/runtime/ops/unary/exp_wgsl.h index 28b6c67cb18..cbf85fd415a 100644 --- a/backends/webgpu/runtime/ops/unary/exp_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/exp_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from exp.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 3171399bc36acf9c1cb2a03c2a31038318203c4c63ab03c4881df7a660346020 inline constexpr const char* kExpWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/hardswish.wgsl b/backends/webgpu/runtime/ops/unary/hardswish.wgsl deleted file mode 100644 index 2278caf3664..00000000000 --- a/backends/webgpu/runtime/ops/unary/hardswish.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = select(select(x * (x + 3.0) / 6.0, x, x >= 3.0), 0.0, x <= -3.0); -} diff --git a/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h b/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h index 0c991547b9a..43f104c7f8f 100644 --- a/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from hardswish.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: c874a15ef6cdaec71187296016cc2a1515f5e7c889b97dfa8fd4b278e6e2c3d5 inline constexpr const char* kHardswishWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/neg.wgsl b/backends/webgpu/runtime/ops/unary/neg.wgsl deleted file mode 100644 index c977957957d..00000000000 --- a/backends/webgpu/runtime/ops/unary/neg.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = -x; -} diff --git a/backends/webgpu/runtime/ops/unary/neg_wgsl.h b/backends/webgpu/runtime/ops/unary/neg_wgsl.h index d528c45fea0..c4c7bb989ac 100644 --- a/backends/webgpu/runtime/ops/unary/neg_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/neg_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from neg.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 8851b9f42d14153f6f04484fee2f8bf67bda26dea892ff48768e09e6ad49cee1 inline constexpr const char* kNegWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/round.wgsl b/backends/webgpu/runtime/ops/unary/round.wgsl deleted file mode 100644 index 2269ca59988..00000000000 --- a/backends/webgpu/runtime/ops/unary/round.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = round(x); -} diff --git a/backends/webgpu/runtime/ops/unary/round_wgsl.h b/backends/webgpu/runtime/ops/unary/round_wgsl.h index 209305bc855..8c805ba4f80 100644 --- a/backends/webgpu/runtime/ops/unary/round_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/round_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from round.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 8f3e0edbeb81aa50f35e691c78554e8057fa8d78fe8a86454f4f42e5e8871452 inline constexpr const char* kRoundWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/rsqrt.wgsl b/backends/webgpu/runtime/ops/unary/rsqrt.wgsl deleted file mode 100644 index 1f50c4f66ac..00000000000 --- a/backends/webgpu/runtime/ops/unary/rsqrt.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = 1.0 / sqrt(x); -} diff --git a/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h b/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h index 83e74a1b8cf..58bd05d010e 100644 --- a/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from rsqrt.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 108765d5a23b87473f34651875d08abf2a5fa8980bd92fc8cbe3617295097747 inline constexpr const char* kRsqrtWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/sin.wgsl b/backends/webgpu/runtime/ops/unary/sin.wgsl deleted file mode 100644 index ffd2a07ea8a..00000000000 --- a/backends/webgpu/runtime/ops/unary/sin.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = sin(x); -} diff --git a/backends/webgpu/runtime/ops/unary/sin_wgsl.h b/backends/webgpu/runtime/ops/unary/sin_wgsl.h index f22229b2342..54184a1ccb8 100644 --- a/backends/webgpu/runtime/ops/unary/sin_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/sin_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from sin.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: e5762804773659d348fddddcef4935807ae6fe7d92c92eb17a2f44aae8f2c5b9 inline constexpr const char* kSinWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/sqrt.wgsl b/backends/webgpu/runtime/ops/unary/sqrt.wgsl deleted file mode 100644 index e34ff440007..00000000000 --- a/backends/webgpu/runtime/ops/unary/sqrt.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = sqrt(x); -} diff --git a/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h b/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h index 260a4ce4266..42dcdb838c0 100644 --- a/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from sqrt.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 008534ae365969f5c180b42e8d6d0b131df78f181e5435abbcafc3ffb8be8aac inline constexpr const char* kSqrtWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/tanh.wgsl b/backends/webgpu/runtime/ops/unary/tanh.wgsl deleted file mode 100644 index 5a541e21699..00000000000 --- a/backends/webgpu/runtime/ops/unary/tanh.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = tanh(clamp(x, -15.0, 15.0)); -} diff --git a/backends/webgpu/runtime/ops/unary/tanh_wgsl.h b/backends/webgpu/runtime/ops/unary/tanh_wgsl.h index eef1ecd91af..51ba4d3919c 100644 --- a/backends/webgpu/runtime/ops/unary/tanh_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/tanh_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from tanh.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 5bd7eb1c6411940d84a9b311884f35b39f15b82103b14bab02902290ed6b0339 inline constexpr const char* kTanhWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/exp.wgsl b/backends/webgpu/runtime/ops/unary/unary.wgsl similarity index 94% rename from backends/webgpu/runtime/ops/unary/exp.wgsl rename to backends/webgpu/runtime/ops/unary/unary.wgsl index b69aa509d8e..d974a2f3319 100644 --- a/backends/webgpu/runtime/ops/unary/exp.wgsl +++ b/backends/webgpu/runtime/ops/unary/unary.wgsl @@ -17,5 +17,5 @@ fn main( return; } let x = input[idx]; - output[idx] = exp(x); + output[idx] = ${OPERATOR}; } diff --git a/backends/webgpu/runtime/ops/unary/unary.yaml b/backends/webgpu/runtime/ops/unary/unary.yaml new file mode 100644 index 00000000000..84c0dd660d7 --- /dev/null +++ b/backends/webgpu/runtime/ops/unary/unary.yaml @@ -0,0 +1,29 @@ +# 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. + +unary: + parameter_names_with_default_values: + OPERATOR: "abs(x)" + shader_variants: + - NAME: abs + - NAME: cos + OPERATOR: "cos(x)" + - NAME: exp + OPERATOR: "exp(x)" + - NAME: hardswish + OPERATOR: "select(select(x * (x + 3.0) / 6.0, x, x >= 3.0), 0.0, x <= -3.0)" + - NAME: neg + OPERATOR: "-x" + - NAME: round + OPERATOR: "round(x)" + - NAME: rsqrt + OPERATOR: "1.0 / sqrt(x)" + - NAME: sin + OPERATOR: "sin(x)" + - NAME: sqrt + OPERATOR: "sqrt(x)" + - NAME: tanh + OPERATOR: "tanh(clamp(x, -15.0, 15.0))" diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index b2b0d028567..6ed7229604b 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -440,6 +440,18 @@ void record_resize_probe(WebGPUGraph&, const ResizeProbeContext& context) { ++*context.calls; } +struct ThrowingResizeContext { + bool* fail; + int* calls; +}; + +void maybe_throw_resize(WebGPUGraph&, const ThrowingResizeContext& context) { + ++*context.calls; + if (*context.fail) { + throw std::runtime_error("resize hook failure"); + } +} + struct GridPickerContext { int tensor_id; uint32_t x_bias; @@ -487,6 +499,25 @@ TEST(WebGPUResizeHooks, TypedRegistrationOwnsContextCopies) { int tensor_calls = 0; int symint_observed = 0; int symint_calls = 0; + using ResizeProbeFn = void (*)(WebGPUGraph&, const ResizeProbeContext&); + EXPECT_THROW( + graph.add_tensor_resize_hook( + kResizeQ, + static_cast(nullptr), + ResizeProbeContext{0, &tensor_observed, &tensor_calls}), + std::runtime_error); + EXPECT_THROW( + graph.add_tensor_resize_hook( + kResizeSymInt, + record_resize_probe, + ResizeProbeContext{0, &tensor_observed, &tensor_calls}), + std::runtime_error); + EXPECT_THROW( + graph.add_resize_hook( + kResizeQ, + record_resize_probe, + ResizeProbeContext{0, &symint_observed, &symint_calls}), + std::runtime_error); { ResizeProbeContext tensor_context = {17, &tensor_observed, &tensor_calls}; ResizeProbeContext symint_context = {23, &symint_observed, &symint_calls}; @@ -510,6 +541,36 @@ TEST(WebGPUResizeHooks, TypedRegistrationOwnsContextCopies) { EXPECT_EQ(symint_calls, 1); } +TEST(WebGPUResizeHooks, RestoresDirtyTriggerWhenHookThrows) { + WebGPUGraph graph; + build_resize_test_graph(graph); + bool hook_fails = true; + int hook_calls = 0; + int picker_calls = 0; + graph.add_tensor_resize_hook( + kResizeQ, + maybe_throw_resize, + ThrowingResizeContext{&hook_fails, &hook_calls}); + const size_t dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_hook_retry"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 1u, 2u, &picker_calls, nullptr}); + picker_calls = 0; + + graph.resize_input(kResizeQ, {4, 3}); + EXPECT_THROW(graph.propagate_resize(), std::runtime_error); + EXPECT_EQ(hook_calls, 1); + EXPECT_EQ(picker_calls, 0); + expect_dispatch_grid(graph, dispatch, 9u, 10u); + + hook_fails = false; + EXPECT_NO_THROW(graph.propagate_resize()); + EXPECT_EQ(hook_calls, 2); + EXPECT_EQ(picker_calls, 1); + expect_dispatch_grid(graph, dispatch, 5u, 5u); +} + TEST(WebGPUDynamicDispatch, InitializesAndIsolatesTriggeredGrids) { WebGPUGraph graph; build_resize_test_graph(graph); @@ -615,12 +676,60 @@ TEST(WebGPUDynamicDispatch, HandlesCascadesAndStagesPickerFailures) { EXPECT_THROW(graph.propagate_resize(), std::runtime_error); expect_dispatch_grid(graph, first, 9u, 10u); expect_dispatch_grid(graph, second, 11u, 12u); + second_fails = false; + EXPECT_NO_THROW(graph.propagate_resize()); + expect_dispatch_grid(graph, first, 5u, 5u); + expect_dispatch_grid(graph, second, 7u, 7u); } TEST(WebGPUDynamicDispatch, RejectsRouteOverlapWithoutPoisoningRegistry) { WebGPUGraph graph; build_resize_test_graph(graph); int calls = 0; + const size_t dispatches_before_invalid_trigger = graph.num_dispatches(); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_negative_trigger"), + -1, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_oob_trigger"), + graph.num_values(), + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_symint_trigger"), + kResizeSymInt, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_EQ(calls, 0); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_invalid_trigger); + + auto pick_zero_grid = [](const WebGPUGraph&, const WebGPUDispatchGrid& grid) { + return grid; + }; + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_zero_x"), + kResizeQ, + +pick_zero_grid, + WebGPUDispatchGrid{0u, 1u}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_zero_y"), + kResizeQ, + +pick_zero_grid, + WebGPUDispatchGrid{1u, 0u}), + std::runtime_error); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_invalid_trigger); + bool initial_fails = true; const size_t dispatches_before_failure = graph.num_dispatches(); EXPECT_THROW( diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index be46d35b520..aa8f000bceb 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -112,7 +112,15 @@ DequantizeConstModule, QuantizeModule, ) -from executorch.backends.webgpu.test.ops.test_reduce import AmaxModule, AminModule +from executorch.backends.webgpu.test.ops.test_reduce import ( + amax_sign_trap_input, + amax_tie_input, + AmaxModule, + amin_sign_trap_input, + amin_tie_input, + AminModule, + EXTREMA_CONFIGS, +) from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_rms_norm import ( _CASES, @@ -489,35 +497,42 @@ def _floor_divide_suite() -> WebGPUTestSuite: ) -def _reduce_suite(module_cls) -> WebGPUTestSuite: - # Last-dim reduction; both keepdim variants over a 2d and a 3d shape. - return WebGPUTestSuite( - module_factory=lambda keepdim: module_cls(keepdim), - cases=[ - Case(name="keepdim_2d", construct={"keepdim": True}, inputs=((M1, M2),)), - Case(name="nodim_2d", construct={"keepdim": False}, inputs=((M1, M2),)), - Case( - name="keepdim_3d", - construct={"keepdim": True}, - inputs=((S, S1, S2),), - ), +def _reduce_suite(module_cls, op: str) -> WebGPUTestSuite: + generators = { + ("amax", "sign_trap"): amax_sign_trap_input, + ("amax", "tie"): amax_tie_input, + ("amin", "sign_trap"): amin_sign_trap_input, + ("amin", "tie"): amin_tie_input, + } + cases = [] + for name, shape, dim, keepdim, input_class in EXTREMA_CONFIGS: + inputs = (shape,) + kwargs = {} + if input_class != "default": + inputs = (InputSpec(shape=shape, gen=generators[(op, input_class)]),) + kwargs = {"atol": 0.0, "rtol": 0.0} + cases.append( Case( - name="nodim_3d", - construct={"keepdim": False}, - inputs=((S, S1, S2),), - ), - ], + name=name, + construct={"keepdim": keepdim, "dim": dim}, + inputs=inputs, + **kwargs, + ) + ) + return WebGPUTestSuite( + module_factory=module_cls, + cases=cases, ) @register_op_test("amax") def _amax_suite() -> WebGPUTestSuite: - return _reduce_suite(AmaxModule) + return _reduce_suite(AmaxModule, "amax") @register_op_test("amin") def _amin_suite() -> WebGPUTestSuite: - return _reduce_suite(AminModule) + return _reduce_suite(AminModule, "amin") @register_op_test("flip") diff --git a/backends/webgpu/test/ops/test_reduce.py b/backends/webgpu/test/ops/test_reduce.py index c7b91d8156c..cdfd7eed96c 100644 --- a/backends/webgpu/test/ops/test_reduce.py +++ b/backends/webgpu/test/ops/test_reduce.py @@ -19,12 +19,14 @@ from __future__ import annotations +import math import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes class ReduceModule(torch.nn.Module): @@ -61,6 +63,70 @@ def _det_input(shape) -> torch.Tensor: return ((flat % 17) - 8).div(16.0).reshape(shape) +# Shared structural and manifest-driven extrema case authority. +EXTREMA_CONFIGS = ( + ("keepdim_2d", (37, 41), -1, True, "default"), + ("nodim_2d", (37, 41), -1, False, "default"), + ("keepdim_3d", (5, 7, 11), -1, True, "default"), + ("nodim_3d", (5, 7, 11), -1, False, "default"), + ("sign_trap_63_drop", (3, 63), -1, False, "sign_trap"), + ("tie_64_keep", (2, 64), -1, True, "tie"), + ("tie_65_posdim_drop", (2, 65), 1, False, "tie"), + ("sign_trap_255_keep", (2, 255), -1, True, "sign_trap"), + ("tie_256_drop", (2, 256), -1, False, "tie"), + ("tie_257_posdim_keep", (2, 257), 1, True, "tie"), +) + + +def _sign_trap_input(shape, *, for_max: bool) -> torch.Tensor: + magnitude = ((torch.arange(math.prod(shape), dtype=torch.float32) % 31) + 1).div( + 8.0 + ) + rows = magnitude.reshape(-1, shape[-1]) + for row_index, row in enumerate(rows): + row.add_(float(row_index)) + values = -magnitude if for_max else magnitude + return values.reshape(shape) + + +def amax_sign_trap_input(shape) -> torch.Tensor: + return _sign_trap_input(shape, for_max=True) + + +def amin_sign_trap_input(shape) -> torch.Tensor: + return _sign_trap_input(shape, for_max=False) + + +def _tie_input(shape, *, for_max: bool) -> torch.Tensor: + values = ((torch.arange(math.prod(shape), dtype=torch.float32) % 29) - 14).div(8.0) + rows = values.reshape(-1, shape[-1]) + for row_index, row in enumerate(rows): + extreme = 16.0 + float(row_index) + if not for_max: + extreme = -extreme + row[-2] = extreme + row[-1] = extreme + return values.reshape(shape) + + +def amax_tie_input(shape) -> torch.Tensor: + return _tie_input(shape, for_max=True) + + +def amin_tie_input(shape) -> torch.Tensor: + return _tie_input(shape, for_max=False) + + +def _extrema_input(op: str, input_class: str, shape) -> torch.Tensor: + if input_class == "default": + return _det_input(shape) + if input_class == "sign_trap": + return ( + amax_sign_trap_input(shape) if op == "amax" else amin_sign_trap_input(shape) + ) + return amax_tie_input(shape) if op == "amax" else amin_tie_input(shape) + + def _export(m: torch.nn.Module, x: torch.Tensor): ep = torch.export.export(m, (x,)) return to_edge_transform_and_lower( @@ -128,21 +194,55 @@ def export_reduce_model( class AmaxModule(torch.nn.Module): - def __init__(self, keepdim: bool) -> None: + def __init__(self, keepdim: bool, dim: int = -1) -> None: super().__init__() self.keepdim = keepdim + self.dim = dim def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.amax(x, dim=-1, keepdim=self.keepdim) + return torch.amax(x, dim=self.dim, keepdim=self.keepdim) class AminModule(torch.nn.Module): - def __init__(self, keepdim: bool) -> None: + def __init__(self, keepdim: bool, dim: int = -1) -> None: super().__init__() self.keepdim = keepdim + self.dim = dim def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.amin(x, dim=-1, keepdim=self.keepdim) + return torch.amin(x, dim=self.dim, keepdim=self.keepdim) + + +class TestExtrema(unittest.TestCase): + def test_config_contract(self) -> None: + self.assertEqual( + EXTREMA_CONFIGS, + ( + ("keepdim_2d", (37, 41), -1, True, "default"), + ("nodim_2d", (37, 41), -1, False, "default"), + ("keepdim_3d", (5, 7, 11), -1, True, "default"), + ("nodim_3d", (5, 7, 11), -1, False, "default"), + ("sign_trap_63_drop", (3, 63), -1, False, "sign_trap"), + ("tie_64_keep", (2, 64), -1, True, "tie"), + ("tie_65_posdim_drop", (2, 65), 1, False, "tie"), + ("sign_trap_255_keep", (2, 255), -1, True, "sign_trap"), + ("tie_256_drop", (2, 256), -1, False, "tie"), + ("tie_257_posdim_keep", (2, 257), 1, True, "tie"), + ), + ) + + def test_exports_fully_delegated(self) -> None: + for op, module_cls in (("amax", AmaxModule), ("amin", AminModule)): + for name, shape, dim, keepdim, input_class in EXTREMA_CONFIGS: + with self.subTest(op=op, config=name): + x = _extrema_input(op, input_class, shape) + ep = torch.export.export(module_cls(keepdim, dim).eval(), (x,)) + edge = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ) + graph = edge.exported_program().graph_module.graph + self.assertEqual(len(get_delegates(graph)), 1) + self.assertEqual(get_non_lowered_nodes(graph), []) if __name__ == "__main__": diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index b8c239483f2..d437a0a42e8 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -213,7 +213,7 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 134) self.assertEqual( digest.hexdigest(), - "ab15be30e7cfa2cb2f6fa7743d3b9f03535f5cc0b88d64d9b6bad8f777efda25", + "ef97dca2336315ee2c8b0f9e896c6aa082834ae94948bda3ccb42b1145f2bd27", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -952,9 +952,95 @@ def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "ce1777820bffe77e7cdda312f86a3fd41a090f8f282a6add66666446d06b1608", + "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", ) + def test_extrema_template_roundtrip_byte_identical(self) -> None: + extrema_dir = g.BACKEND_ROOT / "runtime/ops/extrema" + template_path = extrema_dir / "extrema.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "amax": ( + "max", + "35fc059d7c72caa17f9cb1128823ecfd8f75be4ce24b6cd4f9629a97b52f64c0", + ), + "amin": ( + "min", + "8cb6035ae4d34eb2a6cc973d93d9847905722e967239c96033fccfe3a1943cb2", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (reduce_fn, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["REDUCE_FN"], reduce_fn) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = extrema_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (256, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/extrema/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "amax": "57f929b9f3087dc32403c3587884ce2ed4be2d03c4e80ff7035428b52e7e0e51", + "amin": "5dc947d4781a67df953b9c5970c4ab2119317b29657f983a9d308dfdf123dede", + } + for name, expected_hash in handler_hashes.items(): + handler = g.BACKEND_ROOT / f"runtime/ops/{name}/Reduce.cpp" + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) + + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", + ) + + def test_unary_template_roundtrip_byte_identical(self) -> None: + unary_dir = g.BACKEND_ROOT / "runtime/ops/unary" + template_path = unary_dir / "unary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "abs": "39d3c163fdf6a92286828f4b3217e00294e3ca5634a878ed5fd34e3b1cdf0a27", + "cos": "9df78873e5fae98d347c26db2a02b047ea3d5d2c93f0761cb9ac6995f9a71ab2", + "exp": "3171399bc36acf9c1cb2a03c2a31038318203c4c63ab03c4881df7a660346020", + "hardswish": "c874a15ef6cdaec71187296016cc2a1515f5e7c889b97dfa8fd4b278e6e2c3d5", + "neg": "8851b9f42d14153f6f04484fee2f8bf67bda26dea892ff48768e09e6ad49cee1", + "round": "8f3e0edbeb81aa50f35e691c78554e8057fa8d78fe8a86454f4f42e5e8871452", + "rsqrt": "108765d5a23b87473f34651875d08abf2a5fa8980bd92fc8cbe3617295097747", + "sin": "e5762804773659d348fddddcef4935807ae6fe7d92c92eb17a2f44aae8f2c5b9", + "sqrt": "008534ae365969f5c180b42e8d6d0b131df78f181e5435abbcafc3ffb8be8aac", + "tanh": "5bd7eb1c6411940d84a9b311884f35b39f15b82103b14bab02902290ed6b0339", + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + entries = {entry.name: entry for entry in g.registry_entries()} + for name, expected_hash in expected.items(): + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **variants[name]}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + header = (unary_dir / f"{name}_wgsl.h").read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (256, 1, 1)) + self.assertEqual(entries[name].include, f"runtime/ops/unary/{name}_wgsl.h") + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + self.assertTrue({"clamp", "pow_scalar"}.isdisjoint(variants)) + def test_rms_norm_half_variant_is_type_correct(self) -> None: # A DTYPE=half expansion must emit compilable WGSL: `enable f16;`, an f32 # accumulator, loads widened to f32 for the reduction, and the store From f8110c84e29eef267410d910b1b26468af601f58 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:39 -0700 Subject: [PATCH 10/12] [ExecuTorch][WebGPU] Generate logical and arithmetic binary shader variants Pull Request resolved: https://github.com/pytorch/executorch/pull/21451 **Generate byte-identical logical and arithmetic binary variants from shared WGSL families** Logical AND/OR and four arithmetic binary kernels duplicated shader skeletons and broadcast logic. This consolidates logical AND/OR behind one packed-Boolean family and minimum/pow/floor_divide/mul into the existing binary family, with a permanent mixed-rank broadcast contract. Key changes: - Generate logical AND/OR from one operator-token manifest. - Generate minimum, pow, floor_divide, and mul beside the existing div/sub variants. - Lock same-shape and mixed-rank expressions, exact payloads/workgroups, PTE delegation, and broadcast boundary cases. No runtime C++ dispatch, bindings, pipeline construction, workgroups, or expanded shader payloads change. Four standalone WGSL inputs are removed, and future compatible variants require manifest entries instead of copied kernels. This follows the Vulkan binary-family pattern. Co-authored-with: Claude Code. ghstack-source-id: 411961479 @exported-using-ghexport Differential Revision: [D113979789](https://our.internmc.facebook.com/intern/diff/D113979789/) --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 12 +- .../binary_floor_divide_wgsl.h | 2 +- .../binary_minimum_wgsl.h | 2 +- .../ops/{mul => binary_op}/binary_mul_wgsl.h | 2 +- .../runtime/ops/binary_op/binary_op.wgsl | 25 +++- .../runtime/ops/binary_op/binary_op.yaml | 23 ++++ .../ops/{pow => binary_op}/binary_pow_wgsl.h | 2 +- .../runtime/ops/floor_divide/BinaryOp.cpp | 2 +- .../ops/floor_divide/binary_floor_divide.wgsl | 51 -------- .../runtime/ops/logical_and/LogicalAnd.cpp | 2 +- .../logical_and_wgsl.h | 2 +- .../logical_binary.wgsl} | 7 +- .../ops/logical_binary/logical_binary.yaml | 13 ++ .../logical_or_wgsl.h | 2 +- .../runtime/ops/logical_or/LogicalOr.cpp | 2 +- .../runtime/ops/logical_or/logical_or.wgsl | 25 ---- .../webgpu/runtime/ops/minimum/BinaryOp.cpp | 2 +- .../runtime/ops/minimum/binary_minimum.wgsl | 51 -------- backends/webgpu/runtime/ops/mul/BinaryOp.cpp | 2 +- .../webgpu/runtime/ops/mul/binary_mul.wgsl | 51 -------- backends/webgpu/runtime/ops/pow/BinaryOp.cpp | 2 +- .../webgpu/runtime/ops/pow/binary_pow.wgsl | 51 -------- backends/webgpu/test/op_tests/cases.py | 117 +++++++----------- .../webgpu/test/op_tests/test_generator.py | 115 ++++++++++++++++- backends/webgpu/test/ops/test_bitwise.py | 20 ++- backends/webgpu/test/ops/test_floor_divide.py | 4 +- backends/webgpu/test/ops/test_logical_and.py | 51 ++++---- backends/webgpu/test/ops/test_logical_or.py | 42 ++----- backends/webgpu/test/ops/test_minimum.py | 4 +- backends/webgpu/test/ops/test_pow.py | 4 +- backends/webgpu/test/test_wgsl_codegen.py | 116 +++++++++++++++-- 31 files changed, 402 insertions(+), 404 deletions(-) rename backends/webgpu/runtime/ops/{floor_divide => binary_op}/binary_floor_divide_wgsl.h (97%) rename backends/webgpu/runtime/ops/{minimum => binary_op}/binary_minimum_wgsl.h (97%) rename backends/webgpu/runtime/ops/{mul => binary_op}/binary_mul_wgsl.h (98%) rename backends/webgpu/runtime/ops/{pow => binary_op}/binary_pow_wgsl.h (98%) delete mode 100644 backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl rename backends/webgpu/runtime/ops/{logical_and => logical_binary}/logical_and_wgsl.h (96%) rename backends/webgpu/runtime/ops/{logical_and/logical_and.wgsl => logical_binary/logical_binary.wgsl} (73%) create mode 100644 backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml rename backends/webgpu/runtime/ops/{logical_or => logical_binary}/logical_or_wgsl.h (96%) delete mode 100644 backends/webgpu/runtime/ops/logical_or/logical_or.wgsl delete mode 100644 backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl delete mode 100644 backends/webgpu/runtime/ops/mul/binary_mul.wgsl delete mode 100644 backends/webgpu/runtime/ops/pow/binary_pow.wgsl diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index b475a6cd999..5374390722b 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -17,6 +17,10 @@ #include #include #include +#include +#include +#include +#include #include #include #include @@ -50,7 +54,6 @@ #include #include #include -#include #include #include #include @@ -65,20 +68,17 @@ #include #include #include -#include -#include +#include +#include #include -#include #include #include #include -#include #include #include #include #include #include -#include #include #include #include diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h index fe2315df30f..dcd2b2e46d8 100644 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_floor_divide.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3 inline constexpr const char* kBinaryFloorDivideWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h index 88d9614ba59..c3f5f76c81f 100644 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_minimum.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: 929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972 inline constexpr const char* kBinaryMinimumWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h index c9f60dbd200..68784e82eb2 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_mul.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812 inline constexpr const char* kBinaryMulWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl index 98076ed98b5..4b42665013f 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl @@ -16,11 +16,14 @@ override wg_size: u32 = 64u; $if USE_ALPHA: override alpha: f32 = 1.0; -fn op(a: f32, b: f32) -> f32 { - return ${OP_EXPR}; -} +$if INLINE: + @compute @workgroup_size(wg_size, 1, 1) +$else: + fn op(a: f32, b: f32) -> f32 { + return ${OP_EXPR}; + } -@compute @workgroup_size(wg_size, 1, 1) + @compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { @@ -30,6 +33,8 @@ fn main( return; } + $if INLINE: + // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || @@ -38,10 +43,15 @@ fn main( } } if (same) { - output[idx] = op(input1[idx], input2[idx]); + $if INLINE: + output[idx] = ${SAME_EXPR}; + $else: + output[idx] = op(input1[idx], input2[idx]); return; } + $if INLINE: + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. var rem = idx; var l1: u32 = 0u; var l2: u32 = 0u; @@ -51,5 +61,8 @@ fn main( l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } - output[idx] = op(input1[l1], input2[l2]); + $if INLINE: + output[idx] = ${BROADCAST_EXPR}; + $else: + output[idx] = op(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml index 03040fcbcce..bd4cd452d35 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml @@ -2,6 +2,9 @@ binary_op: parameter_names_with_default_values: OP_EXPR: a + alpha * b USE_ALPHA: 1 + INLINE: 0 + SAME_EXPR: input1[idx] + input2[idx] + BROADCAST_EXPR: input1[l1] + input2[l2] shader_variants: - NAME: binary_div OP_EXPR: a / b @@ -9,3 +12,23 @@ binary_op: - NAME: binary_sub OP_EXPR: a - alpha * b USE_ALPHA: 1 + - NAME: binary_minimum + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: min(input1[idx], input2[idx]) + BROADCAST_EXPR: min(input1[l1], input2[l2]) + - NAME: binary_pow + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: pow(input1[idx], input2[idx]) + BROADCAST_EXPR: pow(input1[l1], input2[l2]) + - NAME: binary_floor_divide + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: floor(input1[idx] / input2[idx]) + BROADCAST_EXPR: floor(input1[l1] / input2[l2]) + - NAME: binary_mul + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: input1[idx] * input2[idx] + BROADCAST_EXPR: input1[l1] * input2[l2] diff --git a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h index 3532c091160..776d4b6693f 100644 --- a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_pow.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c inline constexpr const char* kBinaryPowWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp index a35a7587311..98a76228003 100644 --- a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include #include diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl deleted file mode 100644 index 3b4edd788b0..00000000000 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = floor(input1[idx] / input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = floor(input1[l1] / input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp index b6a7e12830f..358cb29fff0 100644 --- a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h index 6a21a77a687..3c1f861f119 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_and.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20 inline constexpr const char* kLogicalAndWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl similarity index 73% rename from backends/webgpu/runtime/ops/logical_and/logical_and.wgsl rename to backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl index 9acb583f51c..fb5310cfde9 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl @@ -16,10 +16,13 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $if OP == "&": + // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $else: + // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. let w = gid.x + gid.y * (num_workgroups.x * wg_size); if (w >= params.num_words) { return; } - t_out[w] = t_a[w] & t_b[w]; + t_out[w] = t_a[w] ${OP} t_b[w]; } diff --git a/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml new file mode 100644 index 00000000000..65df4465b04 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml @@ -0,0 +1,13 @@ +# 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. + +logical_binary: + parameter_names_with_default_values: + OP: "&" + shader_variants: + - NAME: logical_and + - NAME: logical_or + OP: "|" diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h index d64898cb523..e61317d2ff1 100644 --- a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_or.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: 4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf inline constexpr const char* kLogicalOrWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp index d26f6b486d2..d750af11227 100644 --- a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp +++ b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl b/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl deleted file mode 100644 index d7e6176ba32..00000000000 --- a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl +++ /dev/null @@ -1,25 +0,0 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_a: array; -@group(0) @binding(2) var t_b: array; - -struct Params { - num_words: u32, - pad0: u32, - pad1: u32, - pad2: u32, -} -@group(0) @binding(3) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. - let w = gid.x + gid.y * (num_workgroups.x * wg_size); - if (w >= params.num_words) { - return; - } - t_out[w] = t_a[w] | t_b[w]; -} diff --git a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp index 3c150fc2580..f457bc9dc99 100644 --- a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl b/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl deleted file mode 100644 index e79cb2d2bcc..00000000000 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = min(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = min(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index fdb7984b9e1..27bff334dd3 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -10,7 +10,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl deleted file mode 100644 index f82a16e4b21..00000000000 --- a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = input1[idx] * input2[idx]; - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = input1[l1] * input2[l2]; -} diff --git a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp index d2f923880bf..1ac90015ffe 100644 --- a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl deleted file mode 100644 index 2114ef87fee..00000000000 --- a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = pow(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = pow(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index aa8f000bceb..bbd6e13dc9d 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -38,10 +38,10 @@ ) from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_bitwise import ( + BITWISE_NOT_SHAPES, BitwiseAndModule, BitwiseNotModule, bw_gen_a, - bw_gen_b, ) from executorch.backends.webgpu.test.ops.test_cat import ( CatModule, @@ -71,14 +71,13 @@ make_qcs4w_linear_module, ) from executorch.backends.webgpu.test.ops.test_logical_and import ( - la_gen_a, - la_gen_b, + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, LogicalAndModule, ) from executorch.backends.webgpu.test.ops.test_logical_or import ( BitwiseOrModule, - lo_gen_a, - lo_gen_b, LogicalOrModule, ) from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule @@ -253,7 +252,7 @@ def _rms_norm_suite() -> WebGPUTestSuite: @register_op_test("mul") def _mul_suite() -> WebGPUTestSuite: - # Full numeric coverage incl. broadcast (binary_mul.wgsl over a TensorMeta UBO); fp64 golden. + # Full binary_op-family numeric coverage, including broadcast. return WebGPUTestSuite( module_factory=lambda: MulModule(), cases=[ @@ -278,12 +277,19 @@ def _fn_config_suite(module_cls, configs) -> WebGPUTestSuite: @register_op_test("minimum") def _minimum_suite() -> WebGPUTestSuite: - # Same-shape numeric coverage (flat binary kernel; broadcast stays smoke). + # Same-shape and mixed-rank broadcast numeric coverage. return WebGPUTestSuite( module_factory=lambda: MinimumModule(), cases=[ Case(name="2d", inputs=((M1, M2), (M1, M2))), Case(name="3d", inputs=((S, S1, S2), (S, S1, S2))), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-3.0, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 4.0)), + ), + ), ], ) @@ -334,49 +340,33 @@ def _ge_suite() -> WebGPUTestSuite: return _compare_suite("ge") -@register_op_test("logical_and") -def _logical_and_suite() -> WebGPUTestSuite: - # out = (a>0) && (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), AND'd -> bool. Distinct a/b seeds so the masks - # differ (AND ~25% True, a real mix an OR mutant fails); all shapes numel % - # 4 == 0 (bool packs 4/word). float32 oracle (byte-exact bool golden). +def _logical_binary_suite(module_factory) -> WebGPUTestSuite: + # Every packed u32 receives all four Boolean pairs in byte order. def case(name, shape): return Case( name=name, construct={"shape": shape}, inputs=( - InputSpec(shape=shape, gen=la_gen_a), - InputSpec(shape=shape, gen=la_gen_b), + InputSpec(shape=shape, gen=logical_binary_gen_a), + InputSpec(shape=shape, gen=logical_binary_gen_b), ), ) return WebGPUTestSuite( - module_factory=lambda shape: LogicalAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + module_factory=module_factory, + cases=[case(name, shape) for name, shape in LOGICAL_BINARY_CASES], golden_dtype="float32", ) +@register_op_test("logical_and") +def _logical_and_suite() -> WebGPUTestSuite: + return _logical_binary_suite(lambda shape: LogicalAndModule(shape)) + + @register_op_test("bitwise_and") def _bitwise_and_suite() -> WebGPUTestSuite: - # bool bitwise AND == logical_and for canonical 0/1 (shares the handler). - # Two masks derived on-GPU from float inputs via gt.Tensor (baked zeros), - # distinct a/b seeds (AND ~25% True); all shapes numel % 4 == 0. - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=bw_gen_a), - InputSpec(shape=shape, gen=bw_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseAndModule(shape)) @register_op_test("bitwise_not") @@ -392,52 +382,22 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda shape: BitwiseNotModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case(name, shape) + for name, shape in zip(("2d", "3d", "sq"), BITWISE_NOT_SHAPES) + ], golden_dtype="float32", ) @register_op_test("logical_or") def _logical_or_suite() -> WebGPUTestSuite: - # out = (a>0) || (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), OR'd -> bool. Distinct a/b seeds (~50% each, - # independent -> OR ~75% True, a real mix an AND mutant fails); all shapes - # numel % 4 == 0. float32 oracle (byte-exact bool golden). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: LogicalOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: LogicalOrModule(shape)) @register_op_test("bitwise_or") def _bitwise_or_suite() -> WebGPUTestSuite: - # bool bitwise OR == logical_or for canonical 0/1 (shares the handler). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseOrModule(shape)) @register_op_test("pow") @@ -460,6 +420,13 @@ def _pow_suite() -> WebGPUTestSuite: InputSpec(shape=(S, S1, S2), gen=_unary_lin(-2.0, 3.0)), ), ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(0.1, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 3.0)), + ), + ), ], ) @@ -493,6 +460,14 @@ def _floor_divide_suite() -> WebGPUTestSuite: ), golden_fn=_floor_div_golden, ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-8.0, 8.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(0.5, 4.0)), + ), + golden_fn=_floor_div_golden, + ), ], ) diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index fabf79171c0..65f765812be 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -10,7 +10,15 @@ import torch from executorch.backends.webgpu.test.op_tests import generate_op_tests as g -from executorch.backends.webgpu.test.op_tests.test_suite import op_test_registry +from executorch.backends.webgpu.test.op_tests.test_suite import ( + InputSpec, + op_test_registry, +) +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) def _add_regular_case(): @@ -112,3 +120,108 @@ def test_manifest_schema_roundtrip(tmp_path): gd = e["golden"] assert {"path", "shape", "dtype", "output_index"} <= set(gd) assert gd["output_index"] == 0 + + +def test_logical_binary_case_contract(): + expected_cases = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), + ) + assert LOGICAL_BINARY_CASES == expected_cases + assert logical_binary_gen_a((8,)).tolist() == [ + -1.0, + -1.0, + 1.0, + 1.0, + -1.0, + -1.0, + 1.0, + 1.0, + ] + assert logical_binary_gen_b((8,)).tolist() == [ + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + ] + + for op in ("logical_and", "bitwise_and", "logical_or", "bitwise_or"): + suite = op_test_registry[op] + assert tuple((case.name, case.construct["shape"]) for case in suite.cases) == ( + expected_cases + ) + for case in suite.cases: + assert case.required is True + assert case.heavy is False + assert len(case.inputs) == 2 + assert case.inputs[0].gen is logical_binary_gen_a + assert case.inputs[1].gen is logical_binary_gen_b + + +def test_binary_shader_family_case_contract(): + expected = { + "minimum": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "pow": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "floor_divide": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "mul": ( + ("same", ((8, 32), (8, 32))), + ("bcast_lastdim", ((1, 1, 7, 896), (1, 1, 7, 1))), + ("bcast_firstdim", ((4, 4), (1, 4))), + ("bcast_4d_mixed", ((3, 5, 7, 11), (1, 5, 1, 11))), + ("mixedrank", ((4,), (3, 4))), + ), + } + + for op, cases in expected.items(): + suite = op_test_registry[op] + actual = tuple( + ( + case.name, + tuple( + spec.shape if isinstance(spec, InputSpec) else spec + for spec in case.inputs + ), + ) + for case in suite.cases + ) + assert actual == cases + + ranges = { + "minimum": ((-3.0, 3.0), (-2.0, 4.0)), + "pow": ((0.1, 3.0), (-2.0, 3.0)), + "floor_divide": ((-8.0, 8.0), (0.5, 4.0)), + } + for op, expected_ranges in ranges.items(): + case = next( + c for c in op_test_registry[op].cases if c.name == "broadcast_3d_2d" + ) + for spec, (start, end) in zip(case.inputs, expected_ranges): + assert isinstance(spec, InputSpec) and callable(spec.gen) + values = spec.gen(spec.shape).flatten() + assert torch.isclose(values[0], torch.tensor(start)) + assert torch.isclose(values[-1], torch.tensor(end)) + + assert all( + case.golden_fn is g.cases._floor_div_golden + for case in op_test_registry["floor_divide"].cases + ) diff --git a/backends/webgpu/test/ops/test_bitwise.py b/backends/webgpu/test/ops/test_bitwise.py index b61a37d9990..a8fee7a0095 100644 --- a/backends/webgpu/test/ops/test_bitwise.py +++ b/backends/webgpu/test/ops/test_bitwise.py @@ -19,6 +19,11 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -50,11 +55,10 @@ def g(shape): bw_gen_a = _bw_gen(0) -bw_gen_b = _bw_gen(1) # All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +BITWISE_NOT_SHAPES = ((4, 8), (2, 3, 8), (16, 16)) class BitwiseTest(unittest.TestCase): @@ -75,13 +79,17 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = bw_gen_a(shape) - b = bw_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(op="bitwise_and", case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( BitwiseAndModule(shape), (a, b), "bitwise_and", shape ) + + for shape in BITWISE_NOT_SHAPES: + with self.subTest(op="bitwise_not", shape=shape): + a = bw_gen_a(shape) self._assert_delegates( BitwiseNotModule(shape), (a,), "bitwise_not", shape ) diff --git a/backends/webgpu/test/ops/test_floor_divide.py b/backends/webgpu/test/ops/test_floor_divide.py index 6020480ff2c..141e8314cde 100644 --- a/backends/webgpu/test/ops/test_floor_divide.py +++ b/backends/webgpu/test/ops/test_floor_divide.py @@ -6,8 +6,8 @@ """`aten.div.Tensor_mode` module for the WebGPU op-test framework. -`FloorDivideModule` is imported by `cases.py`. Same-shape elementwise -`div(a, b, rounding_mode="floor")`. The kernel computes `floor(a/b)` mirroring +`FloorDivideModule` is imported by `cases.py`. Same-shape and broadcast +`div(a, b, rounding_mode="floor")` use `floor(a/b)`, mirroring the Vulkan `floor_divide` glsl (`floor(X/Y)`); this differs from torch's own fmod-corrected `div_floor` at rare fp boundaries, so the suite goldens against a `floor(a/b)` `golden_fn` (Vulkan-faithful), not this module's eager output. diff --git a/backends/webgpu/test/ops/test_logical_and.py b/backends/webgpu/test/ops/test_logical_and.py index 4c3939002d0..9cf474fd8de 100644 --- a/backends/webgpu/test/ops/test_logical_and.py +++ b/backends/webgpu/test/ops/test_logical_and.py @@ -4,17 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_and.default` module + configs for the WebGPU op-test framework. - -`LogicalAndModule` derives its two bool operands on-GPU from float inputs -(`a > 0`, `b > 0` via the delegated `gt.Tensor` against a baked zero buffer), so -the only runtime inputs are the two float tensors (the op-test framework is -float-input-only). `a`/`b` use distinct seeds so the two bool masks differ (each -~50% True, independent -> AND ~25% True), a real mix that a wrong op (e.g. OR) -would fail. Output is bool (byte-exact golden). `LogicalAndTest` is the -export-delegation smoke test. -""" +"""Delegation coverage for logical AND with packed truth-table inputs.""" +import math import unittest import torch @@ -32,29 +24,40 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.logical_and(a > self.z, b > self.z) -def _la_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) +LOGICAL_BINARY_CASES = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), +) - return g +def _logical_binary_gen(pattern): + def generate(shape): + numel = math.prod(shape) + if numel == 0 or numel % len(pattern) != 0: + raise ValueError("logical-binary test shapes must have numel % 4 == 0") + return ( + torch.tensor(pattern, dtype=torch.float32) + .repeat(numel // len(pattern)) + .reshape(shape) + ) -la_gen_a = _la_gen(0) -la_gen_b = _la_gen(1) + return generate -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +logical_binary_gen_a = _logical_binary_gen((-1.0, -1.0, 1.0, 1.0)) +logical_binary_gen_b = _logical_binary_gen((-1.0, 1.0, -1.0, 1.0)) class LogicalAndTest(unittest.TestCase): def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = la_gen_a(shape) - b = la_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) ep = torch.export.export(LogicalAndModule(shape).eval(), (a, b)) edge = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] diff --git a/backends/webgpu/test/ops/test_logical_or.py b/backends/webgpu/test/ops/test_logical_or.py index d71c1493ea9..420fee26336 100644 --- a/backends/webgpu/test/ops/test_logical_or.py +++ b/backends/webgpu/test/ops/test_logical_or.py @@ -4,23 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_or.default` / `aten.bitwise_or.Tensor` (bool) modules + configs. - -Mirrors the logical_and/bitwise_and tests: the modules derive their two bool -operands on-GPU from float inputs (`a > 0`, `b > 0` via the delegated `gt.Tensor` -against a baked zero buffer), so the only runtime inputs are the two float -tensors (the op-test framework is float-input-only). `a`/`b` use distinct seeds -so the two bool masks differ (each ~50% True, independent -> OR ~75% True), a -real mix that a wrong op (e.g. AND) would fail. `bitwise_or` on bool is identical -to `logical_or` (shares the handler). Output is bool (byte-exact golden). -`LogicalOrTest` is the export-delegation smoke test. -""" +"""Delegation coverage for logical and bitwise OR truth-table inputs.""" import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -42,23 +37,6 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.bitwise_or(a > self.z, b > self.z) -def _lo_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) - - return g - - -lo_gen_a = _lo_gen(0) -lo_gen_b = _lo_gen(1) - - -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] - - class LogicalOrTest(unittest.TestCase): def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ep = torch.export.export(mod.eval(), inputs) @@ -77,10 +55,10 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = lo_gen_a(shape) - b = lo_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( LogicalOrModule(shape), (a, b), "logical_or", shape ) diff --git a/backends/webgpu/test/ops/test_minimum.py b/backends/webgpu/test/ops/test_minimum.py index f8f8088fef5..ba466317d5b 100644 --- a/backends/webgpu/test/ops/test_minimum.py +++ b/backends/webgpu/test/ops/test_minimum.py @@ -6,8 +6,8 @@ """`aten.minimum.default` module for the WebGPU op-test framework. -`MinimumModule` is imported by `cases.py`. minimum is a same-shape elementwise -binary op mirroring the landed `add`/`mul` pattern (flat 2D-dispatch kernel). +`MinimumModule` is imported by `cases.py`. minimum is an elementwise binary op +with the same-shape fast path and broadcast indexing used by `mul`. """ import torch diff --git a/backends/webgpu/test/ops/test_pow.py b/backends/webgpu/test/ops/test_pow.py index 93596dc5d9f..8498bd4f7aa 100644 --- a/backends/webgpu/test/ops/test_pow.py +++ b/backends/webgpu/test/ops/test_pow.py @@ -6,8 +6,8 @@ """`aten.pow.Tensor_Tensor` module for the WebGPU op-test framework. -`PowModule` is imported by `cases.py`. Same-shape elementwise `pow(a, b)`; the -suite uses a POSITIVE base so `pow(neg, frac)` (NaN) is never exercised. +`PowModule` is imported by `cases.py`. The suite covers same-shape and broadcast +`pow(a, b)` with POSITIVE bases so `pow(neg, frac)` never produces NaN. """ import torch diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index d437a0a42e8..574c3869864 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -213,7 +213,11 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 134) self.assertEqual( digest.hexdigest(), - "ef97dca2336315ee2c8b0f9e896c6aa082834ae94948bda3ccb42b1145f2bd27", + "e502196846f0f8100f468e5d9f8f9c006b67e08df54e1e2e667daa2fc50d8844", + ) + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "492535b396833ad6ebfed29b093057e38f40b1182b5c7d6b3eb2f3577cab024e", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -950,10 +954,6 @@ def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: entries["to_copy_int_to_float"].include, "runtime/ops/to_copy/to_copy_int_to_float_wgsl.h", ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) def test_extrema_template_roundtrip_byte_identical(self) -> None: extrema_dir = g.BACKEND_ROOT / "runtime/ops/extrema" @@ -1004,10 +1004,108 @@ def test_extrema_template_roundtrip_byte_identical(self) -> None: hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) + def test_logical_binary_template_roundtrip_byte_identical(self) -> None: + logical_dir = g.BACKEND_ROOT / "runtime/ops/logical_binary" + template_path = logical_dir / "logical_binary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "logical_and": ( + "&", + "cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20", + ), + "logical_or": ( + "|", + "4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (op, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["OP"], op) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = logical_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/logical_binary/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "logical_and": "eb85a8f97ee7640298a661da49feb08aa79b8c24d3d4458b71d24d3f01bc388d", + "logical_or": "bda18617f7077fee5a812c21cdc495c89542a1688f7e1ef6739ed01da343a66b", + } + for name, expected_hash in handler_hashes.items(): + handler = ( + g.BACKEND_ROOT / f"runtime/ops/{name}/Logical{name[8:].title()}.cpp" + ) + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) + + def test_binary_family_roundtrip_byte_identical(self) -> None: + binary_dir = g.BACKEND_ROOT / "runtime/ops/binary_op" + template_path = binary_dir / "binary_op.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "binary_div": ( + 0, + "e36b560fd623dd5337b9ae57acd8981c9c635b995d6021caf1331c182cd3f0cd", + ), + "binary_sub": ( + 0, + "63209ff70422a21fc340d9aadba0945bc259bba89bdf05db018a6507d01c7ae5", + ), + "binary_minimum": ( + 1, + "929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972", + ), + "binary_pow": ( + 1, + "a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c", + ), + "binary_floor_divide": ( + 1, + "baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3", + ), + "binary_mul": ( + 1, + "d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + entries = {entry.name: entry for entry in g.registry_entries()} + + for name, (inline, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["INLINE"], inline) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header = (binary_dir / f"{name}_wgsl.h").read_text() + literal = header.split('R"(', 1)[1].split(')";', 1)[0] + self.assertEqual(literal, "\n" + expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(expanded), (64, 1, 1)) + self.assertIn(f"k{g.symbol_base(name)}WGSL", header) + self.assertEqual( + entries[name].include, + f"runtime/ops/binary_op/{name}_wgsl.h", + ) def test_unary_template_roundtrip_byte_identical(self) -> None: unary_dir = g.BACKEND_ROOT / "runtime/ops/unary" From 07b93e5e6c8d9598a1407e43891eb233cdc6e0a0 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:41 -0700 Subject: [PATCH 11/12] [ExecuTorch][WebGPU] Preserve Linear and Q4 embedding params across dynamic resize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21483 **Preserve immutable Linear and Q4 embedding shader controls across dynamic resize through typed parameter authorities.** **Problem** The Linear resize hook dropped has_bias, and the Q4 embedding resize path could drop is_linear_weight. Both defects produced correct build-time outputs but silently changed semantics after a live shape update. **Solution** - Before: build and resize populated control words independently, allowing immutable shader state to reset. - After: each build/resize pair calls one typed helper while recomputing only live counts and dispatch. **Implementation** - Linear.cpp — make_linear_params preserves the bias flag across vec4 and tiled routes. - EmbeddingQ4gsw.cpp — EmbeddingLayout and make_embedding_params preserve nibble layout across resize. - Mirrors Vulkan runtime/graph/ops/impl/Linear.cpp and EmbeddingQ4gsw.cpp, which carry immutable shader controls through dynamic dispatch. **Constraints** WGSL bytes, shader selection, bindings, uniform sizes, queue-write counts, dispatch formulas, and pipeline topology are unchanged. Production handler code is smaller because duplicated field population and scalar captures are removed. Co-authored-with: Claude Code. ghstack-source-id: 411961489 @exported-using-ghexport Differential Revision: [D113992326](https://our.internmc.facebook.com/intern/diff/D113992326/) --- .../ops/embedding_q4gsw/EmbeddingQ4gsw.cpp | 58 +++++--- backends/webgpu/runtime/ops/linear/Linear.cpp | 18 +-- .../webgpu/test/native/test_dynamic_shape.cpp | 40 ++++- .../test_dynamic_shape_export.py | 140 +++++++++++++----- 4 files changed, 183 insertions(+), 73 deletions(-) diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index 64956121b17..b5ea0cdca4c 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -37,13 +37,37 @@ static_assert( sizeof(EmbeddingParams) == 32, "EmbeddingParams must be 32 bytes"); +struct EmbeddingLayout { + uint32_t embed_dim; + uint32_t blocks_per_row; + uint32_t group_size; + uint32_t groups_per_row; + uint32_t bytes_per_row; + bool is_linear_weight; +}; + +EmbeddingParams make_embedding_params( + const EmbeddingLayout& layout, + uint32_t num_indices, + uint32_t total_blocks) { + return { + layout.embed_dim, + layout.blocks_per_row, + num_indices, + layout.group_size, + layout.groups_per_row, + layout.bytes_per_row, + total_blocks, + layout.is_linear_weight ? 1u : 0u}; +} + // Resize hook body: recompute counts/dispatch; out = indices dims + // [embed_dim]. void resize_embedding_q4gsw( WebGPUGraph& g, int indices_id, int out_id, - EmbeddingParams params, + const EmbeddingLayout& layout, uint32_t wg_size, size_t dispatch_idx, WGPUBuffer params_buf) { @@ -52,17 +76,17 @@ void resize_embedding_q4gsw( if (ni == 0) { throw std::runtime_error("WebGPU embedding_q4gsw: zero indices"); } - const uint64_t total_blocks = ni * params.blocks_per_row; + const uint64_t total_blocks = ni * layout.blocks_per_row; if (total_blocks > UINT32_MAX) { throw std::runtime_error( "WebGPU embedding_q4gsw: total_blocks exceeds uint32"); } std::vector od = id; - od.push_back(static_cast(params.embed_dim)); + od.push_back(static_cast(layout.embed_dim)); g.set_cur_dims(out_id, od); - params.num_indices = static_cast(ni); - params.total_blocks = static_cast(total_blocks); - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, ¶ms, sizeof(params)); + EmbeddingParams p = make_embedding_params( + layout, static_cast(ni), static_cast(total_blocks)); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); g.dispatch_at(dispatch_idx).workgroup_count_x = utils::compute_1d_workgroup_count( g.device(), @@ -166,15 +190,15 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t workgroup_count = utils::compute_1d_workgroup_count( device, static_cast(total_blocks), wg_size, "embedding_q4gsw"); - EmbeddingParams params = {}; - params.embed_dim = embed_dim; - params.blocks_per_row = blocks_per_row; - params.num_indices = num_indices; // std140 layout only; shader derives it - params.group_size = static_cast(group_size); - params.groups_per_row = groups_per_row; - params.bytes_per_row = bytes_per_row; - params.total_blocks = static_cast(total_blocks); - params.is_linear_weight = is_linear ? 1u : 0u; + const EmbeddingLayout layout = { + embed_dim, + blocks_per_row, + static_cast(group_size), + groups_per_row, + bytes_per_row, + is_linear}; + EmbeddingParams params = make_embedding_params( + layout, num_indices, static_cast(total_blocks)); WGPUBufferDescriptor uniform_desc = {}; uniform_desc.size = sizeof(EmbeddingParams); @@ -223,10 +247,10 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( indices_id, - [indices_id, out_id, params, wg_size, dispatch_idx, params_buf]( + [indices_id, out_id, layout, wg_size, dispatch_idx, params_buf]( WebGPUGraph& g) { resize_embedding_q4gsw( - g, indices_id, out_id, params, wg_size, dispatch_idx, params_buf); + g, indices_id, out_id, layout, wg_size, dispatch_idx, params_buf); }); // Graph owns it so the resize hook can rewrite it; freed in the dtor. diff --git a/backends/webgpu/runtime/ops/linear/Linear.cpp b/backends/webgpu/runtime/ops/linear/Linear.cpp index 45cc27a515d..ba82fcc36a4 100644 --- a/backends/webgpu/runtime/ops/linear/Linear.cpp +++ b/backends/webgpu/runtime/ops/linear/Linear.cpp @@ -33,6 +33,11 @@ static_assert(sizeof(LinearParams) == 16, "LinearParams must be 16 bytes"); constexpr uint32_t kTile = 32u; +LinearParams +make_linear_params(uint32_t M, uint32_t N, uint32_t K, bool has_bias) { + return {M, N, K, has_bias ? 1u : 0u}; +} + // aten.linear (+ optional bias); shared-memory tiled GEMM. void linear_impl(WebGPUGraph& graph, const std::vector& args) { // args: [input, weight, bias?, out]; out is last. bias (arg 2) is a tensor @@ -82,11 +87,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { } } - LinearParams params = {}; - params.M = M; - params.N = N; - params.K = K; - params.has_bias = has_bias ? 1u : 0u; + LinearParams params = make_linear_params(M, N, K, has_bias); // Bias binding (binding 4); a 4-byte dummy satisfies it when None // (WGSL-gated). @@ -138,7 +139,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( in_id, - [in_id, out_id, M, N, K, dispatch_x, dispatch_idx, params_buf]( + [in_id, out_id, M, N, K, has_bias, dispatch_x, dispatch_idx, params_buf]( WebGPUGraph& g) { const auto& d = g.cur_dims(in_id); const uint64_t numel = utils::numel_of(d); @@ -152,10 +153,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "WebGPU linear: live M is 0 or exceeds the build-time max"); } - LinearParams p = {}; - p.M = m; - p.N = N; - p.K = K; + LinearParams p = make_linear_params(m, N, K, has_bias); wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); g.dispatch_at(dispatch_idx).workgroup_count_x = dispatch_x; g.dispatch_at(dispatch_idx).workgroup_count_y = diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 0e72ce6234d..7b3c1d6b0a7 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -133,6 +133,7 @@ constexpr int kLinK = 64; constexpr int kLinAltK = 72; constexpr int kLinN = 128; constexpr int kLinNShmem = 2048; +constexpr int kFp32LinearN = 32; // Run at [m_rows, kLinK] on an already-loaded module (so it can be // reused across M without a fresh load), and compare to the golden. void run_linear( @@ -233,6 +234,14 @@ void check_linear_tiled(int m_rows) { run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); } +void check_fp32_linear_reused(const char* prefix, int k) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load " << prefix << ".pte"; + for (int m_rows : {128, 32, 1, 128}) { + run_linear(module, m_rows, prefix, kFp32LinearN, k, 1e-3f); + } +} + constexpr int kQkvNq = 2048; constexpr int kQkvNk = 512; constexpr int kQkvNv = 512; @@ -941,6 +950,23 @@ TEST(DynamicShape, RmsMul) { } } +// I0: dynamic fp32 linear preserves bias across repeated resizes. +TEST(DynamicShape, Fp32LinearVec4BiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_vec4_bias", 64); +} + +TEST(DynamicShape, Fp32LinearVec4UnbiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_vec4_no_bias", 64); +} + +TEST(DynamicShape, Fp32LinearTiledBiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_tiled_bias", 63); +} + +TEST(DynamicShape, Fp32LinearTiledUnbiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_tiled_no_bias", 63); +} + // I: dynamic 4-bit quantized linear (prefill GEMM) at several M. TEST(DynamicShape, QuantizedLinear) { for (int m_rows : {128, 32, 1}) { @@ -1670,12 +1696,14 @@ TEST(DynamicShape, EmbeddingReusedGraph) { } } -// K3: linear-packed reuse must preserve nibble order across resizes. -TEST(DynamicShape, LinearPackedEmbeddingReusedGraph) { - Module m(g_dir + "/emb_dyn_linear.pte"); - ASSERT_EQ(m.load_forward(), Error::Ok) << "load emb_dyn_linear.pte"; - for (int n : {16, 8, 1, 16}) { - run_embedding(m, n, "emb_dyn_linear"); +// K3: linear/nonlinear-packed reuse must preserve nibble order across resizes. +TEST(DynamicShape, EmbeddingLayoutsReusedGraph) { + for (const char* prefix : {"emb_dyn_linear", "emb_dyn_nonlinear"}) { + Module m(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load " << prefix << ".pte"; + for (int n : {16, 8, 1, 16}) { + run_embedding(m, n, prefix); + } } } 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 a99c9baf46e..b302b3c120d 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 @@ -297,7 +297,10 @@ def export_dynamic_shape_cases(out_dir: str) -> None: ) _write_goldens(rmsmul, "dyn_rmsmul", out_dir, [MAXS, 32, 1]) - # 2d) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM + # 2d) fp32 linear with a dynamic rows (M) dim. + export_dynamic_fp32_linear_cases(out_dir) + + # 2d.1) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM # (register-tiled N=128) + a shmem-GEMM-routed variant (N=2048). _export_dynamic_linear(out_dir) _export_dynamic_linear( @@ -327,7 +330,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: _export_static_sdpa(out_dir, 16, "static_sdpa_s16") # 2f) 4-bit embedding with a DYNAMIC token count (int64 indices). - _export_dynamic_embedding(out_dir) + export_dynamic_embedding_cases(out_dir) # 2g) Interleaved RoPE with a DYNAMIC seq-len S (two outputs xq/xk). _export_dynamic_rope(out_dir) @@ -370,6 +373,10 @@ def export_dynamic_shape_cases(out_dir: str) -> None: LIN_GROUP = 32 LIN_MAXM = 128 +FP32_LINEAR_N = 32 +FP32_LINEAR_MAXM = 128 +FP32_LINEAR_M_VALUES = (FP32_LINEAR_MAXM, 32, 1) + BK64_K = 2048 BK64_N = 2048 BK64_KV_N = 512 @@ -533,6 +540,46 @@ def _export_dynamic_linear( print(f" golden {prefix} M={m}") +def _export_dynamic_fp32_linear_case( + out_dir: str, + *, + k: int, + bias: bool, + prefix: str, +) -> None: + from executorch.backends.webgpu.test.ops.test_linear_fp32 import make_linear + + model = make_linear(k, FP32_LINEAR_N, bias=bias).eval() + x = _ramp((FP32_LINEAR_MAXM, k)) + m_dim = torch.export.Dim("m", min=1, max=FP32_LINEAR_MAXM) + ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + + weight = model.fc.weight.detach().double() + bias_value = model.fc.bias.detach().double() if model.fc.bias is not None else None + for m in FP32_LINEAR_M_VALUES: + xm = _ramp((m, k)) + golden = torch.nn.functional.linear(xm.double(), weight, bias_value) + base = os.path.join(out_dir, f"{prefix}.S{m}") + xm.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) + for route, k in (("vec4", 64), ("tiled", 63)): + for bias in (True, False): + suffix = "bias" if bias else "no_bias" + _export_dynamic_fp32_linear_case( + out_dir, + k=k, + bias=bias, + prefix=f"dyn_linear_fp32_{route}_{suffix}", + ) + + def _make_bk64_model( *, k: int = BK64_K, @@ -1306,52 +1353,64 @@ def _export_static_sdpa(out_dir: str, s: int, prefix: str) -> None: EMB_MAXN = 16 -class _LinearPackedEmbedding(torch.nn.Module): - def __init__(self) -> None: +def _write_embedding_goldens( + out_dir: str, + prefix: str, + weight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + is_linear: bool, +) -> None: + for n in [EMB_MAXN, 8, 1]: + idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB + golden = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, is_linear + ) + idx.detach().numpy().astype(" None: super().__init__() packed = torch.arange(EMB_VOCAB * (EMB_DIM // 2), dtype=torch.int64).reshape( EMB_VOCAB, EMB_DIM // 2 ) self.register_buffer("weight", (packed % 256).to(torch.uint8)) self.register_buffer("scales", torch.ones(EMB_VOCAB, EMB_DIM // EMB_GROUP)) + self.is_linear_weight = is_linear_weight def forward(self, indices: torch.Tensor) -> torch.Tensor: return torch.ops.et_vk.embedding_q4gsw.default( - self.weight, self.scales, EMB_GROUP, indices, True + self.weight, + self.scales, + EMB_GROUP, + indices, + self.is_linear_weight, ) -def _write_embedding_goldens( - out_dir: str, - prefix: str, - weight: torch.Tensor, - scales: torch.Tensor, - group_size: int, - is_linear: bool, +def _write_packed_embedding_goldens( + out_dir: str, prefix: str, model: _PackedEmbedding ) -> None: for n in [EMB_MAXN, 8, 1]: idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB - golden = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, is_linear - ) - if is_linear: - nonlinear_golden = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, False - ) - if torch.equal(golden, nonlinear_golden): - raise RuntimeError( - "emb_dyn_linear fixture does not distinguish nibble packing" - ) + golden = model(idx) idx.detach().numpy().astype(" None: +def export_dynamic_embedding_cases(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import ( _make_quantized_model, _quant_params, @@ -1370,21 +1429,22 @@ def _export_dynamic_embedding(out_dir: str) -> None: weight, scales, group_size = _quant_params(qm) _write_embedding_goldens(out_dir, "emb_dyn", weight, scales, group_size, False) - linear_model = _LinearPackedEmbedding().eval() - _export( - linear_model, - (idx_max,), - ({0: n_dim},), - os.path.join(out_dir, "emb_dyn_linear.pte"), - ) - _write_embedding_goldens( - out_dir, - "emb_dyn_linear", - linear_model.weight, - linear_model.scales, - EMB_GROUP, - True, - ) + packed_models = { + "emb_dyn_linear": _PackedEmbedding(True).eval(), + "emb_dyn_nonlinear": _PackedEmbedding(False).eval(), + } + linear_golden = packed_models["emb_dyn_linear"](idx_max) + nonlinear_golden = packed_models["emb_dyn_nonlinear"](idx_max) + if torch.equal(linear_golden, nonlinear_golden): + raise RuntimeError("embedding layout fixtures must distinguish nibble order") + for prefix, model in packed_models.items(): + _export( + model, + (idx_max,), + ({0: n_dim},), + os.path.join(out_dir, f"{prefix}.pte"), + ) + _write_packed_embedding_goldens(out_dir, prefix, model) # Dynamic RoPE: xq/xk + freqs all share a dynamic seq-len S. From 509ea9534e197da89f434d79a69549450f3d5cca Mon Sep 17 00:00:00 2001 From: pytorchbot Date: Fri, 7 Aug 2026 10:47:11 -0700 Subject: [PATCH 12/12] [ExecuTorch][Vulkan] Preserve persistent buffer mutations (#21662) This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: https://github.com/pytorch/executorch/pull/21597 by @JCNTH ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/204/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/204/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/203/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/204/orig Differential Revision: [D114936148](https://our.internmc.facebook.com/intern/diff/D114936148/) @diff-train-skip-merge cc @SS-JIA @manuelcandales @digantdesai @cbilgin --------- Co-authored-by: Julian Ng-Thow-Hing Co-authored-by: Julian Ng-Thow-Hing <107437036+JCNTH@users.noreply.github.com> --- .../serialization/vulkan_graph_builder.py | 48 +++- backends/vulkan/test/test_serialization.py | 138 +++++++++- .../test/test_vulkan_compile_options.py | 25 +- backends/vulkan/vulkan_preprocess.py | 5 + backends/webgpu/runtime/WebGPUDispatchMath.h | 2 +- backends/webgpu/runtime/WebGPUGraph.cpp | 98 +++++--- backends/webgpu/runtime/WebGPUGraph.h | 2 + .../webgpu/runtime/WebGPUShaderRegistry.cpp | 18 +- .../webgpu/runtime/ops/compare/Compare.cpp | 25 +- .../webgpu/runtime/ops/compare/compare.wgsl | 30 +-- .../webgpu/runtime/ops/compare/compare_wgsl.h | 32 +-- .../webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp | 237 ++++++++++++++++-- .../webgpu/runtime/ops/conv1d_dw/conv1d.wgsl | 53 ++++ .../runtime/ops/conv1d_dw/conv1d_wgsl.h | 77 ++++++ .../runtime/ops/expand_copy/ExpandCopy.cpp | 38 ++- .../runtime/ops/expand_copy/expand_copy.wgsl | 6 +- .../ops/expand_copy/expand_copy_wgsl.h | 8 +- backends/webgpu/runtime/ops/gelu/Gelu.cpp | 52 +++- backends/webgpu/runtime/ops/gelu/gelu.wgsl | 14 +- backends/webgpu/runtime/ops/gelu/gelu_wgsl.h | 16 +- .../webgpu/runtime/ops/to_copy/ToCopy.cpp | 196 ++++++++++++++- .../ops/to_copy/to_copy_bool_to_float.wgsl | 24 ++ .../ops/to_copy/to_copy_bool_to_float_wgsl.h | 48 ++++ .../test/native/test_compute_dispatch.cpp | 170 ++++++++++--- .../webgpu/test/native/test_dynamic_shape.cpp | 108 ++++++++ .../webgpu/test/native/test_webgpu_utils.cpp | 8 + backends/webgpu/test/op_tests/cases.py | 140 ++++++++++- .../webgpu/test/op_tests/generate_op_tests.py | 15 +- .../webgpu/test/op_tests/op_test_driver.cpp | 16 +- .../webgpu/test/op_tests/test_generator.py | 21 ++ backends/webgpu/test/op_tests/test_suite.py | 3 + .../test_dynamic_shape_export.py | 89 +++++++ backends/webgpu/test/ops/test_conv1d_pw.py | 82 +++++- backends/webgpu/test/ops/test_to_copy.py | 60 ++++- backends/webgpu/test/test_wgsl_codegen.py | 16 +- 35 files changed, 1720 insertions(+), 200 deletions(-) create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index ca5ab196dd2..46e01e701b1 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -9,7 +9,7 @@ import logging import operator from types import NoneType -from typing import cast, List, Optional, Union +from typing import cast, Dict, List, Optional, Union import executorch.backends.vulkan.serialization.vulkan_graph_schema as vk_graph_schema import torch @@ -28,6 +28,7 @@ ) from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.utils import DelegateMappingBuilder +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.tensor import TensorSpec from torch._export.utils import get_buffer, get_param, is_buffer, is_param from torch.export import ExportedProgram @@ -49,11 +50,41 @@ def __init__( delegate_mapping_builder: DelegateMappingBuilder, downcast_64_bit: bool = True, force_fp16: bool = False, + alias_buffer_mutations: bool = False, ) -> None: self.program = program self.delegate_mapping_builder = delegate_mapping_builder self.downcast_64_bit = downcast_64_bit self.force_fp16 = force_fp16 + self.buffer_mutation_inputs: Dict[str, Node] = {} + self.buffer_mutation_user_outputs: set[str] = set() + if alias_buffer_mutations: + nodes_by_name = { + node.name: node for node in program.graph_module.graph.nodes + } + buffer_inputs_by_target: Dict[str, Node] = {} + for name, target in program.graph_signature.inputs_to_buffers.items(): + if name not in nodes_by_name: + continue + buffer_input = nodes_by_name[name] + prepack = next( + ( + user + for user in buffer_input.users + if user.op == "call_function" + and user.target == exir_ops.edge.et_vk.prepack.default + ), + None, + ) + buffer_inputs_by_target[target] = prepack or buffer_input + self.buffer_mutation_inputs = { + output_name: buffer_inputs_by_target[target] + for output_name, target in program.graph_signature.buffers_to_mutate.items() + if target in buffer_inputs_by_target + } + self.buffer_mutation_user_outputs = set( + program.graph_signature.user_outputs + ) self.chain = [] self.values = [] self.input_ids = [] @@ -160,6 +191,16 @@ def maybe_add_constant_tensor(self, node: Node) -> int: return constant_id def create_node_value(self, node: Node) -> int: + if node.name in self.buffer_mutation_inputs: + input_node = self.buffer_mutation_inputs[node.name] + if input_node not in self.node_to_value_ids: + raise AssertionError( + "Cannot alias a buffer mutation before its input is serialized" + ) + value_id = self.node_to_value_ids[input_node] + self.node_to_value_ids[node] = value_id + return value_id + # If the node has been marked as a scalar tensor, create a SymInt instead of a tensor if is_symint_node(node) or node.meta.get("etvk_is_scalar_tensor", False): new_id = self.create_symint_value() @@ -448,7 +489,10 @@ def process_output_node(self, node: Node) -> None: ) # Mutable buffers outputs are not included as an output to the # delegate call. Skip marking them as an output. - if is_mutable_buffer_node(out_node, self.program): + if out_node.name in self.buffer_mutation_inputs: + if out_node.name not in self.buffer_mutation_user_outputs: + continue + elif is_mutable_buffer_node(out_node, self.program): continue self.output_ids.append(self.node_to_value_ids[out_node]) diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index c373f5216d2..71a6980635a 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -9,9 +9,14 @@ import ctypes import random import unittest -from typing import List +from types import SimpleNamespace +from typing import List, Tuple +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 import torch +from executorch.backends.vulkan.serialization import ( + vulkan_graph_builder as graph_builder_module, +) from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( IntList, @@ -30,6 +35,137 @@ class TestSerialization(unittest.TestCase): + def _build_mutation_program( + self, prepack: bool, shared_user_output: bool = False + ) -> Tuple[SimpleNamespace, torch.fx.Node, torch.fx.Node, torch.fx.Node]: + graph = torch.fx.Graph() + state = graph.placeholder("state") + user_input = graph.placeholder("user_input") + state.meta["spec"] = graph_builder_module.TensorSpec.from_tensor(torch.zeros(4)) + user_input.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + + state_value = state + if prepack: + state_value = graph.call_function( + graph_builder_module.exir_ops.edge.et_vk.prepack.default, + (state,), + ) + state_value.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.zeros(4) + ) + + mutation = graph.call_function( + torch.ops.aten.add.Tensor, (state_value, user_input) + ) + mutation.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + user_output = mutation + if not shared_user_output: + user_output = graph.call_function( + torch.ops.aten.mul.Tensor, (user_input, 2.0) + ) + user_output.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + graph.output((mutation, user_output)) + + graph_module = torch.fx.GraphModule({}, graph) + signature = SimpleNamespace( + buffers_to_mutate={mutation.name: "state"}, + inputs_to_buffers={state.name: "state"}, + inputs_to_lifted_tensor_constants={}, + inputs_to_parameters={}, + non_persistent_buffers=set(), + user_outputs=(user_output.name,), + ) + program = SimpleNamespace( + constants={}, + graph_module=graph_module, + graph_signature=signature, + state_dict={"state": torch.zeros(4)}, + ) + return program, state_value, mutation, user_output + + def test_alias_buffer_mutations_is_opt_in(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, user_output = ( + self._build_mutation_program(prepack) + ) + + default_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + ) + default_graph = default_builder.build_graph() + self.assertNotEqual( + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + default_graph.output_ids, + [ + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[user_output], + ], + ) + + explicit_false_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=False, + ) + self.assertEqual(default_graph, explicit_false_builder.build_graph()) + + aliasing_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + aliasing_graph = aliasing_builder.build_graph() + self.assertEqual( + aliasing_builder.node_to_value_ids[mutation], + aliasing_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + aliasing_graph.output_ids, + [aliasing_builder.node_to_value_ids[user_output]], + ) + + def test_alias_buffer_mutations_preserves_shared_user_output(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, _ = self._build_mutation_program( + prepack, shared_user_output=True + ) + builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + + graph = builder.build_graph() + + self.assertEqual( + builder.node_to_value_ids[mutation], + builder.node_to_value_ids[state_value], + ) + self.assertEqual( + graph.output_ids, + [builder.node_to_value_ids[mutation]], + ) + def _generate_random_const_tensors(self, num_tensors: int) -> List[torch.Tensor]: """ Helper function to generate `num_tensor` buffers of random sizes and random contents, diff --git a/backends/vulkan/test/test_vulkan_compile_options.py b/backends/vulkan/test/test_vulkan_compile_options.py index f45cfdf12d9..ff507bb3197 100644 --- a/backends/vulkan/test/test_vulkan_compile_options.py +++ b/backends/vulkan/test/test_vulkan_compile_options.py @@ -41,6 +41,10 @@ def test_skip_memory_planning_round_trips(self) -> None: round_tripped = self._round_trip({"skip_memory_planning": True}) self.assertTrue(round_tripped.get("skip_memory_planning")) + def test_alias_buffer_mutations_round_trips(self) -> None: + round_tripped = self._round_trip({"alias_buffer_mutations": True}) + self.assertTrue(round_tripped.get("alias_buffer_mutations")) + def test_force_fp16_round_trips(self) -> None: round_tripped = self._round_trip({"force_fp16": True}) self.assertTrue(round_tripped.get("force_fp16")) @@ -105,15 +109,15 @@ def build_graph(): ), patch( "executorch.backends.vulkan.vulkan_preprocess.VkGraphBuilder", return_value=graph_builder, - ), patch( + ) as graph_builder_factory, patch( "executorch.backends.vulkan.vulkan_preprocess.serialize_vulkan_graph", return_value=b"vk_graph", ): result = VulkanBackend.preprocess(program, parse_compile_options(options)) - return result.data_store_output, externalize_pte_data + return result.data_store_output, externalize_pte_data, graph_builder_factory def test_external_constants_default_keeps_constants_inline(self) -> None: - output, externalize_pte_data = self._preprocess_named_data({}) + output, externalize_pte_data, _ = self._preprocess_named_data({}) self.assertEqual(output.buffers, [b"constant"]) self.assertEqual(output.pte_data, {"constant": DataEntry(0, 16, None)}) @@ -121,7 +125,7 @@ def test_external_constants_default_keeps_constants_inline(self) -> None: externalize_pte_data.assert_not_called() def test_external_constants_option_externalizes_constants(self) -> None: - output, externalize_pte_data = self._preprocess_named_data( + output, externalize_pte_data, _ = self._preprocess_named_data( {"external_constants_max_data_bytes": 16} ) @@ -131,8 +135,21 @@ def test_external_constants_option_externalizes_constants(self) -> None: self.assertEqual(list(next(iter(output.external_data.values()))), ["constant"]) externalize_pte_data.assert_called_once_with(16, "vulkan_constants") + def test_alias_buffer_mutations_reaches_graph_builder(self) -> None: + for options, expected in ( + ({}, False), + ({"alias_buffer_mutations": True}, True), + ): + with self.subTest(options=options): + _, _, graph_builder_factory = self._preprocess_named_data(options) + self.assertIs( + graph_builder_factory.call_args.kwargs["alias_buffer_mutations"], + expected, + ) + def test_unset_options_are_absent(self) -> None: round_tripped = self._round_trip({}) + self.assertNotIn("alias_buffer_mutations", round_tripped) self.assertNotIn("small_texture_limits", round_tripped) self.assertNotIn("skip_memory_planning", round_tripped) self.assertNotIn("external_constants_max_data_bytes", round_tripped) diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index d3954240880..f7d6955ce26 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -130,6 +130,9 @@ def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: if spec.key == "skip_memory_planning": options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "alias_buffer_mutations": + options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "external_constants_max_data_bytes": options[spec.key] = _parse_external_constants_max_data_bytes(spec.value) @@ -172,6 +175,7 @@ def preprocess( # noqa: C901 ) downcast_64_bit = compile_options.get("downcast_64_bit", True) force_fp16 = compile_options.get("force_fp16", False) + alias_buffer_mutations = compile_options.get("alias_buffer_mutations", False) program = unsafe_remove_auto_functionalized_pass(program) @@ -258,6 +262,7 @@ def preprocess( # noqa: C901 DelegateMappingBuilder(generated_identifiers=True), downcast_64_bit=downcast_64_bit, force_fp16=force_fp16, + alias_buffer_mutations=alias_buffer_mutations, ) vk_graph = graph_builder.build_graph() external_constants_max_data_bytes = compile_options.get( diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 561bebc6b87..60638b499bb 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -24,7 +24,7 @@ namespace executorch::backends::webgpu::utils { // Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). template inline T div_up(T a, T b) { - return (a + b - 1) / b; + return a / b + (a % b != 0); } // Product of a tensor's dims; the same accumulation was duplicated per-op. diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 58d7178ddbf..d1e1d625ad7 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -212,6 +212,31 @@ bool vk_datatype_is_int(vkgraph::VkDataType dtype) { } } +size_t storage_buffer_size(size_t nbytes) { + const size_t at_least_four = std::max(nbytes, size_t(4)); + if (at_least_four > std::numeric_limits::max() - 3u) { + throw std::runtime_error("WebGPU: storage buffer size overflows alignment"); + } + return (at_least_four + 3u) & ~size_t(3); +} + +void write_storage_buffer( + WGPUQueue queue, + WGPUBuffer buffer, + const void* data, + size_t nbytes) { + if (nbytes == 0u) { + return; + } + if (nbytes % 4u == 0u) { + wgpuQueueWriteBuffer(queue, buffer, 0, data, nbytes); + return; + } + std::vector padded(storage_buffer_size(nbytes), 0u); + std::memcpy(padded.data(), data, nbytes); + wgpuQueueWriteBuffer(queue, buffer, 0, padded.data(), padded.size()); +} + // Normalize a possibly-negative dim against rank; throws (fail-loud) if OOR. int normalize_dim(int dim, int rank, const char* op) { if (dim < 0) { @@ -276,7 +301,7 @@ WebGPUGraph::WebGPUGraph() = default; WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = nbytes > 0 ? nbytes : 4; + buf_desc.size = storage_buffer_size(nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -286,7 +311,7 @@ WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { } WGPUBuffer WebGPUGraph::acquire_scratch(size_t nbytes) { - nbytes = nbytes > 0 ? nbytes : 4; + nbytes = storage_buffer_size(nbytes); // Best-fit reuse: smallest free slot with size in [nbytes, 2*nbytes] -- the // 2x cap stops a large Cmax-sized buffer from backing a tiny request. Never // reuse an in_use slot (co-live safety). @@ -890,6 +915,7 @@ void WebGPUGraph::build( throw std::runtime_error("WebGPU: tensor byte size overflows"); } tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); + tensor.is_bool = vk_tensor->datatype() == vkgraph::VkDataType::BOOL; tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8; tensor.nbytes = numel * tensor.elem_size; // Live dims start == max (serialized upper bound); resize_input shrinks @@ -911,7 +937,7 @@ void WebGPUGraph::build( tensor.cur_nbytes = tensor.nbytes; tensor_mem_obj_ids_[i] = -1; WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -936,10 +962,9 @@ void WebGPUGraph::build( std::memcpy(&value, src + e * sizeof(float), sizeof(float)); converted[e] = executorch::runtime::etensor::Half(value); } - wgpuQueueWriteBuffer( + write_storage_buffer( queue_, tensor.buffer, - 0, converted.data(), converted.size() * sizeof(converted[0])); }; @@ -1013,7 +1038,7 @@ void WebGPUGraph::build( prepack_src_ids.count(i) != 0 && direct_use_ids.count(i) == 0; if (!defer) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -1110,7 +1135,7 @@ void WebGPUGraph::build( shared_buffers_.resize(shared_buffer_sizes_.size(), nullptr); for (size_t id = 0; id < shared_buffer_sizes_.size(); id++) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(shared_buffer_sizes_[id], size_t(4)); + buf_desc.size = storage_buffer_size(shared_buffer_sizes_[id]); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -1138,7 +1163,7 @@ void WebGPUGraph::build( // Create staging buffer for output readback WGPUBufferDescriptor staging_desc = {}; - staging_desc.size = std::max(tensors_[oid].nbytes, size_t(4)); + staging_desc.size = storage_buffer_size(tensors_[oid].nbytes); staging_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; staging_desc.mappedAtCreation = false; output_staging_buffers_.push_back( @@ -1313,7 +1338,7 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { cs.nbytes, "WebGPU: inline constant exceeds constant data"); if (cs.nbytes != 0) { - wgpuQueueWriteBuffer(queue_, dst, 0, data, cs.nbytes); + write_storage_buffer(queue_, dst, data, cs.nbytes); } } else if (cs.nbytes == 0) { return; @@ -1327,7 +1352,7 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { throw std::runtime_error( "WebGPU: named constant '" + cs.named_key + "' undersized"); } - wgpuQueueWriteBuffer(queue_, dst, 0, buf->data(), cs.nbytes); + write_storage_buffer(queue_, dst, buf->data(), cs.nbytes); buf->Free(); } else { throw std::runtime_error("WebGPU: constant has no source"); @@ -1419,7 +1444,7 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { // Fast path: host and GPU element types match byte-for-byte. if (in.nbytes == live_nbytes) { - wgpuQueueWriteBuffer(queue_, tensor.buffer, 0, in.data, live_nbytes); + write_storage_buffer(queue_, tensor.buffer, in.data, live_nbytes); continue; } @@ -1439,8 +1464,21 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { #endif narrowed[e] = static_cast(src[e]); } - wgpuQueueWriteBuffer( - queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); + continue; + } + + // Require an explicit fp32 host dtype, not merely "not int64": inferring + // the narrow from the 2:1 byte ratio alone would silently reinterpret a + // same-sized non-fp32 host buffer (e.g. a stale int32) as fp32. + if (in.host_is_fp32 && buffer_is_fp16 && in.nbytes == live_nbytes * 2) { + const size_t numel = live_nbytes / sizeof(uint16_t); + const float* src = static_cast(in.data); + std::vector narrowed(numel); + for (size_t e = 0; e < numel; e++) { + narrowed[e] = executorch::runtime::etensor::Half(src[e]); + } + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); continue; } @@ -1684,10 +1722,11 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { } for (size_t i = 0; i < output_copies_.size(); i++) { - const size_t copy_nbytes = tensors_[output_ids_[i]].cur_nbytes; - if (!plan.copy_outputs[i] || copy_nbytes == 0) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { continue; } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); @@ -1715,12 +1754,13 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { return 1; } - // GPU timestamp queries assume one submit; chunked execute is multi-submit. +#ifdef WGPU_BACKEND_ENABLE_PROFILING if (should_timestamp_query()) { throw std::runtime_error( "WebGPU: WEBGPU_TIMESTAMP_QUERY is incompatible with chunked execute " "(multi-submit); disable chunking to use GPU timestamp queries"); } +#endif // WGPU_BACKEND_ENABLE_PROFILING for (size_t chunk_index = 0; chunk_index < plan.dispatch_chunks.size(); chunk_index++) { @@ -1759,10 +1799,11 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { if (chunk_index + 1 == plan.dispatch_chunks.size()) { for (size_t i = 0; i < output_copies_.size(); i++) { - const size_t copy_nbytes = tensors_[output_ids_[i]].cur_nbytes; - if (!plan.copy_outputs[i] || copy_nbytes == 0) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { continue; } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); @@ -1813,13 +1854,13 @@ void WebGPUGraph::copy_outputs( continue; } const auto& tensor = tensors_[output_ids_[i]]; - const size_t map_nbytes = tensor.cur_nbytes; - if (map_nbytes == 0) { + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { continue; } const size_t dst_nbytes = outputs[i].nbytes; const bool is_double_width = - dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; const bool widen_fp16 = is_double_width && !tensor.is_int && tensor.elem_size == 2; const bool widen_int32 = @@ -1832,7 +1873,7 @@ void WebGPUGraph::copy_outputs( if (outputs[i].host_is_fp32 && buffer_is_fp16 && !widen_fp16) { throw std::runtime_error("WebGPU: fp16 output buffer size mismatch"); } - if (dst_nbytes != map_nbytes && !widen_fp16 && !widen_int32) { + if (dst_nbytes != logical_nbytes && !widen_fp16 && !widen_int32) { throw std::runtime_error("WebGPU: output buffer size mismatch"); } } @@ -1842,13 +1883,14 @@ void WebGPUGraph::copy_outputs( continue; } const auto& tensor = tensors_[output_ids_[i]]; - const size_t map_nbytes = tensor.cur_nbytes; - if (map_nbytes == 0) { + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { continue; } + const size_t map_nbytes = storage_buffer_size(logical_nbytes); const size_t dst_nbytes = outputs[i].nbytes; const bool is_double_width = - dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; const bool widen_fp16 = is_double_width && !tensor.is_int && tensor.elem_size == 2; const bool widen_int32 = @@ -1887,7 +1929,7 @@ void WebGPUGraph::copy_outputs( const auto* src = static_cast(mapped); auto* dst = static_cast(outputs[i].data); - const size_t n = map_nbytes / sizeof(*src); + const size_t n = logical_nbytes / sizeof(*src); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } @@ -1895,12 +1937,12 @@ void WebGPUGraph::copy_outputs( // int64 host output backed by an int32 GPU buffer: widen (sign-extend). const int32_t* src = static_cast(mapped); int64_t* dst = static_cast(outputs[i].data); - const size_t n = map_nbytes / sizeof(int32_t); + const size_t n = logical_nbytes / sizeof(int32_t); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } } else { - std::memcpy(outputs[i].data, mapped, map_nbytes); + std::memcpy(outputs[i].data, mapped, logical_nbytes); } wgpuBufferUnmap(output_staging_buffers_[i]); } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index a634ebb4172..23ce9df03ed 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -39,6 +39,8 @@ struct WebGPUTensor { // Serialized (GPU-side) element type, used to narrow wider host inputs. size_t elem_size = 0; bool is_int = false; + // Exact BOOL tag for byte-packed WGSL storage. + bool is_bool = false; // Exactly int8 (not uint8/bool), so int8-only ops can guard their dtype. bool is_int8 = false; }; diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 5374390722b..480944ea93d 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -38,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -124,6 +125,7 @@ #include #include #include +#include #include #include #include @@ -151,7 +153,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -362,6 +364,13 @@ constexpr std::array kShaderRegistry = {{ kConstantPadNdWorkgroupSizeY, kConstantPadNdWorkgroupSizeZ, }, + { + "conv1d", + kConv1dWGSL, + kConv1dWorkgroupSizeX, + kConv1dWorkgroupSizeY, + kConv1dWorkgroupSizeZ, + }, { "conv1d_dw", kConv1dDwWGSL, @@ -1034,6 +1043,13 @@ constexpr std::array kShaderRegistry = {{ kTanhWorkgroupSizeY, kTanhWorkgroupSizeZ, }, + { + "to_copy_bool_to_float", + kToCopyBoolToFloatWGSL, + kToCopyBoolToFloatWorkgroupSizeX, + kToCopyBoolToFloatWorkgroupSizeY, + kToCopyBoolToFloatWorkgroupSizeZ, + }, { "to_copy_float_to_int", kToCopyFloatToIntWGSL, diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp index 99200cb2425..1870ee85213 100644 --- a/backends/webgpu/runtime/ops/compare/Compare.cpp +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -58,13 +58,13 @@ void compare_impl( in2_tensor.elem_size != 4) { throw std::runtime_error("compare: fp32 inputs only"); } - if (!out_tensor.is_int || out_tensor.elem_size != 1) { + if (!out_tensor.is_bool || out_tensor.elem_size != 1) { throw std::runtime_error("compare: output must be a 1-byte bool tensor"); } const uint64_t numel = out_tensor.nbytes; - // out bool packed 4/word (array); numel%4==0 gates the readback map. - if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { - throw std::runtime_error("compare: numel must be a nonzero mult of 4"); + // out bool is byte-packed into ceil(numel / 4) u32 storage words. + if (numel == 0u || numel > UINT32_MAX) { + throw std::runtime_error("compare: numel must be nonzero and fit u32"); } const uint64_t in_numel = in1_tensor.nbytes / sizeof(float); if (in1_tensor.nbytes != in2_tensor.nbytes || in_numel != numel) { @@ -75,7 +75,7 @@ void compare_impl( params.num_elements = static_cast(numel); params.op = op; - const uint32_t words = static_cast(numel / 4u); + const uint32_t words = static_cast((numel + 3u) / 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX); utils::WgCount workgroup_count = @@ -85,9 +85,7 @@ void compare_impl( wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(CompareParams)); - graph.add_uniform_buffer_bytes(sizeof(CompareParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // out (rw storage) + in1/in2 (ro storage) + params (uniform). utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( @@ -97,7 +95,7 @@ void compare_impl( {0, WGPUBufferBindingType_Storage, out_tensor.buffer, - out_tensor.nbytes}, + static_cast(words) * sizeof(uint32_t)}, {1, WGPUBufferBindingType_ReadOnlyStorage, in1_tensor.buffer, @@ -127,9 +125,8 @@ void compare_impl( WebGPUGraph& g) { const auto& d = g.cur_dims(in1_id); const uint64_t n = utils::numel_of(d); - if (n == 0u || n % 4u != 0u || n > UINT32_MAX || - utils::numel_of(g.cur_dims(in2_id)) != n) { - throw std::runtime_error("compare(resize): numel must be a mult of 4"); + if (n == 0u || n > UINT32_MAX || utils::numel_of(g.cur_dims(in2_id)) != n) { + throw std::runtime_error("compare(resize): invalid numel"); } g.set_cur_dims(out_id, d); CompareParams p = {}; @@ -137,14 +134,12 @@ void compare_impl( p.op = op; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), static_cast(n / 4u), wg_size, "compare"); + g.device(), static_cast((n + 3u) / 4u), wg_size, "compare"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }; graph.add_tensor_resize_hook(in1_id, resize); graph.add_tensor_resize_hook(in2_id, resize); - - graph.own_uniform_buffer(uniform_buffer); } void eq_op(WebGPUGraph& graph, const std::vector& args) { diff --git a/backends/webgpu/runtime/ops/compare/compare.wgsl b/backends/webgpu/runtime/ops/compare/compare.wgsl index a16568e5533..a650a6298b5 100644 --- a/backends/webgpu/runtime/ops/compare/compare.wgsl +++ b/backends/webgpu/runtime/ops/compare/compare.wgsl @@ -16,27 +16,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/compare/compare_wgsl.h b/backends/webgpu/runtime/ops/compare/compare_wgsl.h index 672c99b62d8..c1c4ac23e4f 100644 --- a/backends/webgpu/runtime/ops/compare/compare_wgsl.h +++ b/backends/webgpu/runtime/ops/compare/compare_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from compare.wgsl - DO NOT EDIT. -// wgsl-sha256: 241e7e6762b1eded07d28a3767936c970f509f6591e7cbf599d0b1eb61efb181 +// wgsl-sha256: 8f330b5a1e29a1fb8135e64600dadb1dc64c98f8f5371357f8de73a1808b76d6 inline constexpr const char* kCompareWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var input1: array; @@ -33,27 +33,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp index ae86ccffba8..64d8c2d2380 100644 --- a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include @@ -49,8 +50,24 @@ uint32_t conv1d_out_len( int64_t stride, int64_t padding, int64_t dilation) { - return static_cast( - (in_len + 2 * padding - dilation * (k - 1) - 1) / stride + 1); + if (in_len <= 0 || k <= 0 || stride <= 0 || padding < 0 || dilation <= 0) { + throw std::runtime_error("conv1d: invalid geometry parameter"); + } + constexpr int64_t kMaxShaderIndex = std::numeric_limits::max(); + if (in_len > kMaxShaderIndex || k > kMaxShaderIndex || + stride > kMaxShaderIndex || padding > kMaxShaderIndex || + dilation > kMaxShaderIndex) { + throw std::runtime_error("conv1d: geometry parameter exceeds i32"); + } + const int64_t numerator = in_len + 2 * padding - dilation * (k - 1) - 1; + if (numerator < 0) { + throw std::runtime_error("conv1d: kernel exceeds padded input"); + } + const int64_t out_len = numerator / stride + 1; + if (static_cast(out_len) > UINT32_MAX) { + throw std::runtime_error("conv1d: output length exceeds u32"); + } + return static_cast(out_len); } int64_t first_int(const std::vector& v) { @@ -71,6 +88,24 @@ static_assert( sizeof(Conv1dPwParams) == 32, "Conv1dPwParams must match the WGSL Params struct (32 bytes)"); +struct Conv1dParams { + uint32_t in_channels; + uint32_t out_channels; + uint32_t in_len; + uint32_t out_len; + uint32_t kernel_size; + uint32_t stride; + uint32_t padding; + uint32_t dilation; + uint32_t numel; + uint32_t has_bias; +}; +static_assert( + sizeof(Conv1dParams) == 40, + "Conv1dParams must match the WGSL Params struct (40 bytes)"); +constexpr uint64_t kMaxConv1dDispatchElements = + static_cast(std::numeric_limits::max()); + // Pointwise conv1d (K=1, groups=1): a per-position matmul over channels. void add_conv1d_pw_node( WebGPUGraph& graph, @@ -203,6 +238,154 @@ void add_conv1d_pw_node( graph.own_uniform_buffer(params_buf); } +// General groups=1 conv1d. Voxtral uses K=3 with stride 1 then 2. +void add_conv1d_node( + WebGPUGraph& graph, + int in_id, + int weight_id, + int bias_id, + int out_id, + uint32_t stride, + uint32_t padding, + uint32_t dilation) { + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + if (!utils::is_fp32_tensor(in) || !utils::is_fp32_tensor(weight) || + !utils::is_fp32_tensor(out)) { + throw std::runtime_error("conv1d: input, weight, and output must be fp32"); + } + + const uint32_t expected_out_len = conv1d_out_len( + in.dims.at(2), weight.dims.at(2), stride, padding, dilation); + const uint32_t batch = static_cast(in.dims.at(0)); + const uint32_t in_channels = static_cast(in.dims.at(1)); + const uint32_t in_len = static_cast(in.dims.at(2)); + const uint32_t out_channels = static_cast(out.dims.at(1)); + const uint32_t out_len = static_cast(out.dims.at(2)); + const uint32_t kernel_size = static_cast(weight.dims.at(2)); + if (out.dims.at(0) != in.dims.at(0) || out_len != expected_out_len || + weight.dims.at(0) != out.dims.at(1) || + weight.dims.at(1) != in.dims.at(1)) { + throw std::runtime_error("conv1d: shape mismatch"); + } + + const uint64_t in_numel = utils::check_fp32(in, "conv1d", "input"); + const uint64_t out_numel = utils::check_fp32(out, "conv1d", "output"); + const uint64_t weight_numel = utils::check_fp32(weight, "conv1d", "weight"); + if (in_numel != static_cast(batch) * in_channels * in_len || + out_numel != static_cast(batch) * out_channels * out_len || + weight_numel != + static_cast(out_channels) * in_channels * kernel_size || + in_numel > UINT32_MAX || weight_numel > UINT32_MAX || + out_numel > kMaxConv1dDispatchElements) { + throw std::runtime_error("conv1d: fp32 byte-size or u32 mismatch"); + } + if (has_bias) { + const auto& bias = graph.get_tensor(bias_id); + if (!utils::is_fp32_tensor(bias) || bias.dims.size() != 1 || + bias.dims.at(0) != out.dims.at(1) || + utils::check_fp32(bias, "conv1d", "bias") != out_channels) { + throw std::runtime_error("conv1d: bias shape mismatch"); + } + } + + Conv1dParams params = {}; + params.in_channels = in_channels; + params.out_channels = out_channels; + params.in_len = in_len; + params.out_len = out_len; + params.kernel_size = kernel_size; + params.stride = stride; + params.padding = padding; + params.dilation = dilation; + params.numel = static_cast(out_numel); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kConv1dWorkgroupSizeX); + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, params.numel, wg_size, "conv1d"); + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); + WGPUBuffer params_buf = graph.create_params_buffer(params); + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : weight.nbytes; + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConv1dWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(Conv1dParams)}, + }, + &wg_size_constant, + 1); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "conv1d", + workgroup_count.y}); + + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + has_bias, + wg_size, + dispatch_idx, + params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + if (dims.size() != 3 || dims[0] <= 0 || dims[1] <= 0 || dims[2] <= 0 || + dims[1] != static_cast(in_channels)) { + throw std::runtime_error("conv1d(resize): input shape changed"); + } + Conv1dParams p = {}; + p.in_channels = in_channels; + p.out_channels = out_channels; + p.in_len = static_cast(dims[2]); + p.out_len = + conv1d_out_len(dims[2], kernel_size, stride, padding, dilation); + p.kernel_size = kernel_size; + p.stride = stride; + p.padding = padding; + p.dilation = dilation; + const uint64_t input_numel = utils::numel(dims); + const uint64_t numel = utils::numel( + {dims[0], out_channels, static_cast(p.out_len)}); + if (input_numel > UINT32_MAX || numel > kMaxConv1dDispatchElements) { + throw std::runtime_error( + "conv1d(resize): tensor numel exceeds shader index range"); + } + p.numel = static_cast(numel); + p.has_bias = has_bias ? 1u : 0u; + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "conv1d(resize)"); + g.set_cur_dims( + out_id, {dims[0], out_channels, static_cast(p.out_len)}); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); +} + // depthwise-conv1d (groups==C); mirrors Vulkan conv1d_dw (Convolution.cpp:755). void convolution_impl(WebGPUGraph& graph, const std::vector& args) { // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back(). @@ -242,38 +425,44 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { const bool transposed = graph.get_bool(transposed_id); const int64_t groups = graph.get_int(groups_id); - // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. - if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && - first_int(graph.get_int_list(stride_id)) == 1 && - first_int(graph.get_int_list(padding_id)) == 0) { - add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); - return; - } - - // Otherwise only the depthwise config (groups==C, weight [C,1,K]). - if (transposed || groups != static_cast(channels) || - weight_tensor.dims.at(0) != static_cast(channels) || - weight_tensor.dims.at(1) != 1) { - throw std::runtime_error( - "convolution: only depthwise or pointwise conv1d supported"); - } - const int64_t stride_i = first_int(graph.get_int_list(stride_id)); const int64_t padding_i = first_int(graph.get_int_list(padding_id)); const int64_t dilation_i = first_int(graph.get_int_list(dilation_id)); - if (stride_i < 1) { - throw std::runtime_error("convolution: stride must be >= 1"); + if (stride_i < 1 || stride_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: stride must fit positive i32"); } - if (padding_i < 0) { - throw std::runtime_error("convolution: padding must be >= 0"); + if (padding_i < 0 || padding_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: padding must fit nonnegative i32"); } - if (dilation_i < 1) { - throw std::runtime_error("convolution: dilation must be >= 1"); + if (dilation_i < 1 || dilation_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: dilation must fit positive i32"); } const uint32_t stride = static_cast(stride_i); const uint32_t padding = static_cast(padding_i); const uint32_t dilation = static_cast(dilation_i); + // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. + if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && + stride_i == 1 && padding_i == 0) { + add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); + return; + } + + const bool is_depthwise = !transposed && + groups == static_cast(channels) && + weight_tensor.dims.at(0) == static_cast(channels) && + weight_tensor.dims.at(1) == 1; + if (!is_depthwise && !transposed && groups == 1) { + add_conv1d_node( + graph, in_id, weight_id, bias_id, out_id, stride, padding, dilation); + return; + } + + if (!is_depthwise) { + throw std::runtime_error( + "convolution: only depthwise, pointwise, or groups=1 conv1d supported"); + } + uint64_t out_numel = 1; for (int64_t d : out_tensor.dims) { out_numel *= static_cast(d); diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl new file mode 100644 index 00000000000..c51594d2005 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl @@ -0,0 +1,53 @@ +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h new file mode 100644 index 00000000000..e29af80459b --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h @@ -0,0 +1,77 @@ +/* + * 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 conv1d.wgsl - DO NOT EDIT. +// wgsl-sha256: 7bc955b7f43473aab96222e7a1228973b65e38e2e6a19e2b9793e0ed7f3768d9 +inline constexpr const char* kConv1dWGSL = R"( +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} +)"; + +inline constexpr uint32_t kConv1dWorkgroupSizeX = 64; +inline constexpr uint32_t kConv1dWorkgroupSizeY = 1; +inline constexpr uint32_t kConv1dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp index 5bfc0fa3bf1..9f6a4e68ba6 100644 --- a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp +++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp @@ -14,6 +14,7 @@ #include +#include #include namespace executorch::backends::webgpu { @@ -34,6 +35,22 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (graph.get_value_type(args.at(1)) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic target sizes are unsupported"); + } + for (int64_t target_size : graph.get_int_list(args.at(1))) { + if (target_size == -1) { + throw std::runtime_error( + "WebGPU expand_copy: inferred target sizes are unsupported"); + } + } + if (graph.tensor_has_dynamic_dims(in_id) || + graph.tensor_has_dynamic_dims(out_id)) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic shapes are unsupported"); + } + TensorMeta out_meta; TensorMeta in_meta; fill_tensor_meta(out_tensor, &out_meta); @@ -44,21 +61,24 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "expand_copy: non-fp32 operand (nbytes != numel*4)"); } + if (out_meta.numel > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU expand_copy: element count exceeds the flattened 2D dispatch " + "limit"); + } uint32_t wg_size = utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "expand_copy"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer out_meta_buf = - utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); - WGPUBuffer in_meta_buf = - utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); - graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); + WGPUBuffer out_meta_buf = graph.create_params_buffer(out_meta); + WGPUBuffer in_meta_buf = graph.create_params_buffer(in_meta); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -78,10 +98,8 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - wgpuBufferRelease(out_meta_buf); - wgpuBufferRelease(in_meta_buf); + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); } } // namespace diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl index 053311a69f4..fab4df15a90 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl @@ -13,8 +13,10 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h index 83c5881604f..f1449f61793 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from expand_copy.wgsl - DO NOT EDIT. -// wgsl-sha256: 99953670bea89e42bc9c689ab80addfd9a331442c8c8f1a5b0c39dbe11c19370 +// wgsl-sha256: b3c032ab961ffde245fc44289b67df3b5e4ca93eedb9ada2f20a3eaa6f10e9c6 inline constexpr const char* kExpandCopyWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -30,8 +30,10 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/Gelu.cpp b/backends/webgpu/runtime/ops/gelu/Gelu.cpp index 943012515ff..023ed7b5d88 100644 --- a/backends/webgpu/runtime/ops/gelu/Gelu.cpp +++ b/backends/webgpu/runtime/ops/gelu/Gelu.cpp @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -42,13 +43,18 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { const auto& out_tensor = graph.get_tensor(out_id); utils::check_elementwise_fp32_io(in_tensor, out_tensor, "gelu"); - uint32_t num_elements = - static_cast(out_tensor.nbytes / sizeof(float)); + const uint64_t num_elements64 = out_tensor.nbytes / sizeof(float); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu: element count exceeds the flattened 2D dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); // Each thread handles up to 4 elements (vec4 body + scalar-tail idiom). uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kGeluWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, num_vec4_threads, wg_size, "gelu"); WGPUConstantEntry wg_constant = utils::make_wg_size_constant(wg_size); @@ -56,9 +62,7 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { GeluParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(GeluParams)); - graph.add_uniform_buffer_bytes(sizeof(GeluParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // input (read storage) + output (storage) + params. The exact/approximate // choice is baked into the compiled pipeline via the entry point (mirrors @@ -85,10 +89,38 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { 1, exact ? "main_erf" : "main_tanh"); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - // Drop our ref; the bind group keeps the uniform buffer alive until release. - wgpuBufferRelease(uniform_buffer); + const size_t dispatch_idx = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t num_elements64 = utils::numel_of(dims); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu(resize): element count exceeds the flattened 2D " + "dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); + g.set_cur_dims(out_id, dims); + + GeluParams params = {}; + params.num_elements = num_elements; + wgpuQueueWriteBuffer( + g.queue(), params_buf, 0, ¶ms, sizeof(GeluParams)); + + const uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); + const utils::WgCount resized_workgroup_count = + utils::compute_2d_workgroup_count( + g.device(), num_vec4_threads, wg_size, "gelu(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = + resized_workgroup_count.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = + resized_workgroup_count.y; + }); } } // namespace diff --git a/backends/webgpu/runtime/ops/gelu/gelu.wgsl b/backends/webgpu/runtime/ops/gelu/gelu.wgsl index 4f7eb68bc96..9583ef81551 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu.wgsl +++ b/backends/webgpu/runtime/ops/gelu/gelu.wgsl @@ -33,8 +33,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -49,8 +52,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h index 6da12e229af..f8af0f8d2c3 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h +++ b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from gelu.wgsl - DO NOT EDIT. -// wgsl-sha256: 18f4a82d3bad1ef8703397b871c708804140c4cb382451661f7a77367ac2425f +// wgsl-sha256: 96570753688590fa009ee5503f754cf3eb572dcb3dcae6818220fe06fe3139ee inline constexpr const char* kGeluWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -50,8 +50,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -66,8 +69,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp index 05f7ff6fe24..54d80016296 100644 --- a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -10,12 +10,14 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -68,9 +70,7 @@ void add_convert_op( ConvertParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); - graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -113,8 +113,185 @@ void add_convert_op( wg_size, "to_copy(resize)"); }); +} + +// Decode byte-packed bool storage into numeric fp32 values. +void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("to_copy_bool_to_float: null buffer binding"); + } + if (!in_tensor.is_bool || in_tensor.elem_size != 1 || out_tensor.is_int || + out_tensor.elem_size != sizeof(float) || + out_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes / sizeof(float) != in_tensor.nbytes) { + throw std::runtime_error("to_copy_bool_to_float: dtype/numel mismatch"); + } + if (in_tensor.nbytes == 0u || + in_tensor.nbytes > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: numel must be nonzero and fit u32"); + } + + const uint32_t num_elements = static_cast(in_tensor.nbytes); + const uint64_t input_bind_size_u64 = + (static_cast(in_tensor.nbytes) + 3u) & ~uint64_t(3); + if (input_bind_size_u64 > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: input binding size overflows"); + } + const size_t input_bind_size = static_cast(input_bind_size_u64); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_elements, wg_size, "to_copy_bool_to_float"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kToCopyBoolToFloatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + input_bind_size}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(dims); + if (numel == 0u || numel > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float(resize): invalid numel"); + } + g.set_cur_dims(out_id, dims); + ConvertParams p = {}; + p.num_elements = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy_bool_to_float(resize)"); + }); +} + +// Decode byte-packed bool storage into numeric fp32 values. +void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("to_copy_bool_to_float: null buffer binding"); + } + if (!in_tensor.is_bool || in_tensor.elem_size != 1 || out_tensor.is_int || + out_tensor.elem_size != sizeof(float) || + out_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes / sizeof(float) != in_tensor.nbytes) { + throw std::runtime_error("to_copy_bool_to_float: dtype/numel mismatch"); + } + if (in_tensor.nbytes == 0u || + in_tensor.nbytes > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: numel must be nonzero and fit u32"); + } + + const uint32_t num_elements = static_cast(in_tensor.nbytes); + const uint64_t input_bind_size_u64 = + (static_cast(in_tensor.nbytes) + 3u) & ~uint64_t(3); + if (input_bind_size_u64 > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: input binding size overflows"); + } + const size_t input_bind_size = static_cast(input_bind_size_u64); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_elements, wg_size, "to_copy_bool_to_float"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kToCopyBoolToFloatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + input_bind_size}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(dims); + if (numel == 0u || numel > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float(resize): invalid numel"); + } + g.set_cur_dims(out_id, dims); + ConvertParams p = {}; + p.num_elements = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy_bool_to_float(resize)"); + }); - // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } @@ -129,6 +306,12 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.is_bool != out_tensor.is_bool && in_tensor.is_int && + out_tensor.is_int) { + throw std::runtime_error( + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + // Same is_int+width = flat byte copy; unique dtype key in the 32-bit domain. if (in_tensor.is_int == out_tensor.is_int && in_tensor.elem_size == out_tensor.elem_size) { @@ -137,7 +320,10 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { } // int<->float = numeric convert (mirrors Vulkan add_view_copy_convert_node). - if (in_tensor.is_int && !out_tensor.is_int) { + if (in_tensor.is_bool && !out_tensor.is_int && + out_tensor.elem_size == sizeof(float)) { + add_bool_to_float_op(graph, in_id, out_id); + } else if (in_tensor.is_int && !out_tensor.is_int) { add_convert_op( graph, in_id, diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl new file mode 100644 index 00000000000..239730de65d --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl @@ -0,0 +1,24 @@ +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h new file mode 100644 index 00000000000..ef7e40976f3 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h @@ -0,0 +1,48 @@ +/* + * 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 to_copy_bool_to_float.wgsl - DO NOT EDIT. +// wgsl-sha256: 29fd43b2f638489e9b8d72b2cc9140d07174c750cbdc291e63793319a2fa5961 +inline constexpr const char* kToCopyBoolToFloatWGSL = R"( +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} +)"; + +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeX = 256; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 6ed7229604b..949ad3c0de6 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -153,42 +153,90 @@ void expect_dual_q4_topology(WebGPUGraph& graph) { 1); } -TEST(WebGPUShaderRegistry, FindsKnownShaderAndRejectsUnknownName) { - const WebGPUShaderInfo& sigmoid = get_webgpu_shader_info("sigmoid"); - EXPECT_EQ(sigmoid.name, "sigmoid"); - EXPECT_NE(sigmoid.source, nullptr); - EXPECT_GT(sigmoid.workgroup_size_x, 0u); - EXPECT_THROW( - get_webgpu_shader_info("not_a_registered_shader"), std::runtime_error); -} +struct Conv1dRouteCase { + const char* name; + std::vector input_dims; + std::vector weight_dims; + std::vector output_dims; + int64_t stride; + int64_t padding; + int64_t dilation; + int64_t groups; + const char* expected_kernel; +}; + +void build_conv1d_route_graph( + WebGPUGraph& graph, + const Conv1dRouteCase& test_case) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](const std::vector& dims, int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + auto add_int_list = [&](int64_t value) { + const int id = static_cast(values.size()); + const std::vector items = {value}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::IntList, + vk::CreateIntListDirect(fbb, &items).Union())); + return id; + }; + + const int input = add_tensor(test_case.input_dims, 0); + const int weight = add_tensor(test_case.weight_dims, 1); + const int bias = static_cast(values.size()); + values.push_back(vk::CreateVkValue(fbb)); + const int stride = add_int_list(test_case.stride); + const int padding = add_int_list(test_case.padding); + const int dilation = add_int_list(test_case.dilation); + const int transposed = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Bool, vk::CreateBool(fbb, false).Union())); + const int output_padding = add_int_list(0); + const int groups = add_int(test_case.groups); + const int output = add_tensor(test_case.output_dims, 2); + const std::vector args = { + input, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + output}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.convolution.default", &args)); + const std::vector input_ids = { + static_cast(input), static_cast(weight)}; + const std::vector output_ids = {static_cast(output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); -TEST(WebGPUQ4RouteSignal, PreservesStaticAndRecordsBothDynamicSignals) { - WebGPUGraph static_graph; - static_graph.set_device(g_device); - build_q4_route_graph(static_graph, Q4RouteSignal::Static); - const auto static_dispatches = q4_dispatches(static_graph); - ASSERT_EQ(static_dispatches.size(), 1); - EXPECT_EQ(static_graph.num_dispatches(), 1); - EXPECT_NE(static_dispatches[0]->pipeline, nullptr); - EXPECT_NE(static_dispatches[0]->bind_group, nullptr); - EXPECT_FALSE(static_graph.has_dynamic_shapes()); - EXPECT_FALSE(static_graph.config().record_q4gsw_decode_route); - - WebGPUGraph legacy_graph; - legacy_graph.set_device(g_device); - build_q4_route_graph(legacy_graph, Q4RouteSignal::LegacyGraphMarker); - EXPECT_TRUE(legacy_graph.has_dynamic_shapes()); - EXPECT_FALSE(legacy_graph.config().record_q4gsw_decode_route); - EXPECT_EQ(legacy_graph.num_dispatches(), 3); - expect_dual_q4_topology(legacy_graph); - - WebGPUGraph explicit_graph; - explicit_graph.set_device(g_device); - build_q4_route_graph(explicit_graph, Q4RouteSignal::ExplicitOption); - EXPECT_FALSE(explicit_graph.has_dynamic_shapes()); - EXPECT_TRUE(explicit_graph.config().record_q4gsw_decode_route); - EXPECT_EQ(explicit_graph.num_dispatches(), 2); - expect_dual_q4_topology(explicit_graph); + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } TEST(WebGPUComputeDispatch, PipelineKeyCanonicalizesConstants) { @@ -899,6 +947,58 @@ TEST(WebGPURopeValidation, RejectsMalformedGraphsBeforeDispatchAllocation) { } } +TEST(WebGPUToCopyValidation, RejectsBoolAndByteIntegerConversions) { + ASSERT_TRUE(webgpu_operator_registry().has_op("aten._to_copy.default")); + namespace vk = vkgraph; + struct TestCase { + const char* name; + vk::VkDataType input_dtype; + vk::VkDataType output_dtype; + }; + const TestCase cases[] = { + {"bool_to_int8", vk::VkDataType::BOOL, vk::VkDataType::INT8}, + {"bool_to_uint8", vk::VkDataType::BOOL, vk::VkDataType::UINT8}, + {"int8_to_bool", vk::VkDataType::INT8, vk::VkDataType::BOOL}, + {"uint8_to_bool", vk::VkDataType::UINT8, vk::VkDataType::BOOL}, + }; + for (const TestCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + ::flatbuffers::FlatBufferBuilder fbb; + const std::vector dims = {4}; + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.input_dtype, &dims, -1, 0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.output_dtype, &dims, -1, 1) + .Union())); + const std::vector args = {0, 1}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten._to_copy.default", &args)); + const std::vector input_ids = {0}; + const std::vector output_ids = {1}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + FAIL() << test_case.name << " unexpectedly built"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + EXPECT_EQ(graph.memory_stats().num_dispatches, 0); + } +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 7b3c1d6b0a7..cf16cd7e87c 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -18,6 +18,8 @@ // F dyn_rms_chain (rms(rms(x))) at 3 S -> golden (resize CASCADE, DD-4) // G rms+residual H rms*x I dyn_linear J sdpa_dyn K emb_dyn L rope_dyn // M dyn_sigmoid N dyn_select (select_copy(0,-1), dynamic S) +// O ONE dyn_conv1d graph reused across live input lengths +// P ONE dyn_gelu graph reused above -> at -> above the old 1D dispatch cap // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py. // // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else @@ -33,6 +35,7 @@ #include #include +#include #include #include #include @@ -127,6 +130,74 @@ void check_s(Module& m, const std::string& prefix, int s) { << " golden.size=" << golden.size() << ")"; } +constexpr int kConv1dInChannels = 3; +constexpr int kConv1dOutChannels = 4; +constexpr int kConv1dKernel = 3; +constexpr int kConv1dStride = 2; +constexpr int kConv1dPadding = 1; +constexpr int kConv1dDilation = 2; + +void check_conv1d(Module& module, int length) { + const std::string prefix = g_dir + "/dyn_conv1d.S" + std::to_string(length); + auto input = read_bin(prefix + ".input.bin"); + auto golden = read_bin(prefix + ".golden.bin"); + ASSERT_EQ(input.size(), static_cast(kConv1dInChannels * length)); + ASSERT_FALSE(golden.empty()); + auto tensor = + make_tensor_ptr({1, kConv1dInChannels, length}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "conv1d length=" << length << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const int output_length = (length + 2 * kConv1dPadding - + kConv1dDilation * (kConv1dKernel - 1) - 1) / + kConv1dStride + + 1; + ASSERT_EQ(output.dim(), 3); + ASSERT_EQ(output.size(0), 1); + ASSERT_EQ(output.size(1), kConv1dOutChannels); + ASSERT_EQ(output.size(2), output_length); + const size_t numel = static_cast(kConv1dOutChannels * output_length); + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + const float error = max_err(got, golden); + EXPECT_LT(error, 1e-3f) << "conv1d length=" << length << " max_err=" << error; +} + +constexpr int kGeluOld1dDispatchCap = 4 * 64 * 65535; +constexpr int kGelu2dDispatchBoundary = kGeluOld1dDispatchCap + 1; +constexpr int kGeluPatternSize = 257; + +void check_gelu_2d(Module& module, int elements) { + std::array golden = {}; + std::vector input(static_cast(elements)); + for (int i = 0; i < kGeluPatternSize; i++) { + const float value = -4.0f + 8.0f * i / (kGeluPatternSize - 1); + golden[i] = 0.5f * value * (1.0f + std::erf(value * 0.7071067811865476f)); + } + for (int i = 0; i < elements; i++) { + input[i] = -4.0f + 8.0f * (i % kGeluPatternSize) / (kGeluPatternSize - 1); + } + auto tensor = make_tensor_ptr({elements}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "gelu elements=" << elements << " forward failed"; + const auto& output = result.get()[0].toTensor(); + ASSERT_EQ(output.dim(), 1); + ASSERT_EQ(output.size(0), elements); + ASSERT_EQ(output.numel(), elements); + const float* data = output.const_data_ptr(); + float error = 0.0f; + for (int i = 0; i < elements; i++) { + error = std::fmax(error, std::fabs(data[i] - golden[i % kGeluPatternSize])); + } + EXPECT_LT(error, 1e-4f) << "gelu elements=" << elements + << " max_err=" << error; +} + // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; @@ -893,6 +964,43 @@ TEST(DynamicShape, RmsNormReusedGraph) { } } +TEST(DynamicShape, Conv1dReusedGraph) { + Module module(g_dir + "/dyn_conv1d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_conv1d.pte"; + for (int length : {16, 9, 5, 16}) { + check_conv1d(module, length); + } +} + +TEST(DynamicShape, GeluCrosses2dDispatchBoundary) { + if (std::getenv("WEBGPU_TEST_HEAVY") == nullptr) { + GTEST_SKIP() << "WEBGPU_TEST_HEAVY not set"; + } + Module module(g_dir + "/dyn_gelu_2d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_gelu_2d.pte"; + for (int elements : + {kGelu2dDispatchBoundary, + kGeluOld1dDispatchCap, + kGelu2dDispatchBoundary}) { + check_gelu_2d(module, elements); + } +} + +TEST(DynamicShape, ExpandCopyRejectsDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy.pte"; + ASSERT_TRUE(std::ifstream(path).good()) << "missing dyn_expand_copy.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + +TEST(DynamicShape, ExpandCopyRejectsInferredDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy_inferred.pte"; + ASSERT_TRUE(std::ifstream(path).good()) + << "missing dyn_expand_copy_inferred.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + // C2: grow-only reuse — one loaded rms graph run smallest -> largest, so the // FIRST resize grows the dispatch (every other reuse test starts at MAXS and // only shrinks; this catches a hook with a shrink-only short-circuit). diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp index edc0f315294..a839d224b16 100644 --- a/backends/webgpu/test/native/test_webgpu_utils.cpp +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -14,8 +14,16 @@ #include +#include + using namespace executorch::backends::webgpu; +TEST(WebGPUUtils, DivUpDoesNotOverflowAtUint32Max) { + constexpr uint32_t kMax = std::numeric_limits::max(); + EXPECT_EQ(utils::div_up(kMax, 4u), 1073741824u); + EXPECT_EQ(utils::div_up(kMax, kMax), 1u); +} + TEST(WebGPUUtils, DispatchGridStaysOneDimUnderCeiling) { utils::DispatchGrid g = utils::compute_dispatch_grid_from_limits(1000u, 256u, 65535u, "test"); diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index bbd6e13dc9d..6a019ab8e9f 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -53,8 +53,16 @@ CompareModule, ) from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule -from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule +from executorch.backends.webgpu.test.ops.test_conv1d_pw import ( + Conv1dModule, + Conv1dPwModule, + GENERAL_CONFIGS as _CONV1D_CONFIGS, +) from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ConvWithClampModule +from executorch.backends.webgpu.test.ops.test_expand_copy import ( + CONFIGS as _EXPAND_COPY_CONFIGS, + ExpandCopyModule, +) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_grid_priors import GridPriorsModule @@ -153,8 +161,13 @@ ) from executorch.backends.webgpu.test.ops.test_to_copy import ( + bool_tail_input, + compare_to_copy_input_a, + compare_to_copy_input_b, + CompareToCopyBoolToFloatModule, to_copy_float_input, to_copy_int_input, + ToCopyBoolToFloatModule, ToCopyFloatToIntToFloatModule, ToCopyIntToFloatModule, ) @@ -191,6 +204,39 @@ def _add_factory(variant: str = "regular") -> torch.nn.Module: }[variant]() +@register_op_test("to_copy_bool_to_float") +def _to_copy_bool_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=CompareToCopyBoolToFloatModule, + cases=[ + Case( + inputs=( + InputSpec((n,), gen=compare_to_copy_input_a), + InputSpec((n,), gen=compare_to_copy_input_b), + ), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + +@register_op_test("to_copy_bool_input_to_float") +def _to_copy_bool_input_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=ToCopyBoolToFloatModule, + cases=[ + Case( + inputs=(InputSpec((n,), gen=bool_tail_input),), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + @register_op_test("add") def _add_suite() -> WebGPUTestSuite: # Same-shape numeric coverage only: broadcast adds stay export-smoke in @@ -295,10 +341,7 @@ def _minimum_suite() -> WebGPUTestSuite: def _compare_suite(op: str) -> WebGPUTestSuite: - # Elementwise fp32 comparison -> bool (byte-exact golden). The two inputs use - # DIFFERENT discrete-range seeds so a!=b (real lt/gt mix) while colliding - # often (eq/le/ge ties); all shapes have numel % 4 == 0 (bool output packs 4 - # bytes/word). Same-shape only (flat kernel; broadcast=smoke). + # Distinct inputs and tail shapes cover byte-exact packed BOOL output. def case(name, shape): return Case( name=name, @@ -310,7 +353,14 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda: CompareModule(op), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case("tail_1", (1,)), + case("tail_5", (5,)), + case("tail_67", (67,)), + case("2d", (4, 8)), + case("3d", (2, 3, 8)), + case("sq", (16, 16)), + ], golden_dtype="bool", ) @@ -612,6 +662,54 @@ def case(name, C, L, kernel, stride, padding, dilation, bias): case("k3s2p1", 4, 8, 3, 2, 1, 1, True), case("dil2", 3, 10, 3, 1, 2, 2, True), case("k5_nobias", 5, 7, 5, 1, 0, 1, False), + case("single_channel_route", 1, 8, 3, 1, 1, 1, True), + ], + atol=1e-3, + rtol=1e-3, + ) + + +@register_op_test("conv1d") +def _conv1d_suite() -> WebGPUTestSuite: + # General NCL conv1d; neighboring suites cover the retained fast paths. + def case(name, cfg): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + return Case( + name=name, + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + inputs=((n, ic, length),), + ) + + dynamic_cfg = _CONV1D_CONFIGS["voxtral_stride1"] + n, ic, oc, length, kernel, stride, padding, dilation, bias = dynamic_cfg + dynamic_length = torch.export.Dim("manifest_conv1d_length", min=7, max=length) + return WebGPUTestSuite( + module_factory=Conv1dModule, + cases=[ + *[case(name, cfg) for name, cfg in _CONV1D_CONFIGS.items()], + Case( + name="dynamic_length_10_to_7", + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + export_inputs=((n, ic, length),), + inputs=((n, ic, 7),), + dynamic_shapes=({2: dynamic_length},), + ), ], atol=1e-3, rtol=1e-3, @@ -988,12 +1086,36 @@ def _cat_suite() -> WebGPUTestSuite: N as _GELU_N, ) +_GELU_2D_DISPATCH_BOUNDARY = 4 * 64 * 65535 + 1 +_EXPAND_COPY_2D_DISPATCH_BOUNDARY = 64 * 65535 + 1 + def _gelu_full_range(_shape) -> torch.Tensor: # Reuse the deterministic linspace(-6, 6) spanning negatives/zero/positives. return _gelu_det_input() +@register_op_test("expand_copy") +def _expand_copy_suite() -> WebGPUTestSuite: + cases = [ + Case(name=name, construct={"shape": out_shape}, inputs=(in_shape,)) + for name, (in_shape, out_shape) in _EXPAND_COPY_CONFIGS.items() + ] + cases.append( + Case( + name="dispatch_2d_boundary", + construct={"shape": (_EXPAND_COPY_2D_DISPATCH_BOUNDARY,)}, + inputs=((1,),), + heavy=True, + ) + ) + return WebGPUTestSuite( + module_factory=ExpandCopyModule, + cases=cases, + golden_dtype="float32", + ) + + @register_op_test("gelu") def _gelu_suite() -> WebGPUTestSuite: # erf ("none") is the Florence-2/BART + PyTorch default; tanh is the approx. @@ -1015,6 +1137,12 @@ def _gelu_suite() -> WebGPUTestSuite: construct={"approximate": "none"}, inputs=(InputSpec(shape=(_GELU_N,), gen=_gelu_full_range),), ), + Case( + name="erf_dispatch_2d_boundary", + construct={"approximate": "none"}, + inputs=(InputSpec(shape=(_GELU_2D_DISPATCH_BOUNDARY,), gen="ramp"),), + heavy=True, + ), ], atol=1e-4, rtol=1e-3, diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 72f819f94ce..33b7940a9be 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -39,6 +39,8 @@ def _materialize(spec) -> torch.Tensor: shape, gen = spec, "randn" if callable(gen): _t = gen(shape) + if _t.dtype == torch.bool: + return _t return ( _t.to(torch.int32) if not _t.is_floating_point() else _t.to(torch.float32) ) @@ -53,13 +55,17 @@ def _materialize(spec) -> torch.Tensor: def export_case(suite: WebGPUTestSuite, case) -> tuple[torch.nn.Module, tuple, object]: - """Build the module + forward inputs and export to an ExecuTorch program.""" + """Build the module and export it, returning the live runtime inputs.""" module = suite.module_factory(**case.construct) # Seed so an unseeded-randn input is reproducible across generations (the golden uses # the SAME tensor, so this only affects which bytes a case sees, never pass/fail). torch.manual_seed(0) inputs = tuple(_materialize(s) for s in case.inputs) - ep = torch.export.export(module, inputs) + export_inputs = inputs + if case.export_inputs is not None: + torch.manual_seed(0) + export_inputs = tuple(_materialize(s) for s in case.export_inputs) + ep = torch.export.export(module, export_inputs, dynamic_shapes=case.dynamic_shapes) prog = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] ).to_executorch() @@ -169,7 +175,10 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d input_entries: list[dict] = [] for i, t in enumerate(inputs): rel = f"{case_id}.in{i}.bin" - if t.dtype == torch.int32: + if t.dtype == torch.bool: + _write_int8(t.to(torch.int8), os.path.join(out_dir, rel)) + in_dtype = "bool" + elif t.dtype == torch.int32: t.detach().cpu().numpy().astype(" +#include #include #include #include @@ -61,7 +62,15 @@ class OpCase : public ::testing::Test { const size_t n = numel(in.shape); std::vector sizes( in.shape.begin(), in.shape.end()); - if (in.dtype == "int32") { + if (in.dtype == "bool") { + auto data = load_int8_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + std::vector raw(data.begin(), data.end()); + tensors.push_back(make_tensor_ptr( + std::move(sizes), + std::move(raw), + executorch::aten::ScalarType::Bool)); + } else if (in.dtype == "int32") { auto data = load_int32_bin(in.path, n); ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); @@ -96,10 +105,11 @@ class OpCase : public ::testing::Test { auto golden = load_int8_bin(e_.golden.path, gn); ASSERT_FALSE(golden.empty()) << "missing/short golden: " << e_.golden.path; - const bool* out_p = out_tensor.const_data_ptr(); + ASSERT_EQ(out_tensor.scalar_type(), executorch::aten::ScalarType::Bool); + const uint8_t* out_p = out_tensor.const_data_ptr(); int mism = -1; for (size_t i = 0; i < gn; i++) { - if (static_cast(out_p[i]) != golden[i]) { + if (out_p[i] != static_cast(golden[i])) { mism = static_cast(i); break; } diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index 65f765812be..ec4125a2818 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -61,6 +61,27 @@ def test_generate_case_writes_artifacts(tmp_path): assert entry["golden"]["output_index"] == 0 +def test_export_case_separates_upper_bound_from_runtime_inputs(monkeypatch): + suite = op_test_registry["conv1d"] + case = next(c for c in suite.cases if c.name == "dynamic_length_10_to_7") + export_shapes = [] + exported_dynamic_shapes = [] + real_export = torch.export.export + + def capture_export(module, inputs, **kwargs): + export_shapes.append(tuple(inputs[0].shape)) + exported_dynamic_shapes.append(kwargs.get("dynamic_shapes")) + return real_export(module, inputs, **kwargs) + + monkeypatch.setattr(torch.export, "export", capture_export) + _module, runtime_inputs, prog = g.export_case(suite, case) + + assert export_shapes == [(1, 4, 10)] + assert exported_dynamic_shapes == [case.dynamic_shapes] + assert tuple(runtime_inputs[0].shape) == (1, 4, 7) + assert g._has_vulkan_delegate(prog) + + def test_generate_manifest(tmp_path): g.generate(str(tmp_path), ops=["add"]) manifest = tmp_path / "manifest.json" diff --git a/backends/webgpu/test/op_tests/test_suite.py b/backends/webgpu/test/op_tests/test_suite.py index f2714125c84..17542cd2e55 100644 --- a/backends/webgpu/test/op_tests/test_suite.py +++ b/backends/webgpu/test/op_tests/test_suite.py @@ -58,6 +58,9 @@ class Case: required: bool = True heavy: bool = False golden_fn: Callable | None = None + # Optional upper-bound export inputs; `inputs` stay live manifest tensors. + export_inputs: tuple[Input, ...] | None = None + dynamic_shapes: object | None = None def __post_init__(self) -> None: # Mirror kQ4gswConfigs: every heavy config is required=False (export-gated, never FAILs on absence). 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 b302b3c120d..d9f666622ee 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 @@ -19,6 +19,8 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule +from executorch.backends.webgpu.test.ops.test_gelu import GeluModule from executorch.exir import to_edge_transform_and_lower from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes @@ -181,6 +183,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.select(0, -1) +class DynamicExpandCopyModule(torch.nn.Module): + """Dynamic expand_copy is rejected until its TensorMeta can be resized.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, x.shape[1])).clone() + + +class DynamicExpandCopyInferredModule(torch.nn.Module): + """Dynamic expand_copy whose -1 target hides symbolic provenance.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, -1)).clone() + + def _ramp(shape) -> torch.Tensor: n = 1 for d in shape: @@ -252,9 +268,75 @@ def _write_goldens(model, prefix: str, out_dir: str, s_values) -> None: print(f" golden {prefix} S={s}") +def export_dynamic_conv1d_cases(out_dir: str) -> None: + """Write one dynamic Conv1d program and live-length runtime fixtures.""" + os.makedirs(out_dir, exist_ok=True) + max_length = 16 + lengths = (max_length, 9, 5) + model = Conv1dModule( + in_channels=3, + out_channels=4, + kernel_size=3, + stride=2, + padding=1, + dilation=2, + bias=True, + ).eval() + length_dim = torch.export.Dim("conv1d_length", min=5, max=max_length) + _export( + model, + (_ramp((1, 3, max_length)),), + {"x": {2: length_dim}}, + os.path.join(out_dir, "dyn_conv1d.pte"), + ) + for length in lengths: + x = _ramp((1, 3, length)) + with torch.no_grad(): + golden = model(x) + prefix = os.path.join(out_dir, f"dyn_conv1d.S{length}") + x.detach().numpy().astype(" None: + """Write a dynamic GELU fixture crossing the old 1D dispatch cap.""" + os.makedirs(out_dir, exist_ok=True) + max_elements = 4 * 64 * 65535 + 1 + model = GeluModule("none").eval() + elements_dim = torch.export.Dim("gelu_elements", min=1024, max=max_elements) + _export( + model, + (torch.empty((max_elements,), dtype=torch.float32),), + {"x": {0: elements_dim}}, + os.path.join(out_dir, "dyn_gelu_2d.pte"), + ) + + +def export_dynamic_expand_copy_rejection_case(out_dir: str) -> None: + """Write a dynamic expand_copy graph that the runtime must reject at load.""" + model = DynamicExpandCopyModule().eval() + elements_dim = torch.export.Dim("expand_elements", min=1, max=8) + _export( + model, + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy.pte"), + ) + _export( + DynamicExpandCopyInferredModule().eval(), + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy_inferred.pte"), + ) + + def export_dynamic_shape_cases(out_dir: str) -> None: """Write the dynamic + static .pte's and per-S goldens for the native test.""" os.makedirs(out_dir, exist_ok=True) + export_dynamic_conv1d_cases(out_dir) + export_dynamic_expand_copy_rejection_case(out_dir) + if os.environ.get("WEBGPU_TEST_HEAVY"): + export_dynamic_gelu_boundary_cases(out_dir) s_dim = torch.export.Dim("s", min=1, max=MAXS) # 1) Single dynamic rms_norm, graph built at S=MAXS (upper bound). @@ -1536,6 +1618,13 @@ def test_export_dynamic_rms(self) -> None: 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_conv1d.pte", + "dyn_conv1d.S16.input.bin", + "dyn_conv1d.S16.golden.bin", + "dyn_conv1d.S9.input.bin", + "dyn_conv1d.S9.golden.bin", + "dyn_conv1d.S5.input.bin", + "dyn_conv1d.S5.golden.bin", "dyn_linear_bk64.pte", "dyn_linear_bk64.S512.input.bin", "dyn_linear_bk64.S512.golden.bin", diff --git a/backends/webgpu/test/ops/test_conv1d_pw.py b/backends/webgpu/test/ops/test_conv1d_pw.py index a0de8988f81..07036fe4a42 100644 --- a/backends/webgpu/test/ops/test_conv1d_pw.py +++ b/backends/webgpu/test/ops/test_conv1d_pw.py @@ -27,6 +27,16 @@ "batch2": (2, 3, 4, 5, True), } +# name -> N, C_in, C_out, L, K, stride, padding, dilation, bias +GENERAL_CONFIGS = { + "voxtral_stride1": (1, 4, 6, 10, 3, 1, 0, 1, True), + "voxtral_stride2": (1, 6, 5, 10, 3, 2, 0, 1, True), + "no_bias": (1, 3, 2, 9, 3, 1, 0, 1, False), + "padded": (1, 3, 4, 9, 3, 1, 1, 1, True), + "dilated": (1, 2, 3, 11, 3, 1, 2, 2, True), + "batch2": (2, 3, 4, 8, 3, 2, 1, 1, True), +} + class Conv1dPwModule(torch.nn.Module): def __init__(self, in_channels, out_channels, bias) -> None: @@ -42,6 +52,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) +class Conv1dModule(torch.nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + bias, + ) -> None: + super().__init__() + g = torch.Generator().manual_seed(0) + self.conv = torch.nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + with torch.no_grad(): + self.conv.weight.normal_(generator=g) + if bias: + self.conv.bias.normal_(generator=g) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + def _det_input(shape): g = torch.Generator().manual_seed(1) return torch.randn(*shape, generator=g, dtype=torch.float32) @@ -54,6 +95,14 @@ def _lower(cfg): return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) +def _lower_general(cfg, dynamic_shapes=None): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + module = Conv1dModule(ic, oc, kernel, stride, padding, dilation, bias).eval() + inputs = (_det_input((n, ic, length)),) + ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + def _delegated(et) -> bool: return any( d.id == "VulkanBackend" @@ -63,9 +112,17 @@ def _delegated(et) -> bool: def _op_delegated(edge, op_substr: str) -> bool: - # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + # Require the op in a delegate, not merely absent from the host graph. + from executorch.exir.lowered_backend_module import get_lowered_submodules + gm = edge.exported_program().graph_module - return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + if any(op_substr in str(getattr(n, "target", "")) for n in gm.graph.nodes): + return False + return any( + op_substr in str(getattr(dn, "target", "")) + for _, lowered, _ in get_lowered_submodules(gm) + for dn in lowered.original_module.graph_module.graph.nodes + ) class Conv1dPwTest(unittest.TestCase): @@ -82,3 +139,24 @@ def test_export_delegates(self) -> None: _op_delegated(edge, "convolution"), f"conv1d not delegated (fell back to CPU) for {name}", ) + + +class Conv1dTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, cfg in GENERAL_CONFIGS.items(): + with self.subTest(name=name): + edge = _lower_general(cfg) + self.assertTrue( + _delegated(edge.to_executorch()), + f"Expected a VulkanBackend delegate (conv1d {name})", + ) + self.assertTrue( + _op_delegated(edge, "convolution"), + f"conv1d not delegated (fell back to CPU) for {name}", + ) + + def test_dynamic_length_export_delegates(self) -> None: + length = torch.export.Dim("conv1d_length", min=5, max=16) + edge = _lower_general(GENERAL_CONFIGS["padded"], dynamic_shapes=({2: length},)) + self.assertTrue(_delegated(edge.to_executorch())) + self.assertTrue(_op_delegated(edge, "convolution")) diff --git a/backends/webgpu/test/ops/test_to_copy.py b/backends/webgpu/test/ops/test_to_copy.py index 1fa2375f248..54b400ea9ef 100644 --- a/backends/webgpu/test/ops/test_to_copy.py +++ b/backends/webgpu/test/ops/test_to_copy.py @@ -46,6 +46,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.float32, copy=True) +class ToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class ToCopyInt8ToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class CompareToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a > b).to(torch.float32) + + def to_copy_int_input(shape: tuple[int, ...]) -> torch.Tensor: n = math.prod(shape) return (torch.arange(n, dtype=torch.int32) - n // 2).reshape(shape) @@ -61,14 +76,32 @@ def to_copy_float_input(shape: tuple[int, ...]) -> torch.Tensor: return pattern.repeat(repeats)[:n].reshape(shape) -def _lower(model: torch.nn.Module, x: torch.Tensor): - ep = torch.export.export(model.eval(), (x,)) +def bool_tail_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([True, False, True, True, False, False, True]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_a(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_b(shape: tuple[int, ...]) -> torch.Tensor: + return torch.zeros(shape, dtype=torch.float32) + + +def _lower(model: torch.nn.Module, *inputs: torch.Tensor): + ep = torch.export.export(model.eval(), inputs) edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) return ep, edge -def _export(model: torch.nn.Module, x: torch.Tensor): - _, edge = _lower(model, x) +def _export(model: torch.nn.Module, *inputs: torch.Tensor): + _, edge = _lower(model, *inputs) return edge.to_executorch() @@ -142,3 +175,22 @@ def test_float_passthrough_delegates(self) -> None: self.assertTrue( _delegated(et), "Expected a VulkanBackend delegate (to_copy float->float)" ) + + def test_bool_to_float_delegates(self) -> None: + x = bool_tail_input((5,)) + ep, edge = _lower(ToCopyBoolToFloatModule(), x) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_compare_bool_to_float_delegates(self) -> None: + a = compare_to_copy_input_a((5,)) + b = compare_to_copy_input_b((5,)) + ep, edge = _lower(CompareToCopyBoolToFloatModule(), a, b) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_int8_to_float_does_not_delegate(self) -> None: + x = torch.tensor([-2, 0, 3], dtype=torch.int8) + self.assertFalse(_delegated(_export(ToCopyInt8ToFloatModule(), x))) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 574c3869864..9990297b9f0 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -77,6 +77,16 @@ def _function_source(text: str, name: str) -> str: class WgslCodegenTest(unittest.TestCase): + def test_compare_word_count_does_not_overflow_u32(self) -> None: + source = (g.BACKEND_ROOT / "runtime/ops/compare/compare.wgsl").read_text() + expression = "(params.num_elements - 1u) / 4u + 1u" + self.assertIn(expression, source) + for num_elements in (1, 4, 5, (1 << 32) - 3, (1 << 32) - 2, (1 << 32) - 1): + self.assertEqual( + (num_elements - 1) // 4 + 1, + (num_elements + 3) // 4, + ) + def test_registry_entries_match_concrete_headers(self) -> None: entries = g.registry_entries() names = [entry.name for entry in entries] @@ -210,14 +220,14 @@ def test_generated_output_manifest_digest(self) -> None: digest.update(b"\0") digest.update(output.read_bytes()) digest.update(b"\0") - self.assertEqual(len(outputs), 134) + self.assertEqual(len(outputs), 136) self.assertEqual( digest.hexdigest(), - "e502196846f0f8100f468e5d9f8f9c006b67e08df54e1e2e667daa2fc50d8844", + "0512f8d258952e446ffaedcb653b6a3a720eccf8a6b5327d95fd454a912214a3", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "492535b396833ad6ebfed29b093057e38f40b1182b5c7d6b3eb2f3577cab024e", + "28aaa7a8d3e916df43e407120e91d487d0d51cbc5ca93c56bd822d25d109890e", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: